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

Fix Codex image progress heartbeat merge regressions

Browse Source
This commit is contained in:
fawney19
2026-05-10 02:10:23 +08:00
parent eda3738d59 3c1456706a
commit 0f42a6ed82
156 changed files with 11741 additions and 943 deletions
Download Patch File Download Diff File Expand all files Collapse all files

63
crates/aether-usage-runtime/src/request_metadata.rs
View File

@@ -1,3 +1,6 @@
use aether_ai_formats::api::{
sanitize_request_path, sanitize_request_path_and_query, sanitize_request_query_string,
};
use aether_contracts::ExecutionPlan;
use serde_json::{json, Map, Value};
@@ -72,9 +75,13 @@ fn copy_allowed_metadata_fields(source: &Map<String, Value>, target: &mut Map<St
copy_bool(source, target, "client_requested_stream");
copy_bool(source, target, "upstream_is_stream");
copy_bool(source, target, "api_key_is_standalone");
copy_non_empty_string(source, target, "request_path");
copy_non_empty_string(source, target, "request_query_string");
copy_non_empty_string(source, target, "request_path_and_query");
copy_number(source, target, "provider_request_body_base64_bytes");
copy_number(source, target, "provider_response_body_base64_bytes");
copy_number(source, target, "client_response_body_base64_bytes");
copy_number(source, target, "client_response_status_code");
copy_non_null_value(source, target, "billing_snapshot");
copy_non_empty_string(source, target, "billing_snapshot_schema_version");
copy_non_empty_string(source, target, "billing_snapshot_status");
@@ -96,6 +103,7 @@ fn copy_allowed_metadata_fields(source: &Map<String, Value>, target: &mut Map<St
copy_number(source, target, "cache_read_price_per_1m");
copy_number(source, target, "price_per_request");
copy_non_null_value(source, target, "proxy");
sanitize_request_path_metadata_fields(target);
}
fn move_allowed_metadata_fields(mut source: Map<String, Value>, target: &mut Map<String, Value>) {
@@ -105,9 +113,13 @@ fn move_allowed_metadata_fields(mut source: Map<String, Value>, target: &mut Map
remove_bool(&mut source, target, "client_requested_stream");
remove_bool(&mut source, target, "upstream_is_stream");
remove_bool(&mut source, target, "api_key_is_standalone");
remove_non_empty_string(&mut source, target, "request_path");
remove_non_empty_string(&mut source, target, "request_query_string");
remove_non_empty_string(&mut source, target, "request_path_and_query");
remove_number(&mut source, target, "provider_request_body_base64_bytes");
remove_number(&mut source, target, "provider_response_body_base64_bytes");
remove_number(&mut source, target, "client_response_body_base64_bytes");
remove_number(&mut source, target, "client_response_status_code");
remove_non_null_value(&mut source, target, "billing_snapshot");
remove_non_empty_string(&mut source, target, "billing_snapshot_schema_version");
remove_non_empty_string(&mut source, target, "billing_snapshot_status");
@@ -129,6 +141,38 @@ fn move_allowed_metadata_fields(mut source: Map<String, Value>, target: &mut Map
remove_number(&mut source, target, "cache_read_price_per_1m");
remove_number(&mut source, target, "price_per_request");
remove_non_null_value(&mut source, target, "proxy");
sanitize_request_path_metadata_fields(target);
}
fn sanitize_request_path_metadata_fields(target: &mut Map<String, Value>) {
let path = target
.get("request_path")
.and_then(Value::as_str)
.and_then(sanitize_request_path);
let query = target
.get("request_query_string")
.and_then(Value::as_str)
.and_then(sanitize_request_query_string);
let path_and_query = target
.get("request_path_and_query")
.and_then(Value::as_str)
.and_then(|value| sanitize_request_path_and_query(value, None))
.or_else(|| {
path.as_deref()
.and_then(|path| sanitize_request_path_and_query(path, query.as_deref()))
});
apply_optional_string_field(target, "request_path", path.as_deref());
apply_optional_string_field(target, "request_query_string", query.as_deref());
apply_optional_string_field(target, "request_path_and_query", path_and_query.as_deref());
}
fn apply_optional_string_field(target: &mut Map<String, Value>, key: &str, value: Option<&str>) {
if let Some(value) = value {
target.insert(key.to_string(), Value::String(value.to_string()));
} else {
target.remove(key);
}
}
fn copy_non_empty_string(source: &Map<String, Value>, target: &mut Map<String, Value>, key: &str) {
@@ -451,6 +495,25 @@ mod tests {
);
}
#[test]
fn sanitizes_request_path_query_metadata() {
let metadata = sanitize_usage_request_metadata(Some(json!({
"request_path": "/v1beta/models/gemini-2.5-pro:streamGenerateContent?key=secret",
"request_query_string": "key=secret&alt=sse&pageSize=10&token=hidden",
"request_path_and_query": "/v1beta/models/gemini-2.5-pro:streamGenerateContent?key=secret&alt=sse&pageSize=10&token=hidden",
})))
.expect("metadata should remain");
assert_eq!(
metadata,
json!({
"request_path": "/v1beta/models/gemini-2.5-pro:streamGenerateContent",
"request_query_string": "alt=sse&pageSize=10",
"request_path_and_query": "/v1beta/models/gemini-2.5-pro:streamGenerateContent?alt=sse&pageSize=10",
})
);
}
#[test]
fn sanitizes_large_allowed_metadata_values_to_bounded_representations() {
let metadata = sanitize_usage_request_metadata(Some(json!({

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

@@ -529,7 +529,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 = 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);
let provider_api_family = infer_api_family(&provider_contract).map(ToOwned::to_owned);
@@ -734,7 +735,8 @@ pub fn build_sync_terminal_usage_payload_seed(
let context = payload.report_context.as_ref().and_then(Value::as_object);
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 client_response_headers = context_usage_value(context, "client_response_headers")
.or_else(|| headers_to_json(&payload.headers));
SyncTerminalUsagePayloadSeed {
report_kind: payload.report_kind.clone(),
status_code: payload.status_code,
@@ -941,7 +943,7 @@ fn infer_sync_terminal_state(
) -> UsageTerminalState {
if status_code == 499 || report_kind.contains("cancel") {
UsageTerminalState::Cancelled
} else if status_code >= 400
} else if !(200..300).contains(&status_code)
|| provider_response
.and_then(|value| value.get("error"))
.is_some_and(|value| !value.is_null())
@@ -959,7 +961,7 @@ fn infer_stream_terminal_state(
) -> UsageTerminalState {
if cancelled || status_code == 499 || report_kind.contains("cancel") {
UsageTerminalState::Cancelled
} else if status_code >= 400 {
} else if !(200..300).contains(&status_code) {
UsageTerminalState::Failed
} else {
UsageTerminalState::Completed
@@ -1982,6 +1984,8 @@ fn resolve_error_category(status_code: u16, event_type: UsageEventType) -> Optio
UsageEventType::Cancelled => Some("cancelled".to_string()),
UsageEventType::Failed if status_code >= 500 => Some("server_error".to_string()),
UsageEventType::Failed if status_code >= 400 => Some("client_error".to_string()),
UsageEventType::Failed if status_code >= 300 => Some("redirect".to_string()),
UsageEventType::Failed => Some("non_success_status".to_string()),
_ => None,
}
}
@@ -2000,7 +2004,7 @@ fn resolve_error_message(
if explicit_error_message.is_some() {
return explicit_error_message;
}
if status_code < 400 {
if (200..300).contains(&status_code) {
return None;
}
@@ -2798,6 +2802,81 @@ mod tests {
);
}
#[test]
fn stream_terminal_usage_marks_redirect_status_as_failed() {
let plan = ExecutionPlan {
request_id: "req-stream-redirect-usage".to_string(),
candidate_id: Some("cand-stream-redirect-usage".to_string()),
provider_name: Some("ChatGPTWeb".to_string()),
provider_id: "provider-redirect".to_string(),
endpoint_id: "endpoint-redirect".to_string(),
key_id: "key-redirect".to_string(),
method: "POST".to_string(),
url: "https://example.com/v1beta/models/gemini:streamGenerateContent".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: "gemini:generate_content".to_string(),
provider_api_format: "gemini:generate_content".to_string(),
model_name: Some("gemini".to_string()),
proxy: None,
transport_profile: None,
timeouts: None,
};
let client_body = json!({
"error": {
"type": "execution_runtime_non_success_status",
"message": "execution runtime stream returned non-success status 302",
"code": 302,
"upstream_status": 302,
"location": "/"
}
});
let payload = GatewayStreamReportRequest {
trace_id: "trace-stream-redirect-usage".to_string(),
report_kind: "gemini_chat_stream_success".to_string(),
report_context: Some(json!({
"client_api_format": "gemini:generate_content",
"provider_api_format": "gemini:generate_content"
})),
status_code: 302,
headers: BTreeMap::from([
("content-type".to_string(), "application/json".to_string()),
("x-aether-upstream-status".to_string(), "302".to_string()),
]),
provider_body_base64: Some(
base64::engine::general_purpose::STANDARD
.encode(br#"{"error":{"message":"raw redirect body"}}"#),
),
provider_body_state: Some(UsageBodyCaptureState::Inline),
client_body_base64: Some(
base64::engine::general_purpose::STANDARD
.encode(serde_json::to_vec(&client_body).expect("body should encode")),
),
client_body_state: Some(UsageBodyCaptureState::Inline),
terminal_summary: 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::Failed);
assert_eq!(event.data.status_code, Some(302));
assert_eq!(event.data.error_category.as_deref(), Some("redirect"));
assert_eq!(
event.data.error_message.as_deref(),
Some("raw redirect body")
);
}
#[test]
fn builds_stream_terminal_usage_from_terminal_summary_usage_without_decoding_bodies() {
let plan = ExecutionPlan {