mirror of
https://github.com/fawney19/Aether.git
synced 2026-10-07 18:07:47 +08:00
1205 lines
44 KiB
Rust
1205 lines
44 KiB
Rust
use std::collections::BTreeMap;
|
|
|
|
use serde_json::Value;
|
|
|
|
use crate::antigravity::is_antigravity_provider_transport;
|
|
use crate::auth::{
|
|
build_complete_passthrough_headers, build_complete_passthrough_headers_with_auth,
|
|
resolve_local_gemini_auth, resolve_local_openai_bearer_auth, resolve_local_standard_auth,
|
|
};
|
|
use crate::claude_code::build_claude_code_passthrough_headers;
|
|
use crate::claude_code::local_claude_code_transport_unsupported_reason_with_network;
|
|
use crate::gemini_cli::is_gemini_cli_provider_transport;
|
|
use crate::grok::{is_grok_provider_transport, resolve_grok_session_auth};
|
|
use crate::kiro::{
|
|
build_kiro_provider_headers, build_kiro_provider_request_body, is_kiro_provider_transport,
|
|
local_kiro_request_transport_unsupported_reason_with_network, KiroAuthConfig,
|
|
KiroProviderHeadersInput,
|
|
};
|
|
use crate::policy::{
|
|
local_gemini_transport_unsupported_reason_with_network,
|
|
local_standard_transport_unsupported_reason_with_network,
|
|
};
|
|
use crate::rules::{
|
|
apply_local_body_rules_with_request_headers, apply_local_header_rules_with_request_headers,
|
|
};
|
|
use crate::snapshot::GatewayProviderTransportSnapshot;
|
|
use crate::vertex::{
|
|
is_vertex_service_account_transport_context, is_vertex_transport_context,
|
|
local_vertex_gemini_transport_unsupported_reason_with_network,
|
|
};
|
|
use crate::{
|
|
build_transport_request_url_for_request_body, ensure_upstream_auth_header,
|
|
TransportRequestUrlParams,
|
|
};
|
|
|
|
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
|
pub enum SameFormatProviderFamily {
|
|
Standard,
|
|
Gemini,
|
|
}
|
|
|
|
#[derive(Debug, Clone, Copy)]
|
|
pub struct SameFormatProviderRequestBehaviorParams<'a> {
|
|
pub require_streaming: bool,
|
|
pub provider_api_format: &'a str,
|
|
pub report_kind: &'static str,
|
|
}
|
|
|
|
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
|
pub struct SameFormatProviderRequestBehavior {
|
|
pub is_antigravity: bool,
|
|
pub is_gemini_cli: bool,
|
|
pub is_claude_code: bool,
|
|
pub is_vertex: bool,
|
|
pub is_kiro: bool,
|
|
pub upstream_is_stream: bool,
|
|
pub force_body_stream_field: bool,
|
|
pub report_kind: &'static str,
|
|
}
|
|
|
|
#[derive(Debug, Clone, Copy)]
|
|
pub struct SameFormatProviderRequestBodyInput<'a> {
|
|
pub body_json: &'a Value,
|
|
pub mapped_model: &'a str,
|
|
pub client_api_format: &'a str,
|
|
pub provider_api_format: &'a str,
|
|
pub source_model: Option<&'a str>,
|
|
pub family: SameFormatProviderFamily,
|
|
pub body_rules: Option<&'a Value>,
|
|
pub request_headers: Option<&'a http::HeaderMap>,
|
|
pub upstream_is_stream: bool,
|
|
pub force_body_stream_field: bool,
|
|
pub kiro_auth_config: Option<&'a KiroAuthConfig>,
|
|
pub is_claude_code: bool,
|
|
pub enable_model_directives: bool,
|
|
}
|
|
|
|
#[derive(Debug, Clone, Copy)]
|
|
pub struct SameFormatProviderUpstreamUrlParams<'a> {
|
|
pub provider_api_format: &'a str,
|
|
pub mapped_model: &'a str,
|
|
pub upstream_is_stream: bool,
|
|
pub request_query: Option<&'a str>,
|
|
pub kiro_api_region: Option<&'a str>,
|
|
pub provider_request_body: Option<&'a Value>,
|
|
}
|
|
|
|
#[derive(Debug, Clone, Copy)]
|
|
pub struct SameFormatProviderHeadersInput<'a> {
|
|
pub headers: &'a http::HeaderMap,
|
|
pub provider_request_body: &'a Value,
|
|
pub original_request_body: &'a Value,
|
|
pub header_rules: Option<&'a Value>,
|
|
pub behavior: SameFormatProviderRequestBehavior,
|
|
pub auth_header: Option<&'a str>,
|
|
pub auth_value: Option<&'a str>,
|
|
pub extra_headers: &'a BTreeMap<String, String>,
|
|
pub key_fingerprint: Option<&'a Value>,
|
|
pub kiro_auth_config: Option<&'a KiroAuthConfig>,
|
|
pub kiro_machine_id: Option<&'a str>,
|
|
}
|
|
|
|
pub fn classify_same_format_provider_request_behavior(
|
|
transport: &GatewayProviderTransportSnapshot,
|
|
params: SameFormatProviderRequestBehaviorParams<'_>,
|
|
) -> SameFormatProviderRequestBehavior {
|
|
let is_antigravity = is_antigravity_provider_transport(transport);
|
|
let is_gemini_cli = is_gemini_cli_provider_transport(transport);
|
|
let is_claude_code = transport
|
|
.provider
|
|
.provider_type
|
|
.trim()
|
|
.eq_ignore_ascii_case("claude_code");
|
|
let is_vertex = is_vertex_transport_context(transport);
|
|
let is_kiro = is_kiro_provider_transport(transport);
|
|
let gemini_cli_requires_upstream_streaming = is_gemini_cli
|
|
&& crate::gemini_cli::gemini_cli_v1internal_requires_upstream_streaming(
|
|
params.provider_api_format,
|
|
params.require_streaming,
|
|
);
|
|
let upstream_is_stream = aether_ai_formats::resolve_upstream_is_stream_from_endpoint_config(
|
|
transport.endpoint.config.as_ref(),
|
|
params.require_streaming,
|
|
is_kiro
|
|
|| is_antigravity
|
|
|| gemini_cli_requires_upstream_streaming
|
|
|| aether_ai_formats::api::force_upstream_streaming_for_provider(
|
|
transport.provider.provider_type.as_str(),
|
|
params.provider_api_format,
|
|
),
|
|
);
|
|
let force_body_stream_field = aether_ai_formats::endpoint_config_forces_upstream_stream_policy(
|
|
transport.endpoint.config.as_ref(),
|
|
);
|
|
let report_kind = if is_kiro && !params.require_streaming {
|
|
"claude_cli_sync_finalize"
|
|
} else if (is_gemini_cli || is_antigravity) && !params.require_streaming {
|
|
match params.report_kind {
|
|
"gemini_chat_sync_success" => "gemini_chat_sync_finalize",
|
|
"gemini_cli_sync_success" => "gemini_cli_sync_finalize",
|
|
_ => params.report_kind,
|
|
}
|
|
} else {
|
|
params.report_kind
|
|
};
|
|
|
|
SameFormatProviderRequestBehavior {
|
|
is_antigravity,
|
|
is_gemini_cli,
|
|
is_claude_code,
|
|
is_vertex,
|
|
is_kiro,
|
|
upstream_is_stream,
|
|
force_body_stream_field,
|
|
report_kind,
|
|
}
|
|
}
|
|
|
|
pub fn build_same_format_provider_request_body(
|
|
input: SameFormatProviderRequestBodyInput<'_>,
|
|
) -> Option<Value> {
|
|
if let Some(kiro_auth_config) = input.kiro_auth_config {
|
|
return build_kiro_provider_request_body(
|
|
input.body_json,
|
|
input.mapped_model,
|
|
kiro_auth_config,
|
|
input.body_rules,
|
|
input.request_headers,
|
|
);
|
|
}
|
|
|
|
if embedding_multimodal_input_requires_aliyun_provider(
|
|
input.client_api_format,
|
|
input.provider_api_format,
|
|
input.body_json,
|
|
) {
|
|
return None;
|
|
}
|
|
|
|
let mut provider_request_body = if aether_ai_formats::api_format_alias_matches(
|
|
input.client_api_format,
|
|
input.provider_api_format,
|
|
) {
|
|
let request_body_object = input.body_json.as_object()?;
|
|
serde_json::Map::from_iter(
|
|
request_body_object
|
|
.iter()
|
|
.map(|(key, value)| (key.clone(), value.clone())),
|
|
)
|
|
} else {
|
|
aether_ai_formats::convert_request(
|
|
input.client_api_format,
|
|
input.provider_api_format,
|
|
input.body_json,
|
|
&aether_ai_formats::FormatContext::default().with_mapped_model(input.mapped_model),
|
|
)
|
|
.ok()?
|
|
.as_object()?
|
|
.clone()
|
|
};
|
|
match input.family {
|
|
SameFormatProviderFamily::Standard => {
|
|
provider_request_body.insert(
|
|
"model".to_string(),
|
|
Value::String(input.mapped_model.to_string()),
|
|
);
|
|
}
|
|
SameFormatProviderFamily::Gemini => {
|
|
provider_request_body.remove("model");
|
|
}
|
|
}
|
|
let mut provider_request_body = Value::Object(provider_request_body);
|
|
if input.is_claude_code {
|
|
crate::claude_code::sanitize_claude_code_request_body(&mut provider_request_body);
|
|
}
|
|
if input.enable_model_directives {
|
|
if let Some(source_model) = input.source_model {
|
|
aether_ai_formats::apply_model_directive_overrides_from_model(
|
|
&mut provider_request_body,
|
|
input.provider_api_format,
|
|
input.mapped_model,
|
|
source_model,
|
|
);
|
|
}
|
|
}
|
|
if !apply_local_body_rules_with_request_headers(
|
|
&mut provider_request_body,
|
|
input.body_rules,
|
|
Some(input.body_json),
|
|
input.request_headers,
|
|
) {
|
|
return None;
|
|
}
|
|
if matches!(input.family, SameFormatProviderFamily::Gemini)
|
|
&& aether_ai_formats::api_format_alias_matches(
|
|
input.provider_api_format,
|
|
"gemini:generate_content",
|
|
)
|
|
{
|
|
strip_gemini_function_response_ids(&mut provider_request_body);
|
|
}
|
|
let require_body_stream_field = input.force_body_stream_field
|
|
|| input
|
|
.body_json
|
|
.as_object()
|
|
.is_some_and(|object| object.contains_key("stream"));
|
|
aether_ai_formats::enforce_request_body_stream_field(
|
|
&mut provider_request_body,
|
|
input.provider_api_format,
|
|
input.upstream_is_stream,
|
|
require_body_stream_field,
|
|
);
|
|
Some(provider_request_body)
|
|
}
|
|
|
|
fn embedding_multimodal_input_requires_aliyun_provider(
|
|
client_api_format: &str,
|
|
provider_api_format: &str,
|
|
body_json: &Value,
|
|
) -> bool {
|
|
aether_ai_formats::is_embedding_api_format(client_api_format)
|
|
&& embedding_input_is_multimodal(body_json.get("input"))
|
|
&& aether_ai_formats::normalize_api_format_alias(provider_api_format)
|
|
!= "aliyun:multimodal_embedding"
|
|
}
|
|
|
|
fn embedding_input_is_multimodal(value: Option<&Value>) -> bool {
|
|
value
|
|
.and_then(Value::as_array)
|
|
.is_some_and(|items| !items.is_empty() && items.iter().all(embedding_content_is_multimodal))
|
|
}
|
|
|
|
fn embedding_content_is_multimodal(value: &Value) -> bool {
|
|
value.as_object().is_some_and(|object| {
|
|
["text", "image", "video", "multi_images"]
|
|
.iter()
|
|
.any(|key| object.contains_key(*key))
|
|
})
|
|
}
|
|
|
|
fn strip_gemini_function_response_ids(value: &mut Value) {
|
|
match value {
|
|
Value::Object(object) => {
|
|
for key in ["functionResponse", "function_response"] {
|
|
if let Some(function_response) = object.get_mut(key).and_then(Value::as_object_mut)
|
|
{
|
|
function_response.remove("id");
|
|
}
|
|
}
|
|
for child in object.values_mut() {
|
|
strip_gemini_function_response_ids(child);
|
|
}
|
|
}
|
|
Value::Array(items) => {
|
|
for item in items {
|
|
strip_gemini_function_response_ids(item);
|
|
}
|
|
}
|
|
_ => {}
|
|
}
|
|
}
|
|
|
|
pub fn build_same_format_provider_upstream_url(
|
|
transport: &GatewayProviderTransportSnapshot,
|
|
params: SameFormatProviderUpstreamUrlParams<'_>,
|
|
) -> Option<String> {
|
|
build_transport_request_url_for_request_body(
|
|
transport,
|
|
TransportRequestUrlParams {
|
|
provider_api_format: params.provider_api_format,
|
|
mapped_model: Some(params.mapped_model),
|
|
upstream_is_stream: params.upstream_is_stream,
|
|
request_query: params.request_query,
|
|
kiro_api_region: params.kiro_api_region,
|
|
},
|
|
params.provider_request_body,
|
|
)
|
|
}
|
|
|
|
pub fn build_same_format_provider_headers(
|
|
input: SameFormatProviderHeadersInput<'_>,
|
|
) -> Option<BTreeMap<String, String>> {
|
|
if let Some(kiro_auth_config) = input.kiro_auth_config {
|
|
return build_kiro_provider_headers(KiroProviderHeadersInput {
|
|
headers: input.headers,
|
|
provider_request_body: input.provider_request_body,
|
|
original_request_body: input.original_request_body,
|
|
header_rules: input.header_rules,
|
|
auth_header: input.auth_header.unwrap_or_default(),
|
|
auth_value: input.auth_value.unwrap_or_default(),
|
|
auth_config: kiro_auth_config,
|
|
machine_id: input.kiro_machine_id.unwrap_or_default(),
|
|
});
|
|
}
|
|
|
|
let auth_header = input.auth_header.unwrap_or_default();
|
|
let auth_value = input.auth_value.unwrap_or_default();
|
|
let mut provider_request_headers = if input.behavior.is_claude_code {
|
|
build_claude_code_passthrough_headers(
|
|
input.headers,
|
|
auth_header,
|
|
auth_value,
|
|
input.extra_headers,
|
|
input.behavior.upstream_is_stream,
|
|
input.key_fingerprint,
|
|
)
|
|
} else if input.behavior.is_vertex {
|
|
build_complete_passthrough_headers(
|
|
input.headers,
|
|
input.extra_headers,
|
|
Some("application/json"),
|
|
)
|
|
} else {
|
|
build_complete_passthrough_headers_with_auth(
|
|
input.headers,
|
|
auth_header,
|
|
auth_value,
|
|
input.extra_headers,
|
|
Some("application/json"),
|
|
)
|
|
};
|
|
|
|
let protected_headers = input
|
|
.auth_header
|
|
.filter(|value| !value.trim().is_empty())
|
|
.map(|value| vec![value, "content-type"])
|
|
.unwrap_or_else(|| vec!["content-type"]);
|
|
if !apply_local_header_rules_with_request_headers(
|
|
&mut provider_request_headers,
|
|
input.header_rules,
|
|
&protected_headers,
|
|
input.provider_request_body,
|
|
Some(input.original_request_body),
|
|
Some(input.headers),
|
|
) {
|
|
return None;
|
|
}
|
|
if let (Some(auth_header), Some(auth_value)) = (input.auth_header, input.auth_value) {
|
|
ensure_upstream_auth_header(&mut provider_request_headers, auth_header, auth_value);
|
|
}
|
|
if input.behavior.upstream_is_stream {
|
|
provider_request_headers.insert("accept".to_string(), "text/event-stream".to_string());
|
|
}
|
|
Some(provider_request_headers)
|
|
}
|
|
|
|
pub fn same_format_provider_transport_supported(
|
|
behavior: &SameFormatProviderRequestBehavior,
|
|
transport: &GatewayProviderTransportSnapshot,
|
|
family: SameFormatProviderFamily,
|
|
api_format: &str,
|
|
) -> bool {
|
|
same_format_provider_transport_unsupported_reason(behavior, transport, family, api_format)
|
|
.is_none()
|
|
}
|
|
|
|
pub fn same_format_provider_transport_unsupported_reason(
|
|
behavior: &SameFormatProviderRequestBehavior,
|
|
transport: &GatewayProviderTransportSnapshot,
|
|
family: SameFormatProviderFamily,
|
|
api_format: &str,
|
|
) -> Option<&'static str> {
|
|
if is_grok_provider_transport(transport) && matches!(family, SameFormatProviderFamily::Standard)
|
|
{
|
|
return None;
|
|
}
|
|
if behavior.is_kiro {
|
|
local_kiro_request_transport_unsupported_reason_with_network(transport)
|
|
} else if behavior.is_antigravity {
|
|
None
|
|
} else if behavior.is_claude_code {
|
|
local_claude_code_transport_unsupported_reason_with_network(transport, api_format)
|
|
} else if behavior.is_vertex {
|
|
local_vertex_gemini_transport_unsupported_reason_with_network(transport)
|
|
} else {
|
|
match family {
|
|
SameFormatProviderFamily::Standard => {
|
|
local_standard_transport_unsupported_reason_with_network(transport, api_format)
|
|
}
|
|
SameFormatProviderFamily::Gemini => {
|
|
local_gemini_transport_unsupported_reason_with_network(transport, api_format)
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
pub fn same_format_provider_transport_unsupported_reason_for_trace(
|
|
transport: &GatewayProviderTransportSnapshot,
|
|
provider_api_format: &str,
|
|
) -> Option<&'static str> {
|
|
let normalized_api_format =
|
|
match aether_ai_formats::normalize_api_format_alias(provider_api_format).as_str() {
|
|
"openai:chat" => "openai:chat",
|
|
"openai:responses" => "openai:responses",
|
|
"openai:responses:compact" => "openai:responses:compact",
|
|
"claude:messages" => "claude:messages",
|
|
"gemini:generate_content" => "gemini:generate_content",
|
|
_ => return Some("transport_api_format_unsupported"),
|
|
};
|
|
let behavior = classify_same_format_provider_request_behavior(
|
|
transport,
|
|
SameFormatProviderRequestBehaviorParams {
|
|
require_streaming: false,
|
|
provider_api_format: normalized_api_format,
|
|
report_kind: "trace_candidate_metadata",
|
|
},
|
|
);
|
|
if !behavior.is_antigravity
|
|
&& !behavior.is_claude_code
|
|
&& !behavior.is_gemini_cli
|
|
&& !behavior.is_vertex
|
|
&& !behavior.is_kiro
|
|
{
|
|
return None;
|
|
}
|
|
|
|
let family = if normalized_api_format.starts_with("gemini:") {
|
|
SameFormatProviderFamily::Gemini
|
|
} else {
|
|
SameFormatProviderFamily::Standard
|
|
};
|
|
same_format_provider_transport_unsupported_reason(
|
|
&behavior,
|
|
transport,
|
|
family,
|
|
normalized_api_format,
|
|
)
|
|
}
|
|
|
|
pub fn should_try_same_format_provider_oauth_auth(
|
|
behavior: &SameFormatProviderRequestBehavior,
|
|
transport: &GatewayProviderTransportSnapshot,
|
|
family: SameFormatProviderFamily,
|
|
provider_api_format: &str,
|
|
) -> bool {
|
|
let provider_api_format = aether_ai_formats::normalize_api_format_alias(provider_api_format);
|
|
behavior.is_kiro
|
|
|| matches!(family, SameFormatProviderFamily::Standard)
|
|
&& resolve_same_format_standard_direct_auth(transport, provider_api_format.as_str())
|
|
.is_none()
|
|
|| matches!(family, SameFormatProviderFamily::Gemini)
|
|
&& behavior.is_vertex
|
|
&& is_vertex_service_account_transport_context(transport)
|
|
|| matches!(family, SameFormatProviderFamily::Gemini)
|
|
&& !behavior.is_vertex
|
|
&& resolve_local_gemini_auth(transport).is_none()
|
|
}
|
|
|
|
pub fn resolve_same_format_provider_direct_auth(
|
|
behavior: &SameFormatProviderRequestBehavior,
|
|
transport: &GatewayProviderTransportSnapshot,
|
|
family: SameFormatProviderFamily,
|
|
provider_api_format: &str,
|
|
) -> Option<(String, String)> {
|
|
if is_grok_provider_transport(transport) && matches!(family, SameFormatProviderFamily::Standard)
|
|
{
|
|
return resolve_grok_session_auth(transport);
|
|
}
|
|
if behavior.is_vertex {
|
|
None
|
|
} else {
|
|
match family {
|
|
SameFormatProviderFamily::Standard => {
|
|
resolve_same_format_standard_direct_auth(transport, provider_api_format)
|
|
}
|
|
SameFormatProviderFamily::Gemini => resolve_local_gemini_auth(transport),
|
|
}
|
|
}
|
|
}
|
|
|
|
fn resolve_same_format_standard_direct_auth(
|
|
transport: &GatewayProviderTransportSnapshot,
|
|
provider_api_format: &str,
|
|
) -> Option<(String, String)> {
|
|
if aether_ai_formats::api_format_alias_matches(provider_api_format, "openai:embedding") {
|
|
resolve_local_openai_bearer_auth(transport)
|
|
} else {
|
|
resolve_local_standard_auth(transport)
|
|
}
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use std::collections::BTreeMap;
|
|
|
|
use super::*;
|
|
use crate::snapshot::{
|
|
GatewayProviderTransportEndpoint, GatewayProviderTransportKey,
|
|
GatewayProviderTransportProvider,
|
|
};
|
|
use serde_json::json;
|
|
|
|
fn sample_transport(provider_type: &str) -> 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: "openai:chat".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: None,
|
|
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,
|
|
upstream_metadata: None,
|
|
decrypted_api_key: "secret".to_string(),
|
|
decrypted_auth_config: None,
|
|
},
|
|
}
|
|
}
|
|
|
|
#[test]
|
|
fn classifies_streaming_and_report_kind_for_provider_private_transports() {
|
|
let kiro = sample_transport("kiro");
|
|
let behavior = classify_same_format_provider_request_behavior(
|
|
&kiro,
|
|
SameFormatProviderRequestBehaviorParams {
|
|
require_streaming: false,
|
|
provider_api_format: "claude:messages",
|
|
report_kind: "claude_chat_sync_success",
|
|
},
|
|
);
|
|
|
|
assert!(behavior.is_kiro);
|
|
assert!(behavior.upstream_is_stream);
|
|
assert_eq!(behavior.report_kind, "claude_cli_sync_finalize");
|
|
|
|
let antigravity = sample_transport("antigravity");
|
|
let behavior = classify_same_format_provider_request_behavior(
|
|
&antigravity,
|
|
SameFormatProviderRequestBehaviorParams {
|
|
require_streaming: false,
|
|
provider_api_format: "gemini:generate_content",
|
|
report_kind: "gemini_chat_sync_success",
|
|
},
|
|
);
|
|
|
|
assert!(behavior.is_antigravity);
|
|
assert!(behavior.upstream_is_stream);
|
|
assert_eq!(behavior.report_kind, "gemini_chat_sync_finalize");
|
|
|
|
let gemini_cli = sample_transport("gemini_cli");
|
|
let behavior = classify_same_format_provider_request_behavior(
|
|
&gemini_cli,
|
|
SameFormatProviderRequestBehaviorParams {
|
|
require_streaming: false,
|
|
provider_api_format: "gemini:generate_content",
|
|
report_kind: "gemini_cli_sync_success",
|
|
},
|
|
);
|
|
|
|
assert!(!behavior.upstream_is_stream);
|
|
assert_eq!(behavior.report_kind, "gemini_cli_sync_finalize");
|
|
}
|
|
|
|
#[test]
|
|
fn same_format_behavior_resolves_endpoint_stream_policy() {
|
|
let mut force_stream = sample_transport("openai");
|
|
force_stream.endpoint.config = Some(json!({
|
|
"upstream_stream_policy": "force_stream"
|
|
}));
|
|
let behavior = classify_same_format_provider_request_behavior(
|
|
&force_stream,
|
|
SameFormatProviderRequestBehaviorParams {
|
|
require_streaming: false,
|
|
provider_api_format: "openai:chat",
|
|
report_kind: "openai_chat_sync_success",
|
|
},
|
|
);
|
|
assert!(behavior.upstream_is_stream);
|
|
|
|
let mut force_non_stream = sample_transport("openai");
|
|
force_non_stream.endpoint.config = Some(json!({
|
|
"upstreamStreamPolicy": "force_non_stream"
|
|
}));
|
|
let behavior = classify_same_format_provider_request_behavior(
|
|
&force_non_stream,
|
|
SameFormatProviderRequestBehaviorParams {
|
|
require_streaming: true,
|
|
provider_api_format: "openai:chat",
|
|
report_kind: "openai_chat_stream_success",
|
|
},
|
|
);
|
|
assert!(!behavior.upstream_is_stream);
|
|
|
|
let mut auto = sample_transport("openai");
|
|
auto.endpoint.config = Some(json!({
|
|
"upstream_stream": "auto"
|
|
}));
|
|
let stream_behavior = classify_same_format_provider_request_behavior(
|
|
&auto,
|
|
SameFormatProviderRequestBehaviorParams {
|
|
require_streaming: true,
|
|
provider_api_format: "openai:chat",
|
|
report_kind: "openai_chat_stream_success",
|
|
},
|
|
);
|
|
assert!(stream_behavior.upstream_is_stream);
|
|
let sync_behavior = classify_same_format_provider_request_behavior(
|
|
&auto,
|
|
SameFormatProviderRequestBehaviorParams {
|
|
require_streaming: false,
|
|
provider_api_format: "openai:chat",
|
|
report_kind: "openai_chat_sync_success",
|
|
},
|
|
);
|
|
assert!(!sync_behavior.upstream_is_stream);
|
|
}
|
|
|
|
#[test]
|
|
fn same_format_behavior_preserves_hard_streaming_constraint() {
|
|
let mut kiro = sample_transport("kiro");
|
|
kiro.endpoint.config = Some(json!({
|
|
"upstream_stream_policy": "force_non_stream"
|
|
}));
|
|
|
|
let behavior = classify_same_format_provider_request_behavior(
|
|
&kiro,
|
|
SameFormatProviderRequestBehaviorParams {
|
|
require_streaming: true,
|
|
provider_api_format: "claude:messages",
|
|
report_kind: "claude_chat_stream_success",
|
|
},
|
|
);
|
|
|
|
assert!(behavior.upstream_is_stream);
|
|
}
|
|
|
|
#[test]
|
|
fn same_format_behavior_preserves_gemini_cli_streaming_requests() {
|
|
let mut gemini_cli = sample_transport("gemini_cli");
|
|
gemini_cli.endpoint.config = Some(json!({
|
|
"upstream_stream_policy": "force_non_stream"
|
|
}));
|
|
|
|
let stream_behavior = classify_same_format_provider_request_behavior(
|
|
&gemini_cli,
|
|
SameFormatProviderRequestBehaviorParams {
|
|
require_streaming: true,
|
|
provider_api_format: "gemini:generate_content",
|
|
report_kind: "gemini_cli_stream_success",
|
|
},
|
|
);
|
|
assert!(stream_behavior.upstream_is_stream);
|
|
|
|
let sync_behavior = classify_same_format_provider_request_behavior(
|
|
&gemini_cli,
|
|
SameFormatProviderRequestBehaviorParams {
|
|
require_streaming: false,
|
|
provider_api_format: "gemini:generate_content",
|
|
report_kind: "gemini_cli_sync_success",
|
|
},
|
|
);
|
|
assert!(!sync_behavior.upstream_is_stream);
|
|
}
|
|
|
|
#[test]
|
|
fn same_format_policy_resolution_drives_standard_body_stream_field() {
|
|
for (endpoint_config, client_is_stream, expected_stream) in [
|
|
(
|
|
json!({"upstream_stream_policy": "force_stream"}),
|
|
false,
|
|
true,
|
|
),
|
|
(
|
|
json!({"upstreamStreamPolicy": "force_non_stream"}),
|
|
true,
|
|
false,
|
|
),
|
|
(json!({"upstream_stream": "auto"}), true, true),
|
|
(json!({"upstream_stream": "auto"}), false, false),
|
|
] {
|
|
let mut transport = sample_transport("openai");
|
|
transport.endpoint.config = Some(endpoint_config);
|
|
let behavior = classify_same_format_provider_request_behavior(
|
|
&transport,
|
|
SameFormatProviderRequestBehaviorParams {
|
|
require_streaming: client_is_stream,
|
|
provider_api_format: "openai:chat",
|
|
report_kind: "openai_chat_policy_test",
|
|
},
|
|
);
|
|
let body =
|
|
build_same_format_provider_request_body(SameFormatProviderRequestBodyInput {
|
|
body_json: &json!({
|
|
"model": "client-model",
|
|
"messages": [{"role": "user", "content": "hello"}],
|
|
"stream": client_is_stream
|
|
}),
|
|
mapped_model: "upstream-model",
|
|
client_api_format: "openai:chat",
|
|
provider_api_format: "openai:chat",
|
|
source_model: Some("client-model"),
|
|
family: SameFormatProviderFamily::Standard,
|
|
body_rules: None,
|
|
request_headers: None,
|
|
upstream_is_stream: behavior.upstream_is_stream,
|
|
force_body_stream_field: behavior.force_body_stream_field,
|
|
kiro_auth_config: None,
|
|
is_claude_code: false,
|
|
enable_model_directives: false,
|
|
})
|
|
.expect("body should build");
|
|
|
|
assert_eq!(body.get("stream"), Some(&json!(expected_stream)));
|
|
}
|
|
}
|
|
|
|
#[test]
|
|
fn resolves_direct_auth_except_vertex() {
|
|
let transport = sample_transport("openai");
|
|
let behavior = classify_same_format_provider_request_behavior(
|
|
&transport,
|
|
SameFormatProviderRequestBehaviorParams {
|
|
require_streaming: false,
|
|
provider_api_format: "openai:chat",
|
|
report_kind: "openai_chat_sync_success",
|
|
},
|
|
);
|
|
|
|
assert_eq!(
|
|
resolve_same_format_provider_direct_auth(
|
|
&behavior,
|
|
&transport,
|
|
SameFormatProviderFamily::Standard,
|
|
"openai:chat",
|
|
),
|
|
Some(("x-api-key".to_string(), "secret".to_string()))
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
fn resolves_openai_embedding_direct_auth_with_bearer_header() {
|
|
let mut transport = sample_transport("custom");
|
|
transport.endpoint.api_format = "openai:embedding".to_string();
|
|
transport.key.auth_type = "api_key".to_string();
|
|
let behavior = classify_same_format_provider_request_behavior(
|
|
&transport,
|
|
SameFormatProviderRequestBehaviorParams {
|
|
require_streaming: false,
|
|
provider_api_format: "openai:embedding",
|
|
report_kind: "openai_embedding_sync_success",
|
|
},
|
|
);
|
|
|
|
assert_eq!(
|
|
resolve_same_format_provider_direct_auth(
|
|
&behavior,
|
|
&transport,
|
|
SameFormatProviderFamily::Standard,
|
|
"openai:embedding",
|
|
),
|
|
Some(("authorization".to_string(), "Bearer secret".to_string()))
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
fn keeps_claude_same_format_api_key_on_x_api_key_header() {
|
|
let mut transport = sample_transport("custom");
|
|
transport.endpoint.api_format = "claude:messages".to_string();
|
|
transport.key.auth_type = "api_key".to_string();
|
|
let behavior = classify_same_format_provider_request_behavior(
|
|
&transport,
|
|
SameFormatProviderRequestBehaviorParams {
|
|
require_streaming: false,
|
|
provider_api_format: "claude:messages",
|
|
report_kind: "claude_chat_sync_success",
|
|
},
|
|
);
|
|
|
|
assert_eq!(
|
|
resolve_same_format_provider_direct_auth(
|
|
&behavior,
|
|
&transport,
|
|
SameFormatProviderFamily::Standard,
|
|
"claude:messages",
|
|
),
|
|
Some(("x-api-key".to_string(), "secret".to_string()))
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
fn builds_same_format_standard_body_with_mapped_model_and_stream_flag() {
|
|
let body = build_same_format_provider_request_body(SameFormatProviderRequestBodyInput {
|
|
body_json: &json!({
|
|
"model": "client-model",
|
|
"messages": [{"role": "user", "content": "hello"}]
|
|
}),
|
|
mapped_model: "upstream-model",
|
|
client_api_format: "openai:chat",
|
|
provider_api_format: "openai:chat",
|
|
source_model: Some("client-model"),
|
|
family: SameFormatProviderFamily::Standard,
|
|
body_rules: None,
|
|
request_headers: None,
|
|
upstream_is_stream: true,
|
|
force_body_stream_field: false,
|
|
kiro_auth_config: None,
|
|
is_claude_code: false,
|
|
enable_model_directives: false,
|
|
})
|
|
.expect("body should build");
|
|
|
|
assert_eq!(body.get("model"), Some(&json!("upstream-model")));
|
|
assert_eq!(body.get("stream"), Some(&json!(true)));
|
|
}
|
|
|
|
#[test]
|
|
fn same_format_embedding_body_rejects_multimodal_for_openai_like_provider() {
|
|
let body = build_same_format_provider_request_body(SameFormatProviderRequestBodyInput {
|
|
body_json: &json!({
|
|
"model": "qwen3-vl-embedding",
|
|
"input": [
|
|
{"text": "white running shoes"},
|
|
{"image": "https://example.com/shoe.png"}
|
|
]
|
|
}),
|
|
mapped_model: "openai-qwen-fallback",
|
|
client_api_format: "openai:embedding",
|
|
provider_api_format: "openai:embedding",
|
|
source_model: Some("qwen3-vl-embedding"),
|
|
family: SameFormatProviderFamily::Standard,
|
|
body_rules: None,
|
|
request_headers: None,
|
|
upstream_is_stream: false,
|
|
force_body_stream_field: false,
|
|
kiro_auth_config: None,
|
|
is_claude_code: false,
|
|
enable_model_directives: false,
|
|
});
|
|
|
|
assert!(body.is_none());
|
|
}
|
|
|
|
#[test]
|
|
fn same_format_standard_body_overrides_client_stream_for_non_stream_upstream() {
|
|
let body = build_same_format_provider_request_body(SameFormatProviderRequestBodyInput {
|
|
body_json: &json!({
|
|
"model": "client-model",
|
|
"messages": [{"role": "user", "content": "hello"}],
|
|
"stream": true
|
|
}),
|
|
mapped_model: "upstream-model",
|
|
client_api_format: "openai:chat",
|
|
provider_api_format: "openai:chat",
|
|
source_model: Some("client-model"),
|
|
family: SameFormatProviderFamily::Standard,
|
|
body_rules: None,
|
|
request_headers: None,
|
|
upstream_is_stream: false,
|
|
force_body_stream_field: false,
|
|
kiro_auth_config: None,
|
|
is_claude_code: false,
|
|
enable_model_directives: false,
|
|
})
|
|
.expect("body should build");
|
|
|
|
assert_eq!(body.get("model"), Some(&json!("upstream-model")));
|
|
assert_eq!(body.get("stream"), Some(&json!(false)));
|
|
}
|
|
|
|
#[test]
|
|
fn same_format_standard_body_does_not_add_stream_false_for_plain_sync_body() {
|
|
let body = build_same_format_provider_request_body(SameFormatProviderRequestBodyInput {
|
|
body_json: &json!({
|
|
"model": "client-model",
|
|
"messages": [{"role": "user", "content": "hello"}]
|
|
}),
|
|
mapped_model: "upstream-model",
|
|
client_api_format: "openai:chat",
|
|
provider_api_format: "openai:chat",
|
|
source_model: Some("client-model"),
|
|
family: SameFormatProviderFamily::Standard,
|
|
body_rules: None,
|
|
request_headers: None,
|
|
upstream_is_stream: false,
|
|
force_body_stream_field: false,
|
|
kiro_auth_config: None,
|
|
is_claude_code: false,
|
|
enable_model_directives: false,
|
|
})
|
|
.expect("body should build");
|
|
|
|
assert_eq!(body.get("model"), Some(&json!("upstream-model")));
|
|
assert!(body.get("stream").is_none());
|
|
}
|
|
|
|
#[test]
|
|
fn same_format_standard_body_forced_policy_adds_stream_false() {
|
|
let body = build_same_format_provider_request_body(SameFormatProviderRequestBodyInput {
|
|
body_json: &json!({
|
|
"model": "client-model",
|
|
"messages": [{"role": "user", "content": "hello"}]
|
|
}),
|
|
mapped_model: "upstream-model",
|
|
client_api_format: "openai:chat",
|
|
provider_api_format: "openai:chat",
|
|
source_model: Some("client-model"),
|
|
family: SameFormatProviderFamily::Standard,
|
|
body_rules: None,
|
|
request_headers: None,
|
|
upstream_is_stream: false,
|
|
force_body_stream_field: true,
|
|
kiro_auth_config: None,
|
|
is_claude_code: false,
|
|
enable_model_directives: false,
|
|
})
|
|
.expect("body should build");
|
|
|
|
assert_eq!(body.get("model"), Some(&json!("upstream-model")));
|
|
assert_eq!(body.get("stream"), Some(&json!(false)));
|
|
}
|
|
|
|
#[test]
|
|
fn same_format_gemini_body_removes_leaked_client_stream_field() {
|
|
let body = build_same_format_provider_request_body(SameFormatProviderRequestBodyInput {
|
|
body_json: &json!({
|
|
"contents": [{"role": "user", "parts": [{"text": "hello"}]}],
|
|
"stream": true
|
|
}),
|
|
mapped_model: "gemini-upstream",
|
|
client_api_format: "gemini:generate_content",
|
|
provider_api_format: "gemini:generate_content",
|
|
source_model: None,
|
|
family: SameFormatProviderFamily::Gemini,
|
|
body_rules: None,
|
|
request_headers: None,
|
|
upstream_is_stream: false,
|
|
force_body_stream_field: false,
|
|
kiro_auth_config: None,
|
|
is_claude_code: false,
|
|
enable_model_directives: false,
|
|
})
|
|
.expect("body should build");
|
|
|
|
assert!(body.get("stream").is_none());
|
|
}
|
|
|
|
#[test]
|
|
fn same_format_gemini_body_strips_function_response_id_for_upstream() {
|
|
let body = build_same_format_provider_request_body(SameFormatProviderRequestBodyInput {
|
|
body_json: &json!({
|
|
"contents": [
|
|
{
|
|
"role": "user",
|
|
"parts": [
|
|
{
|
|
"functionResponse": {
|
|
"id": "call_123",
|
|
"name": "lookup",
|
|
"response": {
|
|
"result": {
|
|
"id": "keep_result_id",
|
|
"ok": true
|
|
}
|
|
}
|
|
}
|
|
},
|
|
{
|
|
"function_response": {
|
|
"id": "call_456",
|
|
"name": "lookup_snake",
|
|
"response": {"ok": true}
|
|
}
|
|
}
|
|
]
|
|
}
|
|
],
|
|
"stream": true
|
|
}),
|
|
mapped_model: "gemini-upstream",
|
|
client_api_format: "gemini:generate_content",
|
|
provider_api_format: "gemini:generate_content",
|
|
source_model: None,
|
|
family: SameFormatProviderFamily::Gemini,
|
|
body_rules: None,
|
|
request_headers: None,
|
|
upstream_is_stream: true,
|
|
force_body_stream_field: false,
|
|
kiro_auth_config: None,
|
|
is_claude_code: false,
|
|
enable_model_directives: false,
|
|
})
|
|
.expect("body should build");
|
|
|
|
let function_response = &body["contents"][0]["parts"][0]["functionResponse"];
|
|
assert!(function_response.get("id").is_none());
|
|
assert_eq!(function_response["name"], "lookup");
|
|
assert_eq!(
|
|
function_response["response"]["result"]["id"],
|
|
"keep_result_id"
|
|
);
|
|
|
|
let function_response = &body["contents"][0]["parts"][1]["function_response"];
|
|
assert!(function_response.get("id").is_none());
|
|
assert_eq!(function_response["name"], "lookup_snake");
|
|
}
|
|
|
|
#[test]
|
|
fn same_format_stream_policy_wins_after_body_rules() {
|
|
let body_rules = json!([
|
|
{"action":"set","path":"stream","value":true}
|
|
]);
|
|
|
|
let body = build_same_format_provider_request_body(SameFormatProviderRequestBodyInput {
|
|
body_json: &json!({
|
|
"model": "client-model",
|
|
"messages": [{"role": "user", "content": "hello"}],
|
|
"stream": true
|
|
}),
|
|
mapped_model: "upstream-model",
|
|
client_api_format: "openai:chat",
|
|
provider_api_format: "openai:chat",
|
|
source_model: Some("client-model"),
|
|
family: SameFormatProviderFamily::Standard,
|
|
body_rules: Some(&body_rules),
|
|
request_headers: None,
|
|
upstream_is_stream: false,
|
|
force_body_stream_field: false,
|
|
kiro_auth_config: None,
|
|
is_claude_code: false,
|
|
enable_model_directives: false,
|
|
})
|
|
.expect("body should build");
|
|
|
|
assert_eq!(body.get("stream"), Some(&json!(false)));
|
|
}
|
|
|
|
#[test]
|
|
fn same_format_compact_stream_policy_wins_after_body_rules() {
|
|
let body_rules = json!([
|
|
{"action":"set","path":"stream","value":true}
|
|
]);
|
|
|
|
let body = build_same_format_provider_request_body(SameFormatProviderRequestBodyInput {
|
|
body_json: &json!({
|
|
"model": "client-model",
|
|
"input": "hello",
|
|
"stream": true
|
|
}),
|
|
mapped_model: "upstream-model",
|
|
client_api_format: "openai:responses:compact",
|
|
provider_api_format: "openai:responses:compact",
|
|
source_model: Some("client-model"),
|
|
family: SameFormatProviderFamily::Standard,
|
|
body_rules: Some(&body_rules),
|
|
request_headers: None,
|
|
upstream_is_stream: false,
|
|
force_body_stream_field: false,
|
|
kiro_auth_config: None,
|
|
is_claude_code: false,
|
|
enable_model_directives: false,
|
|
})
|
|
.expect("body should build");
|
|
|
|
assert!(body.get("stream").is_none());
|
|
}
|
|
|
|
#[test]
|
|
fn same_format_body_applies_model_directive_before_body_rules() {
|
|
let body = build_same_format_provider_request_body(SameFormatProviderRequestBodyInput {
|
|
body_json: &json!({
|
|
"model": "gpt-5.4-high",
|
|
"messages": [{"role": "user", "content": "hello"}],
|
|
"reasoning_effort": "low"
|
|
}),
|
|
mapped_model: "upstream-model",
|
|
client_api_format: "openai:chat",
|
|
provider_api_format: "openai:chat",
|
|
source_model: Some("gpt-5.4-high"),
|
|
family: SameFormatProviderFamily::Standard,
|
|
body_rules: Some(&json!([
|
|
{"action":"set","path":"metadata.body_rule_seen","value":true}
|
|
])),
|
|
request_headers: None,
|
|
upstream_is_stream: false,
|
|
force_body_stream_field: false,
|
|
kiro_auth_config: None,
|
|
is_claude_code: false,
|
|
enable_model_directives: true,
|
|
})
|
|
.expect("body should build");
|
|
|
|
assert_eq!(body["model"], "upstream-model");
|
|
assert_eq!(body["reasoning_effort"], "high");
|
|
assert_eq!(body["metadata"]["body_rule_seen"], true);
|
|
}
|
|
|
|
#[test]
|
|
fn builds_same_format_headers_with_auth_and_stream_accept() {
|
|
let provider_request_body = json!({"model": "upstream-model"});
|
|
let original_request_body = json!({"model": "client-model"});
|
|
let headers = build_same_format_provider_headers(SameFormatProviderHeadersInput {
|
|
headers: &http::HeaderMap::new(),
|
|
provider_request_body: &provider_request_body,
|
|
original_request_body: &original_request_body,
|
|
header_rules: None,
|
|
behavior: SameFormatProviderRequestBehavior {
|
|
is_antigravity: false,
|
|
is_gemini_cli: false,
|
|
is_claude_code: false,
|
|
is_vertex: false,
|
|
is_kiro: false,
|
|
upstream_is_stream: true,
|
|
force_body_stream_field: false,
|
|
report_kind: "openai_chat_stream_success",
|
|
},
|
|
auth_header: Some("x-api-key"),
|
|
auth_value: Some("secret"),
|
|
extra_headers: &BTreeMap::new(),
|
|
key_fingerprint: None,
|
|
kiro_auth_config: None,
|
|
kiro_machine_id: None,
|
|
})
|
|
.expect("headers should build");
|
|
|
|
assert_eq!(headers.get("x-api-key").map(String::as_str), Some("secret"));
|
|
assert_eq!(
|
|
headers.get("content-type").map(String::as_str),
|
|
Some("application/json")
|
|
);
|
|
assert_eq!(
|
|
headers.get("accept").map(String::as_str),
|
|
Some("text/event-stream")
|
|
);
|
|
}
|
|
}
|