diff --git a/codex-rs/core/config.schema.json b/codex-rs/core/config.schema.json index 246c4abffbbd..040347994929 100644 --- a/codex-rs/core/config.schema.json +++ b/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": [ diff --git a/codex-rs/core/src/config/config_tests.rs b/codex-rs/core/src/config/config_tests.rs index ab004642fc06..a9653d3a200e 100644 --- a/codex-rs/core/src/config/config_tests.rs +++ b/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"); diff --git a/codex-rs/core/src/config/mod.rs b/codex-rs/core/src/config/mod.rs index 1c71afa2dd1a..400c00b255bf 100644 --- a/codex-rs/core/src/config/mod.rs +++ b/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, pub reminder_message_template: String, pub guidance_message: Option, @@ -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, diff --git a/codex-rs/core/src/context/token_budget_context.rs b/codex-rs/core/src/context/token_budget_context.rs index 601042569e06..ea32b437bc1e 100644 --- a/codex-rs/core/src/context/token_budget_context.rs +++ b/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, 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, 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}"), ]; diff --git a/codex-rs/core/src/session/config_lock.rs b/codex-rs/core/src/session/config_lock.rs index ecbe73091e7f..6422757bace2 100644 --- a/codex-rs/core/src/session/config_lock.rs +++ b/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() diff --git a/codex-rs/core/src/session/mod.rs b/codex-rs/core/src/session/mod.rs index db73bdbfa402..b44f343fd551 100644 --- a/codex-rs/core/src/session/mod.rs +++ b/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, diff --git a/codex-rs/core/src/session/tests.rs b/codex-rs/core/src/session/tests.rs index 7c6f2fce4caf..3d6bd8a2ea51 100644 --- a/codex-rs/core/src/session/tests.rs +++ b/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"); diff --git a/codex-rs/core/src/session/token_budget.rs b/codex-rs/core/src/session/token_budget.rs index 80388076cd27..8008d4f2286b 100644 --- a/codex-rs/core/src/session/token_budget.rs +++ b/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()), diff --git a/codex-rs/core/tests/suite/token_budget.rs b/codex-rs/core/tests/suite/token_budget.rs index e6721c38635e..696817ad7ca1 100644 --- a/codex-rs/core/tests/suite/token_budget.rs +++ b/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 { - 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 { fn token_budget_window_ids(text: &str, agent_name: &str) -> (String, Option, 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() diff --git a/codex-rs/features/src/feature_configs.rs b/codex-rs/features/src/feature_configs.rs index bbe072424c36..5776910dac5d 100644 --- a/codex-rs/features/src/feature_configs.rs +++ b/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, + /// Select whether context-window metadata identifies the thread or agent name. + #[serde(skip_serializing_if = "Option::is_none")] + pub mode: Option, /// 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))] diff --git a/codex-rs/features/src/lib.rs b/codex-rs/features/src/lib.rs index 68aafd630a84..1fb5dd6f6cb7 100644 --- a/codex-rs/features/src/lib.rs +++ b/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;