From cff76fa96f70f9f3b63d221446fd02cfd87e6d2e Mon Sep 17 00:00:00 2001 From: pmccrary-oai Date: Wed, 2 Sep 2026 22:13:39 +0000 Subject: [PATCH] Add experimental context management activation (#42385) ## What changed - Add the under-development `features.context_management.experimental_mode` configuration and schema. - For eligible ChatGPT Plus, Pro, and Pro Lite sessions using the Codex backend, enable token-budget context, history notes, and the `new_context` tool when experimental context management is configured. - Keep the feature disabled for custom providers, provider credentials, non-Codex endpoints, and temporary structured threads. ## Testing - Cover subscription eligibility, backend restrictions, configuration resolution, token-budget guidance, and temporary-thread overrides. GitOrigin-RevId: 19402af19c4d69e09a7cdb3687ee7d233f6a562a --- codex-rs/config/src/schema.rs | 9 ++ codex-rs/core/config.schema.json | 26 ++++++ codex-rs/core/src/config/config_tests.rs | 20 +++-- codex-rs/core/src/config/mod.rs | 2 +- codex-rs/core/src/session/mod.rs | 2 + codex-rs/core/src/session/token_budget.rs | 73 ++++++++++++++++ codex-rs/core/tests/suite/token_budget.rs | 84 +++++++++++++++++-- codex-rs/features/src/feature_configs.rs | 14 ++++ codex-rs/features/src/lib.rs | 18 ++++ .../src/app/tests/recap_generation_tests.rs | 3 +- .../tui/src/temporary_structured_request.rs | 1 + 11 files changed, 236 insertions(+), 16 deletions(-) diff --git a/codex-rs/config/src/schema.rs b/codex-rs/config/src/schema.rs index 3198353f38..550f0c3d54 100644 --- a/codex-rs/config/src/schema.rs +++ b/codex-rs/config/src/schema.rs @@ -104,6 +104,15 @@ pub fn features_schema(schema_gen: &mut SchemaGenerator) -> Schema { ); continue; } + if feature.id == codex_features::Feature::ContextManagement { + validation.properties.insert( + feature.key.to_string(), + schema_gen.subschema_for::>(), + ); + continue; + } if feature.id == codex_features::Feature::RolloutBudget { validation.properties.insert( feature.key.to_string(), diff --git a/codex-rs/core/config.schema.json b/codex-rs/core/config.schema.json index 3f95d69f84..33e37a89b1 100644 --- a/codex-rs/core/config.schema.json +++ b/codex-rs/core/config.schema.json @@ -699,6 +699,9 @@ "content_item_kinds": { "type": "boolean" }, + "context_management": { + "$ref": "#/definitions/FeatureToml_for_ContextManagementConfigToml" + }, "current_time_reminder": { "$ref": "#/definitions/FeatureToml_for_CurrentTimeReminderConfigToml" }, @@ -1142,6 +1145,16 @@ }, "type": "object" }, + "ContextManagementConfigToml": { + "additionalProperties": false, + "properties": { + "experimental_mode": { + "description": "Enables experimental context management.", + "type": "boolean" + } + }, + "type": "object" + }, "CurrentTimeReminderConfigToml": { "additionalProperties": false, "properties": { @@ -1255,6 +1268,16 @@ } ] }, + "FeatureToml_for_ContextManagementConfigToml": { + "anyOf": [ + { + "type": "boolean" + }, + { + "$ref": "#/definitions/ContextManagementConfigToml" + } + ] + }, "FeatureToml_for_CurrentTimeReminderConfigToml": { "anyOf": [ { @@ -6014,6 +6037,9 @@ "content_item_kinds": { "type": "boolean" }, + "context_management": { + "$ref": "#/definitions/FeatureToml_for_ContextManagementConfigToml" + }, "current_time_reminder": { "$ref": "#/definitions/FeatureToml_for_CurrentTimeReminderConfigToml" }, diff --git a/codex-rs/core/src/config/config_tests.rs b/codex-rs/core/src/config/config_tests.rs index 75a242f303..6d0d50640c 100644 --- a/codex-rs/core/src/config/config_tests.rs +++ b/codex-rs/core/src/config/config_tests.rs @@ -707,7 +707,11 @@ async fn load_config_resolves_token_budget_config() -> std::io::Result<()> { for (config_toml, expected) in [ ( "[features]\ntoken_budget = true\n", - TokenBudgetConfig::default(), + Some(TokenBudgetConfig::default()), + ), + ( + "features.context_management.experimental_mode = true\n", + None, ), ( r#" @@ -720,14 +724,14 @@ guidance_message = "Preserve important state before compaction." auto_compact_fallback_prompt = " Write notes immediately. " auto_compact_fallback_buffer_tokens = 8000 "#, - TokenBudgetConfig { + Some(TokenBudgetConfig { use_history_notes_extension: true, reminder_threshold_tokens: Some(16_000), reminder_message_template: "Custom reminder: {n_remaining} tokens.".to_string(), guidance_message: Some("Preserve important state before compaction.".to_string()), auto_compact_fallback_prompt: Some("Write notes immediately.".to_string()), auto_compact_fallback_buffer_tokens: Some(8_000), - }, + }), ), ] { let codex_home = tempdir()?; @@ -739,8 +743,14 @@ auto_compact_fallback_buffer_tokens = 8000 ) .await?; - assert!(config.features.enabled(Feature::TokenBudget)); - assert_eq!(config.token_budget, Some(expected)); + assert_eq!( + config.features.enabled(Feature::TokenBudget), + expected.is_some() + ); + if expected.is_none() { + assert!(config.features.enabled(Feature::ContextManagement)); + } + assert_eq!(config.token_budget, expected); } Ok(()) } diff --git a/codex-rs/core/src/config/mod.rs b/codex-rs/core/src/config/mod.rs index ee6faa367c..cd339cd813 100644 --- a/codex-rs/core/src/config/mod.rs +++ b/codex-rs/core/src/config/mod.rs @@ -2766,7 +2766,7 @@ fn resolve_multi_agent_v2_config(config_toml: &ConfigToml) -> MultiAgentV2Config } } -fn resolve_token_budget_config( +pub(crate) fn resolve_token_budget_config( config_toml: &ConfigToml, features: &ManagedFeatures, ) -> std::io::Result> { diff --git a/codex-rs/core/src/session/mod.rs b/codex-rs/core/src/session/mod.rs index 1a81fdb6ca..90fe56ec36 100644 --- a/codex-rs/core/src/session/mod.rs +++ b/codex-rs/core/src/session/mod.rs @@ -682,6 +682,8 @@ impl Session { let model_info = models_manager .get_model_info(model.as_str(), &config.to_models_manager_config()) .await; + let auth = auth_manager.auth_cached(); + token_budget::apply_experimental_context(Arc::make_mut(&mut config), auth.as_ref())?; // Intentionally resolve `enabled` and `use_history_notes_extension` only at // thread startup. Both activation flags stay fixed for this thread runtime, // even if the selected model changes later. diff --git a/codex-rs/core/src/session/token_budget.rs b/codex-rs/core/src/session/token_budget.rs index 179db47caf..0aeb66e4f5 100644 --- a/codex-rs/core/src/session/token_budget.rs +++ b/codex-rs/core/src/session/token_budget.rs @@ -2,10 +2,59 @@ use super::session::Session; use super::turn_context::TurnContext; use crate::config::Config; use crate::config::TokenBudgetConfig; +use crate::config::resolve_token_budget_config; use crate::context::ContextualUserFragment; use codex_features::Feature; +use codex_login::CodexAuth; +use codex_protocol::account::PlanType; +use codex_protocol::auth::AuthMode; use codex_protocol::openai_models::ModelInfo; +fn experimental_context_is_eligible(auth_mode: AuthMode, plan_type: Option) -> bool { + auth_mode == AuthMode::Chatgpt + && matches!( + plan_type, + Some(PlanType::Plus | PlanType::Pro | PlanType::ProLite) + ) +} + +pub(super) fn apply_experimental_context( + config: &mut Config, + auth: Option<&CodexAuth>, +) -> std::io::Result<()> { + let provider = &config.model_provider; + if !config.features.enabled(Feature::ContextManagement) + || !provider.supports_codex_backend_routes() + || !provider.requires_openai_auth + || provider.env_key.is_some() + || provider.experimental_bearer_token.is_some() + || provider.auth.is_some() + || provider.aws.is_some() + || !auth.is_some_and(|auth| { + experimental_context_is_eligible(auth.auth_mode(), auth.account_plan_type()) + }) + || config.features.enable(Feature::TokenBudget).is_err() + || !config.features.enabled(Feature::TokenBudget) + { + return Ok(()); + } + + if config.token_budget.is_none() { + let config_toml = config + .config_layer_stack + .effective_config() + .try_into() + .map_err(|err| std::io::Error::new(std::io::ErrorKind::InvalidData, err))?; + config.token_budget = resolve_token_budget_config(&config_toml, &config.features)?; + } + + config + .token_budget + .get_or_insert_default() + .use_history_notes_extension = true; + Ok(()) +} + /// Detects explicit preferences before model defaults are applied to the turn config. pub(super) fn has_explicit_settings(config: &Config) -> bool { config @@ -163,3 +212,27 @@ pub(super) async fn maybe_record( sess.record_conversation_items(turn_context, std::slice::from_ref(&response_item)) .await; } + +#[cfg(test)] +mod tests { + use super::experimental_context_is_eligible; + use codex_protocol::account::PlanType; + use codex_protocol::auth::AuthMode; + + #[test] + fn experimental_context_requires_eligible_chatgpt_subscription() { + for (auth_mode, plan_type, expected) in [ + (AuthMode::Chatgpt, PlanType::Plus, true), + (AuthMode::Chatgpt, PlanType::Pro, true), + (AuthMode::Chatgpt, PlanType::ProLite, true), + (AuthMode::Chatgpt, PlanType::Free, false), + (AuthMode::Chatgpt, PlanType::Enterprise, false), + (AuthMode::ApiKey, PlanType::Pro, false), + ] { + assert_eq!( + experimental_context_is_eligible(auth_mode, Some(plan_type)), + expected + ); + } + } +} diff --git a/codex-rs/core/tests/suite/token_budget.rs b/codex-rs/core/tests/suite/token_budget.rs index e4d1624ada..0e259ed21c 100644 --- a/codex-rs/core/tests/suite/token_budget.rs +++ b/codex-rs/core/tests/suite/token_budget.rs @@ -6,6 +6,7 @@ use codex_core::config::Config; use codex_core::config::TokenBudgetConfig; use codex_extension_api::ExtensionRegistryBuilder; use codex_features::Feature; +use codex_login::CodexAuth; use codex_model_provider_info::built_in_model_providers; use codex_protocol::config_types::AutoCompactTokenLimitScope; use codex_protocol::items::TurnItem; @@ -235,8 +236,12 @@ async fn token_budget_context_is_only_emitted_with_full_context() -> Result<()> Ok(()) } +#[test_case("features.token_budget.enabled = true"; "token_budget")] +#[test_case("features.context_management.experimental_mode = true"; "experimental_mode")] #[tokio::test(flavor = "multi_thread", worker_threads = 2)] -async fn token_budget_guidance_precedes_standalone_context_window() -> Result<()> { +async fn token_budget_guidance_precedes_standalone_context_window( + activation: &'static str, +) -> Result<()> { skip_if_no_network!(Ok(())); let server = start_mock_server().await; @@ -249,18 +254,26 @@ async fn token_budget_guidance_precedes_standalone_context_window() -> Result<() ) .await; let guidance_message = "Preserve important state before compaction."; + let backend_url = format!("{}/backend-api/codex", server.uri()); let test = test_codex() + .with_auth(CodexAuth::from_external_chatgpt_tokens( + "header.e30.signature", + "account-123", + Some("plus"), + )?) + .with_pre_build_hook(move |home| { + std::fs::write( + home.join("config.toml"), + format!( + "{activation}\nfeatures.token_budget.guidance_message = {guidance_message:?}\n" + ), + ) + .expect("write token-budget guidance configuration"); + }) .with_config(move |config| { + config.model_provider.base_url = Some(backend_url); config.update_plan_enabled = true; config.model_context_window = Some(CONFIGURED_CONTEXT_WINDOW); - config.token_budget = Some(TokenBudgetConfig { - guidance_message: Some(guidance_message.to_string()), - ..TokenBudgetConfig::default() - }); - config - .features - .enable(Feature::TokenBudget) - .expect("test config should allow token budget"); }) .build_with_auto_env(&server) .await?; @@ -291,6 +304,59 @@ async fn token_budget_guidance_precedes_standalone_context_window() -> Result<() Ok(()) } +#[test_case("OpenAI", "/backend-api/codex", None, true; "codex_backend")] +#[test_case("Custom", "/backend-api/codex", None, false; "custom_provider")] +#[test_case("OpenAI", "/v1", None, false; "non_codex_endpoint")] +#[test_case("OpenAI", "/backend-api/codex", Some("test-provider-token"), false; "provider_credentials")] +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn experimental_context_requires_codex_backend( + provider_name: &'static str, + base_path: &'static str, + bearer_token: Option<&'static str>, + expected_enabled: bool, +) -> Result<()> { + skip_if_no_network!(Ok(())); + + let server = start_mock_server().await; + let response = mount_sse_once(&server, sse_completed("resp-1")).await; + let base_url = format!("{}{base_path}", server.uri()); + let test = test_codex() + .with_auth(CodexAuth::from_external_chatgpt_tokens( + "header.e30.signature", + "account-123", + Some("plus"), + )?) + .with_config(move |config| { + config.model_provider.name = provider_name.to_string(); + config.model_provider.base_url = Some(base_url); + config.model_provider.experimental_bearer_token = bearer_token.map(Into::into); + config.model_context_window = Some(CONFIGURED_CONTEXT_WINDOW); + config + .features + .enable(Feature::ContextManagement) + .expect("test config should allow experimental context"); + }) + .build_with_auto_env(&server) + .await?; + + test.submit_turn("inspect experimental context activation") + .await?; + + let request = response.single_request(); + assert_eq!( + tool_names(&request) + .iter() + .any(|name| name == "new_context"), + expected_enabled + ); + assert_eq!( + !token_budget_contexts(&request).is_empty(), + expected_enabled + ); + + Ok(()) +} + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn token_budget_uses_model_message_defaults() -> Result<()> { skip_if_no_network!(Ok(())); diff --git a/codex-rs/features/src/feature_configs.rs b/codex-rs/features/src/feature_configs.rs index 848be62d35..32d3cbaf0e 100644 --- a/codex-rs/features/src/feature_configs.rs +++ b/codex-rs/features/src/feature_configs.rs @@ -293,6 +293,20 @@ impl FeatureConfig for MultiAgentV2ConfigToml { } } +#[derive(Serialize, Deserialize, Debug, Clone, Default, PartialEq, Eq, JsonSchema)] +#[serde(deny_unknown_fields)] +pub struct ContextManagementConfigToml { + /// Enables experimental context management. + #[serde(skip_serializing_if = "Option::is_none")] + pub experimental_mode: Option, +} + +impl FeatureConfig for ContextManagementConfigToml { + fn enabled(&self) -> Option { + self.experimental_mode + } +} + #[derive(Serialize, Deserialize, Debug, Clone, Default, PartialEq, Eq, JsonSchema)] #[serde(deny_unknown_fields)] pub struct TokenBudgetConfigToml { diff --git a/codex-rs/features/src/lib.rs b/codex-rs/features/src/lib.rs index 552792fdf8..6c307e4362 100644 --- a/codex-rs/features/src/lib.rs +++ b/codex-rs/features/src/lib.rs @@ -18,6 +18,7 @@ mod feature_configs; mod legacy; pub use feature_configs::CodeModeConfigToml; pub use feature_configs::CodeModeHostConfigToml; +pub use feature_configs::ContextManagementConfigToml; pub use feature_configs::CurrentTimeReminderConfigToml; pub use feature_configs::CurrentTimeReminderDeliveryMode; pub use feature_configs::CurrentTimeSource; @@ -312,6 +313,8 @@ pub enum Feature { Goals, /// Add current context-window metadata to model-visible context. TokenBudget, + /// Enables experimental context management. + ContextManagement, /// Track and report a shared token budget across a session's agent threads. RolloutBudget, /// Add current-time reminders to model-visible context. @@ -763,6 +766,8 @@ pub struct FeaturesToml { #[serde(default, skip_serializing_if = "Option::is_none")] pub token_budget: Option>, #[serde(default, skip_serializing_if = "Option::is_none")] + pub context_management: Option>, + #[serde(default, skip_serializing_if = "Option::is_none")] pub rollout_budget: Option>, #[serde(default, skip_serializing_if = "Option::is_none")] pub current_time_reminder: Option>, @@ -809,6 +814,13 @@ impl FeaturesToml { if let Some(enabled) = self.token_budget.as_ref().and_then(FeatureToml::enabled) { entries.insert(Feature::TokenBudget.key().to_string(), enabled); } + if let Some(enabled) = self + .context_management + .as_ref() + .and_then(FeatureToml::enabled) + { + entries.insert(Feature::ContextManagement.key().to_string(), enabled); + } if let Some(enabled) = self.rollout_budget.as_ref().and_then(FeatureToml::enabled) { entries.insert(Feature::RolloutBudget.key().to_string(), enabled); } @@ -1549,6 +1561,12 @@ pub const FEATURES: &[FeatureSpec] = &[ stage: Stage::UnderDevelopment, default_enabled: false, }, + FeatureSpec { + id: Feature::ContextManagement, + key: "context_management", + stage: Stage::UnderDevelopment, + default_enabled: false, + }, FeatureSpec { id: Feature::RolloutBudget, key: "rollout_budget", diff --git a/codex-rs/tui/src/app/tests/recap_generation_tests.rs b/codex-rs/tui/src/app/tests/recap_generation_tests.rs index 1988ccc99d..ff3992ff7e 100644 --- a/codex-rs/tui/src/app/tests/recap_generation_tests.rs +++ b/codex-rs/tui/src/app/tests/recap_generation_tests.rs @@ -407,7 +407,7 @@ async fn temporary_recap_threads_disable_memories_and_remote_mcp_servers() -> Re let codex_home = tempdir()?; std::fs::write( codex_home.path().join("config.toml"), - "[mcp_servers.filesystem]\ncommand = 'true'\n", + "features.context_management.experimental_mode = true\n[mcp_servers.filesystem]\ncommand = 'true'\n", )?; app.config.codex_home = codex_home.path().to_path_buf().abs(); app.config.sqlite = codex_state::SqliteConfig::new_for_testing(codex_home.path().abs()); @@ -433,6 +433,7 @@ async fn temporary_recap_threads_disable_memories_and_remote_mcp_servers() -> Re let starts = recorded_params(&requests, "thread/start"); assert_eq!(starts.len(), 1); assert_eq!(starts[0]["config"]["features.memories"], false); + assert_eq!(starts[0]["config"]["features.context_management"], false); assert!(starts[0]["config"].get("features.memory_tool").is_none()); assert_eq!( starts[0]["config"]["mcp_servers"]["filesystem"]["enabled"], diff --git a/codex-rs/tui/src/temporary_structured_request.rs b/codex-rs/tui/src/temporary_structured_request.rs index 7a553e102c..da1926a34b 100644 --- a/codex-rs/tui/src/temporary_structured_request.rs +++ b/codex-rs/tui/src/temporary_structured_request.rs @@ -60,6 +60,7 @@ pub(crate) async fn start_temporary_thread( ("features.apps".to_string(), false.into()), ("features.code_mode".to_string(), false.into()), ("features.code_mode_only".to_string(), false.into()), + ("features.context_management".to_string(), false.into()), ("features.current_time_reminder".to_string(), false.into()), ("features.deferred_executor".to_string(), false.into()), ("features.enable_fanout".to_string(), false.into()),