Preserve raw response usage metadata (#41980)

## What changed

- Include the complete upstream `response.usage` object in
  `ResponseUsageMetadata.metadata` while retaining the existing `amount` field.
- Expose the metadata through `rawResponse/completed` notifications and update
  the generated JSON and TypeScript schemas.

## Testing

- Cover SSE, WebSocket, turn, and compaction completion paths.

GitOrigin-RevId: 6d79b1b4c92ac11abf43ec98f380b4868288a701
This commit is contained in:
sayan-oai
2026-09-01 05:17:11 +00:00
committed by copyberry
parent d038f3448f
commit e017e93ace
14 changed files with 62 additions and 13 deletions
@@ -18732,7 +18732,8 @@
"string",
"null"
]
}
},
"metadata": true
},
"type": "object"
},
@@ -14804,7 +14804,8 @@
"string",
"null"
]
}
},
"metadata": true
},
"type": "object"
},
@@ -9,7 +9,8 @@
"string",
"null"
]
}
},
"metadata": true
},
"type": "object"
},
@@ -1,8 +1,9 @@
// GENERATED CODE! DO NOT MODIFY BY HAND!
// This file was generated by [ts-rs](https://github.com/Aleph-Alpha/ts-rs). Do not edit this file manually.
import type { JsonValue } from "../serde_json/JsonValue";
/**
* Usage metadata reported for one upstream response.
*/
export type ResponseUsageMetadata = { amount: string | null, };
export type ResponseUsageMetadata = { amount: string | null, metadata: JsonValue | null, };
@@ -1856,12 +1856,14 @@ pub struct RawResponseCompletedNotification {
#[ts(export_to = "v2/")]
pub struct ResponseUsageMetadata {
pub amount: Option<String>,
pub metadata: Option<JsonValue>,
}
impl From<codex_protocol::ResponseUsageMetadata> for ResponseUsageMetadata {
fn from(value: codex_protocol::ResponseUsageMetadata) -> Self {
Self {
amount: value.amount,
metadata: value.metadata,
}
}
}
@@ -254,6 +254,8 @@ async fn thread_compact_start_triggers_compaction_and_returns_empty_response() -
]);
let mut completed = responses::ev_completed_with_tokens("r1", /*total_tokens*/ 200);
completed["response"]["usage_metadata"] = serde_json::json!({ "amount": "0.125" });
completed["response"]["usage"]["extra"] = serde_json::json!({ "label": "example" });
let expected_metadata = completed["response"]["usage"].clone();
let sse = responses::sse(vec![
responses::ev_assistant_message("m1", "MANUAL_COMPACT_SUMMARY"),
completed,
@@ -338,6 +340,7 @@ async fn thread_compact_start_triggers_compaction_and_returns_empty_response() -
response_id: "r1".to_string(),
usage_metadata: Some(ResponseUsageMetadata {
amount: Some("0.125".to_string()),
metadata: Some(expected_metadata),
}),
usage: Some(TokenUsageBreakdown {
total_tokens: 200,
@@ -904,19 +904,24 @@ async fn turn_start_sends_service_tier_id_to_model_request() -> Result<()> {
Ok(())
}
#[test_case(None, json!(null); "without_usage_metadata")]
#[test_case(Some(json!({})), json!({ "amount": null }); "without_amount")]
#[test_case(Some(json!({ "amount": null })), json!({ "amount": null }); "null_amount")]
#[test_case(Some(json!({ "amount": "0" })), json!({ "amount": "0" }); "zero_amount")]
#[test_case(None, None, None; "without_usage_metadata")]
#[test_case(Some(json!({})), None, None; "without_amount")]
#[test_case(Some(json!({ "amount": null })), None, None; "null_amount")]
#[test_case(Some(json!({ "amount": "0" })), Some("0"), None; "zero_amount")]
#[test_case(
Some(json!({ "amount": "0.12345678901234567890" })),
json!({ "amount": "0.12345678901234567890" });
Some("0.12345678901234567890"),
Some(json!({ "label": "example", "items": [0, null, true] }));
"exact_amount"
)]
#[test_case(None, None, Some(json!(null)); "null_extra")]
#[test_case(None, None, Some(json!(0)); "zero_extra")]
#[test_case(None, None, Some(json!({ "label": "example" })); "nested_extra")]
#[tokio::test]
async fn turn_start_emits_raw_response_completed_with_upstream_usage(
upstream_metadata: Option<Value>,
expected_metadata: Value,
expected_amount: Option<&str>,
extra: Option<Value>,
) -> Result<()> {
let server = responses::start_mock_server().await;
let mut completed = json!({
@@ -935,6 +940,13 @@ async fn turn_start_emits_raw_response_completed_with_upstream_usage(
if let Some(metadata) = upstream_metadata {
completed["response"]["usage_metadata"] = metadata;
}
if let Some(extra) = extra {
completed["response"]["usage"]["extra"] = extra;
}
let expected_metadata = json!({
"amount": expected_amount,
"metadata": completed["response"]["usage"],
});
let body = responses::sse(vec![
responses::ev_response_created("resp-1"),
responses::ev_assistant_message("msg-1", "Done"),
+8 -1
View File
@@ -478,8 +478,15 @@ pub fn process_responses_event(
}
"response.completed" => {
if let Some(resp_val) = event.response {
let metadata = resp_val
.get("usage")
.filter(|usage| !usage.is_null())
.cloned();
match serde_json::from_value::<ResponseCompleted>(resp_val) {
Ok(resp) => {
Ok(mut resp) => {
if let Some(metadata) = metadata {
resp.usage_metadata.get_or_insert_default().metadata = Some(metadata);
}
return Ok(Some(ResponseEvent::Completed {
response_id: resp.id,
token_usage: resp.usage.map(Into::into),
@@ -118,11 +118,13 @@ async fn responses_stream_parses_items_and_completed_end_to_end() -> Result<()>
"input_tokens": 10,
"output_tokens": 5,
"total_tokens": 15,
"extra": { "label": "example", "items": [0, null, true] },
"codex_rollout_budget_units": 2.5
}
}
});
let expected_metadata = completed["response"]["usage"].clone();
let body = build_responses_body(vec![item1, item2, completed]);
let transport = FixtureSseTransport::new(body);
let client = ResponsesClient::new(transport, provider("openai"), Arc::new(NoAuth));
@@ -174,6 +176,7 @@ async fn responses_stream_parses_items_and_completed_end_to_end() -> Result<()>
usage_metadata,
&Some(codex_protocol::ResponseUsageMetadata {
amount: Some("0.12345678901234567890".to_string()),
metadata: Some(expected_metadata),
})
);
assert_eq!(
+15 -1
View File
@@ -1175,18 +1175,32 @@ mod tests {
}),
usage_metadata: Some(codex_protocol::ResponseUsageMetadata {
amount: Some("0.125".to_string()),
metadata: Some(serde_json::json!({ "extra": { "label": "example" } })),
}),
end_turn: Some(true),
}),
]);
let (sess, turn_context) = crate::session::tests::make_session_and_context().await;
let (sess, turn_context, rx) =
crate::session::tests::make_session_and_context_with_rx().await;
let output = collect_compaction_output(&sess, &turn_context, stream)
.await
.expect("compaction should be collected");
assert_eq!(output.compaction_output, compaction);
assert_eq!(output.response_id, "resp-compact");
let event = rx.recv().await.expect("raw response completion");
let EventMsg::RawResponseCompleted(completed) = event.msg else {
panic!("expected raw response completion, got {:?}", event.msg);
};
assert_eq!(completed.response_id, "resp-compact");
assert_eq!(
completed.usage_metadata,
Some(codex_protocol::ResponseUsageMetadata {
amount: Some("0.125".to_string()),
metadata: Some(serde_json::json!({ "extra": { "label": "example" } })),
})
);
assert_eq!(
output.token_usage,
Some(TokenUsage {
@@ -158,6 +158,7 @@ async fn responses_websocket_preserves_credit_usage_metadata() {
let mut completed = ev_completed("resp-1");
completed["response"]["usage_metadata"] = json!({ "amount": "0.12345678901234567890" });
let expected_metadata = completed["response"]["usage"].clone();
let server =
start_websocket_server(vec![vec![vec![ev_response_created("resp-1"), completed]]]).await;
let harness = websocket_harness_for_codex_backend(&server).await;
@@ -193,6 +194,7 @@ async fn responses_websocket_preserves_credit_usage_metadata() {
usage_metadata,
Some(codex_protocol::ResponseUsageMetadata {
amount: Some("0.12345678901234567890".to_string()),
metadata: Some(expected_metadata),
}),
);
assert_eq!(server.single_connection().len(), 1);
+3 -1
View File
@@ -3,10 +3,12 @@
use schemars::JsonSchema;
use serde::Deserialize;
use serde::Serialize;
use serde_json::Value;
use ts_rs::TS;
/// Usage metadata reported for one upstream response.
#[derive(Debug, Clone, Deserialize, Serialize, PartialEq, Eq, TS, JsonSchema)]
#[derive(Debug, Clone, Default, Deserialize, Serialize, PartialEq, Eq, TS, JsonSchema)]
pub struct ResponseUsageMetadata {
pub amount: Option<String>,
pub metadata: Option<Value>,
}