diff --git a/codex-rs/core/src/rollout_budget.rs b/codex-rs/core/src/rollout_budget.rs index 334e9e73a5..cb4b9f8bd4 100644 --- a/codex-rs/core/src/rollout_budget.rs +++ b/codex-rs/core/src/rollout_budget.rs @@ -1,5 +1,7 @@ use crate::config::RolloutBudgetConfig; use codex_protocol::ThreadId; +use codex_protocol::error::CodexErr; +use codex_protocol::error::Result as CodexResult; use codex_protocol::protocol::TokenUsage; use std::collections::HashMap; use std::sync::Mutex; @@ -41,14 +43,25 @@ impl RolloutBudget { } /// Returns true once the configured budget is exhausted, including on later calls. - pub(crate) fn record_usage(&self, usage: &TokenUsage) -> bool { + pub(crate) fn record_usage(&self, usage: &TokenUsage) -> CodexResult { let Some(mut state) = self.lock() else { - return false; + return Ok(false); }; - state.weighted_tokens_used += usage.output_tokens.max(0) as f64 - * state.config.sampling_token_weight - + usage.non_cached_input() as f64 * state.config.prefill_token_weight; - state.weighted_tokens_used >= state.config.limit_tokens as f64 + let units = if let Some(units) = usage.codex_rollout_budget_units.as_ref() { + let units = units.as_f64().unwrap_or(f64::NAN); + if !units.is_finite() || units < 0.0 { + return Err(CodexErr::Fatal( + "response.completed usage.codex_rollout_budget_units must be finite and non-negative" + .to_string(), + )); + } + units + } else { + usage.output_tokens.max(0) as f64 * state.config.sampling_token_weight + + usage.non_cached_input() as f64 * state.config.prefill_token_weight + }; + state.weighted_tokens_used += units; + Ok(state.weighted_tokens_used >= state.config.limit_tokens as f64) } pub(crate) fn pending_reminder( diff --git a/codex-rs/core/src/session/rollout_budget.rs b/codex-rs/core/src/session/rollout_budget.rs index 1ab5afb0fe..1867d7e3ec 100644 --- a/codex-rs/core/src/session/rollout_budget.rs +++ b/codex-rs/core/src/session/rollout_budget.rs @@ -28,7 +28,7 @@ impl Session { .services .agent_control .rollout_budget() - .record_usage(usage) + .record_usage(usage)? { return Err(CodexErr::SessionBudgetExceeded); } diff --git a/codex-rs/core/tests/suite/rollout_budget.rs b/codex-rs/core/tests/suite/rollout_budget.rs index ba199653e4..9bbf98f1fe 100644 --- a/codex-rs/core/tests/suite/rollout_budget.rs +++ b/codex-rs/core/tests/suite/rollout_budget.rs @@ -55,7 +55,11 @@ fn wire_request_contains(request: &wiremock::Request, text: &str) -> bool { } #[tokio::test(flavor = "multi_thread", worker_threads = 2)] -async fn adds_weighted_initial_and_threshold_reminders() -> Result<()> { +#[test_case(None ; "weighted token usage")] +#[test_case(Some(40.5) ; "provider budget units")] +async fn adds_weighted_initial_and_threshold_reminders( + rollout_budget_units: Option, +) -> Result<()> { skip_if_no_network!(Ok(())); let server = start_mock_server().await; @@ -73,7 +77,8 @@ async fn adds_weighted_initial_and_threshold_reminders() -> Result<()> { "input_tokens_details": { "cached_tokens": 40 }, "output_tokens": 15, "output_tokens_details": null, - "total_tokens": 75 + "total_tokens": 75, + "codex_rollout_budget_units": rollout_budget_units } } }), @@ -105,13 +110,72 @@ async fn adds_weighted_initial_and_threshold_reminders() -> Result<()> { rollout_budget_texts(&requests[1]), vec![ rollout_budget_message(/*remaining_tokens*/ 100), - rollout_budget_message(/*remaining_tokens*/ 60), + rollout_budget_message(if rollout_budget_units.is_some() { + 59 + } else { + 60 + }), ] ); Ok(()) } +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn invalid_provider_rollout_budget_units_fail_without_retry() -> Result<()> { + skip_if_no_network!(Ok(())); + + let server = start_mock_server().await; + let mut completed = ev_completed_with_tokens("invalid-units", /*total_tokens*/ 11); + completed["response"]["usage"]["codex_rollout_budget_units"] = json!(-1.0); + let responses = mount_sse_sequence( + &server, + vec![sse(vec![ev_response_created("invalid-units"), completed])], + ) + .await; + let test = test_codex() + .with_config(|config| { + config.rollout_budget = Some(rollout_budget()); + }) + .build(&server) + .await?; + + test.codex + .submit(Op::UserInput { + items: vec![UserInput::Text { + text: "reject invalid provider budget units".to_string(), + text_elements: Vec::new(), + }], + final_output_json_schema: None, + responsesapi_client_metadata: None, + additional_context: Default::default(), + thread_settings: Default::default(), + }) + .await?; + + let EventMsg::Error(error) = + wait_for_event(&test.codex, |event| matches!(event, EventMsg::Error(_))).await + else { + unreachable!(); + }; + assert_eq!( + error.message, + "Fatal error: response.completed usage.codex_rollout_budget_units must be finite and non-negative" + ); + assert_eq!(error.codex_error_info, Some(CodexErrorInfo::Other)); + wait_for_event(&test.codex, |event| { + matches!(event, EventMsg::TurnComplete(_)) + }) + .await; + assert_eq!( + responses.requests().len(), + 1, + "invalid units should not retry" + ); + + Ok(()) +} + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn subagent_usage_draws_from_the_shared_budget() -> Result<()> { skip_if_no_network!(Ok(())); @@ -278,12 +342,24 @@ async fn exhausted_budget_fails_current_and_later_turns() -> Result<()> { } #[tokio::test(flavor = "multi_thread", worker_threads = 2)] -#[test_case(false ; "local")] -#[test_case(true ; "remote_v2")] -async fn compaction_budget_exhaustion_fails_without_retry(remote_v2: bool) -> Result<()> { +#[test_case(false, false ; "local token usage")] +#[test_case(false, true ; "local provider units")] +#[test_case(true, false ; "remote v2 token usage")] +#[test_case(true, true ; "remote v2 provider units")] +async fn compaction_budget_exhaustion_fails_without_retry( + remote_v2: bool, + provider_units: bool, +) -> Result<()> { skip_if_no_network!(Ok(())); let server = start_mock_server().await; + let mut completed = ev_completed_with_tokens( + "compact", + /*total_tokens*/ if provider_units { 1 } else { 10 }, + ); + if provider_units { + completed["response"]["usage"]["codex_rollout_budget_units"] = json!(10.0); + } let compact_response = if remote_v2 { sse(vec![ json!({ @@ -293,13 +369,13 @@ async fn compaction_budget_exhaustion_fails_without_retry(remote_v2: bool) -> Re "encrypted_content": "encrypted-summary", } }), - ev_completed_with_tokens("compact", /*total_tokens*/ 10), + completed, ]) } else { sse(vec![ ev_response_created("compact"), ev_assistant_message("compact-summary", "compact summary"), - ev_completed_with_tokens("compact", /*total_tokens*/ 10), + completed, ]) }; let responses = mount_sse_sequence(&server, vec![compact_response]).await;