mirror of
https://github.com/fawney19/Aether.git
synced 2026-09-03 09:50:21 +08:00
fix: enforce provider upstream stream policy
This commit is contained in:
@@ -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 {
|
||||
|
||||
@@ -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"),
|
||||
|
||||
@@ -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,
|
||||
})
|
||||
}
|
||||
|
||||
@@ -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,
|
||||
|
||||
Reference in New Issue
Block a user