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

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!({