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

@@ -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::<codex_features::FeatureToml<
+                    codex_features::ContextManagementConfigToml,
+                >>(),
+            );
+            continue;
+        }
         if feature.id == codex_features::Feature::RolloutBudget {
             validation.properties.insert(
                 feature.key.to_string(),

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"
         },

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(())
 }

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<Option<TokenBudgetConfig>> {

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.

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<PlanType>) -> 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
+            );
+        }
+    }
+}

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(()));

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<bool>,
+}
+
+impl FeatureConfig for ContextManagementConfigToml {
+    fn enabled(&self) -> Option<bool> {
+        self.experimental_mode
+    }
+}
+
 #[derive(Serialize, Deserialize, Debug, Clone, Default, PartialEq, Eq, JsonSchema)]
 #[serde(deny_unknown_fields)]
 pub struct TokenBudgetConfigToml {

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<FeatureToml<TokenBudgetConfigToml>>,
     #[serde(default, skip_serializing_if = "Option::is_none")]
+    pub context_management: Option<FeatureToml<ContextManagementConfigToml>>,
+    #[serde(default, skip_serializing_if = "Option::is_none")]
     pub rollout_budget: Option<FeatureToml<RolloutBudgetConfigToml>>,
     #[serde(default, skip_serializing_if = "Option::is_none")]
     pub current_time_reminder: Option<FeatureToml<CurrentTimeReminderConfigToml>>,
@@ -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",

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"],

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()),