Fix Codex image progress heartbeat merge regressions

This commit is contained in:
fawney19
2026-05-10 02:10:23 +08:00
156 changed files with 11741 additions and 943 deletions

View File

@@ -1,8 +1,11 @@
use axum::body::Bytes;
use crate::ai_serving::{
endpoint_config_forces_upstream_stream_policy as endpoint_config_forces_upstream_stream_policy_impl,
enforce_request_body_stream_field as enforce_request_body_stream_field_impl,
force_upstream_streaming_for_provider as force_upstream_streaming_for_provider_impl,
is_json_request, parse_direct_request_body as parse_direct_request_body_impl,
resolve_upstream_is_stream_from_endpoint_config as resolve_upstream_is_stream_from_endpoint_config_impl,
};
pub(crate) use crate::ai_serving::{
CLAUDE_CHAT_STREAM_PLAN_KIND, CLAUDE_CHAT_SYNC_PLAN_KIND, CLAUDE_CLI_STREAM_PLAN_KIND,
@@ -37,6 +40,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);
resolve_upstream_is_stream_from_endpoint_config_impl(
endpoint_config,
client_is_stream,
hard_requires_streaming,
)
}
pub(crate) fn endpoint_config_forces_body_stream_field(
endpoint_config: Option<&serde_json::Value>,
) -> bool {
endpoint_config_forces_upstream_stream_policy_impl(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,
) {
enforce_request_body_stream_field_impl(
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 +105,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 +141,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 =

View File

@@ -112,6 +112,8 @@ pub(crate) async fn maybe_build_local_same_format_provider_decision_payload_for_
provider_request_method: Some(serde_json::Value::Null),
provider_request_headers: Some(&resolved.provider_request_headers),
original_headers: &parts.headers,
request_path: Some(parts.uri.path()),
request_query_string: parts.uri.query(),
request_origin: Some(crate::ai_serving::request_origin_from_parts(parts)),
original_request_body_json: Some(body_json),
original_request_body_base64: None,

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 {

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"),

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,
})
}

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,

View File

@@ -9,7 +9,11 @@ use aether_ai_serving::{
use aether_scheduler_core::{ClientSessionAffinity, SchedulerRankingOutcome};
use serde_json::{Map, Value};
use crate::ai_serving::{request_origin_from_headers, ExecutionRuntimeAuthContext, RequestOrigin};
use crate::ai_serving::{
request_origin_from_headers, request_path_implies_stream_request, sanitize_request_path,
sanitize_request_path_and_query, sanitize_request_query_string, ExecutionRuntimeAuthContext,
RequestOrigin,
};
use crate::client_session_affinity::{
client_session_affinity_report_context_value, CLIENT_SESSION_AFFINITY_REPORT_CONTEXT_FIELD,
};
@@ -40,6 +44,8 @@ pub(crate) struct LocalExecutionReportContextParts<'a> {
pub(crate) provider_request_method: Option<Value>,
pub(crate) provider_request_headers: Option<&'a BTreeMap<String, String>>,
pub(crate) original_headers: &'a http::HeaderMap,
pub(crate) request_path: Option<&'a str>,
pub(crate) request_query_string: Option<&'a str>,
pub(crate) request_origin: Option<RequestOrigin>,
pub(crate) original_request_body_json: Option<&'a Value>,
pub(crate) original_request_body_base64: Option<&'a str>,
@@ -80,6 +86,15 @@ pub(crate) fn build_local_execution_report_context(
{
merge_incoming_tls_fingerprint(&mut extra_fields, incoming_tls);
}
insert_request_path_fields(
&mut extra_fields,
parts.request_path,
parts.request_query_string,
);
let client_requested_stream = parts.client_requested_stream
|| parts
.request_path
.is_some_and(request_path_implies_stream_request);
build_ai_execution_report_context(AiExecutionReportContextParts {
auth_context: parts.auth_context,
@@ -113,7 +128,7 @@ pub(crate) fn build_local_execution_report_context(
client_ip,
user_agent,
},
client_requested_stream: parts.client_requested_stream,
client_requested_stream,
upstream_is_stream: parts.upstream_is_stream,
has_envelope: parts.has_envelope,
needs_conversion: parts.needs_conversion,
@@ -121,6 +136,30 @@ pub(crate) fn build_local_execution_report_context(
})
}
fn insert_request_path_fields(
extra_fields: &mut Map<String, Value>,
request_path: Option<&str>,
request_query_string: Option<&str>,
) {
let Some(path) = request_path.and_then(sanitize_request_path) else {
return;
};
let query = request_query_string.and_then(sanitize_request_query_string);
let path_and_query = sanitize_request_path_and_query(path.as_str(), query.as_deref())
.unwrap_or_else(|| path.clone());
extra_fields
.entry("request_path".to_string())
.or_insert_with(|| Value::String(path.clone()));
if let Some(query) = query.clone() {
extra_fields
.entry("request_query_string".to_string())
.or_insert_with(|| Value::String(query.to_string()));
}
extra_fields
.entry("request_path_and_query".to_string())
.or_insert_with(|| Value::String(path_and_query));
}
pub(crate) fn provider_stream_event_api_format_for_provider_type(
provider_type: &str,
) -> Option<&'static str> {
@@ -226,6 +265,8 @@ mod tests {
provider_request_method: None,
provider_request_headers: Some(&provider_request_headers),
original_headers: &original_headers,
request_path: Some("/v1/chat/completions"),
request_query_string: Some("debug=true&limit=10"),
request_origin: Some(RequestOrigin {
client_ip: Some("203.0.113.8".to_string()),
user_agent: Some("Claude-Code/1.0".to_string()),
@@ -255,6 +296,77 @@ mod tests {
"session_key": "account=account-1;session=session-1"
})
);
assert_eq!(report_context["request_path"], "/v1/chat/completions");
assert_eq!(report_context["request_query_string"], "limit=10");
assert_eq!(
report_context["request_path_and_query"],
"/v1/chat/completions?limit=10"
);
}
#[test]
fn local_execution_report_context_treats_stream_generate_content_path_as_client_stream() {
let auth_context = ExecutionRuntimeAuthContext {
user_id: "user-1".to_string(),
api_key_id: "api-key-1".to_string(),
username: None,
api_key_name: None,
balance_remaining: None,
access_allowed: true,
api_key_is_standalone: false,
};
let original_headers = http::HeaderMap::new();
let provider_request_headers = BTreeMap::new();
let report_context =
build_local_execution_report_context(LocalExecutionReportContextParts {
auth_context: &auth_context,
request_id: "trace-1",
candidate_id: "candidate-1",
attempt_identity: ExecutionAttemptIdentity::new(0, 0),
model: "gemini-3.1-flash-image-preview",
provider_name: "Gemini",
provider_id: "provider-1",
endpoint_id: "endpoint-1",
key_id: "key-1",
key_name: None,
model_id: None,
global_model_id: None,
global_model_name: None,
provider_api_format: "gemini:generate_content",
client_api_format: "gemini:generate_content",
mapped_model: None,
candidate_group_id: None,
ranking: None,
upstream_url: None,
header_rules: None,
body_rules: None,
provider_request_method: None,
provider_request_headers: Some(&provider_request_headers),
original_headers: &original_headers,
request_path: Some(
"/v1beta/models/gemini-3.1-flash-image-preview:streamGenerateContent",
),
request_query_string: Some("key=secret&alt=sse"),
request_origin: None,
original_request_body_json: Some(&json!({
"contents": [{"role": "user", "parts": [{"text": "hi"}]}]
})),
original_request_body_base64: None,
client_session_affinity: None,
client_requested_stream: false,
upstream_is_stream: true,
has_envelope: false,
needs_conversion: false,
extra_fields: Map::new(),
});
assert_eq!(report_context["client_requested_stream"], true);
assert_eq!(report_context["request_query_string"], "alt=sse");
assert_eq!(
report_context["request_path_and_query"],
"/v1beta/models/gemini-3.1-flash-image-preview:streamGenerateContent?alt=sse"
);
}
#[test]
@@ -300,6 +412,8 @@ mod tests {
provider_request_method: None,
provider_request_headers: Some(&provider_request_headers),
original_headers: &original_headers,
request_path: None,
request_query_string: None,
request_origin: None,
original_request_body_json: Some(&json!({"model": "gpt-5"})),
original_request_body_base64: None,

View File

@@ -94,6 +94,8 @@ pub(super) async fn maybe_build_local_gemini_files_decision_payload_for_candidat
provider_request_method: None,
provider_request_headers: None,
original_headers: &parts.headers,
request_path: Some(parts.uri.path()),
request_query_string: parts.uri.query(),
request_origin: Some(crate::ai_serving::request_origin_from_parts(parts)),
original_request_body_json: Some(body_json),
original_request_body_base64: resolved.provider_request_body_base64.as_deref(),

View File

@@ -114,6 +114,8 @@ pub(super) async fn maybe_build_local_openai_image_decision_payload_for_candidat
provider_request_method: Some(serde_json::Value::String(parts.method.to_string())),
provider_request_headers: Some(&resolved.provider_request_headers),
original_headers: &parts.headers,
request_path: Some(parts.uri.path()),
request_query_string: parts.uri.query(),
request_origin: Some(crate::ai_serving::request_origin_from_parts(parts)),
original_request_body_json: Some(body_json),
original_request_body_base64: body_base64,

View File

@@ -75,6 +75,8 @@ pub(super) async fn maybe_build_local_video_create_decision_payload_for_candidat
provider_request_method: None,
provider_request_headers: None,
original_headers: &parts.headers,
request_path: Some(parts.uri.path()),
request_query_string: parts.uri.query(),
request_origin: Some(crate::ai_serving::request_origin_from_parts(parts)),
original_request_body_json: Some(body_json),
original_request_body_base64: None,

View File

@@ -120,6 +120,8 @@ pub(super) async fn maybe_build_local_standard_decision_payload_for_candidate(
provider_request_method: Some(serde_json::Value::Null),
provider_request_headers: Some(&resolved.provider_request_headers),
original_headers: &parts.headers,
request_path: Some(parts.uri.path()),
request_query_string: parts.uri.query(),
request_origin: Some(crate::ai_serving::request_origin_from_parts(parts)),
original_request_body_json: Some(body_json),
original_request_body_base64: None,

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() {

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

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,
};

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)
}

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)
}

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(),

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,
@@ -107,6 +109,8 @@ pub(crate) async fn maybe_build_local_openai_chat_decision_payload_for_candidate
provider_request_method: Some(serde_json::Value::Null),
provider_request_headers: Some(&resolved.provider_request_headers),
original_headers: &parts.headers,
request_path: Some(parts.uri.path()),
request_query_string: parts.uri.query(),
request_origin: Some(crate::ai_serving::request_origin_from_parts(parts)),
original_request_body_json: Some(body_json),
original_request_body_base64: None,
@@ -147,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,

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() {

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,
));
}
}

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 {

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"
));
}
}

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,
@@ -17,6 +18,18 @@ use crate::ai_serving::transport::{
};
use crate::{AiExecutionDecision, GatewayError};
fn effective_stream_accept_mode(
payload_upstream_is_stream: bool,
provider_request_body: &serde_json::Value,
) -> bool {
payload_upstream_is_stream
|| provider_request_body
.as_object()
.and_then(|body| body.get("stream"))
.and_then(serde_json::Value::as_bool)
.unwrap_or(false)
}
pub(crate) fn build_openai_chat_stream_plan_from_decision(
parts: &http::request::Parts,
body_json: &serde_json::Value,
@@ -53,22 +66,34 @@ 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 effective_upstream_is_stream =
effective_stream_accept_mode(payload.upstream_is_stream, &provider_request_body_value);
let extra_headers = std::mem::take(&mut payload.extra_headers);
let mut provider_request_headers =
build_standard_plan_fallback_headers(StandardPlanFallbackHeadersInput {
@@ -82,9 +107,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: effective_upstream_is_stream,
build_from_request_when_empty: true,
accept_policy: StandardPlanFallbackAcceptPolicy::TextEventStreamRequired,
accept_policy: StandardPlanFallbackAcceptPolicy::TextEventStreamIfStreamingOrWildcard,
});
let content_type = payload
.content_type
@@ -157,13 +182,18 @@ 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 effective_upstream_is_stream =
effective_stream_accept_mode(payload.upstream_is_stream, &provider_request_body_value);
let accept_policy = if effective_upstream_is_stream
&& provider_adaptation_requires_eventstream_accept(
envelope_name,
core.provider_api_format.as_str(),
) {
StandardPlanFallbackAcceptPolicy::ProviderEventStreamIfMissing
} else if envelope_name.is_some() {
StandardPlanFallbackAcceptPolicy::TextEventStreamIfStreaming
} else {
StandardPlanFallbackAcceptPolicy::TextEventStreamRequired
StandardPlanFallbackAcceptPolicy::TextEventStreamIfStreamingOrWildcard
};
let mut provider_request_headers =
build_standard_plan_fallback_headers(StandardPlanFallbackHeadersInput {
@@ -177,7 +207,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: effective_upstream_is_stream,
build_from_request_when_empty: false,
accept_policy,
});
@@ -219,7 +249,7 @@ pub(crate) fn build_openai_responses_stream_plan_from_decision(
plan_url = %plan.url,
client_api_format = %plan.client_api_format,
provider_api_format = %plan.provider_api_format,
upstream_is_stream = payload.upstream_is_stream,
upstream_is_stream = effective_upstream_is_stream,
compact,
"gateway built local openai responses stream execution plan"
);
@@ -427,6 +457,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()

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 =

View File

@@ -105,6 +105,8 @@ pub(crate) async fn maybe_build_local_openai_responses_decision_payload_for_cand
provider_request_method: Some(serde_json::Value::Null),
provider_request_headers: Some(&resolved.provider_request_headers),
original_headers: &parts.headers,
request_path: Some(parts.uri.path()),
request_query_string: parts.uri.query(),
request_origin: Some(crate::ai_serving::request_origin_from_parts(parts)),
original_request_body_json: Some(body_json),
original_request_body_base64: None,

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(

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,
});

View File

@@ -40,7 +40,8 @@ pub(crate) use aether_ai_formats::api::{
copy_request_number_field_as, core_error_background_report_kind,
core_error_default_client_api_format, core_success_background_report_kind,
default_model_for_openai_image_operation, encode_done_sse, encode_json_sse,
encode_kiro_sse_events, estimate_kiro_tokens, extract_openai_text_content,
encode_kiro_sse_events, endpoint_config_forces_upstream_stream_policy,
enforce_request_body_stream_field, estimate_kiro_tokens, extract_openai_text_content,
find_kiro_real_thinking_end_tag, find_kiro_real_thinking_end_tag_at_buffer_end,
find_kiro_real_thinking_start_tag, force_upstream_streaming_for_provider,
gemini_request_is_image_generation, implicit_sync_finalize_report_kind,
@@ -70,7 +71,8 @@ pub(crate) use aether_ai_formats::api::{
provider_adaptation_should_unwrap_stream_envelope,
provider_private_response_allows_sync_finalize, request_candidate_api_format_preference,
request_candidate_api_formats, request_conversion_kind,
request_conversion_requires_enable_flag, resolve_claude_stream_spec, resolve_claude_sync_spec,
request_conversion_requires_enable_flag, request_path_implies_stream_request,
resolve_claude_stream_spec, resolve_claude_sync_spec,
resolve_execution_runtime_stream_plan_kind, resolve_execution_runtime_sync_plan_kind,
resolve_finalize_stream_rewrite_mode, resolve_gemini_files_stream_spec,
resolve_gemini_files_sync_spec, resolve_gemini_stream_spec, resolve_gemini_sync_spec,
@@ -79,11 +81,13 @@ pub(crate) use aether_ai_formats::api::{
resolve_local_video_sync_spec, resolve_openai_chat_max_tokens,
resolve_openai_responses_stream_spec, resolve_openai_responses_sync_spec,
resolve_requested_gemini_image_model_for_request,
resolve_requested_openai_image_model_for_request, stream_body_contains_error_event,
supports_stream_execution_decision_kind, supports_sync_execution_decision_kind,
sync_chat_response_conversion_kind, sync_cli_response_conversion_kind,
transform_provider_private_stream_line, value_as_u64, AiControlPlanRequest,
AiSurfaceFinalizeError, AiSurfaceStreamRewriter, CanonicalStreamFrame,
resolve_requested_openai_image_model_for_request,
resolve_upstream_is_stream_from_endpoint_config, sanitize_request_path,
sanitize_request_path_and_query, sanitize_request_query_string,
stream_body_contains_error_event, supports_stream_execution_decision_kind,
supports_sync_execution_decision_kind, sync_chat_response_conversion_kind,
sync_cli_response_conversion_kind, transform_provider_private_stream_line, value_as_u64,
AiControlPlanRequest, AiSurfaceFinalizeError, AiSurfaceStreamRewriter, CanonicalStreamFrame,
ChatGptWebImageRequestError, ClaudeClientEmitter, ClaudeProviderState,
ExecutionRuntimeAuthContext, FinalizeStreamRewriteMode, FormatContext, GeminiClientEmitter,
GeminiImageRequestForOpenAi, GeminiProviderState, KiroToClaudeCliStreamState,