github/codex
4
0
Fork 0
mirror of https://github.com/openai/codex.git synced 2026-10-20 00:18:49 +00:00
Code Issues Packages Projects Releases Wiki Activity

Capture rollout budget units from response usage (#36641)

Browse Source 
## 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
This commit is contained in:
rka-oai
2026-08-03 00:30:57 +00:00
committed by copyberry
parent 5157493c23
commit bb5054fe47
18 changed files with 63 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files
codex-rs/analytics/src/analytics_client_tests.rs
+1
View File
@@ -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,
},
}
}
codex-rs/app-server/src/bespoke_event_handling.rs
+2
View File
@@ -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),
};
codex-rs/app-server/tests/common/rollout.rs
+2
View File
@@ -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),
}),
codex-rs/app-server/tests/suite/v2/thread_resume.rs
+2
View File
@@ -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),
}),
codex-rs/codex-api/src/sse/responses.rs
+6 -1
View File
@@ -126,6 +126,8 @@ struct ResponseCompletedUsage {
output_tokens: i64,
output_tokens_details: Option<ResponseCompletedOutputTokensDetails>,
total_tokens: i64,
#[serde(default)]
codex_rollout_budget_units: Option<serde_json::Number>,
}
impl From<ResponseCompletedUsage> for TokenUsage {
@@ -141,6 +143,7 @@ impl From<ResponseCompletedUsage> 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),
}
);
}
codex-rs/codex-api/tests/sse_end_to_end.rs
+20 -2
View File
@@ -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:?}"),
codex-rs/core/src/compact_remote_v2.rs
+2
View File
@@ -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,
})
);
}
codex-rs/core/src/guardian/metrics.rs
+1
View File
@@ -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()
codex-rs/core/src/guardian/review_session.rs
+4
View File
@@ -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,
}
);
}
codex-rs/core/src/session/mod.rs
+1
View File
@@ -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() {
codex-rs/core/src/session/tests.rs
+7
View File
@@ -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;
codex-rs/core/src/tasks/mod.rs
+1
View File
@@ -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(
codex-rs/ext/goal/src/accounting.rs
+1
View File
@@ -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)
}
codex-rs/ext/goal/tests/accounting.rs
+1
View File
@@ -65,5 +65,6 @@ fn token_usage(
output_tokens,
reasoning_output_tokens,
total_tokens,
codex_rollout_budget_units: None,
}
}
codex-rs/ext/goal/tests/goal_extension_backend.rs
+1
View File
@@ -1462,6 +1462,7 @@ fn token_usage(
output_tokens,
reasoning_output_tokens,
total_tokens,
codex_rollout_budget_units: None,
}
}
codex-rs/memories/write/src/phase1.rs
+3
View File
@@ -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,
})
);
}
codex-rs/protocol/src/protocol.rs
+7
View File
@@ -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<serde_json::Number>,
}
#[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 =
codex-rs/state/src/runtime/threads.rs
+1
View File
@@ -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,