From bb5054fe47abe73ecbbd454751066a28c89f4bb9 Mon Sep 17 00:00:00 2001 From: rka-oai Date: Mon, 3 Aug 2026 00:25:52 +0000 Subject: [PATCH] Capture rollout budget units from response usage (#36641) ## What changed - Parse `codex_rollout_budget_units` from completed Responses API usage into `TokenUsage`. - Keep the provider-only value out of serialized protocol, JSON schema, and TypeScript representations. ## Testing - Cover the value in unit and end-to-end SSE response parsing tests. GitOrigin-RevId: a53d4202beb9c8985d25894cb10e7c01f20b2a44 --- .../analytics/src/analytics_client_tests.rs | 1 + .../app-server/src/bespoke_event_handling.rs | 2 ++ codex-rs/app-server/tests/common/rollout.rs | 2 ++ .../tests/suite/v2/thread_resume.rs | 2 ++ codex-rs/codex-api/src/sse/responses.rs | 7 +++++- codex-rs/codex-api/tests/sse_end_to_end.rs | 22 +++++++++++++++++-- codex-rs/core/src/compact_remote_v2.rs | 2 ++ codex-rs/core/src/guardian/metrics.rs | 1 + codex-rs/core/src/guardian/review_session.rs | 4 ++++ codex-rs/core/src/session/mod.rs | 1 + codex-rs/core/src/session/tests.rs | 7 ++++++ codex-rs/core/src/tasks/mod.rs | 1 + codex-rs/ext/goal/src/accounting.rs | 1 + codex-rs/ext/goal/tests/accounting.rs | 1 + .../ext/goal/tests/goal_extension_backend.rs | 1 + codex-rs/memories/write/src/phase1.rs | 3 +++ codex-rs/protocol/src/protocol.rs | 7 ++++++ codex-rs/state/src/runtime/threads.rs | 1 + 18 files changed, 63 insertions(+), 3 deletions(-) diff --git a/codex-rs/analytics/src/analytics_client_tests.rs b/codex-rs/analytics/src/analytics_client_tests.rs index 42c2eaa192..8ba7faafc3 100644 --- a/codex-rs/analytics/src/analytics_client_tests.rs +++ b/codex-rs/analytics/src/analytics_client_tests.rs @@ -389,6 +389,7 @@ fn sample_turn_token_usage_fact(thread_id: &str, turn_id: &str) -> TurnTokenUsag cache_write_input_tokens: 7, output_tokens: 140, reasoning_output_tokens: 13, + codex_rollout_budget_units: None, }, } } diff --git a/codex-rs/app-server/src/bespoke_event_handling.rs b/codex-rs/app-server/src/bespoke_event_handling.rs index e8ee2e5dc6..ae62eb7018 100644 --- a/codex-rs/app-server/src/bespoke_event_handling.rs +++ b/codex-rs/app-server/src/bespoke_event_handling.rs @@ -3843,6 +3843,7 @@ mod tests { output_tokens: 50, reasoning_output_tokens: 9, total_tokens: 200, + codex_rollout_budget_units: None, }, last_token_usage: TokenUsage { input_tokens: 10, @@ -3851,6 +3852,7 @@ mod tests { output_tokens: 7, reasoning_output_tokens: 1, total_tokens: 23, + codex_rollout_budget_units: None, }, model_context_window: Some(4096), }; diff --git a/codex-rs/app-server/tests/common/rollout.rs b/codex-rs/app-server/tests/common/rollout.rs index 7fae7f9735..d254d6ac84 100644 --- a/codex-rs/app-server/tests/common/rollout.rs +++ b/codex-rs/app-server/tests/common/rollout.rs @@ -122,6 +122,7 @@ pub fn create_fake_rollout_with_token_usage( output_tokens: 30, reasoning_output_tokens: 10, total_tokens: 150, + codex_rollout_budget_units: None, }, last_token_usage: TokenUsage { input_tokens: 70, @@ -130,6 +131,7 @@ pub fn create_fake_rollout_with_token_usage( output_tokens: 20, reasoning_output_tokens: 5, total_tokens: 90, + codex_rollout_budget_units: None, }, model_context_window: Some(200_000), }), diff --git a/codex-rs/app-server/tests/suite/v2/thread_resume.rs b/codex-rs/app-server/tests/suite/v2/thread_resume.rs index e2a21a4ccc..a4410da23b 100644 --- a/codex-rs/app-server/tests/suite/v2/thread_resume.rs +++ b/codex-rs/app-server/tests/suite/v2/thread_resume.rs @@ -2602,6 +2602,7 @@ async fn thread_resume_token_usage_replay_can_belong_to_interrupted_turn() -> Re output_tokens: 50, reasoning_output_tokens: 15, total_tokens: 230, + codex_rollout_budget_units: None, }, last_token_usage: TokenUsage { input_tokens: 90, @@ -2610,6 +2611,7 @@ async fn thread_resume_token_usage_replay_can_belong_to_interrupted_turn() -> Re output_tokens: 40, reasoning_output_tokens: 12, total_tokens: 130, + codex_rollout_budget_units: None, }, model_context_window: Some(200_000), }), diff --git a/codex-rs/codex-api/src/sse/responses.rs b/codex-rs/codex-api/src/sse/responses.rs index 441d7bb3bf..44ed78737b 100644 --- a/codex-rs/codex-api/src/sse/responses.rs +++ b/codex-rs/codex-api/src/sse/responses.rs @@ -126,6 +126,8 @@ struct ResponseCompletedUsage { output_tokens: i64, output_tokens_details: Option, total_tokens: i64, + #[serde(default)] + codex_rollout_budget_units: Option, } impl From for TokenUsage { @@ -141,6 +143,7 @@ impl From for TokenUsage { .map(|d| d.reasoning_tokens) .unwrap_or(0), total_tokens: val.total_tokens, + codex_rollout_budget_units: val.codex_rollout_budget_units, } } } @@ -816,7 +819,8 @@ mod tests { }, "output_tokens": 10, "output_tokens_details": { "reasoning_tokens": 5 }, - "total_tokens": 110 + "total_tokens": 110, + "codex_rollout_budget_units": 2.5 })) .expect("valid response usage"); @@ -829,6 +833,7 @@ mod tests { output_tokens: 10, reasoning_output_tokens: 5, total_tokens: 110, + codex_rollout_budget_units: serde_json::Number::from_f64(2.5), } ); } diff --git a/codex-rs/codex-api/tests/sse_end_to_end.rs b/codex-rs/codex-api/tests/sse_end_to_end.rs index 2526de281d..23871fb984 100644 --- a/codex-rs/codex-api/tests/sse_end_to_end.rs +++ b/codex-rs/codex-api/tests/sse_end_to_end.rs @@ -111,7 +111,15 @@ async fn responses_stream_parses_items_and_completed_end_to_end() -> Result<()> let completed = serde_json::json!({ "type": "response.completed", - "response": { "id": "resp1" } + "response": { + "id": "resp1", + "usage": { + "input_tokens": 10, + "output_tokens": 5, + "total_tokens": 15, + "codex_rollout_budget_units": 2.5 + } + } }); let body = build_responses_body(vec![item1, item2, completed]); @@ -160,7 +168,17 @@ async fn responses_stream_parses_items_and_completed_end_to_end() -> Result<()> end_turn, } => { assert_eq!(response_id, "resp1"); - assert!(token_usage.is_none()); + assert_eq!( + token_usage.as_ref().map(|usage| usage.total_tokens), + Some(15) + ); + assert_eq!( + token_usage + .as_ref() + .and_then(|usage| usage.codex_rollout_budget_units.as_ref()) + .and_then(serde_json::Number::as_f64), + Some(2.5) + ); assert!(end_turn.is_none()); } other => panic!("unexpected third event: {other:?}"), diff --git a/codex-rs/core/src/compact_remote_v2.rs b/codex-rs/core/src/compact_remote_v2.rs index da7502b9b8..34f48513ac 100644 --- a/codex-rs/core/src/compact_remote_v2.rs +++ b/codex-rs/core/src/compact_remote_v2.rs @@ -852,6 +852,7 @@ mod tests { output_tokens: 42, reasoning_output_tokens: 5, total_tokens: 123_498, + codex_rollout_budget_units: None, }), end_turn: Some(true), }), @@ -872,6 +873,7 @@ mod tests { output_tokens: 42, reasoning_output_tokens: 5, total_tokens: 123_498, + codex_rollout_budget_units: None, }) ); } diff --git a/codex-rs/core/src/guardian/metrics.rs b/codex-rs/core/src/guardian/metrics.rs index a0a461b4f4..f648695a4d 100644 --- a/codex-rs/core/src/guardian/metrics.rs +++ b/codex-rs/core/src/guardian/metrics.rs @@ -356,6 +356,7 @@ mod tests { output_tokens: 3, reasoning_output_tokens: 2, total_tokens: 15, + codex_rollout_budget_units: None, }), time_to_first_token_ms: Some(123), ..GuardianReviewAnalyticsResult::without_session() diff --git a/codex-rs/core/src/guardian/review_session.rs b/codex-rs/core/src/guardian/review_session.rs index 3887cc9c69..5b992bfa04 100644 --- a/codex-rs/core/src/guardian/review_session.rs +++ b/codex-rs/core/src/guardian/review_session.rs @@ -135,6 +135,7 @@ fn token_usage_delta(start: &TokenUsage, end: &TokenUsage) -> TokenUsage { reasoning_output_tokens: (end.reasoning_output_tokens - start.reasoning_output_tokens) .max(0), total_tokens: (end.total_tokens - start.total_tokens).max(0), + codex_rollout_budget_units: None, } } @@ -1684,6 +1685,7 @@ mod tests { output_tokens: 6, reasoning_output_tokens: 4, total_tokens: 28, + codex_rollout_budget_units: None, }; let end = TokenUsage { input_tokens: 15, @@ -1692,6 +1694,7 @@ mod tests { output_tokens: 10, reasoning_output_tokens: 2, total_tokens: 34, + codex_rollout_budget_units: None, }; assert_eq!( @@ -1703,6 +1706,7 @@ mod tests { output_tokens: 4, reasoning_output_tokens: 0, total_tokens: 6, + codex_rollout_budget_units: None, } ); } diff --git a/codex-rs/core/src/session/mod.rs b/codex-rs/core/src/session/mod.rs index 9660434f4b..56d7f8befe 100644 --- a/codex-rs/core/src/session/mod.rs +++ b/codex-rs/core/src/session/mod.rs @@ -3827,6 +3827,7 @@ impl Session { output_tokens: 0, reasoning_output_tokens: 0, total_tokens: estimated_total_tokens.max(0), + codex_rollout_budget_units: None, }; if let Some(model_context_window) = turn_context.model_context_window() { diff --git a/codex-rs/core/src/session/tests.rs b/codex-rs/core/src/session/tests.rs index 302e9bcc16..0ec52987a7 100644 --- a/codex-rs/core/src/session/tests.rs +++ b/codex-rs/core/src/session/tests.rs @@ -2572,6 +2572,7 @@ async fn record_initial_history_seeds_token_info_from_rollout() { output_tokens: 20, reasoning_output_tokens: 0, total_tokens: 30, + codex_rollout_budget_units: None, }, last_token_usage: TokenUsage { input_tokens: 3, @@ -2580,6 +2581,7 @@ async fn record_initial_history_seeds_token_info_from_rollout() { output_tokens: 4, reasoning_output_tokens: 0, total_tokens: 7, + codex_rollout_budget_units: None, }, model_context_window: Some(1_000), }; @@ -2591,6 +2593,7 @@ async fn record_initial_history_seeds_token_info_from_rollout() { output_tokens: 200, reasoning_output_tokens: 25, total_tokens: 375, + codex_rollout_budget_units: None, }, last_token_usage: TokenUsage { input_tokens: 10, @@ -2599,6 +2602,7 @@ async fn record_initial_history_seeds_token_info_from_rollout() { output_tokens: 20, reasoning_output_tokens: 5, total_tokens: 35, + codex_rollout_budget_units: None, }, model_context_window: Some(2_000), }; @@ -2768,6 +2772,7 @@ async fn record_token_usage_info_notifies_extension_contributors() { output_tokens: 20, reasoning_output_tokens: 3, total_tokens: 33, + codex_rollout_budget_units: None, }; let second_usage = TokenUsage { input_tokens: 7, @@ -2776,6 +2781,7 @@ async fn record_token_usage_info_notifies_extension_contributors() { output_tokens: 8, reasoning_output_tokens: 5, total_tokens: 20, + codex_rollout_budget_units: None, }; session @@ -2896,6 +2902,7 @@ async fn turn_start_lifecycle_exposes_turn_metadata_and_token_baseline() { output_tokens: 25, reasoning_output_tokens: 5, total_tokens: 130, + codex_rollout_budget_units: None, }; set_total_token_usage(&session, token_usage_at_turn_start.clone()).await; diff --git a/codex-rs/core/src/tasks/mod.rs b/codex-rs/core/src/tasks/mod.rs index e065686b46..bfd18ef09d 100644 --- a/codex-rs/core/src/tasks/mod.rs +++ b/codex-rs/core/src/tasks/mod.rs @@ -678,6 +678,7 @@ impl Session { total_tokens: (total_token_usage.total_tokens - token_usage_at_turn_start.total_tokens) .max(0), + codex_rollout_budget_units: None, }; let current_span = Span::current(); current_span.record( diff --git a/codex-rs/ext/goal/src/accounting.rs b/codex-rs/ext/goal/src/accounting.rs index a414c58e69..ea942ea158 100644 --- a/codex-rs/ext/goal/src/accounting.rs +++ b/codex-rs/ext/goal/src/accounting.rs @@ -324,6 +324,7 @@ fn token_delta_since_last_accounting(last: &TokenUsage, current: &TokenUsage) -> .reasoning_output_tokens .saturating_sub(last.reasoning_output_tokens), total_tokens: current.total_tokens.saturating_sub(last.total_tokens), + codex_rollout_budget_units: None, }; goal_token_delta_for_usage(&delta) } diff --git a/codex-rs/ext/goal/tests/accounting.rs b/codex-rs/ext/goal/tests/accounting.rs index c0485c6185..d5b601c997 100644 --- a/codex-rs/ext/goal/tests/accounting.rs +++ b/codex-rs/ext/goal/tests/accounting.rs @@ -65,5 +65,6 @@ fn token_usage( output_tokens, reasoning_output_tokens, total_tokens, + codex_rollout_budget_units: None, } } diff --git a/codex-rs/ext/goal/tests/goal_extension_backend.rs b/codex-rs/ext/goal/tests/goal_extension_backend.rs index 34fd5a7ade..2ab1b954b7 100644 --- a/codex-rs/ext/goal/tests/goal_extension_backend.rs +++ b/codex-rs/ext/goal/tests/goal_extension_backend.rs @@ -1462,6 +1462,7 @@ fn token_usage( output_tokens, reasoning_output_tokens, total_tokens, + codex_rollout_budget_units: None, } } diff --git a/codex-rs/memories/write/src/phase1.rs b/codex-rs/memories/write/src/phase1.rs index 69fdebdad0..2ba4ef28bb 100644 --- a/codex-rs/memories/write/src/phase1.rs +++ b/codex-rs/memories/write/src/phase1.rs @@ -830,6 +830,7 @@ mod tests { output_tokens: 3, reasoning_output_tokens: 1, total_tokens: 13, + codex_rollout_budget_units: None, }), }, JobResult { @@ -841,6 +842,7 @@ mod tests { output_tokens: 2, reasoning_output_tokens: 0, total_tokens: 9, + codex_rollout_budget_units: None, }), }, JobResult { @@ -862,6 +864,7 @@ mod tests { output_tokens: 5, reasoning_output_tokens: 1, total_tokens: 22, + codex_rollout_budget_units: None, }) ); } diff --git a/codex-rs/protocol/src/protocol.rs b/codex-rs/protocol/src/protocol.rs index 3155da6743..93313c63c2 100644 --- a/codex-rs/protocol/src/protocol.rs +++ b/codex-rs/protocol/src/protocol.rs @@ -2076,6 +2076,11 @@ pub struct TokenUsage { pub reasoning_output_tokens: i64, #[ts(type = "number")] pub total_tokens: i64, + /// Provider-reported units consumed from the shared rollout budget. + #[serde(default, skip_serializing)] + #[schemars(skip)] + #[ts(skip)] + pub codex_rollout_budget_units: Option, } #[derive(Debug, Clone, Deserialize, Serialize, PartialEq, Eq, JsonSchema, TS)] @@ -6309,6 +6314,7 @@ mod tests { output_tokens: 0, reasoning_output_tokens: 0, total_tokens: 10, + codex_rollout_budget_units: None, }); let info = TokenUsageInfo::new_or_append(&initial, &last, Some(128_000)) @@ -6331,6 +6337,7 @@ mod tests { output_tokens: 0, reasoning_output_tokens: 0, total_tokens: 10, + codex_rollout_budget_units: None, }); let info = diff --git a/codex-rs/state/src/runtime/threads.rs b/codex-rs/state/src/runtime/threads.rs index 8ae291fe9c..d65c404070 100644 --- a/codex-rs/state/src/runtime/threads.rs +++ b/codex-rs/state/src/runtime/threads.rs @@ -3217,6 +3217,7 @@ mod tests { output_tokens: 0, reasoning_output_tokens: 0, total_tokens: 321, + codex_rollout_budget_units: None, }, last_token_usage: codex_protocol::protocol::TokenUsage::default(), model_context_window: None,