core: wrap token budget window context (#29494)

Token-budget initial context carries thread and context-window lineage
that the model should treat as one structured context-window block.
Wrapping it in `<context_window>` makes that boundary explicit while
preserving the existing window id content.

Before this change, the window identifiers were injected as an untagged
developer text fragment:

```text
Thread id <THREAD_ID>.
First context window id: <FIRST_WINDOW_ID>
Current context window id: <WINDOW_ID>
Previous context window id: <PREVIOUS_WINDOW_ID>
```

After this change, the same payload is wrapped as a context-window
block:

```text
<context_window>
Thread id: <THREAD_ID>
First context window id: <FIRST_WINDOW_ID>
Current context window id: <WINDOW_ID>
Previous context window id: <PREVIOUS_WINDOW_ID>
</context_window>
```

This adds shared `CONTEXT_WINDOW_*_TAG` protocol constants, updates
`TokenBudgetContext` to render with those 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`
This commit is contained in:
Michael Bolin authored and GitHub committed 2026-06-22 23:37:49 +00:00
1 parent 3b32d861c5
commit f1945de3b7
6 files changed
+35 -10

No files matched your search

@@ -1,5 +1,7 @@
use super::ContextualUserFragment; use super::ContextualUserFragment;
use codex_protocol::ThreadId; use codex_protocol::ThreadId;
use codex_protocol::protocol::CONTEXT_WINDOW_CLOSE_TAG;
use codex_protocol::protocol::CONTEXT_WINDOW_OPEN_TAG;
use uuid::Uuid; use uuid::Uuid;
#[derive(Debug, Clone, PartialEq, Eq)] #[derive(Debug, Clone, PartialEq, Eq)]
@@ -39,7 +41,7 @@ impl ContextualUserFragment for TokenBudgetContext {
} }
fn type_markers() -> (&'static str, &'static str) { fn type_markers() -> (&'static str, &'static str) {
("", "") (CONTEXT_WINDOW_OPEN_TAG, CONTEXT_WINDOW_CLOSE_TAG)
} }
fn body(&self) -> String { fn body(&self) -> String {
@@ -47,7 +49,7 @@ impl ContextualUserFragment for TokenBudgetContext {
let first_window_id = self.first_window_id; let first_window_id = self.first_window_id;
let window_id = self.window_id; let window_id = self.window_id;
let mut lines = vec![ let mut lines = vec![
format!("Thread id {thread_id}."), format!("Thread id: {thread_id}"),
format!("First context window id: {first_window_id}"), format!("First context window id: {first_window_id}"),
format!("Current context window id: {window_id}"), format!("Current context window id: {window_id}"),
]; ];
@@ -57,7 +59,7 @@ impl ContextualUserFragment for TokenBudgetContext {
if let Some(mcp_result) = &self.mcp_result { if let Some(mcp_result) = &self.mcp_result {
lines.push(mcp_result.clone()); lines.push(mcp_result.clone());
} }
lines.join("\n") format!("\n{}\n", lines.join("\n"))
} }
} }
+2
View File
@@ -15,6 +15,7 @@ use codex_protocol::models::is_image_open_tag_text;
use codex_protocol::models::is_local_image_close_tag_text; use codex_protocol::models::is_local_image_close_tag_text;
use codex_protocol::models::is_local_image_open_tag_text; use codex_protocol::models::is_local_image_open_tag_text;
use codex_protocol::protocol::COLLABORATION_MODE_OPEN_TAG; use codex_protocol::protocol::COLLABORATION_MODE_OPEN_TAG;
use codex_protocol::protocol::CONTEXT_WINDOW_OPEN_TAG;
use codex_protocol::protocol::MULTI_AGENT_MODE_OPEN_TAG; use codex_protocol::protocol::MULTI_AGENT_MODE_OPEN_TAG;
use codex_protocol::protocol::REALTIME_CONVERSATION_OPEN_TAG; use codex_protocol::protocol::REALTIME_CONVERSATION_OPEN_TAG;
use codex_protocol::protocol::SKILLS_INSTRUCTIONS_OPEN_TAG; use codex_protocol::protocol::SKILLS_INSTRUCTIONS_OPEN_TAG;
@@ -36,6 +37,7 @@ const CONTEXTUAL_DEVELOPER_PREFIXES: &[&str] = &[
"<personality_spec>", "<personality_spec>",
// Keep recognizing token-budget wrappers persisted by older versions. // Keep recognizing token-budget wrappers persisted by older versions.
"<token_budget>", "<token_budget>",
CONTEXT_WINDOW_OPEN_TAG,
"<rollout_budget>", "<rollout_budget>",
]; ];
+16
View File
@@ -15,6 +15,8 @@ use codex_protocol::models::ReasoningItemContent;
use codex_protocol::models::ReasoningItemReasoningSummary; use codex_protocol::models::ReasoningItemReasoningSummary;
use codex_protocol::models::ResponseItem; use codex_protocol::models::ResponseItem;
use codex_protocol::models::WebSearchAction; use codex_protocol::models::WebSearchAction;
use codex_protocol::protocol::CONTEXT_WINDOW_CLOSE_TAG;
use codex_protocol::protocol::CONTEXT_WINDOW_OPEN_TAG;
use codex_protocol::protocol::SKILLS_INSTRUCTIONS_OPEN_TAG; use codex_protocol::protocol::SKILLS_INSTRUCTIONS_OPEN_TAG;
use codex_protocol::user_input::UserInput; use codex_protocol::user_input::UserInput;
use pretty_assertions::assert_eq; use pretty_assertions::assert_eq;
@@ -39,6 +41,20 @@ fn recognizes_legacy_token_budget_as_contextual_developer_content() {
assert!(!has_non_contextual_dev_message_content(&content)); assert!(!has_non_contextual_dev_message_content(&content));
} }
#[test]
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
{CONTEXT_WINDOW_CLOSE_TAG}"#
),
}];
assert!(is_contextual_dev_message_content(&content));
assert!(!has_non_contextual_dev_message_content(&content));
}
#[test] #[test]
fn parses_user_message_with_text_and_two_images() { fn parses_user_message_with_text_and_two_images() {
let img1 = "https://example.com/one.png".to_string(); let img1 = "https://example.com/one.png".to_string();
@@ -1,6 +1,6 @@
--- ---
source: core/tests/suite/token_budget.rs source: core/tests/suite/token_budget.rs
assertion_line: 588 assertion_line: 545
expression: snapshot expression: snapshot
--- ---
Scenario: New context window tool installs fresh full context before the next follow-up request. 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]: 00:message/developer[3]:
[01] <PERMISSIONS_INSTRUCTIONS> [01] <PERMISSIONS_INSTRUCTIONS>
[02] <SKILLS_INSTRUCTIONS> [02] <SKILLS_INSTRUCTIONS>
[03] Thread id <THREAD_ID>.\nFirst context window id: <FIRST_WINDOW_ID>\nCurrent context window id: <WINDOW_ID>\nPrevious context window id: <FIRST_WINDOW_ID> [03] <context_window>\nThread id: <THREAD_ID>\nFirst context window id: <FIRST_WINDOW_ID>\nCurrent context window id: <WINDOW_ID>\nPrevious context window id: <FIRST_WINDOW_ID>\n</context_window>
01:message/user:<ENVIRONMENT_CONTEXT:cwd=<CWD>> 01:message/user:<ENVIRONMENT_CONTEXT:cwd=<CWD>>
02:function_call/update_plan 02:function_call/update_plan
03:function_call_output:Plan updated 03:function_call_output:Plan updated
+6 -3
View File
@@ -4,6 +4,8 @@ use codex_config::types::McpServerTransportConfig;
use codex_core::config::TokenBudgetConfig; use codex_core::config::TokenBudgetConfig;
use codex_features::Feature; use codex_features::Feature;
use codex_model_provider_info::built_in_model_providers; use codex_model_provider_info::built_in_model_providers;
use codex_protocol::protocol::CONTEXT_WINDOW_CLOSE_TAG;
use codex_protocol::protocol::CONTEXT_WINDOW_OPEN_TAG;
use codex_protocol::protocol::EventMsg; use codex_protocol::protocol::EventMsg;
use codex_protocol::protocol::Op; use codex_protocol::protocol::Op;
use core_test_support::PathBufExt; use core_test_support::PathBufExt;
@@ -34,10 +36,11 @@ use std::time::Duration;
const CONFIGURED_CONTEXT_WINDOW: i64 = 128_000; const CONFIGURED_CONTEXT_WINDOW: i64 = 128_000;
fn token_budget_contexts(request: &ResponsesRequest) -> Vec<String> { fn token_budget_contexts(request: &ResponsesRequest) -> Vec<String> {
let context_window_prefix = format!("{CONTEXT_WINDOW_OPEN_TAG}\nThread id: ");
request request
.message_input_texts("developer") .message_input_texts("developer")
.into_iter() .into_iter()
.filter(|text| text.starts_with("Thread id ")) .filter(|text| text.starts_with(&context_window_prefix))
.collect() .collect()
} }
@@ -47,7 +50,7 @@ fn token_budget_window_ids(
) -> (String, Option<String>, String) { ) -> (String, Option<String>, String) {
let captures = assert_regex_match( let captures = assert_regex_match(
&format!( &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"^{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}$"
), ),
text, text,
); );
@@ -191,7 +194,7 @@ async fn token_budget_context_injects_plain_thread_hint_text() -> Result<()> {
assert_eq!(token_budgets.len(), 1); assert_eq!(token_budgets.len(), 1);
let captures = assert_regex_match( let captures = assert_regex_match(
&format!( &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"^{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], &token_budgets[0],
); );
+4 -2
View File
@@ -92,8 +92,8 @@ use crate::permissions::default_read_only_subpaths_for_writable_root;
pub use crate::request_permissions::RequestPermissionsArgs; pub use crate::request_permissions::RequestPermissionsArgs;
pub use crate::request_user_input::RequestUserInputEvent; pub use crate::request_user_input::RequestUserInputEvent;
/// Open/close tags for special user-input blocks. Used across crates to avoid /// Open/close tags for special context blocks. Used across crates to avoid duplicated hardcoded
/// duplicated hardcoded strings. /// strings.
pub const USER_INSTRUCTIONS_OPEN_TAG: &str = "<user_instructions>"; pub const USER_INSTRUCTIONS_OPEN_TAG: &str = "<user_instructions>";
pub const USER_INSTRUCTIONS_CLOSE_TAG: &str = "</user_instructions>"; pub const USER_INSTRUCTIONS_CLOSE_TAG: &str = "</user_instructions>";
pub const ENVIRONMENT_CONTEXT_OPEN_TAG: &str = "<environment_context>"; pub const ENVIRONMENT_CONTEXT_OPEN_TAG: &str = "<environment_context>";
@@ -110,6 +110,8 @@ pub const MULTI_AGENT_MODE_OPEN_TAG: &str = "<multi_agent_mode>";
pub const MULTI_AGENT_MODE_CLOSE_TAG: &str = "</multi_agent_mode>"; pub const MULTI_AGENT_MODE_CLOSE_TAG: &str = "</multi_agent_mode>";
pub const REALTIME_CONVERSATION_OPEN_TAG: &str = "<realtime_conversation>"; pub const REALTIME_CONVERSATION_OPEN_TAG: &str = "<realtime_conversation>";
pub const REALTIME_CONVERSATION_CLOSE_TAG: &str = "</realtime_conversation>"; pub const REALTIME_CONVERSATION_CLOSE_TAG: &str = "</realtime_conversation>";
pub const CONTEXT_WINDOW_OPEN_TAG: &str = "<context_window>";
pub const CONTEXT_WINDOW_CLOSE_TAG: &str = "</context_window>";
pub const USER_MESSAGE_BEGIN: &str = "## My request for Codex:"; pub const USER_MESSAGE_BEGIN: &str = "## My request for Codex:";
// TODO(anp): Replace `TurnEnvironmentSelection` with `PathUri` once path URIs carry environment // TODO(anp): Replace `TurnEnvironmentSelection` with `PathUri` once path URIs carry environment