fix(ws): harden Responses continuation state

This commit is contained in:
ZheFox
2026-08-20 08:51:16 +08:00
parent bef282cfee
commit 654f798d25
45 changed files with 6403 additions and 493 deletions
@@ -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:?}"),
@@ -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(
@@ -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<ProviderRequestRedaction<'a>, 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::<RedactionSessionSlot>() 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)]
@@ -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<String> {
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://[email protected]/v1"
));
assert!(!is_deepseek_provider(
"custom",
"https://evil.example/[email protected]/v1"
));
assert!(!is_deepseek_provider(
"custom",
"https://evil.example/[email protected]"
));
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::<Vec<_>>();
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"),
@@ -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(
@@ -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,
@@ -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::{
@@ -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<Value> {
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<Value> {
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<Value> {
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,
@@ -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]
@@ -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::<Vec<_>>();
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();
@@ -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,
@@ -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<Option<AiExecutionDecision>, 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<Option<AiExecutionDecision>, 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))
}
@@ -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<Option<LocalOpenAiResponsesCandidatePayloadParts>, 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<Option<LocalOpenAiResponsesCandidatePayloadParts>, 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,
@@ -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<Self> {
pub(crate) fn new(provider_id: &str, endpoint_id: &str, key_id: &str) -> Option<Self> {
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> {
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::<Vec<_>>();
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
{
@@ -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,
};