From ed36afc2167175c25cee97541b11bcdbc644809d Mon Sep 17 00:00:00 2001 From: Rohit Arunachalam Date: Thu, 18 Jun 2026 17:51:37 -0700 Subject: [PATCH] Add rollout budget inspection tool --- codex-rs/core/src/rollout_budget.rs | 16 +++- codex-rs/core/src/tools/handlers/mod.rs | 2 + .../core/src/tools/handlers/rollout_budget.rs | 78 +++++++++++++++++++ codex-rs/core/src/tools/spec_plan.rs | 5 ++ codex-rs/core/tests/suite/rollout_budget.rs | 46 +++++++++++ 5 files changed, 144 insertions(+), 3 deletions(-) create mode 100644 codex-rs/core/src/tools/handlers/rollout_budget.rs diff --git a/codex-rs/core/src/rollout_budget.rs b/codex-rs/core/src/rollout_budget.rs index accabac761..01ccd74da9 100644 --- a/codex-rs/core/src/rollout_budget.rs +++ b/codex-rs/core/src/rollout_budget.rs @@ -24,6 +24,14 @@ struct RolloutBudgetState { deliveries: HashMap, } +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 { + 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, }) } diff --git a/codex-rs/core/src/tools/handlers/mod.rs b/codex-rs/core/src/tools/handlers/mod.rs index d9c338e82d..7d7c8cdf52 100644 --- a/codex-rs/core/src/tools/handlers/mod.rs +++ b/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; diff --git a/codex-rs/core/src/tools/handlers/rollout_budget.rs b/codex-rs/core/src/tools/handlers/rollout_budget.rs new file mode 100644 index 0000000000..a387cdec33 --- /dev/null +++ b/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 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 {} diff --git a/codex-rs/core/src/tools/spec_plan.rs b/codex-rs/core/src/tools/spec_plan.rs index 660f3f4187..279b58338f 100644 --- a/codex-rs/core/src/tools/spec_plan.rs +++ b/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); } diff --git a/codex-rs/core/tests/suite/rollout_budget.rs b/codex-rs/core/tests/suite/rollout_budget.rs index c2a044eb52..86c06a8731 100644 --- a/codex-rs/core/tests/suite/rollout_budget.rs +++ b/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(()) +}