Expose response usage metadata in completion events (#41087)

## What changed

- Parse optional `usage_metadata.amount` values from Responses API completion events without converting their string representation.
- Propagate the metadata through SSE and WebSocket streams, regular turns, and remote compaction.
- Add `usageMetadata` to the app-server `rawResponse/completed` notification and its generated schemas.

## Testing

- Cover missing, null, zero, and high-precision amounts, including completion events emitted during compaction.

GitOrigin-RevId: b16d81e8350eab0dc7ff613d0dc9dbed73041713
This commit is contained in:
Dylan Hurd
2026-08-27 07:50:05 +00:00
committed by copyberry
parent b592a0bfed
commit 2c4a95736b
25 changed files with 250 additions and 18 deletions
@@ -17253,6 +17253,16 @@
"type": "null"
}
]
},
"usageMetadata": {
"anyOf": [
{
"$ref": "#/definitions/v2/ResponseUsageMetadata"
},
{
"type": "null"
}
]
}
},
"required": [
@@ -18501,6 +18511,18 @@
}
]
},
"ResponseUsageMetadata": {
"description": "Usage metadata reported for one upstream response.",
"properties": {
"amount": {
"type": [
"string",
"null"
]
}
},
"type": "object"
},
"ResponsesApiWebSearchAction": {
"oneOf": [
{
@@ -13386,6 +13386,16 @@
"type": "null"
}
]
},
"usageMetadata": {
"anyOf": [
{
"$ref": "#/definitions/ResponseUsageMetadata"
},
{
"type": "null"
}
]
}
},
"required": [
@@ -14634,6 +14644,18 @@
}
]
},
"ResponseUsageMetadata": {
"description": "Usage metadata reported for one upstream response.",
"properties": {
"amount": {
"type": [
"string",
"null"
]
}
},
"type": "object"
},
"ResponsesApiWebSearchAction": {
"oneOf": [
{
@@ -1,6 +1,18 @@
{
"$schema": "http://json-schema.org/draft-07/schema#",
"definitions": {
"ResponseUsageMetadata": {
"description": "Usage metadata reported for one upstream response.",
"properties": {
"amount": {
"type": [
"string",
"null"
]
}
},
"type": "object"
},
"TokenUsageBreakdown": {
"properties": {
"cacheWriteInputTokens": {
@@ -59,6 +71,16 @@
"type": "null"
}
]
},
"usageMetadata": {
"anyOf": [
{
"$ref": "#/definitions/ResponseUsageMetadata"
},
{
"type": "null"
}
]
}
},
"required": [
@@ -1,10 +1,11 @@
// 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 { ResponseUsageMetadata } from "./ResponseUsageMetadata";
import type { TokenUsageBreakdown } from "./TokenUsageBreakdown";
/**
* Internal-only notification containing the exact usage from one upstream
* Responses API completion.
*/
export type RawResponseCompletedNotification = { threadId: string, turnId: string, responseId: string, usage: TokenUsageBreakdown | null, };
export type RawResponseCompletedNotification = { threadId: string, turnId: string, responseId: string, usage: TokenUsageBreakdown | null, usageMetadata: ResponseUsageMetadata | null, };
@@ -0,0 +1,8 @@
// 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.
/**
* Usage metadata reported for one upstream response.
*/
export type ResponseUsageMetadata = { amount: string | null, };
@@ -408,6 +408,7 @@ export type { RemoteControlEnableParams } from "./RemoteControlEnableParams";
export type { RemoteControlStatusChangedNotification } from "./RemoteControlStatusChangedNotification";
export type { RequestPermissionProfile } from "./RequestPermissionProfile";
export type { ResidencyRequirement } from "./ResidencyRequirement";
export type { ResponseUsageMetadata } from "./ResponseUsageMetadata";
export type { ReviewDelivery } from "./ReviewDelivery";
export type { ReviewStartParams } from "./ReviewStartParams";
export type { ReviewStartResponse } from "./ReviewStartResponse";
@@ -1841,6 +1841,23 @@ pub struct RawResponseCompletedNotification {
pub turn_id: String,
pub response_id: String,
pub usage: Option<TokenUsageBreakdown>,
pub usage_metadata: Option<ResponseUsageMetadata>,
}
/// Usage metadata reported for one upstream response.
#[derive(Serialize, Deserialize, Debug, Clone, PartialEq, JsonSchema, TS)]
#[serde(rename_all = "camelCase")]
#[ts(export_to = "v2/")]
pub struct ResponseUsageMetadata {
pub amount: Option<String>,
}
impl From<codex_protocol::ResponseUsageMetadata> for ResponseUsageMetadata {
fn from(value: codex_protocol::ResponseUsageMetadata) -> Self {
Self {
amount: value.amount,
}
}
}
#[derive(Serialize, Deserialize, Debug, Clone, PartialEq, JsonSchema, TS)]
@@ -1118,6 +1118,7 @@ pub(crate) async fn apply_bespoke_event_handling(
turn_id: event_turn_id,
response_id: raw_response_completed_event.response_id,
usage: raw_response_completed_event.token_usage.map(Into::into),
usage_metadata: raw_response_completed_event.usage_metadata.map(Into::into),
};
outgoing
.send_server_notification(ServerNotification::RawResponseCompleted(notification))
@@ -15,6 +15,7 @@ use codex_app_server_protocol::ItemStartedNotification;
use codex_app_server_protocol::JSONRPCError;
use codex_app_server_protocol::RawResponseCompletedNotification;
use codex_app_server_protocol::RequestId;
use codex_app_server_protocol::ResponseUsageMetadata;
use codex_app_server_protocol::ThreadCompactStartParams;
use codex_app_server_protocol::ThreadCompactStartResponse;
use codex_app_server_protocol::ThreadItem;
@@ -244,9 +245,11 @@ async fn thread_compact_start_triggers_compaction_and_returns_empty_response() -
skip_if_no_network!(Ok(()));
let server = responses::start_mock_server().await;
let mut completed = responses::ev_completed_with_tokens("r1", /*total_tokens*/ 200);
completed["response"]["usage_metadata"] = serde_json::json!({ "amount": "0.125" });
let sse = responses::sse(vec![
responses::ev_assistant_message("m1", "MANUAL_COMPACT_SUMMARY"),
responses::ev_completed_with_tokens("r1", /*total_tokens*/ 200),
completed,
]);
responses::mount_sse_sequence(&server, vec![sse]).await;
@@ -300,6 +303,9 @@ async fn thread_compact_start_triggers_compaction_and_returns_empty_response() -
thread_id,
turn_id: started.turn_id,
response_id: "r1".to_string(),
usage_metadata: Some(ResponseUsageMetadata {
amount: Some("0.125".to_string()),
}),
usage: Some(TokenUsageBreakdown {
total_tokens: 200,
input_tokens: 200,
@@ -799,27 +799,43 @@ 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(
Some(json!({ "amount": "0.12345678901234567890" })),
json!({ "amount": "0.12345678901234567890" });
"exact_amount"
)]
#[tokio::test]
async fn turn_start_emits_raw_response_completed_with_upstream_usage() -> Result<()> {
async fn turn_start_emits_raw_response_completed_with_upstream_usage(
upstream_metadata: Option<Value>,
expected_metadata: Value,
) -> Result<()> {
let server = responses::start_mock_server().await;
let mut completed = json!({
"type": "response.completed",
"response": {
"id": "resp-1",
"usage": {
"input_tokens": 30,
"input_tokens_details": { "cached_tokens": 11 },
"output_tokens": 7,
"output_tokens_details": { "reasoning_tokens": 3 },
"total_tokens": 37
}
}
});
if let Some(metadata) = upstream_metadata {
completed["response"]["usage_metadata"] = metadata;
}
let body = responses::sse(vec![
responses::ev_response_created("resp-1"),
responses::ev_assistant_message("msg-1", "Done"),
json!({
"type": "response.completed",
"response": {
"id": "resp-1",
"usage": {
"input_tokens": 30,
"input_tokens_details": { "cached_tokens": 11 },
"output_tokens": 7,
"output_tokens_details": { "reasoning_tokens": 3 },
"total_tokens": 37
}
}
}),
completed,
]);
responses::mount_sse_once(&server, body).await;
let response_mock = responses::mount_sse_once(&server, body).await;
let codex_home = TempDir::new()?;
MockResponsesConfig::new(&server.uri()).write(codex_home.path())?;
@@ -856,6 +872,13 @@ async fn turn_start_emits_raw_response_completed_with_upstream_usage() -> Result
mcp.read_stream_until_notification_message("rawResponse/completed"),
)
.await??;
assert_eq!(
notification
.params
.as_ref()
.and_then(|params| params.get("usageMetadata")),
Some(&expected_metadata),
);
let notification: codex_app_server_protocol::ServerNotification = notification.try_into()?;
let codex_app_server_protocol::ServerNotification::RawResponseCompleted(notification) =
notification
@@ -869,6 +892,7 @@ async fn turn_start_emits_raw_response_completed_with_upstream_usage() -> Result
thread_id: thread.id,
turn_id: turn.id,
response_id: "resp-1".to_string(),
usage_metadata: serde_json::from_value(expected_metadata)?,
usage: Some(TokenUsageBreakdown {
total_tokens: 37,
input_tokens: 30,
@@ -880,6 +904,7 @@ async fn turn_start_emits_raw_response_completed_with_upstream_usage() -> Result
}
);
response_mock.single_request();
Ok(())
}
+2
View File
@@ -1,4 +1,5 @@
use crate::error::ApiError;
use codex_protocol::ResponseUsageMetadata;
use codex_protocol::config_types::ReasoningSummary as ReasoningSummaryConfig;
use codex_protocol::config_types::Verbosity as VerbosityConfig;
use codex_protocol::models::ResponseItem;
@@ -113,6 +114,7 @@ pub enum ResponseEvent {
Completed {
response_id: String,
token_usage: Option<TokenUsage>,
usage_metadata: Option<ResponseUsageMetadata>,
/// Did the model affirmatively end its turn? Some providers do not set this,
/// so we rely on fallback logic when this is `None`.
end_turn: Option<bool>,
+11
View File
@@ -8,6 +8,7 @@ use crate::safety_buffering::treatment_from_headers;
use crate::telemetry::SseTelemetry;
use codex_client::ByteStream;
use codex_client::StreamResponse;
use codex_protocol::ResponseUsageMetadata;
use codex_protocol::models::ResponseItem;
use codex_protocol::protocol::MisalignmentErrorDetails;
use codex_protocol::protocol::ModelVerification;
@@ -118,6 +119,7 @@ struct ResponseCompleted {
id: String,
#[serde(default)]
usage: Option<ResponseCompletedUsage>,
usage_metadata: Option<ResponseUsageMetadata>,
#[serde(default)]
end_turn: Option<bool>,
}
@@ -481,6 +483,7 @@ pub fn process_responses_event(
return Ok(Some(ResponseEvent::Completed {
response_id: resp.id,
token_usage: resp.usage.map(Into::into),
usage_metadata: resp.usage_metadata,
end_turn: resp.end_turn,
}));
}
@@ -864,10 +867,12 @@ mod tests {
Ok(ResponseEvent::Completed {
response_id,
token_usage,
usage_metadata,
end_turn,
}) => {
assert_eq!(response_id, "resp1");
assert!(token_usage.is_none());
assert!(usage_metadata.is_none());
assert!(end_turn.is_none());
}
other => panic!("unexpected third event: {other:?}"),
@@ -1060,10 +1065,12 @@ mod tests {
Ok(ResponseEvent::Completed {
response_id,
token_usage,
usage_metadata,
end_turn,
}) => {
assert_eq!(response_id, "resp1");
assert!(token_usage.is_none());
assert!(usage_metadata.is_none());
assert!(end_turn.is_none());
}
other => panic!("unexpected event: {other:?}"),
@@ -1570,6 +1577,7 @@ mod tests {
ResponseEvent::Completed {
response_id,
token_usage: None,
usage_metadata: None,
end_turn: None,
} if response_id == "resp-1"
);
@@ -1607,6 +1615,7 @@ mod tests {
ResponseEvent::Completed {
response_id,
token_usage: None,
usage_metadata: None,
end_turn: None,
} if response_id == "resp-1"
);
@@ -1642,6 +1651,7 @@ mod tests {
ResponseEvent::Completed {
response_id,
token_usage: None,
usage_metadata: None,
end_turn: None,
} if response_id == "resp-1"
);
@@ -1677,6 +1687,7 @@ mod tests {
ResponseEvent::Completed {
response_id,
token_usage: None,
usage_metadata: None,
end_turn: None,
} if response_id == "resp-1"
);
@@ -113,6 +113,7 @@ async fn responses_stream_parses_items_and_completed_end_to_end() -> Result<()>
"type": "response.completed",
"response": {
"id": "resp1",
"usage_metadata": { "amount": "0.12345678901234567890" },
"usage": {
"input_tokens": 10,
"output_tokens": 5,
@@ -165,9 +166,16 @@ async fn responses_stream_parses_items_and_completed_end_to_end() -> Result<()>
ResponseEvent::Completed {
response_id,
token_usage,
usage_metadata,
end_turn,
} => {
assert_eq!(response_id, "resp1");
assert_eq!(
usage_metadata,
&Some(codex_protocol::ResponseUsageMetadata {
amount: Some("0.12345678901234567890".to_string()),
})
);
assert_eq!(
token_usage.as_ref().map(|usage| usage.total_tokens),
Some(15)
+2
View File
@@ -2184,6 +2184,7 @@ where
Ok(ResponseEvent::Completed {
response_id,
token_usage,
usage_metadata,
end_turn,
}) => {
feedback_tags!(last_model_response_id = &response_id);
@@ -2206,6 +2207,7 @@ where
.send(Ok(ResponseEvent::Completed {
response_id,
token_usage,
usage_metadata,
end_turn,
}))
.await
+1
View File
@@ -728,6 +728,7 @@ async fn response_stream_records_last_model_feedback_ids() {
Ok(ResponseEvent::Completed {
response_id: "resp-123".to_string(),
token_usage: None,
usage_metadata: None,
end_turn: Some(true),
}),
]);
+2
View File
@@ -774,6 +774,7 @@ async fn drain_to_completed(
Ok(ResponseEvent::Completed {
response_id,
token_usage,
usage_metadata,
..
}) => {
sess.send_event(
@@ -781,6 +782,7 @@ async fn drain_to_completed(
EventMsg::RawResponseCompleted(RawResponseCompletedEvent {
response_id,
token_usage: token_usage.clone(),
usage_metadata,
}),
)
.await;
+15
View File
@@ -38,6 +38,7 @@ use codex_context_fragments::to_annotated_content;
use codex_features::Feature;
use codex_history::CodexHarnessMetadata;
use codex_history::ResponseItemEnvelope;
use codex_protocol::ResponseUsageMetadata;
use codex_protocol::error::CodexErr;
use codex_protocol::error::CodexErrorDetails;
use codex_protocol::error::Result as CodexResult;
@@ -361,6 +362,7 @@ struct RemoteCompactionV2Output {
compaction_output: ResponseItem,
response_id: String,
token_usage: Option<TokenUsage>,
usage_metadata: Option<ResponseUsageMetadata>,
}
async fn run_remote_compaction_request_v2(
@@ -422,6 +424,7 @@ async fn collect_compaction_output(
let mut saw_completed = false;
let mut completed_response_id = None;
let mut completed_token_usage = None;
let mut completed_usage_metadata = None;
while let Some(event) = stream.next().await {
match event? {
ResponseEvent::OutputItemDone(item) => {
@@ -436,11 +439,13 @@ async fn collect_compaction_output(
ResponseEvent::Completed {
response_id,
token_usage,
usage_metadata,
..
} => {
saw_completed = true;
completed_response_id = Some(response_id);
completed_token_usage = token_usage;
completed_usage_metadata = usage_metadata;
break;
}
_ => {}
@@ -469,6 +474,7 @@ async fn collect_compaction_output(
compaction_output,
response_id,
token_usage: completed_token_usage,
usage_metadata: completed_usage_metadata,
})
}
@@ -1161,6 +1167,9 @@ mod tests {
total_tokens: 123_498,
codex_rollout_budget_units: None,
}),
usage_metadata: Some(codex_protocol::ResponseUsageMetadata {
amount: Some("0.125".to_string()),
}),
end_turn: Some(true),
}),
]);
@@ -1169,6 +1178,12 @@ mod tests {
.await
.expect("compaction should be collected");
assert_eq!(
output.usage_metadata,
Some(codex_protocol::ResponseUsageMetadata {
amount: Some("0.125".to_string()),
}),
);
assert_eq!(output.compaction_output, compaction);
assert_eq!(output.response_id, "resp-compact");
assert_eq!(
@@ -119,6 +119,7 @@ pub(super) async fn run_remote_compact_v2_attempt(
compaction_output,
response_id,
token_usage,
usage_metadata,
} = compaction_output_result?;
// TODO: Emit this before compaction output validation so malformed completed
// responses still surface their raw upstream usage.
@@ -127,6 +128,7 @@ pub(super) async fn run_remote_compact_v2_attempt(
EventMsg::RawResponseCompleted(RawResponseCompletedEvent {
response_id,
token_usage: token_usage.clone(),
usage_metadata,
}),
)
.await;
+2
View File
@@ -2565,6 +2565,7 @@ async fn try_run_sampling_request(
ResponseEvent::Completed {
response_id,
token_usage,
usage_metadata,
end_turn,
} => {
sess.services
@@ -2589,6 +2590,7 @@ async fn try_run_sampling_request(
EventMsg::RawResponseCompleted(RawResponseCompletedEvent {
response_id,
token_usage: token_usage.clone(),
usage_metadata,
}),
)
.await;
@@ -152,6 +152,53 @@ fn websocket_connection_metadata(harness: &WebsocketTestHarness) -> CodexRespons
)
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn responses_websocket_preserves_credit_usage_metadata() {
skip_if_no_network!();
let mut completed = ev_completed("resp-1");
completed["response"]["usage_metadata"] = json!({ "amount": "0.12345678901234567890" });
let server =
start_websocket_server(vec![vec![vec![ev_response_created("resp-1"), completed]]]).await;
let harness = websocket_harness_for_codex_backend(&server).await;
let mut client_session = harness.client.new_session();
let prompt = prompt_with_input(vec![message_item("hello")]);
let responses_metadata = turn_metadata(&harness, /*turn_id*/ None);
let mut stream = client_session
.stream(
&prompt,
&harness.model_info,
&harness.session_telemetry,
harness.effort.clone(),
harness.summary,
/*service_tier*/ None,
&responses_metadata,
&InferenceTraceContext::disabled(),
)
.await
.expect("websocket stream failed");
let mut usage_metadata = None;
while let Some(event) = stream.next().await {
if let ResponseEvent::Completed {
usage_metadata: metadata,
..
} = event.expect("websocket stream failed")
{
usage_metadata = metadata;
break;
}
}
assert_eq!(
usage_metadata,
Some(codex_protocol::ResponseUsageMetadata {
amount: Some("0.12345678901234567890".to_string()),
}),
);
assert_eq!(server.single_connection().len(), 1);
server.shutdown().await;
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn responses_websocket_streams_request() {
skip_if_no_network!();
+2
View File
@@ -2,6 +2,7 @@ pub mod account;
mod agent_path;
pub mod auth;
mod response_item_id;
mod response_usage;
mod sanitized_git_url;
mod session_id;
mod thread_id;
@@ -10,6 +11,7 @@ pub use agent_path::AgentPath;
pub use permission_profile_intersection::PermissionIntersectionError;
pub use permission_profile_intersection::intersect_effective_permission_profiles;
pub use response_item_id::ResponseItemId;
pub use response_usage::ResponseUsageMetadata;
pub use sanitized_git_url::SanitizedGitUrl;
pub use session_id::SessionId;
pub use thread_id::RolloutId;
+2 -1
View File
@@ -1882,13 +1882,14 @@ pub struct RawResponseItemEvent {
pub item: ResponseItem,
}
/// Exact usage reported by one upstream Responses API completion.
/// Exact usage and metadata reported by one upstream Responses API completion.
///
/// Unlike TokenCountEvent, this is not accumulated, estimated, or replayed.
#[derive(Debug, Clone, Deserialize, Serialize, TS, JsonSchema)]
pub struct RawResponseCompletedEvent {
pub response_id: String,
pub token_usage: Option<TokenUsage>,
pub usage_metadata: Option<crate::ResponseUsageMetadata>,
}
#[derive(Debug, Clone, Deserialize, Serialize, TS, JsonSchema)]
+12
View File
@@ -0,0 +1,12 @@
//! Per-response usage metadata reported by the upstream service, without aggregation.
use schemars::JsonSchema;
use serde::Deserialize;
use serde::Serialize;
use ts_rs::TS;
/// Usage metadata reported for one upstream response.
#[derive(Debug, Clone, Deserialize, Serialize, PartialEq, Eq, TS, JsonSchema)]
pub struct ResponseUsageMetadata {
pub amount: Option<String>,
}