Track prompt cache write token usage (#33454)

## What changed

- Parse `cache_write_tokens` from response input token details and carry it through token usage aggregation.
- Expose `cache_write_input_tokens` in protocol, app-server, exec, and TypeScript SDK usage events, defaulting it to zero for compatibility with older payloads.
- Report cache-write usage in analytics, metrics, traces, and rollout data alongside existing cache-read usage.

## Testing

- Add coverage for response parsing, usage event output, aggregation, analytics, metrics, and telemetry fields.

GitOrigin-RevId: 634e4fbe3086a8923c0ee13b622573e7d9d221bd
This commit is contained in:
Eric Traut
2026-07-16 00:22:53 +00:00
committed by copyberry
parent 78d4a563f9
commit 2edad72de3
46 changed files with 207 additions and 31 deletions
+1 -8
View File
@@ -2002,14 +2002,7 @@ where
}) => {
feedback_tags!(last_model_response_id = &response_id);
if let Some(usage) = &token_usage {
session_telemetry.sse_event_completed(
usage.input_tokens,
usage.output_tokens,
Some(usage.cached_input_tokens),
Some(usage.reasoning_output_tokens),
usage.total_tokens,
ttft_ms,
);
session_telemetry.sse_event_completed(usage, ttft_ms);
}
inference_trace_attempt.record_completed(
&response_id,
+3
View File
@@ -395,6 +395,7 @@ pub(crate) struct CompactionAnalyticsDetails {
pub(crate) retained_image_count: Option<usize>,
pub(crate) compaction_summary_tokens: Option<i64>,
pub(crate) cached_input_tokens: Option<i64>,
pub(crate) cache_write_input_tokens: Option<i64>,
}
impl CompactionAnalyticsAttempt {
@@ -432,6 +433,7 @@ impl CompactionAnalyticsAttempt {
retained_image_count,
compaction_summary_tokens,
cached_input_tokens,
cache_write_input_tokens,
} = details;
let active_context_tokens_before =
active_context_tokens_before.unwrap_or(self.active_context_tokens_before);
@@ -455,6 +457,7 @@ impl CompactionAnalyticsAttempt {
retained_image_count,
compaction_summary_tokens,
cached_input_tokens,
cache_write_input_tokens,
started_at: self.started_at,
completed_at: now_unix_seconds(),
duration_ms: Some(
+3
View File
@@ -277,6 +277,7 @@ async fn run_remote_compact_task_inner_impl(
analytics_details.active_context_tokens_before = Some(token_usage.input_tokens);
analytics_details.compaction_summary_tokens = Some(token_usage.output_tokens);
analytics_details.cached_input_tokens = Some(token_usage.cached_input_tokens);
analytics_details.cache_write_input_tokens = Some(token_usage.cache_write_input_tokens);
}
let (compacted_history, retained_images) =
build_v2_compacted_history(&prompt_input, compaction_output);
@@ -835,6 +836,7 @@ mod tests {
token_usage: Some(TokenUsage {
input_tokens: 123_456,
cached_input_tokens: 7_890,
cache_write_input_tokens: 0,
output_tokens: 42,
reasoning_output_tokens: 5,
total_tokens: 123_498,
@@ -854,6 +856,7 @@ mod tests {
Some(TokenUsage {
input_tokens: 123_456,
cached_input_tokens: 7_890,
cache_write_input_tokens: 0,
output_tokens: 42,
reasoning_output_tokens: 5,
total_tokens: 123_498,
+6
View File
@@ -60,6 +60,10 @@ fn emit_guardian_token_usage_histograms(
("total", token_usage.total_tokens.max(0)),
("input", token_usage.input_tokens.max(0)),
("cached_input", token_usage.cached_input()),
(
"cache_write_input",
token_usage.cache_write_input_tokens.max(0),
),
("non_cached_input", token_usage.non_cached_input()),
("output", token_usage.output_tokens.max(0)),
(
@@ -348,6 +352,7 @@ mod tests {
token_usage: Some(TokenUsage {
input_tokens: 10,
cached_input_tokens: 4,
cache_write_input_tokens: 2,
output_tokens: 3,
reasoning_output_tokens: 2,
total_tokens: 15,
@@ -399,6 +404,7 @@ mod tests {
histogram_sums(&snapshot, GUARDIAN_REVIEW_TOKEN_USAGE_METRIC),
BTreeMap::from([
("cached_input".to_string(), 4),
("cache_write_input".to_string(), 2),
("input".to_string(), 10),
("non_cached_input".to_string(), 6),
("output".to_string(), 3),
@@ -128,6 +128,8 @@ fn token_usage_delta(start: &TokenUsage, end: &TokenUsage) -> TokenUsage {
TokenUsage {
input_tokens: (end.input_tokens - start.input_tokens).max(0),
cached_input_tokens: (end.cached_input_tokens - start.cached_input_tokens).max(0),
cache_write_input_tokens: (end.cache_write_input_tokens - start.cache_write_input_tokens)
.max(0),
output_tokens: (end.output_tokens - start.output_tokens).max(0),
reasoning_output_tokens: (end.reasoning_output_tokens - start.reasoning_output_tokens)
.max(0),
@@ -1636,6 +1638,7 @@ mod tests {
let start = TokenUsage {
input_tokens: 10,
cached_input_tokens: 8,
cache_write_input_tokens: 8,
output_tokens: 6,
reasoning_output_tokens: 4,
total_tokens: 28,
@@ -1643,6 +1646,7 @@ mod tests {
let end = TokenUsage {
input_tokens: 15,
cached_input_tokens: 7,
cache_write_input_tokens: 7,
output_tokens: 10,
reasoning_output_tokens: 2,
total_tokens: 34,
@@ -1653,6 +1657,7 @@ mod tests {
TokenUsage {
input_tokens: 5,
cached_input_tokens: 0,
cache_write_input_tokens: 0,
output_tokens: 4,
reasoning_output_tokens: 0,
total_tokens: 6,
+1
View File
@@ -3713,6 +3713,7 @@ impl Session {
info.last_token_usage = TokenUsage {
input_tokens: 0,
cached_input_tokens: 0,
cache_write_input_tokens: 0,
output_tokens: 0,
reasoning_output_tokens: 0,
total_tokens: estimated_total_tokens.max(0),
+7
View File
@@ -2212,6 +2212,7 @@ async fn record_initial_history_seeds_token_info_from_rollout() {
total_token_usage: TokenUsage {
input_tokens: 10,
cached_input_tokens: 0,
cache_write_input_tokens: 0,
output_tokens: 20,
reasoning_output_tokens: 0,
total_tokens: 30,
@@ -2219,6 +2220,7 @@ async fn record_initial_history_seeds_token_info_from_rollout() {
last_token_usage: TokenUsage {
input_tokens: 3,
cached_input_tokens: 0,
cache_write_input_tokens: 0,
output_tokens: 4,
reasoning_output_tokens: 0,
total_tokens: 7,
@@ -2229,6 +2231,7 @@ async fn record_initial_history_seeds_token_info_from_rollout() {
total_token_usage: TokenUsage {
input_tokens: 100,
cached_input_tokens: 50,
cache_write_input_tokens: 0,
output_tokens: 200,
reasoning_output_tokens: 25,
total_tokens: 375,
@@ -2236,6 +2239,7 @@ async fn record_initial_history_seeds_token_info_from_rollout() {
last_token_usage: TokenUsage {
input_tokens: 10,
cached_input_tokens: 0,
cache_write_input_tokens: 0,
output_tokens: 20,
reasoning_output_tokens: 5,
total_tokens: 35,
@@ -2404,6 +2408,7 @@ async fn record_token_usage_info_notifies_extension_contributors() {
let first_usage = TokenUsage {
input_tokens: 10,
cached_input_tokens: 2,
cache_write_input_tokens: 0,
output_tokens: 20,
reasoning_output_tokens: 3,
total_tokens: 33,
@@ -2411,6 +2416,7 @@ async fn record_token_usage_info_notifies_extension_contributors() {
let second_usage = TokenUsage {
input_tokens: 7,
cached_input_tokens: 1,
cache_write_input_tokens: 0,
output_tokens: 8,
reasoning_output_tokens: 5,
total_tokens: 20,
@@ -2530,6 +2536,7 @@ async fn turn_start_lifecycle_exposes_turn_metadata_and_token_baseline() {
let token_usage_at_turn_start = TokenUsage {
input_tokens: 100,
cached_input_tokens: 40,
cache_write_input_tokens: 0,
output_tokens: 25,
reasoning_output_tokens: 5,
total_tokens: 130,
+1
View File
@@ -2025,6 +2025,7 @@ async fn try_run_sampling_request(
codex.request.reasoning_effort = %reasoning_effort,
gen_ai.usage.input_tokens = field::Empty,
gen_ai.usage.cache_read.input_tokens = field::Empty,
gen_ai.usage.cache_write.input_tokens = field::Empty,
gen_ai.usage.output_tokens = field::Empty,
codex.usage.reasoning_output_tokens = field::Empty,
codex.usage.total_tokens = field::Empty,
+13
View File
@@ -393,6 +393,7 @@ impl Session {
codex.turn.reasoning_effort = %reasoning_effort,
codex.turn.token_usage.input_tokens = field::Empty,
codex.turn.token_usage.cached_input_tokens = field::Empty,
codex.turn.token_usage.cache_write_input_tokens = field::Empty,
codex.turn.token_usage.non_cached_input_tokens = field::Empty,
codex.turn.token_usage.output_tokens = field::Empty,
codex.turn.token_usage.reasoning_output_tokens = field::Empty,
@@ -670,6 +671,9 @@ impl Session {
cached_input_tokens: (total_token_usage.cached_input_tokens
- token_usage_at_turn_start.cached_input_tokens)
.max(0),
cache_write_input_tokens: (total_token_usage.cache_write_input_tokens
- token_usage_at_turn_start.cache_write_input_tokens)
.max(0),
output_tokens: (total_token_usage.output_tokens
- token_usage_at_turn_start.output_tokens)
.max(0),
@@ -689,6 +693,10 @@ impl Session {
"codex.turn.token_usage.cached_input_tokens",
turn_token_usage.cached_input(),
);
current_span.record(
"codex.turn.token_usage.cache_write_input_tokens",
turn_token_usage.cache_write_input_tokens,
);
current_span.record(
"codex.turn.token_usage.non_cached_input_tokens",
turn_token_usage.non_cached_input(),
@@ -727,6 +735,11 @@ impl Session {
turn_token_usage.cached_input(),
&[("token_type", "cached_input"), tmp_mem],
);
self.services.session_telemetry.histogram(
TURN_TOKEN_USAGE_METRIC,
turn_token_usage.cache_write_input_tokens,
&[("token_type", "cache_write_input"), tmp_mem],
);
self.services.session_telemetry.histogram(
TURN_TOKEN_USAGE_METRIC,
turn_token_usage.output_tokens,
+2
View File
@@ -3295,6 +3295,7 @@ async fn token_count_includes_rate_limits_snapshot() {
"total_token_usage": {
"input_tokens": 123,
"cached_input_tokens": 0,
"cache_write_input_tokens": 0,
"output_tokens": 0,
"reasoning_output_tokens": 0,
"total_tokens": 123
@@ -3302,6 +3303,7 @@ async fn token_count_includes_rate_limits_snapshot() {
"last_token_usage": {
"input_tokens": 123,
"cached_input_tokens": 0,
"cache_write_input_tokens": 0,
"output_tokens": 0,
"reasoning_output_tokens": 0,
"total_tokens": 123
+11 -2
View File
@@ -575,7 +575,10 @@ async fn process_sse_emits_completed_telemetry() {
"id": "resp1",
"usage": {
"input_tokens": 3,
"input_tokens_details": { "cached_tokens": 1 },
"input_tokens_details": {
"cached_tokens": 1,
"cache_write_tokens": 2
},
"output_tokens": 5,
"output_tokens_details": { "reasoning_tokens": 2 },
"total_tokens": 9
@@ -613,6 +616,7 @@ async fn process_sse_emits_completed_telemetry() {
&& line.contains("input_token_count=3")
&& line.contains("output_token_count=5")
&& line.contains("cached_token_count=1")
&& line.contains("cache_write_token_count=2")
&& line.contains("reasoning_token_count=2")
&& line.contains("tool_token_count=9")
&& extract_log_field(line, "ttft_ms")
@@ -646,7 +650,10 @@ async fn turn_and_completed_response_spans_record_token_usage() {
"id": "resp1",
"usage": {
"input_tokens": 3,
"input_tokens_details": { "cached_tokens": 1 },
"input_tokens_details": {
"cached_tokens": 1,
"cache_write_tokens": 2
},
"output_tokens": 5,
"output_tokens_details": { "reasoning_tokens": 2 },
"total_tokens": 9
@@ -695,6 +702,7 @@ async fn turn_and_completed_response_spans_record_token_usage() {
&& line.contains("codex.request.reasoning_effort=high")
&& line.contains("gen_ai.usage.input_tokens=3")
&& line.contains("gen_ai.usage.cache_read.input_tokens=1")
&& line.contains("gen_ai.usage.cache_write.input_tokens=2")
&& line.contains("gen_ai.usage.output_tokens=5")
&& line.contains("codex.usage.reasoning_output_tokens=2")
&& line.contains("codex.usage.total_tokens=9")
@@ -707,6 +715,7 @@ async fn turn_and_completed_response_spans_record_token_usage() {
&& line.contains("codex.turn.reasoning_effort=high")
&& line.contains("codex.turn.token_usage.input_tokens=3")
&& line.contains("codex.turn.token_usage.cached_input_tokens=1")
&& line.contains("codex.turn.token_usage.cache_write_input_tokens=2")
&& line.contains("codex.turn.token_usage.non_cached_input_tokens=2")
&& line.contains("codex.turn.token_usage.output_tokens=5")
&& line.contains("codex.turn.token_usage.reasoning_output_tokens=2")