Include thread id in token budget context (#27663)

## Why

The token budget full-context fragment identifies the current context
window, but not the thread that owns that window. Including the thread
id makes the initial context-window metadata self-contained, and
`get_context_remaining` also needs to be usable from Code Mode without
forcing callers to parse the model-facing fragment string.

## What changed

- Include the session thread id in the initial `<token_budget>` context
fragment.
- Expose `get_context_remaining` as a Code Mode nested tool while
keeping `new_context` direct-model-only.
- Keep direct model-facing `get_context_remaining` output as the
existing `<token_budget>` text fragment.
- Return only `tokens_left` from the Code Mode structured result for
`get_context_remaining`.
- Update token-budget integration tests and add Code Mode coverage for
the structured result.

## Verification

- `just test -p codex-core token_budget`
- `just test -p codex-core
code_mode_get_context_remaining_returns_structured_result`
- `just test -p core_test_support redacted_text_mode_normalizes_uuids`

codex-rs/core/src/context/token_budget_context.rs

@@ -1,14 +1,17 @@
 use super::ContextualUserFragment;
+use codex_protocol::ThreadId;
 
 #[derive(Debug, Clone, PartialEq, Eq)]
 pub(crate) struct TokenBudgetContext {
+    thread_id: ThreadId,
     window_id: u64,
     tokens_left: i64,
 }
 
 impl TokenBudgetContext {
-    pub(crate) fn new(window_id: u64, tokens_left: i64) -> Self {
+    pub(crate) fn new(thread_id: ThreadId, window_id: u64, tokens_left: i64) -> Self {
         Self {
+            thread_id,
             window_id,
             tokens_left,
         }
@@ -29,10 +32,11 @@ impl ContextualUserFragment for TokenBudgetContext {
     }
 
     fn body(&self) -> String {
+        let thread_id = self.thread_id;
         let window_id = self.window_id;
         let tokens_left = self.tokens_left;
         format!(
-            "Current context window {window_id}.\nYou have {tokens_left} tokens left in this context window."
+            "Thread id {thread_id}.\nCurrent context window {window_id}.\nYou have {tokens_left} tokens left in this context window."
         )
     }
 }

codex-rs/core/src/session/mod.rs

@@ -2991,6 +2991,7 @@ impl Session {
         {
             developer_sections.push(
                 crate::context::TokenBudgetContext::new(
+                    self.thread_id(),
                     auto_compact_window_id,
                     model_context_window,
                 )

codex-rs/core/src/tools/handlers/get_context_remaining.rs

@@ -2,14 +2,59 @@ use crate::context::ContextualUserFragment;
 use crate::function_tool::FunctionCallError;
 use crate::tools::context::FunctionToolOutput;
 use crate::tools::context::ToolInvocation;
+use crate::tools::context::ToolOutput;
 use crate::tools::context::ToolPayload;
 use crate::tools::context::boxed_tool_output;
 use crate::tools::handlers::get_context_remaining_spec::GET_CONTEXT_REMAINING_TOOL_NAME;
 use crate::tools::handlers::get_context_remaining_spec::create_get_context_remaining_tool;
 use crate::tools::registry::CoreToolRuntime;
 use crate::tools::registry::ToolExecutor;
+use codex_protocol::models::ResponseInputItem;
 use codex_tools::ToolName;
 use codex_tools::ToolSpec;
+use serde_json::Value as JsonValue;
+use serde_json::json;
+
+#[derive(Debug, Clone)]
+struct GetContextRemainingOutput {
+    tokens_left: Option<i64>,
+}
+
+impl GetContextRemainingOutput {
+    fn new(tokens_left: Option<i64>) -> Self {
+        Self { tokens_left }
+    }
+
+    fn fragment(&self) -> String {
+        match self.tokens_left {
+            Some(tokens_left) => {
+                crate::context::TokenBudgetRemainingContext::new(tokens_left).render()
+            }
+            None => crate::context::TokenBudgetRemainingContext::unknown().render(),
+        }
+    }
+}
+
+impl ToolOutput for GetContextRemainingOutput {
+    fn log_preview(&self) -> String {
+        self.fragment()
+    }
+
+    fn success_for_logging(&self) -> bool {
+        true
+    }
+
+    fn to_response_item(&self, call_id: &str, payload: &ToolPayload) -> ResponseInputItem {
+        FunctionToolOutput::from_text(self.fragment(), Some(true))
+            .to_response_item(call_id, payload)
+    }
+
+    fn code_mode_result(&self, _payload: &ToolPayload) -> JsonValue {
+        json!({
+            "tokens_left": self.tokens_left,
+        })
+    }
+}
 
 pub struct GetContextRemainingHandler;
 
@@ -31,22 +76,18 @@ impl ToolExecutor<ToolInvocation> for GetContextRemainingHandler {
             }
 
             let Some(model_context_window) = invocation.turn.model_context_window() else {
-                let fragment = crate::context::TokenBudgetRemainingContext::unknown().render();
-                return Ok(boxed_tool_output(FunctionToolOutput::from_text(
-                    fragment,
-                    Some(true),
+                return Ok(boxed_tool_output(GetContextRemainingOutput::new(
+                    /*tokens_left*/ None,
                 )));
             };
             let active_context_tokens = invocation.session.get_total_token_usage().await.max(0);
             let tokens_left = model_context_window
                 .saturating_sub(active_context_tokens)
                 .max(0);
-            let fragment = crate::context::TokenBudgetRemainingContext::new(tokens_left).render();
 
-            Ok(boxed_tool_output(FunctionToolOutput::from_text(
-                fragment,
-                Some(true),
-            )))
+            Ok(boxed_tool_output(GetContextRemainingOutput::new(Some(
+                tokens_left,
+            ))))
         })
     }
 }

codex-rs/core/src/tools/handlers/get_context_remaining_spec.rs

@@ -1,6 +1,8 @@
 use codex_tools::JsonSchema;
 use codex_tools::ResponsesApiTool;
 use codex_tools::ToolSpec;
+use serde_json::Value;
+use serde_json::json;
 use std::collections::BTreeMap;
 
 pub(crate) const GET_CONTEXT_REMAINING_TOOL_NAME: &str = "get_context_remaining";
@@ -12,6 +14,23 @@ pub fn create_get_context_remaining_tool() -> ToolSpec {
         strict: false,
         defer_loading: None,
         parameters: JsonSchema::object(BTreeMap::new(), /*required*/ None, Some(false.into())),
-        output_schema: None,
+        output_schema: Some(get_context_remaining_output_schema()),
+    })
+}
+
+fn get_context_remaining_output_schema() -> Value {
+    json!({
+        "type": "object",
+        "properties": {
+            "tokens_left": {
+                "anyOf": [
+                    { "type": "integer" },
+                    { "type": "null" }
+                ],
+                "description": "Remaining tokens in the current context window, or null when unavailable."
+            }
+        },
+        "required": ["tokens_left"],
+        "additionalProperties": false
     })
 }

codex-rs/core/src/tools/spec_plan.rs

@@ -654,7 +654,7 @@ fn add_core_utility_tools(context: &CoreToolPlanContext<'_>, planned_tools: &mut
 
     if features.enabled(Feature::TokenBudget) {
         planned_tools.add_with_exposure(NewContextWindowHandler, ToolExposure::DirectModelOnly);
-        planned_tools.add_with_exposure(GetContextRemainingHandler, ToolExposure::DirectModelOnly);
+        planned_tools.add(GetContextRemainingHandler);
     }
 
     if tool_suggest_enabled(turn_context)

codex-rs/core/tests/suite/code_mode.rs

@@ -801,6 +801,47 @@ text(JSON.stringify(result));
     Ok(())
 }
 
+#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
+async fn code_mode_get_context_remaining_returns_structured_result() -> Result<()> {
+    skip_if_no_network!(Ok(()));
+
+    let server = responses::start_mock_server().await;
+    let (_test, second_mock) = run_code_mode_turn_with_config(
+        &server,
+        "use exec to get remaining context",
+        r#"
+const result = await tools.get_context_remaining({});
+text(JSON.stringify(result));
+"#,
+        |config| {
+            config.model_context_window = Some(10_000);
+            config
+                .features
+                .enable(Feature::TokenBudget)
+                .expect("test config should allow token budget");
+        },
+    )
+    .await?;
+
+    let req = second_mock.single_request();
+    let (output, success) = custom_tool_output_body_and_success(&req, "call-1");
+    assert_ne!(
+        success,
+        Some(false),
+        "exec get_context_remaining call failed unexpectedly: {output}"
+    );
+
+    let parsed: Value = serde_json::from_str(&output)?;
+    assert_eq!(
+        parsed,
+        serde_json::json!({
+            "tokens_left": 9500,
+        })
+    );
+
+    Ok(())
+}
+
 #[cfg_attr(windows, ignore = "flaky on windows")]
 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
 async fn code_mode_nested_tool_calls_can_run_in_parallel() -> Result<()> {

codex-rs/core/tests/suite/snapshots/all__suite__token_budget__token_budget_new_context_window_tool_full_context.snap

@@ -1,7 +1,6 @@
 ---
 source: core/tests/suite/token_budget.rs
-assertion_line: 300
-expression: "context_snapshot::format_labeled_requests_snapshot(\"New context window tool installs fresh full context before the next follow-up request.\",\n&[(\"Final Follow-Up Request\", &requests[2])],\n&ContextSnapshotOptions::default(),)"
+expression: snapshot
 ---
 Scenario: New context window tool installs fresh full context before the next follow-up request.
 
@@ -9,7 +8,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] <token_budget>\nCurrent context window 1.\nYou have 121600 tokens left in this context window.\n</token_budget>
+    [03] <token_budget>\nThread id <THREAD_ID>.\nCurrent context window 1.\nYou have 121600 tokens left in this context window.\n</token_budget>
 01:message/user:<ENVIRONMENT_CONTEXT:cwd=<CWD>>
 02:function_call/update_plan
 03:function_call_output:Plan updated

codex-rs/core/tests/suite/token_budget.rs

@@ -79,8 +79,9 @@ 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 expected = vec![format!(
-        "<token_budget>\nCurrent context window 0.\nYou have {EFFECTIVE_CONTEXT_WINDOW} tokens left in this context window.\n</token_budget>"
+        "<token_budget>\nThread id {thread_id}.\nCurrent context window 0.\nYou have {EFFECTIVE_CONTEXT_WINDOW} tokens left in this context window.\n</token_budget>"
     )];
     assert_eq!(
         token_budget_texts(&requests[0]),
@@ -142,8 +143,10 @@ async fn token_budget_remaining_context_emits_on_first_threshold_crossing() -> R
     let requests = responses.requests();
     assert_eq!(requests.len(), 5);
 
-    let full_context = "<token_budget>\nCurrent context window 0.\nYou have 9500 tokens left in this context window.\n</token_budget>"
-        .to_string();
+    let thread_id = test.session_configured.thread_id;
+    let full_context = format!(
+        "<token_budget>\nThread id {thread_id}.\nCurrent context window 0.\nYou have 9500 tokens left in this context window.\n</token_budget>"
+    );
     let threshold_25 =
         "<token_budget>\nYou have 7000 tokens left in this context window.\n</token_budget>"
             .to_string();
@@ -229,8 +232,10 @@ async fn get_context_remaining_returns_token_budget_remaining_fragment() -> Resu
         "get_context_remaining should be exposed when token budget is enabled"
     );
 
-    let full_context = "<token_budget>\nCurrent context window 0.\nYou have 9500 tokens left in this context window.\n</token_budget>"
-        .to_string();
+    let thread_id = test.session_configured.thread_id;
+    let full_context = format!(
+        "<token_budget>\nThread id {thread_id}.\nCurrent context window 0.\nYou have 9500 tokens left in this context window.\n</token_budget>"
+    );
     let remaining_context =
         "<token_budget>\nYou have 7000 tokens left in this context window.\n</token_budget>"
             .to_string();
@@ -356,10 +361,11 @@ async fn token_budget_context_uses_new_window_after_compaction() -> Result<()> {
     let requests = responses.requests();
     assert_eq!(requests.len(), 3);
 
+    let thread_id = test.session_configured.thread_id;
     assert_eq!(
         token_budget_texts(&requests[2]),
         vec![format!(
-            "<token_budget>\nCurrent context window 1.\nYou have {EFFECTIVE_CONTEXT_WINDOW} tokens left in this context window.\n</token_budget>"
+            "<token_budget>\nThread id {thread_id}.\nCurrent context window 1.\nYou have {EFFECTIVE_CONTEXT_WINDOW} tokens left in this context window.\n</token_budget>"
         )],
         "post-compaction full context should report context window 1"
     );
@@ -422,10 +428,11 @@ async fn new_context_tool_starts_new_window_before_follow_up() -> Result<()> {
             .any(|name| name == "new_context"),
         "new_context should be exposed when token budget is enabled"
     );
+    let thread_id = test.session_configured.thread_id;
     assert_eq!(
         token_budget_texts(&requests[2]),
         vec![format!(
-            "<token_budget>\nCurrent context window 1.\nYou have {EFFECTIVE_CONTEXT_WINDOW} tokens left in this context window.\n</token_budget>"
+            "<token_budget>\nThread id {thread_id}.\nCurrent context window 1.\nYou have {EFFECTIVE_CONTEXT_WINDOW} tokens left in this context window.\n</token_budget>"
         )]
     );
     assert!(
@@ -436,13 +443,15 @@ async fn new_context_tool_starts_new_window_before_follow_up() -> Result<()> {
         requests[2].function_call_output_text(continue_call_id),
         Some("Plan updated".to_string())
     );
+    let snapshot = context_snapshot::format_labeled_requests_snapshot(
+        "New context window tool installs fresh full context before the next follow-up request.",
+        &[("Final Follow-Up Request", &requests[2])],
+        &ContextSnapshotOptions::default(),
+    );
+    let snapshot = snapshot.replace(&thread_id.to_string(), "<THREAD_ID>");
     insta::assert_snapshot!(
         "token_budget_new_context_window_tool_full_context",
-        context_snapshot::format_labeled_requests_snapshot(
-            "New context window tool installs fresh full context before the next follow-up request.",
-            &[("Final Follow-Up Request", &requests[2])],
-            &ContextSnapshotOptions::default(),
-        )
+        snapshot
     );
 
     Ok(())