i/Aether
1
0
Fork 0
mirror of https://github.com/fawney19/Aether.git synced 2026-09-02 01:10:23 +08:00
Code Issues Packages Projects Releases Wiki Activity

fix: enforce provider upstream stream policy

Browse Source
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
Download Patch File Download Diff File Expand all files Collapse all files

111
apps/aether-gateway/src/ai_serving/planner/common.rs
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)
}
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> {
aether_ai_serving::extract_ai_standard_requested_model(body_json)
}
@@ -56,8 +102,10 @@ pub(crate) fn extract_requested_model_from_request(
#[cfg(test)]
mod tests {
use super::{
endpoint_config_forces_body_stream_field, enforce_provider_body_stream_policy,
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 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]
fn extracts_standard_requested_model_from_request_body() {
let requested_model =

17
apps/aether-gateway/src/ai_serving/planner/passthrough/provider/family/request.rs
View File

@@ -3,6 +3,9 @@ use std::sync::Arc;
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::{
build_antigravity_safe_v1internal_request, build_antigravity_static_identity_headers,
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(
transport,
provider_api_format,
crate::ai_serving::planner::spec_metadata::LocalExecutionSurfaceSpecMetadata {
api_format: provider_api_format,
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(),
Some(&parts.headers),
prepared.upstream_is_stream,
prepared.force_body_stream_field,
prepared.kiro_auth.as_ref(),
prepared.is_claude_code,
enable_model_directives,
@@ -170,6 +175,18 @@ pub(crate) async fn resolve_local_same_format_provider_candidate_payload_parts(
&mut base_provider_request_body,
&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 {

2
apps/aether-gateway/src/ai_serving/planner/passthrough/provider/family/request/policy.rs
View File

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

8
apps/aether-gateway/src/ai_serving/planner/passthrough/provider/family/request/prepare.rs
View File

@@ -35,6 +35,7 @@ pub(super) struct PreparedSameFormatProviderCandidate {
pub(super) mapped_model: String,
pub(super) report_kind: &'static str,
pub(super) upstream_is_stream: bool,
pub(super) force_body_stream_field: bool,
}
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 transport = Arc::clone(&eligible.transport);
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(
&behavior,
@@ -174,5 +179,6 @@ pub(super) async fn prepare_local_same_format_provider_candidate(
mapped_model,
report_kind: behavior.report_kind,
upstream_is_stream: behavior.upstream_is_stream,
force_body_stream_field: behavior.force_body_stream_field,
})
}

2
apps/aether-gateway/src/ai_serving/planner/passthrough/provider/request/body.rs
View File

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

33
apps/aether-gateway/src/ai_serving/planner/standard/family/request.rs
View File

@@ -7,7 +7,10 @@ use crate::ai_serving::planner::candidate_preparation::{
prepare_header_authenticated_candidate, prepare_header_authenticated_candidate_from_auth,
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::standard::{
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
|| force_upstream_streaming_for_provider(
transport.provider.provider_type.as_str(),
provider_api_format,
);
let upstream_is_stream = resolve_upstream_is_stream_for_provider(
transport.endpoint.config.as_ref(),
transport.provider.provider_type.as_str(),
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 =
crate::system_features::reasoning_model_directive_enabled_for_api_format_and_model(
state,
@@ -225,6 +232,12 @@ pub(crate) async fn resolve_local_standard_candidate_payload_parts(
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) =
crate::system_features::reasoning_model_directive_mapping_for_api_format_and_model(
state,
@@ -237,6 +250,14 @@ pub(crate) async fn resolve_local_standard_candidate_payload_parts(
&mut provider_request_body,
&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() {

4
apps/aether-gateway/src/ai_serving/planner/standard/gemini/plan_builders.rs
View File

@@ -116,9 +116,9 @@ pub(crate) fn build_gemini_stream_plan_from_decision(
content_type: payload.content_type.as_deref(),
provider_api_format: core.provider_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,
accept_policy: StandardPlanFallbackAcceptPolicy::TextEventStreamRequired,
accept_policy: StandardPlanFallbackAcceptPolicy::TextEventStreamIfStreaming,
});
let content_type = payload
.content_type

3
apps/aether-gateway/src/ai_serving/planner/standard/normalize.rs
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_local_openai_responses_upstream_url,
};
pub(super) use crate::ai_serving::planner::common::{
enforce_provider_body_stream_policy, request_requires_body_stream_field,
};

30
apps/aether-gateway/src/ai_serving/planner/standard/normalize/chat.rs
View File

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

16
apps/aether-gateway/src/ai_serving/planner/standard/normalize/responses.rs
View File

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

9
apps/aether-gateway/src/ai_serving/planner/standard/normalize/tests.rs
View File

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

4
apps/aether-gateway/src/ai_serving/planner/standard/openai/chat/decision/payload.rs
View File

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

26
apps/aether-gateway/src/ai_serving/planner/standard/openai/chat/decision/request.rs
View File

@@ -8,7 +8,10 @@ use crate::ai_serving::planner::candidate_preparation::{
OauthPreparationContext,
};
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::{
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,
@@ -72,6 +75,8 @@ pub(crate) async fn resolve_local_openai_chat_candidate_payload_parts(
let candidate = &eligible.candidate;
let provider_api_format = eligible.provider_api_format.as_str();
let transport = &eligible.transport;
let force_body_stream_field =
endpoint_config_forces_body_stream_field(transport.endpoint.config.as_ref());
let enable_model_directives =
crate::system_features::reasoning_model_directive_enabled_for_api_format_and_model(
state,
@@ -128,6 +133,7 @@ pub(crate) async fn resolve_local_openai_chat_candidate_payload_parts(
body_json,
&prepared_candidate.mapped_model,
upstream_is_stream,
force_body_stream_field,
transport.endpoint.body_rules.as_ref(),
&parts.headers,
enable_model_directives,
@@ -214,6 +220,13 @@ pub(crate) async fn resolve_local_openai_chat_candidate_payload_parts(
let (execution_strategy, conversion_mode) =
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 {
auth_header: resolved_headers.auth_header,
auth_value: resolved_headers.auth_value,
@@ -224,7 +237,7 @@ pub(crate) async fn resolve_local_openai_chat_candidate_payload_parts(
upstream_url,
execution_strategy,
conversion_mode,
report_kind: report_kind.to_string(),
report_kind: resolved_report_kind,
envelope_name: None,
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(),
provider_api_format.as_str(),
upstream_is_stream,
force_body_stream_field,
if is_kiro_claude_cli {
None
} else {
@@ -387,6 +401,14 @@ pub(crate) async fn resolve_local_openai_chat_candidate_payload_parts(
&mut provider_request_body,
&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() {

147
apps/aether-gateway/src/ai_serving/planner/standard/openai/chat/plans.rs
View File

@@ -9,6 +9,9 @@ mod stream;
#[path = "plans/sync.rs"]
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::diagnostic::set_local_openai_chat_miss_diagnostic;
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::{
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,
));
}
}

15
apps/aether-gateway/src/ai_serving/planner/standard/openai/chat/plans/stream.rs
View File

@@ -13,6 +13,7 @@ use super::candidates::list_local_openai_chat_candidates;
use super::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 crate::ai_serving::planner::candidate_materialization::LocalExecutionAttemptSource;
use crate::ai_serving::planner::common::OPENAI_CHAT_STREAM_PLAN_KIND;
@@ -120,6 +121,11 @@ impl LocalOpenAiChatStreamAttemptSource<'_> {
&self,
attempt: LocalOpenAiChatCandidateAttempt,
) -> 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(
self.state,
self.parts,
@@ -129,7 +135,7 @@ impl LocalOpenAiChatStreamAttemptSource<'_> {
attempt,
OPENAI_CHAT_STREAM_PLAN_KIND,
"openai_chat_stream_success",
true,
upstream_is_stream,
)
.await
else {
@@ -222,6 +228,11 @@ pub(crate) async fn build_local_openai_chat_stream_plan_and_reports(
let mut plans = Vec::new();
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(
state,
parts,
@@ -231,7 +242,7 @@ pub(crate) async fn build_local_openai_chat_stream_plan_and_reports(
attempt,
OPENAI_CHAT_STREAM_PLAN_KIND,
"openai_chat_stream_success",
true,
upstream_is_stream,
)
.await
else {

47
apps/aether-gateway/src/ai_serving/planner/standard/openai/chat/plans/sync.rs
View File

@@ -13,11 +13,10 @@ use super::candidates::list_local_openai_chat_candidates;
use super::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 crate::ai_serving::planner::candidate_materialization::LocalExecutionAttemptSource;
use crate::ai_serving::planner::common::{
force_upstream_streaming_for_provider, OPENAI_CHAT_SYNC_PLAN_KIND,
};
use crate::ai_serving::planner::common::OPENAI_CHAT_SYNC_PLAN_KIND;
use crate::ai_serving::planner::plan_builders::{
build_openai_chat_sync_plan_from_decision, AiSyncAttempt,
};
@@ -32,13 +31,6 @@ pub(crate) struct LocalOpenAiChatSyncAttemptSource<'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>(
state: &'a AppState,
parts: &'a http::request::Parts,
@@ -129,9 +121,10 @@ impl LocalOpenAiChatSyncAttemptSource<'_> {
&self,
attempt: LocalOpenAiChatCandidateAttempt,
) -> Result<Option<AiSyncAttempt>, GatewayError> {
let upstream_is_stream = openai_chat_sync_upstream_is_stream_for_candidate(
attempt.eligible.transport.provider.provider_type.as_str(),
let upstream_is_stream = openai_chat_upstream_is_stream_for_candidate(
&attempt.eligible.transport,
attempt.eligible.provider_api_format.as_str(),
false,
);
let Some(payload) = maybe_build_local_openai_chat_decision_payload_for_candidate(
self.state,
@@ -235,9 +228,10 @@ pub(crate) async fn build_local_openai_chat_sync_plan_and_reports(
let mut plans = Vec::new();
for attempt in attempts {
let upstream_is_stream = openai_chat_sync_upstream_is_stream_for_candidate(
attempt.eligible.transport.provider.provider_type.as_str(),
let upstream_is_stream = openai_chat_upstream_is_stream_for_candidate(
&attempt.eligible.transport,
attempt.eligible.provider_api_format.as_str(),
false,
);
let Some(payload) = maybe_build_local_openai_chat_decision_payload_for_candidate(
state,
@@ -272,28 +266,3 @@ pub(crate) async fn build_local_openai_chat_sync_plan_and_reports(
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"
));
}
}

134
apps/aether-gateway/src/ai_serving/planner/standard/openai/plan_builders/stream.rs
View File

@@ -9,6 +9,7 @@ use super::super::{
take_ai_upstream_auth_pair, take_non_empty_string, AiExecutionPlanFromDecisionParts,
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::transport::{
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
.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) {
let existing = provider_request_body
let existing = provider_request_object
.get("prompt_cache_key")
.and_then(|value| value.as_str())
.map(str::trim)
.unwrap_or_default();
if existing.is_empty() {
provider_request_body.insert(
provider_request_object.insert(
"prompt_cache_key".to_string(),
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 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(),
provider_api_format: core.provider_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,
accept_policy: StandardPlanFallbackAcceptPolicy::TextEventStreamRequired,
accept_policy: StandardPlanFallbackAcceptPolicy::TextEventStreamIfStreaming,
});
let content_type = payload
.content_type
@@ -157,13 +168,14 @@ pub(crate) fn build_openai_responses_stream_plan_from_decision(
.as_ref()
.and_then(|context| context.get("envelope_name"))
.and_then(serde_json::Value::as_str);
let accept_policy = if provider_adaptation_requires_eventstream_accept(
envelope_name,
core.provider_api_format.as_str(),
) {
let accept_policy = if payload.upstream_is_stream
&& provider_adaptation_requires_eventstream_accept(
envelope_name,
core.provider_api_format.as_str(),
) {
StandardPlanFallbackAcceptPolicy::ProviderEventStreamIfMissing
} else {
StandardPlanFallbackAcceptPolicy::TextEventStreamRequired
StandardPlanFallbackAcceptPolicy::TextEventStreamIfStreaming
};
let mut provider_request_headers =
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(),
provider_api_format: core.provider_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,
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]
fn build_openai_chat_stream_plan_fallback_restores_claude_headers_for_cross_format() {
let parts = http::Request::builder()

21
apps/aether-gateway/src/ai_serving/planner/standard/openai/plan_builders/sync.rs
View File

@@ -9,6 +9,7 @@ use super::super::{
take_ai_upstream_auth_pair, take_non_empty_string, AiExecutionPlanFromDecisionParts,
AiSyncAttempt,
};
use crate::ai_serving::planner::common::enforce_provider_body_stream_policy;
use crate::ai_serving::transport::{
build_standard_plan_fallback_headers, build_standard_plan_fallback_openai_chat_url,
build_standard_plan_fallback_openai_responses_url, StandardPlanFallbackAcceptPolicy,
@@ -51,23 +52,31 @@ pub(crate) fn build_openai_chat_sync_plan_from_decision(
provider_request_body
.insert("model".to_string(), serde_json::Value::String(mapped_model));
}
if payload.upstream_is_stream {
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) {
let existing = provider_request_body
let existing = provider_request_object
.get("prompt_cache_key")
.and_then(|value| value.as_str())
.map(str::trim)
.unwrap_or_default();
if existing.is_empty() {
provider_request_body.insert(
provider_request_object.insert(
"prompt_cache_key".to_string(),
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 mut provider_request_headers =

30
apps/aether-gateway/src/ai_serving/planner/standard/openai/responses/decision/request.rs
View File

@@ -9,7 +9,10 @@ use crate::ai_serving::planner::candidate_preparation::{
OauthPreparationContext,
};
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::standard::{
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;
let needs_bidirectional_conversion = !same_format && conversion_kind.is_some();
let upstream_is_stream = spec_metadata.require_streaming
|| is_antigravity
|| force_upstream_streaming_for_provider(
transport.provider.provider_type.as_str(),
provider_api_format,
);
let upstream_is_stream = resolve_upstream_is_stream_for_provider(
transport.endpoint.config.as_ref(),
transport.provider.provider_type.as_str(),
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 {
build_cross_format_openai_responses_request_body(
body_json,
@@ -237,6 +243,7 @@ pub(crate) async fn resolve_local_openai_responses_candidate_payload_parts(
spec_metadata.api_format,
provider_api_format,
upstream_is_stream,
force_body_stream_field,
transport.provider.provider_type.as_str(),
if is_kiro_claude_cli {
None
@@ -252,6 +259,7 @@ pub(crate) async fn resolve_local_openai_responses_candidate_payload_parts(
body_json,
&mapped_model,
upstream_is_stream,
force_body_stream_field,
transport.provider.provider_type.as_str(),
provider_api_format,
if is_kiro_claude_cli {
@@ -293,6 +301,14 @@ pub(crate) async fn resolve_local_openai_responses_candidate_payload_parts(
&mut base_provider_request_body,
&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 {
match classify_local_antigravity_request_support(

13
apps/aether-gateway/src/ai_serving/planner/standard/plan_builders.rs
View File

@@ -112,13 +112,14 @@ pub(crate) fn build_standard_stream_plan_from_decision(
.as_ref()
.and_then(|context| context.get("envelope_name"))
.and_then(serde_json::Value::as_str);
let accept_policy = if provider_adaptation_requires_eventstream_accept(
envelope_name,
core.provider_api_format.as_str(),
) {
let accept_policy = if payload.upstream_is_stream
&& provider_adaptation_requires_eventstream_accept(
envelope_name,
core.provider_api_format.as_str(),
) {
StandardPlanFallbackAcceptPolicy::ProviderEventStreamIfMissing
} else {
StandardPlanFallbackAcceptPolicy::TextEventStreamRequired
StandardPlanFallbackAcceptPolicy::TextEventStreamIfStreaming
};
let mut provider_request_headers =
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(),
provider_api_format: core.provider_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,
accept_policy,
});