[codex] Add context remaining tool (#27518)

## Why

The token budget feature can inject remaining-context notices into
model-visible context, but the model does not have a direct way to ask
for that same remaining-token fragment on demand.

This PR adds a small model tool for the token budget feature so the
model can request the current remaining context window message without
duplicating the fragment format.

## What changed

- Adds a `get_context_remaining` direct-model tool behind
`Feature::TokenBudget`.
- Renders the tool output through `TokenBudgetRemainingContext`,
matching the existing budget message shape.
- Registers the tool alongside `new_context` in the token budget tool
set.
- Adds integration coverage that verifies the tool is exposed and
returns the same `<token_budget>` remaining fragment already present in
context.

## Validation

- `just test -p codex-core token_budget`

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

@@ -39,12 +39,18 @@ impl ContextualUserFragment for TokenBudgetContext {
 
 #[derive(Debug, Clone, PartialEq, Eq)]
 pub(crate) struct TokenBudgetRemainingContext {
-    tokens_left: i64,
+    tokens_left: Option<i64>,
 }
 
 impl TokenBudgetRemainingContext {
     pub(crate) fn new(tokens_left: i64) -> Self {
-        Self { tokens_left }
+        Self {
+            tokens_left: Some(tokens_left),
+        }
+    }
+
+    pub(crate) fn unknown() -> Self {
+        Self { tokens_left: None }
     }
 }
 
@@ -62,7 +68,11 @@ impl ContextualUserFragment for TokenBudgetRemainingContext {
     }
 
     fn body(&self) -> String {
-        let tokens_left = self.tokens_left;
-        format!("You have {tokens_left} tokens left in this context window.")
+        match self.tokens_left {
+            Some(tokens_left) => {
+                format!("You have {tokens_left} tokens left in this context window.")
+            }
+            None => "You have unknown tokens left in this context window.".to_string(),
+        }
     }
 }

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

@@ -0,0 +1,54 @@
+use crate::context::ContextualUserFragment;
+use crate::function_tool::FunctionCallError;
+use crate::tools::context::FunctionToolOutput;
+use crate::tools::context::ToolInvocation;
+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_tools::ToolName;
+use codex_tools::ToolSpec;
+
+pub struct GetContextRemainingHandler;
+
+impl ToolExecutor<ToolInvocation> for GetContextRemainingHandler {
+    fn tool_name(&self) -> ToolName {
+        ToolName::plain(GET_CONTEXT_REMAINING_TOOL_NAME)
+    }
+
+    fn spec(&self) -> ToolSpec {
+        create_get_context_remaining_tool()
+    }
+
+    fn handle(&self, invocation: ToolInvocation) -> codex_tools::ToolExecutorFuture<'_> {
+        Box::pin(async move {
+            if !matches!(invocation.payload, ToolPayload::Function { .. }) {
+                return Err(FunctionCallError::RespondToModel(
+                    "get_context_remaining handler received unsupported payload".to_string(),
+                ));
+            }
+
+            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),
+                )));
+            };
+            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),
+            )))
+        })
+    }
+}
+
+impl CoreToolRuntime for GetContextRemainingHandler {}

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

@@ -0,0 +1,17 @@
+use codex_tools::JsonSchema;
+use codex_tools::ResponsesApiTool;
+use codex_tools::ToolSpec;
+use std::collections::BTreeMap;
+
+pub(crate) const GET_CONTEXT_REMAINING_TOOL_NAME: &str = "get_context_remaining";
+
+pub fn create_get_context_remaining_tool() -> ToolSpec {
+    ToolSpec::Function(ResponsesApiTool {
+        name: GET_CONTEXT_REMAINING_TOOL_NAME.to_string(),
+        description: "Get the remaining tokens in the current context window.".to_string(),
+        strict: false,
+        defer_loading: None,
+        parameters: JsonSchema::object(BTreeMap::new(), /*required*/ None, Some(false.into())),
+        output_schema: None,
+    })
+}

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

@@ -4,6 +4,8 @@ pub(crate) mod apply_patch;
 pub(crate) mod apply_patch_spec;
 mod dynamic;
 pub(crate) mod extension_tools;
+mod get_context_remaining;
+pub(crate) mod get_context_remaining_spec;
 mod list_available_plugins_to_install;
 pub(crate) mod list_available_plugins_to_install_spec;
 mod mcp;
@@ -53,6 +55,7 @@ pub use apply_patch::ApplyPatchHandler;
 use codex_protocol::models::AdditionalPermissionProfile;
 use codex_protocol::protocol::AskForApproval;
 pub use dynamic::DynamicToolHandler;
+pub use get_context_remaining::GetContextRemainingHandler;
 pub use list_available_plugins_to_install::ListAvailablePluginsToInstallHandler;
 pub use mcp::McpHandler;
 pub use mcp_resource::ListMcpResourceTemplatesHandler;

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

@@ -9,6 +9,7 @@ use crate::tools::handlers::CodeModeWaitHandler;
 use crate::tools::handlers::DynamicToolHandler;
 use crate::tools::handlers::ExecCommandHandler;
 use crate::tools::handlers::ExecCommandHandlerOptions;
+use crate::tools::handlers::GetContextRemainingHandler;
 use crate::tools::handlers::ListAvailablePluginsToInstallHandler;
 use crate::tools::handlers::ListMcpResourceTemplatesHandler;
 use crate::tools::handlers::ListMcpResourcesHandler;
@@ -662,6 +663,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);
     }
 
     if tool_suggest_enabled(turn_context)

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

@@ -179,6 +179,136 @@ async fn token_budget_remaining_context_emits_on_first_threshold_crossing() -> R
     Ok(())
 }
 
+#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
+async fn get_context_remaining_returns_token_budget_remaining_fragment() -> Result<()> {
+    skip_if_no_network!(Ok(()));
+
+    let server = start_mock_server().await;
+    let call_id = "remaining-call";
+    let responses = mount_sse_sequence(
+        &server,
+        vec![
+            sse(vec![
+                ev_response_created("resp-1"),
+                ev_assistant_message("msg-1", "noted"),
+                ev_completed_with_tokens("resp-1", /*total_tokens*/ 2_500),
+            ]),
+            sse(vec![
+                ev_response_created("resp-2"),
+                ev_function_call(call_id, "get_context_remaining", "{}"),
+                ev_completed_with_tokens("resp-2", /*total_tokens*/ 2_500),
+            ]),
+            sse(vec![
+                ev_response_created("resp-3"),
+                ev_assistant_message("msg-3", "done"),
+                ev_completed("resp-3"),
+            ]),
+        ],
+    )
+    .await;
+    let test = test_codex()
+        .with_config(|config| {
+            config.model_context_window = Some(10_000);
+            config
+                .features
+                .enable(Feature::TokenBudget)
+                .expect("test config should allow token budget");
+        })
+        .build(&server)
+        .await?;
+
+    test.submit_turn("spend some tokens").await?;
+    test.submit_turn("check remaining context").await?;
+
+    let requests = responses.requests();
+    assert_eq!(requests.len(), 3);
+    assert!(
+        tool_names(&requests[1])
+            .iter()
+            .any(|name| name == "get_context_remaining"),
+        "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 remaining_context =
+        "<token_budget>\nYou have 7000 tokens left in this context window.\n</token_budget>"
+            .to_string();
+    assert_eq!(
+        token_budget_texts(&requests[1]),
+        vec![full_context, remaining_context.clone()]
+    );
+    assert_eq!(
+        requests[2].function_call_output_content_and_success(call_id),
+        Some((Some(remaining_context), None))
+    );
+
+    Ok(())
+}
+
+#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
+async fn get_context_remaining_returns_unknown_when_window_is_unavailable() -> Result<()> {
+    skip_if_no_network!(Ok(()));
+
+    let server = start_mock_server().await;
+    let call_id = "remaining-call";
+    let responses = mount_sse_sequence(
+        &server,
+        vec![
+            sse(vec![
+                ev_response_created("resp-1"),
+                ev_function_call(call_id, "get_context_remaining", "{}"),
+                ev_completed("resp-1"),
+            ]),
+            sse(vec![
+                ev_response_created("resp-2"),
+                ev_assistant_message("msg-2", "done"),
+                ev_completed("resp-2"),
+            ]),
+        ],
+    )
+    .await;
+    let test = test_codex()
+        .with_model_info_override("gpt-5.2", |model_info| {
+            model_info.context_window = None;
+            model_info.max_context_window = None;
+        })
+        .with_config(|config| {
+            config.model_context_window = None;
+            config
+                .features
+                .enable(Feature::TokenBudget)
+                .expect("test config should allow token budget");
+        })
+        .build(&server)
+        .await?;
+
+    test.submit_turn("check remaining context").await?;
+
+    let requests = responses.requests();
+    assert_eq!(requests.len(), 2);
+    assert!(
+        tool_names(&requests[0])
+            .iter()
+            .any(|name| name == "get_context_remaining"),
+        "get_context_remaining should be exposed when token budget is enabled"
+    );
+
+    assert_eq!(token_budget_texts(&requests[0]), Vec::<String>::new());
+    assert_eq!(
+        requests[1].function_call_output_content_and_success(call_id),
+        Some((
+            Some(
+                "<token_budget>\nYou have unknown tokens left in this context window.\n</token_budget>"
+                    .to_string()
+            ),
+            None,
+        ))
+    );
+
+    Ok(())
+}
+
 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
 async fn token_budget_context_uses_new_window_after_compaction() -> Result<()> {
     skip_if_no_network!(Ok(()));