Make token budget context identity configurable (#36970)

## What changed

- Add a `features.token_budget.mode` setting with `thread` and `name` values.
- Default context-window metadata to the thread ID while allowing configurations to retain the agent name.
- Preserve the selected identity mode when applying model token-budget defaults and replaying locked configuration.

## Testing

- Cover configuration parsing and both identity formats in token-budget context messages.

GitOrigin-RevId: f6e6d2c49e09e70c0352e803bb65784ef77e5bee

codex-rs/core/config.schema.json

@@ -3119,6 +3119,14 @@
           "maxLength": 2000,
           "type": "string"
         },
+        "mode": {
+          "allOf": [
+            {
+              "$ref": "#/definitions/TokenBudgetMode"
+            }
+          ],
+          "description": "Select whether context-window metadata identifies the thread or agent name."
+        },
         "reminder_message_template": {
           "description": "Reminder template. `{n_remaining}` is replaced with the tokens remaining before auto-compaction.",
           "maxLength": 2000,
@@ -3134,6 +3142,14 @@
       },
       "type": "object"
     },
+    "TokenBudgetMode": {
+      "description": "Identity included in the context-window developer message.",
+      "enum": [
+        "thread",
+        "name"
+      ],
+      "type": "string"
+    },
     "ToolExposureSurface": {
       "description": "A model-facing surface on which a tool can be exposed.",
       "oneOf": [

codex-rs/core/src/config/config_tests.rs

@@ -609,6 +609,7 @@ async fn load_config_resolves_token_budget_config() -> std::io::Result<()> {
             r#"
 [features.token_budget]
 enabled = true
+mode = "thread"
 reminder_threshold_tokens = 16000
 reminder_message_template = "Custom reminder: {n_remaining} tokens."
 guidance_message = "Preserve important state before compaction."
@@ -616,6 +617,7 @@ auto_compact_fallback_prompt = "  Write notes immediately.  "
 auto_compact_fallback_buffer_tokens = 8000
 "#,
             TokenBudgetConfig {
+                mode: TokenBudgetMode::Thread,
                 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()),
@@ -623,6 +625,13 @@ auto_compact_fallback_buffer_tokens = 8000
                 auto_compact_fallback_buffer_tokens: Some(8_000),
             },
         ),
+        (
+            "[features.token_budget]\nenabled = true\nmode = \"name\"\n",
+            TokenBudgetConfig {
+                mode: TokenBudgetMode::Name,
+                ..TokenBudgetConfig::default()
+            },
+        ),
     ] {
         let codex_home = tempdir()?;
         let config_toml = toml::from_str(config_toml).expect("TOML should deserialize");

codex-rs/core/src/config/mod.rs

@@ -69,6 +69,7 @@ use codex_features::FeaturesToml;
 use codex_features::MultiAgentV2ConfigToml;
 use codex_features::NetworkProxyConfigToml;
 use codex_features::TokenBudgetConfigToml;
+use codex_features::TokenBudgetMode;
 use codex_git_utils::resolve_root_git_project_for_trust;
 use codex_http_client::HttpClientFactory;
 use codex_http_client::OutboundProxyPolicy;
@@ -1119,6 +1120,7 @@ const AUTO_COMPACT_FALLBACK_PROMPT_MAX_BYTES: usize = 2000;
 
 #[derive(Debug, Clone, PartialEq, Eq, Serialize)]
 pub struct TokenBudgetConfig {
+    pub mode: TokenBudgetMode,
     pub reminder_threshold_tokens: Option<i64>,
     pub reminder_message_template: String,
     pub guidance_message: Option<String>,
@@ -1211,6 +1213,7 @@ impl TokenBudgetConfig {
 impl Default for TokenBudgetConfig {
     fn default() -> Self {
         Self {
+            mode: TokenBudgetMode::default(),
             reminder_threshold_tokens: None,
             reminder_message_template: DEFAULT_TOKEN_BUDGET_REMINDER_MESSAGE_TEMPLATE.to_string(),
             guidance_message: None,
@@ -2757,6 +2760,9 @@ fn resolve_token_budget_config(
     }
 
     let token_budget_config = token_budget_toml_config(config_toml.features.as_ref());
+    let mode = token_budget_config
+        .and_then(|config| config.mode)
+        .unwrap_or_default();
     let reminder_threshold_tokens =
         token_budget_config.and_then(|config| config.reminder_threshold_tokens);
     let reminder_message_template = token_budget_config
@@ -2774,6 +2780,7 @@ fn resolve_token_budget_config(
         token_budget_config.and_then(|config| config.auto_compact_fallback_buffer_tokens);
 
     let token_budget = TokenBudgetConfig {
+        mode,
         reminder_threshold_tokens,
         reminder_message_template,
         guidance_message,

codex-rs/core/src/context/token_budget_context.rs

@@ -1,4 +1,5 @@
 use super::ContextualUserFragment;
+use codex_features::TokenBudgetMode;
 use codex_protocol::AgentPath;
 use codex_protocol::ThreadId;
 use codex_protocol::protocol::CONTEXT_WINDOW_CLOSE_TAG;
@@ -11,6 +12,7 @@ use uuid::Uuid;
 pub(crate) struct TokenBudgetContext {
     thread_id: ThreadId,
     agent_path: AgentPath,
+    mode: TokenBudgetMode,
     first_window_id: Uuid,
     previous_window_id: Option<Uuid>,
     window_id: Uuid,
@@ -21,6 +23,7 @@ impl TokenBudgetContext {
     pub(crate) fn new(
         thread_id: ThreadId,
         agent_path: AgentPath,
+        mode: TokenBudgetMode,
         first_window_id: Uuid,
         previous_window_id: Option<Uuid>,
         window_id: Uuid,
@@ -29,6 +32,7 @@ impl TokenBudgetContext {
         Self {
             thread_id,
             agent_path,
+            mode,
             first_window_id,
             previous_window_id,
             window_id,
@@ -51,11 +55,14 @@ impl ContextualUserFragment for TokenBudgetContext {
     }
 
     fn body(&self) -> String {
-        let agent_path = &self.agent_path;
+        let identity = match self.mode {
+            TokenBudgetMode::Thread => format!("Thread id: {}", self.thread_id),
+            TokenBudgetMode::Name => format!("Agent name: {}", self.agent_path),
+        };
         let first_window_id = self.first_window_id;
         let window_id = self.window_id;
         let mut lines = vec![
-            format!("Agent name: {agent_path}"),
+            identity,
             format!("First context window id: {first_window_id}"),
             format!("Current context window id: {window_id}"),
         ];

codex-rs/core/src/session/config_lock.rs

@@ -272,6 +272,7 @@ mod tests {
         config.multi_agent_v2.subagent_developer_instructions =
             Some("Locked subagent developer instructions.".to_string());
         config.token_budget = Some(crate::config::TokenBudgetConfig {
+            mode: codex_features::TokenBudgetMode::Thread,
             reminder_threshold_tokens: Some(16_000),
             reminder_message_template: "Locked reminder: {n_remaining} tokens.".to_string(),
             guidance_message: Some("Locked context-window guidance.".to_string()),
@@ -371,6 +372,7 @@ mod tests {
             features.token_budget,
             Some(FeatureToml::Config(TokenBudgetConfigToml {
                 enabled: Some(true),
+                mode: Some(codex_features::TokenBudgetMode::Thread),
                 reminder_threshold_tokens: Some(16_000),
                 reminder_message_template: Some(
                     "Locked reminder: {n_remaining} tokens.".to_string()

codex-rs/core/src/session/mod.rs

@@ -3502,6 +3502,12 @@ impl Session {
                     session_source
                         .get_agent_path()
                         .unwrap_or_else(codex_protocol::AgentPath::root),
+                    turn_context
+                        .config
+                        .token_budget
+                        .as_ref()
+                        .map(|config| config.mode)
+                        .unwrap_or_default(),
                     auto_compact_window_ids.first_window_id,
                     auto_compact_window_ids.previous_window_id,
                     auto_compact_window_ids.window_id,

codex-rs/core/src/session/tests.rs

@@ -9105,7 +9105,12 @@ async fn build_initial_context_adds_multi_agent_v2_subagent_usage_hint_as_develo
     let turn_context_mut =
         Arc::get_mut(&mut turn_context).expect("thread settings should not be shared");
     turn_context_mut.session_source = session_source;
-    Arc::make_mut(&mut turn_context_mut.config)
+    let config = Arc::make_mut(&mut turn_context_mut.config);
+    config.token_budget = Some(crate::config::TokenBudgetConfig {
+        mode: codex_features::TokenBudgetMode::Name,
+        ..crate::config::TokenBudgetConfig::default()
+    });
+    config
         .features
         .enable(Feature::TokenBudget)
         .expect("test config should allow token budget");

codex-rs/core/src/session/token_budget.rs

@@ -13,11 +13,14 @@ pub(super) fn has_explicit_settings(config: &Config) -> bool {
         .get("features")
         .and_then(|features| features.get("token_budget"))
         .and_then(|token_budget| token_budget.as_table())
-        .is_some_and(|settings| settings.keys().any(|key| key != "enabled"))
-        || config
-            .token_budget
-            .as_ref()
-            .is_some_and(|token_budget| token_budget != &TokenBudgetConfig::default())
+        .is_some_and(|settings| settings.keys().any(|key| key != "enabled" && key != "mode"))
+        || config.token_budget.as_ref().is_some_and(|token_budget| {
+            token_budget
+                != &TokenBudgetConfig {
+                    mode: token_budget.mode,
+                    ..TokenBudgetConfig::default()
+                }
+        })
 }
 
 pub(super) fn apply_model_defaults(config: &mut Config, model_info: &ModelInfo) {
@@ -34,6 +37,11 @@ pub(super) fn apply_model_defaults(config: &mut Config, model_info: &ModelInfo)
     };
 
     let token_budget = TokenBudgetConfig {
+        mode: config
+            .token_budget
+            .as_ref()
+            .map(|token_budget| token_budget.mode)
+            .unwrap_or_default(),
         reminder_threshold_tokens: Some(model_defaults.reminder_threshold_tokens),
         reminder_message_template: model_defaults.reminder_message_template.clone(),
         guidance_message: Some(model_defaults.guidance_message.clone()),

codex-rs/core/tests/suite/token_budget.rs

@@ -8,6 +8,7 @@ use codex_extension_api::ExtensionRegistryBuilder;
 use codex_features::Feature;
 use codex_features::FeatureToml;
 use codex_features::TokenBudgetConfigToml;
+use codex_features::TokenBudgetMode;
 use codex_model_provider_info::built_in_model_providers;
 use codex_protocol::config_types::AutoCompactTokenLimitScope;
 use codex_protocol::items::TurnItem;
@@ -74,7 +75,7 @@ fn model_token_budget_config() -> ModelTokenBudgetConfig {
 }
 
 fn token_budget_contexts(request: &ResponsesRequest) -> Vec<String> {
-    let context_window_prefix = format!("{CONTEXT_WINDOW_OPEN_TAG}\nAgent name: ");
+    let context_window_prefix = format!("{CONTEXT_WINDOW_OPEN_TAG}\n");
     request
         .message_input_texts("developer")
         .into_iter()
@@ -85,7 +86,7 @@ fn token_budget_contexts(request: &ResponsesRequest) -> Vec<String> {
 fn token_budget_window_ids(text: &str, agent_name: &str) -> (String, Option<String>, String) {
     let captures = assert_regex_match(
         &format!(
-            r"^{CONTEXT_WINDOW_OPEN_TAG}\nAgent name: {agent_name}\nFirst context window id: ([0-9a-f-]{{36}})\nCurrent context window id: ([0-9a-f-]{{36}})(?:\nPrevious context window id: ([0-9a-f-]{{36}}))?\n{CONTEXT_WINDOW_CLOSE_TAG}$"
+            r"^{CONTEXT_WINDOW_OPEN_TAG}\n(?:Thread id: [0-9a-f-]{{36}}|Agent name: {agent_name})\nFirst context window id: ([0-9a-f-]{{36}})\nCurrent context window id: ([0-9a-f-]{{36}})(?:\nPrevious context window id: ([0-9a-f-]{{36}}))?\n{CONTEXT_WINDOW_CLOSE_TAG}$"
         ),
         text,
     );
@@ -290,6 +291,13 @@ async fn token_budget_uses_model_message_defaults() -> Result<()> {
     let model_defaults = model_token_budget_config();
     let expected_guidance = model_defaults.guidance_message.clone();
     let test = test_codex()
+        .with_pre_build_hook(|home| {
+            std::fs::write(
+                home.join("config.toml"),
+                "[features.token_budget]\nenabled = true\nmode = \"name\"\n",
+            )
+            .expect("write agent-name token-budget configuration");
+        })
         .with_model_info_override("gpt-5.2", move |model_info| {
             model_info
                 .model_messages
@@ -310,7 +318,14 @@ async fn token_budget_uses_model_message_defaults() -> Result<()> {
     test.submit_turn("inspect model-owned context guidance")
         .await?;
 
-    let developer_texts = response.single_request().message_input_texts("developer");
+    let request = response.single_request();
+    let token_budget_context = token_budget_contexts(&request);
+    assert_eq!(token_budget_context.len(), 1);
+    assert!(
+        token_budget_context[0]
+            .starts_with(&format!("{CONTEXT_WINDOW_OPEN_TAG}\nAgent name: /root\n"))
+    );
+    let developer_texts = request.message_input_texts("developer");
     assert!(developer_texts.iter().any(|text| {
         text == &format!(
             "{CONTEXT_WINDOW_GUIDANCE_OPEN_TAG}\n{expected_guidance}\n{CONTEXT_WINDOW_GUIDANCE_CLOSE_TAG}"
@@ -414,6 +429,7 @@ async fn token_budget_model_defaults_survive_config_lock_replay() -> Result<()>
             .and_then(|features| features.token_budget.as_ref()),
         Some(&FeatureToml::Config(TokenBudgetConfigToml {
             enabled: Some(true),
+            mode: Some(TokenBudgetMode::Thread),
             reminder_threshold_tokens: Some(6_144),
             reminder_message_template: Some(
                 "Model reminder: {n_remaining} tokens remain.".to_string()
@@ -727,7 +743,7 @@ async fn token_budget_context_injects_plain_thread_hint_text() -> Result<()> {
     assert_eq!(token_budgets.len(), 1);
     let captures = assert_regex_match(
         &format!(
-            r"^{CONTEXT_WINDOW_OPEN_TAG}\nAgent name: /root\nFirst context window id: ([0-9a-f-]{{36}})\nCurrent context window id: ([0-9a-f-]{{36}})\nmanual history hint for thread {thread_id}\nunstructured notes/thread_hint fixture result\n{CONTEXT_WINDOW_CLOSE_TAG}$"
+            r"^{CONTEXT_WINDOW_OPEN_TAG}\nThread id: {thread_id}\nFirst context window id: ([0-9a-f-]{{36}})\nCurrent context window id: ([0-9a-f-]{{36}})\nmanual history hint for thread {thread_id}\nunstructured notes/thread_hint fixture result\n{CONTEXT_WINDOW_CLOSE_TAG}$"
         ),
         &token_budgets[0],
     );
@@ -1492,6 +1508,7 @@ async fn new_context_tool_skips_auto_compact_fallback() -> Result<()> {
         .with_config(|config| {
             config.model_context_window = Some(10_000);
             config.token_budget = Some(TokenBudgetConfig {
+                mode: TokenBudgetMode::Name,
                 auto_compact_fallback_prompt: Some(AUTO_COMPACT_FALLBACK_PROMPT.to_string()),
                 auto_compact_fallback_buffer_tokens: Some(4_000),
                 ..TokenBudgetConfig::default()

codex-rs/features/src/feature_configs.rs

@@ -134,11 +134,23 @@ impl FeatureConfig for MultiAgentV2ConfigToml {
     }
 }
 
+/// Identity included in the context-window developer message.
+#[derive(Serialize, Deserialize, Debug, Clone, Copy, Default, PartialEq, Eq, JsonSchema)]
+#[serde(rename_all = "snake_case")]
+pub enum TokenBudgetMode {
+    #[default]
+    Thread,
+    Name,
+}
+
 #[derive(Serialize, Deserialize, Debug, Clone, Default, PartialEq, Eq, JsonSchema)]
 #[serde(deny_unknown_fields)]
 pub struct TokenBudgetConfigToml {
     #[serde(skip_serializing_if = "Option::is_none")]
     pub enabled: Option<bool>,
+    /// Select whether context-window metadata identifies the thread or agent name.
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub mode: Option<TokenBudgetMode>,
     /// Number of tokens remaining before auto-compaction when the wrap-up reminder is emitted.
     #[serde(skip_serializing_if = "Option::is_none")]
     #[schemars(range(min = 1))]

codex-rs/features/src/lib.rs

@@ -30,6 +30,7 @@ pub use feature_configs::NonPrefixedMcpToolNamesConfigToml;
 use feature_configs::RemovedAppsMcpPathOverrideConfigToml;
 pub use feature_configs::RolloutBudgetConfigToml;
 pub use feature_configs::TokenBudgetConfigToml;
+pub use feature_configs::TokenBudgetMode;
 pub use feature_configs::ToolRegistryConfigToml;
 use legacy::LegacyFeatureToggles;
 pub use legacy::legacy_feature_keys;