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 `<context_window>` makes that boundary explicit while preserving the existing window id content.

This updates `TokenBudgetContext` to render with `<context_window>` 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
2026-06-22 15:39:56 -07:00
parent 372b5ac7c2
commit 60ba7ddcb7
5 changed files with 23 additions and 7 deletions
@@ -2,6 +2,9 @@ use super::ContextualUserFragment;
use codex_protocol::ThreadId;
use uuid::Uuid;
const CONTEXT_WINDOW_OPEN_TAG: &str = "<context_window>";
const CONTEXT_WINDOW_CLOSE_TAG: &str = "</context_window>";
#[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"))
}
}
+1
View File
@@ -36,6 +36,7 @@ const CONTEXTUAL_DEVELOPER_PREFIXES: &[&str] = &[
"<personality_spec>",
// Keep recognizing token-budget wrappers persisted by older versions.
"<token_budget>",
"<context_window>",
"<rollout_budget>",
];
+12
View File
@@ -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:
"<context_window>\nThread id 00000000-0000-0000-0000-000000000000.\n</context_window>"
.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();
@@ -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] <PERMISSIONS_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>>
02:function_call/update_plan
03:function_call_output:Plan updated
+3 -3
View File
@@ -37,7 +37,7 @@ fn token_budget_contexts(request: &ResponsesRequest) -> Vec<String> {
request
.message_input_texts("developer")
.into_iter()
.filter(|text| text.starts_with("Thread id "))
.filter(|text| text.starts_with("<context_window>\nThread id "))
.collect()
}
@@ -47,7 +47,7 @@ fn token_budget_window_ids(
) -> (String, Option<String>, 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"^<context_window>\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>$"
),
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"^<context_window>\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>$"
),
&token_budgets[0],
);