i/Aether
1
0
Fork 0
mirror of https://github.com/fawney19/Aether.git synced 2026-09-01 17:00:21 +08:00
Code Issues Packages Projects Releases Wiki Activity

fix(usage): 流式终端 usage 以更完整值为准，Codex CLI 显式选择 response 解析器

Browse Source 
- 新增 StandardizedUsage::signal_score/is_more_complete_than/choose_more_complete，流式合并与终端落库均按信号完整度择优
- OpenAI Chat/CLI 解析器支持仅 usage 的终结 chunk 与 response.completed usage
- Codex provider 注入 provider_stream_event_api_format=openai:cli，解析器选择改由 report_context 显式决定
- usage_mapper 扩展嵌套 response/message/item 兼容 Claude message_start/message_delta 及 Gemini stream chunks
- usage SQL upsert 在终态（completed/failed/cancelled）时按 GREATEST 写入 token/费用镜像列
This commit is contained in:
fawney19
2026-04-25 00:57:01 +08:00
parent 2c6209277f
commit f30b1f3f6b
14 changed files with 656 additions and 72 deletions
Download Patch File Download Diff File Expand all files Collapse all files

1
crates/aether-ai-pipeline/src/finalize/standard/claude/stream.rs
View File

@@ -788,7 +788,6 @@ impl ClaudeClientEmitter {
)?);
Ok(out)
}
}
fn merge_claude_usage(mut current: CanonicalUsage, next: CanonicalUsage) -> CanonicalUsage {

111
crates/aether-ai-pipeline/src/finalize/standard/openai/stream.rs
View File

@@ -110,6 +110,19 @@ impl OpenAIChatProviderState {
let mut out = Vec::new();
let Some(chunk_choices) = chunk_object.get("choices").and_then(Value::as_array) else {
if let Some(usage) = Self::finish_usage(chunk_object.get("usage")) {
self.ensure_started(report_context, &mut out);
let (id, model) = self.identity(report_context);
out.push(CanonicalStreamFrame {
id,
model,
event: CanonicalStreamEvent::Finish {
finish_reason: self.pending_finish_reason.take(),
usage: Some(usage),
},
});
self.finished = true;
}
return Ok(out);
};
if chunk_choices.is_empty() {
@@ -1838,6 +1851,104 @@ mod tests {
assert_eq!(usage.cache_read_tokens, 19_840);
}
#[test]
fn openai_chat_provider_state_accepts_usage_only_terminal_chunk() {
let mut state = OpenAIChatProviderState::default();
let report_context = json!({});
let _ = state
.push_line(
&report_context,
data_line(json!({
"id": "chatcmpl_123",
"object": "chat.completion.chunk",
"model": "gpt-5.4",
"choices": [{
"index": 0,
"delta": {},
"finish_reason": "stop",
}],
})),
)
.expect("finish chunk should parse");
let frames = state
.push_line(
&report_context,
data_line(json!({
"usage": {
"input_tokens": 26,
"input_tokens_details": {
"cached_tokens": 0,
},
"output_tokens": 144,
"output_tokens_details": {
"reasoning_tokens": 10,
},
"total_tokens": 170,
},
})),
)
.expect("usage-only chunk should parse");
assert!(frames.iter().any(|frame| matches!(
frame.event,
CanonicalStreamEvent::Finish {
finish_reason: Some(ref reason),
usage: Some(CanonicalUsage {
input_tokens: 26,
output_tokens: 144,
cache_read_tokens: 0,
..
}),
} if reason == "stop"
)));
}
#[test]
fn openai_cli_provider_state_extracts_response_completed_usage() {
let mut state = OpenAICliProviderState::default();
let report_context = json!({});
let frames = state
.push_line(
&report_context,
data_line(json!({
"type": "response.completed",
"response": {
"id": "resp_063494bbd780be940169eb8191c4ec8191916347b2080805ee",
"object": "response",
"model": "gpt-5.5",
"status": "completed",
"output": [],
"usage": {
"input_tokens": 26,
"input_tokens_details": {
"cached_tokens": 0,
},
"output_tokens": 137,
"output_tokens_details": {
"reasoning_tokens": 0,
},
"total_tokens": 163,
},
},
"sequence_number": 139,
})),
)
.expect("completed event should parse");
assert!(frames.iter().any(|frame| matches!(
frame.event,
CanonicalStreamEvent::Finish {
usage: Some(CanonicalUsage {
input_tokens: 26,
output_tokens: 137,
cache_read_tokens: 0,
..
}),
..
}
)));
}
#[test]
fn openai_cli_client_emitter_emits_doc_like_text_events() {
let mut emitter = OpenAICliClientEmitter::default();

94
crates/aether-ai-pipeline/src/finalize/standard/stream_core/format_matrix.rs
View File

@@ -238,12 +238,15 @@ enum ClientStreamEmitter {
}
fn provider_api_format_for_context(report_context: &Value) -> String {
report_context
.get("provider_api_format")
.and_then(Value::as_str)
string_context_field(report_context, "provider_stream_event_api_format")
.or_else(|| string_context_field(report_context, "provider_stream_api_format"))
.or_else(|| string_context_field(report_context, "provider_api_format"))
.unwrap_or_default()
.trim()
.to_ascii_lowercase()
}
fn string_context_field(report_context: &Value, key: &str) -> Option<String> {
let value = report_context.get(key)?.as_str()?.trim();
(!value.is_empty()).then(|| value.to_ascii_lowercase())
}
fn client_api_format_for_context(report_context: &Value) -> String {
@@ -317,12 +320,7 @@ impl ClientStreamEmitter {
fn build_client_error_body_for_line(report_context: &Value, line: &[u8]) -> Option<Value> {
let value = decode_json_data_line(line)?;
let provider_api_format = report_context
.get("provider_api_format")
.and_then(Value::as_str)
.unwrap_or_default()
.trim()
.to_ascii_lowercase();
let provider_api_format = provider_api_format_for_context(report_context);
let client_api_format = report_context
.get("client_api_format")
.and_then(Value::as_str)
@@ -801,4 +799,78 @@ mod tests {
assert_eq!(usage.cache_creation_tokens, 42_262);
assert_eq!(usage.cache_read_tokens, 0);
}
#[test]
fn terminal_observer_uses_explicit_provider_stream_event_api_format() {
let mut report_context = report_context("openai:chat", "openai:cli");
report_context["provider_stream_event_api_format"] = json!("openai:cli");
let mut observer = StreamingStandardTerminalObserver::default();
observer
.push_line(
&report_context,
data_line(json!({
"type": "response.completed",
"response": {
"id": "resp_codex_123",
"object": "response",
"model": "gpt-5.5",
"status": "completed",
"output": [],
"usage": {
"input_tokens": 26,
"input_tokens_details": {
"cached_tokens": 0,
},
"output_tokens": 137,
"output_tokens_details": {
"reasoning_tokens": 0,
},
"total_tokens": 163,
},
},
"sequence_number": 139,
})),
)
.expect("response.completed should parse");
let summary = observer
.latest_summary()
.cloned()
.expect("summary should exist");
let usage = summary
.standardized_usage
.expect("standardized usage should exist");
assert_eq!(usage.input_tokens, 26);
assert_eq!(usage.output_tokens, 137);
assert_eq!(usage.cache_read_tokens, 0);
}
#[test]
fn terminal_observer_does_not_infer_provider_stream_event_api_format() {
let report_context = report_context("openai:chat", "openai:cli");
let mut observer = StreamingStandardTerminalObserver::default();
observer
.push_line(
&report_context,
data_line(json!({
"type": "response.completed",
"response": {
"usage": {
"input_tokens": 26,
"output_tokens": 137,
"total_tokens": 163,
},
},
})),
)
.expect("line should be ignored by explicitly selected chat parser");
assert!(
observer.latest_summary().is_none(),
"provider stream parser selection must come from report context, not event sniffing"
);
}
}

65
crates/aether-contracts/src/usage.rs
View File

@@ -83,6 +83,41 @@ impl StandardizedUsage {
}
self
}
pub fn signal_score(&self) -> usize {
[
self.input_tokens,
self.output_tokens,
self.cache_creation_tokens,
self.cache_creation_ephemeral_5m_tokens,
self.cache_creation_ephemeral_1h_tokens,
self.cache_read_tokens,
self.reasoning_tokens,
]
.into_iter()
.filter(|value| *value > 0)
.count()
+ self.dimensions.len()
}
pub fn has_token_signal(&self) -> bool {
self.signal_score() > 0
}
pub fn is_more_complete_than(&self, other: &Self) -> bool {
self.signal_score() > other.signal_score()
}
pub fn choose_more_complete(primary: Option<Self>, candidate: Option<Self>) -> Option<Self> {
match (primary, candidate) {
(Some(primary), Some(candidate)) if candidate.is_more_complete_than(&primary) => {
Some(candidate)
}
(Some(primary), _) => Some(primary),
(None, Some(candidate)) => Some(candidate),
(None, None) => None,
}
}
}
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, Default)]
@@ -116,6 +151,36 @@ fn as_f64(value: &serde_json::Value, default: f64) -> f64 {
mod tests {
use super::StandardizedUsage;
#[test]
fn standardized_usage_prefers_more_complete_candidate() {
let mut output_only = StandardizedUsage::new();
output_only.output_tokens = 131;
let mut complete = StandardizedUsage::new();
complete.input_tokens = 26;
complete.output_tokens = 131;
let selected = StandardizedUsage::choose_more_complete(Some(output_only), Some(complete))
.expect("usage should be selected");
assert_eq!(selected.input_tokens, 26);
assert_eq!(selected.output_tokens, 131);
}
#[test]
fn standardized_usage_keeps_primary_when_candidate_is_not_more_complete() {
let mut primary = StandardizedUsage::new();
primary.input_tokens = 26;
primary.output_tokens = 131;
let mut output_only = StandardizedUsage::new();
output_only.output_tokens = 131;
let selected = StandardizedUsage::choose_more_complete(Some(primary), Some(output_only))
.expect("usage should be selected");
assert_eq!(selected.input_tokens, 26);
assert_eq!(selected.output_tokens, 131);
}
#[test]
fn standardized_usage_reads_and_writes_known_and_extra_fields() {
let mut usage = StandardizedUsage::new();

32
crates/aether-data/src/repository/usage/sql/mod.rs
View File

@@ -6323,18 +6323,28 @@ ORDER BY "usage".user_id ASC
.bind(&usage.provider_endpoint_kind)
.bind(usage.has_format_conversion)
.bind(usage.is_stream)
.bind(None::<i32>)
.bind(None::<i32>)
.bind(None::<i32>)
.bind(None::<i32>)
.bind(None::<i32>)
.bind(None::<i32>)
.bind(None::<i32>)
.bind(None::<f64>)
.bind(None::<f64>)
.bind(None::<f64>)
.bind(None::<f64>)
.bind(usage.input_tokens.map(to_i32).transpose()?)
.bind(usage.output_tokens.map(to_i32).transpose()?)
.bind(usage.total_tokens.map(to_i32).transpose()?)
.bind(usage.cache_creation_input_tokens.map(to_i32).transpose()?)
.bind(
usage
.cache_creation_ephemeral_5m_input_tokens
.map(to_i32)
.transpose()?,
)
.bind(
usage
.cache_creation_ephemeral_1h_input_tokens
.map(to_i32)
.transpose()?,
)
.bind(usage.cache_read_input_tokens.map(to_i32).transpose()?)
.bind(usage.cache_creation_cost_usd)
.bind(usage.cache_read_cost_usd)
.bind(None::<f64>)
.bind(usage.total_cost_usd)
.bind(usage.actual_total_cost_usd)
.bind(usage.status_code.map(i32::from))
.bind(&usage.error_message)
.bind(&usage.error_category)

22
crates/aether-data/src/repository/usage/sql/queries/upsert_sql.sql
View File

@@ -135,18 +135,18 @@ DO UPDATE SET
provider_endpoint_kind = CASE WHEN "usage".billing_status = 'pending' THEN COALESCE(EXCLUDED.provider_endpoint_kind, "usage".provider_endpoint_kind) ELSE "usage".provider_endpoint_kind END,
has_format_conversion = CASE WHEN "usage".billing_status = 'pending' THEN COALESCE(EXCLUDED.has_format_conversion, "usage".has_format_conversion) ELSE "usage".has_format_conversion END,
is_stream = CASE WHEN "usage".billing_status = 'pending' THEN COALESCE(EXCLUDED.is_stream, "usage".is_stream) ELSE "usage".is_stream END,
input_tokens = "usage".input_tokens,
output_tokens = "usage".output_tokens,
total_tokens = "usage".total_tokens,
cache_creation_input_tokens = "usage".cache_creation_input_tokens,
cache_creation_input_tokens_5m = "usage".cache_creation_input_tokens_5m,
cache_creation_input_tokens_1h = "usage".cache_creation_input_tokens_1h,
cache_read_input_tokens = "usage".cache_read_input_tokens,
cache_creation_cost_usd = "usage".cache_creation_cost_usd,
cache_read_cost_usd = "usage".cache_read_cost_usd,
input_tokens = CASE WHEN "usage".billing_status = 'pending' AND EXCLUDED.status IN ('completed', 'failed', 'cancelled') THEN GREATEST("usage".input_tokens, EXCLUDED.input_tokens) ELSE "usage".input_tokens END,
output_tokens = CASE WHEN "usage".billing_status = 'pending' AND EXCLUDED.status IN ('completed', 'failed', 'cancelled') THEN GREATEST("usage".output_tokens, EXCLUDED.output_tokens) ELSE "usage".output_tokens END,
total_tokens = CASE WHEN "usage".billing_status = 'pending' AND EXCLUDED.status IN ('completed', 'failed', 'cancelled') THEN GREATEST("usage".total_tokens, EXCLUDED.total_tokens) ELSE "usage".total_tokens END,
cache_creation_input_tokens = CASE WHEN "usage".billing_status = 'pending' AND EXCLUDED.status IN ('completed', 'failed', 'cancelled') THEN GREATEST("usage".cache_creation_input_tokens, EXCLUDED.cache_creation_input_tokens) ELSE "usage".cache_creation_input_tokens END,
cache_creation_input_tokens_5m = CASE WHEN "usage".billing_status = 'pending' AND EXCLUDED.status IN ('completed', 'failed', 'cancelled') THEN GREATEST("usage".cache_creation_input_tokens_5m, EXCLUDED.cache_creation_input_tokens_5m) ELSE "usage".cache_creation_input_tokens_5m END,
cache_creation_input_tokens_1h = CASE WHEN "usage".billing_status = 'pending' AND EXCLUDED.status IN ('completed', 'failed', 'cancelled') THEN GREATEST("usage".cache_creation_input_tokens_1h, EXCLUDED.cache_creation_input_tokens_1h) ELSE "usage".cache_creation_input_tokens_1h END,
cache_read_input_tokens = CASE WHEN "usage".billing_status = 'pending' AND EXCLUDED.status IN ('completed', 'failed', 'cancelled') THEN GREATEST("usage".cache_read_input_tokens, EXCLUDED.cache_read_input_tokens) ELSE "usage".cache_read_input_tokens END,
cache_creation_cost_usd = CASE WHEN "usage".billing_status = 'pending' AND EXCLUDED.status IN ('completed', 'failed', 'cancelled') THEN GREATEST("usage".cache_creation_cost_usd, EXCLUDED.cache_creation_cost_usd) ELSE "usage".cache_creation_cost_usd END,
cache_read_cost_usd = CASE WHEN "usage".billing_status = 'pending' AND EXCLUDED.status IN ('completed', 'failed', 'cancelled') THEN GREATEST("usage".cache_read_cost_usd, EXCLUDED.cache_read_cost_usd) ELSE "usage".cache_read_cost_usd END,
output_price_per_1m = NULL,
total_cost_usd = "usage".total_cost_usd,
actual_total_cost_usd = "usage".actual_total_cost_usd,
total_cost_usd = CASE WHEN "usage".billing_status = 'pending' AND EXCLUDED.status IN ('completed', 'failed', 'cancelled') THEN GREATEST("usage".total_cost_usd, EXCLUDED.total_cost_usd) ELSE "usage".total_cost_usd END,
actual_total_cost_usd = CASE WHEN "usage".billing_status = 'pending' AND EXCLUDED.status IN ('completed', 'failed', 'cancelled') THEN GREATEST("usage".actual_total_cost_usd, EXCLUDED.actual_total_cost_usd) ELSE "usage".actual_total_cost_usd END,
status_code = CASE WHEN "usage".billing_status = 'pending' THEN CASE
WHEN "usage".status = 'streaming' AND EXCLUDED.status = 'pending' THEN "usage".status_code
WHEN EXCLUDED.status IN ('pending', 'streaming', 'completed', 'cancelled') AND EXCLUDED.status_code IS NULL THEN NULL

49
crates/aether-data/src/repository/usage/sql/tests.rs
View File

@@ -585,27 +585,46 @@ fn usage_sql_writes_usage_settlement_pricing_snapshots() {
}
#[test]
fn usage_sql_does_not_update_deprecated_billing_mirror_columns() {
fn usage_sql_updates_usage_mirror_columns_from_terminal_events_only() {
for assignment in [
"input_tokens = \"usage\".input_tokens",
"output_tokens = \"usage\".output_tokens",
"total_tokens = \"usage\".total_tokens",
"cache_creation_input_tokens = \"usage\".cache_creation_input_tokens",
"cache_creation_input_tokens_5m = \"usage\".cache_creation_input_tokens_5m",
"cache_creation_input_tokens_1h = \"usage\".cache_creation_input_tokens_1h",
"cache_read_input_tokens = \"usage\".cache_read_input_tokens",
"cache_creation_cost_usd = \"usage\".cache_creation_cost_usd",
"cache_read_cost_usd = \"usage\".cache_read_cost_usd",
"total_cost_usd = \"usage\".total_cost_usd",
"actual_total_cost_usd = \"usage\".actual_total_cost_usd",
"input_tokens = CASE WHEN \"usage\".billing_status = 'pending' AND EXCLUDED.status IN ('completed', 'failed', 'cancelled') THEN GREATEST(\"usage\".input_tokens, EXCLUDED.input_tokens) ELSE \"usage\".input_tokens END",
"output_tokens = CASE WHEN \"usage\".billing_status = 'pending' AND EXCLUDED.status IN ('completed', 'failed', 'cancelled') THEN GREATEST(\"usage\".output_tokens, EXCLUDED.output_tokens) ELSE \"usage\".output_tokens END",
"total_tokens = CASE WHEN \"usage\".billing_status = 'pending' AND EXCLUDED.status IN ('completed', 'failed', 'cancelled') THEN GREATEST(\"usage\".total_tokens, EXCLUDED.total_tokens) ELSE \"usage\".total_tokens END",
"cache_creation_input_tokens = CASE WHEN \"usage\".billing_status = 'pending' AND EXCLUDED.status IN ('completed', 'failed', 'cancelled') THEN GREATEST(\"usage\".cache_creation_input_tokens, EXCLUDED.cache_creation_input_tokens) ELSE \"usage\".cache_creation_input_tokens END",
"cache_creation_input_tokens_5m = CASE WHEN \"usage\".billing_status = 'pending' AND EXCLUDED.status IN ('completed', 'failed', 'cancelled') THEN GREATEST(\"usage\".cache_creation_input_tokens_5m, EXCLUDED.cache_creation_input_tokens_5m) ELSE \"usage\".cache_creation_input_tokens_5m END",
"cache_creation_input_tokens_1h = CASE WHEN \"usage\".billing_status = 'pending' AND EXCLUDED.status IN ('completed', 'failed', 'cancelled') THEN GREATEST(\"usage\".cache_creation_input_tokens_1h, EXCLUDED.cache_creation_input_tokens_1h) ELSE \"usage\".cache_creation_input_tokens_1h END",
"cache_read_input_tokens = CASE WHEN \"usage\".billing_status = 'pending' AND EXCLUDED.status IN ('completed', 'failed', 'cancelled') THEN GREATEST(\"usage\".cache_read_input_tokens, EXCLUDED.cache_read_input_tokens) ELSE \"usage\".cache_read_input_tokens END",
"cache_creation_cost_usd = CASE WHEN \"usage\".billing_status = 'pending' AND EXCLUDED.status IN ('completed', 'failed', 'cancelled') THEN GREATEST(\"usage\".cache_creation_cost_usd, EXCLUDED.cache_creation_cost_usd) ELSE \"usage\".cache_creation_cost_usd END",
"cache_read_cost_usd = CASE WHEN \"usage\".billing_status = 'pending' AND EXCLUDED.status IN ('completed', 'failed', 'cancelled') THEN GREATEST(\"usage\".cache_read_cost_usd, EXCLUDED.cache_read_cost_usd) ELSE \"usage\".cache_read_cost_usd END",
"total_cost_usd = CASE WHEN \"usage\".billing_status = 'pending' AND EXCLUDED.status IN ('completed', 'failed', 'cancelled') THEN GREATEST(\"usage\".total_cost_usd, EXCLUDED.total_cost_usd) ELSE \"usage\".total_cost_usd END",
"actual_total_cost_usd = CASE WHEN \"usage\".billing_status = 'pending' AND EXCLUDED.status IN ('completed', 'failed', 'cancelled') THEN GREATEST(\"usage\".actual_total_cost_usd, EXCLUDED.actual_total_cost_usd) ELSE \"usage\".actual_total_cost_usd END",
] {
assert!(
super::UPSERT_SQL.contains(assignment),
"missing deprecated mirror no-op assignment: {assignment}"
"missing terminal mirror assignment: {assignment}"
);
}
}
#[test]
fn usage_sql_binds_usage_mirror_values_for_terminal_upserts() {
let source = include_str!("mod.rs");
for binding in [
".bind(usage.input_tokens.map(to_i32).transpose()?)",
".bind(usage.output_tokens.map(to_i32).transpose()?)",
".bind(usage.total_tokens.map(to_i32).transpose()?)",
".bind(usage.cache_creation_input_tokens.map(to_i32).transpose()?)",
".bind(usage.cache_read_input_tokens.map(to_i32).transpose()?)",
".bind(usage.cache_creation_cost_usd)",
".bind(usage.cache_read_cost_usd)",
".bind(usage.total_cost_usd)",
".bind(usage.actual_total_cost_usd)",
] {
assert!(
source.contains(binding),
"missing usage upsert bind: {binding}"
);
}
assert!(!super::UPSERT_SQL.contains("EXCLUDED.input_tokens"));
assert!(!super::UPSERT_SQL.contains("EXCLUDED.total_cost_usd"));
}
#[test]

102
crates/aether-usage-runtime/src/usage_mapper.rs
View File

@@ -229,9 +229,11 @@ fn resolve_usage_value<'a>(
}
}
if let Some(nested) = response.get("response") {
if let Some(usage) = resolve_usage_value(nested, family) {
return Some(usage);
for nested_key in ["response", "message", "item"] {
if let Some(nested) = response.get(nested_key) {
if let Some(usage) = resolve_usage_value(nested, family) {
return Some(usage);
}
}
}
@@ -418,6 +420,69 @@ mod tests {
assert_eq!(usage.cache_read_tokens, 1);
}
#[test]
fn maps_claude_usage_from_stream_chunks() {
let usage = map_usage_from_response(
&serde_json::json!({
"chunks": [
{
"type": "message_start",
"message": {
"usage": {
"input_tokens": 5,
"cache_creation_input_tokens": 59_573,
"cache_read_input_tokens": 0,
"output_tokens": 0
}
}
},
{
"type": "message_delta",
"usage": {
"input_tokens": 5,
"cache_creation_input_tokens": 59_573,
"cache_read_input_tokens": 0,
"output_tokens": 162
}
}
]
}),
"claude:chat",
);
assert_eq!(usage.input_tokens, 5);
assert_eq!(usage.output_tokens, 162);
assert_eq!(usage.cache_creation_tokens, 59_573);
assert_eq!(usage.cache_read_tokens, 0);
}
#[test]
fn maps_claude_usage_from_message_start_chunk() {
let usage = map_usage_from_response(
&serde_json::json!({
"chunks": [
{
"type": "message_start",
"message": {
"usage": {
"input_tokens": 5,
"cache_creation_input_tokens": 59_573,
"cache_read_input_tokens": 0,
"output_tokens": 0
}
}
}
]
}),
"claude:chat",
);
assert_eq!(usage.input_tokens, 5);
assert_eq!(usage.output_tokens, 0);
assert_eq!(usage.cache_creation_tokens, 59_573);
assert_eq!(usage.cache_read_tokens, 0);
}
#[test]
fn maps_claude_usage_with_large_cache_read_tokens_without_subtracting_input() {
let usage = map_usage(
@@ -507,4 +572,35 @@ mod tests {
assert_eq!(usage.output_tokens, 6);
assert_eq!(usage.cache_read_tokens, 2);
}
#[test]
fn maps_gemini_usage_from_stream_chunks() {
let usage = map_usage_from_response(
&serde_json::json!({
"chunks": [
{
"candidates": [
{
"content": {
"parts": [{ "text": "hello" }]
}
}
]
},
{
"usageMetadata": {
"promptTokenCount": 14,
"candidatesTokenCount": 6,
"cachedContentTokenCount": 2
}
}
]
}),
"gemini:chat",
);
assert_eq!(usage.input_tokens, 14);
assert_eq!(usage.output_tokens, 6);
assert_eq!(usage.cache_read_tokens, 2);
}
}

125
crates/aether-usage-runtime/src/write.rs
View File

@@ -532,18 +532,13 @@ fn build_terminal_usage_event_from_seed_impl(
let endpoint_kind = infer_endpoint_kind(&client_contract).map(ToOwned::to_owned);
let provider_api_family = infer_api_family(&provider_contract).map(ToOwned::to_owned);
let provider_endpoint_kind = infer_endpoint_kind(&provider_contract).map(ToOwned::to_owned);
let derived_standardized_usage = standardized_usage
let derived_standardized_usage = provider_response
.as_ref()
.is_none()
.then(|| {
provider_response
.as_ref()
.filter(|response_body| response_body.is_object())
.map(|response_body| {
map_usage_from_response(response_body, provider_contract.as_str())
})
})
.flatten();
.filter(|response_body| response_body.is_object())
.map(|response_body| map_usage_from_response(response_body, provider_contract.as_str()))
.filter(StandardizedUsage::has_token_signal);
let standardized_usage =
StandardizedUsage::choose_more_complete(standardized_usage, derived_standardized_usage);
let request_metadata = if trusted_request_metadata {
merge_usage_request_metadata_owned(request_metadata, audit_payload)
} else {
@@ -608,9 +603,6 @@ fn build_terminal_usage_event_from_seed_impl(
apply_standardized_usage_seed(usage, &mut data);
}
if let Some(usage) = derived_standardized_usage.as_ref() {
apply_standardized_usage_seed(usage, &mut data);
}
if data.total_tokens.is_none() {
if let Some(tokens) = data
.response_body
@@ -2787,6 +2779,111 @@ mod tests {
assert!(event.data.client_response_body.is_none());
}
#[test]
fn stream_terminal_usage_prefers_more_complete_provider_chunks_usage() {
let plan = ExecutionPlan {
request_id: "req-stream-provider-chunks-usage-1".to_string(),
candidate_id: Some("cand-stream-provider-chunks-usage-1".to_string()),
provider_name: Some("OpenAI".to_string()),
provider_id: "provider-1".to_string(),
endpoint_id: "endpoint-1".to_string(),
key_id: "key-1".to_string(),
method: "POST".to_string(),
url: "https://example.com/v1/responses".to_string(),
headers: BTreeMap::new(),
content_type: None,
content_encoding: None,
body: RequestBody {
json_body: None,
body_bytes_b64: None,
body_ref: None,
},
stream: true,
client_api_format: "openai:cli".to_string(),
provider_api_format: "openai:cli".to_string(),
model_name: Some("gpt-5.5".to_string()),
proxy: None,
tls_profile: None,
timeouts: None,
};
let mut partial_summary_usage = StandardizedUsage::new();
partial_summary_usage.output_tokens = 148;
let provider_body = json!({
"chunks": [
{
"type": "response.created",
"response": {
"id": "resp_123",
"object": "response",
"model": "gpt-5.5",
"status": "in_progress",
"usage": null
}
},
{
"type": "response.completed",
"response": {
"id": "resp_123",
"object": "response",
"model": "gpt-5.5",
"status": "completed",
"usage": {
"input_tokens": 26,
"input_tokens_details": {
"cached_tokens": 0
},
"output_tokens": 148,
"output_tokens_details": {
"reasoning_tokens": 10
},
"total_tokens": 174
}
},
"sequence_number": 141
}
],
"metadata": {
"stream": true,
"total_chunks": 142,
"stored_chunks": 142
}
});
let payload = GatewayStreamReportRequest {
trace_id: "trace-stream-provider-chunks-usage-1".to_string(),
report_kind: "openai_cli_stream_success".to_string(),
report_context: Some(json!({
"client_api_format": "openai:cli",
"provider_api_format": "openai:cli",
})),
status_code: 200,
headers: BTreeMap::new(),
provider_body_base64: Some(
base64::engine::general_purpose::STANDARD.encode(provider_body.to_string()),
),
provider_body_state: Some(UsageBodyCaptureState::Inline),
client_body_base64: None,
client_body_state: Some(UsageBodyCaptureState::None),
terminal_summary: Some(ExecutionStreamTerminalSummary {
standardized_usage: Some(partial_summary_usage),
finish_reason: None,
response_id: Some("resp_123".to_string()),
model: Some("gpt-5.5".to_string()),
observed_finish: true,
parser_error: None,
}),
telemetry: None,
};
let event =
build_stream_terminal_usage_event(&plan, payload.report_context.as_ref(), &payload)
.expect("usage event should build");
assert_eq!(event.data.input_tokens, Some(26));
assert_eq!(event.data.output_tokens, Some(148));
assert_eq!(event.data.total_tokens, Some(174));
assert_eq!(event.data.cache_read_input_tokens, None);
}
#[test]
fn builds_stream_terminal_usage_from_sse_chunks_and_extracts_usage() {
let plan = ExecutionPlan {