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

Merge remote-tracking branch 'origin/pr/536'

Browse Source 
# Conflicts:
#	apps/aether-gateway/src/execution_runtime/stream/execution.rs
This commit is contained in:
fawney19
2026-05-21 23:03:14 +08:00
parent 97997685b5 eeb7995214
commit d6eb41aa78
18 changed files with 1850 additions and 555 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
apps/aether-gateway/src/ai_serving/planner/standard/codex/tests.rs
View File

@@ -85,7 +85,7 @@ fn injects_stable_prompt_cache_key_for_codex_requests() {
assert_eq!(
body["prompt_cache_key"],
"172c39e6-c0a0-5a70-8b63-e0f8e0d185a3"
"53363264-dbb0-5f9d-b9c7-3e92c45c5bdf"
);
}

230
apps/aether-gateway/src/ai_serving/planner/standard/family/request.rs
View File

@@ -65,6 +65,134 @@ fn is_grok_text_provider_api_format(provider_api_format: &str) -> bool {
)
}
fn provider_preserves_claude_thinking_signatures(provider_type: &str, base_url: &str) -> bool {
let provider_type = provider_type.trim().to_ascii_lowercase();
let base_url = base_url.trim().to_ascii_lowercase();
let is_bedrock_runtime_url = base_url.contains("bedrock-runtime")
&& (base_url.contains("amazonaws.com")
|| base_url.contains("amazonaws.com.cn")
|| base_url.contains("api.aws"));
matches!(
provider_type.as_str(),
"anthropic" | "claude_code" | "bedrock" | "aws_bedrock" | "amazon_bedrock"
) || base_url.contains("api.anthropic.com")
|| is_bedrock_runtime_url
}
fn sanitize_claude_thinking_block(block: Value) -> (Option<Value>, bool) {
let Some(object) = block.as_object() else {
return (Some(block), false);
};
let block_type = object
.get("type")
.and_then(Value::as_str)
.map(str::trim)
.unwrap_or_default();
match block_type {
"thinking" => {
let thinking_text = object
.get("thinking")
.and_then(Value::as_str)
.map(str::trim)
.unwrap_or_default();
if thinking_text.is_empty() {
(None, true)
} else {
(
Some(serde_json::json!({
"type": "text",
"text": thinking_text,
})),
true,
)
}
}
"redacted_thinking" => (None, true),
_ => (Some(block), false),
}
}
fn sanitize_claude_message_content_for_non_native_thinking(content: &mut Value) -> bool {
const OMITTED_THINKING_TEXT: &str = "Previous thinking omitted.";
if content.is_object() {
let original = std::mem::take(content);
let (sanitized, changed) = sanitize_claude_thinking_block(original);
if changed {
*content = sanitized.unwrap_or_else(|| {
serde_json::json!({
"type": "text",
"text": OMITTED_THINKING_TEXT,
})
});
}
return changed;
}
let Some(blocks) = content.as_array_mut() else {
return false;
};
let original_blocks = std::mem::take(blocks);
let mut changed = false;
let mut sanitized_blocks = Vec::with_capacity(original_blocks.len());
for block in original_blocks {
let (sanitized, block_changed) = sanitize_claude_thinking_block(block);
changed |= block_changed;
if let Some(sanitized) = sanitized {
sanitized_blocks.push(sanitized);
}
}
if changed && sanitized_blocks.is_empty() {
sanitized_blocks.push(serde_json::json!({
"type": "text",
"text": OMITTED_THINKING_TEXT,
}));
}
*blocks = sanitized_blocks;
changed
}
fn sanitize_claude_request_thinking_signatures_for_non_native(body_json: &mut Value) -> bool {
body_json
.get_mut("messages")
.and_then(Value::as_array_mut)
.map(|messages| {
messages.iter_mut().fold(false, |changed, message| {
let is_assistant = message
.get("role")
.and_then(Value::as_str)
.is_some_and(|role| role.trim().eq_ignore_ascii_case("assistant"));
if !is_assistant {
return changed;
}
let content_changed = message
.get_mut("content")
.is_some_and(sanitize_claude_message_content_for_non_native_thinking);
changed || content_changed
})
})
.unwrap_or(false)
}
fn apply_non_native_claude_thinking_signature_compat(
provider_request_body: &mut Value,
provider_api_format: &str,
transport: &GatewayProviderTransportSnapshot,
) {
if crate::ai_serving::normalize_api_format_alias(provider_api_format) != "claude:messages" {
return;
}
if provider_preserves_claude_thinking_signatures(
transport.provider.provider_type.as_str(),
transport.endpoint.base_url.as_str(),
) {
return;
}
let _ = sanitize_claude_request_thinking_signatures_for_non_native(provider_request_body);
}
pub(crate) async fn resolve_local_standard_candidate_payload_parts(
state: &AppState,
parts: &http::request::Parts,
@@ -390,6 +518,11 @@ pub(crate) async fn resolve_local_standard_candidate_payload_parts(
.await;
return None;
}
apply_non_native_claude_thinking_signature_compat(
&mut provider_request_body,
provider_api_format,
transport,
);
if let Some(mapping) =
crate::system_features::reasoning_model_directive_mapping_for_api_format_and_model(
state,
@@ -433,6 +566,11 @@ pub(crate) async fn resolve_local_standard_candidate_payload_parts(
.await;
return None;
}
apply_non_native_claude_thinking_signature_compat(
&mut provider_request_body,
provider_api_format,
transport,
);
}
if let Some(kiro_auth) = kiro_auth.as_ref() {
@@ -943,3 +1081,95 @@ async fn build_kiro_cross_format_payload_parts(
transport_profile: None,
})
}
#[cfg(test)]
mod tests {
use super::{
provider_preserves_claude_thinking_signatures,
sanitize_claude_request_thinking_signatures_for_non_native,
};
use serde_json::json;
#[test]
fn sanitizes_historical_claude_thinking_for_non_native_relays() {
let mut body = json!({
"model": "claude-opus-4-1",
"messages": [{
"role": "assistant",
"content": [
{
"type": "thinking",
"thinking": "I should keep this short.",
"signature": "sig_123"
},
{
"type": "redacted_thinking",
"data": "opaque"
},
{
"type": "text",
"text": "Done."
}
]
}]
});
assert!(sanitize_claude_request_thinking_signatures_for_non_native(
&mut body
));
assert_eq!(body["messages"][0]["content"][0]["type"], json!("text"));
assert_eq!(
body["messages"][0]["content"][0]["text"],
json!("I should keep this short.")
);
assert_eq!(body["messages"][0]["content"].as_array().unwrap().len(), 2);
assert_eq!(body["messages"][0]["content"][1]["text"], json!("Done."));
}
#[test]
fn inserts_placeholder_when_only_redacted_thinking_would_remain() {
let mut body = json!({
"model": "claude-opus-4-1",
"messages": [{
"role": "assistant",
"content": [{
"type": "redacted_thinking",
"data": "opaque"
}]
}]
});
assert!(sanitize_claude_request_thinking_signatures_for_non_native(
&mut body
));
assert_eq!(body["messages"][0]["content"][0]["type"], json!("text"));
assert_eq!(
body["messages"][0]["content"][0]["text"],
json!("Previous thinking omitted.")
);
}
#[test]
fn official_claude_providers_preserve_thinking_signatures() {
assert!(provider_preserves_claude_thinking_signatures(
"anthropic",
"https://relay.example.com"
));
assert!(provider_preserves_claude_thinking_signatures(
"custom",
"https://api.anthropic.com"
));
assert!(provider_preserves_claude_thinking_signatures(
"aws",
"https://bedrock-runtime.us-east-1.amazonaws.com"
));
assert!(provider_preserves_claude_thinking_signatures(
"amazon_bedrock",
"https://relay.example.com"
));
assert!(!provider_preserves_claude_thinking_signatures(
"openai",
"https://relay.example.com"
));
}
}

6
apps/aether-gateway/src/ai_serving/planner/standard/normalize/tests.rs
View File

@@ -209,7 +209,7 @@ fn local_openai_responses_compact_wrapper_strips_include_for_codex_requests() {
assert_eq!(provider_request_body["instructions"], "");
assert_eq!(
provider_request_body["prompt_cache_key"],
"172c39e6-c0a0-5a70-8b63-e0f8e0d185a3"
"3d2e2842-74cb-55dd-803a-b8940b3500c2"
);
}
@@ -355,7 +355,7 @@ fn injects_codex_prompt_cache_key_for_openai_responses_cross_format_requests() {
assert_eq!(
provider_request_body["prompt_cache_key"],
"172c39e6-c0a0-5a70-8b63-e0f8e0d185a3"
"b4dfeb75-b105-544c-a706-39b92f0bddb0"
);
}
@@ -385,6 +385,6 @@ fn injects_codex_prompt_cache_key_for_openai_chat_cross_format_requests() {
assert_eq!(
provider_request_body["prompt_cache_key"],
"172c39e6-c0a0-5a70-8b63-e0f8e0d185a3"
"4ee6ea6e-3ac6-5a18-8cb8-1f8b956419e5"
);
}

552
apps/aether-gateway/src/execution_runtime/stream/execution.rs
View File

@@ -53,6 +53,7 @@ use crate::ai_serving::api::{
maybe_build_provider_private_stream_normalizer, maybe_build_stream_response_rewriter,
normalize_provider_private_report_context, StreamingStandardTerminalObserver,
};
use crate::ai_serving::is_openai_responses_family_format;
use crate::api::response::{
attach_control_metadata_headers, build_client_response, build_client_response_from_parts,
};
@@ -702,6 +703,84 @@ fn should_replace_stream_usage(
observed.is_more_complete_than(current)
}
fn stream_terminal_summary_missing_observed_finish(
summary: Option<&ExecutionStreamTerminalSummary>,
) -> bool {
summary.is_some_and(|summary| {
!summary.observed_finish
&& !summary
.standardized_usage
.as_ref()
.is_some_and(StandardizedUsage::has_token_signal)
})
}
fn stream_report_context_format_field<'a>(
report_context: Option<&'a Value>,
field: &str,
) -> Option<&'a str> {
report_context
.and_then(Value::as_object)
.and_then(|object| object.get(field))
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
}
fn stream_requires_observed_terminal_event(
provider_api_format: &str,
report_context: Option<&Value>,
) -> bool {
is_openai_responses_family_format(provider_api_format)
|| [
"provider_stream_event_api_format",
"provider_stream_api_format",
"provider_api_format",
]
.into_iter()
.filter_map(|field| stream_report_context_format_field(report_context, field))
.any(is_openai_responses_family_format)
}
fn stream_terminal_summary_missing_observed_finish_with_requirement(
summary: Option<&ExecutionStreamTerminalSummary>,
requires_observed_terminal_event: bool,
) -> bool {
if !requires_observed_terminal_event {
return stream_terminal_summary_missing_observed_finish(summary);
}
summary.is_some_and(|summary| !summary.observed_finish)
}
fn ensure_stream_terminal_summary_for_missing_observed_finish(
summary: &mut Option<ExecutionStreamTerminalSummary>,
requires_observed_terminal_event: bool,
) {
if !requires_observed_terminal_event {
return;
}
let summary = summary.get_or_insert_with(ExecutionStreamTerminalSummary::default);
if !summary.observed_finish && summary.parser_error.is_none() {
summary.parser_error =
Some("execution runtime stream ended before provider terminal event".to_string());
}
}
fn stream_terminal_summary_represents_failure_with_requirement(
summary: Option<&ExecutionStreamTerminalSummary>,
requires_observed_terminal_event: bool,
) -> bool {
summary.is_some_and(|summary| {
summary.parser_error.is_some()
|| stream_terminal_summary_missing_observed_finish_with_requirement(
Some(summary),
requires_observed_terminal_event,
)
})
}
async fn execute_in_process_stream(
state: &AppState,
plan: &ExecutionPlan,
@@ -1476,7 +1555,12 @@ fn stream_chunk_contains_sse_done(chunk: &[u8]) -> bool {
let line = line.trim();
if matches!(
line,
"data: [DONE]" | "event: message_stop" | "event: response.completed"
"data: [DONE]"
| "event: message_stop"
| "event: response.completed"
| "event: response.failed"
| "event: response.incomplete"
| "event: error"
) {
return true;
}
@@ -1489,7 +1573,14 @@ fn stream_chunk_contains_sse_done(chunk: &[u8]) -> bool {
.get("type")
.and_then(serde_json::Value::as_str)
.is_some_and(|event_type| {
matches!(event_type, "message_stop" | "response.completed")
matches!(
event_type,
"message_stop"
| "response.completed"
| "response.failed"
| "response.incomplete"
| "error"
)
})
})
})
@@ -1509,24 +1600,6 @@ where
read_next_frame(lines).await
}
async fn next_stream_frame_until_downstream_closed<R>(
buffered_frames: &mut VecDeque<StreamFrame>,
lines: &mut FramedRead<R, LinesCodec>,
tx: &mpsc::Sender<Result<Bytes, IoError>>,
) -> Result<Option<StreamFrame>, GatewayError>
where
R: tokio::io::AsyncRead + Unpin,
{
if let Some(frame) = buffered_frames.pop_front() {
return Ok(Some(frame));
}
tokio::select! {
frame = read_next_frame(lines) => frame,
() = tx.closed() => Ok(None),
}
}
fn should_refresh_stream_usage_telemetry(
previous: Option<&ExecutionTelemetry>,
next: &ExecutionTelemetry,
@@ -2778,15 +2851,12 @@ async fn execute_stream_from_frame_stream(
image_stream_total_timeout.as_mut()
{
tokio::select! {
result = next_stream_frame_until_downstream_closed(
&mut buffered_frames,
&mut lines,
&tx,
) => result,
() = tx.closed() => {
biased;
_ = tx.closed(), if client_visible_stream_completed => {
downstream_dropped = true;
break;
}
result = next_stream_frame(&mut buffered_frames, &mut lines) => result,
_ = timeout_sleep.as_mut() => {
let timeout_ms = openai_image_stream_total_timeout_ms
.unwrap_or(OPENAI_IMAGE_STREAM_DEFAULT_TOTAL_TIMEOUT_MS);
@@ -2837,8 +2907,14 @@ async fn execute_stream_from_frame_stream(
}
}
} else {
next_stream_frame_until_downstream_closed(&mut buffered_frames, &mut lines, &tx)
.await
tokio::select! {
biased;
_ = tx.closed(), if client_visible_stream_completed => {
downstream_dropped = true;
break;
}
result = next_stream_frame(&mut buffered_frames, &mut lines) => result,
}
};
let next_frame = match next_frame_result {
Ok(frame) => frame,
@@ -3034,6 +3110,9 @@ async fn execute_stream_from_frame_stream(
u64::try_from(rewritten_chunk.len()).unwrap_or(u64::MAX);
let chunk_completed_stream =
stream_chunk_contains_sse_done(&rewritten_chunk);
if downstream_dropped {
continue;
}
if tx.send(Ok(Bytes::from(rewritten_chunk))).await.is_err() {
warn!(
event_name = "stream_execution_downstream_disconnected",
@@ -3041,10 +3120,9 @@ async fn execute_stream_from_frame_stream(
trace_id = %trace_id_owned,
request_id = %request_id_for_report_log,
candidate_id = ?candidate_id_for_report.as_deref(),
"gateway stream downstream dropped; stopping execution runtime stream forwarding"
"gateway stream downstream dropped; continuing to drain execution runtime stream"
);
downstream_dropped = true;
break;
} else {
client_visible_stream_completed |= chunk_completed_stream;
client_stream_bytes.fetch_add(rewritten_chunk_len, Ordering::Relaxed);
@@ -3110,35 +3188,35 @@ async fn execute_stream_from_frame_stream(
}
if downstream_dropped {
drop(lines);
debug!(
event_name = "execution_runtime_stream_flush_skipped",
event_name = "execution_runtime_stream_client_flush_skipped",
log_type = "debug",
debug_context = "redacted",
stream_status = "downstream_disconnected",
trace_id = %trace_id_owned,
"gateway skipped local stream flush after downstream disconnect"
"gateway skipped client stream flush after downstream disconnect"
);
} else {
if let Some(normalizer) = private_stream_normalizer.as_mut() {
match normalizer.finish() {
Ok(normalized_chunk) if !normalized_chunk.is_empty() => {
let provider_private_error_body_json =
extract_provider_private_stream_error_body(
stream_usage_report_context.as_ref(),
&normalized_chunk,
);
if let (Some(observer), Some(report_context)) = (
stream_usage_observer.as_mut(),
}
if let Some(normalizer) = private_stream_normalizer.as_mut() {
match normalizer.finish() {
Ok(normalized_chunk) if !normalized_chunk.is_empty() => {
let provider_private_error_body_json =
extract_provider_private_stream_error_body(
stream_usage_report_context.as_ref(),
) {
observe_stream_usage_bytes(
observer,
report_context,
&mut stream_usage_observer_buffered,
&normalized_chunk,
);
}
&normalized_chunk,
);
if let (Some(observer), Some(report_context)) = (
stream_usage_observer.as_mut(),
stream_usage_report_context.as_ref(),
) {
observe_stream_usage_bytes(
observer,
report_context,
&mut stream_usage_observer_buffered,
&normalized_chunk,
);
}
if !downstream_dropped {
let rewritten_chunk = if let Some(rewriter) = local_stream_rewriter.as_mut()
{
match rewriter.push_chunk(&normalized_chunk) {
@@ -3212,86 +3290,85 @@ async fn execute_stream_from_frame_stream(
});
}
}
}
Ok(_) => {}
Err(err) => {
warn!(
event_name = "stream_execution_normalization_flush_failed",
log_type = "ops",
trace_id = %trace_id_owned,
request_id = %request_id_for_report_log,
candidate_id = ?candidate_id_for_report.as_deref(),
error = ?err,
"gateway failed to flush private stream normalization"
);
terminal_failure.get_or_insert_with(|| {
build_stream_failure_report(
"execution_runtime_stream_rewrite_flush_error",
format!("failed to flush private stream normalization: {err:?}"),
502,
)
});
}
}
}
if !downstream_dropped {
if let Some(rewriter) = local_stream_rewriter.as_mut() {
match rewriter.finish() {
Ok(flushed_chunk) if !flushed_chunk.is_empty() => {
append_stream_capture_bytes(
&mut buffered_body,
&flushed_chunk,
max_stream_body_buffer_bytes,
&mut client_body_truncated,
);
let flushed_chunk_len =
u64::try_from(flushed_chunk.len()).unwrap_or(u64::MAX);
let chunk_completed_stream = stream_chunk_contains_sse_done(&flushed_chunk);
if tx.send(Ok(Bytes::from(flushed_chunk))).await.is_err() {
warn!(
event_name = "stream_execution_downstream_rewrite_flush_disconnected",
log_type = "ops",
trace_id = %trace_id_owned,
request_id = %request_id_for_report_log,
candidate_id = ?candidate_id_for_report.as_deref(),
"gateway stream downstream dropped while flushing local stream rewrite"
);
downstream_dropped = true;
} else {
client_visible_stream_completed |= chunk_completed_stream;
client_stream_bytes.fetch_add(flushed_chunk_len, Ordering::Relaxed);
last_client_chunk_elapsed_ms.store(
stream_started_at_for_report
.elapsed()
.as_millis()
.min(u128::from(u64::MAX))
as u64,
Ordering::Relaxed,
);
}
}
Ok(_) => {}
Err(err) => {
warn!(
event_name = "stream_execution_normalization_flush_failed",
event_name = "stream_execution_rewrite_flush_failed",
log_type = "ops",
trace_id = %trace_id_owned,
request_id = %request_id_for_report_log,
candidate_id = ?candidate_id_for_report.as_deref(),
error = ?err,
"gateway failed to flush private stream normalization"
"gateway failed to flush local stream rewrite"
);
terminal_failure.get_or_insert_with(|| {
build_stream_failure_report(
"execution_runtime_stream_rewrite_flush_error",
format!("failed to flush private stream normalization: {err:?}"),
format!("failed to flush local stream rewrite: {err:?}"),
502,
)
});
}
}
}
if !downstream_dropped {
if let Some(rewriter) = local_stream_rewriter.as_mut() {
match rewriter.finish() {
Ok(flushed_chunk) if !flushed_chunk.is_empty() => {
append_stream_capture_bytes(
&mut buffered_body,
&flushed_chunk,
max_stream_body_buffer_bytes,
&mut client_body_truncated,
);
let flushed_chunk_len =
u64::try_from(flushed_chunk.len()).unwrap_or(u64::MAX);
let chunk_completed_stream =
stream_chunk_contains_sse_done(&flushed_chunk);
if tx.send(Ok(Bytes::from(flushed_chunk))).await.is_err() {
warn!(
event_name = "stream_execution_downstream_rewrite_flush_disconnected",
log_type = "ops",
trace_id = %trace_id_owned,
request_id = %request_id_for_report_log,
candidate_id = ?candidate_id_for_report.as_deref(),
"gateway stream downstream dropped while flushing local stream rewrite"
);
downstream_dropped = true;
} else {
client_visible_stream_completed |= chunk_completed_stream;
client_stream_bytes.fetch_add(flushed_chunk_len, Ordering::Relaxed);
last_client_chunk_elapsed_ms.store(
stream_started_at_for_report
.elapsed()
.as_millis()
.min(u128::from(u64::MAX))
as u64,
Ordering::Relaxed,
);
}
}
Ok(_) => {}
Err(err) => {
warn!(
event_name = "stream_execution_rewrite_flush_failed",
log_type = "ops",
trace_id = %trace_id_owned,
request_id = %request_id_for_report_log,
candidate_id = ?candidate_id_for_report.as_deref(),
error = ?err,
"gateway failed to flush local stream rewrite"
);
terminal_failure.get_or_insert_with(|| {
build_stream_failure_report(
"execution_runtime_stream_rewrite_flush_error",
format!("failed to flush local stream rewrite: {err:?}"),
502,
)
});
}
}
}
}
}
if !downstream_dropped {
@@ -3466,6 +3543,19 @@ async fn execute_stream_from_frame_stream(
report_context_owned.as_ref(),
&mut stream_terminal_summary,
);
let requires_observed_terminal_event = stream_requires_observed_terminal_event(
plan_for_report.provider_api_format.as_str(),
stream_usage_report_context.as_ref(),
);
ensure_stream_terminal_summary_for_missing_observed_finish(
&mut stream_terminal_summary,
requires_observed_terminal_event,
);
let missing_observed_finish =
stream_terminal_summary_missing_observed_finish_with_requirement(
stream_terminal_summary.as_ref(),
requires_observed_terminal_event,
);
let should_submit_report = report_kind_owned.is_some();
let terminal_telemetry = Some(build_terminal_stream_telemetry(
@@ -3474,6 +3564,18 @@ async fn execute_stream_from_frame_stream(
usage_stream_telemetry.as_ref(),
provider_stream_bytes.load(Ordering::Relaxed),
));
let stream_failed = stream_terminal_summary_represents_failure_with_requirement(
stream_terminal_summary.as_ref(),
requires_observed_terminal_event,
);
let stream_terminal_error_message = stream_terminal_summary
.as_ref()
.and_then(|summary| summary.parser_error.clone())
.or_else(|| {
missing_observed_finish.then(|| {
"execution runtime stream ended before provider terminal event".to_string()
})
});
let usage_payload = build_stream_usage_payload(
trace_id_owned.clone(),
report_kind_owned.unwrap_or_default(),
@@ -3487,35 +3589,48 @@ async fn execute_stream_from_frame_stream(
stream_terminal_summary,
terminal_telemetry,
);
apply_local_execution_effect(
&state_for_report,
LocalExecutionEffectContext {
plan: &plan_for_report,
report_context: usage_payload.report_context.as_ref(),
},
LocalExecutionEffect::HealthSuccess(LocalHealthSuccessEffect),
)
.await;
apply_local_execution_effect(
&state_for_report,
LocalExecutionEffectContext {
plan: &plan_for_report,
report_context: usage_payload.report_context.as_ref(),
},
LocalExecutionEffect::AdaptiveSuccess(LocalAdaptiveSuccessEffect),
)
.await;
apply_local_execution_effect(
&state_for_report,
LocalExecutionEffectContext {
plan: &plan_for_report,
report_context: usage_payload.report_context.as_ref(),
},
LocalExecutionEffect::PoolSuccessStream {
payload: &usage_payload,
},
)
.await;
if stream_failed {
warn!(
event_name = "execution_runtime_stream_missing_terminal_event",
log_type = "ops",
trace_id = %trace_id_owned,
request_id = %request_id_for_report_log,
candidate_id = ?candidate_id_for_report.as_deref(),
status_code,
error_message = stream_terminal_error_message.as_deref().unwrap_or_default(),
"gateway stream ended with a failed terminal state"
);
} else {
apply_local_execution_effect(
&state_for_report,
LocalExecutionEffectContext {
plan: &plan_for_report,
report_context: usage_payload.report_context.as_ref(),
},
LocalExecutionEffect::HealthSuccess(LocalHealthSuccessEffect),
)
.await;
apply_local_execution_effect(
&state_for_report,
LocalExecutionEffectContext {
plan: &plan_for_report,
report_context: usage_payload.report_context.as_ref(),
},
LocalExecutionEffect::AdaptiveSuccess(LocalAdaptiveSuccessEffect),
)
.await;
apply_local_execution_effect(
&state_for_report,
LocalExecutionEffectContext {
plan: &plan_for_report,
report_context: usage_payload.report_context.as_ref(),
},
LocalExecutionEffect::PoolSuccessStream {
payload: &usage_payload,
},
)
.await;
}
record_stream_terminal_usage(
&state_for_report,
&plan_for_report,
@@ -3528,10 +3643,24 @@ async fn execute_stream_from_frame_stream(
&plan_for_report,
usage_payload.report_context.as_ref(),
SchedulerRequestCandidateStatusUpdate {
status: RequestCandidateStatus::Success,
status: if stream_failed {
RequestCandidateStatus::Failed
} else {
RequestCandidateStatus::Success
},
status_code: Some(status_code),
error_type: None,
error_message: None,
error_type: if stream_failed {
if missing_observed_finish {
Some("stream_missing_terminal_event".to_string())
} else {
Some("stream_terminal_error".to_string())
}
} else {
None
},
error_message: stream_failed
.then_some(stream_terminal_error_message)
.flatten(),
latency_ms: usage_payload
.telemetry
.as_ref()
@@ -3630,10 +3759,14 @@ mod tests {
use tokio::sync::{mpsc, watch, Notify};
use super::{
build_sse_body_stream, execute_execution_runtime_stream, execute_stream_from_frame_stream,
build_sse_body_stream, ensure_stream_terminal_summary_for_missing_observed_finish,
execute_execution_runtime_stream, execute_stream_from_frame_stream,
maybe_apply_kiro_prompt_cache_usage_to_stream_summary, merge_stream_terminal_summary,
should_limit_direct_finalize_prefetch, should_probe_success_failover_before_stream,
should_skip_direct_finalize_prefetch, stream_chunk_contains_sse_done,
stream_requires_observed_terminal_event, stream_terminal_summary_missing_observed_finish,
stream_terminal_summary_missing_observed_finish_with_requirement,
stream_terminal_summary_represents_failure_with_requirement,
};
use crate::control::GatewayControlDecision;
use crate::tunnel::{tunnel_protocol, TunnelProxyConn};
@@ -3659,6 +3792,9 @@ mod tests {
assert!(stream_chunk_contains_sse_done(
b"event: response.completed\ndata: {\"type\":\"response.completed\",\"response\":{}}\n\n"
));
assert!(stream_chunk_contains_sse_done(
b"event: response.failed\ndata: {\"type\":\"response.failed\",\"response\":{\"status\":\"failed\"}}\n\n"
));
assert!(!stream_chunk_contains_sse_done(
b"event: content_block_delta\ndata: {\"type\":\"content_block_delta\"}\n\n"
));
@@ -3719,6 +3855,98 @@ mod tests {
assert_eq!(merged.unknown_event_count, 3);
}
#[test]
fn detects_missing_observed_finish_only_without_usage_signal() {
assert!(stream_terminal_summary_missing_observed_finish(Some(
&ExecutionStreamTerminalSummary {
response_id: Some("resp_missing_finish".to_string()),
model: Some("gpt-5.5".to_string()),
observed_finish: false,
..ExecutionStreamTerminalSummary::default()
}
)));
let mut usage = StandardizedUsage::new();
usage.output_tokens = 12;
assert!(!stream_terminal_summary_missing_observed_finish(Some(
&ExecutionStreamTerminalSummary {
standardized_usage: Some(usage),
observed_finish: false,
..ExecutionStreamTerminalSummary::default()
}
)));
assert!(!stream_terminal_summary_missing_observed_finish(Some(
&ExecutionStreamTerminalSummary {
observed_finish: true,
..ExecutionStreamTerminalSummary::default()
}
)));
assert!(!stream_terminal_summary_missing_observed_finish(None));
}
#[test]
fn requires_terminal_event_for_openai_responses_streams() {
assert!(stream_requires_observed_terminal_event(
"openai:responses",
None
));
assert!(stream_requires_observed_terminal_event(
"openai:responses:compact",
None
));
assert!(!stream_requires_observed_terminal_event(
"openai:chat",
None
));
assert!(stream_requires_observed_terminal_event(
"openai:chat",
Some(&json!({
"provider_stream_event_api_format": "openai:responses"
}))
));
}
#[test]
fn synthesizes_missing_terminal_summary_for_openai_responses_empty_stream() {
let mut summary = None;
ensure_stream_terminal_summary_for_missing_observed_finish(&mut summary, true);
let summary = summary.expect("summary should be synthesized");
assert!(!summary.observed_finish);
assert_eq!(
summary.parser_error.as_deref(),
Some("execution runtime stream ended before provider terminal event")
);
assert!(
stream_terminal_summary_missing_observed_finish_with_requirement(Some(&summary), true)
);
assert!(stream_terminal_summary_represents_failure_with_requirement(
Some(&summary),
true
));
}
#[test]
fn terminal_required_stream_fails_even_with_usage_without_finish() {
let mut usage = StandardizedUsage::new();
usage.output_tokens = 12;
let mut summary = Some(ExecutionStreamTerminalSummary {
standardized_usage: Some(usage),
observed_finish: false,
..ExecutionStreamTerminalSummary::default()
});
ensure_stream_terminal_summary_for_missing_observed_finish(&mut summary, true);
let summary = summary.as_ref().expect("summary should remain present");
assert!(
stream_terminal_summary_missing_observed_finish_with_requirement(Some(summary), true)
);
assert!(stream_terminal_summary_represents_failure_with_requirement(
Some(summary),
true
));
}
#[test]
fn kiro_stream_summary_applies_prompt_cache_usage_from_original_request() {
let request_body = json!({
@@ -4830,7 +5058,7 @@ mod tests {
}
#[tokio::test]
async fn execute_stream_from_frame_stream_stops_upstream_when_client_drops_body() {
async fn execute_stream_from_frame_stream_drains_upstream_when_client_drops_body() {
let usage_repository = Arc::new(InMemoryUsageReadRepository::default());
let request_candidate_repository = Arc::new(InMemoryRequestCandidateRepository::default());
let state = AppState::new()
@@ -4873,24 +5101,22 @@ mod tests {
transport_profile: None,
timeouts: None,
};
let frame_stream_dropped = Arc::new(Notify::new());
let frame_stream_dropped_for_stream = Arc::clone(&frame_stream_dropped);
let release_terminal = Arc::new(Notify::new());
let terminal_frame_drained = Arc::new(Notify::new());
let release_terminal_for_stream = Arc::clone(&release_terminal);
let terminal_frame_drained_for_stream = Arc::clone(&terminal_frame_drained);
let frame_stream = stream! {
struct NotifyOnDrop(Arc<Notify>);
impl Drop for NotifyOnDrop {
fn drop(&mut self) {
self.0.notify_waiters();
}
}
let _drop_guard = NotifyOnDrop(frame_stream_dropped_for_stream);
yield Ok::<Bytes, std::io::Error>(Bytes::from_static(
b"{\"type\":\"headers\",\"payload\":{\"kind\":\"headers\",\"status_code\":200,\"headers\":{\"content-type\":\"text/event-stream\"}}}\n",
));
yield Ok::<Bytes, std::io::Error>(Bytes::from_static(
b"{\"type\":\"data\",\"payload\":{\"kind\":\"data\",\"text\":\"data: {\\\"id\\\":\\\"first\\\"}\\n\\n\"}}\n",
));
std::future::pending::<()>().await;
release_terminal_for_stream.notified().await;
yield Ok::<Bytes, std::io::Error>(Bytes::from_static(
b"{\"type\":\"data\",\"payload\":{\"kind\":\"data\",\"text\":\"data: {\\\"id\\\":\\\"terminal\\\",\\\"object\\\":\\\"chat.completion.chunk\\\",\\\"model\\\":\\\"gpt-5.4\\\",\\\"choices\\\":[{\\\"index\\\":0,\\\"delta\\\":{},\\\"finish_reason\\\":\\\"stop\\\"}],\\\"usage\\\":{\\\"prompt_tokens\\\":7,\\\"completion_tokens\\\":11,\\\"total_tokens\\\":18}}\\n\\ndata: [DONE]\\n\\n\"}}\n",
));
terminal_frame_drained_for_stream.notify_one();
}
.boxed();
@@ -4936,10 +5162,11 @@ mod tests {
assert_eq!(first.as_ref(), b"data: {\"id\":\"first\"}\n\n");
tokio::time::sleep(Duration::from_millis(30)).await;
drop(body_stream);
release_terminal.notify_one();
tokio::time::timeout(Duration::from_secs(1), frame_stream_dropped.notified())
tokio::time::timeout(Duration::from_secs(1), terminal_frame_drained.notified())
.await
.expect("upstream frame stream should be dropped after client disconnect");
.expect("upstream frame stream should be drained after client disconnect");
let candidates = tokio::time::timeout(Duration::from_secs(1), async {
loop {
let candidates = request_candidate_repository
@@ -4982,6 +5209,9 @@ mod tests {
.expect("usage should be marked cancelled");
assert_eq!(stored_usage.billing_status, "pending");
assert_eq!(stored_usage.status_code, Some(499));
assert_eq!(stored_usage.input_tokens, 7);
assert_eq!(stored_usage.output_tokens, 11);
assert_eq!(stored_usage.total_tokens, 18);
let first_byte_time_ms = stored_usage
.first_byte_time_ms
.expect("cancelled stream should retain first byte time");

2
apps/aether-gateway/src/request_candidate_runtime.rs
View File

@@ -191,7 +191,7 @@ pub(crate) fn snapshot_local_request_candidate_status(
.map(str::trim)
.filter(|value| !value.is_empty())?;
let metadata = parse_request_candidate_report_context(report_context)?;
let candidate_index = metadata.candidate_index?;
let candidate_index = metadata.candidate_index.unwrap_or(0);
Some(LocalRequestCandidateStatusSnapshot {
candidate_id: candidate_id.to_string(),

6
apps/aether-gateway/src/tests/ai_execute/stream/decision.rs
View File

@@ -884,15 +884,15 @@ async fn gateway_executes_openai_chat_stream_via_local_openai_responses_cross_fo
);
assert_eq!(
seen_execution_runtime_request.prompt_cache_key,
"b6741389-8b9e-5c00-bef6-fbce92aee45a"
"bc749eb7-a9e2-5793-8d14-abd659c700b0"
);
assert_eq!(
seen_execution_runtime_request.session_id,
"9fa08f4f14ccba13"
"d1e9b802644e1f52"
);
assert_eq!(
seen_execution_runtime_request.conversation_id,
"9fa08f4f14ccba13"
"d1e9b802644e1f52"
);
assert_eq!(
seen_execution_runtime_request.instructions,

27
apps/aether-gateway/src/usage/reporting/mod.rs
View File

@@ -18,7 +18,7 @@ use context::{report_context_is_locally_actionable, resolve_locally_actionable_r
use aether_usage_runtime::{
is_local_ai_stream_report_kind, is_local_ai_sync_report_kind, report_request_id,
should_handle_local_stream_report, should_handle_local_sync_report,
sync_report_represents_failure,
stream_report_represents_failure, sync_report_represents_failure,
};
pub(crate) use aether_usage_runtime::{GatewayStreamReportRequest, GatewaySyncReportRequest};
@@ -256,14 +256,33 @@ async fn handle_local_stream_report(state: &AppState, payload: &GatewayStreamRep
.telemetry
.as_ref()
.and_then(|telemetry| telemetry.elapsed_ms);
let failed = stream_report_represents_failure(payload);
record_report_request_candidate_status(
state,
payload.report_context.as_ref(),
SchedulerRequestCandidateStatusUpdate {
status: RequestCandidateStatus::Success,
status: if failed {
RequestCandidateStatus::Failed
} else {
RequestCandidateStatus::Success
},
status_code: Some(payload.status_code),
error_type: None,
error_message: None,
error_type: failed.then(|| {
if payload.status_code >= 400 {
"stream_http_error".to_string()
} else {
"stream_terminal_error".to_string()
}
}),
error_message: failed.then(|| {
payload
.terminal_summary
.as_ref()
.and_then(|summary| summary.parser_error.clone())
.unwrap_or_else(|| {
"execution runtime stream ended with a terminal error".to_string()
})
}),
latency_ms,
started_at_unix_ms: None,
finished_at_unix_ms: Some(terminal_unix_ms),

200
crates/aether-ai-formats/src/formats/openai/chat/stream.rs
View File

@@ -1180,6 +1180,23 @@ impl OpenAIResponsesProviderState {
}
}
}
event_type if openai_stream_payload_is_terminal_error(&value) => {
self.finished = true;
let mut payload = value.clone();
if event_type != "response.failed"
&& event_type != "response.incomplete"
&& event_type != "error"
{
payload = openai_stream_terminal_error_body(&value).unwrap_or(payload);
if let Some(object) = payload.as_object_mut() {
object.insert(
"type".to_string(),
Value::String("response.failed".to_string()),
);
}
}
out.push(self.unknown_frame(report_context, payload));
}
"response.completed" => {
let Some(response) = value.get("response").and_then(Value::as_object) else {
return Ok(out);
@@ -1291,6 +1308,8 @@ pub struct OpenAIChatClientEmitter {
model: Option<String>,
started: bool,
finished: bool,
next_tool_call_index: usize,
tool_call_index_by_canonical: BTreeMap<usize, usize>,
}
#[derive(Clone, Default)]
@@ -1375,6 +1394,17 @@ impl OpenAIChatClientEmitter {
)
}
fn chat_tool_call_index(&mut self, canonical_index: usize) -> usize {
if let Some(index) = self.tool_call_index_by_canonical.get(&canonical_index) {
return *index;
}
let index = self.next_tool_call_index;
self.next_tool_call_index += 1;
self.tool_call_index_by_canonical
.insert(canonical_index, index);
index
}
pub fn emit(&mut self, frame: CanonicalStreamFrame) -> Result<Vec<u8>, AiSurfaceFinalizeError> {
self.update_identity(&frame);
match frame.event {
@@ -1483,6 +1513,7 @@ impl OpenAIChatClientEmitter {
name,
} => {
let mut out = self.ensure_started()?;
let chat_index = self.chat_tool_call_index(index);
out.extend(encode_json_sse(
None,
&build_openai_chat_chunk(
@@ -1492,7 +1523,7 @@ impl OpenAIChatClientEmitter {
self.model.as_deref().unwrap_or("unknown"),
String::new(),
Some(vec![json!({
"index": index,
"index": chat_index,
"id": call_id,
"type": "function",
"function": {
@@ -1507,6 +1538,7 @@ impl OpenAIChatClientEmitter {
}
CanonicalStreamEvent::ToolCallArgumentsDelta { index, arguments } => {
let mut out = self.ensure_started()?;
let chat_index = self.chat_tool_call_index(index);
out.extend(encode_json_sse(
None,
&json!({
@@ -1519,7 +1551,7 @@ impl OpenAIChatClientEmitter {
"index": 0,
"delta": {
"tool_calls": [{
"index": index,
"index": chat_index,
"function": {
"arguments": arguments,
}
@@ -1562,6 +1594,13 @@ impl OpenAIChatClientEmitter {
)?);
Ok(out)
}
CanonicalStreamEvent::UnknownEvent(payload)
if openai_stream_terminal_error_body(&payload).is_some() =>
{
self.finished = true;
let error_body = openai_stream_terminal_error_body(&payload).unwrap_or(payload);
encode_json_sse(None, &error_body)
}
CanonicalStreamEvent::UnknownEvent(_) => Ok(Vec::new()),
CanonicalStreamEvent::Finish {
finish_reason,
@@ -2508,6 +2547,29 @@ impl OpenAIResponsesClientEmitter {
}
Ok(out)
}
CanonicalStreamEvent::UnknownEvent(payload)
if openai_stream_terminal_error_body(&payload).is_some() =>
{
self.finished = true;
let raw_event = payload.get("type").and_then(Value::as_str);
let event = raw_event
.filter(|event| {
matches!(*event, "response.failed" | "response.incomplete" | "error")
})
.unwrap_or("response.failed")
.to_string();
let mut payload = if raw_event == Some(event.as_str()) {
payload
} else {
openai_stream_terminal_error_body(&payload).unwrap_or(payload)
};
if payload.get("type").is_none() {
if let Some(object) = payload.as_object_mut() {
object.insert("type".to_string(), Value::String(event.clone()));
}
}
self.encode_response_event(event.as_str(), payload)
}
CanonicalStreamEvent::UnknownEvent(_) => Ok(Vec::new()),
CanonicalStreamEvent::Finish { usage, .. } => {
if self.finished {
@@ -2655,6 +2717,27 @@ mod tests {
parts
}
fn openai_chat_tool_call_indices(sse: &str) -> Vec<u64> {
let mut indices = Vec::new();
for payload in sse.lines().filter_map(|line| line.strip_prefix("data: ")) {
let Ok(value) = serde_json::from_str::<Value>(payload) else {
continue;
};
let Some(tool_calls) = value
.pointer("/choices/0/delta/tool_calls")
.and_then(Value::as_array)
else {
continue;
};
for tool_call in tool_calls {
if let Some(index) = tool_call.get("index").and_then(Value::as_u64) {
indices.push(index);
}
}
}
indices
}
#[test]
fn openai_chat_provider_state_emits_unknown_events_for_unrecognized_deltas() {
let mut state = OpenAIChatProviderState::default();
@@ -2713,6 +2796,40 @@ mod tests {
)));
}
#[test]
fn openai_responses_provider_state_treats_failed_event_as_terminal() {
let mut state = OpenAIResponsesProviderState::default();
let report_context = json!({});
let frames = state
.push_line(
&report_context,
data_line(json!({
"type": "response.failed",
"response": {
"id": "resp_failed_123",
"model": "gpt-5.4",
"status": "failed",
"error": {
"message": "policy failure",
"type": "invalid_request_error",
"code": "cyber_policy"
}
}
})),
)
.expect("failed response event should parse");
assert!(frames.iter().any(|frame| matches!(
frame.event,
CanonicalStreamEvent::UnknownEvent(ref payload)
if payload.get("type").and_then(Value::as_str) == Some("response.failed")
)));
assert!(state
.finish(&report_context)
.expect("terminal failure should not synthesize completion")
.is_empty());
}
#[test]
fn openai_usage_derives_missing_input_tokens_from_total() {
let usage = canonical_usage_from_openai_usage(Some(&json!({
@@ -2879,6 +2996,41 @@ mod tests {
assert_eq!(response_sequence_numbers(&sse), (1..=9).collect::<Vec<_>>());
}
#[test]
fn openai_responses_client_emitter_forwards_failed_unknown_event() {
let mut emitter = OpenAIResponsesClientEmitter::default();
let bytes = emitter
.emit(CanonicalStreamFrame {
id: "resp_failed_123".to_string(),
model: "gpt-5.4".to_string(),
event: CanonicalStreamEvent::UnknownEvent(json!({
"type": "response.failed",
"response": {
"id": "resp_failed_123",
"model": "gpt-5.4",
"status": "failed",
"error": {
"message": "policy failure",
"type": "invalid_request_error",
"code": "cyber_policy"
}
}
})),
})
.expect("failed response event should encode");
let mut all = bytes;
all.extend(
emitter
.finish()
.expect("failed stream should not synthesize completion"),
);
let sse = String::from_utf8(all).expect("sse should be utf8");
assert!(sse.contains("event: response.failed\n"));
assert!(sse.contains("\"message\":\"policy failure\""));
assert!(!sse.contains("event: response.completed\n"));
}
#[test]
fn openai_responses_client_emitter_keeps_text_item_id_stable_after_text_started() {
let mut emitter = OpenAIResponsesClientEmitter::default();
@@ -3367,6 +3519,50 @@ mod tests {
assert!(sse.contains("[Image]"));
}
#[test]
fn openai_chat_client_emitter_normalizes_sparse_tool_call_indices() {
let mut emitter = OpenAIChatClientEmitter::default();
let mut bytes = Vec::new();
for event in [
CanonicalStreamEvent::ToolCallStart {
index: 1,
call_id: "call_first".to_string(),
name: "first_tool".to_string(),
},
CanonicalStreamEvent::ToolCallArgumentsDelta {
index: 1,
arguments: "{\"first\":".to_string(),
},
CanonicalStreamEvent::ToolCallStart {
index: 3,
call_id: "call_second".to_string(),
name: "second_tool".to_string(),
},
CanonicalStreamEvent::ToolCallArgumentsDelta {
index: 3,
arguments: "{\"second\":true}".to_string(),
},
CanonicalStreamEvent::ToolCallArgumentsDelta {
index: 1,
arguments: "true}".to_string(),
},
] {
bytes.extend(
emitter
.emit(CanonicalStreamFrame {
id: "chatcmpl_sparse".to_string(),
model: "claude-opus-4-6".to_string(),
event,
})
.expect("tool event should encode"),
);
}
let sse = String::from_utf8(bytes).expect("sse should be utf8");
assert_eq!(openai_chat_tool_call_indices(&sse), vec![0, 0, 1, 1, 0]);
}
#[test]
fn openai_chat_client_emitter_emits_usage_only_final_chunk() {
let mut emitter = OpenAIChatClientEmitter::default();

536
crates/aether-ai-formats/src/formats/openai/responses/codex.rs
View File

@@ -165,14 +165,25 @@ fn inject_codex_default_variation_prompt(body_object: &mut serde_json::Map<Strin
);
}
fn build_stable_codex_prompt_cache_key(user_api_key_id: &str) -> Option<String> {
let normalized = user_api_key_id.trim();
fn build_stable_codex_prompt_cache_key_from_seed(kind: &str, seed: &str) -> Option<String> {
let normalized = seed.trim();
if normalized.is_empty() {
return None;
}
let normalized_kind = kind
.trim()
.to_ascii_lowercase()
.chars()
.filter(|ch| ch.is_ascii_alphanumeric() || *ch == '_' || *ch == '-')
.collect::<String>();
let normalized_kind = if normalized_kind.is_empty() {
"seed".to_string()
} else {
normalized_kind
};
let namespace = format!(
"aether:codex:prompt-cache:{CODEX_PROMPT_CACHE_NAMESPACE_VERSION}:user:{normalized}"
"aether:codex:prompt-cache:{CODEX_PROMPT_CACHE_NAMESPACE_VERSION}:{normalized_kind}:{normalized}"
);
let mut hasher = Sha1::new();
hasher.update(UUID_NAMESPACE_OID_BYTES);
@@ -186,6 +197,266 @@ fn build_stable_codex_prompt_cache_key(user_api_key_id: &str) -> Option<String>
Some(Uuid::from_bytes(bytes).to_string())
}
fn build_stable_codex_prompt_cache_key(user_api_key_id: &str) -> Option<String> {
build_stable_codex_prompt_cache_key_from_seed("user", user_api_key_id)
}
fn extract_codex_prompt_cache_session_seed(provider_request_body: &Value) -> Option<String> {
fn non_empty_str(value: Option<&Value>) -> Option<&str> {
value
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
}
fn session_seed_from_metadata(metadata: &Value) -> Option<String> {
let object = metadata.as_object()?;
non_empty_str(object.get("session_id"))
.or_else(|| non_empty_str(object.get("sessionId")))
.or_else(|| non_empty_str(object.get("conversation_id")))
.or_else(|| non_empty_str(object.get("conversationId")))
.map(|value| format!("metadata:{value}"))
.or_else(|| {
let user_id = non_empty_str(object.get("user_id"))?;
serde_json::from_str::<Value>(user_id)
.ok()
.and_then(|decoded| {
non_empty_str(decoded.get("session_id"))
.or_else(|| non_empty_str(decoded.get("sessionId")))
.or_else(|| non_empty_str(decoded.get("conversation_id")))
.or_else(|| non_empty_str(decoded.get("conversationId")))
.map(|value| format!("metadata.user_id:{value}"))
})
})
}
let object = provider_request_body.as_object()?;
non_empty_str(object.get("session_id"))
.or_else(|| non_empty_str(object.get("sessionId")))
.or_else(|| non_empty_str(object.get("conversation_id")))
.or_else(|| non_empty_str(object.get("conversationId")))
.map(|value| format!("body:{value}"))
.or_else(|| object.get("metadata").and_then(session_seed_from_metadata))
}
fn sha256_hex(input: &[u8]) -> String {
let digest = Sha256::digest(input);
let mut output = String::with_capacity(digest.len() * 2);
for byte in digest {
let _ = write!(&mut output, "{byte:02x}");
}
output
}
fn stable_json_digest(value: &Value) -> Option<String> {
serde_json::to_vec(value)
.ok()
.map(|serialized| sha256_hex(&serialized))
}
fn compact_prompt_cache_text(value: &str) -> Option<Value> {
const MAX_PROMPT_CACHE_TEXT_CHARS: usize = 4096;
let normalized = value.trim();
if normalized.is_empty() {
return None;
}
let mut text = normalized
.chars()
.take(MAX_PROMPT_CACHE_TEXT_CHARS)
.collect::<String>();
if normalized.chars().count() > MAX_PROMPT_CACHE_TEXT_CHARS {
text.push_str("...");
}
Some(Value::String(text))
}
fn compact_prompt_cache_anchor(value: &Value) -> Value {
match value {
Value::String(text) => compact_prompt_cache_text(text).unwrap_or(Value::Null),
Value::Array(items) => Value::Array(
items
.iter()
.take(16)
.map(compact_prompt_cache_anchor)
.filter(|value| !value.is_null())
.collect(),
),
Value::Object(object) => {
let mut compacted = serde_json::Map::new();
for key in [
"type",
"role",
"id",
"name",
"description",
"text",
"input_text",
"output_text",
"content",
"call_id",
"arguments",
"output",
"parameters",
"strict",
"function",
"effort",
"summary",
] {
let Some(value) = object.get(key) else {
continue;
};
let value = compact_prompt_cache_anchor(value);
if !value.is_null() {
compacted.insert(key.to_string(), value);
}
}
Value::Object(compacted)
}
Value::Null | Value::Bool(_) | Value::Number(_) => value.clone(),
}
}
fn compact_prompt_cache_json_anchor(value: &Value) -> Value {
match value {
Value::String(text) => compact_prompt_cache_text(text).unwrap_or(Value::Null),
Value::Array(items) => Value::Array(
items
.iter()
.take(16)
.map(compact_prompt_cache_json_anchor)
.filter(|value| !value.is_null())
.collect(),
),
Value::Object(object) => {
let mut compacted = serde_json::Map::new();
let mut keys = object.keys().collect::<Vec<_>>();
keys.sort();
for key in keys {
if key == "cache_control" {
continue;
}
let Some(value) = object.get(key) else {
continue;
};
let value = compact_prompt_cache_json_anchor(value);
if !value.is_null() {
compacted.insert(key.clone(), value);
}
}
Value::Object(compacted)
}
Value::Null | Value::Bool(_) | Value::Number(_) => value.clone(),
}
}
fn collect_codex_prompt_cache_control_anchors(value: &Value, anchors: &mut Vec<Value>) {
const MAX_PROMPT_CACHE_CONTROL_ANCHORS: usize = 16;
if anchors.len() >= MAX_PROMPT_CACHE_CONTROL_ANCHORS {
return;
}
match value {
Value::Object(object) => {
if object.contains_key("cache_control") {
let mut anchor = object.clone();
anchor.remove("cache_control");
let anchor = compact_prompt_cache_anchor(&Value::Object(anchor));
if !anchor.is_null() {
anchors.push(anchor);
}
}
for child in object.values() {
if anchors.len() >= MAX_PROMPT_CACHE_CONTROL_ANCHORS {
break;
}
collect_codex_prompt_cache_control_anchors(child, anchors);
}
}
Value::Array(items) => {
for child in items {
if anchors.len() >= MAX_PROMPT_CACHE_CONTROL_ANCHORS {
break;
}
collect_codex_prompt_cache_control_anchors(child, anchors);
}
}
_ => {}
}
}
fn extract_codex_prompt_cache_control_seed(provider_request_body: &Value) -> Option<String> {
let mut anchors = Vec::new();
collect_codex_prompt_cache_control_anchors(provider_request_body, &mut anchors);
if anchors.is_empty() {
return None;
}
let seed = json!({
"model": provider_request_body.get("model"),
"anchors": anchors,
});
stable_json_digest(&seed).map(|digest| format!("cache_control:{digest}"))
}
fn first_responses_input_anchor(input: &Value) -> Option<Value> {
let items = input.as_array()?;
let first_user_message = items.iter().find(|item| {
item.get("type")
.and_then(Value::as_str)
.is_some_and(|value| value == "message")
&& item
.get("role")
.and_then(Value::as_str)
.is_some_and(|value| value == "user")
});
let first_item = first_user_message.or_else(|| items.first())?;
let anchor = compact_prompt_cache_anchor(first_item);
(!anchor.is_null()).then_some(anchor)
}
fn extract_codex_stable_request_prompt_cache_seed(
provider_request_body: &Value,
user_api_key_id: Option<&str>,
) -> Option<String> {
let object = provider_request_body.as_object()?;
let mut seed = serde_json::Map::new();
for key in [
"model",
"instructions",
"reasoning",
"tools",
"tool_choice",
"parallel_tool_calls",
] {
if let Some(value) = object.get(key).filter(|value| !value.is_null()) {
let value = if key == "tools" {
compact_prompt_cache_json_anchor(value)
} else {
compact_prompt_cache_anchor(value)
};
seed.insert(key.to_string(), value);
}
}
if let Some(input_anchor) = object.get("input").and_then(first_responses_input_anchor) {
seed.insert("first_input".to_string(), input_anchor);
}
if let Some(user_api_key_id) = user_api_key_id
.map(str::trim)
.filter(|value| !value.is_empty())
{
seed.insert(
"api_key_id".to_string(),
Value::String(user_api_key_id.to_string()),
);
}
if seed.len() < 2 {
return None;
}
stable_json_digest(&Value::Object(seed)).map(|digest| format!("stable_request:{digest}"))
}
fn build_short_codex_header_id(seed: &str) -> Option<String> {
let normalized = seed.trim();
if normalized.is_empty() {
@@ -261,31 +532,47 @@ fn maybe_insert_default_codex_header(
provider_request_headers.insert(header_name.to_string(), header_value.to_string());
}
fn maybe_inject_codex_prompt_cache_key(
provider_request_body: &mut Value,
fn codex_prompt_cache_key_to_insert(
provider_request_body: &Value,
provider_type: &str,
provider_api_format: &str,
user_api_key_id: Option<&str>,
) {
) -> Option<String> {
if !is_codex_openai_responses_request(provider_type, provider_api_format) {
return;
return None;
}
let Some(body_object) = provider_request_body.as_object_mut() else {
return;
};
let existing = body_object
let existing = provider_request_body
.get("prompt_cache_key")
.and_then(Value::as_str)
.map(str::trim)
.unwrap_or_default();
if !existing.is_empty() {
return;
return None;
}
let Some(prompt_cache_key) = user_api_key_id.and_then(build_stable_codex_prompt_cache_key)
else {
extract_codex_prompt_cache_session_seed(provider_request_body)
.and_then(|seed| build_stable_codex_prompt_cache_key_from_seed("session", &seed))
.or_else(|| {
extract_codex_prompt_cache_control_seed(provider_request_body)
.and_then(|seed| build_stable_codex_prompt_cache_key_from_seed("anchor", &seed))
})
.or_else(|| {
extract_codex_stable_request_prompt_cache_seed(provider_request_body, user_api_key_id)
.and_then(|seed| build_stable_codex_prompt_cache_key_from_seed("request", &seed))
})
.or_else(|| user_api_key_id.and_then(build_stable_codex_prompt_cache_key))
}
fn insert_codex_prompt_cache_key(
provider_request_body: &mut Value,
prompt_cache_key: Option<String>,
) {
let Some(prompt_cache_key) = prompt_cache_key else {
return;
};
let Some(body_object) = provider_request_body.as_object_mut() else {
return;
};
@@ -431,6 +718,13 @@ pub fn apply_codex_openai_responses_special_body_edits(
return;
}
let prompt_cache_key = codex_prompt_cache_key_to_insert(
provider_request_body,
provider_type,
provider_api_format,
user_api_key_id,
);
let Some(body_object) = provider_request_body.as_object_mut() else {
return;
};
@@ -479,12 +773,7 @@ pub fn apply_codex_openai_responses_special_body_edits(
inject_codex_default_variation_prompt(body_object);
}
maybe_inject_codex_prompt_cache_key(
provider_request_body,
provider_type,
provider_api_format,
user_api_key_id,
);
insert_codex_prompt_cache_key(provider_request_body, prompt_cache_key);
}
pub fn apply_codex_openai_responses_chat_body_edits(
@@ -509,6 +798,9 @@ pub fn apply_codex_openai_responses_chat_body_edits(
return;
};
ensure_codex_chat_reasoning_defaults(body_object, provider_api_format, body_rules);
if let Some(prompt_cache_key) = body_object.remove("prompt_cache_key") {
body_object.insert("prompt_cache_key".to_string(), prompt_cache_key);
}
}
pub fn apply_codex_openai_responses_special_headers(
@@ -748,6 +1040,208 @@ mod tests {
);
}
#[test]
fn codex_responses_body_edits_derive_prompt_cache_key_from_session_metadata() {
let mut body_a = json!({
"input": [{"role": "user", "content": "hello"}],
"model": "gpt-5.4",
"metadata": {
"user_id": "{\"session_id\":\"session-a\",\"device_id\":\"device-a\"}"
}
});
let mut body_b = json!({
"input": [{"role": "user", "content": "hello again"}],
"model": "gpt-5.4",
"metadata": {
"user_id": "{\"session_id\":\"session-a\",\"device_id\":\"device-b\"}"
}
});
let mut body_c = json!({
"input": [{"role": "user", "content": "hello"}],
"model": "gpt-5.4",
"metadata": {"session_id": "session-b"}
});
apply_codex_openai_responses_special_body_edits(
&mut body_a,
"codex",
"openai:responses",
None,
Some("key-123"),
);
apply_codex_openai_responses_special_body_edits(
&mut body_b,
"codex",
"openai:responses",
None,
Some("different-key"),
);
apply_codex_openai_responses_special_body_edits(
&mut body_c,
"codex",
"openai:responses",
None,
Some("key-123"),
);
assert_eq!(body_a["prompt_cache_key"], body_b["prompt_cache_key"]);
assert_ne!(body_a["prompt_cache_key"], body_c["prompt_cache_key"]);
assert!(body_a.get("metadata").is_none());
assert!(body_b.get("metadata").is_none());
assert!(body_c.get("metadata").is_none());
}
#[test]
fn codex_responses_body_edits_derive_prompt_cache_key_from_cache_control_anchor() {
let mut body_a = json!({
"input": [{
"type": "message",
"role": "user",
"content": [{
"type": "input_text",
"text": "stable project brief",
"cache_control": {"type": "ephemeral"}
}]
}, {
"type": "message",
"role": "user",
"content": [{"type": "input_text", "text": "new turn A"}]
}],
"model": "gpt-5.4"
});
let mut body_b = json!({
"input": [{
"type": "message",
"role": "user",
"content": [{
"type": "input_text",
"text": "stable project brief",
"cache_control": {"type": "ephemeral"}
}]
}, {
"type": "message",
"role": "user",
"content": [{"type": "input_text", "text": "new turn B"}]
}],
"model": "gpt-5.4"
});
let mut body_c = json!({
"input": [{
"type": "message",
"role": "user",
"content": [{
"type": "input_text",
"text": "different project brief",
"cache_control": {"type": "ephemeral"}
}]
}],
"model": "gpt-5.4"
});
apply_codex_openai_responses_special_body_edits(
&mut body_a,
"codex",
"openai:responses",
None,
Some("key-a"),
);
apply_codex_openai_responses_special_body_edits(
&mut body_b,
"codex",
"openai:responses",
None,
Some("key-b"),
);
apply_codex_openai_responses_special_body_edits(
&mut body_c,
"codex",
"openai:responses",
None,
Some("key-a"),
);
assert_eq!(body_a["prompt_cache_key"], body_b["prompt_cache_key"]);
assert_ne!(body_a["prompt_cache_key"], body_c["prompt_cache_key"]);
}
#[test]
fn codex_responses_body_edits_derive_prompt_cache_key_from_stable_request_anchor() {
let mut body_a = json!({
"input": [{
"type": "message",
"role": "user",
"content": [{"type": "input_text", "text": "open workspace"}]
}, {
"type": "message",
"role": "user",
"content": [{"type": "input_text", "text": "new turn A"}]
}],
"model": "gpt-5.4",
"instructions": "Be concise.",
"tools": [{
"type": "function",
"name": "shell",
"parameters": {"type": "object", "properties": {}}
}],
"reasoning": {"effort": "medium"}
});
let mut body_b = json!({
"input": [{
"type": "message",
"role": "user",
"content": [{"type": "input_text", "text": "open workspace"}]
}, {
"type": "message",
"role": "user",
"content": [{"type": "input_text", "text": "new turn B"}]
}],
"model": "gpt-5.4",
"instructions": "Be concise.",
"tools": [{
"type": "function",
"name": "shell",
"parameters": {"type": "object", "properties": {}}
}],
"reasoning": {"effort": "medium"}
});
let mut body_c = json!({
"input": [{
"type": "message",
"role": "user",
"content": [{"type": "input_text", "text": "open another workspace"}]
}],
"model": "gpt-5.4",
"instructions": "Be concise.",
"tools": [{"type": "function", "name": "shell"}],
"reasoning": {"effort": "medium"}
});
apply_codex_openai_responses_special_body_edits(
&mut body_a,
"codex",
"openai:responses",
None,
Some("key-a"),
);
apply_codex_openai_responses_special_body_edits(
&mut body_b,
"codex",
"openai:responses",
None,
Some("key-a"),
);
apply_codex_openai_responses_special_body_edits(
&mut body_c,
"codex",
"openai:responses",
None,
Some("key-a"),
);
assert_eq!(body_a["prompt_cache_key"], body_b["prompt_cache_key"]);
assert_ne!(body_a["prompt_cache_key"], body_c["prompt_cache_key"]);
}
#[test]
fn compact_body_edits_strip_include_store_and_stream() {
let mut provider_request_body = json!({

101
crates/aether-ai-formats/src/formats/shared/stream_core/common.rs
View File

@@ -108,6 +108,107 @@ pub fn canonical_usage_from_openai_usage(value: Option<&Value>) -> Option<Canoni
})
}
pub fn openai_stream_payload_is_terminal_error(payload: &Value) -> bool {
let event_type = payload
.get("type")
.and_then(Value::as_str)
.unwrap_or_default();
if payload.get("error").is_some() {
return true;
}
if matches!(
event_type,
"error" | "response.failed" | "response.incomplete"
) {
return true;
}
payload
.get("response")
.and_then(Value::as_object)
.and_then(|response| response.get("status"))
.and_then(Value::as_str)
.is_some_and(|status| matches!(status, "failed" | "incomplete"))
}
pub fn openai_stream_terminal_error_body(payload: &Value) -> Option<Value> {
if !openai_stream_payload_is_terminal_error(payload) {
return None;
}
let event_type = payload
.get("type")
.and_then(Value::as_str)
.unwrap_or_default();
let response = payload.get("response").and_then(Value::as_object);
let status = response
.and_then(|response| response.get("status"))
.and_then(Value::as_str);
let raw_error = response
.and_then(|response| response.get("error"))
.or_else(|| payload.get("error"));
let mut error = raw_error
.and_then(Value::as_object)
.cloned()
.unwrap_or_default();
let message = error
.get("message")
.and_then(Value::as_str)
.map(ToOwned::to_owned)
.or_else(|| raw_error.and_then(Value::as_str).map(ToOwned::to_owned))
.or_else(|| {
payload
.get("message")
.and_then(Value::as_str)
.map(ToOwned::to_owned)
})
.or_else(|| {
response
.and_then(|response| response.get("incomplete_details"))
.and_then(|details| details.get("reason"))
.and_then(Value::as_str)
.map(|reason| format!("Response incomplete: {reason}"))
})
.or_else(|| status.map(|status| format!("Response ended with status {status}")))
.unwrap_or_else(|| "Upstream stream ended with an error".to_string());
error
.entry("message".to_string())
.or_insert_with(|| Value::String(message));
error.entry("type".to_string()).or_insert_with(|| {
if event_type == "response.incomplete" || status == Some("incomplete") {
Value::String("incomplete".to_string())
} else {
Value::String("server_error".to_string())
}
});
if !error.contains_key("code") {
if let Some(reason) = response
.and_then(|response| response.get("incomplete_details"))
.and_then(|details| details.get("reason"))
.and_then(Value::as_str)
{
error.insert("code".to_string(), Value::String(reason.to_string()));
}
}
Some(json!({ "error": Value::Object(error) }))
}
pub fn openai_stream_terminal_error_message(payload: &Value) -> Option<String> {
openai_stream_terminal_error_body(payload)
.and_then(|body| body.get("error").cloned())
.and_then(|error| {
error
.get("message")
.and_then(Value::as_str)
.map(ToOwned::to_owned)
})
}
pub fn canonical_usage_from_claude_usage(value: Option<&Value>) -> Option<CanonicalUsage> {
let usage = value?.as_object()?;
let input_tokens = usage

49
crates/aether-ai-formats/src/formats/shared/stream_core/format_matrix.rs
View File

@@ -14,7 +14,8 @@ use crate::formats::shared::error_body::{
};
use crate::formats::shared::sse::encode_json_sse;
use crate::formats::shared::stream_core::common::{
decode_json_data_line, CanonicalStreamEvent, CanonicalStreamFrame, CanonicalUsage,
decode_json_data_line, openai_stream_terminal_error_body, openai_stream_terminal_error_message,
CanonicalStreamEvent, CanonicalStreamFrame, CanonicalUsage,
};
use crate::formats::shared::AiSurfaceFinalizeError;
@@ -197,6 +198,14 @@ impl StreamingStandardTerminalObserver {
summary.model = Some(model);
}
match event {
CanonicalStreamEvent::UnknownEvent(payload)
if openai_stream_terminal_error_body(&payload).is_some() =>
{
summary.unknown_event_count = summary.unknown_event_count.saturating_add(1);
summary.observed_finish = true;
summary.finish_reason = Some("error".to_string());
summary.parser_error = openai_stream_terminal_error_message(&payload);
}
CanonicalStreamEvent::UnknownEvent(_) => {
summary.unknown_event_count = summary.unknown_event_count.saturating_add(1);
}
@@ -410,7 +419,8 @@ fn parse_provider_error(
}
fn parse_openai_error(payload: &Value) -> Option<(String, Option<String>, LocalCoreSyncErrorKind)> {
let error = payload.get("error")?.as_object()?;
let error_body = openai_stream_terminal_error_body(payload)?;
let error = error_body.get("error")?.as_object()?;
let message = error.get("message").and_then(Value::as_str)?.to_string();
let code = error
.get("code")
@@ -973,6 +983,41 @@ mod tests {
assert!(!summary.observed_finish);
}
#[test]
fn terminal_observer_marks_openai_responses_failed_event_as_terminal_error() {
let mut report_context = report_context("openai:chat", "openai:responses");
report_context["provider_stream_event_api_format"] = json!("openai:responses");
let mut observer = StreamingStandardTerminalObserver::default();
observer
.push_line(
&report_context,
data_line(json!({
"type": "response.failed",
"response": {
"id": "resp_failed_123",
"model": "gpt-5.4",
"status": "failed",
"error": {
"message": "policy failure",
"type": "invalid_request_error",
"code": "cyber_policy"
}
}
})),
)
.expect("failed event should be observed");
let summary = observer
.latest_summary()
.cloned()
.expect("summary should exist");
assert!(summary.observed_finish);
assert_eq!(summary.finish_reason.as_deref(), Some("error"));
assert_eq!(summary.parser_error.as_deref(), Some("policy failure"));
assert_eq!(summary.unknown_event_count, 1);
}
#[test]
fn terminal_observer_tracks_openai_image_stream_usage() {
let mut report_context = report_context("openai:image", "openai:chat");

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

@@ -183,23 +183,35 @@ DO UPDATE SET
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 IN ('completed', 'failed', 'cancelled') AND EXCLUDED.status IN ('pending', 'streaming') THEN "usage".status_code
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
ELSE COALESCE(EXCLUDED.status_code, "usage".status_code)
END ELSE "usage".status_code END,
error_message = CASE WHEN "usage".billing_status = 'pending' THEN CASE
WHEN "usage".status IN ('completed', 'failed', 'cancelled') AND EXCLUDED.status IN ('pending', 'streaming') THEN "usage".error_message
WHEN "usage".status = 'streaming' AND EXCLUDED.status = 'pending' THEN "usage".error_message
WHEN EXCLUDED.status IN ('pending', 'streaming', 'completed', 'cancelled') THEN EXCLUDED.error_message
ELSE COALESCE(EXCLUDED.error_message, "usage".error_message)
END ELSE "usage".error_message END,
error_category = CASE WHEN "usage".billing_status = 'pending' THEN CASE
WHEN "usage".status IN ('completed', 'failed', 'cancelled') AND EXCLUDED.status IN ('pending', 'streaming') THEN "usage".error_category
WHEN "usage".status = 'streaming' AND EXCLUDED.status = 'pending' THEN "usage".error_category
WHEN EXCLUDED.status IN ('pending', 'streaming', 'completed', 'cancelled') THEN EXCLUDED.error_category
ELSE COALESCE(EXCLUDED.error_category, "usage".error_category)
END ELSE "usage".error_category END,
response_time_ms = CASE WHEN "usage".billing_status = 'pending' THEN COALESCE(EXCLUDED.response_time_ms, "usage".response_time_ms) ELSE "usage".response_time_ms END,
first_byte_time_ms = CASE WHEN "usage".billing_status = 'pending' THEN COALESCE(EXCLUDED.first_byte_time_ms, "usage".first_byte_time_ms) ELSE "usage".first_byte_time_ms END,
response_time_ms = CASE WHEN "usage".billing_status = 'pending' THEN CASE
WHEN "usage".status IN ('completed', 'failed', 'cancelled') AND EXCLUDED.status IN ('pending', 'streaming') THEN "usage".response_time_ms
WHEN EXCLUDED.response_time_ms IS NULL OR EXCLUDED.response_time_ms = 0 THEN COALESCE("usage".response_time_ms, EXCLUDED.response_time_ms)
ELSE EXCLUDED.response_time_ms
END ELSE "usage".response_time_ms END,
first_byte_time_ms = CASE WHEN "usage".billing_status = 'pending' THEN CASE
WHEN "usage".status IN ('completed', 'failed', 'cancelled') AND EXCLUDED.status IN ('pending', 'streaming') THEN "usage".first_byte_time_ms
WHEN EXCLUDED.first_byte_time_ms IS NULL OR EXCLUDED.first_byte_time_ms = 0 THEN COALESCE("usage".first_byte_time_ms, EXCLUDED.first_byte_time_ms)
ELSE EXCLUDED.first_byte_time_ms
END ELSE "usage".first_byte_time_ms END,
status = CASE WHEN "usage".billing_status = 'pending' THEN CASE
WHEN "usage".status IN ('completed', 'failed', 'cancelled') AND EXCLUDED.status IN ('pending', 'streaming') THEN "usage".status
WHEN "usage".status = 'streaming' AND EXCLUDED.status = 'pending' THEN "usage".status
ELSE EXCLUDED.status
END ELSE "usage".status END,

42
crates/aether-data/src/repository/usage/sqlite.rs
View File

@@ -182,17 +182,41 @@ ON CONFLICT (request_id) DO UPDATE SET
output_price_per_1m = excluded.output_price_per_1m,
total_cost_usd = excluded.total_cost_usd,
actual_total_cost_usd = excluded.actual_total_cost_usd,
status_code = excluded.status_code,
error_message = excluded.error_message,
error_category = excluded.error_category,
response_time_ms = excluded.response_time_ms,
first_byte_time_ms = excluded.first_byte_time_ms,
status = excluded.status,
status_code = CASE
WHEN "usage".status IN ('completed', 'failed', 'cancelled') AND excluded.status IN ('pending', 'streaming') THEN "usage".status_code
WHEN "usage".status = 'streaming' AND excluded.status = 'pending' THEN "usage".status_code
ELSE excluded.status_code
END,
error_message = CASE
WHEN "usage".status IN ('completed', 'failed', 'cancelled') AND excluded.status IN ('pending', 'streaming') THEN "usage".error_message
WHEN "usage".status = 'streaming' AND excluded.status = 'pending' THEN "usage".error_message
ELSE excluded.error_message
END,
error_category = CASE
WHEN "usage".status IN ('completed', 'failed', 'cancelled') AND excluded.status IN ('pending', 'streaming') THEN "usage".error_category
WHEN "usage".status = 'streaming' AND excluded.status = 'pending' THEN "usage".error_category
ELSE excluded.error_category
END,
response_time_ms = CASE
WHEN "usage".status IN ('completed', 'failed', 'cancelled') AND excluded.status IN ('pending', 'streaming') THEN "usage".response_time_ms
WHEN excluded.response_time_ms IS NULL OR excluded.response_time_ms = 0 THEN COALESCE("usage".response_time_ms, excluded.response_time_ms)
ELSE excluded.response_time_ms
END,
first_byte_time_ms = CASE
WHEN "usage".status IN ('completed', 'failed', 'cancelled') AND excluded.status IN ('pending', 'streaming') THEN "usage".first_byte_time_ms
WHEN excluded.first_byte_time_ms IS NULL OR excluded.first_byte_time_ms = 0 THEN COALESCE("usage".first_byte_time_ms, excluded.first_byte_time_ms)
ELSE excluded.first_byte_time_ms
END,
status = CASE
WHEN "usage".status IN ('completed', 'failed', 'cancelled') AND excluded.status IN ('pending', 'streaming') THEN "usage".status
WHEN "usage".status = 'streaming' AND excluded.status = 'pending' THEN "usage".status
ELSE excluded.status
END,
billing_status = excluded.billing_status,
request_metadata = excluded.request_metadata,
candidate_id = excluded.candidate_id,
candidate_index = excluded.candidate_index,
key_name = excluded.key_name,
candidate_id = COALESCE(excluded.candidate_id, "usage".candidate_id),
candidate_index = COALESCE(excluded.candidate_index, "usage".candidate_index),
key_name = COALESCE(excluded.key_name, "usage".key_name),
planner_kind = excluded.planner_kind,
route_family = excluded.route_family,
route_kind = excluded.route_kind,

2
crates/aether-scheduler-core/src/request_candidate.rs
View File

@@ -463,7 +463,7 @@ pub fn build_local_request_candidate_status_record(
.map(str::trim)
.filter(|value| !value.is_empty())?;
let metadata = parse_request_candidate_report_context(report_context)?;
let candidate_index = metadata.candidate_index?;
let candidate_index = metadata.candidate_index.unwrap_or(0);
let extra_data = build_report_candidate_extra_data(ReportCandidateExtraDataInput {
client_api_format: metadata.client_api_format.clone(),
provider_api_format: metadata.provider_api_format.clone(),

6
crates/aether-usage-runtime/src/lib.rs
View File

@@ -27,9 +27,9 @@ pub use report::{
infer_internal_finalize_signature, is_local_ai_stream_report_kind,
is_local_ai_sync_report_kind, normalize_gemini_file_name, report_request_id,
resolve_internal_finalize_route, should_handle_local_stream_report,
should_handle_local_sync_report, sync_report_represents_failure, GatewayStreamReportRequest,
GatewaySyncReportRequest, GeminiFileMappingEntry, InternalFinalizeRoute,
GEMINI_FILE_MAPPING_TTL_SECONDS,
should_handle_local_sync_report, stream_report_represents_failure,
sync_report_represents_failure, GatewayStreamReportRequest, GatewaySyncReportRequest,
GeminiFileMappingEntry, InternalFinalizeRoute, GEMINI_FILE_MAPPING_TTL_SECONDS,
};
pub use report_context::{
build_locally_actionable_report_context_from_request_candidate,

50
crates/aether-usage-runtime/src/report.rs
View File

@@ -292,6 +292,24 @@ pub fn sync_report_represents_failure(
.is_some_and(|value| !value.is_null())
}
fn stream_terminal_summary_represents_failure(summary: &ExecutionStreamTerminalSummary) -> bool {
summary.parser_error.is_some()
|| (!summary.observed_finish
&& !summary
.standardized_usage
.as_ref()
.is_some_and(aether_contracts::StandardizedUsage::has_token_signal))
}
pub fn stream_report_represents_failure(payload: &GatewayStreamReportRequest) -> bool {
payload.status_code >= 400
|| payload.report_kind.contains("error")
|| payload
.terminal_summary
.as_ref()
.is_some_and(stream_terminal_summary_represents_failure)
}
pub fn should_handle_local_sync_report(
report_context: Option<&serde_json::Value>,
report_kind: &str,
@@ -373,6 +391,7 @@ fn content_type_starts_with(headers: &BTreeMap<String, String>, expected_prefix:
mod tests {
use std::collections::BTreeMap;
use aether_contracts::ExecutionStreamTerminalSummary;
use base64::Engine as _;
use serde_json::json;
@@ -381,7 +400,8 @@ mod tests {
infer_internal_finalize_signature, is_local_ai_stream_report_kind,
is_local_ai_sync_report_kind, normalize_gemini_file_name, report_request_id,
resolve_internal_finalize_route, should_handle_local_stream_report,
should_handle_local_sync_report, sync_report_represents_failure, GatewaySyncReportRequest,
should_handle_local_sync_report, stream_report_represents_failure,
sync_report_represents_failure, GatewayStreamReportRequest, GatewaySyncReportRequest,
GeminiFileMappingEntry, InternalFinalizeRoute,
};
@@ -416,6 +436,22 @@ mod tests {
}
}
fn sample_stream_report(report_kind: &str, status_code: u16) -> GatewayStreamReportRequest {
GatewayStreamReportRequest {
trace_id: "trace-stream-123".to_string(),
report_kind: report_kind.to_string(),
report_context: None,
status_code,
headers: BTreeMap::new(),
provider_body_base64: None,
provider_body_state: None,
client_body_base64: None,
client_body_state: None,
terminal_summary: None,
telemetry: None,
}
}
#[test]
fn classifies_local_ai_sync_report_kinds() {
assert!(is_local_ai_sync_report_kind(
@@ -481,6 +517,18 @@ mod tests {
assert!(!sync_report_represents_failure(&success_payload, None));
}
#[test]
fn detects_stream_report_failure_from_terminal_summary_error() {
let mut payload = sample_stream_report("openai_responses_stream_success", 200);
payload.terminal_summary = Some(ExecutionStreamTerminalSummary {
observed_finish: true,
parser_error: Some("policy failure".to_string()),
..ExecutionStreamTerminalSummary::default()
});
assert!(stream_report_represents_failure(&payload));
}
#[test]
fn infers_internal_finalize_signature_from_context_or_report_kind() {
let from_context = sample_sync_report_with_context(

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

@@ -158,6 +158,8 @@ pub struct StreamTerminalUsagePayloadSeed {
pub client_response: Option<Value>,
pub client_response_body_state: Option<UsageBodyCaptureState>,
pub standardized_usage: Option<StandardizedUsage>,
pub observed_stream_finish: Option<bool>,
pub terminal_error_message: Option<String>,
pub capture_metadata: Option<Value>,
}
@@ -183,6 +185,7 @@ pub struct TerminalUsageSeed {
pub has_format_conversion: bool,
pub is_stream: bool,
pub status_code: u16,
pub terminal_error_message: Option<String>,
pub response_time_ms: Option<u64>,
pub first_byte_time_ms: Option<u64>,
pub request_headers: Option<Value>,
@@ -506,6 +509,7 @@ fn build_terminal_usage_event_from_seed_impl(
has_format_conversion,
is_stream,
status_code,
terminal_error_message,
response_time_ms,
first_byte_time_ms,
request_headers,
@@ -530,7 +534,8 @@ fn build_terminal_usage_event_from_seed_impl(
};
let routing = merge_routing_seed_with_metadata_owned(routing, request_metadata.as_ref());
let body_refs = merge_body_refs_seed_with_metadata_owned(body_refs, request_metadata.as_ref());
let error_message = resolve_error_message(status_code, provider_response.as_ref(), None)
let error_message = terminal_error_message
.or_else(|| resolve_error_message(status_code, provider_response.as_ref(), None))
.or_else(|| resolve_error_message(status_code, client_response.as_ref(), None));
let api_family = infer_api_family(&client_contract).map(ToOwned::to_owned);
let endpoint_kind = infer_endpoint_kind(&client_contract).map(ToOwned::to_owned);
@@ -623,10 +628,6 @@ fn build_terminal_usage_event_from_seed_impl(
apply_completed_image_usage_estimate(&mut data);
}
if matches!(event_type, UsageEventType::Cancelled) {
apply_cancelled_usage_estimate(&mut data);
}
let data = if trusted_request_metadata {
sanitize_usage_event_capture_fields_trusted(data)
} else {
@@ -779,6 +780,16 @@ pub fn build_stream_terminal_usage_payload_seed(
let provider_response_headers = context_usage_value(context, "provider_response_headers")
.or_else(|| headers_to_json(&payload.headers));
let client_response_headers = headers_to_json(&payload.headers);
let observed_stream_finish = payload
.terminal_summary
.as_ref()
.map(|summary| summary.observed_finish);
let terminal_error_message = payload
.terminal_summary
.as_ref()
.and_then(|summary| summary.parser_error.clone())
.map(|message| message.trim().to_string())
.filter(|message| !message.is_empty());
StreamTerminalUsagePayloadSeed {
report_kind: payload.report_kind.clone(),
status_code: payload.status_code,
@@ -797,6 +808,8 @@ pub fn build_stream_terminal_usage_payload_seed(
.terminal_summary
.as_ref()
.and_then(|summary| summary.standardized_usage.clone()),
observed_stream_finish,
terminal_error_message,
capture_metadata: build_payload_body_capture_metadata(
payload.provider_body_base64.as_deref(),
payload.client_body_base64.as_deref(),
@@ -853,6 +866,7 @@ pub fn build_sync_terminal_usage_seed(
has_format_conversion: context_seed.has_format_conversion,
is_stream: context_seed.is_stream,
status_code,
terminal_error_message: None,
response_time_ms,
first_byte_time_ms,
request_headers: context_seed.request_headers,
@@ -894,6 +908,8 @@ pub fn build_stream_terminal_usage_seed(
client_response,
client_response_body_state,
standardized_usage,
observed_stream_finish,
terminal_error_message,
capture_metadata,
} = payload_seed;
let standardized_usage = standardized_usage.or_else(|| {
@@ -901,7 +917,28 @@ pub fn build_stream_terminal_usage_seed(
map_usage_from_response(response, context_seed.provider_contract.as_str())
})
});
let terminal_state = infer_stream_terminal_state(report_kind.as_str(), status_code, cancelled);
let missing_observed_finish = matches!(observed_stream_finish, Some(false))
&& !standardized_usage
.as_ref()
.is_some_and(StandardizedUsage::has_token_signal);
let terminal_error_message = terminal_error_message
.or_else(|| {
provider_response_full
.as_ref()
.and_then(extract_explicit_error_message_from_json)
})
.or_else(|| {
client_response
.as_ref()
.and_then(extract_explicit_error_message_from_json)
});
let terminal_state = infer_stream_terminal_state(
report_kind.as_str(),
status_code,
cancelled,
missing_observed_finish,
terminal_error_message.is_some(),
);
TerminalUsageSeed {
terminal_state,
@@ -924,6 +961,7 @@ pub fn build_stream_terminal_usage_seed(
has_format_conversion: context_seed.has_format_conversion,
is_stream: context_seed.is_stream,
status_code,
terminal_error_message,
response_time_ms,
first_byte_time_ms,
request_headers: context_seed.request_headers,
@@ -970,10 +1008,12 @@ fn infer_stream_terminal_state(
report_kind: &str,
status_code: u16,
cancelled: bool,
missing_observed_finish: bool,
terminal_error: bool,
) -> UsageTerminalState {
if cancelled || status_code == 499 || report_kind.contains("cancel") {
UsageTerminalState::Cancelled
} else if !(200..300).contains(&status_code) {
} else if !(200..300).contains(&status_code) || missing_observed_finish || terminal_error {
UsageTerminalState::Failed
} else {
UsageTerminalState::Completed
@@ -2141,6 +2181,31 @@ fn extract_explicit_error_message_from_json(value: &Value) -> Option<String> {
.and_then(|error| error.get("message"))
.and_then(Value::as_str)
.map(ToOwned::to_owned)
.or_else(|| {
value
.get("response")
.and_then(|response| response.get("error"))
.and_then(|error| error.get("message"))
.and_then(Value::as_str)
.map(ToOwned::to_owned)
})
.or_else(|| {
value
.get("response")
.and_then(|response| response.get("incomplete_details"))
.and_then(|details| details.get("reason"))
.and_then(Value::as_str)
.map(|reason| format!("Response incomplete: {reason}"))
})
.or_else(|| extract_stream_error_message_from_chunks(value))
}
fn extract_stream_error_message_from_chunks(value: &Value) -> Option<String> {
value
.get("chunks")
.and_then(Value::as_array)?
.iter()
.find_map(extract_explicit_error_message_from_json)
}
fn extract_generic_error_message_from_json(value: &Value) -> Option<String> {
@@ -2303,48 +2368,6 @@ fn extract_token_counts_from_value(value: &Value) -> Option<(u64, u64, u64)> {
}
}
fn apply_cancelled_usage_estimate(data: &mut UsageEventData) {
let provider_usage_available = data
.response_body
.as_ref()
.and_then(extract_token_counts_from_value)
.is_some();
let request_usage = data
.provider_request_body
.as_ref()
.or(data.request_body.as_ref())
.and_then(estimate_request_usage);
if positive_tokens(data.input_tokens) == 0 {
if let Some(usage) = request_usage.as_ref() {
data.input_tokens = Some(usage.input_tokens);
}
}
if !provider_usage_available {
apply_cancelled_request_cache_estimate(data, request_usage.as_ref());
}
if positive_tokens(data.output_tokens) == 0 {
if let Some(output_tokens) = data
.response_body
.as_ref()
.or(data.client_response_body.as_ref())
.and_then(estimate_response_output_tokens)
{
data.output_tokens = Some(output_tokens);
}
}
if positive_tokens(data.total_tokens) == 0 {
let total_tokens =
positive_tokens(data.input_tokens).saturating_add(positive_tokens(data.output_tokens));
if total_tokens > 0 {
data.total_tokens = Some(total_tokens);
}
}
}
fn apply_completed_image_usage_estimate(data: &mut UsageEventData) {
if !usage_event_data_is_image(data) {
return;
@@ -2369,7 +2392,7 @@ fn apply_completed_image_usage_estimate(data: &mut UsageEventData) {
data.input_tokens = Some(usage.input_tokens);
}
}
apply_cancelled_request_cache_estimate(data, request_usage.as_ref());
apply_request_cache_usage_estimate(data, request_usage.as_ref());
if positive_tokens(data.total_tokens) == 0 {
let total_tokens =
positive_tokens(data.input_tokens).saturating_add(positive_tokens(data.output_tokens));
@@ -2507,7 +2530,7 @@ fn usage_event_data_is_image(data: &UsageEventData) -> bool {
.is_some_and(|value| value.eq_ignore_ascii_case("image"))
}
fn apply_cancelled_request_cache_estimate(
fn apply_request_cache_usage_estimate(
data: &mut UsageEventData,
request_usage: Option<&EstimatedRequestUsage>,
) {
@@ -2674,280 +2697,6 @@ fn estimate_text_tokens(text: &str) -> u64 {
}
}
#[derive(Default)]
struct StreamOutputEstimate {
text: String,
saw_delta: bool,
}
impl StreamOutputEstimate {
fn push_delta(&mut self, text: &str) {
if text.is_empty() {
return;
}
self.saw_delta = true;
self.text.push_str(text);
}
fn push_done(&mut self, text: &str) {
if text.is_empty() || self.saw_delta {
return;
}
self.text.push_str(text);
}
}
fn estimate_response_output_tokens(value: &Value) -> Option<u64> {
let mut estimate = StreamOutputEstimate::default();
collect_stream_output_text(value, &mut estimate);
let tokens = estimate_text_tokens(estimate.text.as_str());
(tokens > 0).then_some(tokens)
}
fn collect_stream_output_text(value: &Value, estimate: &mut StreamOutputEstimate) {
match value {
Value::String(text) => {
for_each_sse_payload(text, |payload| {
if payload == "[DONE]" {
return;
}
if let Ok(json_body) = serde_json::from_str::<Value>(payload) {
collect_stream_output_text(&json_body, estimate);
}
});
}
Value::Array(items) => {
for item in items {
collect_stream_output_text(item, estimate);
}
}
Value::Object(object) => {
if let Some(chunks) = object.get("chunks").and_then(Value::as_array) {
for chunk in chunks {
collect_stream_output_text(chunk, estimate);
}
return;
}
collect_openai_responses_output_text(object, estimate);
collect_openai_chat_output_text(object, estimate);
collect_claude_output_text(object, estimate);
collect_gemini_output_text(object, estimate);
}
_ => {}
}
}
fn collect_openai_responses_output_text(
object: &Map<String, Value>,
estimate: &mut StreamOutputEstimate,
) {
match object
.get("type")
.and_then(Value::as_str)
.unwrap_or_default()
{
"response.output_text.delta" | "response.outtext.delta" => {
if let Some(text) = openai_delta_text(object.get("delta")) {
estimate.push_delta(text.as_str());
}
}
"response.reasoning_summary_text.delta" | "response.function_call_arguments.delta" => {
if let Some(text) = object.get("delta").and_then(Value::as_str) {
estimate.push_delta(text);
}
}
"response.output_text.done" | "response.reasoning_summary_text.done" => {
if let Some(text) = object
.get("text")
.and_then(Value::as_str)
.or_else(|| part_text(object.get("part")))
{
estimate.push_done(text);
}
}
"response.function_call_arguments.done" => {
if let Some(text) = object.get("arguments").and_then(Value::as_str) {
estimate.push_done(text);
}
}
"response.output_item.done" => {
if let Some(item) = object.get("item").and_then(Value::as_object) {
collect_openai_responses_output_item_text(item, estimate);
}
}
"response.completed" => {
if let Some(response) = object.get("response").and_then(Value::as_object) {
collect_openai_responses_completed_text(response, estimate);
}
}
_ => {}
}
}
fn collect_openai_responses_completed_text(
response: &Map<String, Value>,
estimate: &mut StreamOutputEstimate,
) {
for item in response
.get("output")
.and_then(Value::as_array)
.into_iter()
.flatten()
.filter_map(Value::as_object)
{
collect_openai_responses_output_item_text(item, estimate);
}
}
fn collect_openai_responses_output_item_text(
item: &Map<String, Value>,
estimate: &mut StreamOutputEstimate,
) {
match item.get("type").and_then(Value::as_str).unwrap_or_default() {
"message" => {
for content in item
.get("content")
.and_then(Value::as_array)
.into_iter()
.flatten()
.filter_map(Value::as_object)
{
if content.get("type").and_then(Value::as_str) == Some("output_text") {
if let Some(text) = content.get("text").and_then(Value::as_str) {
estimate.push_done(text);
}
}
}
}
"reasoning" => {
for summary in item
.get("summary")
.and_then(Value::as_array)
.into_iter()
.flatten()
.filter_map(Value::as_object)
{
if let Some(text) = summary.get("text").and_then(Value::as_str) {
estimate.push_done(text);
}
}
}
"function_call" => {
if let Some(arguments) = item.get("arguments").and_then(Value::as_str) {
estimate.push_done(arguments);
}
}
_ => {}
}
}
fn collect_openai_chat_output_text(
object: &Map<String, Value>,
estimate: &mut StreamOutputEstimate,
) {
for choice in object
.get("choices")
.and_then(Value::as_array)
.into_iter()
.flatten()
.filter_map(Value::as_object)
{
if let Some(delta) = choice.get("delta").and_then(Value::as_object) {
if let Some(content) = delta.get("content").and_then(Value::as_str) {
estimate.push_delta(content);
}
if let Some(reasoning_content) = delta.get("reasoning_content").and_then(Value::as_str)
{
estimate.push_delta(reasoning_content);
}
for tool_call in delta
.get("tool_calls")
.and_then(Value::as_array)
.into_iter()
.flatten()
.filter_map(Value::as_object)
{
if let Some(arguments) = tool_call
.get("function")
.and_then(Value::as_object)
.and_then(|function| function.get("arguments"))
.and_then(Value::as_str)
{
estimate.push_delta(arguments);
}
}
}
}
}
fn collect_claude_output_text(object: &Map<String, Value>, estimate: &mut StreamOutputEstimate) {
if object.get("type").and_then(Value::as_str) != Some("content_block_delta") {
return;
}
let Some(delta) = object.get("delta").and_then(Value::as_object) else {
return;
};
match delta
.get("type")
.and_then(Value::as_str)
.unwrap_or_default()
{
"text_delta" => {
if let Some(text) = delta.get("text").and_then(Value::as_str) {
estimate.push_delta(text);
}
}
"thinking_delta" => {
if let Some(text) = delta.get("thinking").and_then(Value::as_str) {
estimate.push_delta(text);
}
}
"input_json_delta" => {
if let Some(text) = delta.get("partial_json").and_then(Value::as_str) {
estimate.push_delta(text);
}
}
_ => {}
}
}
fn collect_gemini_output_text(object: &Map<String, Value>, estimate: &mut StreamOutputEstimate) {
for part in object
.get("candidates")
.and_then(Value::as_array)
.into_iter()
.flatten()
.filter_map(|candidate| candidate.get("content"))
.filter_map(Value::as_object)
.filter_map(|content| content.get("parts"))
.filter_map(Value::as_array)
.flatten()
.filter_map(Value::as_object)
{
if let Some(text) = part.get("text").and_then(Value::as_str) {
estimate.push_delta(text);
}
}
}
fn openai_delta_text(value: Option<&Value>) -> Option<String> {
match value {
Some(Value::String(text)) => Some(text.clone()),
Some(Value::Object(object)) => object
.get("text")
.and_then(Value::as_str)
.map(ToOwned::to_owned),
_ => None,
}
}
fn part_text(value: Option<&Value>) -> Option<&str> {
value
.and_then(Value::as_object)
.and_then(|part| part.get("text"))
.and_then(Value::as_str)
}
fn extract_token_counts_from_json(value: &Value) -> Option<(u64, u64, u64)> {
if let Some(usage) = value.get("usage").and_then(Value::as_object) {
let input = usage
@@ -3633,7 +3382,7 @@ mod tests {
}
#[test]
fn cancelled_stream_usage_estimates_tokens_from_request_and_partial_response() {
fn cancelled_stream_usage_does_not_estimate_tokens_from_request_or_partial_response() {
let plan = ExecutionPlan {
request_id: "req-stream-cancelled-estimated-usage-1".to_string(),
candidate_id: Some("cand-stream-cancelled-estimated-usage-1".to_string()),
@@ -3692,16 +3441,14 @@ mod tests {
.expect("usage event should build");
assert_eq!(event.event_type, UsageEventType::Cancelled);
assert!(event.data.input_tokens.unwrap_or_default() > 0);
assert_eq!(event.data.output_tokens, Some(5));
assert_eq!(
event.data.total_tokens,
Some(event.data.input_tokens.unwrap_or_default() + 5)
);
assert_eq!(event.data.input_tokens, None);
assert_eq!(event.data.output_tokens, None);
assert_eq!(event.data.total_tokens, None);
assert_eq!(event.data.cache_read_input_tokens, None);
}
#[test]
fn cancelled_stream_usage_does_not_infer_cache_read_from_prompt_cache_key() {
fn cancelled_stream_usage_does_not_infer_cache_or_token_estimates_from_prompt_cache_key() {
let request_body = json!({
"model": "gpt-5.4",
"input": "Use the cached project context and answer briefly",
@@ -3752,15 +3499,81 @@ mod tests {
let event =
build_stream_terminal_usage_event(&plan, payload.report_context.as_ref(), &payload)
.expect("usage event should build");
let input_tokens = event
.data
.input_tokens
.expect("input estimate should exist");
assert_eq!(event.event_type, UsageEventType::Cancelled);
assert_eq!(event.data.input_tokens, None);
assert_eq!(event.data.output_tokens, None);
assert_eq!(event.data.total_tokens, None);
assert_eq!(event.data.cache_read_input_tokens, None);
assert_eq!(event.data.output_tokens, Some(4));
assert_eq!(event.data.total_tokens, Some(input_tokens + 4));
}
#[test]
fn cancelled_stream_usage_preserves_terminal_summary_usage() {
let plan = ExecutionPlan {
request_id: "req-stream-cancelled-summary-usage-1".to_string(),
candidate_id: Some("cand-stream-cancelled-summary-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: Some("application/json".to_string()),
content_encoding: None,
body: RequestBody::from_json(json!({
"model": "gpt-5.4",
"input": "This cancelled request has terminal upstream usage",
"stream": true
})),
stream: true,
client_api_format: "openai:responses".to_string(),
provider_api_format: "openai:responses".to_string(),
model_name: Some("gpt-5.4".to_string()),
proxy: None,
transport_profile: None,
timeouts: None,
};
let mut standardized_usage = StandardizedUsage::new();
standardized_usage.input_tokens = 13;
standardized_usage.output_tokens = 21;
standardized_usage.cache_creation_tokens = 2;
standardized_usage.cache_read_tokens = 3;
let payload = GatewayStreamReportRequest {
trace_id: "trace-stream-cancelled-summary-usage-1".to_string(),
report_kind: "openai_responses_stream_cancelled".to_string(),
report_context: Some(json!({
"client_api_format": "openai:responses",
"provider_api_format": "openai:responses"
})),
status_code: 499,
headers: BTreeMap::new(),
provider_body_base64: None,
provider_body_state: Some(UsageBodyCaptureState::None),
client_body_base64: None,
client_body_state: Some(UsageBodyCaptureState::None),
terminal_summary: Some(ExecutionStreamTerminalSummary {
standardized_usage: Some(standardized_usage),
finish_reason: None,
response_id: Some("resp_cancel_summary_1".to_string()),
model: Some("gpt-5.4".to_string()),
observed_finish: true,
unknown_event_count: 0,
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.event_type, UsageEventType::Cancelled);
assert_eq!(event.data.input_tokens, Some(13));
assert_eq!(event.data.output_tokens, Some(21));
assert_eq!(event.data.total_tokens, Some(34));
assert_eq!(event.data.cache_creation_input_tokens, Some(2));
assert_eq!(event.data.cache_read_input_tokens, Some(3));
}
#[test]
@@ -3970,6 +3783,69 @@ mod tests {
assert!(event.data.client_response_body.is_none());
}
#[test]
fn stream_terminal_usage_marks_missing_observed_finish_as_failed() {
let plan = ExecutionPlan {
request_id: "req-stream-missing-finish-1".to_string(),
candidate_id: Some("cand-stream-missing-finish-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:responses".to_string(),
provider_api_format: "openai:responses".to_string(),
model_name: Some("gpt-5.5".to_string()),
proxy: None,
transport_profile: None,
timeouts: None,
};
let payload = GatewayStreamReportRequest {
trace_id: "trace-stream-missing-finish-1".to_string(),
report_kind: "openai_responses_stream_success".to_string(),
report_context: Some(json!({
"client_api_format": "openai:responses",
"provider_api_format": "openai:responses"
})),
status_code: 200,
headers: BTreeMap::new(),
provider_body_base64: None,
provider_body_state: Some(UsageBodyCaptureState::None),
client_body_base64: None,
client_body_state: Some(UsageBodyCaptureState::None),
terminal_summary: Some(ExecutionStreamTerminalSummary {
response_id: Some("resp_missing_finish".to_string()),
model: Some("gpt-5.5".to_string()),
observed_finish: false,
..ExecutionStreamTerminalSummary::default()
}),
telemetry: None,
};
let event =
build_stream_terminal_usage_event(&plan, payload.report_context.as_ref(), &payload)
.expect("usage event should build");
assert_eq!(event.event_type, UsageEventType::Failed);
assert_eq!(event.data.status_code, Some(200));
assert_eq!(
event.data.error_category.as_deref(),
Some("non_success_status")
);
assert_eq!(event.data.input_tokens, None);
assert_eq!(event.data.output_tokens, None);
}
#[test]
fn completed_image_usage_estimates_request_tokens_when_provider_usage_is_missing() {
let plan = ExecutionPlan {
@@ -5079,6 +4955,7 @@ mod tests {
..UsageRoutingSeed::default()
},
status_code: 200,
terminal_error_message: None,
response_time_ms: Some(123),
first_byte_time_ms: Some(45),
request_headers: Some(json!({

23
frontend/src/features/usage/composables/useUsageData.ts
View File

@@ -418,6 +418,25 @@ export function useUsageData(options: UseUsageDataOptions) {
}
}
function mergePositiveDurationMs(
existingValue: number | null | undefined,
nextValue: number | null | undefined
): number | null | undefined {
const existingIsPositive = typeof existingValue === 'number' && Number.isFinite(existingValue) && existingValue > 0
const nextIsPositive = typeof nextValue === 'number' && Number.isFinite(nextValue) && nextValue > 0
if (existingIsPositive && nextIsPositive) {
return Math.max(existingValue, nextValue)
}
if (existingIsPositive) {
return existingValue
}
if (nextIsPositive) {
return nextValue
}
return existingValue ?? nextValue
}
function mergeRecordStatus(
current: UsageRecord[],
next: UsageRecord[]
@@ -513,8 +532,8 @@ export function useUsageData(options: UseUsageDataOptions) {
cache_read_input_tokens: existing.cache_read_input_tokens ?? record.cache_read_input_tokens,
cost: existing.cost || record.cost,
actual_cost: existing.actual_cost ?? record.actual_cost,
response_time_ms: existing.response_time_ms ?? record.response_time_ms,
first_byte_time_ms: existing.first_byte_time_ms ?? record.first_byte_time_ms,
response_time_ms: mergePositiveDurationMs(existing.response_time_ms, record.response_time_ms),
first_byte_time_ms: mergePositiveDurationMs(existing.first_byte_time_ms, record.first_byte_time_ms),
is_stream: upstreamIsStream,
upstream_is_stream: upstreamIsStream,
client_requested_stream: clientRequestedStream,