diff --git a/apps/aether-gateway/src/ai_serving/api.rs b/apps/aether-gateway/src/ai_serving/api.rs index a809905df..9ff10edec 100644 --- a/apps/aether-gateway/src/ai_serving/api.rs +++ b/apps/aether-gateway/src/ai_serving/api.rs @@ -50,15 +50,16 @@ pub(crate) use aether_ai_formats::api::{ resolve_claude_stream_spec, resolve_claude_sync_spec, resolve_gemini_stream_spec, resolve_gemini_sync_spec, resolve_local_image_stream_spec, resolve_local_image_sync_spec, resolve_local_same_format_stream_spec, resolve_local_same_format_sync_spec, - resolve_openai_embedding_sync_spec, sanitize_request_path_and_query, AiControlPlanRequest, - CanonicalContentPart, CanonicalStreamEvent, CanonicalStreamFrame, ClaudeClientEmitter, - ExecutionRuntimeAuthContext, LocalCoreSyncErrorKind, LocalOpenAiImageSpec, - LocalSameFormatProviderFamily, LocalSameFormatProviderSpec, LocalStandardSourceFamily, - LocalStandardSourceMode, LocalStandardSpec, OpenAIChatClientEmitter, - OpenAIResponsesClientEmitter, StreamingStandardTerminalObserver, CLAUDE_CHAT_STREAM_PLAN_KIND, - CLAUDE_CLI_STREAM_PLAN_KIND, EXECUTION_RUNTIME_STREAM_DECISION_ACTION, - EXECUTION_RUNTIME_SYNC_DECISION_ACTION, GEMINI_CHAT_STREAM_PLAN_KIND, - GEMINI_CLI_STREAM_PLAN_KIND, GEMINI_EMBEDDING_SYNC_PLAN_KIND, GEMINI_FILES_DOWNLOAD_PLAN_KIND, + resolve_openai_embedding_sync_spec, sanitize_request_path_and_query, + sanitize_request_query_string, AiControlPlanRequest, CanonicalContentPart, + CanonicalStreamEvent, CanonicalStreamFrame, ClaudeClientEmitter, ExecutionRuntimeAuthContext, + LocalCoreSyncErrorKind, LocalOpenAiImageSpec, LocalSameFormatProviderFamily, + LocalSameFormatProviderSpec, LocalStandardSourceFamily, LocalStandardSourceMode, + LocalStandardSpec, OpenAIChatClientEmitter, OpenAIResponsesClientEmitter, + StreamingStandardTerminalObserver, CLAUDE_CHAT_STREAM_PLAN_KIND, CLAUDE_CLI_STREAM_PLAN_KIND, + EXECUTION_RUNTIME_STREAM_DECISION_ACTION, EXECUTION_RUNTIME_SYNC_DECISION_ACTION, + GEMINI_CHAT_STREAM_PLAN_KIND, GEMINI_CLI_STREAM_PLAN_KIND, GEMINI_EMBEDDING_SYNC_PLAN_KIND, + GEMINI_FILES_DOWNLOAD_PLAN_KIND, GEMINI_INTERACTIONS_STREAM_PLAN_KIND, GEMINI_VIDEO_CANCEL_SYNC_PLAN_KIND, OPENAI_CHAT_STREAM_PLAN_KIND, OPENAI_EMBEDDING_SYNC_PLAN_KIND, OPENAI_IMAGE_STREAM_PLAN_KIND, OPENAI_IMAGE_SYNC_FINALIZE_REPORT_KIND, OPENAI_IMAGE_SYNC_PLAN_KIND, @@ -68,10 +69,10 @@ pub(crate) use aether_ai_formats::api::{ OPENAI_VIDEO_REMIX_SYNC_PLAN_KIND, }; pub(crate) use aether_ai_formats::protocol::stream::CanonicalUsage as StreamingCanonicalUsage; -pub(crate) use aether_ai_formats::CODEX_RESPONSES_LITE_HEADER; /// Codex client identity headers re-exported for out-of-crate probe binaries, /// which must reach `aether_ai_formats` through this seam. pub use aether_ai_formats::{CODEX_CLIENT_ORIGINATOR, CODEX_CLIENT_USER_AGENT}; +pub(crate) use aether_ai_formats::{CODEX_RESPONSES_LITE_HEADER, UPSTREAM_IS_STREAM_KEY}; pub(crate) fn parse_direct_request_body( parts: &http::request::Parts, diff --git a/apps/aether-gateway/src/ai_serving/planner/decision_input.rs b/apps/aether-gateway/src/ai_serving/planner/decision_input.rs index 6c4490977..93fbdf708 100644 --- a/apps/aether-gateway/src/ai_serving/planner/decision_input.rs +++ b/apps/aether-gateway/src/ai_serving/planner/decision_input.rs @@ -486,8 +486,7 @@ pub(crate) async fn attach_routing_policy_to_local_requested_model_input( body_json: &Value, client_api_format: &str, ) -> Result<(), GatewayError> { - input.original_client_session_id = routing_header_value_str(&parts.headers, "session-id") - .or_else(|| routing_header_value_str(&parts.headers, "session_id")); + input.original_client_session_id = original_client_session_id_from_headers(&parts.headers); let explicit_group = routing_header_value_str(&parts.headers, ROUTING_GROUP_HEADER); let selected_group = match state.routing_group_read_repository() { Some(repository) => { @@ -738,6 +737,12 @@ pub(crate) async fn attach_routing_policy_to_local_requested_model_input( Ok(()) } +fn original_client_session_id_from_headers(headers: &HeaderMap) -> Option { + routing_header_value_str(headers, "session-id") + .or_else(|| routing_header_value_str(headers, "session_id")) + .or_else(|| routing_header_value_str(headers, "x-session-id")) +} + fn try_attach_static_default_routing_policy_to_input( input: &mut LocalRequestedModelDecisionInput, parts: &http::request::Parts, @@ -1101,6 +1106,38 @@ mod tests { GatewayProviderTransportProvider, }; + #[test] + fn original_client_session_id_accepts_live_header_as_fallback() { + let headers = HeaderMap::from_iter([( + HeaderName::from_static("x-session-id"), + HeaderValue::from_static("live-thread-1"), + )]); + + assert_eq!( + original_client_session_id_from_headers(&headers).as_deref(), + Some("live-thread-1") + ); + } + + #[test] + fn original_client_session_id_prefers_responses_headers_over_live_fallback() { + let headers = HeaderMap::from_iter([ + ( + HeaderName::from_static("session-id"), + HeaderValue::from_static("responses-session"), + ), + ( + HeaderName::from_static("x-session-id"), + HeaderValue::from_static("live-thread"), + ), + ]); + + assert_eq!( + original_client_session_id_from_headers(&headers).as_deref(), + Some("responses-session") + ); + } + #[test] fn explicit_routing_selection_cache_key_is_principal_specific() { let first = routing_group_selection_cache_key( diff --git a/apps/aether-gateway/src/ai_serving/planner/redaction.rs b/apps/aether-gateway/src/ai_serving/planner/redaction.rs index b66ef9ef8..b57b13c9f 100644 --- a/apps/aether-gateway/src/ai_serving/planner/redaction.rs +++ b/apps/aether-gateway/src/ai_serving/planner/redaction.rs @@ -61,6 +61,34 @@ pub(crate) fn request_identity_response_encoding_when_redacted( } } +/// Removes credential-bearing URL components before attaching an upstream URL +/// to a diagnostic event. Endpoint query parameters remain untouched on the +/// wire, but they can contain API keys or signed tokens and must not reach +/// logs. +pub(crate) fn sanitize_upstream_url_for_log(raw: &str) -> String { + if let Ok(mut url) = url::Url::parse(raw) { + if !matches!(url.scheme(), "http" | "https") || url.host_str().is_none() { + return "".to_string(); + } + let _ = url.set_username(""); + let _ = url.set_password(None); + url.set_query(None); + url.set_fragment(None); + return url.to_string(); + } + + let suffix_offset = raw + .char_indices() + .find_map(|(offset, character)| matches!(character, '?' | '#').then_some(offset)) + .unwrap_or(raw.len()); + let path = &raw[..suffix_offset]; + if path.starts_with('/') && !path.starts_with("//") && !path.contains('@') { + path.to_string() + } else { + "".to_string() + } +} + pub(crate) async fn resolve_provider_chat_pii_redaction<'a>( state: &AppState, parts: &http::request::Parts, @@ -240,7 +268,32 @@ fn redaction_mask_error_to_gateway_error(error: RedactionMaskError) -> GatewayEr mod tests { use serde_json::json; - use super::ChatPiiRedactionFeatureSettings; + use super::{sanitize_upstream_url_for_log, ChatPiiRedactionFeatureSettings}; + + #[test] + fn upstream_url_log_projection_removes_all_credential_carriers() { + assert_eq!( + sanitize_upstream_url_for_log( + "https://user:password@api.example.test/v1/responses?api-version=2026-08-01&token=secret#fragment" + ), + "https://api.example.test/v1/responses" + ); + assert_eq!( + sanitize_upstream_url_for_log("/v1/responses?key=secret#fragment"), + "/v1/responses" + ); + for invalid in [ + "https://user:secret@invalid host/v1/responses", + "//user:secret@api.example.test/v1/responses?token=hidden", + "not-a-url?token=hidden", + "data:text/plain,secret", + ] { + assert_eq!( + sanitize_upstream_url_for_log(invalid), + "" + ); + } + } #[test] fn chat_pii_redaction_feature_settings_only_control_enablement() { diff --git a/apps/aether-gateway/src/ai_serving/planner/standard/openai/plan_builders/stream.rs b/apps/aether-gateway/src/ai_serving/planner/standard/openai/plan_builders/stream.rs index e34c0131d..bbda7d62a 100644 --- a/apps/aether-gateway/src/ai_serving/planner/standard/openai/plan_builders/stream.rs +++ b/apps/aether-gateway/src/ai_serving/planner/standard/openai/plan_builders/stream.rs @@ -10,6 +10,7 @@ use super::super::{ AiStreamAttempt, }; use crate::ai_serving::planner::common::enforce_provider_body_stream_policy; +use crate::ai_serving::planner::redaction::sanitize_upstream_url_for_log; use crate::ai_serving::provider_adaptation_requires_eventstream_accept; use crate::ai_serving::transport::{ build_standard_plan_fallback_headers, build_standard_plan_fallback_openai_chat_url, @@ -233,6 +234,19 @@ pub(crate) fn build_openai_responses_stream_plan_from_decision( }, ); + let log_downstream_query = parts + .uri + .query() + .and_then(crate::ai_serving::api::sanitize_request_query_string); + let log_decision_upstream_base_url = payload + .upstream_base_url + .as_deref() + .map(sanitize_upstream_url_for_log); + let log_decision_upstream_url = payload + .upstream_url + .as_deref() + .map(sanitize_upstream_url_for_log); + let log_plan_url = sanitize_upstream_url_for_log(plan.url.as_str()); debug!( event_name = "local_openai_responses_stream_plan_built", log_type = "debug", @@ -242,11 +256,11 @@ pub(crate) fn build_openai_responses_stream_plan_from_decision( endpoint_id = %plan.endpoint_id, key_id = %plan.key_id, downstream_path = %parts.uri.path(), - downstream_query = ?parts.uri.query(), + downstream_query = ?log_downstream_query, url_source, - decision_upstream_base_url = ?payload.upstream_base_url, - decision_upstream_url = ?payload.upstream_url, - plan_url = %plan.url, + decision_upstream_base_url = ?log_decision_upstream_base_url, + decision_upstream_url = ?log_decision_upstream_url, + plan_url = %log_plan_url, client_api_format = %plan.client_api_format, provider_api_format = %plan.provider_api_format, upstream_is_stream = effective_upstream_is_stream, diff --git a/apps/aether-gateway/src/ai_serving/planner/standard/openai/plan_builders/sync.rs b/apps/aether-gateway/src/ai_serving/planner/standard/openai/plan_builders/sync.rs index f1047ad18..44882aa11 100644 --- a/apps/aether-gateway/src/ai_serving/planner/standard/openai/plan_builders/sync.rs +++ b/apps/aether-gateway/src/ai_serving/planner/standard/openai/plan_builders/sync.rs @@ -10,6 +10,7 @@ use super::super::{ AiSyncAttempt, }; use crate::ai_serving::planner::common::enforce_provider_body_stream_policy; +use crate::ai_serving::planner::redaction::sanitize_upstream_url_for_log; use crate::ai_serving::transport::{ build_standard_plan_fallback_headers, build_standard_plan_fallback_openai_chat_url, build_standard_plan_fallback_openai_responses_url, StandardPlanFallbackAcceptPolicy, @@ -200,6 +201,19 @@ pub(crate) fn build_openai_responses_sync_plan_from_decision( }, ); + let log_downstream_query = parts + .uri + .query() + .and_then(crate::ai_serving::api::sanitize_request_query_string); + let log_decision_upstream_base_url = payload + .upstream_base_url + .as_deref() + .map(sanitize_upstream_url_for_log); + let log_decision_upstream_url = payload + .upstream_url + .as_deref() + .map(sanitize_upstream_url_for_log); + let log_plan_url = sanitize_upstream_url_for_log(plan.url.as_str()); debug!( event_name = "local_openai_responses_sync_plan_built", log_type = "debug", @@ -209,11 +223,11 @@ pub(crate) fn build_openai_responses_sync_plan_from_decision( endpoint_id = %plan.endpoint_id, key_id = %plan.key_id, downstream_path = %parts.uri.path(), - downstream_query = ?parts.uri.query(), + downstream_query = ?log_downstream_query, url_source, - decision_upstream_base_url = ?payload.upstream_base_url, - decision_upstream_url = ?payload.upstream_url, - plan_url = %plan.url, + decision_upstream_base_url = ?log_decision_upstream_base_url, + decision_upstream_url = ?log_decision_upstream_url, + plan_url = %log_plan_url, client_api_format = %plan.client_api_format, provider_api_format = %plan.provider_api_format, upstream_is_stream = payload.upstream_is_stream, diff --git a/apps/aether-gateway/src/ai_serving/planner/standard/openai/responses/decision/payload.rs b/apps/aether-gateway/src/ai_serving/planner/standard/openai/responses/decision/payload.rs index 6c953611f..dfc950f2d 100644 --- a/apps/aether-gateway/src/ai_serving/planner/standard/openai/responses/decision/payload.rs +++ b/apps/aether-gateway/src/ai_serving/planner/standard/openai/responses/decision/payload.rs @@ -3,6 +3,7 @@ use tracing::debug; use crate::ai_serving::build_request_trace_proxy_value; use crate::ai_serving::planner::decision_input::apply_provider_request_routing_policy_to_decision_with_websocket_mode; +use crate::ai_serving::planner::redaction::sanitize_upstream_url_for_log; use crate::ai_serving::planner::report_context::{ build_local_execution_report_context, insert_native_client_envelope_name, insert_provider_stream_event_api_format, LocalExecutionReportContextParts, @@ -203,6 +204,12 @@ pub(crate) async fn maybe_build_local_openai_responses_decision_payload_for_cand &resolved.transport, ); + let log_base_url = sanitize_upstream_url_for_log(resolved.transport.endpoint.base_url.as_str()); + let log_request_query = parts + .uri + .query() + .and_then(crate::ai_serving::api::sanitize_request_query_string); + let log_upstream_url = sanitize_upstream_url_for_log(resolved.upstream_url.as_str()); debug!( event_name = "local_openai_responses_decision_payload_built", log_type = "debug", @@ -219,9 +226,9 @@ pub(crate) async fn maybe_build_local_openai_responses_decision_payload_for_cand client_api_format = spec_metadata.api_format, provider_api_format = %resolved.provider_api_format, request_path = %parts.uri.path(), - request_query = ?parts.uri.query(), - upstream_base_url = %resolved.transport.endpoint.base_url, - upstream_url = %resolved.upstream_url, + request_query = ?log_request_query, + upstream_base_url = %log_base_url, + upstream_url = %log_upstream_url, upstream_is_stream = resolved.upstream_is_stream, has_envelope = resolved.envelope_name.is_some(), "gateway built local openai responses decision payload" diff --git a/apps/aether-gateway/src/ai_serving/planner/standard/openai/responses/decision/request.rs b/apps/aether-gateway/src/ai_serving/planner/standard/openai/responses/decision/request.rs index f9da0d939..8d09af2ae 100644 --- a/apps/aether-gateway/src/ai_serving/planner/standard/openai/responses/decision/request.rs +++ b/apps/aether-gateway/src/ai_serving/planner/standard/openai/responses/decision/request.rs @@ -24,6 +24,7 @@ use crate::ai_serving::planner::gemini_cli::{ }; use crate::ai_serving::planner::redaction::{ request_identity_response_encoding_when_redacted, resolve_provider_chat_pii_redaction, + sanitize_upstream_url_for_log, }; use crate::ai_serving::planner::spec_metadata::local_openai_responses_spec_metadata; use crate::ai_serving::planner::standard::{ @@ -865,6 +866,17 @@ pub(crate) async fn resolve_local_openai_responses_candidate_payload_parts_with_ let (execution_strategy, conversion_mode) = ai_local_execution_contract_for_formats(spec_metadata.api_format, provider_api_format); + let log_base_url = sanitize_upstream_url_for_log(transport.endpoint.base_url.as_str()); + let log_custom_path = transport + .endpoint + .custom_path + .as_deref() + .map(sanitize_upstream_url_for_log); + let log_request_query = parts + .uri + .query() + .and_then(crate::ai_serving::api::sanitize_request_query_string); + let log_upstream_url = sanitize_upstream_url_for_log(upstream_url.as_str()); debug!( event_name = "local_openai_responses_upstream_url_resolved", @@ -880,12 +892,12 @@ pub(crate) async fn resolve_local_openai_responses_candidate_payload_parts_with_ provider_api_format = %provider_api_format, execution_strategy = execution_strategy.as_str(), conversion_mode = conversion_mode.as_str(), - base_url = %transport.endpoint.base_url, - custom_path = ?transport.endpoint.custom_path, + base_url = %log_base_url, + custom_path = ?log_custom_path, request_path = %parts.uri.path(), - request_query = ?parts.uri.query(), + request_query = ?log_request_query, mapped_model = %mapped_model, - upstream_url = %upstream_url, + upstream_url = %log_upstream_url, upstream_is_stream, "gateway resolved local openai responses upstream url" ); @@ -1998,6 +2010,7 @@ async fn build_kiro_openai_responses_payload_parts( }; let (execution_strategy, conversion_mode) = ai_local_execution_contract_for_formats(client_api_format, provider_api_format); + let log_upstream_url = sanitize_upstream_url_for_log(upstream_url.as_str()); debug!( event_name = "local_openai_responses_kiro_upstream_url_resolved", @@ -2013,7 +2026,7 @@ async fn build_kiro_openai_responses_payload_parts( provider_api_format = %provider_api_format, execution_strategy = execution_strategy.as_str(), conversion_mode = conversion_mode.as_str(), - upstream_url = %upstream_url, + upstream_url = %log_upstream_url, upstream_is_stream, "gateway resolved local openai responses kiro upstream url" ); diff --git a/apps/aether-gateway/src/ai_serving/planner/standard/openai/responses/mod.rs b/apps/aether-gateway/src/ai_serving/planner/standard/openai/responses/mod.rs index 6666be55a..7b4b3f3d2 100644 --- a/apps/aether-gateway/src/ai_serving/planner/standard/openai/responses/mod.rs +++ b/apps/aether-gateway/src/ai_serving/planner/standard/openai/responses/mod.rs @@ -234,6 +234,10 @@ pub(crate) struct ResponsesWebSocketDecision { pub(crate) execution: AiExecutionDecision, pub(crate) adapter: ResponsesWebSocketAdapter, pub(crate) normalization: ResponsesWebSocketBodyNormalization, + /// Effective key auth after applying the endpoint API-format override. + /// Protocol companions such as Codex Live must not infer this from a URL + /// or from the presence of one particular generated header. + pub(crate) effective_auth_type: String, } /// The scheduler identity a continuation is allowed to reuse. @@ -909,6 +913,10 @@ pub(crate) async fn maybe_build_responses_websocket_decision( // Captured before `attempt` is consumed so a later continuation turn can // reproduce this candidate's body normalization without re-planning. let transport = std::sync::Arc::clone(&attempt.eligible.transport); + let effective_auth_type = + aether_provider_transport::auth::resolve_local_auth_type_for_transport_format( + transport.as_ref(), + ); let candidate_provider_api_format = attempt.eligible.provider_api_format.clone(); let payload = match maybe_build_local_openai_responses_decision_payload_for_candidate_with_websocket_mode( state, @@ -1013,6 +1021,7 @@ pub(crate) async fn maybe_build_responses_websocket_decision( execution: payload, adapter, normalization, + effective_auth_type, }; // The decision report context now carries the lease identity. The // WebSocket ownership layer takes over before any further await. diff --git a/apps/aether-gateway/src/api/ai/registry.rs b/apps/aether-gateway/src/api/ai/registry.rs index b5e46c777..ed674c2b9 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::{proxy_request, responses_websocket}, + handlers::proxy::{live_websocket, proxy_request, responses_websocket}, state::AppState, GatewayError, }; @@ -22,6 +22,7 @@ const AI_POST_ROUTE_PATTERNS: &[&str] = &[ "/v1/rerank", "/v1/responses", "/v1/responses/compact", + "/v1/live", "/v1/alpha/search", "/v1/images/generations", "/v1/images/edits", @@ -57,10 +58,13 @@ pub(crate) fn mount_ai_routes(mut router: Router) -> Router for path in AI_POST_ROUTE_PATTERNS { router = if *path == "/v1/responses" { router.route(path, get(responses_websocket).post(proxy_request)) + } else if *path == "/v1/live" { + router.route(path, get(live_websocket).post(proxy_request)) } else { router.route(path, post(proxy_request)) }; } + router = router.route("/v1/live/{call_id}", get(live_websocket)); for path in CLAUDE_POST_ROUTE_PATTERNS { router = router.route( path, diff --git a/apps/aether-gateway/src/constants.rs b/apps/aether-gateway/src/constants.rs index 44e2f4788..e2741011d 100644 --- a/apps/aether-gateway/src/constants.rs +++ b/apps/aether-gateway/src/constants.rs @@ -125,6 +125,8 @@ pub(crate) const RUST_FRONTDOOR_OWNED_ROUTE_PATTERNS: &[&str] = &[ "/v1/messages/count_tokens", "/v1/responses", "/v1/responses/compact", + "/v1/live", + "/v1/live/{call_id}", "/v1/alpha/search", "/v1/models/{model}:generateContent", "/v1/models/{model}:streamGenerateContent", diff --git a/apps/aether-gateway/src/control/route/ai.rs b/apps/aether-gateway/src/control/route/ai.rs index ba84ad48a..376008b48 100644 --- a/apps/aether-gateway/src/control/route/ai.rs +++ b/apps/aether-gateway/src/control/route/ai.rs @@ -35,6 +35,21 @@ pub(super) fn classify_ai_public_route( "openai:rerank", 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, + )) } else if (method == http::Method::POST || (method == http::Method::GET && normalized_path == "/v1/responses" @@ -291,4 +306,27 @@ mod tests { classify_ai_public_route(&Method::GET, "/v1/responses", &HeaderMap::new()).is_none() ); } + + #[test] + fn classifies_live_http_and_websocket_routes_as_responses_companions() { + 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_kind, "live"); + assert_eq!(post.auth_endpoint_signature, "openai:responses"); + + let mut headers = HeaderMap::new(); + headers.insert(CONNECTION, HeaderValue::from_static("Upgrade")); + headers.insert(UPGRADE, HeaderValue::from_static("websocket")); + 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_kind, "live"); + assert_eq!(route.auth_endpoint_signature, "openai:responses"); + } + + assert!( + classify_ai_public_route(&Method::GET, "/v1/live/rtc_opaque", &HeaderMap::new()) + .is_none() + ); + } } diff --git a/apps/aether-gateway/src/execution_runtime/stream/execution.rs b/apps/aether-gateway/src/execution_runtime/stream/execution.rs index 31b3f059c..e7c90815a 100644 --- a/apps/aether-gateway/src/execution_runtime/stream/execution.rs +++ b/apps/aether-gateway/src/execution_runtime/stream/execution.rs @@ -64,6 +64,11 @@ use crate::ai_serving::api::{ extract_provider_private_stream_error_body, maybe_bridge_standard_sync_json_to_stream, maybe_build_provider_private_stream_normalizer, maybe_build_stream_response_rewriter, normalize_provider_private_report_context, StreamingStandardTerminalObserver, + CLAUDE_CHAT_STREAM_PLAN_KIND, CLAUDE_CLI_STREAM_PLAN_KIND, GEMINI_CHAT_STREAM_PLAN_KIND, + GEMINI_CLI_STREAM_PLAN_KIND, GEMINI_INTERACTIONS_STREAM_PLAN_KIND, + OPENAI_CHAT_STREAM_PLAN_KIND, OPENAI_IMAGE_STREAM_PLAN_KIND, + OPENAI_RESPONSES_COMPACT_STREAM_PLAN_KIND, OPENAI_RESPONSES_STREAM_PLAN_KIND, + UPSTREAM_IS_STREAM_KEY, }; use crate::ai_serving::is_openai_responses_family_format; use crate::api::response::{ @@ -144,7 +149,6 @@ use crate::{ AppState, GatewayError, GEMINI_FILES_DOWNLOAD_PLAN_KIND, OPENAI_VIDEO_CONTENT_PLAN_KIND, }; -const OPENAI_IMAGE_STREAM_PLAN_KIND: &str = "openai_image_stream"; const SSE_KEEPALIVE_INTERVAL: Duration = Duration::from_secs(15); const SSE_KEEPALIVE_BYTES: &[u8] = b": aether-keepalive\n\n"; const SSE_CONTROL_FILTER_MAX_BUFFER_BYTES: usize = 1024 * 1024; @@ -4524,13 +4528,86 @@ fn decode_stream_data_chunk( fn response_headers_indicate_sse(headers: &BTreeMap) -> bool { headers - .get("content-type") - .map(String::as_str) + .iter() + .find(|(name, _)| name.eq_ignore_ascii_case("content-type")) + .map(|(_, value)| value.as_str()) .map(str::trim) .filter(|value| !value.is_empty()) .is_some_and(|value| value.to_ascii_lowercase().contains("text/event-stream")) } +fn report_context_upstream_is_stream(report_context: Option<&Value>) -> bool { + report_context + .and_then(|value| value.get(UPSTREAM_IS_STREAM_KEY)) + .and_then(Value::as_bool) + .unwrap_or(false) +} + +fn response_headers_have_octet_stream_content_type(headers: &BTreeMap) -> bool { + headers + .iter() + .find(|(name, _)| name.eq_ignore_ascii_case("content-type")) + .map(|(_, value)| value.as_str()) + .and_then(|value| value.split(';').next()) + .map(str::trim) + .is_some_and(|value| value.eq_ignore_ascii_case("application/octet-stream")) +} + +fn response_headers_have_only_identity_content_encoding( + headers: &BTreeMap, +) -> bool { + headers + .iter() + .find(|(name, _)| name.eq_ignore_ascii_case("content-encoding")) + .map(|(_, value)| value.as_str()) + .is_none_or(|value| { + value + .split(',') + .map(str::trim) + .all(|coding| coding.is_empty() || coding.eq_ignore_ascii_case("identity")) + }) +} + +fn plan_kind_uses_text_event_stream(plan_kind: &str) -> bool { + matches!( + plan_kind, + OPENAI_CHAT_STREAM_PLAN_KIND + | OPENAI_RESPONSES_STREAM_PLAN_KIND + | OPENAI_RESPONSES_COMPACT_STREAM_PLAN_KIND + | OPENAI_IMAGE_STREAM_PLAN_KIND + | CLAUDE_CHAT_STREAM_PLAN_KIND + | CLAUDE_CLI_STREAM_PLAN_KIND + | GEMINI_CHAT_STREAM_PLAN_KIND + | GEMINI_CLI_STREAM_PLAN_KIND + | GEMINI_INTERACTIONS_STREAM_PLAN_KIND + ) +} + +fn should_normalize_declared_stream_response_headers( + plan_kind: &str, + status_code: u16, + headers: &BTreeMap, + report_context: Option<&Value>, +) -> bool { + plan_kind_uses_text_event_stream(plan_kind) + && (200..300).contains(&status_code) + && report_context_upstream_is_stream(report_context) + && response_headers_have_octet_stream_content_type(headers) + && response_headers_have_only_identity_content_encoding(headers) + && !headers + .keys() + .any(|name| name.eq_ignore_ascii_case("content-length")) +} + +fn normalize_declared_stream_response_headers(headers: &mut BTreeMap) { + headers.retain(|name, _| { + !name.eq_ignore_ascii_case("content-encoding") + && !name.eq_ignore_ascii_case("content-length") + && !name.eq_ignore_ascii_case("content-type") + }); + headers.insert("content-type".to_string(), "text/event-stream".to_string()); +} + fn parse_prefetched_sync_json_body(body: &[u8]) -> Option { let stripped = strip_utf8_bom_and_ws(body); serde_json::from_slice::(stripped).ok() @@ -6042,6 +6119,32 @@ async fn execute_stream_from_frame_stream_with_retry_scope( headers.insert("content-type".to_string(), "text/event-stream".to_string()); } let upstream_content_type = upstream_headers.get("content-type").map(String::as_str); + let normalized_declared_stream_headers = private_stream_normalizer.is_none() + && local_stream_rewriter.is_none() + && should_normalize_declared_stream_response_headers( + plan_kind, + status_code, + &upstream_headers, + report_context.as_ref(), + ); + if normalized_declared_stream_headers { + normalize_declared_stream_response_headers(&mut headers); + debug!( + event_name = "execution_runtime_stream_content_type_corrected", + log_type = "debug", + trace_id = %trace_id, + request_id = %request_id_for_log, + candidate_id = ?candidate_id, + plan_kind, + provider_name, + endpoint_id = %plan.endpoint_id, + key_id = %plan.key_id, + model_name, + candidate_index = candidate_index.as_str(), + upstream_content_type = upstream_content_type.unwrap_or("-"), + "gateway normalized declared upstream stream response headers for the client" + ); + } let prefetch_for_cyber_failover = is_openai_responses_family_format(plan.provider_api_format.as_str()) && cyber_continue_failover_enabled(state).await; @@ -6729,8 +6832,9 @@ async fn execute_stream_from_frame_stream_with_retry_scope( let native_anthropic_stream_for_report = stream_commit_policy.is_native_anthropic(); let plan_for_report = plan; let emit_passthrough_sse_terminal_error = (skip_direct_finalize_prefetch - || stream_commit_policy.is_native_anthropic()) - && response_headers_indicate_sse(&upstream_headers) + || stream_commit_policy.is_native_anthropic() + || normalized_declared_stream_headers) + && (response_headers_indicate_sse(&upstream_headers) || normalized_declared_stream_headers) && !is_openai_image_stream_for_report; let plan_kind_for_report = plan_kind.to_string(); let stream_started_at_for_report = stream_started_at; @@ -8092,20 +8196,22 @@ mod tests { execute_execution_runtime_stream, execute_in_process_stream_with_oauth_retry, execute_stream_from_frame_stream, execute_stream_from_frame_stream_with_retry_scope, maybe_apply_kiro_prompt_cache_usage_to_stream_summary, merge_stream_terminal_summary, - parse_direct_passthrough_mode, prefetch_direct_stream_error_body, - prefetched_openai_responses_body_has_output_boundary, + normalize_declared_stream_response_headers, parse_direct_passthrough_mode, + prefetch_direct_stream_error_body, prefetched_openai_responses_body_has_output_boundary, record_sync_terminal_usage_with_handoff, record_sync_terminal_usage_with_handoff_after_spawn, resolve_provider_stream_error_status_code, select_direct_anthropic_prefetch_wait, - should_limit_direct_finalize_prefetch, should_probe_success_failover_before_stream, - should_skip_direct_finalize_prefetch, stream_chunk_contains_sse_done, - stream_requires_observed_terminal_event, stream_terminal_summary_missing_observed_finish, + should_limit_direct_finalize_prefetch, should_normalize_declared_stream_response_headers, + should_probe_success_failover_before_stream, should_skip_direct_finalize_prefetch, + stream_chunk_contains_sse_done, stream_requires_observed_terminal_event, + stream_terminal_summary_missing_observed_finish, stream_terminal_summary_missing_observed_finish_with_requirement, stream_terminal_summary_represents_failure_with_requirement, ClientVisibleStreamCompletionTracker, DirectPassthroughFinalizer, DirectPassthroughFinalizerCore, DirectPassthroughInlineBodyState, DirectPassthroughMode, PostStopFrameReadBudget, PostStopLimitedStreamReader, ProviderStreamErrorInspection, - ANTHROPIC_POST_STOP_DRAIN_MAX_BYTES, POST_STOP_MAX_EMPTY_CHUNKS_PER_POLL, + ANTHROPIC_POST_STOP_DRAIN_MAX_BYTES, GEMINI_FILES_DOWNLOAD_PLAN_KIND, + OPENAI_CHAT_STREAM_PLAN_KIND, POST_STOP_MAX_EMPTY_CHUNKS_PER_POLL, }; use crate::control::GatewayControlDecision; use crate::stage_metrics::RequestStageTrace; @@ -12059,6 +12165,109 @@ mod tests { )); } + #[test] + fn declared_stream_response_headers_are_normalized_without_body_inspection() { + let mut headers = BTreeMap::from([ + ( + "Content-Type".to_string(), + "Application/Octet-Stream; charset=binary".to_string(), + ), + ("Content-Encoding".to_string(), "identity".to_string()), + ("x-upstream-header".to_string(), "preserved".to_string()), + ]); + assert!(should_normalize_declared_stream_response_headers( + OPENAI_CHAT_STREAM_PLAN_KIND, + 200, + &headers, + Some(&json!({"upstream_is_stream": true})), + )); + headers.insert("Content-Length".to_string(), "4096".to_string()); + normalize_declared_stream_response_headers(&mut headers); + + assert_eq!( + headers.get("content-type").map(String::as_str), + Some("text/event-stream") + ); + assert!(!headers + .keys() + .any(|name| name.eq_ignore_ascii_case("content-encoding"))); + assert!(!headers + .keys() + .any(|name| name.eq_ignore_ascii_case("content-length"))); + assert_eq!( + headers.get("x-upstream-header").map(String::as_str), + Some("preserved") + ); + } + + #[test] + fn declared_stream_header_normalization_requires_success_and_stream_context() { + let headers = BTreeMap::from([( + "content-type".to_string(), + "application/octet-stream".to_string(), + )]); + assert!(!should_normalize_declared_stream_response_headers( + OPENAI_CHAT_STREAM_PLAN_KIND, + 500, + &headers, + Some(&json!({"upstream_is_stream": true})), + )); + assert!(!should_normalize_declared_stream_response_headers( + OPENAI_CHAT_STREAM_PLAN_KIND, + 200, + &headers, + Some(&json!({"upstream_is_stream": false})), + )); + assert!(!should_normalize_declared_stream_response_headers( + OPENAI_CHAT_STREAM_PLAN_KIND, + 200, + &BTreeMap::from([("content-type".to_string(), "text/event-stream".to_string(),)]), + Some(&json!({"upstream_is_stream": true})), + )); + assert!(!should_normalize_declared_stream_response_headers( + OPENAI_CHAT_STREAM_PLAN_KIND, + 200, + &BTreeMap::from([("content-type".to_string(), "application/json".to_string(),)]), + Some(&json!({"upstream_is_stream": true})), + )); + assert!(!should_normalize_declared_stream_response_headers( + OPENAI_CHAT_STREAM_PLAN_KIND, + 200, + &BTreeMap::from([("content-type".to_string(), "text/plain".to_string(),)]), + Some(&json!({"upstream_is_stream": true})), + )); + assert!(!should_normalize_declared_stream_response_headers( + OPENAI_CHAT_STREAM_PLAN_KIND, + 200, + &BTreeMap::from([ + ( + "content-type".to_string(), + "application/octet-stream".to_string(), + ), + ("content-encoding".to_string(), "gzip".to_string()), + ]), + Some(&json!({"upstream_is_stream": true})), + )); + assert!(!should_normalize_declared_stream_response_headers( + OPENAI_CHAT_STREAM_PLAN_KIND, + 200, + &BTreeMap::from([ + ( + "content-type".to_string(), + "application/octet-stream".to_string(), + ), + ("content-length".to_string(), "128".to_string()), + ]), + Some(&json!({"upstream_is_stream": true})), + )); + assert!(!should_normalize_declared_stream_response_headers( + GEMINI_FILES_DOWNLOAD_PLAN_KIND, + 200, + &headers, + Some(&json!({"upstream_is_stream": true})), + )); + } + #[test] fn skips_prefetch_for_event_streams_even_when_cross_format_or_rewritten() { assert!(should_skip_direct_finalize_prefetch( @@ -13791,6 +14000,7 @@ mod tests { Some(json!({ "provider_api_format": "openai:responses", "client_api_format": "openai:responses", + "upstream_is_stream": true, })), ) .await diff --git a/apps/aether-gateway/src/frontdoor_loop_guard.rs b/apps/aether-gateway/src/frontdoor_loop_guard.rs index 827faf7f9..f09119e64 100644 --- a/apps/aether-gateway/src/frontdoor_loop_guard.rs +++ b/apps/aether-gateway/src/frontdoor_loop_guard.rs @@ -41,12 +41,14 @@ pub(crate) fn frontdoor_self_loop_public_ai_path(path: &str) -> bool { | "/v1/rerank" | "/v1/responses" | "/v1/responses/compact" + | "/v1/live" | "/v1/alpha/search" | "/v1beta/files" | "/upload/v1beta/files" | "/v1beta/operations" | "/v1/videos" - ) || path.starts_with("/v1/videos/") + ) || path.starts_with("/v1/live/") + || path.starts_with("/v1/videos/") || path.starts_with("/v1beta/files/") || path.starts_with("/v1beta/operations/") || path.starts_with("/v1internal:") diff --git a/apps/aether-gateway/src/handlers/proxy/mod.rs b/apps/aether-gateway/src/handlers/proxy/mod.rs index e57781864..070059af7 100644 --- a/apps/aether-gateway/src/handlers/proxy/mod.rs +++ b/apps/aether-gateway/src/handlers/proxy/mod.rs @@ -9,6 +9,7 @@ use self::body_buffer::{ 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::responses::responses_websocket; use super::internal::resolve_local_proxy_execution_path; pub(crate) use super::public::matches_model_mapping_for_models; @@ -106,6 +107,7 @@ const AUTH_API_KEY_CONCURRENCY_LIMIT_REACHED_DETAIL: &str = const LOCAL_EXECUTION_PLANNING_TIMEOUT_DETAIL: &str = "当前 AI 请求在本地执行规划阶段超时，请稍后重试"; const EXECUTION_PATH_TUNNEL_AFFINITY_FORWARD: &str = "tunnel_affinity_forward"; +const EXECUTION_PATH_CODEX_LIVE_CALL: &str = "codex_live_call"; const MANAGEMENT_TOKEN_PREFIX: &str = "ae-"; const LEGACY_MANAGEMENT_TOKEN_PREFIX: &str = "ae_"; fn finalize_request_body_buffer_rejection( @@ -939,11 +941,11 @@ pub(crate) async fn proxy_request( ConnectInfo(remote_addr): ConnectInfo, request: Request, ) -> Result, GatewayError> { - crate::request_diagnostics::scope_request_diagnostics(proxy_request_inner( + crate::request_diagnostics::scope_request_diagnostics(Box::pin(proxy_request_inner( state, remote_addr, request, - )) + ))) .await } @@ -1567,6 +1569,26 @@ async fn proxy_request_inner( )); } + if let Some(response) = Box::pin(maybe_handle_live_http( + &state, + &request_context, + &parts, + buffered_body.as_ref(), + &remote_addr, + )) + .await? + { + return Ok(finalize_gateway_response_with_context( + &state, + response, + &remote_addr, + &request_context, + EXECUTION_PATH_CODEX_LIVE_CALL, + &started_at, + request_permit.take(), + )); + } + let local_ai_public_started_at = Instant::now(); let local_ai_public_response = super::public::maybe_build_local_ai_public_response( &state, diff --git a/apps/aether-gateway/src/handlers/proxy/websocket/ingress.rs b/apps/aether-gateway/src/handlers/proxy/websocket/ingress.rs index bd3626fee..7b8907a93 100644 --- a/apps/aether-gateway/src/handlers/proxy/websocket/ingress.rs +++ b/apps/aether-gateway/src/handlers/proxy/websocket/ingress.rs @@ -50,6 +50,82 @@ pub(crate) struct WebSocketIngressSpec { pub(crate) route_unavailable_message: &'static str, } +pub(crate) enum AuthenticatedAiWebSocketUpgradePreparation { + Ready(AuthenticatedAiWebSocketUpgrade), + Rejected(Response), +} + +/// Authenticated HTTP Upgrade state retained while an adapter performs any +/// protocol-specific preflight that must complete before status 101 is sent. +pub(crate) struct AuthenticatedAiWebSocketUpgrade { + state: AppState, + context: WebSocketRequestContext, + request_permit: Option, +} + +impl AuthenticatedAiWebSocketUpgrade { + pub(crate) fn state(&self) -> &AppState { + &self.state + } + + pub(crate) fn context(&self) -> &WebSocketRequestContext { + &self.context + } + + pub(crate) fn rejection_response( + &self, + status: StatusCode, + message: &str, + ) -> Result, GatewayError> { + build_local_http_error_response( + self.context.trace_id.as_str(), + Some(&self.context.decision), + status, + message, + ) + } + + pub(crate) fn into_response( + self, + ws: WebSocketUpgrade, + limits: WebSocketSessionLimits, + run_session: F, + ) -> Response + where + F: FnOnce(WebSocket, AppState, WebSocketRequestContext) -> Fut + Send + 'static, + Fut: Future + Send + 'static, + { + self.into_response_with(ws, limits, (), move |socket, state, context, ()| { + run_session(socket, state, context) + }) + } + + pub(crate) fn into_response_with( + self, + ws: WebSocketUpgrade, + limits: WebSocketSessionLimits, + prepared: P, + run_session: F, + ) -> Response + where + P: Send + 'static, + F: FnOnce(WebSocket, AppState, WebSocketRequestContext, P) -> Fut + Send + 'static, + Fut: Future + Send + 'static, + { + let Self { + state, + context, + request_permit, + } = self; + ws.max_frame_size(limits.max_frame_size) + .max_message_size(limits.max_message_size) + .on_upgrade(move |socket| async move { + drop(request_permit); + run_session(socket, state, context, prepared).await; + }) + } +} + /// Performs the HTTP-only part of an AI WebSocket request. /// /// The ordinary request permit covers only the HTTP Upgrade window. A @@ -69,6 +145,21 @@ where F: FnOnce(WebSocket, AppState, WebSocketRequestContext) -> Fut + Send + 'static, Fut: Future + Send + 'static, { + match prepare_authenticated_ai_websocket(state, remote_addr, headers, uri, spec).await? { + AuthenticatedAiWebSocketUpgradePreparation::Ready(prepared) => { + Ok(prepared.into_response(ws, limits, run_session)) + } + AuthenticatedAiWebSocketUpgradePreparation::Rejected(response) => Ok(response), + } +} + +pub(crate) async fn prepare_authenticated_ai_websocket( + state: AppState, + remote_addr: SocketAddr, + headers: HeaderMap, + uri: Uri, + spec: WebSocketIngressSpec, +) -> Result { let trace_id = extract_or_generate_trace_id(&headers); let client_ip = effective_client_ip(&headers, &remote_addr); if state.admin_security_ip_blacklisted(client_ip).await? { @@ -77,7 +168,8 @@ where None, StatusCode::FORBIDDEN, "当前 IP 已被禁止访问", - ); + ) + .map(AuthenticatedAiWebSocketUpgradePreparation::Rejected); } let request_context = crate::control::resolve_public_request_context( @@ -94,10 +186,12 @@ where None, StatusCode::NOT_FOUND, spec.route_unavailable_message, - ); + ) + .map(AuthenticatedAiWebSocketUpgradePreparation::Rejected); }; if let Some(rejection) = trusted_auth_local_rejection(Some(&decision), &headers) { - return build_local_auth_rejection_response(&trace_id, Some(&decision), &rejection); + return build_local_auth_rejection_response(&trace_id, Some(&decision), &rejection) + .map(AuthenticatedAiWebSocketUpgradePreparation::Rejected); } // Browsers attach cookies to WebSocket handshakes automatically and the // WebSocket API does not let callers add an Authorization header. A @@ -119,14 +213,16 @@ where &trace_id, Some(&decision), &GatewayLocalAuthRejection::InvalidApiKey, - ); + ) + .map(AuthenticatedAiWebSocketUpgradePreparation::Rejected); } let Some(auth_context) = decision.auth_context.as_ref() else { return build_local_auth_rejection_response( &trace_id, Some(&decision), &GatewayLocalAuthRejection::InvalidApiKey, - ); + ) + .map(AuthenticatedAiWebSocketUpgradePreparation::Rejected); }; if !auth_context.access_allowed || auth_context.user_id.trim().is_empty() @@ -136,7 +232,8 @@ where &trace_id, Some(&decision), &GatewayLocalAuthRejection::InvalidApiKey, - ); + ) + .map(AuthenticatedAiWebSocketUpgradePreparation::Rejected); } if !ip_rules_allow(auth_context.ip_rules.as_deref(), client_ip) { return build_local_auth_rejection_response( @@ -145,7 +242,8 @@ where &GatewayLocalAuthRejection::IpNotAllowed { remote_ip: client_ip.to_string(), }, - ); + ) + .map(AuthenticatedAiWebSocketUpgradePreparation::Rejected); } let request_permit = match state.try_acquire_request_permit().await { @@ -157,6 +255,7 @@ where Some(uri.path()), error, ) + .map(AuthenticatedAiWebSocketUpgradePreparation::Rejected) } }; let websocket_connection_permit = match state.try_acquire_websocket_connection_permit().await { @@ -168,6 +267,7 @@ where Some(uri.path()), error, ) + .map(AuthenticatedAiWebSocketUpgradePreparation::Rejected) } }; @@ -188,13 +288,13 @@ where decision, websocket_connection_permit, }; - Ok(ws - .max_frame_size(limits.max_frame_size) - .max_message_size(limits.max_message_size) - .on_upgrade(move |socket| async move { - drop(request_permit); - run_session(socket, state, context).await; - })) + Ok(AuthenticatedAiWebSocketUpgradePreparation::Ready( + AuthenticatedAiWebSocketUpgrade { + state, + context, + request_permit, + }, + )) } fn websocket_credential_carrier_is_allowed(carrier: Option) -> bool { @@ -353,7 +453,7 @@ impl WebSocketConnectionLog { spec, trace_id: context.trace_id.clone(), remote_addr: context.remote_addr, - path: context.uri.path().to_string(), + path: websocket_log_path(context.uri.path()), route_class: context .decision .route_class @@ -392,6 +492,17 @@ impl WebSocketConnectionLog { } } +fn websocket_log_path(path: &str) -> String { + if path + .strip_prefix("/v1/live/") + .is_some_and(|call_id| !call_id.is_empty()) + { + "/v1/live/{call_id}".to_string() + } else { + path.to_string() + } +} + impl Drop for WebSocketConnectionLog { fn drop(&mut self) { info!( @@ -424,7 +535,8 @@ mod tests { use axum::http::{HeaderMap, HeaderValue, Uri}; use super::{ - websocket_credential_carrier_is_allowed, websocket_planning_headers, websocket_planning_uri, + websocket_credential_carrier_is_allowed, websocket_log_path, websocket_planning_headers, + websocket_planning_uri, }; use crate::control::GatewayCredentialCarrier; @@ -524,4 +636,13 @@ mod tests { assert!(websocket_credential_carrier_is_allowed(carrier)); } } + + #[test] + fn live_sideband_access_logs_do_not_retain_the_opaque_call_id() { + assert_eq!(websocket_log_path("/v1/live"), "/v1/live"); + assert_eq!( + websocket_log_path("/v1/live/rtc_secret_opaque"), + "/v1/live/{call_id}" + ); + } } diff --git a/apps/aether-gateway/src/handlers/proxy/websocket/live/http.rs b/apps/aether-gateway/src/handlers/proxy/websocket/live/http.rs new file mode 100644 index 000000000..9340b67b6 --- /dev/null +++ b/apps/aether-gateway/src/handlers/proxy/websocket/live/http.rs @@ -0,0 +1,731 @@ +//! Authenticated WebRTC call creation for Codex Live. + +use std::collections::BTreeMap; +use std::net::SocketAddr; + +use aether_contracts::{ + ExecutionPlan, ExecutionResponseBodyMode, ExecutionResult, EXECUTION_RESPONSE_BODY_MODE_HEADER, +}; +use axum::body::{Body, Bytes}; +use axum::http::{HeaderValue, Response, StatusCode}; +use base64::Engine as _; +use serde_json::json; +use tracing::{info, warn}; + +use crate::ai_serving::build_standard_sync_plan_from_decision; +use crate::api::response::{ + build_client_response_from_parts, build_client_response_from_parts_with_mutator, + build_local_auth_rejection_response, build_local_http_error_response_with_request_path, +}; +use crate::control::{execution_plan_balance_capacity_rejection, GatewayPublicRequestContext}; +use crate::execution_runtime::execute_execution_runtime_sync_plan_with_report_context; +use crate::handlers::proxy::websocket::responses::ResponsesWebSocketTurnAdmission; +use crate::{AppState, GatewayError}; + +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}; +use super::registry::{LiveCallBinding, LiveCallRegistry}; + +const MAX_LIVE_HTTP_BODY_BYTES: usize = 1024 * 1024; + +pub(crate) async fn maybe_handle_live_http( + state: &AppState, + request_context: &GatewayPublicRequestContext, + parts: &http::request::Parts, + body: Option<&Bytes>, + remote_addr: &SocketAddr, +) -> Result>, GatewayError> { + if parts.method != http::Method::POST || request_context.request_path != "/v1/live" { + return Ok(None); + } + let Some(control_decision) = request_context.control_decision.as_ref() else { + return Ok(Some(local_live_error( + request_context, + StatusCode::NOT_FOUND, + "Codex Live route is unavailable", + )?)); + }; + if !live_usage_accounting_is_safe(control_decision) { + return Ok(Some(local_live_error( + request_context, + StatusCode::NOT_IMPLEMENTED, + "Codex Live is unavailable for finite-balance keys until Frameless usage settlement is supported", + )?)); + } + let Some(body) = body else { + return Ok(Some(local_live_error( + request_context, + StatusCode::BAD_REQUEST, + "Codex Live requires a multipart WebRTC offer", + )?)); + }; + if body.len() > MAX_LIVE_HTTP_BODY_BYTES { + return Ok(Some(local_live_error( + request_context, + StatusCode::PAYLOAD_TOO_LARGE, + "Codex Live WebRTC offer exceeds the 1 MiB limit", + )?)); + } + let content_type = parts + .headers + .get(http::header::CONTENT_TYPE) + .and_then(|value| value.to_str().ok()) + .unwrap_or_default(); + let offer = match parse_live_multipart(content_type, body.as_ref()) { + Ok(offer) => offer, + Err(error) => { + return Ok(Some(local_live_error( + request_context, + error.status_code(), + error.client_message(), + )?)) + } + }; + let Some(client_model) = offer + .session + .get("model") + .and_then(serde_json::Value::as_str) + else { + return Ok(Some(local_live_error( + request_context, + StatusCode::BAD_REQUEST, + "Codex Live session.model must be a non-empty model identifier", + )?)); + }; + + let Some(mut candidate) = plan_live_candidate( + state, + request_context.trace_id.as_str(), + control_decision, + &parts.headers, + remote_addr, + client_model, + None, + ) + .await? + else { + return Ok(Some(local_live_error( + request_context, + StatusCode::BAD_GATEWAY, + "No eligible Codex Live provider mapping is available", + )?)); + }; + let lease = LivePoolLeaseGuard::new(state, &candidate); + let binding = LiveCallBinding::from_candidate(&candidate); + let mut provider_session = offer.session.clone(); + provider_session + .as_object_mut() + .expect("validated Live session is a JSON object") + .insert( + "model".to_string(), + serde_json::Value::String(candidate.provider_model.clone()), + ); + let upstream_url = match live_call_url(&candidate) { + Ok(url) => url, + Err(error) => { + lease.release().await; + return Ok(Some(local_live_error( + request_context, + error.status_code(), + error.client_message(), + )?)); + } + }; + + let (provider_content_type, provider_body_base64) = + build_live_call_provider_body(candidate.auth_mode, offer.sdp.as_str(), &provider_session)?; + // The standard plan builder requires a JSON body marker even when the exact wire body is + // carried as bytes. Keep only the mapped model here: retaining the SDP/session projection in + // the decision would unnecessarily widen the surface for future logging or report changes. + let provider_body_marker = json!({"model": candidate.provider_model.clone()}); + candidate.execution.upstream_url = Some(upstream_url); + candidate.execution.provider_request_method = Some("POST".to_string()); + candidate.execution.provider_request_body = Some(provider_body_marker.clone()); + candidate.execution.provider_request_body_base64 = Some(provider_body_base64); + candidate.execution.content_type = Some(provider_content_type.clone()); + candidate.execution.content_encoding = None; + candidate.execution.request_gzip = None; + candidate.execution.upstream_is_stream = false; + prepare_live_call_request_headers( + &mut candidate.execution.provider_request_headers, + provider_content_type.as_str(), + ); + candidate.execution.provider_request_headers.insert( + EXECUTION_RESPONSE_BODY_MODE_HEADER.to_string(), + ExecutionResponseBodyMode::PreserveBytes + .as_str() + .to_string(), + ); + + let Some(attempt) = + build_standard_sync_plan_from_decision(parts, &provider_body_marker, candidate.execution)? + else { + lease.release().await; + return Ok(Some(local_live_error( + request_context, + StatusCode::BAD_GATEWAY, + "Codex Live provider request could not be built", + )?)); + }; + if let Some(rejection) = execution_plan_balance_capacity_rejection( + state, + control_decision, + &attempt.plan, + attempt.report_context.as_ref(), + ) + .await? + { + lease.release().await; + return Ok(Some(build_local_auth_rejection_response( + request_context.trace_id.as_str(), + Some(control_decision), + &rejection, + )?)); + } + let admission = ResponsesWebSocketTurnAdmission::acquire( + state, + &attempt.plan, + request_context.trace_id.as_str(), + ) + .await?; + let result = execute_execution_runtime_sync_plan_with_report_context( + state, + Some(request_context.trace_id.as_str()), + &attempt.plan, + attempt.report_context.as_ref(), + ) + .await; + // These guards intentionally cover only the synchronous call-creation exchange. The + // WebRTC media leg bypasses Aether, and neither the two-hour routing binding nor a sideband + // attachment proves that media is still alive. Holding either guard for a guessed lifetime + // would leak capacity or release it early without an authoritative upstream close signal. + admission.release().await; + let pool_lease_healthy = lease.is_healthy(); + lease.release().await; + let result = result?; + if !(200..300).contains(&result.status_code) { + let response_body = execution_result_body(&result)?; + let downstream_headers = + sanitized_live_response_headers(&result.headers, response_body.preserves_wire_encoding); + warn!( + event_name = "codex_live_call_upstream_failed", + log_type = "ops", + trace_id = %request_context.trace_id, + provider_id = %attempt.plan.provider_id, + endpoint_id = %attempt.plan.endpoint_id, + key_id = %attempt.plan.key_id, + status_code = result.status_code, + elapsed_ms = result.telemetry.as_ref().and_then(|value| value.elapsed_ms), + "Codex Live call creation failed upstream" + ); + return Ok(Some(build_client_response_from_parts( + result.status_code, + &downstream_headers, + Body::from(response_body.bytes), + request_context.trace_id.as_str(), + Some(control_decision), + )?)); + } + if !pool_lease_healthy { + warn_live_call_orphaned( + request_context, + &attempt.plan, + &result, + "pool_lease_lost", + None, + ); + return Ok(Some(local_live_error( + request_context, + StatusCode::SERVICE_UNAVAILABLE, + "Codex Live provider lease expired during call creation", + )?)); + } + let response_body = match execution_result_body(&result) { + Ok(body) => body, + Err(error) => { + warn_live_call_orphaned( + request_context, + &attempt.plan, + &result, + "response_body_unavailable", + None, + ); + return Err(error); + } + }; + let Some(location) = header_value(&result.headers, "location") else { + warn_live_call_orphaned( + request_context, + &attempt.plan, + &result, + "location_missing", + None, + ); + return Ok(Some(local_live_error( + request_context, + StatusCode::BAD_GATEWAY, + "Codex Live upstream response did not include a call location", + )?)); + }; + let call_id = match extract_call_id_from_location(location) { + Ok(call_id) => call_id, + Err(error) => { + warn_live_call_orphaned( + request_context, + &attempt.plan, + &result, + "location_invalid", + Some(error.code()), + ); + return Ok(Some(local_live_error( + request_context, + StatusCode::BAD_GATEWAY, + error.client_message(), + )?)); + } + }; + let Some(auth_context) = control_decision.auth_context.as_ref() else { + warn_live_call_orphaned( + request_context, + &attempt.plan, + &result, + "auth_context_missing", + None, + ); + return Ok(Some(local_live_error( + request_context, + StatusCode::UNAUTHORIZED, + "Codex Live requires an authenticated gateway API key", + )?)); + }; + let registry = LiveCallRegistry::new(std::sync::Arc::clone(&state.runtime_state)); + if let Err(error) = registry + .register( + auth_context.user_id.as_str(), + auth_context.api_key_id.as_str(), + call_id.as_str(), + &binding, + ) + .await + { + warn_live_call_orphaned( + request_context, + &attempt.plan, + &result, + "binding_failed", + Some(error.kind()), + ); + return Ok(Some(local_live_error( + request_context, + StatusCode::SERVICE_UNAVAILABLE, + "Codex Live sideband binding is temporarily unavailable", + )?)); + } + info!( + event_name = "codex_live_call_created", + log_type = "event", + trace_id = %request_context.trace_id, + provider_id = %attempt.plan.provider_id, + endpoint_id = %attempt.plan.endpoint_id, + key_id = %attempt.plan.key_id, + client_model = %binding.client_model(), + status_code = result.status_code, + elapsed_ms = result.telemetry.as_ref().and_then(|value| value.elapsed_ms), + usage_unavailable = true, + "Codex Live created a bound WebRTC call" + ); + let downstream_location = format!("/v1/live/{call_id}"); + let downstream_headers = + sanitized_live_response_headers(&result.headers, response_body.preserves_wire_encoding); + Ok(Some(build_client_response_from_parts_with_mutator( + result.status_code, + &downstream_headers, + Body::from(response_body.bytes), + request_context.trace_id.as_str(), + Some(control_decision), + |headers| { + headers.insert( + http::header::LOCATION, + HeaderValue::from_str(downstream_location.as_str()) + .map_err(|error| GatewayError::Internal(error.to_string()))?, + ); + Ok(()) + }, + )?)) +} + +fn local_live_error( + request_context: &GatewayPublicRequestContext, + status: StatusCode, + message: &str, +) -> Result, GatewayError> { + build_local_http_error_response_with_request_path( + request_context.trace_id.as_str(), + request_context.control_decision.as_ref(), + Some("/v1/live"), + status, + message, + ) +} + +fn remove_headers(headers: &mut BTreeMap, names: &[&str]) { + headers.retain(|candidate, _| { + !names + .iter() + .any(|name| candidate.eq_ignore_ascii_case(name)) + }); +} + +fn header_value<'a>(headers: &'a BTreeMap, name: &str) -> Option<&'a str> { + headers + .iter() + .find(|(candidate, _)| candidate.eq_ignore_ascii_case(name)) + .map(|(_, value)| value.as_str()) +} + +fn prepare_live_call_request_headers(headers: &mut BTreeMap, content_type: &str) { + remove_headers( + headers, + &[ + "content-type", + "content-length", + "content-encoding", + "accept", + "accept-encoding", + ], + ); + headers.insert("content-type".to_string(), content_type.to_string()); + headers.insert("accept".to_string(), "application/sdp".to_string()); + headers.insert("accept-encoding".to_string(), "identity".to_string()); +} + +fn build_live_call_provider_body( + auth_mode: LiveAuthMode, + sdp: &str, + session: &serde_json::Value, +) -> Result<(String, String), GatewayError> { + let (content_type, bytes) = match auth_mode { + LiveAuthMode::ApiKey => { + let (content_type, bytes) = build_live_multipart(sdp, session); + (content_type, bytes) + } + LiveAuthMode::ChatGptOauth => { + let bytes = serde_json::to_vec(&json!({"sdp": sdp, "session": session})) + .map_err(|error| GatewayError::Internal(error.to_string()))?; + ("application/json".to_string(), bytes) + } + }; + Ok(( + content_type, + base64::engine::general_purpose::STANDARD.encode(bytes), + )) +} + +fn sanitized_live_response_headers( + headers: &BTreeMap, + preserves_wire_encoding: bool, +) -> BTreeMap { + let mut sanitized = headers.clone(); + remove_headers(&mut sanitized, &["location", "set-cookie", "set-cookie2"]); + if !preserves_wire_encoding { + remove_headers(&mut sanitized, &["content-length", "content-encoding"]); + } + sanitized +} + +fn warn_live_call_orphaned( + request_context: &GatewayPublicRequestContext, + plan: &ExecutionPlan, + result: &ExecutionResult, + reason: &'static str, + error_kind: Option<&'static str>, +) { + warn!( + event_name = "codex_live_call_orphaned", + log_type = "ops", + trace_id = %request_context.trace_id, + provider_id = %plan.provider_id, + endpoint_id = %plan.endpoint_id, + key_id = %plan.key_id, + status_code = result.status_code, + elapsed_ms = result.telemetry.as_ref().and_then(|value| value.elapsed_ms), + reason, + error_kind = error_kind.unwrap_or("none"), + "Codex Live upstream call succeeded but could not be safely exposed downstream" + ); +} + +struct LiveResponseBody { + bytes: Vec, + preserves_wire_encoding: bool, +} + +fn execution_result_body(result: &ExecutionResult) -> Result { + let Some(body) = result.body.as_ref() else { + return Ok(LiveResponseBody { + bytes: Vec::new(), + preserves_wire_encoding: false, + }); + }; + if let Some(encoded) = body.body_bytes_b64.as_deref() { + let bytes = base64::engine::general_purpose::STANDARD + .decode(encoded) + .map_err(|error| GatewayError::Internal(error.to_string()))?; + return Ok(LiveResponseBody { + bytes, + preserves_wire_encoding: true, + }); + } + let bytes = body + .json_body + .as_ref() + .map(serde_json::to_vec) + .transpose() + .map(|body| body.unwrap_or_default()) + .map_err(|error| GatewayError::Internal(error.to_string()))?; + Ok(LiveResponseBody { + bytes, + preserves_wire_encoding: false, + }) +} + +#[cfg(test)] +mod tests { + use aether_contracts::{ExecutionPlan, ExecutionResult, ResponseBody}; + use axum::body::to_bytes; + + use crate::control::{GatewayControlAuthContext, GatewayControlDecision}; + + use super::*; + + #[test] + fn preserved_wire_bytes_win_over_the_json_projection() { + let result = ExecutionResult { + request_id: "request".to_string(), + candidate_id: None, + status_code: 201, + headers: Default::default(), + response_observation: None, + body: Some(ResponseBody { + json_body: Some(json!({"projected": true})), + body_bytes_b64: Some( + base64::engine::general_purpose::STANDARD.encode(b"raw-sdp-answer"), + ), + }), + telemetry: None, + error: None, + }; + let body = execution_result_body(&result).unwrap(); + assert_eq!(body.bytes, b"raw-sdp-answer"); + assert!(body.preserves_wire_encoding); + } + + #[test] + fn live_call_bodies_are_bytes_only_and_do_not_enter_report_context() { + let session = json!({ + "model": "provider-live-model", + "instructions": "opaque private instructions", + "future_capability": {"enabled": true} + }); + let sdp = "v=0\r\no=private-live-offer"; + let provider_body_marker = json!({"model": "provider-live-model"}); + + for auth_mode in [LiveAuthMode::ApiKey, LiveAuthMode::ChatGptOauth] { + let (content_type, encoded) = + build_live_call_provider_body(auth_mode, sdp, &session).unwrap(); + let report_context = aether_ai_serving::augment_sync_report_context( + Some(json!({"trace_id": "trace-live"})), + &BTreeMap::new(), + &provider_body_marker, + ) + .unwrap() + .unwrap(); + assert!(report_context.get("provider_request_body").is_none()); + assert!(!report_context.to_string().contains("private-live-offer")); + assert!(!report_context + .to_string() + .contains("opaque private instructions")); + + let plan_body = aether_ai_serving::resolve_ai_passthrough_sync_request_body( + Some(provider_body_marker.clone()), + Some(encoded.clone()), + ); + assert!(plan_body.json_body.is_none()); + assert_eq!(plan_body.body_bytes_b64.as_deref(), Some(encoded.as_str())); + let usage_plan = ExecutionPlan { + request_id: "trace-live".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: "POST".to_string(), + url: "https://api.openai.com/v1/live".to_string(), + headers: BTreeMap::new(), + content_type: Some(content_type.clone()), + content_encoding: None, + body: plan_body, + stream: false, + client_api_format: "openai:responses".to_string(), + provider_api_format: "openai:responses".to_string(), + model_name: Some("provider-live-model".to_string()), + proxy: None, + transport_profile: None, + timeouts: None, + }; + let usage_seed = aether_usage_runtime::build_terminal_usage_context_seed( + &usage_plan, + Some(&report_context), + ); + assert!(usage_seed.provider_request.is_none()); + assert!(!usage_seed + .request_metadata + .as_ref() + .map(ToString::to_string) + .unwrap_or_default() + .contains("private-live-offer")); + + let wire = base64::engine::general_purpose::STANDARD + .decode(encoded) + .unwrap(); + match auth_mode { + LiveAuthMode::ApiKey => { + let parsed = parse_live_multipart(content_type.as_str(), wire.as_slice()) + .expect("API-key multipart should round-trip"); + assert_eq!(parsed.sdp, sdp); + assert_eq!(parsed.session, session); + } + LiveAuthMode::ChatGptOauth => { + assert_eq!(content_type, "application/json"); + let decoded: serde_json::Value = serde_json::from_slice(wire.as_slice()) + .expect("OAuth JSON should round-trip"); + assert_eq!(decoded["sdp"], sdp); + assert_eq!(decoded["session"], session); + assert_eq!(decoded["session"]["model"], "provider-live-model"); + assert_eq!( + decoded["session"]["future_capability"], + json!({"enabled": true}) + ); + } + } + } + } + + #[test] + fn live_call_request_headers_replace_stale_body_and_encoding_metadata() { + let mut headers = BTreeMap::from([ + ("Content-Type".to_string(), "stale".to_string()), + ("CONTENT-LENGTH".to_string(), "42".to_string()), + ("Content-Encoding".to_string(), "gzip".to_string()), + ("Accept".to_string(), "application/json".to_string()), + ("ACCEPT-ENCODING".to_string(), "br, gzip".to_string()), + ("x-future".to_string(), "opaque".to_string()), + ]); + prepare_live_call_request_headers(&mut headers, "multipart/form-data; boundary=live-test"); + + assert_eq!( + header_value(&headers, "content-type"), + Some("multipart/form-data; boundary=live-test") + ); + assert_eq!(header_value(&headers, "accept"), Some("application/sdp")); + assert_eq!(header_value(&headers, "accept-encoding"), Some("identity")); + assert_eq!(header_value(&headers, "content-length"), None); + assert_eq!(header_value(&headers, "content-encoding"), None); + assert_eq!(headers.get("x-future").map(String::as_str), Some("opaque")); + } + + #[test] + fn live_response_headers_never_expose_upstream_location_or_cookies() { + let headers = BTreeMap::from([ + ( + "Location".to_string(), + "https://upstream/v1/live/secret".to_string(), + ), + ("SET-COOKIE".to_string(), "session=secret".to_string()), + ("Set-Cookie2".to_string(), "legacy=secret".to_string()), + ("Content-Length".to_string(), "128".to_string()), + ("Content-Encoding".to_string(), "gzip".to_string()), + ("x-future".to_string(), "opaque".to_string()), + ]); + + let sanitized = sanitized_live_response_headers(&headers, true); + assert_eq!(header_value(&sanitized, "location"), None); + assert_eq!(header_value(&sanitized, "set-cookie"), None); + assert_eq!(header_value(&sanitized, "set-cookie2"), None); + assert_eq!(header_value(&sanitized, "content-length"), Some("128")); + assert_eq!(header_value(&sanitized, "content-encoding"), Some("gzip")); + assert_eq!(header_value(&sanitized, "x-future"), Some("opaque")); + } + + #[test] + fn rebuilt_live_response_body_drops_stale_length_and_encoding() { + let headers = BTreeMap::from([ + ("content-length".to_string(), "128".to_string()), + ("content-encoding".to_string(), "gzip".to_string()), + ("content-type".to_string(), "application/json".to_string()), + ]); + + let sanitized = sanitized_live_response_headers(&headers, false); + assert_eq!(header_value(&sanitized, "content-length"), None); + assert_eq!(header_value(&sanitized, "content-encoding"), None); + assert_eq!( + header_value(&sanitized, "content-type"), + Some("application/json") + ); + } + + #[tokio::test] + async fn finite_balance_post_live_fails_before_parsing_or_upstream_execution() { + let mut decision = GatewayControlDecision::synthetic( + "/v1/live", + Some("ai_public".to_string()), + Some("openai".to_string()), + Some("codex_live".to_string()), + Some("openai:responses".to_string()), + ); + decision.auth_context = Some(GatewayControlAuthContext { + user_id: "user-finite".to_string(), + api_key_id: "key-finite".to_string(), + username: Some("finite".to_string()), + api_key_name: Some("finite".to_string()), + balance_remaining: Some(1.25), + access_allowed: true, + user_rate_limit: None, + api_key_rate_limit: None, + api_key_is_standalone: false, + admin_bypass_limits: false, + local_rejection: None, + allowed_models: None, + ip_rules: None, + }); + let request_context = GatewayPublicRequestContext { + trace_id: "trace-live-finite".to_string(), + request_method: http::Method::POST, + request_path: "/v1/live".to_string(), + request_query_string: None, + request_content_type: None, + host_header: None, + control_decision: Some(decision), + }; + let (parts, _) = http::Request::builder() + .method(http::Method::POST) + .uri("/v1/live") + .body(()) + .unwrap() + .into_parts(); + let response = maybe_handle_live_http( + &AppState::new().expect("gateway state should build"), + &request_context, + &parts, + None, + &"127.0.0.1:65000".parse().unwrap(), + ) + .await + .unwrap() + .expect("Live HTTP route must produce a local rejection"); + assert_eq!(response.status(), StatusCode::NOT_IMPLEMENTED); + let body = to_bytes(response.into_body(), usize::MAX).await.unwrap(); + assert!(String::from_utf8_lossy(body.as_ref()).contains("finite-balance")); + } +} diff --git a/apps/aether-gateway/src/handlers/proxy/websocket/live/mod.rs b/apps/aether-gateway/src/handlers/proxy/websocket/live/mod.rs new file mode 100644 index 000000000..90a93640b --- /dev/null +++ b/apps/aether-gateway/src/handlers/proxy/websocket/live/mod.rs @@ -0,0 +1,80 @@ +//! Experimental Codex Frameless Bidi V3 (`/v1/live`) bridge. +//! +//! The public OpenAI Realtime and Responses WebSocket protocols are related +//! transport families, but this Codex protocol has a distinct event grammar. +//! Keeping it in an independent module prevents a `session.update` frame from +//! ever entering the Responses `response.create` state machine. + +mod http; +mod planner; +mod protocol; +mod registry; +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::control::GatewayControlDecision; +use crate::handlers::proxy::websocket::ingress::{ + prepare_authenticated_ai_websocket, AuthenticatedAiWebSocketUpgradePreparation, + WebSocketIngressSpec, +}; +use crate::handlers::proxy::websocket::session::LIVE_WEBSOCKET_SESSION_LIMITS; +use crate::{AppState, GatewayError}; + +pub(crate) use http::maybe_handle_live_http; + +/// Frameless Bidi currently exposes no stable token/cost usage object that can +/// be fed into Aether's settlement pipeline. Fail closed for finite-balance +/// principals instead of silently serving unmetered traffic. Standalone or +/// shared keys backed by an unlimited/no-wallet policy resolve without a +/// finite `balance_remaining` and remain eligible. +fn live_usage_accounting_is_safe(decision: &GatewayControlDecision) -> bool { + decision + .auth_context + .as_ref() + .is_some_and(|auth| auth.balance_remaining.is_none()) +} + +pub(crate) async fn live_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, + LIVE_WEBSOCKET_INGRESS_SPEC, + ) + .await? + { + AuthenticatedAiWebSocketUpgradePreparation::Rejected(response) => Ok(response), + AuthenticatedAiWebSocketUpgradePreparation::Ready(prepared) => { + let live = + match session::prepare_live_websocket(prepared.state(), prepared.context()).await { + Ok(live) => live, + Err(rejection) => { + return prepared.rejection_response(rejection.status(), rejection.message()) + } + }; + Ok(prepared.into_response_with( + ws, + LIVE_WEBSOCKET_SESSION_LIMITS, + live, + session::run_live_websocket, + )) + } + } +} + +const LIVE_WEBSOCKET_INGRESS_SPEC: WebSocketIngressSpec = WebSocketIngressSpec { + route_unavailable_message: "Codex Live WebSocket route is unavailable", +}; diff --git a/apps/aether-gateway/src/handlers/proxy/websocket/live/planner.rs b/apps/aether-gateway/src/handlers/proxy/websocket/live/planner.rs new file mode 100644 index 000000000..189cb148f --- /dev/null +++ b/apps/aether-gateway/src/handlers/proxy/websocket/live/planner.rs @@ -0,0 +1,1036 @@ +//! 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. + +use std::collections::BTreeSet; +use std::net::SocketAddr; +use std::sync::atomic::{AtomicBool, Ordering}; +use std::sync::Arc; +use std::time::Duration; + +use axum::http::header::{AUTHORIZATION, CONNECTION, CONTENT_TYPE, COOKIE, UPGRADE}; +use axum::http::{HeaderMap, HeaderName, HeaderValue, Method}; +use serde_json::json; +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, +}; +use crate::control::GatewayControlDecision; +use crate::headers::request_origin_from_headers_and_remote_addr; +use crate::privacy::RedactionSessionSlot; +use crate::{AppState, GatewayError}; + +use super::protocol::{validate_model, LiveProtocolError}; + +pub(super) const LIVE_ALPHA_HEADER_VALUE: &str = "quicksilver=v2"; +const CHATGPT_ACCOUNT_ID_HEADER: &str = "chatgpt-account-id"; +const CHATGPT_FEDRAMP_HEADER: &str = "x-openai-fedramp"; +const CHATGPT_SESSION_ID_HEADER: &str = "x-session-id"; +const OFFICIAL_CHATGPT_HOST: &str = "chatgpt.com"; +const LIVE_ROUTING_FINGERPRINT_DOMAIN: &[u8] = b"aether-codex-live-routing-v1"; +const MAX_LIVE_MODEL_BYTES: usize = 256; + +#[derive(Debug, Clone, Copy, PartialEq, Eq, serde::Serialize, serde::Deserialize)] +#[serde(rename_all = "snake_case")] +pub(super) enum LiveAuthMode { + ApiKey, + ChatGptOauth, +} + +#[derive(Debug)] +pub(super) struct PlannedLiveCandidate { + pub(super) execution: AiExecutionDecision, + pub(super) pinned_candidate: ResponsesWebSocketPinnedCandidate, + pub(super) client_model: String, + pub(super) provider_model: String, + pub(super) auth_mode: LiveAuthMode, + /// Domain-separated digest of the stable upstream auth/account/origin + /// identity. It deliberately excludes bearer tokens so an OAuth refresh + /// does not invalidate an in-flight WebRTC call. + pub(super) routing_fingerprint: String, +} + +/// Cancellation-safe owner for the scheduler's distributed pool-key lease. +/// Live does not enter the ordinary HTTP/Responses attempt lifecycle, so it +/// must hold and release the lease explicitly for the call or socket lifetime. +pub(super) struct LivePoolLeaseGuard { + state: AppState, + report_context: Option, + renewal_task: Option>, + healthy: Arc, + armed: bool, +} + +impl LivePoolLeaseGuard { + pub(super) fn new(state: &AppState, candidate: &PlannedLiveCandidate) -> Self { + let report_context = candidate.execution.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 LivePoolLeaseGuard { + 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_live_candidate( + state: &AppState, + trace_id: &str, + decision: &GatewayControlDecision, + headers: &HeaderMap, + remote_addr: &SocketAddr, + client_model: &str, + pinned_candidate: Option<&ResponsesWebSocketPinnedCandidate>, +) -> Result, GatewayError> { + if validate_model(client_model).is_err() || client_model.len() > MAX_LIVE_MODEL_BYTES { + return Ok(None); + } + let parts = build_live_planning_parts(headers, remote_addr); + let body = json!({"model": client_model, "input": []}); + let planned = maybe_build_responses_websocket_decision( + state, + &parts, + trace_id, + decision, + None, + &body, + None::<&BTreeSet>, + None::<&BTreeSet>, + pinned_candidate, + ) + .await?; + let Some(planned) = planned else { + 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") + { + crate::orchestration::release_pool_key_lease_from_report_context( + state, + execution.report_context.as_ref(), + ) + .await; + return Ok(None); + } + let Some(pinned_candidate) = ResponsesWebSocketPinnedCandidate::from_decision(&execution) + else { + crate::orchestration::release_pool_key_lease_from_report_context( + state, + execution.report_context.as_ref(), + ) + .await; + return Ok(None); + }; + let provider_model = execution + .mapped_model + .as_deref() + .or(execution.model_name.as_deref()) + .map(str::trim) + .filter(|model| validate_model(model).is_ok()) + .map(str::to_string); + let Some(provider_model) = provider_model else { + crate::orchestration::release_pool_key_lease_from_report_context( + state, + execution.report_context.as_ref(), + ) + .await; + return Ok(None); + }; + let Some(auth_mode) = live_auth_mode(provider_type, effective_auth_type.as_str()) else { + crate::orchestration::release_pool_key_lease_from_report_context( + state, + execution.report_context.as_ref(), + ) + .await; + return Ok(None); + }; + apply_live_headers(&mut execution.provider_request_headers, trace_id); + let routing_fingerprint = + match live_routing_fingerprint(&execution, effective_auth_type.as_str(), auth_mode) { + Ok(fingerprint) => fingerprint, + Err(_) => { + crate::orchestration::release_pool_key_lease_from_report_context( + state, + execution.report_context.as_ref(), + ) + .await; + return Ok(None); + } + }; + Ok(Some(PlannedLiveCandidate { + execution, + pinned_candidate, + client_model: client_model.to_string(), + provider_model, + auth_mode, + routing_fingerprint, + })) +} + +pub(super) fn direct_live_websocket_url( + candidate: &PlannedLiveCandidate, +) -> Result { + if candidate.auth_mode == LiveAuthMode::ChatGptOauth { + return Err(LiveProtocolError::OauthDirectWebSocketUnsupported); + } + replace_responses_suffix( + candidate.execution.upstream_url.as_deref(), + &["live"], + Some(("model", candidate.provider_model.as_str())), + ) +} + +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) + } + LiveAuthMode::ChatGptOauth => { + let source = + validated_official_chatgpt_url(candidate.execution.upstream_url.as_deref())?; + replace_responses_suffix( + Some(source.as_str()), + &["realtime", "calls"], + Some(("intent", "quicksilver")), + ) + } + .and_then(|raw| { + let mut url = + Url::parse(raw.as_str()).map_err(|_| LiveProtocolError::InvalidUpstreamUrl)?; + replace_url_query_pair(&mut url, "architecture", "avas"); + Ok(url.to_string()) + }), + } +} + +pub(super) fn live_sideband_url( + candidate: &PlannedLiveCandidate, + call_id: &str, +) -> Result { + super::protocol::validate_call_id(call_id)?; + match candidate.auth_mode { + LiveAuthMode::ApiKey => replace_responses_suffix( + candidate.execution.upstream_url.as_deref(), + &["live", call_id], + None, + ), + // The Codex ChatGPT call creation endpoint returns an OpenAI Realtime + // call ID. Current Codex connects its sideband to this API origin even + // when call creation used the ChatGPT OAuth backend. + LiveAuthMode::ChatGptOauth => { + validated_official_chatgpt_url(candidate.execution.upstream_url.as_deref())?; + Ok(format!("https://api.openai.com/v1/live/{call_id}")) + } + } +} + +pub(super) fn apply_live_headers( + headers: &mut std::collections::BTreeMap, + seed: &str, +) { + replace_header(headers, "openai-alpha", LIVE_ALPHA_HEADER_VALUE); + let session_id = find_header(headers, "x-session-id") + .or_else(|| find_header(headers, "thread-id")) + .or_else(|| find_header(headers, "session-id")) + .map(str::to_string) + .filter(|value| !value.trim().is_empty()) + .unwrap_or_else(|| seed.to_string()); + replace_header(headers, "x-session-id", session_id.as_str()); +} + +pub(super) fn build_live_stream_admission_attempt( + candidate: &PlannedLiveCandidate, + headers: &HeaderMap, + remote_addr: &SocketAddr, + upstream_url: String, +) -> Result, GatewayError> { + let parts = build_live_planning_parts(headers, remote_addr); + let body = json!({"model": candidate.client_model.as_str(), "input": []}); + let mut execution = candidate.execution.clone(); + execution.upstream_url = Some(upstream_url); + execution.upstream_is_stream = true; + build_standard_stream_plan_from_decision(&parts, &body, execution, false) +} + +fn live_auth_mode(provider_type: &str, effective_auth_type: &str) -> Option { + match effective_auth_type.trim().to_ascii_lowercase().as_str() { + "api_key" | "bearer" => Some(LiveAuthMode::ApiKey), + "oauth" if provider_type.eq_ignore_ascii_case("codex") => Some(LiveAuthMode::ChatGptOauth), + _ => None, + } +} + +fn live_routing_fingerprint( + decision: &AiExecutionDecision, + effective_auth_type: &str, + auth_mode: LiveAuthMode, +) -> Result { + let raw_url = decision + .upstream_url + .as_deref() + .ok_or(LiveProtocolError::MissingUpstreamUrl)?; + let url = Url::parse(raw_url).map_err(|_| LiveProtocolError::InvalidUpstreamUrl)?; + if url.host_str().is_none() + || !url.username().is_empty() + || url.password().is_some() + || url.fragment().is_some() + { + return Err(LiveProtocolError::InvalidUpstreamUrl); + } + let path = url.path().trim_end_matches('/'); + let path_family = path + .strip_suffix("/responses") + .ok_or(LiveProtocolError::InvalidUpstreamUrl)?; + if auth_mode == LiveAuthMode::ChatGptOauth { + validated_official_chatgpt_url(Some(raw_url))?; + } + + let mut hasher = Sha256::new(); + hasher.update(LIVE_ROUTING_FINGERPRINT_DOMAIN); + hasher.update([0]); + for value in [ + effective_auth_type.trim(), + url.scheme(), + url.host_str().unwrap_or_default(), + path_family, + ] { + hasher.update(value.as_bytes()); + hasher.update([0]); + } + hasher.update( + url.port_or_known_default() + .unwrap_or_default() + .to_be_bytes(), + ); + hasher.update([0]); + hasher.update(canonical_live_route_query(&url).as_bytes()); + hasher.update([0]); + + let session_id = find_header( + &decision.provider_request_headers, + CHATGPT_SESSION_ID_HEADER, + ) + .map(str::trim) + .filter(|value| !value.is_empty()) + .ok_or(LiveProtocolError::InvalidUpstreamUrl)?; + hasher.update(session_id.as_bytes()); + hasher.update([0]); + + if auth_mode == LiveAuthMode::ChatGptOauth { + for (name, required) in [ + (CHATGPT_ACCOUNT_ID_HEADER, true), + (CHATGPT_FEDRAMP_HEADER, false), + ] { + let value = find_header(&decision.provider_request_headers, name) + .map(str::trim) + .filter(|value| !value.is_empty()); + if required && value.is_none() { + return Err(LiveProtocolError::InvalidUpstreamUrl); + } + hasher.update(value.unwrap_or_default().as_bytes()); + hasher.update([0]); + } + } + Ok(format!("{:x}", hasher.finalize())) +} + +fn validated_official_chatgpt_url(raw: Option<&str>) -> Result { + let url = Url::parse(raw.ok_or(LiveProtocolError::MissingUpstreamUrl)?) + .map_err(|_| LiveProtocolError::InvalidUpstreamUrl)?; + let official = url.scheme() == "https" + && url + .host_str() + .is_some_and(|host| host.eq_ignore_ascii_case(OFFICIAL_CHATGPT_HOST)) + && url.port_or_known_default() == Some(443) + && url.username().is_empty() + && url.password().is_none() + && url.fragment().is_none() + && url.path().trim_end_matches('/').strip_suffix("/responses") + == Some("/backend-api/codex"); + if !official { + return Err(LiveProtocolError::OauthUpstreamUnsupported); + } + Ok(url) +} + +fn replace_responses_suffix( + raw: Option<&str>, + suffix: &[&str], + query: Option<(&str, &str)>, +) -> Result { + let mut url = Url::parse(raw.ok_or(LiveProtocolError::MissingUpstreamUrl)?) + .map_err(|_| LiveProtocolError::InvalidUpstreamUrl)?; + if url.host_str().is_none() + || !url.username().is_empty() + || url.password().is_some() + || url.fragment().is_some() + { + return Err(LiveProtocolError::InvalidUpstreamUrl); + } + if url.path_segments().and_then(Iterator::last) != Some("responses") { + return Err(LiveProtocolError::InvalidUpstreamUrl); + } + { + let mut path = url + .path_segments_mut() + .map_err(|_| LiveProtocolError::InvalidUpstreamUrl)?; + path.pop_if_empty(); + path.pop(); + for segment in suffix { + path.push(segment); + } + } + if let Some((name, value)) = query { + replace_url_query_pair(&mut url, name, value); + } + Ok(url.to_string()) +} + +fn replace_url_query_pair(url: &mut Url, name: &str, value: &str) { + let retained = url + .query_pairs() + .filter(|(candidate, _)| !candidate.eq_ignore_ascii_case(name)) + .map(|(key, value)| (key.into_owned(), value.into_owned())) + .collect::>(); + url.set_query(None); + let mut query = url.query_pairs_mut(); + for (key, value) in retained { + query.append_pair(key.as_str(), value.as_str()); + } + query.append_pair(name, value); +} + +fn canonical_live_route_query(url: &Url) -> String { + if url.query().is_none() { + return String::new(); + } + let mut pairs = url + .query_pairs() + .map(|(name, value)| { + let value = if live_route_query_key_is_credential(name.as_ref()) { + // Bind the authentication mechanism without pinning a rotating + // credential value into the WebRTC call identity. + String::new() + } else { + value.into_owned() + }; + (name.into_owned(), value) + }) + .collect::>(); + pairs.sort_unstable(); + + let mut serializer = form_urlencoded::Serializer::new(String::new()); + for (name, value) in pairs { + serializer.append_pair(name.as_str(), value.as_str()); + } + serializer.finish() +} + +fn live_route_query_key_is_credential(name: &str) -> bool { + matches!( + name.to_ascii_lowercase().as_str(), + "key" + | "api_key" + | "api-key" + | "x-api-key" + | "x-goog-api-key" + | "access_token" + | "authorization" + | "token" + | "oauth_token" + | "client_secret" + | "secret_key" + | "signature" + | "sig" + ) +} + +fn build_live_planning_parts( + headers: &HeaderMap, + remote_addr: &SocketAddr, +) -> http::request::Parts { + let mut request = http::Request::builder() + .method(Method::POST) + .uri("/v1/responses") + .body(()) + .expect("the fixed Live planning request must be valid"); + *request.headers_mut() = sanitize_live_planning_headers(headers.clone()); + request + .headers_mut() + .insert(CONTENT_TYPE, HeaderValue::from_static("application/json")); + request + .extensions_mut() + .insert(request_origin_from_headers_and_remote_addr( + headers, + remote_addr, + )); + request + .extensions_mut() + .insert(RedactionSessionSlot::default()); + request.into_parts().0 +} + +fn sanitize_live_planning_headers(mut headers: HeaderMap) -> HeaderMap { + let connection_names = headers + .get_all(CONNECTION) + .iter() + .filter_map(|value| value.to_str().ok()) + .flat_map(|value| value.split(',')) + .filter_map(|name| HeaderName::from_bytes(name.trim().as_bytes()).ok()) + .collect::>(); + for name in connection_names { + headers.remove(name); + } + for name in [AUTHORIZATION, CONNECTION, COOKIE, UPGRADE] { + headers.remove(name); + } + for name in [ + "api-key", + "x-api-key", + "x-goog-api-key", + "proxy-authorization", + "proxy-connection", + "keep-alive", + ] { + headers.remove(name); + } + let websocket_headers = headers + .keys() + .filter(|name| name.as_str().starts_with("sec-websocket-")) + .cloned() + .collect::>(); + for name in websocket_headers { + headers.remove(name); + } + headers +} + +fn find_header<'a>( + headers: &'a std::collections::BTreeMap, + name: &str, +) -> Option<&'a str> { + headers + .iter() + .find(|(candidate, _)| candidate.eq_ignore_ascii_case(name)) + .map(|(_, value)| value.as_str()) +} + +fn replace_header( + headers: &mut std::collections::BTreeMap, + name: &str, + value: &str, +) { + headers.retain(|candidate, _| !candidate.eq_ignore_ascii_case(name)); + headers.insert(name.to_string(), value.to_string()); +} + +#[cfg(test)] +mod tests { + use std::collections::BTreeMap; + + use aether_provider_transport::snapshot::{ + GatewayProviderTransportEndpoint, GatewayProviderTransportKey, + GatewayProviderTransportProvider, GatewayProviderTransportSnapshot, + }; + + use super::*; + + fn candidate(url: &str, auth_mode: LiveAuthMode) -> PlannedLiveCandidate { + let execution: AiExecutionDecision = serde_json::from_value(json!({ + "action": "stream", + "provider_id": "provider-1", + "endpoint_id": "endpoint-1", + "key_id": "key-1", + "upstream_url": url, + "provider_type": "codex", + "provider_request_headers": {"authorization": "Bearer secret"} + })) + .expect("decision should deserialize"); + PlannedLiveCandidate { + execution, + pinned_candidate: ResponsesWebSocketPinnedCandidate::new( + "provider-1", + "endpoint-1", + "key-1", + ) + .unwrap(), + client_model: "global-model".to_string(), + provider_model: "provider-model".to_string(), + auth_mode, + routing_fingerprint: "0".repeat(64), + } + } + + fn transport_with_auth_override( + default_auth_type: &str, + auth_type_by_format: Option, + ) -> GatewayProviderTransportSnapshot { + GatewayProviderTransportSnapshot { + provider: GatewayProviderTransportProvider { + id: "provider-1".to_string(), + name: "codex".to_string(), + provider_type: "codex".to_string(), + website: None, + is_active: true, + keep_priority_on_conversion: false, + enable_format_conversion: false, + concurrent_limit: None, + max_retries: None, + proxy: None, + request_timeout_secs: None, + stream_first_byte_timeout_secs: None, + config: None, + }, + 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()), + is_active: true, + base_url: "https://chatgpt.com/backend-api/codex".to_string(), + header_rules: None, + body_rules: None, + max_retries: None, + custom_path: None, + config: None, + format_acceptance_config: None, + proxy: None, + }, + key: GatewayProviderTransportKey { + id: "key-1".to_string(), + provider_id: "provider-1".to_string(), + name: "key".to_string(), + auth_type: default_auth_type.to_string(), + is_active: true, + api_formats: Some(vec!["openai:responses".to_string()]), + auth_type_by_format, + allow_auth_channel_mismatch_formats: None, + allowed_models: None, + capabilities: None, + rate_multipliers: None, + global_priority_by_format: None, + expires_at_unix_secs: None, + proxy: None, + fingerprint: None, + upstream_metadata: None, + decrypted_api_key: "rotating-secret".to_string(), + decrypted_auth_config: None, + }, + } + } + + fn oauth_fingerprint_decision( + token: &str, + account_id: &str, + fedramp: &str, + session_id: &str, + ) -> AiExecutionDecision { + let mut decision = candidate( + "https://chatgpt.com/backend-api/codex/responses", + LiveAuthMode::ChatGptOauth, + ) + .execution; + decision.provider_request_headers = BTreeMap::from([ + ("authorization".to_string(), format!("Bearer {token}")), + ("chatgpt-account-id".to_string(), account_id.to_string()), + ("x-openai-fedramp".to_string(), fedramp.to_string()), + ("x-session-id".to_string(), session_id.to_string()), + ]); + decision + } + + #[test] + fn format_auth_override_selects_the_effective_live_auth_mode() { + let overridden = + transport_with_auth_override("oauth", Some(json!({"openai:responses": "bearer"}))); + let effective = + aether_provider_transport::auth::resolve_local_auth_type_for_transport_format( + &overridden, + ); + assert_eq!(effective, "bearer"); + assert_eq!( + live_auth_mode( + overridden.provider.provider_type.as_str(), + effective.as_str() + ), + Some(LiveAuthMode::ApiKey) + ); + + let oauth = transport_with_auth_override("oauth", None); + let effective = + aether_provider_transport::auth::resolve_local_auth_type_for_transport_format(&oauth); + assert_eq!(effective, "oauth"); + assert_eq!( + live_auth_mode(oauth.provider.provider_type.as_str(), effective.as_str()), + Some(LiveAuthMode::ChatGptOauth) + ); + } + + #[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", + LiveAuthMode::ApiKey, + ); + candidate.provider_model = "upstream/model + future".to_string(); + let direct = Url::parse(direct_live_websocket_url(&candidate).unwrap().as_str()).unwrap(); + assert_eq!(direct.path(), "/v1/live"); + assert_eq!( + direct.query_pairs().collect::>(), + vec![ + ("api-version".into(), "2026-08-01".into()), + ("model".into(), "upstream/model + future".into()), + ] + ); + assert_eq!( + direct + .query_pairs() + .filter(|(name, _)| name.eq_ignore_ascii_case("model")) + .count(), + 1 + ); + assert!(!direct.as_str().contains("global-model")); + assert_eq!( + live_call_url(&candidate).unwrap(), + "https://api.example.test/v1/live?api-version=2026-08-01&model=stale&MODEL=duplicate" + ); + assert_eq!( + live_sideband_url(&candidate, "rtc_abc-123").unwrap(), + "https://api.example.test/v1/live/rtc_abc-123?api-version=2026-08-01&model=stale&MODEL=duplicate" + ); + } + + #[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", + LiveAuthMode::ChatGptOauth, + ); + let call = Url::parse(live_call_url(&candidate).unwrap().as_str()).unwrap(); + assert_eq!(call.path(), "/backend-api/codex/realtime/calls"); + assert_eq!( + call.query_pairs().collect::>(), + vec![ + ("api-version".into(), "2026-08-01".into()), + ("intent".into(), "quicksilver".into()), + ("architecture".into(), "avas".into()), + ] + ); + assert_eq!( + call.query_pairs() + .filter(|(name, _)| name.eq_ignore_ascii_case("intent")) + .count(), + 1 + ); + assert_eq!( + call.query_pairs() + .filter(|(name, _)| name.eq_ignore_ascii_case("architecture")) + .count(), + 1 + ); + assert_eq!( + live_sideband_url(&candidate, "rtc_call_1").unwrap(), + "https://api.openai.com/v1/live/rtc_call_1" + ); + assert_eq!( + direct_live_websocket_url(&candidate), + Err(LiveProtocolError::OauthDirectWebSocketUnsupported) + ); + } + + #[test] + fn chatgpt_oauth_live_fails_closed_for_custom_backend_origins() { + let candidate = candidate( + "https://relay.example/backend-api/codex/responses", + LiveAuthMode::ChatGptOauth, + ); + assert_eq!( + live_call_url(&candidate), + Err(LiveProtocolError::OauthUpstreamUnsupported) + ); + assert_eq!( + live_sideband_url(&candidate, "rtc_custom_backend"), + Err(LiveProtocolError::OauthUpstreamUnsupported) + ); + assert_eq!( + live_routing_fingerprint(&candidate.execution, "oauth", LiveAuthMode::ChatGptOauth), + Err(LiveProtocolError::OauthUpstreamUnsupported) + ); + } + + #[test] + fn routing_fingerprint_ignores_token_rotation_but_binds_oauth_identity() { + let baseline = + oauth_fingerprint_decision("access-token-1", "account-1", "true", "session-1"); + let refreshed = + oauth_fingerprint_decision("access-token-2", "account-1", "true", "session-1"); + let fingerprint = + live_routing_fingerprint(&baseline, "oauth", LiveAuthMode::ChatGptOauth).unwrap(); + assert_eq!( + live_routing_fingerprint(&refreshed, "oauth", LiveAuthMode::ChatGptOauth).unwrap(), + fingerprint, + "access-token refresh must not invalidate an established sideband binding" + ); + + for changed in [ + oauth_fingerprint_decision("access-token-2", "account-2", "true", "session-1"), + oauth_fingerprint_decision("access-token-2", "account-1", "false", "session-1"), + oauth_fingerprint_decision("access-token-2", "account-1", "true", "session-2"), + ] { + assert_ne!( + live_routing_fingerprint(&changed, "oauth", LiveAuthMode::ChatGptOauth).unwrap(), + fingerprint + ); + } + } + + #[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; + first.provider_request_headers.extend([ + ("authorization".to_string(), "Bearer token-1".to_string()), + ("x-session-id".to_string(), "session-1".to_string()), + ]); + let mut refreshed = first.clone(); + refreshed + .provider_request_headers + .insert("authorization".to_string(), "Bearer token-2".to_string()); + assert_eq!( + live_routing_fingerprint(&first, "bearer", LiveAuthMode::ApiKey).unwrap(), + live_routing_fingerprint(&refreshed, "bearer", LiveAuthMode::ApiKey).unwrap() + ); + + let mut changed_origin = + candidate("https://api-b.example/v1/responses", LiveAuthMode::ApiKey).execution; + changed_origin + .provider_request_headers + .insert("x-session-id".to_string(), "session-1".to_string()); + assert_ne!( + live_routing_fingerprint(&first, "bearer", LiveAuthMode::ApiKey).unwrap(), + live_routing_fingerprint(&changed_origin, "bearer", LiveAuthMode::ApiKey).unwrap() + ); + + let mut changed_session = refreshed; + changed_session + .provider_request_headers + .insert("x-session-id".to_string(), "session-2".to_string()); + assert_ne!( + live_routing_fingerprint(&first, "bearer", LiveAuthMode::ApiKey).unwrap(), + live_routing_fingerprint(&changed_session, "bearer", LiveAuthMode::ApiKey).unwrap() + ); + + let missing_session = + candidate("https://api-a.example/v1/responses", LiveAuthMode::ApiKey).execution; + assert_eq!( + live_routing_fingerprint(&missing_session, "bearer", LiveAuthMode::ApiKey), + Err(LiveProtocolError::InvalidUpstreamUrl) + ); + } + + #[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", + LiveAuthMode::ApiKey, + ) + .execution; + baseline + .provider_request_headers + .insert("x-session-id".to_string(), "session-1".to_string()); + let fingerprint = + 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", + LiveAuthMode::ApiKey, + ) + .execution; + reordered + .provider_request_headers + .insert("x-session-id".to_string(), "session-1".to_string()); + assert_eq!( + live_routing_fingerprint(&reordered, "bearer", LiveAuthMode::ApiKey).unwrap(), + fingerprint, + "query order and query credentials must not change the stable route identity" + ); + + 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", + LiveAuthMode::ApiKey, + ) + .execution; + changed_route + .provider_request_headers + .insert("x-session-id".to_string(), "session-1".to_string()); + assert_ne!( + live_routing_fingerprint(&changed_route, "bearer", LiveAuthMode::ApiKey).unwrap(), + fingerprint, + "a non-sensitive endpoint route query change must invalidate the binding" + ); + } + + #[test] + fn live_headers_force_quicksilver_and_preserve_a_stable_session_identity() { + let mut headers = BTreeMap::from([ + ("OpenAI-Alpha".to_string(), "wrong".to_string()), + ("thread-id".to_string(), "thread-stable".to_string()), + ]); + apply_live_headers(&mut headers, "trace-fallback"); + assert_eq!( + headers.get("openai-alpha").map(String::as_str), + Some(LIVE_ALPHA_HEADER_VALUE) + ); + assert_eq!( + headers.get("x-session-id").map(String::as_str), + Some("thread-stable") + ); + assert_eq!( + headers + .keys() + .filter(|name| name.eq_ignore_ascii_case("openai-alpha")) + .count(), + 1 + ); + + let mut legacy_session_header = BTreeMap::from([ + ("Session-Id".to_string(), "legacy-stable".to_string()), + ("chatgpt-account-id".to_string(), "account-1".to_string()), + ]); + apply_live_headers(&mut legacy_session_header, "trace-fallback"); + assert_eq!( + legacy_session_header + .get("x-session-id") + .map(String::as_str), + Some("legacy-stable") + ); + assert_eq!( + legacy_session_header + .get("chatgpt-account-id") + .map(String::as_str), + Some("account-1") + ); + } + + #[test] + fn live_urls_reject_credentials_invalid_suffixes_and_call_ids() { + let credentials = candidate( + "https://token@example.test/v1/responses", + LiveAuthMode::ApiKey, + ); + assert_eq!( + direct_live_websocket_url(&credentials), + Err(LiveProtocolError::InvalidUpstreamUrl) + ); + + let wrong_suffix = candidate( + "https://api.example.test/v1/chat/completions", + LiveAuthMode::ApiKey, + ); + assert_eq!( + live_call_url(&wrong_suffix), + Err(LiveProtocolError::InvalidUpstreamUrl) + ); + assert_eq!( + live_sideband_url(&wrong_suffix, "rtc/escape"), + Err(LiveProtocolError::InvalidCallId) + ); + + let fragment = candidate( + "https://api.example.test/v1/responses#not-sent-upstream", + LiveAuthMode::ApiKey, + ); + assert_eq!( + direct_live_websocket_url(&fragment), + Err(LiveProtocolError::InvalidUpstreamUrl) + ); + assert_eq!( + live_call_url(&fragment), + Err(LiveProtocolError::InvalidUpstreamUrl) + ); + + let mut fragment_fingerprint = fragment.execution; + fragment_fingerprint + .provider_request_headers + .insert("x-session-id".to_string(), "session-1".to_string()); + assert_eq!( + live_routing_fingerprint(&fragment_fingerprint, "bearer", LiveAuthMode::ApiKey), + Err(LiveProtocolError::InvalidUpstreamUrl) + ); + } +} diff --git a/apps/aether-gateway/src/handlers/proxy/websocket/live/protocol.rs b/apps/aether-gateway/src/handlers/proxy/websocket/live/protocol.rs new file mode 100644 index 000000000..cddf33938 --- /dev/null +++ b/apps/aether-gateway/src/handlers/proxy/websocket/live/protocol.rs @@ -0,0 +1,766 @@ +//! Minimal validation at the Codex Live trust boundary. +//! +//! Live events remain opaque after the initial discriminator. This module owns +//! only bounded identifiers, multipart framing and the first `session.update` +//! check; it intentionally does not copy the evolving Codex event schema. + +use axum::http::StatusCode; +use serde_json::Value; + +const MAX_MODEL_BYTES: usize = 256; +const MAX_CALL_ID_BYTES: usize = 256; +const MAX_BOUNDARY_BYTES: usize = 70; +const MAX_MULTIPART_BODY_BYTES: usize = 1024 * 1024; +const MAX_SDP_BYTES: usize = 512 * 1024; +const MAX_SESSION_BYTES: usize = 256 * 1024; +const MAX_PART_HEADERS_BYTES: usize = 8 * 1024; + +#[derive(Debug, Clone, PartialEq, Eq, thiserror::Error)] +pub(super) enum LiveProtocolError { + #[error("missing Live upstream URL")] + MissingUpstreamUrl, + #[error("invalid Live upstream URL")] + InvalidUpstreamUrl, + #[error("ChatGPT OAuth does not support direct Codex Live WebSocket")] + OauthDirectWebSocketUnsupported, + #[error("ChatGPT OAuth Codex Live requires the official backend origin")] + OauthUpstreamUnsupported, + #[error("invalid Live model query")] + InvalidModelQuery, + #[error("invalid Live model")] + InvalidModel, + #[error("invalid Live call ID")] + InvalidCallId, + #[error("unsupported Live media type")] + UnsupportedMediaType, + #[error("invalid Live multipart boundary")] + InvalidBoundary, + #[error("Live multipart body is too large")] + MultipartBodyTooLarge, + #[error("malformed Live multipart body")] + MalformedMultipart, + #[error("unexpected Live multipart part")] + UnexpectedMultipartPart, + #[error("duplicate Live multipart part")] + DuplicateMultipartPart, + #[error("missing Live SDP part")] + MissingSdp, + #[error("invalid Live SDP part")] + InvalidSdp, + #[error("Live SDP is too large")] + SdpTooLarge, + #[error("missing Live session part")] + MissingSession, + #[error("invalid Live session JSON")] + InvalidSession, + #[error("Live session is too large")] + SessionTooLarge, + #[error("invalid initial Live JSON event")] + InvalidInitialEvent, + #[error("initial Live event must be session.update")] + ExpectedSessionUpdate, + #[error("initial Live event must be text")] + InitialEventMustBeText, + #[error("initial Live client read failed")] + InitialClientReadFailed, + #[error("timed out waiting for initial Live session.update")] + InitialSessionUpdateTimeout, + #[error("invalid Live call location")] + InvalidCallLocation, +} + +impl LiveProtocolError { + pub(super) const fn status_code(&self) -> StatusCode { + match self { + Self::MultipartBodyTooLarge | Self::SdpTooLarge | Self::SessionTooLarge => { + StatusCode::PAYLOAD_TOO_LARGE + } + Self::UnsupportedMediaType => StatusCode::UNSUPPORTED_MEDIA_TYPE, + Self::MissingUpstreamUrl + | Self::InvalidUpstreamUrl + | Self::OauthUpstreamUnsupported + | Self::InvalidCallLocation => StatusCode::BAD_GATEWAY, + Self::InitialSessionUpdateTimeout => StatusCode::REQUEST_TIMEOUT, + _ => StatusCode::BAD_REQUEST, + } + } + + pub(super) const fn code(&self) -> &'static str { + match self { + Self::MissingUpstreamUrl => "codex_live_upstream_url_missing", + Self::InvalidUpstreamUrl => "codex_live_upstream_url_invalid", + Self::OauthDirectWebSocketUnsupported => "codex_live_oauth_direct_unsupported", + Self::OauthUpstreamUnsupported => "codex_live_oauth_upstream_unsupported", + Self::InvalidModelQuery => "codex_live_model_query_invalid", + Self::InvalidModel => "codex_live_model_invalid", + Self::InvalidCallId => "codex_live_call_id_invalid", + Self::UnsupportedMediaType => "codex_live_media_type_unsupported", + Self::InvalidBoundary => "codex_live_boundary_invalid", + Self::MultipartBodyTooLarge => "codex_live_body_too_large", + Self::MalformedMultipart => "codex_live_multipart_invalid", + Self::UnexpectedMultipartPart => "codex_live_multipart_part_unexpected", + Self::DuplicateMultipartPart => "codex_live_multipart_part_duplicate", + Self::MissingSdp => "codex_live_sdp_missing", + Self::InvalidSdp => "codex_live_sdp_invalid", + Self::SdpTooLarge => "codex_live_sdp_too_large", + Self::MissingSession => "codex_live_session_missing", + Self::InvalidSession => "codex_live_session_invalid", + Self::SessionTooLarge => "codex_live_session_too_large", + Self::InvalidInitialEvent => "codex_live_initial_event_invalid", + Self::ExpectedSessionUpdate => "codex_live_expected_session_update", + Self::InitialEventMustBeText => "codex_live_initial_event_must_be_text", + Self::InitialClientReadFailed => "codex_live_initial_client_read_failed", + Self::InitialSessionUpdateTimeout => "codex_live_initial_session_update_timeout", + Self::InvalidCallLocation => "codex_live_call_location_invalid", + } + } + + pub(super) const fn client_message(&self) -> &'static str { + match self { + Self::MissingUpstreamUrl | Self::InvalidUpstreamUrl => { + "Codex Live provider URL is invalid" + } + Self::OauthDirectWebSocketUnsupported => { + "Direct Codex Live WebSocket requires an API-key provider; use WebRTC for ChatGPT OAuth" + } + Self::OauthUpstreamUnsupported => { + "ChatGPT OAuth Codex Live requires the official ChatGPT backend" + } + Self::InvalidModelQuery => { + "Codex Live WebSocket requires exactly one model query parameter" + } + Self::InvalidModel => { + "Codex Live model must be a non-empty identifier no longer than 256 bytes" + } + Self::InvalidCallId => "Codex Live call ID is invalid", + Self::UnsupportedMediaType => { + "Codex Live WebRTC call creation requires multipart/form-data" + } + Self::InvalidBoundary => "Codex Live multipart boundary is invalid", + Self::MultipartBodyTooLarge => "Codex Live WebRTC offer exceeds the 1 MiB limit", + Self::MalformedMultipart => "Codex Live multipart body is malformed", + Self::UnexpectedMultipartPart => { + "Codex Live multipart body may contain only sdp and session parts" + } + Self::DuplicateMultipartPart => "Codex Live multipart part is duplicated", + Self::MissingSdp => "Codex Live multipart body is missing the sdp part", + Self::InvalidSdp => "Codex Live sdp part must be non-empty UTF-8", + Self::SdpTooLarge => "Codex Live sdp part exceeds the 512 KiB limit", + Self::MissingSession => "Codex Live multipart body is missing the session part", + Self::InvalidSession => "Codex Live session part must be a JSON object", + Self::SessionTooLarge => "Codex Live session exceeds the 256 KiB limit", + Self::InvalidInitialEvent => { + "The initial Codex Live WebSocket text message must be a JSON object" + } + Self::ExpectedSessionUpdate => { + "Codex Live WebSocket must start with a session.update event" + } + Self::InitialEventMustBeText => { + "The initial Codex Live session.update event must be a text message" + } + Self::InitialClientReadFailed => { + "Failed to read the initial Codex Live session.update event" + } + Self::InitialSessionUpdateTimeout => { + "Timed out waiting for the initial Codex Live session.update event" + } + Self::InvalidCallLocation => { + "Codex Live upstream returned an invalid call location" + } + } + } + + pub(super) const fn is_timeout(&self) -> bool { + matches!(self, Self::InitialSessionUpdateTimeout) + } +} + +#[derive(Debug, Clone, PartialEq)] +pub(super) struct LiveMultipart { + pub(super) sdp: String, + pub(super) session: Value, +} + +pub(super) fn validate_model(model: &str) -> Result<(), LiveProtocolError> { + if model.is_empty() + || model.len() > MAX_MODEL_BYTES + || model.trim() != model + || model.chars().any(char::is_control) + { + return Err(LiveProtocolError::InvalidModel); + } + Ok(()) +} + +pub(super) fn validate_call_id(call_id: &str) -> Result<(), LiveProtocolError> { + if call_id.is_empty() + || call_id.len() > MAX_CALL_ID_BYTES + || matches!(call_id, "." | "..") + || !call_id + .bytes() + .all(|byte| byte.is_ascii_alphanumeric() || matches!(byte, b'_' | b'-' | b'.')) + { + return Err(LiveProtocolError::InvalidCallId); + } + Ok(()) +} + +fn is_live_call_id_segment(call_id: &str) -> bool { + if validate_call_id(call_id).is_err() { + return false; + } + if call_id.starts_with("rtc_") && call_id.len() > "rtc_".len() { + return true; + } + call_id.len() == 36 + && call_id + .bytes() + .enumerate() + .all(|(index, byte)| match index { + 8 | 13 | 18 | 23 => byte == b'-', + _ => byte.is_ascii_hexdigit(), + }) +} + +pub(super) fn direct_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(LiveProtocolError::InvalidModelQuery); + } + validate_model(value.as_ref())?; + model = Some(value.into_owned()); + continue; + } + // Latest Codex preserves provider query parameters while rewriting + // `/v1/realtime` to `/v1/live`. They are downstream transport hints, + // not Aether routing authority, so accept and ignore non-credential + // parameters. Credential values should already have been consumed by + // ingress; rejecting them here keeps this parser safe in isolation. + if live_query_parameter_is_sensitive(name.as_ref()) { + return Err(LiveProtocolError::InvalidModelQuery); + } + } + model.ok_or(LiveProtocolError::InvalidModelQuery) +} + +fn live_query_parameter_is_sensitive(name: &str) -> bool { + matches!( + name.to_ascii_lowercase().as_str(), + "key" + | "api_key" + | "api-key" + | "x-api-key" + | "x-goog-api-key" + | "access_token" + | "authorization" + | "token" + | "oauth_token" + | "client_secret" + | "secret_key" + | "signature" + | "sig" + ) +} + +pub(super) fn call_id_from_path(path: &str) -> Result { + let call_id = path + .strip_prefix("/v1/live/") + .ok_or(LiveProtocolError::InvalidCallId)?; + if call_id.contains('/') { + return Err(LiveProtocolError::InvalidCallId); + } + validate_call_id(call_id)?; + Ok(call_id.to_string()) +} + +pub(super) fn validate_initial_session_update(raw: &str) -> Result<(), LiveProtocolError> { + if raw.len() > MAX_SESSION_BYTES { + return Err(LiveProtocolError::SessionTooLarge); + } + let value: Value = + serde_json::from_str(raw).map_err(|_| LiveProtocolError::InvalidInitialEvent)?; + if !value.is_object() { + return Err(LiveProtocolError::InvalidInitialEvent); + } + if value.get("type").and_then(Value::as_str) != Some("session.update") { + return Err(LiveProtocolError::ExpectedSessionUpdate); + } + Ok(()) +} + +pub(super) fn event_type(raw: &str) -> Option { + serde_json::from_str::(raw) + .ok()? + .get("type")? + .as_str() + .map(str::to_string) +} + +pub(super) fn parse_live_multipart( + content_type: &str, + body: &[u8], +) -> Result { + if body.len() > MAX_MULTIPART_BODY_BYTES { + return Err(LiveProtocolError::MultipartBodyTooLarge); + } + let boundary = multipart_boundary(content_type)?; + let parts = parse_multipart_parts(body, boundary.as_bytes())?; + let mut sdp = None; + let mut session = None; + for part in parts { + match part.name.as_str() { + "sdp" => { + if sdp.is_some() { + return Err(LiveProtocolError::DuplicateMultipartPart); + } + if part.body.len() > MAX_SDP_BYTES { + return Err(LiveProtocolError::SdpTooLarge); + } + let value = + std::str::from_utf8(part.body).map_err(|_| LiveProtocolError::InvalidSdp)?; + if value.trim().is_empty() { + return Err(LiveProtocolError::InvalidSdp); + } + sdp = Some(value.to_string()); + } + "session" => { + if session.is_some() { + return Err(LiveProtocolError::DuplicateMultipartPart); + } + if part.body.len() > MAX_SESSION_BYTES { + return Err(LiveProtocolError::SessionTooLarge); + } + let value: Value = serde_json::from_slice(part.body) + .map_err(|_| LiveProtocolError::InvalidSession)?; + if !value.is_object() { + return Err(LiveProtocolError::InvalidSession); + } + session = Some(value); + } + _ => return Err(LiveProtocolError::UnexpectedMultipartPart), + } + } + Ok(LiveMultipart { + sdp: sdp.ok_or(LiveProtocolError::MissingSdp)?, + session: session.ok_or(LiveProtocolError::MissingSession)?, + }) +} + +pub(super) fn build_live_multipart(sdp: &str, session: &Value) -> (String, Vec) { + let boundary = format!("aether-live-{}", uuid::Uuid::new_v4().simple()); + let session = serde_json::to_vec(session).expect("a JSON value must serialize"); + let mut body = Vec::with_capacity(sdp.len() + session.len() + 320); + append_part( + &mut body, + boundary.as_str(), + "sdp", + "application/sdp", + sdp.as_bytes(), + ); + append_part( + &mut body, + boundary.as_str(), + "session", + "application/json", + session.as_slice(), + ); + body.extend_from_slice(format!("--{boundary}--\r\n").as_bytes()); + (format!("multipart/form-data; boundary={boundary}"), body) +} + +pub(super) fn extract_call_id_from_location(location: &str) -> Result { + let location = location.trim(); + if location.is_empty() { + return Err(LiveProtocolError::InvalidCallLocation); + } + let path = if let Ok(url) = url::Url::parse(location) { + url.path().to_string() + } else { + location + .split_once('?') + .map_or(location, |(path, _)| path) + .to_string() + }; + let call_id = path + .trim_end_matches('/') + .rsplit('/') + .next() + .filter(|value| !value.is_empty()) + .ok_or(LiveProtocolError::InvalidCallLocation)?; + if !is_live_call_id_segment(call_id) { + return Err(LiveProtocolError::InvalidCallLocation); + } + Ok(call_id.to_string()) +} + +struct MultipartPart<'a> { + name: String, + body: &'a [u8], +} + +fn multipart_boundary(content_type: &str) -> Result { + let mut values = content_type.split(';'); + if !values + .next() + .is_some_and(|value| value.trim().eq_ignore_ascii_case("multipart/form-data")) + { + return Err(LiveProtocolError::UnsupportedMediaType); + } + let mut boundary = None; + for parameter in values { + let Some((name, value)) = parameter.trim().split_once('=') else { + continue; + }; + if !name.trim().eq_ignore_ascii_case("boundary") { + continue; + } + if boundary.is_some() { + return Err(LiveProtocolError::InvalidBoundary); + } + let value = value.trim(); + let value = if value.starts_with('"') && value.ends_with('"') && value.len() >= 2 { + &value[1..value.len() - 1] + } else { + value + }; + boundary = Some(value.to_string()); + } + let boundary = boundary.ok_or(LiveProtocolError::InvalidBoundary)?; + if boundary.is_empty() + || boundary.len() > MAX_BOUNDARY_BYTES + || !boundary.bytes().all(|byte| { + byte.is_ascii_alphanumeric() + || matches!( + byte, + b'\'' + | b'(' + | b')' + | b'+' + | b'_' + | b',' + | b'-' + | b'.' + | b'/' + | b':' + | b'=' + | b'?' + ) + }) + { + return Err(LiveProtocolError::InvalidBoundary); + } + Ok(boundary) +} + +fn parse_multipart_parts<'a>( + body: &'a [u8], + boundary: &[u8], +) -> Result>, LiveProtocolError> { + let delimiter = [b"--".as_slice(), boundary].concat(); + if !body.starts_with(delimiter.as_slice()) { + return Err(LiveProtocolError::MalformedMultipart); + } + let mut cursor = delimiter.len(); + let mut parts = Vec::new(); + loop { + if body.get(cursor..cursor + 2) == Some(b"--") { + cursor += 2; + if body + .get(cursor..) + .is_some_and(|tail| tail.is_empty() || tail == b"\r\n") + { + return Ok(parts); + } + return Err(LiveProtocolError::MalformedMultipart); + } + if body.get(cursor..cursor + 2) != Some(b"\r\n") { + return Err(LiveProtocolError::MalformedMultipart); + } + cursor += 2; + let header_end = find_bytes(&body[cursor..], b"\r\n\r\n") + .ok_or(LiveProtocolError::MalformedMultipart)?; + if header_end > MAX_PART_HEADERS_BYTES { + return Err(LiveProtocolError::MalformedMultipart); + } + let headers = &body[cursor..cursor + header_end]; + cursor += header_end + 4; + let marker = [b"\r\n--".as_slice(), boundary].concat(); + let body_end = find_bytes(&body[cursor..], marker.as_slice()) + .ok_or(LiveProtocolError::MalformedMultipart)?; + let part_body = &body[cursor..cursor + body_end]; + let name = multipart_part_name(headers)?; + parts.push(MultipartPart { + name, + body: part_body, + }); + if parts.len() > 2 { + return Err(LiveProtocolError::UnexpectedMultipartPart); + } + cursor += body_end + 2 + delimiter.len(); + } +} + +fn multipart_part_name(headers: &[u8]) -> Result { + let headers = + std::str::from_utf8(headers).map_err(|_| LiveProtocolError::MalformedMultipart)?; + let mut disposition = None; + for line in headers.split("\r\n") { + let Some((name, value)) = line.split_once(':') else { + return Err(LiveProtocolError::MalformedMultipart); + }; + if name.trim().eq_ignore_ascii_case("content-disposition") { + if disposition.is_some() { + return Err(LiveProtocolError::MalformedMultipart); + } + disposition = Some(value.trim()); + } + } + let disposition = disposition.ok_or(LiveProtocolError::MalformedMultipart)?; + let mut parameters = disposition.split(';'); + if !parameters + .next() + .is_some_and(|value| value.trim().eq_ignore_ascii_case("form-data")) + { + return Err(LiveProtocolError::MalformedMultipart); + } + let mut part_name = None; + for parameter in parameters { + let Some((name, value)) = parameter.trim().split_once('=') else { + return Err(LiveProtocolError::MalformedMultipart); + }; + if name.trim().eq_ignore_ascii_case("filename") { + return Err(LiveProtocolError::UnexpectedMultipartPart); + } + if name.trim().eq_ignore_ascii_case("name") { + if part_name.is_some() { + return Err(LiveProtocolError::MalformedMultipart); + } + let value = value.trim(); + if !(value.starts_with('"') && value.ends_with('"') && value.len() >= 2) { + return Err(LiveProtocolError::MalformedMultipart); + } + part_name = Some(value[1..value.len() - 1].to_string()); + } + } + part_name.ok_or(LiveProtocolError::MalformedMultipart) +} + +fn find_bytes(haystack: &[u8], needle: &[u8]) -> Option { + (!needle.is_empty()) + .then(|| { + haystack + .windows(needle.len()) + .position(|value| value == needle) + }) + .flatten() +} + +fn append_part(body: &mut Vec, boundary: &str, name: &str, content_type: &str, value: &[u8]) { + body.extend_from_slice(format!("--{boundary}\r\n").as_bytes()); + body.extend_from_slice( + format!("Content-Disposition: form-data; name=\"{name}\"\r\n").as_bytes(), + ); + body.extend_from_slice(format!("Content-Type: {content_type}\r\n\r\n").as_bytes()); + body.extend_from_slice(value); + body.extend_from_slice(b"\r\n"); +} + +#[cfg(test)] +mod tests { + use serde_json::json; + + use super::*; + + #[test] + fn direct_query_requires_one_bounded_model() { + assert_eq!( + direct_model_from_query(Some("model=gpt-live%2Ffuture")).unwrap(), + "gpt-live/future" + ); + assert_eq!( + direct_model_from_query(Some("foo=bar&model=gpt-live&trace=1")).unwrap(), + "gpt-live" + ); + assert_eq!( + direct_model_from_query(Some("model=a&MODEL=b")), + Err(LiveProtocolError::InvalidModelQuery) + ); + assert_eq!( + direct_model_from_query(Some("model=a&token=secret")), + Err(LiveProtocolError::InvalidModelQuery) + ); + } + + #[test] + fn direct_protocol_starts_with_session_update_not_response_create() { + let opaque = r#"{"type":"session.update","session":{"future_capability":{"version":2}},"future_event_field":[1,2,3]}"#; + validate_initial_session_update(opaque).unwrap(); + assert_eq!( + serde_json::from_str::(opaque).unwrap()["future_event_field"], + json!([1, 2, 3]) + ); + assert_eq!( + validate_initial_session_update(r#"{"type":"response.create","model":"gpt"}"#), + Err(LiveProtocolError::ExpectedSessionUpdate) + ); + assert_eq!( + validate_initial_session_update(r#"["session.update"]"#), + Err(LiveProtocolError::InvalidInitialEvent) + ); + assert_eq!( + validate_initial_session_update("not-json"), + Err(LiveProtocolError::InvalidInitialEvent) + ); + } + + #[test] + fn direct_session_update_uses_the_bounded_session_limit() { + let oversized = format!( + r#"{{"type":"session.update","session":{{"future":"{}"}}}}"#, + "x".repeat(MAX_SESSION_BYTES) + ); + assert_eq!( + validate_initial_session_update(oversized.as_str()), + Err(LiveProtocolError::SessionTooLarge) + ); + } + + #[test] + fn multipart_round_trip_preserves_unknown_session_fields() { + let session = json!({ + "model": "gpt-future-live", + "instructions": "opaque", + "future_capability": {"enabled": true}, + "audio": {"input": {"format": "pcm16"}} + }); + let (content_type, body) = build_live_multipart("v=0\r\no=test", &session); + let parsed = parse_live_multipart(content_type.as_str(), body.as_slice()).unwrap(); + assert_eq!(parsed.sdp, "v=0\r\no=test"); + assert_eq!(parsed.session, session); + } + + #[test] + fn multipart_rejects_duplicate_or_unknown_parts() { + let boundary = "test-boundary"; + let duplicate = format!( + "--{boundary}\r\nContent-Disposition: form-data; name=\"sdp\"\r\n\r\nv=0\r\n--{boundary}\r\nContent-Disposition: form-data; name=\"sdp\"\r\n\r\nv=1\r\n--{boundary}--\r\n" + ); + assert_eq!( + parse_live_multipart( + format!("multipart/form-data; boundary={boundary}").as_str(), + duplicate.as_bytes(), + ), + Err(LiveProtocolError::DuplicateMultipartPart) + ); + + let unknown = format!( + "--{boundary}\r\nContent-Disposition: form-data; name=\"sdp\"\r\n\r\nv=0\r\n--{boundary}\r\nContent-Disposition: form-data; name=\"credentials\"\r\n\r\nsecret\r\n--{boundary}--\r\n" + ); + assert_eq!( + parse_live_multipart( + format!("multipart/form-data; boundary={boundary}").as_str(), + unknown.as_bytes(), + ), + Err(LiveProtocolError::UnexpectedMultipartPart) + ); + } + + #[test] + fn multipart_enforces_total_sdp_and_session_limits() { + let oversized_body = vec![b'x'; MAX_MULTIPART_BODY_BYTES + 1]; + assert_eq!( + parse_live_multipart( + "multipart/form-data; boundary=limit", + oversized_body.as_slice(), + ), + Err(LiveProtocolError::MultipartBodyTooLarge) + ); + + let oversized_sdp = "x".repeat(MAX_SDP_BYTES + 1); + let sdp_body = format!( + "--limit\r\nContent-Disposition: form-data; name=\"sdp\"\r\n\r\n{oversized_sdp}\r\n--limit\r\nContent-Disposition: form-data; name=\"session\"\r\n\r\n{{}}\r\n--limit--\r\n" + ); + assert_eq!( + parse_live_multipart("multipart/form-data; boundary=limit", sdp_body.as_bytes()), + Err(LiveProtocolError::SdpTooLarge) + ); + + let oversized_session = format!(r#"{{"future":"{}"}}"#, "x".repeat(MAX_SESSION_BYTES)); + let session_body = format!( + "--limit\r\nContent-Disposition: form-data; name=\"sdp\"\r\n\r\nv=0\r\n--limit\r\nContent-Disposition: form-data; name=\"session\"\r\n\r\n{oversized_session}\r\n--limit--\r\n" + ); + assert_eq!( + parse_live_multipart( + "multipart/form-data; boundary=limit", + session_body.as_bytes(), + ), + Err(LiveProtocolError::SessionTooLarge) + ); + } + + #[test] + fn multipart_rejects_unbounded_or_ambiguous_boundaries() { + let valid_body = + b"--safe\r\nContent-Disposition: form-data; name=\"sdp\"\r\n\r\nv=0\r\n--safe--\r\n"; + assert_eq!( + parse_live_multipart("application/json", valid_body), + Err(LiveProtocolError::UnsupportedMediaType) + ); + assert_eq!( + parse_live_multipart( + format!( + "multipart/form-data; boundary={}", + "x".repeat(MAX_BOUNDARY_BYTES + 1) + ) + .as_str(), + valid_body, + ), + Err(LiveProtocolError::InvalidBoundary) + ); + assert_eq!( + parse_live_multipart( + "multipart/form-data; boundary=safe; boundary=other", + valid_body, + ), + Err(LiveProtocolError::InvalidBoundary) + ); + } + + #[test] + fn extracts_only_realtime_call_ids_from_location() { + for location in [ + "https://api.openai.com/v1/live/rtc_abc-123", + "/v1/live/550e8400-e29b-41d4-a716-446655440000", + ] { + assert!(extract_call_id_from_location(location).is_ok()); + } + assert_eq!( + extract_call_id_from_location("/v1/live/rtc%2Fescape"), + Err(LiveProtocolError::InvalidCallLocation) + ); + for location in ["/v1/live", "/v1/live/not-a-call-id"] { + assert_eq!( + extract_call_id_from_location(location), + Err(LiveProtocolError::InvalidCallLocation) + ); + } + for dot_segment in [".", ".."] { + assert_eq!( + validate_call_id(dot_segment), + Err(LiveProtocolError::InvalidCallId) + ); + } + } + + #[test] + fn opaque_event_discriminator_does_not_project_unknown_fields() { + let raw = r#"{"type":"delegation.created","unknown":{"nested":[1,2,3]}}"#; + assert_eq!(event_type(raw).as_deref(), Some("delegation.created")); + assert_eq!( + serde_json::from_str::(raw).unwrap()["unknown"]["nested"][2], + 3 + ); + } +} diff --git a/apps/aether-gateway/src/handlers/proxy/websocket/live/registry.rs b/apps/aether-gateway/src/handlers/proxy/websocket/live/registry.rs new file mode 100644 index 000000000..2f738b9e9 --- /dev/null +++ b/apps/aether-gateway/src/handlers/proxy/websocket/live/registry.rs @@ -0,0 +1,1036 @@ +//! Short-lived ownership registry for Codex Live WebRTC sideband calls. +//! +//! Call IDs are opaque references to provider state. The raw ID is never +//! stored: the authenticated downstream principal and call ID are hashed into +//! a RuntimeState key, while the record contains only non-secret routing data. + +use std::sync::Arc; +use std::time::{Duration, SystemTime, UNIX_EPOCH}; + +use aether_runtime_state::{RuntimeLockLease, RuntimeState}; +use serde::{Deserialize, Serialize}; +use sha2::{Digest, Sha256}; +use tokio::sync::{oneshot, watch}; +use tokio::task::JoinHandle; + +use crate::ai_serving::ResponsesWebSocketPinnedCandidate; + +use super::planner::{LiveAuthMode, PlannedLiveCandidate}; +use super::protocol::validate_call_id; + +const SCHEMA_VERSION: u16 = 2; +const RECORD_PREFIX: &str = "codex_live:call:v2:"; +const RECORD_DOMAIN: &[u8] = b"aether-codex-live-call-v2"; +const INDEX_PREFIX: &str = "codex_live:call_index:v2:"; +const INDEX_DOMAIN: &[u8] = b"aether-codex-live-call-index-v2"; +const LOCK_PREFIX: &str = "codex_live:call_lock:v2:"; +const SIDEBAND_LOCK_PREFIX: &str = "codex_live:sideband_lock:v1:"; +const SIDEBAND_LOCK_DOMAIN: &[u8] = b"aether-codex-live-sideband-lock-v1"; +const RECORD_TTL: Duration = Duration::from_secs(2 * 60 * 60); +const EXPIRED_LOOKUP_GRACE: Duration = Duration::from_secs(5 * 60); +const LOCK_TTL: Duration = Duration::from_secs(2); +const SIDEBAND_LOCK_TTL: Duration = Duration::from_secs(30); +const LOCK_ACQUIRE_TIMEOUT: Duration = Duration::from_millis(250); +const COMMIT_VERIFY_TIMEOUT: Duration = Duration::from_millis(250); +const LOCK_INITIAL_RETRY: Duration = Duration::from_millis(5); +const LOCK_MAX_RETRY: Duration = Duration::from_millis(50); +const LOCK_OWNER: &str = "codex_live_call_registry"; +const SIDEBAND_LOCK_OWNER: &str = "codex_live_sideband_attachment"; +const MAX_RECORDS_PER_PRINCIPAL: usize = 64; +const MAX_SERIALIZED_RECORD_BYTES: usize = 4 * 1024; +const MAX_PRINCIPAL_BYTES: usize = 256; +const MAX_RECORD_ID_BYTES: usize = 256; +const LIVE_LOG_TARGET: &str = "aether_gateway::handlers::proxy::codex_live"; + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum RegisterCommitState { + Direct, + VerifiedAfterError, + Uncommitted, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub(super) enum LiveCallLookup { + Found(LiveCallBinding), + Missing, + Expired, +} + +#[derive(Debug, thiserror::Error)] +pub(super) enum LiveCallRegistryError { + #[error("invalid Codex Live call identity: {0}")] + InvalidIdentity(&'static str), + #[error("invalid Codex Live call binding: {0}")] + InvalidRecord(&'static str), + #[error("Codex Live call binding serialization failed")] + Serialization(#[source] serde_json::Error), + #[error("Codex Live call registry contains corrupt data")] + CorruptRecord(#[source] serde_json::Error), + #[error("Codex Live call binding is too large")] + RecordTooLarge, + #[error("Codex Live call ownership conflicts with an existing binding")] + OwnershipConflict, + #[error("Codex Live call capacity lock is busy")] + CapacityLockBusy, + #[error("Codex Live call already has an active sideband attachment")] + SidebandAlreadyAttached, + #[error("Codex Live call registry storage is unavailable")] + Storage(#[source] aether_runtime_state::DataLayerError), +} + +impl LiveCallRegistryError { + pub(super) const fn kind(&self) -> &'static str { + match self { + Self::InvalidIdentity(_) => "invalid_identity", + Self::InvalidRecord(_) => "invalid_record", + Self::Serialization(_) => "serialization_failed", + Self::CorruptRecord(_) => "corrupt_record", + Self::RecordTooLarge => "record_too_large", + Self::OwnershipConflict => "ownership_conflict", + Self::CapacityLockBusy => "capacity_lock_busy", + Self::SidebandAlreadyAttached => "sideband_already_attached", + Self::Storage(_) => "storage_unavailable", + } + } +} + +/// Exclusive, renewable ownership of one authenticated Live sideband call. +/// +/// The runtime lock key contains only a domain-separated digest. Call IDs and +/// downstream principal identifiers are never stored in the lease key. +pub(super) struct LiveSidebandLease { + runtime_state: Arc, + lease: RuntimeLockLease, + renewal_cancel: Option>, + renewal_task: Option>, + health: watch::Receiver, + armed: bool, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum LiveSidebandLeaseHealth { + Healthy, + OwnershipLost, + StorageUnavailable, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(super) enum LiveSidebandLeaseLoss { + OwnershipLost, + StorageUnavailable, +} + +impl LiveSidebandLeaseLoss { + pub(super) const fn kind(self) -> &'static str { + match self { + Self::OwnershipLost => "ownership_lost", + Self::StorageUnavailable => "storage_unavailable", + } + } +} + +impl LiveSidebandLease { + async fn new(runtime_state: Arc, lease: RuntimeLockLease, ttl: Duration) -> Self { + let (renewal_cancel, mut cancel) = oneshot::channel(); + let (renewal_started, started) = oneshot::channel(); + let (health_tx, health) = watch::channel(LiveSidebandLeaseHealth::Healthy); + let renewal_state = Arc::clone(&runtime_state); + let renewal_lease = lease.clone(); + let interval = (ttl / 3).max(Duration::from_millis(1)); + let renewal_task = tokio::spawn(async move { + let mut ticker = + tokio::time::interval_at(tokio::time::Instant::now() + interval, interval); + ticker.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Delay); + let _ = renewal_started.send(()); + loop { + tokio::select! { + _ = &mut cancel => break, + _ = ticker.tick() => { + match renewal_state.lock_renew(&renewal_lease, ttl).await { + Ok(true) => {} + Ok(false) => { + health_tx.send_replace(LiveSidebandLeaseHealth::OwnershipLost); + tracing::warn!( + target: LIVE_LOG_TARGET, + event_name = "codex_live_sideband_lease_auto_renew_failed", + log_type = "ops", + error_kind = "ownership_lost", + "Codex Live sideband ownership was lost during automatic renewal" + ); + break; + } + Err(_) => { + health_tx.send_replace(LiveSidebandLeaseHealth::StorageUnavailable); + tracing::warn!( + target: LIVE_LOG_TARGET, + event_name = "codex_live_sideband_lease_auto_renew_failed", + log_type = "ops", + error_kind = "storage_unavailable", + "Codex Live sideband ownership could not be renewed" + ); + break; + } + } + } + } + } + }); + let sideband_lease = Self { + runtime_state, + lease, + renewal_cancel: Some(renewal_cancel), + renewal_task: Some(renewal_task), + health, + armed: true, + }; + // Do not return ownership to the caller until the renewal task has + // registered its first deadline. This closes the acquire-to-first-poll + // window on single-thread runtimes and heavily loaded executors. + let _ = started.await; + sideband_lease + } + + pub(super) fn loss(&self) -> Option { + match *self.health.borrow() { + LiveSidebandLeaseHealth::Healthy => None, + LiveSidebandLeaseHealth::OwnershipLost => Some(LiveSidebandLeaseLoss::OwnershipLost), + LiveSidebandLeaseHealth::StorageUnavailable => { + Some(LiveSidebandLeaseLoss::StorageUnavailable) + } + } + } + + pub(super) async fn wait_for_loss(&self) -> LiveSidebandLeaseLoss { + let mut health = self.health.clone(); + loop { + match *health.borrow_and_update() { + LiveSidebandLeaseHealth::Healthy => {} + LiveSidebandLeaseHealth::OwnershipLost => { + return LiveSidebandLeaseLoss::OwnershipLost; + } + LiveSidebandLeaseHealth::StorageUnavailable => { + return LiveSidebandLeaseLoss::StorageUnavailable; + } + } + if health.changed().await.is_err() { + return LiveSidebandLeaseLoss::StorageUnavailable; + } + } + } + + pub(super) async fn release(&mut self) -> Result { + self.stop_renewal().await; + let released = self + .runtime_state + .lock_release(&self.lease) + .await + .map_err(LiveCallRegistryError::Storage)?; + self.armed = false; + Ok(released) + } + + async fn stop_renewal(&mut self) { + if let Some(cancel) = self.renewal_cancel.take() { + let _ = cancel.send(()); + } + if let Some(task) = self.renewal_task.take() { + task.abort(); + let _ = task.await; + } + } +} + +impl Drop for LiveSidebandLease { + fn drop(&mut self) { + if !self.armed { + return; + } + if let Some(cancel) = self.renewal_cancel.take() { + let _ = cancel.send(()); + } + let renewal_task = self.renewal_task.take(); + let runtime_state = Arc::clone(&self.runtime_state); + let lease = self.lease.clone(); + if let Ok(runtime) = tokio::runtime::Handle::try_current() { + runtime.spawn(async move { + if let Some(task) = renewal_task { + task.abort(); + let _ = task.await; + } + let _ = runtime_state.lock_release(&lease).await; + }); + } else if let Some(task) = renewal_task { + task.abort(); + } + } +} + +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +#[serde(deny_unknown_fields)] +pub(super) struct LiveCallBinding { + schema_version: u16, + pinned_candidate: ResponsesWebSocketPinnedCandidate, + client_model: String, + provider_model: String, + auth_mode: LiveAuthMode, + routing_fingerprint: String, + created_at_unix_ms: u64, +} + +impl LiveCallBinding { + pub(super) fn from_candidate(candidate: &PlannedLiveCandidate) -> Self { + Self { + schema_version: SCHEMA_VERSION, + pinned_candidate: candidate.pinned_candidate.clone(), + client_model: candidate.client_model.clone(), + provider_model: candidate.provider_model.clone(), + auth_mode: candidate.auth_mode, + routing_fingerprint: candidate.routing_fingerprint.clone(), + created_at_unix_ms: now_unix_ms(), + } + } + + pub(super) fn pinned_candidate(&self) -> &ResponsesWebSocketPinnedCandidate { + &self.pinned_candidate + } + + pub(super) fn client_model(&self) -> &str { + self.client_model.as_str() + } + + pub(super) fn matches_candidate(&self, candidate: &PlannedLiveCandidate) -> bool { + self.pinned_candidate == candidate.pinned_candidate + && self.client_model == candidate.client_model + && self.provider_model == candidate.provider_model + && self.auth_mode == candidate.auth_mode + && self.routing_fingerprint == candidate.routing_fingerprint + } + + fn validate(&self) -> Result<(), LiveCallRegistryError> { + if self.schema_version != SCHEMA_VERSION { + return Err(LiveCallRegistryError::InvalidRecord( + "unsupported_schema_version", + )); + } + for (value, error) in [ + (self.pinned_candidate.provider_id(), "invalid_provider_id"), + (self.pinned_candidate.endpoint_id(), "invalid_endpoint_id"), + (self.pinned_candidate.key_id(), "invalid_key_id"), + (self.client_model.as_str(), "invalid_client_model"), + (self.provider_model.as_str(), "invalid_provider_model"), + ] { + if value.trim().is_empty() || value.len() > MAX_RECORD_ID_BYTES { + return Err(LiveCallRegistryError::InvalidRecord(error)); + } + } + if self.created_at_unix_ms == 0 { + return Err(LiveCallRegistryError::InvalidRecord("invalid_created_at")); + } + if self.routing_fingerprint.len() != 64 + || !self + .routing_fingerprint + .bytes() + .all(|byte| byte.is_ascii_hexdigit()) + { + return Err(LiveCallRegistryError::InvalidRecord( + "invalid_routing_fingerprint", + )); + } + Ok(()) + } +} + +pub(super) struct LiveCallRegistry { + runtime_state: Arc, + ttl: Duration, + max_records_per_principal: usize, +} + +impl LiveCallRegistry { + pub(super) fn new(runtime_state: Arc) -> Self { + Self { + runtime_state, + ttl: RECORD_TTL, + max_records_per_principal: MAX_RECORDS_PER_PRINCIPAL, + } + } + + #[cfg(test)] + fn with_limits(runtime_state: Arc, ttl: Duration, max_records: usize) -> Self { + Self { + runtime_state, + ttl, + max_records_per_principal: max_records, + } + } + + pub(super) async fn register( + &self, + user_id: &str, + api_key_id: &str, + call_id: &str, + binding: &LiveCallBinding, + ) -> Result<(), LiveCallRegistryError> { + binding.validate()?; + let key = record_key(user_id, api_key_id, call_id)?; + let index = index_key(user_id, api_key_id)?; + let lock = lock_key(user_id, api_key_id)?; + let serialized = + serde_json::to_string(binding).map_err(LiveCallRegistryError::Serialization)?; + if serialized.len() > MAX_SERIALIZED_RECORD_BYTES { + return Err(LiveCallRegistryError::RecordTooLarge); + } + let lease = self.acquire_lock(lock.as_str()).await?; + let result = self + .register_locked(key.as_str(), index.as_str(), serialized, binding) + .await; + let exact_binding_committed = if result.is_err() { + self.exact_binding_is_committed(key.as_str(), binding).await + } else { + false + }; + let commit_state = classify_register_commit(&result, exact_binding_committed); + let release = self.runtime_state.lock_release(&lease).await; + + if commit_state == RegisterCommitState::Uncommitted { + return result; + } + + if commit_state == RegisterCommitState::VerifiedAfterError { + tracing::warn!( + target: LIVE_LOG_TARGET, + event_name = "codex_live_call_binding_commit_verified", + log_type = "ops", + register_error_kind = result + .as_ref() + .err() + .map_or("unknown", |error| error.kind()), + "Codex Live accepted a binding after exact commit verification" + ); + } + if release.is_err() { + tracing::warn!( + target: LIVE_LOG_TARGET, + event_name = "codex_live_call_capacity_lock_release_failed", + log_type = "ops", + error_kind = "storage_unavailable", + "Codex Live binding was committed but its short-lived capacity lock could not be released" + ); + } + Ok(()) + } + + pub(super) async fn lookup( + &self, + user_id: &str, + api_key_id: &str, + call_id: &str, + ) -> Result, LiveCallRegistryError> { + Ok( + match self + .lookup_with_status(user_id, api_key_id, call_id) + .await? + { + LiveCallLookup::Found(binding) => Some(binding), + LiveCallLookup::Missing | LiveCallLookup::Expired => None, + }, + ) + } + + pub(super) async fn lookup_with_status( + &self, + user_id: &str, + api_key_id: &str, + call_id: &str, + ) -> Result { + let key = record_key(user_id, api_key_id, call_id)?; + if let Some(serialized) = self + .runtime_state + .kv_get(key.as_str()) + .await + .map_err(LiveCallRegistryError::Storage)? + { + let binding = serde_json::from_str::(serialized.as_str()) + .map_err(LiveCallRegistryError::CorruptRecord)?; + binding.validate()?; + return Ok(LiveCallLookup::Found(binding)); + } + + let index = index_key(user_id, api_key_id)?; + let indexed = self + .runtime_state + .score_range_by_min(index.as_str(), f64::NEG_INFINITY) + .await + .map_err(LiveCallRegistryError::Storage)? + .iter() + .any(|member| member == &key); + Ok(if indexed { + LiveCallLookup::Expired + } else { + LiveCallLookup::Missing + }) + } + + pub(super) async fn acquire_sideband_attachment( + &self, + user_id: &str, + api_key_id: &str, + call_id: &str, + ) -> Result { + self.acquire_sideband_attachment_with_ttl(user_id, api_key_id, call_id, SIDEBAND_LOCK_TTL) + .await + } + + async fn acquire_sideband_attachment_with_ttl( + &self, + user_id: &str, + api_key_id: &str, + call_id: &str, + ttl: Duration, + ) -> Result { + let key = sideband_lock_key(user_id, api_key_id, call_id)?; + let Some(lease) = self + .runtime_state + .lock_try_acquire(key.as_str(), SIDEBAND_LOCK_OWNER, ttl) + .await + .map_err(LiveCallRegistryError::Storage)? + else { + return Err(LiveCallRegistryError::SidebandAlreadyAttached); + }; + Ok(LiveSidebandLease::new(Arc::clone(&self.runtime_state), lease, ttl).await) + } + + async fn acquire_lock(&self, key: &str) -> Result { + let deadline = tokio::time::Instant::now() + LOCK_ACQUIRE_TIMEOUT; + let mut retry = LOCK_INITIAL_RETRY; + loop { + if let Some(lease) = self + .runtime_state + .lock_try_acquire(key, LOCK_OWNER, LOCK_TTL) + .await + .map_err(LiveCallRegistryError::Storage)? + { + return Ok(lease); + } + let now = tokio::time::Instant::now(); + if now >= deadline { + return Err(LiveCallRegistryError::CapacityLockBusy); + } + tokio::time::sleep(retry.min(deadline.saturating_duration_since(now))).await; + retry = retry.saturating_mul(2).min(LOCK_MAX_RETRY); + } + } + + async fn exact_binding_is_committed(&self, key: &str, expected: &LiveCallBinding) -> bool { + let stored = + match tokio::time::timeout(COMMIT_VERIFY_TIMEOUT, self.runtime_state.kv_get(key)).await + { + Ok(Ok(Some(stored))) => stored, + Ok(Ok(None) | Err(_)) | Err(_) => return false, + }; + let Ok(actual) = serde_json::from_str::(stored.as_str()) else { + return false; + }; + actual.validate().is_ok() && actual == *expected + } + + async fn register_locked( + &self, + key: &str, + index: &str, + serialized: String, + binding: &LiveCallBinding, + ) -> Result<(), LiveCallRegistryError> { + if let Some(existing) = self + .runtime_state + .kv_get(key) + .await + .map_err(LiveCallRegistryError::Storage)? + { + let existing = serde_json::from_str::(existing.as_str()) + .map_err(LiveCallRegistryError::CorruptRecord)?; + if existing != *binding { + return Err(LiveCallRegistryError::OwnershipConflict); + } + } + self.runtime_state + .kv_set(key, serialized, Some(self.ttl)) + .await + .map_err(LiveCallRegistryError::Storage)?; + if let Err(error) = self + .runtime_state + .score_set(index, key, now_unix_ms() as f64) + .await + { + let _ = self.runtime_state.kv_delete(key).await; + return Err(LiveCallRegistryError::Storage(error)); + } + if let Err(error) = self + .runtime_state + .key_expire(index, self.ttl.saturating_add(EXPIRED_LOOKUP_GRACE)) + .await + { + let _ = self.runtime_state.score_remove(index, key).await; + let _ = self.runtime_state.kv_delete(key).await; + return Err(LiveCallRegistryError::Storage(error)); + } + let members = self + .runtime_state + .score_range_by_min(index, f64::NEG_INFINITY) + .await + .map_err(LiveCallRegistryError::Storage)?; + let overflow = members.len().saturating_sub(self.max_records_per_principal); + for oldest in members.into_iter().take(overflow) { + self.runtime_state + .kv_delete(oldest.as_str()) + .await + .map_err(LiveCallRegistryError::Storage)?; + self.runtime_state + .score_remove(index, oldest.as_str()) + .await + .map_err(LiveCallRegistryError::Storage)?; + } + Ok(()) + } +} + +fn classify_register_commit( + result: &Result<(), LiveCallRegistryError>, + exact_binding_committed: bool, +) -> RegisterCommitState { + match result { + Ok(()) => RegisterCommitState::Direct, + Err(LiveCallRegistryError::OwnershipConflict) => RegisterCommitState::Uncommitted, + Err(_) if exact_binding_committed => RegisterCommitState::VerifiedAfterError, + Err(_) => RegisterCommitState::Uncommitted, + } +} + +fn record_key( + user_id: &str, + api_key_id: &str, + call_id: &str, +) -> Result { + validate_principal(user_id, "invalid_user_id")?; + validate_principal(api_key_id, "invalid_api_key_id")?; + validate_call_id(call_id) + .map_err(|_| LiveCallRegistryError::InvalidIdentity("invalid_call_id"))?; + Ok(format!( + "{RECORD_PREFIX}{}", + digest(RECORD_DOMAIN, &[user_id, api_key_id, call_id]) + )) +} + +fn index_key(user_id: &str, api_key_id: &str) -> Result { + validate_principal(user_id, "invalid_user_id")?; + validate_principal(api_key_id, "invalid_api_key_id")?; + Ok(format!( + "{INDEX_PREFIX}{}", + digest(INDEX_DOMAIN, &[user_id, api_key_id]) + )) +} + +fn lock_key(user_id: &str, api_key_id: &str) -> Result { + let index = index_key(user_id, api_key_id)?; + Ok(format!( + "{LOCK_PREFIX}{}", + index.strip_prefix(INDEX_PREFIX).unwrap_or(index.as_str()) + )) +} + +fn sideband_lock_key( + user_id: &str, + api_key_id: &str, + call_id: &str, +) -> Result { + validate_principal(user_id, "invalid_user_id")?; + validate_principal(api_key_id, "invalid_api_key_id")?; + validate_call_id(call_id) + .map_err(|_| LiveCallRegistryError::InvalidIdentity("invalid_call_id"))?; + Ok(format!( + "{SIDEBAND_LOCK_PREFIX}{}", + digest(SIDEBAND_LOCK_DOMAIN, &[user_id, api_key_id, call_id]) + )) +} + +fn validate_principal(value: &str, error: &'static str) -> Result<(), LiveCallRegistryError> { + if value.is_empty() || value.len() > MAX_PRINCIPAL_BYTES { + return Err(LiveCallRegistryError::InvalidIdentity(error)); + } + Ok(()) +} + +fn digest(domain: &[u8], components: &[&str]) -> String { + let mut digest = Sha256::new(); + digest.update(domain); + for component in components { + digest.update((component.len() as u64).to_be_bytes()); + digest.update(component.as_bytes()); + } + format!("{:x}", digest.finalize()) +} + +fn now_unix_ms() -> u64 { + SystemTime::now() + .duration_since(UNIX_EPOCH) + .unwrap_or_default() + .as_millis() as u64 +} + +#[cfg(test)] +mod tests { + use aether_runtime_state::{MemoryRuntimeStateConfig, RuntimeState}; + + use super::*; + + fn runtime_state() -> Arc { + Arc::new(RuntimeState::memory(MemoryRuntimeStateConfig::default())) + } + + fn binding(client_model: &str) -> LiveCallBinding { + LiveCallBinding { + schema_version: SCHEMA_VERSION, + pinned_candidate: ResponsesWebSocketPinnedCandidate::new( + "provider-1", + "endpoint-1", + "key-1", + ) + .unwrap(), + client_model: client_model.to_string(), + provider_model: "provider-model".to_string(), + auth_mode: LiveAuthMode::ChatGptOauth, + routing_fingerprint: "a".repeat(64), + created_at_unix_ms: now_unix_ms(), + } + } + + fn candidate_for_binding(binding: &LiveCallBinding) -> PlannedLiveCandidate { + let execution: crate::ai_serving::AiExecutionDecision = + serde_json::from_value(serde_json::json!({ + "action": "stream", + "provider_id": binding.pinned_candidate.provider_id(), + "endpoint_id": binding.pinned_candidate.endpoint_id(), + "key_id": binding.pinned_candidate.key_id(), + "provider_type": "codex", + "upstream_url": "https://chatgpt.com/backend-api/codex/responses" + })) + .unwrap(); + PlannedLiveCandidate { + execution, + pinned_candidate: binding.pinned_candidate.clone(), + client_model: binding.client_model.clone(), + provider_model: binding.provider_model.clone(), + auth_mode: binding.auth_mode, + routing_fingerprint: binding.routing_fingerprint.clone(), + } + } + + #[tokio::test] + async fn binding_is_scoped_to_the_authenticated_principal() { + let state = runtime_state(); + let registry = LiveCallRegistry::new(Arc::clone(&state)); + registry + .register("user-1", "api-key-1", "rtc_secret", &binding("global")) + .await + .unwrap(); + assert!(registry + .lookup("user-1", "api-key-1", "rtc_secret") + .await + .unwrap() + .is_some()); + assert!(registry + .lookup("user-2", "api-key-1", "rtc_secret") + .await + .unwrap() + .is_none()); + assert!(registry + .lookup("user-1", "api-key-2", "rtc_secret") + .await + .unwrap() + .is_none()); + } + + #[tokio::test] + async fn capacity_evicts_the_oldest_binding() { + let state = runtime_state(); + let registry = + LiveCallRegistry::with_limits(Arc::clone(&state), Duration::from_secs(60), 2); + for call_id in ["rtc_1", "rtc_2", "rtc_3"] { + registry + .register("user", "key", call_id, &binding(call_id)) + .await + .unwrap(); + tokio::time::sleep(Duration::from_millis(2)).await; + } + assert!(registry + .lookup("user", "key", "rtc_1") + .await + .unwrap() + .is_none()); + assert!(registry + .lookup("user", "key", "rtc_2") + .await + .unwrap() + .is_some()); + assert!(registry + .lookup("user", "key", "rtc_3") + .await + .unwrap() + .is_some()); + } + + #[tokio::test] + async fn ttl_expiry_removes_a_binding() { + let state = runtime_state(); + let registry = + LiveCallRegistry::with_limits(Arc::clone(&state), Duration::from_millis(5), 2); + registry + .register("user", "key", "rtc_expiring", &binding("global")) + .await + .unwrap(); + tokio::time::sleep(Duration::from_millis(20)).await; + assert!(registry + .lookup("user", "key", "rtc_expiring") + .await + .unwrap() + .is_none()); + assert_eq!( + registry + .lookup_with_status("user", "key", "rtc_expiring") + .await + .unwrap(), + LiveCallLookup::Expired + ); + } + + #[tokio::test] + async fn an_existing_call_id_cannot_be_rebound_to_another_candidate() { + let state = runtime_state(); + let registry = LiveCallRegistry::new(Arc::clone(&state)); + let original = binding("global-a"); + registry + .register("user", "key", "rtc_shared", &original) + .await + .unwrap(); + + assert!(matches!( + registry + .register("user", "key", "rtc_shared", &binding("global-b")) + .await, + Err(LiveCallRegistryError::OwnershipConflict) + )); + assert_eq!( + registry.lookup("user", "key", "rtc_shared").await.unwrap(), + Some(original) + ); + } + + #[test] + fn routing_fingerprint_drift_invalidates_the_pinned_binding() { + let binding = binding("global"); + let mut candidate = candidate_for_binding(&binding); + assert!(binding.matches_candidate(&candidate)); + + candidate.routing_fingerprint = "b".repeat(64); + assert!(!binding.matches_candidate(&candidate)); + } + + #[test] + fn registry_keys_hash_principal_and_raw_call_identifiers() { + let record = record_key("user-private", "key-private", "rtc-private").unwrap(); + let index = index_key("user-private", "key-private").unwrap(); + let sideband = sideband_lock_key("user-private", "key-private", "rtc-private").unwrap(); + for value in [&record, &index, &sideband] { + assert!(!value.contains("user-private")); + assert!(!value.contains("key-private")); + assert!(!value.contains("rtc-private")); + } + assert_eq!(record.len(), RECORD_PREFIX.len() + 64); + assert_eq!(index.len(), INDEX_PREFIX.len() + 64); + assert_eq!(sideband.len(), SIDEBAND_LOCK_PREFIX.len() + 64); + } + + #[tokio::test] + async fn sideband_attachment_is_exclusive_and_release_allows_reconnect() { + let state = runtime_state(); + let registry = LiveCallRegistry::new(Arc::clone(&state)); + let mut first = registry + .acquire_sideband_attachment("user", "key", "rtc_attach") + .await + .unwrap(); + + assert!(matches!( + registry + .acquire_sideband_attachment("user", "key", "rtc_attach") + .await, + Err(LiveCallRegistryError::SidebandAlreadyAttached) + )); + assert!(first.release().await.unwrap()); + + let mut reconnected = registry + .acquire_sideband_attachment("user", "key", "rtc_attach") + .await + .unwrap(); + assert!(reconnected.release().await.unwrap()); + } + + #[tokio::test] + async fn sideband_attachment_auto_renewal_starts_at_acquisition() { + let state = runtime_state(); + let registry = LiveCallRegistry::new(Arc::clone(&state)); + let mut held = registry + .acquire_sideband_attachment_with_ttl( + "user", + "key", + "rtc_renew", + Duration::from_secs(2), + ) + .await + .unwrap(); + + // Keep waking the single-thread test runtime while crossing the + // original TTL. Very small sub-second TTLs are not representative of + // the production 30-second lease and become scheduler-flaky on loaded + // CI hosts. + for _ in 0..25 { + tokio::time::sleep(Duration::from_millis(100)).await; + assert_eq!(held.loss(), None); + } + assert!(matches!( + registry + .acquire_sideband_attachment("user", "key", "rtc_renew") + .await, + Err(LiveCallRegistryError::SidebandAlreadyAttached) + )); + assert!(held.release().await.unwrap()); + } + + #[tokio::test] + async fn expired_owner_cannot_release_a_successor_sideband_lease() { + let state = runtime_state(); + let registry = LiveCallRegistry::new(Arc::clone(&state)); + let mut expired = registry + .acquire_sideband_attachment_with_ttl( + "user", + "key", + "rtc_fenced", + Duration::from_millis(25), + ) + .await + .unwrap(); + expired.stop_renewal().await; + tokio::time::sleep(Duration::from_millis(75)).await; + + let mut successor = registry + .acquire_sideband_attachment("user", "key", "rtc_fenced") + .await + .unwrap(); + assert!(!expired.release().await.unwrap()); + assert!(matches!( + registry + .acquire_sideband_attachment("user", "key", "rtc_fenced") + .await, + Err(LiveCallRegistryError::SidebandAlreadyAttached) + )); + assert!(successor.release().await.unwrap()); + } + + #[tokio::test] + async fn invalid_call_identity_is_rejected_before_storage() { + let state = runtime_state(); + let registry = LiveCallRegistry::new(Arc::clone(&state)); + for call_id in [".", "..", "rtc/escape"] { + assert!(matches!( + registry + .register("user", "key", call_id, &binding("global")) + .await, + Err(LiveCallRegistryError::InvalidIdentity("invalid_call_id")) + )); + assert!(matches!( + registry + .acquire_sideband_attachment("user", "key", call_id) + .await, + Err(LiveCallRegistryError::InvalidIdentity("invalid_call_id")) + )); + } + } + + #[test] + fn register_commit_classification_is_fail_closed() { + let success = Ok(()); + assert_eq!( + classify_register_commit(&success, false), + RegisterCommitState::Direct + ); + + let storage_error = Err(LiveCallRegistryError::Storage( + aether_runtime_state::DataLayerError::UnexpectedValue("injected".to_string()), + )); + assert_eq!( + classify_register_commit(&storage_error, true), + RegisterCommitState::VerifiedAfterError + ); + assert_eq!( + classify_register_commit(&storage_error, false), + RegisterCommitState::Uncommitted + ); + + let ownership_conflict = Err(LiveCallRegistryError::OwnershipConflict); + assert_eq!( + classify_register_commit(&ownership_conflict, true), + RegisterCommitState::Uncommitted + ); + } + + #[tokio::test] + async fn exact_commit_verification_requires_a_valid_identical_binding() { + let state = runtime_state(); + let registry = LiveCallRegistry::new(Arc::clone(&state)); + let key = record_key("user", "key", "rtc_verify").unwrap(); + let expected = binding("global"); + + assert!( + !registry + .exact_binding_is_committed(key.as_str(), &expected) + .await + ); + + state + .kv_set(key.as_str(), "{not-json".to_string(), None) + .await + .unwrap(); + assert!( + !registry + .exact_binding_is_committed(key.as_str(), &expected) + .await + ); + + state + .kv_set( + key.as_str(), + serde_json::to_string(&binding("different")).unwrap(), + None, + ) + .await + .unwrap(); + assert!( + !registry + .exact_binding_is_committed(key.as_str(), &expected) + .await + ); + + state + .kv_set( + key.as_str(), + serde_json::to_string(&expected).unwrap(), + None, + ) + .await + .unwrap(); + assert!( + registry + .exact_binding_is_committed(key.as_str(), &expected) + .await + ); + } +} diff --git a/apps/aether-gateway/src/handlers/proxy/websocket/live/session.rs b/apps/aether-gateway/src/handlers/proxy/websocket/live/session.rs new file mode 100644 index 000000000..34020bba1 --- /dev/null +++ b/apps/aether-gateway/src/handlers/proxy/websocket/live/session.rs @@ -0,0 +1,1268 @@ +//! Opaque direct and WebRTC-sideband WebSocket relay for Codex Live. + +use std::future::Future; +use std::time::{Duration, Instant}; + +use axum::extract::ws::{Message as AxumWsMessage, WebSocket}; +use axum::http::StatusCode; +use futures_util::StreamExt; +use serde_json::{json, Value}; +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::{ + wait_for_optional_deadline, CLOSE_INTERNAL_ERROR, 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, +}; +use crate::{AppState, GatewayError}; + +use super::live_usage_accounting_is_safe; +use super::planner::{ + build_live_stream_admission_attempt, direct_live_websocket_url, live_sideband_url, + plan_live_candidate, LivePoolLeaseGuard, PlannedLiveCandidate, +}; +use super::protocol::{ + call_id_from_path, direct_model_from_query, event_type, validate_initial_session_update, +}; +use super::registry::{ + LiveCallBinding, LiveCallLookup, LiveCallRegistry, LiveCallRegistryError, LiveSidebandLease, + LiveSidebandLeaseLoss, +}; + +const LIVE_LOG_TARGET: &str = "aether_gateway::handlers::proxy::codex_live"; +const SIDEBAND_LOOKUP_TIMEOUT: Duration = Duration::from_millis(500); +const SESSION_CLOSE_DRAIN_TIMEOUT: Duration = Duration::from_secs(5); +const LIVE_CONNECTION_LOG_SPEC: WebSocketConnectionLogSpec = WebSocketConnectionLogSpec { + opened_event_name: "codex_live_websocket_connection_opened", + closed_event_name: "codex_live_websocket_connection_closed", + opened_message: "gateway accepted Codex Live WebSocket connection", + closed_message: "gateway closed Codex Live WebSocket connection", + execution_path: "codex_live_websocket_bridge", + provider_type: "codex_live", +}; +const LIVE_UPSTREAM_ERRORS: UpstreamWebSocketErrorCodes = UpstreamWebSocketErrorCodes { + upstream_url_missing: "codex_live_upstream_url_missing", + upstream_url_invalid: "codex_live_upstream_url_invalid", + frontdoor_self_loop: "codex_live_websocket_frontdoor_self_loop", + headers_invalid: "codex_live_websocket_headers_invalid", + client_build_failed: "codex_live_websocket_client_build_failed", + proxy_invalid: "codex_live_websocket_proxy_invalid", + tunnel_proxy_unsupported: "codex_live_websocket_tunnel_proxy_unsupported", + handshake_failed: "codex_live_websocket_handshake_failed", + upgrade_rejected: "codex_live_websocket_upgrade_rejected", + upgrade_failed: "codex_live_websocket_upgrade_failed", +}; + +#[derive(Debug)] +enum LiveRelayAdmissionError { + PlanUnavailable, + BalanceRejected, + Gateway(GatewayError), +} + +impl LiveRelayAdmissionError { + fn kind(&self) -> &'static str { + match self { + Self::PlanUnavailable => "plan_unavailable", + Self::BalanceRejected => "balance_rejected", + Self::Gateway(error) => gateway_error_kind(error), + } + } + + fn response(&self) -> (u16, &'static str, &'static str, u16, &'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", + ), + } + } +} + +pub(super) enum PreparedLiveWebSocket { + Direct { client_model: String }, + Sideband(PreparedLiveSideband), +} + +pub(super) struct PreparedLiveSideband { + call_id: String, + binding: LiveCallBinding, + lease: LiveSidebandLease, +} + +pub(super) struct LiveWebSocketPreflightRejection { + status: StatusCode, + message: &'static str, +} + +impl LiveWebSocketPreflightRejection { + pub(super) const fn status(&self) -> StatusCode { + self.status + } + + pub(super) const fn message(&self) -> &'static str { + self.message + } +} + +pub(super) async fn prepare_live_websocket( + state: &AppState, + context: &WebSocketRequestContext, +) -> Result { + if !live_usage_accounting_is_safe(&context.decision) { + 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 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", + }); + } + 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(), + } + })?; + return Ok(PreparedLiveWebSocket::Direct { client_model }); + } + 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 registry = LiveCallRegistry::new(std::sync::Arc::clone(&state.runtime_state)); + let lookup = tokio::time::timeout( + SIDEBAND_LOOKUP_TIMEOUT, + registry.lookup_with_status( + auth.user_id.as_str(), + auth.api_key_id.as_str(), + call_id.as_str(), + ), + ) + .await; + let binding = match lookup { + Ok(Ok(LiveCallLookup::Found(binding))) => binding, + Ok(Ok(LiveCallLookup::Missing)) => { + info!( + target: LIVE_LOG_TARGET, + event_name = "codex_live_sideband_binding_miss", + log_type = "event", + transport = WEBSOCKET_LOG_TRANSPORT, + websocket = true, + 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", + }); + } + Ok(Ok(LiveCallLookup::Expired)) => { + info!( + target: LIVE_LOG_TARGET, + event_name = "codex_live_sideband_binding_expired", + log_type = "event", + transport = WEBSOCKET_LOG_TRANSPORT, + websocket = true, + trace_id = %context.trace_id, + "Codex Live sideband binding has expired" + ); + return Err(LiveWebSocketPreflightRejection { + status: StatusCode::GONE, + message: "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", + }); + } + Err(_) => { + return Err(LiveWebSocketPreflightRejection { + status: StatusCode::SERVICE_UNAVAILABLE, + message: "Timed out loading the Codex Live sideband binding", + }); + } + }; + let lease = match tokio::time::timeout( + SIDEBAND_LOOKUP_TIMEOUT, + registry.acquire_sideband_attachment( + auth.user_id.as_str(), + auth.api_key_id.as_str(), + call_id.as_str(), + ), + ) + .await + { + Ok(Ok(lease)) => lease, + Ok(Err(LiveCallRegistryError::SidebandAlreadyAttached)) => { + info!( + target: LIVE_LOG_TARGET, + event_name = "codex_live_sideband_attachment_conflict", + log_type = "event", + transport = WEBSOCKET_LOG_TRANSPORT, + websocket = true, + trace_id = %context.trace_id, + provider_id = %binding.pinned_candidate().provider_id(), + endpoint_id = %binding.pinned_candidate().endpoint_id(), + 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", + }); + } + 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", + }); + } + Err(_) => { + return Err(LiveWebSocketPreflightRejection { + status: StatusCode::SERVICE_UNAVAILABLE, + message: "Timed out acquiring Codex Live sideband ownership", + }); + } + }; + Ok(PreparedLiveWebSocket::Sideband(PreparedLiveSideband { + call_id, + binding, + lease, + })) +} + +pub(super) async fn run_live_websocket( + mut client_socket: WebSocket, + state: AppState, + context: WebSocketRequestContext, + prepared: PreparedLiveWebSocket, +) { + 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::Sideband(prepared) => { + run_sideband(&mut client_socket, &state, &context, prepared).await + } + } +} + +async fn run_direct( + client_socket: &mut WebSocket, + state: &AppState, + context: &WebSocketRequestContext, + client_model: String, +) { + if !live_usage_accounting_is_safe(&context.decision) { + reject_finite_balance_live(client_socket, context).await; + return; + } + let initial = match read_initial_session_update(client_socket).await { + Ok(Some(initial)) => initial, + Ok(None) => return, + Err(error) => { + send_live_error( + client_socket, + error.status_code().as_u16(), + error.code(), + error.client_message(), + ) + .await; + close_client_socket( + client_socket, + if error.is_timeout() { + CLOSE_TRY_AGAIN + } else { + CLOSE_POLICY_VIOLATION + }, + "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" + ); + 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; + return; + } + }; + let initial = + rewrite_live_session_model(initial.as_str(), provider_model.as_str()).unwrap_or(initial); + 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; + close_client_socket(client_socket, CLOSE_TRY_AGAIN, "Live upstream write failed").await; + return; + } + relay_live( + client_socket, + &mut upstream, + context, + "direct", + provider_id.as_str(), + endpoint_id.as_str(), + key_id.as_str(), + provider_model.as_str(), + &lease, + None, + ) + .await; + close_upstream_socket(&mut upstream, None).await; + admission.release().await; + lease.release().await; +} + +async fn run_sideband( + client_socket: &mut WebSocket, + state: &AppState, + context: &WebSocketRequestContext, + prepared: PreparedLiveSideband, +) { + let PreparedLiveSideband { + call_id, + binding, + lease: 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( + client_socket, + &mut upstream, + context, + "sideband", + provider_id.as_str(), + endpoint_id.as_str(), + key_id.as_str(), + provider_model.as_str(), + &pool_lease, + Some(&sideband_lease), + ) + .await; + close_upstream_socket(&mut upstream, None).await; + release_sideband_lease(&mut sideband_lease, context).await; + admission.release().await; + pool_lease.release().await; +} + +async fn acquire_live_relay_admission( + state: &AppState, + context: &WebSocketRequestContext, + candidate: &PlannedLiveCandidate, + upstream_url: String, +) -> Result { + let Some(attempt) = build_live_stream_admission_attempt( + candidate, + &context.headers, + &context.remote_addr, + upstream_url, + ) + .map_err(LiveRelayAdmissionError::Gateway)? + else { + return Err(LiveRelayAdmissionError::PlanUnavailable); + }; + if 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); + } + ResponsesWebSocketTurnAdmission::acquire(state, &attempt.plan, context.trace_id.as_str()) + .await + .map_err(LiveRelayAdmissionError::Gateway) +} + +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, + context: &WebSocketRequestContext, + mode: &'static str, + error: &LiveRelayAdmissionError, +) { + let (status, code, message, close_code, close_reason) = error.response(); + warn!( + target: LIVE_LOG_TARGET, + event_name = "codex_live_relay_admission_failed", + log_type = "ops", + transport = WEBSOCKET_LOG_TRANSPORT, + websocket = true, + trace_id = %context.trace_id, + mode, + status, + error_kind = error.kind(), + "Codex Live relay admission failed" + ); + send_live_error(client_socket, status, code, message).await; + close_client_socket(client_socket, close_code, close_reason).await; +} + +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 reject_lost_pool_lease( + client_socket: &mut WebSocket, + context: &WebSocketRequestContext, + mode: &'static str, + provider_id: &str, + endpoint_id: &str, + key_id: &str, +) { + warn!( + target: LIVE_LOG_TARGET, + event_name = "codex_live_pool_key_lease_lost", + log_type = "ops", + transport = WEBSOCKET_LOG_TRANSPORT, + websocket = true, + trace_id = %context.trace_id, + provider_id, + endpoint_id, + key_id, + mode, + "Codex Live relay stopped after losing its scheduler pool-key lease" + ); + send_live_error( + client_socket, + 503, + "codex_live_pool_key_lease_lost", + "Codex Live provider ownership was lost", + ) + .await; + close_client_socket( + client_socket, + CLOSE_TRY_AGAIN, + "Live provider ownership lost", + ) + .await; +} + +async fn read_initial_session_update( + client_socket: &mut WebSocket, +) -> Result, super::protocol::LiveProtocolError> { + tokio::time::timeout( + LIVE_WEBSOCKET_SESSION_LIMITS.initial_message_timeout, + async { + loop { + let Some(message) = client_socket.next().await else { + return Ok(None); + }; + let message = message + .map_err(|_| super::protocol::LiveProtocolError::InitialClientReadFailed)?; + match message { + AxumWsMessage::Text(text) => { + validate_initial_session_update(text.as_str())?; + return Ok(Some(text.to_string())); + } + AxumWsMessage::Ping(payload) => { + send_client_message(client_socket, AxumWsMessage::Pong(payload)) + .await + .map_err(|_| { + super::protocol::LiveProtocolError::InitialClientReadFailed + })?; + } + AxumWsMessage::Pong(_) => {} + AxumWsMessage::Close(_) => return Ok(None), + AxumWsMessage::Binary(_) => { + return Err(super::protocol::LiveProtocolError::InitialEventMustBeText) + } + } + } + }, + ) + .await + .map_err(|_| super::protocol::LiveProtocolError::InitialSessionUpdateTimeout)? +} + +#[derive(Default)] +struct RelayStats { + client_frames: u64, + client_bytes: u64, + upstream_frames: u64, + upstream_bytes: u64, +} + +async fn relay_live( + client_socket: &mut WebSocket, + upstream: &mut wreq::ws::WebSocket, + context: &WebSocketRequestContext, + mode: &'static str, + provider_id: &str, + endpoint_id: &str, + key_id: &str, + 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); + 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" + } + }; + } + 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 upstream_message = client_message_to_upstream(client); + if send_upstream_message(upstream, upstream_message).await.is_err() { + break "upstream_write_failed"; + } + if is_session_close { + 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"; + } + if is_close { break "upstream_close_frame"; } + } + } + }; + info!( + target: LIVE_LOG_TARGET, + event_name = "codex_live_relay_finished", + log_type = "event", + transport = WEBSOCKET_LOG_TRANSPORT, + websocket = true, + trace_id = %context.trace_id, + provider_id, + endpoint_id, + key_id, + mode, + termination, + client_frames = stats.client_frames, + client_bytes = stats.client_bytes, + upstream_frames = stats.upstream_frames, + upstream_bytes = stats.upstream_bytes, + elapsed_ms = started_at.elapsed().as_millis() as u64, + usage_unavailable = true, + "Codex Live opaque relay finished" + ); +} + +async fn while_sideband_lease_healthy( + lease: &LiveSidebandLease, + operation: F, +) -> Result +where + F: Future, +{ + if let Some(loss) = lease.loss() { + return Err(loss); + } + tokio::select! { + loss = lease.wait_for_loss() => Err(loss), + output = operation => lease.loss().map_or(Ok(output), Err), + } +} + +async fn wait_for_sideband_lease_loss(lease: Option<&LiveSidebandLease>) -> LiveSidebandLeaseLoss { + match lease { + Some(lease) => lease.wait_for_loss().await, + None => std::future::pending().await, + } +} + +async fn reject_sideband_lease_loss( + client_socket: &mut WebSocket, + context: &WebSocketRequestContext, + loss: LiveSidebandLeaseLoss, +) { + warn!( + target: LIVE_LOG_TARGET, + event_name = "codex_live_sideband_lease_lost", + log_type = "ops", + transport = WEBSOCKET_LOG_TRANSPORT, + websocket = true, + trace_id = %context.trace_id, + error_kind = loss.kind(), + "Codex Live sideband attachment lease was lost" + ); + let (code, message, close_reason) = match loss { + LiveSidebandLeaseLoss::OwnershipLost => ( + "codex_live_sideband_lease_lost", + "Codex Live sideband ownership was lost", + "Live sideband ownership lost", + ), + LiveSidebandLeaseLoss::StorageUnavailable => ( + "codex_live_sideband_lease_unavailable", + "Codex Live sideband ownership could not be renewed", + "Live sideband ownership renewal failed", + ), + }; + send_live_error(client_socket, 503, code, message).await; + close_client_socket(client_socket, CLOSE_TRY_AGAIN, close_reason).await; +} + +/// Rewrites only the routing-authoritative model field while keeping the +/// evolving Frameless event schema opaque. Invalid JSON, non-session events, +/// and session updates without an explicit model are forwarded byte-for-byte. +fn rewrite_live_session_model(raw: &str, provider_model: &str) -> Option { + let mut event: Value = serde_json::from_str(raw).ok()?; + if event.get("type").and_then(Value::as_str) != Some("session.update") { + return None; + } + let session = event.get_mut("session")?.as_object_mut()?; + let model = session.get_mut("model")?; + *model = Value::String(provider_model.to_string()); + serde_json::to_string(&event).ok() +} + +fn client_frame_metadata(message: &AxumWsMessage) -> (usize, bool, bool) { + match message { + AxumWsMessage::Text(text) => ( + text.len(), + false, + event_type(text.as_str()).as_deref() == Some("session.close"), + ), + AxumWsMessage::Binary(data) => (data.len(), false, false), + AxumWsMessage::Ping(data) | AxumWsMessage::Pong(data) => (data.len(), false, false), + AxumWsMessage::Close(frame) => ( + frame + .as_ref() + .map_or(0, |frame| 2usize.saturating_add(frame.reason.len())), + true, + false, + ), + } +} + +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; + } + } +} + +async fn send_live_error(client_socket: &mut WebSocket, status: u16, code: &str, message: &str) { + let event = json!({ + "type": "error", + "status": status, + "error": { + "type": "invalid_request_error", + "code": code, + "message": message, + } + }); + let _ = send_client_message(client_socket, AxumWsMessage::Text(event.to_string().into())).await; +} + +fn log_registry_error(context: &WebSocketRequestContext, error: &LiveCallRegistryError) { + warn!( + target: LIVE_LOG_TARGET, + event_name = "codex_live_sideband_binding_lookup_failed", + log_type = "ops", + transport = WEBSOCKET_LOG_TRANSPORT, + websocket = true, + trace_id = %context.trace_id, + error_kind = error.kind(), + "Codex Live sideband binding lookup failed" + ); +} + +fn log_sideband_lease_error( + context: &WebSocketRequestContext, + error: &LiveCallRegistryError, + operation: &'static str, +) { + warn!( + target: LIVE_LOG_TARGET, + event_name = "codex_live_sideband_lease_operation_failed", + log_type = "ops", + transport = WEBSOCKET_LOG_TRANSPORT, + websocket = true, + trace_id = %context.trace_id, + operation, + error_kind = error.kind(), + "Codex Live sideband attachment lease operation failed" + ); +} + +async fn release_sideband_lease(lease: &mut LiveSidebandLease, context: &WebSocketRequestContext) { + match lease.release().await { + Ok(true) => {} + Ok(false) => warn!( + target: LIVE_LOG_TARGET, + event_name = "codex_live_sideband_lease_release_not_owned", + log_type = "ops", + transport = WEBSOCKET_LOG_TRANSPORT, + websocket = true, + trace_id = %context.trace_id, + "Codex Live sideband attachment lease was not owned during release" + ), + Err(error) => log_sideband_lease_error(context, &error, "release"), + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn live_has_a_distinct_frontdoor_self_loop_error() { + assert_eq!( + LIVE_UPSTREAM_ERRORS.frontdoor_self_loop, + "codex_live_websocket_frontdoor_self_loop" + ); + } + + #[test] + fn turn_done_is_opaque_and_does_not_request_connection_close() { + let message = AxumWsMessage::Text(r#"{"type":"turn.done","future":true}"#.into()); + let (_, is_close, is_session_close) = client_frame_metadata(&message); + assert!(!is_close); + assert!(!is_session_close); + } + + #[test] + fn only_session_close_starts_the_bounded_drain() { + let message = AxumWsMessage::Text(r#"{"type":"session.close","future":true}"#.into()); + let (_, is_close, is_session_close) = client_frame_metadata(&message); + assert!(!is_close); + assert!(is_session_close); + } + + #[test] + fn session_update_model_is_pinned_without_dropping_unknown_fields() { + let rewritten = rewrite_live_session_model( + r#"{"type":"session.update","session":{"model":"client-alias","future":true},"event_id":"evt_1","unknown":{"nested":1}}"#, + "provider-model", + ) + .expect("session model should be rewritten"); + let value: Value = serde_json::from_str(&rewritten).expect("rewritten JSON"); + assert_eq!(value["session"]["model"], "provider-model"); + assert_eq!(value["session"]["future"], true); + assert_eq!(value["event_id"], "evt_1"); + assert_eq!(value["unknown"]["nested"], 1); + } + + #[test] + fn non_routing_live_frames_remain_opaque() { + for raw in [ + r#"{"type":"input_audio_buffer.append","audio":"AA=="}"#, + r#"{"type":"session.update","session":{"future":true}}"#, + r#"{"type":"session.update","session":null}"#, + "not-json", + ] { + assert_eq!(rewrite_live_session_model(raw, "provider-model"), None); + } + } +} diff --git a/apps/aether-gateway/src/handlers/proxy/websocket/mod.rs b/apps/aether-gateway/src/handlers/proxy/websocket/mod.rs index 74c794321..e12c4bc18 100644 --- a/apps/aether-gateway/src/handlers/proxy/websocket/mod.rs +++ b/apps/aether-gateway/src/handlers/proxy/websocket/mod.rs @@ -7,6 +7,7 @@ //! decisions. pub(crate) mod ingress; +pub(crate) mod live; pub(crate) mod responses; pub(crate) mod session; pub(crate) mod transport; diff --git a/apps/aether-gateway/src/handlers/proxy/websocket/responses/adapter.rs b/apps/aether-gateway/src/handlers/proxy/websocket/responses/adapter.rs index 3ab4c74b9..3b3d6a2d8 100644 --- a/apps/aether-gateway/src/handlers/proxy/websocket/responses/adapter.rs +++ b/apps/aether-gateway/src/handlers/proxy/websocket/responses/adapter.rs @@ -134,6 +134,7 @@ const STANDARD_UPSTREAM_WEBSOCKET_ERRORS: UpstreamWebSocketErrorCodes = UpstreamWebSocketErrorCodes { upstream_url_missing: "responses_upstream_url_missing", upstream_url_invalid: "responses_upstream_url_invalid", + frontdoor_self_loop: "responses_websocket_frontdoor_self_loop", headers_invalid: "responses_websocket_headers_invalid", client_build_failed: "responses_websocket_client_build_failed", proxy_invalid: "responses_websocket_proxy_invalid", @@ -215,6 +216,10 @@ mod tests { adapter.upstream_errors().handshake_failed, "responses_websocket_handshake_failed" ); + assert_eq!( + adapter.upstream_errors().frontdoor_self_loop, + "responses_websocket_frontdoor_self_loop" + ); } #[test] diff --git a/apps/aether-gateway/src/handlers/proxy/websocket/responses/adapters/codex.rs b/apps/aether-gateway/src/handlers/proxy/websocket/responses/adapters/codex.rs index 059f00559..d422e73fb 100644 --- a/apps/aether-gateway/src/handlers/proxy/websocket/responses/adapters/codex.rs +++ b/apps/aether-gateway/src/handlers/proxy/websocket/responses/adapters/codex.rs @@ -24,6 +24,7 @@ const CODEX_WEBSOCKET_RATE_LIMITS_REPORT_CONTEXT_FIELD: &str = "codex_websocket_ const CODEX_UPSTREAM_WEBSOCKET_ERRORS: UpstreamWebSocketErrorCodes = UpstreamWebSocketErrorCodes { upstream_url_missing: "codex_upstream_url_missing", upstream_url_invalid: "codex_upstream_url_invalid", + frontdoor_self_loop: "codex_websocket_frontdoor_self_loop", headers_invalid: "codex_websocket_headers_invalid", client_build_failed: "codex_websocket_client_build_failed", proxy_invalid: "codex_websocket_proxy_invalid", @@ -259,6 +260,16 @@ mod tests { ResponsesWebSocketRebindSafety, ResponsesWebSocketRelayDirective, }; + #[test] + fn codex_adapter_has_a_distinct_frontdoor_self_loop_error() { + let adapter = CodexResponsesWebSocketAdapter; + + assert_eq!( + adapter.upstream_errors().frontdoor_self_loop, + "codex_websocket_frontdoor_self_loop" + ); + } + #[test] fn codex_rate_limit_chunk_is_kept_for_the_terminal_report() { let adapter = CodexResponsesWebSocketAdapter; diff --git a/apps/aether-gateway/src/handlers/proxy/websocket/responses/admission.rs b/apps/aether-gateway/src/handlers/proxy/websocket/responses/admission.rs index 9c7779ca9..8192f5e90 100644 --- a/apps/aether-gateway/src/handlers/proxy/websocket/responses/admission.rs +++ b/apps/aether-gateway/src/handlers/proxy/websocket/responses/admission.rs @@ -16,7 +16,7 @@ use crate::provider_pool_demand::{ use crate::upstream_admission::UpstreamTargetAdmissionPermit; use crate::{AppState, GatewayError}; -pub(super) struct ResponsesWebSocketTurnAdmission { +pub(crate) struct ResponsesWebSocketTurnAdmission { upstream_execution: Option, upstream_target: Option, provider_pool: Option, @@ -24,7 +24,7 @@ pub(super) struct ResponsesWebSocketTurnAdmission { } impl ResponsesWebSocketTurnAdmission { - pub(super) async fn acquire( + pub(crate) async fn acquire( state: &AppState, plan: &ExecutionPlan, trace_id: &str, @@ -61,7 +61,7 @@ impl ResponsesWebSocketTurnAdmission { /// Release the distributed provider token before the turn's persistence /// work. The remaining permits are local RAII guards and are dropped with /// this value. - pub(super) async fn release(mut self) { + pub(crate) async fn release(mut self) { if let Some(provider_pool) = self.provider_pool.take() { provider_pool.release().await; } diff --git a/apps/aether-gateway/src/handlers/proxy/websocket/responses/mod.rs b/apps/aether-gateway/src/handlers/proxy/websocket/responses/mod.rs index bfe954305..b262513fe 100644 --- a/apps/aether-gateway/src/handlers/proxy/websocket/responses/mod.rs +++ b/apps/aether-gateway/src/handlers/proxy/websocket/responses/mod.rs @@ -29,6 +29,8 @@ mod turn; mod turn_state; mod upstream; +pub(crate) use admission::ResponsesWebSocketTurnAdmission; + use std::net::SocketAddr; use axum::body::Body; diff --git a/apps/aether-gateway/src/handlers/proxy/websocket/session.rs b/apps/aether-gateway/src/handlers/proxy/websocket/session.rs index 1e6549957..0c3b41ad5 100644 --- a/apps/aether-gateway/src/handlers/proxy/websocket/session.rs +++ b/apps/aether-gateway/src/handlers/proxy/websocket/session.rs @@ -20,6 +20,13 @@ pub(crate) const RESPONSES_WEBSOCKET_SESSION_LIMITS: WebSocketSessionLimits = max_connection_duration: Duration::from_secs(60 * 60), }; +pub(crate) const LIVE_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 5b33607f1..bd0352724 100644 --- a/apps/aether-gateway/src/handlers/proxy/websocket/transport.rs +++ b/apps/aether-gateway/src/handlers/proxy/websocket/transport.rs @@ -22,6 +22,7 @@ use crate::ai_serving::AiExecutionDecision; use crate::execution_runtime::transport::{ build_browser_wreq_client, build_request_headers, ExecutionTransportControls, }; +use crate::frontdoor_loop_guard::gateway_frontdoor_self_loop_guard_error; use crate::handlers::proxy::websocket::session::{ WebSocketSessionLimits, RELAY_WRITE_TIMEOUT, TEARDOWN_WRITE_TIMEOUT, }; @@ -30,6 +31,7 @@ use crate::handlers::proxy::websocket::session::{ pub(crate) struct UpstreamWebSocketErrorCodes { pub(crate) upstream_url_missing: &'static str, pub(crate) upstream_url_invalid: &'static str, + pub(crate) frontdoor_self_loop: &'static str, pub(crate) headers_invalid: &'static str, pub(crate) client_build_failed: &'static str, pub(crate) proxy_invalid: &'static str, @@ -53,7 +55,11 @@ pub(crate) async fn connect_upstream_websocket( .upstream_url .as_deref() .ok_or(errors.upstream_url_missing)?; - let upstream_url = websocket_upstream_url(upstream_url, errors.upstream_url_invalid)?; + let upstream_url = guarded_websocket_upstream_url( + upstream_url, + errors.upstream_url_invalid, + errors.frontdoor_self_loop, + )?; let headers = websocket_handshake_headers(&decision.provider_request_headers, errors.headers_invalid)?; let client = build_websocket_client(decision, errors)?; @@ -79,6 +85,18 @@ pub(crate) async fn connect_upstream_websocket( }) } +fn guarded_websocket_upstream_url( + raw: &str, + invalid_code: &'static str, + frontdoor_self_loop_code: &'static str, +) -> Result { + let upstream_url = websocket_upstream_url(raw, invalid_code)?; + if gateway_frontdoor_self_loop_guard_error(upstream_url.as_str()).is_some() { + return Err(frontdoor_self_loop_code); + } + Ok(upstream_url) +} + fn websocket_response_headers(headers: &HeaderMap) -> BTreeMap { headers .iter() @@ -310,6 +328,19 @@ pub(crate) fn upstream_message_to_client(message: WreqWsMessage) -> AxumWsMessag } } +pub(crate) fn client_message_to_upstream(message: AxumWsMessage) -> WreqWsMessage { + match message { + AxumWsMessage::Text(text) => WreqWsMessage::Text(text.to_string().into()), + AxumWsMessage::Binary(data) => WreqWsMessage::Binary(data), + AxumWsMessage::Ping(data) => WreqWsMessage::Ping(data), + AxumWsMessage::Pong(data) => WreqWsMessage::Pong(data), + AxumWsMessage::Close(frame) => WreqWsMessage::Close(frame.map(|frame| WreqCloseFrame { + code: frame.code.into(), + reason: frame.reason.to_string().into(), + })), + } +} + /// Builds a Responses WebSocket error event in the shape understood by the /// official client implementations. The status is part of the event body, /// not the WebSocket handshake, because the connection is already upgraded. @@ -467,11 +498,12 @@ pub(crate) async fn close_client_socket(client_socket: &mut WebSocket, code: u16 #[cfg(test)] mod tests { use super::{ - bounded_send, responses_websocket_error_event, + 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, }; + use crate::frontdoor_loop_guard::configured_gateway_frontdoor_base_url; use axum::http::HeaderMap; use std::collections::BTreeMap; use std::time::Duration; @@ -587,6 +619,39 @@ mod tests { assert!(websocket_upstream_url("https://token@example.test/responses", "invalid").is_err()); } + #[test] + fn rejects_responses_websocket_frontdoor_self_loop_before_connecting() { + let base_url = configured_gateway_frontdoor_base_url(); + let raw_url = format!("{base_url}/v1/responses"); + + assert_eq!( + guarded_websocket_upstream_url( + raw_url.as_str(), + "responses_upstream_url_invalid", + "responses_websocket_frontdoor_self_loop", + ), + Err("responses_websocket_frontdoor_self_loop") + ); + } + + #[test] + fn rejects_live_direct_and_sideband_frontdoor_self_loops_before_connecting() { + let base_url = configured_gateway_frontdoor_base_url(); + + for path in ["/v1/live", "/v1/live/rtc_test"] { + let raw_url = format!("{base_url}{path}"); + assert_eq!( + guarded_websocket_upstream_url( + raw_url.as_str(), + "codex_live_upstream_url_invalid", + "codex_live_websocket_frontdoor_self_loop", + ), + Err("codex_live_websocket_frontdoor_self_loop"), + "{path} must be rejected before an upstream handshake" + ); + } + } + #[test] fn upstream_handshake_keeps_provider_auth_but_drops_transport_managed_headers() { let provider_headers = BTreeMap::from([ diff --git a/apps/aether-gateway/src/tests/ai_execute/stream/decision.rs b/apps/aether-gateway/src/tests/ai_execute/stream/decision.rs index 3f901976d..4759ab296 100644 --- a/apps/aether-gateway/src/tests/ai_execute/stream/decision.rs +++ b/apps/aether-gateway/src/tests/ai_execute/stream/decision.rs @@ -321,15 +321,22 @@ async fn gateway_executes_openai_chat_stream_via_local_decision_gate_without_exe .unwrap_or_default() .to_string(), }); + let response_body = Body::from_stream(async_stream::stream! { + yield Ok::(Bytes::from_static( + b"data: {\"id\":\"chatcmpl-local-123\"}\n\n", + )); + tokio::time::sleep(std::time::Duration::from_millis(250)).await; + yield Ok::(Bytes::from_static( + b"data: [DONE]\n\n", + )); + }); let mut response = Response::builder() .status(StatusCode::OK) - .body(Body::from( - "data: {\"id\":\"chatcmpl-local-123\"}\n\ndata: [DONE]\n\n", - )) + .body(response_body) .expect("response should build"); response.headers_mut().insert( http::header::CONTENT_TYPE, - HeaderValue::from_static("text/event-stream"), + HeaderValue::from_static("application/octet-stream"), ); response } @@ -368,12 +375,14 @@ async fn gateway_executes_openai_chat_stream_via_local_decision_gate_without_exe DEVELOPMENT_ENCRYPTION_KEY, ), ); - let gateway = build_router_with_state(gateway_state); + let gateway = build_router_with_state(gateway_state) + .layer(tower_http::compression::CompressionLayer::new()); let (gateway_url, gateway_handle) = start_server(gateway).await; - let response = reqwest::Client::new() + let mut response = reqwest::Client::new() .post(format!("{gateway_url}/v1/chat/completions")) .header(http::header::CONTENT_TYPE, "application/json") + .header(http::header::ACCEPT_ENCODING, "gzip") .header( http::header::AUTHORIZATION, "Bearer sk-client-openai-local-stream", @@ -392,9 +401,32 @@ async fn gateway_executes_openai_chat_stream_via_local_decision_gate_without_exe .and_then(|value| value.to_str().ok()), Some(EXECUTION_PATH_EXECUTION_RUNTIME_STREAM) ); + assert_eq!( + response + .headers() + .get(http::header::CONTENT_TYPE) + .and_then(|value| value.to_str().ok()), + Some("text/event-stream") + ); + assert!( + response + .headers() + .get(http::header::CONTENT_ENCODING) + .is_none(), + "SSE responses must not be gzip-buffered" + ); + assert_eq!( + tokio::time::timeout( + std::time::Duration::from_millis(100), + super::super::next_non_keepalive_chunk(&mut response), + ) + .await + .expect("first upstream SSE event should reach the client before completion"), + Bytes::from_static(b"data: {\"id\":\"chatcmpl-local-123\"}\n\n") + ); assert_eq!( strip_sse_keepalive_comments(&response.text().await.expect("body should read")), - "data: {\"id\":\"chatcmpl-local-123\"}\n\ndata: [DONE]\n\n" + "data: [DONE]\n\n" ); let seen_upstream_request = seen_upstream diff --git a/apps/aether-gateway/src/tests/ai_execute/stream_cli/direct.rs b/apps/aether-gateway/src/tests/ai_execute/stream_cli/direct.rs index 623750cc8..8ad2bb958 100644 --- a/apps/aether-gateway/src/tests/ai_execute/stream_cli/direct.rs +++ b/apps/aether-gateway/src/tests/ai_execute/stream_cli/direct.rs @@ -504,16 +504,27 @@ async fn gateway_executes_codex_cli_stream_via_local_decision_gate_after_oauth_r .and_then(|value| value.get("json_body")) .is_some_and(|body| body.get("context_management").is_some()), }); - let frames = concat!( - "{\"type\":\"headers\",\"payload\":{\"kind\":\"headers\",\"status_code\":200,\"headers\":{\"content-type\":\"text/event-stream\"}}}\n", - "{\"type\":\"data\",\"payload\":{\"kind\":\"data\",\"text\":\"event: response.output_item.done\\ndata: {\\\"type\\\":\\\"response.output_item.done\\\",\\\"item\\\":{\\\"type\\\":\\\"compaction\\\",\\\"encrypted_content\\\":\\\"ENCRYPTED_CONTEXT_COMPACTION_SUMMARY\\\"}}\\n\\n\"}}\n", - "{\"type\":\"data\",\"payload\":{\"kind\":\"data\",\"text\":\"event: response.completed\\ndata: {\\\"type\\\":\\\"response.completed\\\",\\\"response\\\":{\\\"id\\\":\\\"resp_codex_cli_stream_local_123\\\",\\\"object\\\":\\\"response\\\",\\\"model\\\":\\\"gpt-5.6-sol\\\",\\\"status\\\":\\\"completed\\\",\\\"usage\\\":{\\\"input_tokens\\\":1,\\\"output_tokens\\\":2,\\\"total_tokens\\\":3}}}\\n\\n\"}}\n", - "{\"type\":\"telemetry\",\"payload\":{\"kind\":\"telemetry\",\"telemetry\":{\"elapsed_ms\":41}}}\n", - "{\"type\":\"eof\",\"payload\":{\"kind\":\"eof\"}}\n" - ); + let frames = async_stream::stream! { + yield Ok::(Bytes::from_static( + b"{\"type\":\"headers\",\"payload\":{\"kind\":\"headers\",\"status_code\":200,\"headers\":{\"content-type\":\"application/octet-stream\"}}}\n", + )); + yield Ok::(Bytes::from_static( + b"{\"type\":\"data\",\"payload\":{\"kind\":\"data\",\"text\":\"event: response.output_item.done\\ndata: {\\\"type\\\":\\\"response.output_item.done\\\",\\\"item\\\":{\\\"type\\\":\\\"compaction\\\",\\\"encrypted_content\\\":\\\"ENCRYPTED_CONTEXT_COMPACTION_SUMMARY\\\"}}\\n\\n\"}}\n", + )); + tokio::time::sleep(std::time::Duration::from_millis(250)).await; + yield Ok::(Bytes::from_static( + b"{\"type\":\"data\",\"payload\":{\"kind\":\"data\",\"text\":\"event: response.completed\\ndata: {\\\"type\\\":\\\"response.completed\\\",\\\"response\\\":{\\\"id\\\":\\\"resp_codex_cli_stream_local_123\\\",\\\"object\\\":\\\"response\\\",\\\"model\\\":\\\"gpt-5.6-sol\\\",\\\"status\\\":\\\"completed\\\",\\\"usage\\\":{\\\"input_tokens\\\":1,\\\"output_tokens\\\":2,\\\"total_tokens\\\":3}}}\\n\\n\"}}\n", + )); + yield Ok::(Bytes::from_static( + b"{\"type\":\"telemetry\",\"payload\":{\"kind\":\"telemetry\",\"telemetry\":{\"elapsed_ms\":41}}}\n", + )); + yield Ok::(Bytes::from_static( + b"{\"type\":\"eof\",\"payload\":{\"kind\":\"eof\"}}\n", + )); + }; let mut response = Response::builder() .status(StatusCode::OK) - .body(Body::from(frames)) + .body(Body::from_stream(frames)) .expect("response should build"); response.headers_mut().insert( http::header::CONTENT_TYPE, @@ -574,12 +585,14 @@ async fn gateway_executes_codex_cli_stream_via_local_decision_gate_after_oauth_r enabled: true, ..UsageRuntimeConfig::default() }); - let gateway = build_router_with_state(gateway_state); + let gateway = build_router_with_state(gateway_state) + .layer(tower_http::compression::CompressionLayer::new()); let (gateway_url, gateway_handle) = start_server(gateway).await; - let response = reqwest::Client::new() + let mut response = reqwest::Client::new() .post(format!("{gateway_url}/v1/responses")) .header(http::header::CONTENT_TYPE, "application/json") + .header(http::header::ACCEPT_ENCODING, "gzip") .header( http::header::AUTHORIZATION, format!("Bearer {client_api_key}"), @@ -599,8 +612,34 @@ async fn gateway_executes_codex_cli_stream_via_local_decision_gate_after_oauth_r .expect("request should succeed"); assert_eq!(response.status(), StatusCode::OK); + assert_eq!( + response + .headers() + .get(http::header::CONTENT_TYPE) + .and_then(|value| value.to_str().ok()), + Some("text/event-stream") + ); + assert!( + response + .headers() + .get(http::header::CONTENT_ENCODING) + .is_none(), + "SSE responses must not be gzip-buffered" + ); + let first_event = tokio::time::timeout( + std::time::Duration::from_millis(100), + super::super::next_non_keepalive_chunk(&mut response), + ) + .await + .expect("first upstream SSE event should reach the client before completion"); + assert!( + first_event.starts_with(b"event: response.output_item.done\n"), + "unexpected first event: {}", + String::from_utf8_lossy(&first_event) + ); let response_body = strip_sse_keepalive_comments(&response.text().await.expect("body should read")); + let response_body = format!("{}{}", String::from_utf8_lossy(&first_event), response_body); assert!(response_body.contains("event: response.output_item.done\n")); assert!(response_body.contains("\"type\":\"compaction\"")); assert!(response_body.contains("ENCRYPTED_CONTEXT_COMPACTION_SUMMARY")); diff --git a/apps/aether-gateway/src/tests/frontdoor/ai.rs b/apps/aether-gateway/src/tests/frontdoor/ai.rs index 4cf0a9f27..80a2900ea 100644 --- a/apps/aether-gateway/src/tests/frontdoor/ai.rs +++ b/apps/aether-gateway/src/tests/frontdoor/ai.rs @@ -30,10 +30,18 @@ use aether_data_contracts::repository::provider_catalog::{ StoredProviderCatalogProvider, }; use async_trait::async_trait; +use axum::extract::ws::{Message as AxumWsMessage, WebSocket, WebSocketUpgrade}; +use axum::extract::State; +use axum::http::{HeaderMap, Uri}; use axum::response::IntoResponse; +use axum::routing::get; +use base64::Engine as _; +use futures_util::SinkExt; use std::collections::HashMap; use std::future::pending; use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering}; +use tokio::sync::oneshot; +use wreq::ws::message::Message as WreqWsMessage; fn codex_models_snapshot( api_key_id: &str, @@ -1235,6 +1243,875 @@ async fn run_versioned_codex_model_cards_frontdoor_scenario() { execution_runtime_handle.abort(); } +#[test] +fn gateway_creates_bound_codex_live_oauth_calls_with_opaque_session_fields() { + super::run_frontdoor_async_test( + "codex-live-oauth-frontdoor", + run_codex_live_oauth_frontdoor_scenario(), + ); +} + +async fn run_codex_live_oauth_frontdoor_scenario() { + const PROVIDER_ID: &str = "provider-codex-live"; + const ENDPOINT_ID: &str = "endpoint-provider-codex-live"; + const UPSTREAM_KEY_ID: &str = "key-provider-codex-live"; + const CLIENT_MODEL: &str = "live-future-alias"; + 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); + row.key_allowed_models = Some(vec![PROVIDER_MODEL.to_string()]); + let candidate_repository = + Arc::new(InMemoryMinimalCandidateSelectionReadRepository::seed(vec![ + row, + ])); + 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]), + )])); + + let mut provider = codex_catalog_provider(PROVIDER_ID); + provider.config = Some(json!({ + "responses_websocket": {"enabled": true}, + "codex": {"fingerprint_convergence_enabled": true} + })); + let mut endpoint = codex_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]); + upstream_key.auth_type = "oauth".to_string(); + upstream_key.encrypted_auth_config = Some( + encrypt_python_fernet_plaintext( + DEVELOPMENT_ENCRYPTION_KEY, + r#"{"account_id":"account-live-1","is_fedramp":true}"#, + ) + .expect("Codex Live auth config should encrypt"), + ); + let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed( + vec![provider], + vec![endpoint], + vec![upstream_key], + )); + + let captured_plan = Arc::new(Mutex::new(None::)); + let captured_plan_for_runtime = Arc::clone(&captured_plan); + let execution_runtime = Router::new().route( + "/v1/execute/sync", + any(move |request: Request| { + let captured_plan_for_request = Arc::clone(&captured_plan_for_runtime); + async move { + let (_parts, body) = request.into_parts(); + let raw_body = to_bytes(body, usize::MAX) + .await + .expect("Live execution runtime request body should read"); + let plan: aether_contracts::ExecutionPlan = serde_json::from_slice(&raw_body) + .expect("Live execution runtime plan should parse"); + *captured_plan_for_request + .lock() + .expect("Live plan mutex should lock") = Some(plan.clone()); + Json(ExecutionResult { + request_id: plan.request_id, + candidate_id: plan.candidate_id, + status_code: 201, + headers: std::collections::BTreeMap::from([ + ("Content-Type".to_string(), "application/sdp".to_string()), + ( + "LOCATION".to_string(), + format!("https://api.openai.com/v1/live/{CALL_ID}"), + ), + ]), + response_observation: None, + body: Some(ResponseBody { + json_body: None, + body_bytes_b64: Some( + base64::engine::general_purpose::STANDARD + .encode(b"v=0\r\no=upstream-answer"), + ), + }), + telemetry: Some(ExecutionTelemetry { + ttfb_ms: Some(1), + elapsed_ms: Some(2), + upstream_bytes: Some(24), + }), + error: None, + }) + } + }), + ); + let (execution_runtime_url, execution_runtime_handle) = start_server(execution_runtime).await; + let state = build_state_with_execution_runtime_override(execution_runtime_url) + .with_data_state_for_tests( + crate::data::GatewayDataState::with_minimal_candidate_selection_and_auth_for_tests( + candidate_repository, + auth_repository, + ) + .attach_provider_catalog_repository_for_tests(provider_catalog_repository) + .with_encryption_key_for_tests(DEVELOPMENT_ENCRYPTION_KEY), + ); + let gateway = build_router_with_state(state); + let (gateway_url, gateway_handle) = start_server(gateway).await; + + let boundary = "aether-live-frontdoor"; + let offer_sdp = "v=0\r\no=client-offer"; + let session = json!({ + "model": CLIENT_MODEL, + "instructions": "Keep this opaque", + "future_capability": { + "revision": 7, + "nested": [true, {"mode": "future"}] + } + }); + let body = format!( + "--{boundary}\r\nContent-Disposition: form-data; name=\"sdp\"\r\nContent-Type: application/sdp\r\n\r\n{offer_sdp}\r\n--{boundary}\r\nContent-Disposition: form-data; name=\"session\"\r\nContent-Type: application/json\r\n\r\n{}\r\n--{boundary}--\r\n", + session + ); + let response = reqwest::Client::new() + .post(format!("{gateway_url}/v1/live")) + .header("authorization", "Bearer sk-codex-live") + .header( + "content-type", + format!("multipart/form-data; boundary={boundary}"), + ) + .header("session-id", "client-session-live") + .header("openai-alpha", "client-must-not-control-this") + .body(body) + .send() + .await + .expect("Codex Live call creation should complete"); + assert_eq!(response.status(), StatusCode::CREATED); + let downstream_location = format!("/v1/live/{CALL_ID}"); + assert_eq!( + response + .headers() + .get(http::header::LOCATION) + .and_then(|value| value.to_str().ok()), + Some(downstream_location.as_str()) + ); + assert_eq!( + response + .bytes() + .await + .expect("SDP answer should read") + .as_ref(), + b"v=0\r\no=upstream-answer" + ); + + let plan = captured_plan + .lock() + .expect("Live plan mutex should lock") + .clone() + .expect("Live call must reach the execution runtime"); + let url = url::Url::parse(plan.url.as_str()).expect("Live call URL should parse"); + assert_eq!(url.path(), "/backend-api/codex/realtime/calls"); + assert_eq!( + url.query_pairs().collect::>(), + HashMap::from([ + ("intent".into(), "quicksilver".into()), + ("architecture".into(), "avas".into()), + ]) + ); + assert_eq!(plan.method, "POST"); + assert_eq!(plan.content_type.as_deref(), Some("application/json")); + assert!(!plan.stream); + assert!(plan.body.json_body.is_none()); + let provider_body_bytes = base64::engine::general_purpose::STANDARD + .decode( + plan.body + .body_bytes_b64 + .as_deref() + .expect("OAuth Live call must preserve the exact JSON wire bytes"), + ) + .expect("OAuth Live JSON body should decode"); + let provider_body: serde_json::Value = serde_json::from_slice(&provider_body_bytes) + .expect("OAuth Live call must use the JSON call contract"); + assert_eq!(provider_body["sdp"], offer_sdp); + assert_eq!(provider_body["session"]["model"], PROVIDER_MODEL); + assert_eq!( + provider_body["session"]["future_capability"], + session["future_capability"] + ); + assert_eq!(provider_body["session"]["instructions"], "Keep this opaque"); + assert_eq!( + plan.headers.get("openai-alpha").map(String::as_str), + Some("quicksilver=v2") + ); + assert_eq!( + plan.headers.get("originator").map(String::as_str), + Some("codex_cli_rs") + ); + assert_eq!( + plan.headers.get("chatgpt-account-id").map(String::as_str), + Some("account-live-1") + ); + assert_eq!( + plan.headers.get("x-openai-fedramp").map(String::as_str), + Some("true") + ); + let converged_session = plan + .headers + .get("x-session-id") + .expect("Live must provide a converged session ID"); + assert_ne!(converged_session, "client-session-live"); + assert_eq!(plan.headers.get("thread-id"), Some(converged_session)); + uuid::Uuid::parse_str(converged_session).expect("converged session ID must be a UUID"); + + gateway_handle.abort(); + execution_runtime_handle.abort(); +} + +#[derive(Debug)] +struct ObservedCodexLiveWebSocket { + request_target: String, + authorization: Option, + alpha: Option, + session_id: Option, + initial_event: serde_json::Value, + event_after_turn_done: serde_json::Value, +} + +#[test] +fn gateway_relays_codex_live_api_key_websocket_opaquely() { + super::run_frontdoor_async_test( + "codex-live-api-key-websocket-frontdoor", + run_codex_live_api_key_websocket_frontdoor_scenario(), + ); +} + +async fn run_codex_live_api_key_websocket_frontdoor_scenario() { + const PROVIDER_ID: &str = "provider-codex-live-api-key"; + const ENDPOINT_ID: &str = "endpoint-provider-codex-live-api-key"; + const UPSTREAM_KEY_ID: &str = "key-provider-codex-live-api-key"; + const CLIENT_MODEL: &str = "live-websocket-alias"; + const PROVIDER_MODEL: &str = "gpt-future-live-websocket"; + + let (observed_tx, observed_rx) = oneshot::channel(); + let upstream_state = Arc::new(Mutex::new(Some(observed_tx))); + let upstream = Router::new() + .route("/v1/live", get(mock_codex_live_websocket)) + .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); + row.provider_name = "openai".to_string(); + row.provider_type = "openai".to_string(); + row.key_auth_type = "api_key".to_string(); + row.key_allowed_models = Some(vec![PROVIDER_MODEL.to_string()]); + let candidate_repository = + Arc::new(InMemoryMinimalCandidateSelectionReadRepository::seed(vec![ + row, + ])); + let mut downstream_snapshot = codex_models_snapshot( + "gateway-key-codex-live-websocket", + "user-codex-live-websocket", + &[CLIENT_MODEL], + ); + downstream_snapshot.user_allowed_providers = Some(vec!["openai".to_string()]); + downstream_snapshot.api_key_allowed_providers = Some(vec!["openai".to_string()]); + let auth_repository = Arc::new(InMemoryAuthApiKeySnapshotRepository::seed(vec![( + Some(hash_api_key("sk-codex-live-websocket")), + downstream_snapshot, + )])); + + 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); + endpoint.base_url = format!("{upstream_url}/v1"); + let mut upstream_key = codex_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], + vec![endpoint], + vec![upstream_key], + )); + + let state = AppState::new() + .expect("gateway should build") + .with_data_state_for_tests( + crate::data::GatewayDataState::with_minimal_candidate_selection_and_auth_for_tests( + candidate_repository, + auth_repository, + ) + .attach_provider_catalog_repository_for_tests(provider_catalog_repository) + .with_encryption_key_for_tests(DEVELOPMENT_ENCRYPTION_KEY), + ); + 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-codex-live-websocket"), + ); + handshake_headers.insert( + http::HeaderName::from_static("x-session-id"), + http::HeaderValue::from_static("stable-live-session"), + ); + handshake_headers.insert( + http::HeaderName::from_static("openai-alpha"), + http::HeaderValue::from_static("client-value-must-be-replaced"), + ); + let invalid_model_response = wreq::Client::new() + .websocket(format!( + "{}/v1/live?model={CLIENT_MODEL}&model=second-model", + gateway_url.replacen("http://", "ws://", 1) + )) + .headers(handshake_headers.clone()) + .send() + .await + .expect("invalid Live model handshake should return an HTTP response"); + assert_eq!(invalid_model_response.status(), StatusCode::BAD_REQUEST); + + let websocket_url = format!( + "{}/v1/live?foo=bar&model={CLIENT_MODEL}&trace=1", + gateway_url.replacen("http://", "ws://", 1) + ); + let response = wreq::Client::new() + .websocket(websocket_url) + .headers(handshake_headers) + .send() + .await + .expect("Codex Live gateway WebSocket handshake should complete"); + assert_eq!(response.status(), StatusCode::SWITCHING_PROTOCOLS); + let mut socket = response + .into_websocket() + .await + .expect("Codex Live gateway response should upgrade"); + + let initial_event = json!({ + "type": "session.update", + "session": { + "model": CLIENT_MODEL, + "instructions": "Relay this Live configuration" + }, + "future_client_field": { + "opaque": true, + "revision": 9, + "nested": [1, {"mode": "future"}] + } + }); + socket + .send(WreqWsMessage::text(initial_event.to_string())) + .await + .expect("initial Live session.update should send"); + + let future_event = receive_codex_live_json(&mut socket).await; + assert_eq!( + future_event, + json!({ + "type": "future.live.event", + "future_capability": {"enabled": true, "revision": 11} + }) + ); + let turn_done = receive_codex_live_json(&mut socket).await; + assert_eq!( + turn_done, + json!({ + "type": "turn.done", + "turn": {"id": "turn-live-1"}, + "future_turn_field": "retained" + }) + ); + + let event_after_turn_done = json!({ + "type": "future.client.after_turn_done", + "future_payload": {"still_connected": true} + }); + socket + .send(WreqWsMessage::text(event_after_turn_done.to_string())) + .await + .expect("Live socket should remain writable after turn.done"); + + let observed = tokio::time::timeout(std::time::Duration::from_secs(2), observed_rx) + .await + .expect("mock upstream should observe the post-turn event before timeout") + .expect("mock upstream observation channel should remain open"); + assert_eq!( + observed.request_target, + format!("/v1/live?model={PROVIDER_MODEL}") + ); + assert_eq!( + observed.authorization.as_deref(), + Some("Bearer oauth-upstream-secret") + ); + assert_eq!(observed.alpha.as_deref(), Some("quicksilver=v2")); + assert_eq!(observed.session_id.as_deref(), Some("stable-live-session")); + let mut expected_initial_event = initial_event; + expected_initial_event["session"]["model"] = json!(PROVIDER_MODEL); + assert_eq!(observed.initial_event, expected_initial_event); + assert_eq!(observed.event_after_turn_done, event_after_turn_done); + + drop(socket); + gateway_handle.abort(); + upstream_handle.abort(); +} + +#[derive(Debug)] +struct ObservedCodexLiveSideband { + request_target: String, + authorization: Option, + alpha: Option, + session_id: Option, + first_client_event: serde_json::Value, + session_update: serde_json::Value, +} + +#[test] +fn gateway_creates_and_relays_bound_codex_live_api_key_sideband() { + super::run_frontdoor_async_test( + "codex-live-api-key-sideband-frontdoor", + run_codex_live_api_key_sideband_frontdoor_scenario(), + ); +} + +async fn run_codex_live_api_key_sideband_frontdoor_scenario() { + const PROVIDER_ID: &str = "provider-codex-live-sideband"; + const ENDPOINT_ID: &str = "endpoint-provider-codex-live-sideband"; + const UPSTREAM_KEY_ID: &str = "key-provider-codex-live-sideband"; + const CLIENT_MODEL: &str = "live-sideband-alias"; + const PROVIDER_MODEL: &str = "gpt-future-live-sideband"; + const CALL_ID: &str = "rtc_live_sideband_1"; + + let (sideband_observed_tx, sideband_observed_rx) = oneshot::channel(); + let upstream_state = Arc::new(Mutex::new(Some(sideband_observed_tx))); + let upstream = Router::new() + .route( + "/v1/live/{call_id}", + get(mock_codex_live_sideband_websocket), + ) + .with_state(upstream_state); + let (upstream_url, upstream_handle) = start_server(upstream).await; + + let captured_plan = Arc::new(Mutex::new(None::)); + let captured_plan_for_runtime = Arc::clone(&captured_plan); + let upstream_location = format!("{upstream_url}/v1/live/{CALL_ID}"); + let execution_runtime = Router::new().route( + "/v1/execute/sync", + any(move |request: Request| { + let captured_plan_for_request = Arc::clone(&captured_plan_for_runtime); + let upstream_location = upstream_location.clone(); + async move { + let (_parts, body) = request.into_parts(); + let raw_body = to_bytes(body, usize::MAX) + .await + .expect("Live API-key execution runtime body should read"); + let plan: aether_contracts::ExecutionPlan = serde_json::from_slice(&raw_body) + .expect("Live API-key execution plan should parse"); + *captured_plan_for_request + .lock() + .expect("Live API-key plan mutex should lock") = Some(plan.clone()); + Json(ExecutionResult { + request_id: plan.request_id, + candidate_id: plan.candidate_id, + status_code: 201, + headers: std::collections::BTreeMap::from([ + ("content-type".to_string(), "application/sdp".to_string()), + ("location".to_string(), upstream_location), + ("x-future-live-header".to_string(), "preserved".to_string()), + ]), + response_observation: None, + body: Some(ResponseBody { + json_body: None, + body_bytes_b64: Some( + base64::engine::general_purpose::STANDARD + .encode(b"v=0\r\no=api-key-upstream-answer"), + ), + }), + telemetry: Some(ExecutionTelemetry { + ttfb_ms: Some(1), + elapsed_ms: Some(2), + upstream_bytes: Some(34), + }), + error: None, + }) + } + }), + ); + 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); + row.provider_name = "openai".to_string(); + row.provider_type = "openai".to_string(); + row.key_auth_type = "api_key".to_string(); + row.key_allowed_models = Some(vec![PROVIDER_MODEL.to_string()]); + let candidate_repository = + Arc::new(InMemoryMinimalCandidateSelectionReadRepository::seed(vec![ + row, + ])); + let mut downstream_snapshot = codex_models_snapshot( + "gateway-key-codex-live-sideband", + "user-codex-live-sideband", + &[CLIENT_MODEL], + ); + downstream_snapshot.user_allowed_providers = Some(vec!["openai".to_string()]); + downstream_snapshot.api_key_allowed_providers = Some(vec!["openai".to_string()]); + let auth_repository = Arc::new(InMemoryAuthApiKeySnapshotRepository::seed(vec![( + Some(hash_api_key("sk-codex-live-sideband")), + downstream_snapshot, + )])); + + 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); + endpoint.base_url = format!("{upstream_url}/v1"); + let mut upstream_key = codex_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], + vec![endpoint], + vec![upstream_key], + )); + let state = build_state_with_execution_runtime_override(execution_runtime_url) + .with_data_state_for_tests( + crate::data::GatewayDataState::with_minimal_candidate_selection_and_auth_for_tests( + candidate_repository, + auth_repository, + ) + .attach_provider_catalog_repository_for_tests(provider_catalog_repository) + .with_encryption_key_for_tests(DEVELOPMENT_ENCRYPTION_KEY), + ); + let gateway = build_router_with_state(state); + let (gateway_url, gateway_handle) = start_server(gateway).await; + + let sideband_url = format!( + "{}/v1/live/{CALL_ID}", + gateway_url.replacen("http://", "ws://", 1) + ); + let mut sideband_headers = HeaderMap::new(); + sideband_headers.insert( + http::header::AUTHORIZATION, + http::HeaderValue::from_static("Bearer sk-codex-live-sideband"), + ); + sideband_headers.insert( + http::HeaderName::from_static("x-session-id"), + http::HeaderValue::from_static("stable-live-sideband-session"), + ); + let missing_binding_response = wreq::Client::new() + .websocket(sideband_url.clone()) + .headers(sideband_headers.clone()) + .send() + .await + .expect("missing Live sideband binding should return an HTTP response"); + assert_eq!(missing_binding_response.status(), StatusCode::NOT_FOUND); + + let boundary = "aether-live-api-key-sideband"; + let offer_sdp = "v=0\r\no=api-key-client-offer"; + let session = json!({ + "model": CLIENT_MODEL, + "instructions": "Preserve this API-key Live session", + "future_session_capability": { + "revision": 13, + "nested": [true, {"mode": "opaque"}] + } + }); + let multipart_body = format!( + "--{boundary}\r\nContent-Disposition: form-data; name=\"sdp\"\r\nContent-Type: application/sdp\r\n\r\n{offer_sdp}\r\n--{boundary}\r\nContent-Disposition: form-data; name=\"session\"\r\nContent-Type: application/json\r\n\r\n{}\r\n--{boundary}--\r\n", + session + ); + let response = reqwest::Client::new() + .post(format!("{gateway_url}/v1/live")) + .header("authorization", "Bearer sk-codex-live-sideband") + .header( + "content-type", + format!("multipart/form-data; boundary={boundary}"), + ) + .header("x-session-id", "stable-live-sideband-session") + .header("openai-alpha", "client-value-must-be-replaced") + .body(multipart_body) + .send() + .await + .expect("Codex Live API-key call creation should complete"); + assert_eq!(response.status(), StatusCode::CREATED); + assert_eq!( + response + .headers() + .get(http::header::LOCATION) + .and_then(|value| value.to_str().ok()), + Some(format!("/v1/live/{CALL_ID}").as_str()) + ); + assert_eq!( + response + .headers() + .get("x-future-live-header") + .and_then(|value| value.to_str().ok()), + Some("preserved") + ); + assert_eq!( + response + .bytes() + .await + .expect("Live API-key SDP answer should read") + .as_ref(), + b"v=0\r\no=api-key-upstream-answer" + ); + + let plan = captured_plan + .lock() + .expect("Live API-key plan mutex should lock") + .clone() + .expect("Live API-key call should reach execution runtime"); + let plan_url = url::Url::parse(plan.url.as_str()).expect("Live API-key URL should parse"); + assert_eq!(plan_url.path(), "/v1/live"); + assert!(plan_url.query().is_none()); + assert_eq!(plan.method, "POST"); + assert!(!plan.stream); + assert!(plan + .content_type + .as_deref() + .is_some_and(|value| value.starts_with("multipart/form-data; boundary="))); + assert!(plan.body.json_body.is_none()); + let provider_multipart = base64::engine::general_purpose::STANDARD + .decode( + plan.body + .body_bytes_b64 + .as_deref() + .expect("API-key Live call should preserve multipart wire bytes"), + ) + .expect("provider multipart body should decode"); + let provider_multipart = + String::from_utf8(provider_multipart).expect("provider multipart should be UTF-8"); + assert!(provider_multipart.contains(offer_sdp)); + assert!(provider_multipart.contains(PROVIDER_MODEL)); + assert!(!provider_multipart.contains(CLIENT_MODEL)); + assert!(provider_multipart.contains("future_session_capability")); + assert_eq!( + plan.headers.get("authorization").map(String::as_str), + Some("Bearer oauth-upstream-secret") + ); + assert_eq!( + plan.headers.get("openai-alpha").map(String::as_str), + Some("quicksilver=v2") + ); + assert_eq!( + plan.headers.get("x-session-id").map(String::as_str), + Some("stable-live-sideband-session") + ); + assert_eq!( + plan.headers.get("accept").map(String::as_str), + Some("application/sdp") + ); + + let sideband_response = wreq::Client::new() + .websocket(sideband_url) + .headers(sideband_headers.clone()) + .send() + .await + .expect("Codex Live sideband handshake should complete"); + assert_eq!(sideband_response.status(), StatusCode::SWITCHING_PROTOCOLS); + let mut sideband = sideband_response + .into_websocket() + .await + .expect("Codex Live sideband response should upgrade"); + + // The client deliberately sends nothing before this receive. If sideband + // incorrectly reused the direct-WebSocket session.update bootstrap, this + // event could not arrive. + let ready_event = receive_codex_live_json(&mut sideband).await; + assert_eq!( + ready_event, + json!({ + "type": "future.sideband.ready", + "future_capability": {"opaque": true, "revision": 17} + }) + ); + let conflicting_response = wreq::Client::new() + .websocket(format!( + "{}/v1/live/{CALL_ID}", + gateway_url.replacen("http://", "ws://", 1) + )) + .headers(sideband_headers) + .send() + .await + .expect("duplicate Live sideband attachment should return an HTTP response"); + assert_eq!(conflicting_response.status(), StatusCode::CONFLICT); + + let opaque_command = json!({ + "type": "future.sideband.command", + "future_payload": {"without_session_update": true} + }); + sideband + .send(WreqWsMessage::text(opaque_command.to_string())) + .await + .expect("opaque sideband command should send without session.update"); + let sideband_session_update = json!({ + "type": "session.update", + "session": { + "model": "untrusted-client-model", + "future_session_field": {"opaque": true} + }, + "future_event_field": [1, 2, 3] + }); + sideband + .send(WreqWsMessage::text(sideband_session_update.to_string())) + .await + .expect("sideband session.update should send after an opaque frame"); + + let observed = tokio::time::timeout(std::time::Duration::from_secs(2), sideband_observed_rx) + .await + .expect("mock sideband should observe the opaque command before timeout") + .expect("mock sideband observation channel should remain open"); + assert_eq!(observed.request_target, format!("/v1/live/{CALL_ID}")); + assert_eq!( + observed.authorization.as_deref(), + Some("Bearer oauth-upstream-secret") + ); + assert_eq!(observed.alpha.as_deref(), Some("quicksilver=v2")); + assert_eq!( + observed.session_id.as_deref(), + Some("stable-live-sideband-session") + ); + assert_eq!(observed.first_client_event, opaque_command); + let mut expected_sideband_session_update = sideband_session_update; + expected_sideband_session_update["session"]["model"] = json!(PROVIDER_MODEL); + assert_eq!(observed.session_update, expected_sideband_session_update); + + drop(sideband); + gateway_handle.abort(); + execution_runtime_handle.abort(); + upstream_handle.abort(); +} + +async fn mock_codex_live_sideband_websocket( + State(observed): State>>>>, + uri: Uri, + headers: HeaderMap, + ws: WebSocketUpgrade, +) -> impl IntoResponse { + 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 alpha = headers + .get("openai-alpha") + .and_then(|value| value.to_str().ok()) + .map(str::to_string); + let session_id = headers + .get("x-session-id") + .and_then(|value| value.to_str().ok()) + .map(str::to_string); + ws.on_upgrade(move |mut socket| async move { + socket + .send(AxumWsMessage::Text( + json!({ + "type": "future.sideband.ready", + "future_capability": {"opaque": true, "revision": 17} + }) + .to_string() + .into(), + )) + .await + .expect("mock upstream sideband ready event should send"); + let first_client_event = receive_axum_live_json(&mut socket).await; + let session_update = receive_axum_live_json(&mut socket).await; + let observation = ObservedCodexLiveSideband { + request_target, + authorization, + alpha, + session_id, + first_client_event, + session_update, + }; + if let Some(sender) = observed + .lock() + .expect("mock sideband observation mutex should lock") + .take() + { + let _ = sender.send(observation); + } + }) +} + +async fn mock_codex_live_websocket( + State(observed): State>>>>, + uri: Uri, + headers: HeaderMap, + ws: WebSocketUpgrade, +) -> impl IntoResponse { + 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 alpha = headers + .get("openai-alpha") + .and_then(|value| value.to_str().ok()) + .map(str::to_string); + let session_id = headers + .get("x-session-id") + .and_then(|value| value.to_str().ok()) + .map(str::to_string); + ws.on_upgrade(move |mut socket| async move { + let initial_event = receive_axum_live_json(&mut socket).await; + socket + .send(AxumWsMessage::Text( + json!({ + "type": "future.live.event", + "future_capability": {"enabled": true, "revision": 11} + }) + .to_string() + .into(), + )) + .await + .expect("mock upstream future event should send"); + socket + .send(AxumWsMessage::Text( + json!({ + "type": "turn.done", + "turn": {"id": "turn-live-1"}, + "future_turn_field": "retained" + }) + .to_string() + .into(), + )) + .await + .expect("mock upstream turn.done should send"); + let event_after_turn_done = receive_axum_live_json(&mut socket).await; + let observation = ObservedCodexLiveWebSocket { + request_target, + authorization, + alpha, + session_id, + initial_event, + event_after_turn_done, + }; + if let Some(sender) = observed + .lock() + .expect("mock upstream observation mutex should lock") + .take() + { + let _ = sender.send(observation); + } + }) +} + +async fn receive_axum_live_json(socket: &mut WebSocket) -> serde_json::Value { + let message = tokio::time::timeout(std::time::Duration::from_secs(2), socket.recv()) + .await + .expect("mock upstream should receive a Live event before timeout") + .expect("mock upstream socket should remain open") + .expect("mock upstream Live frame should be readable"); + match message { + AxumWsMessage::Text(text) => { + serde_json::from_str(text.as_str()).expect("mock upstream Live event should be JSON") + } + other => panic!("mock upstream expected text Live event, got {other:?}"), + } +} + +async fn receive_codex_live_json(socket: &mut wreq::ws::WebSocket) -> serde_json::Value { + let message = tokio::time::timeout(std::time::Duration::from_secs(2), socket.recv()) + .await + .expect("Codex Live gateway should send an event before timeout") + .expect("Codex Live gateway socket should remain open") + .expect("Codex Live gateway frame should be readable"); + match message { + WreqWsMessage::Text(text) => serde_json::from_str(text.as_str()) + .expect("Codex Live gateway text event should be JSON"), + other => panic!("Codex Live gateway expected text event, got {other:?}"), + } +} + #[tokio::test] async fn gateway_openai_models_list_drops_disabled_global_model_after_cache_invalidation() { let auth_repository = Arc::new(InMemoryAuthApiKeySnapshotRepository::seed(vec![( diff --git a/crates/aether-ai/formats/src/formats/shared/routing.rs b/crates/aether-ai/formats/src/formats/shared/routing.rs index a7f03a4df..66b47dc8b 100644 --- a/crates/aether-ai/formats/src/formats/shared/routing.rs +++ b/crates/aether-ai/formats/src/formats/shared/routing.rs @@ -411,7 +411,19 @@ pub fn sanitize_request_path(path: &str) -> Option { .map(|(path, _)| path) .unwrap_or_else(|| path.trim()) .trim(); - (!path.is_empty()).then(|| path.to_string()) + if path.is_empty() { + return None; + } + if path + .strip_prefix("/v1/live/") + .is_some_and(|call_id| !call_id.is_empty()) + { + // A Live call id identifies an in-progress WebRTC session. Keep this + // bearer-like capability out of logs and persisted request metadata, + // including for malformed routes that will later be rejected. + return Some("/v1/live/{call_id}".to_string()); + } + Some(path.to_string()) } pub fn sanitize_request_query_string(query: &str) -> Option { @@ -440,12 +452,13 @@ pub fn sanitize_request_path_and_query(path: &str, query: Option<&str>) -> Optio return None; } + let sanitized_path = sanitize_request_path(path)?; let sanitized_query = query .and_then(sanitize_request_query_string) .or_else(|| embedded_query.and_then(sanitize_request_query_string)); Some(match sanitized_query { - Some(query) => format!("{path}?{query}"), - None => path.to_string(), + Some(query) => format!("{sanitized_path}?{query}"), + None => sanitized_path, }) } @@ -871,6 +884,14 @@ mod tests { .as_deref(), Some("/v1beta/models/gemini-2.5-pro:streamGenerateContent?alt=sse") ); + assert_eq!( + sanitize_request_path_and_query( + "/v1/live/rtc_secret_opaque?alt=sse&token=hidden", + None + ) + .as_deref(), + Some("/v1/live/{call_id}?alt=sse") + ); } #[test] diff --git a/crates/aether-gateway/frontdoor/src/middleware/access_log.rs b/crates/aether-gateway/frontdoor/src/middleware/access_log.rs index faed5531d..e7c7d0539 100644 --- a/crates/aether-gateway/frontdoor/src/middleware/access_log.rs +++ b/crates/aether-gateway/frontdoor/src/middleware/access_log.rs @@ -244,6 +244,10 @@ mod tests { ), "/v1beta/models/gemini-3-flash-preview:generateContent?alt=sse&pageSize=10" ); + assert_eq!( + sanitize_access_log_path("/v1/live/rtc_secret_opaque?token=hidden"), + "/v1/live/{call_id}" + ); } #[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 b1b6acee8..16cac5549 100644 --- a/crates/aether-provider/transport/src/codex_fingerprint.rs +++ b/crates/aether-provider/transport/src/codex_fingerprint.rs @@ -138,6 +138,10 @@ fn apply_converged_headers( set_header(headers, "session-id", fingerprint.session_id.clone()); set_header(headers, "session_id", fingerprint.session_id.clone()); set_header(headers, "thread-id", fingerprint.thread_id.clone()); + // Codex Live/Realtime uses `x-session-id` for the thread-scoped session + // identity on the WebSocket upgrade request. Keep it aligned with the + // converged thread identity instead of the account-scoped session value. + set_header(headers, "x-session-id", fingerprint.thread_id.clone()); rewrite_header_turn_metadata(headers, fingerprint); } @@ -356,6 +360,10 @@ mod tests { let transport = sample_transport(); let mut headers = BTreeMap::from([ ("Session-Id".to_string(), "client-session".to_string()), + ( + "X-Session-Id".to_string(), + "client-live-session".to_string(), + ), ( "x-codex-turn-metadata".to_string(), json!({ @@ -410,7 +418,15 @@ mod tests { ); assert_eq!(headers["session_id"], *session_id); assert_eq!(headers["x-client-request-id"], *thread_id); + assert_eq!(headers["x-session-id"], *thread_id); assert_eq!(headers["x-codex-window-id"], format!("{thread_id}:0")); + assert_eq!( + headers + .keys() + .filter(|name| name.eq_ignore_ascii_case("x-session-id")) + .count(), + 1 + ); assert_eq!(body["client_metadata"]["session_id"], *session_id); assert_eq!(body["client_metadata"]["thread_id"], *thread_id); assert_eq!( diff --git a/docs/WebSocket-Mode.md b/docs/WebSocket-Mode.md index 844d54cf7..7c7bd9163 100644 --- a/docs/WebSocket-Mode.md +++ b/docs/WebSocket-Mode.md @@ -4,6 +4,60 @@ The Responses API supports a WebSocket mode for long-running, tool-call-heavy wo WebSocket mode is compatible with both Zero Data Retention (ZDR) and `store=false`. +## Experimental Codex Live bridge + +Aether also exposes the Codex Frameless Bidi V3 transport used by current +Codex clients. It is related to the OpenAI Realtime API, but it is not the +Responses WebSocket protocol and never enters Aether's `response.create` +state machine: + +- Direct WebSocket: `GET /v1/live?model=`. The first client text + frame must be `session.update`; later text, binary, ping, pong, and close + frames are relayed opaquely. +- WebRTC call creation: `POST /v1/live` with bounded `sdp` and `session` + multipart parts. Aether applies the existing global-to-provider model + mapping and rewrites the upstream `Location` to `/v1/live/`. +- WebRTC sideband: `GET /v1/live/`. Frameless sideband attaches to an + 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 + } +} +``` + +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 +origin for its sideband; direct OAuth WebSocket and custom OAuth backend +origins fail closed. The call binding fixes the authenticated downstream +principal, provider/endpoint/key, mapped model, auth mode, account/FedRAMP +identity, session identity, and upstream origin. Raw call IDs are hashed in +RuntimeState keys, records expire after two hours, each principal retains at +most 64 call bindings, and one call permits only one renewable sideband +attachment at a time. The memory RuntimeState backend loses these bindings on +restart. The two-hour binding TTL and 64-record cap bound routing state and +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. + +For the public GA Realtime API's connection and session concepts, see the +[OpenAI Realtime guide](https://developers.openai.com/api/docs/guides/realtime). + OpenAI's current WebSocket service supports named `stream_id` lanes: requests on the same lane are FIFO, while different lanes may run concurrently. Aether's bridge currently exposes only the implicit default lane and deliberately