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