diff --git a/apps/aether-gateway/src/ai_serving/mod.rs b/apps/aether-gateway/src/ai_serving/mod.rs index 483639116..cca447e05 100644 --- a/apps/aether-gateway/src/ai_serving/mod.rs +++ b/apps/aether-gateway/src/ai_serving/mod.rs @@ -53,6 +53,7 @@ pub(crate) use self::planner::{ build_standard_family_sync_plan_and_reports, build_standard_stream_plan_from_decision, build_standard_sync_plan_from_decision, candidate_auth_channel_skip_reason, codex_model_capabilities_for_transport, extract_pool_sticky_session_token, + maybe_build_pinned_stream_local_same_format_provider_decision_payload, maybe_build_responses_websocket_decision, maybe_build_stream_decision_payload, maybe_build_stream_plan_payload, maybe_build_sync_decision_payload, maybe_build_sync_plan_payload, planner_is_matching_stream_request, provider_key_pool_score_id, diff --git a/apps/aether-gateway/src/ai_serving/planner/mod.rs b/apps/aether-gateway/src/ai_serving/planner/mod.rs index 467709e90..66dd53e2b 100644 --- a/apps/aether-gateway/src/ai_serving/planner/mod.rs +++ b/apps/aether-gateway/src/ai_serving/planner/mod.rs @@ -38,6 +38,7 @@ pub(crate) use self::common::resolve_upstream_is_stream_for_provider; pub(crate) use self::passthrough::{ build_local_same_format_stream_attempt_source, build_local_same_format_stream_plan_and_reports, build_local_same_format_sync_attempt_source, build_local_same_format_sync_plan_and_reports, + maybe_build_pinned_stream_local_same_format_provider_decision_payload, }; pub(crate) use self::plan_builders::{ build_gemini_stream_plan_from_decision, build_gemini_sync_plan_from_decision, diff --git a/apps/aether-gateway/src/ai_serving/planner/passthrough/mod.rs b/apps/aether-gateway/src/ai_serving/planner/passthrough/mod.rs index 5b5d22431..c71946cda 100644 --- a/apps/aether-gateway/src/ai_serving/planner/passthrough/mod.rs +++ b/apps/aether-gateway/src/ai_serving/planner/passthrough/mod.rs @@ -8,6 +8,7 @@ pub(crate) use self::provider::{ build_local_sync_attempt_source as build_local_same_format_sync_attempt_source, build_local_sync_plan_and_reports as build_local_same_format_sync_plan_and_reports, maybe_build_local_same_format_provider_decision_payload_for_candidate, + maybe_build_pinned_stream_local_same_format_provider_decision_payload, maybe_build_stream_local_same_format_provider_decision_payload, maybe_build_sync_local_same_format_provider_decision_payload, }; diff --git a/apps/aether-gateway/src/ai_serving/planner/passthrough/provider.rs b/apps/aether-gateway/src/ai_serving/planner/passthrough/provider.rs index 189de54b4..1d9d8d770 100644 --- a/apps/aether-gateway/src/ai_serving/planner/passthrough/provider.rs +++ b/apps/aether-gateway/src/ai_serving/planner/passthrough/provider.rs @@ -69,6 +69,7 @@ pub(crate) use self::family::{ LocalSameFormatProviderFamily, LocalSameFormatProviderSpec, }; pub(crate) use self::family::{ + maybe_build_pinned_stream_local_same_format_provider_decision_payload, maybe_build_stream_local_same_format_provider_decision_payload, maybe_build_sync_local_same_format_provider_decision_payload, }; diff --git a/apps/aether-gateway/src/ai_serving/planner/passthrough/provider/family/build.rs b/apps/aether-gateway/src/ai_serving/planner/passthrough/provider/family/build.rs index cfdd9bbcf..3b5e83495 100644 --- a/apps/aether-gateway/src/ai_serving/planner/passthrough/provider/family/build.rs +++ b/apps/aether-gateway/src/ai_serving/planner/passthrough/provider/family/build.rs @@ -89,6 +89,21 @@ pub(crate) async fn maybe_build_stream_local_same_format_provider_decision_paylo decision: &GatewayControlDecision, body_json: &serde_json::Value, plan_kind: &str, +) -> Result, GatewayError> { + maybe_build_pinned_stream_local_same_format_provider_decision_payload( + state, parts, trace_id, decision, body_json, plan_kind, None, + ) + .await +} + +pub(crate) async fn maybe_build_pinned_stream_local_same_format_provider_decision_payload( + state: &AppState, + parts: &http::request::Parts, + trace_id: &str, + decision: &GatewayControlDecision, + body_json: &serde_json::Value, + plan_kind: &str, + pinned_candidate: Option<(&str, &str, &str)>, ) -> Result, GatewayError> { let Some(spec) = resolve_stream_spec(plan_kind) else { return Ok(None); @@ -141,6 +156,18 @@ pub(crate) async fn maybe_build_stream_local_same_format_provider_decision_paylo ) .await? { + if pinned_candidate.is_some_and(|(provider_id, endpoint_id, key_id)| { + payload.provider_id.as_deref() != Some(provider_id) + || payload.endpoint_id.as_deref() != Some(endpoint_id) + || payload.key_id.as_deref() != Some(key_id) + }) { + crate::orchestration::release_pool_key_lease_from_report_context( + state, + payload.report_context.as_ref(), + ) + .await; + continue; + } return Ok(Some(payload)); } } diff --git a/apps/aether-gateway/src/ai_serving/planner/passthrough/provider/family/mod.rs b/apps/aether-gateway/src/ai_serving/planner/passthrough/provider/family/mod.rs index 38fa23155..044f0d136 100644 --- a/apps/aether-gateway/src/ai_serving/planner/passthrough/provider/family/mod.rs +++ b/apps/aether-gateway/src/ai_serving/planner/passthrough/provider/family/mod.rs @@ -4,6 +4,7 @@ mod payload; mod request; pub(crate) use self::build::{ + maybe_build_pinned_stream_local_same_format_provider_decision_payload, maybe_build_stream_local_same_format_provider_decision_payload, maybe_build_sync_local_same_format_provider_decision_payload, }; diff --git a/apps/aether-gateway/src/ai_serving/pure/mod.rs b/apps/aether-gateway/src/ai_serving/pure/mod.rs index 16fcd2266..5a406a737 100644 --- a/apps/aether-gateway/src/ai_serving/pure/mod.rs +++ b/apps/aether-gateway/src/ai_serving/pure/mod.rs @@ -138,14 +138,14 @@ pub(crate) use aether_ai_formats::api::{ CLAUDE_CLI_STREAM_SUCCESS_REPORT_KIND, CLAUDE_CLI_SYNC_ERROR_REPORT_KIND, CLAUDE_CLI_SYNC_FINALIZE_REPORT_KIND, CLAUDE_CLI_SYNC_PLAN_KIND, CLAUDE_CLI_SYNC_SUCCESS_REPORT_KIND, CLAUDE_COUNT_TOKENS_SYNC_PLAN_KIND, - CODEX_OPENAI_IMAGE_DEFAULT_MODEL, CODEX_OPENAI_IMAGE_DEFAULT_OUTPUT_FORMAT, - CODEX_OPENAI_IMAGE_DEFAULT_VARIATION_MODEL, CODEX_OPENAI_IMAGE_DEFAULT_VARIATION_PROMPT, - CODEX_OPENAI_IMAGE_INTERNAL_MODEL, EXECUTION_RUNTIME_STREAM_ACTION, - EXECUTION_RUNTIME_STREAM_DECISION_ACTION, EXECUTION_RUNTIME_SYNC_ACTION, - EXECUTION_RUNTIME_SYNC_DECISION_ACTION, GEMINI_CHAT_STREAM_PLAN_KIND, - GEMINI_CHAT_STREAM_SUCCESS_REPORT_KIND, GEMINI_CHAT_SYNC_ERROR_REPORT_KIND, - GEMINI_CHAT_SYNC_FINALIZE_REPORT_KIND, GEMINI_CHAT_SYNC_PLAN_KIND, - GEMINI_CHAT_SYNC_SUCCESS_REPORT_KIND, GEMINI_CLI_STREAM_PLAN_KIND, + CODEX_LIVE_STREAM_PLAN_KIND, CODEX_OPENAI_IMAGE_DEFAULT_MODEL, + CODEX_OPENAI_IMAGE_DEFAULT_OUTPUT_FORMAT, CODEX_OPENAI_IMAGE_DEFAULT_VARIATION_MODEL, + CODEX_OPENAI_IMAGE_DEFAULT_VARIATION_PROMPT, CODEX_OPENAI_IMAGE_INTERNAL_MODEL, + EXECUTION_RUNTIME_STREAM_ACTION, EXECUTION_RUNTIME_STREAM_DECISION_ACTION, + EXECUTION_RUNTIME_SYNC_ACTION, EXECUTION_RUNTIME_SYNC_DECISION_ACTION, + GEMINI_CHAT_STREAM_PLAN_KIND, GEMINI_CHAT_STREAM_SUCCESS_REPORT_KIND, + GEMINI_CHAT_SYNC_ERROR_REPORT_KIND, GEMINI_CHAT_SYNC_FINALIZE_REPORT_KIND, + GEMINI_CHAT_SYNC_PLAN_KIND, GEMINI_CHAT_SYNC_SUCCESS_REPORT_KIND, GEMINI_CLI_STREAM_PLAN_KIND, GEMINI_CLI_STREAM_SUCCESS_REPORT_KIND, GEMINI_CLI_SYNC_ERROR_REPORT_KIND, GEMINI_CLI_SYNC_FINALIZE_REPORT_KIND, GEMINI_CLI_SYNC_PLAN_KIND, GEMINI_CLI_SYNC_SUCCESS_REPORT_KIND, GEMINI_CLI_V1INTERNAL_ENVELOPE_NAME, @@ -160,8 +160,8 @@ pub(crate) use aether_ai_formats::api::{ OPENAI_EMBEDDING_SYNC_PLAN_KIND, OPENAI_IMAGE_STREAM_PLAN_KIND, OPENAI_IMAGE_STREAM_SUCCESS_REPORT_KIND, OPENAI_IMAGE_SYNC_FINALIZE_REPORT_KIND, OPENAI_IMAGE_SYNC_PLAN_KIND, OPENAI_IMAGE_SYNC_SUCCESS_REPORT_KIND, - OPENAI_RERANK_SYNC_PLAN_KIND, OPENAI_RESPONSES_COMPACT_STREAM_PLAN_KIND, - OPENAI_RESPONSES_COMPACT_STREAM_SUCCESS_REPORT_KIND, + OPENAI_REALTIME_STREAM_PLAN_KIND, OPENAI_RERANK_SYNC_PLAN_KIND, + OPENAI_RESPONSES_COMPACT_STREAM_PLAN_KIND, OPENAI_RESPONSES_COMPACT_STREAM_SUCCESS_REPORT_KIND, OPENAI_RESPONSES_COMPACT_SYNC_ERROR_REPORT_KIND, OPENAI_RESPONSES_COMPACT_SYNC_FINALIZE_REPORT_KIND, OPENAI_RESPONSES_COMPACT_SYNC_PLAN_KIND, OPENAI_RESPONSES_COMPACT_SYNC_SUCCESS_REPORT_KIND, OPENAI_RESPONSES_STREAM_PLAN_KIND, diff --git a/apps/aether-gateway/src/api/ai/openai.rs b/apps/aether-gateway/src/api/ai/openai.rs index 85d9cbe2e..a3f03effa 100644 --- a/apps/aether-gateway/src/api/ai/openai.rs +++ b/apps/aether-gateway/src/api/ai/openai.rs @@ -5,9 +5,11 @@ pub(crate) fn normalized_signature(api_format: &str) -> Option<&'static str> { "openai:rerank" => Some("openai:rerank"), "openai:responses" => Some("openai:responses"), "openai:responses:compact" => Some("openai:responses:compact"), + "openai:realtime" => Some("openai:realtime"), "openai:search" => Some("openai:search"), "openai:image" => Some("openai:image"), "openai:video" => Some("openai:video"), + "codex:live" => Some("codex:live"), _ => None, } } @@ -19,9 +21,11 @@ pub(crate) fn local_path(api_format: &str) -> Option<&'static str> { "openai:rerank" => Some("/v1/rerank"), "openai:responses" => Some("/v1/responses"), "openai:responses:compact" => Some("/v1/responses/compact"), + "openai:realtime" => Some("/v1/realtime"), "openai:search" => Some("/v1/alpha/search"), "openai:image" => Some("/v1/images/generations"), "openai:video" => Some("/v1/videos"), + "codex:live" => Some("/v1/live"), _ => None, } } diff --git a/apps/aether-gateway/src/api/ai/registry.rs b/apps/aether-gateway/src/api/ai/registry.rs index ed674c2b9..312b44806 100644 --- a/apps/aether-gateway/src/api/ai/registry.rs +++ b/apps/aether-gateway/src/api/ai/registry.rs @@ -8,7 +8,7 @@ use super::{aliyun, claude, doubao, gemini, jina, openai}; use crate::api::response::build_local_http_error_response_with_request_path; use crate::headers::extract_or_generate_trace_id; use crate::{ - handlers::proxy::{live_websocket, proxy_request, responses_websocket}, + handlers::proxy::{live_websocket, proxy_request, realtime_websocket, responses_websocket}, state::AppState, GatewayError, }; @@ -65,6 +65,7 @@ pub(crate) fn mount_ai_routes(mut router: Router) -> Router }; } router = router.route("/v1/live/{call_id}", get(live_websocket)); + router = router.route("/v1/realtime", get(realtime_websocket)); for path in CLAUDE_POST_ROUTE_PATTERNS { router = router.route( path, @@ -160,6 +161,8 @@ mod tests { ), ("openai:rerank", "openai", "rerank", "/v1/rerank"), ("openai:search", "openai", "search", "/v1/alpha/search"), + ("openai:realtime", "openai", "realtime", "/v1/realtime"), + ("codex:live", "codex", "live", "/v1/live"), ("jina:rerank", "jina", "rerank", "/v1/rerank"), ] { assert_eq!( diff --git a/apps/aether-gateway/src/constants.rs b/apps/aether-gateway/src/constants.rs index e2741011d..0ee1460f7 100644 --- a/apps/aether-gateway/src/constants.rs +++ b/apps/aether-gateway/src/constants.rs @@ -125,6 +125,7 @@ pub(crate) const RUST_FRONTDOOR_OWNED_ROUTE_PATTERNS: &[&str] = &[ "/v1/messages/count_tokens", "/v1/responses", "/v1/responses/compact", + "/v1/realtime", "/v1/live", "/v1/live/{call_id}", "/v1/alpha/search", diff --git a/apps/aether-gateway/src/control/auth/credentials.rs b/apps/aether-gateway/src/control/auth/credentials.rs index d8df8378e..f9c0a1825 100644 --- a/apps/aether-gateway/src/control/auth/credentials.rs +++ b/apps/aether-gateway/src/control/auth/credentials.rs @@ -243,7 +243,7 @@ fn select_primary_credential( if signature.starts_with("claude:") { return select_claude_messages_credential(bundle); } - if signature.starts_with("openai:") { + if signature.starts_with("openai:") || signature.starts_with("codex:") { return select_openai_credential(bundle); } if signature.starts_with("aether:") { @@ -491,6 +491,25 @@ mod tests { ); } + #[test] + fn selects_codex_live_bearer_as_provider_api_key() { + let mut headers = http::HeaderMap::new(); + headers.insert( + http::header::AUTHORIZATION, + "Bearer sk-codex-live".parse().unwrap(), + ); + + let extracted = + extract_request_credentials(&headers, &uri("/v1/live?model=gpt-live"), "codex:live"); + assert_eq!( + extracted.primary, + Some(GatewayPrimaryCredential::ProviderApiKey { + raw: "sk-codex-live".to_string(), + carrier: GatewayCredentialCarrier::AuthorizationBearer, + }) + ); + } + #[test] fn prefers_claude_chat_x_api_key_over_bearer() { let mut headers = http::HeaderMap::new(); diff --git a/apps/aether-gateway/src/control/route/ai.rs b/apps/aether-gateway/src/control/route/ai.rs index 376008b48..d07b1b63d 100644 --- a/apps/aether-gateway/src/control/route/ai.rs +++ b/apps/aether-gateway/src/control/route/ai.rs @@ -35,21 +35,25 @@ pub(super) fn classify_ai_public_route( "openai:rerank", true, )) + } else if method == http::Method::GET + && normalized_path == "/v1/realtime" + && is_websocket_upgrade_request(headers) + { + Some(classified( + "ai_public", + "openai", + "realtime", + "openai:realtime", + true, + )) } else if (method == http::Method::POST && normalized_path == "/v1/live") || (method == http::Method::GET && (normalized_path == "/v1/live" || normalized_path.starts_with("/v1/live/")) && is_websocket_upgrade_request(headers)) { - // Codex Live is an experimental companion transport for an existing - // Responses mapping. It deliberately reuses the Responses permission - // surface while its wire protocol is handled by an independent relay. - Some(classified( - "ai_public", - "openai", - "live", - "openai:responses", - true, - )) + // Codex Live has an independent wire contract and permission surface; + // it must never be authorized as an OpenAI Responses request. + Some(classified("ai_public", "codex", "live", "codex:live", true)) } else if (method == http::Method::POST || (method == http::Method::GET && normalized_path == "/v1/responses" @@ -308,11 +312,32 @@ mod tests { } #[test] - fn classifies_live_http_and_websocket_routes_as_responses_companions() { + fn classifies_only_websocket_upgrade_on_realtime_route() { + let mut headers = HeaderMap::new(); + headers.insert(CONNECTION, HeaderValue::from_static("keep-alive, Upgrade")); + headers.insert(UPGRADE, HeaderValue::from_static("websocket")); + + let route = classify_ai_public_route(&Method::GET, "/v1/realtime", &headers) + .expect("Realtime WebSocket should be an AI public route"); + assert_eq!(route.route_class, "ai_public"); + assert_eq!(route.route_family, "openai"); + assert_eq!(route.route_kind, "realtime"); + assert_eq!(route.auth_endpoint_signature, "openai:realtime"); + assert!(route.execution_runtime_candidate); + + assert!( + classify_ai_public_route(&Method::GET, "/v1/realtime", &HeaderMap::new()).is_none() + ); + assert!(classify_ai_public_route(&Method::POST, "/v1/realtime", &headers).is_none()); + } + + #[test] + fn classifies_live_http_and_websocket_routes_as_codex_live() { let post = classify_ai_public_route(&Method::POST, "/v1/live", &HeaderMap::new()) .expect("Live WebRTC call creation should be an AI public route"); + assert_eq!(post.route_family, "codex"); assert_eq!(post.route_kind, "live"); - assert_eq!(post.auth_endpoint_signature, "openai:responses"); + assert_eq!(post.auth_endpoint_signature, "codex:live"); let mut headers = HeaderMap::new(); headers.insert(CONNECTION, HeaderValue::from_static("Upgrade")); @@ -320,8 +345,9 @@ mod tests { for path in ["/v1/live", "/v1/live/rtc_opaque"] { let route = classify_ai_public_route(&Method::GET, path, &headers) .expect("Live WebSocket should be an AI public route"); + assert_eq!(route.route_family, "codex"); assert_eq!(route.route_kind, "live"); - assert_eq!(route.auth_endpoint_signature, "openai:responses"); + assert_eq!(route.auth_endpoint_signature, "codex:live"); } assert!( diff --git a/apps/aether-gateway/src/control/tests/ai.rs b/apps/aether-gateway/src/control/tests/ai.rs index fbfde74bf..b795bd42c 100644 --- a/apps/aether-gateway/src/control/tests/ai.rs +++ b/apps/aether-gateway/src/control/tests/ai.rs @@ -108,6 +108,32 @@ fn classifies_openai_chat_and_responses_separately_from_embedding() { assert_ne!(responses.route_kind.as_deref(), Some("embedding")); } +#[test] +fn classifies_openai_realtime_only_for_websocket_upgrades() { + let websocket_headers = headers(&[ + ("authorization", "Bearer sk-test"), + ("connection", "keep-alive, Upgrade"), + ("upgrade", "websocket"), + ]); + let uri: Uri = "/v1/realtime?model=gpt-realtime" + .parse() + .expect("uri should parse"); + + let decision = classify_control_route(&http::Method::GET, &uri, &websocket_headers) + .expect("Realtime WebSocket route should classify"); + assert_eq!(decision.route_family.as_deref(), Some("openai")); + assert_eq!(decision.route_kind.as_deref(), Some("realtime")); + assert_eq!( + decision.auth_endpoint_signature.as_deref(), + Some("openai:realtime") + ); + assert!(decision.is_execution_runtime_candidate()); + + let plain_headers = headers(&[("authorization", "Bearer sk-test")]); + assert!(classify_control_route(&http::Method::GET, &uri, &plain_headers).is_none()); + assert!(classify_control_route(&http::Method::POST, &uri, &websocket_headers).is_none()); +} + #[test] fn classifies_openai_image_generation_and_edit_but_not_variation() { let headers = headers(&[("authorization", "Bearer sk-test")]); diff --git a/apps/aether-gateway/src/frontdoor_loop_guard.rs b/apps/aether-gateway/src/frontdoor_loop_guard.rs index f09119e64..c7aef4c38 100644 --- a/apps/aether-gateway/src/frontdoor_loop_guard.rs +++ b/apps/aether-gateway/src/frontdoor_loop_guard.rs @@ -41,6 +41,7 @@ pub(crate) fn frontdoor_self_loop_public_ai_path(path: &str) -> bool { | "/v1/rerank" | "/v1/responses" | "/v1/responses/compact" + | "/v1/realtime" | "/v1/live" | "/v1/alpha/search" | "/v1beta/files" @@ -143,3 +144,23 @@ fn normalize_host_for_frontdoor_loop_guard(host: &str) -> String { fn is_loopbackish_host(host: &str) -> bool { matches!(host, "localhost" | "127.0.0.1" | "::1" | "0.0.0.0" | "::") } + +#[cfg(test)] +mod tests { + use super::{ + frontdoor_self_loop_public_ai_path, gateway_frontdoor_self_loop_guard_matches_with_port, + }; + + #[test] + fn realtime_is_protected_from_frontdoor_self_loops() { + assert!(frontdoor_self_loop_public_ai_path("/v1/realtime")); + assert!(gateway_frontdoor_self_loop_guard_matches_with_port( + 8084, + "ws://127.0.0.1:8084/v1/realtime?model=gpt-realtime" + )); + assert!(gateway_frontdoor_self_loop_guard_matches_with_port( + 8084, + "wss://localhost:8084/v1/realtime?model=gpt-realtime" + )); + } +} diff --git a/apps/aether-gateway/src/handlers/admin/observability/usage/summary_routes.rs b/apps/aether-gateway/src/handlers/admin/observability/usage/summary_routes.rs index d31cbf9f0..be83243e4 100644 --- a/apps/aether-gateway/src/handlers/admin/observability/usage/summary_routes.rs +++ b/apps/aether-gateway/src/handlers/admin/observability/usage/summary_routes.rs @@ -52,19 +52,27 @@ fn apply_admin_usage_status_filter(query: &mut UsageAuditListQuery, status: Opti let Some(status) = status .map(str::trim) .filter(|candidate| !candidate.is_empty()) + .map(str::to_ascii_lowercase) else { return; }; - match status { - "stream" => query.is_stream = Some(true), - "standard" => query.is_stream = Some(false), + match status.as_str() { + "stream" => { + query.is_stream = Some(true); + query.is_websocket = Some(false); + } + "standard" => { + query.is_stream = Some(false); + query.is_websocket = Some(false); + } + "websocket" | "ws" => query.is_websocket = Some(true), "error" | "failed" => query.error_only = true, "active" => { query.statuses = Some(vec!["pending".to_string(), "streaming".to_string()]); } "pending" | "streaming" | "completed" | "cancelled" => { - query.statuses = Some(vec![status.to_string()]); + query.statuses = Some(vec![status]); } "has_fallback" | "has_retry" => {} _ => {} @@ -655,6 +663,7 @@ fn build_admin_usage_keyword_search_query( statuses: base_query.statuses.clone(), exclude_status_codes: base_query.exclude_status_codes.clone(), is_stream: base_query.is_stream, + is_websocket: base_query.is_websocket, error_only: base_query.error_only, keywords, matched_user_ids_by_keyword: search_context.matched_user_ids_by_keyword, @@ -1027,7 +1036,10 @@ mod tests { }; use serde_json::json; - use super::admin_usage_terminal_candidate_state_override; + use super::{ + admin_usage_terminal_candidate_state_override, build_admin_usage_keyword_search_query, + build_admin_usage_records_query, AdminUsageSearchContext, + }; fn sample_candidate( candidate_index: i32, @@ -1127,4 +1139,49 @@ mod tests { assert!(payload.is_none()); } + + #[test] + fn admin_usage_transport_statuses_are_disjoint_in_list_and_keyword_queries() { + for status in ["websocket", "ws", "WS"] { + let raw_query = format!("status={status}"); + let list_query = + build_admin_usage_records_query(100, 200, Some(&raw_query), None, None); + + assert_eq!(list_query.is_websocket, Some(true)); + assert_eq!(list_query.is_stream, None); + + let keyword_query = build_admin_usage_keyword_search_query( + &list_query, + vec!["live".to_string()], + None, + AdminUsageSearchContext::default(), + false, + false, + None, + None, + ); + assert_eq!(keyword_query.is_websocket, Some(true)); + } + + for (status, expected_stream) in [("stream", true), ("standard", false)] { + let raw_query = format!("status={status}"); + let list_query = + build_admin_usage_records_query(100, 200, Some(&raw_query), None, None); + assert_eq!(list_query.is_stream, Some(expected_stream)); + assert_eq!(list_query.is_websocket, Some(false)); + + let keyword_query = build_admin_usage_keyword_search_query( + &list_query, + vec!["live".to_string()], + None, + AdminUsageSearchContext::default(), + false, + false, + None, + None, + ); + assert_eq!(keyword_query.is_stream, Some(expected_stream)); + assert_eq!(keyword_query.is_websocket, Some(false)); + } + } } diff --git a/apps/aether-gateway/src/handlers/proxy/mod.rs b/apps/aether-gateway/src/handlers/proxy/mod.rs index 070059af7..0fe8373dc 100644 --- a/apps/aether-gateway/src/handlers/proxy/mod.rs +++ b/apps/aether-gateway/src/handlers/proxy/mod.rs @@ -10,6 +10,7 @@ use self::local::{ maybe_build_local_admin_proxy_response, maybe_build_local_internal_proxy_response, }; pub(crate) use self::websocket::live::{live_websocket, maybe_handle_live_http}; +pub(crate) use self::websocket::realtime::realtime_websocket; pub(crate) use self::websocket::responses::responses_websocket; use super::internal::resolve_local_proxy_execution_path; pub(crate) use super::public::matches_model_mapping_for_models; diff --git a/apps/aether-gateway/src/handlers/proxy/websocket/live/audit.rs b/apps/aether-gateway/src/handlers/proxy/websocket/live/audit.rs new file mode 100644 index 000000000..b8c815d57 --- /dev/null +++ b/apps/aether-gateway/src/handlers/proxy/websocket/live/audit.rs @@ -0,0 +1,395 @@ +//! Session-level audit records for Codex Live transports. +//! +//! Frameless Bidi does not expose an authoritative token/cost usage object. +//! These records therefore capture exactly one bounded lifecycle summary per +//! connection and are explicitly void for billing. They never infer tokens, +//! audio duration, or cost from frame sizes. + +use std::time::Duration; + +use aether_ai_serving::AiStreamAttempt; +use aether_contracts::ExecutionPlan; +use aether_data_contracts::repository::usage::{ + LIVE_SESSION_METADATA_KEY, USAGE_AVAILABLE_METADATA_KEY, USAGE_PRICING_AVAILABLE_METADATA_KEY, + WEBSOCKET_MODE_METADATA_KEY, WEBSOCKET_TRANSPORT_METADATA_KEY, +}; +use aether_usage_runtime::build_usage_event_data_seed; +use serde_json::{json, Map, Value}; +use tracing::warn; + +use crate::usage::{UsageEvent, UsageEventData, UsageEventType}; +use crate::AppState; + +const LIVE_AUDIT_WRITE_WAIT: Duration = Duration::from_secs(5); +const LIVE_AUDIT_SCHEMA_VERSION: &str = "1"; +const LIVE_AUDIT_LOG_TARGET: &str = "aether_gateway::handlers::proxy::codex_live"; + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(super) enum LiveAuditTransport { + WebRtc, + DirectWebSocket, + Sideband, +} + +impl LiveAuditTransport { + const fn transport(self) -> &'static str { + match self { + Self::WebRtc => "webrtc", + Self::DirectWebSocket => "websocket", + Self::Sideband => "sideband", + } + } + + const fn mode(self) -> &'static str { + match self { + Self::WebRtc => "call_create", + Self::DirectWebSocket => "direct", + Self::Sideband => "sideband", + } + } + + const fn websocket_transport(self) -> Option<&'static str> { + match self { + Self::WebRtc => None, + Self::DirectWebSocket => Some("codex_live_direct"), + Self::Sideband => Some("codex_live_sideband"), + } + } +} + +/// Marks the existing synchronous SDP call-create audit row as an unmetered +/// WebRTC control exchange. The media leg bypasses Aether after this request. +pub(super) fn mark_live_call_create_report_context(report_context: &mut Option) { + attach_live_base_metadata(report_context, LiveAuditTransport::WebRtc); +} + +fn attach_live_base_metadata(report_context: &mut Option, transport: LiveAuditTransport) { + let object = report_context_object(report_context); + object.insert(USAGE_AVAILABLE_METADATA_KEY.to_string(), Value::Bool(false)); + object.insert( + USAGE_PRICING_AVAILABLE_METADATA_KEY.to_string(), + Value::Bool(false), + ); + object.insert( + WEBSOCKET_MODE_METADATA_KEY.to_string(), + Value::Bool(transport.websocket_transport().is_some()), + ); + if let Some(websocket_transport) = transport.websocket_transport() { + object.insert( + WEBSOCKET_TRANSPORT_METADATA_KEY.to_string(), + Value::String(websocket_transport.to_string()), + ); + } else { + object.remove(WEBSOCKET_TRANSPORT_METADATA_KEY); + } + object.insert( + LIVE_SESSION_METADATA_KEY.to_string(), + json!({ + "schema_version": LIVE_AUDIT_SCHEMA_VERSION, + "transport": transport.transport(), + "mode": transport.mode(), + "usage_state": "unavailable", + }), + ); +} + +fn report_context_object(report_context: &mut Option) -> &mut Map { + if !matches!(report_context, Some(Value::Object(_))) { + let seed = report_context.take(); + let mut object = Map::new(); + if let Some(seed) = seed.filter(|value| !value.is_null()) { + object.insert("seed".to_string(), seed); + } + *report_context = Some(Value::Object(object)); + } + report_context + .as_mut() + .and_then(Value::as_object_mut) + .expect("Live audit report context was normalized to an object") +} + +pub(super) struct LiveSessionAudit { + plan: ExecutionPlan, + report_context: Option, + transport: LiveAuditTransport, +} + +impl LiveSessionAudit { + pub(super) fn from_attempt(attempt: &AiStreamAttempt, transport: LiveAuditTransport) -> Self { + let mut report_context = attempt.report_context.clone(); + attach_live_base_metadata(&mut report_context, transport); + Self { + plan: attempt.plan.clone(), + report_context, + transport, + } + } + + /// Persists one terminal lifecycle row. The spawned write remains alive if + /// the bounded caller wait elapses, so closing a socket cannot silently + /// cancel the only audit write for that connection. + pub(super) async fn finish(self, state: &AppState, terminal: LiveSessionTerminal) { + let request_id = self.plan.request_id.clone(); + let event = self.build_terminal_event(terminal); + let usage_runtime = std::sync::Arc::clone(&state.usage_runtime); + let usage_data = std::sync::Arc::clone(state.usage_lifecycle_data_state()); + let task = tokio::spawn(async move { + usage_runtime + .record_terminal_event_direct(usage_data.as_ref(), event) + .await; + }); + match tokio::time::timeout(LIVE_AUDIT_WRITE_WAIT, task).await { + Ok(Ok(())) => {} + Ok(Err(error)) => warn!( + target: LIVE_AUDIT_LOG_TARGET, + event_name = "codex_live_session_audit_task_failed", + log_type = "ops", + request_id, + error = %error, + "Codex Live session audit task failed" + ), + Err(_) => warn!( + target: LIVE_AUDIT_LOG_TARGET, + event_name = "codex_live_session_audit_write_slow", + log_type = "ops", + request_id, + wait_ms = LIVE_AUDIT_WRITE_WAIT.as_millis() as u64, + write_detached = true, + "Codex Live stopped waiting for a slow session audit write" + ), + } + } + + fn build_terminal_event(self, terminal: LiveSessionTerminal) -> UsageEvent { + let mut data = build_usage_event_data_seed(&self.plan, self.report_context.as_ref()); + data.request_type = Some("live".to_string()); + data.is_stream = Some(self.transport != LiveAuditTransport::WebRtc); + data.status_code = Some(terminal.status_code); + data.response_time_ms = Some(terminal.elapsed_ms); + data.first_byte_time_ms = terminal.first_upstream_frame_ms; + data.input_tokens = None; + data.output_tokens = None; + data.total_tokens = None; + data.cache_creation_input_tokens = None; + data.cache_creation_ephemeral_5m_input_tokens = None; + data.cache_creation_ephemeral_1h_input_tokens = None; + data.cache_read_input_tokens = None; + data.cache_creation_cost_usd = None; + data.cache_read_cost_usd = None; + data.total_cost_usd = None; + data.actual_total_cost_usd = None; + if terminal.disposition != LiveSessionDisposition::Completed { + data.error_message = Some(terminal.termination.to_string()); + data.error_category = Some(terminal.disposition.error_category().to_string()); + } + data.request_metadata = + attach_terminal_metadata(data.request_metadata, self.transport, &terminal); + UsageEvent::new( + terminal.disposition.event_type(), + self.plan.request_id, + data, + ) + } +} + +fn attach_terminal_metadata( + metadata: Option, + transport: LiveAuditTransport, + terminal: &LiveSessionTerminal, +) -> Option { + let mut object = match metadata { + Some(Value::Object(object)) => object, + _ => Map::new(), + }; + object.insert(USAGE_AVAILABLE_METADATA_KEY.to_string(), Value::Bool(false)); + object.insert( + USAGE_PRICING_AVAILABLE_METADATA_KEY.to_string(), + Value::Bool(false), + ); + object.insert( + WEBSOCKET_MODE_METADATA_KEY.to_string(), + Value::Bool(transport.websocket_transport().is_some()), + ); + if let Some(websocket_transport) = transport.websocket_transport() { + object.insert( + WEBSOCKET_TRANSPORT_METADATA_KEY.to_string(), + Value::String(websocket_transport.to_string()), + ); + } + object.insert( + LIVE_SESSION_METADATA_KEY.to_string(), + json!({ + "schema_version": LIVE_AUDIT_SCHEMA_VERSION, + "transport": transport.transport(), + "mode": transport.mode(), + "state": terminal.disposition.state(), + "termination": terminal.termination, + "elapsed_ms": terminal.elapsed_ms, + "client_frames": terminal.client_frames, + "client_bytes": terminal.client_bytes, + "upstream_frames": terminal.upstream_frames, + "upstream_bytes": terminal.upstream_bytes, + "first_upstream_frame_ms": terminal.first_upstream_frame_ms, + "usage_state": "unavailable", + }), + ); + Some(Value::Object(object)) +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(super) enum LiveSessionDisposition { + Completed, + Failed, + Cancelled, +} + +impl LiveSessionDisposition { + const fn event_type(self) -> UsageEventType { + match self { + Self::Completed => UsageEventType::Completed, + Self::Failed => UsageEventType::Failed, + Self::Cancelled => UsageEventType::Cancelled, + } + } + + const fn state(self) -> &'static str { + match self { + Self::Completed => "closed", + Self::Failed => "failed", + Self::Cancelled => "cancelled", + } + } + + const fn error_category(self) -> &'static str { + match self { + Self::Completed => "none", + Self::Failed => "transport_error", + Self::Cancelled => "client_cancelled", + } + } +} + +#[derive(Debug, Clone, Copy)] +pub(super) struct LiveSessionTerminal { + pub(super) disposition: LiveSessionDisposition, + pub(super) status_code: u16, + pub(super) termination: &'static str, + pub(super) elapsed_ms: u64, + pub(super) first_upstream_frame_ms: Option, + pub(super) client_frames: u64, + pub(super) client_bytes: u64, + pub(super) upstream_frames: u64, + pub(super) upstream_bytes: u64, +} + +impl LiveSessionTerminal { + pub(super) const fn failure( + status_code: u16, + termination: &'static str, + elapsed_ms: u64, + ) -> Self { + Self { + disposition: LiveSessionDisposition::Failed, + status_code, + termination, + elapsed_ms, + first_upstream_frame_ms: None, + client_frames: 0, + client_bytes: 0, + upstream_frames: 0, + upstream_bytes: 0, + } + } +} + +#[cfg(test)] +mod tests { + use std::collections::BTreeMap; + + use aether_contracts::{ExecutionTimeouts, RequestBody}; + + use super::*; + + fn sample_attempt() -> AiStreamAttempt { + AiStreamAttempt { + plan: ExecutionPlan { + request_id: "live-request".to_string(), + candidate_id: Some("candidate-live".to_string()), + provider_name: Some("Codex".to_string()), + provider_id: "provider-live".to_string(), + endpoint_id: "endpoint-live".to_string(), + key_id: "key-live".to_string(), + method: "GET".to_string(), + url: "wss://example.test/v1/live".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: "codex:live".to_string(), + provider_api_format: "codex:live".to_string(), + model_name: Some("gpt-live".to_string()), + proxy: None, + transport_profile: None, + timeouts: Some(ExecutionTimeouts::default()), + }, + report_kind: Some("openai_responses_stream".to_string()), + report_context: Some(json!({ + "user_id": "user-live", + "api_key_id": "gateway-key-live", + "trace_id": "trace-live" + })), + } + } + + #[test] + fn direct_terminal_audit_is_opaque_unmetered_and_void_eligible() { + let audit = + LiveSessionAudit::from_attempt(&sample_attempt(), LiveAuditTransport::DirectWebSocket); + let event = audit.build_terminal_event(LiveSessionTerminal { + disposition: LiveSessionDisposition::Completed, + status_code: 200, + termination: "client_close_frame", + elapsed_ms: 1234, + first_upstream_frame_ms: Some(42), + client_frames: 3, + client_bytes: 128, + upstream_frames: 5, + upstream_bytes: 512, + }); + + assert_eq!(event.event_type, UsageEventType::Completed); + assert_eq!(event.data.input_tokens, None); + assert_eq!(event.data.total_cost_usd, None); + let metadata = event.data.request_metadata.expect("metadata"); + assert_eq!(metadata[USAGE_AVAILABLE_METADATA_KEY], false); + assert_eq!(metadata[USAGE_PRICING_AVAILABLE_METADATA_KEY], false); + assert_eq!(metadata[WEBSOCKET_MODE_METADATA_KEY], true); + assert_eq!( + metadata[WEBSOCKET_TRANSPORT_METADATA_KEY], + "codex_live_direct" + ); + assert_eq!(metadata[LIVE_SESSION_METADATA_KEY]["client_frames"], 3); + assert_eq!( + metadata[LIVE_SESSION_METADATA_KEY]["usage_state"], + "unavailable" + ); + } + + #[test] + fn call_create_is_webrtc_not_websocket() { + let mut context = Some(json!({"trace_id": "trace-live"})); + mark_live_call_create_report_context(&mut context); + let context = context.expect("context"); + + assert_eq!(context[USAGE_AVAILABLE_METADATA_KEY], false); + assert_eq!(context[USAGE_PRICING_AVAILABLE_METADATA_KEY], false); + assert_eq!(context[WEBSOCKET_MODE_METADATA_KEY], false); + assert!(context.get(WEBSOCKET_TRANSPORT_METADATA_KEY).is_none()); + assert_eq!(context[LIVE_SESSION_METADATA_KEY]["transport"], "webrtc"); + } +} diff --git a/apps/aether-gateway/src/handlers/proxy/websocket/live/http.rs b/apps/aether-gateway/src/handlers/proxy/websocket/live/http.rs index 9340b67b6..5ae34f4e7 100644 --- a/apps/aether-gateway/src/handlers/proxy/websocket/live/http.rs +++ b/apps/aether-gateway/src/handlers/proxy/websocket/live/http.rs @@ -22,6 +22,7 @@ use crate::execution_runtime::execute_execution_runtime_sync_plan_with_report_co use crate::handlers::proxy::websocket::responses::ResponsesWebSocketTurnAdmission; use crate::{AppState, GatewayError}; +use super::audit::mark_live_call_create_report_context; use super::live_usage_accounting_is_safe; use super::planner::{live_call_url, plan_live_candidate, LiveAuthMode, LivePoolLeaseGuard}; use super::protocol::{build_live_multipart, extract_call_id_from_location, parse_live_multipart}; @@ -158,7 +159,7 @@ pub(crate) async fn maybe_handle_live_http( .to_string(), ); - let Some(attempt) = + let Some(mut attempt) = build_standard_sync_plan_from_decision(parts, &provider_body_marker, candidate.execution)? else { lease.release().await; @@ -168,6 +169,10 @@ pub(crate) async fn maybe_handle_live_http( "Codex Live provider request could not be built", )?)); }; + // The synchronous SDP exchange has an ordinary request lifecycle, but it + // does not contain the media leg's token/cost usage. Keep the existing row + // while making that boundary explicit and non-billable. + mark_live_call_create_report_context(&mut attempt.report_context); if let Some(rejection) = execution_plan_balance_capacity_rejection( state, control_decision, diff --git a/apps/aether-gateway/src/handlers/proxy/websocket/live/mod.rs b/apps/aether-gateway/src/handlers/proxy/websocket/live/mod.rs index 90a93640b..474129eb8 100644 --- a/apps/aether-gateway/src/handlers/proxy/websocket/live/mod.rs +++ b/apps/aether-gateway/src/handlers/proxy/websocket/live/mod.rs @@ -5,6 +5,7 @@ //! Keeping it in an independent module prevents a `session.update` frame from //! ever entering the Responses `response.create` state machine. +mod audit; mod http; mod planner; mod protocol; diff --git a/apps/aether-gateway/src/handlers/proxy/websocket/live/planner.rs b/apps/aether-gateway/src/handlers/proxy/websocket/live/planner.rs index 189cb148f..ced20ff5e 100644 --- a/apps/aether-gateway/src/handlers/proxy/websocket/live/planner.rs +++ b/apps/aether-gateway/src/handlers/proxy/websocket/live/planner.rs @@ -1,11 +1,10 @@ //! Candidate planning and provider request shaping for Codex Live. //! -//! Live deliberately reuses the existing Responses permission and scheduler -//! surface. Only the selected candidate, model alias and transport identity are -//! reused; Responses body normalization and its WebSocket state machine never -//! see a Live protocol frame. +//! Live has its own endpoint and permission surface. Candidate selection, +//! model aliases and transport policy are shared with the ordinary scheduler, +//! but Responses body normalization and its WebSocket state machine never see +//! a Live protocol frame. -use std::collections::BTreeSet; use std::net::SocketAddr; use std::sync::atomic::{AtomicBool, Ordering}; use std::sync::Arc; @@ -18,8 +17,9 @@ use sha2::{Digest, Sha256}; use url::{form_urlencoded, Url}; use crate::ai_serving::{ - build_standard_stream_plan_from_decision, maybe_build_responses_websocket_decision, - AiExecutionDecision, AiStreamAttempt, ResponsesWebSocketPinnedCandidate, + build_standard_stream_plan_from_decision, + maybe_build_pinned_stream_local_same_format_provider_decision_payload, AiExecutionDecision, + AiStreamAttempt, ResponsesWebSocketPinnedCandidate, }; use crate::control::GatewayControlDecision; use crate::headers::request_origin_from_headers_and_remote_addr; @@ -155,30 +155,59 @@ pub(super) async fn plan_live_candidate( } let parts = build_live_planning_parts(headers, remote_addr); let body = json!({"model": client_model, "input": []}); - let planned = maybe_build_responses_websocket_decision( + let execution = maybe_build_pinned_stream_local_same_format_provider_decision_payload( state, &parts, trace_id, decision, - None, &body, - None::<&BTreeSet>, - None::<&BTreeSet>, - pinned_candidate, + crate::ai_serving::CODEX_LIVE_STREAM_PLAN_KIND, + pinned_candidate + .map(|pinned| (pinned.provider_id(), pinned.endpoint_id(), pinned.key_id())), ) .await?; - let Some(planned) = planned else { + let Some(mut execution) = execution else { + return Ok(None); + }; + if execution + .provider_api_format + .as_deref() + .map(crate::ai_serving::normalize_api_format_alias) + .as_deref() + != Some("codex:live") + { + crate::orchestration::release_pool_key_lease_from_report_context( + state, + execution.report_context.as_ref(), + ) + .await; + return Ok(None); + } + let Some(effective_auth_type) = execution + .report_context + .as_ref() + .and_then(|context| context.get("upstream_credential_mode")) + .and_then(serde_json::Value::as_str) + .map(str::trim) + .filter(|value| !value.is_empty()) + .map(str::to_string) + else { + crate::orchestration::release_pool_key_lease_from_report_context( + state, + execution.report_context.as_ref(), + ) + .await; return Ok(None); }; - let effective_auth_type = planned.effective_auth_type; - let mut execution = planned.execution; let provider_type = execution .provider_type .as_deref() .map(str::trim) .unwrap_or_default(); - if !provider_type.eq_ignore_ascii_case("codex") && !provider_type.eq_ignore_ascii_case("openai") - { + if !matches!( + provider_type.to_ascii_lowercase().as_str(), + "codex" | "openai" | "custom" + ) { crate::orchestration::release_pool_key_lease_from_report_context( state, execution.report_context.as_ref(), @@ -247,7 +276,7 @@ pub(super) fn direct_live_websocket_url( if candidate.auth_mode == LiveAuthMode::ChatGptOauth { return Err(LiveProtocolError::OauthDirectWebSocketUnsupported); } - replace_responses_suffix( + replace_live_suffix( candidate.execution.upstream_url.as_deref(), &["live"], Some(("model", candidate.provider_model.as_str())), @@ -257,12 +286,12 @@ pub(super) fn direct_live_websocket_url( pub(super) fn live_call_url(candidate: &PlannedLiveCandidate) -> Result { match candidate.auth_mode { LiveAuthMode::ApiKey => { - replace_responses_suffix(candidate.execution.upstream_url.as_deref(), &["live"], None) + replace_live_suffix(candidate.execution.upstream_url.as_deref(), &["live"], None) } LiveAuthMode::ChatGptOauth => { let source = validated_official_chatgpt_url(candidate.execution.upstream_url.as_deref())?; - replace_responses_suffix( + replace_live_suffix( Some(source.as_str()), &["realtime", "calls"], Some(("intent", "quicksilver")), @@ -283,7 +312,7 @@ pub(super) fn live_sideband_url( ) -> Result { super::protocol::validate_call_id(call_id)?; match candidate.auth_mode { - LiveAuthMode::ApiKey => replace_responses_suffix( + LiveAuthMode::ApiKey => replace_live_suffix( candidate.execution.upstream_url.as_deref(), &["live", call_id], None, @@ -353,7 +382,7 @@ fn live_routing_fingerprint( } let path = url.path().trim_end_matches('/'); let path_family = path - .strip_suffix("/responses") + .strip_suffix("/live") .ok_or(LiveProtocolError::InvalidUpstreamUrl)?; if auth_mode == LiveAuthMode::ChatGptOauth { validated_official_chatgpt_url(Some(raw_url))?; @@ -419,15 +448,14 @@ fn validated_official_chatgpt_url(raw: Option<&str>) -> Result, suffix: &[&str], query: Option<(&str, &str)>, @@ -441,7 +469,7 @@ fn replace_responses_suffix( { return Err(LiveProtocolError::InvalidUpstreamUrl); } - if url.path_segments().and_then(Iterator::last) != Some("responses") { + if url.path_segments().and_then(Iterator::last) != Some("live") { return Err(LiveProtocolError::InvalidUpstreamUrl); } { @@ -524,8 +552,8 @@ fn build_live_planning_parts( remote_addr: &SocketAddr, ) -> http::request::Parts { let mut request = http::Request::builder() - .method(Method::POST) - .uri("/v1/responses") + .method(Method::GET) + .uri("/v1/live") .body(()) .expect("the fixed Live planning request must be valid"); *request.headers_mut() = sanitize_live_planning_headers(headers.clone()); @@ -658,9 +686,9 @@ mod tests { endpoint: GatewayProviderTransportEndpoint { id: "endpoint-1".to_string(), provider_id: "provider-1".to_string(), - api_format: "openai:responses".to_string(), - api_family: Some("openai".to_string()), - endpoint_kind: Some("responses".to_string()), + api_format: "codex:live".to_string(), + api_family: Some("codex".to_string()), + endpoint_kind: Some("live".to_string()), is_active: true, base_url: "https://chatgpt.com/backend-api/codex".to_string(), header_rules: None, @@ -677,7 +705,7 @@ mod tests { name: "key".to_string(), auth_type: default_auth_type.to_string(), is_active: true, - api_formats: Some(vec!["openai:responses".to_string()]), + api_formats: Some(vec!["codex:live".to_string()]), auth_type_by_format, allow_auth_channel_mismatch_formats: None, allowed_models: None, @@ -701,7 +729,7 @@ mod tests { session_id: &str, ) -> AiExecutionDecision { let mut decision = candidate( - "https://chatgpt.com/backend-api/codex/responses", + "https://chatgpt.com/backend-api/codex/live", LiveAuthMode::ChatGptOauth, ) .execution; @@ -717,7 +745,7 @@ mod tests { #[test] fn format_auth_override_selects_the_effective_live_auth_mode() { let overridden = - transport_with_auth_override("oauth", Some(json!({"openai:responses": "bearer"}))); + transport_with_auth_override("oauth", Some(json!({"codex:live": "bearer"}))); let effective = aether_provider_transport::auth::resolve_local_auth_type_for_transport_format( &overridden, @@ -744,7 +772,7 @@ mod tests { #[test] fn derives_api_key_live_urls_preserves_query_and_replaces_the_mapped_model() { let mut candidate = candidate( - "https://api.example.test/v1/responses?api-version=2026-08-01&model=stale&MODEL=duplicate", + "https://api.example.test/v1/live?api-version=2026-08-01&model=stale&MODEL=duplicate", LiveAuthMode::ApiKey, ); candidate.provider_model = "upstream/model + future".to_string(); @@ -778,7 +806,7 @@ mod tests { #[test] fn derives_chatgpt_call_and_official_sideband_urls() { let candidate = candidate( - "https://chatgpt.com/backend-api/codex/responses?api-version=2026-08-01&intent=stale&INTENT=duplicate&architecture=stale&ARCHITECTURE=duplicate", + "https://chatgpt.com/backend-api/codex/live?api-version=2026-08-01&intent=stale&INTENT=duplicate&architecture=stale&ARCHITECTURE=duplicate", LiveAuthMode::ChatGptOauth, ); let call = Url::parse(live_call_url(&candidate).unwrap().as_str()).unwrap(); @@ -816,7 +844,7 @@ mod tests { #[test] fn chatgpt_oauth_live_fails_closed_for_custom_backend_origins() { let candidate = candidate( - "https://relay.example/backend-api/codex/responses", + "https://relay.example/backend-api/codex/live", LiveAuthMode::ChatGptOauth, ); assert_eq!( @@ -861,8 +889,7 @@ mod tests { #[test] fn routing_fingerprint_binds_api_key_origin_without_hashing_the_token() { - let mut first = - candidate("https://api-a.example/v1/responses", LiveAuthMode::ApiKey).execution; + let mut first = candidate("https://api-a.example/v1/live", LiveAuthMode::ApiKey).execution; first.provider_request_headers.extend([ ("authorization".to_string(), "Bearer token-1".to_string()), ("x-session-id".to_string(), "session-1".to_string()), @@ -877,7 +904,7 @@ mod tests { ); let mut changed_origin = - candidate("https://api-b.example/v1/responses", LiveAuthMode::ApiKey).execution; + candidate("https://api-b.example/v1/live", LiveAuthMode::ApiKey).execution; changed_origin .provider_request_headers .insert("x-session-id".to_string(), "session-1".to_string()); @@ -896,7 +923,7 @@ mod tests { ); let missing_session = - candidate("https://api-a.example/v1/responses", LiveAuthMode::ApiKey).execution; + candidate("https://api-a.example/v1/live", LiveAuthMode::ApiKey).execution; assert_eq!( live_routing_fingerprint(&missing_session, "bearer", LiveAuthMode::ApiKey), Err(LiveProtocolError::InvalidUpstreamUrl) @@ -906,7 +933,7 @@ mod tests { #[test] fn routing_fingerprint_canonicalizes_safe_query_and_ignores_query_credentials() { let mut baseline = candidate( - "https://api-a.example/v1/responses?api-version=2026-08-01&deployment=primary&alt=sse&token=secret-1&key=secret-1", + "https://api-a.example/v1/live?api-version=2026-08-01&deployment=primary&alt=sse&token=secret-1&key=secret-1", LiveAuthMode::ApiKey, ) .execution; @@ -917,7 +944,7 @@ mod tests { live_routing_fingerprint(&baseline, "bearer", LiveAuthMode::ApiKey).unwrap(); let mut reordered = candidate( - "https://api-a.example/v1/responses?key=secret-2&alt=sse&token=secret-2&deployment=primary&api-version=2026-08-01", + "https://api-a.example/v1/live?key=secret-2&alt=sse&token=secret-2&deployment=primary&api-version=2026-08-01", LiveAuthMode::ApiKey, ) .execution; @@ -931,7 +958,7 @@ mod tests { ); let mut changed_route = candidate( - "https://api-a.example/v1/responses?api-version=2026-08-01&deployment=secondary&alt=sse&token=secret-2&key=secret-2", + "https://api-a.example/v1/live?api-version=2026-08-01&deployment=secondary&alt=sse&token=secret-2&key=secret-2", LiveAuthMode::ApiKey, ) .execution; @@ -989,10 +1016,7 @@ mod tests { #[test] fn live_urls_reject_credentials_invalid_suffixes_and_call_ids() { - let credentials = candidate( - "https://token@example.test/v1/responses", - LiveAuthMode::ApiKey, - ); + let credentials = candidate("https://token@example.test/v1/live", LiveAuthMode::ApiKey); assert_eq!( direct_live_websocket_url(&credentials), Err(LiveProtocolError::InvalidUpstreamUrl) @@ -1012,7 +1036,7 @@ mod tests { ); let fragment = candidate( - "https://api.example.test/v1/responses#not-sent-upstream", + "https://api.example.test/v1/live#not-sent-upstream", LiveAuthMode::ApiKey, ); assert_eq!( diff --git a/apps/aether-gateway/src/handlers/proxy/websocket/live/session.rs b/apps/aether-gateway/src/handlers/proxy/websocket/live/session.rs index 34020bba1..c1c5b2fe8 100644 --- a/apps/aether-gateway/src/handlers/proxy/websocket/live/session.rs +++ b/apps/aether-gateway/src/handlers/proxy/websocket/live/session.rs @@ -1,11 +1,12 @@ //! Opaque direct and WebRTC-sideband WebSocket relay for Codex Live. use std::future::Future; +use std::sync::{Arc, Mutex}; use std::time::{Duration, Instant}; use axum::extract::ws::{Message as AxumWsMessage, WebSocket}; use axum::http::StatusCode; -use futures_util::StreamExt; +use futures_util::{SinkExt, StreamExt}; use serde_json::{json, Value}; use tracing::{info, warn}; use wreq::ws::message::Message as WreqWsMessage; @@ -16,16 +17,20 @@ use crate::handlers::proxy::websocket::ingress::{ }; use crate::handlers::proxy::websocket::responses::ResponsesWebSocketTurnAdmission; use crate::handlers::proxy::websocket::session::{ - wait_for_optional_deadline, CLOSE_INTERNAL_ERROR, CLOSE_POLICY_VIOLATION, CLOSE_TRY_AGAIN, + wait_for_optional_deadline, CLOSE_POLICY_VIOLATION, CLOSE_TRY_AGAIN, LIVE_WEBSOCKET_SESSION_LIMITS, WEBSOCKET_LOG_TRANSPORT, }; use crate::handlers::proxy::websocket::transport::{ client_message_to_upstream, close_client_socket, close_upstream_socket, connect_upstream_websocket, send_client_message, send_upstream_message, - upstream_message_to_client, UpstreamWebSocketErrorCodes, + upstream_message_to_client, websocket_relay_frame_queue, UpstreamWebSocketErrorCodes, + WebSocketRelayPumpControl, WebSocketRelayQueueError, WebSocketWriteError, }; use crate::{AppState, GatewayError}; +use super::audit::{ + LiveAuditTransport, LiveSessionAudit, LiveSessionDisposition, LiveSessionTerminal, +}; use super::live_usage_accounting_is_safe; use super::planner::{ build_live_stream_admission_attempt, direct_live_websocket_url, live_sideband_url, @@ -70,72 +75,85 @@ enum LiveRelayAdmissionError { Gateway(GatewayError), } +struct LiveRelayAdmission { + capacity: ResponsesWebSocketTurnAdmission, + audit: LiveSessionAudit, +} + +struct LiveRelayAdmissionFailure { + error: LiveRelayAdmissionError, + audit: Option, +} + impl LiveRelayAdmissionError { - fn kind(&self) -> &'static str { + fn status(&self) -> StatusCode { match self { - Self::PlanUnavailable => "plan_unavailable", - Self::BalanceRejected => "balance_rejected", - Self::Gateway(error) => gateway_error_kind(error), + Self::PlanUnavailable => StatusCode::BAD_GATEWAY, + Self::BalanceRejected | Self::Gateway(GatewayError::AdmissionTimeout { .. }) => { + StatusCode::TOO_MANY_REQUESTS + } + Self::Gateway(GatewayError::Client { status, .. }) => *status, + Self::Gateway(GatewayError::LocalExecutionPlanningTimeout { .. }) => { + StatusCode::GATEWAY_TIMEOUT + } + Self::Gateway(GatewayError::UpstreamUnavailable { .. }) + | Self::Gateway(GatewayError::ControlUnavailable { .. }) => { + StatusCode::SERVICE_UNAVAILABLE + } + Self::Gateway(GatewayError::Internal(_)) => StatusCode::INTERNAL_SERVER_ERROR, } } - fn response(&self) -> (u16, &'static str, &'static str, u16, &'static str) { + fn client_message(&self) -> &'static str { match self { - Self::PlanUnavailable => ( - 502, - "codex_live_admission_plan_unavailable", - "Codex Live provider admission could not be prepared", - CLOSE_INTERNAL_ERROR, - "Live admission plan unavailable", - ), - Self::BalanceRejected => ( - 429, - "codex_live_balance_rejected", - "Codex Live request capacity is unavailable", - CLOSE_POLICY_VIOLATION, - "Live request capacity unavailable", - ), - Self::Gateway(GatewayError::AdmissionTimeout { .. }) => ( - 429, - "codex_live_admission_timeout", - "Gateway capacity is busy; retry this Live connection", - CLOSE_TRY_AGAIN, - "Live admission timeout", - ), - Self::Gateway(GatewayError::Client { status, .. }) => ( - status.as_u16(), - "codex_live_request_rejected", - "Codex Live request was not allowed", - CLOSE_POLICY_VIOLATION, - "Live request rejected", - ), - Self::Gateway(GatewayError::LocalExecutionPlanningTimeout { .. }) => ( - 504, - "codex_live_admission_planning_timeout", - "Codex Live admission planning timed out", - CLOSE_TRY_AGAIN, - "Live admission planning timeout", - ), - Self::Gateway(_) => ( - 500, - "codex_live_admission_failed", - "Gateway could not admit this Codex Live connection", - CLOSE_INTERNAL_ERROR, - "Live admission failed", - ), + Self::PlanUnavailable => "Codex Live provider admission could not be prepared", + Self::BalanceRejected => "Codex Live request capacity is unavailable", + Self::Gateway(GatewayError::AdmissionTimeout { .. }) => { + "Gateway capacity is busy; retry this Live connection" + } + Self::Gateway(GatewayError::Client { .. }) => "Codex Live request was not allowed", + Self::Gateway(GatewayError::LocalExecutionPlanningTimeout { .. }) => { + "Codex Live admission planning timed out" + } + Self::Gateway(_) => "Gateway could not admit this Codex Live connection", + } + } + + fn termination(&self) -> &'static str { + match self { + Self::PlanUnavailable => "admission_plan_unavailable", + Self::BalanceRejected => "balance_rejected", + Self::Gateway(GatewayError::AdmissionTimeout { .. }) => "admission_timeout", + Self::Gateway(GatewayError::Client { .. }) => "request_rejected", + Self::Gateway(GatewayError::LocalExecutionPlanningTimeout { .. }) => { + "admission_planning_timeout" + } + Self::Gateway(GatewayError::UpstreamUnavailable { .. }) => "upstream_unavailable", + Self::Gateway(GatewayError::ControlUnavailable { .. }) => "control_unavailable", + Self::Gateway(GatewayError::Internal(_)) => "admission_failed", } } } pub(super) enum PreparedLiveWebSocket { - Direct { client_model: String }, + Direct(PreparedLiveRelay), Sideband(PreparedLiveSideband), } +pub(super) struct PreparedLiveRelay { + upstream: wreq::ws::WebSocket, + admission: ResponsesWebSocketTurnAdmission, + audit: LiveSessionAudit, + pool_lease: LivePoolLeaseGuard, + provider_id: String, + endpoint_id: String, + key_id: String, + provider_model: String, +} + pub(super) struct PreparedLiveSideband { - call_id: String, - binding: LiveCallBinding, - lease: LiveSidebandLease, + relay: PreparedLiveRelay, + sideband_lease: LiveSidebandLease, } pub(super) struct LiveWebSocketPreflightRejection { @@ -167,33 +185,52 @@ pub(super) async fn prepare_live_websocket( trace_id = %context.trace_id, "Codex Live rejected a finite-balance principal before WebSocket upgrade because Frameless usage is unavailable" ); - return Err(LiveWebSocketPreflightRejection { - status: StatusCode::NOT_IMPLEMENTED, - message: "Codex Live is unavailable for finite-balance keys until Frameless usage settlement is supported", - }); + return Err(preflight_rejection( + context, + "unknown", + StatusCode::NOT_IMPLEMENTED, + "usage_settlement_unavailable", + "Codex Live is unavailable for finite-balance keys until Frameless usage settlement is supported", + )); } if context.uri.path() == "/v1/live" { - let client_model = direct_model_from_query(context.uri.query()).map_err(|error| { - LiveWebSocketPreflightRejection { - status: error.status_code(), - message: error.client_message(), + let client_model = match direct_model_from_query(context.uri.query()) { + Ok(model) => model, + Err(error) => { + return Err(preflight_rejection( + context, + "direct", + error.status_code(), + error.code(), + error.client_message(), + )); } - })?; - return Ok(PreparedLiveWebSocket::Direct { client_model }); + }; + return prepare_direct_live_websocket(state, context, client_model.as_str()) + .await + .map(PreparedLiveWebSocket::Direct); } - let call_id = - call_id_from_path(context.uri.path()).map_err(|error| LiveWebSocketPreflightRejection { - status: error.status_code(), - message: error.client_message(), - })?; - let auth = context - .decision - .auth_context - .as_ref() - .ok_or(LiveWebSocketPreflightRejection { - status: StatusCode::UNAUTHORIZED, - message: "Authentication required", - })?; + let call_id = match call_id_from_path(context.uri.path()) { + Ok(call_id) => call_id, + Err(error) => { + return Err(preflight_rejection( + context, + "sideband", + error.status_code(), + error.code(), + error.client_message(), + )); + } + }; + let Some(auth) = context.decision.auth_context.as_ref() else { + return Err(preflight_rejection( + context, + "sideband", + StatusCode::UNAUTHORIZED, + "authentication_required", + "Authentication required", + )); + }; let registry = LiveCallRegistry::new(std::sync::Arc::clone(&state.runtime_state)); let lookup = tokio::time::timeout( SIDEBAND_LOOKUP_TIMEOUT, @@ -216,10 +253,13 @@ pub(super) async fn prepare_live_websocket( trace_id = %context.trace_id, "Codex Live sideband binding was not found" ); - return Err(LiveWebSocketPreflightRejection { - status: StatusCode::NOT_FOUND, - message: "Codex Live call binding was not found", - }); + return Err(preflight_rejection( + context, + "sideband", + StatusCode::NOT_FOUND, + "sideband_binding_missing", + "Codex Live call binding was not found", + )); } Ok(Ok(LiveCallLookup::Expired)) => { info!( @@ -231,26 +271,236 @@ pub(super) async fn prepare_live_websocket( trace_id = %context.trace_id, "Codex Live sideband binding has expired" ); - return Err(LiveWebSocketPreflightRejection { - status: StatusCode::GONE, - message: "Codex Live call binding has expired", - }); + return Err(preflight_rejection( + context, + "sideband", + StatusCode::GONE, + "sideband_binding_expired", + "Codex Live call binding has expired", + )); } Ok(Err(error)) => { log_registry_error(context, &error); - return Err(LiveWebSocketPreflightRejection { - status: StatusCode::SERVICE_UNAVAILABLE, - message: "Codex Live sideband binding is temporarily unavailable", - }); + return Err(preflight_rejection( + context, + "sideband", + StatusCode::SERVICE_UNAVAILABLE, + "sideband_binding_unavailable", + "Codex Live sideband binding is temporarily unavailable", + )); } Err(_) => { - return Err(LiveWebSocketPreflightRejection { - status: StatusCode::SERVICE_UNAVAILABLE, - message: "Timed out loading the Codex Live sideband binding", - }); + return Err(preflight_rejection( + context, + "sideband", + StatusCode::SERVICE_UNAVAILABLE, + "sideband_binding_lookup_timeout", + "Timed out loading the Codex Live sideband binding", + )); } }; - let lease = match tokio::time::timeout( + prepare_sideband_live_websocket(state, context, call_id, binding) + .await + .map(PreparedLiveWebSocket::Sideband) +} + +async fn prepare_direct_live_websocket( + state: &AppState, + context: &WebSocketRequestContext, + client_model: &str, +) -> Result { + let started_at = Instant::now(); + let candidate = match plan_live_candidate( + state, + context.trace_id.as_str(), + &context.decision, + &context.headers, + &context.remote_addr, + client_model, + None, + ) + .await + { + Ok(Some(candidate)) => candidate, + Ok(None) => { + return Err(preflight_rejection( + context, + "direct", + StatusCode::SERVICE_UNAVAILABLE, + "candidate_unavailable", + "No eligible Codex Live provider mapping is available", + )); + } + Err(error) => { + warn!( + target: LIVE_LOG_TARGET, + event_name = "codex_live_planning_failed", + log_type = "ops", + transport = WEBSOCKET_LOG_TRANSPORT, + websocket = true, + trace_id = %context.trace_id, + mode = "direct", + error_kind = gateway_error_kind(&error), + "Codex Live direct candidate planning failed" + ); + return Err(preflight_rejection( + context, + "direct", + gateway_error_status(&error), + "planning_failed", + "Codex Live provider planning failed", + )); + } + }; + let pool_lease = LivePoolLeaseGuard::new(state, &candidate); + let upstream_url = match direct_live_websocket_url(&candidate) { + Ok(url) => url, + Err(error) => { + pool_lease.release().await; + return Err(preflight_rejection( + context, + "direct", + error.status_code(), + error.code(), + error.client_message(), + )); + } + }; + let LiveRelayAdmission { + capacity: admission, + audit, + } = match acquire_live_relay_admission( + state, + context, + &candidate, + upstream_url.clone(), + LiveAuditTransport::DirectWebSocket, + ) + .await + { + Ok(admission) => admission, + Err(failure) => { + pool_lease.release().await; + let status = failure.error.status(); + let termination = failure.error.termination(); + return Err(audited_preflight_rejection( + state, + context, + "direct", + status, + termination, + failure.error.client_message(), + started_at, + failure.audit, + ) + .await); + } + }; + let provider_id = candidate.execution.provider_id.clone().unwrap_or_default(); + let endpoint_id = candidate.execution.endpoint_id.clone().unwrap_or_default(); + let key_id = candidate.execution.key_id.clone().unwrap_or_default(); + let provider_model = candidate.provider_model.clone(); + if !pool_lease.is_healthy() { + admission.release().await; + pool_lease.release().await; + return Err(audited_preflight_rejection( + state, + context, + "direct", + StatusCode::SERVICE_UNAVAILABLE, + "pool_key_lease_lost", + "Codex Live provider ownership was lost", + started_at, + Some(audit), + ) + .await); + } + let mut execution = candidate.execution; + execution.upstream_url = Some(upstream_url); + let mut upstream = match connect_upstream_websocket( + &execution, + LIVE_WEBSOCKET_SESSION_LIMITS, + LIVE_UPSTREAM_ERRORS, + ) + .await + { + Ok(connection) => connection.socket, + Err(error_code) => { + warn!( + target: LIVE_LOG_TARGET, + event_name = "codex_live_upstream_connect_failed", + log_type = "ops", + transport = WEBSOCKET_LOG_TRANSPORT, + websocket = true, + trace_id = %context.trace_id, + provider_id = %provider_id, + endpoint_id = %endpoint_id, + key_id = %key_id, + mode = "direct", + error_code, + "Codex Live direct upstream connection failed" + ); + admission.release().await; + pool_lease.release().await; + return Err(audited_preflight_rejection( + state, + context, + "direct", + StatusCode::BAD_GATEWAY, + "upstream_connect_failed", + "Codex Live upstream WebSocket connection failed", + started_at, + Some(audit), + ) + .await); + } + }; + if !pool_lease.is_healthy() { + close_upstream_socket(&mut upstream, None).await; + admission.release().await; + pool_lease.release().await; + return Err(audited_preflight_rejection( + state, + context, + "direct", + StatusCode::SERVICE_UNAVAILABLE, + "pool_key_lease_lost", + "Codex Live provider ownership was lost", + started_at, + Some(audit), + ) + .await); + } + Ok(PreparedLiveRelay { + upstream, + admission, + audit, + pool_lease, + provider_id, + endpoint_id, + key_id, + provider_model, + }) +} + +async fn prepare_sideband_live_websocket( + state: &AppState, + context: &WebSocketRequestContext, + call_id: String, + binding: LiveCallBinding, +) -> Result { + let started_at = Instant::now(); + let Some(auth) = context.decision.auth_context.as_ref() else { + return Err(preflight_rejection( + context, + "sideband", + StatusCode::UNAUTHORIZED, + "authentication_required", + "Authentication required", + )); + }; + let registry = LiveCallRegistry::new(std::sync::Arc::clone(&state.runtime_state)); + let mut sideband_lease = match tokio::time::timeout( SIDEBAND_LOOKUP_TIMEOUT, registry.acquire_sideband_attachment( auth.user_id.as_str(), @@ -274,30 +524,275 @@ pub(super) async fn prepare_live_websocket( key_id = %binding.pinned_candidate().key_id(), "Codex Live call already has an active sideband attachment" ); - return Err(LiveWebSocketPreflightRejection { - status: StatusCode::CONFLICT, - message: "Codex Live call already has an active sideband connection", - }); + return Err(preflight_rejection( + context, + "sideband", + StatusCode::CONFLICT, + "sideband_attachment_conflict", + "Codex Live call already has an active sideband connection", + )); } Ok(Err(error)) => { log_sideband_lease_error(context, &error, "acquire"); - return Err(LiveWebSocketPreflightRejection { - status: StatusCode::SERVICE_UNAVAILABLE, - message: "Codex Live sideband ownership is temporarily unavailable", - }); + return Err(preflight_rejection( + context, + "sideband", + StatusCode::SERVICE_UNAVAILABLE, + "sideband_attachment_unavailable", + "Codex Live sideband ownership is temporarily unavailable", + )); } Err(_) => { - return Err(LiveWebSocketPreflightRejection { - status: StatusCode::SERVICE_UNAVAILABLE, - message: "Timed out acquiring Codex Live sideband ownership", - }); + return Err(preflight_rejection( + context, + "sideband", + StatusCode::SERVICE_UNAVAILABLE, + "sideband_attachment_timeout", + "Timed out acquiring Codex Live sideband ownership", + )); } }; - Ok(PreparedLiveWebSocket::Sideband(PreparedLiveSideband { - call_id, - binding, - lease, - })) + + let planned_candidate = while_sideband_lease_healthy( + &sideband_lease, + plan_live_candidate( + state, + context.trace_id.as_str(), + &context.decision, + &context.headers, + &context.remote_addr, + binding.client_model(), + Some(binding.pinned_candidate()), + ), + ) + .await; + let candidate = match planned_candidate { + Err(loss) => { + release_sideband_lease(&mut sideband_lease, context).await; + return Err(preflight_rejection( + context, + "sideband", + StatusCode::SERVICE_UNAVAILABLE, + sideband_loss_termination(loss), + sideband_loss_message(loss), + )); + } + Ok(Ok(Some(candidate))) if binding.matches_candidate(&candidate) => candidate, + Ok(Ok(Some(candidate))) => { + crate::orchestration::release_pool_key_lease_from_report_context( + state, + candidate.execution.report_context.as_ref(), + ) + .await; + release_sideband_lease(&mut sideband_lease, context).await; + return Err(preflight_rejection( + context, + "sideband", + StatusCode::GONE, + "sideband_binding_changed", + "Codex Live call provider binding is no longer valid", + )); + } + Ok(Ok(None)) => { + release_sideband_lease(&mut sideband_lease, context).await; + return Err(preflight_rejection( + context, + "sideband", + StatusCode::GONE, + "sideband_binding_disabled", + "Codex Live call provider key or model is no longer available", + )); + } + Ok(Err(error)) => { + warn!( + target: LIVE_LOG_TARGET, + event_name = "codex_live_sideband_planning_failed", + log_type = "ops", + transport = WEBSOCKET_LOG_TRANSPORT, + websocket = true, + trace_id = %context.trace_id, + error_kind = gateway_error_kind(&error), + "Codex Live sideband pinned candidate validation failed" + ); + release_sideband_lease(&mut sideband_lease, context).await; + return Err(preflight_rejection( + context, + "sideband", + gateway_error_status(&error), + "planning_failed", + "Codex Live provider validation failed", + )); + } + }; + let pool_lease = LivePoolLeaseGuard::new(state, &candidate); + let upstream_url = match live_sideband_url(&candidate, call_id.as_str()) { + Ok(url) => url, + Err(error) => { + release_sideband_lease(&mut sideband_lease, context).await; + pool_lease.release().await; + return Err(preflight_rejection( + context, + "sideband", + error.status_code(), + error.code(), + error.client_message(), + )); + } + }; + let admission_result = while_sideband_lease_healthy( + &sideband_lease, + acquire_live_relay_admission( + state, + context, + &candidate, + upstream_url.clone(), + LiveAuditTransport::Sideband, + ), + ) + .await; + let LiveRelayAdmission { + capacity: admission, + audit, + } = match admission_result { + Err(loss) => { + release_sideband_lease(&mut sideband_lease, context).await; + pool_lease.release().await; + return Err(preflight_rejection( + context, + "sideband", + StatusCode::SERVICE_UNAVAILABLE, + sideband_loss_termination(loss), + sideband_loss_message(loss), + )); + } + Ok(Ok(admission)) => admission, + Ok(Err(failure)) => { + release_sideband_lease(&mut sideband_lease, context).await; + pool_lease.release().await; + let status = failure.error.status(); + let termination = failure.error.termination(); + return Err(audited_preflight_rejection( + state, + context, + "sideband", + status, + termination, + failure.error.client_message(), + started_at, + failure.audit, + ) + .await); + } + }; + let provider_id = candidate.execution.provider_id.clone().unwrap_or_default(); + let endpoint_id = candidate.execution.endpoint_id.clone().unwrap_or_default(); + let key_id = candidate.execution.key_id.clone().unwrap_or_default(); + let provider_model = candidate.provider_model.clone(); + if !pool_lease.is_healthy() { + admission.release().await; + release_sideband_lease(&mut sideband_lease, context).await; + pool_lease.release().await; + return Err(audited_preflight_rejection( + state, + context, + "sideband", + StatusCode::SERVICE_UNAVAILABLE, + "pool_key_lease_lost", + "Codex Live provider ownership was lost", + started_at, + Some(audit), + ) + .await); + } + let mut execution = candidate.execution; + execution.upstream_url = Some(upstream_url); + let upstream_connection = while_sideband_lease_healthy( + &sideband_lease, + connect_upstream_websocket( + &execution, + LIVE_WEBSOCKET_SESSION_LIMITS, + LIVE_UPSTREAM_ERRORS, + ), + ) + .await; + let mut upstream = match upstream_connection { + Err(loss) => { + release_sideband_lease(&mut sideband_lease, context).await; + admission.release().await; + pool_lease.release().await; + return Err(audited_preflight_rejection( + state, + context, + "sideband", + StatusCode::SERVICE_UNAVAILABLE, + sideband_loss_termination(loss), + sideband_loss_message(loss), + started_at, + Some(audit), + ) + .await); + } + Ok(Ok(connection)) => connection.socket, + Ok(Err(error_code)) => { + warn!( + target: LIVE_LOG_TARGET, + event_name = "codex_live_sideband_connect_failed", + log_type = "ops", + transport = WEBSOCKET_LOG_TRANSPORT, + websocket = true, + trace_id = %context.trace_id, + provider_id = %provider_id, + endpoint_id = %endpoint_id, + key_id = %key_id, + error_code, + "Codex Live sideband upstream connection failed" + ); + release_sideband_lease(&mut sideband_lease, context).await; + admission.release().await; + pool_lease.release().await; + return Err(audited_preflight_rejection( + state, + context, + "sideband", + StatusCode::BAD_GATEWAY, + "upstream_connect_failed", + "Codex Live sideband connection failed", + started_at, + Some(audit), + ) + .await); + } + }; + if !pool_lease.is_healthy() { + close_upstream_socket(&mut upstream, None).await; + admission.release().await; + release_sideband_lease(&mut sideband_lease, context).await; + pool_lease.release().await; + return Err(audited_preflight_rejection( + state, + context, + "sideband", + StatusCode::SERVICE_UNAVAILABLE, + "pool_key_lease_lost", + "Codex Live provider ownership was lost", + started_at, + Some(audit), + ) + .await); + } + Ok(PreparedLiveSideband { + relay: PreparedLiveRelay { + upstream, + admission, + audit, + pool_lease, + provider_id, + endpoint_id, + key_id, + provider_model, + }, + sideband_lease, + }) } pub(super) async fn run_live_websocket( @@ -309,8 +804,8 @@ pub(super) async fn run_live_websocket( let connection_log = WebSocketConnectionLog::new(&context, LIVE_CONNECTION_LOG_SPEC); connection_log.log_opened(); match prepared { - PreparedLiveWebSocket::Direct { client_model } => { - run_direct(&mut client_socket, &state, &context, client_model).await + PreparedLiveWebSocket::Direct(prepared) => { + run_direct(&mut client_socket, &state, &context, prepared).await } PreparedLiveWebSocket::Sideband(prepared) => { run_sideband(&mut client_socket, &state, &context, prepared).await @@ -322,15 +817,37 @@ async fn run_direct( client_socket: &mut WebSocket, state: &AppState, context: &WebSocketRequestContext, - client_model: String, + prepared: PreparedLiveRelay, ) { - if !live_usage_accounting_is_safe(&context.decision) { - reject_finite_balance_live(client_socket, context).await; - return; - } + let session_started_at = Instant::now(); + let PreparedLiveRelay { + mut upstream, + admission, + audit, + pool_lease, + provider_id, + endpoint_id, + key_id, + provider_model, + } = prepared; let initial = match read_initial_session_update(client_socket).await { Ok(Some(initial)) => initial, - Ok(None) => return, + Ok(None) => { + close_upstream_socket(&mut upstream, None).await; + admission.release().await; + pool_lease.release().await; + audit + .finish( + state, + live_terminal_from_relay( + "client_closed", + elapsed_ms(session_started_at), + RelayStats::default(), + ), + ) + .await; + return; + } Err(error) => { send_live_error( client_socket, @@ -349,158 +866,44 @@ async fn run_direct( "invalid initial Live event", ) .await; - return; - } - }; - let candidate = match plan_live_candidate( - state, - context.trace_id.as_str(), - &context.decision, - &context.headers, - &context.remote_addr, - client_model.as_str(), - None, - ) - .await - { - Ok(Some(candidate)) => candidate, - Ok(None) => { - send_live_error( - client_socket, - 503, - "codex_live_candidate_unavailable", - "No eligible Codex Live provider mapping is available", - ) - .await; - close_client_socket(client_socket, CLOSE_TRY_AGAIN, "Live provider unavailable").await; - return; - } - Err(error) => { - warn!( - target: LIVE_LOG_TARGET, - event_name = "codex_live_planning_failed", - log_type = "ops", - transport = WEBSOCKET_LOG_TRANSPORT, - websocket = true, - trace_id = %context.trace_id, - error_kind = gateway_error_kind(&error), - "Codex Live direct candidate planning failed" - ); - send_live_error( - client_socket, - 500, - "codex_live_planning_failed", - "Codex Live provider planning failed", - ) - .await; - close_client_socket(client_socket, CLOSE_INTERNAL_ERROR, "Live planning failed").await; - return; - } - }; - let lease = LivePoolLeaseGuard::new(state, &candidate); - let upstream_url = match direct_live_websocket_url(&candidate) { - Ok(url) => url, - Err(error) => { - lease.release().await; - send_live_error( - client_socket, - error.status_code().as_u16(), - error.code(), - error.client_message(), - ) - .await; - close_client_socket( - client_socket, - CLOSE_POLICY_VIOLATION, - "Live auth unsupported", - ) - .await; - return; - } - }; - let admission = match acquire_live_relay_admission( - state, - context, - &candidate, - upstream_url.clone(), - ) - .await - { - Ok(admission) => admission, - Err(error) => { - lease.release().await; - reject_live_relay_admission(client_socket, context, "direct", &error).await; - return; - } - }; - let provider_id = candidate.execution.provider_id.clone().unwrap_or_default(); - let endpoint_id = candidate.execution.endpoint_id.clone().unwrap_or_default(); - let key_id = candidate.execution.key_id.clone().unwrap_or_default(); - let provider_model = candidate.provider_model.clone(); - if !lease.is_healthy() { - admission.release().await; - lease.release().await; - reject_lost_pool_lease( - client_socket, - context, - "direct", - provider_id.as_str(), - endpoint_id.as_str(), - key_id.as_str(), - ) - .await; - return; - } - let mut execution = candidate.execution; - execution.upstream_url = Some(upstream_url); - let mut upstream = match connect_upstream_websocket( - &execution, - LIVE_WEBSOCKET_SESSION_LIMITS, - LIVE_UPSTREAM_ERRORS, - ) - .await - { - Ok(connection) => connection.socket, - Err(error_code) => { - warn!( - target: LIVE_LOG_TARGET, - event_name = "codex_live_upstream_connect_failed", - log_type = "ops", - transport = WEBSOCKET_LOG_TRANSPORT, - websocket = true, - trace_id = %context.trace_id, - provider_id = %provider_id, - endpoint_id = %endpoint_id, - key_id = %key_id, - error_code, - "Codex Live direct upstream connection failed" - ); + close_upstream_socket(&mut upstream, None).await; admission.release().await; - lease.release().await; - send_live_error( - client_socket, - 502, - error_code, - "Codex Live upstream WebSocket connection failed", - ) - .await; - close_client_socket(client_socket, CLOSE_TRY_AGAIN, "Live upstream unavailable").await; + pool_lease.release().await; + audit + .finish( + state, + LiveSessionTerminal::failure( + error.status_code().as_u16(), + error.code(), + elapsed_ms(session_started_at), + ), + ) + .await; return; } }; let initial = rewrite_live_session_model(initial.as_str(), provider_model.as_str()).unwrap_or(initial); + let initial_bytes = initial.len() as u64; if send_upstream_message(&mut upstream, WreqWsMessage::Text(initial.into())) .await .is_err() { close_upstream_socket(&mut upstream, None).await; admission.release().await; - lease.release().await; + pool_lease.release().await; close_client_socket(client_socket, CLOSE_TRY_AGAIN, "Live upstream write failed").await; + let mut terminal = LiveSessionTerminal::failure( + 502, + "initial_upstream_write_failed", + elapsed_ms(session_started_at), + ); + terminal.client_frames = 1; + terminal.client_bytes = initial_bytes; + audit.finish(state, terminal).await; return; } - relay_live( + let terminal = relay_live( client_socket, &mut upstream, context, @@ -509,13 +912,27 @@ async fn run_direct( endpoint_id.as_str(), key_id.as_str(), provider_model.as_str(), - &lease, + &pool_lease, None, + RelayStats { + client_frames: 1, + client_bytes: initial_bytes, + ..RelayStats::default() + }, + session_started_at, ) .await; + let close_client = matches!( + terminal.termination, + "connection_duration_limit" | "connection_admission_lost" + ); close_upstream_socket(&mut upstream, None).await; admission.release().await; - lease.release().await; + pool_lease.release().await; + if close_client { + close_client_socket(client_socket, CLOSE_TRY_AGAIN, terminal.termination).await; + } + audit.finish(state, terminal).await; } async fn run_sideband( @@ -525,212 +942,23 @@ async fn run_sideband( prepared: PreparedLiveSideband, ) { let PreparedLiveSideband { - call_id, - binding, - lease: mut sideband_lease, + relay: + PreparedLiveRelay { + mut upstream, + admission, + audit, + pool_lease, + provider_id, + endpoint_id, + key_id, + provider_model, + }, + mut sideband_lease, } = prepared; - let planned_candidate = match while_sideband_lease_healthy( - &sideband_lease, - plan_live_candidate( - state, - context.trace_id.as_str(), - &context.decision, - &context.headers, - &context.remote_addr, - binding.client_model(), - Some(binding.pinned_candidate()), - ), - ) - .await - { - Ok(result) => result, - Err(loss) => { - reject_sideband_lease_loss(client_socket, context, loss).await; - release_sideband_lease(&mut sideband_lease, context).await; - return; - } - }; - let candidate = match planned_candidate { - Ok(Some(candidate)) if binding.matches_candidate(&candidate) => candidate, - Ok(Some(candidate)) => { - crate::orchestration::release_pool_key_lease_from_report_context( - state, - candidate.execution.report_context.as_ref(), - ) - .await; - release_sideband_lease(&mut sideband_lease, context).await; - send_live_error( - client_socket, - 410, - "codex_live_binding_changed", - "Codex Live call provider binding is no longer valid", - ) - .await; - close_client_socket( - client_socket, - CLOSE_POLICY_VIOLATION, - "Live binding changed", - ) - .await; - return; - } - Ok(None) => { - release_sideband_lease(&mut sideband_lease, context).await; - send_live_error( - client_socket, - 410, - "codex_live_binding_disabled", - "Codex Live call provider key or model is no longer available", - ) - .await; - close_client_socket( - client_socket, - CLOSE_POLICY_VIOLATION, - "Live binding disabled", - ) - .await; - return; - } - Err(error) => { - release_sideband_lease(&mut sideband_lease, context).await; - warn!( - target: LIVE_LOG_TARGET, - event_name = "codex_live_sideband_planning_failed", - log_type = "ops", - transport = WEBSOCKET_LOG_TRANSPORT, - websocket = true, - trace_id = %context.trace_id, - error_kind = gateway_error_kind(&error), - "Codex Live sideband pinned candidate validation failed" - ); - send_live_error( - client_socket, - 500, - "codex_live_planning_failed", - "Codex Live provider validation failed", - ) - .await; - close_client_socket(client_socket, CLOSE_INTERNAL_ERROR, "Live planning failed").await; - return; - } - }; - let pool_lease = LivePoolLeaseGuard::new(state, &candidate); - let upstream_url = match live_sideband_url(&candidate, call_id.as_str()) { - Ok(url) => url, - Err(error) => { - release_sideband_lease(&mut sideband_lease, context).await; - pool_lease.release().await; - send_live_error( - client_socket, - error.status_code().as_u16(), - error.code(), - error.client_message(), - ) - .await; - close_client_socket( - client_socket, - CLOSE_POLICY_VIOLATION, - "invalid Live sideband", - ) - .await; - return; - } - }; - let admission = match while_sideband_lease_healthy( - &sideband_lease, - acquire_live_relay_admission(state, context, &candidate, upstream_url.clone()), - ) - .await - { - Err(loss) => { - reject_sideband_lease_loss(client_socket, context, loss).await; - release_sideband_lease(&mut sideband_lease, context).await; - pool_lease.release().await; - return; - } - Ok(result) => match result { - Ok(admission) => admission, - Err(error) => { - release_sideband_lease(&mut sideband_lease, context).await; - pool_lease.release().await; - reject_live_relay_admission(client_socket, context, "sideband", &error).await; - return; - } - }, - }; - let provider_id = candidate.execution.provider_id.clone().unwrap_or_default(); - let endpoint_id = candidate.execution.endpoint_id.clone().unwrap_or_default(); - let key_id = candidate.execution.key_id.clone().unwrap_or_default(); - let provider_model = candidate.provider_model.clone(); - if !pool_lease.is_healthy() { - admission.release().await; - release_sideband_lease(&mut sideband_lease, context).await; - pool_lease.release().await; - reject_lost_pool_lease( - client_socket, - context, - "sideband", - provider_id.as_str(), - endpoint_id.as_str(), - key_id.as_str(), - ) - .await; - return; - } - let mut execution = candidate.execution; - execution.upstream_url = Some(upstream_url); - let upstream_connection = match while_sideband_lease_healthy( - &sideband_lease, - connect_upstream_websocket( - &execution, - LIVE_WEBSOCKET_SESSION_LIMITS, - LIVE_UPSTREAM_ERRORS, - ), - ) - .await - { - Ok(result) => result, - Err(loss) => { - reject_sideband_lease_loss(client_socket, context, loss).await; - release_sideband_lease(&mut sideband_lease, context).await; - admission.release().await; - pool_lease.release().await; - return; - } - }; - let mut upstream = match upstream_connection { - Ok(connection) => connection.socket, - Err(error_code) => { - warn!( - target: LIVE_LOG_TARGET, - event_name = "codex_live_sideband_connect_failed", - log_type = "ops", - transport = WEBSOCKET_LOG_TRANSPORT, - websocket = true, - trace_id = %context.trace_id, - provider_id = %provider_id, - endpoint_id = %endpoint_id, - key_id = %key_id, - error_code, - "Codex Live sideband upstream connection failed" - ); - release_sideband_lease(&mut sideband_lease, context).await; - admission.release().await; - pool_lease.release().await; - send_live_error( - client_socket, - 502, - error_code, - "Codex Live sideband connection failed", - ) - .await; - close_client_socket(client_socket, CLOSE_TRY_AGAIN, "Live sideband unavailable").await; - return; - } - }; // A sideband attaches to an already-created WebRTC session. Sending a // second synthetic `session.update` here would corrupt the protocol. - relay_live( + let session_started_at = Instant::now(); + let terminal = relay_live( client_socket, &mut upstream, context, @@ -741,12 +969,22 @@ async fn run_sideband( provider_model.as_str(), &pool_lease, Some(&sideband_lease), + RelayStats::default(), + session_started_at, ) .await; + let close_client = matches!( + terminal.termination, + "connection_duration_limit" | "connection_admission_lost" + ); close_upstream_socket(&mut upstream, None).await; release_sideband_lease(&mut sideband_lease, context).await; admission.release().await; pool_lease.release().await; + if close_client { + close_client_socket(client_socket, CLOSE_TRY_AGAIN, terminal.termination).await; + } + audit.finish(state, terminal).await; } async fn acquire_live_relay_admission( @@ -754,83 +992,103 @@ async fn acquire_live_relay_admission( context: &WebSocketRequestContext, candidate: &PlannedLiveCandidate, upstream_url: String, -) -> Result { - let Some(attempt) = build_live_stream_admission_attempt( + transport: LiveAuditTransport, +) -> Result { + let attempt = match build_live_stream_admission_attempt( candidate, &context.headers, &context.remote_addr, upstream_url, - ) - .map_err(LiveRelayAdmissionError::Gateway)? - else { - return Err(LiveRelayAdmissionError::PlanUnavailable); + ) { + Ok(Some(attempt)) => attempt, + Ok(None) => { + return Err(LiveRelayAdmissionFailure { + error: LiveRelayAdmissionError::PlanUnavailable, + audit: None, + }); + } + Err(error) => { + return Err(LiveRelayAdmissionFailure { + error: LiveRelayAdmissionError::Gateway(error), + audit: None, + }); + } }; - if execution_plan_balance_capacity_rejection( + let audit = LiveSessionAudit::from_attempt(&attempt, transport); + let balance_rejection = execution_plan_balance_capacity_rejection( state, &context.decision, &attempt.plan, attempt.report_context.as_ref(), ) - .await - .map_err(LiveRelayAdmissionError::Gateway)? - .is_some() - { - return Err(LiveRelayAdmissionError::BalanceRejected); + .await; + match balance_rejection { + Ok(None) => {} + Ok(Some(_)) => { + return Err(LiveRelayAdmissionFailure { + error: LiveRelayAdmissionError::BalanceRejected, + audit: Some(audit), + }); + } + Err(error) => { + return Err(LiveRelayAdmissionFailure { + error: LiveRelayAdmissionError::Gateway(error), + audit: Some(audit), + }); + } } - ResponsesWebSocketTurnAdmission::acquire(state, &attempt.plan, context.trace_id.as_str()) + match ResponsesWebSocketTurnAdmission::acquire(state, &attempt.plan, context.trace_id.as_str()) .await - .map_err(LiveRelayAdmissionError::Gateway) + { + Ok(capacity) => Ok(LiveRelayAdmission { capacity, audit }), + Err(error) => Err(LiveRelayAdmissionFailure { + error: LiveRelayAdmissionError::Gateway(error), + audit: Some(audit), + }), + } } -async fn reject_finite_balance_live( - client_socket: &mut WebSocket, - context: &WebSocketRequestContext, -) { - warn!( - target: LIVE_LOG_TARGET, - event_name = "codex_live_usage_accounting_unsafe", - log_type = "ops", - transport = WEBSOCKET_LOG_TRANSPORT, - websocket = true, - trace_id = %context.trace_id, - "Codex Live rejected a finite-balance principal because Frameless usage is unavailable" - ); - send_live_error( - client_socket, - 501, - "codex_live_usage_settlement_unavailable", - "Codex Live is unavailable for finite-balance keys until Frameless usage settlement is supported", - ) - .await; - close_client_socket( - client_socket, - CLOSE_POLICY_VIOLATION, - "Live usage settlement unavailable", - ) - .await; -} - -async fn reject_live_relay_admission( - client_socket: &mut WebSocket, +async fn audited_preflight_rejection( + state: &AppState, context: &WebSocketRequestContext, mode: &'static str, - error: &LiveRelayAdmissionError, -) { - let (status, code, message, close_code, close_reason) = error.response(); + status: StatusCode, + termination: &'static str, + message: &'static str, + started_at: Instant, + audit: Option, +) -> LiveWebSocketPreflightRejection { + if let Some(audit) = audit { + audit + .finish( + state, + LiveSessionTerminal::failure(status.as_u16(), termination, elapsed_ms(started_at)), + ) + .await; + } + preflight_rejection(context, mode, status, termination, message) +} + +fn preflight_rejection( + context: &WebSocketRequestContext, + mode: &'static str, + status: StatusCode, + termination: &'static str, + message: &'static str, +) -> LiveWebSocketPreflightRejection { warn!( target: LIVE_LOG_TARGET, - event_name = "codex_live_relay_admission_failed", - log_type = "ops", + event_name = "codex_live_websocket_preflight_rejected", + log_type = "event", transport = WEBSOCKET_LOG_TRANSPORT, websocket = true, trace_id = %context.trace_id, mode, - status, - error_kind = error.kind(), - "Codex Live relay admission failed" + status_code = status.as_u16(), + termination, + "Codex Live WebSocket preflight rejected the HTTP upgrade" ); - send_live_error(client_socket, status, code, message).await; - close_client_socket(client_socket, close_code, close_reason).await; + LiveWebSocketPreflightRejection { status, message } } fn gateway_error_kind(error: &GatewayError) -> &'static str { @@ -844,6 +1102,38 @@ fn gateway_error_kind(error: &GatewayError) -> &'static str { } } +fn gateway_error_status(error: &GatewayError) -> StatusCode { + match error { + GatewayError::UpstreamUnavailable { .. } | GatewayError::ControlUnavailable { .. } => { + StatusCode::SERVICE_UNAVAILABLE + } + GatewayError::LocalExecutionPlanningTimeout { .. } => StatusCode::GATEWAY_TIMEOUT, + GatewayError::AdmissionTimeout { .. } => StatusCode::TOO_MANY_REQUESTS, + GatewayError::Client { status, .. } => *status, + GatewayError::Internal(_) => StatusCode::INTERNAL_SERVER_ERROR, + } +} + +const fn sideband_loss_termination(loss: LiveSidebandLeaseLoss) -> &'static str { + match loss { + LiveSidebandLeaseLoss::OwnershipLost => "sideband_attachment_lease_lost", + LiveSidebandLeaseLoss::StorageUnavailable => "sideband_attachment_lease_renewal_failed", + } +} + +const fn sideband_loss_message(loss: LiveSidebandLeaseLoss) -> &'static str { + match loss { + LiveSidebandLeaseLoss::OwnershipLost => "Codex Live sideband ownership was lost", + LiveSidebandLeaseLoss::StorageUnavailable => { + "Codex Live sideband ownership could not be renewed" + } + } +} + +fn elapsed_ms(started_at: Instant) -> u64 { + started_at.elapsed().as_millis().min(u128::from(u64::MAX)) as u64 +} + async fn reject_lost_pool_lease( client_socket: &mut WebSocket, context: &WebSocketRequestContext, @@ -917,12 +1207,13 @@ async fn read_initial_session_update( .map_err(|_| super::protocol::LiveProtocolError::InitialSessionUpdateTimeout)? } -#[derive(Default)] +#[derive(Clone, Copy, Default)] struct RelayStats { client_frames: u64, client_bytes: u64, upstream_frames: u64, upstream_bytes: u64, + first_upstream_frame_ms: Option, } async fn relay_live( @@ -936,83 +1227,272 @@ async fn relay_live( provider_model: &str, pool_lease: &LivePoolLeaseGuard, sideband_lease: Option<&LiveSidebandLease>, -) { - let started_at = Instant::now(); - let connection_deadline = - tokio::time::sleep(LIVE_WEBSOCKET_SESSION_LIMITS.max_connection_duration); + stats: RelayStats, + started_at: Instant, +) -> LiveSessionTerminal { + let connection_deadline = tokio::time::sleep_until(tokio::time::Instant::from_std( + started_at + LIVE_WEBSOCKET_SESSION_LIMITS.max_connection_duration, + )); tokio::pin!(connection_deadline); - let mut close_deadline = None; let mut pool_lease_health = tokio::time::interval(Duration::from_secs(1)); pool_lease_health.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Delay); - let mut stats = RelayStats::default(); - let termination = loop { - tokio::select! { - _ = &mut connection_deadline => break "connection_duration_limit", - _ = wait_for_connection_permit_loss(context.websocket_connection_permit.as_ref()) => { - break "connection_admission_lost"; - } - _ = pool_lease_health.tick() => { - if !pool_lease.is_healthy() { - reject_lost_pool_lease( - client_socket, - context, - mode, - provider_id, - endpoint_id, - key_id, - ) - .await; - break "pool_key_lease_lost"; - } - } - _ = wait_for_optional_deadline(close_deadline) => break "session_close_drain_timeout", - loss = wait_for_sideband_lease_loss(sideband_lease) => { - reject_sideband_lease_loss(client_socket, context, loss).await; - break match loss { - LiveSidebandLeaseLoss::OwnershipLost => "sideband_attachment_lease_lost", - LiveSidebandLeaseLoss::StorageUnavailable => { - "sideband_attachment_lease_renewal_failed" + let stats = Arc::new(Mutex::new(stats)); + let session_close = Arc::new(tokio::sync::Notify::new()); + let relay_control = WebSocketRelayPumpControl::new(); + + let termination = { + let (mut client_write, mut client_read) = (&mut *client_socket).split(); + let (mut upstream_write, mut upstream_read) = (&mut *upstream).split(); + + let client_to_upstream = { + let control = relay_control.clone(); + let stats = Arc::clone(&stats); + let session_close = Arc::clone(&session_close); + async move { + let (queue_tx, mut queue_rx) = websocket_relay_frame_queue(); + let reader_control = control.clone(); + let reader = async move { + loop { + let client = tokio::select! { + biased; + _ = reader_control.cancelled() => return "relay_cancelled", + client = client_read.next() => client, + }; + let Some(client) = client else { + return "client_closed"; + }; + let Ok(client) = client else { + return "client_read_failed"; + }; + let (bytes, is_close, is_session_close) = client_frame_metadata(&client); + { + let mut stats = stats + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()); + stats.client_frames = stats.client_frames.saturating_add(1); + stats.client_bytes = stats.client_bytes.saturating_add(bytes as u64); + } + let client = match client { + AxumWsMessage::Text(text) => { + rewrite_live_session_model(text.as_str(), provider_model) + .map_or(AxumWsMessage::Text(text), |rewritten| { + AxumWsMessage::Text(rewritten.into()) + }) + } + other => other, + }; + match reader_control + .enqueue(&queue_tx, client_message_to_upstream(client)) + .await + { + Ok(()) => {} + Err(WebSocketRelayQueueError::Cancelled) => { + return "relay_cancelled"; + } + Err(WebSocketRelayQueueError::Closed) => { + return "upstream_write_failed"; + } + } + if is_session_close { + session_close.notify_one(); + } + if is_close { + return "client_close_frame"; + } } }; - } - client = client_socket.next() => { - let Some(client) = client else { break "client_closed"; }; - let Ok(client) = client else { break "client_read_failed"; }; - let (bytes, is_close, is_session_close) = client_frame_metadata(&client); - stats.client_frames = stats.client_frames.saturating_add(1); - stats.client_bytes = stats.client_bytes.saturating_add(bytes as u64); - let client = match client { - AxumWsMessage::Text(text) => rewrite_live_session_model( - text.as_str(), - provider_model, - ) - .map_or(AxumWsMessage::Text(text), |rewritten| { - AxumWsMessage::Text(rewritten.into()) - }), - other => other, + + let writer_control = control; + let writer = async move { + loop { + let message = tokio::select! { + biased; + _ = writer_control.cancelled() => return None, + message = queue_rx.recv() => message, + }; + let Some(message) = message else { + return None; + }; + let result = writer_control + .send(async { upstream_write.send(message).await.map_err(|_| ()) }) + .await; + match result { + Ok(()) => {} + Err(WebSocketWriteError::Cancelled) => return None, + Err(_) => return Some("upstream_write_failed"), + } + } }; - let upstream_message = client_message_to_upstream(client); - if send_upstream_message(upstream, upstream_message).await.is_err() { - break "upstream_write_failed"; + + tokio::pin!(reader, writer); + tokio::select! { + reader_exit = &mut reader => { + writer.await.unwrap_or(reader_exit) + } + writer_exit = &mut writer => { + match writer_exit { + Some(writer_exit) => writer_exit, + None => reader.await, + } + } } - if is_session_close { + } + }; + + let upstream_to_client = { + let control = relay_control.clone(); + let stats = Arc::clone(&stats); + async move { + let (queue_tx, mut queue_rx) = websocket_relay_frame_queue(); + let reader_control = control.clone(); + let reader = async move { + loop { + let provider = tokio::select! { + biased; + _ = reader_control.cancelled() => return "relay_cancelled", + provider = upstream_read.next() => provider, + }; + let Some(provider) = provider else { + return "upstream_closed"; + }; + let Ok(provider) = provider else { + return "upstream_read_failed"; + }; + let (bytes, is_close) = upstream_frame_metadata(&provider); + { + let mut stats = stats + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()); + stats.first_upstream_frame_ms.get_or_insert_with(|| { + started_at.elapsed().as_millis().min(u128::from(u64::MAX)) as u64 + }); + stats.upstream_frames = stats.upstream_frames.saturating_add(1); + stats.upstream_bytes = + stats.upstream_bytes.saturating_add(bytes as u64); + } + match reader_control + .enqueue(&queue_tx, upstream_message_to_client(provider)) + .await + { + Ok(()) => {} + Err(WebSocketRelayQueueError::Cancelled) => { + return "relay_cancelled"; + } + Err(WebSocketRelayQueueError::Closed) => { + return "client_write_failed"; + } + } + if is_close { + return "upstream_close_frame"; + } + } + }; + + let writer_control = control; + let writer = async move { + loop { + let message = tokio::select! { + biased; + _ = writer_control.cancelled() => return None, + message = queue_rx.recv() => message, + }; + let Some(message) = message else { + return None; + }; + let result = writer_control + .send(async { client_write.send(message).await.map_err(|_| ()) }) + .await; + match result { + Ok(()) => {} + Err(WebSocketWriteError::Cancelled) => return None, + Err(_) => return Some("client_write_failed"), + } + } + }; + + tokio::pin!(reader, writer); + tokio::select! { + reader_exit = &mut reader => { + writer.await.unwrap_or(reader_exit) + } + writer_exit = &mut writer => { + match writer_exit { + Some(writer_exit) => writer_exit, + None => reader.await, + } + } + } + } + }; + + tokio::pin!(client_to_upstream, upstream_to_client); + let mut close_deadline = None; + let termination = loop { + tokio::select! { + termination = &mut client_to_upstream => break termination, + termination = &mut upstream_to_client => break termination, + _ = &mut connection_deadline => break "connection_duration_limit", + _ = wait_for_connection_permit_loss(context.websocket_connection_permit.as_ref()) => { + break "connection_admission_lost"; + } + _ = pool_lease_health.tick() => { + if !pool_lease.is_healthy() { + break "pool_key_lease_lost"; + } + } + _ = session_close.notified(), if close_deadline.is_none() => { close_deadline = Some(Instant::now() + SESSION_CLOSE_DRAIN_TIMEOUT); } - if is_close { break "client_close_frame"; } - } - provider = upstream.next() => { - let Some(provider) = provider else { break "upstream_closed"; }; - let Ok(provider) = provider else { break "upstream_read_failed"; }; - let (bytes, is_close) = upstream_frame_metadata(&provider); - stats.upstream_frames = stats.upstream_frames.saturating_add(1); - stats.upstream_bytes = stats.upstream_bytes.saturating_add(bytes as u64); - if send_client_message(client_socket, upstream_message_to_client(provider)).await.is_err() { - break "client_write_failed"; + _ = wait_for_optional_deadline(close_deadline) => { + break "session_close_drain_timeout"; + } + loss = wait_for_sideband_lease_loss(sideband_lease) => { + break match loss { + LiveSidebandLeaseLoss::OwnershipLost => "sideband_attachment_lease_lost", + LiveSidebandLeaseLoss::StorageUnavailable => { + "sideband_attachment_lease_renewal_failed" + } + }; } - if is_close { break "upstream_close_frame"; } } - } + }; + relay_control.cancel(); + termination }; + match termination { + "pool_key_lease_lost" => { + reject_lost_pool_lease( + client_socket, + context, + mode, + provider_id, + endpoint_id, + key_id, + ) + .await; + } + "sideband_attachment_lease_lost" => { + reject_sideband_lease_loss( + client_socket, + context, + LiveSidebandLeaseLoss::OwnershipLost, + ) + .await; + } + "sideband_attachment_lease_renewal_failed" => { + reject_sideband_lease_loss( + client_socket, + context, + LiveSidebandLeaseLoss::StorageUnavailable, + ) + .await; + } + _ => {} + } + let stats = *stats + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()); + let elapsed_ms = started_at.elapsed().as_millis().min(u128::from(u64::MAX)) as u64; info!( target: LIVE_LOG_TARGET, event_name = "codex_live_relay_finished", @@ -1029,10 +1509,46 @@ async fn relay_live( client_bytes = stats.client_bytes, upstream_frames = stats.upstream_frames, upstream_bytes = stats.upstream_bytes, - elapsed_ms = started_at.elapsed().as_millis() as u64, + elapsed_ms, usage_unavailable = true, "Codex Live opaque relay finished" ); + live_terminal_from_relay(termination, elapsed_ms, stats) +} + +fn live_terminal_from_relay( + termination: &'static str, + elapsed_ms: u64, + stats: RelayStats, +) -> LiveSessionTerminal { + let (disposition, status_code) = match termination { + "client_close_frame" | "upstream_close_frame" | "session_close_drain_timeout" => { + (LiveSessionDisposition::Completed, 200) + } + "client_closed" + | "client_read_failed" + | "client_write_failed" + | "connection_duration_limit" => (LiveSessionDisposition::Cancelled, 499), + "pool_key_lease_lost" + | "connection_admission_lost" + | "sideband_attachment_lease_lost" + | "sideband_attachment_lease_renewal_failed" => (LiveSessionDisposition::Failed, 503), + "upstream_closed" | "upstream_read_failed" | "upstream_write_failed" => { + (LiveSessionDisposition::Failed, 502) + } + _ => (LiveSessionDisposition::Failed, 500), + }; + LiveSessionTerminal { + disposition, + status_code, + termination, + elapsed_ms, + first_upstream_frame_ms: stats.first_upstream_frame_ms, + client_frames: stats.client_frames, + client_bytes: stats.client_bytes, + upstream_frames: stats.upstream_frames, + upstream_bytes: stats.upstream_bytes, + } } async fn while_sideband_lease_healthy( diff --git a/apps/aether-gateway/src/handlers/proxy/websocket/mod.rs b/apps/aether-gateway/src/handlers/proxy/websocket/mod.rs index e12c4bc18..e3a1e3060 100644 --- a/apps/aether-gateway/src/handlers/proxy/websocket/mod.rs +++ b/apps/aether-gateway/src/handlers/proxy/websocket/mod.rs @@ -8,6 +8,7 @@ pub(crate) mod ingress; pub(crate) mod live; +pub(crate) mod realtime; pub(crate) mod responses; pub(crate) mod session; pub(crate) mod transport; diff --git a/apps/aether-gateway/src/handlers/proxy/websocket/realtime/audit.rs b/apps/aether-gateway/src/handlers/proxy/websocket/realtime/audit.rs new file mode 100644 index 000000000..0de80776e --- /dev/null +++ b/apps/aether-gateway/src/handlers/proxy/websocket/realtime/audit.rs @@ -0,0 +1,430 @@ +//! One terminal usage/audit row per OpenAI Realtime WebSocket connection. +//! +//! Realtime exposes authoritative token usage on `response.done`. We preserve +//! those counters when present. A connection that closes without any such +//! usage is still visible as a lifecycle row, but is explicitly marked +//! unavailable and cannot participate in billing or balance materialization. + +use std::time::Duration; + +use aether_contracts::ExecutionPlan; +use aether_data_contracts::repository::usage::{ + REALTIME_SESSION_METADATA_KEY, USAGE_AVAILABLE_METADATA_KEY, + USAGE_PRICING_AVAILABLE_METADATA_KEY, WEBSOCKET_MODE_METADATA_KEY, + WEBSOCKET_TRANSPORT_METADATA_KEY, +}; +use aether_usage_runtime::build_usage_event_data_seed; +use serde_json::{json, Map, Value}; +use tracing::warn; + +use crate::usage::{UsageEvent, UsageEventType}; +use crate::AppState; + +use super::protocol::RealtimeUsageTotals; + +const REALTIME_AUDIT_WRITE_WAIT: Duration = Duration::from_secs(5); +const REALTIME_AUDIT_SCHEMA_VERSION: &str = "1"; +const REALTIME_AUDIT_LOG_TARGET: &str = "aether_gateway::handlers::proxy::realtime_ws"; +const REALTIME_WEBSOCKET_TRANSPORT: &str = "openai_realtime"; + +pub(super) struct RealtimeSessionAudit { + plan: ExecutionPlan, + report_context: Option, +} + +impl RealtimeSessionAudit { + pub(super) fn new(plan: &ExecutionPlan, report_context: Option<&Value>) -> Self { + Self { + plan: plan.clone(), + report_context: report_context.cloned(), + } + } + + /// Persist exactly one terminal row. If the bounded caller wait expires, + /// the spawned write remains alive instead of losing the only session row. + pub(super) async fn finish(self, state: &AppState, terminal: RealtimeSessionTerminal) { + let request_id = self.plan.request_id.clone(); + let event = self.build_terminal_event(terminal); + let usage_runtime = std::sync::Arc::clone(&state.usage_runtime); + let usage_data = std::sync::Arc::clone(state.usage_lifecycle_data_state()); + let task = tokio::spawn(async move { + usage_runtime + .record_terminal_event_direct(usage_data.as_ref(), event) + .await; + }); + match tokio::time::timeout(REALTIME_AUDIT_WRITE_WAIT, task).await { + Ok(Ok(())) => {} + Ok(Err(error)) => warn!( + target: REALTIME_AUDIT_LOG_TARGET, + event_name = "openai_realtime_session_audit_task_failed", + log_type = "ops", + request_id, + error = %error, + "OpenAI Realtime session audit task failed" + ), + Err(_) => warn!( + target: REALTIME_AUDIT_LOG_TARGET, + event_name = "openai_realtime_session_audit_write_slow", + log_type = "ops", + request_id, + wait_ms = REALTIME_AUDIT_WRITE_WAIT.as_millis() as u64, + write_detached = true, + "OpenAI Realtime stopped waiting for a slow session audit write" + ), + } + } + + fn build_terminal_event(self, terminal: RealtimeSessionTerminal) -> UsageEvent { + let usage_available = terminal.usage.responses > 0; + let pricing_available = usage_available + && terminal.usage.input_audio_tokens == 0 + && terminal.usage.output_audio_tokens == 0; + let mut data = build_usage_event_data_seed(&self.plan, self.report_context.as_ref()); + data.request_type = Some("realtime".to_string()); + data.is_stream = Some(true); + data.status_code = Some(terminal.status_code); + data.response_time_ms = Some(terminal.elapsed_ms); + data.first_byte_time_ms = terminal.first_upstream_frame_ms; + if usage_available { + data.input_tokens = Some(terminal.usage.input_tokens); + data.output_tokens = Some(terminal.usage.output_tokens); + data.total_tokens = Some(terminal.usage.total_tokens); + data.cache_creation_input_tokens = None; + data.cache_creation_ephemeral_5m_input_tokens = None; + data.cache_creation_ephemeral_1h_input_tokens = None; + data.cache_read_input_tokens = Some(terminal.usage.cached_input_tokens); + } else { + clear_usage_and_cost(&mut data); + } + data.cache_creation_cost_usd = None; + data.cache_read_cost_usd = None; + data.total_cost_usd = None; + data.actual_total_cost_usd = None; + if terminal.disposition != RealtimeSessionDisposition::Completed { + data.error_message = Some(terminal.termination.to_string()); + data.error_category = Some(terminal.disposition.error_category().to_string()); + } + data.request_metadata = attach_terminal_metadata( + data.request_metadata, + usage_available, + pricing_available, + &terminal, + ); + UsageEvent::new( + terminal.disposition.event_type(), + self.plan.request_id, + data, + ) + } +} + +fn clear_usage_and_cost(data: &mut crate::usage::UsageEventData) { + data.input_tokens = None; + data.output_tokens = None; + data.total_tokens = None; + data.cache_creation_input_tokens = None; + data.cache_creation_ephemeral_5m_input_tokens = None; + data.cache_creation_ephemeral_1h_input_tokens = None; + data.cache_read_input_tokens = None; + data.cache_creation_cost_usd = None; + data.cache_read_cost_usd = None; + data.total_cost_usd = None; + data.actual_total_cost_usd = None; +} + +fn attach_terminal_metadata( + metadata: Option, + usage_available: bool, + pricing_available: bool, + terminal: &RealtimeSessionTerminal, +) -> Option { + let mut object = match metadata { + Some(Value::Object(object)) => object, + _ => Map::new(), + }; + object.insert( + USAGE_AVAILABLE_METADATA_KEY.to_string(), + Value::Bool(usage_available), + ); + object.insert( + USAGE_PRICING_AVAILABLE_METADATA_KEY.to_string(), + Value::Bool(pricing_available), + ); + object.insert(WEBSOCKET_MODE_METADATA_KEY.to_string(), Value::Bool(true)); + object.insert( + WEBSOCKET_TRANSPORT_METADATA_KEY.to_string(), + Value::String(REALTIME_WEBSOCKET_TRANSPORT.to_string()), + ); + object.insert( + REALTIME_SESSION_METADATA_KEY.to_string(), + json!({ + "schema_version": REALTIME_AUDIT_SCHEMA_VERSION, + "transport": "websocket", + "state": terminal.disposition.state(), + "termination": terminal.termination, + "elapsed_ms": terminal.elapsed_ms, + "client_frames": terminal.client_frames, + "client_bytes": terminal.client_bytes, + "upstream_frames": terminal.upstream_frames, + "upstream_bytes": terminal.upstream_bytes, + "first_upstream_frame_ms": terminal.first_upstream_frame_ms, + "usage_state": if usage_available { "authoritative" } else { "unavailable" }, + "pricing_state": if !usage_available { + "usage_unavailable" + } else if pricing_available { + "compatible_text_usage" + } else { + "unsupported_audio_breakdown" + }, + "usage_scope": "response_done", + "input_transcription_usage_included": false, + "usage_response_count": terminal.usage.responses, + "cached_input_tokens": terminal.usage.cached_input_tokens, + "input_audio_tokens": terminal.usage.input_audio_tokens, + "output_audio_tokens": terminal.usage.output_audio_tokens, + }), + ); + Some(Value::Object(object)) +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(super) enum RealtimeSessionDisposition { + Completed, + Failed, + Cancelled, +} + +impl RealtimeSessionDisposition { + const fn event_type(self) -> UsageEventType { + match self { + Self::Completed => UsageEventType::Completed, + Self::Failed => UsageEventType::Failed, + Self::Cancelled => UsageEventType::Cancelled, + } + } + + const fn state(self) -> &'static str { + match self { + Self::Completed => "closed", + Self::Failed => "failed", + Self::Cancelled => "cancelled", + } + } + + const fn error_category(self) -> &'static str { + match self { + Self::Completed => "none", + Self::Failed => "transport_error", + Self::Cancelled => "client_cancelled", + } + } +} + +#[derive(Debug, Clone, Copy)] +pub(super) struct RealtimeSessionTerminal { + pub(super) disposition: RealtimeSessionDisposition, + pub(super) status_code: u16, + pub(super) termination: &'static str, + pub(super) elapsed_ms: u64, + pub(super) first_upstream_frame_ms: Option, + pub(super) client_frames: u64, + pub(super) client_bytes: u64, + pub(super) upstream_frames: u64, + pub(super) upstream_bytes: u64, + pub(super) usage: RealtimeUsageTotals, +} + +#[cfg(test)] +mod tests { + use std::collections::BTreeMap; + + use aether_contracts::{ExecutionTimeouts, RequestBody}; + + use super::*; + + fn sample_plan() -> ExecutionPlan { + ExecutionPlan { + request_id: "realtime-request".to_string(), + candidate_id: Some("candidate-realtime".to_string()), + provider_name: Some("OpenAI".to_string()), + provider_id: "provider-realtime".to_string(), + endpoint_id: "endpoint-realtime".to_string(), + key_id: "key-realtime".to_string(), + method: "GET".to_string(), + url: "wss://example.test/v1/realtime?model=gpt-realtime".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:realtime".to_string(), + provider_api_format: "openai:realtime".to_string(), + model_name: Some("gpt-realtime".to_string()), + proxy: None, + transport_profile: None, + timeouts: Some(ExecutionTimeouts::default()), + } + } + + fn terminal(usage: RealtimeUsageTotals) -> RealtimeSessionTerminal { + RealtimeSessionTerminal { + disposition: RealtimeSessionDisposition::Completed, + status_code: 200, + termination: "client_close_frame", + elapsed_ms: 1500, + first_upstream_frame_ms: Some(30), + client_frames: 4, + client_bytes: 600, + upstream_frames: 8, + upstream_bytes: 1200, + usage, + } + } + + #[test] + fn response_done_usage_becomes_authoritative_session_usage() { + let event = RealtimeSessionAudit::new( + &sample_plan(), + Some(&json!({"user_id": "user-1", "api_key_id": "api-key-1"})), + ) + .build_terminal_event(terminal(RealtimeUsageTotals { + responses: 2, + input_tokens: 120, + output_tokens: 40, + total_tokens: 160, + cached_input_tokens: 30, + input_audio_tokens: 20, + output_audio_tokens: 10, + })); + + assert_eq!(event.event_type, UsageEventType::Completed); + assert_eq!(event.data.input_tokens, Some(120)); + assert_eq!(event.data.output_tokens, Some(40)); + assert_eq!(event.data.total_tokens, Some(160)); + assert_eq!(event.data.cache_read_input_tokens, Some(30)); + let metadata = event.data.request_metadata.expect("metadata"); + assert_eq!(metadata[USAGE_AVAILABLE_METADATA_KEY], true); + assert_eq!(metadata[USAGE_PRICING_AVAILABLE_METADATA_KEY], false); + assert_eq!(metadata[WEBSOCKET_MODE_METADATA_KEY], true); + assert_eq!( + metadata[WEBSOCKET_TRANSPORT_METADATA_KEY], + REALTIME_WEBSOCKET_TRANSPORT + ); + assert_eq!( + metadata[REALTIME_SESSION_METADATA_KEY]["usage_state"], + "authoritative" + ); + assert_eq!( + metadata[REALTIME_SESSION_METADATA_KEY]["pricing_state"], + "unsupported_audio_breakdown" + ); + assert_eq!( + metadata[REALTIME_SESSION_METADATA_KEY]["input_audio_tokens"], + 20 + ); + assert_eq!( + metadata[REALTIME_SESSION_METADATA_KEY]["usage_scope"], + "response_done" + ); + assert_eq!( + metadata[REALTIME_SESSION_METADATA_KEY]["input_transcription_usage_included"], + false + ); + } + + #[test] + fn missing_response_done_usage_is_visible_but_unmetered() { + let event = RealtimeSessionAudit::new(&sample_plan(), None) + .build_terminal_event(terminal(RealtimeUsageTotals::default())); + + 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.total_cost_usd, None); + let metadata = event.data.request_metadata.expect("metadata"); + assert_eq!(metadata[USAGE_AVAILABLE_METADATA_KEY], false); + assert_eq!(metadata[USAGE_PRICING_AVAILABLE_METADATA_KEY], false); + assert_eq!( + metadata[REALTIME_SESSION_METADATA_KEY]["usage_state"], + "unavailable" + ); + assert_eq!( + metadata[REALTIME_SESSION_METADATA_KEY]["pricing_state"], + "usage_unavailable" + ); + } + + #[test] + fn text_only_response_done_usage_remains_priceable() { + let event = RealtimeSessionAudit::new(&sample_plan(), None).build_terminal_event(terminal( + RealtimeUsageTotals { + responses: 1, + input_tokens: 12, + output_tokens: 4, + total_tokens: 16, + cached_input_tokens: 2, + input_audio_tokens: 0, + output_audio_tokens: 0, + }, + )); + + let metadata = event.data.request_metadata.expect("metadata"); + assert_eq!(metadata[USAGE_AVAILABLE_METADATA_KEY], true); + assert_eq!(metadata[USAGE_PRICING_AVAILABLE_METADATA_KEY], true); + assert_eq!( + metadata[REALTIME_SESSION_METADATA_KEY]["pricing_state"], + "compatible_text_usage" + ); + } + + #[test] + fn failed_session_preserves_authoritative_usage_without_becoming_billable() { + let mut failed = terminal(RealtimeUsageTotals { + responses: 1, + input_tokens: 18, + output_tokens: 3, + total_tokens: 21, + cached_input_tokens: 4, + input_audio_tokens: 7, + output_audio_tokens: 0, + }); + failed.disposition = RealtimeSessionDisposition::Failed; + failed.status_code = 502; + failed.termination = "upstream_read_failed"; + + let event = RealtimeSessionAudit::new(&sample_plan(), None).build_terminal_event(failed); + + assert_eq!(event.event_type, UsageEventType::Failed); + assert_eq!(event.data.input_tokens, Some(18)); + assert_eq!(event.data.total_tokens, Some(21)); + assert_eq!( + event.data.error_category.as_deref(), + Some("transport_error") + ); + let metadata = event.data.request_metadata.expect("metadata"); + assert_eq!(metadata[USAGE_AVAILABLE_METADATA_KEY], true); + assert_eq!(metadata[USAGE_PRICING_AVAILABLE_METADATA_KEY], false); + assert_eq!(metadata[REALTIME_SESSION_METADATA_KEY]["state"], "failed"); + } + + #[test] + fn failed_session_without_response_usage_is_explicitly_unavailable() { + let mut failed = terminal(RealtimeUsageTotals::default()); + failed.disposition = RealtimeSessionDisposition::Failed; + failed.status_code = 502; + failed.termination = "upstream_closed"; + + let event = RealtimeSessionAudit::new(&sample_plan(), None).build_terminal_event(failed); + + assert_eq!(event.event_type, UsageEventType::Failed); + assert_eq!(event.data.input_tokens, None); + assert_eq!(event.data.total_tokens, None); + let metadata = event.data.request_metadata.expect("metadata"); + assert_eq!(metadata[USAGE_AVAILABLE_METADATA_KEY], false); + assert_eq!(metadata[USAGE_PRICING_AVAILABLE_METADATA_KEY], false); + } +} diff --git a/apps/aether-gateway/src/handlers/proxy/websocket/realtime/mod.rs b/apps/aether-gateway/src/handlers/proxy/websocket/realtime/mod.rs new file mode 100644 index 000000000..fd44f04c5 --- /dev/null +++ b/apps/aether-gateway/src/handlers/proxy/websocket/realtime/mod.rs @@ -0,0 +1,65 @@ +//! Public OpenAI Realtime (`/v1/realtime`) WebSocket bridge. +//! +//! This is intentionally separate from Responses WebSocket mode and Codex +//! Frameless `/v1/live`: all three use WebSocket transport but have different +//! event grammars and lifecycle semantics. + +mod audit; +mod planner; +mod protocol; +mod session; + +use std::net::SocketAddr; + +use axum::body::Body; +use axum::extract::ws::WebSocketUpgrade; +use axum::extract::{ConnectInfo, State}; +use axum::http::{HeaderMap, Response, Uri}; + +use crate::handlers::proxy::websocket::ingress::{ + prepare_authenticated_ai_websocket, AuthenticatedAiWebSocketUpgradePreparation, + WebSocketIngressSpec, +}; +use crate::handlers::proxy::websocket::session::REALTIME_WEBSOCKET_SESSION_LIMITS; +use crate::{AppState, GatewayError}; + +pub(crate) async fn realtime_websocket( + State(state): State, + ConnectInfo(remote_addr): ConnectInfo, + ws: WebSocketUpgrade, + headers: HeaderMap, + uri: Uri, +) -> Result, GatewayError> { + match prepare_authenticated_ai_websocket( + state, + remote_addr, + headers, + uri, + REALTIME_WEBSOCKET_INGRESS_SPEC, + ) + .await? + { + AuthenticatedAiWebSocketUpgradePreparation::Rejected(response) => Ok(response), + AuthenticatedAiWebSocketUpgradePreparation::Ready(prepared) => { + let realtime = + match session::prepare_realtime_websocket(prepared.state(), prepared.context()) + .await + { + Ok(realtime) => realtime, + Err(rejection) => { + return prepared.rejection_response(rejection.status(), rejection.message()) + } + }; + Ok(prepared.into_response_with( + ws, + REALTIME_WEBSOCKET_SESSION_LIMITS, + realtime, + session::run_realtime_websocket, + )) + } + } +} + +const REALTIME_WEBSOCKET_INGRESS_SPEC: WebSocketIngressSpec = WebSocketIngressSpec { + route_unavailable_message: "OpenAI Realtime WebSocket route is unavailable", +}; diff --git a/apps/aether-gateway/src/handlers/proxy/websocket/realtime/planner.rs b/apps/aether-gateway/src/handlers/proxy/websocket/realtime/planner.rs new file mode 100644 index 000000000..d3f6bcafc --- /dev/null +++ b/apps/aether-gateway/src/handlers/proxy/websocket/realtime/planner.rs @@ -0,0 +1,249 @@ +//! Candidate planning for the public OpenAI Realtime WebSocket transport. + +use std::sync::atomic::{AtomicBool, Ordering}; +use std::sync::Arc; +use std::time::Duration; + +use axum::http::{HeaderValue, Method}; +use serde_json::json; + +use crate::ai_serving::{ + build_standard_stream_plan_from_decision, maybe_build_stream_decision_payload, + AiExecutionDecision, +}; +use crate::control::GatewayControlDecision; +use crate::headers::request_origin_from_headers_and_remote_addr; +use crate::privacy::RedactionSessionSlot; +use crate::{AppState, GatewayError}; + +pub(super) struct PlannedRealtimeCandidate { + pub(super) execution: AiExecutionDecision, + pub(super) admission_plan: aether_contracts::ExecutionPlan, + pub(super) provider_id: String, + pub(super) endpoint_id: String, + pub(super) key_id: String, + pub(super) provider_model: String, + pub(super) pool_lease: RealtimePoolLeaseGuard, +} + +pub(super) struct RealtimePoolLeaseGuard { + state: AppState, + report_context: Option, + renewal_task: Option>, + healthy: Arc, + armed: bool, +} + +impl RealtimePoolLeaseGuard { + fn new(state: &AppState, decision: &AiExecutionDecision) -> Self { + let report_context = decision.report_context.clone(); + let lease = crate::orchestration::local_execution_candidate_metadata_from_report_context( + report_context.as_ref(), + ) + .pool_key_lease; + let healthy = Arc::new(AtomicBool::new(true)); + let renewal_task = lease.map(|lease| { + let runtime_state = Arc::clone(&state.runtime_state); + let healthy = Arc::clone(&healthy); + tokio::spawn(async move { + let ttl = Duration::from_millis(lease.ttl_ms); + let interval = Duration::from_millis((lease.ttl_ms / 3).max(1)); + loop { + tokio::time::sleep(interval).await; + match runtime_state.lock_renew(&lease, ttl).await { + Ok(true) => {} + Ok(false) | Err(_) => { + healthy.store(false, Ordering::Release); + return; + } + } + } + }) + }); + Self { + state: state.clone(), + report_context, + renewal_task, + healthy, + armed: true, + } + } + + pub(super) fn is_healthy(&self) -> bool { + self.healthy.load(Ordering::Acquire) + } + + pub(super) async fn release(mut self) { + if let Some(task) = self.renewal_task.take() { + task.abort(); + } + crate::orchestration::release_pool_key_lease_from_report_context( + &self.state, + self.report_context.as_ref(), + ) + .await; + self.armed = false; + } +} + +impl Drop for RealtimePoolLeaseGuard { + fn drop(&mut self) { + if let Some(task) = self.renewal_task.take() { + task.abort(); + } + if !self.armed { + return; + } + let state = self.state.clone(); + let report_context = self.report_context.take(); + if let Ok(runtime) = tokio::runtime::Handle::try_current() { + runtime.spawn(async move { + crate::orchestration::release_pool_key_lease_from_report_context( + &state, + report_context.as_ref(), + ) + .await; + }); + } + } +} + +pub(super) async fn plan_realtime_candidate( + state: &AppState, + context: &crate::handlers::proxy::websocket::ingress::WebSocketRequestContext, + client_model: &str, +) -> Result, GatewayError> { + let parts = realtime_planning_parts(context); + let body = json!({"model": client_model}); + let Some(execution) = maybe_build_stream_decision_payload( + state, + &parts, + context.trace_id.as_str(), + &context.decision, + &body, + None, + ) + .await? + else { + return Ok(None); + }; + if execution + .provider_api_format + .as_deref() + .map(crate::ai_serving::normalize_api_format_alias) + .as_deref() + != Some("openai:realtime") + { + crate::orchestration::release_pool_key_lease_from_report_context( + state, + execution.report_context.as_ref(), + ) + .await; + return Ok(None); + } + + let pool_lease = RealtimePoolLeaseGuard::new(state, &execution); + let Some(attempt) = + build_standard_stream_plan_from_decision(&parts, &body, execution.clone(), false)? + else { + pool_lease.release().await; + return Ok(None); + }; + let provider_id = execution.provider_id.clone().unwrap_or_default(); + let endpoint_id = execution.endpoint_id.clone().unwrap_or_default(); + let key_id = execution.key_id.clone().unwrap_or_default(); + let provider_model = execution + .mapped_model + .clone() + .or_else(|| execution.model_name.clone()) + .unwrap_or_default(); + if provider_id.is_empty() + || endpoint_id.is_empty() + || key_id.is_empty() + || provider_model.trim().is_empty() + || execution.upstream_url.as_deref().is_none_or(str::is_empty) + { + pool_lease.release().await; + return Ok(None); + } + + Ok(Some(PlannedRealtimeCandidate { + execution, + admission_plan: attempt.plan, + provider_id, + endpoint_id, + key_id, + provider_model, + pool_lease, + })) +} + +fn realtime_planning_parts( + context: &crate::handlers::proxy::websocket::ingress::WebSocketRequestContext, +) -> http::request::Parts { + let mut request = http::Request::builder() + .method(Method::GET) + .uri(context.uri.clone()) + .body(()) + .expect("the authenticated Realtime URI must remain valid"); + *request.headers_mut() = context.headers.clone(); + request.headers_mut().insert( + http::header::CONTENT_TYPE, + HeaderValue::from_static("application/json"), + ); + request + .extensions_mut() + .insert(request_origin_from_headers_and_remote_addr( + &context.headers, + &context.remote_addr, + )); + request + .extensions_mut() + .insert(RedactionSessionSlot::default()); + request.into_parts().0 +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn realtime_planning_requires_the_explicit_realtime_route() { + let request = http::Request::builder() + .method(Method::GET) + .uri("/v1/realtime?model=gpt-realtime") + .body(()) + .unwrap(); + let (parts, _) = request.into_parts(); + let decision = GatewayControlDecision { + public_path: "/v1/realtime".to_string(), + public_query_string: Some("model=gpt-realtime".to_string()), + route_class: Some("ai_public".to_string()), + route_family: Some("openai".to_string()), + route_kind: Some("realtime".to_string()), + client_surface: None, + api_operation: None, + gateway_credential_carrier: None, + request_auth_channel: None, + auth_context: None, + admin_principal: None, + auth_endpoint_signature: None, + execution_runtime_candidate: true, + local_auth_rejection: None, + model_directive_policy: Default::default(), + }; + + assert_eq!( + crate::ai_serving::resolve_execution_runtime_stream_plan_kind_with_client_surface( + decision.route_class.as_deref(), + decision.route_family.as_deref(), + decision.route_kind.as_deref(), + decision.client_surface, + decision.request_auth_channel.as_deref(), + &parts.method, + parts.uri.path(), + ), + Some(crate::ai_serving::OPENAI_REALTIME_STREAM_PLAN_KIND) + ); + } +} diff --git a/apps/aether-gateway/src/handlers/proxy/websocket/realtime/protocol.rs b/apps/aether-gateway/src/handlers/proxy/websocket/realtime/protocol.rs new file mode 100644 index 000000000..ee59bcbb6 --- /dev/null +++ b/apps/aether-gateway/src/handlers/proxy/websocket/realtime/protocol.rs @@ -0,0 +1,236 @@ +//! Bounded validation and observation for the public OpenAI Realtime protocol. +//! +//! Realtime events are otherwise relayed as opaque text/binary frames. Keeping +//! this module deliberately small prevents Aether from becoming a schema +//! allowlist for future client and server events. + +use std::collections::BTreeSet; + +use serde_json::{json, Value}; + +const MAX_MODEL_BYTES: usize = 256; +const MAX_OBSERVED_RESPONSE_IDS: usize = 1_024; + +#[derive(Debug, Clone, Copy, PartialEq, Eq, thiserror::Error)] +pub(super) enum RealtimeProtocolError { + #[error("invalid Realtime model query")] + InvalidModelQuery, + #[error("invalid Realtime model")] + InvalidModel, +} + +impl RealtimeProtocolError { + pub(super) const fn client_message(self) -> &'static str { + match self { + Self::InvalidModelQuery => { + "Realtime WebSocket requires exactly one model query parameter" + } + Self::InvalidModel => { + "Realtime model must be a non-empty identifier no longer than 256 bytes" + } + } + } +} + +pub(super) fn model_from_query(query: Option<&str>) -> Result { + let mut model = None; + for (name, value) in url::form_urlencoded::parse(query.unwrap_or_default().as_bytes()) { + if name.eq_ignore_ascii_case("model") { + if model.is_some() { + return Err(RealtimeProtocolError::InvalidModelQuery); + } + validate_model(value.as_ref())?; + model = Some(value.into_owned()); + } else if query_parameter_is_sensitive(name.as_ref()) { + return Err(RealtimeProtocolError::InvalidModelQuery); + } + } + model.ok_or(RealtimeProtocolError::InvalidModelQuery) +} + +fn validate_model(model: &str) -> Result<(), RealtimeProtocolError> { + if model.is_empty() + || model.len() > MAX_MODEL_BYTES + || model.trim() != model + || model.chars().any(char::is_control) + { + return Err(RealtimeProtocolError::InvalidModel); + } + Ok(()) +} + +fn query_parameter_is_sensitive(name: &str) -> bool { + matches!( + name.to_ascii_lowercase().as_str(), + "key" + | "api_key" + | "api-key" + | "x-api-key" + | "access_token" + | "authorization" + | "token" + | "client_secret" + | "secret_key" + | "signature" + | "sig" + ) +} + +pub(super) fn error_event(code: &str, message: &str) -> Value { + json!({ + "type": "error", + "error": { + "type": "server_error", + "code": code, + "message": message, + } + }) +} + +#[derive(Debug, Default, Clone, Copy, PartialEq, Eq)] +pub(super) struct RealtimeUsageTotals { + pub(super) responses: u64, + pub(super) input_tokens: u64, + pub(super) output_tokens: u64, + pub(super) total_tokens: u64, + pub(super) cached_input_tokens: u64, + pub(super) input_audio_tokens: u64, + pub(super) output_audio_tokens: u64, +} + +#[derive(Debug, Default)] +pub(super) struct RealtimeUsageObserver { + totals: RealtimeUsageTotals, + response_ids: BTreeSet, +} + +impl RealtimeUsageObserver { + pub(super) fn observe(&mut self, raw: &str) { + let Ok(event) = serde_json::from_str::(raw) else { + return; + }; + if event.get("type").and_then(Value::as_str) != Some("response.done") { + return; + } + let Some(response) = event.get("response").and_then(Value::as_object) else { + return; + }; + let Some(usage) = response.get("usage").and_then(Value::as_object) else { + return; + }; + if let Some(response_id) = response + .get("id") + .and_then(Value::as_str) + .map(str::trim) + .filter(|value| !value.is_empty()) + { + if self.response_ids.contains(response_id) { + return; + } + if self.response_ids.len() >= MAX_OBSERVED_RESPONSE_IDS { + return; + } + self.response_ids.insert(response_id.to_string()); + } + let input_tokens = json_u64(usage.get("input_tokens")); + let output_tokens = json_u64(usage.get("output_tokens")); + let total_tokens = usage + .get("total_tokens") + .and_then(Value::as_u64) + .unwrap_or_else(|| input_tokens.saturating_add(output_tokens)); + self.totals.responses = self.totals.responses.saturating_add(1); + self.totals.input_tokens = self.totals.input_tokens.saturating_add(input_tokens); + self.totals.output_tokens = self.totals.output_tokens.saturating_add(output_tokens); + self.totals.total_tokens = self.totals.total_tokens.saturating_add(total_tokens); + if let Some(details) = usage.get("input_token_details").and_then(Value::as_object) { + self.totals.cached_input_tokens = self + .totals + .cached_input_tokens + .saturating_add(json_u64(details.get("cached_tokens"))); + self.totals.input_audio_tokens = self + .totals + .input_audio_tokens + .saturating_add(json_u64(details.get("audio_tokens"))); + } + if let Some(details) = usage.get("output_token_details").and_then(Value::as_object) { + self.totals.output_audio_tokens = self + .totals + .output_audio_tokens + .saturating_add(json_u64(details.get("audio_tokens"))); + } + } + + pub(super) const fn totals(&self) -> RealtimeUsageTotals { + self.totals + } +} + +fn json_u64(value: Option<&Value>) -> u64 { + value.and_then(Value::as_u64).unwrap_or(0) +} + +#[cfg(test)] +mod tests { + use super::{model_from_query, RealtimeUsageObserver}; + + #[test] + fn model_query_requires_one_bounded_model_and_ignores_safe_hints() { + assert_eq!( + model_from_query(Some("trace=1&model=gpt-realtime-client")), + Ok("gpt-realtime-client".to_string()) + ); + assert!(model_from_query(None).is_err()); + assert!(model_from_query(Some("model=a&MODEL=b")).is_err()); + assert!(model_from_query(Some("model=a&key=secret")).is_err()); + assert!(model_from_query(Some(format!("model={}", "x".repeat(257)).as_str())).is_err()); + } + + #[test] + fn response_done_usage_is_observed_once_without_reconstructing_events() { + let event = serde_json::json!({ + "type": "response.done", + "future_server_field": {"opaque": true}, + "response": { + "id": "resp_1", + "usage": { + "input_tokens": 12, + "output_tokens": 7, + "total_tokens": 19, + "input_token_details": {"cached_tokens": 4, "audio_tokens": 3}, + "output_token_details": {"audio_tokens": 2} + } + } + }) + .to_string(); + let mut observer = RealtimeUsageObserver::default(); + observer.observe(event.as_str()); + observer.observe(event.as_str()); + + let totals = observer.totals(); + assert_eq!(totals.responses, 1); + assert_eq!(totals.input_tokens, 12); + assert_eq!(totals.output_tokens, 7); + assert_eq!(totals.total_tokens, 19); + assert_eq!(totals.cached_input_tokens, 4); + assert_eq!(totals.input_audio_tokens, 3); + assert_eq!(totals.output_audio_tokens, 2); + } + + #[test] + fn missing_total_tokens_uses_authoritative_component_sum() { + let mut observer = RealtimeUsageObserver::default(); + observer.observe( + serde_json::json!({ + "type": "response.done", + "response": { + "id": "resp_without_total", + "usage": {"input_tokens": 9, "output_tokens": 4} + } + }) + .to_string() + .as_str(), + ); + + assert_eq!(observer.totals().total_tokens, 13); + } +} diff --git a/apps/aether-gateway/src/handlers/proxy/websocket/realtime/session.rs b/apps/aether-gateway/src/handlers/proxy/websocket/realtime/session.rs new file mode 100644 index 000000000..4845244c9 --- /dev/null +++ b/apps/aether-gateway/src/handlers/proxy/websocket/realtime/session.rs @@ -0,0 +1,652 @@ +//! Opaque bidirectional relay for the public OpenAI Realtime WebSocket API. + +use std::sync::{Arc, Mutex}; +use std::time::{Duration, Instant}; + +use axum::extract::ws::{Message as AxumWsMessage, WebSocket}; +use axum::http::StatusCode; +use futures_util::{SinkExt, StreamExt}; +use tracing::{info, warn}; +use wreq::ws::message::Message as WreqWsMessage; + +use crate::control::execution_plan_balance_capacity_rejection; +use crate::handlers::proxy::websocket::ingress::{ + WebSocketConnectionLog, WebSocketConnectionLogSpec, WebSocketRequestContext, +}; +use crate::handlers::proxy::websocket::responses::ResponsesWebSocketTurnAdmission; +use crate::handlers::proxy::websocket::session::{ + CLOSE_INTERNAL_ERROR, CLOSE_POLICY_VIOLATION, CLOSE_TRY_AGAIN, + REALTIME_WEBSOCKET_SESSION_LIMITS, WEBSOCKET_LOG_TRANSPORT, +}; +use crate::handlers::proxy::websocket::transport::{ + client_message_to_upstream, close_client_socket, close_upstream_socket, + connect_upstream_websocket, send_client_message, upstream_message_to_client, + websocket_relay_frame_queue, UpstreamWebSocketErrorCodes, WebSocketRelayPumpControl, + WebSocketRelayQueueError, WebSocketWriteError, +}; +use crate::{AppState, GatewayError}; + +use super::audit::{RealtimeSessionAudit, RealtimeSessionDisposition, RealtimeSessionTerminal}; +use super::planner::{plan_realtime_candidate, PlannedRealtimeCandidate}; +use super::protocol::{error_event, model_from_query, RealtimeUsageObserver}; + +const REALTIME_LOG_TARGET: &str = "aether_gateway::handlers::proxy::realtime_ws"; +const REALTIME_CONNECTION_LOG_SPEC: WebSocketConnectionLogSpec = WebSocketConnectionLogSpec { + opened_event_name: "openai_realtime_websocket_connection_opened", + closed_event_name: "openai_realtime_websocket_connection_closed", + opened_message: "gateway accepted OpenAI Realtime WebSocket connection", + closed_message: "gateway closed OpenAI Realtime WebSocket connection", + execution_path: "openai_realtime_websocket_bridge", + provider_type: "openai_realtime", +}; +const REALTIME_UPSTREAM_ERRORS: UpstreamWebSocketErrorCodes = UpstreamWebSocketErrorCodes { + upstream_url_missing: "openai_realtime_upstream_url_missing", + upstream_url_invalid: "openai_realtime_upstream_url_invalid", + frontdoor_self_loop: "openai_realtime_websocket_frontdoor_self_loop", + headers_invalid: "openai_realtime_websocket_headers_invalid", + client_build_failed: "openai_realtime_websocket_client_build_failed", + proxy_invalid: "openai_realtime_websocket_proxy_invalid", + tunnel_proxy_unsupported: "openai_realtime_websocket_tunnel_proxy_unsupported", + handshake_failed: "openai_realtime_websocket_handshake_failed", + upgrade_rejected: "openai_realtime_websocket_upgrade_rejected", + upgrade_failed: "openai_realtime_websocket_upgrade_failed", +}; + +pub(super) struct PreparedRealtimeWebSocket { + upstream: wreq::ws::WebSocket, + admission: ResponsesWebSocketTurnAdmission, + candidate: PlannedRealtimeCandidate, +} + +pub(super) struct RealtimeWebSocketPreflightRejection { + status: StatusCode, + message: String, +} + +impl RealtimeWebSocketPreflightRejection { + pub(super) const fn status(&self) -> StatusCode { + self.status + } + + pub(super) fn message(&self) -> &str { + self.message.as_str() + } +} + +pub(super) async fn prepare_realtime_websocket( + state: &AppState, + context: &WebSocketRequestContext, +) -> Result { + if !realtime_usage_accounting_is_safe(context) { + return Err(rejection( + StatusCode::NOT_IMPLEMENTED, + "Realtime WebSocket is unavailable for finite-balance keys until session usage settlement is enabled", + )); + } + let client_model = model_from_query(context.uri.query()) + .map_err(|error| rejection(StatusCode::BAD_REQUEST, error.client_message()))?; + let candidate = plan_realtime_candidate(state, context, client_model.as_str()) + .await + .map_err(|error| { + warn!( + target: REALTIME_LOG_TARGET, + event_name = "openai_realtime_planning_failed", + log_type = "ops", + transport = WEBSOCKET_LOG_TRANSPORT, + websocket = true, + trace_id = %context.trace_id, + error_kind = gateway_error_kind(&error), + "OpenAI Realtime candidate planning failed" + ); + rejection( + StatusCode::INTERNAL_SERVER_ERROR, + "Realtime provider planning failed", + ) + })? + .ok_or_else(|| { + rejection( + StatusCode::SERVICE_UNAVAILABLE, + "No eligible OpenAI Realtime provider mapping is available", + ) + })?; + + if execution_plan_balance_capacity_rejection( + state, + &context.decision, + &candidate.admission_plan, + candidate.execution.report_context.as_ref(), + ) + .await + .map_err(|_| { + rejection( + StatusCode::INTERNAL_SERVER_ERROR, + "Realtime balance admission failed", + ) + })? + .is_some() + { + candidate.pool_lease.release().await; + return Err(rejection( + StatusCode::TOO_MANY_REQUESTS, + "Realtime request capacity is unavailable", + )); + } + + let admission = match ResponsesWebSocketTurnAdmission::acquire( + state, + &candidate.admission_plan, + context.trace_id.as_str(), + ) + .await + { + Ok(admission) => admission, + Err(error) => { + candidate.pool_lease.release().await; + return Err(rejection( + admission_error_status(&error), + "Realtime connection admission failed", + )); + } + }; + if !candidate.pool_lease.is_healthy() { + admission.release().await; + candidate.pool_lease.release().await; + return Err(rejection( + StatusCode::SERVICE_UNAVAILABLE, + "Realtime provider ownership was lost", + )); + } + let mut upstream = match connect_upstream_websocket( + &candidate.execution, + REALTIME_WEBSOCKET_SESSION_LIMITS, + REALTIME_UPSTREAM_ERRORS, + ) + .await + { + Ok(connection) => connection.socket, + Err(error_code) => { + warn!( + target: REALTIME_LOG_TARGET, + event_name = "openai_realtime_upstream_connect_failed", + log_type = "ops", + transport = WEBSOCKET_LOG_TRANSPORT, + websocket = true, + trace_id = %context.trace_id, + provider_id = %candidate.provider_id, + endpoint_id = %candidate.endpoint_id, + key_id = %candidate.key_id, + error_code, + "OpenAI Realtime upstream connection failed" + ); + admission.release().await; + candidate.pool_lease.release().await; + return Err(rejection( + StatusCode::BAD_GATEWAY, + "Realtime upstream WebSocket connection failed", + )); + } + }; + // The scheduler lease can expire while the upstream WebSocket handshake + // is in flight. Re-check it after the handshake so an invalid provider + // candidate is rejected before the downstream HTTP 101 is committed. + if !candidate.pool_lease.is_healthy() { + close_upstream_socket(&mut upstream, None).await; + admission.release().await; + candidate.pool_lease.release().await; + return Err(rejection( + StatusCode::SERVICE_UNAVAILABLE, + "Realtime provider ownership was lost", + )); + } + + Ok(PreparedRealtimeWebSocket { + upstream, + admission, + candidate, + }) +} + +pub(super) async fn run_realtime_websocket( + mut client_socket: WebSocket, + state: AppState, + context: WebSocketRequestContext, + prepared: PreparedRealtimeWebSocket, +) { + let connection_log = WebSocketConnectionLog::new(&context, REALTIME_CONNECTION_LOG_SPEC); + connection_log.log_opened(); + let PreparedRealtimeWebSocket { + mut upstream, + admission, + candidate, + } = prepared; + let audit = RealtimeSessionAudit::new( + &candidate.admission_plan, + candidate.execution.report_context.as_ref(), + ); + let terminal = relay_realtime(&mut client_socket, &mut upstream, &context, &candidate).await; + close_upstream_socket(&mut upstream, None).await; + admission.release().await; + candidate.pool_lease.release().await; + if matches!( + terminal.termination, + "connection_duration_limit" | "connection_admission_lost" | "pool_key_lease_lost" + ) { + close_client_socket(&mut client_socket, CLOSE_TRY_AGAIN, terminal.termination).await; + } + audit.finish(&state, terminal).await; +} + +async fn relay_realtime( + client_socket: &mut WebSocket, + upstream: &mut wreq::ws::WebSocket, + context: &WebSocketRequestContext, + candidate: &PlannedRealtimeCandidate, +) -> RealtimeSessionTerminal { + let started_at = Instant::now(); + let connection_deadline = + tokio::time::sleep(REALTIME_WEBSOCKET_SESSION_LIMITS.max_connection_duration); + tokio::pin!(connection_deadline); + let mut lease_health = tokio::time::interval(Duration::from_secs(1)); + lease_health.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Delay); + let stats = Arc::new(Mutex::new(RelayStats::default())); + let usage = Arc::new(Mutex::new(RealtimeUsageObserver::default())); + let relay_control = WebSocketRelayPumpControl::new(); + + let termination = { + let (mut client_write, mut client_read) = (&mut *client_socket).split(); + let (mut upstream_write, mut upstream_read) = (&mut *upstream).split(); + + let client_to_upstream = { + let control = relay_control.clone(); + let stats = Arc::clone(&stats); + async move { + let (queue_tx, mut queue_rx) = websocket_relay_frame_queue(); + let reader_control = control.clone(); + let reader = async move { + loop { + let client = tokio::select! { + biased; + _ = reader_control.cancelled() => return "relay_cancelled", + client = client_read.next() => client, + }; + let Some(client) = client else { + return "client_closed"; + }; + let Ok(client) = client else { + return "client_read_failed"; + }; + let (bytes, is_close) = client_frame_metadata(&client); + { + let mut stats = stats + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()); + stats.client_frames = stats.client_frames.saturating_add(1); + stats.client_bytes = stats.client_bytes.saturating_add(bytes as u64); + } + match reader_control + .enqueue(&queue_tx, client_message_to_upstream(client)) + .await + { + Ok(()) => {} + Err(WebSocketRelayQueueError::Cancelled) => { + return "relay_cancelled"; + } + Err(WebSocketRelayQueueError::Closed) => { + return "upstream_write_failed"; + } + } + if is_close { + return "client_close_frame"; + } + } + }; + + let writer_control = control; + let writer = async move { + loop { + let message = tokio::select! { + biased; + _ = writer_control.cancelled() => return None, + message = queue_rx.recv() => message, + }; + let Some(message) = message else { + return None; + }; + let result = writer_control + .send(async { upstream_write.send(message).await.map_err(|_| ()) }) + .await; + match result { + Ok(()) => {} + Err(WebSocketWriteError::Cancelled) => return None, + Err(_) => return Some("upstream_write_failed"), + } + } + }; + + tokio::pin!(reader, writer); + tokio::select! { + reader_exit = &mut reader => { + writer.await.unwrap_or(reader_exit) + } + writer_exit = &mut writer => { + match writer_exit { + Some(writer_exit) => writer_exit, + None => reader.await, + } + } + } + } + }; + + let upstream_to_client = { + let control = relay_control.clone(); + let stats = Arc::clone(&stats); + let usage = Arc::clone(&usage); + async move { + let (queue_tx, mut queue_rx) = websocket_relay_frame_queue(); + let reader_control = control.clone(); + let reader = async move { + loop { + let provider = tokio::select! { + biased; + _ = reader_control.cancelled() => return "relay_cancelled", + provider = upstream_read.next() => provider, + }; + let Some(provider) = provider else { + return "upstream_closed"; + }; + let Ok(provider) = provider else { + return "upstream_read_failed"; + }; + let (bytes, is_close) = upstream_frame_metadata(&provider); + { + let mut stats = stats + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()); + stats.first_upstream_frame_ms.get_or_insert_with(|| { + started_at.elapsed().as_millis().min(u128::from(u64::MAX)) as u64 + }); + stats.upstream_frames = stats.upstream_frames.saturating_add(1); + stats.upstream_bytes = + stats.upstream_bytes.saturating_add(bytes as u64); + } + if let WreqWsMessage::Text(text) = &provider { + usage + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()) + .observe(text.as_str()); + } + match reader_control + .enqueue(&queue_tx, upstream_message_to_client(provider)) + .await + { + Ok(()) => {} + Err(WebSocketRelayQueueError::Cancelled) => { + return "relay_cancelled"; + } + Err(WebSocketRelayQueueError::Closed) => { + return "client_write_failed"; + } + } + if is_close { + return "upstream_close_frame"; + } + } + }; + + let writer_control = control; + let writer = async move { + loop { + let message = tokio::select! { + biased; + _ = writer_control.cancelled() => return None, + message = queue_rx.recv() => message, + }; + let Some(message) = message else { + return None; + }; + let result = writer_control + .send(async { client_write.send(message).await.map_err(|_| ()) }) + .await; + match result { + Ok(()) => {} + Err(WebSocketWriteError::Cancelled) => return None, + Err(_) => return Some("client_write_failed"), + } + } + }; + + tokio::pin!(reader, writer); + tokio::select! { + reader_exit = &mut reader => { + writer.await.unwrap_or(reader_exit) + } + writer_exit = &mut writer => { + match writer_exit { + Some(writer_exit) => writer_exit, + None => reader.await, + } + } + } + } + }; + + tokio::pin!(client_to_upstream, upstream_to_client); + let termination = loop { + tokio::select! { + termination = &mut client_to_upstream => break termination, + termination = &mut upstream_to_client => break termination, + _ = &mut connection_deadline => break "connection_duration_limit", + _ = wait_for_connection_permit_loss(context.websocket_connection_permit.as_ref()) => { + break "connection_admission_lost"; + } + _ = lease_health.tick() => { + if !candidate.pool_lease.is_healthy() { + break "pool_key_lease_lost"; + } + } + } + }; + relay_control.cancel(); + termination + }; + if termination == "pool_key_lease_lost" { + send_realtime_error( + client_socket, + "openai_realtime_pool_key_lease_lost", + "Realtime provider ownership was lost", + ) + .await; + } + let stats = *stats + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()); + let totals = usage + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()) + .totals(); + let elapsed_ms = started_at.elapsed().as_millis().min(u128::from(u64::MAX)) as u64; + info!( + target: REALTIME_LOG_TARGET, + event_name = "openai_realtime_relay_finished", + log_type = "event", + transport = WEBSOCKET_LOG_TRANSPORT, + websocket = true, + trace_id = %context.trace_id, + provider_id = %candidate.provider_id, + endpoint_id = %candidate.endpoint_id, + key_id = %candidate.key_id, + model = %candidate.provider_model, + termination, + client_frames = stats.client_frames, + client_bytes = stats.client_bytes, + upstream_frames = stats.upstream_frames, + upstream_bytes = stats.upstream_bytes, + response_count = totals.responses, + input_tokens = totals.input_tokens, + output_tokens = totals.output_tokens, + total_tokens = totals.total_tokens, + cached_input_tokens = totals.cached_input_tokens, + input_audio_tokens = totals.input_audio_tokens, + output_audio_tokens = totals.output_audio_tokens, + elapsed_ms, + "OpenAI Realtime opaque relay finished" + ); + realtime_terminal_from_relay(termination, elapsed_ms, stats, totals) +} + +fn realtime_terminal_from_relay( + termination: &'static str, + elapsed_ms: u64, + stats: RelayStats, + usage: super::protocol::RealtimeUsageTotals, +) -> RealtimeSessionTerminal { + let (disposition, status_code) = match termination { + "client_close_frame" | "upstream_close_frame" => { + (RealtimeSessionDisposition::Completed, 200) + } + "client_closed" + | "client_read_failed" + | "client_write_failed" + | "connection_duration_limit" => (RealtimeSessionDisposition::Cancelled, 499), + "pool_key_lease_lost" | "connection_admission_lost" => { + (RealtimeSessionDisposition::Failed, 503) + } + "upstream_closed" | "upstream_read_failed" | "upstream_write_failed" => { + (RealtimeSessionDisposition::Failed, 502) + } + _ => (RealtimeSessionDisposition::Failed, 500), + }; + RealtimeSessionTerminal { + disposition, + status_code, + termination, + elapsed_ms, + first_upstream_frame_ms: stats.first_upstream_frame_ms, + client_frames: stats.client_frames, + client_bytes: stats.client_bytes, + upstream_frames: stats.upstream_frames, + upstream_bytes: stats.upstream_bytes, + usage, + } +} + +fn realtime_usage_accounting_is_safe(context: &WebSocketRequestContext) -> bool { + context + .decision + .auth_context + .as_ref() + .is_some_and(|auth| auth.balance_remaining.is_none()) +} + +fn rejection( + status: StatusCode, + message: impl Into, +) -> RealtimeWebSocketPreflightRejection { + RealtimeWebSocketPreflightRejection { + status, + message: message.into(), + } +} + +fn admission_error_status(error: &GatewayError) -> StatusCode { + match error { + GatewayError::AdmissionTimeout { .. } => StatusCode::TOO_MANY_REQUESTS, + GatewayError::Client { status, .. } => *status, + GatewayError::LocalExecutionPlanningTimeout { .. } => StatusCode::GATEWAY_TIMEOUT, + _ => StatusCode::INTERNAL_SERVER_ERROR, + } +} + +fn gateway_error_kind(error: &GatewayError) -> &'static str { + match error { + GatewayError::UpstreamUnavailable { .. } => "upstream_unavailable", + GatewayError::ControlUnavailable { .. } => "control_unavailable", + GatewayError::LocalExecutionPlanningTimeout { .. } => "planning_timeout", + GatewayError::AdmissionTimeout { .. } => "admission_timeout", + GatewayError::Client { .. } => "client_error", + GatewayError::Internal(_) => "internal_error", + } +} + +async fn send_realtime_error(client_socket: &mut WebSocket, code: &str, message: &str) { + let event = error_event(code, message).to_string(); + let _ = send_client_message(client_socket, AxumWsMessage::Text(event.into())).await; +} + +#[derive(Clone, Copy, Default)] +struct RelayStats { + client_frames: u64, + client_bytes: u64, + upstream_frames: u64, + upstream_bytes: u64, + first_upstream_frame_ms: Option, +} + +fn client_frame_metadata(message: &AxumWsMessage) -> (usize, bool) { + match message { + AxumWsMessage::Text(text) => (text.len(), false), + AxumWsMessage::Binary(data) | AxumWsMessage::Ping(data) | AxumWsMessage::Pong(data) => { + (data.len(), false) + } + AxumWsMessage::Close(frame) => ( + frame + .as_ref() + .map_or(0, |frame| 2usize.saturating_add(frame.reason.len())), + true, + ), + } +} + +fn upstream_frame_metadata(message: &WreqWsMessage) -> (usize, bool) { + match message { + WreqWsMessage::Text(text) => (text.len(), false), + WreqWsMessage::Binary(data) | WreqWsMessage::Ping(data) | WreqWsMessage::Pong(data) => { + (data.len(), false) + } + WreqWsMessage::Close(frame) => ( + frame + .as_ref() + .map_or(0, |frame| 2usize.saturating_add(frame.reason.len())), + true, + ), + } +} + +async fn wait_for_connection_permit_loss(permit: Option<&aether_runtime::AdmissionPermit>) { + let Some(permit) = permit else { + std::future::pending::<()>().await; + return; + }; + let mut health = tokio::time::interval(Duration::from_secs(1)); + health.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Delay); + loop { + health.tick().await; + if !permit.is_healthy() { + return; + } + } +} + +#[cfg(test)] +mod tests { + use super::{client_frame_metadata, upstream_frame_metadata}; + use axum::extract::ws::Message as AxumWsMessage; + use wreq::ws::message::Message as WreqWsMessage; + + #[test] + fn opaque_frame_accounting_does_not_coalesce_audio_or_json_messages() { + assert_eq!( + client_frame_metadata(&AxumWsMessage::Text("{\"type\":\"session.update\"}".into())), + (25, false) + ); + assert_eq!( + client_frame_metadata(&AxumWsMessage::Binary(vec![1, 2, 3].into())), + (3, false) + ); + assert_eq!( + upstream_frame_metadata(&WreqWsMessage::Text("delta".into())), + (5, false) + ); + } +} diff --git a/apps/aether-gateway/src/handlers/proxy/websocket/session.rs b/apps/aether-gateway/src/handlers/proxy/websocket/session.rs index 0c3b41ad5..e3499c547 100644 --- a/apps/aether-gateway/src/handlers/proxy/websocket/session.rs +++ b/apps/aether-gateway/src/handlers/proxy/websocket/session.rs @@ -27,6 +27,14 @@ pub(crate) const LIVE_WEBSOCKET_SESSION_LIMITS: WebSocketSessionLimits = WebSock max_connection_duration: Duration::from_secs(60 * 60), }; +pub(crate) const REALTIME_WEBSOCKET_SESSION_LIMITS: WebSocketSessionLimits = + WebSocketSessionLimits { + max_frame_size: 16 << 20, + max_message_size: 16 << 20, + initial_message_timeout: Duration::from_secs(60), + max_connection_duration: Duration::from_secs(60 * 60), + }; + /// A peer that stops draining its receive window must not be able to pin the /// relay loop. Session loops await socket writes inside a `tokio::select!`, /// so an unbounded write also suspends the connection and per-turn deadlines diff --git a/apps/aether-gateway/src/handlers/proxy/websocket/transport.rs b/apps/aether-gateway/src/handlers/proxy/websocket/transport.rs index bd0352724..4fe1f30aa 100644 --- a/apps/aether-gateway/src/handlers/proxy/websocket/transport.rs +++ b/apps/aether-gateway/src/handlers/proxy/websocket/transport.rs @@ -15,6 +15,8 @@ use axum::http::header::{ use axum::http::{HeaderMap, HeaderName}; use futures_util::{SinkExt, TryFutureExt}; use serde_json::json; +use tokio::sync::mpsc; +use tokio_util::sync::CancellationToken; use url::Url; use wreq::ws::message::{CloseFrame as WreqCloseFrame, Message as WreqWsMessage}; @@ -255,6 +257,7 @@ pub(crate) fn websocket_timeouts( pub(crate) enum WebSocketWriteError { Failed, TimedOut, + Cancelled, } impl WebSocketWriteError { @@ -262,10 +265,76 @@ impl WebSocketWriteError { match self { Self::Failed => "write_failed", Self::TimedOut => "write_timeout", + Self::Cancelled => "write_cancelled", } } } +/// A small per-direction buffer keeps a slow reader from blocking the opposite +/// WebSocket direction while still applying bounded backpressure. At the Live +/// audio cadence this is deliberately only a short burst buffer, not a place +/// where a session can accumulate unbounded media. +pub(crate) const RELAY_FRAME_QUEUE_CAPACITY: usize = 16; + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(crate) enum WebSocketRelayQueueError { + Closed, + Cancelled, +} + +/// Shared cancellation for both read/write halves of a bidirectional relay. +/// +/// Queue admission and socket writes both observe this token, so a connection +/// deadline or lease loss can interrupt a full queue and an in-flight slow +/// write immediately instead of waiting for [`RELAY_WRITE_TIMEOUT`]. +#[derive(Clone, Default)] +pub(crate) struct WebSocketRelayPumpControl { + cancellation: CancellationToken, +} + +impl WebSocketRelayPumpControl { + pub(crate) fn new() -> Self { + Self::default() + } + + pub(crate) fn cancel(&self) { + self.cancellation.cancel(); + } + + pub(crate) async fn cancelled(&self) { + self.cancellation.cancelled().await; + } + + pub(crate) async fn enqueue( + &self, + sender: &mpsc::Sender, + message: T, + ) -> Result<(), WebSocketRelayQueueError> { + tokio::select! { + biased; + _ = self.cancellation.cancelled() => Err(WebSocketRelayQueueError::Cancelled), + result = sender.send(message) => { + result.map_err(|_| WebSocketRelayQueueError::Closed) + } + } + } + + pub(crate) async fn send(&self, write: F) -> Result<(), WebSocketWriteError> + where + F: std::future::Future>, + { + tokio::select! { + biased; + _ = self.cancellation.cancelled() => Err(WebSocketWriteError::Cancelled), + result = bounded_send(RELAY_WRITE_TIMEOUT, write) => result, + } + } +} + +pub(crate) fn websocket_relay_frame_queue() -> (mpsc::Sender, mpsc::Receiver) { + mpsc::channel(RELAY_FRAME_QUEUE_CAPACITY) +} + /// Relays one frame to the client under [`RELAY_WRITE_TIMEOUT`]. pub(crate) async fn send_client_message( client_socket: &mut WebSocket, @@ -500,8 +569,9 @@ mod tests { use super::{ bounded_send, guarded_websocket_upstream_url, responses_websocket_error_event, responses_websocket_error_event_with_stream_id, websocket_handshake_headers, - websocket_response_headers, websocket_upstream_url, WebSocketWriteError, - RELAY_WRITE_TIMEOUT, TEARDOWN_WRITE_TIMEOUT, + websocket_relay_frame_queue, websocket_response_headers, websocket_upstream_url, + WebSocketRelayPumpControl, WebSocketRelayQueueError, WebSocketWriteError, + RELAY_FRAME_QUEUE_CAPACITY, RELAY_WRITE_TIMEOUT, TEARDOWN_WRITE_TIMEOUT, }; use crate::frontdoor_loop_guard::configured_gateway_frontdoor_base_url; use axum::http::HeaderMap; @@ -524,6 +594,64 @@ mod tests { assert_eq!(outcome, Err(WebSocketWriteError::Failed)); assert_eq!(WebSocketWriteError::Failed.as_str(), "write_failed"); assert_eq!(WebSocketWriteError::TimedOut.as_str(), "write_timeout"); + assert_eq!(WebSocketWriteError::Cancelled.as_str(), "write_cancelled"); + } + + #[tokio::test] + async fn relay_frame_queue_is_bounded_and_fifo() { + let (sender, mut receiver) = websocket_relay_frame_queue(); + for frame in 0..RELAY_FRAME_QUEUE_CAPACITY { + sender + .try_send(frame) + .expect("the configured burst buffer should accept this frame"); + } + assert!(matches!( + sender.try_send(RELAY_FRAME_QUEUE_CAPACITY), + Err(tokio::sync::mpsc::error::TrySendError::Full(_)) + )); + for expected in 0..RELAY_FRAME_QUEUE_CAPACITY { + assert_eq!(receiver.recv().await, Some(expected)); + } + } + + #[tokio::test] + async fn relay_cancellation_interrupts_a_full_queue_without_waiting_for_capacity() { + let control = WebSocketRelayPumpControl::new(); + let (sender, _receiver) = websocket_relay_frame_queue(); + for frame in 0..RELAY_FRAME_QUEUE_CAPACITY { + sender.try_send(frame).expect("queue should fill exactly"); + } + let enqueue = control.enqueue(&sender, RELAY_FRAME_QUEUE_CAPACITY); + tokio::pin!(enqueue); + assert!(tokio::time::timeout(Duration::from_millis(5), &mut enqueue) + .await + .is_err()); + + control.cancel(); + assert_eq!( + tokio::time::timeout(Duration::from_millis(100), enqueue) + .await + .expect("cancellation should wake a blocked producer"), + Err(WebSocketRelayQueueError::Cancelled) + ); + } + + #[tokio::test] + async fn relay_cancellation_interrupts_a_stalled_socket_write() { + let control = WebSocketRelayPumpControl::new(); + let write = control.send(std::future::pending::>()); + tokio::pin!(write); + assert!(tokio::time::timeout(Duration::from_millis(5), &mut write) + .await + .is_err()); + + control.cancel(); + assert_eq!( + tokio::time::timeout(Duration::from_millis(100), write) + .await + .expect("cancellation should wake a stalled writer"), + Err(WebSocketWriteError::Cancelled) + ); } #[tokio::test] diff --git a/apps/aether-gateway/src/handlers/public/support/user_me_usage.rs b/apps/aether-gateway/src/handlers/public/support/user_me_usage.rs index ddb641217..8b35cb0e4 100644 --- a/apps/aether-gateway/src/handlers/public/support/user_me_usage.rs +++ b/apps/aether-gateway/src/handlers/public/support/user_me_usage.rs @@ -65,6 +65,45 @@ fn parse_users_me_usage_offset(query: Option<&str>) -> Result { } } +#[derive(Clone, Debug, Default, PartialEq, Eq)] +struct UsersMeUsageRecordFilter { + statuses: Option>, + is_stream: Option, + is_websocket: Option, + error_only: bool, +} + +fn parse_users_me_usage_record_filter(query: Option<&str>) -> UsersMeUsageRecordFilter { + let Some(status) = query_param_value(query, "status") + .map(|value| value.trim().to_ascii_lowercase()) + .filter(|value| !value.is_empty()) + else { + return UsersMeUsageRecordFilter::default(); + }; + + let mut filter = UsersMeUsageRecordFilter::default(); + match status.as_str() { + "stream" => { + filter.is_stream = Some(true); + filter.is_websocket = Some(false); + } + "standard" => { + filter.is_stream = Some(false); + filter.is_websocket = Some(false); + } + "websocket" | "ws" => filter.is_websocket = Some(true), + "error" | "failed" => filter.error_only = true, + "active" => { + filter.statuses = Some(vec!["pending".to_string(), "streaming".to_string()]); + } + "pending" | "streaming" | "completed" | "cancelled" => { + filter.statuses = Some(vec![status]); + } + _ => {} + } + filter +} + fn parse_users_me_usage_hours(query: Option<&str>) -> Result { match query_param_value(query, "hours") { Some(value) => parse_bounded_u32("hours", &value, 1, 720), @@ -480,6 +519,11 @@ fn build_users_me_usage_record_payload( "first_byte_time_ms": item.first_byte_time_ms, "is_stream": item.is_stream, "is_websocket": item.is_websocket(), + "websocket_transport": item.websocket_transport(), + "usage_available": item.usage_available(), + "usage_pricing_available": item.usage_pricing_available(), + "input_audio_tokens": item.realtime_input_audio_tokens(), + "output_audio_tokens": item.realtime_output_audio_tokens(), "upstream_is_stream": upstream_is_stream, "client_requested_stream": client_is_stream, "client_is_stream": client_is_stream, @@ -564,6 +608,11 @@ fn build_users_me_usage_active_payload(item: &StoredRequestUsageAudit) -> serde_ "endpoint_api_format": item.endpoint_api_format, "is_stream": item.is_stream, "is_websocket": item.is_websocket(), + "websocket_transport": item.websocket_transport(), + "usage_available": item.usage_available(), + "usage_pricing_available": item.usage_pricing_available(), + "input_audio_tokens": item.realtime_input_audio_tokens(), + "output_audio_tokens": item.realtime_output_audio_tokens(), "upstream_is_stream": upstream_is_stream, "client_requested_stream": client_is_stream, "client_is_stream": client_is_stream, @@ -947,6 +996,7 @@ pub(super) async fn handle_users_me_usage_get( Ok(value) => value, Err(detail) => return admin_stats_bad_request_response(detail), }; + let record_filter = parse_users_me_usage_record_filter(query); // When no time range is specified, default to 7 days to avoid full-table scans. let effective_time_range = time_range.or_else(|| { @@ -1087,10 +1137,11 @@ pub(super) async fn handle_users_me_usage_get( api_format: None, client_family: None, exclude_unknown_model_or_provider: false, - statuses: None, + statuses: record_filter.statuses.clone(), exclude_status_codes: Vec::new(), - is_stream: None, - error_only: false, + is_stream: record_filter.is_stream, + is_websocket: record_filter.is_websocket, + error_only: record_filter.error_only, keywords, matched_user_ids_by_keyword: Vec::new(), auth_user_reader_available: false, @@ -1143,10 +1194,11 @@ pub(super) async fn handle_users_me_usage_get( api_format: None, client_family: None, exclude_unknown_model_or_provider: false, - statuses: None, + statuses: record_filter.statuses.clone(), exclude_status_codes: Vec::new(), - is_stream: None, - error_only: false, + is_stream: record_filter.is_stream, + is_websocket: record_filter.is_websocket, + error_only: record_filter.error_only, limit: None, offset: None, newest_first: true, @@ -1172,10 +1224,11 @@ pub(super) async fn handle_users_me_usage_get( api_format: None, client_family: None, exclude_unknown_model_or_provider: false, - statuses: None, + statuses: record_filter.statuses.clone(), exclude_status_codes: Vec::new(), - is_stream: None, - error_only: false, + is_stream: record_filter.is_stream, + is_websocket: record_filter.is_websocket, + error_only: record_filter.error_only, limit: Some(limit), offset: Some(offset), newest_first: true, @@ -1314,6 +1367,7 @@ pub(super) async fn handle_users_me_usage_active_get( statuses: Some(vec!["pending".to_string(), "streaming".to_string()]), exclude_status_codes: Vec::new(), is_stream: None, + is_websocket: None, error_only: false, limit: Some(50), offset: None, @@ -1568,10 +1622,34 @@ mod tests { use super::{ build_users_me_usage_active_payload, build_users_me_usage_record_payload, - users_me_usage_client_is_stream, users_me_usage_is_failed, - users_me_usage_terminal_candidate_state_override, users_me_usage_upstream_is_stream, + parse_users_me_usage_record_filter, users_me_usage_client_is_stream, + users_me_usage_is_failed, users_me_usage_terminal_candidate_state_override, + users_me_usage_upstream_is_stream, }; + #[test] + fn users_me_usage_transport_statuses_are_disjoint_server_side_filters() { + for status in ["websocket", "ws", "WS"] { + let filter = parse_users_me_usage_record_filter(Some( + format!("limit=20&status={status}").as_str(), + )); + assert_eq!(filter.is_websocket, Some(true)); + assert_eq!(filter.is_stream, None); + assert_eq!(filter.statuses, None); + assert!(!filter.error_only); + } + + for (status, expected_stream) in [("stream", true), ("standard", false)] { + let filter = parse_users_me_usage_record_filter(Some( + format!("limit=20&status={status}").as_str(), + )); + assert_eq!(filter.is_stream, Some(expected_stream)); + assert_eq!(filter.is_websocket, Some(false)); + assert_eq!(filter.statuses, None); + assert!(!filter.error_only); + } + } + fn sample_usage(status: &str) -> StoredRequestUsageAudit { StoredRequestUsageAudit::new( "usage-1".to_string(), @@ -1731,6 +1809,12 @@ mod tests { request_metadata: Some(json!({ "websocket_mode": true, "websocket_transport": "responses", + "usage_available": false, + "usage_pricing_available": false, + "realtime_session": { + "input_audio_tokens": 7, + "output_audio_tokens": 3, + }, })), ..sample_usage("completed") }; @@ -1740,6 +1824,16 @@ mod tests { assert_eq!(record["is_websocket"], true); assert_eq!(active["is_websocket"], true); + assert_eq!(record["websocket_transport"], "responses"); + assert_eq!(active["websocket_transport"], "responses"); + assert_eq!(record["usage_available"], false); + assert_eq!(active["usage_available"], false); + assert_eq!(record["usage_pricing_available"], false); + assert_eq!(active["usage_pricing_available"], false); + assert_eq!(record["input_audio_tokens"], 7); + assert_eq!(active["input_audio_tokens"], 7); + assert_eq!(record["output_audio_tokens"], 3); + assert_eq!(active["output_audio_tokens"], 3); } #[test] diff --git a/apps/aether-gateway/src/maintenance/runtime/fixed_provider_reconciliation.rs b/apps/aether-gateway/src/maintenance/runtime/fixed_provider_reconciliation.rs index b1beefbde..c1a38d466 100644 --- a/apps/aether-gateway/src/maintenance/runtime/fixed_provider_reconciliation.rs +++ b/apps/aether-gateway/src/maintenance/runtime/fixed_provider_reconciliation.rs @@ -210,6 +210,28 @@ mod tests { .expect("endpoint transport should build"); responses.updated_at_unix_secs = Some(100); + let mut live = StoredProviderCatalogEndpoint::new( + "endpoint-codex-live".to_string(), + provider.id.clone(), + "codex:live".to_string(), + Some("codex".to_string()), + Some("live".to_string()), + false, + ) + .expect("Live endpoint should build") + .with_transport_fields( + "https://voice-proxy.internal/v1".to_string(), + None, + None, + Some(6), + Some("/socket/live".to_string()), + Some(json!({"custom_transport_option": true})), + None, + Some(json!({"url": "http://voice-proxy.internal:8080"})), + ) + .expect("Live endpoint transport should build"); + live.updated_at_unix_secs = Some(100); + let mut key = StoredProviderCatalogKey::new( "key-codex".to_string(), provider.id.clone(), @@ -228,10 +250,17 @@ mod tests { "custom".to_string(), ) .expect("unrelated provider should build"); + let fresh_codex_provider = StoredProviderCatalogProvider::new( + "provider-codex-fresh".to_string(), + "Fresh Codex".to_string(), + None, + "codex".to_string(), + ) + .expect("fresh Codex provider should build"); let repository = Arc::new(InMemoryProviderCatalogReadRepository::seed( - vec![provider, unrelated_provider], - vec![responses], + vec![provider, fresh_codex_provider, unrelated_provider], + vec![responses, live], vec![key], )); let state = AppState::new() @@ -247,7 +276,7 @@ mod tests { .list_endpoints_by_provider_ids(&["provider-codex".to_string()]) .await .expect("endpoints should list"); - assert_eq!(first_endpoints.len(), 4); + assert_eq!(first_endpoints.len(), 5); let responses = first_endpoints .iter() .find(|endpoint| endpoint.api_format == "openai:responses") @@ -272,6 +301,26 @@ mod tests { assert!(first_endpoints .iter() .any(|endpoint| endpoint.api_format == "openai:search")); + let live = first_endpoints + .iter() + .find(|endpoint| endpoint.api_format == "codex:live") + .expect("Codex Live endpoint should be reconciled from the v2 template"); + assert_eq!(live.api_family.as_deref(), Some("codex")); + assert_eq!(live.endpoint_kind.as_deref(), Some("live")); + assert_eq!(live.base_url, "https://voice-proxy.internal/v1"); + assert_eq!(live.custom_path.as_deref(), Some("/socket/live")); + assert_eq!(live.max_retries, Some(6)); + assert_eq!( + live.proxy, + Some(json!({"url": "http://voice-proxy.internal:8080"})) + ); + assert_eq!( + live.config + .as_ref() + .and_then(|value| value.get("custom_transport_option")), + Some(&json!(true)) + ); + assert!(!live.is_active); let keys = repository .list_keys_by_provider_ids(&["provider-codex".to_string()]) .await @@ -283,6 +332,17 @@ mod tests { .await .expect("unrelated endpoints should list") .is_empty()); + let fresh_endpoints = repository + .list_endpoints_by_provider_ids(&["provider-codex-fresh".to_string()]) + .await + .expect("fresh Codex endpoints should list"); + assert_eq!(fresh_endpoints.len(), 5); + let fresh_live = fresh_endpoints + .iter() + .find(|endpoint| endpoint.api_format == "codex:live") + .expect("v2 reconciliation should add a Codex Live endpoint"); + assert_eq!(fresh_live.base_url, "https://chatgpt.com/backend-api/codex"); + assert!(fresh_live.is_active); assert!(perform_fixed_provider_reconciliation_once(&state) .await @@ -292,6 +352,11 @@ mod tests { .await .expect("endpoints should list again"); assert_eq!(second_endpoints, first_endpoints); + let second_fresh_endpoints = repository + .list_endpoints_by_provider_ids(&["provider-codex-fresh".to_string()]) + .await + .expect("fresh Codex endpoints should list again"); + assert_eq!(second_fresh_endpoints, fresh_endpoints); } #[tokio::test] diff --git a/apps/aether-gateway/src/tests/control/admin/providers.rs b/apps/aether-gateway/src/tests/control/admin/providers.rs index 60919b9b1..e65167bd8 100644 --- a/apps/aether-gateway/src/tests/control/admin/providers.rs +++ b/apps/aether-gateway/src/tests/control/admin/providers.rs @@ -1188,7 +1188,7 @@ async fn gateway_creates_admin_provider_locally_with_trusted_admin_principal() { .list_endpoints_by_provider_ids(std::slice::from_ref(&created.id)) .await .expect("endpoints should list"); - assert_eq!(endpoints.len(), 4); + assert_eq!(endpoints.len(), 5); let responses_endpoint = endpoints .iter() .find(|endpoint| endpoint.api_format == "openai:responses") @@ -1205,6 +1205,10 @@ async fn gateway_creates_admin_provider_locally_with_trusted_admin_principal() { .iter() .find(|endpoint| endpoint.api_format == "openai:image") .expect("image endpoint should exist"); + let live_endpoint = endpoints + .iter() + .find(|endpoint| endpoint.api_format == "codex:live") + .expect("Codex Live endpoint should exist"); assert_eq!( responses_endpoint.base_url, "https://chatgpt.com/backend-api/codex" @@ -1225,6 +1229,14 @@ async fn gateway_creates_admin_provider_locally_with_trusted_admin_principal() { assert_eq!(compact_endpoint.max_retries, Some(7)); assert_eq!(search_endpoint.max_retries, Some(7)); assert_eq!(image_endpoint.max_retries, Some(7)); + assert_eq!(live_endpoint.api_family.as_deref(), Some("codex")); + assert_eq!(live_endpoint.endpoint_kind.as_deref(), Some("live")); + assert_eq!( + crate::api::ai::public_api_format_local_path(&live_endpoint.api_format), + "/v1/live" + ); + assert!(live_endpoint.custom_path.is_none()); + assert_eq!(live_endpoint.max_retries, Some(7)); assert_eq!( responses_endpoint .config @@ -1253,6 +1265,7 @@ async fn gateway_creates_admin_provider_locally_with_trusted_admin_principal() { assert!(compact_endpoint.body_rules.is_none()); assert!(search_endpoint.body_rules.is_none()); assert!(image_endpoint.body_rules.is_none()); + assert!(live_endpoint.body_rules.is_none()); assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0); gateway_handle.abort(); @@ -1296,6 +1309,14 @@ async fn gateway_updates_fixed_provider_and_reconciles_template_managed_endpoint ); cli_endpoint.max_retries = Some(2); cli_endpoint.config = Some(json!({"upstream_stream_policy": "force_stream"})); + let mut live_endpoint = sample_endpoint( + "endpoint-codex-live", + "provider-codex", + "codex:live", + "https://chatgpt.com/backend-api/codex", + ); + live_endpoint.max_retries = Some(2); + live_endpoint.custom_path = Some("/custom/live".to_string()); let mut key = sample_key( "key-codex-oauth", "provider-codex", @@ -1307,7 +1328,7 @@ async fn gateway_updates_fixed_provider_and_reconciles_template_managed_endpoint let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed( vec![provider], - vec![cli_endpoint], + vec![cli_endpoint, live_endpoint], vec![key], )); @@ -1367,7 +1388,7 @@ async fn gateway_updates_fixed_provider_and_reconciles_template_managed_endpoint .list_endpoints_by_provider_ids(&["provider-codex".to_string()]) .await .expect("endpoints should list"); - assert_eq!(endpoints.len(), 4); + assert_eq!(endpoints.len(), 5); let responses_endpoint = endpoints .iter() .find(|endpoint| endpoint.api_format == "openai:responses") @@ -1384,11 +1405,23 @@ async fn gateway_updates_fixed_provider_and_reconciles_template_managed_endpoint .iter() .find(|endpoint| endpoint.api_format == "openai:image") .expect("image endpoint should exist"); + let live_endpoint = endpoints + .iter() + .find(|endpoint| endpoint.api_format == "codex:live") + .expect("Codex Live endpoint should exist"); assert_eq!(responses_endpoint.max_retries, Some(9)); assert_eq!(compact_endpoint.max_retries, Some(9)); assert_eq!(search_endpoint.max_retries, Some(9)); assert_eq!(image_endpoint.max_retries, Some(9)); + assert_eq!(live_endpoint.api_family.as_deref(), Some("codex")); + assert_eq!(live_endpoint.endpoint_kind.as_deref(), Some("live")); + assert_eq!( + crate::api::ai::public_api_format_local_path(&live_endpoint.api_format), + "/v1/live" + ); + assert_eq!(live_endpoint.custom_path.as_deref(), Some("/custom/live")); + assert_eq!(live_endpoint.max_retries, Some(9)); assert_eq!( responses_endpoint .config diff --git a/apps/aether-gateway/src/tests/frontdoor/ai.rs b/apps/aether-gateway/src/tests/frontdoor/ai.rs index 80a2900ea..eb829133d 100644 --- a/apps/aether-gateway/src/tests/frontdoor/ai.rs +++ b/apps/aether-gateway/src/tests/frontdoor/ai.rs @@ -1,6 +1,7 @@ use super::{ - hash_api_key, sample_models_candidate_row, unrestricted_models_snapshot, - InMemoryAuthApiKeySnapshotRepository, InMemoryMinimalCandidateSelectionReadRepository, + hash_api_key, sample_endpoint, sample_key, sample_models_candidate_row, sample_provider, + unrestricted_models_snapshot, InMemoryAuthApiKeySnapshotRepository, + InMemoryMinimalCandidateSelectionReadRepository, InMemoryRequestCandidateRepository, InMemoryVideoTaskRepository, StoredAuthApiKeySnapshot, UpsertVideoTask, VideoTaskLookupKey, VideoTaskReadRepository, VideoTaskStatus, VideoTaskWriteRepository, DEVELOPMENT_ENCRYPTION_KEY, }; @@ -15,6 +16,7 @@ use aether_contracts::{ExecutionResult, ExecutionTelemetry, ResponseBody}; use aether_crypto::encrypt_python_fernet_plaintext; use aether_data::repository::global_models::InMemoryGlobalModelReadRepository; use aether_data::repository::provider_catalog::InMemoryProviderCatalogReadRepository; +use aether_data::repository::usage::InMemoryUsageReadRepository; use aether_data::DataLayerError; use aether_data_contracts::repository::auth::AuthApiKeyWriteRepository; use aether_data_contracts::repository::candidate_selection::{ @@ -29,6 +31,7 @@ use aether_data_contracts::repository::provider_catalog::{ ProviderCatalogReadRepository, StoredProviderCatalogEndpoint, StoredProviderCatalogKey, StoredProviderCatalogProvider, }; +use aether_data_contracts::repository::usage::{UsageAuditListQuery, UsageReadRepository}; use async_trait::async_trait; use axum::extract::ws::{Message as AxumWsMessage, WebSocket, WebSocketUpgrade}; use axum::extract::State; @@ -74,6 +77,17 @@ fn codex_models_snapshot( .expect("Codex models auth snapshot should build") } +fn codex_live_snapshot( + api_key_id: &str, + user_id: &str, + allowed_models: &[&str], +) -> StoredAuthApiKeySnapshot { + let mut snapshot = codex_models_snapshot(api_key_id, user_id, allowed_models); + snapshot.user_allowed_api_formats = Some(vec!["codex:live".to_string()]); + snapshot.api_key_allowed_api_formats = Some(vec!["codex:live".to_string()]); + snapshot +} + fn sample_codex_models_candidate_row( provider_id: &str, global_model_name: &str, @@ -101,6 +115,25 @@ fn sample_codex_models_candidate_row( row } +fn sample_codex_live_candidate_row( + provider_id: &str, + global_model_name: &str, + source_model_name: &str, +) -> StoredMinimalCandidateSelectionRow { + let mut row = + sample_codex_models_candidate_row(provider_id, global_model_name, source_model_name); + row.endpoint_api_format = "codex:live".to_string(); + row.endpoint_api_family = Some("codex".to_string()); + row.endpoint_kind = Some("live".to_string()); + row.key_api_formats = Some(vec!["codex:live".to_string()]); + if let Some(mappings) = row.model_provider_model_mappings.as_mut() { + for mapping in mappings { + mapping.api_formats = Some(vec!["codex:live".to_string()]); + } + } + row +} + fn complete_codex_model_card(source_model_name: &str) -> serde_json::Value { json!({ "id": source_model_name, @@ -170,6 +203,32 @@ fn codex_catalog_endpoint(provider_id: &str, endpoint_id: &str) -> StoredProvide .expect("Codex endpoint transport should build") } +fn codex_live_catalog_endpoint( + provider_id: &str, + endpoint_id: &str, +) -> StoredProviderCatalogEndpoint { + StoredProviderCatalogEndpoint::new( + endpoint_id.to_string(), + provider_id.to_string(), + "codex:live".to_string(), + Some("codex".to_string()), + Some("live".to_string()), + true, + ) + .expect("Codex Live endpoint should build") + .with_transport_fields( + "https://chatgpt.example/backend-api/codex".to_string(), + None, + None, + None, + None, + None, + None, + None, + ) + .expect("Codex Live endpoint transport should build") +} + fn codex_catalog_key( provider_id: &str, key_id: &str, @@ -204,6 +263,16 @@ fn codex_catalog_key( key } +fn codex_live_catalog_key( + provider_id: &str, + key_id: &str, + allowed_models: &[&str], +) -> StoredProviderCatalogKey { + let mut key = codex_catalog_key(provider_id, key_id, allowed_models); + key.api_formats = Some(json!(["codex:live"])); + key +} + fn codex_catalog_execution_result( plan: &aether_contracts::ExecutionPlan, status_code: u16, @@ -1243,6 +1312,401 @@ async fn run_versioned_codex_model_cards_frontdoor_scenario() { execution_runtime_handle.abort(); } +#[derive(Debug)] +struct ObservedOpenAiRealtimeWebSocket { + request_target: String, + authorization: Option, + route_header: Option, + session_update: serde_json::Value, + audio_append: serde_json::Value, + binary_frame: Vec, +} + +#[test] +fn gateway_relays_openai_realtime_audio_and_future_events_opaquely() { + super::run_frontdoor_async_test( + "openai-realtime-websocket-frontdoor", + run_openai_realtime_websocket_frontdoor_scenario(), + ); +} + +async fn run_openai_realtime_websocket_frontdoor_scenario() { + const PROVIDER_ID: &str = "provider-openai-realtime"; + const ENDPOINT_ID: &str = "endpoint-provider-openai-realtime"; + const UPSTREAM_KEY_ID: &str = "key-provider-openai-realtime"; + const CLIENT_MODEL: &str = "realtime-client-alias"; + const PROVIDER_MODEL: &str = "gpt-realtime-future"; + + let (observed_tx, observed_rx) = oneshot::channel(); + let upstream_state = Arc::new(Mutex::new(Some(observed_tx))); + let upstream = Router::new() + .route("/v1/realtime", get(mock_openai_realtime_websocket)) + .with_state(upstream_state); + let (upstream_url, upstream_handle) = start_server(upstream).await; + + let mut row = + sample_models_candidate_row(PROVIDER_ID, "openai", "openai:realtime", CLIENT_MODEL, 10); + row.endpoint_api_family = Some("openai".to_string()); + row.endpoint_kind = Some("realtime".to_string()); + row.key_allowed_models = Some(vec![PROVIDER_MODEL.to_string()]); + row.model_provider_model_name = PROVIDER_MODEL.to_string(); + row.model_provider_model_mappings = Some(vec![ + aether_data_contracts::repository::candidate_selection::StoredProviderModelMapping { + name: PROVIDER_MODEL.to_string(), + priority: 1, + api_formats: Some(vec!["openai:realtime".to_string()]), + endpoint_ids: None, + operations: None, + }, + ]); + let candidate_repository = + Arc::new(InMemoryMinimalCandidateSelectionReadRepository::seed(vec![ + row, + ])); + + let mut downstream_snapshot = + unrestricted_models_snapshot("gateway-key-openai-realtime", "user-openai-realtime"); + downstream_snapshot.user_allowed_providers = Some(vec!["openai".to_string()]); + downstream_snapshot.api_key_allowed_providers = Some(vec!["openai".to_string()]); + downstream_snapshot.user_allowed_api_formats = Some(vec!["openai:realtime".to_string()]); + downstream_snapshot.api_key_allowed_api_formats = Some(vec!["openai:realtime".to_string()]); + downstream_snapshot.user_allowed_models = Some(vec![CLIENT_MODEL.to_string()]); + downstream_snapshot.api_key_allowed_models = Some(vec![CLIENT_MODEL.to_string()]); + let auth_repository = Arc::new(InMemoryAuthApiKeySnapshotRepository::seed(vec![( + Some(hash_api_key("sk-openai-realtime")), + downstream_snapshot, + )])); + + let provider = sample_provider(PROVIDER_ID, "openai", 10); + let mut endpoint = sample_endpoint( + ENDPOINT_ID, + PROVIDER_ID, + "openai:realtime", + format!("{upstream_url}/v1").as_str(), + ); + endpoint.api_family = Some("openai".to_string()); + endpoint.endpoint_kind = Some("realtime".to_string()); + endpoint.header_rules = Some(json!([ + {"action": "set", "key": "x-upstream-realtime-route", "value": "opaque"} + ])); + let mut upstream_key = sample_key( + UPSTREAM_KEY_ID, + PROVIDER_ID, + "openai:realtime", + "realtime-upstream-secret", + ); + upstream_key.allowed_models = Some(json!([PROVIDER_MODEL])); + let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed( + vec![provider], + vec![endpoint], + vec![upstream_key], + )); + let request_candidate_repository = Arc::new(InMemoryRequestCandidateRepository::default()); + let usage_repository = Arc::new(InMemoryUsageReadRepository::default()); + + let state = AppState::new() + .expect("gateway should build") + .with_data_state_for_tests( + crate::data::GatewayDataState::with_auth_candidate_selection_provider_catalog_request_candidates_and_usage_for_tests( + auth_repository, + candidate_repository, + provider_catalog_repository, + request_candidate_repository, + Arc::clone(&usage_repository), + DEVELOPMENT_ENCRYPTION_KEY, + ), + ) + .with_usage_runtime_for_tests(crate::usage::UsageRuntimeConfig { + enabled: true, + ..crate::usage::UsageRuntimeConfig::default() + }); + let gateway = build_router_with_state(state); + let (gateway_url, gateway_handle) = start_server(gateway).await; + + let mut handshake_headers = HeaderMap::new(); + handshake_headers.insert( + http::header::AUTHORIZATION, + http::HeaderValue::from_static("Bearer sk-openai-realtime"), + ); + let invalid_model_response = wreq::Client::new() + .websocket(format!( + "{}/v1/realtime?model={CLIENT_MODEL}&model=duplicate", + gateway_url.replacen("http://", "ws://", 1) + )) + .headers(handshake_headers.clone()) + .send() + .await + .expect("invalid Realtime model query should return an HTTP response"); + assert_eq!(invalid_model_response.status(), StatusCode::BAD_REQUEST); + + let rejected_upstream_response = wreq::Client::new() + .websocket(format!( + "{}/v1/realtime?upstream_reject=1&model={CLIENT_MODEL}", + gateway_url.replacen("http://", "ws://", 1) + )) + .headers(handshake_headers.clone()) + .send() + .await + .expect("rejected upstream Realtime handshake should stay an HTTP response"); + assert_eq!(rejected_upstream_response.status(), StatusCode::BAD_GATEWAY); + + let response = wreq::Client::new() + .websocket(format!( + "{}/v1/realtime?trace=opaque&model={CLIENT_MODEL}", + gateway_url.replacen("http://", "ws://", 1) + )) + .headers(handshake_headers) + .send() + .await + .expect("Realtime gateway WebSocket handshake should complete"); + assert_eq!(response.status(), StatusCode::SWITCHING_PROTOCOLS); + let mut socket = response + .into_websocket() + .await + .expect("Realtime gateway response should upgrade"); + + let session_update = json!({ + "type": "session.update", + "session": { + "modalities": ["audio", "text"], + "future_session_capability": {"opaque": true, "revision": 23} + }, + "future_event_field": [1, {"nested": true}] + }); + let audio_append = json!({ + "type": "input_audio_buffer.append", + "audio": "AQIDBA==", + "future_audio_field": {"codec_revision": 7} + }); + socket + .send(WreqWsMessage::Text(session_update.to_string().into())) + .await + .expect("Realtime session.update should send"); + socket + .send(WreqWsMessage::Text(audio_append.to_string().into())) + .await + .expect("Realtime audio append should send"); + socket + .send(WreqWsMessage::Binary(vec![0, 1, 2, 255].into())) + .await + .expect("Realtime binary frame should send"); + + let audio_delta = receive_realtime_message(&mut socket).await; + let WreqWsMessage::Text(audio_delta) = audio_delta else { + panic!("Realtime audio delta should remain a text frame"); + }; + assert_eq!( + serde_json::from_str::(audio_delta.as_str()) + .expect("Realtime audio delta should remain valid JSON"), + json!({ + "type": "response.audio.delta", + "delta": "BQYHCA==", + "future_server_field": {"opaque": true, "revision": 29} + }) + ); + match receive_realtime_message(&mut socket).await { + WreqWsMessage::Binary(data) => assert_eq!(data.as_ref(), &[9, 8, 7]), + other => panic!("Realtime binary response changed frame type: {other:?}"), + } + let response_done = receive_realtime_message(&mut socket).await; + let WreqWsMessage::Text(response_done) = response_done else { + panic!("Realtime response.done should remain a text frame"); + }; + let response_done: serde_json::Value = serde_json::from_str(response_done.as_str()) + .expect("Realtime response.done should remain valid JSON"); + assert_eq!(response_done["type"], "response.done"); + assert_eq!(response_done["future_done_field"]["opaque"], true); + assert_eq!(response_done["response"]["usage"]["input_tokens"], 12); + assert_eq!( + response_done["response"]["usage"]["output_token_details"]["audio_tokens"], + 3 + ); + match receive_realtime_message(&mut socket).await { + WreqWsMessage::Close(_) => {} + other => panic!("Realtime upstream close changed frame type: {other:?}"), + } + + let observed = tokio::time::timeout(std::time::Duration::from_secs(2), observed_rx) + .await + .expect("mock Realtime upstream should report before timeout") + .expect("mock Realtime observation channel should remain open"); + assert_eq!( + observed.request_target, + format!("/v1/realtime?trace=opaque&model={PROVIDER_MODEL}") + ); + assert_eq!( + observed.authorization.as_deref(), + Some("Bearer realtime-upstream-secret") + ); + assert_eq!(observed.route_header.as_deref(), Some("opaque")); + assert_eq!(observed.session_update, session_update); + assert_eq!(observed.audio_append, audio_append); + assert_eq!(observed.binary_frame, vec![0, 1, 2, 255]); + + let realtime_usage = tokio::time::timeout(std::time::Duration::from_secs(5), async { + loop { + let records = usage_repository + .list_usage_audits(&UsageAuditListQuery::default()) + .await + .expect("Realtime usage audit list should load"); + if let Some(record) = records + .into_iter() + .find(|record| record.request_type.as_deref() == Some("realtime")) + { + break record; + } + tokio::task::yield_now().await; + } + }) + .await + .expect("Realtime session usage audit should be persisted before timeout"); + assert_eq!(realtime_usage.status, "completed"); + assert_eq!( + realtime_usage.api_format.as_deref(), + Some("openai:realtime") + ); + assert_eq!( + realtime_usage.endpoint_api_format.as_deref(), + Some("openai:realtime") + ); + assert!(realtime_usage.is_websocket()); + assert_eq!( + realtime_usage.websocket_transport(), + Some("openai_realtime") + ); + assert!(realtime_usage.usage_available()); + assert!(!realtime_usage.usage_pricing_available()); + assert_eq!(realtime_usage.billing_status, "void"); + assert_eq!(realtime_usage.input_tokens, 12); + assert_eq!(realtime_usage.output_tokens, 7); + assert_eq!(realtime_usage.total_tokens, 19); + assert_eq!(realtime_usage.cache_read_input_tokens, 4); + assert_eq!(realtime_usage.total_cost_usd, 0.0); + assert_eq!(realtime_usage.actual_total_cost_usd, 0.0); + let realtime_metadata = realtime_usage + .request_metadata + .as_ref() + .expect("Realtime usage metadata should be present"); + assert_eq!( + realtime_metadata["realtime_session"]["usage_scope"], + "response_done" + ); + assert_eq!( + realtime_metadata["realtime_session"]["input_audio_tokens"], + 5 + ); + assert_eq!( + realtime_metadata["realtime_session"]["output_audio_tokens"], + 3 + ); + + gateway_handle.abort(); + upstream_handle.abort(); +} + +async fn mock_openai_realtime_websocket( + State(observed): State>>>>, + uri: Uri, + headers: HeaderMap, + ws: WebSocketUpgrade, +) -> axum::response::Response { + if uri.query().is_some_and(|query| { + url::form_urlencoded::parse(query.as_bytes()) + .any(|(name, value)| name == "upstream_reject" && value == "1") + }) { + return StatusCode::UNAUTHORIZED.into_response(); + } + let request_target = uri.to_string(); + let authorization = headers + .get(http::header::AUTHORIZATION) + .and_then(|value| value.to_str().ok()) + .map(str::to_string); + let route_header = headers + .get("x-upstream-realtime-route") + .and_then(|value| value.to_str().ok()) + .map(str::to_string); + ws.on_upgrade(move |mut socket| async move { + let session_update = receive_axum_live_json(&mut socket).await; + let audio_append = receive_axum_live_json(&mut socket).await; + let binary_frame = + match tokio::time::timeout(std::time::Duration::from_secs(2), socket.recv()) + .await + .expect("mock Realtime upstream should receive binary frame before timeout") + .expect("mock Realtime upstream should remain open") + .expect("mock Realtime binary frame should be readable") + { + AxumWsMessage::Binary(data) => data.to_vec(), + other => panic!("mock Realtime upstream expected binary frame, got {other:?}"), + }; + + socket + .send(AxumWsMessage::Text( + json!({ + "type": "response.audio.delta", + "delta": "BQYHCA==", + "future_server_field": {"opaque": true, "revision": 29} + }) + .to_string() + .into(), + )) + .await + .expect("mock Realtime audio delta should send"); + socket + .send(AxumWsMessage::Binary(vec![9, 8, 7].into())) + .await + .expect("mock Realtime binary response should send"); + socket + .send(AxumWsMessage::Text( + json!({ + "type": "response.done", + "future_done_field": {"opaque": true}, + "response": { + "id": "resp_realtime_frontdoor", + "usage": { + "input_tokens": 12, + "output_tokens": 7, + "total_tokens": 19, + "input_token_details": {"cached_tokens": 4, "audio_tokens": 5}, + "output_token_details": {"audio_tokens": 3} + } + } + }) + .to_string() + .into(), + )) + .await + .expect("mock Realtime response.done should send"); + socket + .send(AxumWsMessage::Close(None)) + .await + .expect("mock Realtime close should send"); + + if let Some(sender) = observed + .lock() + .expect("mock Realtime observation mutex should lock") + .take() + { + let _ = sender.send(ObservedOpenAiRealtimeWebSocket { + request_target, + authorization, + route_header, + session_update, + audio_append, + binary_frame, + }); + } + }) + .into_response() +} + +async fn receive_realtime_message(socket: &mut wreq::ws::WebSocket) -> WreqWsMessage { + tokio::time::timeout(std::time::Duration::from_secs(2), socket.recv()) + .await + .expect("Realtime gateway should send a frame before timeout") + .expect("Realtime gateway socket should remain open") + .expect("Realtime gateway frame should be readable") +} + #[test] fn gateway_creates_bound_codex_live_oauth_calls_with_opaque_session_fields() { super::run_frontdoor_async_test( @@ -1259,7 +1723,7 @@ async fn run_codex_live_oauth_frontdoor_scenario() { const PROVIDER_MODEL: &str = "gpt-future-live"; const CALL_ID: &str = "rtc_frontdoor_live"; - let mut row = sample_codex_models_candidate_row(PROVIDER_ID, CLIENT_MODEL, PROVIDER_MODEL); + let mut row = sample_codex_live_candidate_row(PROVIDER_ID, CLIENT_MODEL, PROVIDER_MODEL); row.key_allowed_models = Some(vec![PROVIDER_MODEL.to_string()]); let candidate_repository = Arc::new(InMemoryMinimalCandidateSelectionReadRepository::seed(vec![ @@ -1267,7 +1731,7 @@ async fn run_codex_live_oauth_frontdoor_scenario() { ])); let auth_repository = Arc::new(InMemoryAuthApiKeySnapshotRepository::seed(vec![( Some(hash_api_key("sk-codex-live")), - codex_models_snapshot("gateway-key-codex-live", "user-codex-live", &[CLIENT_MODEL]), + codex_live_snapshot("gateway-key-codex-live", "user-codex-live", &[CLIENT_MODEL]), )])); let mut provider = codex_catalog_provider(PROVIDER_ID); @@ -1275,9 +1739,9 @@ async fn run_codex_live_oauth_frontdoor_scenario() { "responses_websocket": {"enabled": true}, "codex": {"fingerprint_convergence_enabled": true} })); - let mut endpoint = codex_catalog_endpoint(PROVIDER_ID, ENDPOINT_ID); + let mut endpoint = codex_live_catalog_endpoint(PROVIDER_ID, ENDPOINT_ID); endpoint.base_url = "https://chatgpt.com/backend-api/codex".to_string(); - let mut upstream_key = codex_catalog_key(PROVIDER_ID, UPSTREAM_KEY_ID, &[PROVIDER_MODEL]); + let mut upstream_key = codex_live_catalog_key(PROVIDER_ID, UPSTREAM_KEY_ID, &[PROVIDER_MODEL]); upstream_key.auth_type = "oauth".to_string(); upstream_key.encrypted_auth_config = Some( encrypt_python_fernet_plaintext( @@ -1490,7 +1954,7 @@ async fn run_codex_live_api_key_websocket_frontdoor_scenario() { .with_state(upstream_state); let (upstream_url, upstream_handle) = start_server(upstream).await; - let mut row = sample_codex_models_candidate_row(PROVIDER_ID, CLIENT_MODEL, PROVIDER_MODEL); + let mut row = sample_codex_live_candidate_row(PROVIDER_ID, CLIENT_MODEL, PROVIDER_MODEL); row.provider_name = "openai".to_string(); row.provider_type = "openai".to_string(); row.key_auth_type = "api_key".to_string(); @@ -1499,7 +1963,7 @@ async fn run_codex_live_api_key_websocket_frontdoor_scenario() { Arc::new(InMemoryMinimalCandidateSelectionReadRepository::seed(vec![ row, ])); - let mut downstream_snapshot = codex_models_snapshot( + let mut downstream_snapshot = codex_live_snapshot( "gateway-key-codex-live-websocket", "user-codex-live-websocket", &[CLIENT_MODEL], @@ -1514,9 +1978,9 @@ async fn run_codex_live_api_key_websocket_frontdoor_scenario() { let mut provider = codex_catalog_provider(PROVIDER_ID); provider.provider_type = "openai".to_string(); provider.config = Some(json!({"responses_websocket": {"enabled": true}})); - let mut endpoint = codex_catalog_endpoint(PROVIDER_ID, ENDPOINT_ID); + let mut endpoint = codex_live_catalog_endpoint(PROVIDER_ID, ENDPOINT_ID); endpoint.base_url = format!("{upstream_url}/v1"); - let mut upstream_key = codex_catalog_key(PROVIDER_ID, UPSTREAM_KEY_ID, &[PROVIDER_MODEL]); + let mut upstream_key = codex_live_catalog_key(PROVIDER_ID, UPSTREAM_KEY_ID, &[PROVIDER_MODEL]); upstream_key.auth_type = "api_key".to_string(); let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed( vec![provider], @@ -1728,7 +2192,7 @@ async fn run_codex_live_api_key_sideband_frontdoor_scenario() { ); let (execution_runtime_url, execution_runtime_handle) = start_server(execution_runtime).await; - let mut row = sample_codex_models_candidate_row(PROVIDER_ID, CLIENT_MODEL, PROVIDER_MODEL); + let mut row = sample_codex_live_candidate_row(PROVIDER_ID, CLIENT_MODEL, PROVIDER_MODEL); row.provider_name = "openai".to_string(); row.provider_type = "openai".to_string(); row.key_auth_type = "api_key".to_string(); @@ -1737,7 +2201,7 @@ async fn run_codex_live_api_key_sideband_frontdoor_scenario() { Arc::new(InMemoryMinimalCandidateSelectionReadRepository::seed(vec![ row, ])); - let mut downstream_snapshot = codex_models_snapshot( + let mut downstream_snapshot = codex_live_snapshot( "gateway-key-codex-live-sideband", "user-codex-live-sideband", &[CLIENT_MODEL], @@ -1752,9 +2216,9 @@ async fn run_codex_live_api_key_sideband_frontdoor_scenario() { let mut provider = codex_catalog_provider(PROVIDER_ID); provider.provider_type = "openai".to_string(); provider.config = Some(json!({"responses_websocket": {"enabled": true}})); - let mut endpoint = codex_catalog_endpoint(PROVIDER_ID, ENDPOINT_ID); + let mut endpoint = codex_live_catalog_endpoint(PROVIDER_ID, ENDPOINT_ID); endpoint.base_url = format!("{upstream_url}/v1"); - let mut upstream_key = codex_catalog_key(PROVIDER_ID, UPSTREAM_KEY_ID, &[PROVIDER_MODEL]); + let mut upstream_key = codex_live_catalog_key(PROVIDER_ID, UPSTREAM_KEY_ID, &[PROVIDER_MODEL]); upstream_key.auth_type = "api_key".to_string(); let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed( vec![provider], diff --git a/apps/aether-gateway/src/tests/frontdoor/ops.rs b/apps/aether-gateway/src/tests/frontdoor/ops.rs index e072e16f9..81fd14e2b 100644 --- a/apps/aether-gateway/src/tests/frontdoor/ops.rs +++ b/apps/aether-gateway/src/tests/frontdoor/ops.rs @@ -70,6 +70,7 @@ async fn gateway_exposes_frontdoor_manifest_without_proxying_upstream() { assert!(owned_routes .iter() .any(|value| value == "/v1/responses/compact")); + assert!(owned_routes.iter().any(|value| value == "/v1/realtime")); assert!(owned_routes.iter().any(|value| value == "/v1/alpha/search")); assert!(owned_routes.iter().any(|value| value == "/health")); assert!(owned_routes.iter().any(|value| value == "/v1/health")); diff --git a/crates/aether-admin/src/observability/usage.rs b/crates/aether-admin/src/observability/usage.rs index 7c861affd..5c4023d31 100644 --- a/crates/aether-admin/src/observability/usage.rs +++ b/crates/aether-admin/src/observability/usage.rs @@ -263,12 +263,17 @@ pub fn admin_usage_has_fallback(item: &StoredRequestUsageAudit) -> bool { } pub fn admin_usage_matches_status(item: &StoredRequestUsageAudit, status: Option<&str>) -> bool { - let Some(status) = status.map(str::trim).filter(|value| !value.is_empty()) else { + let Some(status) = status + .map(str::trim) + .filter(|value| !value.is_empty()) + .map(str::to_ascii_lowercase) + else { return true; }; - match status { - "stream" => item.is_stream, - "standard" => !item.is_stream, + match status.as_str() { + "stream" => item.is_stream && !item.is_websocket(), + "standard" => !item.is_stream && !item.is_websocket(), + "websocket" | "ws" => item.is_websocket(), "error" => { item.status_code .is_some_and(|value| !(200..300).contains(&value)) @@ -1217,6 +1222,11 @@ fn admin_usage_active_request_json( "provider_key_name": provider_key_name, "is_stream": item.is_stream, "is_websocket": item.is_websocket(), + "websocket_transport": item.websocket_transport(), + "usage_available": item.usage_available(), + "usage_pricing_available": item.usage_pricing_available(), + "input_audio_tokens": item.realtime_input_audio_tokens(), + "output_audio_tokens": item.realtime_output_audio_tokens(), "upstream_is_stream": upstream_is_stream, "client_requested_stream": client_is_stream, "client_is_stream": client_is_stream, @@ -1350,6 +1360,23 @@ pub fn admin_usage_record_json( )), ); object.insert("is_websocket".to_string(), json!(item.is_websocket())); + object.insert( + "websocket_transport".to_string(), + json!(item.websocket_transport()), + ); + object.insert("usage_available".to_string(), json!(item.usage_available())); + object.insert( + "usage_pricing_available".to_string(), + json!(item.usage_pricing_available()), + ); + object.insert( + "input_audio_tokens".to_string(), + json!(item.realtime_input_audio_tokens()), + ); + object.insert( + "output_audio_tokens".to_string(), + json!(item.realtime_output_audio_tokens()), + ); object.insert("is_stream".to_string(), json!(item.is_stream)); object.insert( UPSTREAM_IS_STREAM_KEY.to_string(), @@ -2705,6 +2732,12 @@ mod tests { request_metadata: Some(json!({ "websocket_mode": true, "websocket_transport": "responses", + "usage_available": false, + "usage_pricing_available": false, + "realtime_session": { + "input_audio_tokens": 7, + "output_audio_tokens": 3, + }, })), ..sample_usage("completed", Some(200), None) }; @@ -2721,6 +2754,21 @@ mod tests { assert_eq!(record["is_websocket"], true); assert_eq!(active["is_websocket"], true); + assert_eq!(record["websocket_transport"], "responses"); + assert_eq!(active["websocket_transport"], "responses"); + assert_eq!(record["usage_available"], false); + assert_eq!(active["usage_available"], false); + assert_eq!(record["usage_pricing_available"], false); + assert_eq!(active["usage_pricing_available"], false); + assert_eq!(record["input_audio_tokens"], 7); + assert_eq!(active["input_audio_tokens"], 7); + assert_eq!(record["output_audio_tokens"], 3); + assert_eq!(active["output_audio_tokens"], 3); + assert!(admin_usage_matches_status(&item, Some("websocket"))); + assert!(admin_usage_matches_status(&item, Some("ws"))); + assert!(admin_usage_matches_status(&item, Some("WS"))); + assert!(!admin_usage_matches_status(&item, Some("standard"))); + assert!(!admin_usage_matches_status(&item, Some("stream"))); } #[test] diff --git a/crates/aether-admin/src/system.rs b/crates/aether-admin/src/system.rs index b204a8d95..ddd653dff 100644 --- a/crates/aether-admin/src/system.rs +++ b/crates/aether-admin/src/system.rs @@ -769,6 +769,12 @@ const ADMIN_API_FORMAT_DEFINITIONS: &[AdminApiFormatDefinition] = &[ default_path: "/v1/responses/compact", aliases: &["responses_compact"], }, + AdminApiFormatDefinition { + value: "openai:realtime", + label: "OpenAI Realtime", + default_path: "/v1/realtime", + aliases: &["openai_realtime", "realtime"], + }, AdminApiFormatDefinition { value: "openai:search", label: "OpenAI Search", @@ -799,6 +805,12 @@ const ADMIN_API_FORMAT_DEFINITIONS: &[AdminApiFormatDefinition] = &[ default_path: "/v1/videos", aliases: &["openai_video", "sora"], }, + AdminApiFormatDefinition { + value: "codex:live", + label: "Codex Live", + default_path: "/v1/live", + aliases: &["codex_live", "live"], + }, AdminApiFormatDefinition { value: "claude:messages", label: "Claude Messages", @@ -3121,6 +3133,32 @@ fn mask_admin_proxy_node_password(password: Option<&str>) -> Option { mod tests { use super::*; + #[test] + fn api_formats_payload_exposes_realtime_and_codex_live_separately() { + let payload = build_admin_api_formats_payload(); + let formats = payload["formats"] + .as_array() + .expect("formats payload should be an array"); + let realtime = formats + .iter() + .find(|format| format["value"] == "openai:realtime") + .expect("OpenAI Realtime format should be registered"); + let live = formats + .iter() + .find(|format| format["value"] == "codex:live") + .expect("Codex Live format should be registered"); + + assert_eq!(realtime["label"], "OpenAI Realtime"); + assert_eq!(realtime["default_path"], "/v1/realtime"); + assert_eq!( + realtime["aliases"], + serde_json::json!(["openai_realtime", "realtime"]) + ); + assert_eq!(live["label"], "Codex Live"); + assert_eq!(live["default_path"], "/v1/live"); + assert_eq!(live["aliases"], serde_json::json!(["codex_live", "live"])); + } + #[test] fn build_admin_system_check_update_payload_reports_available_release() { let payload = build_admin_system_check_update_payload_with_release( diff --git a/crates/aether-ai/formats/src/api.rs b/crates/aether-ai/formats/src/api.rs index 52111d069..e75172785 100644 --- a/crates/aether-ai/formats/src/api.rs +++ b/crates/aether-ai/formats/src/api.rs @@ -9,12 +9,12 @@ pub use crate::contracts::{ CLAUDE_CLI_STREAM_SUCCESS_REPORT_KIND, CLAUDE_CLI_SYNC_ERROR_REPORT_KIND, CLAUDE_CLI_SYNC_FINALIZE_REPORT_KIND, CLAUDE_CLI_SYNC_PLAN_KIND, CLAUDE_CLI_SYNC_SUCCESS_REPORT_KIND, CLAUDE_COUNT_TOKENS_SYNC_PLAN_KIND, - CLAUDE_COUNT_TOKENS_SYNC_SUCCESS_REPORT_KIND, EXECUTION_RUNTIME_STREAM_ACTION, - EXECUTION_RUNTIME_STREAM_DECISION_ACTION, EXECUTION_RUNTIME_SYNC_ACTION, - EXECUTION_RUNTIME_SYNC_DECISION_ACTION, GEMINI_CHAT_STREAM_PLAN_KIND, - GEMINI_CHAT_STREAM_SUCCESS_REPORT_KIND, GEMINI_CHAT_SYNC_ERROR_REPORT_KIND, - GEMINI_CHAT_SYNC_FINALIZE_REPORT_KIND, GEMINI_CHAT_SYNC_PLAN_KIND, - GEMINI_CHAT_SYNC_SUCCESS_REPORT_KIND, GEMINI_CLI_STREAM_PLAN_KIND, + CLAUDE_COUNT_TOKENS_SYNC_SUCCESS_REPORT_KIND, CODEX_LIVE_STREAM_PLAN_KIND, + EXECUTION_RUNTIME_STREAM_ACTION, EXECUTION_RUNTIME_STREAM_DECISION_ACTION, + EXECUTION_RUNTIME_SYNC_ACTION, EXECUTION_RUNTIME_SYNC_DECISION_ACTION, + GEMINI_CHAT_STREAM_PLAN_KIND, GEMINI_CHAT_STREAM_SUCCESS_REPORT_KIND, + GEMINI_CHAT_SYNC_ERROR_REPORT_KIND, GEMINI_CHAT_SYNC_FINALIZE_REPORT_KIND, + GEMINI_CHAT_SYNC_PLAN_KIND, GEMINI_CHAT_SYNC_SUCCESS_REPORT_KIND, GEMINI_CLI_STREAM_PLAN_KIND, GEMINI_CLI_STREAM_SUCCESS_REPORT_KIND, GEMINI_CLI_SYNC_ERROR_REPORT_KIND, GEMINI_CLI_SYNC_FINALIZE_REPORT_KIND, GEMINI_CLI_SYNC_PLAN_KIND, GEMINI_CLI_SYNC_SUCCESS_REPORT_KIND, GEMINI_EMBEDDING_SYNC_PLAN_KIND, @@ -33,8 +33,8 @@ pub use crate::contracts::{ OPENAI_IMAGE_STREAM_PLAN_KIND, OPENAI_IMAGE_STREAM_SUCCESS_REPORT_KIND, OPENAI_IMAGE_SYNC_ERROR_REPORT_KIND, OPENAI_IMAGE_SYNC_FINALIZE_REPORT_KIND, OPENAI_IMAGE_SYNC_PLAN_KIND, OPENAI_IMAGE_SYNC_SUCCESS_REPORT_KIND, - OPENAI_RERANK_SYNC_PLAN_KIND, OPENAI_RESPONSES_COMPACT_STREAM_PLAN_KIND, - OPENAI_RESPONSES_COMPACT_STREAM_SUCCESS_REPORT_KIND, + OPENAI_REALTIME_STREAM_PLAN_KIND, OPENAI_RERANK_SYNC_PLAN_KIND, + OPENAI_RESPONSES_COMPACT_STREAM_PLAN_KIND, OPENAI_RESPONSES_COMPACT_STREAM_SUCCESS_REPORT_KIND, OPENAI_RESPONSES_COMPACT_SYNC_ERROR_REPORT_KIND, OPENAI_RESPONSES_COMPACT_SYNC_FINALIZE_REPORT_KIND, OPENAI_RESPONSES_COMPACT_SYNC_PLAN_KIND, OPENAI_RESPONSES_COMPACT_SYNC_SUCCESS_REPORT_KIND, OPENAI_RESPONSES_STREAM_PLAN_KIND, diff --git a/crates/aether-ai/formats/src/contracts/mod.rs b/crates/aether-ai/formats/src/contracts/mod.rs index 2c26e0f93..900c0afd9 100644 --- a/crates/aether-ai/formats/src/contracts/mod.rs +++ b/crates/aether-ai/formats/src/contracts/mod.rs @@ -14,20 +14,20 @@ pub use control_payloads::{build_ai_control_plan_request, AiControlPlanRequest}; pub use plan_kinds::{ is_openai_responses_stream_plan_kind, is_openai_responses_sync_plan_kind, CLAUDE_CHAT_STREAM_PLAN_KIND, CLAUDE_CHAT_SYNC_PLAN_KIND, CLAUDE_CLI_STREAM_PLAN_KIND, - CLAUDE_CLI_SYNC_PLAN_KIND, CLAUDE_COUNT_TOKENS_SYNC_PLAN_KIND, GEMINI_CHAT_STREAM_PLAN_KIND, - GEMINI_CHAT_SYNC_PLAN_KIND, GEMINI_CLI_STREAM_PLAN_KIND, GEMINI_CLI_SYNC_PLAN_KIND, - GEMINI_EMBEDDING_SYNC_PLAN_KIND, GEMINI_FILES_DELETE_PLAN_KIND, + CLAUDE_CLI_SYNC_PLAN_KIND, CLAUDE_COUNT_TOKENS_SYNC_PLAN_KIND, CODEX_LIVE_STREAM_PLAN_KIND, + GEMINI_CHAT_STREAM_PLAN_KIND, GEMINI_CHAT_SYNC_PLAN_KIND, GEMINI_CLI_STREAM_PLAN_KIND, + GEMINI_CLI_SYNC_PLAN_KIND, GEMINI_EMBEDDING_SYNC_PLAN_KIND, GEMINI_FILES_DELETE_PLAN_KIND, GEMINI_FILES_DOWNLOAD_PLAN_KIND, GEMINI_FILES_GET_PLAN_KIND, GEMINI_FILES_LIST_PLAN_KIND, GEMINI_FILES_UPLOAD_PLAN_KIND, GEMINI_INTERACTIONS_STREAM_PLAN_KIND, GEMINI_INTERACTIONS_SYNC_PLAN_KIND, GEMINI_VIDEO_CANCEL_SYNC_PLAN_KIND, GEMINI_VIDEO_CREATE_SYNC_PLAN_KIND, OPENAI_CHAT_STREAM_PLAN_KIND, OPENAI_CHAT_SYNC_PLAN_KIND, OPENAI_EMBEDDING_SYNC_PLAN_KIND, OPENAI_IMAGE_STREAM_PLAN_KIND, OPENAI_IMAGE_SYNC_PLAN_KIND, - OPENAI_RERANK_SYNC_PLAN_KIND, OPENAI_RESPONSES_COMPACT_STREAM_PLAN_KIND, - OPENAI_RESPONSES_COMPACT_SYNC_PLAN_KIND, OPENAI_RESPONSES_STREAM_PLAN_KIND, - OPENAI_RESPONSES_SYNC_PLAN_KIND, OPENAI_SEARCH_SYNC_PLAN_KIND, - OPENAI_VIDEO_CANCEL_SYNC_PLAN_KIND, OPENAI_VIDEO_CONTENT_PLAN_KIND, - OPENAI_VIDEO_CREATE_SYNC_PLAN_KIND, OPENAI_VIDEO_DELETE_SYNC_PLAN_KIND, - OPENAI_VIDEO_REMIX_SYNC_PLAN_KIND, + OPENAI_REALTIME_STREAM_PLAN_KIND, OPENAI_RERANK_SYNC_PLAN_KIND, + OPENAI_RESPONSES_COMPACT_STREAM_PLAN_KIND, OPENAI_RESPONSES_COMPACT_SYNC_PLAN_KIND, + OPENAI_RESPONSES_STREAM_PLAN_KIND, OPENAI_RESPONSES_SYNC_PLAN_KIND, + OPENAI_SEARCH_SYNC_PLAN_KIND, OPENAI_VIDEO_CANCEL_SYNC_PLAN_KIND, + OPENAI_VIDEO_CONTENT_PLAN_KIND, OPENAI_VIDEO_CREATE_SYNC_PLAN_KIND, + OPENAI_VIDEO_DELETE_SYNC_PLAN_KIND, OPENAI_VIDEO_REMIX_SYNC_PLAN_KIND, }; pub use report_kinds::{ core_error_background_report_kind, core_error_default_client_api_format, diff --git a/crates/aether-ai/formats/src/contracts/plan_kinds.rs b/crates/aether-ai/formats/src/contracts/plan_kinds.rs index 8684eff10..0da34706f 100644 --- a/crates/aether-ai/formats/src/contracts/plan_kinds.rs +++ b/crates/aether-ai/formats/src/contracts/plan_kinds.rs @@ -17,6 +17,8 @@ pub const CLAUDE_COUNT_TOKENS_SYNC_PLAN_KIND: &str = "claude_count_tokens_sync"; pub const GEMINI_CHAT_STREAM_PLAN_KIND: &str = "gemini_chat_stream"; pub const GEMINI_INTERACTIONS_STREAM_PLAN_KIND: &str = "gemini_interactions_stream"; pub const OPENAI_RESPONSES_STREAM_PLAN_KIND: &str = "openai_responses_stream"; +pub const OPENAI_REALTIME_STREAM_PLAN_KIND: &str = "openai_realtime_stream"; +pub const CODEX_LIVE_STREAM_PLAN_KIND: &str = "codex_live_stream"; pub const OPENAI_RESPONSES_COMPACT_STREAM_PLAN_KIND: &str = "openai_responses_compact_stream"; pub const CLAUDE_CLI_STREAM_PLAN_KIND: &str = "claude_cli_stream"; pub const GEMINI_CLI_STREAM_PLAN_KIND: &str = "gemini_cli_stream"; diff --git a/crates/aether-ai/formats/src/formats/id.rs b/crates/aether-ai/formats/src/formats/id.rs index c49c682d1..1bac0971f 100644 --- a/crates/aether-ai/formats/src/formats/id.rs +++ b/crates/aether-ai/formats/src/formats/id.rs @@ -5,6 +5,7 @@ use std::{fmt, str::FromStr}; #[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)] pub enum FormatFamily { OpenAi, + Codex, Claude, Gemini, Jina, @@ -23,9 +24,11 @@ pub enum FormatId { OpenAiChat, OpenAiResponses, OpenAiResponsesCompact, + OpenAiRealtime, OpenAiSearch, OpenAiEmbedding, OpenAiRerank, + CodexLive, ClaudeMessages, GeminiGenerateContent, GeminiInteractions, @@ -50,9 +53,11 @@ impl FormatId { Self::OpenAiChat | Self::OpenAiResponses | Self::OpenAiResponsesCompact + | Self::OpenAiRealtime | Self::OpenAiSearch | Self::OpenAiEmbedding | Self::OpenAiRerank => FormatFamily::OpenAi, + Self::CodexLive => FormatFamily::Codex, Self::ClaudeMessages => FormatFamily::Claude, Self::GeminiGenerateContent | Self::GeminiInteractions | Self::GeminiEmbedding => { FormatFamily::Gemini @@ -75,9 +80,11 @@ impl FormatId { Self::OpenAiChat => "openai:chat", Self::OpenAiResponses => "openai:responses", Self::OpenAiResponsesCompact => "openai:responses:compact", + Self::OpenAiRealtime => "openai:realtime", Self::OpenAiSearch => "openai:search", Self::OpenAiEmbedding => "openai:embedding", Self::OpenAiRerank => "openai:rerank", + Self::CodexLive => "codex:live", Self::ClaudeMessages => "claude:messages", Self::GeminiGenerateContent => "gemini:generate_content", Self::GeminiInteractions => "gemini:interactions", @@ -106,11 +113,15 @@ impl FromStr for FormatId { "openai:responses:compact" | "/v1/responses/compact" => { Ok(Self::OpenAiResponsesCompact) } + "openai:realtime" | "openai_realtime" | "realtime" | "/v1/realtime" => { + Ok(Self::OpenAiRealtime) + } "openai:search" | "openai_search" | "search" | "/v1/alpha/search" => { Ok(Self::OpenAiSearch) } "openai:embedding" | "/v1/embeddings" => Ok(Self::OpenAiEmbedding), "openai:rerank" | "/v1/rerank" => Ok(Self::OpenAiRerank), + "codex:live" | "codex_live" | "live" | "/v1/live" => Ok(Self::CodexLive), "claude:messages" | "/v1/messages" => Ok(Self::ClaudeMessages), "gemini:generate_content" => Ok(Self::GeminiGenerateContent), "gemini:interactions" @@ -272,6 +283,59 @@ mod tests { assert_eq!(FormatId::parse("gemini:cli"), None); } + #[test] + fn realtime_and_codex_live_are_distinct_first_class_formats() { + use super::{FormatFamily, FormatProfile}; + + for alias in [ + "openai:realtime", + "OPENAI_REALTIME", + "realtime", + "/v1/realtime", + ] { + assert_eq!(FormatId::parse(alias), Some(FormatId::OpenAiRealtime)); + assert_eq!(normalize_api_format_alias(alias), "openai:realtime"); + } + for alias in ["codex:live", "CODEX_LIVE", "live", "/v1/live"] { + assert_eq!(FormatId::parse(alias), Some(FormatId::CodexLive)); + assert_eq!(normalize_api_format_alias(alias), "codex:live"); + } + + assert_eq!(FormatId::OpenAiRealtime.family(), FormatFamily::OpenAi); + assert_eq!(FormatId::CodexLive.family(), FormatFamily::Codex); + assert_eq!(FormatId::OpenAiRealtime.profile(), FormatProfile::Default); + assert_eq!(FormatId::CodexLive.profile(), FormatProfile::Default); + assert!(!api_format_uses_body_stream_field("openai:realtime")); + assert!(!api_format_uses_body_stream_field("codex:live")); + } + + #[test] + fn realtime_live_and_responses_permissions_do_not_cover_each_other() { + for (allowed, requested) in [ + ("openai:responses", "openai:realtime"), + ("openai:responses", "codex:live"), + ("openai:realtime", "openai:responses"), + ("openai:realtime", "codex:live"), + ("codex:live", "openai:responses"), + ("codex:live", "openai:realtime"), + ] { + assert!(!api_format_permission_covers(allowed, requested)); + } + assert!(api_format_permission_covers( + "OPENAI_REALTIME", + "/v1/realtime" + )); + assert!(api_format_permission_covers("CODEX_LIVE", "/v1/live")); + assert_eq!( + api_format_permission_storage_aliases("openai:realtime"), + vec!["openai:realtime".to_string()] + ); + assert_eq!( + api_format_permission_storage_aliases("codex:live"), + vec!["codex:live".to_string()] + ); + } + #[test] fn responses_permission_covers_its_companion_endpoints() { assert!(api_format_permission_covers( diff --git a/crates/aether-ai/formats/src/formats/matrix.rs b/crates/aether-ai/formats/src/formats/matrix.rs index 0c2568e54..72b03f45b 100644 --- a/crates/aether-ai/formats/src/formats/matrix.rs +++ b/crates/aether-ai/formats/src/formats/matrix.rs @@ -57,6 +57,9 @@ pub fn request_candidate_api_format_preference( if client_api_format == "openai:responses:compact" { return (provider_api_format == "openai:responses:compact").then_some((0, 0)); } + if matches!(client_api_format.as_str(), "openai:realtime" | "codex:live") { + return (provider_api_format == client_api_format).then_some((0, 0)); + } if client_api_format == "openai:search" { return (provider_api_format == "openai:search").then_some((0, 0)); } @@ -112,6 +115,12 @@ pub fn request_candidate_api_formats( if client_api_format == "openai:responses:compact" { return vec!["openai:responses:compact"]; } + if client_api_format == "openai:realtime" { + return vec!["openai:realtime"]; + } + if client_api_format == "codex:live" { + return vec!["codex:live"]; + } if client_api_format == "openai:search" { return vec!["openai:search"]; } @@ -746,6 +755,34 @@ mod tests { ); } + #[test] + fn websocket_transports_only_select_their_exact_wire_format() { + assert_eq!( + request_candidate_api_formats("/v1/realtime", true), + vec!["openai:realtime"] + ); + assert_eq!( + request_candidate_api_formats("/v1/live", true), + vec!["codex:live"] + ); + assert_eq!( + request_candidate_api_format_preference("openai:realtime", "openai:realtime"), + Some((0, 0)) + ); + assert_eq!( + request_candidate_api_format_preference("codex:live", "codex:live"), + Some((0, 0)) + ); + assert_eq!( + request_candidate_api_format_preference("openai:realtime", "openai:responses"), + None + ); + assert_eq!( + request_candidate_api_format_preference("codex:live", "openai:responses"), + None + ); + } + #[test] fn request_conversion_enable_flag_only_applies_to_real_data_format_conversions() { assert!(!request_conversion_requires_enable_flag( diff --git a/crates/aether-ai/formats/src/formats/openai/responses/codex.rs b/crates/aether-ai/formats/src/formats/openai/responses/codex.rs index 357523782..14ad409f2 100644 --- a/crates/aether-ai/formats/src/formats/openai/responses/codex.rs +++ b/crates/aether-ai/formats/src/formats/openai/responses/codex.rs @@ -69,6 +69,7 @@ enum CodexOpenAiEndpointKind { Compact, Search, Images, + Live, } #[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)] @@ -652,6 +653,8 @@ fn codex_openai_endpoint_kind( Some(CodexOpenAiEndpointKind::Search) } else if is_openai_image_request(provider_api_format) { Some(CodexOpenAiEndpointKind::Images) + } else if aether_ai_formats::api_format_alias_matches(provider_api_format, "codex:live") { + Some(CodexOpenAiEndpointKind::Live) } else { None } @@ -2079,7 +2082,13 @@ pub fn apply_codex_openai_special_headers( } return; } - if endpoint_kind == CodexOpenAiEndpointKind::Images { + if matches!( + endpoint_kind, + CodexOpenAiEndpointKind::Images | CodexOpenAiEndpointKind::Live + ) { + if endpoint_kind == CodexOpenAiEndpointKind::Live { + remove_btree_header(provider_request_headers, CODEX_RESPONSES_LITE_HEADER); + } return; } @@ -2873,6 +2882,42 @@ mod tests { assert!(!headers.contains_key("accept")); } + #[test] + fn codex_live_uses_account_identity_without_responses_lite_headers() { + let mut headers = std::collections::BTreeMap::from([( + CODEX_RESPONSES_LITE_HEADER.to_string(), + "true".to_string(), + )]); + + apply_codex_openai_special_headers( + &mut headers, + &json!({"model": "gpt-live"}), + &http::HeaderMap::new(), + "codex", + "codex:live", + Some("request-live"), + Some(r#"{"account_id":"account-live","is_fedramp":true}"#), + ); + + assert_eq!( + headers.get("chatgpt-account-id").map(String::as_str), + Some("account-live") + ); + assert_eq!( + headers.get("x-openai-fedramp").map(String::as_str), + Some("true") + ); + assert_eq!( + headers.get("user-agent").map(String::as_str), + Some(CODEX_CLIENT_USER_AGENT) + ); + assert_eq!( + headers.get("originator").map(String::as_str), + Some(CODEX_CLIENT_ORIGINATOR) + ); + assert!(!headers.contains_key(CODEX_RESPONSES_LITE_HEADER)); + } + #[test] fn standard_codex_models_do_not_send_the_responses_lite_header() { let mut headers = std::collections::BTreeMap::from([( diff --git a/crates/aether-ai/formats/src/formats/registry.rs b/crates/aether-ai/formats/src/formats/registry.rs index ba9846a8e..bffce5342 100644 --- a/crates/aether-ai/formats/src/formats/registry.rs +++ b/crates/aether-ai/formats/src/formats/registry.rs @@ -38,7 +38,10 @@ pub fn parse_request( FormatId::GeminiEmbedding => gemini::embedding::request::from(body, ctx), FormatId::DoubaoEmbedding => doubao::embedding::request::from(body, ctx), FormatId::AliyunMultimodalEmbedding => aliyun::embedding::request::from(body, ctx), - FormatId::OpenAiSearch | FormatId::GeminiInteractions => None, + FormatId::OpenAiRealtime + | FormatId::OpenAiSearch + | FormatId::CodexLive + | FormatId::GeminiInteractions => None, } .ok_or_else(|| FormatError::RequestParseFailed { format: source.as_str().to_string(), @@ -72,7 +75,10 @@ fn emit_request_inner( FormatId::GeminiEmbedding => gemini::embedding::request::to(request, ctx), FormatId::DoubaoEmbedding => doubao::embedding::request::to(request, ctx), FormatId::AliyunMultimodalEmbedding => aliyun::embedding::request::to(request, ctx), - FormatId::OpenAiSearch | FormatId::GeminiInteractions => None, + FormatId::OpenAiRealtime + | FormatId::OpenAiSearch + | FormatId::CodexLive + | FormatId::GeminiInteractions => None, } .ok_or_else(|| FormatError::RequestEmitFailed { format: target.as_str().to_string(), @@ -296,6 +302,7 @@ pub fn parse_response( FormatId::ClaudeMessages => claude_messages::response::from(body, ctx), FormatId::GeminiGenerateContent => gemini_generate_content::response::from(body, ctx), FormatId::OpenAiEmbedding + | FormatId::OpenAiRealtime | FormatId::OpenAiSearch | FormatId::JinaEmbedding | FormatId::OpenAiRerank @@ -303,6 +310,7 @@ pub fn parse_response( | FormatId::GeminiEmbedding | FormatId::DoubaoEmbedding | FormatId::AliyunMultimodalEmbedding + | FormatId::CodexLive | FormatId::GeminiInteractions => None, } .ok_or_else(|| FormatError::ResponseParseFailed { @@ -335,6 +343,7 @@ fn emit_response_inner( FormatId::ClaudeMessages => claude_messages::response::to(response, ctx), FormatId::GeminiGenerateContent => gemini_generate_content::response::to(response, ctx), FormatId::OpenAiEmbedding + | FormatId::OpenAiRealtime | FormatId::OpenAiSearch | FormatId::JinaEmbedding | FormatId::OpenAiRerank @@ -342,6 +351,7 @@ fn emit_response_inner( | FormatId::GeminiEmbedding | FormatId::DoubaoEmbedding | FormatId::AliyunMultimodalEmbedding + | FormatId::CodexLive | FormatId::GeminiInteractions => None, } .ok_or_else(|| FormatError::ResponseEmitFailed { @@ -1086,13 +1096,15 @@ fn standard_request_root_field_is_audited(source: FormatId, key: &str) -> bool { | "tools" ), FormatId::OpenAiEmbedding + | FormatId::OpenAiRealtime | FormatId::OpenAiSearch | FormatId::OpenAiRerank | FormatId::GeminiEmbedding | FormatId::JinaEmbedding | FormatId::JinaRerank | FormatId::DoubaoEmbedding - | FormatId::AliyunMultimodalEmbedding => true, + | FormatId::AliyunMultimodalEmbedding + | FormatId::CodexLive => true, } } @@ -1575,13 +1587,15 @@ fn validate_source_response_stop_enums( FormatId::GeminiGenerateContent => validate_gemini_response_finish_reasons(body, target), FormatId::GeminiInteractions => Ok(()), FormatId::OpenAiEmbedding + | FormatId::OpenAiRealtime | FormatId::OpenAiSearch | FormatId::OpenAiRerank | FormatId::GeminiEmbedding | FormatId::JinaEmbedding | FormatId::JinaRerank | FormatId::DoubaoEmbedding - | FormatId::AliyunMultimodalEmbedding => Ok(()), + | FormatId::AliyunMultimodalEmbedding + | FormatId::CodexLive => Ok(()), } } diff --git a/crates/aether-ai/formats/src/formats/shared/passthrough.rs b/crates/aether-ai/formats/src/formats/shared/passthrough.rs index f1b7ccbad..2d8c8ed82 100644 --- a/crates/aether-ai/formats/src/formats/shared/passthrough.rs +++ b/crates/aether-ai/formats/src/formats/shared/passthrough.rs @@ -1,12 +1,13 @@ use crate::contracts::{ ApiOperation, CLAUDE_CHAT_STREAM_PLAN_KIND, CLAUDE_CHAT_SYNC_PLAN_KIND, CLAUDE_CLI_STREAM_PLAN_KIND, CLAUDE_CLI_SYNC_PLAN_KIND, CLAUDE_COUNT_TOKENS_SYNC_PLAN_KIND, - CLAUDE_COUNT_TOKENS_SYNC_SUCCESS_REPORT_KIND, GEMINI_CHAT_STREAM_PLAN_KIND, - GEMINI_CHAT_SYNC_PLAN_KIND, GEMINI_CLI_STREAM_PLAN_KIND, GEMINI_CLI_SYNC_PLAN_KIND, - GEMINI_EMBEDDING_SYNC_PLAN_KIND, GEMINI_EMBEDDING_SYNC_SUCCESS_REPORT_KIND, - GEMINI_INTERACTIONS_STREAM_PLAN_KIND, GEMINI_INTERACTIONS_STREAM_SUCCESS_REPORT_KIND, - GEMINI_INTERACTIONS_SYNC_PLAN_KIND, GEMINI_INTERACTIONS_SYNC_SUCCESS_REPORT_KIND, - OPENAI_EMBEDDING_SYNC_PLAN_KIND, OPENAI_RERANK_SYNC_PLAN_KIND, OPENAI_SEARCH_SYNC_PLAN_KIND, + CLAUDE_COUNT_TOKENS_SYNC_SUCCESS_REPORT_KIND, CODEX_LIVE_STREAM_PLAN_KIND, + GEMINI_CHAT_STREAM_PLAN_KIND, GEMINI_CHAT_SYNC_PLAN_KIND, GEMINI_CLI_STREAM_PLAN_KIND, + GEMINI_CLI_SYNC_PLAN_KIND, GEMINI_EMBEDDING_SYNC_PLAN_KIND, + GEMINI_EMBEDDING_SYNC_SUCCESS_REPORT_KIND, GEMINI_INTERACTIONS_STREAM_PLAN_KIND, + GEMINI_INTERACTIONS_STREAM_SUCCESS_REPORT_KIND, GEMINI_INTERACTIONS_SYNC_PLAN_KIND, + GEMINI_INTERACTIONS_SYNC_SUCCESS_REPORT_KIND, OPENAI_EMBEDDING_SYNC_PLAN_KIND, + OPENAI_REALTIME_STREAM_PLAN_KIND, OPENAI_RERANK_SYNC_PLAN_KIND, OPENAI_SEARCH_SYNC_PLAN_KIND, OPENAI_SEARCH_SYNC_SUCCESS_REPORT_KIND, }; @@ -114,6 +115,22 @@ pub fn resolve_sync_spec(plan_kind: &str) -> Option pub fn resolve_stream_spec(plan_kind: &str) -> Option { match plan_kind { + CODEX_LIVE_STREAM_PLAN_KIND => Some(LocalSameFormatProviderSpec { + api_format: "codex:live", + decision_kind: CODEX_LIVE_STREAM_PLAN_KIND, + report_kind: "codex_live_websocket_success", + family: LocalSameFormatProviderFamily::Standard, + require_streaming: true, + operation: None, + }), + OPENAI_REALTIME_STREAM_PLAN_KIND => Some(LocalSameFormatProviderSpec { + api_format: "openai:realtime", + decision_kind: OPENAI_REALTIME_STREAM_PLAN_KIND, + report_kind: "openai_realtime_websocket_success", + family: LocalSameFormatProviderFamily::Standard, + require_streaming: true, + operation: None, + }), CLAUDE_CHAT_STREAM_PLAN_KIND => Some(LocalSameFormatProviderSpec { api_format: "claude:messages", decision_kind: CLAUDE_CHAT_STREAM_PLAN_KIND, @@ -178,6 +195,24 @@ mod tests { assert!(spec.require_streaming); } + #[test] + fn resolves_openai_realtime_websocket_same_format_spec() { + let spec = resolve_stream_spec("openai_realtime_stream").expect("spec"); + assert_eq!(spec.api_format, "openai:realtime"); + assert_eq!(spec.report_kind, "openai_realtime_websocket_success"); + assert!(spec.require_streaming); + assert_eq!(spec.family, super::LocalSameFormatProviderFamily::Standard); + } + + #[test] + fn resolves_codex_live_websocket_same_format_spec() { + let spec = resolve_stream_spec("codex_live_stream").expect("spec"); + assert_eq!(spec.api_format, "codex:live"); + assert_eq!(spec.report_kind, "codex_live_websocket_success"); + assert!(spec.require_streaming); + assert_eq!(spec.family, super::LocalSameFormatProviderFamily::Standard); + } + #[test] fn resolves_openai_embedding_sync_same_format_spec() { let spec = resolve_sync_spec("openai_embedding_sync").expect("spec"); diff --git a/crates/aether-ai/formats/src/formats/shared/routing.rs b/crates/aether-ai/formats/src/formats/shared/routing.rs index 66b47dc8b..48dfd6e7a 100644 --- a/crates/aether-ai/formats/src/formats/shared/routing.rs +++ b/crates/aether-ai/formats/src/formats/shared/routing.rs @@ -4,13 +4,14 @@ use url::form_urlencoded; use crate::contracts::{ ClientSurface, CLAUDE_CHAT_STREAM_PLAN_KIND, CLAUDE_CHAT_SYNC_PLAN_KIND, CLAUDE_CLI_STREAM_PLAN_KIND, CLAUDE_CLI_SYNC_PLAN_KIND, CLAUDE_COUNT_TOKENS_SYNC_PLAN_KIND, - GEMINI_CHAT_STREAM_PLAN_KIND, GEMINI_CHAT_SYNC_PLAN_KIND, GEMINI_CLI_STREAM_PLAN_KIND, - GEMINI_CLI_SYNC_PLAN_KIND, GEMINI_EMBEDDING_SYNC_PLAN_KIND, GEMINI_FILES_DELETE_PLAN_KIND, - GEMINI_FILES_DOWNLOAD_PLAN_KIND, GEMINI_FILES_GET_PLAN_KIND, GEMINI_FILES_LIST_PLAN_KIND, - GEMINI_FILES_UPLOAD_PLAN_KIND, GEMINI_INTERACTIONS_STREAM_PLAN_KIND, - GEMINI_INTERACTIONS_SYNC_PLAN_KIND, GEMINI_VIDEO_CANCEL_SYNC_PLAN_KIND, - GEMINI_VIDEO_CREATE_SYNC_PLAN_KIND, OPENAI_CHAT_STREAM_PLAN_KIND, OPENAI_CHAT_SYNC_PLAN_KIND, - OPENAI_EMBEDDING_SYNC_PLAN_KIND, OPENAI_IMAGE_STREAM_PLAN_KIND, OPENAI_IMAGE_SYNC_PLAN_KIND, + CODEX_LIVE_STREAM_PLAN_KIND, GEMINI_CHAT_STREAM_PLAN_KIND, GEMINI_CHAT_SYNC_PLAN_KIND, + GEMINI_CLI_STREAM_PLAN_KIND, GEMINI_CLI_SYNC_PLAN_KIND, GEMINI_EMBEDDING_SYNC_PLAN_KIND, + GEMINI_FILES_DELETE_PLAN_KIND, GEMINI_FILES_DOWNLOAD_PLAN_KIND, GEMINI_FILES_GET_PLAN_KIND, + GEMINI_FILES_LIST_PLAN_KIND, GEMINI_FILES_UPLOAD_PLAN_KIND, + GEMINI_INTERACTIONS_STREAM_PLAN_KIND, GEMINI_INTERACTIONS_SYNC_PLAN_KIND, + GEMINI_VIDEO_CANCEL_SYNC_PLAN_KIND, GEMINI_VIDEO_CREATE_SYNC_PLAN_KIND, + OPENAI_CHAT_STREAM_PLAN_KIND, OPENAI_CHAT_SYNC_PLAN_KIND, OPENAI_EMBEDDING_SYNC_PLAN_KIND, + OPENAI_IMAGE_STREAM_PLAN_KIND, OPENAI_IMAGE_SYNC_PLAN_KIND, OPENAI_REALTIME_STREAM_PLAN_KIND, OPENAI_RERANK_SYNC_PLAN_KIND, OPENAI_RESPONSES_COMPACT_STREAM_PLAN_KIND, OPENAI_RESPONSES_COMPACT_SYNC_PLAN_KIND, OPENAI_RESPONSES_STREAM_PLAN_KIND, OPENAI_RESPONSES_SYNC_PLAN_KIND, OPENAI_SEARCH_SYNC_PLAN_KIND, @@ -116,6 +117,22 @@ pub fn resolve_execution_runtime_stream_plan_kind_with_client_surface( return Some(OPENAI_RESPONSES_STREAM_PLAN_KIND); } + if route_family == Some("openai") + && route_kind == Some("realtime") + && *method == Method::GET + && path == "/v1/realtime" + { + return Some(OPENAI_REALTIME_STREAM_PLAN_KIND); + } + + if route_family == Some("codex") + && route_kind == Some("live") + && ((*method == Method::GET && (path == "/v1/live" || path.starts_with("/v1/live/"))) + || (*method == Method::POST && path == "/v1/live")) + { + return Some(CODEX_LIVE_STREAM_PLAN_KIND); + } + if route_family == Some("openai") && route_kind == Some("image") && *method == Method::POST @@ -465,7 +482,7 @@ pub fn sanitize_request_path_and_query(path: &str, query: Option<&str>) -> Optio fn request_query_key_is_safe_to_trace(key: &str) -> bool { matches!( key.to_ascii_lowercase().as_str(), - "alt" | "view" | "pagesize" | "page_size" | "limit" | "offset" + "alt" | "view" | "pagesize" | "page_size" | "limit" | "offset" | "model" ) } @@ -541,6 +558,8 @@ pub fn supports_stream_execution_decision_kind(plan_kind: &str) -> bool { matches!( plan_kind, OPENAI_CHAT_STREAM_PLAN_KIND + | CODEX_LIVE_STREAM_PLAN_KIND + | OPENAI_REALTIME_STREAM_PLAN_KIND | CLAUDE_CHAT_STREAM_PLAN_KIND | GEMINI_CHAT_STREAM_PLAN_KIND | OPENAI_RESPONSES_STREAM_PLAN_KIND @@ -567,14 +586,15 @@ mod tests { }; use crate::contracts::{ CLAUDE_CHAT_STREAM_PLAN_KIND, CLAUDE_CHAT_SYNC_PLAN_KIND, CLAUDE_CLI_STREAM_PLAN_KIND, - CLAUDE_CLI_SYNC_PLAN_KIND, GEMINI_CHAT_STREAM_PLAN_KIND, GEMINI_CHAT_SYNC_PLAN_KIND, - GEMINI_CLI_STREAM_PLAN_KIND, GEMINI_CLI_SYNC_PLAN_KIND, GEMINI_EMBEDDING_SYNC_PLAN_KIND, - GEMINI_INTERACTIONS_STREAM_PLAN_KIND, GEMINI_INTERACTIONS_SYNC_PLAN_KIND, - OPENAI_CHAT_STREAM_PLAN_KIND, OPENAI_CHAT_SYNC_PLAN_KIND, OPENAI_EMBEDDING_SYNC_PLAN_KIND, - OPENAI_IMAGE_STREAM_PLAN_KIND, OPENAI_IMAGE_SYNC_PLAN_KIND, OPENAI_RERANK_SYNC_PLAN_KIND, - OPENAI_RESPONSES_COMPACT_STREAM_PLAN_KIND, OPENAI_RESPONSES_COMPACT_SYNC_PLAN_KIND, - OPENAI_RESPONSES_STREAM_PLAN_KIND, OPENAI_RESPONSES_SYNC_PLAN_KIND, - OPENAI_SEARCH_SYNC_PLAN_KIND, + CLAUDE_CLI_SYNC_PLAN_KIND, CODEX_LIVE_STREAM_PLAN_KIND, GEMINI_CHAT_STREAM_PLAN_KIND, + GEMINI_CHAT_SYNC_PLAN_KIND, GEMINI_CLI_STREAM_PLAN_KIND, GEMINI_CLI_SYNC_PLAN_KIND, + GEMINI_EMBEDDING_SYNC_PLAN_KIND, GEMINI_INTERACTIONS_STREAM_PLAN_KIND, + GEMINI_INTERACTIONS_SYNC_PLAN_KIND, OPENAI_CHAT_STREAM_PLAN_KIND, + OPENAI_CHAT_SYNC_PLAN_KIND, OPENAI_EMBEDDING_SYNC_PLAN_KIND, OPENAI_IMAGE_STREAM_PLAN_KIND, + OPENAI_IMAGE_SYNC_PLAN_KIND, OPENAI_REALTIME_STREAM_PLAN_KIND, + OPENAI_RERANK_SYNC_PLAN_KIND, OPENAI_RESPONSES_COMPACT_STREAM_PLAN_KIND, + OPENAI_RESPONSES_COMPACT_SYNC_PLAN_KIND, OPENAI_RESPONSES_STREAM_PLAN_KIND, + OPENAI_RESPONSES_SYNC_PLAN_KIND, OPENAI_SEARCH_SYNC_PLAN_KIND, }; #[test] @@ -646,6 +666,59 @@ mod tests { )); } + #[test] + fn resolves_openai_realtime_as_websocket_stream_only() { + assert_eq!( + resolve_execution_runtime_stream_plan_kind( + Some("ai_public"), + Some("openai"), + Some("realtime"), + None, + &Method::GET, + "/v1/realtime", + ), + Some(OPENAI_REALTIME_STREAM_PLAN_KIND) + ); + assert_eq!( + resolve_execution_runtime_sync_plan_kind( + Some("ai_public"), + Some("openai"), + Some("realtime"), + None, + &Method::GET, + "/v1/realtime", + ), + None + ); + assert!(supports_stream_execution_decision_kind( + OPENAI_REALTIME_STREAM_PLAN_KIND + )); + } + + #[test] + fn resolves_codex_live_call_create_and_websocket_routes() { + for (method, path) in [ + (Method::POST, "/v1/live"), + (Method::GET, "/v1/live"), + (Method::GET, "/v1/live/rtc_opaque"), + ] { + assert_eq!( + resolve_execution_runtime_stream_plan_kind( + Some("ai_public"), + Some("codex"), + Some("live"), + None, + &method, + path, + ), + Some(CODEX_LIVE_STREAM_PLAN_KIND) + ); + } + assert!(supports_stream_execution_decision_kind( + CODEX_LIVE_STREAM_PLAN_KIND + )); + } + #[test] fn resolves_openai_responses_compact_plan_kinds() { assert_eq!( @@ -892,6 +965,14 @@ mod tests { .as_deref(), Some("/v1/live/{call_id}?alt=sse") ); + assert_eq!( + sanitize_request_path_and_query( + "/v1/realtime?model=gpt-realtime-2.1&api_key=secret&token=hidden", + None + ) + .as_deref(), + Some("/v1/realtime?model=gpt-realtime-2.1") + ); } #[test] diff --git a/crates/aether-ai/formats/src/formats/shared/stream_core/format_matrix.rs b/crates/aether-ai/formats/src/formats/shared/stream_core/format_matrix.rs index 62c71740e..601750dee 100644 --- a/crates/aether-ai/formats/src/formats/shared/stream_core/format_matrix.rs +++ b/crates/aether-ai/formats/src/formats/shared/stream_core/format_matrix.rs @@ -418,6 +418,7 @@ impl ProviderStreamParser { FormatId::ClaudeMessages => Self::Claude(ClaudeProviderState::default()), FormatId::GeminiGenerateContent => Self::Gemini(GeminiProviderState::default()), FormatId::OpenAiEmbedding + | FormatId::OpenAiRealtime | FormatId::OpenAiSearch | FormatId::OpenAiRerank | FormatId::GeminiEmbedding @@ -425,7 +426,8 @@ impl ProviderStreamParser { | FormatId::JinaEmbedding | FormatId::JinaRerank | FormatId::DoubaoEmbedding - | FormatId::AliyunMultimodalEmbedding => return None, + | FormatId::AliyunMultimodalEmbedding + | FormatId::CodexLive => return None, }) } @@ -537,6 +539,7 @@ impl ClientStreamEmitter { FormatId::ClaudeMessages => Self::Claude(ClaudeClientEmitter::default()), FormatId::GeminiGenerateContent => Self::Gemini(GeminiClientEmitter::default()), FormatId::OpenAiEmbedding + | FormatId::OpenAiRealtime | FormatId::OpenAiSearch | FormatId::OpenAiRerank | FormatId::GeminiEmbedding @@ -544,7 +547,8 @@ impl ClientStreamEmitter { | FormatId::JinaEmbedding | FormatId::JinaRerank | FormatId::DoubaoEmbedding - | FormatId::AliyunMultimodalEmbedding => return None, + | FormatId::AliyunMultimodalEmbedding + | FormatId::CodexLive => return None, }) } @@ -660,13 +664,15 @@ fn parse_provider_error( parse_gemini_error(payload) } FormatId::OpenAiEmbedding + | FormatId::OpenAiRealtime | FormatId::OpenAiSearch | FormatId::OpenAiRerank | FormatId::GeminiEmbedding | FormatId::JinaEmbedding | FormatId::JinaRerank | FormatId::DoubaoEmbedding - | FormatId::AliyunMultimodalEmbedding => None, + | FormatId::AliyunMultimodalEmbedding + | FormatId::CodexLive => None, } } diff --git a/crates/aether-billing/src/event_enrichment.rs b/crates/aether-billing/src/event_enrichment.rs index 0b93326ec..a6dd1ccd1 100644 --- a/crates/aether-billing/src/event_enrichment.rs +++ b/crates/aether-billing/src/event_enrichment.rs @@ -1,7 +1,8 @@ use aether_data_contracts::repository::billing::StoredBillingModelContext; use aether_data_contracts::repository::usage::{ extract_provider_cache_ttl_minutes_from_metadata, resolve_provider_cache_ttl_minutes, - resolve_provider_service_tier_from_request_capture, + resolve_provider_service_tier_from_request_capture, USAGE_AVAILABLE_METADATA_KEY, + USAGE_PRICING_AVAILABLE_METADATA_KEY, }; use aether_data_contracts::DataLayerError; use aether_usage_runtime::{UsageEvent, UsageEventType}; @@ -39,6 +40,31 @@ pub async fn enrich_usage_event_with_billing( data: &dyn BillingModelContextLookup, event: &mut UsageEvent, ) -> Result<(), DataLayerError> { + // Session transports such as Codex Live expose lifecycle telemetry but no + // authoritative token/cost object. Do not run request-based pricing with + // zero default tokens: that would turn "unknown" into a fabricated charge. + if usage_is_explicitly_unavailable(event) { + event.data.input_tokens = None; + event.data.output_tokens = None; + event.data.total_tokens = None; + event.data.cache_creation_input_tokens = None; + event.data.cache_creation_ephemeral_5m_input_tokens = None; + event.data.cache_creation_ephemeral_1h_input_tokens = None; + event.data.cache_read_input_tokens = None; + event.data.cache_creation_cost_usd = None; + event.data.cache_read_cost_usd = None; + event.data.total_cost_usd = None; + event.data.actual_total_cost_usd = None; + return Ok(()); + } + // Some protocols expose authoritative token totals with dimensions that + // Aether's pricing schema cannot safely express yet (Realtime audio is the + // first example). Preserve those tokens for observability, but never apply + // ordinary text token prices to them. + if usage_pricing_is_explicitly_unavailable(event) { + clear_usage_costs(event); + return Ok(()); + } if !matches!(event.event_type, UsageEventType::Completed) { event.data.total_cost_usd = Some(0.0); event.data.actual_total_cost_usd = Some(0.0); @@ -108,6 +134,35 @@ pub async fn enrich_usage_event_with_billing( Ok(()) } +fn usage_is_explicitly_unavailable(event: &UsageEvent) -> bool { + event + .data + .request_metadata + .as_ref() + .and_then(Value::as_object) + .and_then(|metadata| metadata.get(USAGE_AVAILABLE_METADATA_KEY)) + .and_then(Value::as_bool) + == Some(false) +} + +fn usage_pricing_is_explicitly_unavailable(event: &UsageEvent) -> bool { + event + .data + .request_metadata + .as_ref() + .and_then(Value::as_object) + .and_then(|metadata| metadata.get(USAGE_PRICING_AVAILABLE_METADATA_KEY)) + .and_then(Value::as_bool) + == Some(false) +} + +fn clear_usage_costs(event: &mut UsageEvent) { + event.data.cache_creation_cost_usd = None; + event.data.cache_read_cost_usd = None; + event.data.total_cost_usd = None; + event.data.actual_total_cost_usd = None; +} + fn billing_model_lookup_names(data: &aether_usage_runtime::UsageEventData) -> Vec<&str> { let mut names = Vec::new(); for value in [data.target_model.as_deref(), Some(data.model.as_str())] @@ -433,6 +488,75 @@ mod tests { } } + #[tokio::test] + async fn unmetered_session_audit_does_not_fabricate_tokens_or_request_cost() { + let lookup = TestLookup { + name_context: None, + model_id_context: None, + }; + let mut event = UsageEvent::new( + UsageEventType::Completed, + "req-live-unmetered", + UsageEventData { + provider_name: "OpenAI".to_string(), + provider_id: Some("provider-live".to_string()), + model: "gpt-live".to_string(), + input_tokens: Some(0), + output_tokens: Some(0), + total_tokens: Some(0), + total_cost_usd: Some(0.0), + actual_total_cost_usd: Some(0.0), + request_metadata: Some(json!({"usage_available": false})), + ..UsageEventData::default() + }, + ); + + enrich_usage_event_with_billing(&lookup, &mut event) + .await + .expect("unmetered session audit should be accepted"); + + 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.total_cost_usd, None); + assert_eq!(event.data.actual_total_cost_usd, None); + } + + #[tokio::test] + async fn authoritative_realtime_audio_tokens_are_visible_but_not_text_priced() { + let lookup = TestLookup { + name_context: None, + model_id_context: None, + }; + let mut event = UsageEvent::new( + UsageEventType::Completed, + "req-realtime-audio", + UsageEventData { + provider_name: "OpenAI".to_string(), + provider_id: Some("provider-realtime".to_string()), + model: "gpt-realtime".to_string(), + input_tokens: Some(120), + output_tokens: Some(40), + total_tokens: Some(160), + request_metadata: Some(json!({ + "usage_available": true, + "usage_pricing_available": false, + })), + ..UsageEventData::default() + }, + ); + + enrich_usage_event_with_billing(&lookup, &mut event) + .await + .expect("unsupported audio pricing should not fail enrichment"); + + assert_eq!(event.data.input_tokens, Some(120)); + assert_eq!(event.data.output_tokens, Some(40)); + assert_eq!(event.data.total_tokens, Some(160)); + assert_eq!(event.data.total_cost_usd, None); + assert_eq!(event.data.actual_total_cost_usd, None); + } + #[test] fn processing_tier_fact_uses_final_provider_request_body() { let data = UsageEventData { diff --git a/crates/aether-data/adapters/mysql/migrations/20260821000000_migrate_legacy_codex_live_permissions.sql b/crates/aether-data/adapters/mysql/migrations/20260821000000_migrate_legacy_codex_live_permissions.sql new file mode 100644 index 000000000..7f000be64 --- /dev/null +++ b/crates/aether-data/adapters/mysql/migrations/20260821000000_migrate_legacy_codex_live_permissions.sql @@ -0,0 +1,203 @@ +-- Preserve the explicit access granted by the original #741 implementation, +-- which represented Codex Live as openai:responses. Runtime permission +-- matching remains strict after this one-time data upgrade. + +UPDATE users +SET allowed_api_formats = JSON_ARRAY_APPEND(allowed_api_formats, '$', 'codex:live') +WHERE LOWER(TRIM(COALESCE(allowed_api_formats_mode, ''))) = 'specific' + AND JSON_VALID(allowed_api_formats) + AND JSON_TYPE(IF(JSON_VALID(allowed_api_formats), allowed_api_formats, NULL)) = 'ARRAY' + AND JSON_CONTAINS( + IF(JSON_VALID(allowed_api_formats), allowed_api_formats, JSON_ARRAY()), + JSON_QUOTE('openai:responses'), '$' + ) = 1 + AND JSON_CONTAINS( + IF(JSON_VALID(allowed_api_formats), allowed_api_formats, JSON_ARRAY()), + JSON_QUOTE('codex:live'), '$' + ) = 0; + +UPDATE user_groups +SET allowed_api_formats = JSON_ARRAY_APPEND(allowed_api_formats, '$', 'codex:live') +WHERE LOWER(TRIM(COALESCE(allowed_api_formats_mode, ''))) = 'specific' + AND JSON_VALID(allowed_api_formats) + AND JSON_TYPE(IF(JSON_VALID(allowed_api_formats), allowed_api_formats, NULL)) = 'ARRAY' + AND JSON_CONTAINS( + IF(JSON_VALID(allowed_api_formats), allowed_api_formats, JSON_ARRAY()), + JSON_QUOTE('openai:responses'), '$' + ) = 1 + AND JSON_CONTAINS( + IF(JSON_VALID(allowed_api_formats), allowed_api_formats, JSON_ARRAY()), + JSON_QUOTE('codex:live'), '$' + ) = 0; + +UPDATE api_keys +SET allowed_api_formats = JSON_ARRAY_APPEND(allowed_api_formats, '$', 'codex:live') +WHERE JSON_VALID(allowed_api_formats) + AND JSON_TYPE(IF(JSON_VALID(allowed_api_formats), allowed_api_formats, NULL)) = 'ARRAY' + AND JSON_CONTAINS( + IF(JSON_VALID(allowed_api_formats), allowed_api_formats, JSON_ARRAY()), + JSON_QUOTE('openai:responses'), '$' + ) = 1 + AND JSON_CONTAINS( + IF(JSON_VALID(allowed_api_formats), allowed_api_formats, JSON_ARRAY()), + JSON_QUOTE('codex:live'), '$' + ) = 0; + +UPDATE provider_api_keys AS provider_key +INNER JOIN providers AS provider ON provider.id = provider_key.provider_id +SET + provider_key.api_formats = JSON_ARRAY_APPEND(provider_key.api_formats, '$', 'codex:live'), + provider_key.updated_at = UNIX_TIMESTAMP() +WHERE LOWER(TRIM(COALESCE(provider.provider_type, ''))) = 'codex' + AND JSON_VALID(provider_key.api_formats) + AND JSON_TYPE( + IF(JSON_VALID(provider_key.api_formats), provider_key.api_formats, NULL) + ) = 'ARRAY' + AND JSON_CONTAINS( + IF(JSON_VALID(provider_key.api_formats), provider_key.api_formats, JSON_ARRAY()), + JSON_QUOTE('openai:responses'), '$' + ) = 1 + AND JSON_CONTAINS( + IF(JSON_VALID(provider_key.api_formats), provider_key.api_formats, JSON_ARRAY()), + JSON_QUOTE('codex:live'), '$' + ) = 0; + +UPDATE provider_api_keys AS provider_key +INNER JOIN providers AS provider ON provider.id = provider_key.provider_id +SET + provider_key.auth_type_by_format = JSON_SET( + provider_key.auth_type_by_format, + '$."codex:live"', + JSON_EXTRACT(provider_key.auth_type_by_format, '$."openai:responses"') + ), + provider_key.updated_at = UNIX_TIMESTAMP() +WHERE LOWER(TRIM(COALESCE(provider.provider_type, ''))) = 'codex' + AND JSON_VALID(provider_key.auth_type_by_format) + AND JSON_TYPE( + IF(JSON_VALID(provider_key.auth_type_by_format), provider_key.auth_type_by_format, NULL) + ) = 'OBJECT' + AND JSON_CONTAINS_PATH( + IF( + JSON_VALID(provider_key.auth_type_by_format), + provider_key.auth_type_by_format, + JSON_OBJECT() + ), + 'one', + '$."openai:responses"' + ) = 1 + AND JSON_CONTAINS_PATH( + IF( + JSON_VALID(provider_key.auth_type_by_format), + provider_key.auth_type_by_format, + JSON_OBJECT() + ), + 'one', + '$."codex:live"' + ) = 0; + +UPDATE provider_api_keys AS provider_key +INNER JOIN providers AS provider ON provider.id = provider_key.provider_id +SET + provider_key.allow_auth_channel_mismatch_formats = JSON_ARRAY_APPEND( + provider_key.allow_auth_channel_mismatch_formats, + '$', + 'codex:live' + ), + provider_key.updated_at = UNIX_TIMESTAMP() +WHERE LOWER(TRIM(COALESCE(provider.provider_type, ''))) = 'codex' + AND JSON_VALID(provider_key.allow_auth_channel_mismatch_formats) + AND JSON_TYPE( + IF( + JSON_VALID(provider_key.allow_auth_channel_mismatch_formats), + provider_key.allow_auth_channel_mismatch_formats, + NULL + ) + ) = 'ARRAY' + AND JSON_CONTAINS( + IF( + JSON_VALID(provider_key.allow_auth_channel_mismatch_formats), + provider_key.allow_auth_channel_mismatch_formats, + JSON_ARRAY() + ), + JSON_QUOTE('openai:responses'), + '$' + ) = 1 + AND JSON_CONTAINS( + IF( + JSON_VALID(provider_key.allow_auth_channel_mismatch_formats), + provider_key.allow_auth_channel_mismatch_formats, + JSON_ARRAY() + ), + JSON_QUOTE('codex:live'), + '$' + ) = 0; + +UPDATE provider_api_keys AS provider_key +INNER JOIN providers AS provider ON provider.id = provider_key.provider_id +SET + provider_key.rate_multipliers = JSON_SET( + provider_key.rate_multipliers, + '$."codex:live"', + JSON_EXTRACT(provider_key.rate_multipliers, '$."openai:responses"') + ), + provider_key.updated_at = UNIX_TIMESTAMP() +WHERE LOWER(TRIM(COALESCE(provider.provider_type, ''))) = 'codex' + AND JSON_VALID(provider_key.rate_multipliers) + AND JSON_TYPE( + IF(JSON_VALID(provider_key.rate_multipliers), provider_key.rate_multipliers, NULL) + ) = 'OBJECT' + AND JSON_CONTAINS_PATH( + IF( + JSON_VALID(provider_key.rate_multipliers), + provider_key.rate_multipliers, + JSON_OBJECT() + ), + 'one', + '$."openai:responses"' + ) = 1 + AND JSON_CONTAINS_PATH( + IF( + JSON_VALID(provider_key.rate_multipliers), + provider_key.rate_multipliers, + JSON_OBJECT() + ), + 'one', + '$."codex:live"' + ) = 0; + +UPDATE provider_api_keys AS provider_key +INNER JOIN providers AS provider ON provider.id = provider_key.provider_id +SET + provider_key.global_priority_by_format = JSON_SET( + provider_key.global_priority_by_format, + '$."codex:live"', + JSON_EXTRACT(provider_key.global_priority_by_format, '$."openai:responses"') + ), + provider_key.updated_at = UNIX_TIMESTAMP() +WHERE LOWER(TRIM(COALESCE(provider.provider_type, ''))) = 'codex' + AND JSON_VALID(provider_key.global_priority_by_format) + AND JSON_TYPE( + IF( + JSON_VALID(provider_key.global_priority_by_format), + provider_key.global_priority_by_format, + NULL + ) + ) = 'OBJECT' + AND JSON_CONTAINS_PATH( + IF( + JSON_VALID(provider_key.global_priority_by_format), + provider_key.global_priority_by_format, + JSON_OBJECT() + ), + 'one', + '$."openai:responses"' + ) = 1 + AND JSON_CONTAINS_PATH( + IF( + JSON_VALID(provider_key.global_priority_by_format), + provider_key.global_priority_by_format, + JSON_OBJECT() + ), + 'one', + '$."codex:live"' + ) = 0; diff --git a/crates/aether-data/adapters/mysql/src/candidate_selection.rs b/crates/aether-data/adapters/mysql/src/candidate_selection.rs index 1967af316..9b45fb690 100644 --- a/crates/aether-data/adapters/mysql/src/candidate_selection.rs +++ b/crates/aether-data/adapters/mysql/src/candidate_selection.rs @@ -4,10 +4,11 @@ use async_trait::async_trait; use sqlx::{mysql::MySqlRow, MySql, QueryBuilder, Row}; use aether_data_contracts::repository::candidate_selection::{ - MinimalCandidateSelectionReadRepository, StoredApiFormatCandidateRowsQuery, - StoredMinimalCandidateSelectionRow, StoredPoolKeyCandidateOrder, - StoredPoolKeyCandidateRowsByKeyIdsQuery, StoredPoolKeyCandidateRowsQuery, - StoredProviderModelMapping, StoredRequestedModelCandidateRowsQuery, + provider_model_mapping_api_format_covers, MinimalCandidateSelectionReadRepository, + StoredApiFormatCandidateRowsQuery, StoredMinimalCandidateSelectionRow, + StoredPoolKeyCandidateOrder, StoredPoolKeyCandidateRowsByKeyIdsQuery, + StoredPoolKeyCandidateRowsQuery, StoredProviderModelMapping, + StoredRequestedModelCandidateRowsQuery, }; use aether_data_contracts::DataLayerError; @@ -675,9 +676,9 @@ fn mapping_scope_matches( api_format: &str, ) -> bool { mapping.api_formats.as_ref().is_none_or(|formats| { - formats - .iter() - .any(|value| api_format_scope_covers(value, api_format)) + formats.iter().any(|value| { + provider_model_mapping_api_format_covers(&row.provider_type, value, api_format) + }) }) && mapping.endpoint_ids.as_ref().is_none_or(|endpoint_ids| { endpoint_ids .iter() @@ -690,7 +691,7 @@ fn push_key_auth_channel_filter(builder: &mut QueryBuilder<'_, MySql>, api_forma builder.push(" AND LOWER(TRIM(pak.auth_type)) = 'oauth' AND "); builder.push_bind(api_format.to_string()); builder.push( - " IN ('openai:responses', 'openai:responses:compact', 'openai:search', 'openai:image'))", + " IN ('openai:responses', 'openai:responses:compact', 'openai:search', 'openai:image', 'codex:live'))", ); builder.push(" OR (LOWER(TRIM(p.provider_type)) = 'chatgpt_web'"); @@ -749,6 +750,7 @@ fn key_auth_channel_matches(row: &CandidateSelectionRow, api_format: &str) -> bo | "openai:responses:compact" | "openai:search" | "openai:image" + | "codex:live" ) } "chatgpt_web" => { @@ -1096,10 +1098,6 @@ fn api_format_matches(left: &str, right: &str) -> bool { aether_ai_formats::api_format_alias_matches(left, right) } -fn api_format_scope_covers(allowed: &str, requested: &str) -> bool { - aether_ai_formats::api_format_permission_covers(allowed, requested) -} - fn sql_match_aliases(api_formats: &[String]) -> Vec { api_formats .iter() @@ -1111,6 +1109,7 @@ fn sql_match_aliases(api_formats: &[String]) -> Vec { mod tests { use super::{ api_format_page_query, pool_key_group_by_key_ids_query, pool_key_group_query, + provider_model_mapping_api_format_covers, push_key_auth_channel_filter, requested_model_page_query, vertex_key_auth_channel_matches, ExactPageAccumulator, MysqlMinimalCandidateSelectionReadRepository, REQUESTED_MODEL_RAW_SCAN_LIMIT, }; @@ -1147,6 +1146,47 @@ mod tests { assert!(sql.contains("LIMIT ? OFFSET ?")); } + #[test] + fn codex_auth_sql_allows_live_for_oauth_keys() { + let mut builder = sqlx::QueryBuilder::::new("SELECT 1 WHERE 1 = 1"); + push_key_auth_channel_filter(&mut builder, "codex:live"); + let sql = builder.sql(); + let codex_clause = sql + .split_once("LOWER(TRIM(p.provider_type)) = 'codex'") + .and_then(|(_, suffix)| { + suffix.split_once("LOWER(TRIM(p.provider_type)) = 'chatgpt_web'") + }) + .map(|(clause, _)| clause) + .expect("Codex auth clause should exist"); + + assert!(codex_clause.contains("LOWER(TRIM(pak.auth_type)) = 'oauth'")); + assert!(codex_clause.contains("'codex:live'")); + } + + #[test] + fn mysql_mapping_scope_keeps_legacy_responses_compatibility_codex_only() { + assert!(provider_model_mapping_api_format_covers( + "codex", + "openai:responses", + "codex:live" + )); + assert!(!provider_model_mapping_api_format_covers( + "openai", + "openai:responses", + "codex:live" + )); + assert!(!provider_model_mapping_api_format_covers( + "custom", + "openai:responses", + "codex:live" + )); + assert!(!provider_model_mapping_api_format_covers( + "codex", + "openai:chat", + "codex:live" + )); + } + #[test] fn exact_page_accumulator_continues_after_coarse_false_positives() { let mut accumulator = ExactPageAccumulator::new(1, 2); diff --git a/crates/aether-data/adapters/mysql/src/migrations.rs b/crates/aether-data/adapters/mysql/src/migrations.rs index 68400ae74..da8046f5d 100644 --- a/crates/aether-data/adapters/mysql/src/migrations.rs +++ b/crates/aether-data/adapters/mysql/src/migrations.rs @@ -100,6 +100,30 @@ mod tests { assert!(versions.windows(2).all(|pair| pair[0] < pair[1])); } + #[test] + fn embeds_scoped_codex_live_permission_migration() { + let migration = MIGRATOR + .iter() + .find(|migration| migration.version == 20260821000000) + .expect("Codex Live permission migration should be embedded"); + let sql = migration.sql.as_ref(); + + for required_fragment in [ + "UPDATE users", + "UPDATE user_groups", + "UPDATE api_keys", + "UPDATE provider_api_keys", + "provider.provider_type", + "openai:responses", + "codex:live", + ] { + assert!( + sql.contains(required_fragment), + "Codex Live permission migration is missing {required_fragment}" + ); + } + } + #[test] fn embeds_cross_driver_schema_parity_migration() { let migration = MIGRATOR diff --git a/crates/aether-data/adapters/mysql/src/usage/read.rs b/crates/aether-data/adapters/mysql/src/usage/read.rs index bca791beb..9aac4e1d3 100644 --- a/crates/aether-data/adapters/mysql/src/usage/read.rs +++ b/crates/aether-data/adapters/mysql/src/usage/read.rs @@ -583,6 +583,12 @@ AND LOWER(TRIM(COALESCE(`usage`.provider_name, ''))) NOT IN ('unknown', 'unknow' push_where(builder, has_where); builder.push("`usage`.is_stream = ").push_bind(is_stream); } + if let Some(is_websocket) = query.is_websocket { + push_where(builder, has_where); + builder + .push("COALESCE(JSON_UNQUOTE(JSON_EXTRACT(`usage`.request_metadata, '$.websocket_mode')), 'false') = ") + .push_bind(if is_websocket { "true" } else { "false" }); + } if query.error_only { push_where(builder, has_where); builder.push( @@ -613,6 +619,7 @@ fn push_keyword_filters( statuses: query.statuses.clone(), exclude_status_codes: query.exclude_status_codes.clone(), is_stream: query.is_stream, + is_websocket: query.is_websocket, error_only: query.error_only, limit: None, offset: None, @@ -808,4 +815,38 @@ mod tests { .sql() .contains("LOWER(COALESCE(`usage`.api_key_name, '')) LIKE ?")); } + + #[test] + fn websocket_filter_is_applied_to_list_and_keyword_queries() { + let mut list_query = QueryBuilder::::new(USAGE_COLUMNS); + let mut has_where = false; + push_list_filters( + &mut list_query, + &UsageAuditListQuery { + is_websocket: Some(true), + ..UsageAuditListQuery::default() + }, + &mut has_where, + ) + .expect("WebSocket list query should build"); + assert!(list_query + .sql() + .contains("JSON_UNQUOTE(JSON_EXTRACT(`usage`.request_metadata, '$.websocket_mode'))")); + + let mut keyword_query = QueryBuilder::::new(USAGE_COLUMNS); + let mut has_where = false; + push_keyword_filters( + &mut keyword_query, + &UsageAuditKeywordSearchQuery { + is_websocket: Some(true), + keywords: vec!["live".to_string()], + ..UsageAuditKeywordSearchQuery::default() + }, + &mut has_where, + ) + .expect("WebSocket keyword query should build"); + assert!(keyword_query + .sql() + .contains("JSON_UNQUOTE(JSON_EXTRACT(`usage`.request_metadata, '$.websocket_mode'))")); + } } diff --git a/crates/aether-data/adapters/postgres/migrations/20260821000000_migrate_legacy_codex_live_permissions.sql b/crates/aether-data/adapters/postgres/migrations/20260821000000_migrate_legacy_codex_live_permissions.sql new file mode 100644 index 000000000..4309e4a79 --- /dev/null +++ b/crates/aether-data/adapters/postgres/migrations/20260821000000_migrate_legacy_codex_live_permissions.sql @@ -0,0 +1,99 @@ +-- #741 exposed Codex Frameless Live through the openai:responses permission +-- bucket. Codex Live is now a first-class format, so preserve the access that +-- administrators had already granted while keeping the runtime permission +-- relationship strict after this one-time migration. + +UPDATE public.users +SET allowed_api_formats = (allowed_api_formats::jsonb || '["codex:live"]'::jsonb)::json +WHERE LOWER(BTRIM(COALESCE(allowed_api_formats_mode, ''))) = 'specific' + AND json_typeof(allowed_api_formats) = 'array' + AND allowed_api_formats::jsonb ? 'openai:responses' + AND NOT (allowed_api_formats::jsonb ? 'codex:live'); + +UPDATE public.user_groups +SET allowed_api_formats = (allowed_api_formats::jsonb || '["codex:live"]'::jsonb)::json +WHERE LOWER(BTRIM(COALESCE(allowed_api_formats_mode, ''))) = 'specific' + AND json_typeof(allowed_api_formats) = 'array' + AND allowed_api_formats::jsonb ? 'openai:responses' + AND NOT (allowed_api_formats::jsonb ? 'codex:live'); + +UPDATE public.api_keys +SET allowed_api_formats = (allowed_api_formats::jsonb || '["codex:live"]'::jsonb)::json +WHERE json_typeof(allowed_api_formats) = 'array' + AND allowed_api_formats::jsonb ? 'openai:responses' + AND NOT (allowed_api_formats::jsonb ? 'codex:live'); + +UPDATE public.provider_api_keys AS provider_key +SET + api_formats = (provider_key.api_formats::jsonb || '["codex:live"]'::jsonb)::json, + updated_at = NOW() +FROM public.providers AS provider +WHERE provider.id = provider_key.provider_id + AND LOWER(BTRIM(COALESCE(provider.provider_type, ''))) = 'codex' + AND json_typeof(provider_key.api_formats) = 'array' + AND provider_key.api_formats::jsonb ? 'openai:responses' + AND NOT (provider_key.api_formats::jsonb ? 'codex:live'); + +UPDATE public.provider_api_keys AS provider_key +SET + auth_type_by_format = ( + provider_key.auth_type_by_format::jsonb + || jsonb_build_object( + 'codex:live', + provider_key.auth_type_by_format::jsonb -> 'openai:responses' + ) + )::json, + updated_at = NOW() +FROM public.providers AS provider +WHERE provider.id = provider_key.provider_id + AND LOWER(BTRIM(COALESCE(provider.provider_type, ''))) = 'codex' + AND json_typeof(provider_key.auth_type_by_format) = 'object' + AND provider_key.auth_type_by_format::jsonb ? 'openai:responses' + AND NOT (provider_key.auth_type_by_format::jsonb ? 'codex:live'); + +UPDATE public.provider_api_keys AS provider_key +SET + allow_auth_channel_mismatch_formats = ( + provider_key.allow_auth_channel_mismatch_formats::jsonb || '["codex:live"]'::jsonb + )::json, + updated_at = NOW() +FROM public.providers AS provider +WHERE provider.id = provider_key.provider_id + AND LOWER(BTRIM(COALESCE(provider.provider_type, ''))) = 'codex' + AND json_typeof(provider_key.allow_auth_channel_mismatch_formats) = 'array' + AND provider_key.allow_auth_channel_mismatch_formats::jsonb ? 'openai:responses' + AND NOT (provider_key.allow_auth_channel_mismatch_formats::jsonb ? 'codex:live'); + +UPDATE public.provider_api_keys AS provider_key +SET + rate_multipliers = ( + provider_key.rate_multipliers::jsonb + || jsonb_build_object( + 'codex:live', + provider_key.rate_multipliers::jsonb -> 'openai:responses' + ) + )::json, + updated_at = NOW() +FROM public.providers AS provider +WHERE provider.id = provider_key.provider_id + AND LOWER(BTRIM(COALESCE(provider.provider_type, ''))) = 'codex' + AND json_typeof(provider_key.rate_multipliers) = 'object' + AND provider_key.rate_multipliers::jsonb ? 'openai:responses' + AND NOT (provider_key.rate_multipliers::jsonb ? 'codex:live'); + +UPDATE public.provider_api_keys AS provider_key +SET + global_priority_by_format = ( + provider_key.global_priority_by_format::jsonb + || jsonb_build_object( + 'codex:live', + provider_key.global_priority_by_format::jsonb -> 'openai:responses' + ) + )::json, + updated_at = NOW() +FROM public.providers AS provider +WHERE provider.id = provider_key.provider_id + AND LOWER(BTRIM(COALESCE(provider.provider_type, ''))) = 'codex' + AND json_typeof(provider_key.global_priority_by_format) = 'object' + AND provider_key.global_priority_by_format::jsonb ? 'openai:responses' + AND NOT (provider_key.global_priority_by_format::jsonb ? 'codex:live'); diff --git a/crates/aether-data/adapters/postgres/src/candidate_selection.rs b/crates/aether-data/adapters/postgres/src/candidate_selection.rs index 319de03e8..96bed87d1 100644 --- a/crates/aether-data/adapters/postgres/src/candidate_selection.rs +++ b/crates/aether-data/adapters/postgres/src/candidate_selection.rs @@ -73,7 +73,7 @@ INNER JOIN LATERAL ( ( LOWER(BTRIM(p.provider_type)) = 'codex' AND LOWER(BTRIM(pak.auth_type)) = 'oauth' - AND LOWER($3) IN ('openai:responses', 'openai:responses:compact', 'openai:search', 'openai:image') + AND LOWER($3) IN ('openai:responses', 'openai:responses:compact', 'openai:search', 'openai:image', 'codex:live') ) OR ( LOWER(BTRIM(p.provider_type)) = 'chatgpt_web' @@ -158,7 +158,7 @@ WHERE p.is_active = TRUE ( LOWER(BTRIM(p.provider_type)) = 'codex' AND LOWER(BTRIM(pak.auth_type)) = 'oauth' - AND LOWER($3) IN ('openai:responses', 'openai:responses:compact', 'openai:search', 'openai:image') + AND LOWER($3) IN ('openai:responses', 'openai:responses:compact', 'openai:search', 'openai:image', 'codex:live') ) OR ( LOWER(BTRIM(p.provider_type)) = 'chatgpt_web' @@ -336,7 +336,7 @@ INNER JOIN LATERAL ( ( LOWER(BTRIM(p.provider_type)) = 'codex' AND LOWER(BTRIM(pak.auth_type)) = 'oauth' - AND LOWER($4) IN ('openai:responses', 'openai:responses:compact', 'openai:search', 'openai:image') + AND LOWER($4) IN ('openai:responses', 'openai:responses:compact', 'openai:search', 'openai:image', 'codex:live') ) OR ( LOWER(BTRIM(p.provider_type)) = 'chatgpt_web' @@ -422,7 +422,7 @@ WHERE p.is_active = TRUE ( LOWER(BTRIM(p.provider_type)) = 'codex' AND LOWER(BTRIM(pak.auth_type)) = 'oauth' - AND LOWER($4) IN ('openai:responses', 'openai:responses:compact', 'openai:search', 'openai:image') + AND LOWER($4) IN ('openai:responses', 'openai:responses:compact', 'openai:search', 'openai:image', 'codex:live') ) OR ( LOWER(BTRIM(p.provider_type)) = 'chatgpt_web' @@ -608,7 +608,7 @@ WHERE p.is_active = TRUE ( LOWER(BTRIM(p.provider_type)) = 'codex' AND LOWER(BTRIM(pak.auth_type)) = 'oauth' - AND LOWER($6) IN ('openai:responses', 'openai:responses:compact', 'openai:search', 'openai:image') + AND LOWER($6) IN ('openai:responses', 'openai:responses:compact', 'openai:search', 'openai:image', 'codex:live') ) OR ( LOWER(BTRIM(p.provider_type)) = 'chatgpt_web' @@ -971,6 +971,17 @@ impl SqlxMinimalCandidateSelectionReadRepository { } } +const PROVIDER_MODEL_MAPPING_API_FORMAT_MATCH_MARKER: &str = + "__AETHER_PROVIDER_MODEL_MAPPING_API_FORMAT_MATCH__"; +const PROVIDER_MODEL_MAPPING_API_FORMAT_MATCH_SQL: &str = r#"( + LOWER(BTRIM(fmt.value)) = ANY($3::text[]) + OR ( + LOWER(BTRIM(p.provider_type)) = 'codex' + AND LOWER($4) = 'codex:live' + AND LOWER(BTRIM(fmt.value)) = 'openai:responses' + ) +)"#; + fn requested_model_selection_sql() -> String { LIST_FOR_EXACT_API_FORMAT_AND_GLOBAL_MODEL_SQL .replace( @@ -996,7 +1007,7 @@ fn requested_model_selection_sql() -> String { OR EXISTS ( SELECT 1 FROM jsonb_array_elements_text(mapping.value -> 'api_formats') AS fmt(value) - WHERE LOWER(BTRIM(fmt.value)) = ANY($3::text[]) + WHERE __AETHER_PROVIDER_MODEL_MAPPING_API_FORMAT_MATCH__ ) ) AND ( @@ -1054,7 +1065,7 @@ fn requested_model_selection_sql() -> String { OR EXISTS ( SELECT 1 FROM jsonb_array_elements_text(mapping.value -> 'api_formats') AS fmt(value) - WHERE LOWER(BTRIM(fmt.value)) = ANY($3::text[]) + WHERE __AETHER_PROVIDER_MODEL_MAPPING_API_FORMAT_MATCH__ ) ) AND ( @@ -1087,7 +1098,7 @@ fn requested_model_selection_sql() -> String { OR EXISTS ( SELECT 1 FROM jsonb_array_elements_text(mapping.value -> 'api_formats') AS fmt(value) - WHERE LOWER(BTRIM(fmt.value)) = ANY($3::text[]) + WHERE __AETHER_PROVIDER_MODEL_MAPPING_API_FORMAT_MATCH__ ) ) AND ( @@ -1103,6 +1114,10 @@ fn requested_model_selection_sql() -> String { ) )"#, ) + .replace( + PROVIDER_MODEL_MAPPING_API_FORMAT_MATCH_MARKER, + PROVIDER_MODEL_MAPPING_API_FORMAT_MATCH_SQL, + ) .replace( "ORDER BY\n provider_priority ASC,", "ORDER BY\n global_model_name ASC,\n provider_priority ASC,", @@ -1494,7 +1509,8 @@ mod tests { requested_model_selection_page_sql, requested_model_selection_sql, SqlxMinimalCandidateSelectionReadRepository, LIST_FOR_EXACT_API_FORMAT_AND_GLOBAL_MODEL_SQL, LIST_FOR_EXACT_API_FORMAT_SQL, - LIST_POOL_KEYS_FOR_GROUP_SQL, + LIST_POOL_KEYS_FOR_GROUP_SQL, PROVIDER_MODEL_MAPPING_API_FORMAT_MATCH_MARKER, + PROVIDER_MODEL_MAPPING_API_FORMAT_MATCH_SQL, }; use crate::{PostgresPoolConfig, PostgresPoolFactory}; use aether_data_contracts::repository::candidate_selection::{ @@ -1566,6 +1582,36 @@ mod tests { } } + #[test] + fn candidate_selection_sql_allows_codex_live_oauth_auth() { + let requested_model_sql = requested_model_selection_sql(); + for sql in [ + LIST_FOR_EXACT_API_FORMAT_SQL, + LIST_FOR_EXACT_API_FORMAT_AND_GLOBAL_MODEL_SQL, + LIST_POOL_KEYS_FOR_GROUP_SQL, + requested_model_sql.as_str(), + ] { + assert!(sql.contains("LOWER(BTRIM(p.provider_type)) = 'codex'")); + assert!(sql.contains("LOWER(BTRIM(pak.auth_type)) = 'oauth'")); + assert!(sql.contains("'codex:live'")); + } + } + + #[test] + fn requested_model_sql_scopes_legacy_responses_mapping_to_codex_live() { + let sql = requested_model_selection_sql(); + let compatibility = PROVIDER_MODEL_MAPPING_API_FORMAT_MATCH_SQL; + + assert_eq!(sql.matches(compatibility).count(), 3); + assert!(!sql.contains(PROVIDER_MODEL_MAPPING_API_FORMAT_MATCH_MARKER)); + assert!(compatibility.contains("LOWER(BTRIM(p.provider_type)) = 'codex'")); + assert!(compatibility.contains("LOWER($4) = 'codex:live'")); + assert!(compatibility.contains("LOWER(BTRIM(fmt.value)) = 'openai:responses'")); + assert!(!LIST_FOR_EXACT_API_FORMAT_SQL.contains(compatibility)); + assert!(!LIST_FOR_EXACT_API_FORMAT_AND_GLOBAL_MODEL_SQL.contains(compatibility)); + assert!(!LIST_POOL_KEYS_FOR_GROUP_SQL.contains(compatibility)); + } + #[test] fn candidate_selection_sql_allows_grok_oauth_chat_auth() { let requested_model_sql = requested_model_selection_sql(); diff --git a/crates/aether-data/adapters/postgres/src/migrations.rs b/crates/aether-data/adapters/postgres/src/migrations.rs index 6c4759b2f..5ffa3ea9b 100644 --- a/crates/aether-data/adapters/postgres/src/migrations.rs +++ b/crates/aether-data/adapters/postgres/src/migrations.rs @@ -382,6 +382,30 @@ mod tests { assert_eq!(pending_migrations_from_applied(&[]), all_up_migrations()); } + #[test] + fn embeds_scoped_codex_live_permission_migration() { + let migration = POSTGRES_MIGRATOR + .iter() + .find(|migration| migration.version == 20260821000000) + .expect("Codex Live permission migration should be embedded"); + let sql = migration.sql.as_ref(); + + for required_fragment in [ + "UPDATE public.users", + "UPDATE public.user_groups", + "UPDATE public.api_keys", + "UPDATE public.provider_api_keys", + "provider.provider_type", + "openai:responses", + "codex:live", + ] { + assert!( + sql.contains(required_fragment), + "Codex Live permission migration is missing {required_fragment}" + ); + } + } + #[test] fn concurrent_index_migrations_opt_out_of_transactions() { for version in [20260715000000, 20260715130000, 20260720000000] { diff --git a/crates/aether-data/adapters/postgres/src/usage/mod.rs b/crates/aether-data/adapters/postgres/src/usage/mod.rs index a95db63fe..bda75d59c 100644 --- a/crates/aether-data/adapters/postgres/src/usage/mod.rs +++ b/crates/aether-data/adapters/postgres/src/usage/mod.rs @@ -1640,6 +1640,21 @@ fn push_postgres_usage_client_family_filter( .push_bind(client_family.to_ascii_lowercase()); } +fn push_postgres_usage_websocket_filter( + builder: &mut QueryBuilder<'_, Postgres>, + has_where: &mut bool, + is_websocket: Option, +) { + let Some(is_websocket) = is_websocket else { + return; + }; + + push_postgres_usage_where(builder, has_where); + builder + .push("LOWER(COALESCE(\"usage\".request_metadata->>'websocket_mode', 'false')) = ") + .push_bind(if is_websocket { "true" } else { "false" }); +} + fn push_postgres_usage_exclude_unknown_filter( builder: &mut QueryBuilder<'_, Postgres>, has_where: &mut bool, @@ -3015,6 +3030,7 @@ ORDER BY request_count DESC, "usage".provider_name ASC has_where = true; builder.push("\"usage\".is_stream = ").push_bind(is_stream); } + push_postgres_usage_websocket_filter(&mut builder, &mut has_where, query.is_websocket); if query.error_only { builder.push(if has_where { " AND " } else { " WHERE " }); builder.push( @@ -3127,6 +3143,7 @@ OR (\"usage\".error_message IS NOT NULL AND BTRIM(\"usage\".error_message) <> '' has_where = true; builder.push("\"usage\".is_stream = ").push_bind(is_stream); } + push_postgres_usage_websocket_filter(&mut builder, &mut has_where, query.is_websocket); if query.error_only { builder.push(if has_where { " AND " } else { " WHERE " }); has_where = true; @@ -3320,6 +3337,7 @@ OR (\"usage\".error_message IS NOT NULL AND BTRIM(\"usage\".error_message) <> '' has_where = true; builder.push("\"usage\".is_stream = ").push_bind(is_stream); } + push_postgres_usage_websocket_filter(&mut builder, &mut has_where, query.is_websocket); if query.error_only { builder.push(if has_where { " AND " } else { " WHERE " }); builder.push( @@ -3421,6 +3439,7 @@ OR (\"usage\".error_message IS NOT NULL AND BTRIM(\"usage\".error_message) <> '' has_where = true; builder.push("\"usage\".is_stream = ").push_bind(is_stream); } + push_postgres_usage_websocket_filter(&mut builder, &mut has_where, query.is_websocket); if query.error_only { builder.push(if has_where { " AND " } else { " WHERE " }); has_where = true; diff --git a/crates/aether-data/adapters/postgres/src/usage/tests.rs b/crates/aether-data/adapters/postgres/src/usage/tests.rs index 93f87a5be..22d14b690 100644 --- a/crates/aether-data/adapters/postgres/src/usage/tests.rs +++ b/crates/aether-data/adapters/postgres/src/usage/tests.rs @@ -1,6 +1,6 @@ use chrono::{TimeZone, Utc}; use serde_json::json; -use sqlx::Row; +use sqlx::{Postgres, QueryBuilder, Row}; use std::sync::Arc; use super::{ @@ -8,14 +8,15 @@ use super::{ attach_usage_routing_snapshot_metadata, attach_usage_settlement_pricing_snapshot_metadata, clear_previous_request_body_facts, inflate_usage_json_value, prepare_request_metadata_for_body_storage, prepare_usage_body_storage, - prepare_usage_upsert_context, request_body_capture_replaces_derived_facts, - resolved_read_usage_body_ref, resolved_write_usage_body_ref, - split_dashboard_daily_aggregate_range, split_dashboard_hourly_aggregate_range, - usage_body_capture_state_for_storage, usage_body_ref, usage_capture_update_allowed, - usage_effective_input_tokens, usage_http_audit_body_refs, usage_http_audit_capture_mode, - usage_routing_snapshot_from_usage, usage_settlement_pricing_snapshot_from_usage, - usage_total_input_context, AggregateRangeSplit, SqlxUsageReadRepository, UsageHttpAuditRefs, - UsageRoutingSnapshot, UsageSettlementPricingSnapshot, MAX_INLINE_USAGE_BODY_BYTES, + prepare_usage_upsert_context, push_postgres_usage_websocket_filter, + request_body_capture_replaces_derived_facts, resolved_read_usage_body_ref, + resolved_write_usage_body_ref, split_dashboard_daily_aggregate_range, + split_dashboard_hourly_aggregate_range, usage_body_capture_state_for_storage, usage_body_ref, + usage_capture_update_allowed, usage_effective_input_tokens, usage_http_audit_body_refs, + usage_http_audit_capture_mode, usage_routing_snapshot_from_usage, + usage_settlement_pricing_snapshot_from_usage, usage_total_input_context, AggregateRangeSplit, + SqlxUsageReadRepository, UsageHttpAuditRefs, UsageRoutingSnapshot, + UsageSettlementPricingSnapshot, MAX_INLINE_USAGE_BODY_BYTES, SELECT_STALE_PENDING_USAGE_BATCH_SQL, }; use crate::{PostgresPoolConfig, PostgresPoolFactory}; @@ -3247,6 +3248,26 @@ fn usage_sql_admin_record_filters_are_pushed_into_postgres_queries() { assert!(source.contains("request_metadata->>'client_family'")); assert!(source.contains("exclude_unknown_model_or_provider")); assert!(source.contains("NOT IN ('unknown', 'unknow')")); + assert_eq!( + source + .matches("push_postgres_usage_websocket_filter(") + .count(), + 5, + "the WebSocket filter must cover list/count and keyword list/count" + ); +} + +#[test] +fn usage_sql_websocket_filter_compares_json_metadata_without_boolean_casts() { + let mut builder = QueryBuilder::::new("SELECT * FROM usage"); + let mut has_where = false; + push_postgres_usage_websocket_filter(&mut builder, &mut has_where, Some(true)); + + assert!(has_where); + assert!(builder + .sql() + .contains("LOWER(COALESCE(\"usage\".request_metadata->>'websocket_mode', 'false'))")); + assert!(!builder.sql().contains("::boolean")); } #[test] diff --git a/crates/aether-data/adapters/sqlite/migrations/20260821000000_migrate_legacy_codex_live_permissions.sql b/crates/aether-data/adapters/sqlite/migrations/20260821000000_migrate_legacy_codex_live_permissions.sql new file mode 100644 index 000000000..1ba4907fd --- /dev/null +++ b/crates/aether-data/adapters/sqlite/migrations/20260821000000_migrate_legacy_codex_live_permissions.sql @@ -0,0 +1,190 @@ +-- Preserve the explicit access granted by the original #741 implementation, +-- which represented Codex Live as openai:responses. Runtime permission +-- matching remains strict after this one-time data upgrade. + +UPDATE users +SET allowed_api_formats = json_insert(allowed_api_formats, '$[#]', 'codex:live') +WHERE LOWER(TRIM(COALESCE(allowed_api_formats_mode, ''))) = 'specific' + AND json_valid(allowed_api_formats) + AND json_type(CASE WHEN json_valid(allowed_api_formats) THEN allowed_api_formats END) = 'array' + AND EXISTS ( + SELECT 1 FROM json_each( + CASE WHEN json_valid(users.allowed_api_formats) THEN users.allowed_api_formats ELSE '[]' END + ) + WHERE value = 'openai:responses' + ) + AND NOT EXISTS ( + SELECT 1 FROM json_each( + CASE WHEN json_valid(users.allowed_api_formats) THEN users.allowed_api_formats ELSE '[]' END + ) + WHERE value = 'codex:live' + ); + +UPDATE user_groups +SET allowed_api_formats = json_insert(allowed_api_formats, '$[#]', 'codex:live') +WHERE LOWER(TRIM(COALESCE(allowed_api_formats_mode, ''))) = 'specific' + AND json_valid(allowed_api_formats) + AND json_type(CASE WHEN json_valid(allowed_api_formats) THEN allowed_api_formats END) = 'array' + AND EXISTS ( + SELECT 1 FROM json_each( + CASE WHEN json_valid(user_groups.allowed_api_formats) THEN user_groups.allowed_api_formats ELSE '[]' END + ) + WHERE value = 'openai:responses' + ) + AND NOT EXISTS ( + SELECT 1 FROM json_each( + CASE WHEN json_valid(user_groups.allowed_api_formats) THEN user_groups.allowed_api_formats ELSE '[]' END + ) + WHERE value = 'codex:live' + ); + +UPDATE api_keys +SET allowed_api_formats = json_insert(allowed_api_formats, '$[#]', 'codex:live') +WHERE json_valid(allowed_api_formats) + AND json_type(CASE WHEN json_valid(allowed_api_formats) THEN allowed_api_formats END) = 'array' + AND EXISTS ( + SELECT 1 FROM json_each( + CASE WHEN json_valid(api_keys.allowed_api_formats) THEN api_keys.allowed_api_formats ELSE '[]' END + ) + WHERE value = 'openai:responses' + ) + AND NOT EXISTS ( + SELECT 1 FROM json_each( + CASE WHEN json_valid(api_keys.allowed_api_formats) THEN api_keys.allowed_api_formats ELSE '[]' END + ) + WHERE value = 'codex:live' + ); + +UPDATE provider_api_keys +SET + api_formats = json_insert(api_formats, '$[#]', 'codex:live'), + updated_at = CAST(strftime('%s', 'now') AS INTEGER) +WHERE provider_id IN ( + SELECT id FROM providers + WHERE LOWER(TRIM(COALESCE(provider_type, ''))) = 'codex' + ) + AND json_valid(api_formats) + AND json_type(CASE WHEN json_valid(api_formats) THEN api_formats END) = 'array' + AND EXISTS ( + SELECT 1 FROM json_each( + CASE WHEN json_valid(provider_api_keys.api_formats) THEN provider_api_keys.api_formats ELSE '[]' END + ) + WHERE value = 'openai:responses' + ) + AND NOT EXISTS ( + SELECT 1 FROM json_each( + CASE WHEN json_valid(provider_api_keys.api_formats) THEN provider_api_keys.api_formats ELSE '[]' END + ) + WHERE value = 'codex:live' + ); + +UPDATE provider_api_keys +SET + auth_type_by_format = json_set( + auth_type_by_format, + '$."codex:live"', + json_extract(auth_type_by_format, '$."openai:responses"') + ), + updated_at = CAST(strftime('%s', 'now') AS INTEGER) +WHERE provider_id IN ( + SELECT id FROM providers + WHERE LOWER(TRIM(COALESCE(provider_type, ''))) = 'codex' + ) + AND json_valid(auth_type_by_format) + AND json_type(CASE WHEN json_valid(auth_type_by_format) THEN auth_type_by_format END) = 'object' + AND json_type( + CASE WHEN json_valid(auth_type_by_format) THEN auth_type_by_format END, + '$."openai:responses"' + ) IS NOT NULL + AND json_type( + CASE WHEN json_valid(auth_type_by_format) THEN auth_type_by_format END, + '$."codex:live"' + ) IS NULL; + +UPDATE provider_api_keys +SET + allow_auth_channel_mismatch_formats = json_insert( + allow_auth_channel_mismatch_formats, + '$[#]', + 'codex:live' + ), + updated_at = CAST(strftime('%s', 'now') AS INTEGER) +WHERE provider_id IN ( + SELECT id FROM providers + WHERE LOWER(TRIM(COALESCE(provider_type, ''))) = 'codex' + ) + AND json_valid(allow_auth_channel_mismatch_formats) + AND json_type( + CASE + WHEN json_valid(allow_auth_channel_mismatch_formats) + THEN allow_auth_channel_mismatch_formats + END + ) = 'array' + AND EXISTS ( + SELECT 1 FROM json_each( + CASE + WHEN json_valid(provider_api_keys.allow_auth_channel_mismatch_formats) + THEN provider_api_keys.allow_auth_channel_mismatch_formats + ELSE '[]' + END + ) + WHERE value = 'openai:responses' + ) + AND NOT EXISTS ( + SELECT 1 FROM json_each( + CASE + WHEN json_valid(provider_api_keys.allow_auth_channel_mismatch_formats) + THEN provider_api_keys.allow_auth_channel_mismatch_formats + ELSE '[]' + END + ) + WHERE value = 'codex:live' + ); + +UPDATE provider_api_keys +SET + rate_multipliers = json_set( + rate_multipliers, + '$."codex:live"', + json_extract(rate_multipliers, '$."openai:responses"') + ), + updated_at = CAST(strftime('%s', 'now') AS INTEGER) +WHERE provider_id IN ( + SELECT id FROM providers + WHERE LOWER(TRIM(COALESCE(provider_type, ''))) = 'codex' + ) + AND json_valid(rate_multipliers) + AND json_type(CASE WHEN json_valid(rate_multipliers) THEN rate_multipliers END) = 'object' + AND json_type( + CASE WHEN json_valid(rate_multipliers) THEN rate_multipliers END, + '$."openai:responses"' + ) IS NOT NULL + AND json_type( + CASE WHEN json_valid(rate_multipliers) THEN rate_multipliers END, + '$."codex:live"' + ) IS NULL; + +UPDATE provider_api_keys +SET + global_priority_by_format = json_set( + global_priority_by_format, + '$."codex:live"', + json_extract(global_priority_by_format, '$."openai:responses"') + ), + updated_at = CAST(strftime('%s', 'now') AS INTEGER) +WHERE provider_id IN ( + SELECT id FROM providers + WHERE LOWER(TRIM(COALESCE(provider_type, ''))) = 'codex' + ) + AND json_valid(global_priority_by_format) + AND json_type( + CASE WHEN json_valid(global_priority_by_format) THEN global_priority_by_format END + ) = 'object' + AND json_type( + CASE WHEN json_valid(global_priority_by_format) THEN global_priority_by_format END, + '$."openai:responses"' + ) IS NOT NULL + AND json_type( + CASE WHEN json_valid(global_priority_by_format) THEN global_priority_by_format END, + '$."codex:live"' + ) IS NULL; diff --git a/crates/aether-data/adapters/sqlite/src/candidate_selection.rs b/crates/aether-data/adapters/sqlite/src/candidate_selection.rs index e6937aa79..54b5dbc64 100644 --- a/crates/aether-data/adapters/sqlite/src/candidate_selection.rs +++ b/crates/aether-data/adapters/sqlite/src/candidate_selection.rs @@ -4,10 +4,11 @@ use async_trait::async_trait; use sqlx::{sqlite::SqliteRow, QueryBuilder, Row, Sqlite}; use aether_data_contracts::repository::candidate_selection::{ - MinimalCandidateSelectionReadRepository, StoredApiFormatCandidateRowsQuery, - StoredMinimalCandidateSelectionRow, StoredPoolKeyCandidateOrder, - StoredPoolKeyCandidateRowsByKeyIdsQuery, StoredPoolKeyCandidateRowsQuery, - StoredProviderModelMapping, StoredRequestedModelCandidateRowsQuery, + provider_model_mapping_api_format_covers, MinimalCandidateSelectionReadRepository, + StoredApiFormatCandidateRowsQuery, StoredMinimalCandidateSelectionRow, + StoredPoolKeyCandidateOrder, StoredPoolKeyCandidateRowsByKeyIdsQuery, + StoredPoolKeyCandidateRowsQuery, StoredProviderModelMapping, + StoredRequestedModelCandidateRowsQuery, }; use aether_data_contracts::DataLayerError; @@ -605,7 +606,7 @@ fn push_key_auth_channel_sql_filter( ); builder.push_bind(api_format.clone()); builder.push( - r#" IN ('openai:responses', 'openai:responses:compact', 'openai:search', 'openai:image') + r#" IN ('openai:responses', 'openai:responses:compact', 'openai:search', 'openai:image', 'codex:live') ) OR ( LOWER(TRIM(p.provider_type)) = 'chatgpt_web' @@ -945,9 +946,9 @@ fn mapping_scope_matches( api_format: &str, ) -> bool { mapping.api_formats.as_ref().is_none_or(|formats| { - formats - .iter() - .any(|value| api_format_scope_covers(value, api_format)) + formats.iter().any(|value| { + provider_model_mapping_api_format_covers(&row.provider_type, value, api_format) + }) }) && mapping.endpoint_ids.as_ref().is_none_or(|endpoint_ids| { endpoint_ids .iter() @@ -968,6 +969,7 @@ fn key_auth_channel_matches(row: &CandidateSelectionRow, api_format: &str) -> bo | "openai:responses:compact" | "openai:search" | "openai:image" + | "codex:live" ) } "chatgpt_web" => { @@ -1326,10 +1328,6 @@ fn api_format_matches(left: &str, right: &str) -> bool { aether_ai_formats::api_format_alias_matches(left, right) } -fn api_format_scope_covers(allowed: &str, requested: &str) -> bool { - aether_ai_formats::api_format_permission_covers(allowed, requested) -} - fn sql_match_aliases(api_formats: &[String]) -> Vec { api_formats .iter() @@ -1340,9 +1338,9 @@ fn sql_match_aliases(api_formats: &[String]) -> Vec { #[cfg(test)] mod tests { use super::{ - push_key_auth_channel_sql_filter, push_pool_key_order, vertex_key_auth_channel_matches, - ExactPageAccumulator, SqliteMinimalCandidateSelectionReadRepository, - REQUESTED_MODEL_RAW_SCAN_LIMIT, + provider_model_mapping_api_format_covers, push_key_auth_channel_sql_filter, + push_pool_key_order, vertex_key_auth_channel_matches, ExactPageAccumulator, + SqliteMinimalCandidateSelectionReadRepository, REQUESTED_MODEL_RAW_SCAN_LIMIT, }; use crate::run_migrations; use aether_data_contracts::repository::candidate_selection::{ @@ -1380,6 +1378,47 @@ mod tests { assert!(vertex_clause.contains("gemini:embedding")); } + #[test] + fn codex_auth_sql_allows_live_for_oauth_keys() { + let mut builder = sqlx::QueryBuilder::::new("SELECT 1 WHERE 1 = 1"); + push_key_auth_channel_sql_filter(&mut builder, "codex:live"); + let sql = builder.sql(); + let codex_clause = sql + .split_once("LOWER(TRIM(p.provider_type)) = 'codex'") + .and_then(|(_, suffix)| { + suffix.split_once("LOWER(TRIM(p.provider_type)) = 'chatgpt_web'") + }) + .map(|(clause, _)| clause) + .expect("Codex auth clause should exist"); + + assert!(codex_clause.contains("LOWER(TRIM(pak.auth_type)) = 'oauth'")); + assert!(codex_clause.contains("'codex:live'")); + } + + #[test] + fn sqlite_mapping_scope_keeps_legacy_responses_compatibility_codex_only() { + assert!(provider_model_mapping_api_format_covers( + "codex", + "openai:responses", + "codex:live" + )); + assert!(!provider_model_mapping_api_format_covers( + "openai", + "openai:responses", + "codex:live" + )); + assert!(!provider_model_mapping_api_format_covers( + "custom", + "openai:responses", + "codex:live" + )); + assert!(!provider_model_mapping_api_format_covers( + "codex", + "openai:chat", + "codex:live" + )); + } + #[test] fn exact_page_accumulator_continues_after_coarse_false_positives() { let mut accumulator = ExactPageAccumulator::new(1, 2); diff --git a/crates/aether-data/adapters/sqlite/src/migrations.rs b/crates/aether-data/adapters/sqlite/src/migrations.rs index aedab37e9..89c07fd66 100644 --- a/crates/aether-data/adapters/sqlite/src/migrations.rs +++ b/crates/aether-data/adapters/sqlite/src/migrations.rs @@ -678,6 +678,158 @@ ORDER BY id ); } + #[tokio::test] + async fn codex_live_permission_migration_is_scoped_and_idempotent() { + const MIGRATION_VERSION: i64 = 20260821000000; + + let pool = sqlx::sqlite::SqlitePoolOptions::new() + .max_connections(1) + .connect("sqlite::memory:") + .await + .expect("in-memory sqlite pool"); + run_migrations(&pool).await.expect("run sqlite migrations"); + + sqlx::raw_sql( + r#" +INSERT INTO users ( + id, username, allowed_api_formats, allowed_api_formats_mode, created_at, updated_at +) VALUES + ('user-specific', 'user-specific', '["openai:responses"]', 'specific', 1, 1), + ('user-unrestricted', 'user-unrestricted', '["openai:responses"]', 'unrestricted', 1, 1); + +INSERT INTO api_keys ( + id, user_id, key_hash, allowed_api_formats, created_at, updated_at +) VALUES ( + 'api-key-specific', 'user-specific', 'api-key-specific-hash', + '["openai:responses"]', 1, 1 +); + +INSERT INTO user_groups ( + id, name, normalized_name, allowed_api_formats, allowed_api_formats_mode, + created_at, updated_at +) VALUES ( + 'group-specific', 'Specific', 'specific', '["openai:responses"]', 'specific', 1, 1 +); + +INSERT INTO providers (id, name, provider_type, created_at, updated_at) VALUES + ('provider-codex', 'Codex migration fixture', 'codex', 1, 1), + ('provider-openai', 'OpenAI migration fixture', 'openai', 1, 1); + +INSERT INTO provider_api_keys ( + id, provider_id, name, api_formats, auth_type_by_format, + allow_auth_channel_mismatch_formats, rate_multipliers, + global_priority_by_format, created_at, updated_at +) VALUES + ( + 'provider-key-codex', 'provider-codex', 'Codex key', + '["openai:responses"]', '{"openai:responses":"oauth"}', + '["openai:responses"]', '{"openai:responses":1.25}', + '{"openai:responses":17}', 1, 1 + ), + ( + 'provider-key-openai', 'provider-openai', 'OpenAI key', + '["openai:responses"]', '{"openai:responses":"api_key"}', + '["openai:responses"]', '{"openai:responses":2.0}', + '{"openai:responses":23}', 1, 1 + ); +"#, + ) + .execute(&pool) + .await + .expect("legacy Codex Live fixtures should insert"); + + let migration = MIGRATOR + .iter() + .find(|migration| migration.version == MIGRATION_VERSION) + .expect("Codex Live permission migration should be embedded"); + for _ in 0..2 { + sqlx::raw_sql(migration.sql.as_ref()) + .execute(&pool) + .await + .expect("Codex Live permission migration should run idempotently"); + } + + let user_specific: String = + sqlx::query_scalar("SELECT allowed_api_formats FROM users WHERE id = 'user-specific'") + .fetch_one(&pool) + .await + .expect("specific user formats should load"); + assert_eq!( + serde_json::from_str::(&user_specific).unwrap(), + serde_json::json!(["openai:responses", "codex:live"]) + ); + + let user_unrestricted: String = sqlx::query_scalar( + "SELECT allowed_api_formats FROM users WHERE id = 'user-unrestricted'", + ) + .fetch_one(&pool) + .await + .expect("unrestricted user formats should load"); + assert_eq!( + serde_json::from_str::(&user_unrestricted).unwrap(), + serde_json::json!(["openai:responses"]) + ); + + for (table, id_column, id) in [ + ("api_keys", "id", "api-key-specific"), + ("user_groups", "id", "group-specific"), + ] { + let sql = format!("SELECT allowed_api_formats FROM {table} WHERE {id_column} = ?"); + let formats: String = sqlx::query_scalar(sql.as_str()) + .bind(id) + .fetch_one(&pool) + .await + .expect("migrated downstream formats should load"); + assert_eq!( + serde_json::from_str::(&formats).unwrap(), + serde_json::json!(["openai:responses", "codex:live"]) + ); + } + + let codex_key = sqlx::query_as::<_, (String, String, String, String, String)>( + r#" +SELECT api_formats, auth_type_by_format, allow_auth_channel_mismatch_formats, + rate_multipliers, global_priority_by_format +FROM provider_api_keys +WHERE id = 'provider-key-codex' +"#, + ) + .fetch_one(&pool) + .await + .expect("migrated Codex key should load"); + assert_eq!( + serde_json::from_str::(&codex_key.0).unwrap(), + serde_json::json!(["openai:responses", "codex:live"]) + ); + assert_eq!( + serde_json::from_str::(&codex_key.1).unwrap(), + serde_json::json!({"openai:responses": "oauth", "codex:live": "oauth"}) + ); + assert_eq!( + serde_json::from_str::(&codex_key.2).unwrap(), + serde_json::json!(["openai:responses", "codex:live"]) + ); + assert_eq!( + serde_json::from_str::(&codex_key.3).unwrap(), + serde_json::json!({"openai:responses": 1.25, "codex:live": 1.25}) + ); + assert_eq!( + serde_json::from_str::(&codex_key.4).unwrap(), + serde_json::json!({"openai:responses": 17, "codex:live": 17}) + ); + + let openai_formats: String = sqlx::query_scalar( + "SELECT api_formats FROM provider_api_keys WHERE id = 'provider-key-openai'", + ) + .fetch_one(&pool) + .await + .expect("non-Codex key formats should load"); + assert_eq!( + serde_json::from_str::(&openai_formats).unwrap(), + serde_json::json!(["openai:responses"]) + ); + } + #[tokio::test] async fn pending_and_startup_preparation_reject_dirty_migration_state() { let pool = sqlx::sqlite::SqlitePoolOptions::new() diff --git a/crates/aether-data/adapters/sqlite/src/usage.rs b/crates/aether-data/adapters/sqlite/src/usage.rs index 6030a894d..64b4d3534 100644 --- a/crates/aether-data/adapters/sqlite/src/usage.rs +++ b/crates/aether-data/adapters/sqlite/src/usage.rs @@ -885,6 +885,12 @@ AND LOWER(TRIM(COALESCE(provider_name, ''))) NOT IN ('unknown', 'unknow'))", .push("is_stream = ") .push_bind(if is_stream { 1_i64 } else { 0_i64 }); } + if let Some(is_websocket) = query.is_websocket { + push_sqlite_usage_where(builder, has_where); + builder + .push("COALESCE(CAST(json_extract(request_metadata, '$.websocket_mode') AS INTEGER), 0) = ") + .push_bind(if is_websocket { 1_i64 } else { 0_i64 }); + } if query.error_only { push_sqlite_usage_where(builder, has_where); builder.push( @@ -931,6 +937,7 @@ fn push_sqlite_usage_keyword_filters( statuses: query.statuses.clone(), exclude_status_codes: query.exclude_status_codes.clone(), is_stream: query.is_stream, + is_websocket: query.is_websocket, error_only: query.error_only, limit: None, offset: None, diff --git a/crates/aether-data/adapters/sqlite/src/usage/tests.rs b/crates/aether-data/adapters/sqlite/src/usage/tests.rs index b0782b8b6..425da833d 100644 --- a/crates/aether-data/adapters/sqlite/src/usage/tests.rs +++ b/crates/aether-data/adapters/sqlite/src/usage/tests.rs @@ -2,11 +2,12 @@ use super::{SqliteUsageReadRepository, SqliteUsageWriteRepository}; use crate::run_migrations; use aether_data_contracts::repository::usage::{ ProviderApiKeyWindowUsageRequest, UpsertUsageRecord, UsageAuditAggregationGroupBy, - UsageAuditAggregationQuery, UsageAuditListQuery, UsageAuditSummaryQuery, UsageBodyCaptureState, - UsageBreakdownGroupBy, UsageBreakdownSummaryQuery, UsageCleanupExecutionMode, - UsageCleanupTargets, UsageCleanupWindow, UsageDailyHeatmapQuery, - UsageDashboardDailyBreakdownQuery, UsageDashboardSummaryQuery, UsageProviderPerformanceQuery, - UsageReadRepository, UsageTimeSeriesGranularity, UsageWriteRepository, + UsageAuditAggregationQuery, UsageAuditKeywordSearchQuery, UsageAuditListQuery, + UsageAuditSummaryQuery, UsageBodyCaptureState, UsageBreakdownGroupBy, + UsageBreakdownSummaryQuery, UsageCleanupExecutionMode, UsageCleanupTargets, UsageCleanupWindow, + UsageDailyHeatmapQuery, UsageDashboardDailyBreakdownQuery, UsageDashboardSummaryQuery, + UsageProviderPerformanceQuery, UsageReadRepository, UsageTimeSeriesGranularity, + UsageWriteRepository, }; use chrono::{DateTime, Utc}; @@ -1745,6 +1746,105 @@ WHERE request_id = 'request-1' assert_eq!(summary.total_tokens, 10); } +#[tokio::test] +async fn sqlite_usage_websocket_filter_applies_to_list_count_and_keyword_search() { + let pool = sqlx::sqlite::SqlitePoolOptions::new() + .max_connections(1) + .connect("sqlite::memory:") + .await + .expect("sqlite pool should connect"); + run_migrations(&pool) + .await + .expect("sqlite migrations should run"); + seed_stats_targets(&pool).await; + + let writer = SqliteUsageWriteRepository::new(pool.clone()); + writer + .upsert(sample_usage("request-http", "completed", "settled", 1_000)) + .await + .expect("HTTP usage should upsert"); + let mut websocket = sample_usage("request-ws", "completed", "void", 1_001); + websocket.request_metadata = Some(serde_json::json!({ + "websocket_mode": true, + "websocket_transport": "codex_live_direct", + "usage_available": false, + })); + websocket.input_tokens = None; + websocket.output_tokens = None; + websocket.total_tokens = None; + websocket.cache_creation_input_tokens = None; + websocket.cache_creation_ephemeral_5m_input_tokens = None; + websocket.cache_creation_ephemeral_1h_input_tokens = None; + websocket.cache_read_input_tokens = None; + websocket.cache_creation_cost_usd = None; + websocket.cache_read_cost_usd = None; + websocket.total_cost_usd = None; + websocket.actual_total_cost_usd = None; + writer + .upsert(websocket) + .await + .expect("WebSocket usage should upsert"); + + let reader = SqliteUsageReadRepository::new(pool); + let list_query = UsageAuditListQuery { + is_websocket: Some(true), + ..UsageAuditListQuery::default() + }; + let listed = reader + .list_usage_audits(&list_query) + .await + .expect("WebSocket list should load"); + assert_eq!(listed.len(), 1); + assert_eq!(listed[0].request_id, "request-ws"); + assert_eq!( + reader + .count_usage_audits(&list_query) + .await + .expect("WebSocket count should load"), + 1 + ); + + let keyword_query = UsageAuditKeywordSearchQuery { + is_websocket: Some(true), + keywords: vec!["model-1".to_string()], + ..UsageAuditKeywordSearchQuery::default() + }; + let keyword_matches = reader + .list_usage_audits_by_keyword_search(&keyword_query) + .await + .expect("WebSocket keyword list should load"); + assert_eq!(keyword_matches.len(), 1); + assert_eq!(keyword_matches[0].request_id, "request-ws"); + assert_eq!( + reader + .count_usage_audits_by_keyword_search(&keyword_query) + .await + .expect("WebSocket keyword count should load"), + 1 + ); + + let summary = reader + .summarize_usage_audits(&UsageAuditSummaryQuery { + created_from_unix_secs: 0, + created_until_unix_secs: 2_000, + ..UsageAuditSummaryQuery::default() + }) + .await + .expect("lifecycle summary should load"); + assert_eq!(summary.total_requests, 2); + assert_eq!(summary.recorded_total_tokens, 5); + + let provider_key_summaries = reader + .summarize_usage_by_provider_api_key_ids(&["provider-key-1".to_string()]) + .await + .expect("provider key lifecycle summary should load"); + let provider_key_summary = provider_key_summaries + .get("provider-key-1") + .expect("provider key summary"); + assert_eq!(provider_key_summary.request_count, 2); + assert_eq!(provider_key_summary.total_tokens, 5); +} + #[tokio::test] async fn sqlite_usage_daily_heatmap_reads_imported_daily_aggregates() { let pool = sqlx::sqlite::SqlitePoolOptions::new() diff --git a/crates/aether-data/contracts/src/repository/candidate_selection/mod.rs b/crates/aether-data/contracts/src/repository/candidate_selection/mod.rs index bc46fb39d..7b130be47 100644 --- a/crates/aether-data/contracts/src/repository/candidate_selection/mod.rs +++ b/crates/aether-data/contracts/src/repository/candidate_selection/mod.rs @@ -1,9 +1,9 @@ mod types; pub use types::{ - MinimalCandidateSelectionReadRepository, MinimalCandidateSelectionRepository, - StoredApiFormatCandidateRowsQuery, StoredMinimalCandidateSelectionRow, - StoredPoolKeyCandidateOrder, StoredPoolKeyCandidateRowsByKeyIdsQuery, - StoredPoolKeyCandidateRowsQuery, StoredProviderModelMapping, - StoredRequestedModelCandidateRowsQuery, + provider_model_mapping_api_format_covers, MinimalCandidateSelectionReadRepository, + MinimalCandidateSelectionRepository, StoredApiFormatCandidateRowsQuery, + StoredMinimalCandidateSelectionRow, StoredPoolKeyCandidateOrder, + StoredPoolKeyCandidateRowsByKeyIdsQuery, StoredPoolKeyCandidateRowsQuery, + StoredProviderModelMapping, StoredRequestedModelCandidateRowsQuery, }; diff --git a/crates/aether-data/contracts/src/repository/candidate_selection/types.rs b/crates/aether-data/contracts/src/repository/candidate_selection/types.rs index 48e74b3db..d6fde1a61 100644 --- a/crates/aether-data/contracts/src/repository/candidate_selection/types.rs +++ b/crates/aether-data/contracts/src/repository/candidate_selection/types.rs @@ -113,6 +113,27 @@ impl StoredMinimalCandidateSelectionRow { } } +/// Evaluates the API-format scope on a provider-model mapping. +/// +/// Codex Live was introduced after existing Codex model associations had +/// already stored their source-model scope as `openai:responses`. Preserve +/// those associations for the same Codex provider without treating the two +/// formats as globally interchangeable. Endpoint and key permissions remain +/// independently scoped to `codex:live`. +pub fn provider_model_mapping_api_format_covers( + provider_type: &str, + mapping_api_format: &str, + requested_api_format: &str, +) -> bool { + if aether_ai_formats::api_format_permission_covers(mapping_api_format, requested_api_format) { + return true; + } + + provider_type.trim().eq_ignore_ascii_case("codex") + && aether_ai_formats::normalize_api_format_alias(requested_api_format) == "codex:live" + && aether_ai_formats::normalize_api_format_alias(mapping_api_format) == "openai:responses" +} + fn api_format_permission_covers(allowed: &str, requested: &str) -> bool { aether_ai_formats::api_format_permission_covers(allowed, requested) } @@ -176,3 +197,37 @@ impl MinimalCandidateSelectionRepository for T where T: MinimalCandidateSelectionReadRepository + Send + Sync { } + +#[cfg(test)] +mod tests { + use super::provider_model_mapping_api_format_covers; + + #[test] + fn legacy_responses_mapping_is_only_compatible_with_codex_live() { + assert!(provider_model_mapping_api_format_covers( + "codex", + "openai:responses", + "codex:live" + )); + assert!(provider_model_mapping_api_format_covers( + " CoDeX ", + "/v1/responses", + "codex:live" + )); + + for provider_type in ["openai", "custom", "chatgpt_web"] { + assert!(!provider_model_mapping_api_format_covers( + provider_type, + "openai:responses", + "codex:live" + )); + } + for requested_api_format in ["openai:chat", "claude:messages", "openai:image"] { + assert!(!provider_model_mapping_api_format_covers( + "codex", + "openai:responses", + requested_api_format + )); + } + } +} diff --git a/crates/aether-data/contracts/src/repository/usage/mod.rs b/crates/aether-data/contracts/src/repository/usage/mod.rs index f55975e34..00329640d 100644 --- a/crates/aether-data/contracts/src/repository/usage/mod.rs +++ b/crates/aether-data/contracts/src/repository/usage/mod.rs @@ -33,9 +33,11 @@ pub use types::{ UsageLeaderboardGroupBy, UsageLeaderboardQuery, UsageMonitoringErrorCountQuery, UsageMonitoringErrorListQuery, UsagePerformancePercentilesQuery, UsageProviderPerformanceQuery, UsageReadRepository, UsageRepository, UsageSettledCostSummaryQuery, UsageTimeSeriesGranularity, - UsageTimeSeriesQuery, UsageWriteRepository, PROVIDER_ACTUAL_SERVICE_TIER_METADATA_KEY, - PROVIDER_CACHE_TTL_MINUTES_METADATA_KEY, PROVIDER_REASONING_EFFORT_METADATA_KEY, - PROVIDER_SERVICE_TIER_METADATA_KEY, REQUESTED_REASONING_EFFORT_METADATA_KEY, + UsageTimeSeriesQuery, UsageWriteRepository, LIVE_SESSION_METADATA_KEY, + PROVIDER_ACTUAL_SERVICE_TIER_METADATA_KEY, PROVIDER_CACHE_TTL_MINUTES_METADATA_KEY, + PROVIDER_REASONING_EFFORT_METADATA_KEY, PROVIDER_SERVICE_TIER_METADATA_KEY, + REALTIME_SESSION_METADATA_KEY, REQUESTED_REASONING_EFFORT_METADATA_KEY, ROUTING_CANDIDATE_SKIP_REASON_METADATA_KEY, ROUTING_FAILURE_DIAGNOSTIC_METADATA_KEY, + USAGE_AVAILABLE_METADATA_KEY, USAGE_PRICING_AVAILABLE_METADATA_KEY, WEBSOCKET_MODE_METADATA_KEY, WEBSOCKET_TRANSPORT_METADATA_KEY, }; diff --git a/crates/aether-data/contracts/src/repository/usage/policy.rs b/crates/aether-data/contracts/src/repository/usage/policy.rs index e45ba3f0d..c8780a420 100644 --- a/crates/aether-data/contracts/src/repository/usage/policy.rs +++ b/crates/aether-data/contracts/src/repository/usage/policy.rs @@ -254,12 +254,15 @@ pub fn provider_api_key_usage_contribution( request_count: 1, success_count: i64::from(is_success), error_count: i64::from(is_error), - total_tokens: if is_in_flight { + total_tokens: if is_in_flight || !usage.usage_available() { 0 } else { i64::try_from(usage.total_tokens).unwrap_or(i64::MAX) }, - total_cost_usd: if is_in_flight { + total_cost_usd: if is_in_flight + || !usage.usage_available() + || !usage.usage_pricing_available() + { 0.0 } else if usage.total_cost_usd.is_finite() { usage.total_cost_usd.max(0.0) @@ -308,8 +311,14 @@ pub fn api_key_usage_contribution( Some(ApiKeyUsageContribution { api_key_id, total_requests: 1, - total_tokens: i64::try_from(usage.total_tokens).unwrap_or(i64::MAX), - total_cost_usd: if usage.total_cost_usd.is_finite() { + total_tokens: if usage.usage_available() { + i64::try_from(usage.total_tokens).unwrap_or(i64::MAX) + } else { + 0 + }, + total_cost_usd: if !usage.usage_available() || !usage.usage_pricing_available() { + 0.0 + } else if usage.total_cost_usd.is_finite() { usage.total_cost_usd.max(0.0) } else { 0.0 @@ -325,3 +334,80 @@ fn newer_last_used_at(before: Option, after: Option) -> Option { _ => None, } } + +#[cfg(test)] +mod tests { + use serde_json::json; + + use super::{api_key_usage_contribution, provider_api_key_usage_contribution}; + use crate::repository::usage::StoredRequestUsageAudit; + + fn authoritative_unpriced_usage() -> StoredRequestUsageAudit { + let mut usage = StoredRequestUsageAudit::new( + "usage-realtime".to_string(), + "request-realtime".to_string(), + None, + Some("downstream-key".to_string()), + None, + None, + "OpenAI".to_string(), + "gpt-realtime".to_string(), + None, + Some("provider-realtime".to_string()), + Some("endpoint-realtime".to_string()), + Some("provider-key-realtime".to_string()), + Some("realtime".to_string()), + Some("openai:realtime".to_string()), + Some("openai".to_string()), + Some("realtime".to_string()), + Some("openai:realtime".to_string()), + Some("openai".to_string()), + Some("realtime".to_string()), + false, + true, + 120, + 40, + 160, + 9.75, + 11.25, + Some(200), + None, + None, + Some(250), + Some(30), + "completed".to_string(), + "void".to_string(), + 100, + 101, + Some(102), + ) + .expect("usage should build"); + usage.request_metadata = Some(json!({ + "usage_available": true, + "usage_pricing_available": false, + "realtime_session": { + "input_audio_tokens": 20, + "output_audio_tokens": 10, + } + })); + usage + } + + #[test] + fn authoritative_unpriced_usage_contributes_tokens_but_never_cost() { + let usage = authoritative_unpriced_usage(); + + let provider = provider_api_key_usage_contribution(&usage) + .expect("provider key contribution should exist"); + assert_eq!(provider.request_count, 1); + assert_eq!(provider.success_count, 1); + assert_eq!(provider.total_tokens, 160); + assert_eq!(provider.total_cost_usd, 0.0); + + let downstream = + api_key_usage_contribution(&usage).expect("API key contribution should exist"); + assert_eq!(downstream.total_requests, 1); + assert_eq!(downstream.total_tokens, 160); + assert_eq!(downstream.total_cost_usd, 0.0); + } +} diff --git a/crates/aether-data/contracts/src/repository/usage/types.rs b/crates/aether-data/contracts/src/repository/usage/types.rs index 3c2585b8d..f0f1cea61 100644 --- a/crates/aether-data/contracts/src/repository/usage/types.rs +++ b/crates/aether-data/contracts/src/repository/usage/types.rs @@ -11,6 +11,24 @@ pub const ROUTING_CANDIDATE_SKIP_REASON_METADATA_KEY: &str = "routing_candidate_ pub const ROUTING_FAILURE_DIAGNOSTIC_METADATA_KEY: &str = "routing_failure_diagnostic"; pub const WEBSOCKET_MODE_METADATA_KEY: &str = "websocket_mode"; pub const WEBSOCKET_TRANSPORT_METADATA_KEY: &str = "websocket_transport"; +/// Whether token/cost usage is authoritative for this audit row. +/// +/// The field is absent for legacy and normally-metered requests. An explicit +/// `false` marks a transport/session audit whose lifecycle is known while the +/// upstream protocol exposes no trustworthy token/cost usage. Such rows still +/// count as requests and retain status/latency; only token and cost accounting +/// is unavailable. +pub const USAGE_AVAILABLE_METADATA_KEY: &str = "usage_available"; +/// Whether Aether has a compatible pricing model for the authoritative usage +/// dimensions on this row. An explicit `false` keeps token telemetry visible +/// while preventing those tokens from being priced with an incompatible rule. +pub const USAGE_PRICING_AVAILABLE_METADATA_KEY: &str = "usage_pricing_available"; +/// Bounded session-level telemetry for transports that do not expose token +/// usage (for example Codex Live direct/sideband WebSockets). +pub const LIVE_SESSION_METADATA_KEY: &str = "live_session"; +/// Bounded lifecycle and authoritative usage facts for an OpenAI Realtime +/// WebSocket connection. Audio payloads themselves are never stored here. +pub const REALTIME_SESSION_METADATA_KEY: &str = "realtime_session"; pub fn extract_provider_reasoning_effort_from_body(value: Option<&Value>) -> Option { let object = value.and_then(Value::as_object)?; @@ -546,6 +564,40 @@ impl StoredRequestUsageAudit { .unwrap_or(false) } + pub fn websocket_transport(&self) -> Option<&str> { + self.request_metadata_string(WEBSOCKET_TRANSPORT_METADATA_KEY) + } + + /// Returns whether this row may participate in token/cost accounting. + /// Missing metadata is treated as available for backward compatibility. + pub fn usage_available(&self) -> bool { + self.request_metadata_bool(USAGE_AVAILABLE_METADATA_KEY) + .unwrap_or(true) + } + + /// Returns whether token usage can be safely converted into cost. Missing + /// metadata remains eligible for backward compatibility. + pub fn usage_pricing_available(&self) -> bool { + self.request_metadata_bool(USAGE_PRICING_AVAILABLE_METADATA_KEY) + .unwrap_or(true) + } + + pub fn realtime_input_audio_tokens(&self) -> Option { + self.request_metadata_object() + .and_then(|metadata| metadata.get(REALTIME_SESSION_METADATA_KEY)) + .and_then(Value::as_object) + .and_then(|session| session.get("input_audio_tokens")) + .and_then(Value::as_u64) + } + + pub fn realtime_output_audio_tokens(&self) -> Option { + self.request_metadata_object() + .and_then(|metadata| metadata.get(REALTIME_SESSION_METADATA_KEY)) + .and_then(Value::as_object) + .and_then(|session| session.get("output_audio_tokens")) + .and_then(Value::as_u64) + } + fn billing_snapshot_resolved_number(&self, key: &str) -> Option { self.request_metadata_object() .and_then(|metadata| metadata.get("billing_snapshot")) @@ -953,6 +1005,7 @@ pub struct UsageAuditListQuery { pub statuses: Option>, pub exclude_status_codes: Vec, pub is_stream: Option, + pub is_websocket: Option, pub error_only: bool, pub limit: Option, pub offset: Option, @@ -972,6 +1025,7 @@ pub struct UsageAuditKeywordSearchQuery { pub statuses: Option>, pub exclude_status_codes: Vec, pub is_stream: Option, + pub is_websocket: Option, pub error_only: bool, pub keywords: Vec, pub matched_user_ids_by_keyword: Vec>, @@ -2401,8 +2455,9 @@ mod tests { extract_provider_actual_service_tier_from_response, extract_provider_service_tier_from_body, resolve_provider_cache_ttl_minutes, StoredRequestUsageAudit, UpsertUsageRecord, UsageBodyCaptureState, UsageBodyCaptureStorage, - UsageBodyField, UsageProviderPerformanceQuery, WEBSOCKET_MODE_METADATA_KEY, - WEBSOCKET_TRANSPORT_METADATA_KEY, + UsageBodyField, UsageProviderPerformanceQuery, REALTIME_SESSION_METADATA_KEY, + USAGE_AVAILABLE_METADATA_KEY, USAGE_PRICING_AVAILABLE_METADATA_KEY, + WEBSOCKET_MODE_METADATA_KEY, WEBSOCKET_TRANSPORT_METADATA_KEY, }; use serde_json::{json, Value}; @@ -2675,13 +2730,26 @@ mod tests { fn websocket_transport_uses_typed_request_metadata() { let mut usage = sample_usage(); assert!(!usage.is_websocket()); + assert!(usage.usage_available()); + assert!(usage.usage_pricing_available()); usage.request_metadata = Some(json!({ - WEBSOCKET_MODE_METADATA_KEY: true, - WEBSOCKET_TRANSPORT_METADATA_KEY: "responses", + (WEBSOCKET_MODE_METADATA_KEY): true, + (WEBSOCKET_TRANSPORT_METADATA_KEY): "responses", + (USAGE_AVAILABLE_METADATA_KEY): false, + (USAGE_PRICING_AVAILABLE_METADATA_KEY): false, + (REALTIME_SESSION_METADATA_KEY): { + "input_audio_tokens": 7, + "output_audio_tokens": 3, + }, })); assert!(usage.is_websocket()); + assert_eq!(usage.websocket_transport(), Some("responses")); + assert!(!usage.usage_available()); + assert!(!usage.usage_pricing_available()); + assert_eq!(usage.realtime_input_audio_tokens(), Some(7)); + assert_eq!(usage.realtime_output_audio_tokens(), Some(3)); } #[test] diff --git a/crates/aether-data/runtime/src/repository/candidate_selection/memory.rs b/crates/aether-data/runtime/src/repository/candidate_selection/memory.rs index 673e8e384..665803064 100644 --- a/crates/aether-data/runtime/src/repository/candidate_selection/memory.rs +++ b/crates/aether-data/runtime/src/repository/candidate_selection/memory.rs @@ -1,5 +1,6 @@ use std::sync::RwLock; +use aether_data_contracts::repository::candidate_selection::provider_model_mapping_api_format_covers; use async_trait::async_trait; use super::{ @@ -242,9 +243,13 @@ fn row_matches_requested_model( .is_some_and(|mappings| { mappings.iter().any(|mapping| { mapping.api_formats.as_ref().is_none_or(|formats| { - formats - .iter() - .any(|value| api_format_scope_covers(value, api_format)) + formats.iter().any(|value| { + provider_model_mapping_api_format_covers( + &row.provider_type, + value, + api_format, + ) + }) }) && mapping.endpoint_ids.as_ref().is_none_or(|endpoint_ids| { endpoint_ids .iter() @@ -298,9 +303,9 @@ fn mapping_scope_matches( api_format: &str, ) -> bool { mapping.api_formats.as_ref().is_none_or(|formats| { - formats - .iter() - .any(|value| api_format_scope_covers(value, api_format)) + formats.iter().any(|value| { + provider_model_mapping_api_format_covers(&row.provider_type, value, api_format) + }) }) && mapping.endpoint_ids.as_ref().is_none_or(|endpoint_ids| { endpoint_ids .iter() @@ -308,10 +313,6 @@ fn mapping_scope_matches( }) } -fn api_format_scope_covers(allowed: &str, requested: &str) -> bool { - aether_ai_formats::api_format_permission_covers(allowed, requested) -} - fn key_auth_channel_matches(row: &StoredMinimalCandidateSelectionRow, api_format: &str) -> bool { let provider_type = row.provider_type.trim().to_ascii_lowercase(); let auth_type = row.key_auth_type.trim().to_ascii_lowercase(); @@ -325,6 +326,7 @@ fn key_auth_channel_matches(row: &StoredMinimalCandidateSelectionRow, api_format | "openai:responses:compact" | "openai:search" | "openai:image" + | "codex:live" ) } "chatgpt_web" => { @@ -496,6 +498,70 @@ mod tests { ); } + #[tokio::test] + async fn codex_live_accepts_oauth_keys_without_accepting_api_keys() { + let mut oauth = sample_row("codex-live-oauth", "codex:live", "gpt-live", 10); + oauth.provider_type = "codex".to_string(); + oauth.key_auth_type = "oauth".to_string(); + + let mut api_key = oauth.clone(); + api_key.provider_id = "codex-live-api-key".to_string(); + api_key.endpoint_id = "endpoint-codex-live-api-key".to_string(); + api_key.key_id = "key-codex-live-api-key".to_string(); + api_key.model_id = "model-codex-live-api-key".to_string(); + api_key.key_auth_type = "api_key".to_string(); + api_key.provider_priority = 20; + + let repository = + InMemoryMinimalCandidateSelectionReadRepository::seed(vec![api_key, oauth]); + let rows = repository + .list_for_exact_api_format_and_requested_model("codex:live", "gpt-live") + .await + .expect("Codex Live candidate should load"); + + assert_eq!(rows.len(), 1); + assert_eq!(rows[0].provider_id, "codex-live-oauth"); + assert_eq!(rows[0].key_auth_type, "oauth"); + } + + #[tokio::test] + async fn codex_live_reuses_only_codex_legacy_responses_model_mappings() { + let legacy_row = |provider_id: &str, provider_type: &str, auth_type: &str| { + let mut row = sample_row( + provider_id, + "codex:live", + &format!("global-{provider_id}"), + 10, + ); + row.provider_type = provider_type.to_string(); + row.key_auth_type = auth_type.to_string(); + row.key_api_formats = Some(vec!["codex:live".to_string()]); + row.model_provider_model_name = format!("upstream-{provider_id}"); + row.model_provider_model_mappings = Some(vec![StoredProviderModelMapping { + name: "legacy-live-alias".to_string(), + priority: 1, + api_formats: Some(vec!["openai:responses".to_string()]), + endpoint_ids: None, + operations: None, + }]); + row + }; + let repository = InMemoryMinimalCandidateSelectionReadRepository::seed(vec![ + legacy_row("codex-provider", "codex", "oauth"), + legacy_row("openai-provider", "openai", "api_key"), + legacy_row("custom-provider", "custom", "api_key"), + ]); + + let rows = repository + .list_for_exact_api_format_and_requested_model("codex:live", "legacy-live-alias") + .await + .expect("legacy Codex Live mapping should resolve"); + + assert_eq!(rows.len(), 1); + assert_eq!(rows[0].provider_id, "codex-provider"); + assert_eq!(rows[0].provider_type, "codex"); + } + #[tokio::test] async fn includes_grok_oauth_rows_for_chat_models() { let mut row = sample_row( diff --git a/crates/aether-data/runtime/src/repository/usage/memory.rs b/crates/aether-data/runtime/src/repository/usage/memory.rs index 2db8a4133..3134441a8 100644 --- a/crates/aether-data/runtime/src/repository/usage/memory.rs +++ b/crates/aether-data/runtime/src/repository/usage/memory.rs @@ -310,6 +310,11 @@ fn usage_matches_list_query(item: &StoredRequestUsageAudit, query: &UsageAuditLi return false; } } + if let Some(is_websocket) = query.is_websocket { + if item.is_websocket() != is_websocket { + return false; + } + } if query.error_only && item.status != "failed" && item.status_code.unwrap_or_default() < 400 @@ -391,6 +396,11 @@ fn usage_matches_keyword_search_query( return false; } } + if let Some(is_websocket) = query.is_websocket { + if item.is_websocket() != is_websocket { + return false; + } + } if query.error_only && item.status != "failed" && item.status_code.unwrap_or_default() < 400 @@ -2396,7 +2406,9 @@ impl UsageReadRepository for InMemoryUsageReadRepository { continue; } let entry = totals.entry(api_key_id.to_string()).or_insert(0); - *entry = (*entry).saturating_add(item.total_tokens); + if item.usage_available() { + *entry = (*entry).saturating_add(item.total_tokens); + } } Ok(totals) } @@ -2436,7 +2448,9 @@ impl UsageReadRepository for InMemoryUsageReadRepository { ..Default::default() }); entry.request_count = entry.request_count.saturating_add(1); - entry.total_tokens = entry.total_tokens.saturating_add(item.total_tokens); + if item.usage_available() { + entry.total_tokens = entry.total_tokens.saturating_add(item.total_tokens); + } } Ok(totals.into_values().collect()) } @@ -2469,8 +2483,10 @@ impl UsageReadRepository for InMemoryUsageReadRepository { ..StoredProviderApiKeyUsageSummary::default() }); entry.request_count = entry.request_count.saturating_add(1); - entry.total_tokens = entry.total_tokens.saturating_add(item.total_tokens); - entry.total_cost_usd += item.total_cost_usd; + if item.usage_available() { + entry.total_tokens = entry.total_tokens.saturating_add(item.total_tokens); + entry.total_cost_usd += item.total_cost_usd; + } entry.last_used_at_unix_secs = Some( entry .last_used_at_unix_secs @@ -2524,8 +2540,10 @@ impl UsageReadRepository for InMemoryUsageReadRepository { } summary.request_count = summary.request_count.saturating_add(1); - summary.total_tokens = summary.total_tokens.saturating_add(item.total_tokens); - summary.total_cost_usd += item.total_cost_usd; + if item.usage_available() { + summary.total_tokens = summary.total_tokens.saturating_add(item.total_tokens); + summary.total_cost_usd += item.total_cost_usd; + } } summaries.push(summary); @@ -2598,6 +2616,9 @@ impl UsageReadRepository for InMemoryUsageReadRepository { let date_key = dt.date_naive().to_string(); let entry = daily.entry(date_key).or_insert((0, 0, 0.0, 0.0)); entry.0 += 1; + if !item.usage_available() { + continue; + } let cache_creation = if item.cache_creation_input_tokens == 0 && (item.cache_creation_ephemeral_5m_input_tokens + item.cache_creation_ephemeral_1h_input_tokens) diff --git a/crates/aether-data/runtime/src/repository/usage/memory/tests.rs b/crates/aether-data/runtime/src/repository/usage/memory/tests.rs index 1c9f76e08..c4c561fcc 100644 --- a/crates/aether-data/runtime/src/repository/usage/memory/tests.rs +++ b/crates/aether-data/runtime/src/repository/usage/memory/tests.rs @@ -15,7 +15,8 @@ use crate::repository::usage::{ }; use aether_data_contracts::repository::usage::{ usage_body_ref, ProviderApiKeyWindowUsageRequest, UsageAuditAggregationGroupBy, - UsageAuditAggregationQuery, UsageBodyCaptureState, UsageBodyField, UsageDashboardSummaryQuery, + UsageAuditAggregationQuery, UsageAuditKeywordSearchQuery, UsageAuditListQuery, + UsageAuditSummaryQuery, UsageBodyCaptureState, UsageBodyField, UsageDashboardSummaryQuery, UsageLeaderboardGroupBy, UsageLeaderboardQuery, UsageProviderPerformanceQuery, UsageTimeSeriesGranularity, }; @@ -332,6 +333,79 @@ async fn provider_aggregation_skips_unknown_provider_labels() { ); } +#[tokio::test] +async fn unmetered_session_audit_counts_lifecycle_without_token_or_cost_contribution() { + let metered = sample_usage("req-metered", 100); + let mut live = sample_usage("req-live", 200); + live.request_metadata = Some(json!({ + "usage_available": false, + "websocket_mode": true, + "websocket_transport": "codex_live_direct", + })); + live.billing_status = "void".to_string(); + live.input_tokens = 0; + live.output_tokens = 0; + live.total_tokens = 0; + live.cache_creation_input_tokens = 0; + live.cache_creation_ephemeral_5m_input_tokens = 0; + live.cache_creation_ephemeral_1h_input_tokens = 0; + live.cache_read_input_tokens = 0; + live.total_cost_usd = 0.0; + live.actual_total_cost_usd = 0.0; + let repository = InMemoryUsageReadRepository::seed(vec![metered, live]); + + let listed = repository + .list_usage_audits(&UsageAuditListQuery { + created_from_unix_secs: Some(0), + created_until_unix_secs: Some(1_000), + newest_first: true, + ..UsageAuditListQuery::default() + }) + .await + .expect("audit list should succeed"); + assert_eq!(listed.len(), 2); + assert!(!listed + .iter() + .find(|item| item.request_id == "req-live") + .expect("Live row should remain visible") + .usage_available()); + + let aggregate = repository + .aggregate_usage_audits(&UsageAuditAggregationQuery { + created_from_unix_secs: 0, + created_until_unix_secs: 1_000, + group_by: UsageAuditAggregationGroupBy::Model, + limit: 10, + exclude_reserved_provider_labels: false, + }) + .await + .expect("aggregate should succeed"); + assert_eq!(aggregate.len(), 1); + assert_eq!(aggregate[0].request_count, 2); + assert_eq!(aggregate[0].total_tokens, 150); + + let summary = repository + .summarize_usage_audits(&UsageAuditSummaryQuery { + created_from_unix_secs: 0, + created_until_unix_secs: 1_000, + ..UsageAuditSummaryQuery::default() + }) + .await + .expect("summary should succeed"); + assert_eq!(summary.total_requests, 2); + assert_eq!(summary.recorded_total_tokens, 150); + + let provider_key_summaries = repository + .summarize_usage_by_provider_api_key_ids(&["provider-key-1".to_string()]) + .await + .expect("provider key lifecycle summary should succeed"); + let provider_key_summary = provider_key_summaries + .get("provider-key-1") + .expect("provider key summary"); + assert_eq!(provider_key_summary.request_count, 2); + assert_eq!(provider_key_summary.total_tokens, 150); +} + #[tokio::test] async fn aggregation_can_skip_unknown_provider_records_for_model_and_api_format() { let mut unknown = sample_usage("req-unknown-provider", 100); @@ -1910,6 +1984,54 @@ async fn list_usage_audits_applies_second_based_time_filters() { assert_eq!(items[0].request_id, "req-2"); } +#[tokio::test] +async fn usage_audit_websocket_filter_applies_to_list_count_and_keyword_search() { + let mut websocket = sample_usage("req-ws", 2); + websocket.request_metadata = Some(json!({ + "websocket_mode": true, + "websocket_transport": "codex_live_direct", + })); + let repository = + InMemoryUsageReadRepository::seed(vec![sample_usage("req-http", 1), websocket]); + + let list_query = crate::repository::usage::UsageAuditListQuery { + is_websocket: Some(true), + ..Default::default() + }; + let listed = repository + .list_usage_audits(&list_query) + .await + .expect("WebSocket list should succeed"); + assert_eq!(listed.len(), 1); + assert_eq!(listed[0].request_id, "req-ws"); + assert_eq!( + repository + .count_usage_audits(&list_query) + .await + .expect("WebSocket count should succeed"), + 1 + ); + + let keyword_query = UsageAuditKeywordSearchQuery { + is_websocket: Some(true), + keywords: vec!["gpt-4.1".to_string()], + ..Default::default() + }; + let keyword_matches = repository + .list_usage_audits_by_keyword_search(&keyword_query) + .await + .expect("WebSocket keyword list should succeed"); + assert_eq!(keyword_matches.len(), 1); + assert_eq!(keyword_matches[0].request_id, "req-ws"); + assert_eq!( + repository + .count_usage_audits_by_keyword_search(&keyword_query) + .await + .expect("WebSocket keyword count should succeed"), + 1 + ); +} + #[tokio::test] async fn dashboard_and_leaderboard_total_tokens_use_effective_cache_aware_tokens() { let mut item = sample_usage("req-cache-aware-total", 1_711_000_000); diff --git a/crates/aether-gateway/frontdoor/src/middleware/access_log.rs b/crates/aether-gateway/frontdoor/src/middleware/access_log.rs index e7c7d0539..7effde668 100644 --- a/crates/aether-gateway/frontdoor/src/middleware/access_log.rs +++ b/crates/aether-gateway/frontdoor/src/middleware/access_log.rs @@ -248,6 +248,12 @@ mod tests { sanitize_access_log_path("/v1/live/rtc_secret_opaque?token=hidden"), "/v1/live/{call_id}" ); + assert_eq!( + sanitize_access_log_path( + "/v1/realtime?model=gpt-realtime-2.1&api_key=secret&token=hidden" + ), + "/v1/realtime?model=gpt-realtime-2.1" + ); } #[tokio::test(flavor = "current_thread")] diff --git a/crates/aether-provider/transport/src/codex_fingerprint.rs b/crates/aether-provider/transport/src/codex_fingerprint.rs index 16cac5549..32c56ffb3 100644 --- a/crates/aether-provider/transport/src/codex_fingerprint.rs +++ b/crates/aether-provider/transport/src/codex_fingerprint.rs @@ -41,6 +41,8 @@ pub fn apply_codex_oauth_fingerprint_convergence( provider_request_headers: &mut BTreeMap, provider_request_body: &mut Value, ) -> bool { + let is_responses = aether_ai_formats::is_openai_responses_format(provider_api_format); + let is_live = aether_ai_formats::api_format_alias_matches(provider_api_format, "codex:live"); if !transport .provider .provider_type @@ -48,11 +50,12 @@ pub fn apply_codex_oauth_fingerprint_convergence( .eq_ignore_ascii_case("codex") || !transport.key.auth_type.trim().eq_ignore_ascii_case("oauth") || crate::agent_identity::is_codex_agent_identity_transport(transport) - || !aether_ai_formats::is_openai_responses_format(provider_api_format) - || aether_ai_formats::openai_responses_request_operation( - provider_api_format, - provider_request_body, - ) == Some(aether_ai_formats::OPENAI_RESPONSES_OPERATION_COMPACT) + || (!is_responses && !is_live) + || is_responses + && aether_ai_formats::openai_responses_request_operation( + provider_api_format, + provider_request_body, + ) == Some(aether_ai_formats::OPENAI_RESPONSES_OPERATION_COMPACT) || !codex_fingerprint_convergence_enabled( transport.provider.provider_type.as_str(), transport.provider.config.as_ref(), @@ -74,7 +77,12 @@ pub fn apply_codex_oauth_fingerprint_convergence( let fingerprint = resolve_converged_fingerprint(account_seed, original_client_session_id); apply_converged_headers(provider_request_headers, &fingerprint); - apply_converged_client_metadata(provider_request_body, &fingerprint); + // Live uses the converged identity on the WebSocket/call-control headers. + // Its event/session payload is an independent opaque protocol and must not + // receive Responses-only `client_metadata` fields. + if is_responses { + apply_converged_client_metadata(provider_request_body, &fingerprint); + } true } @@ -477,6 +485,27 @@ mod tests { assert_ne!(first.installation_id, other_account.installation_id); } + #[test] + fn live_convergence_sets_the_websocket_identity_without_mutating_the_payload() { + let transport = sample_transport(); + let original_body = json!({"model": "gpt-live", "future_live_field": true}); + let mut body = original_body.clone(); + let mut headers = BTreeMap::new(); + + assert!(apply_codex_oauth_fingerprint_convergence( + &transport, + "codex:live", + Some("client-live-session"), + &mut headers, + &mut body, + )); + + assert_eq!(body, original_body); + assert_eq!(headers.get("x-session-id"), headers.get("thread-id")); + assert!(headers.contains_key("x-codex-installation-id")); + assert!(headers.contains_key("x-codex-window-id")); + } + #[test] fn disabled_or_out_of_scope_requests_are_unchanged() { let mut transport = sample_transport(); diff --git a/crates/aether-provider/transport/src/provider_types.rs b/crates/aether-provider/transport/src/provider_types.rs index 98e4a4c2c..b18f43ecf 100644 --- a/crates/aether-provider/transport/src/provider_types.rs +++ b/crates/aether-provider/transport/src/provider_types.rs @@ -291,7 +291,7 @@ const CLAUDE_CODE_FIXED_PROVIDER_TEMPLATE: FixedProviderTemplate = FixedProvider const CODEX_FIXED_PROVIDER_TEMPLATE: FixedProviderTemplate = FixedProviderTemplate { provider_type: "codex", - version: 1, + version: 2, base_url: "https://chatgpt.com/backend-api/codex", endpoints: &[ FixedProviderEndpointTemplate { @@ -318,6 +318,12 @@ const CODEX_FIXED_PROVIDER_TEMPLATE: FixedProviderTemplate = FixedProviderTempla custom_path: None, config_defaults: EMPTY_ENDPOINT_CONFIG_DEFAULTS, }, + FixedProviderEndpointTemplate { + item_key: "codex:live", + api_format: "codex:live", + custom_path: None, + config_defaults: EMPTY_ENDPOINT_CONFIG_DEFAULTS, + }, ], runtime_policy: CODEX_RUNTIME_POLICY, }; @@ -656,7 +662,7 @@ mod tests { fn codex_fixed_provider_template_includes_codex_companion_endpoints() { let template = fixed_provider_template("codex").expect("codex template should exist"); assert_eq!(template.base_url, "https://chatgpt.com/backend-api/codex"); - assert_eq!(template.version, 1); + assert_eq!(template.version, 2); assert_eq!( template .endpoints @@ -667,7 +673,8 @@ mod tests { "openai:responses", "openai:responses:compact", "openai:search", - "openai:image" + "openai:image", + "codex:live" ] ); @@ -680,6 +687,12 @@ mod tests { fixed_provider_endpoint_template_by_api_format("codex", "openai:search") .expect("codex search endpoint should exist"); assert!(search_template.config_defaults.is_empty()); + + let live_template = fixed_provider_endpoint_template_by_api_format("codex", "codex:live") + .expect("codex Live endpoint should exist"); + assert_eq!(live_template.item_key, "codex:live"); + assert_eq!(live_template.custom_path, None); + assert!(live_template.config_defaults.is_empty()); } #[test] diff --git a/crates/aether-provider/transport/src/request_url/mod.rs b/crates/aether-provider/transport/src/request_url/mod.rs index 3590bb6fa..00d58ef93 100644 --- a/crates/aether-provider/transport/src/request_url/mod.rs +++ b/crates/aether-provider/transport/src/request_url/mod.rs @@ -4,7 +4,7 @@ use std::sync::OnceLock; use aether_ai_formats::ApiOperation; use regex::Regex; use serde_json::Value; -use url::form_urlencoded; +use url::{form_urlencoded, Url}; use crate::antigravity::{ build_antigravity_v1internal_url, is_antigravity_provider_transport, @@ -128,7 +128,10 @@ fn build_transport_request_url_inner( .ends_with("/messages/count_tokens"); let blocked_keys = if normalized_provider_api_format.starts_with("gemini:") || normalized_provider_api_format == "claude:messages" - { + || matches!( + normalized_provider_api_format.as_str(), + "openai:realtime" | "codex:live" + ) { GATEWAY_CREDENTIAL_QUERY_KEYS } else { &[][..] @@ -153,6 +156,11 @@ fn build_transport_request_url_inner( { url = build_claude_messages_url(&url, None); } + let url = if normalized_provider_api_format == "openai:realtime" { + replace_realtime_model_query(url, params.mapped_model?)? + } else { + url + }; return Some(maybe_add_gemini_stream_alt_sse( url, &provider_api_format, @@ -179,6 +187,19 @@ fn build_transport_request_url_inner( &transport.endpoint.base_url, params.request_query, )), + "openai:realtime" => build_passthrough_path_url( + &transport.endpoint.base_url, + "/v1/realtime", + params.request_query, + GATEWAY_CREDENTIAL_QUERY_KEYS, + ) + .and_then(|url| replace_realtime_model_query(url, params.mapped_model?)), + "codex:live" => build_passthrough_path_url( + &transport.endpoint.base_url, + "/live", + params.request_query, + GATEWAY_CREDENTIAL_QUERY_KEYS, + ), "openai:embedding" | "jina:embedding" => { build_provider_embedding_v1_url(&transport.endpoint.base_url, params.request_query) } @@ -228,6 +249,24 @@ fn build_transport_request_url_inner( )) } +fn replace_realtime_model_query(raw_url: String, mapped_model: &str) -> Option { + let mut url = Url::parse(raw_url.as_str()).ok()?; + let retained = url + .query_pairs() + .filter(|(name, _)| !name.eq_ignore_ascii_case("model")) + .map(|(name, value)| (name.into_owned(), value.into_owned())) + .collect::>(); + url.set_query(None); + { + let mut query = url.query_pairs_mut(); + for (name, value) in retained { + query.append_pair(name.as_str(), value.as_str()); + } + query.append_pair("model", mapped_model); + } + Some(url.to_string()) +} + pub fn build_local_openai_chat_upstream_url( transport: &GatewayProviderTransportSnapshot, request_query: Option<&str>, @@ -786,6 +825,118 @@ mod tests { } } + #[test] + fn realtime_url_uses_mapped_model_and_drops_downstream_credentials() { + let transport = sample_transport( + "custom", + "openai:realtime", + "https://api.openai.com/v1", + None, + ); + + let url = build_transport_request_url( + &transport, + TransportRequestUrlParams { + provider_api_format: "openai:realtime", + mapped_model: Some("gpt-realtime-provider"), + upstream_is_stream: true, + request_query: Some("model=client-alias&trace=1&key=downstream-secret"), + kiro_api_region: None, + api_operation: None, + }, + ) + .expect("realtime URL"); + + assert_eq!( + url, + "https://api.openai.com/v1/realtime?trace=1&model=gpt-realtime-provider" + ); + assert!(!url.contains("downstream-secret")); + assert!(!url.contains("client-alias")); + } + + #[test] + fn realtime_custom_path_still_rewrites_the_model_query() { + let transport = sample_transport( + "custom", + "openai:realtime", + "https://voice.example.test/api", + Some("/socket/realtime?deployment=future"), + ); + + let url = build_transport_request_url( + &transport, + TransportRequestUrlParams { + provider_api_format: "openai:realtime", + mapped_model: Some("provider-model"), + upstream_is_stream: true, + request_query: Some("model=client-model&tenant=alpha"), + kiro_api_region: None, + api_operation: None, + }, + ) + .expect("custom realtime URL"); + + assert_eq!( + url, + "https://voice.example.test/api/socket/realtime?deployment=future&tenant=alpha&model=provider-model" + ); + } + + #[test] + fn codex_live_url_uses_the_dedicated_endpoint_without_forwarding_credentials() { + let transport = sample_transport( + "custom", + "codex:live", + "https://voice.example.test/v1", + None, + ); + + let url = build_transport_request_url( + &transport, + TransportRequestUrlParams { + provider_api_format: "codex:live", + mapped_model: Some("provider-model"), + upstream_is_stream: true, + request_query: Some("trace=1&key=downstream-secret"), + kiro_api_region: None, + api_operation: None, + }, + ) + .expect("Codex Live URL"); + + assert_eq!(url, "https://voice.example.test/v1/live?trace=1"); + assert!(!url.contains("downstream-secret")); + } + + #[test] + fn codex_live_custom_path_remains_an_explicit_protocol_endpoint() { + let transport = sample_transport( + "custom", + "codex:live", + "https://voice.example.test/v1", + Some("/socket/live?deployment=future"), + ); + + let url = build_transport_request_url( + &transport, + TransportRequestUrlParams { + provider_api_format: "codex:live", + mapped_model: Some("provider-model"), + upstream_is_stream: true, + request_query: Some("tenant=alpha&key=downstream-secret"), + kiro_api_region: None, + api_operation: None, + }, + ) + .expect("custom Codex Live URL"); + + assert_eq!( + url, + "https://voice.example.test/v1/socket/live?deployment=future&tenant=alpha" + ); + } + #[test] fn uses_vertex_hook_before_custom_path_for_custom_aiplatform_transport() { let transport = sample_transport( diff --git a/crates/aether-provider/transport/src/same_format_provider/mod.rs b/crates/aether-provider/transport/src/same_format_provider/mod.rs index fa6dd1fa9..828b87c2e 100644 --- a/crates/aether-provider/transport/src/same_format_provider/mod.rs +++ b/crates/aether-provider/transport/src/same_format_provider/mod.rs @@ -835,6 +835,8 @@ fn resolve_same_format_standard_direct_auth( ) -> Option<(String, String)> { if aether_ai_formats::api_format_alias_matches(provider_api_format, "openai:embedding") || aether_ai_formats::api_format_alias_matches(provider_api_format, "openai:search") + || aether_ai_formats::api_format_alias_matches(provider_api_format, "openai:realtime") + || aether_ai_formats::api_format_alias_matches(provider_api_format, "codex:live") { resolve_local_openai_bearer_auth(transport) } else { @@ -1431,6 +1433,56 @@ mod tests { ); } + #[test] + fn resolves_openai_realtime_direct_auth_with_bearer_header() { + let mut transport = sample_transport("openai"); + transport.endpoint.api_format = "openai:realtime".to_string(); + transport.key.auth_type = "api_key".to_string(); + let behavior = classify_same_format_provider_request_behavior( + &transport, + SameFormatProviderRequestBehaviorParams { + require_streaming: true, + provider_api_format: "openai:realtime", + report_kind: "openai_realtime_websocket_success", + }, + ); + + assert_eq!( + resolve_same_format_provider_direct_auth( + &behavior, + &transport, + SameFormatProviderFamily::Standard, + "openai:realtime", + ), + Some(("authorization".to_string(), "Bearer secret".to_string())) + ); + } + + #[test] + fn resolves_codex_live_direct_auth_with_bearer_header() { + let mut transport = sample_transport("openai"); + transport.endpoint.api_format = "codex:live".to_string(); + transport.key.auth_type = "api_key".to_string(); + let behavior = classify_same_format_provider_request_behavior( + &transport, + SameFormatProviderRequestBehaviorParams { + require_streaming: true, + provider_api_format: "codex:live", + report_kind: "codex_live_websocket_success", + }, + ); + + assert_eq!( + resolve_same_format_provider_direct_auth( + &behavior, + &transport, + SameFormatProviderFamily::Standard, + "codex:live", + ), + Some(("authorization".to_string(), "Bearer secret".to_string())) + ); + } + #[test] fn keeps_claude_same_format_api_key_on_x_api_key_header() { let mut transport = sample_transport("custom"); diff --git a/crates/aether-usage/runtime/src/record.rs b/crates/aether-usage/runtime/src/record.rs index 697a80204..62b75acba 100644 --- a/crates/aether-usage/runtime/src/record.rs +++ b/crates/aether-usage/runtime/src/record.rs @@ -1,4 +1,6 @@ -use aether_data_contracts::repository::usage::UpsertUsageRecord; +use aether_data_contracts::repository::usage::{ + UpsertUsageRecord, USAGE_AVAILABLE_METADATA_KEY, USAGE_PRICING_AVAILABLE_METADATA_KEY, +}; use aether_data_contracts::DataLayerError; use crate::request_metadata::{ @@ -33,7 +35,8 @@ fn metadata_u64(metadata: Option<&serde_json::Value>, key: &str) -> Option pub fn build_upsert_usage_record_from_event( event: &UsageEvent, ) -> Result { - let (status, billing_status) = lifecycle_status_and_billing(event.event_type); + let (status, billing_status) = + lifecycle_status_and_billing(event.event_type, event.data.request_metadata.as_ref()); let finalized_at_unix_secs = match event.event_type { UsageEventType::Pending | UsageEventType::Streaming => None, UsageEventType::Completed | UsageEventType::Failed | UsageEventType::Cancelled => { @@ -41,6 +44,11 @@ pub fn build_upsert_usage_record_from_event( } }; let mut data = event.data.clone(); + if usage_is_explicitly_unavailable(data.request_metadata.as_ref()) { + clear_unavailable_usage_fields(&mut data); + } else if usage_pricing_is_explicitly_unavailable(data.request_metadata.as_ref()) { + clear_unavailable_pricing_fields(&mut data); + } // Request-derived facts are captured before body capture policy is applied. Do not let a // truncation/disabled placeholder clear those facts while converting the queued event into a // database record. Inline (or ref-loaded) bodies remain authoritative and may clear stale @@ -177,16 +185,62 @@ pub fn build_upsert_usage_record_from_event( }) } -fn lifecycle_status_and_billing(event_type: UsageEventType) -> (&'static str, &'static str) { +fn clear_unavailable_usage_fields(data: &mut crate::UsageEventData) { + data.input_tokens = None; + data.output_tokens = None; + data.total_tokens = None; + data.cache_creation_input_tokens = None; + data.cache_creation_ephemeral_5m_input_tokens = None; + data.cache_creation_ephemeral_1h_input_tokens = None; + data.cache_read_input_tokens = None; + data.cache_creation_cost_usd = None; + data.cache_read_cost_usd = None; + data.total_cost_usd = None; + data.actual_total_cost_usd = None; +} + +fn clear_unavailable_pricing_fields(data: &mut crate::UsageEventData) { + data.cache_creation_cost_usd = None; + data.cache_read_cost_usd = None; + data.total_cost_usd = None; + data.actual_total_cost_usd = None; +} + +fn lifecycle_status_and_billing( + event_type: UsageEventType, + request_metadata: Option<&serde_json::Value>, +) -> (&'static str, &'static str) { match event_type { UsageEventType::Pending => ("pending", "pending"), UsageEventType::Streaming => ("streaming", "pending"), + UsageEventType::Completed + if usage_is_explicitly_unavailable(request_metadata) + || usage_pricing_is_explicitly_unavailable(request_metadata) => + { + ("completed", "void") + } UsageEventType::Completed => ("completed", "pending"), UsageEventType::Failed => ("failed", "void"), UsageEventType::Cancelled => ("cancelled", "void"), } } +fn usage_is_explicitly_unavailable(request_metadata: Option<&serde_json::Value>) -> bool { + request_metadata + .and_then(serde_json::Value::as_object) + .and_then(|metadata| metadata.get(USAGE_AVAILABLE_METADATA_KEY)) + .and_then(serde_json::Value::as_bool) + == Some(false) +} + +fn usage_pricing_is_explicitly_unavailable(request_metadata: Option<&serde_json::Value>) -> bool { + request_metadata + .and_then(serde_json::Value::as_object) + .and_then(|metadata| metadata.get(USAGE_PRICING_AVAILABLE_METADATA_KEY)) + .and_then(serde_json::Value::as_bool) + == Some(false) +} + fn empty_to_none(value: Option) -> Option { value .map(|value| value.trim().to_string()) @@ -437,6 +491,97 @@ mod tests { assert_eq!(record.first_byte_time_ms, Some(50)); } + #[test] + fn completed_unmetered_session_audit_is_void_without_fabricated_usage() { + let record = build_upsert_usage_record_from_event(&UsageEvent { + event_type: UsageEventType::Completed, + request_id: "req-live-session".to_string(), + timestamp_ms: 1_700_000_000_000, + data: UsageEventData { + provider_name: "OpenAI".to_string(), + model: "gpt-live".to_string(), + status_code: Some(200), + input_tokens: Some(100), + output_tokens: Some(20), + total_tokens: Some(120), + total_cost_usd: Some(1.25), + actual_total_cost_usd: Some(1.25), + request_metadata: Some(serde_json::json!({ + "usage_available": false, + "websocket_mode": true, + "websocket_transport": "codex_live_direct", + })), + ..UsageEventData::default() + }, + }) + .expect("record should build"); + + assert_eq!(record.status, "completed"); + assert_eq!(record.billing_status, "void"); + assert_eq!(record.input_tokens, None); + assert_eq!(record.output_tokens, None); + assert_eq!(record.total_tokens, None); + assert_eq!(record.total_cost_usd, None); + assert_eq!(record.actual_total_cost_usd, None); + assert_eq!( + record + .request_metadata + .as_ref() + .and_then(|metadata| metadata.get("usage_available")) + .and_then(serde_json::Value::as_bool), + Some(false) + ); + } + + #[test] + fn completed_authoritative_unpriced_session_is_void_but_keeps_tokens() { + let record = build_upsert_usage_record_from_event(&UsageEvent { + event_type: UsageEventType::Completed, + request_id: "req-realtime-audio".to_string(), + timestamp_ms: 1_700_000_000_000, + data: UsageEventData { + provider_name: "OpenAI".to_string(), + model: "gpt-realtime".to_string(), + status_code: Some(200), + input_tokens: Some(120), + output_tokens: Some(40), + total_tokens: Some(160), + cache_read_input_tokens: Some(30), + cache_creation_cost_usd: Some(0.25), + cache_read_cost_usd: Some(0.1), + total_cost_usd: Some(1.25), + actual_total_cost_usd: Some(1.5), + request_metadata: Some(serde_json::json!({ + "usage_available": true, + "usage_pricing_available": false, + "websocket_mode": true, + "websocket_transport": "openai_realtime", + })), + ..UsageEventData::default() + }, + }) + .expect("record should build"); + + assert_eq!(record.status, "completed"); + assert_eq!(record.billing_status, "void"); + assert_eq!(record.input_tokens, Some(120)); + assert_eq!(record.output_tokens, Some(40)); + assert_eq!(record.total_tokens, Some(160)); + assert_eq!(record.cache_read_input_tokens, Some(30)); + assert_eq!(record.cache_creation_cost_usd, None); + assert_eq!(record.cache_read_cost_usd, None); + assert_eq!(record.total_cost_usd, None); + assert_eq!(record.actual_total_cost_usd, None); + assert_eq!( + record + .request_metadata + .as_ref() + .and_then(|metadata| metadata.get("usage_pricing_available")) + .and_then(serde_json::Value::as_bool), + Some(false) + ); + } + #[test] fn sanitizes_request_metadata_before_building_upsert_record() { let record = build_upsert_usage_record_from_event(&UsageEvent { diff --git a/crates/aether-usage/runtime/src/request_metadata.rs b/crates/aether-usage/runtime/src/request_metadata.rs index 946af70af..7d45d4136 100644 --- a/crates/aether-usage/runtime/src/request_metadata.rs +++ b/crates/aether-usage/runtime/src/request_metadata.rs @@ -7,10 +7,12 @@ use aether_data_contracts::repository::usage::{ extract_provider_actual_service_tier_from_response, extract_provider_reasoning_effort_from_body, extract_provider_service_tier_from_body, normalize_provider_service_tier, resolve_provider_cache_ttl_minutes, UsageBodyCaptureState, - PROVIDER_ACTUAL_SERVICE_TIER_METADATA_KEY, PROVIDER_CACHE_TTL_MINUTES_METADATA_KEY, - PROVIDER_REASONING_EFFORT_METADATA_KEY, PROVIDER_SERVICE_TIER_METADATA_KEY, + LIVE_SESSION_METADATA_KEY, PROVIDER_ACTUAL_SERVICE_TIER_METADATA_KEY, + PROVIDER_CACHE_TTL_MINUTES_METADATA_KEY, PROVIDER_REASONING_EFFORT_METADATA_KEY, + PROVIDER_SERVICE_TIER_METADATA_KEY, REALTIME_SESSION_METADATA_KEY, REQUESTED_REASONING_EFFORT_METADATA_KEY, ROUTING_CANDIDATE_SKIP_REASON_METADATA_KEY, - ROUTING_FAILURE_DIAGNOSTIC_METADATA_KEY, WEBSOCKET_MODE_METADATA_KEY, + ROUTING_FAILURE_DIAGNOSTIC_METADATA_KEY, USAGE_AVAILABLE_METADATA_KEY, + USAGE_PRICING_AVAILABLE_METADATA_KEY, WEBSOCKET_MODE_METADATA_KEY, WEBSOCKET_TRANSPORT_METADATA_KEY, }; use serde_json::{json, Map, Value}; @@ -132,6 +134,12 @@ pub(crate) fn retain_first_byte_request_metadata(value: Option) -> Option | "upstream_is_stream" | "client_session_affinity" | "api_key_is_standalone" + | "websocket_mode" + | "websocket_transport" + | "usage_available" + | "usage_pricing_available" + | "live_session" + | "realtime_session" | "request_path" | "request_query_string" | "request_path_and_query" @@ -351,6 +359,10 @@ fn copy_allowed_metadata_fields(source: &Map, target: &mut Map, target: &mut Map remove_bool(&mut source, target, "api_key_is_standalone"); remove_bool(&mut source, target, WEBSOCKET_MODE_METADATA_KEY); remove_non_empty_string(&mut source, target, WEBSOCKET_TRANSPORT_METADATA_KEY); + remove_bool(&mut source, target, USAGE_AVAILABLE_METADATA_KEY); + remove_bool(&mut source, target, USAGE_PRICING_AVAILABLE_METADATA_KEY); + remove_non_null_value(&mut source, target, LIVE_SESSION_METADATA_KEY); + remove_non_null_value(&mut source, target, REALTIME_SESSION_METADATA_KEY); 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"); diff --git a/crates/aether-usage/runtime/src/write.rs b/crates/aether-usage/runtime/src/write.rs index 78bd2f16f..ba332fdab 100644 --- a/crates/aether-usage/runtime/src/write.rs +++ b/crates/aether-usage/runtime/src/write.rs @@ -3,8 +3,9 @@ use std::collections::BTreeMap; use aether_ai_formats::UPSTREAM_IS_STREAM_KEY; use aether_contracts::{ExecutionPlan, ExecutionTelemetry}; use aether_data_contracts::repository::usage::{ - UpsertUsageRecord, UsageBodyCaptureState, WEBSOCKET_MODE_METADATA_KEY, - WEBSOCKET_TRANSPORT_METADATA_KEY, + UpsertUsageRecord, UsageBodyCaptureState, LIVE_SESSION_METADATA_KEY, + USAGE_AVAILABLE_METADATA_KEY, USAGE_PRICING_AVAILABLE_METADATA_KEY, + WEBSOCKET_MODE_METADATA_KEY, WEBSOCKET_TRANSPORT_METADATA_KEY, }; use aether_data_contracts::DataLayerError; use base64::Engine as _; @@ -2129,6 +2130,21 @@ fn build_runtime_request_metadata_seed_from_parts( Value::String(websocket_transport), ); } + if let Some(usage_available) = context_bool(context, USAGE_AVAILABLE_METADATA_KEY) { + metadata.insert( + USAGE_AVAILABLE_METADATA_KEY.to_string(), + Value::Bool(usage_available), + ); + } + if let Some(pricing_available) = context_bool(context, USAGE_PRICING_AVAILABLE_METADATA_KEY) { + metadata.insert( + USAGE_PRICING_AVAILABLE_METADATA_KEY.to_string(), + Value::Bool(pricing_available), + ); + } + if let Some(live_session) = context_value_ref(context, LIVE_SESSION_METADATA_KEY) { + metadata.insert(LIVE_SESSION_METADATA_KEY.to_string(), live_session.clone()); + } let provider_source_bytes = provider_request_body_base64.and_then(decoded_base64_len_hint); append_runtime_body_capture_metadata( &mut metadata, diff --git a/docs/WebSocket-Mode.md b/docs/WebSocket-Mode.md index 7c7bd9163..93b2ac11e 100644 --- a/docs/WebSocket-Mode.md +++ b/docs/WebSocket-Mode.md @@ -1,9 +1,60 @@ -# WebSocket Mode +# WebSocket transports + +Aether exposes three independent WebSocket surfaces. They share transport +machinery, but not request schemas or continuation state: + +| Public route | API format | Protocol | +| --- | --- | --- | +| `GET /v1/responses` | `openai:responses` | Responses WebSocket mode; every turn starts with `response.create`. | +| `GET /v1/realtime?model=...` | `openai:realtime` | OpenAI Realtime JSON events, including Base64 audio events. | +| `GET /v1/live[/{call_id}]` | `codex:live` | Codex Frameless/Live direct and WebRTC-sideband transport. | + +Do not point one surface at an endpoint configured for another. In particular, +a Realtime or Live event is not passed through the Responses +`response.create` state machine. + +## Responses WebSocket mode The Responses API supports a WebSocket mode for long-running, tool-call-heavy workflows. In this mode, you keep a persistent connection to `/v1/responses` and continue each turn by sending only new input items plus `previous_response_id`. WebSocket mode is compatible with both Zero Data Retention (ZDR) and `store=false`. +## OpenAI Realtime WebSocket bridge + +Configure an active `openai:realtime` provider endpoint, then connect to: + +```text +wss:///v1/realtime?model= +``` + +Aether authenticates and plans the request before returning the downstream +WebSocket upgrade. The global model alias is replaced in the upstream query, +while safe non-credential query parameters, provider authentication, +`header_rules`, and proxy settings continue to apply. Client credentials in +the query string are rejected or removed rather than forwarded upstream. + +After the handshake, Aether relays text, binary, ping, pong, and close frames +one at a time. JSON events, Base64 audio payloads, and unknown future fields are +not rebuilt or coalesced. When an upstream `response.done` contains an +authoritative `response.usage`, its text/audio token counters are accumulated +for the connection's usage record. A session that closes without authoritative +usage is recorded as `usage_available=false`; Aether does not estimate token +counts, audio duration, or cost from frame sizes. + +`response.done` covers Realtime Response usage. Optional input transcription +is reported by a different event and can use a different transcription model; +it is not folded into the Response model's session row or priced as if it used +that model. Finite-balance Realtime access therefore remains fail-closed until +multi-event, multi-model settlement is implemented. + +The upstream handshake is completed before Aether sends HTTP 101 to the +client. A provider authentication, TLS, proxy, or upgrade failure therefore +returns an ordinary bounded HTTP error instead of opening a socket that fails +immediately. + +See the official [OpenAI Realtime WebSocket guide](https://developers.openai.com/api/docs/guides/realtime-websocket) +for the current event contract. + ## Experimental Codex Live bridge Aether also exposes the Codex Frameless Bidi V3 transport used by current @@ -21,16 +72,11 @@ state machine: already initialized call, so Aether neither waits for nor sends a second `session.update` frame. -The provider must expose an `openai:responses` endpoint and explicitly enable -the existing provider-scoped WebSocket capability: - -```json -{ - "responses_websocket": { - "enabled": true - } -} -``` +The provider must expose an active, dedicated `codex:live` endpoint. Fixed +Codex providers receive this endpoint from the managed provider template; +custom providers can add it in the endpoint editor. The +`responses_websocket.enabled` provider option belongs only to +`openai:responses` WebSocket mode and is not reused as the Live permission. API-key and bearer providers can use direct WebSocket or WebRTC. ChatGPT OAuth uses the official Codex backend for WebRTC call creation and the OpenAI Live @@ -47,13 +93,20 @@ abuse; they are not provider-concurrency reservations. Frameless V3 currently has no stable usage object that Aether can settle into its wallet pipeline. Aether therefore enables Live only for principals without a finite `balance_remaining`; finite-balance keys receive an explicit local -error instead of unmetered service. Aether-relayed direct and sideband -WebSocket connections are limited to 60 minutes; the WebRTC media leg itself -does not traverse Aether after call creation. The provider-pool and admission -leases therefore cover only the synchronous HTTP call-creation exchange and -are released after its SDP response. Aether cannot infer media lifetime from -the binding TTL or sideband lifetime, so a created call that never attaches a -sideband is not held against provider concurrency after call creation. +error instead of unmetered service. Aether writes one lifecycle record for each +relayed direct or sideband WebSocket connection, with frame/byte counts and +`usage_available=false`; it does not create one database row per audio frame. +The synchronous WebRTC call-creation exchange keeps its ordinary HTTP record +and is also marked usage-unavailable. The WebRTC media leg itself does not +traverse Aether after call creation, so Aether cannot observe or invent a +separate audio-session usage record, token count, duration, or cost for it. + +Aether-relayed direct and sideband WebSocket connections are limited to 60 +minutes. The provider-pool and admission leases cover only the synchronous +HTTP call-creation exchange and are released after its SDP response. Aether +cannot infer media lifetime from the binding TTL or sideband lifetime, so a +created call that never attaches a sideband is not held against provider +concurrency after call creation. For the public GA Realtime API's connection and session concepts, see the [OpenAI Realtime guide](https://developers.openai.com/api/docs/guides/realtime). diff --git a/frontend/src/api/__tests__/usage-contract.spec.ts b/frontend/src/api/__tests__/usage-contract.spec.ts index bd1b12834..8d2566228 100644 --- a/frontend/src/api/__tests__/usage-contract.spec.ts +++ b/frontend/src/api/__tests__/usage-contract.spec.ts @@ -153,4 +153,29 @@ describe('usageApi contract alignment', () => { timeout: 120000, }) }) + + it('passes the WebSocket usage type filter to the admin records endpoint', async () => { + getMock.mockResolvedValueOnce({ + data: { + records: [], + total: 0, + limit: 20, + offset: 0, + }, + }) + + await usageApi.getAllUsageRecords({ + status: 'websocket', + limit: 20, + offset: 0, + }) + + expect(getMock).toHaveBeenCalledWith('/api/admin/usage/records', { + params: { + status: 'websocket', + limit: 20, + offset: 0, + }, + }) + }) }) diff --git a/frontend/src/api/dashboard.ts b/frontend/src/api/dashboard.ts index e83b5857b..b56cc7175 100644 --- a/frontend/src/api/dashboard.ts +++ b/frontend/src/api/dashboard.ts @@ -251,6 +251,12 @@ export interface RequestDetail { price_per_request?: number // 按次计费价格 request_type: string is_stream: boolean + is_websocket?: boolean + websocket_transport?: string | null + usage_available?: boolean + usage_pricing_available?: boolean + input_audio_tokens?: number | null + output_audio_tokens?: number | null upstream_is_stream?: boolean client_requested_stream?: boolean client_is_stream?: boolean diff --git a/frontend/src/api/endpoints/types/__tests__/api-format.spec.ts b/frontend/src/api/endpoints/types/__tests__/api-format.spec.ts index 74f6e6ade..c3d0e6c86 100644 --- a/frontend/src/api/endpoints/types/__tests__/api-format.spec.ts +++ b/frontend/src/api/endpoints/types/__tests__/api-format.spec.ts @@ -15,6 +15,10 @@ describe('api format display helpers', () => { expect(normalizeApiFormatAlias('CLAUDE_MESSAGES')).toBe(API_FORMATS.CLAUDE_MESSAGES) expect(normalizeApiFormatAlias('OPENAI_RESPONSES')).toBe(API_FORMATS.OPENAI_RESPONSES) expect(normalizeApiFormatAlias('OPENAI_RESPONSES_COMPACT')).toBe(API_FORMATS.OPENAI_RESPONSES_COMPACT) + expect(normalizeApiFormatAlias('OPENAI_REALTIME')).toBe(API_FORMATS.OPENAI_REALTIME) + expect(normalizeApiFormatAlias('REALTIME')).toBe(API_FORMATS.OPENAI_REALTIME) + expect(normalizeApiFormatAlias('CODEX_LIVE')).toBe(API_FORMATS.CODEX_LIVE) + expect(normalizeApiFormatAlias('LIVE')).toBe(API_FORMATS.CODEX_LIVE) expect(normalizeApiFormatAlias('OPENAI_SEARCH')).toBe(API_FORMATS.OPENAI_SEARCH) expect(normalizeApiFormatAlias('SEARCH')).toBe(API_FORMATS.OPENAI_SEARCH) expect(normalizeApiFormatAlias('GEMINI_GENERATE_CONTENT')).toBe(API_FORMATS.GEMINI_GENERATE_CONTENT) @@ -57,6 +61,26 @@ describe('api format display helpers', () => { expect(apiFormatPermissionCovers('openai:responses:compact', 'openai:responses')).toBe(false) }) + it('keeps public Realtime and private Codex Live formats distinct', () => { + expect(formatApiFormat(API_FORMATS.OPENAI_REALTIME)).toBe('OpenAI Realtime') + expect(formatApiFormatShort(API_FORMATS.OPENAI_REALTIME)).toBe('ORT') + expect(formatApiFormat(API_FORMATS.CODEX_LIVE)).toBe('Codex Live') + expect(formatApiFormatShort(API_FORMATS.CODEX_LIVE)).toBe('CL') + expect(apiFormatPermissionCovers('openai:realtime', 'openai:realtime')).toBe(true) + expect(apiFormatPermissionCovers('codex:live', 'codex:live')).toBe(true) + expect(apiFormatPermissionCovers('openai:responses', 'openai:realtime')).toBe(false) + expect(apiFormatPermissionCovers('openai:responses', 'codex:live')).toBe(false) + expect(apiFormatPermissionCovers('openai:realtime', 'codex:live')).toBe(false) + expect(apiFormatPermissionCovers('codex:live', 'openai:realtime')).toBe(false) + expect(groupApiFormats([ + API_FORMATS.CODEX_LIVE, + API_FORMATS.OPENAI_REALTIME, + ])).toEqual([ + { family: 'openai', label: 'OpenAI', formats: [API_FORMATS.OPENAI_REALTIME] }, + { family: 'codex', label: 'Codex', formats: [API_FORMATS.CODEX_LIVE] }, + ]) + }) + it('formats embedding api format ids distinctly from chat formats', () => { expect(formatApiFormat(API_FORMATS.GEMINI_INTERACTIONS)).toBe('Gemini Interactions') expect(formatApiFormatShort(API_FORMATS.GEMINI_INTERACTIONS)).toBe('GI') diff --git a/frontend/src/api/endpoints/types/api-format.ts b/frontend/src/api/endpoints/types/api-format.ts index 177c1be79..6fb0c56df 100644 --- a/frontend/src/api/endpoints/types/api-format.ts +++ b/frontend/src/api/endpoints/types/api-format.ts @@ -6,11 +6,13 @@ export const API_FORMATS = { OPENAI: 'openai:chat', OPENAI_RESPONSES: 'openai:responses', OPENAI_RESPONSES_COMPACT: 'openai:responses:compact', + OPENAI_REALTIME: 'openai:realtime', OPENAI_SEARCH: 'openai:search', OPENAI_IMAGE: 'openai:image', OPENAI_VIDEO: 'openai:video', OPENAI_EMBEDDING: 'openai:embedding', OPENAI_RERANK: 'openai:rerank', + CODEX_LIVE: 'codex:live', GEMINI: 'gemini:generate_content', GEMINI_GENERATE_CONTENT: 'gemini:generate_content', GEMINI_INTERACTIONS: 'gemini:interactions', @@ -31,11 +33,13 @@ export const API_FORMAT_LABELS: Record = { [API_FORMATS.OPENAI]: 'OpenAI Chat', [API_FORMATS.OPENAI_RESPONSES]: 'OpenAI Responses', [API_FORMATS.OPENAI_RESPONSES_COMPACT]: 'OpenAI Responses Compact', + [API_FORMATS.OPENAI_REALTIME]: 'OpenAI Realtime', [API_FORMATS.OPENAI_SEARCH]: 'OpenAI Search', [API_FORMATS.OPENAI_IMAGE]: 'OpenAI Image', [API_FORMATS.OPENAI_VIDEO]: 'OpenAI Video', [API_FORMATS.OPENAI_EMBEDDING]: 'OpenAI Embedding', [API_FORMATS.OPENAI_RERANK]: 'OpenAI Rerank', + [API_FORMATS.CODEX_LIVE]: 'Codex Live', [API_FORMATS.GEMINI_GENERATE_CONTENT]: 'Gemini Generate Content', [API_FORMATS.GEMINI_INTERACTIONS]: 'Gemini Interactions', [API_FORMATS.GEMINI_VIDEO]: 'Gemini Video', @@ -50,11 +54,15 @@ export const API_FORMAT_LABELS: Record = { OPENAI: 'OpenAI Chat', OPENAI_RESPONSES: 'OpenAI Responses', OPENAI_RESPONSES_COMPACT: 'OpenAI Responses Compact', + OPENAI_REALTIME: 'OpenAI Realtime', + REALTIME: 'OpenAI Realtime', OPENAI_SEARCH: 'OpenAI Search', OPENAI_IMAGE: 'OpenAI Image', OPENAI_VIDEO: 'OpenAI Video', OPENAI_EMBEDDING: 'OpenAI Embedding', OPENAI_RERANK: 'OpenAI Rerank', + CODEX_LIVE: 'Codex Live', + LIVE: 'Codex Live', GEMINI: 'Gemini Generate Content', GEMINI_GENERATE_CONTENT: 'Gemini Generate Content', GEMINI_INTERACTIONS: 'Gemini Interactions', @@ -72,11 +80,13 @@ export const API_FORMAT_SHORT: Record = { [API_FORMATS.OPENAI]: 'O', [API_FORMATS.OPENAI_RESPONSES]: 'OR', [API_FORMATS.OPENAI_RESPONSES_COMPACT]: 'ORC', + [API_FORMATS.OPENAI_REALTIME]: 'ORT', [API_FORMATS.OPENAI_SEARCH]: 'OS', [API_FORMATS.OPENAI_IMAGE]: 'OI', [API_FORMATS.OPENAI_VIDEO]: 'OV', [API_FORMATS.OPENAI_EMBEDDING]: 'OE', [API_FORMATS.OPENAI_RERANK]: 'ORR', + [API_FORMATS.CODEX_LIVE]: 'CL', [API_FORMATS.CLAUDE_MESSAGES]: 'CM', [API_FORMATS.GEMINI_GENERATE_CONTENT]: 'G', [API_FORMATS.GEMINI_INTERACTIONS]: 'GI', @@ -90,11 +100,15 @@ export const API_FORMAT_SHORT: Record = { OPENAI: 'O', OPENAI_RESPONSES: 'OR', OPENAI_RESPONSES_COMPACT: 'ORC', + OPENAI_REALTIME: 'ORT', + REALTIME: 'ORT', OPENAI_SEARCH: 'OS', OPENAI_IMAGE: 'OI', OPENAI_VIDEO: 'OV', OPENAI_EMBEDDING: 'OE', OPENAI_RERANK: 'ORR', + CODEX_LIVE: 'CL', + LIVE: 'CL', CLAUDE: 'CM', CLAUDE_MESSAGES: 'CM', GEMINI: 'G', @@ -114,11 +128,13 @@ export const API_FORMAT_ORDER: string[] = [ API_FORMATS.OPENAI, API_FORMATS.OPENAI_RESPONSES, API_FORMATS.OPENAI_RESPONSES_COMPACT, + API_FORMATS.OPENAI_REALTIME, API_FORMATS.OPENAI_SEARCH, API_FORMATS.OPENAI_EMBEDDING, API_FORMATS.OPENAI_RERANK, API_FORMATS.OPENAI_IMAGE, API_FORMATS.OPENAI_VIDEO, + API_FORMATS.CODEX_LIVE, API_FORMATS.CLAUDE_MESSAGES, API_FORMATS.GEMINI_GENERATE_CONTENT, API_FORMATS.GEMINI_INTERACTIONS, @@ -134,6 +150,7 @@ export const API_FORMAT_ORDER: string[] = [ // Family 显示名称映射 export const API_FORMAT_FAMILY_LABELS: Record = { openai: 'OpenAI', + codex: 'Codex', claude: 'Claude', gemini: 'Gemini', jina: 'Jina', @@ -146,6 +163,8 @@ export const API_FORMAT_KIND_LABELS: Record = { chat: 'Chat', responses: 'Responses', 'responses:compact': 'Responses Compact', + realtime: 'Realtime', + live: 'Live', search: 'Search', messages: 'Messages', generate_content: 'Generate Content', @@ -158,7 +177,7 @@ export const API_FORMAT_KIND_LABELS: Record = { } // Family 排序顺序 -const FAMILY_ORDER = ['openai', 'claude', 'gemini', 'jina', 'doubao', 'aliyun'] +const FAMILY_ORDER = ['openai', 'codex', 'claude', 'gemini', 'jina', 'doubao', 'aliyun'] // 工具函数：从 API 格式中提取 family 和 kind export function parseApiFormat(format: string): { family: string; kind: string } { @@ -181,6 +200,9 @@ export function normalizeApiFormatAlias(format: string | null | undefined): stri return API_FORMATS.OPENAI_RESPONSES case 'OPENAI_RESPONSES_COMPACT': return API_FORMATS.OPENAI_RESPONSES_COMPACT + case 'OPENAI_REALTIME': + case 'REALTIME': + return API_FORMATS.OPENAI_REALTIME case 'OPENAI_SEARCH': case 'SEARCH': return API_FORMATS.OPENAI_SEARCH @@ -192,6 +214,9 @@ export function normalizeApiFormatAlias(format: string | null | undefined): stri return API_FORMATS.OPENAI_EMBEDDING case 'OPENAI_RERANK': return API_FORMATS.OPENAI_RERANK + case 'CODEX_LIVE': + case 'LIVE': + return API_FORMATS.CODEX_LIVE case 'GEMINI': case 'GEMINI_GENERATE_CONTENT': return API_FORMATS.GEMINI_GENERATE_CONTENT diff --git a/frontend/src/api/me.ts b/frontend/src/api/me.ts index 108fb8ffc..8999123d6 100644 --- a/frontend/src/api/me.ts +++ b/frontend/src/api/me.ts @@ -74,6 +74,11 @@ export interface UsageRecordDetail { response_time_updated_at?: string | null is_stream: boolean is_websocket?: boolean + websocket_transport?: string | null + usage_available?: boolean + usage_pricing_available?: boolean + input_audio_tokens?: number | null + output_audio_tokens?: number | null upstream_is_stream?: boolean client_requested_stream?: boolean client_is_stream?: boolean @@ -338,6 +343,7 @@ export const meApi = { timezone?: string tz_offset_minutes?: number search?: string // 通用搜索：密钥名、模型名 + status?: string limit?: number offset?: number }): Promise { @@ -372,6 +378,11 @@ export const meApi = { endpoint_api_format?: string | null is_stream?: boolean | null is_websocket?: boolean | null + websocket_transport?: string | null + usage_available?: boolean | null + usage_pricing_available?: boolean | null + input_audio_tokens?: number | null + output_audio_tokens?: number | null upstream_is_stream?: boolean | null client_requested_stream?: boolean | null client_is_stream?: boolean | null diff --git a/frontend/src/api/usage.ts b/frontend/src/api/usage.ts index 0028146e7..fbd3b20ed 100644 --- a/frontend/src/api/usage.ts +++ b/frontend/src/api/usage.ts @@ -34,6 +34,11 @@ export interface UsageRecord { end_to_end_time_ms?: number | null end_to_end_first_byte_time_ms?: number | null is_websocket?: boolean + websocket_transport?: string | null + usage_available?: boolean + usage_pricing_available?: boolean + input_audio_tokens?: number | null + output_audio_tokens?: number | null created_at: string updated_at?: string | null response_time_updated_at?: string | null @@ -139,6 +144,21 @@ export interface UsageRequestOptions { skipCache?: boolean } +export type UsageRecordStatusFilter = + | 'stream' + | 'standard' + | 'websocket' + | 'ws' + | 'active' + | 'pending' + | 'streaming' + | 'completed' + | 'failed' + | 'error' + | 'cancelled' + | 'has_fallback' + | 'has_retry' + type UsageListResponse = { records?: unknown pagination?: { @@ -483,7 +503,7 @@ export const usageApi = { model?: string provider?: string api_format?: string // API 格式筛选（如 openai:chat, claude:messages） - status?: string // 'stream' | 'standard' | 'error' + status?: UsageRecordStatusFilter client_family?: string hide_unknown?: boolean include_total?: boolean @@ -516,7 +536,7 @@ export const usageApi = { model?: string provider?: string api_format?: string - status?: string + status?: UsageRecordStatusFilter client_family?: string hide_unknown?: boolean }): Promise { @@ -572,6 +592,11 @@ export const usageApi = { endpoint_api_format?: string | null is_stream?: boolean | null is_websocket?: boolean | null + websocket_transport?: string | null + usage_available?: boolean | null + usage_pricing_available?: boolean | null + input_audio_tokens?: number | null + output_audio_tokens?: number | null upstream_is_stream?: boolean | null client_requested_stream?: boolean | null client_is_stream?: boolean | null diff --git a/frontend/src/features/providers/components/EndpointFormDialog.vue b/frontend/src/features/providers/components/EndpointFormDialog.vue index f7a9f344f..a1dc646a2 100644 --- a/frontend/src/features/providers/components/EndpointFormDialog.vue +++ b/frontend/src/features/providers/components/EndpointFormDialog.vue @@ -42,13 +42,17 @@