diff --git a/codex-rs/core/src/context/token_budget_context.rs b/codex-rs/core/src/context/token_budget_context.rs index f5a4409821..601042569e 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_protocol::AgentPath; use codex_protocol::ThreadId; use codex_protocol::protocol::CONTEXT_WINDOW_CLOSE_TAG; use codex_protocol::protocol::CONTEXT_WINDOW_GUIDANCE_CLOSE_TAG; @@ -9,6 +10,7 @@ use uuid::Uuid; #[derive(Debug, Clone, PartialEq, Eq)] pub(crate) struct TokenBudgetContext { thread_id: ThreadId, + agent_path: AgentPath, first_window_id: Uuid, previous_window_id: Option, window_id: Uuid, @@ -18,6 +20,7 @@ pub(crate) struct TokenBudgetContext { impl TokenBudgetContext { pub(crate) fn new( thread_id: ThreadId, + agent_path: AgentPath, first_window_id: Uuid, previous_window_id: Option, window_id: Uuid, @@ -25,6 +28,7 @@ impl TokenBudgetContext { ) -> Self { Self { thread_id, + agent_path, first_window_id, previous_window_id, window_id, @@ -47,11 +51,11 @@ impl ContextualUserFragment for TokenBudgetContext { } fn body(&self) -> String { - let thread_id = self.thread_id; + let agent_path = &self.agent_path; let first_window_id = self.first_window_id; let window_id = self.window_id; let mut lines = vec![ - format!("Thread id: {thread_id}"), + format!("Agent name: {agent_path}"), format!("First context window id: {first_window_id}"), format!("Current context window id: {window_id}"), ]; diff --git a/codex-rs/core/src/event_mapping_tests.rs b/codex-rs/core/src/event_mapping_tests.rs index ac8ed66c5e..972df5e400 100644 --- a/codex-rs/core/src/event_mapping_tests.rs +++ b/codex-rs/core/src/event_mapping_tests.rs @@ -49,7 +49,7 @@ fn recognizes_context_window_as_contextual_developer_content() { let content = vec![ContentItem::InputText { text: format!( r#"{CONTEXT_WINDOW_OPEN_TAG} -Thread id: 00000000-0000-0000-0000-000000000000 +Agent name: /root {CONTEXT_WINDOW_CLOSE_TAG}"# ), }]; diff --git a/codex-rs/core/src/session/mod.rs b/codex-rs/core/src/session/mod.rs index df667895c4..f515fa2787 100644 --- a/codex-rs/core/src/session/mod.rs +++ b/codex-rs/core/src/session/mod.rs @@ -3493,6 +3493,9 @@ impl Session { developer_sections.push( crate::context::TokenBudgetContext::new( self.thread_id(), + session_source + .get_agent_path() + .unwrap_or_else(codex_protocol::AgentPath::root), 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 d2970beb45..040138acc0 100644 --- a/codex-rs/core/src/session/tests.rs +++ b/codex-rs/core/src/session/tests.rs @@ -9066,13 +9066,24 @@ async fn build_initial_context_adds_multi_agent_v2_subagent_usage_hint_as_develo .await .session_configuration .session_source = session_source.clone(); - Arc::get_mut(&mut turn_context) - .expect("thread settings should not be shared") - .session_source = session_source; + 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) + .features + .enable(Feature::TokenBudget) + .expect("test config should allow token budget"); let initial_context = build_initial_context(&session, &turn_context).await; let developer_messages = developer_message_texts(&initial_context); + assert!( + developer_messages + .iter() + .flatten() + .any(|text| text.contains("\nAgent name: /root/worker\n")), + "expected subagent context window to include its canonical name, got {developer_messages:?}" + ); assert!( developer_messages .iter() 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 648353a1a4..46e90eeecd 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 @@ -6,7 +6,7 @@ Scenario: New context window tool installs fresh full context before the next fo ## Final Follow-Up Request 00:message/developer[3]: - [01] \nThread id: \nFirst context window id: \nCurrent context window id: \nPrevious context window id: \n + [01] \nAgent name: /root\nFirst context window id: \nCurrent context window id: \nPrevious context window id: \n [02] [03] 01:message/user:> diff --git a/codex-rs/core/tests/suite/token_budget.rs b/codex-rs/core/tests/suite/token_budget.rs index 088aedb562..e6721c3863 100644 --- a/codex-rs/core/tests/suite/token_budget.rs +++ b/codex-rs/core/tests/suite/token_budget.rs @@ -74,7 +74,7 @@ fn model_token_budget_config() -> ModelTokenBudgetConfig { } fn token_budget_contexts(request: &ResponsesRequest) -> Vec { - let context_window_prefix = format!("{CONTEXT_WINDOW_OPEN_TAG}\nThread id: "); + let context_window_prefix = format!("{CONTEXT_WINDOW_OPEN_TAG}\nAgent name: "); request .message_input_texts("developer") .into_iter() @@ -82,13 +82,10 @@ fn token_budget_contexts(request: &ResponsesRequest) -> Vec { .collect() } -fn token_budget_window_ids( - text: &str, - thread_id: codex_protocol::ThreadId, -) -> (String, Option, String) { +fn token_budget_window_ids(text: &str, agent_name: &str) -> (String, Option, String) { 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}})(?:\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, ); @@ -223,11 +220,10 @@ async fn token_budget_context_is_only_emitted_with_full_context() -> Result<()> let requests = responses.requests(); assert_eq!(requests.len(), 2); - let thread_id = test.session_configured.thread_id; let initial_token_budget = token_budget_contexts(&requests[0]); assert_eq!(initial_token_budget.len(), 1); let (first_window_id, previous_window_id, window_id) = - token_budget_window_ids(&initial_token_budget[0], thread_id); + token_budget_window_ids(&initial_token_budget[0], "/root"); assert_eq!(previous_window_id, None); assert_eq!(first_window_id, window_id); assert_eq!( @@ -731,7 +727,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], ); @@ -915,11 +911,10 @@ async fn get_context_remaining_returns_token_budget_remaining_fragment() -> Resu "get_context_remaining should be exposed when token budget is enabled" ); - let thread_id = test.session_configured.thread_id; let remaining_context = "You have 6500 tokens left in this context window.".to_string(); let token_budgets = token_budget_contexts(&requests[1]); assert_eq!(token_budgets.len(), 1); - token_budget_window_ids(&token_budgets[0], thread_id); + token_budget_window_ids(&token_budgets[0], "/root"); assert_eq!( requests[2].function_call_output_content_and_success(call_id), Some((Some(remaining_context), None)) @@ -1099,18 +1094,17 @@ async fn token_budget_context_uses_new_window_after_compaction() -> Result<()> { "token budget compaction should not call server-side compaction" ); - let thread_id = test.session_configured.thread_id; let initial_token_budget = token_budget_contexts(&requests[0]); assert_eq!(initial_token_budget.len(), 1); let (initial_first_window_id, initial_previous_window_id, initial_window_id) = - token_budget_window_ids(&initial_token_budget[0], thread_id); + token_budget_window_ids(&initial_token_budget[0], "/root"); let post_compaction_token_budget = token_budget_contexts(&requests[1]); assert_eq!(post_compaction_token_budget.len(), 1); let ( post_compaction_first_window_id, post_compaction_previous_window_id, post_compaction_window_id, - ) = token_budget_window_ids(&post_compaction_token_budget[0], thread_id); + ) = token_budget_window_ids(&post_compaction_token_budget[0], "/root"); assert_eq!(initial_previous_window_id, None); assert_eq!(initial_first_window_id, initial_window_id); assert_eq!(post_compaction_first_window_id, initial_first_window_id); @@ -1264,15 +1258,14 @@ async fn token_budget_mid_turn_auto_compaction_resets_before_active_follow_up() "fresh token-budget windows should drop active tool output with the prior history" ); - let thread_id = test.session_configured.thread_id; let initial_token_budget = token_budget_contexts(&requests[0]); assert_eq!(initial_token_budget.len(), 1); let (initial_first_window_id, _, initial_window_id) = - token_budget_window_ids(&initial_token_budget[0], thread_id); + token_budget_window_ids(&initial_token_budget[0], "/root"); let follow_up_token_budget = token_budget_contexts(&requests[1]); assert_eq!(follow_up_token_budget.len(), 1); let (follow_up_first_window_id, follow_up_previous_window_id, follow_up_window_id) = - token_budget_window_ids(&follow_up_token_budget[0], thread_id); + token_budget_window_ids(&follow_up_token_budget[0], "/root"); assert_eq!(follow_up_first_window_id, initial_first_window_id); assert_eq!( follow_up_previous_window_id.as_deref(), @@ -1364,16 +1357,15 @@ async fn token_budget_auto_compact_fallback_uses_buffer_until_new_context() -> R fallback_request.function_call_output_text(trigger_call_id), Some("You have 0 tokens left in this context window.".to_string()) ); - let thread_id = test.session_configured.thread_id; let initial_context = token_budget_contexts(&requests[0]); assert_eq!(token_budget_contexts(fallback_request), initial_context); assert_eq!(token_budget_contexts(&requests[2]), initial_context); assert!(requests[2].body_contains_text(AUTO_COMPACT_FALLBACK_PROMPT)); assert!(requests[2].body_contains_text(fallback_call_id)); - let (_, _, initial_window_id) = token_budget_window_ids(&initial_context[0], thread_id); + let (_, _, initial_window_id) = token_budget_window_ids(&initial_context[0], "/root"); let follow_up_context = token_budget_contexts(&requests[3]); let (_, previous_window_id, follow_up_window_id) = - token_budget_window_ids(&follow_up_context[0], thread_id); + token_budget_window_ids(&follow_up_context[0], "/root"); assert_eq!( previous_window_id.as_deref(), Some(initial_window_id.as_str()) @@ -1523,15 +1515,14 @@ async fn new_context_tool_skips_auto_compact_fallback() -> Result<()> { .any(|name| name == "new_context"), "new_context should be exposed when token budget is enabled" ); - let thread_id = test.session_configured.thread_id; let initial_token_budget = token_budget_contexts(&requests[0]); assert_eq!(initial_token_budget.len(), 1); let (initial_first_window_id, _, initial_window_id) = - token_budget_window_ids(&initial_token_budget[0], thread_id); + token_budget_window_ids(&initial_token_budget[0], "/root"); let new_window_token_budget = token_budget_contexts(&requests[2]); assert_eq!(new_window_token_budget.len(), 1); let (new_first_window_id, new_previous_window_id, new_window_id) = - token_budget_window_ids(&new_window_token_budget[0], thread_id); + token_budget_window_ids(&new_window_token_budget[0], "/root"); assert_eq!(new_first_window_id, initial_first_window_id); assert_eq!( new_previous_window_id.as_deref(), @@ -1552,7 +1543,6 @@ async fn new_context_tool_skips_auto_compact_fallback() -> Result<()> { &ContextSnapshotOptions::default(), ); let snapshot = snapshot - .replace(&thread_id.to_string(), "") .replace(&new_first_window_id, "") .replace(&new_window_id, ""); insta::assert_snapshot!(