fix: enforce provider upstream stream policy

This commit is contained in:
stabey
2026-05-08 21:00:44 +08:00
parent 8958bf5e08
commit 7e5a08e09d
29 changed files with 1518 additions and 118 deletions

View File

@@ -37,6 +37,52 @@ pub(crate) fn force_upstream_streaming_for_provider(
force_upstream_streaming_for_provider_impl(provider_type, provider_api_format) force_upstream_streaming_for_provider_impl(provider_type, provider_api_format)
} }
pub(crate) fn resolve_upstream_is_stream_for_provider(
endpoint_config: Option<&serde_json::Value>,
provider_type: &str,
provider_api_format: &str,
client_is_stream: bool,
hard_requires_streaming: bool,
) -> bool {
let hard_requires_streaming = hard_requires_streaming
|| force_upstream_streaming_for_provider(provider_type, provider_api_format);
aether_ai_formats::resolve_upstream_is_stream_from_endpoint_config(
endpoint_config,
client_is_stream,
hard_requires_streaming,
)
}
pub(crate) fn endpoint_config_forces_body_stream_field(
endpoint_config: Option<&serde_json::Value>,
) -> bool {
aether_ai_formats::endpoint_config_forces_upstream_stream_policy(endpoint_config)
}
pub(crate) fn request_requires_body_stream_field(
body_json: &serde_json::Value,
force_body_stream_field: bool,
) -> bool {
force_body_stream_field
|| body_json
.as_object()
.is_some_and(|object| object.contains_key("stream"))
}
pub(crate) fn enforce_provider_body_stream_policy(
provider_request_body: &mut serde_json::Value,
provider_api_format: &str,
upstream_is_stream: bool,
require_body_stream_field: bool,
) {
aether_ai_formats::enforce_request_body_stream_field(
provider_request_body,
provider_api_format,
upstream_is_stream,
require_body_stream_field,
);
}
pub(crate) fn extract_standard_requested_model(body_json: &serde_json::Value) -> Option<String> { pub(crate) fn extract_standard_requested_model(body_json: &serde_json::Value) -> Option<String> {
aether_ai_serving::extract_ai_standard_requested_model(body_json) aether_ai_serving::extract_ai_standard_requested_model(body_json)
} }
@@ -56,8 +102,10 @@ pub(crate) fn extract_requested_model_from_request(
#[cfg(test)] #[cfg(test)]
mod tests { mod tests {
use super::{ use super::{
endpoint_config_forces_body_stream_field, enforce_provider_body_stream_policy,
extract_requested_model_from_request, extract_standard_requested_model, extract_requested_model_from_request, extract_standard_requested_model,
force_upstream_streaming_for_provider, RequestedModelFamily, force_upstream_streaming_for_provider, resolve_upstream_is_stream_for_provider,
RequestedModelFamily,
}; };
use axum::http::Request; use axum::http::Request;
use serde_json::json; use serde_json::json;
@@ -90,6 +138,67 @@ mod tests {
)); ));
} }
#[test]
fn resolves_endpoint_upstream_stream_policy_with_provider_hard_constraints() {
assert!(resolve_upstream_is_stream_for_provider(
Some(&json!({"upstream_stream_policy": "force_stream"})),
"openai",
"openai:chat",
false,
false,
));
assert!(!resolve_upstream_is_stream_for_provider(
Some(&json!({"upstream_stream_policy": "force_non_stream"})),
"openai",
"openai:chat",
true,
false,
));
assert!(resolve_upstream_is_stream_for_provider(
Some(&json!({"upstream_stream_policy": "auto"})),
"openai",
"openai:chat",
true,
false,
));
assert!(resolve_upstream_is_stream_for_provider(
Some(&json!({"upstream_stream_policy": "force_non_stream"})),
"codex",
"openai:responses",
true,
false,
));
}
#[test]
fn enforces_provider_body_stream_policy_for_body_and_streamless_formats() {
let mut openai_chat = json!({"stream": true});
enforce_provider_body_stream_policy(&mut openai_chat, "openai:chat", false, false);
assert_eq!(openai_chat.get("stream"), Some(&json!(false)));
let mut ordinary_sync = json!({"messages": []});
enforce_provider_body_stream_policy(&mut ordinary_sync, "openai:chat", false, false);
assert!(ordinary_sync.get("stream").is_none());
let mut compact = json!({"stream": true});
enforce_provider_body_stream_policy(&mut compact, "openai:responses:compact", true, true);
assert!(compact.get("stream").is_none());
}
#[test]
fn detects_endpoint_configs_that_force_body_stream_field() {
assert!(endpoint_config_forces_body_stream_field(Some(
&json!({"upstream_stream_policy": "force_stream"})
)));
assert!(endpoint_config_forces_body_stream_field(Some(
&json!({"upstream_stream_policy": "force_non_stream"})
)));
assert!(!endpoint_config_forces_body_stream_field(Some(
&json!({"upstream_stream_policy": "auto"})
)));
assert!(!endpoint_config_forces_body_stream_field(None));
}
#[test] #[test]
fn extracts_standard_requested_model_from_request_body() { fn extracts_standard_requested_model_from_request_body() {
let requested_model = let requested_model =

View File

@@ -3,6 +3,9 @@ use std::sync::Arc;
use serde_json::Value; use serde_json::Value;
use crate::ai_serving::planner::common::{
enforce_provider_body_stream_policy, request_requires_body_stream_field,
};
use crate::ai_serving::transport::antigravity::{ use crate::ai_serving::transport::antigravity::{
build_antigravity_safe_v1internal_request, build_antigravity_static_identity_headers, build_antigravity_safe_v1internal_request, build_antigravity_static_identity_headers,
classify_local_antigravity_request_support, AntigravityEnvelopeRequestType, classify_local_antigravity_request_support, AntigravityEnvelopeRequestType,
@@ -50,6 +53,7 @@ pub(crate) fn resolve_same_format_provider_transport_unsupported_reason_for_trac
}; };
let behavior = policy::classify_same_format_provider_request_behavior( let behavior = policy::classify_same_format_provider_request_behavior(
transport, transport,
provider_api_format,
crate::ai_serving::planner::spec_metadata::LocalExecutionSurfaceSpecMetadata { crate::ai_serving::planner::spec_metadata::LocalExecutionSurfaceSpecMetadata {
api_format: provider_api_format, api_format: provider_api_format,
require_streaming: false, require_streaming: false,
@@ -131,6 +135,7 @@ pub(crate) async fn resolve_local_same_format_provider_candidate_payload_parts(
prepared.transport.endpoint.body_rules.as_ref(), prepared.transport.endpoint.body_rules.as_ref(),
Some(&parts.headers), Some(&parts.headers),
prepared.upstream_is_stream, prepared.upstream_is_stream,
prepared.force_body_stream_field,
prepared.kiro_auth.as_ref(), prepared.kiro_auth.as_ref(),
prepared.is_claude_code, prepared.is_claude_code,
enable_model_directives, enable_model_directives,
@@ -170,6 +175,18 @@ pub(crate) async fn resolve_local_same_format_provider_candidate_payload_parts(
&mut base_provider_request_body, &mut base_provider_request_body,
&mapping, &mapping,
); );
// Directive mapping is a deep-merge patch and may overwrite/add `stream`;
// re-enforce stream-field policy afterward.
// Kiro behavior classification already hard-requires upstream streaming,
// and the Kiro envelope does not use a top-level body stream field.
if prepared.kiro_auth.is_none() {
enforce_provider_body_stream_policy(
&mut base_provider_request_body,
prepared.provider_api_format.as_str(),
prepared.upstream_is_stream,
request_requires_body_stream_field(body_json, prepared.force_body_stream_field),
);
}
} }
let antigravity_auth = if prepared.is_antigravity { let antigravity_auth = if prepared.is_antigravity {

View File

@@ -13,12 +13,14 @@ use super::super::LocalSameFormatProviderFamily;
pub(super) fn classify_same_format_provider_request_behavior( pub(super) fn classify_same_format_provider_request_behavior(
transport: &GatewayProviderTransportSnapshot, transport: &GatewayProviderTransportSnapshot,
provider_api_format: &str,
spec_metadata: LocalExecutionSurfaceSpecMetadata, spec_metadata: LocalExecutionSurfaceSpecMetadata,
) -> SameFormatProviderRequestBehavior { ) -> SameFormatProviderRequestBehavior {
classify_same_format_provider_request_behavior_impl( classify_same_format_provider_request_behavior_impl(
transport, transport,
SameFormatProviderRequestBehaviorParams { SameFormatProviderRequestBehaviorParams {
require_streaming: spec_metadata.require_streaming, require_streaming: spec_metadata.require_streaming,
provider_api_format,
report_kind: spec_metadata report_kind: spec_metadata
.report_kind .report_kind
.expect("same-format provider specs should declare report kind"), .expect("same-format provider specs should declare report kind"),

View File

@@ -35,6 +35,7 @@ pub(super) struct PreparedSameFormatProviderCandidate {
pub(super) mapped_model: String, pub(super) mapped_model: String,
pub(super) report_kind: &'static str, pub(super) report_kind: &'static str,
pub(super) upstream_is_stream: bool, pub(super) upstream_is_stream: bool,
pub(super) force_body_stream_field: bool,
} }
pub(super) async fn prepare_local_same_format_provider_candidate( pub(super) async fn prepare_local_same_format_provider_candidate(
@@ -51,7 +52,11 @@ pub(super) async fn prepare_local_same_format_provider_candidate(
let candidate = &eligible.candidate; let candidate = &eligible.candidate;
let transport = Arc::clone(&eligible.transport); let transport = Arc::clone(&eligible.transport);
let provider_api_format = eligible.provider_api_format.as_str(); let provider_api_format = eligible.provider_api_format.as_str();
let behavior = classify_same_format_provider_request_behavior(&transport, spec_metadata); let behavior = classify_same_format_provider_request_behavior(
&transport,
provider_api_format,
spec_metadata,
);
if !same_format_provider_transport_supported( if !same_format_provider_transport_supported(
&behavior, &behavior,
@@ -174,5 +179,6 @@ pub(super) async fn prepare_local_same_format_provider_candidate(
mapped_model, mapped_model,
report_kind: behavior.report_kind, report_kind: behavior.report_kind,
upstream_is_stream: behavior.upstream_is_stream, upstream_is_stream: behavior.upstream_is_stream,
force_body_stream_field: behavior.force_body_stream_field,
}) })
} }

View File

@@ -14,6 +14,7 @@ pub(crate) fn build_same_format_provider_request_body(
body_rules: Option<&Value>, body_rules: Option<&Value>,
request_headers: Option<&http::HeaderMap>, request_headers: Option<&http::HeaderMap>,
upstream_is_stream: bool, upstream_is_stream: bool,
force_body_stream_field: bool,
kiro_auth: Option<&crate::ai_serving::transport::kiro::KiroRequestAuth>, kiro_auth: Option<&crate::ai_serving::transport::kiro::KiroRequestAuth>,
is_claude_code: bool, is_claude_code: bool,
enable_model_directives: bool, enable_model_directives: bool,
@@ -28,6 +29,7 @@ pub(crate) fn build_same_format_provider_request_body(
body_rules, body_rules,
request_headers, request_headers,
upstream_is_stream, upstream_is_stream,
force_body_stream_field,
kiro_auth_config: kiro_auth.map(|auth| &auth.auth_config), kiro_auth_config: kiro_auth.map(|auth| &auth.auth_config),
is_claude_code, is_claude_code,
enable_model_directives, enable_model_directives,

View File

@@ -7,7 +7,10 @@ use crate::ai_serving::planner::candidate_preparation::{
prepare_header_authenticated_candidate, prepare_header_authenticated_candidate_from_auth, prepare_header_authenticated_candidate, prepare_header_authenticated_candidate_from_auth,
OauthPreparationContext, OauthPreparationContext,
}; };
use crate::ai_serving::planner::common::force_upstream_streaming_for_provider; use crate::ai_serving::planner::common::{
endpoint_config_forces_body_stream_field, enforce_provider_body_stream_policy,
request_requires_body_stream_field, resolve_upstream_is_stream_for_provider,
};
use crate::ai_serving::planner::spec_metadata::local_standard_spec_metadata; use crate::ai_serving::planner::spec_metadata::local_standard_spec_metadata;
use crate::ai_serving::planner::standard::{ use crate::ai_serving::planner::standard::{
apply_codex_openai_responses_special_headers, request_body_build_failure_extra_data, apply_codex_openai_responses_special_headers, request_body_build_failure_extra_data,
@@ -175,11 +178,15 @@ pub(crate) async fn resolve_local_standard_candidate_payload_parts(
} }
}; };
let upstream_is_stream = spec_metadata.require_streaming let upstream_is_stream = resolve_upstream_is_stream_for_provider(
|| force_upstream_streaming_for_provider( transport.endpoint.config.as_ref(),
transport.provider.provider_type.as_str(), transport.provider.provider_type.as_str(),
provider_api_format, provider_api_format,
spec_metadata.require_streaming,
is_kiro_claude_cli,
); );
let force_body_stream_field =
endpoint_config_forces_body_stream_field(transport.endpoint.config.as_ref());
let enable_model_directives = let enable_model_directives =
crate::system_features::reasoning_model_directive_enabled_for_api_format_and_model( crate::system_features::reasoning_model_directive_enabled_for_api_format_and_model(
state, state,
@@ -225,6 +232,12 @@ pub(crate) async fn resolve_local_standard_candidate_payload_parts(
return None; return None;
} }
}; };
enforce_provider_body_stream_policy(
&mut provider_request_body,
provider_api_format,
upstream_is_stream,
request_requires_body_stream_field(body_json, force_body_stream_field),
);
if let Some(mapping) = if let Some(mapping) =
crate::system_features::reasoning_model_directive_mapping_for_api_format_and_model( crate::system_features::reasoning_model_directive_mapping_for_api_format_and_model(
state, state,
@@ -237,6 +250,14 @@ pub(crate) async fn resolve_local_standard_candidate_payload_parts(
&mut provider_request_body, &mut provider_request_body,
&mapping, &mapping,
); );
// Directive mapping is a deep-merge patch and may overwrite/add `stream`;
// re-enforce stream-field policy afterward.
enforce_provider_body_stream_policy(
&mut provider_request_body,
provider_api_format,
upstream_is_stream,
request_requires_body_stream_field(body_json, force_body_stream_field),
);
} }
if let Some(kiro_auth) = kiro_auth.as_ref() { if let Some(kiro_auth) = kiro_auth.as_ref() {

View File

@@ -116,9 +116,9 @@ pub(crate) fn build_gemini_stream_plan_from_decision(
content_type: payload.content_type.as_deref(), content_type: payload.content_type.as_deref(),
provider_api_format: core.provider_api_format.as_str(), provider_api_format: core.provider_api_format.as_str(),
client_api_format: core.client_api_format.as_str(), client_api_format: core.client_api_format.as_str(),
upstream_is_stream: true, upstream_is_stream: payload.upstream_is_stream,
build_from_request_when_empty: false, build_from_request_when_empty: false,
accept_policy: StandardPlanFallbackAcceptPolicy::TextEventStreamRequired, accept_policy: StandardPlanFallbackAcceptPolicy::TextEventStreamIfStreaming,
}); });
let content_type = payload let content_type = payload
.content_type .content_type

View File

@@ -15,3 +15,6 @@ pub(crate) use self::responses::{
build_cross_format_openai_responses_upstream_url, build_local_openai_responses_request_body, build_cross_format_openai_responses_upstream_url, build_local_openai_responses_request_body,
build_local_openai_responses_upstream_url, build_local_openai_responses_upstream_url,
}; };
pub(super) use crate::ai_serving::planner::common::{
enforce_provider_body_stream_policy, request_requires_body_stream_field,
};

View File

@@ -9,10 +9,13 @@ use crate::ai_serving::{
GatewayProviderTransportSnapshot, GatewayProviderTransportSnapshot,
}; };
use super::{enforce_provider_body_stream_policy, request_requires_body_stream_field};
pub(crate) fn build_local_openai_chat_request_body( pub(crate) fn build_local_openai_chat_request_body(
body_json: &Value, body_json: &Value,
mapped_model: &str, mapped_model: &str,
upstream_is_stream: bool, upstream_is_stream: bool,
force_body_stream_field: bool,
body_rules: Option<&Value>, body_rules: Option<&Value>,
request_headers: &http::HeaderMap, request_headers: &http::HeaderMap,
enable_model_directives: bool, enable_model_directives: bool,
@@ -23,12 +26,20 @@ pub(crate) fn build_local_openai_chat_request_body(
upstream_is_stream, upstream_is_stream,
enable_model_directives, enable_model_directives,
)?; )?;
let mut provider_request_body =
apply_standard_provider_request_body_rules_with_request_headers( apply_standard_provider_request_body_rules_with_request_headers(
provider_request_body, provider_request_body,
body_rules, body_rules,
body_json, body_json,
request_headers, request_headers,
) )?;
enforce_provider_body_stream_policy(
&mut provider_request_body,
"openai:chat",
upstream_is_stream,
request_requires_body_stream_field(body_json, force_body_stream_field),
);
Some(provider_request_body)
} }
pub(crate) fn build_local_openai_chat_upstream_url( pub(crate) fn build_local_openai_chat_upstream_url(
@@ -44,6 +55,7 @@ pub(crate) fn build_cross_format_openai_chat_request_body(
provider_type: &str, provider_type: &str,
provider_api_format: &str, provider_api_format: &str,
upstream_is_stream: bool, upstream_is_stream: bool,
force_body_stream_field: bool,
body_rules: Option<&Value>, body_rules: Option<&Value>,
user_api_key_id: Option<&str>, user_api_key_id: Option<&str>,
request_headers: &http::HeaderMap, request_headers: &http::HeaderMap,
@@ -74,6 +86,12 @@ pub(crate) fn build_cross_format_openai_chat_request_body(
&mut provider_request_body, &mut provider_request_body,
provider_api_format, provider_api_format,
); );
enforce_provider_body_stream_policy(
&mut provider_request_body,
provider_api_format,
upstream_is_stream,
request_requires_body_stream_field(body_json, force_body_stream_field),
);
Some(provider_request_body) Some(provider_request_body)
} }

View File

@@ -9,10 +9,13 @@ use crate::ai_serving::{
GatewayProviderTransportSnapshot, GatewayProviderTransportSnapshot,
}; };
use super::{enforce_provider_body_stream_policy, request_requires_body_stream_field};
pub(crate) fn build_local_openai_responses_request_body( pub(crate) fn build_local_openai_responses_request_body(
body_json: &Value, body_json: &Value,
mapped_model: &str, mapped_model: &str,
require_streaming: bool, require_streaming: bool,
force_body_stream_field: bool,
provider_type: &str, provider_type: &str,
provider_api_format: &str, provider_api_format: &str,
body_rules: Option<&Value>, body_rules: Option<&Value>,
@@ -44,6 +47,12 @@ pub(crate) fn build_local_openai_responses_request_body(
&mut provider_request_body, &mut provider_request_body,
provider_api_format, provider_api_format,
); );
enforce_provider_body_stream_policy(
&mut provider_request_body,
provider_api_format,
require_streaming,
request_requires_body_stream_field(body_json, force_body_stream_field),
);
Some(provider_request_body) Some(provider_request_body)
} }
@@ -53,6 +62,7 @@ pub(crate) fn build_cross_format_openai_responses_request_body(
client_api_format: &str, client_api_format: &str,
provider_api_format: &str, provider_api_format: &str,
upstream_is_stream: bool, upstream_is_stream: bool,
force_body_stream_field: bool,
provider_type: &str, provider_type: &str,
body_rules: Option<&Value>, body_rules: Option<&Value>,
user_api_key_id: Option<&str>, user_api_key_id: Option<&str>,
@@ -85,6 +95,12 @@ pub(crate) fn build_cross_format_openai_responses_request_body(
&mut provider_request_body, &mut provider_request_body,
provider_api_format, provider_api_format,
); );
enforce_provider_body_stream_policy(
&mut provider_request_body,
provider_api_format,
upstream_is_stream,
request_requires_body_stream_field(body_json, force_body_stream_field),
);
Some(provider_request_body) Some(provider_request_body)
} }

View File

@@ -87,6 +87,7 @@ fn builds_openai_chat_cross_format_request_body_from_openai_responses_source() {
"openai:responses", "openai:responses",
"openai:chat", "openai:chat",
false, false,
false,
"openai", "openai",
None, None,
None, None,
@@ -121,6 +122,7 @@ fn local_openai_responses_wrapper_preserves_body_order_after_edits() {
&body_json, &body_json,
"gpt-5.4", "gpt-5.4",
true, true,
false,
"codex", "codex",
"openai:responses", "openai:responses",
None, None,
@@ -160,6 +162,7 @@ fn local_openai_responses_compact_wrapper_strips_store_for_same_format_requests(
&body_json, &body_json,
"gpt-5.4", "gpt-5.4",
false, false,
false,
"openai", "openai",
"openai:responses:compact", "openai:responses:compact",
None, None,
@@ -170,6 +173,7 @@ fn local_openai_responses_compact_wrapper_strips_store_for_same_format_requests(
.expect("local openai compact body should build"); .expect("local openai compact body should build");
assert!(provider_request_body.get("store").is_none()); assert!(provider_request_body.get("store").is_none());
assert!(provider_request_body.get("stream").is_none());
} }
#[test] #[test]
@@ -187,6 +191,7 @@ fn local_openai_responses_wrapper_applies_model_directive_before_body_rules() {
&body_json, &body_json,
"gpt-5.4", "gpt-5.4",
false, false,
false,
"openai", "openai",
"openai:responses", "openai:responses",
Some(&body_rules), Some(&body_rules),
@@ -237,6 +242,7 @@ fn strips_metadata_for_codex_openai_responses_requests() {
"claude:messages", "claude:messages",
"openai:responses", "openai:responses",
true, true,
false,
"codex", "codex",
None, None,
None, None,
@@ -271,6 +277,7 @@ fn applies_codex_defaults_unless_body_rules_handle_the_field() {
"claude:messages", "claude:messages",
"openai:responses", "openai:responses",
true, true,
false,
"codex", "codex",
Some(&body_rules), Some(&body_rules),
None, None,
@@ -300,6 +307,7 @@ fn injects_codex_prompt_cache_key_for_openai_responses_cross_format_requests() {
"claude:messages", "claude:messages",
"openai:responses", "openai:responses",
true, true,
false,
"codex", "codex",
None, None,
Some("key-123"), Some("key-123"),
@@ -330,6 +338,7 @@ fn injects_codex_prompt_cache_key_for_openai_chat_cross_format_requests() {
"codex", "codex",
"openai:responses", "openai:responses",
false, false,
false,
None, None,
Some("key-123"), Some("key-123"),
&http::HeaderMap::new(), &http::HeaderMap::new(),

View File

@@ -1,4 +1,5 @@
use crate::ai_serving::build_request_trace_proxy_value; use crate::ai_serving::build_request_trace_proxy_value;
use crate::ai_serving::planner::common::OPENAI_CHAT_STREAM_PLAN_KIND;
use crate::ai_serving::planner::report_context::{ use crate::ai_serving::planner::report_context::{
build_local_execution_report_context, insert_provider_stream_event_api_format, build_local_execution_report_context, insert_provider_stream_event_api_format,
LocalExecutionReportContextParts, LocalExecutionReportContextParts,
@@ -29,6 +30,7 @@ pub(crate) async fn maybe_build_local_openai_chat_decision_payload_for_candidate
report_kind: &str, report_kind: &str,
upstream_is_stream: bool, upstream_is_stream: bool,
) -> Option<AiExecutionDecision> { ) -> Option<AiExecutionDecision> {
let decision_is_stream = decision_kind == OPENAI_CHAT_STREAM_PLAN_KIND;
let attempt_identity = attempt.attempt_identity(); let attempt_identity = attempt.attempt_identity();
let LocalOpenAiChatCandidateAttempt { let LocalOpenAiChatCandidateAttempt {
eligible, eligible,
@@ -149,7 +151,7 @@ pub(crate) async fn maybe_build_local_openai_chat_decision_payload_for_candidate
Some(build_ai_execution_decision_response( Some(build_ai_execution_decision_response(
AiExecutionDecisionResponseParts { AiExecutionDecisionResponseParts {
decision_is_stream: upstream_is_stream, decision_is_stream,
decision_kind: decision_kind.to_string(), decision_kind: decision_kind.to_string(),
execution_strategy, execution_strategy,
conversion_mode, conversion_mode,

View File

@@ -8,7 +8,10 @@ use crate::ai_serving::planner::candidate_preparation::{
OauthPreparationContext, OauthPreparationContext,
}; };
use crate::ai_serving::planner::candidate_resolution::EligibleLocalExecutionCandidate; use crate::ai_serving::planner::candidate_resolution::EligibleLocalExecutionCandidate;
use crate::ai_serving::planner::common::OPENAI_CHAT_STREAM_PLAN_KIND; use crate::ai_serving::planner::common::{
endpoint_config_forces_body_stream_field, enforce_provider_body_stream_policy,
request_requires_body_stream_field, OPENAI_CHAT_STREAM_PLAN_KIND,
};
use crate::ai_serving::planner::standard::{ use crate::ai_serving::planner::standard::{
apply_codex_openai_responses_special_headers, build_cross_format_openai_chat_request_body, apply_codex_openai_responses_special_headers, build_cross_format_openai_chat_request_body,
build_cross_format_openai_chat_upstream_url, build_local_openai_chat_request_body, build_cross_format_openai_chat_upstream_url, build_local_openai_chat_request_body,
@@ -72,6 +75,8 @@ pub(crate) async fn resolve_local_openai_chat_candidate_payload_parts(
let candidate = &eligible.candidate; let candidate = &eligible.candidate;
let provider_api_format = eligible.provider_api_format.as_str(); let provider_api_format = eligible.provider_api_format.as_str();
let transport = &eligible.transport; let transport = &eligible.transport;
let force_body_stream_field =
endpoint_config_forces_body_stream_field(transport.endpoint.config.as_ref());
let enable_model_directives = let enable_model_directives =
crate::system_features::reasoning_model_directive_enabled_for_api_format_and_model( crate::system_features::reasoning_model_directive_enabled_for_api_format_and_model(
state, state,
@@ -128,6 +133,7 @@ pub(crate) async fn resolve_local_openai_chat_candidate_payload_parts(
body_json, body_json,
&prepared_candidate.mapped_model, &prepared_candidate.mapped_model,
upstream_is_stream, upstream_is_stream,
force_body_stream_field,
transport.endpoint.body_rules.as_ref(), transport.endpoint.body_rules.as_ref(),
&parts.headers, &parts.headers,
enable_model_directives, enable_model_directives,
@@ -214,6 +220,13 @@ pub(crate) async fn resolve_local_openai_chat_candidate_payload_parts(
let (execution_strategy, conversion_mode) = let (execution_strategy, conversion_mode) =
ai_local_execution_contract_for_formats("openai:chat", "openai:chat"); ai_local_execution_contract_for_formats("openai:chat", "openai:chat");
let resolved_report_kind =
if decision_kind == OPENAI_CHAT_STREAM_PLAN_KIND || !upstream_is_stream {
report_kind.to_string()
} else {
"openai_chat_sync_finalize".to_string()
};
return Some(LocalOpenAiChatCandidatePayloadParts { return Some(LocalOpenAiChatCandidatePayloadParts {
auth_header: resolved_headers.auth_header, auth_header: resolved_headers.auth_header,
auth_value: resolved_headers.auth_value, auth_value: resolved_headers.auth_value,
@@ -224,7 +237,7 @@ pub(crate) async fn resolve_local_openai_chat_candidate_payload_parts(
upstream_url, upstream_url,
execution_strategy, execution_strategy,
conversion_mode, conversion_mode,
report_kind: report_kind.to_string(), report_kind: resolved_report_kind,
envelope_name: None, envelope_name: None,
transport: Arc::clone(transport), transport: Arc::clone(transport),
}); });
@@ -349,6 +362,7 @@ pub(crate) async fn resolve_local_openai_chat_candidate_payload_parts(
transport.provider.provider_type.as_str(), transport.provider.provider_type.as_str(),
provider_api_format.as_str(), provider_api_format.as_str(),
upstream_is_stream, upstream_is_stream,
force_body_stream_field,
if is_kiro_claude_cli { if is_kiro_claude_cli {
None None
} else { } else {
@@ -387,6 +401,14 @@ pub(crate) async fn resolve_local_openai_chat_candidate_payload_parts(
&mut provider_request_body, &mut provider_request_body,
&mapping, &mapping,
); );
// Directive mapping is a deep-merge patch and may overwrite/add `stream`;
// re-enforce stream-field policy afterward.
enforce_provider_body_stream_policy(
&mut provider_request_body,
provider_api_format.as_str(),
upstream_is_stream,
request_requires_body_stream_field(body_json, force_body_stream_field),
);
} }
if let Some(kiro_auth) = kiro_auth.as_ref() { if let Some(kiro_auth) = kiro_auth.as_ref() {

View File

@@ -9,6 +9,9 @@ mod stream;
#[path = "plans/sync.rs"] #[path = "plans/sync.rs"]
mod sync; mod sync;
use crate::ai_serving::planner::common::resolve_upstream_is_stream_for_provider;
use crate::ai_serving::GatewayProviderTransportSnapshot;
pub(super) use self::candidates::list_local_openai_chat_candidates; pub(super) use self::candidates::list_local_openai_chat_candidates;
pub(super) use self::diagnostic::set_local_openai_chat_miss_diagnostic; pub(super) use self::diagnostic::set_local_openai_chat_miss_diagnostic;
pub(super) use self::resolve::resolve_local_openai_chat_decision_input; pub(super) use self::resolve::resolve_local_openai_chat_decision_input;
@@ -18,3 +21,147 @@ pub(super) use self::stream::{
pub(super) use self::sync::{ pub(super) use self::sync::{
build_local_openai_chat_sync_attempt_source, build_local_openai_chat_sync_plan_and_reports, build_local_openai_chat_sync_attempt_source, build_local_openai_chat_sync_plan_and_reports,
}; };
fn openai_chat_upstream_is_stream_for_candidate(
transport: &GatewayProviderTransportSnapshot,
provider_api_format: &str,
client_is_stream: bool,
) -> bool {
let hard_requires_streaming =
crate::ai_serving::transport::kiro::is_kiro_claude_messages_transport(
transport,
provider_api_format,
);
resolve_upstream_is_stream_for_provider(
transport.endpoint.config.as_ref(),
transport.provider.provider_type.as_str(),
provider_api_format,
client_is_stream,
hard_requires_streaming,
)
}
#[cfg(test)]
mod tests {
use super::openai_chat_upstream_is_stream_for_candidate;
use aether_provider_transport::snapshot::{
GatewayProviderTransportEndpoint, GatewayProviderTransportKey,
GatewayProviderTransportProvider, GatewayProviderTransportSnapshot,
};
use serde_json::{json, Value};
fn sample_transport(
provider_type: &str,
api_format: &str,
endpoint_config: Option<Value>,
) -> GatewayProviderTransportSnapshot {
GatewayProviderTransportSnapshot {
provider: GatewayProviderTransportProvider {
id: "provider-1".to_string(),
name: "Provider".to_string(),
provider_type: provider_type.to_string(),
website: None,
is_active: true,
keep_priority_on_conversion: false,
enable_format_conversion: true,
concurrent_limit: None,
max_retries: None,
proxy: None,
request_timeout_secs: None,
stream_first_byte_timeout_secs: None,
config: None,
},
endpoint: GatewayProviderTransportEndpoint {
id: "endpoint-1".to_string(),
provider_id: "provider-1".to_string(),
api_format: api_format.to_string(),
api_family: None,
endpoint_kind: None,
is_active: true,
base_url: "https://api.example.test".to_string(),
header_rules: None,
body_rules: None,
max_retries: None,
custom_path: None,
config: endpoint_config,
format_acceptance_config: None,
proxy: None,
},
key: GatewayProviderTransportKey {
id: "key-1".to_string(),
provider_id: "provider-1".to_string(),
name: "key".to_string(),
auth_type: "api_key".to_string(),
is_active: true,
api_formats: None,
auth_type_by_format: None,
allow_auth_channel_mismatch_formats: None,
allowed_models: None,
capabilities: None,
rate_multipliers: None,
global_priority_by_format: None,
expires_at_unix_secs: None,
proxy: None,
fingerprint: None,
decrypted_api_key: "secret".to_string(),
decrypted_auth_config: None,
},
}
}
#[test]
fn openai_chat_policy_resolver_supports_force_stream_force_non_stream_and_auto() {
let force_stream = sample_transport(
"openai",
"openai:chat",
Some(json!({"upstream_stream_policy": "force_stream"})),
);
assert!(openai_chat_upstream_is_stream_for_candidate(
&force_stream,
"openai:chat",
false,
));
let force_non_stream = sample_transport(
"openai",
"openai:chat",
Some(json!({"upstream_stream_policy": "force_non_stream"})),
);
assert!(!openai_chat_upstream_is_stream_for_candidate(
&force_non_stream,
"openai:chat",
true,
));
let auto = sample_transport(
"openai",
"openai:chat",
Some(json!({"upstream_stream_policy": "auto"})),
);
assert!(openai_chat_upstream_is_stream_for_candidate(
&auto,
"openai:chat",
true,
));
assert!(!openai_chat_upstream_is_stream_for_candidate(
&auto,
"openai:chat",
false,
));
}
#[test]
fn openai_chat_policy_resolver_preserves_provider_hard_streaming() {
let codex = sample_transport(
"codex",
"openai:responses",
Some(json!({"upstream_stream_policy": "force_non_stream"})),
);
assert!(openai_chat_upstream_is_stream_for_candidate(
&codex,
"openai:responses",
false,
));
}
}

View File

@@ -13,6 +13,7 @@ use super::candidates::list_local_openai_chat_candidates;
use super::diagnostic::{ use super::diagnostic::{
set_local_openai_chat_candidate_evaluation_diagnostic, set_local_openai_chat_miss_diagnostic, set_local_openai_chat_candidate_evaluation_diagnostic, set_local_openai_chat_miss_diagnostic,
}; };
use super::openai_chat_upstream_is_stream_for_candidate;
use super::resolve::resolve_local_openai_chat_decision_input; use super::resolve::resolve_local_openai_chat_decision_input;
use crate::ai_serving::planner::candidate_materialization::LocalExecutionAttemptSource; use crate::ai_serving::planner::candidate_materialization::LocalExecutionAttemptSource;
use crate::ai_serving::planner::common::OPENAI_CHAT_STREAM_PLAN_KIND; use crate::ai_serving::planner::common::OPENAI_CHAT_STREAM_PLAN_KIND;
@@ -120,6 +121,11 @@ impl LocalOpenAiChatStreamAttemptSource<'_> {
&self, &self,
attempt: LocalOpenAiChatCandidateAttempt, attempt: LocalOpenAiChatCandidateAttempt,
) -> Result<Option<AiStreamAttempt>, GatewayError> { ) -> Result<Option<AiStreamAttempt>, GatewayError> {
let upstream_is_stream = openai_chat_upstream_is_stream_for_candidate(
&attempt.eligible.transport,
attempt.eligible.provider_api_format.as_str(),
true,
);
let Some(payload) = maybe_build_local_openai_chat_decision_payload_for_candidate( let Some(payload) = maybe_build_local_openai_chat_decision_payload_for_candidate(
self.state, self.state,
self.parts, self.parts,
@@ -129,7 +135,7 @@ impl LocalOpenAiChatStreamAttemptSource<'_> {
attempt, attempt,
OPENAI_CHAT_STREAM_PLAN_KIND, OPENAI_CHAT_STREAM_PLAN_KIND,
"openai_chat_stream_success", "openai_chat_stream_success",
true, upstream_is_stream,
) )
.await .await
else { else {
@@ -222,6 +228,11 @@ pub(crate) async fn build_local_openai_chat_stream_plan_and_reports(
let mut plans = Vec::new(); let mut plans = Vec::new();
for attempt in attempts { for attempt in attempts {
let upstream_is_stream = openai_chat_upstream_is_stream_for_candidate(
&attempt.eligible.transport,
attempt.eligible.provider_api_format.as_str(),
true,
);
let Some(payload) = maybe_build_local_openai_chat_decision_payload_for_candidate( let Some(payload) = maybe_build_local_openai_chat_decision_payload_for_candidate(
state, state,
parts, parts,
@@ -231,7 +242,7 @@ pub(crate) async fn build_local_openai_chat_stream_plan_and_reports(
attempt, attempt,
OPENAI_CHAT_STREAM_PLAN_KIND, OPENAI_CHAT_STREAM_PLAN_KIND,
"openai_chat_stream_success", "openai_chat_stream_success",
true, upstream_is_stream,
) )
.await .await
else { else {

View File

@@ -13,11 +13,10 @@ use super::candidates::list_local_openai_chat_candidates;
use super::diagnostic::{ use super::diagnostic::{
set_local_openai_chat_candidate_evaluation_diagnostic, set_local_openai_chat_miss_diagnostic, set_local_openai_chat_candidate_evaluation_diagnostic, set_local_openai_chat_miss_diagnostic,
}; };
use super::openai_chat_upstream_is_stream_for_candidate;
use super::resolve::resolve_local_openai_chat_decision_input; use super::resolve::resolve_local_openai_chat_decision_input;
use crate::ai_serving::planner::candidate_materialization::LocalExecutionAttemptSource; use crate::ai_serving::planner::candidate_materialization::LocalExecutionAttemptSource;
use crate::ai_serving::planner::common::{ use crate::ai_serving::planner::common::OPENAI_CHAT_SYNC_PLAN_KIND;
force_upstream_streaming_for_provider, OPENAI_CHAT_SYNC_PLAN_KIND,
};
use crate::ai_serving::planner::plan_builders::{ use crate::ai_serving::planner::plan_builders::{
build_openai_chat_sync_plan_from_decision, AiSyncAttempt, build_openai_chat_sync_plan_from_decision, AiSyncAttempt,
}; };
@@ -32,13 +31,6 @@ pub(crate) struct LocalOpenAiChatSyncAttemptSource<'a> {
candidates: LocalOpenAiChatCandidateAttemptSource<'a>, candidates: LocalOpenAiChatCandidateAttemptSource<'a>,
} }
fn openai_chat_sync_upstream_is_stream_for_candidate(
provider_type: &str,
provider_api_format: &str,
) -> bool {
force_upstream_streaming_for_provider(provider_type, provider_api_format)
}
pub(crate) async fn build_local_openai_chat_sync_attempt_source<'a>( pub(crate) async fn build_local_openai_chat_sync_attempt_source<'a>(
state: &'a AppState, state: &'a AppState,
parts: &'a http::request::Parts, parts: &'a http::request::Parts,
@@ -129,9 +121,10 @@ impl LocalOpenAiChatSyncAttemptSource<'_> {
&self, &self,
attempt: LocalOpenAiChatCandidateAttempt, attempt: LocalOpenAiChatCandidateAttempt,
) -> Result<Option<AiSyncAttempt>, GatewayError> { ) -> Result<Option<AiSyncAttempt>, GatewayError> {
let upstream_is_stream = openai_chat_sync_upstream_is_stream_for_candidate( let upstream_is_stream = openai_chat_upstream_is_stream_for_candidate(
attempt.eligible.transport.provider.provider_type.as_str(), &attempt.eligible.transport,
attempt.eligible.provider_api_format.as_str(), attempt.eligible.provider_api_format.as_str(),
false,
); );
let Some(payload) = maybe_build_local_openai_chat_decision_payload_for_candidate( let Some(payload) = maybe_build_local_openai_chat_decision_payload_for_candidate(
self.state, self.state,
@@ -235,9 +228,10 @@ pub(crate) async fn build_local_openai_chat_sync_plan_and_reports(
let mut plans = Vec::new(); let mut plans = Vec::new();
for attempt in attempts { for attempt in attempts {
let upstream_is_stream = openai_chat_sync_upstream_is_stream_for_candidate( let upstream_is_stream = openai_chat_upstream_is_stream_for_candidate(
attempt.eligible.transport.provider.provider_type.as_str(), &attempt.eligible.transport,
attempt.eligible.provider_api_format.as_str(), attempt.eligible.provider_api_format.as_str(),
false,
); );
let Some(payload) = maybe_build_local_openai_chat_decision_payload_for_candidate( let Some(payload) = maybe_build_local_openai_chat_decision_payload_for_candidate(
state, state,
@@ -272,28 +266,3 @@ pub(crate) async fn build_local_openai_chat_sync_plan_and_reports(
Ok(plans) Ok(plans)
} }
#[cfg(test)]
mod tests {
use super::openai_chat_sync_upstream_is_stream_for_candidate;
#[test]
fn openai_chat_sync_forces_streaming_for_codex_openai_responses_candidates() {
assert!(openai_chat_sync_upstream_is_stream_for_candidate(
"codex",
"openai:responses"
));
assert!(!openai_chat_sync_upstream_is_stream_for_candidate(
"codex",
"openai:responses:compact"
));
assert!(!openai_chat_sync_upstream_is_stream_for_candidate(
"openai",
"openai:responses"
));
assert!(!openai_chat_sync_upstream_is_stream_for_candidate(
"codex",
"openai:chat"
));
}
}

View File

@@ -9,6 +9,7 @@ use super::super::{
take_ai_upstream_auth_pair, take_non_empty_string, AiExecutionPlanFromDecisionParts, take_ai_upstream_auth_pair, take_non_empty_string, AiExecutionPlanFromDecisionParts,
AiStreamAttempt, AiStreamAttempt,
}; };
use crate::ai_serving::planner::common::enforce_provider_body_stream_policy;
use crate::ai_serving::provider_adaptation_requires_eventstream_accept; use crate::ai_serving::provider_adaptation_requires_eventstream_accept;
use crate::ai_serving::transport::{ use crate::ai_serving::transport::{
build_standard_plan_fallback_headers, build_standard_plan_fallback_openai_chat_url, build_standard_plan_fallback_headers, build_standard_plan_fallback_openai_chat_url,
@@ -53,21 +54,31 @@ pub(crate) fn build_openai_chat_stream_plan_from_decision(
provider_request_body provider_request_body
.insert("model".to_string(), serde_json::Value::String(mapped_model)); .insert("model".to_string(), serde_json::Value::String(mapped_model));
} }
provider_request_body.insert("stream".to_string(), serde_json::Value::Bool(true)); let require_body_stream_field = provider_request_body.contains_key("stream");
let mut provider_request_body = serde_json::Value::Object(provider_request_body);
enforce_provider_body_stream_policy(
&mut provider_request_body,
core.provider_api_format.as_str(),
payload.upstream_is_stream,
require_body_stream_field,
);
let Some(provider_request_object) = provider_request_body.as_object_mut() else {
return Ok(None);
};
if let Some(prompt_cache_key) = take_non_empty_string(&mut payload.prompt_cache_key) { if let Some(prompt_cache_key) = take_non_empty_string(&mut payload.prompt_cache_key) {
let existing = provider_request_body let existing = provider_request_object
.get("prompt_cache_key") .get("prompt_cache_key")
.and_then(|value| value.as_str()) .and_then(|value| value.as_str())
.map(str::trim) .map(str::trim)
.unwrap_or_default(); .unwrap_or_default();
if existing.is_empty() { if existing.is_empty() {
provider_request_body.insert( provider_request_object.insert(
"prompt_cache_key".to_string(), "prompt_cache_key".to_string(),
serde_json::Value::String(prompt_cache_key), serde_json::Value::String(prompt_cache_key),
); );
} }
} }
serde_json::Value::Object(provider_request_body) provider_request_body
}; };
let extra_headers = std::mem::take(&mut payload.extra_headers); let extra_headers = std::mem::take(&mut payload.extra_headers);
let mut provider_request_headers = let mut provider_request_headers =
@@ -82,9 +93,9 @@ pub(crate) fn build_openai_chat_stream_plan_from_decision(
content_type: payload.content_type.as_deref(), content_type: payload.content_type.as_deref(),
provider_api_format: core.provider_api_format.as_str(), provider_api_format: core.provider_api_format.as_str(),
client_api_format: core.client_api_format.as_str(), client_api_format: core.client_api_format.as_str(),
upstream_is_stream: true, upstream_is_stream: payload.upstream_is_stream,
build_from_request_when_empty: true, build_from_request_when_empty: true,
accept_policy: StandardPlanFallbackAcceptPolicy::TextEventStreamRequired, accept_policy: StandardPlanFallbackAcceptPolicy::TextEventStreamIfStreaming,
}); });
let content_type = payload let content_type = payload
.content_type .content_type
@@ -157,13 +168,14 @@ pub(crate) fn build_openai_responses_stream_plan_from_decision(
.as_ref() .as_ref()
.and_then(|context| context.get("envelope_name")) .and_then(|context| context.get("envelope_name"))
.and_then(serde_json::Value::as_str); .and_then(serde_json::Value::as_str);
let accept_policy = if provider_adaptation_requires_eventstream_accept( let accept_policy = if payload.upstream_is_stream
&& provider_adaptation_requires_eventstream_accept(
envelope_name, envelope_name,
core.provider_api_format.as_str(), core.provider_api_format.as_str(),
) { ) {
StandardPlanFallbackAcceptPolicy::ProviderEventStreamIfMissing StandardPlanFallbackAcceptPolicy::ProviderEventStreamIfMissing
} else { } else {
StandardPlanFallbackAcceptPolicy::TextEventStreamRequired StandardPlanFallbackAcceptPolicy::TextEventStreamIfStreaming
}; };
let mut provider_request_headers = let mut provider_request_headers =
build_standard_plan_fallback_headers(StandardPlanFallbackHeadersInput { build_standard_plan_fallback_headers(StandardPlanFallbackHeadersInput {
@@ -177,7 +189,7 @@ pub(crate) fn build_openai_responses_stream_plan_from_decision(
content_type: payload.content_type.as_deref(), content_type: payload.content_type.as_deref(),
provider_api_format: core.provider_api_format.as_str(), provider_api_format: core.provider_api_format.as_str(),
client_api_format: core.client_api_format.as_str(), client_api_format: core.client_api_format.as_str(),
upstream_is_stream: true, upstream_is_stream: payload.upstream_is_stream,
build_from_request_when_empty: false, build_from_request_when_empty: false,
accept_policy, accept_policy,
}); });
@@ -427,6 +439,104 @@ mod tests {
); );
} }
#[test]
fn build_openai_chat_stream_plan_keeps_downstream_stream_for_force_non_stream_upstream() {
fn force_non_stream_payload(provider_request_body: Option<Value>) -> AiExecutionDecision {
AiExecutionDecision {
action: "stream".to_string(),
decision_kind: Some("openai_chat_stream".to_string()),
execution_strategy: None,
conversion_mode: None,
request_id: Some("req_force_non_stream".to_string()),
candidate_id: Some("cand_force_non_stream".to_string()),
provider_name: Some("OpenAI".to_string()),
provider_id: Some("prov_force_non_stream".to_string()),
endpoint_id: Some("ep_force_non_stream".to_string()),
key_id: Some("key_force_non_stream".to_string()),
upstream_base_url: Some("https://example.com".to_string()),
upstream_url: Some("https://example.com/v1/chat/completions".to_string()),
provider_request_method: None,
auth_header: Some("authorization".to_string()),
auth_value: Some("Bearer upstream-token".to_string()),
provider_api_format: Some("openai:chat".to_string()),
client_api_format: Some("openai:chat".to_string()),
provider_contract: Some("openai:chat".to_string()),
client_contract: Some("openai:chat".to_string()),
model_name: Some("gpt-5.4".to_string()),
mapped_model: Some("gpt-5.4".to_string()),
prompt_cache_key: None,
extra_headers: BTreeMap::new(),
provider_request_headers: BTreeMap::new(),
provider_request_body,
provider_request_body_base64: None,
content_type: Some("application/json".to_string()),
proxy: None,
transport_profile: None,
timeouts: None,
upstream_is_stream: false,
report_kind: Some("openai_chat_stream_success".to_string()),
report_context: Some(json!({})),
auth_context: None,
}
}
let parts = http::Request::builder()
.uri("http://localhost/v1/chat/completions")
.body(())
.expect("request should build")
.into_parts()
.0;
let built = build_openai_chat_stream_plan_from_decision(
&parts,
&json!({}),
force_non_stream_payload(Some(json!({
"model": "gpt-5.4",
"messages": [],
"stream": false
}))),
)
.expect("plan build should succeed")
.expect("plan should be produced");
assert!(built.plan.stream);
assert_eq!(
built
.plan
.body
.json_body
.as_ref()
.and_then(|body| body.get("stream"))
.and_then(Value::as_bool),
Some(false)
);
let fallback_body = json!({
"model": "client-model",
"messages": [],
"stream": true
});
let built = build_openai_chat_stream_plan_from_decision(
&parts,
&fallback_body,
force_non_stream_payload(None),
)
.expect("fallback plan build should succeed")
.expect("fallback plan should be produced");
assert!(built.plan.stream);
assert_eq!(
built
.plan
.body
.json_body
.as_ref()
.and_then(|body| body.get("stream"))
.and_then(Value::as_bool),
Some(false)
);
}
#[test] #[test]
fn build_openai_chat_stream_plan_fallback_restores_claude_headers_for_cross_format() { fn build_openai_chat_stream_plan_fallback_restores_claude_headers_for_cross_format() {
let parts = http::Request::builder() let parts = http::Request::builder()

View File

@@ -9,6 +9,7 @@ use super::super::{
take_ai_upstream_auth_pair, take_non_empty_string, AiExecutionPlanFromDecisionParts, take_ai_upstream_auth_pair, take_non_empty_string, AiExecutionPlanFromDecisionParts,
AiSyncAttempt, AiSyncAttempt,
}; };
use crate::ai_serving::planner::common::enforce_provider_body_stream_policy;
use crate::ai_serving::transport::{ use crate::ai_serving::transport::{
build_standard_plan_fallback_headers, build_standard_plan_fallback_openai_chat_url, build_standard_plan_fallback_headers, build_standard_plan_fallback_openai_chat_url,
build_standard_plan_fallback_openai_responses_url, StandardPlanFallbackAcceptPolicy, build_standard_plan_fallback_openai_responses_url, StandardPlanFallbackAcceptPolicy,
@@ -51,23 +52,31 @@ pub(crate) fn build_openai_chat_sync_plan_from_decision(
provider_request_body provider_request_body
.insert("model".to_string(), serde_json::Value::String(mapped_model)); .insert("model".to_string(), serde_json::Value::String(mapped_model));
} }
if payload.upstream_is_stream { let require_body_stream_field = provider_request_body.contains_key("stream");
provider_request_body.insert("stream".to_string(), serde_json::Value::Bool(true)); let mut provider_request_body = serde_json::Value::Object(provider_request_body);
} enforce_provider_body_stream_policy(
&mut provider_request_body,
core.provider_api_format.as_str(),
payload.upstream_is_stream,
require_body_stream_field,
);
let Some(provider_request_object) = provider_request_body.as_object_mut() else {
return Ok(None);
};
if let Some(prompt_cache_key) = take_non_empty_string(&mut payload.prompt_cache_key) { if let Some(prompt_cache_key) = take_non_empty_string(&mut payload.prompt_cache_key) {
let existing = provider_request_body let existing = provider_request_object
.get("prompt_cache_key") .get("prompt_cache_key")
.and_then(|value| value.as_str()) .and_then(|value| value.as_str())
.map(str::trim) .map(str::trim)
.unwrap_or_default(); .unwrap_or_default();
if existing.is_empty() { if existing.is_empty() {
provider_request_body.insert( provider_request_object.insert(
"prompt_cache_key".to_string(), "prompt_cache_key".to_string(),
serde_json::Value::String(prompt_cache_key), serde_json::Value::String(prompt_cache_key),
); );
} }
} }
serde_json::Value::Object(provider_request_body) provider_request_body
}; };
let extra_headers = std::mem::take(&mut payload.extra_headers); let extra_headers = std::mem::take(&mut payload.extra_headers);
let mut provider_request_headers = let mut provider_request_headers =

View File

@@ -9,7 +9,10 @@ use crate::ai_serving::planner::candidate_preparation::{
OauthPreparationContext, OauthPreparationContext,
}; };
use crate::ai_serving::planner::candidate_resolution::EligibleLocalExecutionCandidate; use crate::ai_serving::planner::candidate_resolution::EligibleLocalExecutionCandidate;
use crate::ai_serving::planner::common::force_upstream_streaming_for_provider; use crate::ai_serving::planner::common::{
endpoint_config_forces_body_stream_field, enforce_provider_body_stream_policy,
request_requires_body_stream_field, resolve_upstream_is_stream_for_provider,
};
use crate::ai_serving::planner::spec_metadata::local_openai_responses_spec_metadata; use crate::ai_serving::planner::spec_metadata::local_openai_responses_spec_metadata;
use crate::ai_serving::planner::standard::{ use crate::ai_serving::planner::standard::{
apply_codex_openai_responses_special_headers, build_cross_format_openai_responses_request_body, apply_codex_openai_responses_special_headers, build_cross_format_openai_responses_request_body,
@@ -224,12 +227,15 @@ pub(crate) async fn resolve_local_openai_responses_candidate_payload_parts(
.await; .await;
let needs_bidirectional_conversion = !same_format && conversion_kind.is_some(); let needs_bidirectional_conversion = !same_format && conversion_kind.is_some();
let upstream_is_stream = spec_metadata.require_streaming let upstream_is_stream = resolve_upstream_is_stream_for_provider(
|| is_antigravity transport.endpoint.config.as_ref(),
|| force_upstream_streaming_for_provider(
transport.provider.provider_type.as_str(), transport.provider.provider_type.as_str(),
provider_api_format, provider_api_format,
spec_metadata.require_streaming,
is_antigravity || is_kiro_claude_cli,
); );
let force_body_stream_field =
endpoint_config_forces_body_stream_field(transport.endpoint.config.as_ref());
let Some(mut base_provider_request_body) = (if needs_bidirectional_conversion { let Some(mut base_provider_request_body) = (if needs_bidirectional_conversion {
build_cross_format_openai_responses_request_body( build_cross_format_openai_responses_request_body(
body_json, body_json,
@@ -237,6 +243,7 @@ pub(crate) async fn resolve_local_openai_responses_candidate_payload_parts(
spec_metadata.api_format, spec_metadata.api_format,
provider_api_format, provider_api_format,
upstream_is_stream, upstream_is_stream,
force_body_stream_field,
transport.provider.provider_type.as_str(), transport.provider.provider_type.as_str(),
if is_kiro_claude_cli { if is_kiro_claude_cli {
None None
@@ -252,6 +259,7 @@ pub(crate) async fn resolve_local_openai_responses_candidate_payload_parts(
body_json, body_json,
&mapped_model, &mapped_model,
upstream_is_stream, upstream_is_stream,
force_body_stream_field,
transport.provider.provider_type.as_str(), transport.provider.provider_type.as_str(),
provider_api_format, provider_api_format,
if is_kiro_claude_cli { if is_kiro_claude_cli {
@@ -293,6 +301,14 @@ pub(crate) async fn resolve_local_openai_responses_candidate_payload_parts(
&mut base_provider_request_body, &mut base_provider_request_body,
&mapping, &mapping,
); );
// Directive mapping is a deep-merge patch and may overwrite/add `stream`;
// re-enforce stream-field policy afterward.
enforce_provider_body_stream_policy(
&mut base_provider_request_body,
provider_api_format,
upstream_is_stream,
request_requires_body_stream_field(body_json, force_body_stream_field),
);
} }
let antigravity_auth = if is_antigravity { let antigravity_auth = if is_antigravity {
match classify_local_antigravity_request_support( match classify_local_antigravity_request_support(

View File

@@ -112,13 +112,14 @@ pub(crate) fn build_standard_stream_plan_from_decision(
.as_ref() .as_ref()
.and_then(|context| context.get("envelope_name")) .and_then(|context| context.get("envelope_name"))
.and_then(serde_json::Value::as_str); .and_then(serde_json::Value::as_str);
let accept_policy = if provider_adaptation_requires_eventstream_accept( let accept_policy = if payload.upstream_is_stream
&& provider_adaptation_requires_eventstream_accept(
envelope_name, envelope_name,
core.provider_api_format.as_str(), core.provider_api_format.as_str(),
) { ) {
StandardPlanFallbackAcceptPolicy::ProviderEventStreamIfMissing StandardPlanFallbackAcceptPolicy::ProviderEventStreamIfMissing
} else { } else {
StandardPlanFallbackAcceptPolicy::TextEventStreamRequired StandardPlanFallbackAcceptPolicy::TextEventStreamIfStreaming
}; };
let mut provider_request_headers = let mut provider_request_headers =
build_standard_plan_fallback_headers(StandardPlanFallbackHeadersInput { build_standard_plan_fallback_headers(StandardPlanFallbackHeadersInput {
@@ -132,7 +133,7 @@ pub(crate) fn build_standard_stream_plan_from_decision(
content_type: payload.content_type.as_deref(), content_type: payload.content_type.as_deref(),
provider_api_format: core.provider_api_format.as_str(), provider_api_format: core.provider_api_format.as_str(),
client_api_format: core.client_api_format.as_str(), client_api_format: core.client_api_format.as_str(),
upstream_is_stream: true, upstream_is_stream: payload.upstream_is_stream,
build_from_request_when_empty: false, build_from_request_when_empty: false,
accept_policy, accept_policy,
}); });

View File

@@ -82,7 +82,9 @@ pub use crate::formats::shared::passthrough::{
LocalSameFormatProviderSpec, LocalSameFormatProviderSpec,
}; };
pub use crate::formats::shared::request::{ pub use crate::formats::shared::request::{
endpoint_config_forces_upstream_stream_policy, enforce_request_body_stream_field,
force_upstream_streaming_for_provider, parse_direct_request_body, force_upstream_streaming_for_provider, parse_direct_request_body,
resolve_upstream_is_stream_from_endpoint_config,
}; };
pub use crate::formats::shared::request_matrix::{ pub use crate::formats::shared::request_matrix::{
build_standard_request_body_from_canonical, build_standard_request_body_from_canonical,

View File

@@ -135,10 +135,18 @@ pub fn is_openai_responses_family_format(value: &str) -> bool {
) )
} }
pub fn api_format_uses_body_stream_field(value: &str) -> bool {
matches!(
FormatId::parse(value).map(FormatId::canonical),
Some(FormatId::OpenAiChat | FormatId::OpenAiResponses | FormatId::ClaudeMessages)
)
}
#[cfg(test)] #[cfg(test)]
mod tests { mod tests {
use super::{ use super::{
api_format_alias_matches, api_format_storage_aliases, normalize_api_format_alias, FormatId, api_format_alias_matches, api_format_storage_aliases, api_format_uses_body_stream_field,
normalize_api_format_alias, FormatId,
}; };
#[test] #[test]
@@ -304,4 +312,22 @@ mod tests {
vec!["doubao:embedding".to_string()] vec!["doubao:embedding".to_string()]
); );
} }
#[test]
fn body_stream_field_support_matches_provider_wire_formats() {
assert!(api_format_uses_body_stream_field("openai:chat"));
assert!(api_format_uses_body_stream_field("/v1/chat/completions"));
assert!(api_format_uses_body_stream_field("openai:responses"));
assert!(api_format_uses_body_stream_field("/v1/responses"));
assert!(api_format_uses_body_stream_field("claude:messages"));
assert!(api_format_uses_body_stream_field("/v1/messages"));
assert!(!api_format_uses_body_stream_field(
"openai:responses:compact"
));
assert!(!api_format_uses_body_stream_field("/v1/responses/compact"));
assert!(!api_format_uses_body_stream_field(
"gemini:generate_content"
));
assert!(!api_format_uses_body_stream_field("openai:embedding"));
}
} }

View File

@@ -281,8 +281,9 @@ pub fn apply_openai_responses_compact_special_body_edits(
return; return;
}; };
// `/v1/responses/compact` does not accept `store`. // `/v1/responses/compact` does not accept `store` or body-level `stream`.
body_object.remove("store"); body_object.remove("store");
body_object.remove("stream");
} }
pub fn apply_codex_openai_responses_special_body_edits( pub fn apply_codex_openai_responses_special_body_edits(

View File

@@ -188,6 +188,9 @@ pub fn to_raw(
OPENAI_RESPONSES_LEGACY_EXTENSION_NAMESPACE, OPENAI_RESPONSES_LEGACY_EXTENSION_NAMESPACE,
&output, &output,
)); ));
if compact {
output.remove("stream");
}
output.remove("verbosity"); output.remove("verbosity");
Some(Value::Object(output)) Some(Value::Object(output))
} }

View File

@@ -1,5 +1,14 @@
use base64::Engine as _; use base64::Engine as _;
use crate::formats::id::api_format_uses_body_stream_field;
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub(crate) enum UpstreamStreamPolicy {
Auto,
ForceStream,
ForceNonStream,
}
pub fn parse_direct_request_body( pub fn parse_direct_request_body(
is_json_request: bool, is_json_request: bool,
body_bytes: &[u8], body_bytes: &[u8],
@@ -29,9 +38,130 @@ pub fn force_upstream_streaming_for_provider(
&& aether_ai_formats::is_openai_responses_format(provider_api_format) && aether_ai_formats::is_openai_responses_format(provider_api_format)
} }
pub(crate) fn parse_upstream_stream_policy(
value: Option<&serde_json::Value>,
) -> UpstreamStreamPolicy {
let Some(value) = value else {
return UpstreamStreamPolicy::Auto;
};
if let Some(value) = value.as_bool() {
return if value {
UpstreamStreamPolicy::ForceStream
} else {
UpstreamStreamPolicy::ForceNonStream
};
}
let serde_json::Value::String(value) = value else {
return UpstreamStreamPolicy::Auto;
};
let raw = value.trim().to_ascii_lowercase();
match raw.as_str() {
"" | "auto" | "follow" | "client" | "default" => UpstreamStreamPolicy::Auto,
"force_stream" | "stream" | "sse" | "true" | "1" | "yes" => {
UpstreamStreamPolicy::ForceStream
}
"force_non_stream" | "force_sync" | "non_stream" | "sync" | "false" | "0" | "no" => {
UpstreamStreamPolicy::ForceNonStream
}
_ => UpstreamStreamPolicy::Auto,
}
}
pub(crate) fn upstream_stream_policy_from_endpoint_config(
endpoint_config: Option<&serde_json::Value>,
) -> UpstreamStreamPolicy {
let Some(config) = endpoint_config.and_then(serde_json::Value::as_object) else {
return UpstreamStreamPolicy::Auto;
};
for key in [
"upstream_stream_policy",
"upstreamStreamPolicy",
"upstream_stream",
] {
if let Some(value) = config.get(key) {
return parse_upstream_stream_policy(Some(value));
}
}
UpstreamStreamPolicy::Auto
}
pub fn endpoint_config_forces_upstream_stream_policy(
endpoint_config: Option<&serde_json::Value>,
) -> bool {
matches!(
upstream_stream_policy_from_endpoint_config(endpoint_config),
UpstreamStreamPolicy::ForceStream | UpstreamStreamPolicy::ForceNonStream
)
}
/// Resolves the upstream provider transport mode.
///
/// `client_is_stream` means the request landed on a streaming surface or should
/// be treated as streaming; the original JSON body may not have had
/// `"stream": true`.
pub(crate) fn resolve_upstream_is_stream(
client_is_stream: bool,
hard_requires_streaming: bool,
policy: UpstreamStreamPolicy,
) -> bool {
// ForceStream is unconditional, while ForceNonStream yields to hard
// stream-only constraints such as Kiro or Codex OpenAI Responses.
match policy {
UpstreamStreamPolicy::ForceStream => true,
UpstreamStreamPolicy::ForceNonStream => hard_requires_streaming,
UpstreamStreamPolicy::Auto => hard_requires_streaming || client_is_stream,
}
}
pub fn enforce_request_body_stream_field(
body: &mut serde_json::Value,
provider_api_format: &str,
upstream_is_stream: bool,
require_body_stream_field: bool,
) {
let Some(body_object) = body.as_object_mut() else {
return;
};
if !api_format_uses_body_stream_field(provider_api_format) {
body_object.remove("stream");
return;
}
// Final-body fallback catches body rules, directive patches, and other
// provider-body mutations that introduce `stream`.
if upstream_is_stream || require_body_stream_field || body_object.contains_key("stream") {
body_object.insert(
"stream".to_string(),
serde_json::Value::Bool(upstream_is_stream),
);
} else {
body_object.remove("stream");
}
}
pub fn resolve_upstream_is_stream_from_endpoint_config(
endpoint_config: Option<&serde_json::Value>,
client_is_stream: bool,
hard_requires_streaming: bool,
) -> bool {
resolve_upstream_is_stream(
client_is_stream,
hard_requires_streaming,
upstream_stream_policy_from_endpoint_config(endpoint_config),
)
}
#[cfg(test)] #[cfg(test)]
mod tests { mod tests {
use super::{force_upstream_streaming_for_provider, parse_direct_request_body}; use super::{
endpoint_config_forces_upstream_stream_policy, enforce_request_body_stream_field,
force_upstream_streaming_for_provider, parse_direct_request_body,
parse_upstream_stream_policy, resolve_upstream_is_stream,
resolve_upstream_is_stream_from_endpoint_config,
upstream_stream_policy_from_endpoint_config, UpstreamStreamPolicy,
};
use serde_json::json;
#[test] #[test]
fn parses_empty_json_body_as_empty_object() { fn parses_empty_json_body_as_empty_object() {
@@ -77,4 +207,177 @@ mod tests {
"openai:responses" "openai:responses"
)); ));
} }
#[test]
fn parses_python_compatible_upstream_stream_policy_values() {
assert_eq!(
parse_upstream_stream_policy(None),
UpstreamStreamPolicy::Auto
);
for value in [
json!(""),
json!("auto"),
json!("follow"),
json!("client"),
json!("default"),
json!("unknown"),
] {
assert_eq!(
parse_upstream_stream_policy(Some(&value)),
UpstreamStreamPolicy::Auto
);
}
for value in [
json!(true),
json!("force_stream"),
json!("stream"),
json!("sse"),
json!("true"),
json!("1"),
json!("yes"),
] {
assert_eq!(
parse_upstream_stream_policy(Some(&value)),
UpstreamStreamPolicy::ForceStream
);
}
for value in [
json!(false),
json!("force_non_stream"),
json!("force_sync"),
json!("non_stream"),
json!("sync"),
json!("false"),
json!("0"),
json!("no"),
] {
assert_eq!(
parse_upstream_stream_policy(Some(&value)),
UpstreamStreamPolicy::ForceNonStream
);
}
}
#[test]
fn parses_non_string_non_bool_policy_values_as_auto() {
for value in [json!(1), json!(0), json!(null), json!({}), json!([])] {
assert_eq!(
parse_upstream_stream_policy(Some(&value)),
UpstreamStreamPolicy::Auto
);
}
}
#[test]
fn enforces_request_body_stream_field_for_stream_and_streamless_formats() {
let mut openai_chat = json!({"stream": true});
enforce_request_body_stream_field(&mut openai_chat, "openai:chat", false, false);
assert_eq!(openai_chat.get("stream"), Some(&json!(false)));
let mut ordinary_sync = json!({"messages": []});
enforce_request_body_stream_field(&mut ordinary_sync, "openai:chat", false, false);
assert!(ordinary_sync.get("stream").is_none());
let mut forced_sync = json!({"messages": []});
enforce_request_body_stream_field(&mut forced_sync, "openai:chat", false, true);
assert_eq!(forced_sync.get("stream"), Some(&json!(false)));
let mut compact = json!({"stream": true});
enforce_request_body_stream_field(&mut compact, "openai:responses:compact", true, true);
assert!(compact.get("stream").is_none());
}
#[test]
fn reads_endpoint_policy_keys_in_python_compatible_order() {
assert_eq!(
upstream_stream_policy_from_endpoint_config(Some(&json!({
"upstream_stream_policy": "force_non_stream",
"upstreamStreamPolicy": "force_stream",
"upstream_stream": "force_stream"
}))),
UpstreamStreamPolicy::ForceNonStream
);
assert_eq!(
upstream_stream_policy_from_endpoint_config(Some(&json!({
"upstreamStreamPolicy": "force_stream"
}))),
UpstreamStreamPolicy::ForceStream
);
assert_eq!(
upstream_stream_policy_from_endpoint_config(Some(&json!({
"upstream_stream": false
}))),
UpstreamStreamPolicy::ForceNonStream
);
}
#[test]
fn detects_forced_endpoint_policy_values() {
assert!(endpoint_config_forces_upstream_stream_policy(Some(
&json!({"upstream_stream_policy": "force_stream"})
)));
assert!(endpoint_config_forces_upstream_stream_policy(Some(
&json!({"upstream_stream_policy": "force_non_stream"})
)));
assert!(!endpoint_config_forces_upstream_stream_policy(Some(
&json!({"upstream_stream_policy": "auto"})
)));
assert!(!endpoint_config_forces_upstream_stream_policy(None));
}
#[test]
fn resolves_upstream_stream_policy_against_client_mode_and_hard_constraints() {
assert!(resolve_upstream_is_stream(
false,
false,
UpstreamStreamPolicy::ForceStream
));
assert!(!resolve_upstream_is_stream(
true,
false,
UpstreamStreamPolicy::ForceNonStream
));
assert!(resolve_upstream_is_stream(
true,
true,
UpstreamStreamPolicy::ForceNonStream
));
assert!(!resolve_upstream_is_stream(
false,
false,
UpstreamStreamPolicy::Auto
));
assert!(resolve_upstream_is_stream(
true,
false,
UpstreamStreamPolicy::Auto
));
assert!(resolve_upstream_is_stream(
false,
true,
UpstreamStreamPolicy::Auto
));
}
#[test]
fn resolves_endpoint_policy_config_to_upstream_mode() {
assert!(resolve_upstream_is_stream_from_endpoint_config(
Some(&json!({"upstream_stream_policy": "force_stream"})),
false,
false,
));
assert!(!resolve_upstream_is_stream_from_endpoint_config(
Some(&json!({"upstream_stream_policy": "force_non_stream"})),
true,
false,
));
assert!(resolve_upstream_is_stream_from_endpoint_config(
Some(&json!({"upstream_stream_policy": "auto"})),
true,
false,
));
assert!(!resolve_upstream_is_stream_from_endpoint_config(
None, false, false,
));
}
} }

View File

@@ -128,6 +128,18 @@ pub fn build_standard_request_body_with_model_directives_and_request_headers(
&mut provider_request_body, &mut provider_request_body,
provider_api_format, provider_api_format,
); );
let require_body_stream_field = body_json
.as_object()
.is_some_and(|object| object.contains_key("stream"))
|| provider_request_body
.as_object()
.is_some_and(|object| object.contains_key("stream"));
crate::formats::shared::request::enforce_request_body_stream_field(
&mut provider_request_body,
provider_api_format,
upstream_is_stream,
require_body_stream_field,
);
Some(provider_request_body) Some(provider_request_body)
} }
@@ -302,13 +314,23 @@ mod tests {
fn assert_stream_flag(provider_api_format: &str, upstream_is_stream: bool, converted: &Value) { fn assert_stream_flag(provider_api_format: &str, upstream_is_stream: bool, converted: &Value) {
match provider_api_format { match provider_api_format {
"openai:chat" | "openai:responses" | "claude:messages" => { "openai:chat" | "openai:responses" | "claude:messages" => {
if upstream_is_stream {
assert_eq!( assert_eq!(
converted converted.get("stream").and_then(Value::as_bool),
.get("stream") Some(true),
.and_then(Value::as_bool) "{provider_api_format} stream flag should be true for upstream streaming"
.unwrap_or(false), );
upstream_is_stream, } else {
"{provider_api_format} stream flag should follow upstream_is_stream" assert!(
converted.get("stream").is_none(),
"{provider_api_format} should not gain stream:false for ordinary sync requests"
);
}
}
"openai:responses:compact" => {
assert!(
converted.get("stream").is_none(),
"openai responses compact keeps stream out of the request body"
); );
} }
"gemini:generate_content" => { "gemini:generate_content" => {
@@ -321,6 +343,93 @@ mod tests {
} }
} }
fn assert_explicit_stream_flag(
provider_api_format: &str,
upstream_is_stream: bool,
converted: &Value,
) {
match provider_api_format {
"openai:chat" | "openai:responses" | "claude:messages" => {
assert_eq!(
converted.get("stream").and_then(Value::as_bool),
Some(upstream_is_stream),
"{provider_api_format} stream flag should follow upstream_is_stream"
);
}
"openai:responses:compact" | "gemini:generate_content" => {
assert!(converted.get("stream").is_none());
}
other => panic!("unexpected provider api format: {other}"),
}
}
#[test]
fn standard_request_body_overrides_client_stream_true_for_non_stream_upstream() {
let request = json!({
"model": "source-model",
"messages": [{"role": "user", "content": "hello"}],
"stream": true
});
for provider_api_format in [
"openai:chat",
"openai:responses",
"openai:responses:compact",
"claude:messages",
"gemini:generate_content",
] {
let converted = build_standard_request_body(
&request,
"openai:chat",
"mapped-model",
"custom",
provider_api_format,
"/v1/chat/completions",
false,
None,
None,
)
.unwrap_or_else(|| panic!("openai:chat -> {provider_api_format} should build"));
assert_explicit_stream_flag(provider_api_format, false, &converted);
}
}
#[test]
fn standard_request_body_stream_policy_wins_after_body_rules() {
let request = json!({
"model": "source-model",
"messages": [{"role": "user", "content": "hello"}],
"stream": true
});
let body_rules = json!([
{"action":"set","path":"stream","value":true}
]);
for provider_api_format in [
"openai:chat",
"openai:responses",
"openai:responses:compact",
"claude:messages",
"gemini:generate_content",
] {
let converted = build_standard_request_body(
&request,
"openai:chat",
"mapped-model",
"custom",
provider_api_format,
"/v1/chat/completions",
false,
Some(&body_rules),
None,
)
.unwrap_or_else(|| panic!("openai:chat -> {provider_api_format} should build"));
assert_explicit_stream_flag(provider_api_format, false, &converted);
}
}
fn codex_default_body_rules() -> Value { fn codex_default_body_rules() -> Value {
json!([ json!([
{"action":"drop","path":"max_output_tokens"}, {"action":"drop","path":"max_output_tokens"},

View File

@@ -50,14 +50,24 @@ pub fn build_local_openai_chat_request_body_with_model_directives(
} }
} }
} }
Some(with_model_directive_overrides( let mut provider_request_body = with_model_directive_overrides(
Value::Object(provider_request_body), Value::Object(provider_request_body),
"openai:chat", "openai:chat",
mapped_model, mapped_model,
body_json, body_json,
None, None,
enable_model_directives, enable_model_directives,
)) );
let require_body_stream_field = body_json
.as_object()
.is_some_and(|object| object.contains_key("stream"));
crate::formats::shared::request::enforce_request_body_stream_field(
&mut provider_request_body,
"openai:chat",
upstream_is_stream,
require_body_stream_field,
);
Some(provider_request_body)
} }
pub fn build_cross_format_openai_chat_request_body( pub fn build_cross_format_openai_chat_request_body(
@@ -104,14 +114,24 @@ pub fn build_cross_format_openai_chat_request_body_with_model_directives(
} }
_ => return None, _ => return None,
}; };
Some(with_model_directive_overrides( let mut provider_request_body = with_model_directive_overrides(
provider_request_body, provider_request_body,
provider_api_format, provider_api_format,
mapped_model, mapped_model,
body_json, body_json,
None, None,
enable_model_directives, enable_model_directives,
)) );
let require_body_stream_field = body_json
.as_object()
.is_some_and(|object| object.contains_key("stream"));
crate::formats::shared::request::enforce_request_body_stream_field(
&mut provider_request_body,
provider_api_format,
upstream_is_stream,
require_body_stream_field,
);
Some(provider_request_body)
} }
pub fn build_local_openai_responses_request_body( pub fn build_local_openai_responses_request_body(
@@ -143,14 +163,24 @@ pub fn build_local_openai_responses_request_body_with_model_directives(
if require_streaming { if require_streaming {
provider_request_body.insert("stream".to_string(), Value::Bool(true)); provider_request_body.insert("stream".to_string(), Value::Bool(true));
} }
Some(with_model_directive_overrides( let mut provider_request_body = with_model_directive_overrides(
Value::Object(provider_request_body), Value::Object(provider_request_body),
"openai:responses", "openai:responses",
mapped_model, mapped_model,
body_json, body_json,
None, None,
enable_model_directives, enable_model_directives,
)) );
let require_body_stream_field = body_json
.as_object()
.is_some_and(|object| object.contains_key("stream"));
crate::formats::shared::request::enforce_request_body_stream_field(
&mut provider_request_body,
"openai:responses",
require_streaming,
require_body_stream_field,
);
Some(provider_request_body)
} }
pub fn build_cross_format_openai_responses_request_body( pub fn build_cross_format_openai_responses_request_body(
@@ -208,14 +238,24 @@ pub fn build_cross_format_openai_responses_request_body_with_model_directives(
upstream_is_stream, upstream_is_stream,
)?, )?,
}; };
Some(with_model_directive_overrides( let mut provider_request_body = with_model_directive_overrides(
provider_request_body, provider_request_body,
provider_api_format, provider_api_format,
mapped_model, mapped_model,
body_json, body_json,
None, None,
enable_model_directives, enable_model_directives,
)) );
let require_body_stream_field = body_json
.as_object()
.is_some_and(|object| object.contains_key("stream"));
crate::formats::shared::request::enforce_request_body_stream_field(
&mut provider_request_body,
provider_api_format,
upstream_is_stream,
require_body_stream_field,
);
Some(provider_request_body)
} }
fn with_model_directive_overrides( fn with_model_directive_overrides(
@@ -324,6 +364,108 @@ mod tests {
); );
} }
#[test]
fn local_openai_chat_request_body_overrides_client_stream_for_non_stream_upstream() {
let body_json = json!({
"model": "gpt-5",
"messages": [{
"role": "user",
"content": "hello"
}],
"stream": true
});
let provider_request_body =
build_local_openai_chat_request_body(&body_json, "gpt-5-upstream", false)
.expect("openai chat body should build");
assert_eq!(provider_request_body["model"], "gpt-5-upstream");
assert_eq!(provider_request_body["stream"], false);
}
#[test]
fn local_openai_responses_request_body_overrides_client_stream_for_non_stream_upstream() {
let body_json = json!({
"model": "gpt-5",
"input": "hello",
"stream": true
});
let provider_request_body =
build_local_openai_responses_request_body(&body_json, "gpt-5-upstream", false)
.expect("openai responses body should build");
assert_eq!(provider_request_body["model"], "gpt-5-upstream");
assert_eq!(provider_request_body["stream"], false);
}
#[test]
fn cross_format_openai_chat_request_body_overrides_client_stream_for_non_stream_upstream() {
let body_json = json!({
"model": "gpt-5",
"messages": [{"role": "user", "content": "hello"}],
"stream": true
});
let claude = build_cross_format_openai_chat_request_body_with_model_directives(
&body_json,
"claude-sonnet-4-5",
"claude:messages",
false,
false,
)
.expect("claude body should build");
assert_eq!(claude["stream"], false);
let responses = build_cross_format_openai_chat_request_body_with_model_directives(
&body_json,
"gpt-5-upstream",
"openai:responses",
false,
false,
)
.expect("responses body should build");
assert_eq!(responses["stream"], false);
}
#[test]
fn cross_format_openai_chat_request_body_does_not_add_stream_false_for_plain_sync_body() {
let body_json = json!({
"model": "gpt-5",
"messages": [{"role": "user", "content": "hello"}]
});
let claude = build_cross_format_openai_chat_request_body_with_model_directives(
&body_json,
"claude-sonnet-4-5",
"claude:messages",
false,
false,
)
.expect("claude body should build");
assert!(claude.get("stream").is_none());
}
#[test]
fn cross_format_openai_responses_body_overrides_client_stream_for_non_stream_upstream() {
let body_json = json!({
"model": "gpt-5",
"input": "hello",
"stream": true
});
let provider_request_body = build_cross_format_openai_responses_request_body(
&body_json,
"claude-sonnet-4-5",
"openai:responses",
"claude:messages",
false,
)
.expect("claude body should build");
assert_eq!(provider_request_body["stream"], false);
}
#[test] #[test]
fn local_openai_chat_request_body_applies_reasoning_effort_suffix() { fn local_openai_chat_request_body_applies_reasoning_effort_suffix() {
let body_json = json!({ let body_json = json!({

View File

@@ -8,9 +8,9 @@ pub mod provider_compat;
pub use formats::context::{FormatContext, FormatError}; pub use formats::context::{FormatContext, FormatError};
pub use formats::id::{ pub use formats::id::{
api_format_alias_matches, api_format_storage_aliases, is_openai_responses_compact_format, api_format_alias_matches, api_format_storage_aliases, api_format_uses_body_stream_field,
is_openai_responses_family_format, is_openai_responses_format, normalize_api_format_alias, is_openai_responses_compact_format, is_openai_responses_family_format,
FormatFamily, FormatId, FormatProfile, is_openai_responses_format, normalize_api_format_alias, FormatFamily, FormatId, FormatProfile,
}; };
pub use formats::matrix::{ pub use formats::matrix::{
is_embedding_api_format, is_rerank_api_format, request_candidate_api_format_preference, is_embedding_api_format, is_rerank_api_format, request_candidate_api_format_preference,
@@ -27,6 +27,10 @@ pub use formats::shared::model_directives::{
normalize_model_directive_model, parse_model_directive, ModelDirective, ModelOverride, normalize_model_directive_model, parse_model_directive, ModelDirective, ModelOverride,
ReasoningEffort, ReasoningEffort,
}; };
pub use formats::shared::request::{
endpoint_config_forces_upstream_stream_policy, enforce_request_body_stream_field,
resolve_upstream_is_stream_from_endpoint_config,
};
pub use protocol::canonical::{ pub use protocol::canonical::{
canonical_request_unknown_block_count, canonical_response_unknown_block_count, canonical_request_unknown_block_count, canonical_response_unknown_block_count,
canonical_to_claude_request, canonical_to_claude_response, canonical_to_embedding_response, canonical_to_claude_request, canonical_to_claude_response, canonical_to_embedding_response,

View File

@@ -35,8 +35,9 @@ pub enum SameFormatProviderFamily {
} }
#[derive(Debug, Clone, Copy)] #[derive(Debug, Clone, Copy)]
pub struct SameFormatProviderRequestBehaviorParams { pub struct SameFormatProviderRequestBehaviorParams<'a> {
pub require_streaming: bool, pub require_streaming: bool,
pub provider_api_format: &'a str,
pub report_kind: &'static str, pub report_kind: &'static str,
} }
@@ -47,6 +48,7 @@ pub struct SameFormatProviderRequestBehavior {
pub is_vertex: bool, pub is_vertex: bool,
pub is_kiro: bool, pub is_kiro: bool,
pub upstream_is_stream: bool, pub upstream_is_stream: bool,
pub force_body_stream_field: bool,
pub report_kind: &'static str, pub report_kind: &'static str,
} }
@@ -61,6 +63,7 @@ pub struct SameFormatProviderRequestBodyInput<'a> {
pub body_rules: Option<&'a Value>, pub body_rules: Option<&'a Value>,
pub request_headers: Option<&'a http::HeaderMap>, pub request_headers: Option<&'a http::HeaderMap>,
pub upstream_is_stream: bool, pub upstream_is_stream: bool,
pub force_body_stream_field: bool,
pub kiro_auth_config: Option<&'a KiroAuthConfig>, pub kiro_auth_config: Option<&'a KiroAuthConfig>,
pub is_claude_code: bool, pub is_claude_code: bool,
pub enable_model_directives: bool, pub enable_model_directives: bool,
@@ -92,7 +95,7 @@ pub struct SameFormatProviderHeadersInput<'a> {
pub fn classify_same_format_provider_request_behavior( pub fn classify_same_format_provider_request_behavior(
transport: &GatewayProviderTransportSnapshot, transport: &GatewayProviderTransportSnapshot,
params: SameFormatProviderRequestBehaviorParams, params: SameFormatProviderRequestBehaviorParams<'_>,
) -> SameFormatProviderRequestBehavior { ) -> SameFormatProviderRequestBehavior {
let is_antigravity = is_antigravity_provider_transport(transport); let is_antigravity = is_antigravity_provider_transport(transport);
let is_claude_code = transport let is_claude_code = transport
@@ -102,7 +105,19 @@ pub fn classify_same_format_provider_request_behavior(
.eq_ignore_ascii_case("claude_code"); .eq_ignore_ascii_case("claude_code");
let is_vertex = is_vertex_api_key_transport_context(transport); let is_vertex = is_vertex_api_key_transport_context(transport);
let is_kiro = is_kiro_provider_transport(transport); let is_kiro = is_kiro_provider_transport(transport);
let upstream_is_stream = is_kiro || is_antigravity || params.require_streaming; let upstream_is_stream = aether_ai_formats::resolve_upstream_is_stream_from_endpoint_config(
transport.endpoint.config.as_ref(),
params.require_streaming,
is_kiro
|| is_antigravity
|| aether_ai_formats::api::force_upstream_streaming_for_provider(
transport.provider.provider_type.as_str(),
params.provider_api_format,
),
);
let force_body_stream_field = aether_ai_formats::endpoint_config_forces_upstream_stream_policy(
transport.endpoint.config.as_ref(),
);
let report_kind = if is_kiro && !params.require_streaming { let report_kind = if is_kiro && !params.require_streaming {
"claude_cli_sync_finalize" "claude_cli_sync_finalize"
} else if is_antigravity && !params.require_streaming { } else if is_antigravity && !params.require_streaming {
@@ -121,6 +136,7 @@ pub fn classify_same_format_provider_request_behavior(
is_vertex, is_vertex,
is_kiro, is_kiro,
upstream_is_stream, upstream_is_stream,
force_body_stream_field,
report_kind, report_kind,
} }
} }
@@ -165,9 +181,6 @@ pub fn build_same_format_provider_request_body(
"model".to_string(), "model".to_string(),
Value::String(input.mapped_model.to_string()), Value::String(input.mapped_model.to_string()),
); );
if input.upstream_is_stream {
provider_request_body.insert("stream".to_string(), Value::Bool(true));
}
} }
SameFormatProviderFamily::Gemini => { SameFormatProviderFamily::Gemini => {
provider_request_body.remove("model"); provider_request_body.remove("model");
@@ -195,6 +208,17 @@ pub fn build_same_format_provider_request_body(
) { ) {
return None; return None;
} }
let require_body_stream_field = input.force_body_stream_field
|| input
.body_json
.as_object()
.is_some_and(|object| object.contains_key("stream"));
aether_ai_formats::enforce_request_body_stream_field(
&mut provider_request_body,
input.provider_api_format,
input.upstream_is_stream,
require_body_stream_field,
);
Some(provider_request_body) Some(provider_request_body)
} }
@@ -334,6 +358,7 @@ pub fn same_format_provider_transport_unsupported_reason_for_trace(
transport, transport,
SameFormatProviderRequestBehaviorParams { SameFormatProviderRequestBehaviorParams {
require_streaming: false, require_streaming: false,
provider_api_format: normalized_api_format,
report_kind: "trace_candidate_metadata", report_kind: "trace_candidate_metadata",
}, },
); );
@@ -459,6 +484,7 @@ mod tests {
&kiro, &kiro,
SameFormatProviderRequestBehaviorParams { SameFormatProviderRequestBehaviorParams {
require_streaming: false, require_streaming: false,
provider_api_format: "claude:messages",
report_kind: "claude_chat_sync_success", report_kind: "claude_chat_sync_success",
}, },
); );
@@ -472,6 +498,7 @@ mod tests {
&antigravity, &antigravity,
SameFormatProviderRequestBehaviorParams { SameFormatProviderRequestBehaviorParams {
require_streaming: false, require_streaming: false,
provider_api_format: "gemini:generate_content",
report_kind: "gemini_chat_sync_success", report_kind: "gemini_chat_sync_success",
}, },
); );
@@ -481,6 +508,131 @@ mod tests {
assert_eq!(behavior.report_kind, "gemini_chat_sync_finalize"); assert_eq!(behavior.report_kind, "gemini_chat_sync_finalize");
} }
#[test]
fn same_format_behavior_resolves_endpoint_stream_policy() {
let mut force_stream = sample_transport("openai");
force_stream.endpoint.config = Some(json!({
"upstream_stream_policy": "force_stream"
}));
let behavior = classify_same_format_provider_request_behavior(
&force_stream,
SameFormatProviderRequestBehaviorParams {
require_streaming: false,
provider_api_format: "openai:chat",
report_kind: "openai_chat_sync_success",
},
);
assert!(behavior.upstream_is_stream);
let mut force_non_stream = sample_transport("openai");
force_non_stream.endpoint.config = Some(json!({
"upstreamStreamPolicy": "force_non_stream"
}));
let behavior = classify_same_format_provider_request_behavior(
&force_non_stream,
SameFormatProviderRequestBehaviorParams {
require_streaming: true,
provider_api_format: "openai:chat",
report_kind: "openai_chat_stream_success",
},
);
assert!(!behavior.upstream_is_stream);
let mut auto = sample_transport("openai");
auto.endpoint.config = Some(json!({
"upstream_stream": "auto"
}));
let stream_behavior = classify_same_format_provider_request_behavior(
&auto,
SameFormatProviderRequestBehaviorParams {
require_streaming: true,
provider_api_format: "openai:chat",
report_kind: "openai_chat_stream_success",
},
);
assert!(stream_behavior.upstream_is_stream);
let sync_behavior = classify_same_format_provider_request_behavior(
&auto,
SameFormatProviderRequestBehaviorParams {
require_streaming: false,
provider_api_format: "openai:chat",
report_kind: "openai_chat_sync_success",
},
);
assert!(!sync_behavior.upstream_is_stream);
}
#[test]
fn same_format_behavior_preserves_hard_streaming_constraint() {
let mut kiro = sample_transport("kiro");
kiro.endpoint.config = Some(json!({
"upstream_stream_policy": "force_non_stream"
}));
let behavior = classify_same_format_provider_request_behavior(
&kiro,
SameFormatProviderRequestBehaviorParams {
require_streaming: true,
provider_api_format: "claude:messages",
report_kind: "claude_chat_stream_success",
},
);
assert!(behavior.upstream_is_stream);
}
#[test]
fn same_format_policy_resolution_drives_standard_body_stream_field() {
for (endpoint_config, client_is_stream, expected_stream) in [
(
json!({"upstream_stream_policy": "force_stream"}),
false,
true,
),
(
json!({"upstreamStreamPolicy": "force_non_stream"}),
true,
false,
),
(json!({"upstream_stream": "auto"}), true, true),
(json!({"upstream_stream": "auto"}), false, false),
] {
let mut transport = sample_transport("openai");
transport.endpoint.config = Some(endpoint_config);
let behavior = classify_same_format_provider_request_behavior(
&transport,
SameFormatProviderRequestBehaviorParams {
require_streaming: client_is_stream,
provider_api_format: "openai:chat",
report_kind: "openai_chat_policy_test",
},
);
let body =
build_same_format_provider_request_body(SameFormatProviderRequestBodyInput {
body_json: &json!({
"model": "client-model",
"messages": [{"role": "user", "content": "hello"}],
"stream": client_is_stream
}),
mapped_model: "upstream-model",
client_api_format: "openai:chat",
provider_api_format: "openai:chat",
source_model: Some("client-model"),
family: SameFormatProviderFamily::Standard,
body_rules: None,
request_headers: None,
upstream_is_stream: behavior.upstream_is_stream,
force_body_stream_field: behavior.force_body_stream_field,
kiro_auth_config: None,
is_claude_code: false,
enable_model_directives: false,
})
.expect("body should build");
assert_eq!(body.get("stream"), Some(&json!(expected_stream)));
}
}
#[test] #[test]
fn resolves_direct_auth_except_vertex() { fn resolves_direct_auth_except_vertex() {
let transport = sample_transport("openai"); let transport = sample_transport("openai");
@@ -488,6 +640,7 @@ mod tests {
&transport, &transport,
SameFormatProviderRequestBehaviorParams { SameFormatProviderRequestBehaviorParams {
require_streaming: false, require_streaming: false,
provider_api_format: "openai:chat",
report_kind: "openai_chat_sync_success", report_kind: "openai_chat_sync_success",
}, },
); );
@@ -517,6 +670,7 @@ mod tests {
body_rules: None, body_rules: None,
request_headers: None, request_headers: None,
upstream_is_stream: true, upstream_is_stream: true,
force_body_stream_field: false,
kiro_auth_config: None, kiro_auth_config: None,
is_claude_code: false, is_claude_code: false,
enable_model_directives: false, enable_model_directives: false,
@@ -527,6 +681,170 @@ mod tests {
assert_eq!(body.get("stream"), Some(&json!(true))); assert_eq!(body.get("stream"), Some(&json!(true)));
} }
#[test]
fn same_format_standard_body_overrides_client_stream_for_non_stream_upstream() {
let body = build_same_format_provider_request_body(SameFormatProviderRequestBodyInput {
body_json: &json!({
"model": "client-model",
"messages": [{"role": "user", "content": "hello"}],
"stream": true
}),
mapped_model: "upstream-model",
client_api_format: "openai:chat",
provider_api_format: "openai:chat",
source_model: Some("client-model"),
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,
})
.expect("body should build");
assert_eq!(body.get("model"), Some(&json!("upstream-model")));
assert_eq!(body.get("stream"), Some(&json!(false)));
}
#[test]
fn same_format_standard_body_does_not_add_stream_false_for_plain_sync_body() {
let body = build_same_format_provider_request_body(SameFormatProviderRequestBodyInput {
body_json: &json!({
"model": "client-model",
"messages": [{"role": "user", "content": "hello"}]
}),
mapped_model: "upstream-model",
client_api_format: "openai:chat",
provider_api_format: "openai:chat",
source_model: Some("client-model"),
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,
})
.expect("body should build");
assert_eq!(body.get("model"), Some(&json!("upstream-model")));
assert!(body.get("stream").is_none());
}
#[test]
fn same_format_standard_body_forced_policy_adds_stream_false() {
let body = build_same_format_provider_request_body(SameFormatProviderRequestBodyInput {
body_json: &json!({
"model": "client-model",
"messages": [{"role": "user", "content": "hello"}]
}),
mapped_model: "upstream-model",
client_api_format: "openai:chat",
provider_api_format: "openai:chat",
source_model: Some("client-model"),
family: SameFormatProviderFamily::Standard,
body_rules: None,
request_headers: None,
upstream_is_stream: false,
force_body_stream_field: true,
kiro_auth_config: None,
is_claude_code: false,
enable_model_directives: false,
})
.expect("body should build");
assert_eq!(body.get("model"), Some(&json!("upstream-model")));
assert_eq!(body.get("stream"), Some(&json!(false)));
}
#[test]
fn same_format_gemini_body_removes_leaked_client_stream_field() {
let body = build_same_format_provider_request_body(SameFormatProviderRequestBodyInput {
body_json: &json!({
"contents": [{"role": "user", "parts": [{"text": "hello"}]}],
"stream": true
}),
mapped_model: "gemini-upstream",
client_api_format: "gemini:generate_content",
provider_api_format: "gemini:generate_content",
source_model: None,
family: SameFormatProviderFamily::Gemini,
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,
})
.expect("body should build");
assert!(body.get("stream").is_none());
}
#[test]
fn same_format_stream_policy_wins_after_body_rules() {
let body_rules = json!([
{"action":"set","path":"stream","value":true}
]);
let body = build_same_format_provider_request_body(SameFormatProviderRequestBodyInput {
body_json: &json!({
"model": "client-model",
"messages": [{"role": "user", "content": "hello"}],
"stream": true
}),
mapped_model: "upstream-model",
client_api_format: "openai:chat",
provider_api_format: "openai:chat",
source_model: Some("client-model"),
family: SameFormatProviderFamily::Standard,
body_rules: Some(&body_rules),
request_headers: None,
upstream_is_stream: false,
force_body_stream_field: false,
kiro_auth_config: None,
is_claude_code: false,
enable_model_directives: false,
})
.expect("body should build");
assert_eq!(body.get("stream"), Some(&json!(false)));
}
#[test]
fn same_format_compact_stream_policy_wins_after_body_rules() {
let body_rules = json!([
{"action":"set","path":"stream","value":true}
]);
let body = build_same_format_provider_request_body(SameFormatProviderRequestBodyInput {
body_json: &json!({
"model": "client-model",
"input": "hello",
"stream": true
}),
mapped_model: "upstream-model",
client_api_format: "openai:responses:compact",
provider_api_format: "openai:responses:compact",
source_model: Some("client-model"),
family: SameFormatProviderFamily::Standard,
body_rules: Some(&body_rules),
request_headers: None,
upstream_is_stream: false,
force_body_stream_field: false,
kiro_auth_config: None,
is_claude_code: false,
enable_model_directives: false,
})
.expect("body should build");
assert!(body.get("stream").is_none());
}
#[test] #[test]
fn same_format_body_applies_model_directive_before_body_rules() { fn same_format_body_applies_model_directive_before_body_rules() {
let body = build_same_format_provider_request_body(SameFormatProviderRequestBodyInput { let body = build_same_format_provider_request_body(SameFormatProviderRequestBodyInput {
@@ -545,6 +863,7 @@ mod tests {
])), ])),
request_headers: None, request_headers: None,
upstream_is_stream: false, upstream_is_stream: false,
force_body_stream_field: false,
kiro_auth_config: None, kiro_auth_config: None,
is_claude_code: false, is_claude_code: false,
enable_model_directives: true, enable_model_directives: true,
@@ -571,6 +890,7 @@ mod tests {
is_vertex: false, is_vertex: false,
is_kiro: false, is_kiro: false,
upstream_is_stream: true, upstream_is_stream: true,
force_body_stream_field: false,
report_kind: "openai_chat_stream_success", report_kind: "openai_chat_stream_success",
}, },
auth_header: Some("x-api-key"), auth_header: Some("x-api-key"),