Skip to content

Surface Codex reasoning tokens and rate limits in Anthropic usage #149

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
sommaht wants to merge 1 commit into
base: main
Choose a base branch
from
Open

Surface Codex reasoning tokens and rate limits in Anthropic usage #149

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
29 changes: 27 additions & 2 deletions src/providers/codex/translate/live_stream.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -10,7 +10,8 @@ use super::IncompleteResponsePolicy;
use super::read_rewrite::sanitize_read_args;
use super::reasoning_signature::{PendingReasoning, encode_reasoning_signature};
use super::reducer::{
CodexUsage, STOP_END_TURN, STOP_MAX_TOKENS, STOP_TOOL_USE, map_codex_usage_to_anthropic,
CodexRateLimits, CodexUsage, STOP_END_TURN, STOP_MAX_TOKENS, STOP_TOOL_USE,
map_codex_usage_to_anthropic, parse_codex_rate_limits,
};

const BUFFERED_READ_REPAIR_TRAILING_WHITESPACE_BYTES: usize = 1_024;
Expand Down Expand Up @@ -72,6 +73,9 @@ pub struct LiveStreamTranslator {
estimated_input_tokens: u64,
incomplete_response_policy: IncompleteResponsePolicy,
finished: bool,
// Latest `codex.rate_limits` event seen on this stream; handed to the
// usage mapper at finish so the terminal usage block carries the meter.
rate_limits: Option<CodexRateLimits>,
}

impl LiveStreamTranslator {
Expand Down Expand Up @@ -102,6 +106,7 @@ impl LiveStreamTranslator {
estimated_input_tokens,
incomplete_response_policy: IncompleteResponsePolicy::Error,
finished: false,
rate_limits: None,
}
}

Expand Down Expand Up @@ -134,6 +139,7 @@ impl LiveStreamTranslator {

match kind {
"codex.rate_limits" => {
self.rate_limits = parse_codex_rate_limits(payload).or(self.rate_limits.take());
self.emit_ping(traffic, &mut out);
}
"keepalive" | "response.created" | "response.in_progress" => {
Expand Down Expand Up @@ -912,7 +918,10 @@ impl LiveStreamTranslator {
self.close_open_blocks(traffic, out);
self.emit_web_searches(traffic, out);
self.ensure_message_start(traffic, out);
let usage = payload.get("response").map(parse_codex_usage);
let mut usage = payload.get("response").map(parse_codex_usage);
if let Some(usage) = usage.as_mut() {
usage.rate_limits = self.rate_limits.clone();
}
let stop_reason = if self.incomplete_response_policy
== IncompleteResponsePolicy::AllowMaxOutputTokens
&& is_standard_max_output_tokens_incomplete(payload)
Expand Down Expand Up @@ -1111,6 +1120,7 @@ fn parse_codex_usage(response: &serde_json::Value) -> CodexUsage {
.get("output_tokens_details")
.and_then(|d| d.get("reasoning_tokens"))
.and_then(|v| v.as_u64()),
rate_limits: None,
}
}

Expand Down Expand Up @@ -1624,6 +1634,21 @@ mod tests {
let out = String::from_utf8(out).unwrap();
assert!(out.contains("event: ping"));
assert!(!out.contains("event: error"));

let finish = translator
.accept(
&json!({
"type": "response.completed",
"response": {"id": "resp_1", "status": "completed", "usage": {"input_tokens": 5, "output_tokens": 2}}
}),
None,
)
.unwrap();
let finish = String::from_utf8(finish).unwrap();
assert!(
finish.contains(r#""codex_rate_limits":{"limit_reached":true}"#),
"{finish}"
);
}

#[test]
Expand Down
129 changes: 127 additions & 2 deletions src/providers/codex/translate/reducer.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -54,6 +54,56 @@ pub struct CodexUsage {
pub output_tokens: Option<u64>,
pub input_tokens_details_cached: Option<u64>,
pub output_tokens_details_reasoning: Option<u64>,
/// Latest `codex.rate_limits` snapshot seen on the stream that produced
/// this usage. Codex sends it as its own event before any usage exists,
/// so the translator carries it here for the mapper.
pub rate_limits: Option<CodexRateLimits>,
}

/// Subscription meter summary from a `codex.rate_limits` stream event.
#[derive(Debug, Clone, Default, PartialEq, serde::Serialize, serde::Deserialize)]
pub struct CodexRateLimits {
#[serde(default, skip_serializing_if = "Option::is_none")]
pub plan_type: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub limit_reached: Option<bool>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub primary: Option<CodexRateLimitWindow>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub secondary: Option<CodexRateLimitWindow>,
}

#[derive(Debug, Clone, Default, PartialEq, serde::Serialize, serde::Deserialize)]
pub struct CodexRateLimitWindow {
#[serde(default, skip_serializing_if = "Option::is_none")]
pub used_percent: Option<f64>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub window_minutes: Option<u64>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub reset_at: Option<u64>,
}

/// Parse a `codex.rate_limits` event payload. Returns `None` when the event
/// carries no `rate_limits` object.
pub fn parse_codex_rate_limits(payload: &serde_json::Value) -> Option<CodexRateLimits> {
let limits = payload.get("rate_limits")?.as_object()?;
let window = |name: &str| -> Option<CodexRateLimitWindow> {
let w = limits.get(name)?.as_object()?;
Some(CodexRateLimitWindow {
used_percent: w.get("used_percent").and_then(|v| v.as_f64()),
window_minutes: w.get("window_minutes").and_then(|v| v.as_u64()),
reset_at: w.get("reset_at").and_then(|v| v.as_u64()),
})
};
Some(CodexRateLimits {
plan_type: payload
.get("plan_type")
.and_then(|v| v.as_str())
.map(str::to_string),
limit_reached: limits.get("limit_reached").and_then(|v| v.as_bool()),
primary: window("primary"),
secondary: window("secondary"),
})
}

pub type StopReason = &'static str;
Expand Down Expand Up @@ -271,6 +321,7 @@ pub(crate) fn reduce_upstream_bytes_with_policy(
let mut active_thinking: Option<ActiveThinking> = None;
let mut saw_tool_use = false;
let mut final_usage: Option<CodexUsage> = None;
let mut latest_rate_limits: Option<CodexRateLimits> = None;
let mut response_id: Option<String> = None;
let mut terminal_type: Option<String> = None;
let mut continuation_eligible = false;
Expand Down Expand Up @@ -394,6 +445,7 @@ pub(crate) fn reduce_upstream_bytes_with_policy(
}

if t == "codex.rate_limits" {
latest_rate_limits = parse_codex_rate_limits(&p).or(latest_rate_limits);
out.push(ReducerEvent::Progress);
continue;
}
Expand Down Expand Up @@ -804,6 +856,9 @@ pub(crate) fn reduce_upstream_bytes_with_policy(
.and_then(|v| v.as_str())
.map(|s| s.to_string());
final_usage = p.get("response").map(parse_codex_usage);
if let Some(usage) = final_usage.as_mut() {
usage.rate_limits = latest_rate_limits.clone();
}
incomplete = response_is_incomplete_terminal(&p);
continuation_eligible =
(t == "response.completed" || t == "response.done") && !incomplete;
Expand Down Expand Up @@ -913,6 +968,7 @@ fn parse_codex_usage(response: &serde_json::Value) -> CodexUsage {
.get("output_tokens_details")
.and_then(|d| d.get("reasoning_tokens"))
.and_then(|v| v.as_u64()),
rate_limits: None,
}
}

Expand Down Expand Up @@ -1038,6 +1094,13 @@ pub fn map_codex_usage_to_anthropic(
cache_creation_input_tokens: 0,
cache_read_input_tokens: cached,
server_tool_use: None,
output_tokens_details: usage
.output_tokens_details_reasoning
.map(|reasoning_tokens| OutputTokensDetails {
reasoning_tokens,
thinking_tokens: reasoning_tokens,
}),
codex_rate_limits: usage.rate_limits.clone(),
};

if let Some(requests) = web_search_requests
Expand All @@ -1059,6 +1122,22 @@ pub struct AnthropicUsage {
pub cache_read_input_tokens: u64,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub server_tool_use: Option<WebSearchUsage>,
/// Reasoning tokens as reported by Codex. `thinking_tokens` mirrors the
/// field name Claude Code records for Anthropic models so per-turn
/// accounting in transcripts stays comparable across providers.
#[serde(default, skip_serializing_if = "Option::is_none")]
pub output_tokens_details: Option<OutputTokensDetails>,
/// Subscription meter as of this response, from the `codex.rate_limits`
/// event, so a downstream proxy can account per turn. Claude Code ignores
/// the unknown key.
#[serde(default, skip_serializing_if = "Option::is_none")]
pub codex_rate_limits: Option<CodexRateLimits>,
}

#[derive(Debug, Clone, Default, serde::Serialize, serde::Deserialize)]
pub struct OutputTokensDetails {
pub reasoning_tokens: u64,
pub thinking_tokens: u64,
}

#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
Expand Down Expand Up @@ -1180,7 +1259,7 @@ mod tests {
"{}{}",
sse(
"codex.rate_limits",
json!({"rate_limits":{"limit_reached":true,"primary":{"reset_after_seconds":30}}}),
json!({"plan_type":"prolite","rate_limits":{"limit_reached":true,"primary":{"reset_after_seconds":30,"used_percent":42}}}),
),
sse(
"response.completed",
Expand All @@ -1189,7 +1268,19 @@ mod tests {
);
let out = reduce_upstream_bytes(upstream.as_bytes()).unwrap();
assert!(matches!(out.first(), Some(ReducerEvent::Progress)));
assert!(matches!(out.last(), Some(ReducerEvent::Finish { .. })));
let Some(ReducerEvent::Finish { usage, .. }) = out.last() else {
panic!("expected Finish");
};
let limits = usage
.as_ref()
.and_then(|usage| usage.rate_limits.as_ref())
.expect("rate limits carried into finish usage");
assert_eq!(limits.plan_type.as_deref(), Some("prolite"));
assert_eq!(limits.limit_reached, Some(true));
assert_eq!(
limits.primary.as_ref().and_then(|w| w.used_percent),
Some(42.0)
);
}

#[test]
Expand Down Expand Up @@ -1607,11 +1698,45 @@ mod tests {
output_tokens: Some(50),
input_tokens_details_cached: Some(20),
output_tokens_details_reasoning: None,
rate_limits: None,
};
let mapped = map_codex_usage_to_anthropic(&Some(usage), None);
assert_eq!(mapped.input_tokens, 80);
assert_eq!(mapped.output_tokens, 50);
assert_eq!(mapped.cache_read_input_tokens, 20);
assert!(mapped.output_tokens_details.is_none());
assert!(mapped.codex_rate_limits.is_none());
}

#[test]
fn map_usage_reports_reasoning_and_rate_limits() {
let usage = CodexUsage {
input_tokens: Some(100),
output_tokens: Some(50),
input_tokens_details_cached: None,
output_tokens_details_reasoning: Some(30),
rate_limits: Some(CodexRateLimits {
plan_type: Some("prolite".into()),
limit_reached: Some(false),
primary: Some(CodexRateLimitWindow {
used_percent: Some(8.0),
window_minutes: Some(10080),
reset_at: Some(1_789_765_755),
}),
secondary: None,
}),
};
let mapped = map_codex_usage_to_anthropic(&Some(usage), None);
let details = mapped
.output_tokens_details
.as_ref()
.expect("reasoning details");
assert_eq!(details.reasoning_tokens, 30);
assert_eq!(details.thinking_tokens, 30);
let json = serde_json::to_value(&mapped).unwrap();
assert_eq!(json["codex_rate_limits"]["primary"]["used_percent"], 8.0);
assert_eq!(json["codex_rate_limits"]["plan_type"], "prolite");
assert!(json["codex_rate_limits"].get("secondary").is_none());
}

#[test]
Expand Down