feat(gateway): harden provider request execution

Preserve exact request payloads and model client surface and API operation explicitly.

Add Anthropic compatibility profiles, bounded stream commitment, and scoped OAuth retry behavior across provider transports.
This commit is contained in:
elky
2026-07-27 09:36:31 +08:00
parent 79b70f7b5c
commit 531cf11025
152 changed files with 13984 additions and 2075 deletions
+24 -16
View File
@@ -87,28 +87,36 @@ pub(crate) fn resolve_execution_runtime_stream_plan_kind(
parts: &http::request::Parts,
decision: &GatewayControlDecision,
) -> Option<&'static str> {
aether_ai_formats::api::resolve_execution_runtime_stream_plan_kind(
decision.route_class.as_deref(),
decision.route_family.as_deref(),
decision.route_kind.as_deref(),
decision.request_auth_channel.as_deref(),
&parts.method,
parts.uri.path(),
)
let plan_kind =
aether_ai_formats::api::resolve_execution_runtime_stream_plan_kind_with_client_surface(
decision.route_class.as_deref(),
decision.route_family.as_deref(),
decision.route_kind.as_deref(),
decision.client_surface,
decision.request_auth_channel.as_deref(),
&parts.method,
parts.uri.path(),
)?;
crate::ai_serving::plan_kind_matches_api_operation(plan_kind, true, decision.api_operation)
.then_some(plan_kind)
}
pub(crate) fn resolve_execution_runtime_sync_plan_kind(
parts: &http::request::Parts,
decision: &GatewayControlDecision,
) -> Option<&'static str> {
aether_ai_formats::api::resolve_execution_runtime_sync_plan_kind(
decision.route_class.as_deref(),
decision.route_family.as_deref(),
decision.route_kind.as_deref(),
decision.request_auth_channel.as_deref(),
&parts.method,
parts.uri.path(),
)
let plan_kind =
aether_ai_formats::api::resolve_execution_runtime_sync_plan_kind_with_client_surface(
decision.route_class.as_deref(),
decision.route_family.as_deref(),
decision.route_kind.as_deref(),
decision.client_surface,
decision.request_auth_channel.as_deref(),
&parts.method,
parts.uri.path(),
)?;
crate::ai_serving::plan_kind_matches_api_operation(plan_kind, false, decision.api_operation)
.then_some(plan_kind)
}
pub(crate) fn is_matching_stream_request(
@@ -13,6 +13,7 @@ fn same_format_claude_local_stream_rewriter_sanitizes_read_input_json_delta() {
let report_context = json!({
"provider_api_format": "claude:messages",
"client_api_format": "claude:messages",
"anthropic_compatibility_profile": "claude_code_legacy",
"needs_conversion": false,
});
let mut rewriter =
@@ -25,6 +25,9 @@ fn test_decision() -> GatewayControlDecision {
route_class: Some("ai_public".to_string()),
route_family: Some("openai".to_string()),
route_kind: Some("compact".to_string()),
client_surface: None,
api_operation: None,
gateway_credential_carrier: None,
request_auth_channel: None,
auth_endpoint_signature: Some("openai:responses:compact".to_string()),
execution_runtime_candidate: true,
@@ -1923,6 +1926,9 @@ fn local_finalize_handles_claude_chat_cross_format_sync_response_from_openai_cha
route_class: Some("ai_public".to_string()),
route_family: Some("claude".to_string()),
route_kind: Some("chat".to_string()),
client_surface: None,
api_operation: None,
gateway_credential_carrier: None,
request_auth_channel: None,
auth_endpoint_signature: Some("claude:messages".to_string()),
execution_runtime_candidate: true,
@@ -1991,6 +1997,9 @@ fn local_finalize_handles_gemini_cli_cross_format_sync_response_from_claude_cli(
route_class: Some("ai_public".to_string()),
route_family: Some("gemini".to_string()),
route_kind: Some("cli".to_string()),
client_surface: None,
api_operation: None,
gateway_credential_carrier: None,
request_auth_channel: None,
auth_endpoint_signature: Some("gemini:generate_content".to_string()),
execution_runtime_candidate: true,
+5 -1
View File
@@ -71,7 +71,7 @@ pub(crate) use self::transport::{
request_pair_allowed_for_transport, request_pair_direct_auth,
request_pair_transport_unsupported_reason, CandidateTransportPolicyFacts,
};
pub(crate) use crate::control::GatewayControlDecision;
pub(crate) use crate::control::{GatewayControlDecision, GatewayCredentialCarrier};
pub(crate) use crate::execution_runtime::{ConversionMode, ExecutionStrategy};
pub(crate) use crate::headers::RequestOrigin;
pub(crate) use aether_ai_serving::{
@@ -89,6 +89,7 @@ pub(crate) fn build_provider_transport_request_url(
upstream_is_stream: bool,
request_query: Option<&str>,
kiro_api_region: Option<&str>,
api_operation: Option<ApiOperation>,
) -> Option<String> {
self::transport::build_transport_request_url(
transport,
@@ -98,6 +99,7 @@ pub(crate) fn build_provider_transport_request_url(
upstream_is_stream,
request_query,
kiro_api_region,
api_operation,
},
)
}
@@ -109,6 +111,7 @@ pub(crate) fn build_provider_transport_request_url_for_request_body(
upstream_is_stream: bool,
request_query: Option<&str>,
kiro_api_region: Option<&str>,
api_operation: Option<ApiOperation>,
provider_request_body: Option<&serde_json::Value>,
) -> Option<String> {
self::transport::build_transport_request_url_for_request_body(
@@ -119,6 +122,7 @@ pub(crate) fn build_provider_transport_request_url_for_request_body(
upstream_is_stream,
request_query,
kiro_api_region,
api_operation,
},
provider_request_body,
)
@@ -67,6 +67,7 @@ pub(crate) async fn build_antigravity_v1internal_provider_request(
input.upstream_is_stream,
input.parts.uri.query(),
None,
None,
Some(&payload.body),
)
.ok_or(AntigravityV1InternalRequestError::UpstreamUrlUnavailable)?;
@@ -68,6 +68,8 @@ pub(crate) struct LocalExecutionCandidateAttempt {
pub(crate) struct LocalExecutionCandidateAttemptSource<'a> {
items: VecDeque<LocalExecutionCandidateAttemptSourceItem<'a>>,
skipped_provider_ids: BTreeSet<String>,
skipped_endpoint_ids: BTreeSet<String>,
skipped_credential_ids: BTreeSet<String>,
}
type DecorateSkippedCandidateFn<'a> = Arc<
@@ -80,6 +82,10 @@ pub(crate) trait LocalExecutionAttemptSource<T>: Send {
async fn drain_execution_attempts(&mut self) -> Result<Vec<T>, GatewayError>;
async fn skip_credential(&mut self, key_id: &str) -> Result<(), GatewayError>;
async fn skip_endpoint(&mut self, endpoint_id: &str) -> Result<(), GatewayError>;
async fn skip_provider(&mut self, provider_id: &str) -> Result<(), GatewayError>;
}
@@ -111,6 +117,8 @@ impl<'a> LocalExecutionCandidateAttemptSource<'a> {
Self {
items,
skipped_provider_ids: BTreeSet::new(),
skipped_endpoint_ids: BTreeSet::new(),
skipped_credential_ids: BTreeSet::new(),
}
}
@@ -123,7 +131,12 @@ impl<'a> LocalExecutionCandidateAttemptSource<'a> {
};
match front {
LocalExecutionCandidateAttemptSourceItem::Static { attempts } => {
if dispatch_sequence_provider_is_skipped(attempts, &self.skipped_provider_ids) {
if dispatch_sequence_candidate_is_skipped(
attempts,
&self.skipped_provider_ids,
&self.skipped_endpoint_ids,
&self.skipped_credential_ids,
) {
self.items.pop_front();
continue;
}
@@ -141,10 +154,20 @@ impl<'a> LocalExecutionCandidateAttemptSource<'a> {
pending_attempts,
pool_exhaustion_persistence,
} => {
if self.skipped_provider_ids.contains(cursor.provider_id()) {
if self.skipped_provider_ids.contains(cursor.provider_id())
|| self.skipped_endpoint_ids.contains(cursor.endpoint_id())
{
self.items.pop_front();
continue;
}
if dispatch_sequence_candidate_is_skipped(
pending_attempts,
&self.skipped_provider_ids,
&self.skipped_endpoint_ids,
&self.skipped_credential_ids,
) {
*pending_attempts = DispatchSequence::new(Vec::new());
}
if let Some(attempt) = next_attempt_from_dispatch_sequence(pending_attempts) {
return Ok(Some(attempt));
}
@@ -162,6 +185,14 @@ impl<'a> LocalExecutionCandidateAttemptSource<'a> {
self.items.pop_front();
continue;
};
if candidate_is_skipped(
&candidate,
&self.skipped_provider_ids,
&self.skipped_endpoint_ids,
&self.skipped_credential_ids,
) {
continue;
}
*pending_attempts = dispatch_sequence_from_attempts(
build_unpersisted_local_execution_candidate_attempts(
candidate,
@@ -174,6 +205,12 @@ impl<'a> LocalExecutionCandidateAttemptSource<'a> {
for provider_id in &self.skipped_provider_ids {
cursor.skip_provider(provider_id);
}
for endpoint_id in &self.skipped_endpoint_ids {
cursor.skip_endpoint(endpoint_id);
}
for key_id in &self.skipped_credential_ids {
cursor.skip_credential(key_id);
}
let Some(attempt) = cursor.next_attempt().await? else {
self.items.pop_front();
continue;
@@ -201,6 +238,32 @@ impl<'a> LocalExecutionCandidateAttemptSource<'a> {
}
}
}
pub(crate) fn skip_endpoint(&mut self, endpoint_id: &str) {
let endpoint_id = endpoint_id.trim();
if endpoint_id.is_empty() {
return;
}
self.skipped_endpoint_ids.insert(endpoint_id.to_string());
for item in &mut self.items {
if let LocalExecutionCandidateAttemptSourceItem::RequestedModelPage { cursor } = item {
cursor.skip_endpoint(endpoint_id);
}
}
}
pub(crate) fn skip_credential(&mut self, key_id: &str) {
let key_id = key_id.trim();
if key_id.is_empty() {
return;
}
self.skipped_credential_ids.insert(key_id.to_string());
for item in &mut self.items {
if let LocalExecutionCandidateAttemptSourceItem::RequestedModelPage { cursor } = item {
cursor.skip_credential(key_id);
}
}
}
}
impl LocalExecutionCandidateAttempt {
@@ -668,6 +731,8 @@ where
LocalExecutionCandidateAttemptSource {
items,
skipped_provider_ids: BTreeSet::new(),
skipped_endpoint_ids: BTreeSet::new(),
skipped_credential_ids: BTreeSet::new(),
},
candidate_count,
)
@@ -802,6 +867,8 @@ where
page_cursor,
pending_items: VecDeque::new(),
skipped_provider_ids: BTreeSet::new(),
skipped_endpoint_ids: BTreeSet::new(),
skipped_credential_ids: BTreeSet::new(),
candidate_count: 0,
next_candidate_index: 0,
remembered_affinity: false,
@@ -825,6 +892,8 @@ where
LocalExecutionCandidateAttemptSource {
items,
skipped_provider_ids: BTreeSet::new(),
skipped_endpoint_ids: BTreeSet::new(),
skipped_credential_ids: BTreeSet::new(),
},
candidate_count,
)
@@ -851,6 +920,8 @@ struct RequestedModelAttemptPageCursor<'a> {
page_cursor: LocalCandidatePreselectionPageCursor<'a>,
pending_items: VecDeque<LocalExecutionCandidateAttemptSourceItem<'a>>,
skipped_provider_ids: BTreeSet<String>,
skipped_endpoint_ids: BTreeSet<String>,
skipped_credential_ids: BTreeSet<String>,
candidate_count: usize,
next_candidate_index: u32,
remembered_affinity: bool,
@@ -864,6 +935,14 @@ impl<'a> RequestedModelAttemptPageCursor<'a> {
self.skipped_provider_ids.insert(provider_id.to_string());
}
fn skip_endpoint(&mut self, endpoint_id: &str) {
self.skipped_endpoint_ids.insert(endpoint_id.to_string());
}
fn skip_credential(&mut self, key_id: &str) {
self.skipped_credential_ids.insert(key_id.to_string());
}
async fn next_attempt(
&mut self,
) -> Result<Option<LocalExecutionCandidateAttempt>, GatewayError> {
@@ -871,8 +950,13 @@ impl<'a> RequestedModelAttemptPageCursor<'a> {
return Err(error);
}
loop {
if let Some(attempt) =
pop_attempt_from_items(&mut self.pending_items, &self.skipped_provider_ids).await
if let Some(attempt) = pop_attempt_from_items(
&mut self.pending_items,
&self.skipped_provider_ids,
&self.skipped_endpoint_ids,
&self.skipped_credential_ids,
)
.await
{
return Ok(Some(attempt));
}
@@ -1075,12 +1159,19 @@ fn page_is_exact_auth_api_key_concurrency_limited(
async fn pop_attempt_from_items(
items: &mut VecDeque<LocalExecutionCandidateAttemptSourceItem<'_>>,
skipped_provider_ids: &BTreeSet<String>,
skipped_endpoint_ids: &BTreeSet<String>,
skipped_credential_ids: &BTreeSet<String>,
) -> Option<LocalExecutionCandidateAttempt> {
loop {
let front = items.front_mut()?;
match front {
LocalExecutionCandidateAttemptSourceItem::Static { attempts } => {
if dispatch_sequence_provider_is_skipped(attempts, skipped_provider_ids) {
if dispatch_sequence_candidate_is_skipped(
attempts,
skipped_provider_ids,
skipped_endpoint_ids,
skipped_credential_ids,
) {
items.pop_front();
continue;
}
@@ -1098,10 +1189,20 @@ async fn pop_attempt_from_items(
pending_attempts,
pool_exhaustion_persistence,
} => {
if skipped_provider_ids.contains(cursor.provider_id()) {
if skipped_provider_ids.contains(cursor.provider_id())
|| skipped_endpoint_ids.contains(cursor.endpoint_id())
{
items.pop_front();
continue;
}
if dispatch_sequence_candidate_is_skipped(
pending_attempts,
skipped_provider_ids,
skipped_endpoint_ids,
skipped_credential_ids,
) {
*pending_attempts = DispatchSequence::new(Vec::new());
}
if let Some(attempt) = next_attempt_from_dispatch_sequence(pending_attempts) {
return Some(attempt);
}
@@ -1119,6 +1220,14 @@ async fn pop_attempt_from_items(
items.pop_front();
continue;
};
if candidate_is_skipped(
&candidate,
skipped_provider_ids,
skipped_endpoint_ids,
skipped_credential_ids,
) {
continue;
}
*pending_attempts = dispatch_sequence_from_attempts(
build_unpersisted_local_execution_candidate_attempts(
candidate,
@@ -1780,15 +1889,33 @@ fn next_attempt_from_dispatch_sequence(
Some(attempt)
}
fn dispatch_sequence_provider_is_skipped(
fn dispatch_sequence_candidate_is_skipped(
sequence: &DispatchSequence<LocalExecutionCandidateAttempt>,
skipped_provider_ids: &BTreeSet<String>,
skipped_endpoint_ids: &BTreeSet<String>,
skipped_credential_ids: &BTreeSet<String>,
) -> bool {
sequence.peek_current().is_some_and(|item| {
skipped_provider_ids.contains(&item.candidate.eligible.candidate.provider_id)
candidate_is_skipped(
&item.candidate.eligible,
skipped_provider_ids,
skipped_endpoint_ids,
skipped_credential_ids,
)
})
}
fn candidate_is_skipped(
candidate: &EligibleLocalExecutionCandidate,
skipped_provider_ids: &BTreeSet<String>,
skipped_endpoint_ids: &BTreeSet<String>,
skipped_credential_ids: &BTreeSet<String>,
) -> bool {
skipped_provider_ids.contains(&candidate.candidate.provider_id)
|| skipped_endpoint_ids.contains(&candidate.candidate.endpoint_id)
|| skipped_credential_ids.contains(&candidate.candidate.key_id)
}
fn dispatch_sequence_exhausted(
sequence: &mut DispatchSequence<LocalExecutionCandidateAttempt>,
) -> bool {
@@ -2341,6 +2468,8 @@ mod tests {
page_cursor,
pending_items: VecDeque::new(),
skipped_provider_ids: BTreeSet::new(),
skipped_endpoint_ids: BTreeSet::new(),
skipped_credential_ids: BTreeSet::new(),
candidate_count: 0,
next_candidate_index: 0,
remembered_affinity: false,
@@ -2436,6 +2565,8 @@ mod tests {
page_cursor,
pending_items: VecDeque::new(),
skipped_provider_ids: BTreeSet::new(),
skipped_endpoint_ids: BTreeSet::new(),
skipped_credential_ids: BTreeSet::new(),
candidate_count: 0,
next_candidate_index: 0,
remembered_affinity: false,
@@ -2634,6 +2765,8 @@ mod tests {
),
}]),
skipped_provider_ids: BTreeSet::new(),
skipped_endpoint_ids: BTreeSet::new(),
skipped_credential_ids: BTreeSet::new(),
};
let first = source
@@ -2652,6 +2785,111 @@ mod tests {
.is_none());
}
#[tokio::test]
async fn dynamic_attempt_source_skips_credentials_and_endpoints_across_static_candidates() {
let key_a = sample_eligible("key-a", None);
let key_b = sample_eligible("key-b", None);
let mut key_c = sample_eligible("key-c", None);
key_c.candidate.endpoint_id = "endpoint-2".to_string();
Arc::make_mut(&mut key_c.transport).endpoint.id = "endpoint-2".to_string();
let static_item =
|candidate, candidate_index| LocalExecutionCandidateAttemptSourceItem::Static {
attempts: dispatch_sequence_from_attempts(
build_unpersisted_local_execution_candidate_attempts(
candidate,
candidate_index,
)
.into(),
),
};
let mut source = LocalExecutionCandidateAttemptSource {
items: VecDeque::from([
static_item(key_a, 0),
static_item(key_b, 1),
static_item(key_c, 2),
]),
skipped_provider_ids: BTreeSet::new(),
skipped_endpoint_ids: BTreeSet::new(),
skipped_credential_ids: BTreeSet::new(),
};
source.skip_credential("key-a");
let key_b_attempt = source
.next_attempt()
.await
.expect("candidate source should succeed")
.expect("a different credential should remain");
assert_eq!(key_b_attempt.eligible.candidate.key_id, "key-b");
source.skip_endpoint("endpoint-1");
let endpoint_2_attempt = source
.next_attempt()
.await
.expect("candidate source should succeed")
.expect("a different endpoint should remain");
assert_eq!(endpoint_2_attempt.eligible.candidate.key_id, "key-c");
assert_eq!(
endpoint_2_attempt.eligible.candidate.endpoint_id,
"endpoint-2"
);
}
#[tokio::test]
async fn dynamic_attempt_source_filters_skipped_pool_pending_credential() {
let app = AppState::new().expect("state should build");
let mut pool_group = sample_eligible("pool-group", None);
pool_group.kind = LocalExecutionCandidateKind::PoolGroup;
pool_group.transport = sample_transport("pool-group", Some(json!({ "pool_advanced": {} })));
let pool_cursor = PoolKeyCursor::new(
PlannerAppState::new(&app),
pool_group,
None,
Some("gpt-5"),
None,
);
let pool_key_attempts = dispatch_sequence_from_attempts(
build_unpersisted_local_execution_candidate_attempts(
sample_eligible("pool-key-a", None),
0,
)
.into(),
);
let mut fallback = sample_eligible("fallback-key", None);
fallback.candidate.provider_id = "provider-b".to_string();
Arc::make_mut(&mut fallback.transport).provider.id = "provider-b".to_string();
Arc::make_mut(&mut fallback.transport).key.provider_id = "provider-b".to_string();
let fallback_attempts = dispatch_sequence_from_attempts(
build_unpersisted_local_execution_candidate_attempts(fallback, 1).into(),
);
let mut source = LocalExecutionCandidateAttemptSource {
items: VecDeque::from([
LocalExecutionCandidateAttemptSourceItem::Pool {
cursor: pool_cursor,
candidate_index: 0,
pending_attempts: pool_key_attempts,
pool_exhaustion_persistence: None,
},
LocalExecutionCandidateAttemptSourceItem::Static {
attempts: fallback_attempts,
},
]),
skipped_provider_ids: BTreeSet::new(),
skipped_endpoint_ids: BTreeSet::new(),
skipped_credential_ids: BTreeSet::new(),
};
source.skip_credential("pool-key-a");
let attempt = source
.next_attempt()
.await
.expect("candidate source should succeed")
.expect("fallback credential should remain");
assert_eq!(attempt.eligible.candidate.provider_id, "provider-b");
assert_eq!(attempt.eligible.candidate.key_id, "fallback-key");
}
#[tokio::test]
async fn skipped_provider_discards_pool_cursor_and_continues_with_next_provider() {
let app = AppState::new().expect("state should build");
@@ -2686,6 +2924,8 @@ mod tests {
},
]),
skipped_provider_ids: BTreeSet::new(),
skipped_endpoint_ids: BTreeSet::new(),
skipped_credential_ids: BTreeSet::new(),
};
source.skip_provider("provider-1");
@@ -2752,6 +2992,8 @@ mod tests {
pool_exhaustion_persistence: Some(pool_exhaustion_persistence),
}]),
skipped_provider_ids: BTreeSet::new(),
skipped_endpoint_ids: BTreeSet::new(),
skipped_credential_ids: BTreeSet::new(),
};
assert!(source
@@ -10,7 +10,7 @@ use crate::ai_serving::{
};
pub(crate) use crate::ai_serving::{
CLAUDE_CHAT_STREAM_PLAN_KIND, CLAUDE_CHAT_SYNC_PLAN_KIND, CLAUDE_CLI_STREAM_PLAN_KIND,
CLAUDE_CLI_SYNC_PLAN_KIND, EXECUTION_RUNTIME_STREAM_ACTION,
CLAUDE_CLI_SYNC_PLAN_KIND, CLAUDE_COUNT_TOKENS_SYNC_PLAN_KIND, EXECUTION_RUNTIME_STREAM_ACTION,
EXECUTION_RUNTIME_STREAM_DECISION_ACTION, EXECUTION_RUNTIME_SYNC_ACTION,
EXECUTION_RUNTIME_SYNC_DECISION_ACTION, GEMINI_CHAT_STREAM_PLAN_KIND,
GEMINI_CHAT_SYNC_PLAN_KIND, GEMINI_CLI_STREAM_PLAN_KIND, GEMINI_CLI_SYNC_PLAN_KIND,
@@ -1,17 +1,17 @@
use crate::ai_serving::planner::common::{
CLAUDE_CHAT_STREAM_PLAN_KIND, CLAUDE_CHAT_SYNC_PLAN_KIND, CLAUDE_CLI_STREAM_PLAN_KIND,
CLAUDE_CLI_SYNC_PLAN_KIND, GEMINI_CHAT_STREAM_PLAN_KIND, GEMINI_CHAT_SYNC_PLAN_KIND,
GEMINI_CLI_STREAM_PLAN_KIND, GEMINI_CLI_SYNC_PLAN_KIND, GEMINI_EMBEDDING_SYNC_PLAN_KIND,
GEMINI_FILES_DELETE_PLAN_KIND, GEMINI_FILES_DOWNLOAD_PLAN_KIND, GEMINI_FILES_GET_PLAN_KIND,
GEMINI_FILES_LIST_PLAN_KIND, GEMINI_VIDEO_CANCEL_SYNC_PLAN_KIND,
GEMINI_VIDEO_CREATE_SYNC_PLAN_KIND, OPENAI_CHAT_STREAM_PLAN_KIND, OPENAI_CHAT_SYNC_PLAN_KIND,
OPENAI_EMBEDDING_SYNC_PLAN_KIND, OPENAI_IMAGE_STREAM_PLAN_KIND, OPENAI_IMAGE_SYNC_PLAN_KIND,
OPENAI_RERANK_SYNC_PLAN_KIND, OPENAI_RESPONSES_COMPACT_STREAM_PLAN_KIND,
OPENAI_RESPONSES_COMPACT_SYNC_PLAN_KIND, OPENAI_RESPONSES_STREAM_PLAN_KIND,
OPENAI_RESPONSES_SYNC_PLAN_KIND, OPENAI_SEARCH_SYNC_PLAN_KIND,
OPENAI_VIDEO_CANCEL_SYNC_PLAN_KIND, OPENAI_VIDEO_CONTENT_PLAN_KIND,
OPENAI_VIDEO_CREATE_SYNC_PLAN_KIND, OPENAI_VIDEO_DELETE_SYNC_PLAN_KIND,
OPENAI_VIDEO_REMIX_SYNC_PLAN_KIND,
CLAUDE_CLI_SYNC_PLAN_KIND, CLAUDE_COUNT_TOKENS_SYNC_PLAN_KIND, GEMINI_CHAT_STREAM_PLAN_KIND,
GEMINI_CHAT_SYNC_PLAN_KIND, GEMINI_CLI_STREAM_PLAN_KIND, GEMINI_CLI_SYNC_PLAN_KIND,
GEMINI_EMBEDDING_SYNC_PLAN_KIND, GEMINI_FILES_DELETE_PLAN_KIND,
GEMINI_FILES_DOWNLOAD_PLAN_KIND, GEMINI_FILES_GET_PLAN_KIND, GEMINI_FILES_LIST_PLAN_KIND,
GEMINI_VIDEO_CANCEL_SYNC_PLAN_KIND, GEMINI_VIDEO_CREATE_SYNC_PLAN_KIND,
OPENAI_CHAT_STREAM_PLAN_KIND, OPENAI_CHAT_SYNC_PLAN_KIND, OPENAI_EMBEDDING_SYNC_PLAN_KIND,
OPENAI_IMAGE_STREAM_PLAN_KIND, OPENAI_IMAGE_SYNC_PLAN_KIND, OPENAI_RERANK_SYNC_PLAN_KIND,
OPENAI_RESPONSES_COMPACT_STREAM_PLAN_KIND, OPENAI_RESPONSES_COMPACT_SYNC_PLAN_KIND,
OPENAI_RESPONSES_STREAM_PLAN_KIND, OPENAI_RESPONSES_SYNC_PLAN_KIND,
OPENAI_SEARCH_SYNC_PLAN_KIND, OPENAI_VIDEO_CANCEL_SYNC_PLAN_KIND,
OPENAI_VIDEO_CONTENT_PLAN_KIND, OPENAI_VIDEO_CREATE_SYNC_PLAN_KIND,
OPENAI_VIDEO_DELETE_SYNC_PLAN_KIND, OPENAI_VIDEO_REMIX_SYNC_PLAN_KIND,
};
use crate::ai_serving::planner::plan_builders::{
build_gemini_stream_plan_from_decision, build_gemini_sync_plan_from_decision,
@@ -109,6 +109,7 @@ fn build_sync_plan_payload_from_decision(
}
CLAUDE_CHAT_SYNC_PLAN_KIND
| CLAUDE_CLI_SYNC_PLAN_KIND
| CLAUDE_COUNT_TOKENS_SYNC_PLAN_KIND
| OPENAI_EMBEDDING_SYNC_PLAN_KIND
| OPENAI_RERANK_SYNC_PLAN_KIND => {
build_standard_sync_plan_from_decision(parts, body_json, payload)?
@@ -15,8 +15,9 @@ use tracing::warn;
use crate::ai_serving::planner::common::extract_standard_requested_model;
use crate::ai_serving::{
ExecutionRuntimeAuthContext, GatewayAuthApiKeySnapshot, GatewayProviderTransportSnapshot,
PlannerAppState, CODEX_RESPONSES_LITE_HEADER,
ClientSurface, ExecutionRuntimeAuthContext, GatewayAuthApiKeySnapshot,
GatewayCredentialCarrier, GatewayProviderTransportSnapshot, PlannerAppState,
CODEX_RESPONSES_LITE_HEADER,
};
use crate::cache::CacheLoadObserver;
use crate::client_session_affinity::client_session_affinity_from_api_request;
@@ -52,6 +53,8 @@ pub(crate) struct LocalRequestedModelDecisionInput {
pub(crate) auth_snapshot: GatewayAuthApiKeySnapshot,
pub(crate) required_capabilities: Option<serde_json::Value>,
pub(crate) request_auth_channel: Option<String>,
pub(crate) client_surface: Option<ClientSurface>,
pub(crate) gateway_credential_carrier: Option<GatewayCredentialCarrier>,
pub(crate) client_session_affinity: Option<ClientSessionAffinity>,
pub(crate) routing_policy: Option<ResolvedRoutingPolicy>,
pub(crate) routing_trace_seed: Option<RoutingDecisionTrace>,
@@ -378,6 +381,8 @@ pub(crate) fn build_local_requested_model_decision_input(
auth_snapshot: resolved_input.auth_snapshot,
required_capabilities: resolved_input.required_capabilities,
request_auth_channel: None,
client_surface: None,
gateway_credential_carrier: None,
client_session_affinity: None,
routing_policy: None,
routing_trace_seed: None,
@@ -1128,6 +1133,8 @@ mod tests {
auth_snapshot: sample_auth_snapshot(),
required_capabilities: None,
request_auth_channel: None,
client_surface: None,
gateway_credential_carrier: None,
client_session_affinity: None,
routing_policy: None,
routing_trace_seed: None,
@@ -1323,6 +1330,8 @@ mod tests {
auth_snapshot: sample_auth_snapshot(),
required_capabilities: None,
request_auth_channel: None,
client_surface: None,
gateway_credential_carrier: None,
client_session_affinity: None,
routing_policy: None,
routing_trace_seed: None,
@@ -1390,6 +1399,8 @@ mod tests {
auth_snapshot: sample_auth_snapshot(),
required_capabilities: None,
request_auth_channel: None,
client_surface: None,
gateway_credential_carrier: None,
client_session_affinity: None,
routing_policy: None,
routing_trace_seed: None,
@@ -1459,6 +1470,59 @@ mod tests {
);
}
#[test]
fn provider_request_routing_policy_cannot_restore_credentials_or_aether_internal_headers() {
for header_name in [
"authorization",
"proxy-authorization",
"api-key",
"x-api-key",
"x-goog-api-key",
"cookie",
"cookie2",
"set-cookie",
"x-aether-auth-user-id",
"x-aether-control-future",
] {
let mut input = sample_decision_input();
set_provider_request_rules(
&mut input,
&["gpt-5"],
json!([{
"type": "patch_headers",
"patch": [{
"op": "set",
"name": header_name,
"value": "must-not-reach-upstream"
}]
}]),
);
let mut decision = sample_decision();
let error =
apply_provider_request_routing_policy_to_decision(&input, &mut decision, None)
.expect_err("reserved provider header mutation should fail closed");
assert!(
matches!(
&error,
GatewayError::Client {
status: StatusCode::BAD_REQUEST,
..
}
),
"unexpected error for {header_name}: {error:?}"
);
assert!(
!decision
.provider_request_headers
.keys()
.any(|name| name.eq_ignore_ascii_case(header_name)),
"reserved header reached the provider decision: {header_name}"
);
}
}
#[test]
fn codex_prompt_cache_identity_headers_are_terminal_after_routing_mutations() {
let mut input = sample_decision_input();
@@ -59,6 +59,7 @@ pub(crate) async fn build_gemini_cli_v1internal_provider_request(
input.upstream_is_stream,
input.parts.uri.query(),
None,
None,
Some(&payload.body),
)
.ok_or(GeminiCliV1InternalRequestError::UpstreamUrlUnavailable)?;
@@ -81,6 +81,8 @@ pub(crate) async fn resolve_local_same_format_provider_decision_input(
let mut input = build_local_requested_model_decision_input(resolved_input, requested_model);
input.request_auth_channel = decision.request_auth_channel.clone();
input.client_surface = decision.client_surface;
input.gateway_credential_carrier = decision.gateway_credential_carrier;
input.client_session_affinity = client_session_affinity_from_api_request(
spec_metadata.api_format,
&parts.headers,
@@ -128,7 +130,7 @@ pub(crate) async fn materialize_local_same_format_provider_candidate_attempts(
.base_model()
.unwrap_or(&input.requested_model);
let (candidates, preselection_skipped) = planner_state
.list_selectable_candidates_with_skip_reasons(
.list_selectable_candidates_with_skip_reasons_for_request_operation(
spec_metadata.api_format,
routing_model,
spec_metadata.require_streaming,
@@ -137,6 +139,7 @@ pub(crate) async fn materialize_local_same_format_provider_candidate_attempts(
input.client_session_affinity.as_ref(),
current_unix_secs(),
false,
spec.operation.map(|operation| operation.as_str()),
)
.await?;
let outcome = materialize_local_execution_candidates_with_serving(
@@ -232,7 +235,7 @@ pub(crate) async fn build_local_same_format_provider_candidate_attempt_source<'a
.base_model()
.unwrap_or(&input.requested_model);
let (candidates, preselection_skipped) = planner_state
.list_selectable_candidates_with_skip_reasons(
.list_selectable_candidates_with_skip_reasons_for_request_operation(
spec_metadata.api_format,
routing_model,
spec_metadata.require_streaming,
@@ -241,6 +244,7 @@ pub(crate) async fn build_local_same_format_provider_candidate_attempt_source<'a
input.client_session_affinity.as_ref(),
current_unix_secs(),
false,
spec.operation.map(|operation| operation.as_str()),
)
.await?;
@@ -1,5 +1,8 @@
use serde_json::json;
use aether_ai_serving::{AdaptationMode, AiRequestGzipPolicy, OriginalRequestPayload};
use aether_contracts::{ExecutionResponseBodyMode, EXECUTION_RESPONSE_BODY_MODE_HEADER};
use crate::ai_serving::ai_local_execution_contract_for_formats;
use crate::ai_serving::build_request_trace_proxy_value;
use crate::ai_serving::planner::candidate_materialization::{
@@ -61,6 +64,8 @@ pub(crate) async fn maybe_build_local_same_format_provider_decision_payload_for_
else {
return Ok(None);
};
let request_redacted = resolved.request_redacted;
let compatibility_edits_empty = resolved.compatibility_edits.is_empty();
let original_request_body_json = if resolved.request_redacted {
Some(&resolved.provider_request_body)
} else {
@@ -82,6 +87,51 @@ pub(crate) async fn maybe_build_local_same_format_provider_decision_payload_for_
.clone()
.or_else(|| resolve_transport_profile(&resolved.transport));
let mut extra_fields = serde_json::Map::new();
extra_fields.insert(
"provider_type".to_string(),
json!(resolved.transport.provider.provider_type.as_str()),
);
if let Some(operation) = spec.operation {
extra_fields.insert("api_operation".to_string(), json!(operation.as_str()));
}
if let Some(client_surface) = input.client_surface {
extra_fields.insert("client_surface".to_string(), json!(client_surface.as_str()));
}
if let Some(carrier) = input.gateway_credential_carrier {
extra_fields.insert(
"gateway_credential_carrier".to_string(),
json!(carrier.as_str()),
);
}
extra_fields.insert(
"upstream_credential_mode".to_string(),
json!(resolved.transport.key.auth_type.trim().to_ascii_lowercase()),
);
let mut adaptation_mode = if resolved.compatibility_edits.is_empty() {
AdaptationMode::NativeTransparent
} else {
AdaptationMode::SameFormatCompat
};
if crate::ai_serving::normalize_api_format_alias(&resolved.provider_api_format)
== "claude:messages"
{
let compatibility_profile =
crate::ai_serving::transport::resolve_anthropic_compatibility_profile(
&resolved.transport,
&resolved.provider_api_format,
);
extra_fields.insert(
"anthropic_compatibility_profile".to_string(),
json!(compatibility_profile.as_str()),
);
if compatibility_profile.uses_claude_code_compatibility() {
adaptation_mode = AdaptationMode::SameFormatCompat;
}
}
extra_fields.insert(
"adaptation_mode".to_string(),
json!(adaptation_mode.as_str()),
);
if let Some(proxy_value) =
build_request_trace_proxy_value(Some(&resolved.transport), proxy.as_ref())
{
@@ -227,9 +277,79 @@ pub(crate) async fn maybe_build_local_same_format_provider_decision_payload_for_
&mut decision,
Some(transport.as_ref()),
)?;
enforce_provider_api_operation_invariants(
spec.operation,
decision.provider_request_body.as_mut(),
&mut decision.provider_request_headers,
);
decision.provider_request_body_base64 = original_request_body_base64(
parts,
decision.provider_request_body.as_ref(),
adaptation_mode,
request_redacted,
compatibility_edits_empty,
decision.content_encoding.as_deref(),
decision.request_gzip.as_ref(),
);
decision
.provider_request_headers
.retain(|name, _| !name.eq_ignore_ascii_case(EXECUTION_RESPONSE_BODY_MODE_HEADER));
if !spec_metadata.require_streaming && decision.provider_request_body_base64.is_some() {
decision.provider_request_headers.insert(
EXECUTION_RESPONSE_BODY_MODE_HEADER.to_string(),
ExecutionResponseBodyMode::PreserveBytes
.as_str()
.to_string(),
);
}
Ok(Some(decision))
}
fn enforce_provider_api_operation_invariants(
operation: Option<crate::ai_serving::ApiOperation>,
provider_request_body: Option<&mut serde_json::Value>,
provider_request_headers: &mut std::collections::BTreeMap<String, String>,
) {
if operation != Some(crate::ai_serving::ApiOperation::ClaudeCountTokens) {
return;
}
if let Some(provider_request_body) = provider_request_body {
crate::ai_serving::transport::enforce_same_format_provider_api_operation_body_policy(
provider_request_body,
operation,
);
}
for header_name in ["accept", "content-type"] {
provider_request_headers.retain(|name, _| !name.eq_ignore_ascii_case(header_name));
provider_request_headers.insert(header_name.to_string(), "application/json".to_string());
}
}
fn original_request_body_base64(
parts: &http::request::Parts,
provider_request_body: Option<&serde_json::Value>,
adaptation_mode: AdaptationMode,
request_redacted: bool,
compatibility_edits_empty: bool,
content_encoding: Option<&str>,
request_gzip: Option<&AiRequestGzipPolicy>,
) -> Option<String> {
if adaptation_mode != AdaptationMode::NativeTransparent
|| request_redacted
|| !compatibility_edits_empty
|| content_encoding.is_some_and(|value| !value.trim().is_empty())
|| request_gzip.is_some_and(|policy| policy.enabled != Some(false))
{
return None;
}
parts
.extensions
.get::<OriginalRequestPayload>()?
.body_bytes_base64_if_unchanged(provider_request_body?)
}
pub(super) async fn mark_skipped_local_same_format_provider_candidate(
state: &AppState,
input: &LocalSameFormatProviderDecisionInput,
@@ -313,3 +433,177 @@ pub(super) async fn mark_skipped_local_same_format_provider_candidate_with_failu
)
.await;
}
#[cfg(test)]
mod tests {
use std::collections::BTreeMap;
use base64::Engine as _;
use super::{
enforce_provider_api_operation_invariants, original_request_body_base64, AdaptationMode,
AiRequestGzipPolicy, OriginalRequestPayload,
};
use crate::ai_serving::ApiOperation;
fn request_parts_with_original_payload(
body_json: serde_json::Value,
body_bytes: &[u8],
) -> http::request::Parts {
let (mut parts, ()) = http::Request::new(()).into_parts();
parts
.extensions
.insert(OriginalRequestPayload::from_parsed_json(
body_json, body_bytes,
));
parts
}
#[test]
fn count_tokens_invariants_win_after_provider_routing_mutations() {
let mut body = serde_json::json!({
"model": "claude-sonnet-4",
"messages": [],
"stream": true
});
let mut headers = BTreeMap::from([
("Accept".to_string(), "text/event-stream".to_string()),
("Content-Type".to_string(), "text/plain".to_string()),
("x-provider-route".to_string(), "kept".to_string()),
]);
enforce_provider_api_operation_invariants(
Some(ApiOperation::ClaudeCountTokens),
Some(&mut body),
&mut headers,
);
assert!(body.get("stream").is_none());
assert_eq!(
headers.get("accept").map(String::as_str),
Some("application/json")
);
assert_eq!(
headers.get("content-type").map(String::as_str),
Some("application/json")
);
assert_eq!(
headers.get("x-provider-route").map(String::as_str),
Some("kept")
);
assert_eq!(
headers
.keys()
.filter(|name| name.eq_ignore_ascii_case("accept"))
.count(),
1
);
assert_eq!(
headers
.keys()
.filter(|name| name.eq_ignore_ascii_case("content-type"))
.count(),
1
);
}
#[test]
fn unchanged_same_format_body_preserves_original_json_bytes() {
let raw = br#"{ "unknown": {"enabled":true}, "messages": [], "model": "claude-sonnet-4" }"#;
let body_json: serde_json::Value = serde_json::from_slice(raw).expect("body should parse");
let parts = request_parts_with_original_payload(body_json.clone(), raw);
let encoded = original_request_body_base64(
&parts,
Some(&body_json),
AdaptationMode::NativeTransparent,
false,
true,
None,
None,
)
.expect("unchanged request should retain exact bytes");
assert_eq!(
base64::engine::general_purpose::STANDARD
.decode(encoded)
.expect("body should decode"),
raw
);
}
#[test]
fn request_edits_or_encoding_disable_original_json_bytes() {
let raw = br#"{"model":"claude-sonnet-4","messages":[]}"#;
let body_json: serde_json::Value = serde_json::from_slice(raw).expect("body should parse");
let parts = request_parts_with_original_payload(body_json.clone(), raw);
let changed_body = serde_json::json!({
"model": "claude-sonnet-4-5",
"messages": []
});
assert!(original_request_body_base64(
&parts,
Some(&changed_body),
AdaptationMode::NativeTransparent,
false,
true,
None,
None,
)
.is_none());
assert!(original_request_body_base64(
&parts,
Some(&body_json),
AdaptationMode::NativeTransparent,
true,
true,
None,
None,
)
.is_none());
assert!(original_request_body_base64(
&parts,
Some(&body_json),
AdaptationMode::NativeTransparent,
false,
false,
None,
None,
)
.is_none());
assert!(original_request_body_base64(
&parts,
Some(&body_json),
AdaptationMode::SameFormatCompat,
false,
true,
None,
None,
)
.is_none());
assert!(original_request_body_base64(
&parts,
Some(&body_json),
AdaptationMode::NativeTransparent,
false,
true,
Some("gzip"),
None,
)
.is_none());
assert!(original_request_body_base64(
&parts,
Some(&body_json),
AdaptationMode::NativeTransparent,
false,
true,
None,
Some(&AiRequestGzipPolicy {
enabled: Some(true),
min_bytes: Some(1),
}),
)
.is_none());
}
}
@@ -75,9 +75,10 @@ pub(crate) fn resolve_same_format_provider_transport_unsupported_reason_for_trac
decision_kind: "trace_candidate_metadata",
report_kind: Some("trace_candidate_metadata"),
},
None,
);
if !behavior.is_antigravity
&& !behavior.is_claude_code
&& !behavior.is_claude_code_transport
&& !behavior.is_gemini_cli
&& !behavior.is_vertex
&& !behavior.is_kiro
@@ -127,6 +128,23 @@ pub(crate) async fn resolve_local_same_format_provider_candidate_payload_parts(
spec: LocalSameFormatProviderSpec,
) -> Result<Option<LocalSameFormatProviderCandidatePayloadParts>, GatewayError> {
let candidate = &attempt.eligible.candidate;
if let Some(skip_reason) = same_format_provider_operation_skip_reason(
&attempt.eligible.transport,
attempt.eligible.provider_api_format.as_str(),
spec.operation,
) {
mark_skipped_local_same_format_provider_candidate(
state,
input,
trace_id,
candidate,
attempt.candidate_index,
&attempt.candidate_id,
skip_reason,
)
.await;
return Ok(None);
}
let Some(prepared) = prepare_local_same_format_provider_candidate(
state,
trace_id,
@@ -364,7 +382,7 @@ pub(crate) async fn resolve_local_same_format_provider_candidate_payload_parts(
} else {
None
};
let provider_request_body = if let Some(antigravity_auth) = antigravity_auth.as_ref() {
let mut provider_request_body = if let Some(antigravity_auth) = antigravity_auth.as_ref() {
match build_antigravity_safe_v1internal_request(
antigravity_auth,
trace_id,
@@ -424,6 +442,16 @@ pub(crate) async fn resolve_local_same_format_provider_candidate_payload_parts(
} else {
base_provider_request_body
};
if crate::ai_serving::transport::enforce_same_format_provider_api_operation_body_policy(
&mut provider_request_body,
spec.operation,
) {
compatibility_edits.push(SameFormatProviderCompatibilityEdit {
field: "stream".to_string(),
action: SameFormatProviderCompatibilityEditAction::RuntimeRewrite,
detail: "removed stream field for non-streaming API operation".to_string(),
});
}
let is_grok = prepared
.transport
@@ -490,10 +518,10 @@ pub(crate) async fn resolve_local_same_format_provider_candidate_payload_parts(
original_request_body: body_json,
header_rules: transport.endpoint.header_rules.as_ref(),
behavior: prepared.behavior,
api_operation: spec.operation,
auth_header: prepared.auth_header.as_deref(),
auth_value: prepared.auth_value.as_deref(),
extra_headers: &extra_headers,
key_fingerprint: transport.key.fingerprint.as_ref(),
kiro_auth_config: prepared.kiro_auth.as_ref().map(|auth| &auth.auth_config),
kiro_machine_id: prepared
.kiro_auth
@@ -564,3 +592,98 @@ pub(crate) async fn resolve_local_same_format_provider_candidate_payload_parts(
request_redacted: redaction.redacted,
}))
}
fn same_format_provider_operation_skip_reason(
transport: &GatewayProviderTransportSnapshot,
provider_api_format: &str,
operation: Option<crate::ai_serving::ApiOperation>,
) -> Option<&'static str> {
(!crate::ai_serving::transport::transport_supports_api_operation(
transport,
provider_api_format,
operation,
))
.then_some("transport_operation_unsupported")
}
#[cfg(test)]
mod tests {
use super::same_format_provider_operation_skip_reason;
use crate::ai_serving::transport::snapshot::{
GatewayProviderTransportEndpoint, GatewayProviderTransportKey,
GatewayProviderTransportProvider,
};
use crate::ai_serving::{ApiOperation, GatewayProviderTransportSnapshot};
fn private_adapter_transport(provider_type: &str) -> GatewayProviderTransportSnapshot {
GatewayProviderTransportSnapshot {
provider: GatewayProviderTransportProvider {
id: "provider-1".to_string(),
name: provider_type.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: "claude:messages".to_string(),
api_family: Some("claude".to_string()),
endpoint_kind: Some("chat".to_string()),
is_active: true,
base_url: "https://private.example".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: "oauth".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: String::new(),
decrypted_auth_config: None,
},
}
}
#[test]
fn private_adapter_count_tokens_is_rejected_by_pre_auth_operation_gate() {
for provider_type in ["kiro", "grok"] {
let transport = private_adapter_transport(provider_type);
assert_eq!(
same_format_provider_operation_skip_reason(
&transport,
"claude:messages",
Some(ApiOperation::ClaudeCountTokens),
),
Some("transport_operation_unsupported"),
"provider_type={provider_type}"
);
}
}
}
@@ -1,6 +1,6 @@
use crate::ai_serving::planner::spec_metadata::LocalExecutionSurfaceSpecMetadata;
use crate::ai_serving::transport::{
classify_same_format_provider_request_behavior as classify_same_format_provider_request_behavior_impl,
classify_same_format_provider_request_behavior_for_operation as classify_same_format_provider_request_behavior_impl,
resolve_same_format_provider_direct_auth as resolve_same_format_provider_direct_auth_impl,
same_format_provider_transport_supported as same_format_provider_transport_supported_impl,
same_format_provider_transport_unsupported_reason as same_format_provider_transport_unsupported_reason_impl,
@@ -15,6 +15,7 @@ pub(super) fn classify_same_format_provider_request_behavior(
transport: &GatewayProviderTransportSnapshot,
provider_api_format: &str,
spec_metadata: LocalExecutionSurfaceSpecMetadata,
api_operation: Option<crate::ai_serving::ApiOperation>,
) -> SameFormatProviderRequestBehavior {
classify_same_format_provider_request_behavior_impl(
transport,
@@ -25,6 +26,7 @@ pub(super) fn classify_same_format_provider_request_behavior(
.report_kind
.expect("same-format provider specs should declare report kind"),
},
api_operation,
)
}
@@ -59,6 +59,7 @@ pub(super) async fn prepare_local_same_format_provider_candidate(
&transport,
provider_api_format,
spec_metadata,
spec.operation,
);
if !same_format_provider_transport_supported(
@@ -214,6 +214,16 @@ impl LocalExecutionAttemptSource<AiSyncAttempt> for LocalSameFormatProviderSyncA
Ok(drained)
}
async fn skip_credential(&mut self, key_id: &str) -> Result<(), GatewayError> {
self.candidates.skip_credential(key_id);
Ok(())
}
async fn skip_endpoint(&mut self, endpoint_id: &str) -> Result<(), GatewayError> {
self.candidates.skip_endpoint(endpoint_id);
Ok(())
}
async fn skip_provider(&mut self, provider_id: &str) -> Result<(), GatewayError> {
self.candidates.skip_provider(provider_id);
Ok(())
@@ -249,6 +259,16 @@ impl LocalExecutionAttemptSource<AiStreamAttempt>
Ok(drained)
}
async fn skip_credential(&mut self, key_id: &str) -> Result<(), GatewayError> {
self.candidates.skip_credential(key_id);
Ok(())
}
async fn skip_endpoint(&mut self, endpoint_id: &str) -> Result<(), GatewayError> {
self.candidates.skip_endpoint(endpoint_id);
Ok(())
}
async fn skip_provider(&mut self, provider_id: &str) -> Result<(), GatewayError> {
self.candidates.skip_provider(provider_id);
Ok(())
@@ -24,6 +24,7 @@ pub(crate) fn build_same_format_upstream_url(
upstream_is_stream,
request_query: parts.uri.query(),
kiro_api_region: kiro_auth.map(|auth| auth.auth_config.effective_api_region()),
api_operation: spec.operation,
provider_request_body,
},
)
@@ -1,8 +1,8 @@
use crate::ai_serving::GatewayControlDecision;
use crate::ai_serving::{
is_matching_stream_http_request as is_matching_stream_http_request_impl,
resolve_execution_runtime_stream_plan_kind as resolve_execution_runtime_stream_plan_kind_impl,
resolve_execution_runtime_sync_plan_kind as resolve_execution_runtime_sync_plan_kind_impl,
resolve_execution_runtime_stream_plan_kind_with_client_surface as resolve_execution_runtime_stream_plan_kind_impl,
resolve_execution_runtime_sync_plan_kind_with_client_surface as resolve_execution_runtime_sync_plan_kind_impl,
supports_stream_execution_decision_kind as supports_stream_execution_decision_kind_impl,
supports_sync_execution_decision_kind as supports_sync_execution_decision_kind_impl,
};
@@ -11,28 +11,34 @@ pub(crate) fn resolve_execution_runtime_stream_plan_kind(
parts: &http::request::Parts,
decision: &GatewayControlDecision,
) -> Option<&'static str> {
resolve_execution_runtime_stream_plan_kind_impl(
let plan_kind = resolve_execution_runtime_stream_plan_kind_impl(
decision.route_class.as_deref(),
decision.route_family.as_deref(),
decision.route_kind.as_deref(),
decision.client_surface,
decision.request_auth_channel.as_deref(),
&parts.method,
parts.uri.path(),
)
)?;
crate::ai_serving::plan_kind_matches_api_operation(plan_kind, true, decision.api_operation)
.then_some(plan_kind)
}
pub(crate) fn resolve_execution_runtime_sync_plan_kind(
parts: &http::request::Parts,
decision: &GatewayControlDecision,
) -> Option<&'static str> {
resolve_execution_runtime_sync_plan_kind_impl(
let plan_kind = resolve_execution_runtime_sync_plan_kind_impl(
decision.route_class.as_deref(),
decision.route_family.as_deref(),
decision.route_kind.as_deref(),
decision.client_surface,
decision.request_auth_channel.as_deref(),
&parts.method,
parts.uri.path(),
)
)?;
crate::ai_serving::plan_kind_matches_api_operation(plan_kind, false, decision.api_operation)
.then_some(plan_kind)
}
pub(crate) fn is_matching_stream_request(
@@ -62,7 +68,7 @@ mod tests {
resolve_execution_runtime_sync_plan_kind, supports_stream_execution_decision_kind,
supports_sync_execution_decision_kind,
};
use crate::ai_serving::GatewayControlDecision;
use crate::ai_serving::{ApiOperation, ClientSurface, GatewayControlDecision};
fn sample_decision(route_family: &str, route_kind: &str) -> GatewayControlDecision {
GatewayControlDecision {
@@ -71,6 +77,9 @@ mod tests {
route_class: Some("ai_public".to_string()),
route_family: Some(route_family.to_string()),
route_kind: Some(route_kind.to_string()),
client_surface: None,
api_operation: None,
gateway_credential_carrier: None,
request_auth_channel: None,
auth_context: None,
admin_principal: None,
@@ -121,7 +130,9 @@ mod tests {
let (claude_parts, _) = claude_request.into_parts();
let claude_api_key = sample_decision_with_auth_channel("claude", "messages", "api_key");
let claude_bearer = sample_decision_with_auth_channel("claude", "messages", "bearer_like");
let mut claude_bearer =
sample_decision_with_auth_channel("claude", "messages", "bearer_like");
claude_bearer.client_surface = Some(ClientSurface::ClaudeCode);
assert_eq!(
resolve_execution_runtime_sync_plan_kind(&claude_parts, &claude_api_key),
Some("claude_chat_sync")
@@ -131,6 +142,13 @@ mod tests {
Some("claude_cli_stream")
);
let claude_sdk_bearer =
sample_decision_with_auth_channel("claude", "messages", "bearer_like");
assert_eq!(
resolve_execution_runtime_sync_plan_kind(&claude_parts, &claude_sdk_bearer),
Some("claude_chat_sync")
);
let gemini_request = Request::builder()
.method(Method::POST)
.uri("/v1beta/models/gemini-2.5-pro:generateContent")
@@ -152,6 +170,36 @@ mod tests {
);
}
#[test]
fn resolves_claude_count_tokens_as_native_sync_operation() {
let request = Request::builder()
.method(Method::POST)
.uri("/v1/messages/count_tokens")
.body(())
.expect("request should build");
let (parts, _) = request.into_parts();
let mut decision = sample_decision("claude", "count_tokens");
decision.api_operation = Some(ApiOperation::ClaudeCountTokens);
assert_eq!(
resolve_execution_runtime_sync_plan_kind(&parts, &decision),
Some("claude_count_tokens_sync")
);
assert!(supports_sync_execution_decision_kind(
"claude_count_tokens_sync"
));
decision.api_operation = Some(ApiOperation::ClaudeMessagesCreate);
assert_eq!(
resolve_execution_runtime_sync_plan_kind(&parts, &decision),
None
);
assert_eq!(
resolve_execution_runtime_stream_plan_kind(&parts, &decision),
None
);
}
#[test]
fn stream_matching_uses_surface_route_logic() {
let request = Request::builder()
@@ -194,6 +194,16 @@ impl LocalExecutionAttemptSource<AiSyncAttempt> for LocalGeminiFilesSyncAttemptS
Ok(drained)
}
async fn skip_credential(&mut self, key_id: &str) -> Result<(), GatewayError> {
self.candidates.skip_credential(key_id);
Ok(())
}
async fn skip_endpoint(&mut self, endpoint_id: &str) -> Result<(), GatewayError> {
self.candidates.skip_endpoint(endpoint_id);
Ok(())
}
async fn skip_provider(&mut self, provider_id: &str) -> Result<(), GatewayError> {
self.candidates.skip_provider(provider_id);
Ok(())
@@ -222,6 +232,16 @@ impl LocalExecutionAttemptSource<AiStreamAttempt> for LocalGeminiFilesStreamAtte
Ok(drained)
}
async fn skip_credential(&mut self, key_id: &str) -> Result<(), GatewayError> {
self.candidates.skip_credential(key_id);
Ok(())
}
async fn skip_endpoint(&mut self, endpoint_id: &str) -> Result<(), GatewayError> {
self.candidates.skip_endpoint(endpoint_id);
Ok(())
}
async fn skip_provider(&mut self, provider_id: &str) -> Result<(), GatewayError> {
self.candidates.skip_provider(provider_id);
Ok(())
@@ -272,6 +272,16 @@ impl LocalExecutionAttemptSource<AiSyncAttempt> for LocalOpenAiImageSyncAttemptS
Ok(drained)
}
async fn skip_credential(&mut self, key_id: &str) -> Result<(), GatewayError> {
self.candidates.skip_credential(key_id);
Ok(())
}
async fn skip_endpoint(&mut self, endpoint_id: &str) -> Result<(), GatewayError> {
self.candidates.skip_endpoint(endpoint_id);
Ok(())
}
async fn skip_provider(&mut self, provider_id: &str) -> Result<(), GatewayError> {
self.candidates.skip_provider(provider_id);
Ok(())
@@ -300,6 +310,16 @@ impl LocalExecutionAttemptSource<AiStreamAttempt> for LocalOpenAiImageStreamAtte
Ok(drained)
}
async fn skip_credential(&mut self, key_id: &str) -> Result<(), GatewayError> {
self.candidates.skip_credential(key_id);
Ok(())
}
async fn skip_endpoint(&mut self, endpoint_id: &str) -> Result<(), GatewayError> {
self.candidates.skip_endpoint(endpoint_id);
Ok(())
}
async fn skip_provider(&mut self, provider_id: &str) -> Result<(), GatewayError> {
self.candidates.skip_provider(provider_id);
Ok(())
@@ -124,6 +124,16 @@ impl LocalExecutionAttemptSource<AiSyncAttempt> for LocalVideoCreateSyncAttemptS
Ok(drained)
}
async fn skip_credential(&mut self, key_id: &str) -> Result<(), GatewayError> {
self.candidates.skip_credential(key_id);
Ok(())
}
async fn skip_endpoint(&mut self, endpoint_id: &str) -> Result<(), GatewayError> {
self.candidates.skip_endpoint(endpoint_id);
Ok(())
}
async fn skip_provider(&mut self, provider_id: &str) -> Result<(), GatewayError> {
self.candidates.skip_provider(provider_id);
Ok(())
@@ -202,6 +202,16 @@ impl LocalExecutionAttemptSource<AiSyncAttempt> for LocalStandardSyncAttemptSour
Ok(drained)
}
async fn skip_credential(&mut self, key_id: &str) -> Result<(), GatewayError> {
self.candidates.skip_credential(key_id);
Ok(())
}
async fn skip_endpoint(&mut self, endpoint_id: &str) -> Result<(), GatewayError> {
self.candidates.skip_endpoint(endpoint_id);
Ok(())
}
async fn skip_provider(&mut self, provider_id: &str) -> Result<(), GatewayError> {
self.candidates.skip_provider(provider_id);
Ok(())
@@ -235,6 +245,16 @@ impl LocalExecutionAttemptSource<AiStreamAttempt> for LocalStandardStreamAttempt
Ok(drained)
}
async fn skip_credential(&mut self, key_id: &str) -> Result<(), GatewayError> {
self.candidates.skip_credential(key_id);
Ok(())
}
async fn skip_endpoint(&mut self, endpoint_id: &str) -> Result<(), GatewayError> {
self.candidates.skip_endpoint(endpoint_id);
Ok(())
}
async fn skip_provider(&mut self, provider_id: &str) -> Result<(), GatewayError> {
self.candidates.skip_provider(provider_id);
Ok(())
@@ -373,6 +373,8 @@ mod tests {
auth_snapshot: sample_auth_snapshot(),
required_capabilities: None,
request_auth_channel: None,
client_surface: None,
gateway_credential_carrier: None,
client_session_affinity: None,
routing_policy: None,
routing_trace_seed: None,
@@ -1,12 +1,10 @@
use std::collections::BTreeMap;
use aether_contracts::RequestBody;
use super::{
augment_sync_report_context, build_ai_execution_plan_from_decision,
generic_decision_missing_exact_provider_request, take_ai_decision_plan_core,
take_ai_upstream_auth_pair, take_non_empty_string, AiExecutionPlanFromDecisionParts,
AiStreamAttempt, AiSyncAttempt,
generic_decision_missing_exact_provider_request, resolve_ai_passthrough_sync_request_body,
take_ai_decision_plan_core, take_ai_upstream_auth_pair, take_non_empty_string,
AiExecutionPlanFromDecisionParts, AiStreamAttempt, AiSyncAttempt,
};
use crate::ai_serving::transport::{
build_standard_plan_fallback_headers, StandardPlanFallbackAcceptPolicy,
@@ -61,6 +59,10 @@ pub(crate) fn build_gemini_sync_plan_from_decision(
&provider_request_headers,
&provider_request_body_value,
)?;
let request_body = resolve_ai_passthrough_sync_request_body(
Some(provider_request_body_value),
payload.provider_request_body_base64.take(),
);
let stream = payload.upstream_is_stream;
let plan = build_ai_execution_plan_from_decision(
&mut payload,
@@ -70,7 +72,7 @@ pub(crate) fn build_gemini_sync_plan_from_decision(
url,
headers: std::mem::take(&mut provider_request_headers),
content_type,
body: RequestBody::from_json(provider_request_body_value),
body: request_body,
stream,
},
);
@@ -129,6 +131,10 @@ pub(crate) fn build_gemini_stream_plan_from_decision(
&provider_request_headers,
&provider_request_body_value,
)?;
let request_body = resolve_ai_passthrough_sync_request_body(
Some(provider_request_body_value),
payload.provider_request_body_base64.take(),
);
let plan = build_ai_execution_plan_from_decision(
&mut payload,
AiExecutionPlanFromDecisionParts {
@@ -137,7 +143,7 @@ pub(crate) fn build_gemini_stream_plan_from_decision(
url,
headers: std::mem::take(&mut provider_request_headers),
content_type,
body: RequestBody::from_json(provider_request_body_value),
body: request_body,
stream: true,
},
);
@@ -84,6 +84,7 @@ pub(crate) fn build_standard_upstream_url(
upstream_is_stream,
parts.uri.query(),
None,
None,
provider_request_body,
)
}
@@ -2199,6 +2199,8 @@ mod tests {
auth_snapshot: sample_auth_snapshot(),
required_capabilities: None,
request_auth_channel: None,
client_surface: None,
gateway_credential_carrier: None,
client_session_affinity: None,
routing_policy: None,
routing_trace_seed: None,
@@ -139,6 +139,20 @@ impl LocalExecutionAttemptSource<AiStreamAttempt> for LocalOpenAiChatStreamAttem
Ok(drained)
}
async fn skip_credential(&mut self, key_id: &str) -> Result<(), GatewayError> {
self.prefetched_attempts
.retain(|attempt| attempt.eligible.candidate.key_id != key_id);
self.candidates.skip_credential(key_id);
Ok(())
}
async fn skip_endpoint(&mut self, endpoint_id: &str) -> Result<(), GatewayError> {
self.prefetched_attempts
.retain(|attempt| attempt.eligible.candidate.endpoint_id != endpoint_id);
self.candidates.skip_endpoint(endpoint_id);
Ok(())
}
async fn skip_provider(&mut self, provider_id: &str) -> Result<(), GatewayError> {
self.prefetched_attempts
.retain(|attempt| attempt.eligible.candidate.provider_id != provider_id);
@@ -117,6 +117,16 @@ impl LocalExecutionAttemptSource<AiSyncAttempt> for LocalOpenAiChatSyncAttemptSo
Ok(drained)
}
async fn skip_credential(&mut self, key_id: &str) -> Result<(), GatewayError> {
self.candidates.skip_credential(key_id);
Ok(())
}
async fn skip_endpoint(&mut self, endpoint_id: &str) -> Result<(), GatewayError> {
self.candidates.skip_endpoint(endpoint_id);
Ok(())
}
async fn skip_provider(&mut self, provider_id: &str) -> Result<(), GatewayError> {
self.candidates.skip_provider(provider_id);
Ok(())
@@ -186,6 +186,16 @@ impl LocalExecutionAttemptSource<AiSyncAttempt> for LocalOpenAiResponsesSyncAtte
Ok(drained)
}
async fn skip_credential(&mut self, key_id: &str) -> Result<(), GatewayError> {
self.candidates.skip_credential(key_id);
Ok(())
}
async fn skip_endpoint(&mut self, endpoint_id: &str) -> Result<(), GatewayError> {
self.candidates.skip_endpoint(endpoint_id);
Ok(())
}
async fn skip_provider(&mut self, provider_id: &str) -> Result<(), GatewayError> {
self.candidates.skip_provider(provider_id);
Ok(())
@@ -219,6 +229,16 @@ impl LocalExecutionAttemptSource<AiStreamAttempt> for LocalOpenAiResponsesStream
Ok(drained)
}
async fn skip_credential(&mut self, key_id: &str) -> Result<(), GatewayError> {
self.candidates.skip_credential(key_id);
Ok(())
}
async fn skip_endpoint(&mut self, endpoint_id: &str) -> Result<(), GatewayError> {
self.candidates.skip_endpoint(endpoint_id);
Ok(())
}
async fn skip_provider(&mut self, provider_id: &str) -> Result<(), GatewayError> {
self.candidates.skip_provider(provider_id);
Ok(())
@@ -1,9 +1,8 @@
use std::collections::BTreeMap;
use aether_contracts::RequestBody;
use super::{
augment_sync_report_context, build_ai_execution_plan_from_decision, take_ai_decision_plan_core,
augment_sync_report_context, build_ai_execution_plan_from_decision,
resolve_ai_passthrough_sync_request_body, take_ai_decision_plan_core,
take_ai_upstream_auth_pair, take_non_empty_string, AiExecutionPlanFromDecisionParts,
AiStreamAttempt, AiSyncAttempt,
};
@@ -63,6 +62,10 @@ pub(crate) fn build_standard_sync_plan_from_decision(
&provider_request_headers,
&provider_request_body_value,
)?;
let request_body = resolve_ai_passthrough_sync_request_body(
Some(provider_request_body_value),
payload.provider_request_body_base64.take(),
);
let stream = payload.upstream_is_stream;
let plan = build_ai_execution_plan_from_decision(
&mut payload,
@@ -72,7 +75,7 @@ pub(crate) fn build_standard_sync_plan_from_decision(
url,
headers: std::mem::take(&mut provider_request_headers),
content_type,
body: RequestBody::from_json(provider_request_body_value),
body: request_body,
stream,
},
);
@@ -146,6 +149,10 @@ pub(crate) fn build_standard_stream_plan_from_decision(
&provider_request_headers,
&provider_request_body_value,
)?;
let request_body = resolve_ai_passthrough_sync_request_body(
Some(provider_request_body_value),
payload.provider_request_body_base64.take(),
);
let stream = payload.upstream_is_stream;
let plan = build_ai_execution_plan_from_decision(
&mut payload,
@@ -155,7 +162,7 @@ pub(crate) fn build_standard_stream_plan_from_decision(
url,
headers: std::mem::take(&mut provider_request_headers),
content_type,
body: RequestBody::from_json(provider_request_body_value),
body: request_body,
stream,
},
);
@@ -166,3 +173,88 @@ pub(crate) fn build_standard_stream_plan_from_decision(
report_context,
}))
}
#[cfg(test)]
mod tests {
use aether_contracts::{ExecutionResponseBodyMode, EXECUTION_RESPONSE_BODY_MODE_HEADER};
use serde_json::json;
use super::{
build_standard_stream_plan_from_decision, build_standard_sync_plan_from_decision,
AiExecutionDecision,
};
fn decision_with_raw_body(upstream_is_stream: bool) -> AiExecutionDecision {
serde_json::from_value(json!({
"action": if upstream_is_stream { "stream" } else { "sync" },
"request_id": "req-raw",
"provider_id": "provider-raw",
"endpoint_id": "endpoint-raw",
"key_id": "key-raw",
"upstream_url": "https://api.anthropic.test/v1/messages",
"provider_api_format": "claude:messages",
"client_api_format": "claude:messages",
"provider_request_headers": {
"content-type": "application/json",
(EXECUTION_RESPONSE_BODY_MODE_HEADER): ExecutionResponseBodyMode::PreserveBytes.as_str()
},
"provider_request_body": {
"model": "claude-sonnet-4",
"messages": []
},
"provider_request_body_base64": "eyAibW9kZWwiOiAiY2xhdWRlLXNvbm5ldC00IiwgIm1lc3NhZ2VzIjogW10gfQ==",
"content_type": "application/json",
"upstream_is_stream": upstream_is_stream
}))
.expect("decision should deserialize")
}
fn request_parts() -> http::request::Parts {
http::Request::builder()
.uri("http://localhost/v1/messages")
.body(())
.expect("request should build")
.into_parts()
.0
}
#[test]
fn standard_sync_plan_prefers_exact_request_body_bytes() {
let built = build_standard_sync_plan_from_decision(
&request_parts(),
&json!({}),
decision_with_raw_body(false),
)
.expect("plan should build")
.expect("plan should exist");
assert!(built.plan.body.json_body.is_none());
assert_eq!(
built.plan.body.body_bytes_b64.as_deref(),
Some("eyAibW9kZWwiOiAiY2xhdWRlLXNvbm5ldC00IiwgIm1lc3NhZ2VzIjogW10gfQ==")
);
assert_eq!(
built
.plan
.headers
.get(EXECUTION_RESPONSE_BODY_MODE_HEADER)
.map(String::as_str),
Some(ExecutionResponseBodyMode::PreserveBytes.as_str())
);
}
#[test]
fn standard_stream_plan_prefers_exact_request_body_bytes() {
let built = build_standard_stream_plan_from_decision(
&request_parts(),
&json!({}),
decision_with_raw_body(true),
false,
)
.expect("plan should build")
.expect("plan should exist");
assert!(built.plan.body.json_body.is_none());
assert!(built.plan.body.body_bytes_b64.is_some());
}
}
+39 -12
View File
@@ -92,9 +92,12 @@ pub(crate) use aether_ai_formats::api::{
request_conversion_requires_enable_flag, request_path_implies_stream_request,
resolve_claude_stream_spec, resolve_claude_sync_spec,
resolve_codex_responses_model_capabilities, 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, resolve_local_image_stream_spec, resolve_local_image_sync_spec,
resolve_execution_runtime_stream_plan_kind_with_client_surface,
resolve_execution_runtime_sync_plan_kind,
resolve_execution_runtime_sync_plan_kind_with_client_surface,
resolve_finalize_stream_rewrite_mode, resolve_gemini_files_stream_spec,
resolve_gemini_files_sync_spec, resolve_gemini_stream_spec, resolve_gemini_sync_spec,
resolve_local_image_stream_spec, resolve_local_image_sync_spec,
resolve_local_same_format_stream_spec, resolve_local_same_format_sync_spec,
resolve_local_video_sync_spec, resolve_openai_chat_max_tokens,
resolve_openai_embedding_sync_spec, resolve_openai_responses_stream_spec,
@@ -129,14 +132,15 @@ pub(crate) use aether_ai_formats::api::{
CLAUDE_CHAT_SYNC_SUCCESS_REPORT_KIND, CLAUDE_CLI_STREAM_PLAN_KIND,
CLAUDE_CLI_STREAM_SUCCESS_REPORT_KIND, CLAUDE_CLI_SYNC_ERROR_REPORT_KIND,
CLAUDE_CLI_SYNC_FINALIZE_REPORT_KIND, CLAUDE_CLI_SYNC_PLAN_KIND,
CLAUDE_CLI_SYNC_SUCCESS_REPORT_KIND, CODEX_OPENAI_IMAGE_DEFAULT_MODEL,
CODEX_OPENAI_IMAGE_DEFAULT_OUTPUT_FORMAT, CODEX_OPENAI_IMAGE_DEFAULT_VARIATION_MODEL,
CODEX_OPENAI_IMAGE_DEFAULT_VARIATION_PROMPT, CODEX_OPENAI_IMAGE_INTERNAL_MODEL,
EXECUTION_RUNTIME_STREAM_ACTION, EXECUTION_RUNTIME_STREAM_DECISION_ACTION,
EXECUTION_RUNTIME_SYNC_ACTION, EXECUTION_RUNTIME_SYNC_DECISION_ACTION,
GEMINI_CHAT_STREAM_PLAN_KIND, GEMINI_CHAT_STREAM_SUCCESS_REPORT_KIND,
GEMINI_CHAT_SYNC_ERROR_REPORT_KIND, GEMINI_CHAT_SYNC_FINALIZE_REPORT_KIND,
GEMINI_CHAT_SYNC_PLAN_KIND, GEMINI_CHAT_SYNC_SUCCESS_REPORT_KIND, GEMINI_CLI_STREAM_PLAN_KIND,
CLAUDE_CLI_SYNC_SUCCESS_REPORT_KIND, CLAUDE_COUNT_TOKENS_SYNC_PLAN_KIND,
CODEX_OPENAI_IMAGE_DEFAULT_MODEL, CODEX_OPENAI_IMAGE_DEFAULT_OUTPUT_FORMAT,
CODEX_OPENAI_IMAGE_DEFAULT_VARIATION_MODEL, CODEX_OPENAI_IMAGE_DEFAULT_VARIATION_PROMPT,
CODEX_OPENAI_IMAGE_INTERNAL_MODEL, EXECUTION_RUNTIME_STREAM_ACTION,
EXECUTION_RUNTIME_STREAM_DECISION_ACTION, EXECUTION_RUNTIME_SYNC_ACTION,
EXECUTION_RUNTIME_SYNC_DECISION_ACTION, GEMINI_CHAT_STREAM_PLAN_KIND,
GEMINI_CHAT_STREAM_SUCCESS_REPORT_KIND, GEMINI_CHAT_SYNC_ERROR_REPORT_KIND,
GEMINI_CHAT_SYNC_FINALIZE_REPORT_KIND, GEMINI_CHAT_SYNC_PLAN_KIND,
GEMINI_CHAT_SYNC_SUCCESS_REPORT_KIND, GEMINI_CLI_STREAM_PLAN_KIND,
GEMINI_CLI_STREAM_SUCCESS_REPORT_KIND, GEMINI_CLI_SYNC_ERROR_REPORT_KIND,
GEMINI_CLI_SYNC_FINALIZE_REPORT_KIND, GEMINI_CLI_SYNC_PLAN_KIND,
GEMINI_CLI_SYNC_SUCCESS_REPORT_KIND, GEMINI_CLI_V1INTERNAL_ENVELOPE_NAME,
@@ -167,5 +171,28 @@ pub(crate) use aether_ai_formats::api::{
pub(crate) use aether_ai_formats::{
api_format_defaults_to_client_error_failover, api_format_defaults_to_non_stream,
api_format_permission_covers, intersect_api_format_allowed_lists, is_embedding_api_format,
is_rerank_api_format, openai_responses_request_operation,
is_rerank_api_format, openai_responses_request_operation, ApiOperation, ClientSurface,
};
pub(crate) fn plan_kind_matches_api_operation(
plan_kind: &str,
require_streaming: bool,
expected_operation: Option<ApiOperation>,
) -> bool {
let Some(expected_operation) = expected_operation else {
return true;
};
if expected_operation == ApiOperation::OpenAiResponsesCompact {
return if require_streaming {
plan_kind == OPENAI_RESPONSES_COMPACT_STREAM_PLAN_KIND
} else {
plan_kind == OPENAI_RESPONSES_COMPACT_SYNC_PLAN_KIND
};
}
let resolved_operation = if require_streaming {
resolve_local_same_format_stream_spec(plan_kind).and_then(|spec| spec.operation)
} else {
resolve_local_same_format_sync_spec(plan_kind).and_then(|spec| spec.operation)
};
resolved_operation == Some(expected_operation)
}
+19 -18
View File
@@ -82,10 +82,11 @@ pub(crate) use aether_provider_transport::{
build_windsurf_cascade_headers, build_windsurf_cascade_request_body,
build_windsurf_cascade_upstream_url, candidate_common_transport_skip_reason,
candidate_transport_pair_skip_reason, classify_same_format_provider_request_behavior,
ensure_upstream_auth_header, gemini_files_transport_unsupported_reason,
header_rules_are_locally_supported, header_rules_have_enabled_rules,
is_gemini_cli_provider_transport, is_windsurf_provider_transport,
local_gemini_transport_unsupported_reason_with_network,
classify_same_format_provider_request_behavior_for_operation,
enforce_same_format_provider_api_operation_body_policy, ensure_upstream_auth_header,
gemini_files_transport_unsupported_reason, header_rules_are_locally_supported,
header_rules_have_enabled_rules, is_gemini_cli_provider_transport,
is_windsurf_provider_transport, local_gemini_transport_unsupported_reason_with_network,
local_openai_chat_transport_unsupported_reason,
local_standard_transport_unsupported_reason_with_network,
local_windsurf_request_transport_unsupported_reason_with_network,
@@ -93,22 +94,22 @@ pub(crate) use aether_provider_transport::{
request_conversion_enabled_for_transport, request_conversion_transport_supported,
request_conversion_transport_unsupported_reason, request_pair_allowed_for_transport,
request_pair_direct_auth, request_pair_transport_unsupported_reason,
resolve_gemini_cli_project_id, resolve_gemini_files_auth, resolve_grok_session_auth,
resolve_local_gemini_cli_request_auth, resolve_openai_image_auth,
resolve_same_format_provider_direct_auth, resolve_transport_execution_timeouts,
resolve_transport_profile, resolve_transport_proxy_snapshot,
resolve_transport_proxy_snapshot_with_tunnel_affinity, resolve_video_create_auth,
same_format_provider_transport_supported, same_format_provider_transport_unsupported_reason,
should_skip_upstream_passthrough_header, should_try_same_format_provider_oauth_auth,
supports_local_gemini_transport_with_network,
resolve_anthropic_compatibility_profile, resolve_gemini_cli_project_id,
resolve_gemini_files_auth, resolve_grok_session_auth, resolve_local_gemini_cli_request_auth,
resolve_openai_image_auth, resolve_same_format_provider_direct_auth,
resolve_transport_execution_timeouts, resolve_transport_profile,
resolve_transport_proxy_snapshot, resolve_transport_proxy_snapshot_with_tunnel_affinity,
resolve_video_create_auth, same_format_provider_transport_supported,
same_format_provider_transport_unsupported_reason, should_skip_upstream_passthrough_header,
should_try_same_format_provider_oauth_auth, supports_local_gemini_transport_with_network,
supports_local_generic_oauth_request_auth_resolution,
supports_local_oauth_request_auth_resolution, transport_proxy_is_locally_supported,
video_create_transport_unsupported_reason, CandidateTransportPolicyFacts,
GatewayProviderTransportSnapshot, GeminiCliRequestAuth, GeminiCliRequestAuthSupport,
GeminiCliRequestAuthUnsupportedReason, GeminiCliRequestEnvelopeSupport,
GeminiFilesHeadersInput, GeminiFilesRequestBodyError, GeminiFilesRequestBodyParts,
GrokHeaderInput, LocalResolvedOAuthRequestAuth, ProviderOpenAiImageHeadersInput,
ProviderVideoCreateFamily, ProviderVideoCreateHeadersInput,
transport_supports_api_operation, video_create_transport_unsupported_reason,
AnthropicCompatibilityProfile, CandidateTransportPolicyFacts, GatewayProviderTransportSnapshot,
GeminiCliRequestAuth, GeminiCliRequestAuthSupport, GeminiCliRequestAuthUnsupportedReason,
GeminiCliRequestEnvelopeSupport, GeminiFilesHeadersInput, GeminiFilesRequestBodyError,
GeminiFilesRequestBodyParts, GrokHeaderInput, LocalResolvedOAuthRequestAuth,
ProviderOpenAiImageHeadersInput, ProviderVideoCreateFamily, ProviderVideoCreateHeadersInput,
SameFormatProviderCompatibilityEdit, SameFormatProviderCompatibilityEditAction,
SameFormatProviderFamily, SameFormatProviderHeadersInput, SameFormatProviderRequestBehavior,
SameFormatProviderRequestBehaviorParams, SameFormatProviderRequestBodyInput,