use std::collections::BTreeMap; use aether_contracts::{ExecutionStreamTerminalSummary, ExecutionTelemetry}; use aether_data_contracts::repository::usage::UsageBodyCaptureState; use base64::Engine as _; use serde::{Deserialize, Serialize}; use serde_json::Value; pub const GEMINI_FILE_MAPPING_TTL_SECONDS: u64 = 60 * 60 * 48; const GEMINI_FILE_MAPPING_CACHE_PREFIX: &str = "gemini_files:key"; #[derive(Debug, Clone, Deserialize, Serialize)] pub struct GatewaySyncReportRequest { pub trace_id: String, pub report_kind: String, #[serde(default, skip_serializing_if = "Option::is_none")] pub report_context: Option, pub status_code: u16, pub headers: BTreeMap, #[serde(default, skip_serializing_if = "Option::is_none")] pub body_json: Option, #[serde(default, skip_serializing_if = "Option::is_none")] pub client_body_json: Option, #[serde(default, skip_serializing_if = "Option::is_none")] pub body_base64: Option, #[serde(default, skip_serializing_if = "Option::is_none")] pub telemetry: Option, } #[derive(Debug, Clone, Deserialize, Serialize)] pub struct GatewayStreamReportRequest { pub trace_id: String, pub report_kind: String, #[serde(default, skip_serializing_if = "Option::is_none")] pub report_context: Option, pub status_code: u16, pub headers: BTreeMap, #[serde(default, skip_serializing_if = "Option::is_none")] pub provider_body_base64: Option, #[serde(default, skip_serializing_if = "Option::is_none")] pub provider_body_state: Option, #[serde(default, skip_serializing_if = "Option::is_none")] pub client_body_base64: Option, #[serde(default, skip_serializing_if = "Option::is_none")] pub client_body_state: Option, #[serde(default, skip_serializing_if = "Option::is_none")] pub terminal_summary: Option, #[serde(default, skip_serializing_if = "Option::is_none")] pub telemetry: Option, } #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub struct InternalFinalizeRoute { pub public_path: &'static str, pub route_family: &'static str, pub route_kind: &'static str, } #[derive(Debug, Clone, PartialEq, Eq)] pub struct GeminiFileMappingEntry { pub file_name: String, pub display_name: Option, pub mime_type: Option, } pub fn infer_internal_finalize_signature(payload: &GatewaySyncReportRequest) -> Option { let report_context = payload.report_context.as_ref()?; let from_context = report_context .get("client_api_format") .and_then(serde_json::Value::as_str) .or_else(|| { report_context .get("provider_api_format") .and_then(serde_json::Value::as_str) }) .map(str::trim) .filter(|value| !value.is_empty()) .map(ToOwned::to_owned); if from_context.is_some() { return from_context; } let report_kind = payload.report_kind.trim().to_ascii_lowercase(); if report_kind.starts_with("openai_chat_") { return Some("openai:chat".to_string()); } if report_kind.starts_with("openai_compact_") { return Some("openai:responses:compact".to_string()); } if report_kind.starts_with("openai_responses_compact_") { return Some("openai:responses:compact".to_string()); } if report_kind.starts_with("openai_responses_") { return Some("openai:responses".to_string()); } if report_kind.starts_with("openai_image_") { return Some("openai:image".to_string()); } if report_kind.starts_with("openai_cli_") { return Some("openai:responses".to_string()); } if report_kind.starts_with("openai_video_") { return Some("openai:video".to_string()); } if report_kind.starts_with("claude_chat_") { return Some("claude:messages".to_string()); } if report_kind.starts_with("claude_cli_") { return Some("claude:messages".to_string()); } if report_kind.starts_with("gemini_chat_") { return Some("gemini:generate_content".to_string()); } if report_kind.starts_with("gemini_cli_") { return Some("gemini:generate_content".to_string()); } if report_kind.starts_with("gemini_video_") { return Some("gemini:video".to_string()); } None } pub fn resolve_internal_finalize_route(signature: &str) -> Option { match aether_ai_formats::normalize_api_format_alias(signature).as_str() { "openai:chat" => Some(InternalFinalizeRoute { public_path: "/v1/chat/completions", route_family: "openai", route_kind: "chat", }), "openai:responses" => Some(InternalFinalizeRoute { public_path: "/v1/responses", route_family: "openai", route_kind: "responses", }), "openai:responses:compact" => Some(InternalFinalizeRoute { public_path: "/v1/responses/compact", route_family: "openai", route_kind: "responses:compact", }), "openai:image" => Some(InternalFinalizeRoute { public_path: "/v1/images/generations", route_family: "openai", route_kind: "image", }), "openai:video" => Some(InternalFinalizeRoute { public_path: "/v1/videos", route_family: "openai", route_kind: "video", }), "claude:messages" => Some(InternalFinalizeRoute { public_path: "/v1/messages", route_family: "claude", route_kind: "messages", }), "gemini:generate_content" => Some(InternalFinalizeRoute { public_path: "/v1beta/models", route_family: "gemini", route_kind: "generate_content", }), "gemini:video" => Some(InternalFinalizeRoute { public_path: "/v1beta/models", route_family: "gemini", route_kind: "video", }), _ => None, } } pub fn normalize_gemini_file_name(file_name: &str) -> Option { let file_name = file_name.trim(); if file_name.is_empty() { return None; } if file_name.starts_with("files/") { Some(file_name.to_string()) } else { Some(format!("files/{file_name}")) } } pub fn gemini_file_mapping_cache_key(file_name: &str) -> String { format!("{GEMINI_FILE_MAPPING_CACHE_PREFIX}:{file_name}") } pub fn extract_gemini_file_mapping_entries( payload: &GatewaySyncReportRequest, ) -> Vec { let Some(body) = extract_sync_report_body_json(payload) else { return Vec::new(); }; let Some(object) = body.as_object() else { return Vec::new(); }; let mut entries = Vec::new(); maybe_push_gemini_file_mapping_entry(&mut entries, object); if let Some(file_object) = object.get("file").and_then(Value::as_object) { maybe_push_gemini_file_mapping_entry(&mut entries, file_object); } if let Some(files) = object.get("files").and_then(Value::as_array) { for item in files { if let Some(file_object) = item.as_object() { maybe_push_gemini_file_mapping_entry(&mut entries, file_object); } } } entries } pub fn report_request_id(report_context: Option<&serde_json::Value>) -> &str { report_context .and_then(|context| context.get("request_id")) .and_then(Value::as_str) .filter(|value| !value.trim().is_empty()) .unwrap_or("-") } pub fn is_local_ai_sync_report_kind(report_kind: &str) -> bool { matches!( report_kind, "openai_chat_sync_success" | "claude_chat_sync_success" | "gemini_chat_sync_success" | "openai_chat_sync_error" | "claude_chat_sync_error" | "gemini_chat_sync_error" | "openai_responses_sync_success" | "openai_responses_compact_sync_success" | "openai_responses_sync_error" | "openai_responses_compact_sync_error" | "openai_cli_sync_success" | "openai_image_sync_success" | "openai_image_sync_error" | "claude_cli_sync_success" | "gemini_cli_sync_success" | "openai_cli_sync_error" | "openai_compact_sync_error" | "claude_cli_sync_error" | "gemini_cli_sync_error" | "openai_video_create_sync_success" | "openai_video_remix_sync_success" | "gemini_video_create_sync_success" | "openai_video_delete_sync_success" | "openai_video_cancel_sync_success" | "gemini_video_cancel_sync_success" | "openai_video_create_sync_error" | "openai_video_remix_sync_error" | "gemini_video_create_sync_error" | "openai_video_delete_sync_error" | "openai_video_cancel_sync_error" | "gemini_video_cancel_sync_error" | "gemini_files_store_mapping" | "gemini_files_delete_mapping" ) } pub fn is_local_ai_stream_report_kind(report_kind: &str) -> bool { matches!( report_kind, "openai_chat_stream_success" | "claude_chat_stream_success" | "gemini_chat_stream_success" | "openai_responses_stream_success" | "openai_responses_compact_stream_success" | "openai_cli_stream_success" | "claude_cli_stream_success" | "gemini_cli_stream_success" ) } pub fn sync_report_represents_failure( payload: &GatewaySyncReportRequest, error_type: Option<&str>, ) -> bool { if payload.report_kind == "openai_video_delete_sync_success" && payload.status_code == 404 { return false; } payload.status_code >= 400 || payload.report_kind.contains("error") || error_type.is_some() || payload .body_json .as_ref() .and_then(|body| body.get("error")) .is_some_and(|value| !value.is_null()) } pub fn should_handle_local_sync_report( report_context: Option<&serde_json::Value>, report_kind: &str, ) -> bool { crate::report_context::report_context_is_locally_actionable(report_context) && is_local_ai_sync_report_kind(report_kind) } pub fn should_handle_local_stream_report( report_context: Option<&serde_json::Value>, report_kind: &str, ) -> bool { crate::report_context::report_context_is_locally_actionable(report_context) && is_local_ai_stream_report_kind(report_kind) } fn maybe_push_gemini_file_mapping_entry( entries: &mut Vec, object: &serde_json::Map, ) { let file_name = object .get("name") .and_then(Value::as_str) .and_then(normalize_gemini_file_name); let Some(file_name) = file_name else { return; }; if entries.iter().any(|entry| entry.file_name == file_name) { return; } entries.push(GeminiFileMappingEntry { file_name, display_name: object .get("displayName") .or_else(|| object.get("display_name")) .and_then(Value::as_str) .map(str::trim) .filter(|value| !value.is_empty()) .map(ToOwned::to_owned), mime_type: object .get("mimeType") .or_else(|| object.get("mime_type")) .and_then(Value::as_str) .map(str::trim) .filter(|value| !value.is_empty()) .map(ToOwned::to_owned), }); } fn extract_sync_report_body_json(payload: &GatewaySyncReportRequest) -> Option { if let Some(body_json) = payload.body_json.as_ref() { return Some(body_json.clone()); } if let Some(client_body_json) = payload.client_body_json.as_ref() { return Some(client_body_json.clone()); } if !content_type_starts_with(&payload.headers, "application/json") { return None; } let body_base64 = payload.body_base64.as_deref()?; let bytes = base64::engine::general_purpose::STANDARD .decode(body_base64) .ok()?; serde_json::from_slice(&bytes).ok() } fn content_type_starts_with(headers: &BTreeMap, expected_prefix: &str) -> bool { headers .iter() .find(|(key, _)| key.eq_ignore_ascii_case("content-type")) .map(|(_, value)| value.trim().to_ascii_lowercase()) .is_some_and(|value| value.starts_with(expected_prefix)) } #[cfg(test)] mod tests { use std::collections::BTreeMap; use base64::Engine as _; use serde_json::json; use super::{ extract_gemini_file_mapping_entries, gemini_file_mapping_cache_key, 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, GeminiFileMappingEntry, InternalFinalizeRoute, }; fn sample_sync_report(report_kind: &str, status_code: u16) -> GatewaySyncReportRequest { GatewaySyncReportRequest { trace_id: "trace-123".to_string(), report_kind: report_kind.to_string(), report_context: None, status_code, headers: BTreeMap::new(), body_json: None, client_body_json: None, body_base64: None, telemetry: None, } } fn sample_sync_report_with_context( report_kind: &str, report_context: serde_json::Value, ) -> GatewaySyncReportRequest { GatewaySyncReportRequest { trace_id: "trace-123".to_string(), report_kind: report_kind.to_string(), report_context: Some(report_context), status_code: 200, headers: BTreeMap::new(), body_json: None, client_body_json: None, body_base64: None, telemetry: None, } } #[test] fn classifies_local_ai_sync_report_kinds() { assert!(is_local_ai_sync_report_kind( "openai_video_create_sync_success" )); assert!(is_local_ai_sync_report_kind( "openai_responses_compact_sync_success" )); assert!(is_local_ai_sync_report_kind( "openai_responses_compact_sync_error" )); assert!(is_local_ai_sync_report_kind("openai_image_sync_success")); assert!(is_local_ai_sync_report_kind("openai_image_sync_error")); assert!(is_local_ai_sync_report_kind("gemini_files_delete_mapping")); assert!(!is_local_ai_sync_report_kind("unknown_sync_kind")); } #[test] fn classifies_local_ai_stream_report_kinds() { assert!(is_local_ai_stream_report_kind("openai_chat_stream_success")); assert!(is_local_ai_stream_report_kind( "openai_responses_compact_stream_success" )); assert!(!is_local_ai_stream_report_kind("openai_chat_stream_error")); } #[test] fn treats_openai_video_delete_404_success_as_non_failure() { let payload = sample_sync_report("openai_video_delete_sync_success", 404); assert!(!sync_report_represents_failure(&payload, None)); } #[test] fn detects_sync_report_failure_from_status_kind_error_type_or_body() { let status_payload = sample_sync_report("openai_chat_sync_success", 500); assert!(sync_report_represents_failure(&status_payload, None)); let kind_payload = sample_sync_report("openai_chat_sync_error", 200); assert!(sync_report_represents_failure(&kind_payload, None)); let error_type_payload = sample_sync_report("openai_chat_sync_success", 200); assert!(sync_report_represents_failure( &error_type_payload, Some("authentication_error") )); let mut error_body_payload = sample_sync_report("openai_chat_sync_success", 200); error_body_payload.body_json = Some(json!({"error": {"message": "bad request"}})); assert!(sync_report_represents_failure(&error_body_payload, None)); let mut null_error_payload = sample_sync_report("openai_chat_sync_success", 200); null_error_payload.body_json = Some(json!({"error": null})); assert!(!sync_report_represents_failure(&null_error_payload, None)); let success_payload = sample_sync_report("openai_chat_sync_success", 200); assert!(!sync_report_represents_failure(&success_payload, None)); } #[test] fn infers_internal_finalize_signature_from_context_or_report_kind() { let from_context = sample_sync_report_with_context( "unknown_sync_finalize", json!({"client_api_format": "gemini:video"}), ); assert_eq!( infer_internal_finalize_signature(&from_context), Some("gemini:video".to_string()) ); let from_report_kind = sample_sync_report_with_context("openai_video_create_sync_finalize", json!({})); assert_eq!( infer_internal_finalize_signature(&from_report_kind), Some("openai:video".to_string()) ); let from_image_report_kind = sample_sync_report_with_context("openai_image_sync_finalize", json!({})); assert_eq!( infer_internal_finalize_signature(&from_image_report_kind), Some("openai:image".to_string()) ); let from_compact_report_kind = sample_sync_report_with_context("openai_responses_compact_sync_finalize", json!({})); assert_eq!( infer_internal_finalize_signature(&from_compact_report_kind), Some("openai:responses:compact".to_string()) ); let unknown = sample_sync_report("unknown_sync_finalize", 200); assert_eq!(infer_internal_finalize_signature(&unknown), None); } #[test] fn resolves_internal_finalize_route_for_supported_signatures() { assert_eq!( resolve_internal_finalize_route("openai:responses:compact"), Some(InternalFinalizeRoute { public_path: "/v1/responses/compact", route_family: "openai", route_kind: "responses:compact", }) ); assert_eq!(resolve_internal_finalize_route("openai:compact"), None); assert_eq!( resolve_internal_finalize_route("gemini:video"), Some(InternalFinalizeRoute { public_path: "/v1beta/models", route_family: "gemini", route_kind: "video", }) ); assert_eq!( resolve_internal_finalize_route("openai:image"), Some(InternalFinalizeRoute { public_path: "/v1/images/generations", route_family: "openai", route_kind: "image", }) ); assert_eq!(resolve_internal_finalize_route("unknown:kind"), None); } #[test] fn normalizes_gemini_file_names() { assert_eq!( normalize_gemini_file_name("abc123"), Some("files/abc123".to_string()) ); assert_eq!( normalize_gemini_file_name("files/abc123"), Some("files/abc123".to_string()) ); assert_eq!(normalize_gemini_file_name(" "), None); } #[test] fn builds_gemini_file_mapping_cache_keys() { assert_eq!( gemini_file_mapping_cache_key("files/abc123"), "gemini_files:key:files/abc123" ); } #[test] fn extracts_gemini_file_mapping_entries_from_supported_shapes() { let payload = GatewaySyncReportRequest { trace_id: "trace-123".to_string(), report_kind: "gemini_files_store_mapping".to_string(), report_context: None, status_code: 200, headers: BTreeMap::new(), body_json: Some(json!({ "name": "abc123", "displayName": "root-name", "file": { "name": "files/def456", "mimeType": "image/png" }, "files": [ { "name": "abc123", "display_name": "deduped" }, { "name": "ghi789", "display_name": "third" } ] })), client_body_json: None, body_base64: None, telemetry: None, }; let entries = extract_gemini_file_mapping_entries(&payload); assert_eq!( entries, vec![ GeminiFileMappingEntry { file_name: "files/abc123".to_string(), display_name: Some("root-name".to_string()), mime_type: None, }, GeminiFileMappingEntry { file_name: "files/def456".to_string(), display_name: None, mime_type: Some("image/png".to_string()), }, GeminiFileMappingEntry { file_name: "files/ghi789".to_string(), display_name: Some("third".to_string()), mime_type: None, } ] ); } #[test] fn extracts_gemini_file_mapping_entries_from_base64_json_body() { let encoded_body = base64::engine::general_purpose::STANDARD.encode( serde_json::to_vec(&json!({ "name": "base64-file" })) .expect("json should encode"), ); let payload = GatewaySyncReportRequest { trace_id: "trace-123".to_string(), report_kind: "gemini_files_store_mapping".to_string(), report_context: None, status_code: 200, headers: BTreeMap::from([( "content-type".to_string(), "application/json; charset=utf-8".to_string(), )]), body_json: None, client_body_json: None, body_base64: Some(encoded_body), telemetry: None, }; let entries = extract_gemini_file_mapping_entries(&payload); assert_eq!(entries.len(), 1); assert_eq!(entries[0].file_name, "files/base64-file"); } #[test] fn reads_report_request_id_from_context() { assert_eq!( report_request_id(Some(&json!({"request_id": "req-123"}))), "req-123" ); assert_eq!(report_request_id(Some(&json!({"request_id": " "}))), "-"); assert_eq!(report_request_id(None), "-"); } #[test] fn decides_when_local_sync_report_should_be_handled() { assert!(should_handle_local_sync_report( Some(&json!({ "request_id": "req-123", "provider_id": "provider-123", "endpoint_id": "endpoint-123", "key_id": "key-123" })), "openai_chat_sync_success" )); assert!(should_handle_local_sync_report( Some(&json!({ "request_id": "req-123", "provider_id": "provider-123", "endpoint_id": "endpoint-123", "key_id": "key-123" })), "openai_image_sync_success" )); assert!(!should_handle_local_sync_report( Some(&json!({"request_id": "req-123"})), "openai_chat_sync_success" )); assert!(!should_handle_local_sync_report( Some(&json!({ "request_id": "req-123", "provider_id": "provider-123", "endpoint_id": "endpoint-123", "key_id": "key-123" })), "unknown_sync_kind" )); } #[test] fn decides_when_local_stream_report_should_be_handled() { assert!(should_handle_local_stream_report( Some(&json!({ "request_id": "req-123", "provider_id": "provider-123", "endpoint_id": "endpoint-123", "key_id": "key-123" })), "openai_chat_stream_success" )); assert!(!should_handle_local_stream_report( Some(&json!({ "request_id": "req-123", "provider_id": "provider-123", "endpoint_id": "endpoint-123", "key_id": "key-123" })), "openai_chat_stream_error" )); } }