Use provider-reported rollout budget units (#36715)

## What changed

- Charge `codex_rollout_budget_units` against the shared rollout budget when the provider includes it in response usage.
- Fall back to weighted input and output token accounting when provider units are absent.
- Reject non-finite or negative provider units as a fatal response error.

## Testing

- Cover provider units in reminder thresholds and local and remote compaction budget exhaustion.
- Verify invalid units fail without retrying the response.

GitOrigin-RevId: b452403e365985854d16f298d1ba46383e9892c4

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<bool> {
         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(

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);
         }

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<f64>,
+) -> 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;