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 871500eb4..6c4490977 100644 --- a/apps/aether-gateway/src/ai_serving/planner/decision_input.rs +++ b/apps/aether-gateway/src/ai_serving/planner/decision_input.rs @@ -101,6 +101,22 @@ pub(crate) fn apply_provider_request_routing_policy_to_decision( input: &LocalRequestedModelDecisionInput, decision: &mut AiExecutionDecision, transport: Option<&GatewayProviderTransportSnapshot>, +) -> Result<(), GatewayError> { + apply_provider_request_routing_policy_to_decision_with_websocket_mode( + input, decision, transport, false, + ) +} + +/// Applies provider-request routing mutations while retaining the transport +/// boundary of a pinned Responses WebSocket continuation. Routing rules may +/// mutate the body and therefore require a second provider-contract pass; the +/// pass must use the same explicit continuation mode as the first pass rather +/// than guessing from JSON fields. +pub(crate) fn apply_provider_request_routing_policy_to_decision_with_websocket_mode( + input: &LocalRequestedModelDecisionInput, + decision: &mut AiExecutionDecision, + transport: Option<&GatewayProviderTransportSnapshot>, + websocket_continuation: bool, ) -> Result<(), GatewayError> { let provider_api_format = decision .provider_api_format @@ -257,9 +273,8 @@ pub(crate) fn apply_provider_request_routing_policy_to_decision( input.requested_model.as_str(), ) }); - crate::ai_serving::finalize_openai_provider_request_with_codex_model_capabilities_and_reasoning_replay_policy( - &mut provider_request_body, - crate::ai_serving::OpenAiProviderRequestFinalization { + { + let finalization = crate::ai_serving::OpenAiProviderRequestFinalization { source_api_format: context.client_api_format.as_str(), provider_api_format: provider_api_format.as_str(), provider_type: provider_type.as_str(), @@ -270,17 +285,31 @@ pub(crate) fn apply_provider_request_routing_policy_to_decision( require_body_stream_field: original_provider_request_body .as_ref() .is_some_and(|body| body.get("stream").is_some()), - }, - model_capabilities.as_ref(), - transport + }; + let reasoning_replay_policy = transport .map(|transport| { crate::ai_serving::openai_responses_reasoning_replay_policy( transport.provider.provider_type.as_str(), transport.endpoint.base_url.as_str(), ) }) - .unwrap_or_default(), - ) + .unwrap_or_default(); + if websocket_continuation { + crate::ai_serving::finalize_openai_provider_request_with_codex_model_capabilities_and_reasoning_replay_policy_for_websocket_continuation( + &mut provider_request_body, + finalization, + model_capabilities.as_ref(), + reasoning_replay_policy, + ) + } else { + crate::ai_serving::finalize_openai_provider_request_with_codex_model_capabilities_and_reasoning_replay_policy( + &mut provider_request_body, + finalization, + model_capabilities.as_ref(), + reasoning_replay_policy, + ) + } + } .map_err(|violation| GatewayError::Client { status: StatusCode::BAD_REQUEST, message: format!("routing provider_request violates provider contract: {violation:?}"), diff --git a/apps/aether-gateway/src/ai_serving/planner/passthrough/provider/family/request.rs b/apps/aether-gateway/src/ai_serving/planner/passthrough/provider/family/request.rs index 921690452..fbedcef7b 100644 --- a/apps/aether-gateway/src/ai_serving/planner/passthrough/provider/family/request.rs +++ b/apps/aether-gateway/src/ai_serving/planner/passthrough/provider/family/request.rs @@ -180,12 +180,17 @@ pub(crate) async fn resolve_local_same_format_provider_candidate_payload_parts( } }; let effective_headers = input.effective_headers(&parts.headers); + let reasoning_replay_policy = openai_responses_reasoning_replay_policy( + prepared.transport.provider.provider_type.as_str(), + prepared.transport.endpoint.base_url.as_str(), + ); let redaction = resolve_provider_chat_pii_redaction( state, parts, body_json, &input.auth_context, spec.api_format, + reasoning_replay_policy, &attempt.candidate_id, ) .await?; @@ -205,10 +210,7 @@ pub(crate) async fn resolve_local_same_format_provider_candidate_payload_parts( prepared.kiro_auth.as_ref(), prepared.is_claude_code, false, - openai_responses_reasoning_replay_policy( - prepared.transport.provider.provider_type.as_str(), - prepared.transport.endpoint.base_url.as_str(), - ), + reasoning_replay_policy, ) else { mark_skipped_local_same_format_provider_candidate_with_extra_data( @@ -289,10 +291,7 @@ pub(crate) async fn resolve_local_same_format_provider_candidate_payload_parts( ), }, codex_model_capabilities.as_ref(), - openai_responses_reasoning_replay_policy( - transport.provider.provider_type.as_str(), - transport.endpoint.base_url.as_str(), - ), + reasoning_replay_policy, ) { mark_skipped_local_same_format_provider_candidate_with_extra_data( diff --git a/apps/aether-gateway/src/ai_serving/planner/redaction.rs b/apps/aether-gateway/src/ai_serving/planner/redaction.rs index 4b4814d21..b66ef9ef8 100644 --- a/apps/aether-gateway/src/ai_serving/planner/redaction.rs +++ b/apps/aether-gateway/src/ai_serving/planner/redaction.rs @@ -67,6 +67,7 @@ pub(crate) async fn resolve_provider_chat_pii_redaction<'a>( body_json: &'a Value, auth_context: &ExecutionRuntimeAuthContext, client_api_format: &str, + reasoning_replay_policy: crate::ai_serving::OpenAiResponsesReasoningReplayPolicy, candidate_id: &str, ) -> Result, GatewayError> { let Some(format) = ChatPiiRedactionRequestFormat::from_api_format(client_api_format) else { @@ -75,7 +76,7 @@ pub(crate) async fn resolve_provider_chat_pii_redaction<'a>( let Some(slot) = parts.extensions.get::() else { return Ok(ProviderRequestRedaction::disabled(body_json)); }; - let request_cache_key = request_redaction_cache_key(format, body_json); + let request_cache_key = request_redaction_cache_key(format, reasoning_replay_policy, body_json); if let Some(cached) = slot.cached_request_redaction(&request_cache_key) { crate::stage_metrics::record_chat_pii_redaction_request_cache_hit(); observe_gateway_stage_ms("chat_pii_redaction_request_cache_hit", 0); @@ -132,7 +133,7 @@ pub(crate) async fn resolve_provider_chat_pii_redaction<'a>( body_json, format, build_redaction_session_config(hmac_key, &runtime_config, now_unix_secs), - MaskChatRequestOptions::runtime(), + MaskChatRequestOptions::runtime().with_reasoning_replay_policy(reasoning_replay_policy), Some(&cache), ) .await @@ -165,8 +166,16 @@ pub(crate) async fn resolve_provider_chat_pii_redaction<'a>( }) } -fn request_redaction_cache_key(format: ChatPiiRedactionRequestFormat, body_json: &Value) -> String { - format!("{format:?}:{:p}", body_json) +fn request_redaction_cache_key( + format: ChatPiiRedactionRequestFormat, + reasoning_replay_policy: crate::ai_serving::OpenAiResponsesReasoningReplayPolicy, + body_json: &Value, +) -> String { + // A request may be attempted against providers with different replay + // contracts. Reusing a cached DeepSeek opaque decision for an ordinary + // Responses candidate (or vice versa) would either bypass masking or + // corrupt provider-owned continuation state. + format!("{format:?}:{reasoning_replay_policy:?}:{:p}", body_json) } fn provider_redaction_from_cached<'a>( @@ -216,7 +225,15 @@ async fn resolve_chat_pii_redaction_feature_settings( } fn redaction_mask_error_to_gateway_error(error: RedactionMaskError) -> GatewayError { - match error {} + match error { + RedactionMaskError::SensitiveOpaqueReasoningState => { + warn!("gateway rejected provider-bound reasoning state containing sensitive text"); + GatewayError::Client { + status: http::StatusCode::BAD_REQUEST, + message: "provider-bound reasoning state contains sensitive text and cannot be safely replayed while chat PII redaction is enabled".to_string(), + } + } + } } #[cfg(test)] diff --git a/apps/aether-gateway/src/ai_serving/planner/standard/deepseek.rs b/apps/aether-gateway/src/ai_serving/planner/standard/deepseek.rs index 119f6aac3..14389089a 100644 --- a/apps/aether-gateway/src/ai_serving/planner/standard/deepseek.rs +++ b/apps/aether-gateway/src/ai_serving/planner/standard/deepseek.rs @@ -9,7 +9,9 @@ pub(crate) fn is_deepseek_provider(provider_type: &str, base_url: &str) -> bool return true; } - let host = base_url_host(base_url); + let Some(host) = base_url_host(base_url) else { + return false; + }; host == "deepseek.com" || host.ends_with(".deepseek.com") } @@ -47,24 +49,33 @@ pub(crate) fn apply_deepseek_tool_call_thinking_compat( } } -fn base_url_host(base_url: &str) -> String { - let lower = base_url.trim().to_ascii_lowercase(); - let without_scheme = lower - .split_once("://") - .map(|(_, rest)| rest) - .unwrap_or(lower.as_str()); - let without_userinfo = without_scheme - .rsplit_once('@') - .map(|(_, host)| host) - .unwrap_or(without_scheme); - without_userinfo - .split(['/', '?', '#']) - .next() - .unwrap_or_default() - .split(':') - .next() - .unwrap_or_default() - .to_string() +fn base_url_host(base_url: &str) -> Option { + let base_url = base_url.trim(); + if base_url.is_empty() { + return None; + } + + // Provider configuration historically accepted both absolute URLs and a + // bare authority/path. Use a real URL parser for both forms: hand-parsing + // userinfo with `rsplit_once('@')` can mistake an `@` in the path or query + // for the authority delimiter and classify an attacker-controlled host as + // `api.deepseek.com`. + if let Ok(parsed) = url::Url::parse(base_url) { + if let Some(host) = parsed + .host_str() + .filter(|_| matches!(parsed.scheme(), "http" | "https" | "ws" | "wss")) + { + return Some(host.to_ascii_lowercase()); + } + if base_url.contains("://") { + return None; + } + } + + url::Url::parse(&format!("https://{base_url}")) + .ok()? + .host_str() + .map(str::to_ascii_lowercase) } fn source_disables_thinking( @@ -267,10 +278,29 @@ mod tests { "custom", "https://api.deepseek.com/v1" )); + assert!(is_deepseek_provider("custom", "api.deepseek.com/v1")); + assert!(is_deepseek_provider("custom", "api.deepseek.com:443/v1")); assert!(!is_deepseek_provider( "custom", "https://example.com/deepseek" )); + assert!(!is_deepseek_provider( + "custom", + "https://api.deepseek.com.evil.example/v1" + )); + assert!(!is_deepseek_provider( + "custom", + "https://api.deepseek.com@evil.example/v1" + )); + assert!(!is_deepseek_provider( + "custom", + "https://evil.example/path@api.deepseek.com/v1" + )); + assert!(!is_deepseek_provider( + "custom", + "https://evil.example/?relay=@api.deepseek.com" + )); + assert!(!is_deepseek_provider("custom", "ftp://api.deepseek.com/v1")); assert_eq!( openai_responses_reasoning_replay_policy("custom", "https://api.deepseek.com/v1"), crate::ai_serving::OpenAiResponsesReasoningReplayPolicy::DeepSeekOpaque @@ -295,11 +325,52 @@ mod tests { }) }) .collect::>(); + let request = json!({ + "model": "deepseek-v4-flash", + "input": reasoning_items.clone(), + "future_request_field": {"preserve": true} + }); + let replay_policy = + openai_responses_reasoning_replay_policy("custom", "https://api.deepseek.com/v1"); + let mut provider_body = crate::ai_serving::build_standard_request_body_with_model_directives_and_request_headers_and_reasoning_replay_policy( + &request, + "openai:responses", + "deepseek-v4-flash", + "custom", + "openai:responses", + "/v1/responses", + false, + None, + None, + None, + false, + replay_policy, + ) + .expect("custom DeepSeek Responses body should build"); + crate::ai_serving::finalize_openai_provider_request_with_codex_model_capabilities_and_reasoning_replay_policy( + &mut provider_body, + crate::ai_serving::OpenAiProviderRequestFinalization { + source_api_format: "openai:responses", + provider_api_format: "openai:responses", + provider_type: "custom", + provider_model: "deepseek-v4-flash", + source_model: "deepseek-v4-flash", + body_rules: None, + upstream_is_stream: false, + require_body_stream_field: false, + }, + None, + replay_policy, + ) + .expect("custom DeepSeek finalization should accept opaque reasoning replay"); + assert_eq!(provider_body["input"].as_array().map(Vec::len), Some(66)); + assert_eq!(provider_body["future_request_field"]["preserve"], true); + let mut deepseek = json!({"input": reasoning_items.clone()}); let mut openai = json!({"input": reasoning_items}); assert_eq!( - aether_ai_formats::strip_incompatible_openai_responses_reasoning_items_with_policy( + crate::ai_serving::strip_incompatible_openai_responses_reasoning_items_with_policy( &mut deepseek, "openai:responses", openai_responses_reasoning_replay_policy("custom", "https://api.deepseek.com/v1"), @@ -309,7 +380,7 @@ mod tests { assert_eq!(deepseek["input"].as_array().map(Vec::len), Some(66)); assert_eq!( - aether_ai_formats::strip_incompatible_openai_responses_reasoning_items_with_policy( + crate::ai_serving::strip_incompatible_openai_responses_reasoning_items_with_policy( &mut openai, "openai:responses", openai_responses_reasoning_replay_policy("openai", "https://api.openai.com/v1"), diff --git a/apps/aether-gateway/src/ai_serving/planner/standard/family/request.rs b/apps/aether-gateway/src/ai_serving/planner/standard/family/request.rs index 37fcc9d99..7b3f128c9 100644 --- a/apps/aether-gateway/src/ai_serving/planner/standard/family/request.rs +++ b/apps/aether-gateway/src/ai_serving/planner/standard/family/request.rs @@ -364,6 +364,7 @@ pub(crate) async fn resolve_local_standard_candidate_payload_parts( body_json, &input.auth_context, spec_metadata.api_format, + crate::ai_serving::OpenAiResponsesReasoningReplayPolicy::OpenAiItemIds, &attempt.candidate_id, ) .await?; @@ -593,18 +594,23 @@ pub(crate) async fn resolve_local_standard_candidate_payload_parts( input.auth_context.api_key_id.as_str(), ) .await?; + let reasoning_replay_policy = openai_responses_reasoning_replay_policy( + transport.provider.provider_type.as_str(), + transport.endpoint.base_url.as_str(), + ); let redaction = resolve_provider_chat_pii_redaction( state, parts, body_json, &input.auth_context, spec_metadata.api_format, + reasoning_replay_policy, &attempt.candidate_id, ) .await?; let body_json = redaction.body_json.as_ref(); let mut provider_request_body = - match crate::ai_serving::planner::standard::build_standard_request_body_with_model_directives_and_request_headers( + match crate::ai_serving::planner::standard::build_standard_request_body_with_model_directives_and_request_headers_and_reasoning_replay_policy( body_json, spec_metadata.api_format, &prepared_candidate.mapped_model, @@ -620,10 +626,7 @@ pub(crate) async fn resolve_local_standard_candidate_payload_parts( Some(input.auth_context.api_key_id.as_str()), Some(effective_headers), false, - openai_responses_reasoning_replay_policy( - transport.provider.provider_type.as_str(), - transport.endpoint.base_url.as_str(), - ), + reasoning_replay_policy, ) { Some(body) => body, None => { @@ -770,10 +773,7 @@ pub(crate) async fn resolve_local_standard_candidate_payload_parts( ), }, codex_model_capabilities.as_ref(), - openai_responses_reasoning_replay_policy( - transport.provider.provider_type.as_str(), - transport.endpoint.base_url.as_str(), - ), + reasoning_replay_policy, ) { mark_skipped_local_standard_candidate_with_extra_data( diff --git a/apps/aether-gateway/src/ai_serving/planner/standard/mod.rs b/apps/aether-gateway/src/ai_serving/planner/standard/mod.rs index f8c3f57bc..732345602 100644 --- a/apps/aether-gateway/src/ai_serving/planner/standard/mod.rs +++ b/apps/aether-gateway/src/ai_serving/planner/standard/mod.rs @@ -33,6 +33,7 @@ pub(crate) use self::normalize::{ build_cross_format_openai_responses_upstream_url, build_local_openai_chat_request_body, build_local_openai_chat_upstream_url, build_local_openai_responses_request_body, build_local_openai_responses_request_body_with_codex_model_capabilities, + build_local_openai_responses_request_body_with_codex_model_capabilities_for_websocket_continuation, build_local_openai_responses_upstream_url, validate_final_openai_provider_request, }; pub(crate) use self::openai::{ @@ -64,7 +65,7 @@ pub(crate) use crate::ai_serving::{ }; pub(crate) use crate::ai_serving::{ build_standard_request_body, build_standard_request_body_with_model_directives, - build_standard_request_body_with_model_directives_and_request_headers, + build_standard_request_body_with_model_directives_and_request_headers_and_reasoning_replay_policy, convert_openai_chat_request_to_claude_request, convert_openai_chat_request_to_gemini_request, convert_openai_chat_request_to_openai_responses_request, extract_openai_text_content, normalize_openai_responses_request_to_openai_chat_request, parse_openai_tool_result_content, diff --git a/apps/aether-gateway/src/ai_serving/planner/standard/normalize.rs b/apps/aether-gateway/src/ai_serving/planner/standard/normalize.rs index 898497435..11c44ecea 100644 --- a/apps/aether-gateway/src/ai_serving/planner/standard/normalize.rs +++ b/apps/aether-gateway/src/ai_serving/planner/standard/normalize.rs @@ -15,6 +15,7 @@ pub(crate) use self::responses::{ build_cross_format_openai_responses_request_body_with_codex_model_capabilities, build_cross_format_openai_responses_upstream_url, build_local_openai_responses_request_body, build_local_openai_responses_request_body_with_codex_model_capabilities, + build_local_openai_responses_request_body_with_codex_model_capabilities_for_websocket_continuation, build_local_openai_responses_upstream_url, }; pub(super) use crate::ai_serving::planner::common::{ diff --git a/apps/aether-gateway/src/ai_serving/planner/standard/normalize/responses.rs b/apps/aether-gateway/src/ai_serving/planner/standard/normalize/responses.rs index 6ac9b14cc..47c91f1b9 100644 --- a/apps/aether-gateway/src/ai_serving/planner/standard/normalize/responses.rs +++ b/apps/aether-gateway/src/ai_serving/planner/standard/normalize/responses.rs @@ -50,6 +50,69 @@ pub(crate) fn build_local_openai_responses_request_body_with_codex_model_capabil request_headers: &http::HeaderMap, model_capabilities: Option<&crate::ai_serving::CodexResponsesModelCapabilities>, enable_model_directives: bool, +) -> Option { + build_local_openai_responses_request_body_with_codex_model_capabilities_and_websocket_mode( + body_json, + mapped_model, + require_streaming, + force_body_stream_field, + provider_type, + provider_api_format, + body_rules, + request_headers, + model_capabilities, + enable_model_directives, + false, + ) +} + +/// Builds a Responses body for a pinned WebSocket continuation. +/// +/// This is intentionally an additive variant of the ordinary HTTP builder. +/// The WebSocket framing layer, rather than a JSON-body heuristic, tells the +/// Codex compatibility pass that `previous_response_id` is transport state and +/// that Responses Lite `tools`/`instructions` must not be materialized into a +/// second historical input prefix. +pub(crate) fn build_local_openai_responses_request_body_with_codex_model_capabilities_for_websocket_continuation( + body_json: &Value, + mapped_model: &str, + require_streaming: bool, + force_body_stream_field: bool, + provider_type: &str, + provider_api_format: &str, + body_rules: Option<&Value>, + request_headers: &http::HeaderMap, + model_capabilities: Option<&crate::ai_serving::CodexResponsesModelCapabilities>, + enable_model_directives: bool, +) -> Option { + build_local_openai_responses_request_body_with_codex_model_capabilities_and_websocket_mode( + body_json, + mapped_model, + require_streaming, + force_body_stream_field, + provider_type, + provider_api_format, + body_rules, + request_headers, + model_capabilities, + enable_model_directives, + true, + ) +} + +#[allow(clippy::too_many_arguments)] +fn build_local_openai_responses_request_body_with_codex_model_capabilities_and_websocket_mode( + body_json: &Value, + mapped_model: &str, + require_streaming: bool, + force_body_stream_field: bool, + provider_type: &str, + provider_api_format: &str, + body_rules: Option<&Value>, + request_headers: &http::HeaderMap, + model_capabilities: Option<&crate::ai_serving::CodexResponsesModelCapabilities>, + enable_model_directives: bool, + websocket_continuation: bool, ) -> Option { let provider_request_body = surface_build_local_openai_responses_request_body( body_json, @@ -68,15 +131,27 @@ pub(crate) fn build_local_openai_responses_request_body_with_codex_model_capabil .get("model") .and_then(Value::as_str) .unwrap_or(mapped_model); - crate::ai_serving::apply_codex_openai_responses_special_body_edits_with_source_model_and_capabilities( - &mut provider_request_body, - provider_type, - provider_api_format, - mapped_model, - source_model, - model_capabilities, - body_rules, - ); + if websocket_continuation { + crate::ai_serving::apply_codex_openai_responses_websocket_continuation_body_edits_with_source_model_and_capabilities( + &mut provider_request_body, + provider_type, + provider_api_format, + mapped_model, + source_model, + model_capabilities, + body_rules, + ); + } else { + crate::ai_serving::apply_codex_openai_responses_special_body_edits_with_source_model_and_capabilities( + &mut provider_request_body, + provider_type, + provider_api_format, + mapped_model, + source_model, + model_capabilities, + body_rules, + ); + } apply_openai_responses_compact_special_body_edits( &mut provider_request_body, provider_api_format, diff --git a/apps/aether-gateway/src/ai_serving/planner/standard/normalize/tests.rs b/apps/aether-gateway/src/ai_serving/planner/standard/normalize/tests.rs index 2529b107a..b3890dae0 100644 --- a/apps/aether-gateway/src/ai_serving/planner/standard/normalize/tests.rs +++ b/apps/aether-gateway/src/ai_serving/planner/standard/normalize/tests.rs @@ -153,7 +153,7 @@ fn local_openai_responses_wrapper_preserves_body_order_after_edits() { } #[test] -fn local_openai_responses_wrapper_strips_foreign_reasoning_item_ids() { +fn local_openai_responses_wrapper_defers_reasoning_replay_filtering() { let body_json = json!({ "model": "gpt-5.4", "input": [ @@ -184,9 +184,14 @@ fn local_openai_responses_wrapper_strips_foreign_reasoning_item_ids() { let input = provider_request_body["input"] .as_array() .expect("input array"); - assert_eq!(input.len(), 2); + // This provider-agnostic normalization layer cannot decide whether an + // id-less/foreign reasoning item is opaque state required by DeepSeek. + // The provider-aware finalization pass applies the strict or DeepSeek + // replay policy once the selected upstream base URL is known. + assert_eq!(input.len(), 3); assert_eq!(input[0]["id"], "rs_provider_123"); - assert_eq!(input[1]["type"], "message"); + assert_eq!(input[1]["id"], "item_72d3bd8d367d01977ace23f1"); + assert_eq!(input[2]["type"], "message"); } #[test] diff --git a/apps/aether-gateway/src/ai_serving/planner/standard/openai/chat/decision/request.rs b/apps/aether-gateway/src/ai_serving/planner/standard/openai/chat/decision/request.rs index e6c4dc14c..599a1e763 100644 --- a/apps/aether-gateway/src/ai_serving/planner/standard/openai/chat/decision/request.rs +++ b/apps/aether-gateway/src/ai_serving/planner/standard/openai/chat/decision/request.rs @@ -29,8 +29,8 @@ use crate::ai_serving::planner::standard::{ apply_deepseek_tool_call_thinking_compat, build_cross_format_openai_chat_request_body, build_cross_format_openai_chat_upstream_url, build_local_openai_chat_request_body, build_local_openai_chat_upstream_url, codex_model_capabilities_for_transport, - openai_provider_request_contract_failure_extra_data, request_body_build_failure_extra_data, - request_conversion_failure_extra_data, + openai_provider_request_contract_failure_extra_data, openai_responses_reasoning_replay_policy, + request_body_build_failure_extra_data, request_conversion_failure_extra_data, }; use crate::ai_serving::transport::antigravity::is_antigravity_provider_transport; use crate::ai_serving::transport::auth::resolve_local_openai_bearer_auth; @@ -140,7 +140,7 @@ fn finalize_openai_chat_provider_request_body( mapped_model, source_model, ); - crate::ai_serving::finalize_openai_provider_request_with_codex_model_capabilities( + crate::ai_serving::finalize_openai_provider_request_with_codex_model_capabilities_and_reasoning_replay_policy( provider_request_body, crate::ai_serving::OpenAiProviderRequestFinalization { source_api_format: "openai:chat", @@ -156,6 +156,10 @@ fn finalize_openai_chat_provider_request_body( ), }, codex_model_capabilities.as_ref(), + openai_responses_reasoning_replay_policy( + transport.provider.provider_type.as_str(), + transport.endpoint.base_url.as_str(), + ), ) .err() .map(|violation| { @@ -206,6 +210,7 @@ pub(crate) async fn resolve_local_openai_chat_candidate_payload_parts( body_json, &input.auth_context, "openai:chat", + crate::ai_serving::OpenAiResponsesReasoningReplayPolicy::OpenAiItemIds, candidate_id, ) .await?; @@ -2347,6 +2352,23 @@ mod tests { eligible } + fn sample_custom_deepseek_responses_transport() -> GatewayProviderTransportSnapshot { + let mut transport = sample_gemini_cli_transport(); + transport.provider.name = "deepseek".to_string(); + transport.provider.provider_type = "custom".to_string(); + transport.endpoint.api_format = "openai:responses".to_string(); + transport.endpoint.api_family = Some("openai".to_string()); + transport.endpoint.endpoint_kind = Some("responses".to_string()); + transport.endpoint.base_url = "https://api.deepseek.com/v1".to_string(); + transport.endpoint.custom_path = None; + transport.key.api_formats = Some(vec!["openai:responses".to_string()]); + transport.key.auth_type = "bearer".to_string(); + transport.key.decrypted_api_key = "test-api-key".to_string(); + transport.key.decrypted_auth_config = None; + transport.key.upstream_metadata = None; + transport + } + fn sample_custom_directive_input() -> LocalOpenAiChatDecisionInput { let mut input = sample_input(); input.requested_model = "gpt-5.6-sol-high".to_string(); @@ -2372,6 +2394,58 @@ mod tests { input } + #[test] + fn responses_shaped_chat_request_preserves_deepseek_opaque_reasoning_at_finalization() { + let reasoning_items = (0..66) + .map(|index| { + json!({ + "type": "reasoning", + "encrypted_content": format!("opaque-deepseek-state-{index}"), + "content": [{ + "type": "reasoning_text", + "text": format!("provider thinking state {index}") + }], + "future_capability": {"preserve": true} + }) + }) + .collect::>(); + let original_body = json!({ + "model": "deepseek-v4-flash", + "input": reasoning_items + }); + let transport = sample_custom_deepseek_responses_transport(); + let mut provider_body = build_cross_format_openai_chat_request_body( + &original_body, + "deepseek-v4-flash", + "custom", + "openai:responses", + false, + false, + None, + None, + &http::HeaderMap::new(), + false, + ) + .expect("Responses-shaped chat request should build"); + + assert!(finalize_openai_chat_provider_request_body( + &mut provider_body, + None, + "openai:responses", + false, + false, + &original_body, + &transport, + "deepseek-v4-flash", + ) + .is_none()); + let input = provider_body["input"].as_array().expect("provider input"); + assert_eq!(input.len(), 66); + assert_eq!(input[0]["type"], "reasoning"); + assert_eq!(input[0]["content"][0]["type"], "reasoning_text"); + assert_eq!(input[0]["future_capability"]["preserve"], true); + } + fn sample_alias_max_directive_input() -> LocalOpenAiChatDecisionInput { let mut input = sample_input(); input.requested_model = "deployment-alias-max".to_string(); diff --git a/apps/aether-gateway/src/ai_serving/planner/standard/openai/responses/decision.rs b/apps/aether-gateway/src/ai_serving/planner/standard/openai/responses/decision.rs index e3c4a1524..077d95d25 100644 --- a/apps/aether-gateway/src/ai_serving/planner/standard/openai/responses/decision.rs +++ b/apps/aether-gateway/src/ai_serving/planner/standard/openai/responses/decision.rs @@ -5,7 +5,10 @@ mod request; #[path = "decision/support.rs"] mod support; -pub(super) use self::payload::maybe_build_local_openai_responses_decision_payload_for_candidate; +pub(super) use self::payload::{ + maybe_build_local_openai_responses_decision_payload_for_candidate, + maybe_build_local_openai_responses_decision_payload_for_candidate_with_websocket_mode, +}; pub(super) use self::support::{ build_local_openai_responses_candidate_attempt_source, materialize_local_openai_responses_candidate_attempts, 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 ae0f4f4c3..6c953611f 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 @@ -2,7 +2,7 @@ use serde_json::json; 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; +use crate::ai_serving::planner::decision_input::apply_provider_request_routing_policy_to_decision_with_websocket_mode; use crate::ai_serving::planner::report_context::{ build_local_execution_report_context, insert_native_client_envelope_name, insert_provider_stream_event_api_format, LocalExecutionReportContextParts, @@ -20,7 +20,10 @@ use crate::{ AiExecutionDecision, AppState, GatewayError, }; -use super::request::resolve_local_openai_responses_candidate_payload_parts; +use super::request::{ + resolve_local_openai_responses_candidate_payload_parts, + resolve_local_openai_responses_candidate_payload_parts_with_websocket_mode, +}; use super::support::{LocalOpenAiResponsesCandidateAttempt, LocalOpenAiResponsesDecisionInput}; use super::LocalOpenAiResponsesSpec; @@ -32,6 +35,26 @@ pub(crate) async fn maybe_build_local_openai_responses_decision_payload_for_cand input: &LocalOpenAiResponsesDecisionInput, attempt: LocalOpenAiResponsesCandidateAttempt, spec: LocalOpenAiResponsesSpec, +) -> Result, GatewayError> { + maybe_build_local_openai_responses_decision_payload_for_candidate_with_websocket_mode( + state, parts, trace_id, body_json, input, attempt, spec, false, + ) + .await +} + +/// Builds a candidate payload for a pinned WebSocket turn without changing +/// the ordinary HTTP/plan-builder path. The explicit mode is carried all the +/// way to body normalization because a JSON `type` field is not a reliable +/// transport discriminator once body rules and conversions have run. +pub(crate) async fn maybe_build_local_openai_responses_decision_payload_for_candidate_with_websocket_mode( + state: &AppState, + parts: &http::request::Parts, + trace_id: &str, + body_json: &serde_json::Value, + input: &LocalOpenAiResponsesDecisionInput, + attempt: LocalOpenAiResponsesCandidateAttempt, + spec: LocalOpenAiResponsesSpec, + websocket_continuation: bool, ) -> Result, GatewayError> { let spec_metadata = local_openai_responses_spec_metadata(spec); let attempt_identity = attempt.attempt_identity(); @@ -41,19 +64,35 @@ pub(crate) async fn maybe_build_local_openai_responses_decision_payload_for_cand candidate_id, .. } = attempt; - let Some(resolved) = resolve_local_openai_responses_candidate_payload_parts( - state, - parts, - trace_id, - body_json, - input, - &eligible, - candidate_index, - &candidate_id, - spec, - ) - .await? - else { + let resolved = if websocket_continuation { + resolve_local_openai_responses_candidate_payload_parts_with_websocket_mode( + state, + parts, + trace_id, + body_json, + input, + &eligible, + candidate_index, + &candidate_id, + spec, + true, + ) + .await? + } else { + resolve_local_openai_responses_candidate_payload_parts( + state, + parts, + trace_id, + body_json, + input, + &eligible, + candidate_index, + &candidate_id, + spec, + ) + .await? + }; + let Some(resolved) = resolved else { return Ok(None); }; let candidate = &eligible.candidate; @@ -243,10 +282,11 @@ pub(crate) async fn maybe_build_local_openai_responses_decision_payload_for_cand report_context: Some(report_context), auth_context: input.auth_context.clone(), }); - apply_provider_request_routing_policy_to_decision( + apply_provider_request_routing_policy_to_decision_with_websocket_mode( input, &mut decision, Some(transport.as_ref()), + websocket_continuation, )?; Ok(Some(decision)) } 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 b2290eb0c..f9da0d939 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 @@ -31,6 +31,7 @@ use crate::ai_serving::planner::standard::{ build_cross_format_openai_responses_request_body_with_codex_model_capabilities, build_cross_format_openai_responses_upstream_url, build_local_openai_responses_request_body_with_codex_model_capabilities, + build_local_openai_responses_request_body_with_codex_model_capabilities_for_websocket_continuation, build_local_openai_responses_upstream_url, codex_model_capabilities_for_transport, openai_provider_request_contract_failure_extra_data, openai_responses_reasoning_replay_policy, request_body_build_failure_extra_data, request_conversion_failure_extra_data, @@ -203,6 +204,34 @@ pub(crate) async fn resolve_local_openai_responses_candidate_payload_parts( candidate_index: u32, candidate_id: &str, spec: LocalOpenAiResponsesSpec, +) -> Result, GatewayError> { + resolve_local_openai_responses_candidate_payload_parts_with_websocket_mode( + state, + parts, + trace_id, + body_json, + input, + eligible, + candidate_index, + candidate_id, + spec, + false, + ) + .await +} + +#[allow(clippy::too_many_arguments)] +pub(crate) async fn resolve_local_openai_responses_candidate_payload_parts_with_websocket_mode( + state: &AppState, + parts: &http::request::Parts, + trace_id: &str, + body_json: &serde_json::Value, + input: &LocalOpenAiResponsesDecisionInput, + eligible: &EligibleLocalExecutionCandidate, + candidate_index: u32, + candidate_id: &str, + spec: LocalOpenAiResponsesSpec, + websocket_continuation: bool, ) -> Result, GatewayError> { let spec_metadata = local_openai_responses_spec_metadata(spec); let client_api_format = spec_metadata.api_format.trim().to_ascii_lowercase(); @@ -405,12 +434,17 @@ pub(crate) async fn resolve_local_openai_responses_candidate_payload_parts( input.auth_context.api_key_id.as_str(), ) .await?; + let reasoning_replay_policy = openai_responses_reasoning_replay_policy( + transport.provider.provider_type.as_str(), + transport.endpoint.base_url.as_str(), + ); let redaction = resolve_provider_chat_pii_redaction( state, parts, body_json, &input.auth_context, spec_metadata.api_format, + reasoning_replay_policy, candidate_id, ) .await?; @@ -437,41 +471,42 @@ pub(crate) async fn resolve_local_openai_responses_candidate_payload_parts( mapped_model.as_str(), source_model, ); - let Some(mut base_provider_request_body) = - (if is_grok && is_grok_text_provider_api_format(provider_api_format) { - build_local_openai_responses_request_body_with_codex_model_capabilities( - body_json, - &mapped_model, - upstream_is_stream, - force_body_stream_field, - transport.provider.provider_type.as_str(), - spec_metadata.api_format, - transport.endpoint.body_rules.as_ref(), - effective_headers, - codex_model_capabilities.as_ref(), - false, - ) - } else if needs_bidirectional_conversion { - build_cross_format_openai_responses_request_body_with_codex_model_capabilities( - body_json, - &mapped_model, - spec_metadata.api_format, - provider_api_format, - upstream_is_stream, - force_body_stream_field, - transport.provider.provider_type.as_str(), - if is_kiro_claude_cli || is_windsurf_cascade { - None - } else { - transport.endpoint.body_rules.as_ref() - }, - effective_headers, - Some(input.auth_context.api_key_id.as_str()), - codex_model_capabilities.as_ref(), - false, - ) - } else { - build_local_openai_responses_request_body_with_codex_model_capabilities( + let Some(mut base_provider_request_body) = (if is_grok + && is_grok_text_provider_api_format(provider_api_format) + { + build_local_openai_responses_request_body_with_codex_model_capabilities( + body_json, + &mapped_model, + upstream_is_stream, + force_body_stream_field, + transport.provider.provider_type.as_str(), + spec_metadata.api_format, + transport.endpoint.body_rules.as_ref(), + effective_headers, + codex_model_capabilities.as_ref(), + false, + ) + } else if needs_bidirectional_conversion { + build_cross_format_openai_responses_request_body_with_codex_model_capabilities( + body_json, + &mapped_model, + spec_metadata.api_format, + provider_api_format, + upstream_is_stream, + force_body_stream_field, + transport.provider.provider_type.as_str(), + if is_kiro_claude_cli || is_windsurf_cascade { + None + } else { + transport.endpoint.body_rules.as_ref() + }, + effective_headers, + Some(input.auth_context.api_key_id.as_str()), + codex_model_capabilities.as_ref(), + false, + ) + } else if websocket_continuation { + build_local_openai_responses_request_body_with_codex_model_capabilities_for_websocket_continuation( body_json, &mapped_model, upstream_is_stream, @@ -487,8 +522,24 @@ pub(crate) async fn resolve_local_openai_responses_candidate_payload_parts( codex_model_capabilities.as_ref(), false, ) - }) - else { + } else { + build_local_openai_responses_request_body_with_codex_model_capabilities( + body_json, + &mapped_model, + upstream_is_stream, + force_body_stream_field, + transport.provider.provider_type.as_str(), + provider_api_format, + if is_kiro_claude_cli || is_windsurf_cascade { + None + } else { + transport.endpoint.body_rules.as_ref() + }, + effective_headers, + codex_model_capabilities.as_ref(), + false, + ) + }) else { mark_skipped_local_openai_responses_candidate_with_extra_data( state, input, @@ -531,29 +582,35 @@ pub(crate) async fn resolve_local_openai_responses_candidate_payload_parts( provider_api_format, Some(body_json), ); - if let Err(violation) = + let finalization = crate::ai_serving::OpenAiProviderRequestFinalization { + source_api_format: spec_metadata.api_format, + provider_api_format, + provider_type: transport.provider.provider_type.as_str(), + provider_model: mapped_model.as_str(), + source_model, + body_rules: transport.endpoint.body_rules.as_ref(), + upstream_is_stream, + require_body_stream_field: request_requires_body_stream_field( + body_json, + force_body_stream_field, + ), + }; + let finalization_result = if websocket_continuation { + crate::ai_serving::finalize_openai_provider_request_with_codex_model_capabilities_and_reasoning_replay_policy_for_websocket_continuation( + &mut base_provider_request_body, + finalization, + codex_model_capabilities.as_ref(), + reasoning_replay_policy, + ) + } else { crate::ai_serving::finalize_openai_provider_request_with_codex_model_capabilities_and_reasoning_replay_policy( &mut base_provider_request_body, - crate::ai_serving::OpenAiProviderRequestFinalization { - source_api_format: spec_metadata.api_format, - provider_api_format, - provider_type: transport.provider.provider_type.as_str(), - provider_model: mapped_model.as_str(), - source_model, - body_rules: transport.endpoint.body_rules.as_ref(), - upstream_is_stream, - require_body_stream_field: request_requires_body_stream_field( - body_json, - force_body_stream_field, - ), - }, + finalization, codex_model_capabilities.as_ref(), - openai_responses_reasoning_replay_policy( - transport.provider.provider_type.as_str(), - transport.endpoint.base_url.as_str(), - ), + reasoning_replay_policy, ) - { + }; + if let Err(violation) = finalization_result { mark_skipped_local_openai_responses_candidate_with_extra_data( state, input, 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 716a89aa4..6666be55a 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 @@ -2,7 +2,10 @@ use crate::ai_serving::planner::common::endpoint_config_forces_body_stream_field use crate::ai_serving::planner::plan_builders::{AiStreamAttempt, AiSyncAttempt}; use crate::ai_serving::planner::spec_metadata::local_openai_responses_spec_metadata; use crate::ai_serving::planner::standard::codex::codex_model_capabilities_for_transport; -use crate::ai_serving::planner::standard::normalize::build_local_openai_responses_request_body_with_codex_model_capabilities; +use crate::ai_serving::planner::standard::normalize::{ + build_local_openai_responses_request_body_with_codex_model_capabilities, + build_local_openai_responses_request_body_with_codex_model_capabilities_for_websocket_continuation, +}; use crate::ai_serving::planner::standard::openai_responses_reasoning_replay_policy; use crate::ai_serving::GatewayControlDecision; use crate::orchestration::{ @@ -62,6 +65,7 @@ mod plans; use self::decision::{ build_local_openai_responses_candidate_attempt_source, maybe_build_local_openai_responses_decision_payload_for_candidate, + maybe_build_local_openai_responses_decision_payload_for_candidate_with_websocket_mode, resolve_local_openai_responses_decision_input, resolve_local_openai_responses_decision_input_with_snapshot, }; @@ -238,7 +242,8 @@ pub(crate) struct ResponsesWebSocketDecision { /// but it still has to pass the current scheduler runtime checks on every /// turn. The planner uses this identity as a filter rather than selecting an /// arbitrary eligible replacement. -#[derive(Debug, Clone, PartialEq, Eq)] +#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)] +#[serde(deny_unknown_fields)] pub(crate) struct ResponsesWebSocketPinnedCandidate { provider_id: String, endpoint_id: String, @@ -246,14 +251,34 @@ pub(crate) struct ResponsesWebSocketPinnedCandidate { } impl ResponsesWebSocketPinnedCandidate { - pub(crate) fn from_decision(decision: &AiExecutionDecision) -> Option { + pub(crate) fn new(provider_id: &str, endpoint_id: &str, key_id: &str) -> Option { Some(Self { - provider_id: non_empty_decision_identity(decision.provider_id.as_deref())?, - endpoint_id: non_empty_decision_identity(decision.endpoint_id.as_deref())?, - key_id: non_empty_decision_identity(decision.key_id.as_deref())?, + provider_id: non_empty_decision_identity(Some(provider_id))?, + endpoint_id: non_empty_decision_identity(Some(endpoint_id))?, + key_id: non_empty_decision_identity(Some(key_id))?, }) } + pub(crate) fn from_decision(decision: &AiExecutionDecision) -> Option { + Self::new( + decision.provider_id.as_deref()?, + decision.endpoint_id.as_deref()?, + decision.key_id.as_deref()?, + ) + } + + pub(crate) fn provider_id(&self) -> &str { + self.provider_id.as_str() + } + + pub(crate) fn endpoint_id(&self) -> &str { + self.endpoint_id.as_str() + } + + pub(crate) fn key_id(&self) -> &str { + self.key_id.as_str() + } + fn matches( &self, candidate: &aether_scheduler_core::SchedulerMinimalCandidateSelectionCandidate, @@ -327,12 +352,143 @@ impl ResponsesWebSocketBodyNormalization { self } + #[cfg(test)] + pub(crate) fn with_body_rules_for_tests(mut self, body_rules: serde_json::Value) -> Self { + self.body_rules = Some(body_rules); + self + } + + #[cfg(test)] + pub(crate) fn with_reasoning_replay_policy_for_tests( + mut self, + reasoning_replay_policy: crate::ai_serving::OpenAiResponsesReasoningReplayPolicy, + ) -> Self { + self.reasoning_replay_policy = reasoning_replay_policy; + self + } + #[cfg(test)] pub(crate) fn with_model_directive_patch_for_tests(mut self, patch: serde_json::Value) -> Self { self.model_directive_patch = Some(patch); self } + pub(crate) fn uses_codex_responses_lite(&self) -> bool { + if !self.provider_type.trim().eq_ignore_ascii_case("codex") + || !crate::ai_serving::is_openai_responses_family_format( + self.provider_api_format.as_str(), + ) + { + return false; + } + self.codex_model_capabilities + .clone() + .unwrap_or_else(|| { + crate::ai_serving::resolve_codex_responses_model_capabilities( + self.mapped_model.as_str(), + self.requested_model.as_str(), + None, + ) + }) + .use_responses_lite + } + + pub(crate) fn reasoning_replay_policy( + &self, + ) -> crate::ai_serving::OpenAiResponsesReasoningReplayPolicy { + self.reasoning_replay_policy + } + + /// Returns whether an enabled endpoint body rule that applies to this + /// request owns the final value of a non-lineage WebSocket framing field. + /// + /// Codex's HTTP-shaped normalization intentionally removes or rewrites a + /// few WebSocket-only fields. The framing layer may restore a value from + /// the raw client event only when an administrator rule did not handle + /// that path; otherwise the restore would silently undo the endpoint + /// policy after all request finalization had completed. Opaque lineage + /// (`previous_response_id`) is deliberately excluded by the framing layer: + /// its final value must remain the authenticated client value. + pub(crate) fn body_rules_handle_websocket_field( + &self, + client_event: &serde_json::Value, + field: &str, + ) -> bool { + let Some(mut body_before_rules) = + crate::ai_serving::build_local_openai_responses_request_body_with_model_directives( + client_event, + self.mapped_model.as_str(), + self.upstream_is_stream, + false, + ) + else { + // Normalization will reject the same malformed event. Keep the + // framing pass fail closed if this method is ever called alone. + return true; + }; + crate::ai_serving::transport::rules::apply_local_body_rules_with_request_headers_and_track_path( + &mut body_before_rules, + self.body_rules.as_ref(), + Some(client_event), + Some(&self.request_headers), + field, + ) + .unwrap_or(true) + } + + pub(crate) fn has_same_responses_lite_static_contract(&self, other: &Self) -> bool { + self.provider_type + .trim() + .eq_ignore_ascii_case(other.provider_type.trim()) + && crate::ai_serving::api_format_alias_matches( + self.provider_api_format.as_str(), + other.provider_api_format.as_str(), + ) + && self.mapped_model == other.mapped_model + && self.requested_model == other.requested_model + && self.body_rules == other.body_rules + && self.codex_model_capabilities == other.codex_model_capabilities + && self.model_directive_patch == other.model_directive_patch + && self.uses_codex_responses_lite() == other.uses_codex_responses_lite() + } + + /// Produces a versioned digest of the complete body-normalization + /// contract. A continuation registry stores only this digest so a new + /// socket can fail closed when endpoint rules, model capabilities or + /// header-dependent normalization has changed, without persisting request + /// headers or other sensitive configuration. + pub(crate) fn continuation_fingerprint(&self) -> [u8; 32] { + use sha2::Digest as _; + + let mut digest = sha2::Sha256::new(); + digest.update(b"aether-responses-websocket-normalization-v1"); + update_normalization_string_digest(&mut digest, self.provider_type.as_str()); + update_normalization_string_digest(&mut digest, self.provider_api_format.as_str()); + update_normalization_string_digest(&mut digest, self.client_api_format.as_str()); + update_normalization_string_digest(&mut digest, self.mapped_model.as_str()); + update_normalization_string_digest(&mut digest, self.requested_model.as_str()); + digest.update([ + u8::from(self.upstream_is_stream), + u8::from(self.force_body_stream_field), + ]); + update_normalization_optional_json_digest(&mut digest, self.body_rules.as_ref()); + update_normalization_body_rule_headers_digest( + &mut digest, + &self.request_headers, + self.body_rules.as_ref(), + ); + update_normalization_codex_capabilities_digest( + &mut digest, + self.codex_model_capabilities.as_ref(), + ); + digest.update([match self.reasoning_replay_policy { + crate::ai_serving::OpenAiResponsesReasoningReplayPolicy::OpenAiItemIds => 0, + crate::ai_serving::OpenAiResponsesReasoningReplayPolicy::DeepSeekOpaque => 1, + }]); + update_normalization_optional_json_digest(&mut digest, self.model_directive_patch.as_ref()); + digest.finalize().into() + } + /// Applies the same body transformations the planner applied on the turn /// that bound this upstream. /// @@ -342,9 +498,9 @@ impl ResponsesWebSocketBodyNormalization { /// here: the WebSocket planner only returns candidates whose provider API /// format is `openai:responses`. /// - /// Returns `None` when normalization fails, leaving the caller to fall back - /// to the unnormalized event — a continuation cannot re-select a candidate, - /// so failing the turn outright would be worse than sending it as-is. + /// Returns `None` when normalization fails. The WebSocket caller rejects + /// that turn rather than sending an unnormalized event that bypasses body + /// rules or replays a Responses Lite static prefix. pub(crate) fn normalize_response_create( &self, client_event: &serde_json::Value, @@ -357,20 +513,44 @@ impl ResponsesWebSocketBodyNormalization { .get("model") .and_then(serde_json::Value::as_str) .unwrap_or(self.requested_model.as_str()); + // The first response.create on a socket is a normal Responses request. + // Only a non-empty previous_response_id denotes a continuation whose + // stored history already contains the synthetic Responses Lite + // tools/instructions prefix. Keep this discriminator explicit instead + // of applying continuation edits to every socket turn. + let websocket_continuation = client_event + .get("previous_response_id") + .and_then(serde_json::Value::as_str) + .is_some_and(|value| !value.trim().is_empty()); let require_body_stream_field = request_requires_body_stream_field(client_event, self.force_body_stream_field); - let mut body = build_local_openai_responses_request_body_with_codex_model_capabilities( - client_event, - &self.mapped_model, - self.upstream_is_stream, - self.force_body_stream_field, - self.provider_type.as_str(), - self.provider_api_format.as_str(), - self.body_rules.as_ref(), - &self.request_headers, - self.codex_model_capabilities.as_ref(), - false, - )?; + let mut body = if websocket_continuation { + build_local_openai_responses_request_body_with_codex_model_capabilities_for_websocket_continuation( + client_event, + &self.mapped_model, + self.upstream_is_stream, + self.force_body_stream_field, + self.provider_type.as_str(), + self.provider_api_format.as_str(), + self.body_rules.as_ref(), + &self.request_headers, + self.codex_model_capabilities.as_ref(), + false, + ) + } else { + build_local_openai_responses_request_body_with_codex_model_capabilities( + client_event, + &self.mapped_model, + self.upstream_is_stream, + self.force_body_stream_field, + self.provider_type.as_str(), + self.provider_api_format.as_str(), + self.body_rules.as_ref(), + &self.request_headers, + self.codex_model_capabilities.as_ref(), + false, + ) + }?; if let Some(patch) = self.model_directive_patch.as_ref() { crate::ai_serving::apply_model_directive_mapping_patch(&mut body, patch); // The patch is a deep merge and may reintroduce `stream`. @@ -381,26 +561,281 @@ impl ResponsesWebSocketBodyNormalization { require_body_stream_field, ); } - crate::ai_serving::finalize_openai_provider_request_with_codex_model_capabilities_and_reasoning_replay_policy( - &mut body, - crate::ai_serving::OpenAiProviderRequestFinalization { - source_api_format: self.client_api_format.as_str(), - provider_api_format: self.provider_api_format.as_str(), - provider_type: self.provider_type.as_str(), - provider_model: self.mapped_model.as_str(), - source_model, - body_rules: self.body_rules.as_ref(), - upstream_is_stream: self.upstream_is_stream, - require_body_stream_field, - }, - self.codex_model_capabilities.as_ref(), - self.reasoning_replay_policy, - ) - .ok()?; + let finalization = crate::ai_serving::OpenAiProviderRequestFinalization { + source_api_format: self.client_api_format.as_str(), + provider_api_format: self.provider_api_format.as_str(), + provider_type: self.provider_type.as_str(), + provider_model: self.mapped_model.as_str(), + source_model, + body_rules: self.body_rules.as_ref(), + upstream_is_stream: self.upstream_is_stream, + require_body_stream_field, + }; + let finalized = if websocket_continuation { + crate::ai_serving::finalize_openai_provider_request_with_codex_model_capabilities_and_reasoning_replay_policy_for_websocket_continuation( + &mut body, + finalization, + self.codex_model_capabilities.as_ref(), + self.reasoning_replay_policy, + ) + } else { + crate::ai_serving::finalize_openai_provider_request_with_codex_model_capabilities_and_reasoning_replay_policy( + &mut body, + finalization, + self.codex_model_capabilities.as_ref(), + self.reasoning_replay_policy, + ) + }; + finalized.ok()?; Some(body) } } +fn update_normalization_bytes_digest(digest: &mut sha2::Sha256, value: &[u8]) { + use sha2::Digest as _; + + digest.update((value.len() as u64).to_be_bytes()); + digest.update(value); +} + +fn update_normalization_string_digest(digest: &mut sha2::Sha256, value: &str) { + update_normalization_bytes_digest(digest, value.as_bytes()); +} + +fn update_normalization_optional_string_digest(digest: &mut sha2::Sha256, value: Option<&str>) { + use sha2::Digest as _; + + match value { + Some(value) => { + digest.update([1]); + update_normalization_string_digest(digest, value); + } + None => digest.update([0]), + } +} + +fn update_normalization_string_vec_digest(digest: &mut sha2::Sha256, values: &[String]) { + use sha2::Digest as _; + + digest.update((values.len() as u64).to_be_bytes()); + for value in values { + update_normalization_string_digest(digest, value); + } +} + +fn update_normalization_optional_json_digest( + digest: &mut sha2::Sha256, + value: Option<&serde_json::Value>, +) { + use sha2::Digest as _; + + match value { + Some(value) => { + digest.update([1]); + update_normalization_json_digest(digest, value); + } + None => digest.update([0]), + } +} + +fn update_normalization_json_digest(digest: &mut sha2::Sha256, value: &serde_json::Value) { + use serde_json::Value; + use sha2::Digest as _; + + match value { + Value::Null => digest.update(b"n"), + Value::Bool(value) => digest.update(if *value { b"t" } else { b"f" }), + Value::Number(value) => { + digest.update(b"d"); + update_normalization_string_digest(digest, value.to_string().as_str()); + } + Value::String(value) => { + digest.update(b"s"); + update_normalization_string_digest(digest, value); + } + Value::Array(values) => { + digest.update(b"["); + digest.update((values.len() as u64).to_be_bytes()); + for value in values { + update_normalization_json_digest(digest, value); + } + digest.update(b"]"); + } + Value::Object(values) => { + digest.update(b"{"); + digest.update((values.len() as u64).to_be_bytes()); + let mut keys = values.keys().collect::>(); + keys.sort_unstable(); + for key in keys { + update_normalization_string_digest(digest, key); + update_normalization_json_digest(digest, &values[key]); + } + digest.update(b"}"); + } + } +} + +fn update_normalization_body_rule_headers_digest( + digest: &mut sha2::Sha256, + headers: &http::HeaderMap, + body_rules: Option<&serde_json::Value>, +) { + use sha2::Digest as _; + + let dependencies = + crate::ai_serving::transport::rules::body_rules_request_header_dependencies(body_rules); + digest.update((dependencies.len() as u64).to_be_bytes()); + for name in dependencies { + update_normalization_string_digest(digest, name.as_str()); + let value = headers + .get(name.as_str()) + .and_then(|value| value.to_str().ok()) + .map(str::trim); + update_normalization_optional_string_digest(digest, value); + } +} + +fn update_normalization_codex_capabilities_digest( + digest: &mut sha2::Sha256, + capabilities: Option<&crate::ai_serving::CodexResponsesModelCapabilities>, +) { + use sha2::Digest as _; + + let Some(capabilities) = capabilities else { + digest.update([0]); + return; + }; + digest.update([1]); + digest.update([ + u8::from(capabilities.use_responses_lite), + u8::from(capabilities.supports_reasoning_summary_parameter), + u8::from(capabilities.supports_parallel_tool_calls), + u8::from(capabilities.support_verbosity), + ]); + update_normalization_optional_string_digest( + digest, + capabilities.default_reasoning_effort.as_deref(), + ); + update_normalization_optional_string_digest( + digest, + capabilities.default_reasoning_summary.as_deref(), + ); + update_normalization_string_vec_digest(digest, &capabilities.supported_reasoning_efforts); + update_normalization_optional_string_digest(digest, capabilities.default_verbosity.as_deref()); + update_normalization_string_vec_digest(digest, &capabilities.supported_service_tiers); +} + +#[cfg(test)] +mod continuation_fingerprint_tests { + use http::HeaderValue; + use serde_json::json; + + use super::ResponsesWebSocketBodyNormalization; + use crate::ai_serving::OpenAiResponsesReasoningReplayPolicy; + + #[test] + fn normalization_fingerprint_is_stable_for_json_object_key_order() { + let first = ResponsesWebSocketBodyNormalization::for_tests("provider-model") + .with_model_directive_patch_for_tests(json!({"reasoning": {"effort": "high"}, "x": 1})); + let second = ResponsesWebSocketBodyNormalization::for_tests("provider-model") + .with_model_directive_patch_for_tests(json!({"x": 1, "reasoning": {"effort": "high"}})); + assert_eq!( + first.continuation_fingerprint(), + second.continuation_fingerprint() + ); + } + + #[test] + fn normalization_fingerprint_changes_with_effective_contract() { + let base = ResponsesWebSocketBodyNormalization::for_tests("provider-model"); + let changed_policy = base.clone().with_reasoning_replay_policy_for_tests( + OpenAiResponsesReasoningReplayPolicy::DeepSeekOpaque, + ); + assert_ne!( + base.continuation_fingerprint(), + changed_policy.continuation_fingerprint() + ); + + let changed_patch = base + .clone() + .with_model_directive_patch_for_tests(json!({"reasoning": {"effort": "low"}})); + assert_ne!( + base.continuation_fingerprint(), + changed_patch.continuation_fingerprint() + ); + } + + #[test] + fn normalization_fingerprint_ignores_unrelated_volatile_request_headers() { + let body_rules = json!([{ + "action": "set", + "path": "store", + "value": false, + "condition": { + "source": "request_headers", + "path": "x-contract", + "op": "eq", + "value": "enabled" + } + }]); + let mut first = ResponsesWebSocketBodyNormalization::for_tests("provider-model") + .with_body_rules_for_tests(body_rules); + first + .request_headers + .insert("x-contract", HeaderValue::from_static("enabled")); + first + .request_headers + .insert("x-request-id", HeaderValue::from_static("request-1")); + first + .request_headers + .insert("cf-ray", HeaderValue::from_static("edge-1")); + let mut second = first.clone(); + second + .request_headers + .insert("x-request-id", HeaderValue::from_static("request-2")); + second + .request_headers + .insert("cf-ray", HeaderValue::from_static("edge-2")); + + assert_eq!( + first.continuation_fingerprint(), + second.continuation_fingerprint(), + "headers that no body-rule condition reads must not invalidate a persisted continuation" + ); + } + + #[test] + fn normalization_fingerprint_tracks_headers_used_by_body_rule_conditions() { + let body_rules = json!([{ + "action": "set", + "path": "store", + "value": false, + "condition": { + "source": "request_headers", + "path": "X-Contract", + "op": "eq", + "value": "enabled" + } + }]); + let mut first = ResponsesWebSocketBodyNormalization::for_tests("provider-model") + .with_body_rules_for_tests(body_rules.clone()); + first + .request_headers + .insert("x-contract", HeaderValue::from_static("enabled")); + let mut second = ResponsesWebSocketBodyNormalization::for_tests("provider-model") + .with_body_rules_for_tests(body_rules); + second + .request_headers + .insert("x-contract", HeaderValue::from_static("disabled")); + + assert_ne!( + first.continuation_fingerprint(), + second.continuation_fingerprint(), + "a header that controls an effective body-rule condition remains part of the contract" + ); + } +} + /// Builds one upstream decision for a Responses WebSocket turn. The session /// reuses this decision for same-model turns and invokes the planner again when /// a later `response.create` changes the public model. @@ -432,6 +867,14 @@ pub(crate) async fn maybe_build_responses_websocket_decision( else { return Ok(None); }; + // The continuation discriminator belongs to the WebSocket protocol, not + // provider body rules/redaction. Capture it before the planner creates its + // effective body so a rule cannot accidentally turn a valid chain into a + // first-turn Lite normalization pass. + let websocket_continuation = body_json + .get("previous_response_id") + .and_then(serde_json::Value::as_str) + .is_some_and(|value| !value.trim().is_empty()); let body_json = input.effective_body_json(body_json); let (mut source, _) = build_local_openai_responses_candidate_attempt_source( state, trace_id, &input, body_json, spec, @@ -467,8 +910,15 @@ pub(crate) async fn maybe_build_responses_websocket_decision( // reproduce this candidate's body normalization without re-planning. let transport = std::sync::Arc::clone(&attempt.eligible.transport); let candidate_provider_api_format = attempt.eligible.provider_api_format.clone(); - let payload = match maybe_build_local_openai_responses_decision_payload_for_candidate( - state, parts, trace_id, body_json, &input, attempt, spec, + let payload = match maybe_build_local_openai_responses_decision_payload_for_candidate_with_websocket_mode( + state, + parts, + trace_id, + body_json, + &input, + attempt, + spec, + websocket_continuation, ) .await { diff --git a/apps/aether-gateway/src/ai_serving/pure/mod.rs b/apps/aether-gateway/src/ai_serving/pure/mod.rs index e6ca938fc..16fcd2266 100644 --- a/apps/aether-gateway/src/ai_serving/pure/mod.rs +++ b/apps/aether-gateway/src/ai_serving/pure/mod.rs @@ -9,6 +9,7 @@ pub(crate) use aether_ai_formats::api::{ apply_codex_openai_responses_lite_header_with_capabilities, apply_codex_openai_responses_special_body_edits, apply_codex_openai_responses_special_body_edits_with_source_model_and_capabilities, + apply_codex_openai_responses_websocket_continuation_body_edits_with_source_model_and_capabilities, apply_codex_openai_special_headers, apply_model_directive_mapping_patch, apply_model_directive_overrides_from_model, apply_model_directive_overrides_from_request, apply_openai_responses_compact_special_body_edits, build_chatgpt_web_image_request_body, @@ -37,6 +38,7 @@ pub(crate) use aether_ai_formats::api::{ build_standard_request_body_from_canonical_with_model_directives, build_standard_request_body_with_model_directives, build_standard_request_body_with_model_directives_and_request_headers, + build_standard_request_body_with_model_directives_and_request_headers_and_reasoning_replay_policy, calculate_kiro_context_input_tokens, canonicalize_tool_arguments, convert_claude_chat_response_to_openai_chat, convert_claude_response_to_openai_responses, convert_gemini_chat_response_to_openai_chat, convert_gemini_response_to_openai_responses, @@ -55,6 +57,7 @@ pub(crate) use aether_ai_formats::api::{ finalize_openai_provider_request, finalize_openai_provider_request_with_codex_model_capabilities, finalize_openai_provider_request_with_codex_model_capabilities_and_reasoning_replay_policy, + finalize_openai_provider_request_with_codex_model_capabilities_and_reasoning_replay_policy_for_websocket_continuation, find_kiro_real_thinking_end_tag, find_kiro_real_thinking_end_tag_at_buffer_end, find_kiro_real_thinking_start_tag, forbid_upstream_streaming_for_provider, force_upstream_streaming_for_provider, gemini_request_is_image_generation, @@ -172,10 +175,11 @@ pub(crate) use aether_ai_formats::api::{ }; pub(crate) use aether_ai_formats::{ api_format_defaults_to_client_error_failover, api_format_defaults_to_non_stream, - api_format_permission_covers, intersect_api_format_allowed_lists, is_embedding_api_format, - is_rerank_api_format, openai_responses_request_operation, - openai_responses_synthetic_reasoning_item_id, - strip_incompatible_openai_responses_reasoning_items, ApiOperation, ClientSurface, + api_format_permission_covers, codex_responses_lite_tool_is_client_executed, + intersect_api_format_allowed_lists, is_embedding_api_format, is_rerank_api_format, + openai_responses_request_operation, openai_responses_synthetic_reasoning_item_id, + strip_incompatible_openai_responses_reasoning_items, + strip_incompatible_openai_responses_reasoning_items_with_policy, ApiOperation, ClientSurface, CODEX_CLIENT_VERSION, }; diff --git a/apps/aether-gateway/src/handlers/admin/provider/query/models/model_test.rs b/apps/aether-gateway/src/handlers/admin/provider/query/models/model_test.rs index 1210e89f0..ea6ff99d5 100644 --- a/apps/aether-gateway/src/handlers/admin/provider/query/models/model_test.rs +++ b/apps/aether-gateway/src/handlers/admin/provider/query/models/model_test.rs @@ -3162,7 +3162,7 @@ async fn provider_query_execute_standard_test_candidate( | "openai:rerank" | "jina:rerank" => { let Some(mut provider_request_body) = - crate::ai_serving::build_standard_request_body_with_model_directives_and_request_headers( + crate::ai_serving::build_standard_request_body_with_model_directives_and_request_headers_and_reasoning_replay_policy( &request_body, client_api_format, request_model, diff --git a/apps/aether-gateway/src/handlers/proxy/websocket/responses/binding.rs b/apps/aether-gateway/src/handlers/proxy/websocket/responses/binding.rs index 62ab8af7f..05e75b1b3 100644 --- a/apps/aether-gateway/src/handlers/proxy/websocket/responses/binding.rs +++ b/apps/aether-gateway/src/handlers/proxy/websocket/responses/binding.rs @@ -143,17 +143,162 @@ impl UpstreamBindingIdentity { } changed } + + /// Returns a versioned, one-way identity suitable for a continuation + /// registry. The digest covers every field used by `PartialEq`, including + /// effective credential and transport state, without persisting header or + /// credential values themselves. + pub(super) fn continuation_fingerprint(&self) -> [u8; 32] { + let mut digest = Sha256::new(); + digest.update(b"aether-responses-websocket-binding-v1"); + digest.update([match self.adapter_kind { + ResponsesWebSocketAdapter::Standard => 0, + ResponsesWebSocketAdapter::Codex => 1, + }]); + update_optional_string_digest(&mut digest, self.provider_id.as_deref()); + update_optional_string_digest(&mut digest, self.endpoint_id.as_deref()); + update_optional_string_digest(&mut digest, self.key_id.as_deref()); + update_string_digest(&mut digest, self.upstream_url.as_str()); + update_string_map_digest(&mut digest, &self.handshake_headers); + digest.update(self.credential_fingerprint); + update_proxy_digest(&mut digest, self.proxy.as_ref()); + update_transport_profile_digest(&mut digest, self.transport_profile.as_ref()); + digest.finalize().into() + } + + pub(super) fn adapter_kind(&self) -> ResponsesWebSocketAdapter { + self.adapter_kind + } } -/// `x-codex-turn-metadata` describes one logical response turn. Fingerprint -/// convergence intentionally rewrites its `turn_id` and timestamp for every -/// `response.create`, including a continuation on an already-upgraded socket. -/// It therefore cannot identify the physical handshake that owns a -/// `previous_response_id`; the per-turn copy in `client_metadata` still travels -/// in the response.create body. +fn update_bytes_digest(digest: &mut Sha256, value: &[u8]) { + digest.update((value.len() as u64).to_be_bytes()); + digest.update(value); +} + +fn update_string_digest(digest: &mut Sha256, value: &str) { + update_bytes_digest(digest, value.as_bytes()); +} + +fn update_optional_string_digest(digest: &mut Sha256, value: Option<&str>) { + match value { + Some(value) => { + digest.update([1]); + update_string_digest(digest, value); + } + None => digest.update([0]), + } +} + +fn update_optional_bool_digest(digest: &mut Sha256, value: Option) { + digest.update([match value { + None => 0, + Some(false) => 1, + Some(true) => 2, + }]); +} + +fn update_string_map_digest(digest: &mut Sha256, values: &BTreeMap) { + digest.update((values.len() as u64).to_be_bytes()); + for (name, value) in values { + update_string_digest(digest, name); + update_string_digest(digest, value); + } +} + +fn update_optional_json_digest(digest: &mut Sha256, value: Option<&serde_json::Value>) { + match value { + Some(value) => { + digest.update([1]); + update_json_digest(digest, value); + } + None => digest.update([0]), + } +} + +fn update_json_digest(digest: &mut Sha256, value: &serde_json::Value) { + use serde_json::Value; + + match value { + Value::Null => digest.update(b"n"), + Value::Bool(value) => digest.update(if *value { b"t" } else { b"f" }), + Value::Number(value) => { + digest.update(b"d"); + update_string_digest(digest, value.to_string().as_str()); + } + Value::String(value) => { + digest.update(b"s"); + update_string_digest(digest, value); + } + Value::Array(values) => { + digest.update(b"["); + digest.update((values.len() as u64).to_be_bytes()); + for value in values { + update_json_digest(digest, value); + } + digest.update(b"]"); + } + Value::Object(values) => { + digest.update(b"{"); + digest.update((values.len() as u64).to_be_bytes()); + let mut keys = values.keys().collect::>(); + keys.sort_unstable(); + for key in keys { + update_string_digest(digest, key); + update_json_digest(digest, &values[key]); + } + digest.update(b"}"); + } + } +} + +fn update_proxy_digest(digest: &mut Sha256, proxy: Option<&ProxySnapshot>) { + let Some(proxy) = proxy else { + digest.update([0]); + return; + }; + digest.update([1]); + update_optional_bool_digest(digest, proxy.enabled); + update_optional_string_digest(digest, proxy.mode.as_deref()); + update_optional_string_digest(digest, proxy.node_id.as_deref()); + update_optional_string_digest(digest, proxy.label.as_deref()); + update_optional_string_digest(digest, proxy.url.as_deref()); + update_optional_json_digest(digest, proxy.extra.as_ref()); +} + +fn update_transport_profile_digest( + digest: &mut Sha256, + profile: Option<&ResolvedTransportProfile>, +) { + let Some(profile) = profile else { + digest.update([0]); + return; + }; + digest.update([1]); + update_string_digest(digest, profile.profile_id.as_str()); + update_string_digest(digest, profile.backend.as_str()); + update_string_digest(digest, profile.http_mode.as_str()); + update_string_digest(digest, profile.pool_scope.as_str()); + update_optional_json_digest(digest, profile.header_fingerprint.as_ref()); + update_optional_json_digest(digest, profile.extra.as_ref()); +} + +/// Excludes headers whose values identify one downstream request/turn rather +/// than the provider connection contract. +/// +/// `x-trace-id` is generated by Aether's access middleware for every +/// downstream request. A cross-socket continuation therefore receives a new +/// value even when it revalidates to the exact same provider binding. +/// +/// `x-codex-turn-metadata` likewise describes one logical response turn. +/// Fingerprint convergence intentionally rewrites its `turn_id` and timestamp +/// for every `response.create`, including a continuation on an already-upgraded +/// socket. The per-turn copy in `client_metadata` still travels in the +/// response.create body. fn is_turn_scoped_handshake_header(adapter_kind: ResponsesWebSocketAdapter, name: &str) -> bool { - adapter_kind == ResponsesWebSocketAdapter::Codex - && name.eq_ignore_ascii_case("x-codex-turn-metadata") + name.eq_ignore_ascii_case(crate::constants::TRACE_ID_HEADER) + || (adapter_kind == ResponsesWebSocketAdapter::Codex + && name.eq_ignore_ascii_case("x-codex-turn-metadata")) } /// Header names that carry credentials in the provider handshake. The @@ -386,6 +531,51 @@ mod tests { ); } + #[test] + fn request_trace_id_does_not_change_standard_or_codex_binding_identity() { + for adapter_kind in [ + ResponsesWebSocketAdapter::Standard, + ResponsesWebSocketAdapter::Codex, + ] { + let adapter = resolve_responses_websocket_adapter(adapter_kind); + let mut first = decision(); + if adapter_kind == ResponsesWebSocketAdapter::Codex { + first.provider_type = Some("codex".to_string()); + first.report_context = Some(json!({ + "codex_credential_generation": "credential-generation-1" + })); + } + first + .provider_request_headers + .insert("x-trace-id".to_string(), "request-1".to_string()); + let first_identity = UpstreamBindingIdentity::from_decision(adapter, &first).unwrap(); + assert!(!first_identity.handshake_headers.contains_key("x-trace-id")); + + let mut continuation = first; + continuation + .provider_request_headers + .insert("x-trace-id".to_string(), "request-2".to_string()); + let continuation_identity = + UpstreamBindingIdentity::from_decision(adapter, &continuation).unwrap(); + assert_eq!(first_identity, continuation_identity); + assert_eq!( + first_identity.continuation_fingerprint(), + continuation_identity.continuation_fingerprint() + ); + + continuation + .provider_request_headers + .insert("x-correlation-id".to_string(), "changed".to_string()); + let changed_identity = + UpstreamBindingIdentity::from_decision(adapter, &continuation).unwrap(); + assert_ne!(first_identity, changed_identity); + assert_eq!( + first_identity.changed_field_names(&changed_identity), + vec!["handshake_header:x-correlation-id".to_string()] + ); + } + } + #[test] fn identity_changes_when_physical_binding_changes() { let adapter = resolve_responses_websocket_adapter(ResponsesWebSocketAdapter::Standard); @@ -441,6 +631,25 @@ mod tests { ); } + #[test] + fn continuation_fingerprint_changes_with_binding_and_is_not_secret_text() { + let adapter = resolve_responses_websocket_adapter(ResponsesWebSocketAdapter::Standard); + let base = decision(); + let identity = UpstreamBindingIdentity::from_decision(adapter, &base).unwrap(); + let mut changed = base; + changed + .provider_request_headers + .insert("X-Client".to_string(), "different-client".to_string()); + let changed_identity = UpstreamBindingIdentity::from_decision(adapter, &changed).unwrap(); + assert_ne!( + identity.continuation_fingerprint(), + changed_identity.continuation_fingerprint() + ); + let digest = format!("{:?}", identity.continuation_fingerprint()); + assert!(!digest.contains("secret")); + assert!(!digest.contains("api.example")); + } + #[test] fn stable_key_identity_rejects_custom_static_auth_value_rotation() { let adapter = resolve_responses_websocket_adapter(ResponsesWebSocketAdapter::Standard); @@ -528,7 +737,7 @@ mod tests { } #[test] - fn only_codex_turn_metadata_is_excluded_from_binding_headers() { + fn unknown_codex_headers_remain_part_of_the_binding_identity() { let codex_adapter = resolve_responses_websocket_adapter(ResponsesWebSocketAdapter::Codex); let mut first = decision(); first.provider_type = Some("codex".to_string()); diff --git a/apps/aether-gateway/src/handlers/proxy/websocket/responses/client.rs b/apps/aether-gateway/src/handlers/proxy/websocket/responses/client.rs index f7140f84c..a7a978a6a 100644 --- a/apps/aether-gateway/src/handlers/proxy/websocket/responses/client.rs +++ b/apps/aether-gateway/src/handlers/proxy/websocket/responses/client.rs @@ -17,11 +17,14 @@ use super::ownership::{ spawn_owned_responses_websocket_plan, OwnedResponsesWebSocketDecision, }; use super::quota::mark_active_response_retry_unsafe; -use super::redaction::redact_responses_websocket_client_event; +use super::redaction::redact_responses_websocket_client_event_with_reasoning_replay_policy; use super::request::{ - build_planning_parts, changed_followup_response_create_model, planned_response_create_event, - provider_model_from_decision, response_create_has_previous_response_id, - response_create_model_or_current, + build_planning_parts, changed_followup_response_create_model, + planned_request_uses_codex_responses_lite, planned_response_create_event, + prepare_responses_lite_continuation, provider_model_from_decision, + response_create_has_previous_response_id, response_create_model_or_current, + validate_response_create_previous_response_id, validate_response_create_stream_id_support, + validated_named_stream_id, ResponsesLiteStaticConfig, }; use super::state::BoundResponsesConnection; use super::turn::{ @@ -39,8 +42,10 @@ use crate::handlers::proxy::websocket::ingress::WebSocketRequestContext; use crate::handlers::proxy::websocket::session::{CLOSE_INTERNAL_ERROR, WEBSOCKET_LOG_TRANSPORT}; use crate::handlers::proxy::websocket::transport::{ close_client_socket, close_upstream_socket, send_client_message, send_gateway_error, - send_gateway_error_with_status, send_upstream_message, + send_gateway_error_with_status, send_gateway_error_with_stream_id, + send_responses_websocket_error_with_param, send_upstream_message, }; +use crate::privacy::RedactionSession; use crate::rate_limit::FrontdoorUserRpmOutcome; use crate::AppState; @@ -80,20 +85,51 @@ pub(super) fn adapter_drain_ready( )) } -fn parse_response_create_event(text: &str) -> Result { - let event = serde_json::from_str::(text).map_err(|_| "invalid_response_create")?; +#[derive(Debug, Clone, PartialEq, Eq)] +struct ResponseCreateParseError { + code: &'static str, + stream_id: Option, +} + +impl ResponseCreateParseError { + fn new(code: &'static str, event: Option<&Value>) -> Self { + // A syntactically valid named lane is safe to reflect on every + // request-scoped error, even when another response.create field fails + // validation first. Invalid lane values never pass this helper and + // therefore cannot be reflected. + let stream_id = event + .and_then(validated_named_stream_id) + .map(str::to_string); + Self { code, stream_id } + } +} + +fn parse_response_create_event(text: &str) -> Result { + let event = serde_json::from_str::(text) + .map_err(|_| ResponseCreateParseError::new("invalid_response_create", None))?; if event.as_object().is_none() { - return Err("invalid_response_create"); + return Err(ResponseCreateParseError::new( + "invalid_response_create", + Some(&event), + )); } if event.get("type").and_then(Value::as_str) != Some("response.create") { - return Err("expected_response_create"); + return Err(ResponseCreateParseError::new( + "expected_response_create", + Some(&event), + )); } + validate_response_create_previous_response_id(&event) + .map_err(|code| ResponseCreateParseError::new(code, Some(&event)))?; + validate_response_create_stream_id_support(&event) + .map_err(|code| ResponseCreateParseError::new(code, Some(&event)))?; Ok(event) } #[derive(Debug, Clone, Copy, PartialEq, Eq)] enum ContinuationConstraint { Pinned, + UnknownResponseId, UpstreamUnavailable, ModelChangeUnsupported, } @@ -106,10 +142,14 @@ fn continuation_constraint( event: &Value, current_client_model: &str, upstream_available: bool, + response_id_owned_by_connection: bool, ) -> Result, &'static str> { if !response_create_has_previous_response_id(event) { return Ok(None); } + if !response_id_owned_by_connection { + return Ok(Some(ContinuationConstraint::UnknownResponseId)); + } if changed_followup_response_create_model(event, current_client_model)?.is_some() { return Ok(Some(ContinuationConstraint::ModelChangeUnsupported)); } @@ -151,11 +191,24 @@ pub(super) async fn forward_client_message( let text = text.to_string(); let mut client_event = match parse_response_create_event(&text) { Ok(event) => event, - Err(code) => { - send_gateway_error( + Err(error) => { + let message = match error.code { + "invalid_response_create_previous_response_id" => { + "response.create.previous_response_id must be null or a non-empty string" + } + "invalid_response_create_stream_id" => { + "response.create.stream_id must be 1-256 ASCII letters, numbers, underscores, hyphens, or periods" + } + "responses_websocket_named_stream_unsupported" => { + "Aether currently supports only the implicit default WebSocket lane; omit response.create.stream_id" + } + _ => "WebSocket client text events must be response.create JSON objects", + }; + send_gateway_error_with_stream_id( client_socket, - code, - "WebSocket client text events must be response.create JSON objects", + error.code, + message, + error.stream_id.as_deref(), ) .await; return RelayDisposition::Continue; @@ -256,25 +309,116 @@ pub(super) async fn forward_client_message( } } - // 这一轮的 planning Parts 只构造一次（它携带 per-turn 的 - // RedactionSessionSlot），并且客户端事件也只在这里脱敏一次： - // 复用已绑定 upstream 的 continuation 根本不进 planner，只靠 planner - // 内部脱敏拦不住它。之后 re-plan / continuation / 配额重试都只看脱敏 - // 后的事件，上游请求体与审计 original_request_body 因此一致。 - let redacted_client_event = redact_responses_websocket_client_event( - state, - &planning_parts, - &turn_control.decision, + let response_id_owned_by_connection = client_event + .get("previous_response_id") + .and_then(Value::as_str) + .is_none_or(|response_id| bound.continuation_response_ids.contains(response_id)); + let is_pinned_continuation = match continuation_constraint( &client_event, - ) - .await; - let client_event = match redacted_client_event { - Ok(Some(redaction)) => { - // 这一轮的映射登记到连接上，响应帧才能在最后一跳还原回真实值。 - bound.redaction_restorer.register(redaction.session); - redaction.client_event + &bound.client_model, + bound.upstream.is_some(), + response_id_owned_by_connection, + ) { + Ok(Some(ContinuationConstraint::Pinned)) => true, + Ok(Some(ContinuationConstraint::UnknownResponseId)) => { + send_responses_websocket_error_with_param( + client_socket, + 400, + "invalid_request_error", + "previous_response_not_found", + "The previous response is unavailable on this authenticated WebSocket connection", + "previous_response_id", + ) + .await; + return RelayDisposition::Continue; } - Ok(None) => client_event, + Ok(Some(ContinuationConstraint::UpstreamUnavailable)) => { + send_gateway_error_with_status( + client_socket, + 503, + "responses_continuation_provider_unavailable", + "The bound provider connection is unavailable for this continuation", + ) + .await; + return RelayDisposition::Continue; + } + Ok(Some(ContinuationConstraint::ModelChangeUnsupported)) => { + send_gateway_error_with_status( + client_socket, + 409, + "responses_continuation_model_change_unsupported", + "A continuation cannot change models on the bound provider connection", + ) + .await; + return RelayDisposition::Continue; + } + Ok(None) => false, + Err(code) => { + send_gateway_error( + client_socket, + code, + "response.create.model must be a non-empty string", + ) + .await; + return RelayDisposition::Continue; + } + }; + + // Static Responses Lite configuration belongs to the response + // chain, not to a redaction TTL bucket. Compare and strip it from + // continuation input while it is still the raw client value; only + // the incremental event is redacted below. Independent turns keep + // a raw hash so a later sentinel rotation cannot look like a tools + // or instructions change. + let raw_responses_lite_static_config = (!is_pinned_continuation) + .then(|| ResponsesLiteStaticConfig::from_response_create(&client_event)); + let client_event = if is_pinned_continuation { + if let Some(static_config) = bound.responses_lite_static_config.as_ref() { + match prepare_responses_lite_continuation(&client_event, static_config) { + Ok(event) => event, + Err("responses_lite_continuation_static_config_changed") => { + send_gateway_error_with_status( + client_socket, + 409, + "responses_lite_continuation_static_config_changed", + "Responses Lite tools or instructions changed; start a new response without previous_response_id", + ) + .await; + return RelayDisposition::Continue; + } + Err(code) => { + send_gateway_error( + client_socket, + code, + "Gateway could not validate the Responses Lite continuation", + ) + .await; + return RelayDisposition::Continue; + } + } + } else { + client_event + } + } else { + client_event + }; + + // 这一轮的 planning Parts 只构造一次（它携带 per-turn 的 + // RedactionSessionSlot），并且客户端事件也只在这里脱敏一次。 + // continuation 已经去掉继承的 static prefix，所以只 mask 新增 input； + // 后续 re-plan / continuation / 配额重试都只看脱敏后的增量事件。 + let redacted_client_event = + redact_responses_websocket_client_event_with_reasoning_replay_policy( + state, + &planning_parts, + &turn_control.decision, + &client_event, + bound.body_normalization.reasoning_replay_policy(), + ) + .await; + let (client_event, turn_redaction_session) = match redacted_client_event { + Ok(Some(redaction)) => (redaction.client_event, Some(redaction.session)), + Ok(None) => (client_event, None), Err(error) => { warn!( event_name = "responses_websocket_followup_redaction_failed", @@ -301,53 +445,18 @@ pub(super) async fn forward_client_message( return RelayDisposition::Close; } }; - match continuation_constraint( - &client_event, - &bound.client_model, - bound.upstream.is_some(), - ) { - Ok(Some(ContinuationConstraint::Pinned)) => { - return forward_pinned_continuation( - bound, - client_socket, - state, - context, - planning_parts, - client_event, - turn_control, - ) - .await; - } - Ok(Some(ContinuationConstraint::UpstreamUnavailable)) => { - send_gateway_error_with_status( - client_socket, - 503, - "responses_continuation_provider_unavailable", - "The bound provider connection is unavailable for this continuation", - ) - .await; - return RelayDisposition::Continue; - } - Ok(Some(ContinuationConstraint::ModelChangeUnsupported)) => { - send_gateway_error_with_status( - client_socket, - 409, - "responses_continuation_model_change_unsupported", - "A continuation cannot change models on the bound provider connection", - ) - .await; - return RelayDisposition::Continue; - } - Ok(None) => {} - Err(code) => { - send_gateway_error( - client_socket, - code, - "response.create.model must be a non-empty string", - ) - .await; - return RelayDisposition::Continue; - } + if is_pinned_continuation { + return forward_pinned_continuation( + bound, + client_socket, + state, + context, + planning_parts, + client_event, + turn_control, + turn_redaction_session, + ) + .await; } forward_replanned_response_create( bound, @@ -358,6 +467,9 @@ pub(super) async fn forward_client_message( client_event, requested_model, turn_control, + raw_responses_lite_static_config + .expect("independent turns always retain their raw static config"), + turn_redaction_session, ) .await } @@ -402,6 +514,7 @@ async fn forward_pinned_continuation( planning_parts: http::request::Parts, client_event: Value, turn_control: ResponsesWebSocketTurnControl, + turn_redaction_session: Option, ) -> RelayDisposition { let Some(pinned_candidate) = ResponsesWebSocketPinnedCandidate::from_decision(&bound.decision_template) @@ -426,8 +539,8 @@ async fn forward_pinned_continuation( // The public protocol allows follow-ups to omit `model`. The planner still // needs the effective public model to enumerate the pinned mapping; this - // injected copy never replaces the opaque client event kept for audit and - // protocol-field restoration. + // injected copy never replaces the already de-duplicated, redacted client + // event kept for audit and protocol-field restoration. let planning_event = match pinned_continuation_planning_event(&client_event, bound.client_model.as_str()) { Ok(event) => event, @@ -497,6 +610,25 @@ async fn forward_pinned_continuation( let adapter = resolve_responses_websocket_adapter(planned.adapter); let normalization = planned.normalization; let decision = planned.execution; + let bound_uses_responses_lite = bound.responses_lite_static_config.is_some(); + let planned_uses_responses_lite = + planned_request_uses_codex_responses_lite(&decision, &normalization); + if bound_uses_responses_lite != planned_uses_responses_lite + || (bound_uses_responses_lite + && !bound + .body_normalization + .has_same_responses_lite_static_contract(&normalization)) + { + planned_lease.release().await; + send_gateway_error_with_status( + client_socket, + 409, + "responses_lite_continuation_contract_changed", + "The Responses Lite provider contract changed; start a new response without previous_response_id", + ) + .await; + return RelayDisposition::Continue; + } let planned_provider_model = provider_model_from_decision(&decision); let reuses_bound_upstream = decision_reuses_bound_upstream(bound, adapter, &decision); let provider_model_changed = @@ -530,10 +662,12 @@ async fn forward_pinned_continuation( } let provider_event = - match planned_response_create_event(&decision, &client_event).and_then(|event| { - serde_json::from_str::(&event) - .map_err(|_| "response_create_serialization_failed") - }) { + match planned_response_create_event(&decision, &normalization, &client_event).and_then( + |event| { + serde_json::from_str::(&event) + .map_err(|_| "response_create_serialization_failed") + }, + ) { Ok(event) => event, Err(code) => { planned_lease.release().await; @@ -624,11 +758,16 @@ async fn forward_pinned_continuation( turn.mark_upstream_request_sent(); turn.set_provider_response_headers(bound.upstream_response_headers.clone()); + if let Some(session) = turn_redaction_session { + bound.redaction_restorer.register(session); + } bound.adapter = adapter; bound.decision_template = decision; bound.body_normalization = normalization; bound.turn_state.begin( - LogicalTurn::new(client_event, turn_index, logical_turn_id).with_turn_control(turn_control), + LogicalTurn::new(client_event, turn_index, logical_turn_id) + .with_provider_store(provider_event.get("store") == Some(&Value::Bool(true))) + .with_turn_control(turn_control), turn, ); bound.next_turn_index = bound.next_turn_index.saturating_add(1); @@ -662,6 +801,8 @@ async fn forward_replanned_response_create( client_event: Value, requested_model: String, turn_control: ResponsesWebSocketTurnControl, + raw_responses_lite_static_config: ResponsesLiteStaticConfig, + turn_redaction_session: Option, ) -> RelayDisposition { let turn_request_id = Uuid::new_v4().to_string(); let logical_turn_id = Uuid::new_v4().to_string(); @@ -726,10 +867,12 @@ async fn forward_replanned_response_create( let decision = planned.execution; let reuses_bound_upstream = decision_reuses_bound_upstream(bound, adapter, &decision); let provider_event = - match planned_response_create_event(&decision, &client_event).and_then(|event| { - serde_json::from_str::(&event) - .map_err(|_| "response_create_serialization_failed") - }) { + match planned_response_create_event(&decision, &normalization, &client_event).and_then( + |event| { + serde_json::from_str::(&event) + .map_err(|_| "response_create_serialization_failed") + }, + ) { Ok(event) => event, Err(code) => { planned_lease.release().await; @@ -826,18 +969,30 @@ async fn forward_replanned_response_create( return RelayDisposition::UpstreamError("responses_websocket_send_failed"); } + // A response.create without previous_response_id starts a new chain. + // IDs from the preceding chain must not be accepted merely because + // the planner can reuse the same physical provider socket. + bound.continuation_response_ids.clear(); + bound + .redaction_restorer + .start_new_chain(turn_redaction_session); turn.mark_upstream_request_sent(); turn.set_provider_response_headers(bound.upstream_response_headers.clone()); let provider_model = provider_model_from_decision(&decision).unwrap_or_else(|| bound.provider_model.clone()); let previous_client_model = std::mem::replace(&mut bound.client_model, requested_model); let previous_provider_model = std::mem::replace(&mut bound.provider_model, provider_model); + let uses_responses_lite = + planned_request_uses_codex_responses_lite(&decision, &normalization); bound.decision_template = decision; // The re-plan keeps this upstream but resolved a new model, so later // continuations must normalize against the new plan, not the old one. + bound.responses_lite_static_config = + uses_responses_lite.then_some(raw_responses_lite_static_config); bound.body_normalization = normalization; bound.turn_state.begin( LogicalTurn::new(client_event.clone(), turn_index, logical_turn_id.clone()) + .with_provider_store(provider_event.get("store") == Some(&Value::Bool(true))) .with_turn_control(turn_control), turn, ); @@ -891,6 +1046,9 @@ async fn forward_replanned_response_create( return RelayDisposition::Continue; } }; + if replacement.responses_lite_static_config.is_some() { + replacement.responses_lite_static_config = Some(raw_responses_lite_static_config); + } turn.mark_upstream_request_sent(); turn.set_provider_response_headers(replacement.upstream_response_headers.clone()); @@ -903,14 +1061,23 @@ async fn forward_replanned_response_create( if let Some(mut previous_upstream) = bound.upstream.replace(replacement_upstream) { close_upstream_socket(&mut previous_upstream, None).await; } + // Provider connection-local response state cannot survive a physical + // rebind, and this request starts an independent chain in any case. + bound.continuation_response_ids.clear(); + bound + .redaction_restorer + .start_new_chain(turn_redaction_session); bound.adapter = replacement.adapter; bound.client_model = replacement.client_model; bound.provider_model = replacement.provider_model; bound.decision_template = replacement.decision_template; bound.body_normalization = replacement.body_normalization; + bound.responses_lite_static_config = replacement.responses_lite_static_config; bound.binding_identity = replacement.binding_identity; bound.turn_state.begin( - LogicalTurn::new(client_event, turn_index, logical_turn_id).with_turn_control(turn_control), + LogicalTurn::new(client_event, turn_index, logical_turn_id) + .with_provider_store(provider_event.get("store") == Some(&Value::Bool(true))) + .with_turn_control(turn_control), turn, ); bound.next_turn_index = bound.next_turn_index.saturating_add(1); @@ -965,17 +1132,66 @@ mod tests { #[test] fn invalid_client_text_does_not_poison_the_next_response_create() { assert_eq!( - parse_response_create_event("not-json"), - Err("invalid_response_create") + parse_response_create_event("not-json") + .expect_err("invalid JSON") + .code, + "invalid_response_create" ); assert_eq!( - parse_response_create_event("[]"), - Err("invalid_response_create") + parse_response_create_event("[]") + .expect_err("non-object JSON") + .code, + "invalid_response_create" ); assert_eq!( - parse_response_create_event(r#"{"type":"response.cancel"}"#), - Err("expected_response_create") + parse_response_create_event(r#"{"type":"response.cancel"}"#) + .expect_err("wrong event type") + .code, + "expected_response_create" ); + for invalid_previous_response_id in [ + r#"{"type":"response.create","previous_response_id":""}"#, + r#"{"type":"response.create","previous_response_id":42}"#, + r#"{"type":"response.create","previous_response_id":{"id":"resp_1"}}"#, + ] { + assert_eq!( + parse_response_create_event(invalid_previous_response_id) + .expect_err("invalid previous_response_id") + .code, + "invalid_response_create_previous_response_id" + ); + } + let invalid_previous_on_named_lane = parse_response_create_event( + r#"{"type":"response.create","stream_id":"main","previous_response_id":""}"#, + ) + .expect_err("invalid previous_response_id must retain a valid lane identity"); + assert_eq!( + invalid_previous_on_named_lane.code, + "invalid_response_create_previous_response_id" + ); + assert_eq!( + invalid_previous_on_named_lane.stream_id.as_deref(), + Some("main") + ); + let named_stream = parse_response_create_event( + r#"{"type":"response.create","stream_id":"main-lane_1.test"}"#, + ) + .expect_err("named streams are not implemented"); + assert_eq!( + named_stream.code, + "responses_websocket_named_stream_unsupported" + ); + assert_eq!(named_stream.stream_id.as_deref(), Some("main-lane_1.test")); + for invalid_stream_id in [ + r#"{"type":"response.create","stream_id":null}"#, + r#"{"type":"response.create","stream_id":""}"#, + r#"{"type":"response.create","stream_id":"not/a/lane"}"#, + ] { + let error = parse_response_create_event(invalid_stream_id) + .expect_err("invalid stream_id must be rejected"); + assert_eq!(error.code, "invalid_response_create_stream_id"); + assert_eq!(error.stream_id, None); + } let valid = parse_response_create_event( r#"{"type":"response.create","model":"gpt-test","store":true}"#, @@ -993,11 +1209,11 @@ mod tests { }); assert_eq!( - continuation_constraint(&continuation, "gpt-current", false), + continuation_constraint(&continuation, "gpt-current", false, true), Ok(Some(ContinuationConstraint::UpstreamUnavailable)) ); assert_eq!( - continuation_constraint(&continuation, "gpt-current", true), + continuation_constraint(&continuation, "gpt-current", true, true), Ok(Some(ContinuationConstraint::Pinned)) ); assert_eq!( @@ -1005,6 +1221,7 @@ mod tests { &json!({"type": "response.create", "model": "gpt-current"}), "gpt-current", false, + false, ), Ok(None) ); @@ -1019,11 +1236,25 @@ mod tests { }); assert_eq!( - continuation_constraint(&continuation, "gpt-current", true), + continuation_constraint(&continuation, "gpt-current", true, true), Ok(Some(ContinuationConstraint::ModelChangeUnsupported)) ); } + #[test] + fn unknown_same_socket_response_id_never_reaches_the_pinned_provider() { + let continuation = json!({ + "type": "response.create", + "model": "gpt-current", + "previous_response_id": "resp_from_another_principal", + }); + + assert_eq!( + continuation_constraint(&continuation, "gpt-current", true, false), + Ok(Some(ContinuationConstraint::UnknownResponseId)) + ); + } + #[test] fn pinned_planning_uses_the_canonical_bound_model() { let client_event = json!({ diff --git a/apps/aether-gateway/src/handlers/proxy/websocket/responses/connection.rs b/apps/aether-gateway/src/handlers/proxy/websocket/responses/connection.rs index dd380a731..64ca00a2f 100644 --- a/apps/aether-gateway/src/handlers/proxy/websocket/responses/connection.rs +++ b/apps/aether-gateway/src/handlers/proxy/websocket/responses/connection.rs @@ -9,6 +9,9 @@ use wreq::ws::message::Message as WreqWsMessage; use super::adapter::ResponsesWebSocketRelayDirective; use super::client::{adapter_drain_ready, forward_client_message, RelayDisposition}; +use super::continuation::{ + ResponsesWebSocketContinuationRecord, ResponsesWebSocketContinuationRegistry, +}; use super::frame::{encode_opaque_websocket_event, ParsedResponsesWebSocketFrame}; use super::lifecycle::{ await_pending_adapter_observation, finalize_active_turn, queue_turn_finalization, @@ -26,6 +29,7 @@ use super::state::BoundResponsesConnection; use super::turn::{ ResponsesProviderAttempt, ResponsesWebSocketTurnObservation, ResponsesWebSocketTurnOutcome, }; +use super::turn_state::LogicalTurn; use super::upstream::{close_bound_upstream, receive_optional_upstream}; use crate::handlers::proxy::websocket::ingress::WebSocketRequestContext; use crate::handlers::proxy::websocket::session::{ @@ -38,6 +42,7 @@ use crate::handlers::proxy::websocket::transport::{ use crate::AppState; const LOG_TARGET: &str = "aether_gateway::handlers::proxy::responses_ws"; +const CONTINUATION_REGISTRATION_TIMEOUT: Duration = Duration::from_millis(500); /// 写客户端 socket 失败时记录的投递失败原因。刻意不说「客户端在终态前断开」： /// 供应商的终态可能已经到达，只是最后一跳没送出去。 @@ -207,6 +212,55 @@ pub(super) async fn relay_bound_connection( } _ => None, }; + if parsed_upstream_frame + .as_ref() + .is_some_and(ParsedResponsesWebSocketFrame::carries_stream_id) + { + // The session currently owns only the implicit default + // lane. Reject a provider-side named identity before the + // adapter, usage observer, continuation cache, PII + // restorer, or logical-turn state can attribute an + // interleaved event to the sole default-lane attempt. + let policy = fatal_relay_policy( + FatalRelaySignal::UnexpectedUpstreamStreamId, + ); + let frame = parsed_upstream_frame + .as_ref() + .expect("the stream-id guard requires a parsed frame"); + warn!( + event_name = "responses_websocket_unexpected_upstream_stream_id", + log_type = "ops", + transport = WEBSOCKET_LOG_TRANSPORT, + websocket = true, + trace_id = %context.trace_id, + event_type = %frame.event_type_for_log(), + frame_bytes = frame.raw_text().len(), + chunked = frame.is_chunked(), + "gateway rejected a named-lane provider event on a default-lane Responses WebSocket" + ); + finalize_active_turn( + bound, + state, + ResponsesWebSocketTurnOutcome::upstream_receive_failed(), + ) + .await; + send_responses_websocket_error( + client_socket, + policy.status_code, + "server_error", + policy.error_code, + policy.client_message, + ) + .await; + close_bound_upstream(bound).await; + close_client_socket( + client_socket, + policy.close_code, + policy.close_reason, + ) + .await; + break; + } let parsed_upstream_event = parsed_upstream_frame .as_ref() .map(ParsedResponsesWebSocketFrame::event); @@ -300,6 +354,46 @@ pub(super) async fn relay_bound_connection( Some(ResponsesWebSocketTurnObservation::Terminal(outcome)) => Some(outcome), _ => None, }; + if let Some(frame) = parsed_upstream_frame.as_ref() { + if let Some(response_id) = evicted_default_lane_continuation_response_id( + bound.turn_state.logical(), + frame, + ) { + // A 4xx/5xx continuation terminal evicts the referenced + // ID from the provider's implicit default-lane cache. + // Do not keep claiming local ownership and replay it. + bound + .continuation_response_ids + .forget_connection_local(response_id); + } + if let Some(response_id) = connection_local_terminal_response_id( + bound.turn_state.logical(), + frame, + ) { + // Remember every successful response on this physical + // socket, including store=false responses that exist + // only in the provider's connection-local cache. + bound + .continuation_response_ids + .remember_connection_local(response_id); + } + if let Some(registration) = + prepare_persisted_continuation_registration(bound, context, frame) + { + if let Some(response_id) = + register_persisted_continuation_before_terminal_delivery( + state, + context, + registration, + ) + .await + { + bound + .continuation_response_ids + .remember_persisted(response_id.as_str()); + } + } + } if matches!(&upstream_message, WreqWsMessage::Text(_)) && parsed_upstream_frame.is_none() { @@ -605,6 +699,203 @@ pub(super) async fn relay_bound_connection( } } +struct PendingContinuationRegistration { + user_id: String, + api_key_id: String, + response_id: String, + record: ResponsesWebSocketContinuationRecord, +} + +fn prepare_persisted_continuation_registration( + bound: &BoundResponsesConnection, + context: &WebSocketRequestContext, + frame: &ParsedResponsesWebSocketFrame<'_>, +) -> Option { + let logical = bound.turn_state.logical(); + let Some(response_id) = persistable_terminal_response_id(logical, frame).map(str::to_string) + else { + return None; + }; + let logical = logical.expect("a persistable terminal requires an active logical turn"); + let Some(auth_context) = logical + .turn_control + .as_ref() + .and_then(|control| control.decision.auth_context.as_ref()) + else { + warn!( + event_name = "responses_websocket_continuation_registration_skipped", + log_type = "ops", + transport = WEBSOCKET_LOG_TRANSPORT, + websocket = true, + trace_id = %context.trace_id, + reason = "missing_live_auth_context", + "gateway did not register a persisted Responses continuation" + ); + return None; + }; + let Some(pinned_candidate) = + crate::ai_serving::ResponsesWebSocketPinnedCandidate::from_decision( + &bound.decision_template, + ) + else { + warn!( + event_name = "responses_websocket_continuation_registration_skipped", + log_type = "ops", + transport = WEBSOCKET_LOG_TRANSPORT, + websocket = true, + trace_id = %context.trace_id, + reason = "missing_binding_identity", + "gateway did not register a persisted Responses continuation" + ); + return None; + }; + let record = match ResponsesWebSocketContinuationRecord::from_binding( + pinned_candidate, + bound.client_model.as_str(), + bound.provider_model.as_str(), + &bound.binding_identity, + &bound.body_normalization, + bound.redaction_restorer.has_sessions(), + bound.responses_lite_static_config.clone(), + ) { + Ok(record) => record, + Err(error) => { + warn!( + event_name = "responses_websocket_continuation_registration_skipped", + log_type = "ops", + transport = WEBSOCKET_LOG_TRANSPORT, + websocket = true, + trace_id = %context.trace_id, + provider_id = %bound.decision_template.provider_id.as_deref().unwrap_or("-"), + endpoint_id = %bound.decision_template.endpoint_id.as_deref().unwrap_or("-"), + key_id = %bound.decision_template.key_id.as_deref().unwrap_or("-"), + reason = error.kind(), + "gateway could not build a persisted Responses continuation record" + ); + return None; + } + }; + Some(PendingContinuationRegistration { + user_id: auth_context.user_id.clone(), + api_key_id: auth_context.api_key_id.clone(), + response_id, + record, + }) +} + +fn persistable_terminal_response_id<'a>( + logical: Option<&LogicalTurn>, + frame: &'a ParsedResponsesWebSocketFrame<'_>, +) -> Option<&'a str> { + let logical = logical?; + // `provider_store` is derived from the final framed provider event. False + // or absent is ZDR/connection-local and must never create a 24-hour KV + // record, even if the provider happens to return an ID. + if !logical.provider_store { + return None; + } + connection_local_terminal_response_id(Some(logical), frame) +} + +fn connection_local_terminal_response_id<'a>( + logical: Option<&LogicalTurn>, + frame: &'a ParsedResponsesWebSocketFrame<'_>, +) -> Option<&'a str> { + logical?; + frame.continuation_response_id() +} + +fn evicted_default_lane_continuation_response_id<'a>( + logical: Option<&'a LogicalTurn>, + frame: &ParsedResponsesWebSocketFrame<'_>, +) -> Option<&'a str> { + let logical = logical?; + let terminal = frame.terminal()?; + if terminal.status_code < 400 || terminal.cancelled { + return None; + } + logical + .client_event + .get("previous_response_id") + .and_then(Value::as_str) + .filter(|response_id| !response_id.trim().is_empty()) +} + +async fn register_persisted_continuation_before_terminal_delivery( + state: &AppState, + context: &WebSocketRequestContext, + registration: PendingContinuationRegistration, +) -> Option { + let PendingContinuationRegistration { + user_id, + api_key_id, + response_id, + record, + } = registration; + let registry = ResponsesWebSocketContinuationRegistry::new(state.runtime_state.as_ref()); + match tokio::time::timeout( + CONTINUATION_REGISTRATION_TIMEOUT, + registry.register( + user_id.as_str(), + api_key_id.as_str(), + response_id.as_str(), + &record, + ), + ) + .await + { + Ok(Ok(())) => { + debug!( + event_name = "responses_websocket_continuation_registered", + log_type = "event", + transport = WEBSOCKET_LOG_TRANSPORT, + websocket = true, + trace_id = %context.trace_id, + user_id = %user_id, + api_key_id = %api_key_id, + provider_id = %record.pinned_candidate().provider_id(), + endpoint_id = %record.pinned_candidate().endpoint_id(), + key_id = %record.pinned_candidate().key_id(), + client_model = %record.client_model(), + provider_model = %record.provider_model(), + "gateway registered a persisted Responses continuation before terminal delivery" + ); + Some(response_id) + } + Ok(Err(error)) => { + warn!( + event_name = "responses_websocket_continuation_registration_failed", + log_type = "ops", + transport = WEBSOCKET_LOG_TRANSPORT, + websocket = true, + trace_id = %context.trace_id, + provider_id = %record.pinned_candidate().provider_id(), + endpoint_id = %record.pinned_candidate().endpoint_id(), + key_id = %record.pinned_candidate().key_id(), + reason = error.kind(), + "gateway failed to register a persisted Responses continuation" + ); + None + } + Err(_) => { + warn!( + event_name = "responses_websocket_continuation_registration_failed", + log_type = "ops", + transport = WEBSOCKET_LOG_TRANSPORT, + websocket = true, + trace_id = %context.trace_id, + provider_id = %record.pinned_candidate().provider_id(), + endpoint_id = %record.pinned_candidate().endpoint_id(), + key_id = %record.pinned_candidate().key_id(), + reason = "timeout", + timeout_ms = CONTINUATION_REGISTRATION_TIMEOUT.as_millis() as u64, + "gateway timed out registering a persisted Responses continuation" + ); + None + } + } +} + pub(super) async fn wait_for_connection_permit_loss( permit: Option<&aether_runtime::AdmissionPermit>, ) { @@ -621,3 +912,79 @@ pub(super) async fn wait_for_connection_permit_loss( } } } + +#[cfg(test)] +mod tests { + use serde_json::json; + + use super::{ + connection_local_terminal_response_id, evicted_default_lane_continuation_response_id, + persistable_terminal_response_id, LogicalTurn, ParsedResponsesWebSocketFrame, + }; + + #[test] + fn cross_connection_registration_requires_explicit_store_true_and_a_success_terminal() { + let completed = ParsedResponsesWebSocketFrame::parse( + r#"{"type":"response.completed","response":{"id":"resp_persisted"}}"#, + ) + .expect("valid completed event"); + let failed = ParsedResponsesWebSocketFrame::parse( + r#"{"type":"response.failed","response":{"id":"resp_failed"}}"#, + ) + .expect("valid failed event"); + + let omitted_or_false = LogicalTurn::new( + json!({"type": "response.create", "store": false}), + 1, + "logical-local".to_string(), + ); + assert_eq!( + persistable_terminal_response_id(Some(&omitted_or_false), &completed), + None, + "store=false or an omitted provider-side store must remain connection-local" + ); + assert_eq!(persistable_terminal_response_id(None, &completed), None); + assert_eq!( + connection_local_terminal_response_id(Some(&omitted_or_false), &completed), + Some("resp_persisted"), + "store=false continuations remain valid on the same physical socket" + ); + assert_eq!( + connection_local_terminal_response_id(None, &completed), + None + ); + + let persisted = LogicalTurn::new( + json!({"type": "response.create", "store": true}), + 1, + "logical-persisted".to_string(), + ) + .with_provider_store(true); + assert_eq!( + persistable_terminal_response_id(Some(&persisted), &completed), + Some("resp_persisted") + ); + assert_eq!( + persistable_terminal_response_id(Some(&persisted), &failed), + None, + "a failed terminal must never establish cross-connection ownership" + ); + + let continuation = LogicalTurn::new( + json!({ + "type": "response.create", + "previous_response_id": "resp_parent" + }), + 2, + "logical-continuation".to_string(), + ); + assert_eq!( + evicted_default_lane_continuation_response_id(Some(&continuation), &failed), + Some("resp_parent") + ); + assert_eq!( + evicted_default_lane_continuation_response_id(Some(&continuation), &completed), + None + ); + } +} diff --git a/apps/aether-gateway/src/handlers/proxy/websocket/responses/continuation.rs b/apps/aether-gateway/src/handlers/proxy/websocket/responses/continuation.rs new file mode 100644 index 000000000..d8de44134 --- /dev/null +++ b/apps/aether-gateway/src/handlers/proxy/websocket/responses/continuation.rs @@ -0,0 +1,781 @@ +//! Short-lived ownership registry for Responses WebSocket continuations. +//! +//! OpenAI response IDs are opaque bearer-like references to provider state. A +//! response created on one physical provider binding must never be resumed on +//! a scheduler-selected replacement. This registry stores only non-secret +//! routing metadata and one-way contract fingerprints. Raw response IDs, +//! downstream credentials and upstream credentials are never persisted. + +use std::time::{Duration, SystemTime, UNIX_EPOCH}; + +use aether_runtime_state::{RuntimeLockLease, RuntimeState}; +use serde::{Deserialize, Serialize}; +use sha2::{Digest, Sha256}; + +use super::binding::UpstreamBindingIdentity; +use super::request::{ResponsesLiteStaticConfig, MAX_RESPONSES_WEBSOCKET_RESPONSE_ID_BYTES}; +use crate::ai_serving::{ResponsesWebSocketBodyNormalization, ResponsesWebSocketPinnedCandidate}; +use crate::orchestration::ResponsesWebSocketAdapter; + +const CONTINUATION_RECORD_SCHEMA_VERSION: u16 = 1; +const CONTINUATION_KEY_PREFIX: &str = "responses_ws:continuation:v1:"; +const CONTINUATION_KEY_DOMAIN: &[u8] = b"aether-responses-websocket-continuation-key-v1"; +const CONTINUATION_INDEX_PREFIX: &str = "responses_ws:continuation_index:v1:"; +const CONTINUATION_INDEX_DOMAIN: &[u8] = b"aether-responses-websocket-continuation-index-v1"; +const CONTINUATION_LOCK_PREFIX: &str = "responses_ws:continuation_lock:v1:"; +const CONTINUATION_RECORD_TTL: Duration = Duration::from_secs(24 * 60 * 60); +const CONTINUATION_INDEX_LOCK_TTL: Duration = Duration::from_secs(2); +const CONTINUATION_INDEX_LOCK_ACQUIRE_TIMEOUT: Duration = Duration::from_millis(250); +const CONTINUATION_INDEX_LOCK_INITIAL_RETRY_DELAY: Duration = Duration::from_millis(5); +const CONTINUATION_INDEX_LOCK_MAX_RETRY_DELAY: Duration = Duration::from_millis(50); +const CONTINUATION_INDEX_LOCK_OWNER: &str = "responses_ws_continuation_registry"; +const MAX_CONTINUATION_RECORDS_PER_PRINCIPAL: usize = 1_024; +const MAX_SERIALIZED_CONTINUATION_RECORD_BYTES: usize = 16 * 1024; +const MAX_CONTINUATION_PRINCIPAL_BYTES: usize = 256; +const MAX_CONTINUATION_RECORD_ID_BYTES: usize = 256; + +#[derive(Debug, thiserror::Error)] +pub(super) enum ResponsesWebSocketContinuationRegistryError { + #[error("invalid Responses WebSocket continuation identity: {0}")] + InvalidIdentity(&'static str), + #[error("invalid Responses WebSocket continuation record: {0}")] + InvalidRecord(&'static str), + #[error("Responses WebSocket continuation registry serialization failed")] + Serialization(#[source] serde_json::Error), + #[error("Responses WebSocket continuation registry contains a corrupt record")] + CorruptRecord(#[source] serde_json::Error), + #[error("Responses WebSocket continuation registry record is too large")] + RecordTooLarge, + #[error("Responses WebSocket continuation registry ownership conflict")] + OwnershipConflict, + #[error("Responses WebSocket continuation registry capacity lock is busy")] + CapacityLockBusy, + #[error("Responses WebSocket continuation registry storage is unavailable")] + Storage(#[source] aether_runtime_state::DataLayerError), +} + +impl ResponsesWebSocketContinuationRegistryError { + 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::Storage(_) => "storage_unavailable", + } + } +} + +/// Non-secret metadata required to prove ownership of a persisted response. +/// +/// The record deliberately contains no raw response ID. Its RuntimeState key +/// is derived from the live authenticated principal plus a SHA-256 digest of +/// the opaque response ID. +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +#[serde(deny_unknown_fields)] +pub(super) struct ResponsesWebSocketContinuationRecord { + schema_version: u16, + pinned_candidate: ResponsesWebSocketPinnedCandidate, + client_model: String, + provider_model: String, + adapter: ResponsesWebSocketAdapter, + binding_fingerprint: [u8; 32], + normalization_fingerprint: [u8; 32], + /// Server-derived replay contract for a provider whose id-less reasoning + /// state must remain byte-identical. This is trusted only because the + /// record is created after binding an authenticated provider candidate; + /// request JSON can never set it. + #[serde(default)] + deepseek_opaque_reasoning_replay: bool, + /// A prior turn stored PII sentinels whose restore mapping exists only on + /// the original downstream socket. Such a chain cannot safely resume on a + /// new socket without leaking sentinels, so lookup succeeds but bootstrap + /// rejects it before contacting the provider. + #[serde(default)] + has_connection_local_redaction: bool, + #[serde(default, skip_serializing_if = "Option::is_none")] + responses_lite_static_config: Option, +} + +impl ResponsesWebSocketContinuationRecord { + pub(super) fn from_binding( + pinned_candidate: ResponsesWebSocketPinnedCandidate, + client_model: &str, + provider_model: &str, + binding: &UpstreamBindingIdentity, + normalization: &ResponsesWebSocketBodyNormalization, + has_connection_local_redaction: bool, + responses_lite_static_config: Option, + ) -> Result { + let record = Self { + schema_version: CONTINUATION_RECORD_SCHEMA_VERSION, + pinned_candidate, + client_model: client_model.to_string(), + provider_model: provider_model.to_string(), + adapter: binding.adapter_kind(), + binding_fingerprint: binding.continuation_fingerprint(), + normalization_fingerprint: normalization.continuation_fingerprint(), + deepseek_opaque_reasoning_replay: matches!( + normalization.reasoning_replay_policy(), + crate::ai_serving::OpenAiResponsesReasoningReplayPolicy::DeepSeekOpaque + ), + has_connection_local_redaction, + responses_lite_static_config, + }; + record.validate()?; + Ok(record) + } + + 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 provider_model(&self) -> &str { + self.provider_model.as_str() + } + + pub(super) fn adapter(&self) -> ResponsesWebSocketAdapter { + self.adapter + } + + pub(super) fn responses_lite_static_config(&self) -> Option<&ResponsesLiteStaticConfig> { + self.responses_lite_static_config.as_ref() + } + + pub(super) fn has_connection_local_redaction(&self) -> bool { + self.has_connection_local_redaction + } + + pub(super) fn reasoning_replay_policy( + &self, + ) -> crate::ai_serving::OpenAiResponsesReasoningReplayPolicy { + if self.deepseek_opaque_reasoning_replay { + crate::ai_serving::OpenAiResponsesReasoningReplayPolicy::DeepSeekOpaque + } else { + crate::ai_serving::OpenAiResponsesReasoningReplayPolicy::OpenAiItemIds + } + } + + pub(super) fn matches_contract( + &self, + binding: &UpstreamBindingIdentity, + normalization: &ResponsesWebSocketBodyNormalization, + ) -> bool { + self.adapter == binding.adapter_kind() + && self.binding_fingerprint == binding.continuation_fingerprint() + && self.normalization_fingerprint == normalization.continuation_fingerprint() + } + + fn validate(&self) -> Result<(), ResponsesWebSocketContinuationRegistryError> { + if self.schema_version != CONTINUATION_RECORD_SCHEMA_VERSION { + return Err(ResponsesWebSocketContinuationRegistryError::InvalidRecord( + "unsupported_schema_version", + )); + } + validate_record_identifier(self.pinned_candidate.provider_id(), "invalid_provider_id")?; + validate_record_identifier(self.pinned_candidate.endpoint_id(), "invalid_endpoint_id")?; + validate_record_identifier(self.pinned_candidate.key_id(), "invalid_key_id")?; + validate_record_identifier(self.client_model.as_str(), "invalid_client_model")?; + validate_record_identifier(self.provider_model.as_str(), "invalid_provider_model")?; + Ok(()) + } +} + +pub(super) struct ResponsesWebSocketContinuationRegistry<'a> { + runtime_state: &'a RuntimeState, + ttl: Duration, + max_records_per_principal: usize, +} + +impl<'a> ResponsesWebSocketContinuationRegistry<'a> { + pub(super) fn new(runtime_state: &'a RuntimeState) -> Self { + Self { + runtime_state, + ttl: CONTINUATION_RECORD_TTL, + max_records_per_principal: MAX_CONTINUATION_RECORDS_PER_PRINCIPAL, + } + } + + #[cfg(test)] + fn with_limits(runtime_state: &'a RuntimeState, 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, + response_id: &str, + record: &ResponsesWebSocketContinuationRecord, + ) -> Result<(), ResponsesWebSocketContinuationRegistryError> { + let key = continuation_registry_key(user_id, api_key_id, response_id)?; + let index_key = continuation_registry_index_key(user_id, api_key_id)?; + let lock_key = continuation_registry_lock_key(user_id, api_key_id)?; + record.validate()?; + let serialized = serde_json::to_string(record) + .map_err(ResponsesWebSocketContinuationRegistryError::Serialization)?; + if serialized.len() > MAX_SERIALIZED_CONTINUATION_RECORD_BYTES { + return Err(ResponsesWebSocketContinuationRegistryError::RecordTooLarge); + } + let lease = self.acquire_capacity_lock(&lock_key).await?; + let result = self + .register_under_capacity_lock(&key, &index_key, serialized, record) + .await; + let release_result = self.runtime_state.lock_release(&lease).await; + match (result, release_result) { + (Err(error), _) => Err(error), + (Ok(()), Ok(_)) => Ok(()), + (Ok(()), Err(error)) => { + Err(ResponsesWebSocketContinuationRegistryError::Storage(error)) + } + } + } + + async fn acquire_capacity_lock( + &self, + lock_key: &str, + ) -> Result { + let deadline = tokio::time::Instant::now() + CONTINUATION_INDEX_LOCK_ACQUIRE_TIMEOUT; + let mut retry_delay = CONTINUATION_INDEX_LOCK_INITIAL_RETRY_DELAY; + loop { + if let Some(lease) = self + .runtime_state + .lock_try_acquire( + lock_key, + CONTINUATION_INDEX_LOCK_OWNER, + CONTINUATION_INDEX_LOCK_TTL, + ) + .await + .map_err(ResponsesWebSocketContinuationRegistryError::Storage)? + { + return Ok(lease); + } + + let now = tokio::time::Instant::now(); + if now >= deadline { + return Err(ResponsesWebSocketContinuationRegistryError::CapacityLockBusy); + } + tokio::time::sleep(retry_delay.min(deadline.saturating_duration_since(now))).await; + retry_delay = retry_delay + .saturating_mul(2) + .min(CONTINUATION_INDEX_LOCK_MAX_RETRY_DELAY); + } + } + + async fn register_under_capacity_lock( + &self, + key: &str, + index_key: &str, + serialized: String, + record: &ResponsesWebSocketContinuationRecord, + ) -> Result<(), ResponsesWebSocketContinuationRegistryError> { + if let Some(existing) = self + .runtime_state + .kv_get(key) + .await + .map_err(ResponsesWebSocketContinuationRegistryError::Storage)? + { + let existing = serde_json::from_str::(&existing) + .map_err(ResponsesWebSocketContinuationRegistryError::CorruptRecord)?; + if existing != *record { + return Err(ResponsesWebSocketContinuationRegistryError::OwnershipConflict); + } + } + self.runtime_state + .kv_set(key, serialized, Some(self.ttl)) + .await + .map_err(ResponsesWebSocketContinuationRegistryError::Storage)?; + let score = SystemTime::now() + .duration_since(UNIX_EPOCH) + .unwrap_or_default() + .as_millis() as f64; + if let Err(error) = self.runtime_state.score_set(index_key, key, score).await { + let _ = self.runtime_state.kv_delete(key).await; + return Err(ResponsesWebSocketContinuationRegistryError::Storage(error)); + } + if let Err(error) = self.runtime_state.key_expire(index_key, self.ttl).await { + let _ = self.runtime_state.score_remove(index_key, key).await; + let _ = self.runtime_state.kv_delete(key).await; + return Err(ResponsesWebSocketContinuationRegistryError::Storage(error)); + } + let members = self + .runtime_state + .score_range_by_min(index_key, f64::NEG_INFINITY) + .await + .map_err(ResponsesWebSocketContinuationRegistryError::Storage)?; + let overflow = members.len().saturating_sub(self.max_records_per_principal); + for oldest_key in members.into_iter().take(overflow) { + self.runtime_state + .kv_delete(&oldest_key) + .await + .map_err(ResponsesWebSocketContinuationRegistryError::Storage)?; + self.runtime_state + .score_remove(index_key, &oldest_key) + .await + .map_err(ResponsesWebSocketContinuationRegistryError::Storage)?; + } + Ok(()) + } + + pub(super) async fn lookup( + &self, + user_id: &str, + api_key_id: &str, + response_id: &str, + ) -> Result< + Option, + ResponsesWebSocketContinuationRegistryError, + > { + let key = continuation_registry_key(user_id, api_key_id, response_id)?; + let Some(serialized) = self + .runtime_state + .kv_get(&key) + .await + .map_err(ResponsesWebSocketContinuationRegistryError::Storage)? + else { + return Ok(None); + }; + let record = serde_json::from_str::(&serialized) + .map_err(ResponsesWebSocketContinuationRegistryError::CorruptRecord)?; + record.validate()?; + Ok(Some(record)) + } +} + +fn validate_record_identifier( + value: &str, + error: &'static str, +) -> Result<(), ResponsesWebSocketContinuationRegistryError> { + if value.trim().is_empty() || value.len() > MAX_CONTINUATION_RECORD_ID_BYTES { + return Err(ResponsesWebSocketContinuationRegistryError::InvalidRecord( + error, + )); + } + Ok(()) +} + +fn validate_key_component( + value: &str, + max_bytes: usize, + error: &'static str, +) -> Result<(), ResponsesWebSocketContinuationRegistryError> { + if value.is_empty() || value.len() > max_bytes { + return Err(ResponsesWebSocketContinuationRegistryError::InvalidIdentity(error)); + } + Ok(()) +} + +fn continuation_registry_key( + user_id: &str, + api_key_id: &str, + response_id: &str, +) -> Result { + validate_key_component(user_id, MAX_CONTINUATION_PRINCIPAL_BYTES, "invalid_user_id")?; + validate_key_component( + api_key_id, + MAX_CONTINUATION_PRINCIPAL_BYTES, + "invalid_api_key_id", + )?; + validate_key_component( + response_id, + MAX_RESPONSES_WEBSOCKET_RESPONSE_ID_BYTES, + "invalid_response_id", + )?; + + let digest = digest_key_components( + CONTINUATION_KEY_DOMAIN, + [ + user_id.as_bytes(), + api_key_id.as_bytes(), + response_id.as_bytes(), + ], + ); + Ok(format!("{CONTINUATION_KEY_PREFIX}{digest}")) +} + +fn continuation_registry_index_key( + user_id: &str, + api_key_id: &str, +) -> Result { + validate_key_component(user_id, MAX_CONTINUATION_PRINCIPAL_BYTES, "invalid_user_id")?; + validate_key_component( + api_key_id, + MAX_CONTINUATION_PRINCIPAL_BYTES, + "invalid_api_key_id", + )?; + let digest = digest_key_components( + CONTINUATION_INDEX_DOMAIN, + [user_id.as_bytes(), api_key_id.as_bytes()], + ); + Ok(format!("{CONTINUATION_INDEX_PREFIX}{digest}")) +} + +fn continuation_registry_lock_key( + user_id: &str, + api_key_id: &str, +) -> Result { + let index = continuation_registry_index_key(user_id, api_key_id)?; + Ok(format!( + "{CONTINUATION_LOCK_PREFIX}{}", + index + .strip_prefix(CONTINUATION_INDEX_PREFIX) + .unwrap_or(index.as_str()) + )) +} + +fn digest_key_components(domain: &[u8], components: [&[u8]; N]) -> 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); + } + format!("{:x}", digest.finalize()) +} + +#[cfg(test)] +mod tests { + use std::time::Duration; + + use aether_runtime_state::{MemoryRuntimeStateConfig, RuntimeState}; + use serde_json::json; + + use super::*; + + const USER_ID: &str = "user-live-0198"; + const API_KEY_ID: &str = "api-key-live-0198"; + const RESPONSE_ID: &str = "resp_opaque_super_secret_reference"; + + fn runtime_state() -> RuntimeState { + RuntimeState::memory(MemoryRuntimeStateConfig::default()) + } + + fn record() -> ResponsesWebSocketContinuationRecord { + ResponsesWebSocketContinuationRecord { + schema_version: CONTINUATION_RECORD_SCHEMA_VERSION, + pinned_candidate: ResponsesWebSocketPinnedCandidate::new( + "provider-1", + "endpoint-1", + "key-1", + ) + .expect("candidate"), + client_model: "public-model".to_string(), + provider_model: "provider-model".to_string(), + adapter: ResponsesWebSocketAdapter::Codex, + binding_fingerprint: [7; 32], + normalization_fingerprint: [9; 32], + deepseek_opaque_reasoning_replay: false, + has_connection_local_redaction: false, + responses_lite_static_config: Some(ResponsesLiteStaticConfig::from_response_create( + &json!({ + "tools": [{"type": "function", "name": "lookup"}], + "instructions": "Do not persist this plaintext" + }), + )), + } + } + + #[tokio::test] + async fn same_principal_gets_the_exact_pinned_candidate_and_other_principals_miss() { + let runtime = runtime_state(); + let registry = ResponsesWebSocketContinuationRegistry::new(&runtime); + let expected = record(); + registry + .register(USER_ID, API_KEY_ID, RESPONSE_ID, &expected) + .await + .expect("register"); + + let found = registry + .lookup(USER_ID, API_KEY_ID, RESPONSE_ID) + .await + .expect("lookup") + .expect("same authenticated principal must find its record"); + assert_eq!(found, expected); + assert_eq!(found.pinned_candidate(), expected.pinned_candidate()); + for (user_id, api_key_id, response_id) in [ + ("other-user", API_KEY_ID, RESPONSE_ID), + (USER_ID, "other-api-key", RESPONSE_ID), + (USER_ID, API_KEY_ID, "resp_other"), + ] { + assert_eq!( + registry + .lookup(user_id, api_key_id, response_id) + .await + .expect("isolated lookup"), + None + ); + } + } + + #[tokio::test] + async fn corrupt_or_unsupported_records_fail_closed() { + let runtime = runtime_state(); + let registry = ResponsesWebSocketContinuationRegistry::new(&runtime); + let key = continuation_registry_key(USER_ID, API_KEY_ID, RESPONSE_ID).expect("key"); + runtime + .kv_set(&key, "not-json", Some(Duration::from_secs(60))) + .await + .expect("seed corrupt record"); + assert!(matches!( + registry.lookup(USER_ID, API_KEY_ID, RESPONSE_ID).await, + Err(ResponsesWebSocketContinuationRegistryError::CorruptRecord( + _ + )) + )); + + let mut unsupported = serde_json::to_value(record()).expect("serialize record"); + unsupported["schema_version"] = json!(CONTINUATION_RECORD_SCHEMA_VERSION + 1); + runtime + .kv_set(&key, unsupported.to_string(), Some(Duration::from_secs(60))) + .await + .expect("seed unsupported record"); + assert!(matches!( + registry.lookup(USER_ID, API_KEY_ID, RESPONSE_ID).await, + Err(ResponsesWebSocketContinuationRegistryError::InvalidRecord( + "unsupported_schema_version" + )) + )); + } + + #[tokio::test] + async fn expired_records_are_not_returned() { + let runtime = runtime_state(); + let registry = ResponsesWebSocketContinuationRegistry::with_limits( + &runtime, + Duration::from_millis(5), + MAX_CONTINUATION_RECORDS_PER_PRINCIPAL, + ); + registry + .register(USER_ID, API_KEY_ID, RESPONSE_ID, &record()) + .await + .expect("register"); + tokio::time::sleep(Duration::from_millis(20)).await; + assert_eq!( + registry + .lookup(USER_ID, API_KEY_ID, RESPONSE_ID) + .await + .expect("lookup"), + None + ); + } + + #[tokio::test] + async fn registry_evicts_the_oldest_record_per_authenticated_principal() { + let runtime = runtime_state(); + let registry = ResponsesWebSocketContinuationRegistry::with_limits( + &runtime, + Duration::from_secs(60), + 2, + ); + for response_id in ["resp_oldest", "resp_middle", "resp_newest"] { + registry + .register(USER_ID, API_KEY_ID, response_id, &record()) + .await + .expect("register"); + tokio::time::sleep(Duration::from_millis(2)).await; + } + assert_eq!( + registry + .lookup(USER_ID, API_KEY_ID, "resp_oldest") + .await + .expect("lookup"), + None + ); + for response_id in ["resp_middle", "resp_newest"] { + assert_eq!( + registry + .lookup(USER_ID, API_KEY_ID, response_id) + .await + .expect("lookup"), + Some(record()) + ); + } + let index_key = continuation_registry_index_key(USER_ID, API_KEY_ID).expect("index key"); + assert_eq!(runtime.score_len(&index_key).await.expect("index len"), 2); + } + + #[tokio::test] + async fn registry_retries_a_briefly_contended_principal_capacity_lock() { + let runtime = runtime_state(); + let registry = ResponsesWebSocketContinuationRegistry::new(&runtime); + let lock_key = continuation_registry_lock_key(USER_ID, API_KEY_ID).expect("lock key"); + let held = runtime + .lock_try_acquire( + &lock_key, + "continuation-registry-contention-test", + CONTINUATION_INDEX_LOCK_TTL, + ) + .await + .expect("acquire test lock") + .expect("test lock should be uncontended"); + + let release_held_lock = async { + tokio::time::sleep(Duration::from_millis(20)).await; + assert!(runtime + .lock_release(&held) + .await + .expect("release test lock")); + }; + let expected = record(); + let (registration, ()) = tokio::join!( + registry.register(USER_ID, API_KEY_ID, RESPONSE_ID, &expected), + release_held_lock, + ); + + registration.expect("registration should retry until the short contention clears"); + assert!(registry + .lookup(USER_ID, API_KEY_ID, RESPONSE_ID) + .await + .expect("lookup") + .is_some()); + } + + #[tokio::test] + async fn same_response_id_cannot_be_rebound_to_a_different_owner_record() { + let runtime = runtime_state(); + let registry = ResponsesWebSocketContinuationRegistry::new(&runtime); + let original = record(); + registry + .register(USER_ID, API_KEY_ID, RESPONSE_ID, &original) + .await + .expect("register"); + let mut replacement = original.clone(); + replacement.provider_model = "different-provider-model".to_string(); + assert!(matches!( + registry + .register(USER_ID, API_KEY_ID, RESPONSE_ID, &replacement) + .await, + Err(ResponsesWebSocketContinuationRegistryError::OwnershipConflict) + )); + assert_eq!( + registry + .lookup(USER_ID, API_KEY_ID, RESPONSE_ID) + .await + .expect("lookup"), + Some(original) + ); + } + + #[test] + fn registry_key_is_length_delimited_hashed_and_contains_no_plaintext() { + let key = continuation_registry_key(USER_ID, API_KEY_ID, RESPONSE_ID).expect("key"); + assert!(key.starts_with(CONTINUATION_KEY_PREFIX)); + assert_eq!(key.len(), CONTINUATION_KEY_PREFIX.len() + 64); + for secret in [USER_ID, API_KEY_ID, RESPONSE_ID, "super_secret_reference"] { + assert!(!key.contains(secret)); + } + let index_key = continuation_registry_index_key(USER_ID, API_KEY_ID).expect("index key"); + for secret in [USER_ID, API_KEY_ID] { + assert!(!index_key.contains(secret)); + } + assert_ne!( + continuation_registry_key("ab", "c", "d").expect("key"), + continuation_registry_key("a", "bc", "d").expect("key") + ); + } + + #[test] + fn invalid_or_oversized_identity_never_produces_a_cache_key() { + for (user_id, api_key_id, response_id) in [ + ("", API_KEY_ID, RESPONSE_ID), + (USER_ID, "", RESPONSE_ID), + (USER_ID, API_KEY_ID, ""), + ] { + assert!(continuation_registry_key(user_id, api_key_id, response_id).is_err()); + } + let oversized = "x".repeat(MAX_RESPONSES_WEBSOCKET_RESPONSE_ID_BYTES + 1); + assert!(continuation_registry_key(USER_ID, API_KEY_ID, &oversized).is_err()); + } + + #[test] + fn serialized_record_contains_only_digests_for_static_and_binding_state() { + let serialized = serde_json::to_string(&record()).expect("serialize"); + for plaintext in [ + RESPONSE_ID, + "Do not persist this plaintext", + "lookup", + "upstream-oauth-token", + ] { + assert!(!serialized.contains(plaintext)); + } + let decoded: ResponsesWebSocketContinuationRecord = + serde_json::from_str(&serialized).expect("deserialize"); + assert_eq!(decoded, record()); + } + + #[test] + fn serialized_record_preserves_only_the_server_derived_reasoning_replay_policy_bit() { + let mut expected = record(); + expected.deepseek_opaque_reasoning_replay = true; + + let serialized = serde_json::to_string(&expected).expect("serialize"); + let decoded: ResponsesWebSocketContinuationRecord = + serde_json::from_str(&serialized).expect("deserialize"); + assert_eq!( + decoded.reasoning_replay_policy(), + crate::ai_serving::OpenAiResponsesReasoningReplayPolicy::DeepSeekOpaque + ); + + let mut legacy = serde_json::to_value(&expected).expect("serialize legacy fixture"); + legacy + .as_object_mut() + .expect("record is an object") + .remove("deepseek_opaque_reasoning_replay"); + let legacy: ResponsesWebSocketContinuationRecord = + serde_json::from_value(legacy).expect("legacy record should remain readable"); + assert_eq!( + legacy.reasoning_replay_policy(), + crate::ai_serving::OpenAiResponsesReasoningReplayPolicy::OpenAiItemIds, + "old records must fail closed instead of inferring trust from request shape" + ); + } + + #[test] + fn record_derives_deepseek_replay_policy_from_the_bound_normalization() { + let decision: crate::ai_serving::AiExecutionDecision = serde_json::from_value(json!({ + "action": "local", + "provider_id": "provider-1", + "endpoint_id": "endpoint-1", + "key_id": "key-1", + "upstream_url": "https://api.deepseek.com/v1/responses", + "provider_request_headers": {} + })) + .expect("minimal decision"); + let adapter = super::super::adapter::resolve_responses_websocket_adapter( + ResponsesWebSocketAdapter::Standard, + ); + let binding = + UpstreamBindingIdentity::from_decision(adapter, &decision).expect("binding identity"); + let normalization = ResponsesWebSocketBodyNormalization::for_tests("deepseek-reasoner") + .with_reasoning_replay_policy_for_tests( + crate::ai_serving::OpenAiResponsesReasoningReplayPolicy::DeepSeekOpaque, + ); + let record = ResponsesWebSocketContinuationRecord::from_binding( + ResponsesWebSocketPinnedCandidate::new("provider-1", "endpoint-1", "key-1") + .expect("pinned candidate"), + "public-model", + "deepseek-reasoner", + &binding, + &normalization, + false, + None, + ) + .expect("continuation record"); + + assert_eq!( + record.reasoning_replay_policy(), + crate::ai_serving::OpenAiResponsesReasoningReplayPolicy::DeepSeekOpaque + ); + } +} diff --git a/apps/aether-gateway/src/handlers/proxy/websocket/responses/frame.rs b/apps/aether-gateway/src/handlers/proxy/websocket/responses/frame.rs index 16e0b1d2f..bc5fbb9d3 100644 --- a/apps/aether-gateway/src/handlers/proxy/websocket/responses/frame.rs +++ b/apps/aether-gateway/src/handlers/proxy/websocket/responses/frame.rs @@ -7,6 +7,8 @@ use serde_json::Value; +use super::request::MAX_RESPONSES_WEBSOCKET_RESPONSE_ID_BYTES; + #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub(super) struct ResponsesWebSocketFrameTerminal { pub(super) status_code: u16, @@ -89,6 +91,23 @@ impl<'a> ParsedResponsesWebSocketFrame<'a> { &self.event } + /// Whether the provider attached a named-lane identity to this frame. + /// + /// Aether's current relay owns only the implicit default lane. OpenAI's + /// named-lane contract places `stream_id` directly on each public event; + /// Codex may additionally batch those events under `chunks`. Inspect both + /// wire positions without looking through arbitrary response payloads so + /// a model-produced field named `stream_id` is not mistaken for transport + /// framing. + pub(super) fn carries_stream_id(&self) -> bool { + self.event.get("stream_id").is_some() + || self + .event + .get("chunks") + .and_then(Value::as_array) + .is_some_and(|chunks| chunks.iter().any(|event| event.get("stream_id").is_some())) + } + pub(super) fn event_type(&self) -> Option<&str> { self.event_type.as_deref() } @@ -109,6 +128,30 @@ impl<'a> ParsedResponsesWebSocketFrame<'a> { self.terminal } + /// Returns the opaque response ID only for a successfully observed, + /// provider-persistable terminal. Failed, cancelled, malformed and + /// oversized terminal IDs must never establish continuation ownership. + pub(super) fn continuation_response_id(&self) -> Option<&str> { + let terminal = self.terminal?; + if terminal.status_code >= 400 || terminal.cancelled { + return None; + } + let event = self.terminal_event.as_ref()?; + if !matches!( + event_type_of(event), + Some("response.completed" | "response.done" | "response.incomplete") + ) { + return None; + } + let response_id = event + .pointer("/response/id") + .or_else(|| event.get("response_id")) + .and_then(Value::as_str)?; + (!response_id.trim().is_empty() + && response_id.len() <= MAX_RESPONSES_WEBSOCKET_RESPONSE_ID_BYTES) + .then_some(response_id) + } + /// Return a bounded label suitable for structured logs. Event payloads /// are never inserted directly into a log field. pub(super) fn event_type_for_log(&self) -> String { @@ -206,7 +249,7 @@ fn responses_incomplete_default_status(event: &Value) -> u16 { fn terminal_for_event(event: &Value) -> Option { match event_type_of(event).unwrap_or_default() { - "response.completed" => Some(ResponsesWebSocketFrameTerminal { + "response.completed" | "response.done" => Some(ResponsesWebSocketFrameTerminal { status_code: websocket_event_status_code(event, 200), cancelled: false, }), @@ -307,6 +350,29 @@ mod tests { assert_eq!(frame.event_type_for_log(), "response.in_progress"); } + #[test] + fn continuation_id_requires_a_successful_persistable_terminal() { + for event_type in ["response.completed", "response.done"] { + let raw = format!(r#"{{"type":"{event_type}","response":{{"id":"resp_persisted"}}}}"#); + let frame = ParsedResponsesWebSocketFrame::parse(&raw).expect("valid terminal"); + assert_eq!(frame.continuation_response_id(), Some("resp_persisted")); + } + let incomplete = ParsedResponsesWebSocketFrame::parse( + r#"{"type":"response.incomplete","response":{"id":"resp_partial","incomplete_details":{"reason":"max_output_tokens"}}}"#, + ) + .expect("valid incomplete terminal"); + assert_eq!(incomplete.continuation_response_id(), Some("resp_partial")); + + for raw in [ + r#"{"type":"response.failed","response":{"id":"resp_failed"}}"#, + r#"{"type":"response.cancelled","response":{"id":"resp_cancelled"}}"#, + r#"{"type":"response.completed","status_code":500,"response":{"id":"resp_error"}}"#, + ] { + let frame = ParsedResponsesWebSocketFrame::parse(raw).expect("valid terminal"); + assert_eq!(frame.continuation_response_id(), None); + } + } + #[test] fn future_response_event_keeps_its_exact_original_text_and_unknown_fields() { let raw = "{ \n \"future_top_level\": {\"nested\": [1, true, null]}, \n \"type\": \"response.future_capability.delta\", \n \"delta\": {\"new_wire_shape\": \"opaque\"}\n}"; @@ -338,6 +404,36 @@ mod tests { assert_eq!(round_trip["response"]["future_usage"]["novel_tokens"], 7); } + #[test] + fn detects_a_top_level_named_lane_identity() { + let frame = ParsedResponsesWebSocketFrame::parse( + r#"{"type":"response.in_progress","stream_id":"planner","response":{"id":"resp_1"}}"#, + ) + .expect("valid named-lane event"); + + assert!(frame.carries_stream_id()); + } + + #[test] + fn detects_a_named_lane_identity_inside_a_batch() { + let frame = ParsedResponsesWebSocketFrame::parse( + r#"{"type":"codex.response.metadata","chunks":[{"type":"codex.rate_limits"},{"type":"response.completed","stream_id":"planner","response":{"id":"resp_1"}}]}"#, + ) + .expect("valid named-lane batch"); + + assert!(frame.carries_stream_id()); + } + + #[test] + fn response_payload_fields_are_not_mistaken_for_lane_framing() { + let frame = ParsedResponsesWebSocketFrame::parse( + r#"{"type":"response.completed","response":{"id":"resp_1","metadata":{"stream_id":"model-owned"}}}"#, + ) + .expect("valid default-lane event"); + + assert!(!frame.carries_stream_id()); + } + #[test] fn classifies_terminal_status_and_cancellation() { let completed = ParsedResponsesWebSocketFrame::parse( 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 2aa903e78..bfe954305 100644 --- a/apps/aether-gateway/src/handlers/proxy/websocket/responses/mod.rs +++ b/apps/aether-gateway/src/handlers/proxy/websocket/responses/mod.rs @@ -12,6 +12,7 @@ mod admission; mod binding; mod client; mod connection; +mod continuation; mod control; mod frame; mod lifecycle; diff --git a/apps/aether-gateway/src/handlers/proxy/websocket/responses/quota.rs b/apps/aether-gateway/src/handlers/proxy/websocket/responses/quota.rs index b43fec3ba..4ad15bcda 100644 --- a/apps/aether-gateway/src/handlers/proxy/websocket/responses/quota.rs +++ b/apps/aether-gateway/src/handlers/proxy/websocket/responses/quota.rs @@ -13,7 +13,9 @@ use super::ownership::{ await_owned_responses_websocket_plan, begin_responses_websocket_turn_with_planned_lease, spawn_owned_responses_websocket_plan, OwnedResponsesWebSocketDecision, }; -use super::request::{build_planning_parts, planned_response_create_event}; +use super::request::{ + build_planning_parts, planned_response_create_event, ResponsesLiteStaticConfig, +}; use super::state::BoundResponsesConnection; use super::turn::{prepare_responses_websocket_turn_decision, ResponsesWebSocketTurnOutcome}; use super::upstream::{bind_responses_upstream, close_bound_upstream}; @@ -102,6 +104,14 @@ pub(super) async fn retry_active_turn_after_quota_exhaustion( context: &WebSocketRequestContext, _previous_settled: PreviousAttemptSettled, ) -> bool { + // `LogicalTurn::client_event` is intentionally redacted before it is + // retained for replay. The binding, however, keeps the hash of the raw + // client-side Responses Lite tools/instructions so a later continuation + // can compare the client's plaintext configuration before redaction. + // Preserve that chain identity across this transparent rebind instead of + // replacing it with the hash `bind_responses_upstream` derives from the + // redacted replay event. + let responses_lite_static_config = bound.responses_lite_static_config.clone(); let Some(active) = bound.turn_state.logical_mut() else { return false; }; @@ -213,12 +223,14 @@ pub(super) async fn retry_active_turn_after_quota_exhaustion( ); return false; } - let provider_event = match planned_response_create_event(&decision, &client_event).and_then( - |event| { - serde_json::from_str::(&event) - .map_err(|_| "response_create_serialization_failed") - }, - ) { + let provider_event = match planned_response_create_event( + &decision, + &normalization, + &client_event, + ) + .and_then(|event| { + serde_json::from_str::(&event).map_err(|_| "response_create_serialization_failed") + }) { Ok(event) => event, Err(code) => { planned_lease.release().await; @@ -234,6 +246,7 @@ pub(super) async fn retry_active_turn_after_quota_exhaustion( return false; } }; + let replacement_provider_store = provider_event.get("store") == Some(&Value::Bool(true)); let turn_decision = prepare_responses_websocket_turn_decision( &decision, turn_request_id, @@ -309,12 +322,20 @@ pub(super) async fn retry_active_turn_after_quota_exhaustion( if let Some(mut previous_upstream) = bound.upstream.replace(replacement_upstream) { close_upstream_socket(&mut previous_upstream, None).await; } + // The replacement socket has no access to response IDs cached only on + // the exhausted physical connection. Continuation turns are never quota + // replayed, so clearing here cannot discard the active turn's parent. + bound.continuation_response_ids.clear(); let previous_key_id = bound.decision_template.key_id.clone(); bound.adapter = replacement.adapter; bound.client_model = replacement.client_model; bound.provider_model = replacement.provider_model; bound.decision_template = replacement.decision_template; bound.body_normalization = replacement.body_normalization; + bound.responses_lite_static_config = responses_lite_static_config_after_rebind( + responses_lite_static_config, + replacement.responses_lite_static_config, + ); bound.binding_identity = replacement.binding_identity; // 同一个 logical turn 的下一个 attempt 就位。状态不符时把 attempt 交回 // drop guard 结算并让调用方走「透明重试失败」分支，不静默丢弃一条已经写了 @@ -323,6 +344,9 @@ pub(super) async fn retry_active_turn_after_quota_exhaustion( drop(orphan); return false; } + if let Some(logical) = bound.turn_state.logical_mut() { + logical.provider_store = replacement_provider_store; + } bound.upstream_response_headers = replacement.upstream_response_headers; bound.pending_adapter_drain = None; debug!( @@ -341,6 +365,13 @@ pub(super) async fn retry_active_turn_after_quota_exhaustion( true } +fn responses_lite_static_config_after_rebind( + previous: Option, + replacement: Option, +) -> Option { + replacement.map(|replacement| previous.unwrap_or(replacement)) +} + pub(super) fn is_usage_limit_error_event(event: &Value) -> bool { let is_error = |value: &Value| { value.get("type").and_then(Value::as_str) == Some("error") @@ -375,3 +406,97 @@ pub(super) fn mark_active_response_retry_unsafe( active.mark_retry_unsafe(reason); } } + +#[cfg(test)] +mod tests { + use serde_json::json; + + use super::responses_lite_static_config_after_rebind; + use crate::handlers::proxy::websocket::responses::request::{ + prepare_responses_lite_continuation, ResponsesLiteStaticConfig, + }; + + #[test] + fn quota_rebind_preserves_the_raw_responses_lite_static_hash() { + let raw = json!({ + "type": "response.create", + "model": "gpt-5.6-sol", + "instructions": "Contact alice@example.com before using the tool", + "tools": [{ + "type": "function", + "name": "lookup", + "description": "Look up alice@example.com", + "parameters": {"type": "object", "properties": {}} + }], + "input": [{"role": "user", "content": "hello"}] + }); + let redacted_replay = json!({ + "type": "response.create", + "model": "gpt-5.6-sol", + "instructions": "Contact before using the tool", + "tools": [{ + "type": "function", + "name": "lookup", + "description": "Look up ", + "parameters": {"type": "object", "properties": {}} + }], + "input": [{"role": "user", "content": "hello"}] + }); + let raw_static_config = ResponsesLiteStaticConfig::from_response_create(&raw); + let redacted_static_config = + ResponsesLiteStaticConfig::from_response_create(&redacted_replay); + assert_ne!(raw_static_config, redacted_static_config); + + let rebound_static_config = responses_lite_static_config_after_rebind( + Some(raw_static_config.clone()), + Some(redacted_static_config), + ) + .expect("the replacement still uses Responses Lite"); + assert_eq!(rebound_static_config, raw_static_config); + + let continuation = json!({ + "type": "response.create", + "model": "gpt-5.6-sol", + "previous_response_id": "resp_after_quota_retry", + "instructions": "Contact alice@example.com before using the tool", + "tools": [{ + "type": "function", + "name": "lookup", + "description": "Look up alice@example.com", + "parameters": {"type": "object", "properties": {}} + }], + "input": [{ + "type": "function_call_output", + "call_id": "call_1", + "output": "ok" + }] + }); + let prepared = prepare_responses_lite_continuation(&continuation, &rebound_static_config) + .expect("an unchanged plaintext continuation must survive a quota retry"); + assert!(prepared.get("tools").is_none()); + assert!(prepared.get("instructions").is_none()); + } + + #[test] + fn quota_rebind_still_tracks_the_replacement_contract() { + let raw = ResponsesLiteStaticConfig::from_response_create(&json!({ + "type": "response.create", + "tools": [{"type": "function", "name": "lookup", "parameters": {}}] + })); + let replacement = ResponsesLiteStaticConfig::from_response_create(&json!({ + "type": "response.create", + "instructions": "replacement" + })); + + assert_eq!( + responses_lite_static_config_after_rebind(Some(raw), None), + None, + "a non-Lite replacement must clear the Lite chain marker" + ); + assert_eq!( + responses_lite_static_config_after_rebind(None, Some(replacement.clone())), + Some(replacement), + "a newly selected Lite contract has no earlier raw hash to preserve" + ); + } +} diff --git a/apps/aether-gateway/src/handlers/proxy/websocket/responses/redaction.rs b/apps/aether-gateway/src/handlers/proxy/websocket/responses/redaction.rs index 7a698dfad..6312c81a0 100644 --- a/apps/aether-gateway/src/handlers/proxy/websocket/responses/redaction.rs +++ b/apps/aether-gateway/src/handlers/proxy/websocket/responses/redaction.rs @@ -29,12 +29,14 @@ //! （"你刚才给我的邮箱是……"），本轮 session 里没有这条映射，占位符就漏给客户端。 //! HTTP 不会漏，是因为它每次都重发整段历史，重新 mask 同一个值会派生出同一个 //! sentinel（HMAC over 规则 + bucket + 值），所以映射天然齐备。 -//! * 挂在连接上（当前实现）：每轮仍然各自 mask、各自持有独立 session +//! * 挂在当前 response chain 上（当前实现）：每轮仍然各自 mask、各自持有独立 session //! （per-turn 语义不变），连接只是把最近若干轮的 session 留下来一起参与还原， //! 凑出的映射集合正好等于「等价 HTTP 请求会拥有的那一份」。 //! -//! 选后者。代价是每帧最多对 [`MAX_RETAINED_TURN_REDACTION_SESSIONS`] 个 session -//! 各扫一遍，以及这些 session 的映射会驻留到连接结束；用有界 FIFO 兜住上限。 +//! 选后者。省略（或置空）`previous_response_id` 会开始一条独立 response chain， +//! 此时必须丢弃旧链的映射，否则新链里偶然出现的旧 sentinel 会被还原成旧链 PII。 +//! 代价是每帧最多对 [`MAX_RETAINED_TURN_REDACTION_SESSIONS`] 个 session 各扫一遍， +//! 以及这些 session 的映射会驻留到当前链结束；用有界 FIFO 兜住上限。 //! 窗口不够用或每帧成本变高时，正确的下一步是在 `privacy` 侧提供跨 session 的 //! 合并匹配器，而不是把这个窗口调大。 @@ -89,6 +91,27 @@ pub(super) async fn redact_responses_websocket_client_event( parts: &http::request::Parts, control_decision: &GatewayControlDecision, client_event: &Value, +) -> Result, GatewayError> { + redact_responses_websocket_client_event_with_reasoning_replay_policy( + state, + parts, + control_decision, + client_event, + crate::ai_serving::OpenAiResponsesReasoningReplayPolicy::OpenAiItemIds, + ) + .await +} + +/// Variant used only after the gateway has selected and authenticated the +/// provider binding. The replay policy comes from that trusted binding, never +/// from client JSON, so a forged reasoning-item shape cannot opt itself into +/// byte-opaque PII handling. +pub(super) async fn redact_responses_websocket_client_event_with_reasoning_replay_policy( + state: &AppState, + parts: &http::request::Parts, + control_decision: &GatewayControlDecision, + client_event: &Value, + reasoning_replay_policy: crate::ai_serving::OpenAiResponsesReasoningReplayPolicy, ) -> Result, GatewayError> { let Some(auth_context) = resolve_local_decision_execution_runtime_auth_context(control_decision) @@ -101,6 +124,7 @@ pub(super) async fn redact_responses_websocket_client_event( client_event, &auth_context, RESPONSES_WEBSOCKET_CLIENT_API_FORMAT, + reasoning_replay_policy, WEBSOCKET_TURN_REDACTION_CANDIDATE_ID, ) .await?; @@ -126,17 +150,22 @@ pub(super) async fn redact_responses_websocket_client_event( })) } -/// 一条连接上「我们 mask 过哪些映射」的留存集合，供响应侧还原使用。 +/// 当前 response chain 上「我们 mask 过哪些映射」的留存集合，供响应侧还原使用。 /// -/// 每轮一个独立 session（per-turn mask 语义不变），连接按 FIFO 留最近 -/// [`MAX_RETAINED_TURN_REDACTION_SESSIONS`] 轮。上游重绑不清空：客户端仍在同一段 -/// 对话里，旧占位符可能随重发的输入再次出现。 +/// 每轮一个独立 session（per-turn mask 语义不变），当前链按 FIFO 留最近 +/// [`MAX_RETAINED_TURN_REDACTION_SESSIONS`] 轮。物理上游重绑本身不决定生命周期； +/// `previous_response_id` 决定是否延续旧链。独立请求成功发出时由调用方通过 +/// [`Self::start_new_chain`] 原子替换为新链的首轮 session。 #[derive(Default)] pub(super) struct ResponsesWebSocketRedactionRestorer { sessions: VecDeque, } impl ResponsesWebSocketRedactionRestorer { + pub(super) fn has_sessions(&self) -> bool { + !self.sessions.is_empty() + } + /// 登记这一轮的 mask session。 pub(super) fn register(&mut self, session: RedactionSession) { if session.mapping_count() == 0 { @@ -148,6 +177,18 @@ impl ResponsesWebSocketRedactionRestorer { } } + /// Commits a successfully started independent response chain. + /// + /// Keep this transition next to the successful upstream send/bind. A + /// rejected independent request has not replaced the active chain and + /// therefore must not discard the old chain's restore mappings. + pub(super) fn start_new_chain(&mut self, session: Option) { + self.sessions.clear(); + if let Some(session) = session { + self.register(session); + } + } + /// 把一帧 provider 事件里的占位符换回真实值，返回要发给客户端的帧文本。 /// /// `None` 表示这一帧没有任何东西要还原，调用方必须原样转发上游字节：未启用 @@ -497,8 +538,9 @@ mod tests { seed_report_context_with_raw_pii(), ); // 首轮实际发上游的事件由 decision.provider_request_body 派生。 + let normalization = ResponsesWebSocketBodyNormalization::for_tests("provider-model"); let provider_event: Value = serde_json::from_str( - &planned_response_create_event(&template, &effective_event) + &planned_response_create_event(&template, &normalization, &effective_event) .expect("first provider event should serialize"), ) .expect("first provider event should parse"); @@ -602,8 +644,9 @@ mod tests { provider_body_from(&active.client_event), seed_report_context_with_raw_pii(), ); + let normalization = ResponsesWebSocketBodyNormalization::for_tests("provider-model"); let provider_event: Value = serde_json::from_str( - &planned_response_create_event(&template, &active.client_event) + &planned_response_create_event(&template, &normalization, &active.client_event) .expect("retry provider event should serialize"), ) .expect("retry provider event should parse"); @@ -814,6 +857,53 @@ mod tests { assert!(!restored.contains(&second_sentinel), "{restored}"); } + /// Omitting `previous_response_id` starts a new response chain. Restore + /// mappings from the prior chain must not leak into that independent + /// response, while the new chain's first-turn mapping remains available. + #[tokio::test] + async fn an_independent_chain_replaces_prior_restore_mappings() { + let state = redaction_enabled_state(); + let decision = control_decision(); + let prior = turn_redaction(&state, &decision, TEST_EMAIL).await; + let current = turn_redaction(&state, &decision, OTHER_TEST_EMAIL).await; + let prior_sentinel = sentinel_for(&prior, TEST_EMAIL); + let current_sentinel = sentinel_for(¤t, OTHER_TEST_EMAIL); + + let mut restorer = ResponsesWebSocketRedactionRestorer::default(); + restorer.register(prior.session); + restorer.start_new_chain(Some(current.session)); + + assert!( + restorer + .restore_provider_frame_text(&provider_delta_frame(&prior_sentinel)) + .is_none(), + "an independent chain must not restore PII from its predecessor" + ); + let restored = restorer + .restore_provider_frame_text(&provider_delta_frame(¤t_sentinel)) + .expect("the new chain's first-turn mapping must remain available"); + assert!(restored.contains(OTHER_TEST_EMAIL), "{restored}"); + assert!(!restored.contains(¤t_sentinel), "{restored}"); + } + + #[tokio::test] + async fn an_unredacted_independent_chain_clears_prior_restore_mappings() { + let state = redaction_enabled_state(); + let prior = turn_redaction(&state, &control_decision(), TEST_EMAIL).await; + let prior_sentinel = sentinel_for(&prior, TEST_EMAIL); + + let mut restorer = ResponsesWebSocketRedactionRestorer::default(); + restorer.register(prior.session); + assert!(restorer.has_sessions()); + + restorer.start_new_chain(None); + + assert!(!restorer.has_sessions()); + assert!(restorer + .restore_provider_frame_text(&provider_delta_frame(&prior_sentinel)) + .is_none()); + } + /// 留存窗口是有界的：长连接不能无限累积映射，代价是更早的轮次会退回 /// 「占位符原样透传」而不是被错误还原成别的值。 #[tokio::test] diff --git a/apps/aether-gateway/src/handlers/proxy/websocket/responses/relay_policy.rs b/apps/aether-gateway/src/handlers/proxy/websocket/responses/relay_policy.rs index 99b974b4a..7f342c13d 100644 --- a/apps/aether-gateway/src/handlers/proxy/websocket/responses/relay_policy.rs +++ b/apps/aether-gateway/src/handlers/proxy/websocket/responses/relay_policy.rs @@ -10,6 +10,7 @@ pub enum FatalRelaySignal { ConnectionAdmissionLost, InvalidUpstreamText, + UnexpectedUpstreamStreamId, } #[derive(Debug, Clone, Copy, PartialEq, Eq)] @@ -40,6 +41,13 @@ pub const fn fatal_relay_policy(signal: FatalRelaySignal) -> FatalRelayPolicy { client_message: "Provider returned an invalid WebSocket event", close_reason: "invalid_upstream_event", }, + FatalRelaySignal::UnexpectedUpstreamStreamId => FatalRelayPolicy { + status_code: 502, + close_code: 1011, + error_code: "responses_websocket_unexpected_upstream_stream_id", + client_message: "Provider returned a named-lane event on a default-lane connection", + close_reason: "unexpected_upstream_stream_id", + }, } } @@ -230,6 +238,20 @@ mod tests { assert_eq!(upstream.next(), None); } + #[test] + fn unexpected_named_lane_is_a_terminal_provider_protocol_error() { + assert_eq!( + fatal_relay_policy(FatalRelaySignal::UnexpectedUpstreamStreamId), + FatalRelayPolicy { + status_code: 502, + close_code: 1011, + error_code: "responses_websocket_unexpected_upstream_stream_id", + client_message: "Provider returned a named-lane event on a default-lane connection", + close_reason: "unexpected_upstream_stream_id", + } + ); + } + #[test] fn quota_snapshot_without_definitive_error_does_not_trigger_retry() { assert_eq!( diff --git a/apps/aether-gateway/src/handlers/proxy/websocket/responses/request.rs b/apps/aether-gateway/src/handlers/proxy/websocket/responses/request.rs index a3530fcc0..68ddb8635 100644 --- a/apps/aether-gateway/src/handlers/proxy/websocket/responses/request.rs +++ b/apps/aether-gateway/src/handlers/proxy/websocket/responses/request.rs @@ -7,6 +7,7 @@ use axum::http::header::{AUTHORIZATION, CONNECTION, CONTENT_TYPE, UPGRADE}; use axum::http::Method; use serde_json::Value; +use sha2::Digest as _; use crate::ai_serving::{AiExecutionDecision, ResponsesWebSocketBodyNormalization}; use crate::handlers::proxy::websocket::ingress::WebSocketRequestContext; @@ -17,6 +18,260 @@ use crate::privacy::RedactionSessionSlot; /// any planning/logging so a single 16 MiB WebSocket frame cannot amplify into /// repeated multi-megabyte log records. pub(super) const MAX_RESPONSES_WEBSOCKET_MODEL_BYTES: usize = 256; +/// Response IDs are opaque, but bounding their wire size prevents an +/// untrusted first frame from creating unbounded registry/hash work. +pub(super) const MAX_RESPONSES_WEBSOCKET_RESPONSE_ID_BYTES: usize = 256; + +/// The static Responses Lite prefix already stored in a response chain. +/// +/// Codex normally sends only the incremental input on a WebSocket +/// continuation. Some compatible clients repeat the current top-level tools +/// and instructions, while others repeat the Lite synthetic input prefix. We +/// retain the first turn's effective configuration so repeated copies can be +/// removed without silently discarding an actual configuration change. +#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)] +pub(super) struct ResponsesLiteStaticConfig { + tools_sha256: [u8; 32], + instructions_sha256: [u8; 32], +} + +impl Default for ResponsesLiteStaticConfig { + fn default() -> Self { + Self { + tools_sha256: responses_lite_static_value_sha256(&Value::Array(Vec::new())), + instructions_sha256: responses_lite_static_bytes_sha256(b""), + } + } +} + +#[derive(Debug, Default)] +struct ResponsesLiteStaticConfigObservation { + tools: Option, + instructions: Option, + leading_input_items: usize, +} + +impl ResponsesLiteStaticConfig { + pub(super) fn from_response_create(event: &Value) -> Self { + let observation = observe_responses_lite_static_config(event).unwrap_or_default(); + let tools = observation + .tools + .unwrap_or_else(|| Value::Array(Vec::new())); + let instructions = observation.instructions.unwrap_or_default(); + Self { + tools_sha256: responses_lite_static_value_sha256(&tools), + instructions_sha256: responses_lite_static_bytes_sha256(instructions.as_bytes()), + } + } +} + +fn responses_lite_static_bytes_sha256(value: &[u8]) -> [u8; 32] { + sha2::Sha256::digest(value).into() +} + +fn responses_lite_static_value_sha256(value: &Value) -> [u8; 32] { + let mut digest = sha2::Sha256::new(); + update_responses_lite_static_digest(&mut digest, value); + digest.finalize().into() +} + +fn update_responses_lite_static_length(digest: &mut sha2::Sha256, length: usize) { + let length = u64::try_from(length).expect("JSON values cannot exceed the u64 hash domain"); + digest.update(length.to_be_bytes()); +} + +fn update_responses_lite_static_digest(digest: &mut sha2::Sha256, value: &Value) { + match value { + Value::Null => digest.update(b"n"), + Value::Bool(value) => digest.update(if *value { b"t" } else { b"f" }), + Value::Number(value) => { + digest.update(b"d"); + digest.update(value.to_string().as_bytes()); + digest.update(b";"); + } + Value::String(value) => { + digest.update(b"s"); + update_responses_lite_static_length(digest, value.len()); + digest.update(value.as_bytes()); + } + Value::Array(values) => { + digest.update(b"["); + update_responses_lite_static_length(digest, values.len()); + for value in values { + update_responses_lite_static_digest(digest, value); + } + digest.update(b"]"); + } + Value::Object(values) => { + digest.update(b"{"); + update_responses_lite_static_length(digest, values.len()); + let mut keys = values.keys().collect::>(); + keys.sort_unstable(); + for key in keys { + update_responses_lite_static_length(digest, key.len()); + digest.update(key.as_bytes()); + update_responses_lite_static_digest(digest, &values[key]); + } + digest.update(b"}"); + } + } +} + +fn is_responses_lite_additional_tools_item(item: &Value) -> bool { + item.get("type").and_then(Value::as_str) == Some("additional_tools") + && item.get("role").and_then(Value::as_str) == Some("developer") + && item.get("tools").is_some_and(Value::is_array) +} + +fn responses_lite_instruction_text(item: &Value) -> Option<&str> { + (item.get("type").and_then(Value::as_str) == Some("message") + && item.get("role").and_then(Value::as_str) == Some("developer")) + .then(|| { + item.get("content") + .and_then(Value::as_array) + .filter(|content| content.len() == 1) + .and_then(|content| content[0].as_object()) + .filter(|content| content.get("type").and_then(Value::as_str) == Some("input_text")) + .and_then(|content| content.get("text")) + .and_then(Value::as_str) + }) + .flatten() +} + +fn normalize_responses_lite_tools(value: &Value) -> Result { + match value { + Value::Null => Ok(Value::Array(Vec::new())), + Value::Array(tools) => Ok(Value::Array( + tools + .iter() + .filter(|tool| { + crate::ai_serving::codex_responses_lite_tool_is_client_executed(tool) + }) + .cloned() + .collect(), + )), + _ => Err("invalid_response_create_tools"), + } +} + +fn normalize_responses_lite_instructions(value: &Value) -> Result { + match value { + Value::Null => Ok(String::new()), + Value::String(value) => Ok(value.clone()), + _ => Err("invalid_response_create_instructions"), + } +} + +fn observe_responses_lite_static_config( + event: &Value, +) -> Result { + let object = event.as_object().ok_or("invalid_response_create")?; + let mut observation = ResponsesLiteStaticConfigObservation::default(); + + if let Some(tools) = object.get("tools") { + observation.tools = Some(normalize_responses_lite_tools(tools)?); + } + if let Some(instructions) = object.get("instructions") { + observation.instructions = Some(normalize_responses_lite_instructions(instructions)?); + } + + let Some(input) = object.get("input").and_then(Value::as_array) else { + return Ok(observation); + }; + let mut consumed_leading_instruction = false; + while input + .get(observation.leading_input_items) + .is_some_and(is_responses_lite_additional_tools_item) + { + let additional_tools = &input[observation.leading_input_items]; + observation.leading_input_items += 1; + let tools = normalize_responses_lite_tools(&additional_tools["tools"])?; + if observation.tools.as_ref().is_some_and(|existing| { + responses_lite_static_value_sha256(existing) + != responses_lite_static_value_sha256(&tools) + }) { + return Err("responses_lite_static_tools_conflict"); + } + observation.tools = Some(tools); + if let Some(instructions) = input + .get(observation.leading_input_items) + .and_then(responses_lite_instruction_text) + { + observation.leading_input_items += 1; + consumed_leading_instruction = true; + if observation + .instructions + .as_deref() + .is_some_and(|existing| existing != instructions) + { + return Err("responses_lite_static_instructions_conflict"); + } + observation.instructions = Some(instructions.to_string()); + } + } + // A Lite request may contain instructions without any client-executed + // tools. Its normalized prefix then starts directly with the synthetic + // developer message, which must be inherited just like the two-item + // additional_tools + instructions prefix. + if !consumed_leading_instruction { + if let Some(instructions) = input + .get(observation.leading_input_items) + .and_then(responses_lite_instruction_text) + { + observation.leading_input_items += 1; + if observation + .instructions + .as_deref() + .is_some_and(|existing| existing != instructions) + { + return Err("responses_lite_static_instructions_conflict"); + } + observation.instructions = Some(instructions.to_string()); + } + } + Ok(observation) +} + +/// Removes a repeated Responses Lite static prefix from a continuation. +/// +/// A changed prefix cannot safely be appended to the stored history: doing so +/// is the context-growth bug this path prevents, and the Lite backend has no +/// request primitive that erases the previous synthetic item. Match Codex's +/// own transport behavior by requiring a new response chain when request +/// properties change instead of silently retaining stale tools/instructions. +pub(super) fn prepare_responses_lite_continuation( + event: &Value, + stored: &ResponsesLiteStaticConfig, +) -> Result { + let observation = observe_responses_lite_static_config(event)?; + if observation + .tools + .as_ref() + .is_some_and(|tools| responses_lite_static_value_sha256(tools) != stored.tools_sha256) + || observation + .instructions + .as_ref() + .is_some_and(|instructions| { + responses_lite_static_bytes_sha256(instructions.as_bytes()) + != stored.instructions_sha256 + }) + { + return Err("responses_lite_continuation_static_config_changed"); + } + + let mut prepared = event.clone(); + let object = prepared.as_object_mut().ok_or("invalid_response_create")?; + object.remove("tools"); + object.remove("instructions"); + if observation.leading_input_items > 0 { + let input = object + .get_mut("input") + .and_then(Value::as_array_mut) + .ok_or("invalid_response_create_input")?; + input.drain(..observation.leading_input_items); + } + Ok(prepared) +} pub(super) fn validated_response_create_model(value: &Value) -> Result<&str, &'static str> { let Some(model) = value @@ -80,13 +335,14 @@ pub(super) fn build_planning_parts(context: &WebSocketRequestContext) -> http::r pub(super) fn planned_response_create_event( decision: &AiExecutionDecision, + normalization: &ResponsesWebSocketBodyNormalization, fallback: &Value, ) -> Result { let event = decision .provider_request_body .clone() .unwrap_or_else(|| fallback.clone()); - finish_response_create_event(event, fallback) + finish_response_create_event(event, fallback, normalization) } /// Restores the WebSocket protocol framing that provider-body normalization is @@ -94,13 +350,19 @@ pub(super) fn planned_response_create_event( /// /// `previous_response_id` is on the Codex unsupported-field list, Codex HTTP /// normalization may force `store`, and `generate` is not an HTTP body option -/// at all. Those fields are WebSocket protocol state, so an explicitly supplied -/// value (including `null`) must be re-grafted verbatim from the client event. +/// at all. Those fields are WebSocket protocol state. `store` and `generate` +/// may still be owned by an endpoint body rule, but lineage is security state: +/// the final `previous_response_id` must be exactly the client value whose +/// ownership the session validated. A body rule may neither replace nor inject +/// an opaque response ID after that check. /// `stream`/`background` go the other way: the normalizer inserts `stream`, and -/// the WebSocket protocol has no use for it. +/// the WebSocket protocol has no use for it. Named `stream_id` lanes are not +/// yet exposed by Aether, so provider rules cannot inject one behind the +/// default-lane validator. fn finish_response_create_event( mut event: Value, client_event: &Value, + normalization: &ResponsesWebSocketBodyNormalization, ) -> Result { let object = event .as_object_mut() @@ -109,20 +371,78 @@ fn finish_response_create_event( "type".to_string(), Value::String("response.create".to_string()), ); - for field in ["store", "previous_response_id", "generate"] { - if let Some(value) = client_event.get(field) { - object.insert(field.to_string(), value.clone()); + for field in ["store", "generate"] { + if !normalization.body_rules_handle_websocket_field(client_event, field) { + if let Some(value) = client_event.get(field) { + object.insert(field.to_string(), value.clone()); + } + } + } + match client_event.get("previous_response_id") { + Some(value) => { + object.insert("previous_response_id".to_string(), value.clone()); + } + None => { + object.remove("previous_response_id"); } } object.remove("stream"); object.remove("background"); + object.remove("stream_id"); serde_json::to_string(&event).map_err(|_| "responses_websocket_request_invalid") } pub(super) fn response_create_has_previous_response_id(event: &Value) -> bool { event .get("previous_response_id") - .is_some_and(|value| !value.is_null()) + .and_then(Value::as_str) + .is_some_and(|value| !value.trim().is_empty()) +} + +pub(super) fn validate_response_create_previous_response_id( + event: &Value, +) -> Result<(), &'static str> { + match event.get("previous_response_id") { + None | Some(Value::Null) => Ok(()), + Some(Value::String(value)) + if !value.trim().is_empty() + && value.len() <= MAX_RESPONSES_WEBSOCKET_RESPONSE_ID_BYTES => + { + Ok(()) + } + Some(_) => Err("invalid_response_create_previous_response_id"), + } +} + +/// Returns a named lane only after the complete public `stream_id` grammar has +/// been checked. This is the sole source for reflecting a request-scoped lane +/// into gateway-generated error events. +pub(super) fn validated_named_stream_id(event: &Value) -> Option<&str> { + let stream_id = event.get("stream_id")?.as_str()?; + (!stream_id.is_empty() + && stream_id.len() <= 256 + && stream_id + .bytes() + .all(|byte| byte.is_ascii_alphanumeric() || matches!(byte, b'_' | b'-' | b'.'))) + .then_some(stream_id) +} + +/// Aether currently owns one logical turn for each physical Responses socket. +/// +/// OpenAI's named `stream_id` lanes require per-lane turn state, FIFO queues, +/// usage settlement, timeouts, and gateway-generated error routing. Silently +/// dropping the field would merge independent conversations into the default +/// lane, so fail closed until that complete multiplexing contract is present. +pub(super) fn validate_response_create_stream_id_support( + event: &Value, +) -> Result<(), &'static str> { + if event.get("stream_id").is_none() { + return Ok(()); + } + if validated_named_stream_id(event).is_none() { + return Err("invalid_response_create_stream_id"); + } + Err("responses_websocket_named_stream_unsupported") } pub(super) fn changed_followup_response_create_model( @@ -175,6 +495,41 @@ pub(super) fn provider_model_from_decision(decision: &AiExecutionDecision) -> Op .map(str::to_string) } +/// Returns the effective Codex contract after body rules, routing mutations, +/// model directives, and terminal header convergence have all run. +/// +/// Model capability alone is insufficient: a non-null `context_management` +/// request deliberately disables Responses Lite for that turn. The protected +/// Lite header is emitted from the final provider body, so it is the stable +/// contract marker retained by a WebSocket binding. +pub(super) fn planned_request_uses_codex_responses_lite( + decision: &AiExecutionDecision, + normalization: &ResponsesWebSocketBodyNormalization, +) -> bool { + let decision_is_codex_responses = decision + .provider_type + .as_deref() + .is_some_and(|value| value.trim().eq_ignore_ascii_case("codex")) + && decision + .provider_api_format + .as_deref() + .is_some_and(crate::ai_serving::is_openai_responses_family_format); + let final_body_supports_lite = decision + .provider_request_body + .as_ref() + .is_none_or(|body| body.get("context_management").is_none_or(Value::is_null)); + decision_is_codex_responses + && normalization.uses_codex_responses_lite() + && final_body_supports_lite + && decision + .provider_request_headers + .iter() + .any(|(name, value)| { + name.eq_ignore_ascii_case(crate::ai_serving::CODEX_RESPONSES_LITE_HEADER) + && value.trim().eq_ignore_ascii_case("true") + }) +} + /// Prepares a continuation `response.create` for the already-bound upstream. /// /// The turn cannot be re-planned without risking a different provider key, so @@ -192,12 +547,13 @@ pub(super) fn normalize_followup_response_create( if event.get("type").and_then(Value::as_str) != Some("response.create") { return Err("invalid_response_create"); } - // Normalization is best-effort here: a continuation cannot fall back to - // another candidate, so a body the contract rejects is still better sent - // than dropped. + validate_response_create_previous_response_id(event)?; + // Never fall back to the raw event. That would bypass endpoint rules and + // the Responses Lite de-duplication exactly when normalization rejected a + // malformed request, potentially re-appending static configuration. let mut normalized = normalization .normalize_response_create(event) - .unwrap_or_else(|| event.clone()); + .ok_or("responses_websocket_request_normalization_failed")?; let Some(object) = normalized.as_object_mut() else { return Err("invalid_response_create"); }; @@ -207,7 +563,7 @@ pub(super) fn normalize_followup_response_create( "model".to_string(), Value::String(provider_model.to_string()), ); - finish_response_create_event(normalized, event) + finish_response_create_event(normalized, event, normalization) .map_err(|_| "response_create_serialization_failed") } @@ -220,7 +576,11 @@ mod tests { use super::{ build_planning_parts, normalize_followup_response_create, - response_create_has_previous_response_id, + planned_request_uses_codex_responses_lite, planned_response_create_event, + prepare_responses_lite_continuation, response_create_has_previous_response_id, + validate_response_create_previous_response_id, validate_response_create_stream_id_support, + validated_named_stream_id, ResponsesLiteStaticConfig, + MAX_RESPONSES_WEBSOCKET_RESPONSE_ID_BYTES, }; use crate::ai_serving::ResponsesWebSocketBodyNormalization; use crate::control::GatewayControlDecision; @@ -311,12 +671,45 @@ mod tests { } #[test] - fn explicit_store_and_previous_response_id_are_forwarded_opaquely() { + fn websocket_first_response_create_keeps_responses_lite_static_config() { + // A socket's first event has no previous_response_id, so its tools and + // instructions still need the normal Responses Lite projection. The + // continuation-only pass must not accidentally suppress this prefix. + let event = json!({ + "type": "response.create", + "model": "gpt-5.6-sol", + "instructions": "developer instructions", + "tools": [{ + "type": "function", + "name": "lookup", + "parameters": {"type": "object"} + }], + "input": [{ + "role": "user", + "content": [{"type": "input_text", "text": "hello"}] + }] + }); + + let normalization = ResponsesWebSocketBodyNormalization::for_tests("gpt-5.6-sol") + .with_provider_type_for_tests("codex"); + let normalized = normalized_continuation(&event, &normalization); + let input = normalized["input"].as_array().expect("normalized input"); + assert_eq!(input[0]["type"], "additional_tools"); + assert_eq!(input[0]["tools"][0]["name"], "lookup"); + assert_eq!(input[1]["role"], "developer"); + assert_eq!(input[1]["content"][0]["text"], "developer instructions"); + assert!(normalized.get("tools").is_none()); + assert!(normalized.get("instructions").is_none()); + assert!(normalized.get("previous_response_id").is_none()); + } + + #[test] + fn explicit_store_and_previous_response_id_are_forwarded() { let event = json!({ "type": "response.create", "model": "public-model", "store": true, - "previous_response_id": {"future": "opaque"}, + "previous_response_id": "resp_123", "input": [], }); @@ -330,10 +723,126 @@ mod tests { // `previous_response_id`. WebSocket framing restores exactly what the // client sent so the upstream owns validation and continuation lookup. assert_eq!(normalized["store"], true); - assert_eq!( - normalized["previous_response_id"], - json!({"future": "opaque"}) + assert_eq!(normalized["previous_response_id"], "resp_123"); + } + + #[test] + fn endpoint_body_rules_cannot_replace_validated_websocket_lineage() { + let event = json!({ + "type": "response.create", + "model": "public-model", + "store": true, + "previous_response_id": "resp_client", + "generate": false, + "input": [], + }); + let normalization = ResponsesWebSocketBodyNormalization::for_tests("provider-model") + .with_provider_type_for_tests("codex") + .with_body_rules_for_tests(json!([ + {"action": "set", "path": "store", "value": false}, + { + "action": "set", + "path": "previous_response_id", + "value": "resp_admin" + }, + {"action": "set", "path": "generate", "value": true} + ])); + + let normalized = normalized_continuation(&event, &normalization); + + // WebSocket framing runs after provider-body finalization. Store and + // generate remain administrator-owned, but the opaque lineage ID must + // be the exact value validated against the authenticated connection. + assert_eq!(normalized["store"], false); + assert_eq!(normalized["previous_response_id"], "resp_client"); + assert_eq!(normalized["generate"], true); + } + + #[test] + fn conditional_body_rules_own_store_and_generate_only_when_their_values_apply() { + let normalization = ResponsesWebSocketBodyNormalization::for_tests("provider-model") + .with_provider_type_for_tests("codex") + .with_body_rules_for_tests(json!([ + { + "action": "set", + "path": "store", + "value": false, + "condition": {"path": "metadata.mode", "op": "eq", "value": "enforce"} + }, + { + "action": "set", + "path": "generate", + "value": true, + "condition": {"path": "metadata.mode", "op": "eq", "value": "enforce"} + } + ])); + + let skipped = normalized_continuation( + &json!({ + "type": "response.create", + "model": "public-model", + "store": true, + "generate": false, + "metadata": {"mode": "observe"}, + "input": [] + }), + &normalization, ); + assert_eq!(skipped["store"], true); + assert_eq!(skipped["generate"], false); + + let applied = normalized_continuation( + &json!({ + "type": "response.create", + "model": "public-model", + "store": true, + "generate": false, + "metadata": {"mode": "enforce"}, + "input": [] + }), + &normalization, + ); + assert_eq!(applied["store"], false); + assert_eq!(applied["generate"], true); + } + + #[test] + fn endpoint_body_rules_cannot_inject_websocket_lineage() { + let event = json!({ + "type": "response.create", + "model": "public-model", + "input": [], + }); + let normalization = ResponsesWebSocketBodyNormalization::for_tests("provider-model") + .with_provider_type_for_tests("codex") + .with_body_rules_for_tests(json!([{ + "action": "set", + "path": "previous_response_id", + "value": "resp_admin" + }])); + + let normalized = normalized_continuation(&event, &normalization); + + assert!(normalized.get("previous_response_id").is_none()); + } + + #[test] + fn endpoint_body_rules_cannot_inject_an_untracked_named_lane() { + let event = json!({ + "type": "response.create", + "model": "public-model", + "input": [], + }); + let normalization = ResponsesWebSocketBodyNormalization::for_tests("provider-model") + .with_body_rules_for_tests(json!([{ + "action": "set", + "path": "stream_id", + "value": "admin-injected-lane" + }])); + + let normalized = normalized_continuation(&event, &normalization); + + assert!(normalized.get("stream_id").is_none()); } #[test] @@ -436,6 +945,51 @@ mod tests { } } + #[test] + fn deepseek_continuations_preserve_idless_opaque_reasoning_state() { + let normalization = ResponsesWebSocketBodyNormalization::for_tests("deepseek-v4-flash") + .with_provider_type_for_tests("custom") + .with_reasoning_replay_policy_for_tests( + crate::ai_serving::OpenAiResponsesReasoningReplayPolicy::DeepSeekOpaque, + ); + let opaque_reasoning = json!({ + "type": "reasoning", + "encrypted_content": "opaque-deepseek-state", + "content": [{ + "type": "reasoning_text", + "text": "provider-owned thinking state" + }], + "future_capability": {"preserve": true} + }); + let event = json!({ + "type": "response.create", + "model": "deepseek-v4-flash", + "previous_response_id": "resp_deepseek_1", + "input": [ + opaque_reasoning.clone(), + { + "type": "function_call_output", + "call_id": "call_1", + "output": "result" + } + ] + }); + + let normalized = normalized_continuation(&event, &normalization); + assert_eq!(normalized["input"][0], opaque_reasoning); + assert_eq!(normalized["input"][1]["type"], "function_call_output"); + assert_eq!(normalized["previous_response_id"], "resp_deepseek_1"); + + let strict = normalized_continuation( + &event, + &ResponsesWebSocketBodyNormalization::for_tests("deepseek-v4-flash") + .with_provider_type_for_tests("custom"), + ); + let strict_input = strict["input"].as_array().expect("strict provider input"); + assert_eq!(strict_input.len(), 1); + assert_eq!(strict_input[0]["type"], "function_call_output"); + } + #[test] fn continuation_keeps_a_warmup_generate_flag() { let event = json!({ @@ -509,13 +1063,416 @@ mod tests { } #[test] - fn previous_response_id_is_protocol_state_even_when_not_a_string() { + fn previous_response_id_requires_a_non_empty_string() { assert!(response_create_has_previous_response_id( + &json!({"previous_response_id": "resp_123"}) + )); + assert!(!response_create_has_previous_response_id( + &json!({"previous_response_id": " "}) + )); + assert!(!response_create_has_previous_response_id( &json!({"previous_response_id": 42}) )); + assert!(!response_create_has_previous_response_id( + &json!({"previous_response_id": {"id": "resp_123"}}) + )); assert!(!response_create_has_previous_response_id( &json!({"previous_response_id": null}) )); assert!(!response_create_has_previous_response_id(&json!({}))); + + assert!(validate_response_create_previous_response_id( + &json!({"previous_response_id": "resp_123"}) + ) + .is_ok()); + assert!(validate_response_create_previous_response_id( + &json!({"previous_response_id": null}) + ) + .is_ok()); + for invalid in [json!(" "), json!(42), json!({"id": "resp_123"})] { + assert!(validate_response_create_previous_response_id( + &json!({"previous_response_id": invalid}) + ) + .is_err()); + } + assert!(validate_response_create_previous_response_id(&json!({ + "previous_response_id": "x".repeat(MAX_RESPONSES_WEBSOCKET_RESPONSE_ID_BYTES + 1) + })) + .is_err()); + } + + #[test] + fn named_streams_fail_closed_until_lane_multiplexing_is_implemented() { + assert!(validate_response_create_stream_id_support(&json!({})).is_ok()); + for stream_id in [ + json!(null), + json!(""), + json!("invalid/lane"), + json!("x".repeat(257)), + json!(42), + ] { + assert_eq!( + validate_response_create_stream_id_support(&json!({"stream_id": stream_id})), + Err("invalid_response_create_stream_id") + ); + } + assert_eq!( + validate_response_create_stream_id_support(&json!({"stream_id": "main-lane_1.test"})), + Err("responses_websocket_named_stream_unsupported") + ); + assert_eq!( + validated_named_stream_id(&json!({"stream_id": "main-lane_1.test"})), + Some("main-lane_1.test") + ); + for invalid in [ + json!(null), + json!(""), + json!("invalid/lane"), + json!("x".repeat(257)), + json!(42), + ] { + assert_eq!( + validated_named_stream_id(&json!({"stream_id": invalid})), + None + ); + } + } + + #[test] + fn effective_lite_contract_uses_the_converged_provider_header() { + let mut decision: crate::ai_serving::AiExecutionDecision = serde_json::from_value(json!({ + "action": "local", + "provider_type": "codex", + "provider_api_format": "openai:responses", + "provider_request_headers": {} + })) + .expect("minimal decision"); + let normalization = ResponsesWebSocketBodyNormalization::for_tests("gpt-5.6-sol") + .with_provider_type_for_tests("codex"); + assert!(!planned_request_uses_codex_responses_lite( + &decision, + &normalization + )); + + decision.provider_request_headers.insert( + crate::ai_serving::CODEX_RESPONSES_LITE_HEADER.to_string(), + "false".to_string(), + ); + assert!(!planned_request_uses_codex_responses_lite( + &decision, + &normalization + )); + + decision.provider_request_headers.clear(); + decision.provider_request_headers.insert( + crate::ai_serving::CODEX_RESPONSES_LITE_HEADER.to_ascii_uppercase(), + "TRUE".to_string(), + ); + assert!(planned_request_uses_codex_responses_lite( + &decision, + &normalization + )); + + decision.provider_request_body = Some(json!({ + "model": "gpt-5.6-sol", + "context_management": {"compact_threshold": 1000} + })); + assert!(!planned_request_uses_codex_responses_lite( + &decision, + &normalization + )); + decision.provider_request_body = None; + + // Header rules on a custom provider must not be able to spoof the + // internal Codex contract marker and enable Lite de-duplication. + decision.provider_type = Some("custom".to_string()); + assert!(!planned_request_uses_codex_responses_lite( + &decision, + &normalization + )); + decision.provider_type = Some("codex".to_string()); + let custom_normalization = ResponsesWebSocketBodyNormalization::for_tests("gpt-5.6-sol"); + assert!(!planned_request_uses_codex_responses_lite( + &decision, + &custom_normalization + )); + } + + #[test] + fn responses_lite_continuation_deduplicates_only_matching_static_config() { + let first = json!({ + "type": "response.create", + "model": "gpt-5.6-sol", + "instructions": "developer instructions", + "tools": [{"type": "function", "name": "lookup", "parameters": {}}], + "input": [{"role": "user", "content": "hello"}] + }); + let stored = ResponsesLiteStaticConfig::from_response_create(&first); + let continuation = json!({ + "type": "response.create", + "model": "gpt-5.6-sol", + "previous_response_id": "resp_1", + "instructions": "developer instructions", + "tools": [{"type": "function", "name": "lookup", "parameters": {}}], + "input": [{"type": "function_call_output", "call_id": "call_1", "output": "ok"}] + }); + + let prepared = prepare_responses_lite_continuation(&continuation, &stored) + .expect("matching static config should be inherited"); + assert!(prepared.get("tools").is_none()); + assert!(prepared.get("instructions").is_none()); + assert_eq!(prepared["input"].as_array().map(Vec::len), Some(1)); + + let changed_instructions = json!({ + "type": "response.create", + "model": "gpt-5.6-sol", + "previous_response_id": "resp_1", + "instructions": "different instructions", + "input": [] + }); + assert_eq!( + prepare_responses_lite_continuation(&changed_instructions, &stored), + Err("responses_lite_continuation_static_config_changed") + ); + let changed_tools = json!({ + "type": "response.create", + "model": "gpt-5.6-sol", + "previous_response_id": "resp_1", + "tools": [{"type": "function", "name": "other", "parameters": {}}], + "input": [] + }); + assert_eq!( + prepare_responses_lite_continuation(&changed_tools, &stored), + Err("responses_lite_continuation_static_config_changed") + ); + } + + #[test] + fn responses_lite_static_tools_match_the_actual_client_executed_synthetic_subset() { + let function = json!({"type": "function", "name": "lookup", "parameters": {}}); + let custom = json!({"type": "custom", "name": "shell", "format": {}}); + let namespace = json!({"type": "namespace", "name": "browser", "tools": []}); + let client_search = json!({"type": "tool_search", "execution": "client"}); + let first = json!({ + "type": "response.create", + "model": "gpt-5.6-sol", + "instructions": "developer instructions", + "tools": [ + function.clone(), + {"type": "web_search"}, + {"type": "image_generation"}, + custom.clone(), + namespace.clone(), + client_search.clone(), + {"type": "tool_search", "execution": "server"}, + {"type": "tool_search"}, + {"type": "future_hosted_tool"} + ], + "input": [{"role": "user", "content": "hello"}] + }); + let stored = ResponsesLiteStaticConfig::from_response_create(&first); + let continuation = json!({ + "type": "response.create", + "model": "gpt-5.6-sol", + "previous_response_id": "resp_1", + "input": [ + { + "type": "additional_tools", + "role": "developer", + "tools": [function, custom, namespace, client_search] + }, + { + "type": "message", + "role": "developer", + "content": [{"type": "input_text", "text": "developer instructions"}] + }, + {"type": "function_call_output", "call_id": "call_1", "output": "ok"} + ] + }); + + let prepared = prepare_responses_lite_continuation(&continuation, &stored) + .expect("hosted tools are not part of the stored Lite synthetic prefix"); + let input = prepared["input"].as_array().expect("prepared input"); + assert_eq!(input.len(), 1); + assert_eq!(input[0]["type"], "function_call_output"); + } + + #[test] + fn responses_lite_static_identity_must_be_retained_before_redaction() { + let raw_first = json!({ + "type": "response.create", + "model": "gpt-5.6-sol", + "instructions": "customer secret", + "tools": [{ + "type": "function", + "name": "lookup", + "description": "customer secret" + }], + "input": [] + }); + let redacted_first = json!({ + "type": "response.create", + "model": "gpt-5.6-sol", + "instructions": "", + "tools": [{ + "type": "function", + "name": "lookup", + "description": "" + }], + "input": [] + }); + let raw_continuation = json!({ + "type": "response.create", + "model": "gpt-5.6-sol", + "previous_response_id": "resp_1", + "instructions": "customer secret", + "tools": [{ + "type": "function", + "name": "lookup", + "description": "customer secret" + }], + "input": [{"type": "function_call_output", "call_id": "call_1", "output": "ok"}] + }); + + let raw_identity = ResponsesLiteStaticConfig::from_response_create(&raw_first); + let prepared = prepare_responses_lite_continuation(&raw_continuation, &raw_identity) + .expect("the same plaintext configuration should be inherited"); + assert!(prepared.get("tools").is_none()); + assert!(prepared.get("instructions").is_none()); + + let redacted_identity = ResponsesLiteStaticConfig::from_response_create(&redacted_first); + assert_eq!( + prepare_responses_lite_continuation(&raw_continuation, &redacted_identity), + Err("responses_lite_continuation_static_config_changed") + ); + } + + #[test] + fn responses_lite_static_hash_is_stable_across_object_key_order() { + let first = json!({ + "type": "response.create", + "tools": [{ + "type": "function", + "name": "lookup", + "parameters": {"type": "object", "properties": {"b": {}, "a": {}}} + }], + "instructions": "same" + }); + let reordered = json!({ + "instructions": "same", + "tools": [{ + "parameters": {"properties": {"a": {}, "b": {}}, "type": "object"}, + "name": "lookup", + "type": "function" + }], + "type": "response.create" + }); + + assert_eq!( + ResponsesLiteStaticConfig::from_response_create(&first), + ResponsesLiteStaticConfig::from_response_create(&reordered) + ); + } + + #[test] + fn decision_without_a_provider_body_uses_the_prepared_continuation_fallback() { + let first = json!({ + "type": "response.create", + "model": "gpt-5.6-sol", + "instructions": "developer instructions", + "tools": [{"type": "function", "name": "lookup", "parameters": {}}], + "input": [{"role": "user", "content": "hello"}] + }); + let stored = ResponsesLiteStaticConfig::from_response_create(&first); + let continuation = json!({ + "type": "response.create", + "model": "gpt-5.6-sol", + "previous_response_id": "resp_1", + "instructions": "developer instructions", + "tools": [{"type": "function", "name": "lookup", "parameters": {}}], + "input": [{"type": "function_call_output", "call_id": "call_1", "output": "ok"}] + }); + let prepared = prepare_responses_lite_continuation(&continuation, &stored) + .expect("matching configuration"); + let decision: crate::ai_serving::AiExecutionDecision = serde_json::from_value(json!({ + "action": "local" + })) + .expect("minimal decision"); + + let normalization = ResponsesWebSocketBodyNormalization::for_tests("gpt-5.6-sol") + .with_provider_type_for_tests("codex"); + let outbound = planned_response_create_event(&decision, &normalization, &prepared) + .expect("prepared fallback should serialize"); + let outbound: Value = serde_json::from_str(&outbound).expect("provider event"); + assert!(outbound.get("tools").is_none()); + assert!(outbound.get("instructions").is_none()); + assert_eq!(outbound["previous_response_id"], "resp_1"); + assert_eq!(outbound["input"].as_array().map(Vec::len), Some(1)); + } + + #[test] + fn responses_lite_continuation_removes_a_repeated_synthetic_prefix() { + let first = json!({ + "type": "response.create", + "model": "gpt-5.6-sol", + "input": [ + {"type": "additional_tools", "role": "developer", "tools": [ + {"type": "function", "name": "lookup", "parameters": {}} + ]}, + {"type": "message", "role": "developer", "content": [ + {"type": "input_text", "text": "developer instructions"} + ]}, + {"role": "user", "content": "hello"} + ] + }); + let stored = ResponsesLiteStaticConfig::from_response_create(&first); + let continuation = json!({ + "type": "response.create", + "model": "gpt-5.6-sol", + "previous_response_id": "resp_1", + "input": [ + {"type": "additional_tools", "role": "developer", "tools": [ + {"type": "function", "name": "lookup", "parameters": {}} + ]}, + {"type": "message", "role": "developer", "content": [ + {"type": "input_text", "text": "developer instructions"} + ]}, + {"type": "function_call_output", "call_id": "call_1", "output": "ok"} + ] + }); + + let prepared = prepare_responses_lite_continuation(&continuation, &stored) + .expect("matching synthetic prefix should be inherited"); + let input = prepared["input"].as_array().expect("prepared input"); + assert_eq!(input.len(), 1); + assert_eq!(input[0]["type"], "function_call_output"); + } + + #[test] + fn responses_lite_continuation_removes_an_instructions_only_synthetic_prefix() { + let first = json!({ + "type": "response.create", + "model": "gpt-5.6-sol", + "instructions": "developer instructions", + "input": [{"role": "user", "content": "hello"}] + }); + let stored = ResponsesLiteStaticConfig::from_response_create(&first); + let continuation = json!({ + "type": "response.create", + "model": "gpt-5.6-sol", + "previous_response_id": "resp_1", + "input": [ + {"type": "message", "role": "developer", "content": [ + {"type": "input_text", "text": "developer instructions"} + ]}, + {"type": "function_call_output", "call_id": "call_1", "output": "ok"} + ] + }); + + let prepared = prepare_responses_lite_continuation(&continuation, &stored) + .expect("matching instructions-only prefix should be inherited"); + let input = prepared["input"].as_array().expect("prepared input"); + assert_eq!(input.len(), 1); + assert_eq!(input[0]["type"], "function_call_output"); } } diff --git a/apps/aether-gateway/src/handlers/proxy/websocket/responses/session.rs b/apps/aether-gateway/src/handlers/proxy/websocket/responses/session.rs index 7d7f75c0d..266142c48 100644 --- a/apps/aether-gateway/src/handlers/proxy/websocket/responses/session.rs +++ b/apps/aether-gateway/src/handlers/proxy/websocket/responses/session.rs @@ -14,8 +14,12 @@ use serde_json::Value; use uuid::Uuid; use super::adapter::resolve_responses_websocket_adapter; +use super::binding::UpstreamBindingIdentity; use super::client::consume_response_create_rate_limit; use super::connection::{relay_bound_connection, wait_for_connection_permit_loss}; +use super::continuation::{ + ResponsesWebSocketContinuationRecord, ResponsesWebSocketContinuationRegistry, +}; use super::control::resolve_responses_websocket_turn_control; use super::lifecycle::{ await_pending_adapter_observation, await_pending_turn_finalization, @@ -26,16 +30,20 @@ use super::ownership::{ await_owned_responses_websocket_plan, begin_responses_websocket_turn_with_planned_lease, spawn_owned_responses_websocket_plan, OwnedResponsesWebSocketDecision, }; -use super::redaction::redact_responses_websocket_client_event; +use super::redaction::redact_responses_websocket_client_event_with_reasoning_replay_policy; use super::relay_policy::{fatal_relay_policy, FatalRelaySignal}; use super::request::{ - build_planning_parts, planned_response_create_event, validated_response_create_model, + build_planning_parts, planned_request_uses_codex_responses_lite, planned_response_create_event, + prepare_responses_lite_continuation, validate_response_create_previous_response_id, + validate_response_create_stream_id_support, validated_named_stream_id, + validated_response_create_model, ResponsesLiteStaticConfig, }; use super::state::BoundResponsesConnection; use super::turn::{prepare_responses_websocket_turn_decision, ResponsesWebSocketTurnOutcome}; use super::turn_state::LogicalTurn; use super::upstream::{bind_responses_upstream, close_bound_upstream}; +use crate::ai_serving::ResponsesWebSocketPinnedCandidate; use crate::handlers::proxy::websocket::ingress::{ WebSocketConnectionLog, WebSocketConnectionLogSpec, WebSocketRequestContext, }; @@ -45,10 +53,13 @@ use crate::handlers::proxy::websocket::session::{ }; use crate::handlers::proxy::websocket::transport::{ close_client_socket, send_gateway_error, send_gateway_error_with_status, + send_gateway_error_with_stream_id, send_responses_websocket_error_with_param, }; +use crate::privacy::RedactionSession; use crate::AppState; const RESPONSES_WEBSOCKET_LOG_TARGET: &str = "aether_gateway::handlers::proxy::responses_ws"; +const CONTINUATION_LOOKUP_TIMEOUT: std::time::Duration = std::time::Duration::from_millis(500); const RESPONSES_CONNECTION_LOG_SPEC: WebSocketConnectionLogSpec = WebSocketConnectionLogSpec { opened_event_name: "responses_websocket_connection_opened", closed_event_name: "responses_websocket_connection_closed", @@ -74,6 +85,9 @@ enum InitialMessageError { MissingResponseCreate, MissingModel, InvalidModel, + InvalidPreviousResponseId, + InvalidStreamId, + UnsupportedStreamId, } #[derive(Debug, Clone, Copy, PartialEq, Eq)] @@ -113,10 +127,11 @@ impl InitialMessageFrameMetadata { } } -#[derive(Debug, Clone, Copy, PartialEq, Eq)] +#[derive(Debug, Clone, PartialEq, Eq)] struct InitialMessageFailure { error: InitialMessageError, last_frame: Option, + stream_id: Option, } impl InitialMessageFailure { @@ -124,7 +139,27 @@ impl InitialMessageFailure { error: InitialMessageError, last_frame: Option, ) -> Self { - Self { error, last_frame } + Self { + error, + last_frame, + stream_id: None, + } + } + + fn for_event( + error: InitialMessageError, + last_frame: Option, + event: &Value, + ) -> Self { + // Preserve a valid lane identity for any request-scoped validation + // error. Invalid `stream_id` values are rejected by the grammar helper + // and are never reflected back to the client. + let stream_id = validated_named_stream_id(event).map(str::to_string); + Self { + error, + last_frame, + stream_id, + } } } @@ -156,6 +191,9 @@ impl InitialMessageError { Self::MissingResponseCreate => "expected_response_create", Self::MissingModel => "response_create_model_required", Self::InvalidModel => "invalid_response_create_model", + Self::InvalidPreviousResponseId => "invalid_response_create_previous_response_id", + Self::InvalidStreamId => "invalid_response_create_stream_id", + Self::UnsupportedStreamId => "responses_websocket_named_stream_unsupported", } } @@ -167,6 +205,9 @@ impl InitialMessageError { Self::MissingResponseCreate | Self::MissingModel | Self::InvalidModel => { CLOSE_POLICY_VIOLATION } + Self::InvalidPreviousResponseId => CLOSE_POLICY_VIOLATION, + Self::InvalidStreamId => CLOSE_POLICY_VIOLATION, + Self::UnsupportedStreamId => CLOSE_POLICY_VIOLATION, } } @@ -190,6 +231,15 @@ impl InitialMessageError { Self::InvalidModel => { Some("response.create.model must be a non-empty string no longer than 256 bytes") } + Self::InvalidPreviousResponseId => { + Some("response.create.previous_response_id must be null or a non-empty string") + } + Self::InvalidStreamId => Some( + "response.create.stream_id must be 1-256 ASCII letters, numbers, underscores, hyphens, or periods", + ), + Self::UnsupportedStreamId => Some( + "Aether currently supports only the implicit default WebSocket lane; omit response.create.stream_id", + ), } } @@ -203,6 +253,9 @@ impl InitialMessageError { Self::MissingResponseCreate => "unexpected_event_type", Self::MissingModel => "missing_model", Self::InvalidModel => "invalid_model", + Self::InvalidPreviousResponseId => "invalid_previous_response_id", + Self::InvalidStreamId => "invalid_stream_id", + Self::UnsupportedStreamId => "unsupported_stream_id", } } @@ -211,7 +264,7 @@ impl InitialMessageError { } } -fn initial_message_diagnostic(failure: InitialMessageFailure) -> InitialMessageDiagnostic { +fn initial_message_diagnostic(failure: &InitialMessageFailure) -> InitialMessageDiagnostic { InitialMessageDiagnostic { error_code: failure.error.code(), error_kind: failure.error.kind(), @@ -225,7 +278,7 @@ fn initial_message_diagnostic(failure: InitialMessageFailure) -> InitialMessageD fn log_initial_message_failure( context: &WebSocketRequestContext, - failure: InitialMessageFailure, + failure: &InitialMessageFailure, upgraded_at: std::time::Instant, ) { let diagnostic = initial_message_diagnostic(failure); @@ -360,8 +413,14 @@ async fn bootstrap_responses_websocket( Ok(value) => value, Err(failure) => { if let Some(client_message) = failure.error.client_message() { - log_initial_message_failure(context, failure, upgraded_at); - send_gateway_error(client_socket, failure.error.code(), client_message).await; + log_initial_message_failure(context, &failure, upgraded_at); + send_gateway_error_with_stream_id( + client_socket, + failure.error.code(), + client_message, + failure.stream_id.as_deref(), + ) + .await; close_client_socket( client_socket, failure.error.close_code(), @@ -373,6 +432,18 @@ async fn bootstrap_responses_websocket( } }; + let initial_previous_response_id = first_event + .get("previous_response_id") + .and_then(Value::as_str) + .filter(|value| !value.trim().is_empty()) + .map(str::to_string); + + // Keep the response-chain identity on the client's raw configuration. + // Redaction sentinels may rotate between turns, but that must not turn + // identical plaintext tools/instructions into a synthetic config change. + let raw_responses_lite_static_config = + ResponsesLiteStaticConfig::from_response_create(&first_event); + let planning_parts = build_planning_parts(context); let turn_control = match resolve_responses_websocket_turn_control( &state, @@ -418,7 +489,7 @@ async fn bootstrap_responses_websocket( close_client_socket(client_socket, CLOSE_TRY_AGAIN, "rate_limit_exceeded").await; return None; } - Err(()) => { + Err(_) => { warn!( event_name = "responses_websocket_rate_limit_check_failed", log_type = "ops", @@ -444,16 +515,152 @@ async fn bootstrap_responses_websocket( } } + // A cross-socket response chain must be owned by this exact live + // authenticated principal. Missing, expired, corrupt or unavailable state + // fails closed; allowing the normal scheduler to choose a provider/key + // would disclose an opaque response ID to an unrelated account. + let continuation_record = if let Some(previous_response_id) = + initial_previous_response_id.as_deref() + { + let Some(auth_context) = turn_control.decision.auth_context.as_ref() else { + reject_initial_previous_response(client_socket).await; + return None; + }; + let registry = ResponsesWebSocketContinuationRegistry::new(state.runtime_state.as_ref()); + match tokio::time::timeout( + CONTINUATION_LOOKUP_TIMEOUT, + registry.lookup( + auth_context.user_id.as_str(), + auth_context.api_key_id.as_str(), + previous_response_id, + ), + ) + .await + { + Ok(Ok(Some(record))) + if record.client_model() + == first_event + .get("model") + .and_then(Value::as_str) + .map(str::trim) + .unwrap_or_default() + && !record.has_connection_local_redaction() => + { + Some(record) + } + Ok(Ok(Some(record))) => { + warn!( + event_name = "responses_websocket_continuation_registry_rejected", + log_type = "ops", + transport = WEBSOCKET_LOG_TRANSPORT, + websocket = true, + trace_id = %context.trace_id, + user_id = %auth_context.user_id, + api_key_id = %auth_context.api_key_id, + provider_id = %record.pinned_candidate().provider_id(), + endpoint_id = %record.pinned_candidate().endpoint_id(), + key_id = %record.pinned_candidate().key_id(), + reason = if record.has_connection_local_redaction() { + "connection_local_redaction_state_unavailable" + } else { + "client_model_mismatch" + }, + "gateway rejected a cross-socket Responses continuation" + ); + reject_initial_previous_response(client_socket).await; + return None; + } + Ok(Ok(None)) => { + warn!( + event_name = "responses_websocket_continuation_registry_miss", + log_type = "ops", + transport = WEBSOCKET_LOG_TRANSPORT, + websocket = true, + trace_id = %context.trace_id, + user_id = %auth_context.user_id, + api_key_id = %auth_context.api_key_id, + reason = "not_found_or_expired", + "gateway could not prove ownership of a cross-socket Responses continuation" + ); + reject_initial_previous_response(client_socket).await; + return None; + } + Ok(Err(error)) => { + warn!( + event_name = "responses_websocket_continuation_registry_lookup_failed", + log_type = "ops", + transport = WEBSOCKET_LOG_TRANSPORT, + websocket = true, + trace_id = %context.trace_id, + user_id = %auth_context.user_id, + api_key_id = %auth_context.api_key_id, + reason = error.kind(), + "gateway failed closed while looking up a cross-socket Responses continuation" + ); + reject_initial_previous_response(client_socket).await; + return None; + } + Err(_) => { + warn!( + event_name = "responses_websocket_continuation_registry_lookup_failed", + log_type = "ops", + transport = WEBSOCKET_LOG_TRANSPORT, + websocket = true, + trace_id = %context.trace_id, + user_id = %auth_context.user_id, + api_key_id = %auth_context.api_key_id, + reason = "timeout", + timeout_ms = CONTINUATION_LOOKUP_TIMEOUT.as_millis() as u64, + "gateway timed out looking up a cross-socket Responses continuation" + ); + reject_initial_previous_response(client_socket).await; + return None; + } + } + } else { + None + }; + + // Validate and remove any repeated Lite static prefix against the stored + // chain identity before PII redaction rotates per-turn sentinels. + let first_event = match continuation_record + .as_ref() + .and_then(ResponsesWebSocketContinuationRecord::responses_lite_static_config) + { + Some(stored) => match prepare_responses_lite_continuation(&first_event, stored) { + Ok(prepared) => prepared, + Err(_) => { + warn!( + event_name = "responses_websocket_continuation_static_contract_rejected", + log_type = "ops", + transport = WEBSOCKET_LOG_TRANSPORT, + websocket = true, + trace_id = %context.trace_id, + "gateway rejected changed Responses Lite static configuration on a continuation" + ); + reject_initial_previous_response(client_socket).await; + return None; + } + }, + None => first_event, + }; + // 请求侧脱敏必须在规划之前完成，而且这一轮只在这里做一次：planner 会把这份 // body 写进 upstream 请求体和审计 original_request_body，绑定上游的首条 // response.create 也从它派生。脱敏失败时直接断开，绝不退回原文发上游。 - let redacted_first_event = redact_responses_websocket_client_event( - &state, - &planning_parts, - &turn_control.decision, - &first_event, - ) - .await; + let reasoning_replay_policy = continuation_record + .as_ref() + .map(ResponsesWebSocketContinuationRecord::reasoning_replay_policy) + .unwrap_or_default(); + let redacted_first_event = + redact_responses_websocket_client_event_with_reasoning_replay_policy( + &state, + &planning_parts, + &turn_control.decision, + &first_event, + reasoning_replay_policy, + ) + .await; // 首轮的 mask session 要活到响应帧还原，但连接此刻还没绑定，只能先接住， // 等 `bind_responses_upstream` 之后登记到连接上。 let (first_event, first_turn_redaction_session) = match redacted_first_event { @@ -486,6 +693,21 @@ async fn bootstrap_responses_websocket( } }; + let pinned_candidate = match initial_continuation_planning_candidate( + initial_previous_response_id.is_some(), + continuation_record + .as_ref() + .map(|record| record.pinned_candidate().clone()), + ) { + Ok(candidate) => candidate, + Err(_) => { + // Keep this invariant next to the planner boundary as a second + // fail-closed guard: an unproved response ID must never enter the + // ordinary scheduler and land on an unrelated provider/key. + reject_initial_previous_response(client_socket).await; + return None; + } + }; let planned = match await_owned_responses_websocket_plan(spawn_owned_responses_websocket_plan( state.clone(), planning_parts, @@ -495,12 +717,24 @@ async fn bootstrap_responses_websocket( first_event.clone(), None, None, - None, + pinned_candidate, )) .await { Ok(Some(decision)) => decision, Ok(None) => { + if continuation_record.is_some() { + warn!( + event_name = "responses_websocket_continuation_pinned_candidate_unavailable", + log_type = "ops", + transport = WEBSOCKET_LOG_TRANSPORT, + websocket = true, + trace_id = %context.trace_id, + "gateway could not revalidate the registered Responses continuation binding" + ); + reject_initial_previous_response(client_socket).await; + return None; + } send_gateway_error_with_status( client_socket, 503, @@ -547,35 +781,85 @@ async fn bootstrap_responses_websocket( planning_parts, planned_lease, } = planned; - let adapter = resolve_responses_websocket_adapter(planned.adapter); + let adapter_kind = planned.adapter; + let adapter = resolve_responses_websocket_adapter(adapter_kind); let normalization = planned.normalization; let decision = planned.execution; - let first_provider_event = match planned_response_create_event(&decision, &first_event) - .and_then(|event| { - serde_json::from_str::(&event).map_err(|_| "responses_websocket_request_invalid") - }) { - Ok(event) => event, - Err(code) => { + if let Some(record) = continuation_record.as_ref() { + let planned_candidate = ResponsesWebSocketPinnedCandidate::from_decision(&decision); + let planned_provider_model = decision + .provider_request_body + .as_ref() + .and_then(|body| body.get("model")) + .and_then(Value::as_str) + .map(str::trim) + .filter(|value| !value.is_empty()) + .or_else(|| { + decision + .mapped_model + .as_deref() + .map(str::trim) + .filter(|value| !value.is_empty()) + }); + let planned_binding = UpstreamBindingIdentity::from_decision(adapter, &decision).ok(); + let planned_uses_responses_lite = + planned_request_uses_codex_responses_lite(&decision, &normalization); + let matches_record = record.adapter() == adapter_kind + && planned_candidate.as_ref() == Some(record.pinned_candidate()) + && planned_provider_model == Some(record.provider_model()) + && responses_lite_contract_modes_match( + record.responses_lite_static_config().is_some(), + planned_uses_responses_lite, + ) + && planned_binding + .as_ref() + .is_some_and(|binding| record.matches_contract(binding, &normalization)); + if !matches_record { planned_lease.release().await; warn!( - event_name = "responses_websocket_initial_event_normalization_failed", + event_name = "responses_websocket_continuation_binding_rejected", log_type = "ops", transport = WEBSOCKET_LOG_TRANSPORT, websocket = true, trace_id = %context.trace_id, - error_code = code, - "gateway could not normalize the initial Responses WebSocket event" + provider_id = %record.pinned_candidate().provider_id(), + endpoint_id = %record.pinned_candidate().endpoint_id(), + key_id = %record.pinned_candidate().key_id(), + "gateway rejected a cross-socket continuation after pinned planning changed its contract" ); - send_gateway_error( - client_socket, - code, - "Gateway could not prepare the Responses response.create event", - ) - .await; - close_client_socket(client_socket, CLOSE_POLICY_VIOLATION, code).await; + reject_initial_previous_response(client_socket).await; return None; } - }; + } + let first_provider_event = + match planned_response_create_event(&decision, &normalization, &first_event).and_then( + |event| { + serde_json::from_str::(&event) + .map_err(|_| "responses_websocket_request_invalid") + }, + ) { + Ok(event) => event, + Err(code) => { + planned_lease.release().await; + warn!( + event_name = "responses_websocket_initial_event_normalization_failed", + log_type = "ops", + transport = WEBSOCKET_LOG_TRANSPORT, + websocket = true, + trace_id = %context.trace_id, + error_code = code, + "gateway could not normalize the initial Responses WebSocket event" + ); + send_gateway_error( + client_socket, + code, + "Gateway could not prepare the Responses response.create event", + ) + .await; + close_client_socket(client_socket, CLOSE_POLICY_VIOLATION, code).await; + return None; + } + }; let first_logical_turn_id = Uuid::new_v4().to_string(); let first_turn_decision = prepare_responses_websocket_turn_decision( &decision, @@ -647,19 +931,83 @@ async fn bootstrap_responses_websocket( return None; } }; + if bound.responses_lite_static_config.is_some() { + bound.responses_lite_static_config = continuation_record + .as_ref() + .and_then(ResponsesWebSocketContinuationRecord::responses_lite_static_config) + .cloned() + .or(Some(raw_responses_lite_static_config)); + } + if let Some(previous_response_id) = initial_previous_response_id.as_deref() { + // Reaching this point means the principal-scoped registry record and + // the newly planned physical binding were both proved above. Keep the + // parent as persisted ownership: the new physical socket may hydrate + // it, but a failed attempt can evict only its connection-local copy. + bound + .continuation_response_ids + .remember_persisted(previous_response_id); + } first_turn.mark_upstream_request_sent(); first_turn.set_provider_response_headers(bound.upstream_response_headers.clone()); if let Some(session) = first_turn_redaction_session { - bound.redaction_restorer.register(session); + register_initial_redaction_session(&mut bound, session); } bound.turn_state.begin( - LogicalTurn::new(first_event, 1, first_logical_turn_id).with_turn_control(turn_control), + LogicalTurn::new(first_event, 1, first_logical_turn_id) + .with_provider_store(first_provider_event.get("store") == Some(&Value::Bool(true))) + .with_turn_control(turn_control), first_turn, ); Some(bound) } +fn register_initial_redaction_session( + bound: &mut BoundResponsesConnection, + mut session: RedactionSession, +) { + session.set_reasoning_replay_policy(bound.body_normalization.reasoning_replay_policy()); + bound.redaction_restorer.register(session); +} + +fn responses_lite_contract_modes_match( + stored_chain_uses_responses_lite: bool, + planned_request_uses_responses_lite: bool, +) -> bool { + stored_chain_uses_responses_lite == planned_request_uses_responses_lite +} + +fn initial_continuation_planning_candidate( + has_previous_response_id: bool, + registered_candidate: Option, +) -> Result, &'static str> { + match (has_previous_response_id, registered_candidate) { + (false, None) => Ok(None), + (true, Some(candidate)) => Ok(Some(candidate)), + // A registry miss/corruption and an impossible stray record both fail + // closed. Neither state is allowed to turn into an unpinned plan. + (true, None) | (false, Some(_)) => Err("previous_response_not_found"), + } +} + +async fn reject_initial_previous_response(client_socket: &mut WebSocket) { + send_responses_websocket_error_with_param( + client_socket, + 400, + "invalid_request_error", + "previous_response_not_found", + "The previous response is unavailable for this authenticated WebSocket connection", + "previous_response_id", + ) + .await; + close_client_socket( + client_socket, + CLOSE_POLICY_VIOLATION, + "previous_response_not_found", + ) + .await; +} + async fn close_terminated_bootstrap( client_socket: &mut WebSocket, context: &WebSocketRequestContext, @@ -828,7 +1176,7 @@ where InitialMessageFailure::new(InitialMessageError::InvalidJson, last_frame) })?; validate_initial_response_create(&event) - .map_err(|error| InitialMessageFailure::new(error, last_frame))?; + .map_err(|error| InitialMessageFailure::for_event(error, last_frame, &event))?; return Ok((text, event)); } } @@ -844,6 +1192,15 @@ fn validate_initial_response_create(event: &Value) -> Result<(), InitialMessageE .get("model") .ok_or(InitialMessageError::MissingModel)?; validated_response_create_model(model).map_err(|_| InitialMessageError::InvalidModel)?; + validate_response_create_previous_response_id(event) + .map_err(|_| InitialMessageError::InvalidPreviousResponseId)?; + match validate_response_create_stream_id_support(event) { + Ok(()) => {} + Err("invalid_response_create_stream_id") => { + return Err(InitialMessageError::InvalidStreamId); + } + Err(_) => return Err(InitialMessageError::UnsupportedStreamId), + } Ok(()) } @@ -875,11 +1232,15 @@ mod tests { }; use super::super::turn_state::{LogicalTurn, ResponsesTurnState}; use super::super::upstream::bind_responses_upstream; - use crate::ai_serving::{AiExecutionDecision, ResponsesWebSocketBodyNormalization}; + use crate::ai_serving::{ + AiExecutionDecision, OpenAiResponsesReasoningReplayPolicy, + ResponsesWebSocketBodyNormalization, + }; use crate::handlers::proxy::websocket::session::wait_for_optional_deadline; use crate::handlers::proxy::websocket::transport::{ websocket_handshake_headers, websocket_timeouts, websocket_upstream_url, }; + use crate::privacy::{RedactionSession, RedactionSessionConfig}; use axum::extract::ws::{Message, WebSocket, WebSocketUpgrade}; use axum::extract::State; use axum::http::header::{AUTHORIZATION, CONTENT_TYPE}; @@ -987,10 +1348,38 @@ mod tests { 1008, false, ), + ( + InitialMessageError::InvalidPreviousResponseId, + "invalid_response_create_previous_response_id", + "invalid_previous_response_id", + Some("response.create.previous_response_id must be null or a non-empty string"), + 1008, + false, + ), + ( + InitialMessageError::InvalidStreamId, + "invalid_response_create_stream_id", + "invalid_stream_id", + Some( + "response.create.stream_id must be 1-256 ASCII letters, numbers, underscores, hyphens, or periods", + ), + 1008, + false, + ), + ( + InitialMessageError::UnsupportedStreamId, + "responses_websocket_named_stream_unsupported", + "unsupported_stream_id", + Some( + "Aether currently supports only the implicit default WebSocket lane; omit response.create.stream_id", + ), + 1008, + false, + ), ]; for (error, error_code, error_kind, client_message, close_code, timed_out) in cases { - let diagnostic = initial_message_diagnostic(InitialMessageFailure::new(error, None)); + let diagnostic = initial_message_diagnostic(&InitialMessageFailure::new(error, None)); assert_eq!(diagnostic.error_code, error_code); assert_eq!(diagnostic.error_kind, error_kind); assert_eq!(diagnostic.client_message, client_message); @@ -1011,7 +1400,7 @@ mod tests { let secret_body = r#"{"type":"not-response.create","token":"must-not-log"}"#; let frame = Message::Text(secret_body.to_string().into()); let metadata = InitialMessageFrameMetadata::from_message(&frame); - let diagnostic = initial_message_diagnostic(InitialMessageFailure::new( + let diagnostic = initial_message_diagnostic(&InitialMessageFailure::new( InitialMessageError::MissingResponseCreate, Some(metadata), )); @@ -1211,8 +1600,10 @@ mod tests { "stream": true, "background": true, })); + let normalization = ResponsesWebSocketBodyNormalization::for_tests("provider-model"); let event = planned_response_create_event( &decision, + &normalization, &json!({ "type": "response.create", "model": "public-model", @@ -1309,6 +1700,168 @@ mod tests { ); } + #[test] + fn malformed_initial_previous_response_id_is_rejected_before_planning() { + for previous_response_id in [json!(""), json!(42), json!({"id": "resp_1"})] { + let initial = json!({ + "type": "response.create", + "model": "gpt-5.6-sol", + "previous_response_id": previous_response_id, + "input": [], + }); + assert!(matches!( + super::validate_initial_response_create(&initial), + Err(super::InitialMessageError::InvalidPreviousResponseId) + )); + } + + let named = json!({ + "type": "response.create", + "model": "gpt-5.6-sol", + "stream_id": "main", + "previous_response_id": "", + "input": [], + }); + let failure = super::InitialMessageFailure::for_event( + super::InitialMessageError::InvalidPreviousResponseId, + None, + &named, + ); + assert_eq!(failure.stream_id.as_deref(), Some("main")); + } + + #[test] + fn valid_previous_response_can_start_a_new_socket() { + let initial = json!({ + "type": "response.create", + "model": "gpt-5.6-sol", + "previous_response_id": "resp_existing", + "input": [{"type": "function_call_output", "call_id": "call_1", "output": "ok"}], + }); + assert!(super::validate_initial_response_create(&initial).is_ok()); + + let null_previous = json!({ + "type": "response.create", + "model": "gpt-5.6-sol", + "previous_response_id": null, + "input": [{"role": "user", "content": "new chain"}], + }); + assert!(super::validate_initial_response_create(&null_previous).is_ok()); + } + + #[test] + fn cross_socket_registry_miss_never_falls_through_to_random_provider_planning() { + let pinned = crate::ai_serving::ResponsesWebSocketPinnedCandidate::new( + "provider-original", + "endpoint-original", + "key-original", + ) + .expect("valid pinned candidate"); + + assert_eq!( + super::initial_continuation_planning_candidate(false, None), + Ok(None), + "a genuinely new response may use the ordinary planner" + ); + assert_eq!( + super::initial_continuation_planning_candidate(true, Some(pinned.clone())), + Ok(Some(pinned)), + "a proved continuation must retain its exact provider/endpoint/key" + ); + assert_eq!( + super::initial_continuation_planning_candidate(true, None), + Err("previous_response_not_found"), + "a registry miss must be rejected before an unpinned planner can choose another key" + ); + } + + #[test] + fn cross_socket_continuation_rejects_an_effective_lite_mode_change() { + let mut decision: AiExecutionDecision = serde_json::from_value(json!({ + "action": "local", + "provider_type": "codex", + "provider_api_format": "openai:responses", + "provider_request_headers": {} + })) + .expect("minimal Codex decision"); + decision.provider_request_headers.insert( + crate::ai_serving::CODEX_RESPONSES_LITE_HEADER.to_string(), + "true".to_string(), + ); + let normalization = ResponsesWebSocketBodyNormalization::for_tests("gpt-5.6-sol") + .with_provider_type_for_tests("codex"); + + let effective_lite = + super::planned_request_uses_codex_responses_lite(&decision, &normalization); + assert!(effective_lite); + assert!(super::responses_lite_contract_modes_match( + true, + effective_lite + )); + + // A non-null context_management object suppresses the converged Lite + // contract/header even though the model capability remains enabled. + // A chain whose stored prefix used Lite must not cross that boundary. + decision.provider_request_body = Some(json!({ + "model": "gpt-5.6-sol", + "context_management": {"compact_threshold": 1_000} + })); + let effective_lite = + super::planned_request_uses_codex_responses_lite(&decision, &normalization); + assert!(!effective_lite); + assert!(!super::responses_lite_contract_modes_match( + true, + effective_lite + )); + assert!(super::responses_lite_contract_modes_match( + false, + effective_lite + )); + } + + #[test] + fn initial_named_stream_is_rejected_before_planning() { + let initial = json!({ + "type": "response.create", + "model": "gpt-5.6-sol", + "stream_id": "main", + "input": [], + }); + assert!(matches!( + super::validate_initial_response_create(&initial), + Err(super::InitialMessageError::UnsupportedStreamId) + )); + + let failure = super::InitialMessageFailure::for_event( + super::InitialMessageError::UnsupportedStreamId, + None, + &initial, + ); + assert_eq!(failure.stream_id.as_deref(), Some("main")); + } + + #[test] + fn malformed_initial_stream_id_is_rejected_before_planning() { + for stream_id in [json!(null), json!(""), json!("not/a/lane"), json!(42)] { + let initial = json!({ + "type": "response.create", + "model": "gpt-5.6-sol", + "stream_id": stream_id, + "input": [], + }); + assert!(matches!( + super::validate_initial_response_create(&initial), + Err(super::InitialMessageError::InvalidStreamId) + )); + let failure = super::InitialMessageFailure::for_event( + super::InitialMessageError::InvalidStreamId, + None, + &initial, + ); + assert_eq!(failure.stream_id, None); + } + } + #[test] fn model_at_the_identifier_limit_remains_valid() { let model = "m".repeat(256); @@ -1669,7 +2222,9 @@ mod tests { provider_model: "gpt-5.6-sol".to_string(), decision_template: decision, body_normalization: ResponsesWebSocketBodyNormalization::for_tests("gpt-5.6-sol"), + responses_lite_static_config: None, binding_identity, + continuation_response_ids: Default::default(), // Replanning：logical turn 在、attempt 不在。重放安全与配额排除都只看 // logical turn，所以这些用例不需要真实 socket 或真实 attempt。 turn_state: ResponsesTurnState::Replanning { @@ -1689,6 +2244,57 @@ mod tests { } } + #[test] + fn initial_deepseek_binding_upgrades_redaction_restore_policy() { + let mut session = RedactionSession::new(RedactionSessionConfig::new( + b"initial-deepseek-redaction-test".to_vec(), + 300, + 600, + )); + let sentinel = session.redact_text("alice@example.com").text; + let provider_event = json!({ + "type": "response.completed", + "response": { + "output": [{ + "type": "reasoning", + "encrypted_content": "provider-owned-state", + "content": [{ + "type": "reasoning_text", + "text": format!("opaque replay {sentinel}") + }] + }] + } + }); + + let mut ordinary_bound = sample_bound_for_rebind_safety(); + super::register_initial_redaction_session(&mut ordinary_bound, session.clone()); + let restored = ordinary_bound + .redaction_restorer + .restore_provider_frame_text(&provider_event) + .expect("ordinary OpenAI replay policy should restore response text"); + let restored: serde_json::Value = + serde_json::from_str(&restored).expect("restored provider event should remain JSON"); + assert_eq!( + restored["response"]["output"][0]["content"][0]["text"], + "opaque replay alice@example.com" + ); + + let mut deepseek_bound = sample_bound_for_rebind_safety(); + deepseek_bound.body_normalization = + ResponsesWebSocketBodyNormalization::for_tests("deepseek-reasoner") + .with_reasoning_replay_policy_for_tests( + OpenAiResponsesReasoningReplayPolicy::DeepSeekOpaque, + ); + super::register_initial_redaction_session(&mut deepseek_bound, session); + assert!( + deepseek_bound + .redaction_restorer + .restore_provider_frame_text(&provider_event) + .is_none(), + "the authenticated DeepSeek binding must keep opaque reasoning state byte-identical" + ); + } + /// 用 mpsc 驱动的 FakeSocket，实现 Stream + Sink 两个 trait。 /// 测试侧通过 tx 注入消息，通过 pong_rx 观察 Pong 回包。 struct FakeSocket { diff --git a/apps/aether-gateway/src/handlers/proxy/websocket/responses/state.rs b/apps/aether-gateway/src/handlers/proxy/websocket/responses/state.rs index 941948c13..7025550e1 100644 --- a/apps/aether-gateway/src/handlers/proxy/websocket/responses/state.rs +++ b/apps/aether-gateway/src/handlers/proxy/websocket/responses/state.rs @@ -4,16 +4,93 @@ //! connection may survive many `response.create` turns, while the turn //! lifecycle and upstream binding are replaced independently. -use std::collections::{BTreeMap, BTreeSet}; +use std::collections::{BTreeMap, BTreeSet, VecDeque}; use tokio::task::JoinHandle; use super::adapter::{ResponsesWebSocketDrainDirective, ResponsesWebSocketProtocolAdapter}; use super::binding::UpstreamBindingIdentity; use super::redaction::ResponsesWebSocketRedactionRestorer; +use super::request::ResponsesLiteStaticConfig; use super::turn_state::ResponsesTurnState; use crate::ai_serving::{AiExecutionDecision, ResponsesWebSocketBodyNormalization}; const EXHAUSTED_KEY_EXCLUSION_FALLBACK_SECONDS: u64 = 300; +const MAX_CONNECTION_LOCAL_CONTINUATION_RESPONSE_IDS: usize = 1_024; + +#[derive(Debug, Default)] +struct BoundedContinuationResponseIds { + ids: BTreeSet, + insertion_order: VecDeque, +} + +impl BoundedContinuationResponseIds { + fn contains(&self, response_id: &str) -> bool { + self.ids.contains(response_id) + } + + fn remember(&mut self, response_id: &str) { + if !self.ids.insert(response_id.to_string()) { + return; + } + self.insertion_order.push_back(response_id.to_string()); + while self.ids.len() > MAX_CONNECTION_LOCAL_CONTINUATION_RESPONSE_IDS { + let Some(oldest) = self.insertion_order.pop_front() else { + break; + }; + self.ids.remove(oldest.as_str()); + } + } + + fn forget(&mut self, response_id: &str) { + if self.ids.remove(response_id) { + self.insertion_order + .retain(|remembered| remembered != response_id); + } + } + + fn clear(&mut self) { + self.ids.clear(); + self.insertion_order.clear(); + } +} + +/// Principal-proved response IDs for the currently bound response chain. +/// +/// Connection-local IDs prove that the provider's current physical socket can +/// resolve a parent, including `store=false` responses. Persisted IDs are kept +/// separately after a successful principal-scoped registry write (or a proved +/// cross-socket bootstrap), so a 4xx/5xx eviction of the provider's local cache +/// does not suppress the provider's documented `store=true` hydration fallback. +/// Starting an independent chain or replacing the physical upstream clears +/// both bounded sets. +#[derive(Debug, Default)] +pub(super) struct ContinuationResponseIds { + connection_local: BoundedContinuationResponseIds, + persisted: BoundedContinuationResponseIds, +} + +impl ContinuationResponseIds { + pub(super) fn contains(&self, response_id: &str) -> bool { + self.connection_local.contains(response_id) || self.persisted.contains(response_id) + } + + pub(super) fn remember_connection_local(&mut self, response_id: &str) { + self.connection_local.remember(response_id); + } + + pub(super) fn remember_persisted(&mut self, response_id: &str) { + self.persisted.remember(response_id); + } + + pub(super) fn forget_connection_local(&mut self, response_id: &str) { + self.connection_local.forget(response_id); + } + + pub(super) fn clear(&mut self) { + self.connection_local.clear(); + self.persisted.clear(); + } +} /// All mutable state associated with the physical upstream connection. pub(super) struct BoundResponsesConnection { @@ -26,7 +103,15 @@ pub(super) struct BoundResponsesConnection { /// turns, which must not re-enter the planner. Replaced whenever the /// binding or its decision is replaced. pub(super) body_normalization: ResponsesWebSocketBodyNormalization, + /// Static Responses Lite configuration already represented in the current + /// response chain. A continuation may repeat it, but may not append a + /// changed synthetic prefix to the inherited history. + pub(super) responses_lite_static_config: Option, pub(super) binding_identity: UpstreamBindingIdentity, + /// IDs observed on the current physical upstream and current independent + /// response chain. A continuation must reference one of these IDs; the + /// cross-socket bootstrap path seeds the already registry-proved parent. + pub(super) continuation_response_ids: ContinuationResponseIds, /// 这条连接上「有没有正在进行的 logical turn」的唯一事实来源。 pub(super) turn_state: ResponsesTurnState, /// 这条连接迄今 mask 出来的映射，用于把 provider 事件里的占位符换回真实值。 @@ -104,3 +189,63 @@ impl ExhaustedResponsesWebSocketExclusions { .retain(|_, expires_at| *expires_at > now_unix_secs); } } + +#[cfg(test)] +mod tests { + use super::{ContinuationResponseIds, MAX_CONNECTION_LOCAL_CONTINUATION_RESPONSE_IDS}; + + #[test] + fn continuation_ids_are_bounded_and_clear_with_the_chain() { + let mut ids = ContinuationResponseIds::default(); + for index in 0..=MAX_CONNECTION_LOCAL_CONTINUATION_RESPONSE_IDS { + ids.remember_connection_local(format!("resp_{index}").as_str()); + } + + assert!(!ids.contains("resp_0")); + assert!(ids.contains("resp_1")); + assert!( + ids.contains(format!("resp_{MAX_CONNECTION_LOCAL_CONTINUATION_RESPONSE_IDS}").as_str()) + ); + assert_eq!( + ids.connection_local.ids.len(), + MAX_CONNECTION_LOCAL_CONTINUATION_RESPONSE_IDS + ); + assert_eq!( + ids.connection_local.insertion_order.len(), + MAX_CONNECTION_LOCAL_CONTINUATION_RESPONSE_IDS + ); + + ids.forget_connection_local("resp_1"); + assert!(!ids.contains("resp_1")); + assert_eq!( + ids.connection_local.insertion_order.len(), + MAX_CONNECTION_LOCAL_CONTINUATION_RESPONSE_IDS - 1 + ); + + ids.clear(); + assert!(ids.connection_local.ids.is_empty()); + assert!(ids.connection_local.insertion_order.is_empty()); + assert!(ids.persisted.ids.is_empty()); + assert!(ids.persisted.insertion_order.is_empty()); + } + + #[test] + fn local_eviction_keeps_only_registered_persisted_ownership() { + let mut ids = ContinuationResponseIds::default(); + ids.remember_connection_local("resp_store_false"); + ids.remember_connection_local("resp_store_true"); + ids.remember_persisted("resp_store_true"); + + ids.forget_connection_local("resp_store_false"); + ids.forget_connection_local("resp_store_true"); + + assert!( + !ids.contains("resp_store_false"), + "store=false has no persisted hydration fallback after local eviction" + ); + assert!( + ids.contains("resp_store_true"), + "a registered store=true parent remains eligible for persisted hydration" + ); + } +} diff --git a/apps/aether-gateway/src/handlers/proxy/websocket/responses/turn_state.rs b/apps/aether-gateway/src/handlers/proxy/websocket/responses/turn_state.rs index 859512cec..8fbb69e05 100644 --- a/apps/aether-gateway/src/handlers/proxy/websocket/responses/turn_state.rs +++ b/apps/aether-gateway/src/handlers/proxy/websocket/responses/turn_state.rs @@ -20,6 +20,10 @@ use super::request::response_create_has_previous_response_id; #[derive(Debug, Clone)] pub(super) struct LogicalTurn { pub(super) client_event: Value, + /// Effective `store` after provider body rules and WebSocket framing. Only + /// an explicit provider-side `true` permits cross-connection registry + /// state; false or absent remains ZDR/connection-local. + pub(super) provider_store: bool, pub(super) turn_index: u64, pub(super) logical_turn_id: String, pub(super) turn_attempt: u32, @@ -35,6 +39,7 @@ impl LogicalTurn { pub(super) fn new(client_event: Value, turn_index: u64, logical_turn_id: String) -> Self { Self { client_event, + provider_store: false, turn_index, logical_turn_id, turn_attempt: 1, @@ -49,6 +54,11 @@ impl LogicalTurn { self } + pub(super) fn with_provider_store(mut self, provider_store: bool) -> Self { + self.provider_store = provider_store; + self + } + pub(super) fn quota_retry_block_reason(&self) -> Option<&'static str> { if self.retry_attempted { Some("quota_retry_already_attempted") diff --git a/apps/aether-gateway/src/handlers/proxy/websocket/responses/upstream.rs b/apps/aether-gateway/src/handlers/proxy/websocket/responses/upstream.rs index 31f7384df..8644588c6 100644 --- a/apps/aether-gateway/src/handlers/proxy/websocket/responses/upstream.rs +++ b/apps/aether-gateway/src/handlers/proxy/websocket/responses/upstream.rs @@ -8,8 +8,10 @@ use wreq::ws::message::Message as WreqWsMessage; use super::adapter::ResponsesWebSocketProtocolAdapter; use super::binding::{UpstreamBindingIdentity, UpstreamBindingIdentityError}; use super::redaction::ResponsesWebSocketRedactionRestorer; -use super::request::planned_response_create_event; -use super::state::{BoundResponsesConnection, ExhaustedResponsesWebSocketExclusions}; +use super::request::{planned_request_uses_codex_responses_lite, planned_response_create_event}; +use super::state::{ + BoundResponsesConnection, ContinuationResponseIds, ExhaustedResponsesWebSocketExclusions, +}; use super::turn_state::ResponsesTurnState; use crate::ai_serving::{AiExecutionDecision, ResponsesWebSocketBodyNormalization}; use crate::handlers::proxy::websocket::session::RESPONSES_WEBSOCKET_SESSION_LIMITS; @@ -79,7 +81,7 @@ async fn bind_responses_upstream_inner( adapter.upstream_errors(), ) .await?; - let first_event = planned_response_create_event(decision, initial_event)?; + let first_event = planned_response_create_event(decision, &normalization, initial_event)?; send_upstream_message(&mut upstream.socket, WreqWsMessage::text(first_event)) .await .map_err(|_| "responses_websocket_initial_send_failed")?; @@ -108,6 +110,11 @@ async fn bind_responses_upstream_inner( .ok_or("responses_websocket_mapped_model_missing")? .to_string(); + let responses_lite_static_config = + planned_request_uses_codex_responses_lite(decision, &normalization).then(|| { + super::request::ResponsesLiteStaticConfig::from_response_create(initial_event) + }); + Ok(BoundResponsesConnection { upstream: Some(upstream.socket), adapter, @@ -115,7 +122,9 @@ async fn bind_responses_upstream_inner( provider_model, decision_template: decision.clone(), body_normalization: normalization, + responses_lite_static_config, binding_identity, + continuation_response_ids: ContinuationResponseIds::default(), // 首条 response.create 已经发出，但这一轮的 logical turn 和 attempt 由调用方 // 通过 `ResponsesTurnState::begin` 装上：绑定本身不持有记账状态。 turn_state: ResponsesTurnState::Idle, diff --git a/apps/aether-gateway/src/handlers/proxy/websocket/transport.rs b/apps/aether-gateway/src/handlers/proxy/websocket/transport.rs index 19ebd25f1..5b33607f1 100644 --- a/apps/aether-gateway/src/handlers/proxy/websocket/transport.rs +++ b/apps/aether-gateway/src/handlers/proxy/websocket/transport.rs @@ -82,6 +82,7 @@ pub(crate) async fn connect_upstream_websocket( fn websocket_response_headers(headers: &HeaderMap) -> BTreeMap { headers .iter() + .filter(|(name, _)| websocket_response_header_is_safe_to_retain(name)) .filter_map(|(name, value)| { value .to_str() @@ -91,6 +92,24 @@ fn websocket_response_headers(headers: &HeaderMap) -> BTreeMap { .collect() } +fn websocket_response_header_is_safe_to_retain(name: &HeaderName) -> bool { + !matches!( + name.as_str(), + "authorization" + | "proxy-authorization" + | "www-authenticate" + | "proxy-authenticate" + | "authentication-info" + | "proxy-authentication-info" + | "cookie" + | "set-cookie" + | "set-cookie2" + | "x-api-key" + | "api-key" + | "x-goog-api-key" + ) +} + pub(crate) fn websocket_upstream_url( raw: &str, invalid_code: &'static str, @@ -300,7 +319,20 @@ pub(crate) fn responses_websocket_error_event( code: &str, message: &str, ) -> serde_json::Value { - json!({ + responses_websocket_error_event_with_stream_id(status, error_type, code, message, None) +} + +/// Builds a request-scoped Responses error. Callers must supply `stream_id` +/// only after validating the protocol's named-lane grammar; untrusted or +/// malformed identifiers must never be reflected into a provider event. +pub(crate) fn responses_websocket_error_event_with_stream_id( + status: u16, + error_type: &str, + code: &str, + message: &str, + stream_id: Option<&str>, +) -> serde_json::Value { + let mut event = json!({ "type": "error", "status": status, "error": { @@ -308,7 +340,17 @@ pub(crate) fn responses_websocket_error_event( "code": code, "message": message, }, - }) + }); + if let Some(stream_id) = stream_id { + event + .as_object_mut() + .expect("Responses error events are JSON objects") + .insert( + "stream_id".to_string(), + serde_json::Value::String(stream_id.to_string()), + ); + } + event } pub(crate) async fn send_responses_websocket_error( @@ -318,7 +360,49 @@ pub(crate) async fn send_responses_websocket_error( code: &str, message: &str, ) { - let event = responses_websocket_error_event(status, error_type, code, message); + send_responses_websocket_error_with_stream_id( + client_socket, + status, + error_type, + code, + message, + None, + ) + .await; +} + +/// Sends a standard invalid-request error with a bounded, server-owned +/// parameter name. This is used for protocol fields such as +/// `previous_response_id`; no untrusted value is reflected. +pub(crate) async fn send_responses_websocket_error_with_param( + client_socket: &mut WebSocket, + status: u16, + error_type: &str, + code: &str, + message: &str, + param: &'static str, +) { + let mut event = responses_websocket_error_event(status, error_type, code, message); + event["error"]["param"] = serde_json::Value::String(param.to_string()); + send_teardown_message( + client_socket + .send(AxumWsMessage::Text(event.to_string().into())) + .map_err(|_| ()), + ) + .await; +} + +pub(crate) async fn send_responses_websocket_error_with_stream_id( + client_socket: &mut WebSocket, + status: u16, + error_type: &str, + code: &str, + message: &str, + stream_id: Option<&str>, +) { + let event = responses_websocket_error_event_with_stream_id( + status, error_type, code, message, stream_id, + ); send_teardown_message( client_socket .send(AxumWsMessage::Text(event.to_string().into())) @@ -331,13 +415,41 @@ pub(crate) async fn send_gateway_error(client_socket: &mut WebSocket, code: &str send_gateway_error_with_status(client_socket, 400, code, message).await; } +pub(crate) async fn send_gateway_error_with_stream_id( + client_socket: &mut WebSocket, + code: &str, + message: &str, + stream_id: Option<&str>, +) { + send_gateway_error_with_status_and_stream_id(client_socket, 400, code, message, stream_id) + .await; +} + pub(crate) async fn send_gateway_error_with_status( client_socket: &mut WebSocket, status: u16, code: &str, message: &str, ) { - send_responses_websocket_error(client_socket, status, "gateway_error", code, message).await; + send_gateway_error_with_status_and_stream_id(client_socket, status, code, message, None).await; +} + +pub(crate) async fn send_gateway_error_with_status_and_stream_id( + client_socket: &mut WebSocket, + status: u16, + code: &str, + message: &str, + stream_id: Option<&str>, +) { + send_responses_websocket_error_with_stream_id( + client_socket, + status, + "gateway_error", + code, + message, + stream_id, + ) + .await; } pub(crate) async fn close_client_socket(client_socket: &mut WebSocket, code: u16, reason: &str) { @@ -355,9 +467,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, websocket_handshake_headers, - websocket_upstream_url, WebSocketWriteError, RELAY_WRITE_TIMEOUT, TEARDOWN_WRITE_TIMEOUT, + bounded_send, 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 axum::http::HeaderMap; use std::collections::BTreeMap; use std::time::Duration; @@ -391,6 +506,32 @@ mod tests { assert!(TEARDOWN_WRITE_TIMEOUT < RELAY_WRITE_TIMEOUT); } + #[test] + fn upstream_handshake_observability_drops_credential_bearing_headers() { + let mut headers = HeaderMap::new(); + headers.insert("x-codex-primary-used-percent", "10".parse().unwrap()); + headers.insert("x-request-id", "request-123".parse().unwrap()); + headers.insert("set-cookie", "session=secret".parse().unwrap()); + headers.insert("www-authenticate", "Bearer secret".parse().unwrap()); + headers.insert("authentication-info", "nextnonce=secret".parse().unwrap()); + + let retained = websocket_response_headers(&headers); + + assert_eq!( + retained + .get("x-codex-primary-used-percent") + .map(String::as_str), + Some("10") + ); + assert_eq!( + retained.get("x-request-id").map(String::as_str), + Some("request-123") + ); + assert!(!retained.contains_key("set-cookie")); + assert!(!retained.contains_key("www-authenticate")); + assert!(!retained.contains_key("authentication-info")); + } + #[test] fn builds_a_client_compatible_responses_error_event() { let event = responses_websocket_error_event( @@ -408,6 +549,24 @@ mod tests { event["error"]["message"], "Previous response was not found." ); + assert!(event.get("stream_id").is_none()); + } + + #[test] + fn request_scoped_responses_errors_include_the_validated_named_stream() { + let event = responses_websocket_error_event_with_stream_id( + 400, + "gateway_error", + "responses_websocket_named_stream_unsupported", + "Named streams are not supported.", + Some("main-lane_1.test"), + ); + + assert_eq!(event["stream_id"], "main-lane_1.test"); + assert_eq!( + event["error"]["code"], + "responses_websocket_named_stream_unsupported" + ); } #[test] diff --git a/apps/aether-gateway/src/handlers/public/support/models/route.rs b/apps/aether-gateway/src/handlers/public/support/models/route.rs index 39b795646..e874a00fb 100644 --- a/apps/aether-gateway/src/handlers/public/support/models/route.rs +++ b/apps/aether-gateway/src/handlers/public/support/models/route.rs @@ -146,9 +146,8 @@ async fn load_codex_model_cards( event_name = "codex_catalog_aggregate_incomplete", client_version = %client_version.as_str(), target_count = targets.len(), - "Codex catalog aggregation was incomplete; returning an empty remote catalog so the client can use its bundled fallback" + "Codex catalog aggregation was incomplete; serving cards from available last-known-good snapshots" ); - return (Vec::new(), None); } let mut seen_global_models = BTreeSet::new(); let possible_inference_catalogs = rows @@ -210,9 +209,8 @@ async fn load_codex_model_cards( expected_model_count = expected_global_models.len(), projected_model_count = cards.len(), missing_model_count, - "Codex upstream catalogs omitted authorized mappings; returning an empty remote catalog so the client can use its bundled fallback" + "Codex upstream catalogs omitted authorized mappings; serving the available cards without fabricating missing model metadata" ); - return (Vec::new(), None); } if !codex_projected_catalog_fits_response_limits(&cards) { warn!( diff --git a/apps/aether-gateway/src/orchestration/policy.rs b/apps/aether-gateway/src/orchestration/policy.rs index 7c94570d2..f1d3731a3 100644 --- a/apps/aether-gateway/src/orchestration/policy.rs +++ b/apps/aether-gateway/src/orchestration/policy.rs @@ -302,7 +302,8 @@ pub(crate) fn codex_cyber_flag_passthrough_enabled( /// WebSocket upstream. Provider-scoped feature switches remain the source of /// truth; this enum only identifies provider-specific extensions around the /// otherwise standard Responses WebSocket protocol. -#[derive(Debug, Clone, Copy, PartialEq, Eq)] +#[derive(Debug, Clone, Copy, PartialEq, Eq, serde::Serialize, serde::Deserialize)] +#[serde(rename_all = "snake_case")] pub(crate) enum ResponsesWebSocketAdapter { /// A provider that speaks the standard OpenAI Responses WebSocket protocol. Standard, diff --git a/apps/aether-gateway/src/privacy/mod.rs b/apps/aether-gateway/src/privacy/mod.rs index b00c8de01..3d66c47ec 100644 --- a/apps/aether-gateway/src/privacy/mod.rs +++ b/apps/aether-gateway/src/privacy/mod.rs @@ -244,7 +244,19 @@ impl RedactionSessionConfig { } } -pub(crate) type RedactionMaskError = Infallible; +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub(crate) enum RedactionMaskError { + /// Provider-owned reasoning text cannot be rewritten independently from + /// its encrypted continuation state. Reject the request instead of either + /// leaking detected sensitive text or corrupting the provider binding. + SensitiveOpaqueReasoningState, +} + +impl From for RedactionMaskError { + fn from(value: Infallible) -> Self { + match value {} + } +} #[derive(Clone, Copy, Default)] struct RedactionScanState; @@ -307,6 +319,7 @@ pub(crate) struct RedactionSession { mappings: HashMap, sentinel_index: HashMap, collision_corpus: Vec, + preserve_deepseek_opaque_reasoning_state: bool, } impl RedactionSession { @@ -316,13 +329,49 @@ impl RedactionSession { mappings: HashMap::new(), sentinel_index: HashMap::new(), collision_corpus: Vec::new(), + preserve_deepseek_opaque_reasoning_state: false, } } + fn apply_mask_options(&mut self, options: MaskChatRequestOptions) { + self.preserve_deepseek_opaque_reasoning_state = + options.preserve_deepseek_opaque_reasoning_state; + } + + /// Updates the provider-owned reasoning replay policy after the initial + /// request has been planned and its upstream binding authenticated. + /// + /// A new WebSocket chain has to redact its first event before planning, so + /// it starts with the conservative OpenAI item-id policy. Once the + /// planner has selected the provider, response restoration must use that + /// trusted binding's policy as well; otherwise a DeepSeek opaque reasoning + /// item could have a PII sentinel restored inside provider-owned state. + pub(crate) fn set_reasoning_replay_policy( + &mut self, + policy: crate::ai_serving::OpenAiResponsesReasoningReplayPolicy, + ) { + self.apply_mask_options( + MaskChatRequestOptions::runtime().with_reasoning_replay_policy(policy), + ); + } + + fn preserves_deepseek_opaque_reasoning_state(&self) -> bool { + self.preserve_deepseek_opaque_reasoning_state + } + fn set_collision_corpus(&mut self, collision_corpus: Vec) { self.collision_corpus = collision_corpus; } + fn text_has_redaction_candidate(&self, input: &str) -> bool { + !select_non_overlapping(detect_candidates_for_session_config( + input, + &self.config, + None, + )) + .is_empty() + } + pub(crate) fn redact_text(&mut self, input: &str) -> RedactedText { self.redact_text_internal(input, None) .expect("unlimited redaction scan should not fail") @@ -630,6 +679,10 @@ impl fmt::Debug for RedactionSession { .field("bucket", &self.config.bucket()) .field("mapping_count", &self.mappings.len()) .field("type_counts", &counts) + .field( + "preserve_deepseek_opaque_reasoning_state", + &self.preserve_deepseek_opaque_reasoning_state, + ) .finish() } } @@ -972,12 +1025,29 @@ impl ChatPiiRedactionRuntimeConfigCache { } } -#[derive(Clone, Copy, Default)] -pub(crate) struct MaskChatRequestOptions; +#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)] +pub(crate) struct MaskChatRequestOptions { + /// This bit is derived from the selected provider configuration, never + /// from request JSON. It permits byte-identical handling only for the + /// provider-owned, id-less continuation state used by DeepSeek's + /// Responses contract. + preserve_deepseek_opaque_reasoning_state: bool, +} impl MaskChatRequestOptions { pub(crate) fn runtime() -> Self { - Self + Self::default() + } + + pub(crate) fn with_reasoning_replay_policy( + mut self, + policy: crate::ai_serving::OpenAiResponsesReasoningReplayPolicy, + ) -> Self { + self.preserve_deepseek_opaque_reasoning_state = matches!( + policy, + crate::ai_serving::OpenAiResponsesReasoningReplayPolicy::DeepSeekOpaque + ); + self } } @@ -1387,7 +1457,7 @@ pub(crate) fn try_mask_chat_request_json_with_options( body: &[u8], config: RedactionSessionConfig, options: MaskChatRequestOptions, -) -> Result { +) -> Result { try_mask_chat_pii_request_json_with_options( body, ChatPiiRedactionRequestFormat::OpenAiChat, @@ -1417,8 +1487,9 @@ pub(crate) fn try_mask_chat_pii_request_json_with_options( format: ChatPiiRedactionRequestFormat, config: RedactionSessionConfig, options: MaskChatRequestOptions, -) -> Result { +) -> Result { let mut session = RedactionSession::new(config); + session.apply_mask_options(options); let Ok(mut value) = serde_json::from_slice::(body) else { return Ok(MaskedChatRequest { body: body.to_vec(), @@ -1428,8 +1499,10 @@ pub(crate) fn try_mask_chat_pii_request_json_with_options( }; session.set_collision_corpus(request_collision_corpus(format, &value)); + reject_sensitive_opaque_reasoning_state(format, &value, &session)?; let mut scan_state = RedactionScanState; - let redacted = mask_request_value(format, &mut value, &mut session, &mut scan_state, options)?; + let redacted = mask_request_value(format, &mut value, &mut session, &mut scan_state, options) + .unwrap_or_else(|never| match never {}); if !redacted { return Ok(MaskedChatRequest { @@ -1455,6 +1528,7 @@ pub(crate) async fn try_mask_chat_pii_request_json_with_cache_options( cache: Option<&RedisRedactionMappingCache<'_>>, ) -> Result { let mut session = RedactionSession::new(config); + session.apply_mask_options(options); let Ok(mut value) = serde_json::from_slice::(body) else { return Ok(MaskedChatRequest { body: body.to_vec(), @@ -1464,6 +1538,7 @@ pub(crate) async fn try_mask_chat_pii_request_json_with_cache_options( }; session.set_collision_corpus(request_collision_corpus(format, &value)); + reject_sensitive_opaque_reasoning_state(format, &value, &session)?; let mut scan_state = RedactionScanState; let redacted = mask_request_value_async( format, @@ -1499,9 +1574,11 @@ pub(crate) async fn try_mask_chat_pii_request_value_with_cache_options( cache: Option<&RedisRedactionMappingCache<'_>>, ) -> Result { let mut session = RedactionSession::new(config); + session.apply_mask_options(options); let mut value = body_json.clone(); session.set_collision_corpus(request_collision_corpus(format, &value)); + reject_sensitive_opaque_reasoning_state(format, &value, &session)?; let mut scan_state = RedactionScanState; let redacted = mask_request_value_async( format, @@ -1533,6 +1610,62 @@ fn request_collision_corpus(format: ChatPiiRedactionRequestFormat, value: &Value } } +fn reject_sensitive_opaque_reasoning_state( + format: ChatPiiRedactionRequestFormat, + value: &Value, + session: &RedactionSession, +) -> Result<(), RedactionMaskError> { + if !session.preserves_deepseek_opaque_reasoning_state() { + return Ok(()); + } + if !matches!( + format, + ChatPiiRedactionRequestFormat::OpenAiResponses + | ChatPiiRedactionRequestFormat::OpenAiSearch + ) { + return Ok(()); + } + let Some(items) = value.get("input").and_then(Value::as_array) else { + return Ok(()); + }; + if items.iter().any(|item| { + item.as_object().is_some_and(|object| { + openai_responses_object_is_opaque_reasoning_state(object) + && opaque_reasoning_state_has_sensitive_text(item, session) + }) + }) { + return Err(RedactionMaskError::SensitiveOpaqueReasoningState); + } + Ok(()) +} + +fn opaque_reasoning_state_has_sensitive_text(value: &Value, session: &RedactionSession) -> bool { + let Some(object) = value.as_object() else { + return false; + }; + object.iter().any(|(key, value)| { + // The provider token is ciphertext/high-entropy state. Scanning it as + // user text would reject essentially every valid continuation (the + // generic API-key detector intentionally matches high-entropy data). + // Only this exact top-level protocol field is excluded; every other + // known or future string in the opaque item is scanned fail-closed. + key != "encrypted_content" && json_value_has_redaction_candidate(value, session) + }) +} + +fn json_value_has_redaction_candidate(value: &Value, session: &RedactionSession) -> bool { + match value { + Value::String(text) => session.text_has_redaction_candidate(text), + Value::Array(values) => values + .iter() + .any(|value| json_value_has_redaction_candidate(value, session)), + Value::Object(values) => values + .values() + .any(|value| json_value_has_redaction_candidate(value, session)), + _ => false, + } +} + fn mask_request_value( format: ChatPiiRedactionRequestFormat, value: &mut Value, @@ -1824,7 +1957,7 @@ fn collect_openai_responses_content_part_collision_text(part: &Value, corpus: &m fn response_textish_type(raw_type: Option<&str>) -> bool { matches!( raw_type, - Some("text" | "input_text" | "output_text" | "summary_text") + Some("text" | "input_text" | "output_text" | "summary_text" | "reasoning_text") ) } @@ -2032,6 +2165,16 @@ fn mask_openai_responses_input_item( let Some(item) = item.as_object_mut() else { return Ok(false); }; + // Some Responses-compatible providers bind the clear-text reasoning + // payload to sibling opaque continuation state. The top-level validation + // rejects such an item when its text contains sensitive data; a safe item + // can then remain byte-identical for provider replay. Unbound reasoning + // text still follows the ordinary PII policy below. + if session.preserves_deepseek_opaque_reasoning_state() + && openai_responses_object_is_opaque_reasoning_state(item) + { + return Ok(false); + } let mut redacted = false; if let Some(content) = item.get_mut("content") { redacted |= mask_openai_responses_content_value(content, session, scan_state)?; @@ -2348,6 +2491,11 @@ async fn mask_openai_responses_input_item_async( let Some(item) = item.as_object_mut() else { return Ok(false); }; + if session.preserves_deepseek_opaque_reasoning_state() + && openai_responses_object_is_opaque_reasoning_state(item) + { + return Ok(false); + } let mut redacted = false; if let Some(content) = item.get_mut("content") { redacted |= @@ -2530,6 +2678,12 @@ fn restore_json_response_body( /// 两边因此保持同一套还原语义：未映射的占位符原样保留，`type` / `model` / `id` /// 这类协议字段虽然也被遍历，但它们不可能包含本 session 派生出的 sentinel， /// 所以不会被改写。 +/// +/// Provider-owned Responses reasoning state is the exception. DeepSeek-style +/// items bind `reasoning_text` to `encrypted_content` and require both to be +/// replayed unchanged. Restoring a sentinel in only the text half would make +/// the client return a different continuation item, so those objects/events +/// stay opaque even when another response field is restored. pub(crate) fn restore_json_strings(value: &mut Value, session: &RedactionSession) -> bool { match value { Value::String(text) => { @@ -2548,6 +2702,11 @@ pub(crate) fn restore_json_strings(value: &mut Value, session: &RedactionSession restored } Value::Object(values) => { + if session.preserves_deepseek_opaque_reasoning_state() + && openai_responses_object_is_opaque_reasoning_state(values) + { + return false; + } let mut restored = false; for value in values.values_mut() { restored = restore_json_strings(value, session) || restored; @@ -2558,6 +2717,38 @@ pub(crate) fn restore_json_strings(value: &mut Value, session: &RedactionSession } } +fn openai_responses_object_is_opaque_reasoning_state( + object: &serde_json::Map, +) -> bool { + let Some(item_type) = object.get("type").and_then(Value::as_str) else { + return false; + }; + // Never treat a `reasoning_text` content part or delta/done event name by + // itself as proof of opaque state: ordinary OpenAI reasoning text may + // contain mask placeholders that still need client-side restoration. The + // binding evidence is the parent reasoning item carrying provider-owned + // encrypted continuation state. Returning here keeps that whole object, + // including its nested reasoning_text parts, value-for-value unchanged. + let id_is_absent_or_empty = object + .get("id") + .is_none_or(|id| id.is_null() || id.as_str().is_some_and(|value| value.trim().is_empty())); + item_type == "reasoning" + && id_is_absent_or_empty + && object + .get("encrypted_content") + .and_then(Value::as_str) + .is_some_and(|state| !state.trim().is_empty()) + && object + .get("content") + .and_then(Value::as_array) + .is_some_and(|content| { + content.iter().any(|part| { + part.get("type").and_then(Value::as_str) == Some("reasoning_text") + && part.get("text").is_some_and(Value::is_string) + }) + }) +} + fn restore_text_response_body(body: &[u8], session: &RedactionSession) -> RestoredSyncResponseBody { let Ok(text) = std::str::from_utf8(body) else { return RestoredSyncResponseBody { @@ -4091,14 +4282,14 @@ fn redacted_sentinel_debug(sentinel: &str) -> String { mod tests { use super::{ build_redaction_session_config, detect_candidates_with_probe, mask_chat_request_json, - mask_chat_request_json_with_options, parse_chat_pii_redaction_rules, + mask_chat_request_json_with_options, parse_chat_pii_redaction_rules, restore_json_strings, restore_sync_response_body, try_mask_chat_pii_request_json_with_options, try_mask_chat_pii_request_value_with_cache_options, try_mask_chat_request_json_with_cache_options, try_mask_chat_request_json_with_options, ChatPiiRedactionRequestFormat, ChatPiiRedactionRuntimeConfig, DetectorProbe, MappingKey, - MaskChatRequestOptions, RedactionKind, RedactionMapping, RedactionSession, - RedactionSessionConfig, RedactionSessionSlot, RedisRedactionMappingCache, SentinelMatcher, - StreamingResponseRestorer, + MaskChatRequestOptions, RedactionKind, RedactionMapping, RedactionMaskError, + RedactionSession, RedactionSessionConfig, RedactionSessionSlot, RedisRedactionMappingCache, + SentinelMatcher, StreamingResponseRestorer, }; use std::collections::BTreeMap; use std::time::Duration; @@ -4772,6 +4963,349 @@ mod tests { .contains("secretValueABCDEF1234567890abcdef")); } + #[test] + fn pii_redaction_preserves_opaque_reasoning_state_while_masking_user_input() { + let request = json!({ + "model": "deepseek-reasoner", + "input": [ + { + "type": "reasoning", + "encrypted_content": "opaque-state-must-remain-byte-identical", + "future_state": {"version": 2}, + "content": [ + { + "type": "reasoning_text", + "text": "Provider reasoning state must remain byte-identical" + }, + { + "type": "summary_text", + "text": "Bound future text must also remain byte-identical" + } + ] + }, + { + "type": "message", + "role": "user", + "content": [{ + "type": "input_text", + "text": "Please contact bob@example.net" + }] + } + ] + }); + let raw = serde_json::to_vec(&request).expect("request should serialize"); + + let masked = try_mask_chat_pii_request_json_with_options( + &raw, + ChatPiiRedactionRequestFormat::OpenAiResponses, + test_config(), + deepseek_opaque_replay_options(), + ) + .expect("reasoning request should mask"); + + assert!(masked.redacted); + let masked_json: serde_json::Value = + serde_json::from_slice(&masked.body).expect("masked request should stay valid JSON"); + assert_eq!( + masked_json["input"][0]["encrypted_content"], + "opaque-state-must-remain-byte-identical" + ); + assert_eq!( + masked_json["input"][0]["future_state"], + json!({"version": 2}) + ); + assert_eq!( + masked_json["input"][0]["content"][0]["text"], + "Provider reasoning state must remain byte-identical" + ); + assert_eq!( + masked_json["input"][0]["content"][1]["text"], + "Bound future text must also remain byte-identical", + "the complete provider-bound item must remain unchanged" + ); + let user_text = masked_json["input"][1]["content"][0]["text"] + .as_str() + .expect("user text should remain a string"); + assert!(!user_text.contains("bob@example.net")); + assert!(user_text.contains(" RedactionSessionConfig { RedactionSessionConfig::new(b"redaction-test-key".to_vec(), 300, 600) } + + fn deepseek_opaque_replay_options() -> MaskChatRequestOptions { + MaskChatRequestOptions::runtime().with_reasoning_replay_policy( + crate::ai_serving::OpenAiResponsesReasoningReplayPolicy::DeepSeekOpaque, + ) + } } diff --git a/apps/aether-gateway/src/tests/frontdoor/ai.rs b/apps/aether-gateway/src/tests/frontdoor/ai.rs index 72d88382b..4cf0a9f27 100644 --- a/apps/aether-gateway/src/tests/frontdoor/ai.rs +++ b/apps/aether-gateway/src/tests/frontdoor/ai.rs @@ -723,6 +723,14 @@ async fn run_versioned_codex_model_cards_frontdoor_scenario() { &["future-alias", "hidden-alias"], ), ), + ( + Some(hash_api_key("sk-codex-hidden-only")), + codex_models_snapshot( + "key-codex-hidden-only", + "user-codex-hidden-only", + &["hidden-alias"], + ), + ), ( Some(hash_api_key("sk-codex-second-mixed")), codex_models_snapshot( @@ -1015,10 +1023,13 @@ async fn run_versioned_codex_model_cards_frontdoor_scenario() { .await .expect("incomplete authorized Codex catalog request should succeed"); assert_eq!(incomplete_authorized_response.status(), StatusCode::OK); - assert!(incomplete_authorized_response - .headers() - .get(http::header::ETAG) - .is_none()); + assert_eq!( + incomplete_authorized_response + .headers() + .get(http::header::ETAG) + .and_then(|value| value.to_str().ok()), + Some("\"catalog-etag-v1\"") + ); let incomplete_authorized_payload: serde_json::Value = incomplete_authorized_response .json() .await @@ -1026,9 +1037,34 @@ async fn run_versioned_codex_model_cards_frontdoor_scenario() { assert_eq!( incomplete_authorized_payload["models"] .as_array() - .map(Vec::len), + .map(|models| models + .iter() + .map(|model| model["slug"].as_str().unwrap_or_default()) + .collect::>()), + Some(vec!["future-alias"]), + "one hidden or not-yet-described mapping must not erase valid dynamic cards" + ); + assert_eq!(catalog_hits.load(Ordering::SeqCst), 1); + + let hidden_only_response = client + .get(format!("{gateway_url}/v1/models?client_version=0.145.2")) + .header("authorization", "Bearer sk-codex-hidden-only") + .send() + .await + .expect("hidden-only authorized Codex catalog request should succeed"); + assert_eq!(hidden_only_response.status(), StatusCode::OK); + assert!(hidden_only_response + .headers() + .get(http::header::ETAG) + .is_none()); + let hidden_only_payload: serde_json::Value = hidden_only_response + .json() + .await + .expect("hidden-only authorized Codex body should parse"); + assert_eq!( + hidden_only_payload["models"].as_array().map(Vec::len), Some(0), - "a partial non-empty remote catalog would hide the client's bundled fallback models" + "a model absent from the authoritative upstream catalog must not receive a fabricated card" ); assert_eq!(catalog_hits.load(Ordering::SeqCst), 1); @@ -1044,8 +1080,13 @@ async fn run_versioned_codex_model_cards_frontdoor_scenario() { .await .expect("not-yet-published authorized model body should parse"); assert_eq!( - pending_second_payload["models"].as_array().map(Vec::len), - Some(0) + pending_second_payload["models"] + .as_array() + .map(|models| models + .iter() + .map(|model| model["slug"].as_str().unwrap_or_default()) + .collect::>()), + Some(vec!["future-alias"]) ); assert_eq!(catalog_hits.load(Ordering::SeqCst), 1); diff --git a/crates/aether-ai/formats/src/api.rs b/crates/aether-ai/formats/src/api.rs index 283e001f6..52111d069 100644 --- a/crates/aether-ai/formats/src/api.rs +++ b/crates/aether-ai/formats/src/api.rs @@ -73,6 +73,7 @@ pub use crate::formats::openai::{ finalize_openai_provider_request, finalize_openai_provider_request_with_codex_model_capabilities, finalize_openai_provider_request_with_codex_model_capabilities_and_reasoning_replay_policy, + finalize_openai_provider_request_with_codex_model_capabilities_and_reasoning_replay_policy_for_websocket_continuation, validate_openai_provider_request_contract, OpenAiProviderRequestContractViolation, OpenAiProviderRequestFinalization, }, @@ -188,6 +189,7 @@ pub use crate::formats::{ apply_codex_openai_responses_lite_header_with_capabilities, apply_codex_openai_responses_special_body_edits, apply_codex_openai_responses_special_body_edits_with_source_model_and_capabilities, + apply_codex_openai_responses_websocket_continuation_body_edits_with_source_model_and_capabilities, apply_codex_openai_special_headers, apply_openai_responses_compact_special_body_edits, build_codex_model_catalog_metadata, parse_codex_auth_identity, @@ -212,6 +214,7 @@ pub use crate::formats::{ standard_matrix::{ build_standard_request_body, build_standard_request_body_with_model_directives, build_standard_request_body_with_model_directives_and_request_headers, + build_standard_request_body_with_model_directives_and_request_headers_and_reasoning_replay_policy, }, standard_normalize::{ build_cross_format_openai_chat_request_body, diff --git a/crates/aether-ai/formats/src/formats/openai/request_contract.rs b/crates/aether-ai/formats/src/formats/openai/request_contract.rs index ffad9d38f..ebac54bd6 100644 --- a/crates/aether-ai/formats/src/formats/openai/request_contract.rs +++ b/crates/aether-ai/formats/src/formats/openai/request_contract.rs @@ -53,6 +53,37 @@ pub fn finalize_openai_provider_request_with_codex_model_capabilities_and_reason finalization: OpenAiProviderRequestFinalization<'_>, model_capabilities: Option<&super::responses::codex::CodexResponsesModelCapabilities>, reasoning_replay_policy: super::responses::OpenAiResponsesReasoningReplayPolicy, +) -> Result<(), OpenAiProviderRequestContractViolation> { + finalize_openai_provider_request_with_codex_model_capabilities_and_reasoning_replay_policy_inner( + body, + finalization, + model_capabilities, + reasoning_replay_policy, + false, + ) +} + +pub fn finalize_openai_provider_request_with_codex_model_capabilities_and_reasoning_replay_policy_for_websocket_continuation( + body: &mut Value, + finalization: OpenAiProviderRequestFinalization<'_>, + model_capabilities: Option<&super::responses::codex::CodexResponsesModelCapabilities>, + reasoning_replay_policy: super::responses::OpenAiResponsesReasoningReplayPolicy, +) -> Result<(), OpenAiProviderRequestContractViolation> { + finalize_openai_provider_request_with_codex_model_capabilities_and_reasoning_replay_policy_inner( + body, + finalization, + model_capabilities, + reasoning_replay_policy, + true, + ) +} + +fn finalize_openai_provider_request_with_codex_model_capabilities_and_reasoning_replay_policy_inner( + body: &mut Value, + finalization: OpenAiProviderRequestFinalization<'_>, + model_capabilities: Option<&super::responses::codex::CodexResponsesModelCapabilities>, + reasoning_replay_policy: super::responses::OpenAiResponsesReasoningReplayPolicy, + websocket_continuation: bool, ) -> Result<(), OpenAiProviderRequestContractViolation> { let is_codex_reasoning_endpoint = finalization .provider_type @@ -73,15 +104,27 @@ pub fn finalize_openai_provider_request_with_codex_model_capabilities_and_reason .flatten(); match crate::normalize_api_format_alias(finalization.source_api_format).as_str() { "openai:responses" | "openai:responses:compact" => { - super::responses::codex::apply_codex_openai_responses_special_body_edits_with_source_model_and_capabilities( - body, - finalization.provider_type, - finalization.provider_api_format, - finalization.provider_model, - finalization.source_model, - model_capabilities, - finalization.body_rules, - ); + if websocket_continuation { + super::responses::codex::apply_codex_openai_responses_websocket_continuation_body_edits_with_source_model_and_capabilities( + body, + finalization.provider_type, + finalization.provider_api_format, + finalization.provider_model, + finalization.source_model, + model_capabilities, + finalization.body_rules, + ); + } else { + super::responses::codex::apply_codex_openai_responses_special_body_edits_with_source_model_and_capabilities( + body, + finalization.provider_type, + finalization.provider_api_format, + finalization.provider_model, + finalization.source_model, + model_capabilities, + finalization.body_rules, + ); + } } _ => { super::responses::codex::apply_codex_openai_responses_chat_body_edits_with_source_model_and_capabilities( diff --git a/crates/aether-ai/formats/src/formats/openai/responses/codex.rs b/crates/aether-ai/formats/src/formats/openai/responses/codex.rs index d53ec50a6..357523782 100644 --- a/crates/aether-ai/formats/src/formats/openai/responses/codex.rs +++ b/crates/aether-ai/formats/src/formats/openai/responses/codex.rs @@ -1258,11 +1258,23 @@ fn is_codex_responses_lite_additional_tools_item(value: &Value) -> bool { .is_some_and(|item_type| item_type == "additional_tools") } +fn is_codex_responses_lite_static_additional_tools_item(value: &Value) -> bool { + is_codex_responses_lite_additional_tools_item(value) + && value.get("role").and_then(Value::as_str) == Some("developer") + && value.get("tools").is_some_and(Value::is_array) +} + fn codex_tool_type_accepts_top_level_name(tool_type: &str) -> bool { matches!(tool_type, "function" | "custom" | "namespace") } -fn is_codex_client_executed_tool(tool: &Value) -> bool { +/// Returns whether a tool is represented in the Responses Lite synthetic +/// `additional_tools` item and therefore becomes part of stored history. +/// +/// Keep callers that compare raw client configuration with the synthetic wire +/// prefix on this predicate so hosted/server-executed tools do not create a +/// false configuration change. +pub fn codex_responses_lite_tool_is_client_executed(tool: &Value) -> bool { match tool.get("type").and_then(Value::as_str) { Some("function" | "custom" | "namespace") => true, Some("tool_search") => tool.get("execution").and_then(Value::as_str) == Some("client"), @@ -1277,7 +1289,7 @@ fn retain_codex_client_executed_tools(additional_tools: &mut Value) { else { return; }; - tools.retain(is_codex_client_executed_tool); + tools.retain(codex_responses_lite_tool_is_client_executed); } fn is_codex_responses_lite_instruction_item(value: &Value) -> bool { @@ -1370,7 +1382,7 @@ fn apply_codex_responses_lite_body_contract( .map(|tools| { tools .into_iter() - .filter(is_codex_client_executed_tool) + .filter(codex_responses_lite_tool_is_client_executed) .collect::>() }); let top_level_instructions = body_object @@ -1389,10 +1401,28 @@ fn apply_codex_responses_lite_body_contract( // the Lite wire requirements below, but do not synthesize static config // for a continuation. if websocket_continuation { - for item in input - .iter_mut() - .filter(|item| !is_codex_responses_lite_additional_tools_item(item)) + // A few compatible clients send the already-normalized Lite prefix + // instead of top-level fields. It is still inherited through + // `previous_response_id`, so remove every repeated leading prefix + // item. Limit this to the prefix: a later developer message can be + // genuine incremental input and must not be deleted by shape alone. + let mut repeated_prefix_len = 0; + while input + .get(repeated_prefix_len) + .is_some_and(is_codex_responses_lite_static_additional_tools_item) { + repeated_prefix_len += 1; + if input + .get(repeated_prefix_len) + .is_some_and(is_codex_responses_lite_instruction_item) + { + repeated_prefix_len += 1; + } + } + if repeated_prefix_len > 0 { + input.drain(..repeated_prefix_len); + } + for item in input.iter_mut() { strip_codex_responses_lite_image_details(item); } body_object.insert("parallel_tool_calls".to_string(), json!(false)); @@ -1719,6 +1749,50 @@ pub fn apply_codex_openai_responses_special_body_edits_with_source_model_and_cap source_model: &str, model_capabilities: Option<&CodexResponsesModelCapabilities>, body_rules: Option<&Value>, +) { + apply_codex_openai_responses_special_body_edits_with_source_model_and_capabilities_inner( + provider_request_body, + provider_type, + provider_api_format, + provider_model, + source_model, + model_capabilities, + body_rules, + false, + ); +} + +pub fn apply_codex_openai_responses_websocket_continuation_body_edits_with_source_model_and_capabilities( + provider_request_body: &mut Value, + provider_type: &str, + provider_api_format: &str, + provider_model: &str, + source_model: &str, + model_capabilities: Option<&CodexResponsesModelCapabilities>, + body_rules: Option<&Value>, +) { + apply_codex_openai_responses_special_body_edits_with_source_model_and_capabilities_inner( + provider_request_body, + provider_type, + provider_api_format, + provider_model, + source_model, + model_capabilities, + body_rules, + true, + ); +} + +#[allow(clippy::too_many_arguments)] +fn apply_codex_openai_responses_special_body_edits_with_source_model_and_capabilities_inner( + provider_request_body: &mut Value, + provider_type: &str, + provider_api_format: &str, + provider_model: &str, + source_model: &str, + model_capabilities: Option<&CodexResponsesModelCapabilities>, + body_rules: Option<&Value>, + websocket_continuation: bool, ) { if !is_codex_openai_responses_request(provider_type, provider_api_format) { return; @@ -1728,21 +1802,8 @@ pub fn apply_codex_openai_responses_special_body_edits_with_source_model_and_cap return; }; - // HTTP Responses bodies do not carry an event `type`. Preserve this - // marker only for an actual WebSocket response.create continuation so it - // survives both normalization/finalization passes. The framing layer owns - // the field and will forward it verbatim to the already-bound upstream. - let websocket_continuation = body_object.get("type").and_then(Value::as_str) - == Some("response.create") - && body_object - .get("previous_response_id") - .is_some_and(|value| !value.is_null()); - wrap_codex_responses_string_input_for_backend(body_object, provider_api_format); for field in CODEX_OPENAI_RESPONSES_UNSUPPORTED_BODY_FIELDS { - if *field == "previous_response_id" && websocket_continuation { - continue; - } if !body_rules_handle_path(body_rules, field) { body_object.remove(*field); } @@ -2052,6 +2113,7 @@ mod tests { apply_codex_openai_responses_lite_header_with_capabilities, apply_codex_openai_responses_special_body_edits, apply_codex_openai_responses_special_body_edits_with_source_model_and_capabilities, + apply_codex_openai_responses_websocket_continuation_body_edits_with_source_model_and_capabilities, apply_codex_openai_special_headers, apply_openai_responses_compact_special_body_edits, build_codex_model_catalog_metadata, bundled_codex_model_cards, effective_codex_model_cards, project_codex_catalog_model_card, resolve_codex_responses_model_capabilities, @@ -3257,21 +3319,89 @@ mod tests { "input": incremental_input.clone() }); - apply_codex_openai_responses_special_body_edits( + apply_codex_openai_responses_websocket_continuation_body_edits_with_source_model_and_capabilities( &mut provider_request_body, "codex", "openai:responses", + "gpt-5.6-sol", + "gpt-5.6-sol", None, None, ); - assert_eq!(provider_request_body["previous_response_id"], "resp_123"); + assert!(provider_request_body.get("previous_response_id").is_none()); assert!(provider_request_body.get("instructions").is_none()); assert!(provider_request_body.get("tools").is_none()); assert_eq!(provider_request_body["input"], incremental_input); assert_eq!(provider_request_body["parallel_tool_calls"], false); let once = provider_request_body.clone(); + apply_codex_openai_responses_websocket_continuation_body_edits_with_source_model_and_capabilities( + &mut provider_request_body, + "codex", + "openai:responses", + "gpt-5.6-sol", + "gpt-5.6-sol", + None, + None, + ); + assert_eq!(provider_request_body, once); + } + + #[test] + fn codex_responses_lite_websocket_continuation_removes_only_the_static_prefix() { + let later_developer_message = json!({ + "type": "message", + "role": "developer", + "content": [{"type": "input_text", "text": "genuine incremental input"}] + }); + let mut provider_request_body = json!({ + "model": "gpt-5.6-sol", + "input": [ + { + "type": "additional_tools", + "role": "developer", + "tools": [{"type": "function", "name": "lookup", "parameters": {}}] + }, + { + "type": "message", + "role": "developer", + "content": [{"type": "input_text", "text": "static instructions"}] + }, + {"type": "function_call_output", "call_id": "call_1", "output": "ok"}, + later_developer_message.clone() + ] + }); + + apply_codex_openai_responses_websocket_continuation_body_edits_with_source_model_and_capabilities( + &mut provider_request_body, + "codex", + "openai:responses", + "gpt-5.6-sol", + "gpt-5.6-sol", + None, + None, + ); + + let input = provider_request_body["input"] + .as_array() + .expect("continuation input"); + assert_eq!(input.len(), 2); + assert_eq!(input[0]["type"], "function_call_output"); + assert_eq!(input[1], later_developer_message); + } + + #[test] + fn codex_responses_lite_http_body_cannot_opt_into_websocket_continuation_edits() { + let mut provider_request_body = json!({ + "type": "response.create", + "model": "gpt-5.6-sol", + "previous_response_id": "resp_123", + "instructions": "Static developer instructions.", + "tools": [{"type": "function", "name": "lookup", "parameters": {}}], + "input": [{"type": "message", "role": "user", "content": []}] + }); + apply_codex_openai_responses_special_body_edits( &mut provider_request_body, "codex", @@ -3279,7 +3409,13 @@ mod tests { None, None, ); - assert_eq!(provider_request_body, once); + + assert!(provider_request_body.get("previous_response_id").is_none()); + assert_eq!( + provider_request_body["input"][0]["type"], + "additional_tools" + ); + assert_eq!(provider_request_body["input"][1]["role"], "developer"); } #[test] diff --git a/crates/aether-ai/formats/src/formats/openai/responses/mod.rs b/crates/aether-ai/formats/src/formats/openai/responses/mod.rs index 68f6e8a46..5be47f610 100644 --- a/crates/aether-ai/formats/src/formats/openai/responses/mod.rs +++ b/crates/aether-ai/formats/src/formats/openai/responses/mod.rs @@ -63,6 +63,20 @@ pub fn strip_incompatible_openai_responses_reasoning_items_with_policy( if !aether_ai_formats::is_openai_responses_family_format(provider_api_format) { return 0; } + // DeepSeek's id-less opaque state is valid only on the normal Responses + // continuation contract. Both the legacy Compact endpoint and the current + // `compaction_trigger` operation must retain the strict OpenAI item-id + // replay rules even when the same provider key serves ordinary Responses. + let normal_responses = + aether_ai_formats::normalize_api_format_alias(provider_api_format) == "openai:responses"; + let compact_operation = openai_responses_request_operation(provider_api_format, body).is_some(); + let policy = if policy == OpenAiResponsesReasoningReplayPolicy::DeepSeekOpaque + && (!normal_responses || compact_operation) + { + OpenAiResponsesReasoningReplayPolicy::OpenAiItemIds + } else { + policy + }; let Some(items) = body.get_mut("input").and_then(Value::as_array_mut) else { return 0; }; @@ -121,10 +135,7 @@ fn deepseek_opaque_reasoning_item_is_replayable(object: &serde_json::Map, request_headers: Option<&http::HeaderMap>, enable_model_directives: bool, +) -> Option { + build_standard_request_body_with_model_directives_and_request_headers_and_reasoning_replay_policy( + body_json, + client_api_format, + mapped_model, + provider_type, + provider_api_format, + request_path, + upstream_is_stream, + body_rules, + user_api_key_id, + request_headers, + enable_model_directives, + crate::formats::openai::responses::OpenAiResponsesReasoningReplayPolicy::OpenAiItemIds, + ) +} + +#[allow(clippy::too_many_arguments)] +pub fn build_standard_request_body_with_model_directives_and_request_headers_and_reasoning_replay_policy( + body_json: &Value, + client_api_format: &str, + mapped_model: &str, + provider_type: &str, + provider_api_format: &str, + request_path: &str, + upstream_is_stream: bool, + body_rules: Option<&Value>, + user_api_key_id: Option<&str>, + request_headers: Option<&http::HeaderMap>, + enable_model_directives: bool, reasoning_replay_policy: crate::formats::openai::responses::OpenAiResponsesReasoningReplayPolicy, ) -> Option { let mut format_context = FormatContext::default() @@ -104,26 +133,29 @@ pub fn build_standard_request_body_with_model_directives_and_request_headers( client_api_format, provider_api_format, ); - // A same-family Responses hop is a wire-preserving route. Parsing through the - // canonical request model here would intentionally discard provider-owned - // input item fields (for example DeepSeek's id-less `reasoning_text` and - // future opaque capability fields), even though no format conversion is - // required. Keep the original object and only rewrite the routing model; - // the normal provider-contract and compatibility passes below still apply. - let mut provider_request_body = - if is_same_openai_responses_family(source_api_format.as_ref(), provider_api_format) { - let mut object = body_json.as_object()?.clone(); - object.insert("model".to_string(), Value::String(mapped_model.to_string())); - Value::Object(object) - } else { - convert_request( - source_api_format.as_ref(), - provider_api_format, - body_json, - &format_context, - ) - .ok()? - }; + // DeepSeek's Responses continuation state is opaque. Parsing a same-wire-format + // request through the canonical model would discard its id-less `reasoning_text` + // items and future provider-owned fields even though no conversion is required. + // Keep that provider-specific route wire-preserving, while retaining canonical + // normalization for ordinary OpenAI Responses and for Responses/Compact + // cross-format conversions. + let mut provider_request_body = if is_wire_preserving_deepseek_responses_hop( + source_api_format.as_ref(), + provider_api_format, + reasoning_replay_policy, + ) { + let mut object = body_json.as_object()?.clone(); + object.insert("model".to_string(), Value::String(mapped_model.to_string())); + Value::Object(object) + } else { + convert_request( + source_api_format.as_ref(), + provider_api_format, + body_json, + &format_context, + ) + .ok()? + }; if enable_model_directives { apply_model_directive_overrides_from_request( @@ -192,14 +224,19 @@ pub fn build_standard_request_body_with_model_directives_and_request_headers( Some(provider_request_body) } -fn is_same_openai_responses_family(source_api_format: &str, provider_api_format: &str) -> bool { - matches!( - aether_ai_formats::normalize_api_format_alias(source_api_format).as_str(), - "openai:responses" | "openai:responses:compact" - ) && matches!( - aether_ai_formats::normalize_api_format_alias(provider_api_format).as_str(), - "openai:responses" | "openai:responses:compact" - ) +fn is_wire_preserving_deepseek_responses_hop( + source_api_format: &str, + provider_api_format: &str, + reasoning_replay_policy: crate::formats::openai::responses::OpenAiResponsesReasoningReplayPolicy, +) -> bool { + if reasoning_replay_policy + != crate::formats::openai::responses::OpenAiResponsesReasoningReplayPolicy::DeepSeekOpaque + { + return false; + } + let source_api_format = aether_ai_formats::normalize_api_format_alias(source_api_format); + let provider_api_format = aether_ai_formats::normalize_api_format_alias(provider_api_format); + source_api_format == "openai:responses" && source_api_format == provider_api_format } fn compatible_source_format_for_standard_request<'a>( @@ -378,7 +415,7 @@ mod tests { use super::{ build_standard_request_body, build_standard_request_body_from_canonical, build_standard_request_body_with_model_directives, - build_standard_request_body_with_model_directives_and_request_headers, + build_standard_request_body_with_model_directives_and_request_headers_and_reasoning_replay_policy, normalize_standard_request_to_openai_chat_request, }; use crate::formats::openai::responses::OpenAiResponsesReasoningReplayPolicy; @@ -549,7 +586,7 @@ mod tests { "future_request_field": {"preserve": true} }); - let preserved = build_standard_request_body_with_model_directives_and_request_headers( + let preserved = build_standard_request_body_with_model_directives_and_request_headers_and_reasoning_replay_policy( &request, "openai:responses", "deepseek-v4-flash", @@ -580,7 +617,7 @@ mod tests { "dynamic" ); - let strict = build_standard_request_body_with_model_directives_and_request_headers( + let strict = build_standard_request_body_with_model_directives_and_request_headers_and_reasoning_replay_policy( &request, "openai:responses", "deepseek-v4-flash", @@ -604,6 +641,32 @@ mod tests { .count(), 0 ); + + let compact = build_standard_request_body_with_model_directives_and_request_headers_and_reasoning_replay_policy( + &request, + "openai:responses:compact", + "deepseek-v4-flash", + "custom", + "openai:responses:compact", + "/v1/responses/compact", + false, + None, + None, + None, + false, + OpenAiResponsesReasoningReplayPolicy::DeepSeekOpaque, + ) + .expect("Compact body should continue through canonical normalization"); + assert_eq!( + compact["input"] + .as_array() + .expect("compact input array") + .iter() + .filter(|item| item["type"] == "reasoning") + .count(), + 0, + "compact must not use the DeepSeek opaque replay policy" + ); } #[test] diff --git a/crates/aether-ai/formats/src/formats/shared/sync_products.rs b/crates/aether-ai/formats/src/formats/shared/sync_products.rs index c13bbd3f9..57d512eb6 100644 --- a/crates/aether-ai/formats/src/formats/shared/sync_products.rs +++ b/crates/aether-ai/formats/src/formats/shared/sync_products.rs @@ -3184,6 +3184,13 @@ fn materialize_openai_responses_reasoning_item( state: OpenAIResponsesSyncReasoningState, ) -> Value { let mut item = state.item; + let has_provider_opaque_state = item + .get("encrypted_content") + .and_then(Value::as_str) + .is_some_and(|value| !value.trim().is_empty()); + if has_provider_opaque_state { + return Value::Object(item); + } item.entry("type".to_string()) .or_insert_with(|| Value::String("reasoning".to_string())); item.entry("id".to_string()).or_insert_with(|| { @@ -3939,7 +3946,7 @@ mod tests { aggregate_claude_stream_sync_response, aggregate_gemini_stream_sync_response, aggregate_openai_chat_stream_sync_response, aggregate_openai_responses_stream_sync_response, convert_standard_chat_response, - convert_standard_cli_response, + convert_standard_cli_response, materialize_openai_responses_reasoning_item, maybe_build_openai_chat_cross_format_sync_product_from_normalized_payload, maybe_build_openai_responses_cross_format_sync_product_from_normalized_payload, maybe_build_openai_responses_same_family_sync_body_from_normalized_payload, @@ -3947,7 +3954,9 @@ mod tests { maybe_build_standard_cross_format_sync_product_from_normalized_payload, maybe_build_standard_same_format_sync_body_from_normalized_payload, maybe_build_standard_sync_finalize_product_from_normalized_payload, - try_aggregate_openai_responses_stream_sync_response, StandardSyncFinalizeNormalizedProduct, + openai_responses_synthetic_reasoning_item_id, + try_aggregate_openai_responses_stream_sync_response, OpenAIResponsesSyncReasoningState, + StandardSyncFinalizeNormalizedProduct, }; use aether_ai_formats::formats::conversion::response::{ convert_claude_chat_response_to_openai_chat, convert_gemini_chat_response_to_openai_chat, @@ -5336,6 +5345,85 @@ mod tests { assert!(result["completed_at"].as_i64().is_some()); } + #[test] + fn preserves_idless_provider_opaque_reasoning_item_during_materialization() { + let original = json!({ + "type": "reasoning", + "encrypted_content": "opaque-provider-state", + "content": [{ + "type": "reasoning_text", + "text": "private chain of thought" + }], + "summary": [{ + "type": "provider_summary", + "text": "provider-owned summary" + }], + "future_provider_field": {"version": 2} + }); + let state = OpenAIResponsesSyncReasoningState { + item: original + .as_object() + .expect("reasoning item should be an object") + .clone(), + summary_text: "must not replace provider-owned state".to_string(), + }; + + let materialized = materialize_openai_responses_reasoning_item("resp_opaque_123", state); + + assert_eq!(materialized, original); + assert!(materialized.get("id").is_none()); + assert!(materialized.get("status").is_none()); + } + + #[test] + fn aggregates_authoritative_provider_opaque_reasoning_item_without_mutation() { + let body = concat!( + "event: response.output_item.added\n", + "data: {\"type\":\"response.output_item.added\",\"output_index\":0,\"item\":{\"type\":\"reasoning\",\"encrypted_content\":\"opaque-provider-state\",\"content\":[]}}\n\n", + "event: response.reasoning_text.done\n", + "data: {\"type\":\"response.reasoning_text.done\",\"output_index\":0,\"text\":\"provider reasoning\"}\n\n", + "event: response.output_item.done\n", + "data: {\"type\":\"response.output_item.done\",\"output_index\":0,\"item\":{\"type\":\"reasoning\",\"encrypted_content\":\"opaque-provider-state\",\"content\":[{\"type\":\"reasoning_text\",\"text\":\"provider reasoning\"}],\"future_provider_field\":{\"version\":2}}}\n\n", + "event: response.completed\n", + "data: {\"type\":\"response.completed\",\"response\":{\"id\":\"resp_opaque_stream_123\",\"object\":\"response\",\"model\":\"deepseek-reasoner\",\"status\":\"completed\",\"output\":[]}}\n\n", + ); + + let result = aggregate_openai_responses_stream_sync_response(body.as_bytes()) + .expect("provider opaque reasoning stream should aggregate"); + + assert_eq!( + result["output"][0], + json!({ + "type": "reasoning", + "encrypted_content": "opaque-provider-state", + "content": [{ + "type": "reasoning_text", + "text": "provider reasoning" + }], + "future_provider_field": {"version": 2} + }) + ); + } + + #[test] + fn synthesizes_wire_compatible_id_for_local_reasoning_summary() { + let state = OpenAIResponsesSyncReasoningState { + item: json!({"type": "reasoning"}) + .as_object() + .expect("reasoning item should be an object") + .clone(), + summary_text: "Need care".to_string(), + }; + + let materialized = materialize_openai_responses_reasoning_item("resp_summary_123", state); + + assert_eq!( + materialized["id"], + openai_responses_synthetic_reasoning_item_id("resp_summary_123", 0) + ); + assert_eq!(materialized["summary"][0]["text"], "Need care"); + } + #[test] fn accepts_openai_responses_same_family_stream_when_needs_conversion_is_true() { let body = concat!( diff --git a/crates/aether-ai/formats/src/lib.rs b/crates/aether-ai/formats/src/lib.rs index dbcb445e1..fbb8a3f1e 100644 --- a/crates/aether-ai/formats/src/lib.rs +++ b/crates/aether-ai/formats/src/lib.rs @@ -40,11 +40,13 @@ pub use formats::openai::request_contract::{ finalize_openai_provider_request, finalize_openai_provider_request_with_codex_model_capabilities, finalize_openai_provider_request_with_codex_model_capabilities_and_reasoning_replay_policy, + finalize_openai_provider_request_with_codex_model_capabilities_and_reasoning_replay_policy_for_websocket_continuation, validate_openai_provider_request_contract, OpenAiProviderRequestContractViolation, OpenAiProviderRequestFinalization, }; pub use formats::openai::responses::codex::{ - build_codex_model_catalog_metadata, bundled_codex_model_cards, effective_codex_model_cards, + build_codex_model_catalog_metadata, bundled_codex_model_cards, + codex_responses_lite_tool_is_client_executed, effective_codex_model_cards, parse_codex_auth_identity, project_codex_catalog_model_card, resolve_codex_responses_model_capabilities, CodexAuthIdentity, CodexResponsesModelCapabilities, CODEX_CLIENT_ORIGINATOR, CODEX_CLIENT_USER_AGENT, CODEX_CLIENT_VERSION, diff --git a/crates/aether-ai/formats/src/provider_compat/proxy/rules.rs b/crates/aether-ai/formats/src/provider_compat/proxy/rules.rs index c0f942c95..4830609cd 100644 --- a/crates/aether-ai/formats/src/provider_compat/proxy/rules.rs +++ b/crates/aether-ai/formats/src/provider_compat/proxy/rules.rs @@ -1,5 +1,5 @@ use std::{ - collections::{BTreeMap, HashSet}, + collections::{BTreeMap, BTreeSet, HashSet}, sync::OnceLock, }; @@ -193,6 +193,73 @@ pub fn body_rules_have_enabled_rules(rules: Option<&Value>) -> bool { .any(|rule| rule.as_object().is_some_and(body_rule_is_enabled)) } +/// Returns the normalized request-header names that can affect enabled body +/// rules. +/// +/// Callers that persist a body-normalization identity should bind only these +/// values, rather than every incidental ingress header. Dynamic proxy headers +/// such as request IDs and forwarding hops do not change the provider body +/// unless an enabled, locally supported condition actually reads them. +pub fn body_rules_request_header_dependencies(rules: Option<&Value>) -> BTreeSet { + let Some(rules) = rules.and_then(Value::as_array) else { + return BTreeSet::new(); + }; + let mut dependencies = BTreeSet::new(); + for rule in rules { + let Some(rule) = rule.as_object() else { + continue; + }; + if !body_rule_is_enabled(rule) { + continue; + } + let Some(condition) = rule.get("condition").filter(|value| !value.is_null()) else { + continue; + }; + if condition_is_locally_supported(condition) { + collect_condition_request_header_dependencies(condition, &mut dependencies); + } + } + dependencies +} + +fn collect_condition_request_header_dependencies( + condition: &Value, + dependencies: &mut BTreeSet, +) { + let Some(condition) = condition.as_object() else { + return; + }; + if let Some(children) = condition.get("all").and_then(Value::as_array) { + for child in children { + collect_condition_request_header_dependencies(child, dependencies); + } + return; + } + if let Some(children) = condition.get("any").and_then(Value::as_array) { + for child in children { + collect_condition_request_header_dependencies(child, dependencies); + } + return; + } + + let source = condition + .get("source") + .and_then(Value::as_str) + .map(str::trim) + .unwrap_or("body"); + if !condition_source_is_headers(source) { + return; + } + if let Some(path) = condition + .get("path") + .and_then(Value::as_str) + .map(str::trim) + .filter(|path| !path.is_empty()) + { + dependencies.insert(path.to_ascii_lowercase()); + } +} + pub fn body_rules_handle_path(rules: Option<&Value>, path: &str) -> bool { let Some(target_path) = parse_body_path(path) else { return false; @@ -245,7 +312,7 @@ pub fn apply_local_body_rules( rules: Option<&Value>, original_body: Option<&Value>, ) -> bool { - apply_local_body_rules_inner(body, rules, original_body, None) + apply_local_body_rules_inner(body, rules, original_body, None, None) } pub fn apply_local_body_rules_with_request_headers( @@ -259,14 +326,56 @@ pub fn apply_local_body_rules_with_request_headers( rules, original_body, request_headers.map(ConditionHeaders::Request), + None, ) } +/// Applies body rules and reports whether an enabled, condition-matching rule +/// actually handled `path`. +/// +/// This differs from [`body_rules_handle_path`], which intentionally answers a +/// static configuration question and therefore cannot account for per-request +/// conditions. WebSocket framing uses this variant before deciding whether it +/// may restore a client protocol field after provider normalization. +pub fn apply_local_body_rules_with_request_headers_and_track_path( + body: &mut Value, + rules: Option<&Value>, + original_body: Option<&Value>, + request_headers: Option<&http::HeaderMap>, + path: &str, +) -> Option { + let target = parse_body_path(path)?; + let mut tracker = AppliedBodyRulePath { + target, + handled: false, + }; + apply_local_body_rules_inner( + body, + rules, + original_body, + request_headers.map(ConditionHeaders::Request), + Some(&mut tracker), + ) + .then_some(tracker.handled) +} + +struct AppliedBodyRulePath { + target: Vec, + handled: bool, +} + +impl AppliedBodyRulePath { + fn observe(&mut self, path: &[BodyPathSegment]) { + self.handled |= path_matches(path, &self.target); + } +} + fn apply_local_body_rules_inner( body: &mut Value, rules: Option<&Value>, original_body: Option<&Value>, request_headers: Option>, + mut tracker: Option<&mut AppliedBodyRulePath>, ) -> bool { let Some(rules) = rules else { return true; @@ -331,7 +440,11 @@ fn apply_local_body_rules_inner( } else { value_template.clone() }; - let _ = set_nested_value(body, &target_path, value); + if set_nested_value(body, &target_path, value) { + if let Some(tracker) = tracker.as_deref_mut() { + tracker.observe(&target_path); + } + } } } Some("drop") => { @@ -354,7 +467,11 @@ fn apply_local_body_rules_inner( reverse: true, }, ) { - let _ = delete_nested_value(body, &target_path); + if delete_nested_value(body, &target_path) { + if let Some(tracker) = tracker.as_deref_mut() { + tracker.observe(&target_path); + } + } } } Some("rename") => { @@ -375,7 +492,12 @@ fn apply_local_body_rules_inner( if has_wildcard(&from) || has_wildcard(&to) { continue; } - let _ = rename_nested_value(body, &from, &to); + if rename_nested_value(body, &from, &to) { + if let Some(tracker) = tracker.as_deref_mut() { + tracker.observe(&from); + tracker.observe(&to); + } + } } Some("append") => { let Some(path) = rule @@ -401,6 +523,9 @@ fn apply_local_body_rules_inner( if let Some(target) = get_nested_value_mut(body, &target_path) { if let Some(values) = target.as_array_mut() { values.push(value.clone()); + if let Some(tracker) = tracker.as_deref_mut() { + tracker.observe(&target_path); + } } } } @@ -424,6 +549,9 @@ fn apply_local_body_rules_inner( if let Some(values) = target.as_array_mut() { let insert_index = normalize_insert_index(values.len(), index); values.insert(insert_index, value); + if let Some(tracker) = tracker.as_deref_mut() { + tracker.observe(&path); + } } } } @@ -474,6 +602,9 @@ fn apply_local_body_rules_inner( pattern.replacen(¤t, count, replacement).to_string() }; *target = Value::String(replaced); + if let Some(tracker) = tracker.as_deref_mut() { + tracker.observe(&target_path); + } } } } @@ -1331,9 +1462,10 @@ fn rename_nested_value( mod tests { use super::{ apply_local_body_rules, apply_local_body_rules_with_request_headers, - apply_local_header_rules, apply_local_header_rules_with_request_headers, - body_rules_are_locally_supported, body_rules_handle_path, body_rules_have_enabled_rules, - header_rules_are_locally_supported, header_rules_have_enabled_rules, + apply_local_body_rules_with_request_headers_and_track_path, apply_local_header_rules, + apply_local_header_rules_with_request_headers, body_rules_are_locally_supported, + body_rules_handle_path, body_rules_have_enabled_rules, header_rules_are_locally_supported, + header_rules_have_enabled_rules, }; #[test] @@ -1782,4 +1914,48 @@ mod tests { assert!(!body_rules_handle_path(Some(&rules), "tools[3].kind")); assert!(!body_rules_handle_path(Some(&rules), "instructions")); } + + #[test] + fn body_rule_path_tracking_requires_the_rule_condition_to_apply() { + let rules = serde_json::json!([{ + "action": "set", + "path": "store", + "value": false, + "condition": {"path": "metadata.mode", "op": "eq", "value": "enforce"} + }]); + + let skipped_original = serde_json::json!({ + "store": true, + "metadata": {"mode": "observe"} + }); + let mut skipped = skipped_original.clone(); + assert_eq!( + apply_local_body_rules_with_request_headers_and_track_path( + &mut skipped, + Some(&rules), + Some(&skipped_original), + None, + "store", + ), + Some(false) + ); + assert_eq!(skipped["store"], true); + + let applied_original = serde_json::json!({ + "store": false, + "metadata": {"mode": "enforce"} + }); + let mut applied = applied_original.clone(); + assert_eq!( + apply_local_body_rules_with_request_headers_and_track_path( + &mut applied, + Some(&rules), + Some(&applied_original), + None, + "store", + ), + Some(true), + "setting the existing value still counts as rule ownership" + ); + } } diff --git a/crates/aether-provider/transport/src/same_format_provider/mod.rs b/crates/aether-provider/transport/src/same_format_provider/mod.rs index 99a61fddd..fa6dd1fa9 100644 --- a/crates/aether-provider/transport/src/same_format_provider/mod.rs +++ b/crates/aether-provider/transport/src/same_format_provider/mod.rs @@ -2188,6 +2188,44 @@ mod tests { ); } + #[test] + fn same_format_compact_keeps_strict_reasoning_replay_under_deepseek_policy() { + let request_body = json!({ + "model": "deepseek-v4-flash", + "input": [{ + "type": "reasoning", + "encrypted_content": "opaque-deepseek-state", + "content": [{"type": "reasoning_text", "text": "thinking"}] + }] + }); + let output = build_same_format_provider_request_body_with_compatibility_report_and_reasoning_replay_policy( + SameFormatProviderRequestBodyInput { + body_json: &request_body, + mapped_model: "deepseek-v4-flash", + client_api_format: "openai:responses:compact", + provider_api_format: "openai:responses:compact", + source_model: Some("deepseek-v4-flash"), + family: SameFormatProviderFamily::Standard, + body_rules: None, + request_headers: None, + upstream_is_stream: false, + force_body_stream_field: false, + kiro_auth_config: None, + is_claude_code: false, + enable_model_directives: false, + }, + aether_ai_formats::OpenAiResponsesReasoningReplayPolicy::DeepSeekOpaque, + ) + .expect("Compact body should keep the strict OpenAI replay contract"); + + assert_eq!(output.body["input"].as_array().map(Vec::len), Some(0)); + assert!(output.compatibility_edits.iter().any(|edit| { + edit.field == "input[].id" + && edit.action + == SameFormatProviderCompatibilityEditAction::ProviderCompatibilityRewrite + })); + } + #[test] fn same_format_stream_policy_wins_after_body_rules() { let body_rules = json!([ diff --git a/docs/WebSocket-Mode.md b/docs/WebSocket-Mode.md index d605a871d..844d54cf7 100644 --- a/docs/WebSocket-Mode.md +++ b/docs/WebSocket-Mode.md @@ -4,6 +4,16 @@ 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`. +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 +rejects `response.create.stream_id` until per-lane binding, ordering, timeout, +usage, and error routing are implemented end to end. Use separate WebSocket +connections for parallel runs through Aether. A syntactically valid named +`stream_id` is rejected with `responses_websocket_named_stream_unsupported`; +the error event echoes the validated ID so the client can associate the error +with its attempted lane. Invalid or untrusted IDs are not echoed. + ## Why use WebSocket mode WebSocket mode is most useful when a workflow involves many model-tool round trips (for example, agentic coding or orchestration loops with repeated tool calls). @@ -86,14 +96,68 @@ ws.send( WebSocket mode uses the same `previous_response_id` chaining semantics as HTTP mode, but it adds a lower-latency continuation path on the active socket. -On an active WebSocket connection, the service keeps one previous-response state in a connection-local in-memory cache (the most recent response). Continuing from that most recent response is fast because the service can reuse connection-local state. Because the previous-response state is retained only in memory and is not written to disk, you can use WebSocket mode in a way that is compatible with `store=false` and Zero Data Retention (ZDR). +On an active Aether WebSocket connection, the selected upstream keeps the +previous-response state for the single default lane in its connection-local +cache. Continuing from that most recent response is fast because the service +can reuse connection-local state. Because the previous-response state is +retained only in memory and is not written to disk, you can use WebSocket mode +in a way that is compatible with `store=false` and Zero Data Retention (ZDR). -If a `previous_response_id` is not in the in-memory cache, behavior depends on whether you store responses: +If a `previous_response_id` is not in the upstream connection's in-memory +cache, behavior depends on whether the upstream stored the response: -- With `store=true`, the service may hydrate older response IDs from persisted state when available. Continuation can still work, but it usually loses the in-memory latency benefit. +- With `store=true`, the upstream service may hydrate older response IDs from its persisted state when available. Continuation can still work, but it usually loses the in-memory latency benefit. - With `store=false` (including ZDR), there is no persisted fallback. If the ID is uncached, the request returns `previous_response_not_found`. -If a turn fails (`4xx` or `5xx`), the service evicts the referenced `previous_response_id` from the connection-local cache. This prevents reusing stale cached state for that failed continuation. +For a new downstream WebSocket connection, Aether also has to prove that the +response belongs to the currently authenticated user/API key and to the exact +provider endpoint, key, credential generation, transport, adapter, model, and +normalization contract. Aether records this ownership only when the effective +provider `response.create`, after Aether's body rules and framing, explicitly +has `store=true` and a successful +`response.completed`, `response.done`, or non-error `response.incomplete` +terminal supplies a valid response ID. `store=false`, an omitted/overridden +`store`, failures, cancellations, malformed IDs, and ZDR turns never create +this registry state. This explicit-true rule is intentionally conservative: +Aether does not infer a provider default for an omitted `store` field. + +The RuntimeState key contains a SHA-256 digest over length-delimited live +`user_id`, `api_key_id`, and the opaque response ID; raw response IDs and +credentials are not stored in keys or values. Records expire after 24 hours, +are capped at the 1,024 most recently registered IDs per user/API-key pair, +and are bounded in serialized size. The registry stores ownership/routing +proof and contract digests only; it does not store response contents and is +not a replacement for the upstream's `store=true` persistence. A registry +write failure does not turn a successful provider response into a failure, so +that terminal can reach the client but cannot later resume on a new socket. + +With the Redis RuntimeState backend, ownership is shared across gateway +instances for the record TTL, subject to that Redis deployment's own +availability and persistence configuration. The memory backend is +process-local, is not shared between instances, and loses the registry on +restart. Expiry, per-principal eviction, a RuntimeState outage, or a memory +backend restart causes the first continuation on a new connection to fail +closed with `previous_response_not_found`, even if the upstream might still +retain the response. Aether never falls back to the ordinary scheduler for +such a miss and never sends the opaque response ID to a different provider or +key. + +PII-redaction restore mappings intentionally remain connection-local and are +not persisted. If a stored response chain contains Aether PII sentinels, Aether +rejects cross-connection continuation rather than risk exposing those +sentinels without the original restore mapping. Start a new response with the +complete required context in that case. + +If a continuation on the same lane fails (`4xx` or `5xx`), the service evicts +the referenced `previous_response_id` from the connection-local cache. Aether +only supports the implicit default lane, so this same-lane rule applies to all +continuations it currently accepts. The upstream service preserves a shared +parent when a cross-lane fork fails, but Aether does not yet expose that named +lane behavior. + +The continuation must keep the model selected for the response chain. Aether +rejects a model change with status `409` and code +`responses_continuation_model_change_unsupported`. ## Compaction and creating new responses @@ -105,7 +169,7 @@ When you enable server-side compaction (`context_management` with `compact_thres ### Standalone `/responses/compact` -The standalone [`/responses/compact` endpoint](https://developers.openai.com/api/docs/api-reference/responses/compact) returns a new compacted input window, not a response ID. After compaction, create a new response on your WebSocket connection using the compacted window as `input` (plus the next user/tool items). +The standalone [`/responses/compact` endpoint](https://developers.openai.com/api/reference/resources/responses/methods/compact) returns a new compacted input window, not a response ID. After compaction, create a new response on your WebSocket connection using the compacted window as `input` (plus the next user/tool items). Start a new chain by omitting `previous_response_id` or setting it to `null`. Pass the compacted output as-is; do not prune the returned window. @@ -141,8 +205,9 @@ ws.send( ## Connection behavior and limits - Server events and ordering match the existing Responses streaming event model. -- A single WebSocket connection can receive multiple `response.create` messages, but it runs them sequentially (one in-flight response at a time). -- No multiplexing support today. Use multiple connections if you need parallel runs. +- A single Aether WebSocket connection can receive multiple `response.create` messages over its lifetime, but the client must wait for a terminal event before sending the next one. Aether does not queue overlapping creates and returns `response_already_in_progress` while a turn is active. +- Named `stream_id` multiplexing is not exposed by Aether yet. Use multiple connections if you need parallel runs. +- The upstream OpenAI service allows at most 16 active/in-flight responses on one connection; additional `response.create` events are queued. It also allows at most 32 distinct named `stream_id` values per connection, and the implicit default lane does not count toward that 32-lane limit. These describe upstream multiplexing limits, not capabilities exposed by Aether's current single-lane bridge. - Connection duration is limited to 60 minutes. Reconnect when the limit is reached. - Aether binds each upstream WebSocket to one selected provider key. A provider must explicitly enable the standard Responses WebSocket capability and expose an `openai:responses` endpoint before it is eligible for this bridge. - The Codex adapter additionally watches Codex quota events. A `usage_limit_reached` terminal error immediately marks the bound account unavailable. If the client has not received a standard `response.*` event and the request has no `previous_response_id`, Aether retries that one turn once on another eligible key without closing the public socket. @@ -151,10 +216,11 @@ ws.send( ## Reconnect and recover -When a connection closes (or hits the 60-minute limit), open a new WebSocket connection and continue with one of these patterns: +When a connection closes (or hits the 60-minute limit), open a new WebSocket +connection and continue with one of these patterns: -1. If your prior response is persisted (`store=true`) and you have a valid response ID, continue with `previous_response_id` and new input items. -2. If you cannot continue the chain (for example, `store=false`/ZDR or `previous_response_not_found`), start a new response by setting `previous_response_id` to `null` (or omitting it) and send the full input context for the next turn. +1. If the prior response is persisted (`store=true`) and its response ID remains valid, continue with `previous_response_id` and only the new input items. +2. If the chain cannot be hydrated (for example, `store=false`/ZDR or `previous_response_not_found`), start a new response by setting `previous_response_id` to `null` (or omitting it) and send the complete input context needed for the next turn. 3. If you compacted context with `/responses/compact`, use the returned compacted window as the base `input` for that new response, then append the latest user/tool items. ## Errors to handle @@ -166,6 +232,7 @@ When a connection closes (or hits the 60-minute limit), open a new WebSocket con "type": "error", "status": 400, "error": { + "type": "invalid_request_error", "code": "previous_response_not_found", "message": "Previous response with id 'resp_abc' not found.", "param": "previous_response_id" @@ -191,4 +258,5 @@ When a connection closes (or hits the 60-minute limit), open a new WebSocket con - [Conversation state](https://developers.openai.com/api/docs/guides/conversation-state) - [Streaming API responses](https://developers.openai.com/api/docs/guides/streaming-responses) -- [Responses streaming events reference](https://developers.openai.com/api/docs/api-reference/responses-streaming) +- [Responses streaming events reference](https://developers.openai.com/api/reference/resources/responses) +- [Responses WebSocket events reference](https://developers.openai.com/api/reference/resources/responses/websocket-events)