From 60ba7ddcb76456ae8cdd628a1a4d7e401c168ce9 Mon Sep 17 00:00:00 2001 From: Michael Bolin Date: Mon, 22 Jun 2026 15:39:56 -0700 Subject: [PATCH] core: wrap token budget window context Token-budget initial context carries thread and context-window lineage that the model should treat as one structured context-window block. Wrapping it in `` makes that boundary explicit while preserving the existing window id content. This updates `TokenBudgetContext` to render with `` markers, treats the new wrapper as contextual developer content when mapping history, and refreshes the token-budget request-shape assertions and snapshot. Verification: - `just test -p codex-core token_budget` - `just test -p codex-core recognizes_context_window_as_contextual_developer_content` --- codex-rs/core/src/context/token_budget_context.rs | 7 +++++-- codex-rs/core/src/event_mapping.rs | 1 + codex-rs/core/src/event_mapping_tests.rs | 12 ++++++++++++ ..._budget_new_context_window_tool_full_context.snap | 4 ++-- codex-rs/core/tests/suite/token_budget.rs | 6 +++--- 5 files changed, 23 insertions(+), 7 deletions(-) diff --git a/codex-rs/core/src/context/token_budget_context.rs b/codex-rs/core/src/context/token_budget_context.rs index 55a94b2b11..054e746a19 100644 --- a/codex-rs/core/src/context/token_budget_context.rs +++ b/codex-rs/core/src/context/token_budget_context.rs @@ -2,6 +2,9 @@ use super::ContextualUserFragment; use codex_protocol::ThreadId; use uuid::Uuid; +const CONTEXT_WINDOW_OPEN_TAG: &str = ""; +const CONTEXT_WINDOW_CLOSE_TAG: &str = ""; + #[derive(Debug, Clone, PartialEq, Eq)] pub(crate) struct TokenBudgetContext { thread_id: ThreadId, @@ -39,7 +42,7 @@ impl ContextualUserFragment for TokenBudgetContext { } fn type_markers() -> (&'static str, &'static str) { - ("", "") + (CONTEXT_WINDOW_OPEN_TAG, CONTEXT_WINDOW_CLOSE_TAG) } fn body(&self) -> String { @@ -57,7 +60,7 @@ impl ContextualUserFragment for TokenBudgetContext { if let Some(mcp_result) = &self.mcp_result { lines.push(mcp_result.clone()); } - lines.join("\n") + format!("\n{}\n", lines.join("\n")) } } diff --git a/codex-rs/core/src/event_mapping.rs b/codex-rs/core/src/event_mapping.rs index 4c57c5cadc..c4c07e95b4 100644 --- a/codex-rs/core/src/event_mapping.rs +++ b/codex-rs/core/src/event_mapping.rs @@ -36,6 +36,7 @@ const CONTEXTUAL_DEVELOPER_PREFIXES: &[&str] = &[ "", // Keep recognizing token-budget wrappers persisted by older versions. "", + "", "", ]; diff --git a/codex-rs/core/src/event_mapping_tests.rs b/codex-rs/core/src/event_mapping_tests.rs index f4c3ddd3a6..cb37c73bad 100644 --- a/codex-rs/core/src/event_mapping_tests.rs +++ b/codex-rs/core/src/event_mapping_tests.rs @@ -39,6 +39,18 @@ fn recognizes_legacy_token_budget_as_contextual_developer_content() { assert!(!has_non_contextual_dev_message_content(&content)); } +#[test] +fn recognizes_context_window_as_contextual_developer_content() { + let content = vec![ContentItem::InputText { + text: + "\nThread id 00000000-0000-0000-0000-000000000000.\n" + .to_string(), + }]; + + assert!(is_contextual_dev_message_content(&content)); + assert!(!has_non_contextual_dev_message_content(&content)); +} + #[test] fn parses_user_message_with_text_and_two_images() { let img1 = "https://example.com/one.png".to_string(); 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 c0c905625f..f4296217e8 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 @@ -1,6 +1,6 @@ --- source: core/tests/suite/token_budget.rs -assertion_line: 588 +assertion_line: 545 expression: snapshot --- Scenario: New context window tool installs fresh full context before the next follow-up request. @@ -9,7 +9,7 @@ Scenario: New context window tool installs fresh full context before the next fo 00:message/developer[3]: [01] [02] - [03] Thread id .\nFirst context window id: \nCurrent context window id: \nPrevious context window id: + [03] \nThread id .\nFirst context window id: \nCurrent context window id: \nPrevious context window id: \n 01:message/user:> 02:function_call/update_plan 03:function_call_output:Plan updated diff --git a/codex-rs/core/tests/suite/token_budget.rs b/codex-rs/core/tests/suite/token_budget.rs index ec86ce5946..e883958dad 100644 --- a/codex-rs/core/tests/suite/token_budget.rs +++ b/codex-rs/core/tests/suite/token_budget.rs @@ -37,7 +37,7 @@ fn token_budget_contexts(request: &ResponsesRequest) -> Vec { request .message_input_texts("developer") .into_iter() - .filter(|text| text.starts_with("Thread id ")) + .filter(|text| text.starts_with("\nThread id ")) .collect() } @@ -47,7 +47,7 @@ fn token_budget_window_ids( ) -> (String, Option, String) { let captures = assert_regex_match( &format!( - r"^Thread 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}}))?$" + r"^\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$" ), text, ); @@ -191,7 +191,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"^Thread 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$" + r"^\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$" ), &token_budgets[0], );