Add rollout budget inspection tool

codex-rs/core/src/rollout_budget.rs

@@ -24,6 +24,14 @@ struct RolloutBudgetState {
     deliveries: HashMap<ThreadId, ThreadBudgetDelivery>,
 }
 
+impl RolloutBudgetState {
+    fn remaining_tokens(&self) -> i64 {
+        (self.config.limit_tokens as f64 - self.weighted_tokens_used)
+            .max(0.0)
+            .floor() as i64
+    }
+}
+
 struct ThreadBudgetDelivery {
     window_id: String,
     reminder_index: i64,
@@ -49,6 +57,10 @@ impl RolloutBudget {
             + usage.non_cached_input() as f64 * state.config.prefill_token_weight;
     }
 
+    pub(crate) fn remaining_tokens(&self) -> Option<i64> {
+        self.lock().map(|state| state.remaining_tokens())
+    }
+
     pub(crate) fn pending_reminder(
         &self,
         thread_id: ThreadId,
@@ -64,9 +76,7 @@ impl RolloutBudget {
             return None;
         }
         Some(RolloutBudgetReminder {
-            remaining_tokens: (state.config.limit_tokens as f64 - state.weighted_tokens_used)
-                .max(0.0)
-                .floor() as i64,
+            remaining_tokens: state.remaining_tokens(),
             reminder_index,
         })
     }

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

@@ -24,6 +24,7 @@ mod request_plugin_install;
 pub(crate) mod request_plugin_install_spec;
 mod request_user_input;
 pub(crate) mod request_user_input_spec;
+mod rollout_budget;
 mod shell;
 pub(crate) mod shell_spec;
 mod sleep;
@@ -67,6 +68,7 @@ pub use plan::PlanHandler;
 pub use request_permissions::RequestPermissionsHandler;
 pub use request_plugin_install::RequestPluginInstallHandler;
 pub use request_user_input::RequestUserInputHandler;
+pub use rollout_budget::RolloutBudgetHandler;
 pub use shell::ShellCommandHandler;
 pub(crate) use shell::ShellCommandHandlerOptions;
 pub use sleep::SleepHandler;

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

@@ -0,0 +1,78 @@
+use crate::context::ContextualUserFragment;
+use crate::context::RolloutBudgetContext;
+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::registry::CoreToolRuntime;
+use crate::tools::registry::ToolExecutor;
+use codex_tools::JsonSchema;
+use codex_tools::ResponsesApiNamespace;
+use codex_tools::ResponsesApiNamespaceTool;
+use codex_tools::ResponsesApiTool;
+use codex_tools::ToolName;
+use codex_tools::ToolSpec;
+use std::collections::BTreeMap;
+
+const NAMESPACE: &str = "rollout";
+const TOOL_NAME: &str = "remaining_budget";
+
+pub struct RolloutBudgetHandler;
+
+impl ToolExecutor<ToolInvocation> for RolloutBudgetHandler {
+    fn tool_name(&self) -> ToolName {
+        ToolName::namespaced(NAMESPACE, TOOL_NAME)
+    }
+
+    fn spec(&self) -> ToolSpec {
+        ToolSpec::Namespace(ResponsesApiNamespace {
+            name: NAMESPACE.to_string(),
+            description: "Tools for inspecting the current rollout state.".to_string(),
+            tools: vec![ResponsesApiNamespaceTool::Function(ResponsesApiTool {
+                name: TOOL_NAME.to_string(),
+                description:
+                    "Return the weighted tokens remaining in the shared session token budget."
+                        .to_string(),
+                strict: false,
+                defer_loading: None,
+                parameters: JsonSchema::object(
+                    BTreeMap::new(),
+                    /*required*/ None,
+                    /*additional_properties*/ Some(false.into()),
+                ),
+                output_schema: None,
+            })],
+        })
+    }
+
+    fn handle(&self, invocation: ToolInvocation) -> codex_tools::ToolExecutorFuture<'_> {
+        Box::pin(async move {
+            if !matches!(invocation.payload, ToolPayload::Function { .. }) {
+                return Err(FunctionCallError::RespondToModel(format!(
+                    "{TOOL_NAME} handler received unsupported payload"
+                )));
+            }
+
+            let Some(remaining_tokens) = invocation
+                .session
+                .services
+                .agent_control
+                .rollout_budget()
+                .remaining_tokens()
+            else {
+                return Err(FunctionCallError::RespondToModel(
+                    "rollout budget is not configured".to_string(),
+                ));
+            };
+            let output = RolloutBudgetContext { remaining_tokens }.render();
+
+            Ok(boxed_tool_output(FunctionToolOutput::from_text(
+                output,
+                /*success*/ Some(true),
+            )))
+        })
+    }
+}
+
+impl CoreToolRuntime for RolloutBudgetHandler {}

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

@@ -21,6 +21,7 @@ use crate::tools::handlers::ReadMcpResourceHandler;
 use crate::tools::handlers::RequestPermissionsHandler;
 use crate::tools::handlers::RequestPluginInstallHandler;
 use crate::tools::handlers::RequestUserInputHandler;
+use crate::tools::handlers::RolloutBudgetHandler;
 use crate::tools::handlers::ShellCommandHandler;
 use crate::tools::handlers::ShellCommandHandlerOptions;
 use crate::tools::handlers::SleepHandler;
@@ -711,6 +712,10 @@ fn add_core_utility_tools(context: &CoreToolPlanContext<'_>, planned_tools: &mut
         planned_tools.add(GetContextRemainingHandler);
     }
 
+    if features.enabled(Feature::RolloutBudget) {
+        planned_tools.add(RolloutBudgetHandler);
+    }
+
     if features.enabled(Feature::SleepTool) {
         planned_tools.add(SleepHandler);
     }

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

@@ -9,6 +9,7 @@ use core_test_support::responses::ev_assistant_message;
 use core_test_support::responses::ev_completed;
 use core_test_support::responses::ev_completed_with_tokens;
 use core_test_support::responses::ev_function_call;
+use core_test_support::responses::ev_function_call_with_namespace;
 use core_test_support::responses::ev_response_created;
 use core_test_support::responses::mount_sse_once_match;
 use core_test_support::responses::mount_sse_sequence;
@@ -302,3 +303,48 @@ async fn restates_the_current_remainder_after_rollback() -> Result<()> {
 
     Ok(())
 }
+
+#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
+async fn remaining_budget_tool_returns_the_current_remainder() -> Result<()> {
+    skip_if_no_network!(Ok(()));
+
+    const CALL_ID: &str = "remaining-budget";
+
+    let server = start_mock_server().await;
+    let responses = mount_sse_sequence(
+        &server,
+        vec![
+            sse(vec![
+                ev_response_created("resp-1"),
+                ev_completed_with_tokens("resp-1", /*total_tokens*/ 30),
+            ]),
+            sse(vec![
+                ev_response_created("resp-2"),
+                ev_function_call_with_namespace(CALL_ID, "rollout", "remaining_budget", "{}"),
+                ev_completed("resp-2"),
+            ]),
+            sse(vec![ev_response_created("resp-3"), ev_completed("resp-3")]),
+        ],
+    )
+    .await;
+    let test = test_codex()
+        .with_config(|config| {
+            config
+                .features
+                .enable(Feature::RolloutBudget)
+                .expect("test config should allow rollout budgets");
+            config.rollout_budget = Some(ROLLOUT_BUDGET);
+        })
+        .build(&server)
+        .await?;
+
+    test.submit_turn("use some budget").await?;
+    test.submit_turn("check the remaining budget").await?;
+
+    assert_eq!(
+        responses.requests()[2].function_call_output_text(CALL_ID),
+        Some(rollout_budget_message(/*remaining_tokens*/ 70))
+    );
+
+    Ok(())
+}