openai/codex
1
0
Fork 0
mirror of https://github.com/openai/codex.git synced 2026-10-19 17:28:46 +00:00
Code Issues Packages Projects Releases Wiki Activity

Identify agents by name in token budget context (#36815)

Browse Source 
## What changed

- Replace the thread ID in `<context_window>` metadata with the session's canonical agent path.
- Default sessions without an agent path to `/root`.
- Verify that root sessions emit `/root` and subagent sessions emit their own path, such as `/root/worker`.

GitOrigin-RevId: bcf057842ed31f93d554b7de063b6c03403a2594
This commit is contained in:
Boyang Niu
2026-08-04 00:51:53 +00:00
committed by copyberry
parent 60c722e075
commit 7431f10d0d
6 changed files with 39 additions and 31 deletions
Download Patch File Download Diff File Expand all files Collapse all files
codex-rs/core/src/context/token_budget_context.rs
+6 -2
View File
@@ -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<Uuid>,
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<Uuid>,
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}"),
];
codex-rs/core/src/event_mapping_tests.rs
+1 -1
View File
@@ -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}"#
),
}];
codex-rs/core/src/session/mod.rs
+3
View File
@@ -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,
codex-rs/core/src/session/tests.rs
+14 -3
View File
@@ -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("<context_window>\nAgent name: /root/worker\n")),
"expected subagent context window to include its canonical name, got {developer_messages:?}"
);
assert!(
developer_messages
.iter()
codex-rs/core/tests/suite/snapshots/all__suite__token_budget__token_budget_new_context_window_tool_full_context.snap
+1 -1
View File
@@ -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] <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] <context_window>\nAgent name: /root\nFirst context window id: <FIRST_WINDOW_ID>\nCurrent context window id: <WINDOW_ID>\nPrevious context window id: <FIRST_WINDOW_ID>\n</context_window>
[02] <SKILLS_INSTRUCTIONS>
[03] <PERMISSIONS_INSTRUCTIONS>
01:message/user:<ENVIRONMENT_CONTEXT:cwd=<CWD>>
codex-rs/core/tests/suite/token_budget.rs
+14 -24
View File
@@ -74,7 +74,7 @@ fn model_token_budget_config() -> ModelTokenBudgetConfig {
}
fn token_budget_contexts(request: &ResponsesRequest) -> Vec<String> {
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<String> {
.collect()
}
fn token_budget_window_ids(
text: &str,
thread_id: codex_protocol::ThreadId,
) -> (String, Option<String>, String) {
fn token_budget_window_ids(text: &str, agent_name: &str) -> (String, Option<String>, 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(), "<THREAD_ID>")
.replace(&new_first_window_id, "<FIRST_WINDOW_ID>")
.replace(&new_window_id, "<WINDOW_ID>");
insta::assert_snapshot!(