From 69b6152c17943e681941484ad3de31e0c212e814 Mon Sep 17 00:00:00 2001 From: Boyang Niu Date: Fri, 7 Aug 2026 00:24:14 +0000 Subject: [PATCH] Track context windows per agent (#37347) ## Why Forked subagents can inherit a parent's compacted history, but their context-window metadata needs to identify the child and start a distinct window lineage. ## What changed - Always identify context windows by agent name and remove the `features.token_budget.mode` setting. - Persist the agent path in world state and emit context-window metadata as a standalone developer message when the identity changes. - Reset inherited compaction metadata to the forked subagent's initial context window and restore that window in session state. ## Testing Extend the full-history subagent test to verify parent and child world-state snapshots, agent identities, distinct window IDs, and the reset child checkpoint. GitOrigin-RevId: 6584201e57b3d2fba0ce626b89a8b3b0b3c752db --- codex-rs/core/config.schema.json | 16 --- codex-rs/core/src/config/config_tests.rs | 9 -- codex-rs/core/src/config/mod.rs | 7 -- .../core/src/context/token_budget_context.rs | 37 +++--- codex-rs/core/src/session/config_lock.rs | 2 - codex-rs/core/src/session/mod.rs | 9 +- codex-rs/core/src/session/session.rs | 20 ++++ codex-rs/core/src/session/tests.rs | 5 +- codex-rs/core/src/session/token_budget.rs | 18 +-- codex-rs/core/src/session/world_state.rs | 20 +++- ..._new_context_window_tool_full_context.snap | 14 +-- .../tests/suite/subagent_notifications.rs | 108 +++++++++++++++++- codex-rs/core/tests/suite/token_budget.rs | 28 ++--- codex-rs/features/src/feature_configs.rs | 12 -- codex-rs/features/src/lib.rs | 1 - 15 files changed, 195 insertions(+), 111 deletions(-) diff --git a/codex-rs/core/config.schema.json b/codex-rs/core/config.schema.json index 9d67ab7a84..8c0a10169a 100644 --- a/codex-rs/core/config.schema.json +++ b/codex-rs/core/config.schema.json @@ -3125,14 +3125,6 @@ "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, @@ -3148,14 +3140,6 @@ }, "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 e3e91a670a..19ab4f3346 100644 --- a/codex-rs/core/src/config/config_tests.rs +++ b/codex-rs/core/src/config/config_tests.rs @@ -609,7 +609,6 @@ 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." @@ -617,7 +616,6 @@ 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()), @@ -625,13 +623,6 @@ 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 fc3b7f54de..61cf285bca 100644 --- a/codex-rs/core/src/config/mod.rs +++ b/codex-rs/core/src/config/mod.rs @@ -70,7 +70,6 @@ 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; @@ -1136,7 +1135,6 @@ 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, @@ -1229,7 +1227,6 @@ 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, @@ -2813,9 +2810,6 @@ 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 @@ -2833,7 +2827,6 @@ 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 ea32b437bc..b74db69cf3 100644 --- a/codex-rs/core/src/context/token_budget_context.rs +++ b/codex-rs/core/src/context/token_budget_context.rs @@ -1,7 +1,7 @@ use super::ContextualUserFragment; -use codex_features::TokenBudgetMode; +use super::world_state::PreviousSectionState; +use super::world_state::WorldStateSection; use codex_protocol::AgentPath; -use codex_protocol::ThreadId; use codex_protocol::protocol::CONTEXT_WINDOW_CLOSE_TAG; use codex_protocol::protocol::CONTEXT_WINDOW_GUIDANCE_CLOSE_TAG; use codex_protocol::protocol::CONTEXT_WINDOW_GUIDANCE_OPEN_TAG; @@ -10,9 +10,7 @@ use uuid::Uuid; #[derive(Debug, Clone, PartialEq, Eq)] pub(crate) struct TokenBudgetContext { - thread_id: ThreadId, agent_path: AgentPath, - mode: TokenBudgetMode, first_window_id: Uuid, previous_window_id: Option, window_id: Uuid, @@ -21,18 +19,14 @@ pub(crate) struct TokenBudgetContext { 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, mcp_result: Option, ) -> Self { Self { - thread_id, agent_path, - mode, first_window_id, previous_window_id, window_id, @@ -46,6 +40,10 @@ impl ContextualUserFragment for TokenBudgetContext { "developer" } + fn requires_separate_message(&self) -> bool { + true + } + fn markers(&self) -> (&'static str, &'static str) { Self::type_markers() } @@ -55,14 +53,10 @@ impl ContextualUserFragment for TokenBudgetContext { } fn body(&self) -> String { - 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![ - identity, + format!("Agent name: {}", self.agent_path), format!("First context window id: {first_window_id}"), format!("Current context window id: {window_id}"), ]; @@ -76,6 +70,23 @@ impl ContextualUserFragment for TokenBudgetContext { } } +impl WorldStateSection for TokenBudgetContext { + const ID: &'static str = "context_window"; + type Snapshot = AgentPath; + + fn snapshot(&self) -> Self::Snapshot { + self.agent_path.clone() + } + + fn render_diff( + &self, + previous: PreviousSectionState<'_, Self::Snapshot>, + ) -> Option> { + matches!(previous, PreviousSectionState::Known(agent_path) if agent_path != &self.agent_path) + .then(|| Box::new(self.clone()) as Box) + } +} + #[derive(Debug, Clone, PartialEq, Eq)] pub(crate) struct ContextWindowGuidance { message: String, diff --git a/codex-rs/core/src/session/config_lock.rs b/codex-rs/core/src/session/config_lock.rs index 6422757bac..ecbe73091e 100644 --- a/codex-rs/core/src/session/config_lock.rs +++ b/codex-rs/core/src/session/config_lock.rs @@ -272,7 +272,6 @@ 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()), @@ -372,7 +371,6 @@ 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 4db2ba4dac..de33ce20d0 100644 --- a/codex-rs/core/src/session/mod.rs +++ b/codex-rs/core/src/session/mod.rs @@ -3542,18 +3542,11 @@ impl Session { .join("\n"); (!text.is_empty()).then_some(text) }); - developer_sections.push( + separate_developer_sections.push( crate::context::TokenBudgetContext::new( - self.thread_id(), 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/session.rs b/codex-rs/core/src/session/session.rs index ed99e8d01f..265765eaa5 100644 --- a/codex-rs/core/src/session/session.rs +++ b/codex-rs/core/src/session/session.rs @@ -595,6 +595,20 @@ impl Session { } }); let initial_auto_compact_window_ids = AutoCompactWindowIds::new_initial(); + let restore_child_window = matches!(&initial_history, InitialHistory::Forked(_)) + && session_configuration.session_source.is_non_root_agent() + && config.features.enabled(Feature::TokenBudget); + if restore_child_window && let InitialHistory::Forked(items) = &mut initial_history { + let child_window_id = initial_auto_compact_window_ids.window_id.to_string(); + for item in items { + if let RolloutItem::Compacted(checkpoint) = item { + checkpoint.window_number = Some(0); + checkpoint.first_window_id = Some(child_window_id.clone()); + checkpoint.previous_window_id = None; + checkpoint.window_id = Some(child_window_id.clone()); + } + } + } let agent_control = agent_control.with_session_id( session_id, config @@ -1313,6 +1327,12 @@ impl Session { // record_initial_history can emit events. We record only after the SessionConfiguredEvent is emitted. Box::pin(sess.record_initial_history(initial_history)).await; + if restore_child_window { + sess.state.lock().await.restore_auto_compact_window( + /*window_number*/ 0, + initial_auto_compact_window_ids, + ); + } if matches!(&sess.fork_persistence, ForkPersistence::Referenced { .. }) { // Keep the source reserved until the child's history reference is durable. sess.try_ensure_rollout_materialized().await?; diff --git a/codex-rs/core/src/session/tests.rs b/codex-rs/core/src/session/tests.rs index 8048bf525e..5ceb7da2cd 100644 --- a/codex-rs/core/src/session/tests.rs +++ b/codex-rs/core/src/session/tests.rs @@ -9299,10 +9299,7 @@ async fn build_initial_context_adds_multi_agent_v2_subagent_usage_hint_as_develo Arc::get_mut(&mut turn_context).expect("thread settings should not be shared"); turn_context_mut.session_source = session_source; 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.token_budget = Some(crate::config::TokenBudgetConfig::default()); config .features .enable(Feature::TokenBudget) diff --git a/codex-rs/core/src/session/token_budget.rs b/codex-rs/core/src/session/token_budget.rs index 8008d4f228..80388076cd 100644 --- a/codex-rs/core/src/session/token_budget.rs +++ b/codex-rs/core/src/session/token_budget.rs @@ -13,14 +13,11 @@ 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" && key != "mode")) - || config.token_budget.as_ref().is_some_and(|token_budget| { - token_budget - != &TokenBudgetConfig { - mode: token_budget.mode, - ..TokenBudgetConfig::default() - } - }) + .is_some_and(|settings| settings.keys().any(|key| key != "enabled")) + || config + .token_budget + .as_ref() + .is_some_and(|token_budget| token_budget != &TokenBudgetConfig::default()) } pub(super) fn apply_model_defaults(config: &mut Config, model_info: &ModelInfo) { @@ -37,11 +34,6 @@ 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/src/session/world_state.rs b/codex-rs/core/src/session/world_state.rs index 14ac8fc656..60f1829f40 100644 --- a/codex-rs/core/src/session/world_state.rs +++ b/codex-rs/core/src/session/world_state.rs @@ -4,6 +4,7 @@ use super::session::Session; use super::step_context::StepContext; use crate::connectors; use crate::context::ApprovalPromptContext; +use crate::context::TokenBudgetContext; use crate::context::world_state::AgentsMdState; use crate::context::world_state::AppsInstructionsState; use crate::context::world_state::CollaborationModeState; @@ -99,14 +100,27 @@ impl Session { } if turn_context.config.features.enabled(Feature::TokenBudget) && turn_context.model_context_window().is_some() - && let Some(guidance) = turn_context + { + let window_ids = self.state.lock().await.auto_compact_window_ids(); + world_state.add_section(TokenBudgetContext::new( + turn_context + .session_source + .get_agent_path() + .unwrap_or_else(codex_protocol::AgentPath::root), + window_ids.first_window_id, + window_ids.previous_window_id, + window_ids.window_id, + /*mcp_result*/ None, + )); + if let Some(guidance) = turn_context .config .token_budget .as_ref() .and_then(|config| config.guidance_message.as_deref()) .filter(|message| !message.trim().is_empty()) - { - world_state.add_section(ContextWindowGuidanceState::new(guidance)); + { + world_state.add_section(ContextWindowGuidanceState::new(guidance)); + } } let realtime_mode_instructions = self.conversation.mode_instructions().await; world_state.add_section(RealtimeState::new( diff --git a/codex-rs/core/tests/suite/snapshots/all__suite__token_budget__token_budget_new_context_window_tool_full_context.snap b/codex-rs/core/tests/suite/snapshots/all__suite__token_budget__token_budget_new_context_window_tool_full_context.snap index 46e90eeecd..1b9cda1acd 100644 --- a/codex-rs/core/tests/suite/snapshots/all__suite__token_budget__token_budget_new_context_window_tool_full_context.snap +++ b/codex-rs/core/tests/suite/snapshots/all__suite__token_budget__token_budget_new_context_window_tool_full_context.snap @@ -5,10 +5,10 @@ expression: snapshot Scenario: New context window tool installs fresh full context before the next follow-up request. ## Final Follow-Up Request -00:message/developer[3]: - [01] \nAgent name: /root\nFirst context window id: \nCurrent context window id: \nPrevious context window id: \n - [02] - [03] -01:message/user:> -02:function_call/update_plan -03:function_call_output:Plan updated +00:message/developer[2]: + [01] + [02] +01:message/developer:\nAgent name: /root\nFirst context window id: \nCurrent context window id: \nPrevious context window id: \n +02:message/user:> +03:function_call/update_plan +04:function_call_output:Plan updated diff --git a/codex-rs/core/tests/suite/subagent_notifications.rs b/codex-rs/core/tests/suite/subagent_notifications.rs index 899a325acf..55547c97d2 100644 --- a/codex-rs/core/tests/suite/subagent_notifications.rs +++ b/codex-rs/core/tests/suite/subagent_notifications.rs @@ -12,6 +12,7 @@ use codex_protocol::protocol::AgentStatus; use codex_protocol::protocol::AskForApproval; use codex_protocol::protocol::EventMsg; use codex_protocol::protocol::Op; +use codex_protocol::protocol::RolloutItem; use codex_protocol::protocol::SessionSource; use codex_protocol::protocol::SubAgentSource; use codex_protocol::protocol::ThreadHistoryMode; @@ -38,6 +39,7 @@ use core_test_support::test_codex::TestCodex; use core_test_support::test_codex::local_selections; use core_test_support::test_codex::test_codex; use core_test_support::test_codex::turn_permission_fields; +use core_test_support::wait_for_event; use core_test_support::wait_for_event_match; use pretty_assertions::assert_eq; use serde_json::Value; @@ -1039,10 +1041,12 @@ async fn spawned_child_receives_forked_parent_context( enum FullHistoryV2ModelSelection { ConfiguredDefault, ExplicitOverride, + WorldStateIdentity, } #[test_case(FullHistoryV2ModelSelection::ConfiguredDefault; "configured default with omitted fork_turns")] #[test_case(FullHistoryV2ModelSelection::ExplicitOverride; "explicit override with fork_turns all")] +#[test_case(FullHistoryV2ModelSelection::WorldStateIdentity; "world state appends context window when agent identity changes")] #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn spawned_full_history_v2_child_uses_model_precedence_without_dropping_context( selection: FullHistoryV2ModelSelection, @@ -1061,7 +1065,8 @@ async fn spawned_full_history_v2_child_uses_model_precedence_without_dropping_co ) .await; let (spawn_args, expected_model, expected_reasoning_effort) = match selection { - FullHistoryV2ModelSelection::ConfiguredDefault => ( + FullHistoryV2ModelSelection::ConfiguredDefault + | FullHistoryV2ModelSelection::WorldStateIdentity => ( json!({ "message": CHILD_PROMPT, "task_name": "worker", @@ -1119,7 +1124,7 @@ async fn spawned_full_history_v2_child_uses_model_precedence_without_dropping_co ]), ) .await; - let mut builder = test_codex().with_config(|config| { + let mut builder = test_codex().with_config(move |config| { config .features .enable(Feature::Collab) @@ -1128,12 +1133,29 @@ async fn spawned_full_history_v2_child_uses_model_precedence_without_dropping_co .features .enable(Feature::MultiAgentV2) .expect("test config should allow feature update"); + if matches!(selection, FullHistoryV2ModelSelection::WorldStateIdentity) { + config + .features + .enable(Feature::TokenBudget) + .expect("test config should allow feature update"); + config.model_context_window = Some(128_000); + } config.model = Some(INHERITED_MODEL.to_string()); config.model_reasoning_effort = Some(INHERITED_REASONING_EFFORT); config.agent_default_subagent_model = Some(V2_DEFAULT_MODEL.to_string()); config.agent_default_subagent_reasoning_effort = Some(V2_DEFAULT_REASONING_EFFORT); }); + if matches!(selection, FullHistoryV2ModelSelection::WorldStateIdentity) { + builder = builder.with_history_mode(ThreadHistoryMode::Paginated); + } let test = builder.build(&server).await?; + if matches!(selection, FullHistoryV2ModelSelection::WorldStateIdentity) { + test.codex.submit(Op::Compact).await?; + wait_for_event(&test.codex, |event| { + matches!(event, EventMsg::TurnComplete(_)) + }) + .await; + } test.submit_turn(TURN_0_FORK_PROMPT).await?; let _ = seed_turn.single_request(); @@ -1143,6 +1165,88 @@ async fn spawned_full_history_v2_child_uses_model_precedence_without_dropping_co let child_request = wait_for_request_with_model(&child_request_log, expected_model).await?; assert!(child_request.body_contains_text(TURN_0_FORK_PROMPT)); let child_body = child_request.body_json(); + if matches!(selection, FullHistoryV2ModelSelection::WorldStateIdentity) { + let child_thread_id = ThreadId::from_string( + child_body["client_metadata"]["thread_id"] + .as_str() + .expect("child thread id"), + )?; + let child_thread = test.thread_manager.get_thread(child_thread_id).await?; + child_thread.flush_rollout().await?; + let child_rollout = codex_rollout::RolloutRecorder::get_rollout_history( + &child_thread + .rollout_path() + .expect("child rollout should exist"), + ) + .await?; + let context_window_snapshots = child_rollout + .get_rollout_items() + .iter() + .filter_map(|item| match item { + RolloutItem::WorldState(world_state) => { + world_state.state.get("context_window").cloned() + } + _ => None, + }) + .collect::>(); + assert_eq!( + context_window_snapshots, + vec![json!("/root"), json!("/root/worker")] + ); + let context_windows = child_request + .message_input_texts("developer") + .into_iter() + .filter(|text| text.starts_with("\n")) + .collect::>(); + let identities = context_windows + .iter() + .map(|text| text.lines().nth(1).expect("agent identity")) + .collect::>(); + assert_eq!( + identities, + ["Agent name: /root", "Agent name: /root/worker"] + ); + let window_ids = context_windows + .iter() + .map(|text| { + text.lines() + .find_map(|line| line.strip_prefix("Current context window id: ")) + .expect("context window id") + }) + .collect::>(); + assert_ne!(window_ids[0], window_ids[1]); + let checkpoint = child_rollout + .get_rollout_items() + .iter() + .find_map(|item| match item { + RolloutItem::Compacted(checkpoint) => Some(checkpoint), + _ => None, + }) + .expect("inherited compaction checkpoint"); + assert!( + checkpoint + .replacement_history + .as_ref() + .is_some_and(|history| !history.is_empty()) + ); + assert_eq!( + ( + checkpoint.window_number, + checkpoint.first_window_id.as_deref(), + checkpoint.previous_window_id.as_deref(), + checkpoint.window_id.as_deref(), + ), + (Some(0), Some(window_ids[1]), None, Some(window_ids[1])) + ); + assert!( + child_request.has_message_with_input_texts("developer", |message| { + matches!( + message, + [text] if text.starts_with("\nAgent name: /root/worker\n") + ) + }) + ); + } assert_eq!( ( child_body["model"].clone(), diff --git a/codex-rs/core/tests/suite/token_budget.rs b/codex-rs/core/tests/suite/token_budget.rs index 696817ad7c..43c11a193e 100644 --- a/codex-rs/core/tests/suite/token_budget.rs +++ b/codex-rs/core/tests/suite/token_budget.rs @@ -8,7 +8,6 @@ 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; @@ -86,7 +85,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}\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}$" + 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}$" ), text, ); @@ -237,7 +236,7 @@ async fn token_budget_context_is_only_emitted_with_full_context() -> Result<()> } #[tokio::test(flavor = "multi_thread", worker_threads = 2)] -async fn token_budget_guidance_follows_context_window() -> Result<()> { +async fn token_budget_guidance_precedes_standalone_context_window() -> Result<()> { skip_if_no_network!(Ok(())); let server = start_mock_server().await; @@ -272,12 +271,15 @@ async fn token_budget_guidance_follows_context_window() -> Result<()> { .iter() .position(|text| text.starts_with(CONTEXT_WINDOW_OPEN_TAG)) .expect("context-window metadata should be present"); - assert_eq!( - developer_texts.get(context_window_index + 1), - Some(&format!( - "{CONTEXT_WINDOW_GUIDANCE_OPEN_TAG}\n{guidance_message}\n{CONTEXT_WINDOW_GUIDANCE_CLOSE_TAG}" - )) - ); + let guidance_index = developer_texts + .iter() + .position(|text| { + text == &format!( + "{CONTEXT_WINDOW_GUIDANCE_OPEN_TAG}\n{guidance_message}\n{CONTEXT_WINDOW_GUIDANCE_CLOSE_TAG}" + ) + }) + .expect("context-window guidance should be present"); + assert!(guidance_index < context_window_index); Ok(()) } @@ -294,9 +296,9 @@ async fn token_budget_uses_model_message_defaults() -> Result<()> { .with_pre_build_hook(|home| { std::fs::write( home.join("config.toml"), - "[features.token_budget]\nenabled = true\nmode = \"name\"\n", + "[features.token_budget]\nenabled = true\n", ) - .expect("write agent-name token-budget configuration"); + .expect("write token-budget configuration"); }) .with_model_info_override("gpt-5.2", move |model_info| { model_info @@ -429,7 +431,6 @@ 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() @@ -743,7 +744,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}\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}$" + 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}$" ), &token_budgets[0], ); @@ -1508,7 +1509,6 @@ 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 5776910dac..bbe072424c 100644 --- a/codex-rs/features/src/feature_configs.rs +++ b/codex-rs/features/src/feature_configs.rs @@ -134,23 +134,11 @@ 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 74881e5637..51f77749cd 100644 --- a/codex-rs/features/src/lib.rs +++ b/codex-rs/features/src/lib.rs @@ -30,7 +30,6 @@ 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;