feat(codex): align Search and execution protocol

This commit is contained in:
MMEXA
2026-07-12 03:04:15 +08:00
parent 59d37ae1dd
commit 2316df5c9a
121 changed files with 5141 additions and 543 deletions
@@ -21,9 +21,10 @@ pub(crate) use crate::ai_serving::{
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_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,
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,
};
pub(crate) use aether_ai_serving::AiRequestedModelFamily as RequestedModelFamily;
@@ -8,9 +8,10 @@ use crate::ai_serving::planner::common::{
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_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,
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,
@@ -100,7 +101,9 @@ fn build_sync_plan_payload_from_decision(
OPENAI_RESPONSES_SYNC_PLAN_KIND => {
build_openai_responses_sync_plan_from_decision(parts, body_json, payload, false)?
}
OPENAI_IMAGE_SYNC_PLAN_KIND => build_passthrough_sync_plan_from_decision(parts, payload)?,
OPENAI_IMAGE_SYNC_PLAN_KIND | OPENAI_SEARCH_SYNC_PLAN_KIND => {
build_passthrough_sync_plan_from_decision(parts, payload)?
}
OPENAI_RESPONSES_COMPACT_SYNC_PLAN_KIND => {
build_openai_responses_sync_plan_from_decision(parts, body_json, payload, true)?
}
@@ -18,7 +18,7 @@ use crate::ai_serving::{
ExecutionRuntimeAuthContext, GatewayAuthApiKeySnapshot, GatewayProviderTransportSnapshot,
PlannerAppState, CODEX_RESPONSES_LITE_HEADER,
};
use crate::client_session_affinity::client_session_affinity_from_request;
use crate::client_session_affinity::client_session_affinity_from_api_request;
use crate::clock::current_unix_secs;
use crate::routing::{
apply_routing_mutation_plan, build_routing_trace_seed, resolve_gateway_routing_policy,
@@ -449,8 +449,11 @@ pub(crate) async fn attach_routing_policy_to_local_requested_model_input(
let Some((group_id, group_version, group_config_json, selection_source)) = selected_group
else {
input.client_session_affinity =
client_session_affinity_from_request(&parts.headers, Some(body_json));
input.client_session_affinity = client_session_affinity_from_api_request(
client_api_format,
&parts.headers,
Some(body_json),
);
input.routing_policy = None;
input.routing_trace_seed = None;
input.routing_context = None;
@@ -527,8 +530,11 @@ pub(crate) async fn attach_routing_policy_to_local_requested_model_input(
}
let effective_headers_json = headers_to_routing_value(&effective_headers);
input.client_session_affinity =
client_session_affinity_from_request(&effective_headers, Some(&effective_body_json));
input.client_session_affinity = client_session_affinity_from_api_request(
client_api_format,
&effective_headers,
Some(&effective_body_json),
);
let final_policy_resolve_started_at = std::time::Instant::now();
let mut final_policy = resolve_gateway_routing_policy(GatewayRoutingPolicyInput {
group_id: group_id.as_deref(),
@@ -595,8 +601,11 @@ fn try_attach_static_default_routing_policy_to_input(
static_policy_resolve_started_at.elapsed().as_millis() as u64,
);
input.client_session_affinity =
client_session_affinity_from_request(&parts.headers, Some(body_json));
input.client_session_affinity = client_session_affinity_from_api_request(
client_api_format,
&parts.headers,
Some(body_json),
);
input.routing_trace_seed = Some(build_routing_trace_seed(&policy, client_api_format));
input.routing_policy = Some(policy);
input.routing_context = None;
@@ -24,7 +24,7 @@ use crate::ai_serving::{
ai_local_execution_contract_for_formats, extract_pool_sticky_session_token,
resolve_local_decision_execution_runtime_auth_context, GatewayControlDecision, PlannerAppState,
};
use crate::client_session_affinity::client_session_affinity_from_parts;
use crate::client_session_affinity::client_session_affinity_from_api_request;
use crate::clock::current_unix_secs;
use crate::{AppState, GatewayError};
@@ -81,7 +81,11 @@ 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_session_affinity = client_session_affinity_from_parts(parts, Some(body_json));
input.client_session_affinity = client_session_affinity_from_api_request(
spec_metadata.api_format,
&parts.headers,
Some(body_json),
);
if let Err(err) = attach_routing_policy_to_local_requested_model_input(
state,
parts,
@@ -53,6 +53,7 @@ pub(crate) fn resolve_same_format_provider_transport_unsupported_reason_for_trac
"openai:chat" => "openai:chat",
"openai:responses" => "openai:responses",
"openai:responses:compact" => "openai:responses:compact",
"openai:search" => "openai:search",
"openai:embedding" => "openai:embedding",
"openai:rerank" => "openai:rerank",
"claude:messages" => "claude:messages",
@@ -12,20 +12,33 @@ pub(crate) fn codex_model_capabilities_for_transport(
provider_model: &str,
source_model: &str,
) -> Option<crate::ai_serving::CodexResponsesModelCapabilities> {
if !transport
.provider
.provider_type
.trim()
.eq_ignore_ascii_case("codex")
|| !crate::ai_serving::is_openai_responses_family_format(provider_api_format)
{
codex_model_capabilities(
&transport.provider.provider_type,
provider_api_format,
provider_model,
source_model,
transport.key.upstream_metadata.as_ref(),
)
}
fn codex_model_capabilities(
provider_type: &str,
provider_api_format: &str,
provider_model: &str,
source_model: &str,
upstream_metadata: Option<&serde_json::Value>,
) -> Option<crate::ai_serving::CodexResponsesModelCapabilities> {
let uses_codex_model_catalog =
crate::ai_serving::is_openai_responses_family_format(provider_api_format)
|| crate::ai_serving::api_format_alias_matches(provider_api_format, "openai:search");
if !provider_type.trim().eq_ignore_ascii_case("codex") || !uses_codex_model_catalog {
return None;
}
Some(
crate::ai_serving::resolve_codex_responses_model_capabilities(
provider_model,
source_model,
transport.key.upstream_metadata.as_ref(),
upstream_metadata,
),
)
}
@@ -1,10 +1,52 @@
use std::collections::BTreeMap;
use super::{apply_codex_openai_responses_special_body_edits, apply_codex_openai_special_headers};
use super::{
apply_codex_openai_responses_special_body_edits, apply_codex_openai_special_headers,
codex_model_capabilities,
};
use crate::ai_serving::planner::standard::build_local_openai_responses_request_body;
use http::{HeaderMap, HeaderValue};
use serde_json::json;
#[test]
fn search_uses_live_codex_model_catalog_capabilities() {
let metadata = crate::ai_serving::build_codex_model_catalog_metadata(&[json!({
"slug": "gpt-search-custom",
"default_reasoning_level": "low",
"supported_reasoning_levels": [
{"effort": "low"},
{"effort": "max"}
],
"supports_parallel_tool_calls": true
})]);
let capabilities = codex_model_capabilities(
"codex",
"openai:search",
"gpt-search-custom",
"gpt-search-custom",
Some(&metadata),
)
.expect("Search should resolve capabilities from the Codex model catalog");
assert_eq!(
capabilities.default_reasoning_effort.as_deref(),
Some("low")
);
assert_eq!(
capabilities.supported_reasoning_efforts,
vec!["low".to_string(), "max".to_string()]
);
assert!(codex_model_capabilities(
"codex",
"openai:chat",
"gpt-search-custom",
"gpt-search-custom",
Some(&metadata),
)
.is_none());
}
#[test]
fn applies_codex_defaults_when_body_rules_do_not_handle_fields() {
let mut body = json!({
@@ -156,7 +156,8 @@ pub(crate) use aether_ai_formats::api::{
OPENAI_RESPONSES_COMPACT_SYNC_SUCCESS_REPORT_KIND, OPENAI_RESPONSES_STREAM_PLAN_KIND,
OPENAI_RESPONSES_STREAM_SUCCESS_REPORT_KIND, OPENAI_RESPONSES_SYNC_ERROR_REPORT_KIND,
OPENAI_RESPONSES_SYNC_FINALIZE_REPORT_KIND, OPENAI_RESPONSES_SYNC_PLAN_KIND,
OPENAI_RESPONSES_SYNC_SUCCESS_REPORT_KIND, OPENAI_VIDEO_CANCEL_SYNC_PLAN_KIND,
OPENAI_RESPONSES_SYNC_SUCCESS_REPORT_KIND, OPENAI_SEARCH_SYNC_PLAN_KIND,
OPENAI_SEARCH_SYNC_SUCCESS_REPORT_KIND, OPENAI_VIDEO_CANCEL_SYNC_PLAN_KIND,
OPENAI_VIDEO_CONTENT_PLAN_KIND, OPENAI_VIDEO_CREATE_SYNC_FINALIZE_REPORT_KIND,
OPENAI_VIDEO_CREATE_SYNC_PLAN_KIND, OPENAI_VIDEO_DELETE_SYNC_PLAN_KIND,
OPENAI_VIDEO_REMIX_SYNC_PLAN_KIND,
+2
View File
@@ -5,6 +5,7 @@ pub(crate) fn normalized_signature(api_format: &str) -> Option<&'static str> {
"openai:rerank" => Some("openai:rerank"),
"openai:responses" => Some("openai:responses"),
"openai:responses:compact" => Some("openai:responses:compact"),
"openai:search" => Some("openai:search"),
"openai:image" => Some("openai:image"),
"openai:video" => Some("openai:video"),
_ => None,
@@ -18,6 +19,7 @@ pub(crate) fn local_path(api_format: &str) -> Option<&'static str> {
"openai:rerank" => Some("/v1/rerank"),
"openai:responses" => Some("/v1/responses"),
"openai:responses:compact" => Some("/v1/responses/compact"),
"openai:search" => Some("/v1/alpha/search"),
"openai:image" => Some("/v1/images/generations"),
"openai:video" => Some("/v1/videos"),
_ => None,
@@ -15,6 +15,7 @@ const AI_POST_ROUTE_PATTERNS: &[&str] = &[
"/v1/messages/count_tokens",
"/v1/responses",
"/v1/responses/compact",
"/v1/alpha/search",
"/v1/images/generations",
"/v1/images/edits",
"/v1/interactions",
@@ -120,6 +121,7 @@ mod tests {
"/api/v1/services/embeddings/multimodal-embedding/multimodal-embedding",
),
("openai:rerank", "openai", "rerank", "/v1/rerank"),
("openai:search", "openai", "search", "/v1/alpha/search"),
("jina:rerank", "jina", "rerank", "/v1/rerank"),
] {
assert_eq!(
@@ -101,6 +101,29 @@ pub(crate) fn client_session_affinity_from_request(
client_session_scope_from_request(headers, body_json)?.scheduler_affinity()
}
pub(crate) fn client_session_affinity_from_api_request(
api_format: &str,
headers: &http::HeaderMap,
body_json: Option<&Value>,
) -> Option<ClientSessionAffinity> {
client_session_scope_from_api_request(api_format, headers, body_json)?.scheduler_affinity()
}
fn client_session_scope_from_api_request(
api_format: &str,
headers: &http::HeaderMap,
body_json: Option<&Value>,
) -> Option<ClientSessionScope> {
if api_format.trim().eq_ignore_ascii_case("openai:search") {
let request = ClientSessionRequest { headers, body_json };
return explicit_aether_session_scope(&request, CodexSessionScopeAdapter.family())
.or_else(|| codex_search_session_scope(&request))
.or_else(|| client_session_scope_from_request(headers, body_json));
}
client_session_scope_from_request(headers, body_json)
}
pub(crate) fn client_session_scope_from_request(
headers: &http::HeaderMap,
body_json: Option<&Value>,
@@ -113,6 +136,22 @@ pub(crate) fn client_session_scope_from_request(
.or_else(|| extract_scope_from_other_specific_adapters(&request, client_family.as_str()))
}
fn codex_search_session_scope(request: &ClientSessionRequest<'_>) -> Option<ClientSessionScope> {
let session_id = request
.body_json?
.get("id")?
.as_str()
.map(str::trim)
.filter(|value| !value.is_empty())?;
Some(ClientSessionScope::new(
CodexSessionScopeAdapter.family(),
session_id,
None,
header_value_str(request.headers, "chatgpt-account-id"),
ClientSessionSignalSource::Body,
))
}
pub(crate) fn client_session_affinity_from_parts(
parts: &http::request::Parts,
body_json: Option<&Value>,
@@ -789,6 +828,7 @@ fn has_header_with_prefix(headers: &http::HeaderMap, prefix: &str) -> bool {
#[cfg(test)]
mod tests {
use super::{
client_session_affinity_from_api_request,
client_session_affinity_from_report_context_value, client_session_affinity_from_request,
client_session_affinity_report_context_value, client_session_scope_from_request,
ClientSessionSignalSource, AETHER_AGENT_ID_HEADER, AETHER_SESSION_ID_HEADER,
@@ -1125,4 +1165,28 @@ mod tests {
assert!(client_session_affinity_from_request(&headers, Some(&body)).is_none());
}
#[test]
fn codex_search_uses_request_id_as_session_affinity() {
let mut headers = HeaderMap::new();
headers.insert("chatgpt-account-id", HeaderValue::from_static("account-1"));
let body = json!({"id": "codex-session-1", "model": "gpt-5.6"});
let affinity =
client_session_affinity_from_api_request("openai:search", &headers, Some(&body))
.expect("search affinity should build");
assert_eq!(affinity.client_family.as_deref(), Some("codex"));
assert_eq!(
affinity.session_key.as_deref(),
Some("account=account-1;session=codex-session-1")
);
}
#[test]
fn top_level_request_id_is_not_a_generic_session_signal() {
let body = json!({"id": "request-id", "model": "gpt-5.6"});
assert!(client_session_affinity_from_request(&HeaderMap::new(), Some(&body)).is_none());
}
}
+1
View File
@@ -125,6 +125,7 @@ pub(crate) const RUST_FRONTDOOR_OWNED_ROUTE_PATTERNS: &[&str] = &[
"/v1/messages/count_tokens",
"/v1/responses",
"/v1/responses/compact",
"/v1/alpha/search",
"/v1/models/{model}:generateContent",
"/v1/models/{model}:streamGenerateContent",
"/v1/models/{model}:predictLongRunning",
@@ -1115,10 +1115,9 @@ fn normalize_api_format_alias(value: &str) -> String {
fn auth_gate_api_format(auth_endpoint_signature: &str) -> String {
let normalized = normalize_api_format_alias(auth_endpoint_signature);
if normalized == "antigravity:v1internal" {
"gemini:generate_content".to_string()
} else {
normalized
match normalized.as_str() {
"antigravity:v1internal" => "gemini:generate_content".to_string(),
_ => normalized,
}
}
@@ -54,6 +54,14 @@ pub(super) fn classify_ai_public_route(
true,
))
}
} else if method == http::Method::POST && normalized_path == "/v1/alpha/search" {
Some(classified(
"ai_public",
"openai",
"search",
"openai:search",
true,
))
} else if method == http::Method::POST
&& matches!(
normalized_path,
@@ -56,6 +56,28 @@ fn classifies_openai_rerank_as_rerank_not_chat() {
assert!(decision.is_execution_runtime_candidate());
}
#[test]
fn classifies_openai_search_as_its_own_sync_endpoint() {
let headers = headers(&[("authorization", "Bearer sk-test")]);
let uri: Uri = "/v1/alpha/search".parse().expect("uri should parse");
let decision =
classify_control_route(&http::Method::POST, &uri, &headers).expect("route should classify");
assert_eq!(decision.route_family.as_deref(), Some("openai"));
assert_eq!(decision.route_kind.as_deref(), Some("search"));
assert_eq!(
decision.auth_endpoint_signature.as_deref(),
Some("openai:search")
);
assert!(decision.is_execution_runtime_candidate());
assert!(classify_control_route(&http::Method::GET, &uri, &headers).is_none());
let upstream_uri: Uri = "/backend-api/codex/alpha/search"
.parse()
.expect("uri should parse");
assert!(classify_control_route(&http::Method::POST, &upstream_uri, &headers).is_none());
}
#[test]
fn classifies_openai_chat_and_responses_separately_from_embedding() {
let headers = headers(&[("authorization", "Bearer sk-test")]);
+127 -3
View File
@@ -1816,7 +1816,7 @@ impl GatewayDataState {
)
});
let mut allowed_api_formats =
resolve_effective_list_policy(None, "unrestricted", &groups, |group| {
resolve_effective_api_format_policy(None, "unrestricted", &groups, |group| {
(
&group.allowed_api_formats_mode,
group.allowed_api_formats.clone(),
@@ -1832,7 +1832,7 @@ impl GatewayDataState {
&mut allowed_providers,
&mut snapshot.api_key_allowed_providers,
);
constrain_api_key_list_policy_to_user_policy(
constrain_api_key_api_format_policy_to_user_policy(
&mut allowed_api_formats,
&mut snapshot.api_key_allowed_api_formats,
);
@@ -1875,7 +1875,7 @@ impl GatewayDataState {
)
},
),
allowed_api_formats: resolve_effective_list_policy(
allowed_api_formats: resolve_effective_api_format_policy(
user.allowed_api_formats.clone(),
&user.allowed_api_formats_mode,
&groups,
@@ -1995,6 +1995,19 @@ fn resolve_effective_list_policy(
intersect_list_policies(group_policy, user_policy)
}
fn resolve_effective_api_format_policy(
user_values: Option<Vec<String>>,
user_mode: &str,
groups: &[aether_data::repository::users::StoredUserGroup],
group_field: impl Fn(
&aether_data::repository::users::StoredUserGroup,
) -> (&str, Option<Vec<String>>),
) -> Option<Vec<String>> {
let group_policy = union_group_list_policies(groups, group_field);
let user_policy = list_restriction_from_mode(user_mode, user_values);
intersect_api_format_list_policies(group_policy, user_policy)
}
fn union_group_list_policies(
groups: &[aether_data::repository::users::StoredUserGroup],
group_field: impl Fn(
@@ -2089,6 +2102,19 @@ fn intersect_list_policies(
}
}
fn intersect_api_format_list_policies(
left: Option<Vec<String>>,
right: Option<Vec<String>>,
) -> Option<Vec<String>> {
match (left, right) {
(None, None) => None,
(Some(values), None) | (None, Some(values)) => Some(values),
(Some(left_values), Some(right_values)) => Some(
aether_ai_formats::intersect_api_format_allowed_lists(&left_values, &right_values),
),
}
}
fn intersect_rate_limit_policies(
left: Option<RateLimitRestriction>,
right: Option<RateLimitRestriction>,
@@ -2149,6 +2175,22 @@ fn constrain_api_key_list_policy_to_user_policy(
*api_key_policy = Some(effective);
}
fn constrain_api_key_api_format_policy_to_user_policy(
user_policy: &mut Option<Vec<String>>,
api_key_policy: &mut Option<Vec<String>>,
) {
let Some(api_key_values) = api_key_policy.as_ref().filter(|values| !values.is_empty()) else {
return;
};
let Some(user_values) = user_policy.as_ref() else {
return;
};
let effective =
aether_ai_formats::intersect_api_format_allowed_lists(api_key_values, user_values);
*user_policy = Some(effective.clone());
*api_key_policy = Some(effective);
}
#[cfg(test)]
mod tests {
use std::sync::Arc;
@@ -2279,6 +2321,41 @@ mod tests {
);
}
#[test]
fn api_format_policy_intersection_preserves_search_companion_scope() {
let mut responses_group =
sample_group("responses", 10, None, "unrestricted", None, "system");
responses_group.allowed_api_formats = Some(vec!["openai:responses".to_string()]);
responses_group.allowed_api_formats_mode = "specific".to_string();
let search_policy = resolve_effective_api_format_policy(
Some(vec!["openai:search".to_string()]),
"specific",
std::slice::from_ref(&responses_group),
|group| {
(
&group.allowed_api_formats_mode,
group.allowed_api_formats.clone(),
)
},
);
assert_eq!(search_policy, Some(vec!["openai:search".to_string()]));
responses_group.allowed_api_formats = Some(vec!["openai:search".to_string()]);
let responses_policy = resolve_effective_api_format_policy(
Some(vec!["openai:responses".to_string()]),
"specific",
&[responses_group],
|group| {
(
&group.allowed_api_formats_mode,
group.allowed_api_formats.clone(),
)
},
);
assert_eq!(responses_policy, Some(vec!["openai:search".to_string()]));
}
#[test]
fn list_policy_unions_multiple_group_restrictions_legacy_case() {
let groups = vec![
@@ -2577,6 +2654,53 @@ mod tests {
assert_eq!(resolved.user_rate_limit, Some(30));
}
#[tokio::test]
async fn group_responses_permission_and_key_search_scope_resolve_to_search() {
let mut snapshot = sample_snapshot("key-search", "user-search");
snapshot.api_key_allowed_api_formats = Some(vec!["openai:search".to_string()]);
let auth_repository = Arc::new(InMemoryAuthApiKeySnapshotRepository::seed(vec![(
Some("hash-search".to_string()),
snapshot,
)]));
let user_repository = Arc::new(InMemoryUserReadRepository::seed_auth_users(vec![
sample_auth_user("user-search", "user"),
]));
let group = user_repository
.create_user_group(UpsertUserGroupRecord {
name: "Responses".to_string(),
description: None,
priority: 10,
allowed_providers: None,
allowed_providers_mode: "unrestricted".to_string(),
allowed_api_formats: Some(vec!["openai:responses".to_string()]),
allowed_api_formats_mode: "specific".to_string(),
allowed_models: None,
allowed_models_mode: "unrestricted".to_string(),
rate_limit: None,
rate_limit_mode: "system".to_string(),
})
.await
.expect("group should create")
.expect("group should exist");
user_repository
.add_user_to_group(&group.id, "user-search")
.await
.expect("group membership should create");
let state = GatewayDataState::with_auth_api_key_reader_for_tests(auth_repository)
.with_user_reader(user_repository);
let resolved = state
.read_auth_api_key_snapshot_by_key_hash("hash-search", 100)
.await
.expect("snapshot should resolve")
.expect("snapshot should exist");
assert_eq!(
resolved.effective_allowed_api_formats(),
Some(&["openai:search".to_string()][..])
);
}
#[tokio::test]
async fn data_state_lists_auth_api_key_export_records() {
let repository = Arc::new(
@@ -941,6 +941,7 @@ mod tests {
success_failover_patterns: Vec::new(),
error_stop_patterns: Vec::new(),
stop_cyber_policy_errors: false,
retry_client_errors_by_default: true,
}
);
}
@@ -59,7 +59,6 @@ const DEFAULT_STREAM_FIRST_BYTE_TIMEOUT_MS: u64 = 30_000;
const DEFAULT_NON_STREAM_TOTAL_TIMEOUT_MS: u64 = 300_000;
const DEFAULT_CODEX_COMPACT_TOTAL_TIMEOUT_MS: u64 = 1_200_000;
const MIN_TUNNEL_TIMEOUT_SECS: u64 = 1;
const MAX_TUNNEL_TIMEOUT_SECS: u64 = 1_200;
const DIRECT_REQWEST_H2_CLIENT_SHARDS_ENV: &str = "AETHER_GATEWAY_DIRECT_REQWEST_H2_CLIENT_SHARDS";
const DIRECT_REQWEST_CLIENT_SHARDS_ENV: &str = "AETHER_GATEWAY_DIRECT_REQWEST_CLIENT_SHARDS";
const DIRECT_REQWEST_H2_TARGET_STREAMS_PER_CLIENT_ENV: &str =
@@ -2216,36 +2215,51 @@ fn resolve_tunnel_first_byte_timeout(plan: &ExecutionPlan) -> Option<Duration> {
})
}
fn resolve_non_stream_total_timeout(plan: &ExecutionPlan) -> Option<Duration> {
if plan.stream {
pub(crate) fn resolve_non_stream_total_timeout_for_request(
is_stream: bool,
provider_api_format: &str,
timeouts: Option<&aether_contracts::ExecutionTimeouts>,
) -> Option<Duration> {
if is_stream {
return None;
}
let default_timeout_ms =
if crate::ai_serving::is_openai_responses_compact_format(&plan.provider_api_format) {
if crate::ai_serving::is_openai_responses_compact_format(provider_api_format) {
DEFAULT_CODEX_COMPACT_TOTAL_TIMEOUT_MS
} else {
DEFAULT_NON_STREAM_TOTAL_TIMEOUT_MS
};
let timeout_ms = plan
.timeouts
.as_ref()
let timeout_ms = timeouts
.and_then(|timeouts| timeouts.total_ms)
.unwrap_or(default_timeout_ms);
Some(Duration::from_millis(timeout_ms.max(1)))
}
pub(crate) fn resolve_stream_first_byte_timeout(plan: &ExecutionPlan) -> Option<Duration> {
if !plan.stream {
fn resolve_non_stream_total_timeout(plan: &ExecutionPlan) -> Option<Duration> {
resolve_non_stream_total_timeout_for_request(
plan.stream,
&plan.provider_api_format,
plan.timeouts.as_ref(),
)
}
pub(crate) fn resolve_stream_first_byte_timeout_for_request(
is_stream: bool,
timeouts: Option<&aether_contracts::ExecutionTimeouts>,
) -> Option<Duration> {
if !is_stream {
return None;
}
let timeout_ms = plan
.timeouts
.as_ref()
let timeout_ms = timeouts
.and_then(|timeouts| timeouts.first_byte_ms)
.unwrap_or(DEFAULT_STREAM_FIRST_BYTE_TIMEOUT_MS);
Some(Duration::from_millis(timeout_ms.max(1)))
}
pub(crate) fn resolve_stream_first_byte_timeout(plan: &ExecutionPlan) -> Option<Duration> {
resolve_stream_first_byte_timeout_for_request(plan.stream, plan.timeouts.as_ref())
}
pub(crate) async fn with_non_stream_total_timeout<T, F>(
plan: &ExecutionPlan,
future: F,
@@ -2363,7 +2377,10 @@ fn resolve_tunnel_timeout_metadata(plan: &ExecutionPlan) -> TunnelTimeoutMetadat
fn timeout_ms_to_secs(ms: u64) -> u64 {
let secs = ms.div_ceil(1_000);
secs.clamp(MIN_TUNNEL_TIMEOUT_SECS, MAX_TUNNEL_TIMEOUT_SECS)
secs.clamp(
MIN_TUNNEL_TIMEOUT_SECS,
aether_contracts::MAX_EXECUTION_REQUEST_TIMEOUT_SECS,
)
}
fn resolve_tunnel_node_id(proxy: Option<&ProxySnapshot>) -> Option<String> {
@@ -41,6 +41,7 @@ pub(crate) fn frontdoor_self_loop_public_ai_path(path: &str) -> bool {
| "/v1/rerank"
| "/v1/responses"
| "/v1/responses/compact"
| "/v1/alpha/search"
| "/v1beta/files"
| "/upload/v1beta/files"
| "/v1beta/operations"
@@ -104,6 +104,7 @@ fn is_known_admin_monitoring_api_format(value: &str) -> bool {
"openai:chat"
| "openai:responses"
| "openai:responses:compact"
| "openai:search"
| "openai:image"
| "openai:video"
| "openai:embedding"
@@ -492,6 +493,20 @@ mod tests {
);
}
#[test]
fn parses_v2_scheduler_affinity_key_for_openai_search() {
let parsed = parse_admin_monitoring_scheduler_affinity_key(
"scheduler_affinity:v2:user-key-1:openai:search:gpt-5.6-sol:codex:sessionhash",
)
.expect("Search scheduler key should parse");
assert_eq!(parsed.affinity_key, "user-key-1");
assert_eq!(parsed.api_format, "openai:search");
assert_eq!(parsed.model_name, "gpt-5.6-sol");
assert_eq!(parsed.client_family.as_deref(), Some("codex"));
assert_eq!(parsed.session_hash.as_deref(), Some("sessionhash"));
}
#[test]
fn parses_v2_scheduler_affinity_key_with_three_segment_api_format() {
let parsed = parse_admin_monitoring_scheduler_affinity_key(
@@ -6,7 +6,8 @@ use crate::handlers::admin::provider::shared::payloads::{
AdminProviderCreateRequest, AdminProviderUpdatePatch,
};
use crate::handlers::admin::provider::write::provider::{
reconcile_admin_fixed_provider_template_endpoints, reconcile_admin_fixed_provider_template_keys,
reconcile_admin_fixed_provider_template_endpoints,
reconcile_admin_fixed_provider_template_endpoints_after_update,
};
use crate::handlers::admin::request::{AdminAppState, AdminRequestContext};
use crate::handlers::admin::shared::attach_admin_audit_response;
@@ -80,7 +81,6 @@ pub(crate) async fn maybe_build_local_admin_provider_writes_response(
.is_some()
{
reconcile_admin_fixed_provider_template_endpoints(state, &created_provider).await?;
reconcile_admin_fixed_provider_template_keys(state, &created_provider).await?;
}
return Ok(Some(attach_admin_audit_response(
Json(json!({
@@ -161,8 +161,12 @@ pub(crate) async fn maybe_build_local_admin_provider_writes_response(
.fixed_provider_template(&updated_record.provider_type)
.is_some()
{
reconcile_admin_fixed_provider_template_endpoints(state, &updated_record).await?;
reconcile_admin_fixed_provider_template_keys(state, &updated_record).await?;
reconcile_admin_fixed_provider_template_endpoints_after_update(
state,
&existing_provider,
&updated_record,
)
.await?;
}
return Ok(Some(
match state
@@ -29,10 +29,7 @@ use crate::handlers::shared::{
provider_key_status_snapshot_payload,
};
use crate::model_fetch::ModelFetchRuntimeState;
use crate::provider_key_auth::{
provider_key_auth_semantics, provider_key_configured_api_formats,
provider_key_inherits_provider_api_formats,
};
use crate::provider_key_auth::provider_key_auth_semantics;
use crate::provider_transport::antigravity::{
build_antigravity_safe_v1internal_request, build_antigravity_static_identity_headers,
classify_local_antigravity_request_support, AntigravityEnvelopeRequestType,
@@ -543,6 +540,22 @@ fn provider_query_build_test_request_body_for_api_format(
model: &str,
route_path: &str,
client_api_format: &str,
) -> Value {
provider_query_build_test_request_body_for_api_format_with_search_session(
payload,
model,
route_path,
client_api_format,
None,
)
}
fn provider_query_build_test_request_body_for_api_format_with_search_session(
payload: &Value,
model: &str,
route_path: &str,
client_api_format: &str,
search_session_id: Option<&str>,
) -> Value {
let client_api_format = provider_query_normalize_api_format_alias(client_api_format);
let override_custom_model = route_path.ends_with("/test-model-failover")
@@ -568,7 +581,7 @@ fn provider_query_build_test_request_body_for_api_format(
);
} else if matches!(
client_api_format.as_str(),
"openai:responses" | "openai:responses:compact"
"openai:responses" | "openai:responses:compact" | "openai:search"
) && !value_has_non_empty_text(object.get("input"))
{
if let Some(prompt) = object
@@ -580,7 +593,7 @@ fn provider_query_build_test_request_body_for_api_format(
}
if matches!(
client_api_format.as_str(),
"openai:responses" | "openai:responses:compact"
"openai:responses" | "openai:responses:compact" | "openai:search"
) && value_has_non_empty_text(object.get("input"))
{
object.remove("prompt");
@@ -590,6 +603,9 @@ fn provider_query_build_test_request_body_for_api_format(
{
object.remove("messages");
}
if client_api_format == "openai:search" {
provider_query_ensure_search_test_fields(object, payload, search_session_id);
}
}
return body;
}
@@ -620,6 +636,15 @@ fn provider_query_build_test_request_body_for_api_format(
"temperature": 0.7,
"stream": true,
}),
"openai:search" => json!({
"id": provider_query_search_test_session_id(search_session_id),
"model": model,
"input": message,
"commands": {
"search_query": [{"q": message}]
},
"max_output_tokens": 256,
}),
"claude:messages" => json!({
"model": model,
"messages": [{
@@ -680,7 +705,7 @@ fn provider_query_insert_default_test_conversation(
.entry("top_n".to_string())
.or_insert_with(|| Value::from(4_u64));
}
"openai:responses" | "openai:responses:compact" => {
"openai:responses" | "openai:responses:compact" | "openai:search" => {
object.insert("input".to_string(), Value::String(message));
}
"claude:messages" => {
@@ -698,6 +723,48 @@ fn provider_query_insert_default_test_conversation(
}
}
fn provider_query_search_test_session_id(search_session_id: Option<&str>) -> String {
search_session_id
.map(str::trim)
.filter(|value| !value.is_empty())
.map(|value| format!("aether-model-test-{value}"))
.unwrap_or_else(|| format!("aether-model-test-{}", Uuid::new_v4().simple()))
}
fn provider_query_ensure_search_test_fields(
object: &mut Map<String, Value>,
payload: &Value,
search_session_id: Option<&str>,
) {
let query = object
.get("input")
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
.map(ToOwned::to_owned)
.or_else(|| provider_query_extract_message(payload))
.unwrap_or_else(|| DEFAULT_PROVIDER_QUERY_TEST_MESSAGE.to_string());
let has_session_id = object
.get("id")
.and_then(Value::as_str)
.is_some_and(|value| !value.trim().is_empty());
if !has_session_id {
object.insert(
"id".to_string(),
provider_query_search_test_session_id(search_session_id).into(),
);
}
object
.entry("input".to_string())
.or_insert_with(|| Value::String(query.clone()));
object
.entry("commands".to_string())
.or_insert_with(|| json!({"search_query": [{"q": query}]}));
object
.entry("max_output_tokens".to_string())
.or_insert_with(|| Value::from(256_u64));
}
fn provider_query_grok_test_client_api_format(provider_api_format: &str) -> &'static str {
match provider_query_normalize_api_format_alias(provider_api_format).as_str() {
"openai:responses" | "openai:responses:compact" => "openai:responses",
@@ -767,7 +834,7 @@ fn provider_query_request_body_has_conversation_for_api_format(
client_api_format: &str,
) -> bool {
match provider_query_normalize_api_format_alias(client_api_format).as_str() {
"openai:responses" | "openai:responses:compact" => {
"openai:responses" | "openai:responses:compact" | "openai:search" => {
value_has_non_empty_text(body.get("input"))
|| value_has_non_empty_text(body.get("prompt"))
}
@@ -874,15 +941,11 @@ fn provider_query_key_supports_endpoint(
provider_type: &str,
endpoint_api_format: &str,
) -> bool {
if provider_key_inherits_provider_api_formats(key, provider_type) {
return true;
}
let formats = provider_key_configured_api_formats(key);
let endpoint_api_format = provider_query_normalize_api_format_alias(endpoint_api_format);
formats.is_empty()
|| formats
.iter()
.any(|value| provider_query_normalize_api_format_alias(value) == endpoint_api_format)
crate::handlers::shared::provider_catalog_key_supports_format(
key,
provider_type,
endpoint_api_format,
)
}
async fn provider_query_select_preferred_non_kiro_endpoint(
@@ -1624,6 +1687,15 @@ fn provider_query_standard_execution_response_body(
{
return None;
}
if result.status_code < 400
&& provider_query_normalize_api_format_alias(provider_api_format) == "openai:search"
&& !body
.get("output")
.and_then(Value::as_str)
.is_some_and(|value| !value.trim().is_empty())
{
return None;
}
Some(body)
}
@@ -2864,12 +2936,14 @@ async fn provider_query_execute_standard_test_candidate(
crate::ai_serving::normalize_api_format_alias(provider_api_format);
let client_api_format =
provider_query_standard_test_client_api_format(normalized_provider_api_format.as_str());
let original_request_body = provider_query_build_test_request_body_for_api_format(
payload,
&candidate.effective_model,
route_path,
client_api_format,
);
let original_request_body =
provider_query_build_test_request_body_for_api_format_with_search_session(
payload,
&candidate.effective_model,
route_path,
client_api_format,
Some(trace_id),
);
if crate::provider_transport::is_windsurf_provider_transport(&transport)
&& provider_query_normalize_api_format_alias(candidate.endpoint.api_format.as_str())
== "openai:chat"
@@ -3014,6 +3088,45 @@ async fn provider_query_execute_standard_test_candidate(
);
provider_request_body
}
"openai:search" => {
let Some(mut provider_request_body) =
crate::provider_transport::build_same_format_provider_request_body(
crate::provider_transport::SameFormatProviderRequestBodyInput {
body_json: &request_body,
mapped_model: request_model,
client_api_format,
provider_api_format,
source_model: request_body.get("model").and_then(Value::as_str),
family: crate::provider_transport::SameFormatProviderFamily::Standard,
body_rules: transport.endpoint.body_rules.as_ref(),
request_headers: Some(&incoming_request_headers),
upstream_is_stream,
force_body_stream_field: require_body_stream_field,
kiro_auth_config: None,
is_claude_code: false,
enable_model_directives: false,
},
)
else {
return Ok(provider_query_skipped_execution_outcome(
request_body.clone(),
format!("Provider request body could not be built for {provider_api_format}"),
));
};
if let Err(err) = crate::provider_transport::apply_transport_request_body_semantics(
&mut provider_request_body,
&transport,
normalized_provider_api_format.as_str(),
) {
return Ok(provider_query_skipped_execution_outcome(
provider_request_body,
format!(
"Provider request body is not compatible with transport semantics: {err}"
),
));
}
provider_request_body
}
"openai:embedding"
| "gemini:embedding"
| "jina:embedding"
@@ -3077,7 +3190,7 @@ async fn provider_query_execute_standard_test_candidate(
);
if matches!(
normalized_provider_api_format.as_str(),
"openai:chat" | "openai:responses" | "openai:responses:compact"
"openai:chat" | "openai:responses" | "openai:responses:compact" | "openai:search"
) && crate::ai_serving::finalize_openai_provider_request_with_codex_model_capabilities(
&mut provider_request_body,
crate::ai_serving::OpenAiProviderRequestFinalization {
@@ -3175,6 +3288,7 @@ async fn provider_query_execute_standard_test_candidate(
"openai:chat"
| "openai:responses"
| "openai:responses:compact"
| "openai:search"
| "claude:messages"
| "gemini:generate_content"
| "openai:embedding"
@@ -3190,6 +3304,7 @@ async fn provider_query_execute_standard_test_candidate(
"openai:chat"
| "openai:responses"
| "openai:responses:compact"
| "openai:search"
| "openai:embedding"
| "jina:embedding"
| "doubao:embedding"
@@ -3255,7 +3370,7 @@ async fn provider_query_execute_standard_test_candidate(
&BTreeMap::new(),
Some("application/json"),
),
"openai:responses" | "openai:responses:compact" => {
"openai:responses" | "openai:responses:compact" | "openai:search" => {
crate::provider_transport::auth::build_complete_passthrough_headers_with_auth(
&parts.headers,
auth_header.as_deref().unwrap_or_default(),
@@ -3321,7 +3436,9 @@ async fn provider_query_execute_standard_test_candidate(
response_body: None,
});
}
if crate::ai_serving::is_openai_responses_family_format(provider_api_format) {
if crate::ai_serving::is_openai_responses_family_format(provider_api_format)
|| crate::ai_serving::api_format_alias_matches(provider_api_format, "openai:search")
{
crate::ai_serving::apply_codex_openai_special_headers(
&mut request_headers,
&provider_request_body,
@@ -33,6 +33,8 @@ pub(super) fn provider_query_standard_test_client_api_format(
let normalized_api_format = crate::ai_serving::normalize_api_format_alias(provider_api_format);
if normalized_api_format == "openai:responses:compact" {
"openai:responses:compact"
} else if normalized_api_format == "openai:search" {
"openai:search"
} else if crate::ai_serving::is_embedding_api_format(&normalized_api_format) {
"openai:embedding"
} else if crate::ai_serving::is_rerank_api_format(&normalized_api_format) {
@@ -71,6 +73,7 @@ pub(super) fn provider_query_standard_test_unsupported_reason(
}
"openai:responses"
| "openai:responses:compact"
| "openai:search"
| "claude:messages"
| "openai:embedding"
| "jina:embedding"
@@ -261,6 +264,7 @@ pub(super) fn provider_query_test_adapter_for_provider_api_format(
"openai:chat"
| "openai:responses"
| "openai:responses:compact"
| "openai:search"
| "claude:messages"
| "gemini:generate_content"
| "gemini:interactions"
@@ -365,6 +369,7 @@ pub(super) fn provider_query_transport_supports_model_test_execution(
}
"openai:responses"
| "openai:responses:compact"
| "openai:search"
| "openai:embedding"
| "jina:embedding"
| "doubao:embedding"
@@ -101,8 +101,7 @@ fn provider_query_model_mapping_matches_endpoint(
) -> bool {
let api_format_matches = mapping.api_formats.as_ref().is_none_or(|api_formats| {
api_formats.iter().any(|value| {
aether_scheduler_core::normalize_api_format(value)
== aether_scheduler_core::normalize_api_format(&endpoint.api_format)
aether_ai_formats::api_format_permission_covers(value, &endpoint.api_format)
})
});
if !api_format_matches {
@@ -367,3 +366,41 @@ fn provider_query_parse_mapping_string_list_array(
}
Ok(parsed)
}
#[cfg(test)]
mod tests {
use super::*;
fn sample_endpoint(api_format: &str) -> StoredProviderCatalogEndpoint {
StoredProviderCatalogEndpoint::new(
"endpoint-1".to_string(),
"provider-1".to_string(),
api_format.to_string(),
None,
None,
true,
)
.expect("endpoint should build")
}
fn sample_mapping(api_format: &str) -> StoredProviderModelMapping {
StoredProviderModelMapping {
name: "gpt-5.6-luna".to_string(),
priority: 1,
api_formats: Some(vec![api_format.to_string()]),
endpoint_ids: None,
}
}
#[test]
fn responses_mapping_scope_covers_search_in_one_direction() {
assert!(provider_query_model_mapping_matches_endpoint(
&sample_mapping("openai:responses"),
&sample_endpoint("openai:search"),
));
assert!(!provider_query_model_mapping_matches_endpoint(
&sample_mapping("openai:search"),
&sample_endpoint("openai:responses"),
));
}
}
@@ -475,6 +475,118 @@ fn provider_query_compact_test_request_body_defaults_to_responses_input() {
assert!(body.get("messages").is_none());
}
#[test]
fn provider_query_search_test_request_body_defaults_to_typed_search_input() {
let payload = json!({"message": "find current documentation"});
let client_api_format = provider_query_standard_test_client_api_format("openai:search");
let body = provider_query_build_test_request_body_for_api_format(
&payload,
"gpt-5.6-sol",
"/api/admin/provider-query/test-model",
client_api_format,
);
assert_eq!(client_api_format, "openai:search");
assert!(body["id"]
.as_str()
.is_some_and(|value| value.starts_with("aether-model-test-")));
assert_eq!(body["model"], json!("gpt-5.6-sol"));
assert_eq!(body["input"], json!("find current documentation"));
assert_eq!(
body["commands"]["search_query"][0]["q"],
json!("find current documentation")
);
assert_eq!(body["max_output_tokens"], json!(256));
assert!(body.get("messages").is_none());
assert!(body.get("stream").is_none());
}
#[test]
fn provider_query_search_test_completes_missing_protocol_fields() {
let payload = json!({
"request_body": {
"model": "gpt-5.6-luna",
"input": "find release notes"
}
});
let body = provider_query_build_test_request_body_for_api_format_with_search_session(
&payload,
"gpt-5.6-luna",
"/api/admin/provider-query/test-model",
"openai:search",
Some("trace-model-test-1"),
);
assert_eq!(body["id"], json!("aether-model-test-trace-model-test-1"));
assert_eq!(
body["commands"]["search_query"][0]["q"],
json!("find release notes")
);
assert_eq!(body["max_output_tokens"], json!(256));
}
#[test]
fn provider_query_search_test_preserves_a_non_empty_client_session_id() {
let payload = json!({
"request_body": {
"id": "client-search-session",
"model": "gpt-5.6-luna",
"input": "find release notes"
}
});
let body = provider_query_build_test_request_body_for_api_format_with_search_session(
&payload,
"gpt-5.6-luna",
"/api/admin/provider-query/test-model",
"openai:search",
Some("trace-model-test-1"),
);
assert_eq!(body["id"], json!("client-search-session"));
}
#[test]
fn provider_query_search_success_requires_non_empty_output() {
fn result(body: Value) -> aether_contracts::ExecutionResult {
aether_contracts::ExecutionResult {
request_id: "provider-search-test".to_string(),
candidate_id: Some("candidate-0".to_string()),
status_code: 200,
headers: BTreeMap::new(),
body: Some(aether_contracts::ResponseBody {
json_body: Some(body),
body_bytes_b64: None,
}),
telemetry: None,
error: None,
}
}
assert!(provider_query_standard_execution_response_body(
"openai:search",
&result(json!({})),
None,
)
.is_none());
assert!(provider_query_standard_execution_response_body(
"openai:search",
&result(json!({"output": " "})),
None,
)
.is_none());
assert_eq!(
provider_query_standard_execution_response_body(
"openai:search",
&result(json!({"output": "search result", "encrypted_output": "ciphertext"})),
None,
),
Some(json!({"output": "search result", "encrypted_output": "ciphertext"}))
);
}
#[test]
fn provider_query_embedding_test_request_body_defaults_to_embedding_input() {
let payload = json!({"message": "hello from embedding"});
@@ -637,6 +749,10 @@ fn provider_query_test_adapter_routes_fixed_provider_endpoint_types() {
provider_query_test_adapter_for_provider_api_format("codex", "openai:responses:compact"),
Some(ProviderQueryTestAdapter::Standard)
);
assert_eq!(
provider_query_test_adapter_for_provider_api_format("codex", "openai:search"),
Some(ProviderQueryTestAdapter::Standard)
);
assert_eq!(
provider_query_test_adapter_for_provider_api_format("chatgpt_web", "openai:image"),
Some(ProviderQueryTestAdapter::OpenAiImage)
@@ -722,6 +838,10 @@ fn provider_query_endpoint_priority_prefers_text_before_cli_and_image() {
provider_query_model_test_endpoint_priority("codex", "openai:responses:compact"),
Some(1)
);
assert_eq!(
provider_query_model_test_endpoint_priority("codex", "openai:search"),
Some(1)
);
assert_eq!(
provider_query_model_test_endpoint_priority("chatgpt_web", "openai:image"),
Some(2)
@@ -3,11 +3,79 @@ mod endpoint;
mod template;
mod update;
fn normalize_provider_request_timeout(value: Option<f64>) -> Result<Option<f64>, String> {
let max = aether_contracts::MAX_EXECUTION_REQUEST_TIMEOUT_SECS as f64;
match value {
Some(value) if (1.0..=max).contains(&value) => Ok(Some(value)),
Some(_) => Err(format!(
"request_timeout 必须是 1 到 {} 之间的数字",
aether_contracts::MAX_EXECUTION_REQUEST_TIMEOUT_SECS
)),
None => Ok(None),
}
}
fn normalize_provider_stream_first_byte_timeout(value: Option<f64>) -> Result<Option<f64>, String> {
let max = aether_contracts::MAX_EXECUTION_STREAM_FIRST_BYTE_TIMEOUT_SECS as f64;
match value {
Some(value) if (1.0..=max).contains(&value) => Ok(Some(value)),
Some(_) => Err(format!(
"stream_first_byte_timeout 必须是 1 到 {} 之间的数字",
aether_contracts::MAX_EXECUTION_STREAM_FIRST_BYTE_TIMEOUT_SECS
)),
None => Ok(None),
}
}
pub(crate) use self::create::build_admin_create_provider_record;
pub(crate) use self::endpoint::build_admin_fixed_provider_endpoint_record;
pub(crate) use self::template::{
apply_admin_fixed_provider_endpoint_template_overrides,
reconcile_admin_fixed_provider_template_endpoints,
reconcile_admin_fixed_provider_template_keys,
reconcile_admin_fixed_provider_template_endpoints_after_update,
};
pub(crate) use self::update::build_admin_update_provider_record;
#[cfg(test)]
mod tests {
use super::{normalize_provider_request_timeout, normalize_provider_stream_first_byte_timeout};
#[test]
fn provider_request_timeout_accepts_the_execution_protocol_range() {
assert_eq!(normalize_provider_request_timeout(Some(1.0)), Ok(Some(1.0)));
assert_eq!(
normalize_provider_request_timeout(Some(
aether_contracts::MAX_EXECUTION_REQUEST_TIMEOUT_SECS as f64,
)),
Ok(Some(
aether_contracts::MAX_EXECUTION_REQUEST_TIMEOUT_SECS as f64,
))
);
assert!(normalize_provider_request_timeout(Some(0.0)).is_err());
assert!(normalize_provider_request_timeout(Some(
aether_contracts::MAX_EXECUTION_REQUEST_TIMEOUT_SECS as f64 + 1.0,
))
.is_err());
assert!(normalize_provider_request_timeout(Some(f64::NAN)).is_err());
}
#[test]
fn provider_stream_first_byte_timeout_keeps_its_protocol_range() {
assert_eq!(
normalize_provider_stream_first_byte_timeout(Some(1.0)),
Ok(Some(1.0))
);
assert_eq!(
normalize_provider_stream_first_byte_timeout(Some(
aether_contracts::MAX_EXECUTION_STREAM_FIRST_BYTE_TIMEOUT_SECS as f64,
)),
Ok(Some(
aether_contracts::MAX_EXECUTION_STREAM_FIRST_BYTE_TIMEOUT_SECS as f64,
))
);
assert!(normalize_provider_stream_first_byte_timeout(Some(
aether_contracts::MAX_EXECUTION_STREAM_FIRST_BYTE_TIMEOUT_SECS as f64 + 1.0,
))
.is_err());
}
}
@@ -107,16 +107,9 @@ pub(crate) async fn build_admin_create_provider_record(
None => Some(2),
};
let proxy = normalize_json_object(payload.proxy, "proxy")?;
let stream_first_byte_timeout_secs = match payload.stream_first_byte_timeout {
Some(value) if (1.0..=300.0).contains(&value) => Some(value),
Some(_) => return Err("stream_first_byte_timeout 必须是 1 到 300 之间的数字".to_string()),
None => None,
};
let request_timeout_secs = match payload.request_timeout {
Some(value) if (1.0..=600.0).contains(&value) => Some(value),
Some(_) => return Err("request_timeout 必须是 1 到 600 之间的数字".to_string()),
None => None,
};
let stream_first_byte_timeout_secs =
super::normalize_provider_stream_first_byte_timeout(payload.stream_first_byte_timeout)?;
let request_timeout_secs = super::normalize_provider_request_timeout(payload.request_timeout)?;
let mut config_map = normalize_json_object(payload.config, "config")?
.and_then(|value| value.as_object().cloned())
@@ -1,11 +1,11 @@
use super::endpoint::{
build_admin_fixed_provider_endpoint_defaults, build_admin_fixed_provider_endpoint_record,
AdminFixedProviderEndpointDefaults,
};
use crate::handlers::admin::request::AdminAppState;
use crate::provider_key_auth::provider_key_is_oauth_managed;
use crate::GatewayError;
use aether_data_contracts::repository::provider_catalog::{
StoredProviderCatalogEndpoint, StoredProviderCatalogKey, StoredProviderCatalogProvider,
StoredProviderCatalogEndpoint, StoredProviderCatalogProvider,
};
use aether_provider_transport::provider_types::{
fixed_provider_template, FixedProviderEndpointTemplate, FixedProviderTemplate,
@@ -36,6 +36,30 @@ struct FixedProviderEndpointMetadata {
pub(crate) async fn reconcile_admin_fixed_provider_template_endpoints(
state: &AdminAppState<'_>,
provider: &StoredProviderCatalogProvider,
) -> Result<(), GatewayError> {
reconcile_admin_fixed_provider_template_endpoints_with_adoption_provider(
state, provider, provider,
)
.await
}
pub(crate) async fn reconcile_admin_fixed_provider_template_endpoints_after_update(
state: &AdminAppState<'_>,
previous_provider: &StoredProviderCatalogProvider,
updated_provider: &StoredProviderCatalogProvider,
) -> Result<(), GatewayError> {
reconcile_admin_fixed_provider_template_endpoints_with_adoption_provider(
state,
updated_provider,
previous_provider,
)
.await
}
async fn reconcile_admin_fixed_provider_template_endpoints_with_adoption_provider(
state: &AdminAppState<'_>,
provider: &StoredProviderCatalogProvider,
adoption_provider: &StoredProviderCatalogProvider,
) -> Result<(), GatewayError> {
let Some(template) = state.fixed_provider_template(&provider.provider_type) else {
return Ok(());
@@ -53,8 +77,9 @@ pub(crate) async fn reconcile_admin_fixed_provider_template_endpoints(
match existing_endpoint {
Some(existing_endpoint) => {
matched_endpoint_ids.insert(existing_endpoint.id.clone());
let updated = reconcile_fixed_provider_endpoint(
let updated = reconcile_fixed_provider_endpoint_with_adoption_provider(
provider,
adoption_provider,
existing_endpoint,
template,
endpoint_template,
@@ -115,31 +140,6 @@ pub(crate) async fn reconcile_admin_fixed_provider_template_endpoints(
Ok(())
}
pub(crate) async fn reconcile_admin_fixed_provider_template_keys(
state: &AdminAppState<'_>,
provider: &StoredProviderCatalogProvider,
) -> Result<(), GatewayError> {
let Some(_) = state.fixed_provider_template(&provider.provider_type) else {
return Ok(());
};
let existing_keys = state
.list_provider_catalog_keys_by_provider_ids(std::slice::from_ref(&provider.id))
.await?;
for existing_key in existing_keys {
let Some(updated_key) = reconcile_fixed_provider_key(provider, &existing_key) else {
continue;
};
let Some(_) = state.update_provider_catalog_key(&updated_key).await? else {
return Err(GatewayError::Internal(
"provider catalog key writer unavailable".to_string(),
));
};
}
Ok(())
}
pub(crate) fn apply_admin_fixed_provider_endpoint_template_overrides(
provider: &StoredProviderCatalogProvider,
existing_endpoint: &StoredProviderCatalogEndpoint,
@@ -156,8 +156,14 @@ pub(crate) fn apply_admin_fixed_provider_endpoint_template_overrides(
let defaults =
build_admin_fixed_provider_endpoint_defaults(provider, template, endpoint_template)?;
let mut metadata = fixed_provider_endpoint_metadata(existing_endpoint)
.unwrap_or_else(|| managed_fixed_provider_endpoint_metadata(template, endpoint_template));
let mut metadata = fixed_provider_endpoint_metadata(existing_endpoint).unwrap_or_else(|| {
adopt_fixed_provider_endpoint_metadata(
existing_endpoint,
&defaults,
template,
endpoint_template,
)
});
let mut overrides = metadata.overrides.clone();
sync_override_if_changed(
@@ -250,17 +256,48 @@ pub(crate) fn apply_admin_fixed_provider_endpoint_template_overrides(
Ok(())
}
#[cfg(test)]
fn reconcile_fixed_provider_endpoint(
provider: &StoredProviderCatalogProvider,
existing_endpoint: &StoredProviderCatalogEndpoint,
template: &FixedProviderTemplate,
endpoint_template: &FixedProviderEndpointTemplate,
) -> Result<StoredProviderCatalogEndpoint, String> {
reconcile_fixed_provider_endpoint_with_adoption_provider(
provider,
provider,
existing_endpoint,
template,
endpoint_template,
)
}
fn reconcile_fixed_provider_endpoint_with_adoption_provider(
provider: &StoredProviderCatalogProvider,
adoption_provider: &StoredProviderCatalogProvider,
existing_endpoint: &StoredProviderCatalogEndpoint,
template: &FixedProviderTemplate,
endpoint_template: &FixedProviderEndpointTemplate,
) -> Result<StoredProviderCatalogEndpoint, String> {
let defaults =
build_admin_fixed_provider_endpoint_defaults(provider, template, endpoint_template)?;
let mut updated = existing_endpoint.clone();
let metadata = fixed_provider_endpoint_metadata(existing_endpoint)
.unwrap_or_else(|| managed_fixed_provider_endpoint_metadata(template, endpoint_template));
let metadata = match fixed_provider_endpoint_metadata(existing_endpoint) {
Some(metadata) => metadata,
None => {
let adoption_defaults = build_admin_fixed_provider_endpoint_defaults(
adoption_provider,
template,
endpoint_template,
)?;
adopt_fixed_provider_endpoint_metadata(
existing_endpoint,
&adoption_defaults,
template,
endpoint_template,
)
}
};
updated.api_format = defaults.api_format.clone();
updated.api_family = Some(defaults.api_family.clone());
@@ -455,6 +492,77 @@ fn managed_fixed_provider_endpoint_metadata(
}
}
fn adopt_fixed_provider_endpoint_metadata(
existing_endpoint: &StoredProviderCatalogEndpoint,
defaults: &AdminFixedProviderEndpointDefaults,
template: &FixedProviderTemplate,
endpoint_template: &FixedProviderEndpointTemplate,
) -> FixedProviderEndpointMetadata {
let mut metadata = managed_fixed_provider_endpoint_metadata(template, endpoint_template);
sync_override(
&mut metadata.overrides,
OVERRIDE_BASE_URL,
&existing_endpoint.base_url,
&defaults.base_url,
);
sync_override(
&mut metadata.overrides,
OVERRIDE_CUSTOM_PATH,
&existing_endpoint.custom_path,
&defaults.custom_path,
);
sync_override(
&mut metadata.overrides,
OVERRIDE_HEADER_RULES,
&existing_endpoint.header_rules,
&defaults.header_rules,
);
sync_override(
&mut metadata.overrides,
OVERRIDE_BODY_RULES,
&existing_endpoint.body_rules,
&defaults.body_rules,
);
sync_override(
&mut metadata.overrides,
OVERRIDE_MAX_RETRIES,
&existing_endpoint.max_retries,
&defaults.max_retries,
);
sync_override(
&mut metadata.overrides,
OVERRIDE_IS_ACTIVE,
&existing_endpoint.is_active,
&defaults.is_active,
);
sync_override(
&mut metadata.overrides,
OVERRIDE_PROXY,
&existing_endpoint.proxy,
&defaults.proxy,
);
sync_override(
&mut metadata.overrides,
OVERRIDE_FORMAT_ACCEPTANCE_CONFIG,
&existing_endpoint.format_acceptance_config,
&defaults.format_acceptance_config,
);
let existing_config = endpoint_config_without_metadata(existing_endpoint.config.as_ref());
for (key, desired) in fixed_provider_endpoint_config_defaults(endpoint_template) {
let Some(actual) = existing_config.get(&key) else {
continue;
};
sync_override(
&mut metadata.overrides,
&config_override_key(&key),
actual,
&desired,
);
}
metadata
}
fn upsert_fixed_provider_endpoint_metadata(
endpoint: &mut StoredProviderCatalogEndpoint,
metadata: &FixedProviderEndpointMetadata,
@@ -513,22 +621,6 @@ fn current_unix_secs() -> u64 {
.unwrap_or(0)
}
fn reconcile_fixed_provider_key(
provider: &StoredProviderCatalogProvider,
existing_key: &StoredProviderCatalogKey,
) -> Option<StoredProviderCatalogKey> {
if !provider_key_is_oauth_managed(existing_key, &provider.provider_type)
|| existing_key.api_formats.is_none()
{
return None;
}
let mut updated = existing_key.clone();
updated.api_formats = None;
updated.updated_at_unix_secs = Some(current_unix_secs());
Some(updated)
}
fn sync_override<T>(overrides: &mut BTreeSet<String>, key: &str, actual: &T, desired: &T)
where
T: PartialEq,
@@ -624,4 +716,68 @@ mod tests {
.expect("endpoint should reconcile");
assert_eq!(reconciled.base_url, "http://127.0.0.1:18181/v1");
}
#[test]
fn fixed_provider_endpoint_reconcile_adopts_existing_customization() {
let provider = sample_codex_provider();
let template = fixed_provider_template("codex").expect("codex template should exist");
let endpoint_template = template
.endpoints
.iter()
.find(|endpoint| endpoint.api_format == "openai:responses")
.expect("responses endpoint template should exist");
let mut existing = sample_codex_endpoint("http://127.0.0.1:18181/backend-api/codex");
existing.is_active = false;
existing.max_retries = Some(9);
existing.proxy = Some(serde_json::json!({"url": "http://proxy.internal:8080"}));
existing.config = Some(serde_json::json!({
"upstream_stream_policy": "force_non_stream",
"custom_transport_option": true
}));
let reconciled =
reconcile_fixed_provider_endpoint(&provider, &existing, template, endpoint_template)
.expect("endpoint should reconcile");
assert_eq!(
reconciled.base_url,
"http://127.0.0.1:18181/backend-api/codex"
);
assert!(!reconciled.is_active);
assert_eq!(reconciled.max_retries, Some(9));
assert_eq!(
reconciled.proxy,
Some(serde_json::json!({"url": "http://proxy.internal:8080"}))
);
assert_eq!(
reconciled
.config
.as_ref()
.and_then(|value| value.get("upstream_stream_policy")),
Some(&serde_json::json!("force_non_stream"))
);
assert_eq!(
reconciled
.config
.as_ref()
.and_then(|value| value.get("custom_transport_option")),
Some(&serde_json::json!(true))
);
let metadata = fixed_provider_endpoint_metadata(&reconciled)
.expect("fixed provider metadata should exist");
for key in [
"base_url",
"is_active",
"max_retries",
"proxy",
"config.upstream_stream_policy",
] {
assert!(metadata.overrides.contains(key), "missing override {key}");
}
let second =
reconcile_fixed_provider_endpoint(&provider, &reconciled, template, endpoint_template)
.expect("endpoint should reconcile idempotently");
assert_eq!(second, reconciled);
}
}
@@ -202,21 +202,13 @@ pub(crate) async fn build_admin_update_provider_record(
}
if fields.contains("stream_first_byte_timeout") {
updated.stream_first_byte_timeout_secs = match payload.stream_first_byte_timeout {
Some(value) if (1.0..=300.0).contains(&value) => Some(value),
Some(_) => {
return Err("stream_first_byte_timeout 必须是 1 到 300 之间的数字".to_string());
}
None => None,
};
updated.stream_first_byte_timeout_secs =
super::normalize_provider_stream_first_byte_timeout(payload.stream_first_byte_timeout)?;
}
if fields.contains("request_timeout") {
updated.request_timeout_secs = match payload.request_timeout {
Some(value) if (1.0..=600.0).contains(&value) => Some(value),
Some(_) => return Err("request_timeout 必须是 1 到 600 之间的数字".to_string()),
None => None,
};
updated.request_timeout_secs =
super::normalize_provider_request_timeout(payload.request_timeout)?;
}
if fields.contains("enable_format_conversion") {
@@ -1,5 +1,26 @@
use super::*;
fn validate_admin_endpoint_stream_policy(
api_format: &str,
config: Option<&serde_json::Value>,
) -> Result<(), String> {
if !crate::ai_serving::api_format_alias_matches(api_format, "openai:search") {
return Ok(());
}
let requested = config
.and_then(serde_json::Value::as_object)
.and_then(|config| {
config
.get("upstream_stream_policy")
.or_else(|| config.get("upstreamStreamPolicy"))
.or_else(|| config.get("upstream_stream"))
});
if requested.is_some_and(crate::handlers::public::admin_requested_force_stream) {
return Err("OpenAI Search 端点仅支持非流式上游请求".to_string());
}
Ok(())
}
impl<'a> AdminAppState<'a> {
pub(crate) async fn build_admin_keys_grouped_by_format_payload(
&self,
@@ -232,6 +253,7 @@ impl<'a> AdminAppState<'a> {
let (normalized_api_format, api_family, endpoint_kind) =
admin_endpoint_signature_parts(&payload.api_format)
.ok_or_else(|| format!("无效的 api_format: {}", payload.api_format))?;
validate_admin_endpoint_stream_policy(normalized_api_format, payload.config.as_ref())?;
let base_url = normalize_admin_base_url(&payload.base_url)?;
let existing_endpoints = self
@@ -337,6 +359,13 @@ impl<'a> AdminAppState<'a> {
&update_fields,
)?;
if fields.contains("config") {
validate_admin_endpoint_stream_policy(
existing_endpoint.api_format.as_str(),
updated.config.as_ref(),
)?;
}
if provider_type == "codex"
&& crate::ai_serving::is_openai_responses_format(&existing_endpoint.api_format)
{
@@ -389,3 +418,40 @@ impl<'a> AdminAppState<'a> {
Ok(updated)
}
}
#[cfg(test)]
mod tests {
use serde_json::json;
use super::validate_admin_endpoint_stream_policy;
#[test]
fn search_endpoint_rejects_explicit_streaming_policy_for_all_config_keys() {
for (api_format, key, value) in [
(
"openai:search",
"upstream_stream_policy",
json!("force_stream"),
),
("openai:search", "upstreamStreamPolicy", json!(true)),
("/v1/alpha/search", "upstream_stream", json!("sse")),
] {
let config = json!({(key): value});
assert!(validate_admin_endpoint_stream_policy(api_format, Some(&config),).is_err());
}
}
#[test]
fn search_endpoint_accepts_non_streaming_and_unrelated_config() {
assert!(validate_admin_endpoint_stream_policy(
"openai:search",
Some(&json!({"upstream_stream_policy": "force_non_stream"})),
)
.is_ok());
assert!(validate_admin_endpoint_stream_policy(
"openai:responses",
Some(&json!({"upstream_stream_policy": "force_stream"})),
)
.is_ok());
}
}
+177 -14
View File
@@ -83,6 +83,8 @@ const OPENAI_RESPONSES_LOCAL_EXECUTION_RUNTIME_MISS_DETAIL: &str =
"当前 OpenAI Responses 请求无法在本地执行:没有匹配到可用的执行路径";
const OPENAI_RESPONSES_COMPACT_LOCAL_EXECUTION_RUNTIME_MISS_DETAIL: &str =
"当前 OpenAI Responses Compact 请求无法在本地执行:没有匹配到可用的执行路径";
const OPENAI_SEARCH_LOCAL_EXECUTION_RUNTIME_MISS_DETAIL: &str =
"当前 OpenAI Search 请求无法在本地执行:没有匹配到可用的执行路径";
const OPENAI_VIDEO_LOCAL_EXECUTION_RUNTIME_MISS_DETAIL: &str =
"当前 OpenAI Video 请求无法在本地执行:没有匹配到可用的执行路径";
const CLAUDE_MESSAGES_LOCAL_EXECUTION_RUNTIME_MISS_DETAIL: &str =
@@ -565,8 +567,9 @@ async fn maybe_forward_public_request_to_tunnel_owner(
serde_json::from_slice::<serde_json::Value>(body.as_ref()).ok()
});
let client_session_affinity =
crate::client_session_affinity::client_session_affinity_from_parts(
parts,
crate::client_session_affinity::client_session_affinity_from_api_request(
api_format,
&parts.headers,
body_json.as_ref(),
);
let Some(target) = crate::scheduler::affinity::read_cached_scheduler_affinity_target(
@@ -636,7 +639,27 @@ async fn maybe_forward_public_request_to_tunnel_owner(
owner.relay_base_url.trim_end_matches('/'),
request_context.request_path_and_query()
);
let mut upstream_request = state.client.request(parts.method.clone(), owner_url);
let is_stream =
owner_forward_request_is_stream(parts, decision, buffered_body.unwrap_or(&empty_body));
let transport_timeouts =
crate::provider_transport::resolve_transport_execution_timeouts(&transport);
let non_stream_timeout =
crate::execution_runtime::transport::resolve_non_stream_total_timeout_for_request(
is_stream,
&transport.endpoint.api_format,
transport_timeouts.as_ref(),
);
let stream_first_byte_timeout =
crate::execution_runtime::transport::resolve_stream_first_byte_timeout_for_request(
is_stream,
transport_timeouts.as_ref(),
);
let mut upstream_request = state
.owner_forward_client
.request(parts.method.clone(), owner_url);
if let Some(timeout) = non_stream_timeout {
upstream_request = upstream_request.timeout(timeout);
}
for (name, value) in &parts.headers {
if should_skip_request_header(name.as_str()) || name == http::header::HOST {
continue;
@@ -682,14 +705,15 @@ async fn maybe_forward_public_request_to_tunnel_owner(
upstream_request.header(TRUSTED_AUTH_BALANCE_HEADER, balance_remaining.to_string());
}
let upstream_response = upstream_request
.body(buffered_body.cloned().unwrap_or_default())
.send()
.await
.map_err(|err| GatewayError::UpstreamUnavailable {
trace_id: request_context.trace_id.clone(),
message: format!("owner gateway affinity forward failed: {err}"),
})?;
let upstream_response = crate::tunnel::send_owner_forward_request(
upstream_request.body(buffered_body.cloned().unwrap_or_default()),
stream_first_byte_timeout,
)
.await
.map_err(|message| GatewayError::UpstreamUnavailable {
trace_id: request_context.trace_id.clone(),
message: format!("owner gateway affinity forward failed: {message}"),
})?;
let mut response = build_sync_aware_affinity_forward_response(
request_context,
@@ -707,6 +731,29 @@ async fn maybe_forward_public_request_to_tunnel_owner(
Ok(Some(response))
}
fn owner_forward_request_is_stream(
parts: &http::request::Parts,
decision: &GatewayControlDecision,
body_bytes: &Bytes,
) -> bool {
let Some(plan_kind) =
crate::ai_serving::api::resolve_execution_runtime_stream_plan_kind(parts, decision)
else {
return false;
};
let Some((body_json, body_base64)) =
crate::ai_serving::api::parse_direct_request_body(parts, body_bytes)
else {
return false;
};
crate::ai_serving::api::is_matching_stream_request(
plan_kind,
parts,
&body_json,
body_base64.as_deref(),
)
}
fn upstream_response_is_sse(headers: &reqwest::header::HeaderMap) -> bool {
headers
.get(http::header::CONTENT_TYPE)
@@ -2372,6 +2419,7 @@ fn local_execution_runtime_miss_route_label(
"/v1/chat/completions" => "OpenAI Chat Completions",
"/v1/responses" => "OpenAI Responses",
"/v1/responses/compact" => "OpenAI Responses Compact",
"/v1/alpha/search" => "OpenAI Search",
"/v1/messages" => "Claude Messages",
path if path.starts_with("/v1/videos") => "OpenAI Video",
path if path.starts_with("/upload/v1beta/files") || path.starts_with("/v1beta/files") => {
@@ -2414,6 +2462,7 @@ fn local_execution_runtime_miss_route_detail(
"/v1/responses/compact" => {
Some(OPENAI_RESPONSES_COMPACT_LOCAL_EXECUTION_RUNTIME_MISS_DETAIL)
}
"/v1/alpha/search" => Some(OPENAI_SEARCH_LOCAL_EXECUTION_RUNTIME_MISS_DETAIL),
"/v1/messages" => Some(CLAUDE_MESSAGES_LOCAL_EXECUTION_RUNTIME_MISS_DETAIL),
path if path.starts_with("/v1/videos") => {
Some(OPENAI_VIDEO_LOCAL_EXECUTION_RUNTIME_MISS_DETAIL)
@@ -2437,14 +2486,128 @@ mod tests {
use super::{
api_key_remote_ip_allowed, buffer_and_normalize_request_body,
diagnostic_is_auth_api_key_concurrency_limited, local_execution_runtime_miss_detail,
restore_redacted_stream_execution_response, restore_redacted_sync_execution_response,
GatewayControlDecision, LocalExecutionRuntimeMissDiagnostic, RequestBodyBufferError,
RequestBodyBufferPolicy,
owner_forward_request_is_stream, restore_redacted_stream_execution_response,
restore_redacted_sync_execution_response, GatewayControlDecision,
LocalExecutionRuntimeMissDiagnostic, RequestBodyBufferError, RequestBodyBufferPolicy,
};
use axum::body::{to_bytes, Body, Bytes};
use axum::http::{header, HeaderMap, Method, Response};
use serde_json::json;
#[test]
fn owner_forward_uses_search_protocol_timeout_semantics() {
let request = http::Request::builder()
.method(Method::POST)
.uri("/v1/alpha/search")
.header(header::CONTENT_TYPE, "application/json")
.body(())
.expect("request should build");
let (parts, _) = request.into_parts();
let decision = GatewayControlDecision::synthetic(
"/v1/alpha/search",
Some("ai_public".to_string()),
Some("openai".to_string()),
Some("search".to_string()),
Some("openai:search".to_string()),
);
let body =
Bytes::from_static(br#"{"model":"gpt-5.6-sol","input":"find docs","stream":true}"#);
let is_stream = owner_forward_request_is_stream(&parts, &decision, &body);
let timeouts = aether_contracts::ExecutionTimeouts {
total_ms: Some(aether_contracts::MAX_EXECUTION_REQUEST_TIMEOUT_MS),
first_byte_ms: Some(10),
..aether_contracts::ExecutionTimeouts::default()
};
assert!(!is_stream);
assert_eq!(
crate::execution_runtime::transport::resolve_non_stream_total_timeout_for_request(
is_stream,
"openai:search",
Some(&timeouts),
),
Some(Duration::from_millis(
aether_contracts::MAX_EXECUTION_REQUEST_TIMEOUT_MS
))
);
assert_eq!(
crate::execution_runtime::transport::resolve_stream_first_byte_timeout_for_request(
is_stream,
Some(&timeouts),
),
None
);
}
#[test]
fn owner_forward_keeps_streaming_for_stream_capable_protocols() {
let chat_request = http::Request::builder()
.method(Method::POST)
.uri("/v1/chat/completions")
.header(header::CONTENT_TYPE, "application/json")
.body(())
.expect("request should build");
let (chat_parts, _) = chat_request.into_parts();
let chat_decision = GatewayControlDecision::synthetic(
"/v1/chat/completions",
Some("ai_public".to_string()),
Some("openai".to_string()),
Some("chat".to_string()),
Some("openai:chat".to_string()),
);
assert!(owner_forward_request_is_stream(
&chat_parts,
&chat_decision,
&Bytes::from_static(br#"{"model":"gpt-5.6-sol","stream":true}"#),
));
assert!(!owner_forward_request_is_stream(
&chat_parts,
&chat_decision,
&Bytes::from_static(br#"{"model":"gpt-5.6-sol","stream":false}"#),
));
let image_request = http::Request::builder()
.method(Method::POST)
.uri("/v1/images/generations")
.header(header::CONTENT_TYPE, "application/json")
.body(())
.expect("request should build");
let (image_parts, _) = image_request.into_parts();
let image_decision = GatewayControlDecision::synthetic(
"/v1/images/generations",
Some("ai_public".to_string()),
Some("openai".to_string()),
Some("image".to_string()),
Some("openai:image".to_string()),
);
assert!(owner_forward_request_is_stream(
&image_parts,
&image_decision,
&Bytes::from_static(br#"{"model":"gpt-image-1","stream":true}"#),
));
let compact_request = http::Request::builder()
.method(Method::POST)
.uri("/v1/responses/compact")
.header(header::CONTENT_TYPE, "application/json")
.body(())
.expect("request should build");
let (compact_parts, _) = compact_request.into_parts();
let compact_decision = GatewayControlDecision::synthetic(
"/v1/responses/compact",
Some("ai_public".to_string()),
Some("openai".to_string()),
Some("responses:compact".to_string()),
Some("openai:responses:compact".to_string()),
);
assert!(!owner_forward_request_is_stream(
&compact_parts,
&compact_decision,
&Bytes::from_static(br#"{"model":"gpt-5.6-sol","stream":true}"#),
));
}
#[test]
fn api_key_remote_ip_allows_unrestricted_keys() {
let remote_ip = "203.0.113.10".parse().expect("valid ip");
@@ -220,14 +220,7 @@ fn users_me_usage_api_format_defaults_to_non_stream(item: &StoredRequestUsageAud
let Some(value) = api_format else {
return false;
};
matches!(
crate::ai_serving::normalize_api_format_alias(value).as_str(),
"openai:chat"
| "openai:responses"
| "openai:responses:compact"
| "openai:image"
| "claude:messages"
)
aether_ai_formats::api_format_defaults_to_non_stream(value)
}
fn users_me_usage_request_body_implies_default_non_stream(item: &StoredRequestUsageAudit) -> bool {
@@ -1846,6 +1839,27 @@ mod tests {
assert_eq!(active_payload["client_is_stream"], false);
}
#[test]
fn user_usage_stream_defaults_to_non_stream_for_openai_search() {
let item = StoredRequestUsageAudit {
is_stream: false,
api_format: Some("openai:search".to_string()),
request_body: Some(json!({
"id": "session-search-1",
"model": "gpt-5.6-sol",
"input": "current documentation"
})),
..sample_usage("completed")
};
assert!(!users_me_usage_client_is_stream(&item));
let record_payload =
build_users_me_usage_record_payload(&item, false, &BTreeMap::new(), false);
assert_eq!(record_payload["client_requested_stream"], false);
assert_eq!(record_payload["client_is_stream"], false);
}
#[test]
fn user_usage_upstream_stream_prefers_request_metadata_flag() {
let item = StoredRequestUsageAudit {
@@ -38,7 +38,7 @@ pub(crate) fn provider_catalog_key_supports_format(
}
formats
.iter()
.any(|candidate| crate::ai_serving::api_format_alias_matches(candidate, api_format))
.any(|candidate| aether_ai_formats::api_format_permission_covers(candidate, api_format))
}
pub(crate) fn decrypt_catalog_secret_with_fallbacks(
@@ -2732,6 +2732,25 @@ mod tests {
.expect("key transport should build")
}
#[test]
fn responses_key_scope_covers_search_in_one_direction() {
let mut responses_key = sample_catalog_key();
responses_key.api_formats = Some(json!(["openai:responses"]));
assert!(provider_catalog_key_supports_format(
&responses_key,
"codex",
"openai:search",
));
let mut search_key = sample_catalog_key();
search_key.api_formats = Some(json!(["openai:search"]));
assert!(!provider_catalog_key_supports_format(
&search_key,
"codex",
"openai:responses",
));
}
#[test]
fn masked_catalog_api_key_handles_unicode_plaintext_without_panicking() {
let state = AppState::new().expect("gateway should build");
+2 -2
View File
@@ -12,8 +12,8 @@ pub(crate) use runtime::{
restore_proxy_upgrade_rollout_skipped_nodes, retry_proxy_upgrade_rollout_node,
run_admin_system_cleanup_once, run_manual_usage_cleanup_once, skip_proxy_upgrade_rollout_node,
spawn_account_self_check_worker, spawn_audit_cleanup_worker, spawn_db_maintenance_worker,
spawn_gemini_file_mapping_cleanup_worker, spawn_oauth_token_refresh_worker,
spawn_pending_cleanup_worker, spawn_pool_monitor_worker,
spawn_fixed_provider_reconciliation_task, spawn_gemini_file_mapping_cleanup_worker,
spawn_oauth_token_refresh_worker, spawn_pending_cleanup_worker, spawn_pool_monitor_worker,
spawn_pool_quota_probe_replenish_for_request, spawn_pool_quota_probe_worker,
spawn_pool_score_rebuild_worker, spawn_provider_checkin_worker,
spawn_provider_quota_alert_worker, spawn_proxy_node_metrics_cleanup_worker,
@@ -16,6 +16,8 @@ mod cleanup_runs;
mod config;
#[path = "runtime/db_maintenance.rs"]
mod db_maintenance;
#[path = "runtime/fixed_provider_reconciliation.rs"]
mod fixed_provider_reconciliation;
#[path = "runtime/oauth_token_refresh.rs"]
mod oauth_token_refresh;
#[path = "runtime/pending_cleanup.rs"]
@@ -71,6 +73,9 @@ pub(crate) use cleanup_runs::{
};
use config::*;
use db_maintenance::*;
pub(crate) use fixed_provider_reconciliation::{
perform_fixed_provider_reconciliation_once, spawn_fixed_provider_reconciliation_task,
};
pub(crate) use oauth_token_refresh::{
perform_oauth_token_refresh_once, OAuthTokenRefreshRunSummary,
};
@@ -0,0 +1,301 @@
use std::time::Duration;
use tracing::warn;
use crate::admin_api::{reconcile_admin_fixed_provider_template_endpoints, AdminAppState};
use crate::task_runtime::{
spawn_fire_and_forget, task_definition, TASK_KEY_FIXED_PROVIDER_RECONCILIATION,
};
use crate::{AppState, GatewayError};
const FIXED_PROVIDER_RECONCILIATION_LOCK_KEY: &str =
"task_runtime:lock:maintenance.provider.fixed_template.reconcile";
const FIXED_PROVIDER_RECONCILIATION_LOCK_TTL: Duration = Duration::from_secs(10 * 60);
const FIXED_PROVIDER_RECONCILIATION_RETRY_DELAY: Duration = Duration::from_secs(2);
const RECONCILED_PROVIDER_TYPE: &str = "codex";
pub(crate) async fn perform_fixed_provider_reconciliation_once(
state: &AppState,
) -> Result<bool, GatewayError> {
if !state.has_provider_catalog_data_reader() || !state.has_provider_catalog_data_writer() {
return Ok(false);
}
let Some(lock) = state
.runtime_state
.lock_try_acquire(
FIXED_PROVIDER_RECONCILIATION_LOCK_KEY,
state.tunnel.local_instance_id(),
FIXED_PROVIDER_RECONCILIATION_LOCK_TTL,
)
.await
.map_err(|error| GatewayError::Internal(error.to_string()))?
else {
return Ok(false);
};
let result = reconcile_fixed_provider_templates(state).await;
if let Err(error) = state.runtime_state.lock_release(&lock).await {
warn!(
event_name = "fixed_provider_reconciliation_lock_release_failed",
log_type = "ops",
error = ?error,
"gateway fixed provider reconciliation lock release failed"
);
}
result.map(|()| true)
}
async fn reconcile_fixed_provider_templates(state: &AppState) -> Result<(), GatewayError> {
let providers = state.list_provider_catalog_providers(false).await?;
let admin_state = AdminAppState::new(state);
let mut failures = Vec::new();
for provider in &providers {
if !provider
.provider_type
.trim()
.eq_ignore_ascii_case(RECONCILED_PROVIDER_TYPE)
|| admin_state
.fixed_provider_template(&provider.provider_type)
.is_none()
{
continue;
}
if let Err(error) =
reconcile_admin_fixed_provider_template_endpoints(&admin_state, provider).await
{
failures.push(format!(
"provider {} endpoint reconciliation failed: {error:?}",
provider.id,
));
continue;
}
}
if !failures.is_empty() {
return Err(GatewayError::Internal(failures.join("; ")));
}
Ok(())
}
async fn perform_fixed_provider_reconciliation_with_retry(state: &AppState) {
let max_attempts = task_definition(TASK_KEY_FIXED_PROVIDER_RECONCILIATION)
.map(|definition| definition.retry_policy.max_attempts)
.unwrap_or(1)
.max(1);
for attempt in 1..=max_attempts {
match perform_fixed_provider_reconciliation_once(state).await {
Ok(_) => return,
Err(error) if attempt < max_attempts => {
warn!(
event_name = "fixed_provider_reconciliation_retrying",
log_type = "ops",
attempt,
max_attempts,
error = ?error,
"gateway fixed provider reconciliation will retry"
);
tokio::time::sleep(FIXED_PROVIDER_RECONCILIATION_RETRY_DELAY).await;
}
Err(error) => {
warn!(
event_name = "fixed_provider_reconciliation_failed",
log_type = "ops",
attempt,
max_attempts,
error = ?error,
"gateway fixed provider reconciliation failed"
);
return;
}
}
}
}
pub(crate) fn spawn_fixed_provider_reconciliation_task(
state: AppState,
) -> Option<tokio::task::JoinHandle<()>> {
if !state.has_provider_catalog_data_reader() || !state.has_provider_catalog_data_writer() {
return None;
}
Some(spawn_fire_and_forget(
TASK_KEY_FIXED_PROVIDER_RECONCILIATION,
async move {
perform_fixed_provider_reconciliation_with_retry(&state).await;
},
))
}
#[cfg(test)]
mod tests {
use std::sync::Arc;
use aether_data::repository::provider_catalog::InMemoryProviderCatalogReadRepository;
use aether_data_contracts::repository::provider_catalog::{
ProviderCatalogReadRepository, StoredProviderCatalogEndpoint, StoredProviderCatalogKey,
StoredProviderCatalogProvider,
};
use serde_json::json;
use super::{
perform_fixed_provider_reconciliation_once, FIXED_PROVIDER_RECONCILIATION_LOCK_KEY,
FIXED_PROVIDER_RECONCILIATION_LOCK_TTL,
};
use crate::data::GatewayDataState;
use crate::AppState;
#[tokio::test]
async fn fixed_provider_reconciliation_respects_runtime_singleton_lock() {
let repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![],
vec![],
vec![],
));
let state = AppState::new()
.expect("gateway state should build")
.with_data_state_for_tests(
GatewayDataState::with_provider_catalog_repository_for_tests(repository),
);
let lock = state
.runtime_state
.lock_try_acquire(
FIXED_PROVIDER_RECONCILIATION_LOCK_KEY,
"another-gateway",
FIXED_PROVIDER_RECONCILIATION_LOCK_TTL,
)
.await
.expect("runtime lock should be available")
.expect("runtime lock should be acquired");
assert!(!perform_fixed_provider_reconciliation_once(&state)
.await
.expect("locked reconciliation should skip"));
assert!(state
.runtime_state
.lock_release(&lock)
.await
.expect("runtime lock should release"));
assert!(perform_fixed_provider_reconciliation_once(&state)
.await
.expect("unlocked reconciliation should run"));
}
#[tokio::test]
async fn fixed_provider_reconciliation_preserves_existing_endpoint_and_is_idempotent() {
let mut provider = StoredProviderCatalogProvider::new(
"provider-codex".to_string(),
"Codex".to_string(),
None,
"codex".to_string(),
)
.expect("provider should build");
provider.is_active = false;
provider.max_retries = Some(2);
let mut responses = StoredProviderCatalogEndpoint::new(
"endpoint-codex-responses".to_string(),
provider.id.clone(),
"openai:responses".to_string(),
Some("openai".to_string()),
Some("cli".to_string()),
false,
)
.expect("endpoint should build")
.with_transport_fields(
"http://127.0.0.1:18181/backend-api/codex".to_string(),
None,
None,
Some(9),
None,
Some(json!({"upstream_stream_policy": "force_non_stream"})),
None,
Some(json!({"url": "http://proxy.internal:8080"})),
)
.expect("endpoint transport should build");
responses.updated_at_unix_secs = Some(100);
let mut key = StoredProviderCatalogKey::new(
"key-codex".to_string(),
provider.id.clone(),
"oauth".to_string(),
"oauth".to_string(),
None,
true,
)
.expect("key should build");
key.api_formats = Some(json!(["openai:responses"]));
let unrelated_fixed_provider = StoredProviderCatalogProvider::new(
"provider-claude-code".to_string(),
"Claude Code".to_string(),
None,
"claude_code".to_string(),
)
.expect("unrelated fixed provider should build");
let repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![provider, unrelated_fixed_provider],
vec![responses],
vec![key],
));
let state = AppState::new()
.expect("gateway state should build")
.with_data_state_for_tests(
GatewayDataState::with_provider_catalog_repository_for_tests(repository.clone()),
);
assert!(perform_fixed_provider_reconciliation_once(&state)
.await
.expect("reconciliation should run"));
let first_endpoints = repository
.list_endpoints_by_provider_ids(&["provider-codex".to_string()])
.await
.expect("endpoints should list");
assert_eq!(first_endpoints.len(), 4);
let responses = first_endpoints
.iter()
.find(|endpoint| endpoint.api_format == "openai:responses")
.expect("responses endpoint should exist");
assert_eq!(
responses.base_url,
"http://127.0.0.1:18181/backend-api/codex"
);
assert!(!responses.is_active);
assert_eq!(responses.max_retries, Some(9));
assert_eq!(
responses.proxy,
Some(json!({"url": "http://proxy.internal:8080"}))
);
assert_eq!(
responses
.config
.as_ref()
.and_then(|value| value.get("upstream_stream_policy")),
Some(&json!("force_non_stream"))
);
assert!(first_endpoints
.iter()
.any(|endpoint| endpoint.api_format == "openai:search"));
let keys = repository
.list_keys_by_provider_ids(&["provider-codex".to_string()])
.await
.expect("keys should list");
assert_eq!(keys.len(), 1);
assert_eq!(keys[0].api_formats, Some(json!(["openai:responses"])));
assert!(repository
.list_endpoints_by_provider_ids(&["provider-claude-code".to_string()])
.await
.expect("unrelated endpoints should list")
.is_empty());
assert!(perform_fixed_provider_reconciliation_once(&state)
.await
.expect("second reconciliation should run"));
let second_endpoints = repository
.list_endpoints_by_provider_ids(&["provider-codex".to_string()])
.await
.expect("endpoints should list again");
assert_eq!(second_endpoints, first_endpoints);
}
}
@@ -21,7 +21,8 @@ use super::{
pending_cleanup_timeout_minutes, plan_pending_cleanup_batch, provider_checkin_schedule,
proxy_node_metrics_cleanup_settings, record_proxy_upgrade_traffic_success,
run_db_maintenance_with, run_proxy_upgrade_rollout_once, spawn_account_self_check_worker,
spawn_audit_cleanup_worker, spawn_db_maintenance_worker, spawn_oauth_token_refresh_worker,
spawn_audit_cleanup_worker, spawn_db_maintenance_worker,
spawn_fixed_provider_reconciliation_task, spawn_oauth_token_refresh_worker,
spawn_pending_cleanup_worker, spawn_pool_monitor_worker, spawn_pool_quota_probe_worker,
spawn_provider_checkin_worker, spawn_proxy_node_stale_cleanup_worker,
spawn_proxy_upgrade_rollout_worker, spawn_stats_aggregation_worker,
@@ -74,6 +75,14 @@ async fn spawn_oauth_token_refresh_worker_skips_when_provider_catalog_unavailabl
assert!(spawn_oauth_token_refresh_worker(state).is_none());
}
#[tokio::test]
async fn spawn_fixed_provider_reconciliation_task_skips_when_provider_catalog_unavailable() {
let state = AppState::new()
.expect("gateway state should build")
.with_data_state_for_tests(GatewayDataState::disabled());
assert!(spawn_fixed_provider_reconciliation_task(state).is_none());
}
#[tokio::test]
async fn spawn_proxy_upgrade_rollout_worker_skips_when_system_config_unavailable() {
let repository = Arc::new(InMemoryProxyNodeRepository::seed(vec![]));
@@ -92,7 +92,10 @@ pub(crate) fn classify_local_failover(
return LocalFailoverClassification::RetryStatusCode;
}
if should_failover_local_upstream_status(input.status_code) {
if should_failover_local_upstream_status(
input.status_code,
policy.retry_client_errors_by_default,
) {
return LocalFailoverClassification::RetryUpstreamFailure;
}
@@ -109,8 +112,11 @@ pub(crate) fn local_failover_error_message(response_text: Option<&str>) -> Optio
.filter(|value| !value.is_empty())
}
fn should_failover_local_upstream_status(status_code: u16) -> bool {
status_code >= 400
fn should_failover_local_upstream_status(
status_code: u16,
retry_client_errors_by_default: bool,
) -> bool {
status_code >= 500 || status_code >= 400 && retry_client_errors_by_default
}
fn local_error_response_has_cyber_policy_code(response_text: Option<&str>) -> bool {
@@ -474,6 +480,39 @@ mod tests {
}
}
#[test]
fn classifier_passes_through_client_errors_when_protocol_default_disables_failover() {
let policy = LocalFailoverPolicy {
retry_client_errors_by_default: false,
..LocalFailoverPolicy::default()
};
for status_code in [400, 401, 429, 499] {
assert_eq!(
classify_local_failover(&policy, LocalFailoverInput::new(status_code, None)),
LocalFailoverClassification::UseDefault
);
}
assert_eq!(
classify_local_failover(&policy, LocalFailoverInput::new(500, None)),
LocalFailoverClassification::RetryUpstreamFailure
);
}
#[test]
fn classifier_explicit_continue_rule_overrides_protocol_client_error_default() {
let policy = LocalFailoverPolicy {
continue_status_codes: [429].into_iter().collect(),
retry_client_errors_by_default: false,
..LocalFailoverPolicy::default()
};
assert_eq!(
classify_local_failover(&policy, LocalFailoverInput::new(429, None)),
LocalFailoverClassification::RetryStatusCode
);
}
#[test]
fn classifier_keeps_embedded_rate_limit_error_in_success_response_on_default_path() {
assert_eq!(
@@ -243,7 +243,11 @@ fn local_client_session_affinity(report_context: Option<&Value>) -> Option<Clien
.get("original_request_body")
.filter(|value| !value.is_null());
crate::client_session_affinity::client_session_affinity_from_request(&headers, body_json)
crate::client_session_affinity::client_session_affinity_from_api_request(
report_context_string_field(Some(report_context), "client_api_format").unwrap_or_default(),
&headers,
body_json,
)
}
fn header_map_from_report_context(headers: Option<&Value>) -> http::HeaderMap {
@@ -7,7 +7,7 @@ use tracing::debug;
use crate::provider_transport::GatewayProviderTransportSnapshot;
use crate::AppState;
#[derive(Debug, Clone, Default, PartialEq, Eq)]
#[derive(Debug, Clone, PartialEq, Eq)]
pub(crate) struct LocalFailoverPolicy {
pub(crate) max_retries: Option<u64>,
pub(crate) stop_status_codes: BTreeSet<u16>,
@@ -15,6 +15,21 @@ pub(crate) struct LocalFailoverPolicy {
pub(crate) success_failover_patterns: Vec<LocalFailoverRegexRule>,
pub(crate) error_stop_patterns: Vec<LocalFailoverRegexRule>,
pub(crate) stop_cyber_policy_errors: bool,
pub(crate) retry_client_errors_by_default: bool,
}
impl Default for LocalFailoverPolicy {
fn default() -> Self {
Self {
max_retries: None,
stop_status_codes: BTreeSet::new(),
continue_status_codes: BTreeSet::new(),
success_failover_patterns: Vec::new(),
error_stop_patterns: Vec::new(),
stop_cyber_policy_errors: false,
retry_client_errors_by_default: true,
}
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
@@ -81,6 +96,10 @@ pub(crate) fn local_failover_policy_from_transport(
LocalFailoverPolicy {
max_retries,
retry_client_errors_by_default:
aether_ai_formats::api_format_defaults_to_client_error_failover(
&transport.endpoint.api_format,
),
stop_cyber_policy_errors: codex_cyber_flag_passthrough_enabled(
&transport.provider.provider_type,
transport.provider.config.as_ref(),
@@ -144,6 +163,10 @@ pub(crate) fn local_failover_policy_from_report_context(
.get("stop_cyber_policy_errors")
.and_then(Value::as_bool)
.unwrap_or(false),
retry_client_errors_by_default: object
.get("retry_client_errors_by_default")
.and_then(Value::as_bool)
.unwrap_or(true),
})
}
@@ -178,6 +201,7 @@ fn local_failover_policy_to_value(policy: &LocalFailoverPolicy) -> Value {
"success_failover_patterns": policy.success_failover_patterns.iter().map(local_failover_regex_rule_to_value).collect::<Vec<_>>(),
"error_stop_patterns": policy.error_stop_patterns.iter().map(local_failover_regex_rule_to_value).collect::<Vec<_>>(),
"stop_cyber_policy_errors": policy.stop_cyber_policy_errors,
"retry_client_errors_by_default": policy.retry_client_errors_by_default,
})
}
@@ -375,10 +399,26 @@ mod tests {
status_codes: [422].into_iter().collect(),
}],
stop_cyber_policy_errors: false,
retry_client_errors_by_default: true,
})
);
}
#[test]
fn search_transport_disables_default_client_error_failover() {
let mut transport = sample_transport(None, None, None);
transport.endpoint.api_format = "openai:search".to_string();
let policy = local_failover_policy_from_transport(&transport);
assert!(!policy.retry_client_errors_by_default);
let report_context = append_local_failover_policy_to_value(json!({}), &transport);
assert_eq!(
local_failover_policy_from_report_context(Some(&report_context))
.map(|policy| policy.retry_client_errors_by_default),
Some(false)
);
}
#[test]
fn codex_cyber_policy_passthrough_defaults_on_and_can_be_disabled() {
let mut transport = sample_transport(None, None, None);
+180 -1
View File
@@ -1063,6 +1063,7 @@ impl MaskChatRequestOptions {
pub(crate) enum ChatPiiRedactionRequestFormat {
OpenAiChat,
OpenAiResponses,
OpenAiSearch,
ClaudeMessages,
}
@@ -1071,6 +1072,7 @@ impl ChatPiiRedactionRequestFormat {
match api_format.trim().to_ascii_lowercase().as_str() {
"openai:chat" => Some(Self::OpenAiChat),
"openai:responses" | "openai:responses:compact" => Some(Self::OpenAiResponses),
"openai:search" => Some(Self::OpenAiSearch),
"claude:messages" => Some(Self::ClaudeMessages),
_ => None,
}
@@ -1604,6 +1606,7 @@ fn request_collision_corpus(format: ChatPiiRedactionRequestFormat, value: &Value
.map(|messages| chat_message_collision_corpus(messages))
.unwrap_or_default(),
ChatPiiRedactionRequestFormat::OpenAiResponses => openai_responses_collision_corpus(value),
ChatPiiRedactionRequestFormat::OpenAiSearch => openai_search_collision_corpus(value),
ChatPiiRedactionRequestFormat::ClaudeMessages => claude_messages_collision_corpus(value),
}
}
@@ -1622,6 +1625,9 @@ fn mask_request_value(
ChatPiiRedactionRequestFormat::OpenAiResponses => {
mask_openai_responses_request_value(value, session, scan_state)
}
ChatPiiRedactionRequestFormat::OpenAiSearch => {
mask_openai_search_request_value(value, session, scan_state)
}
ChatPiiRedactionRequestFormat::ClaudeMessages => {
mask_claude_messages_request_value(value, session, scan_state)
}
@@ -1643,6 +1649,9 @@ async fn mask_request_value_async(
ChatPiiRedactionRequestFormat::OpenAiResponses => {
mask_openai_responses_request_value_async(value, session, scan_state, cache).await
}
ChatPiiRedactionRequestFormat::OpenAiSearch => {
mask_openai_search_request_value_async(value, session, scan_state, cache).await
}
ChatPiiRedactionRequestFormat::ClaudeMessages => {
mask_claude_messages_request_value_async(value, session, scan_state, cache).await
}
@@ -1810,6 +1819,31 @@ fn openai_responses_collision_corpus(value: &Value) -> Vec<String> {
corpus
}
const OPENAI_SEARCH_COMMAND_TEXT_FIELDS: [(&str, &str); 4] = [
("search_query", "q"),
("image_query", "q"),
("find", "pattern"),
("weather", "location"),
];
fn openai_search_collision_corpus(value: &Value) -> Vec<String> {
let mut corpus = openai_responses_collision_corpus(value);
let Some(commands) = value.get("commands").and_then(Value::as_object) else {
return corpus;
};
for (command, field) in OPENAI_SEARCH_COMMAND_TEXT_FIELDS {
let Some(entries) = commands.get(command).and_then(Value::as_array) else {
continue;
};
for entry in entries {
if let Some(text) = entry.get(field).and_then(Value::as_str) {
corpus.push(text.to_string());
}
}
}
corpus
}
fn collect_openai_responses_input_collision_text(value: &Value, corpus: &mut Vec<String>) {
match value {
Value::String(text) => corpus.push(text.clone()),
@@ -2025,6 +2059,31 @@ fn mask_openai_responses_request_value(
Ok(redacted)
}
fn mask_openai_search_request_value(
value: &mut Value,
session: &mut RedactionSession,
scan_state: &mut RedactionScanState,
) -> Result<bool, RedactionLimitError> {
let mut redacted = false;
if let Some(input) = value.get_mut("input") {
redacted |= mask_openai_responses_input_value(input, session, scan_state)?;
}
let Some(commands) = value.get_mut("commands").and_then(Value::as_object_mut) else {
return Ok(redacted);
};
for (command, field) in OPENAI_SEARCH_COMMAND_TEXT_FIELDS {
let Some(entries) = commands.get_mut(command).and_then(Value::as_array_mut) else {
continue;
};
for entry in entries {
if let Some(Value::String(text)) = entry.get_mut(field) {
redacted |= mask_json_string(text, session, scan_state)?;
}
}
}
Ok(redacted)
}
fn mask_openai_responses_input_value(
value: &mut Value,
session: &mut RedactionSession,
@@ -2310,6 +2369,33 @@ async fn mask_openai_responses_request_value_async(
Ok(redacted)
}
async fn mask_openai_search_request_value_async(
value: &mut Value,
session: &mut RedactionSession,
scan_state: &mut RedactionScanState,
cache: Option<&RedisRedactionMappingCache<'_>>,
) -> Result<bool, RedactionMaskError> {
let mut redacted = false;
if let Some(input) = value.get_mut("input") {
redacted |=
mask_openai_responses_input_value_async(input, session, scan_state, cache).await?;
}
let Some(commands) = value.get_mut("commands").and_then(Value::as_object_mut) else {
return Ok(redacted);
};
for (command, field) in OPENAI_SEARCH_COMMAND_TEXT_FIELDS {
let Some(entries) = commands.get_mut(command).and_then(Value::as_array_mut) else {
continue;
};
for entry in entries {
if let Some(Value::String(text)) = entry.get_mut(field) {
redacted |= mask_json_string_async(text, session, scan_state, cache).await?;
}
}
}
Ok(redacted)
}
async fn mask_openai_responses_input_value_async(
value: &mut Value,
session: &mut RedactionSession,
@@ -4072,6 +4158,7 @@ mod tests {
build_redaction_session_config, detect_candidates_with_probe, mask_chat_request_json,
mask_chat_request_json_with_options, parse_chat_pii_redaction_rules,
restore_sync_response_body, try_mask_chat_pii_request_json_with_options,
try_mask_chat_pii_request_value_with_cache_options,
try_mask_chat_request_json_with_cache_options, try_mask_chat_request_json_with_options,
ChatPiiRedactionRequestFormat, ChatPiiRedactionRuntimeConfig, DetectorProbe, MappingKey,
MaskChatRequestOptions, RedactionKind, RedactionLimitError, RedactionMapping,
@@ -4083,7 +4170,7 @@ mod tests {
use aether_runtime_state::{RedisClientConfig, RuntimeState};
use aether_testkit::ManagedRedisServer;
use serde_json::json;
use serde_json::{json, Value};
fn assert_debug_surface_hides_values(debug: &str, originals: &[&str], sentinels: &[String]) {
for original in originals {
@@ -4734,6 +4821,98 @@ mod tests {
.contains("secretValueABCDEF1234567890abcdef"));
}
#[test]
fn pii_redaction_request_masks_openai_search_text_fields() {
let request = json!({
"id": "session-1",
"model": "gpt-5.6",
"input": "Find alice@example.com",
"commands": {
"search_query": [{"q": "Phone +14155552671"}],
"image_query": [{"q": "Image for bob@example.com"}],
"find": [{"ref_id": "https://example.com/alice@example.com", "pattern": "secret_key=secretValueABCDEF1234567890abcdef"}],
"weather": [{"location": "Contact carol@example.com"}],
"open": [{"ref_id": "https://example.com/alice@example.com"}]
}
});
let raw = serde_json::to_vec(&request).expect("request should serialize");
let masked = try_mask_chat_pii_request_json_with_options(
&raw,
ChatPiiRedactionRequestFormat::OpenAiSearch,
test_config(),
MaskChatRequestOptions::runtime(),
)
.expect("search request should mask");
let masked_json: Value =
serde_json::from_slice(&masked.body).expect("masked request should parse");
assert!(masked.redacted);
assert!(!masked_json["input"]
.as_str()
.unwrap()
.contains("alice@example.com"));
assert!(!masked_json["commands"]["search_query"][0]["q"]
.as_str()
.unwrap()
.contains("+14155552671"));
assert!(!masked_json["commands"]["image_query"][0]["q"]
.as_str()
.unwrap()
.contains("bob@example.com"));
assert!(!masked_json["commands"]["find"][0]["pattern"]
.as_str()
.unwrap()
.contains("secretValueABCDEF1234567890abcdef"));
assert!(!masked_json["commands"]["weather"][0]["location"]
.as_str()
.unwrap()
.contains("carol@example.com"));
assert_eq!(
masked_json["commands"]["open"][0]["ref_id"],
"https://example.com/alice@example.com"
);
}
#[tokio::test]
async fn pii_redaction_async_request_masks_openai_search_text_fields() {
let request = json!({
"id": "session-1",
"model": "gpt-5.6",
"input": "Find alice@example.com",
"commands": {
"search_query": [{"q": "Phone +14155552671"}],
"find": [{"ref_id": "turn0search0", "pattern": "bob@example.com"}]
}
});
let masked = try_mask_chat_pii_request_value_with_cache_options(
&request,
ChatPiiRedactionRequestFormat::OpenAiSearch,
test_config(),
MaskChatRequestOptions::runtime(),
None,
)
.await
.expect("search request should mask");
let masked_json = masked.body_json.expect("masked body should be present");
assert!(masked.redacted);
assert!(!masked_json["input"]
.as_str()
.unwrap()
.contains("alice@example.com"));
assert!(!masked_json["commands"]["search_query"][0]["q"]
.as_str()
.unwrap()
.contains("+14155552671"));
assert!(!masked_json["commands"]["find"][0]["pattern"]
.as_str()
.unwrap()
.contains("bob@example.com"));
assert_eq!(masked_json["commands"]["find"][0]["ref_id"], "turn0search0");
}
#[test]
fn pii_redaction_request_avoids_cross_message_and_tool_argument_sentinel_collisions() {
let mut probe = session_at(600);
+1
View File
@@ -344,6 +344,7 @@ pub struct AppState {
pub(crate) upstream_target_admission: Arc<crate::upstream_admission::UpstreamTargetAdmission>,
pub(crate) distributed_request_gate: Option<Arc<RuntimeSemaphore>>,
pub(crate) client: reqwest::Client,
pub(crate) owner_forward_client: reqwest::Client,
pub(crate) auth_context_cache: Arc<AuthContextCache>,
pub(crate) auth_snapshot_cache: Arc<AuthSnapshotCache>,
pub(crate) user_model_capability_settings_cache: Arc<JsonValueCache<String>>,
+14
View File
@@ -51,6 +51,7 @@ use super::super::{provider_transport, usage};
use crate::maintenance::spawn_account_self_check_worker;
use crate::maintenance::spawn_audit_cleanup_worker;
use crate::maintenance::spawn_db_maintenance_worker;
use crate::maintenance::spawn_fixed_provider_reconciliation_task;
use crate::maintenance::spawn_gemini_file_mapping_cleanup_worker;
use crate::maintenance::spawn_oauth_token_refresh_worker;
use crate::maintenance::spawn_pending_cleanup_worker;
@@ -250,6 +251,11 @@ impl AppState {
http2_adaptive_window: true,
..HttpClientConfig::default()
})?;
let owner_forward_client = build_http_client(&HttpClientConfig {
connect_timeout_ms: Some(10_000),
http2_adaptive_window: true,
..HttpClientConfig::default()
})?;
let frontdoor_runtime_guards = Arc::new(FrontdoorRuntimeGuardConfig::from_env());
Ok(Self {
#[cfg(test)]
@@ -284,6 +290,7 @@ impl AppState {
),
distributed_request_gate: None,
client,
owner_forward_client,
auth_context_cache: Arc::new(AuthContextCache::default()),
auth_snapshot_cache: Arc::new(AuthSnapshotCache::default()),
user_model_capability_settings_cache: Arc::new(JsonValueCache::default()),
@@ -1463,6 +1470,13 @@ impl AppState {
record_boot(crate::task_runtime::TASK_KEY_USAGE_QUEUE_WORKER);
}
if let Some(handle) = spawn_fixed_provider_reconciliation_task(self.clone()) {
supervisor.supervise_handle(
crate::task_runtime::TASK_KEY_FIXED_PROVIDER_RECONCILIATION,
handle,
);
}
let mut supervise_worker =
|task_key: &'static str, handle: Option<tokio::task::JoinHandle<()>>| {
if let Some(handle) = handle {
@@ -32,6 +32,8 @@ pub(crate) const TASK_KEY_DB_MAINTENANCE: &str = "maintenance.database";
pub(crate) const TASK_KEY_PENDING_CLEANUP: &str = "maintenance.pending.cleanup";
pub(crate) const TASK_KEY_REQUEST_CANDIDATE_CLEANUP: &str = "maintenance.request.candidate.cleanup";
pub(crate) const TASK_KEY_GEMINI_FILES_CLEANUP: &str = "maintenance.gemini.files.cleanup";
pub(crate) const TASK_KEY_FIXED_PROVIDER_RECONCILIATION: &str =
"maintenance.provider.fixed_template.reconcile";
pub(crate) const TASK_KEY_OAUTH_TOKEN_REFRESH: &str = "maintenance.oauth.token.refresh";
pub(crate) const TASK_KEY_PROXY_NODE_STALE_CLEANUP: &str = "maintenance.proxy.node.stale.cleanup";
pub(crate) const TASK_KEY_PROXY_NODE_METRICS_CLEANUP: &str =
@@ -49,6 +51,7 @@ pub(crate) const TASK_KEY_PROVIDER_BALANCE_REFRESH: &str = "provider.ops.balance
const PROVIDER_DELETE_LOCK_TTL_SECS: u64 = 60 * 60 * 6;
const RETRY_ONCE: RetryPolicy = RetryPolicy { max_attempts: 1 };
const RETRY_THREE: RetryPolicy = RetryPolicy { max_attempts: 3 };
const TASK_DEFINITIONS: &[TaskDefinition] = &[
TaskDefinition::new(
@@ -187,6 +190,14 @@ const TASK_DEFINITIONS: &[TaskDefinition] = &[
true,
RETRY_ONCE,
),
TaskDefinition::new(
TASK_KEY_FIXED_PROVIDER_RECONCILIATION,
TaskKind::FireAndForget,
"startup",
true,
false,
RETRY_THREE,
),
TaskDefinition::new(
TASK_KEY_OAUTH_TOKEN_REFRESH,
TaskKind::Scheduled,
@@ -29,3 +29,4 @@ mod cli;
mod gemini;
mod image;
mod pii_redaction_formats;
mod search;
@@ -0,0 +1,597 @@
use super::{
any, build_router_with_state, build_state_with_execution_runtime_override, json, start_server,
to_bytes, Arc, Body, Json, Mutex, Request, Router, StatusCode,
EXECUTION_PATH_EXECUTION_RUNTIME_SYNC, EXECUTION_PATH_HEADER, TRACE_ID_HEADER,
};
use aether_crypto::{encrypt_python_fernet_plaintext, DEVELOPMENT_ENCRYPTION_KEY};
use aether_data::repository::auth::{
InMemoryAuthApiKeySnapshotRepository, StoredAuthApiKeySnapshot,
};
use aether_data::repository::candidate_selection::InMemoryMinimalCandidateSelectionReadRepository;
use aether_data::repository::candidates::InMemoryRequestCandidateRepository;
use aether_data::repository::provider_catalog::InMemoryProviderCatalogReadRepository;
use aether_data_contracts::repository::candidate_selection::{
StoredMinimalCandidateSelectionRow, StoredProviderModelMapping,
};
use aether_data_contracts::repository::candidates::{
RequestCandidateReadRepository, RequestCandidateStatus,
};
use aether_data_contracts::repository::provider_catalog::{
StoredProviderCatalogEndpoint, StoredProviderCatalogKey, StoredProviderCatalogProvider,
};
use sha2::{Digest, Sha256};
const SEARCH_SYNC_TEST_STACK_BYTES: usize = 16 * 1024 * 1024;
fn run_search_sync_test<F, Fut>(test_name: &'static str, make_future: F)
where
F: FnOnce() -> Fut + Send + 'static,
Fut: std::future::Future<Output = ()> + 'static,
{
let handle = std::thread::Builder::new()
.name(test_name.to_string())
.stack_size(SEARCH_SYNC_TEST_STACK_BYTES)
.spawn(move || {
let runtime = tokio::runtime::Builder::new_current_thread()
.enable_all()
.build()
.expect("test runtime should build");
runtime.block_on(make_future());
})
.expect("search sync test thread should spawn");
if let Err(payload) = handle.join() {
std::panic::resume_unwind(payload);
}
}
#[test]
fn gateway_executes_codex_search_with_responses_permission_and_search_contract() {
run_search_sync_test(
"gateway_executes_codex_search_with_responses_permission_and_search_contract",
gateway_executes_codex_search_with_responses_permission_and_search_contract_impl,
);
}
async fn gateway_executes_codex_search_with_responses_permission_and_search_contract_impl() {
fn hash_api_key(value: &str) -> String {
let mut hasher = Sha256::new();
hasher.update(value.as_bytes());
format!("{:x}", hasher.finalize())
}
fn auth_snapshot() -> StoredAuthApiKeySnapshot {
StoredAuthApiKeySnapshot::new(
"user-search-1".to_string(),
"alice".to_string(),
Some("alice@example.com".to_string()),
"user".to_string(),
"local".to_string(),
true,
false,
Some(json!(["openai", "codex"])),
Some(json!(["openai:responses"])),
None,
"api-key-search-1".to_string(),
Some("search-client".to_string()),
true,
false,
false,
Some(60),
Some(5),
Some(4_102_444_800_i64),
Some(json!(["openai", "codex"])),
Some(json!(["openai:responses"])),
None,
)
.expect("auth snapshot should build")
}
fn candidate_row() -> StoredMinimalCandidateSelectionRow {
StoredMinimalCandidateSelectionRow {
provider_id: "provider-codex-search-1".to_string(),
provider_name: "codex".to_string(),
provider_type: "codex".to_string(),
provider_priority: 10,
provider_is_active: true,
endpoint_id: "endpoint-codex-search-1".to_string(),
endpoint_api_format: "openai:search".to_string(),
endpoint_api_family: Some("openai".to_string()),
endpoint_kind: Some("search".to_string()),
endpoint_is_active: true,
key_id: "key-codex-search-1".to_string(),
key_name: "oauth".to_string(),
key_auth_type: "oauth".to_string(),
key_is_active: true,
key_api_formats: Some(vec!["openai:responses".to_string()]),
key_allowed_models: None,
key_capabilities: None,
key_internal_priority: 5,
key_global_priority_by_format: Some(json!({"openai:search": 1})),
model_id: "model-codex-search-1".to_string(),
global_model_id: "global-model-codex-search-1".to_string(),
global_model_name: "gpt-5.6-sol".to_string(),
global_model_mappings: None,
global_model_supports_streaming: Some(false),
model_provider_model_name: "gpt-5.6-sol".to_string(),
model_provider_model_mappings: Some(vec![StoredProviderModelMapping {
name: "gpt-5.6-sol".to_string(),
priority: 1,
api_formats: Some(vec!["openai:responses".to_string()]),
endpoint_ids: None,
}]),
model_supports_streaming: Some(false),
model_is_active: true,
model_is_available: true,
}
}
fn provider() -> StoredProviderCatalogProvider {
StoredProviderCatalogProvider::new(
"provider-codex-search-1".to_string(),
"codex".to_string(),
Some("https://chatgpt.com".to_string()),
"codex".to_string(),
)
.expect("provider should build")
.with_transport_fields(
true,
false,
false,
None,
Some(2),
None,
Some(900.0),
None,
None,
)
}
fn endpoint() -> StoredProviderCatalogEndpoint {
StoredProviderCatalogEndpoint::new(
"endpoint-codex-search-1".to_string(),
"provider-codex-search-1".to_string(),
"openai:search".to_string(),
Some("openai".to_string()),
Some("search".to_string()),
true,
)
.expect("endpoint should build")
.with_transport_fields(
"https://chatgpt.com/backend-api/codex".to_string(),
None,
None,
Some(2),
None,
None,
None,
None,
)
.expect("endpoint transport should build")
}
fn key() -> StoredProviderCatalogKey {
let auth_config = encrypt_python_fernet_plaintext(
DEVELOPMENT_ENCRYPTION_KEY,
r#"{"provider_type":"codex","account_id":"account-search-1","is_fedramp":true}"#,
)
.expect("auth config should encrypt");
StoredProviderCatalogKey::new(
"key-codex-search-1".to_string(),
"provider-codex-search-1".to_string(),
"oauth".to_string(),
"oauth".to_string(),
None,
true,
)
.expect("key should build")
.with_transport_fields(
Some(json!(["openai:responses"])),
encrypt_python_fernet_plaintext(
DEVELOPMENT_ENCRYPTION_KEY,
"codex-search-access-token",
)
.expect("access token should encrypt"),
Some(auth_config),
None,
Some(json!({"openai:search": 1})),
None,
Some(4_102_444_800),
None,
None,
)
.expect("key transport should build")
}
let seen_plans = Arc::new(Mutex::new(Vec::<serde_json::Value>::new()));
let seen_plans_clone = Arc::clone(&seen_plans);
let execution_runtime = Router::new().route(
"/v1/execute/sync",
any(move |request: Request| {
let seen_plans_inner = Arc::clone(&seen_plans_clone);
async move {
let (_, body) = request.into_parts();
let bytes = to_bytes(body, usize::MAX).await.expect("body should read");
let payload: serde_json::Value =
serde_json::from_slice(&bytes).expect("execution payload should parse");
let request_id = payload["request_id"]
.as_str()
.unwrap_or_default()
.to_string();
let provider_id = payload["provider_id"]
.as_str()
.unwrap_or_default()
.to_string();
seen_plans_inner
.lock()
.expect("mutex should lock")
.push(payload);
let execution_result = if request_id == "trace-search-error-1" {
json!({
"request_id": request_id,
"status_code": 400,
"headers": {
"content-type": "application/json",
"x-search-upstream": "rate-limited"
},
"body": {
"json_body": {
"error": {
"type": "rate_limit_error",
"message": "Search capacity reached",
"param": null,
"code": "rate_limit_exceeded"
},
"future_error_field": {"retryable": true}
}
},
"telemetry": {"elapsed_ms": 17}
})
} else if request_id == "trace-search-failover-1"
&& provider_id == "provider-codex-search-1"
{
json!({
"request_id": request_id,
"status_code": 500,
"headers": {
"content-type": "application/json",
"x-search-upstream": "primary"
},
"body": {
"json_body": {
"error": {
"type": "server_error",
"message": "Search backend unavailable"
}
}
},
"telemetry": {"elapsed_ms": 11}
})
} else if request_id == "trace-search-failover-1" {
json!({
"request_id": request_id,
"status_code": 200,
"headers": {
"content-type": "application/json",
"x-search-upstream": "backup"
},
"body": {
"json_body": {
"output": "search fallback result"
}
},
"telemetry": {"elapsed_ms": 23}
})
} else {
json!({
"request_id": request_id,
"status_code": 201,
"headers": {
"content-type": "application/json",
"x-search-upstream": "alpha"
},
"body": {
"json_body": {
"output": "search result",
"encrypted_output": "encrypted-search-result",
"future_response_field": {"enabled": true}
}
},
"telemetry": {"elapsed_ms": 42}
})
};
(
StatusCode::OK,
[("x-search-source", "codex-alpha")],
Json(execution_result),
)
}
}),
);
let client_api_key = "sk-client-search";
let auth_repository = Arc::new(InMemoryAuthApiKeySnapshotRepository::seed(vec![(
Some(hash_api_key(client_api_key)),
auth_snapshot(),
)]));
let candidate_repository = Arc::new(InMemoryMinimalCandidateSelectionReadRepository::seed({
let primary = candidate_row();
let mut backup = primary.clone();
backup.provider_id = "provider-codex-search-2".to_string();
backup.provider_name = "codex-backup".to_string();
backup.provider_priority = 20;
backup.endpoint_id = "endpoint-codex-search-2".to_string();
backup.key_id = "key-codex-search-2".to_string();
backup.key_name = "oauth-backup".to_string();
backup.key_internal_priority = 6;
backup.key_global_priority_by_format = Some(json!({"openai:search": 2}));
backup.model_id = "model-codex-search-2".to_string();
vec![primary, backup]
}));
let catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
{
let primary = provider();
let mut backup = primary.clone();
backup.id = "provider-codex-search-2".to_string();
backup.name = "codex-backup".to_string();
vec![primary, backup]
},
{
let primary = endpoint();
let mut backup = primary.clone();
backup.id = "endpoint-codex-search-2".to_string();
backup.provider_id = "provider-codex-search-2".to_string();
vec![primary, backup]
},
{
let primary = key();
let mut backup = primary.clone();
backup.id = "key-codex-search-2".to_string();
backup.provider_id = "provider-codex-search-2".to_string();
backup.name = "oauth-backup".to_string();
backup.global_priority_by_format = Some(json!({"openai:search": 2}));
vec![primary, backup]
},
));
let request_candidates = Arc::new(InMemoryRequestCandidateRepository::default());
let (execution_runtime_url, execution_runtime_handle) = start_server(execution_runtime).await;
let data_state =
crate::data::GatewayDataState::with_auth_candidate_selection_provider_catalog_and_request_candidate_repository_for_tests(
auth_repository,
candidate_repository,
catalog_repository,
Arc::clone(&request_candidates),
DEVELOPMENT_ENCRYPTION_KEY,
)
.with_system_config_values_for_tests([(
crate::system_features::ENABLE_MODEL_DIRECTIVES_CONFIG_KEY.to_string(),
json!(true),
)]);
let state = build_state_with_execution_runtime_override(execution_runtime_url)
.with_data_state_for_tests(data_state);
let gateway = build_router_with_state(state);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
.post(format!("{gateway_url}/v1/alpha/search"))
.header(http::header::CONTENT_TYPE, "application/json")
.header(
http::header::AUTHORIZATION,
format!("Bearer {client_api_key}"),
)
.header(TRACE_ID_HEADER, "trace-search-1")
.json(&json!({
"id": "session-search-1",
"model": "gpt-5.6-sol-ultra-fast",
"reasoning": {"effort": "low", "summary": "auto"},
"input": "find current OpenAI documentation",
"commands": {
"search_query": [{"q": "OpenAI Codex search"}],
"open": [{"ref_id": "turn0search0"}]
},
"settings": {
"search_context_size": "high",
"allowed_callers": ["direct"]
},
"max_output_tokens": 4096,
"store": false,
"stream": true,
"future_request_field": {"enabled": true}
}))
.send()
.await
.expect("search request should succeed");
if response.status() != StatusCode::CREATED {
let status = response.status();
let body = response.text().await.expect("error response should read");
panic!("Search request returned {status}: {body}");
}
assert_eq!(
response
.headers()
.get("x-search-upstream")
.and_then(|value| value.to_str().ok()),
Some("alpha")
);
assert_eq!(
response
.headers()
.get(EXECUTION_PATH_HEADER)
.and_then(|value| value.to_str().ok()),
Some(EXECUTION_PATH_EXECUTION_RUNTIME_SYNC)
);
let response_json: serde_json::Value = response.json().await.expect("response should parse");
assert_eq!(response_json["output"], "search result");
assert_eq!(response_json["encrypted_output"], "encrypted-search-result");
assert_eq!(response_json["future_response_field"]["enabled"], true);
let plan = seen_plans
.lock()
.expect("mutex should lock")
.first()
.cloned()
.expect("execution plan should be captured");
assert_eq!(
plan["url"],
"https://chatgpt.com/backend-api/codex/alpha/search"
);
assert_eq!(plan["client_api_format"], "openai:search");
assert_eq!(plan["provider_api_format"], "openai:search");
assert_eq!(plan["stream"], false);
assert_eq!(plan["timeouts"]["total_ms"], 900_000);
assert_eq!(
plan["headers"]["authorization"],
"Bearer codex-search-access-token"
);
assert_eq!(plan["headers"]["chatgpt-account-id"], "account-search-1");
assert_eq!(plan["headers"]["x-openai-fedramp"], "true");
assert_eq!(plan["headers"]["originator"], "codex_cli_rs");
assert!(plan["headers"]["user-agent"]
.as_str()
.is_some_and(|value| value.starts_with("codex_cli_rs/")));
assert!(plan["headers"].get("openai-beta").is_none());
assert!(plan["headers"]
.get("x-openai-internal-codex-responses-lite")
.is_none());
assert_ne!(plan["headers"]["accept"], "text/event-stream");
let body = &plan["body"]["json_body"];
assert_eq!(body["id"], "session-search-1");
assert_eq!(body["model"], "gpt-5.6-sol");
assert_eq!(body["reasoning"]["effort"], "max");
assert_eq!(body["reasoning"]["summary"], "auto");
assert_eq!(
body["commands"]["search_query"][0]["q"],
"OpenAI Codex search"
);
assert_eq!(body["commands"]["open"][0]["ref_id"], "turn0search0");
assert_eq!(body["settings"]["search_context_size"], "high");
assert_eq!(body["max_output_tokens"], 4096);
assert!(body.get("store").is_none());
assert!(body.get("future_request_field").is_none());
assert!(body.get("stream").is_none());
assert!(body.get("service_tier").is_none());
let candidates = request_candidates
.list_by_request_id("trace-search-1")
.await
.expect("request candidates should read");
assert_eq!(candidates.len(), 1);
assert_eq!(candidates[0].status, RequestCandidateStatus::Success);
let expected_error_body = json!({
"error": {
"type": "rate_limit_error",
"message": "Search capacity reached",
"param": null,
"code": "rate_limit_exceeded"
},
"future_error_field": {"retryable": true}
});
let error_response = reqwest::Client::new()
.post(format!("{gateway_url}/v1/alpha/search"))
.header(http::header::CONTENT_TYPE, "application/json")
.header(
http::header::AUTHORIZATION,
format!("Bearer {client_api_key}"),
)
.header(TRACE_ID_HEADER, "trace-search-error-1")
.json(&json!({
"id": "session-search-error-1",
"model": "gpt-5.6-sol",
"input": "find current OpenAI documentation",
"commands": {"search_query": [{"q": "OpenAI documentation"}]}
}))
.send()
.await
.expect("search error response should return");
assert_eq!(error_response.status(), StatusCode::BAD_REQUEST);
assert_eq!(
error_response
.headers()
.get("x-search-upstream")
.and_then(|value| value.to_str().ok()),
Some("rate-limited")
);
assert_eq!(
error_response
.json::<serde_json::Value>()
.await
.expect("error response should parse"),
expected_error_body
);
let error_candidates = request_candidates
.list_by_request_id("trace-search-error-1")
.await
.expect("error request candidates should read");
assert_eq!(error_candidates.len(), 1);
assert_eq!(error_candidates[0].status, RequestCandidateStatus::Failed);
let failover_response = reqwest::Client::new()
.post(format!("{gateway_url}/v1/alpha/search"))
.header(http::header::CONTENT_TYPE, "application/json")
.header(
http::header::AUTHORIZATION,
format!("Bearer {client_api_key}"),
)
.header(TRACE_ID_HEADER, "trace-search-failover-1")
.json(&json!({
"id": "session-search-failover-1",
"model": "gpt-5.6-sol",
"input": "find current OpenAI documentation",
"commands": {"search_query": [{"q": "OpenAI documentation"}]}
}))
.send()
.await
.expect("search failover response should return");
assert_eq!(failover_response.status(), StatusCode::OK);
assert_eq!(
failover_response
.headers()
.get("x-search-upstream")
.and_then(|value| value.to_str().ok()),
Some("backup")
);
assert_eq!(
failover_response
.json::<serde_json::Value>()
.await
.expect("failover response should parse")["output"],
"search fallback result"
);
let failover_plans = seen_plans
.lock()
.expect("mutex should lock")
.iter()
.filter(|plan| plan["request_id"] == "trace-search-failover-1")
.map(|plan| plan["provider_id"].clone())
.collect::<Vec<_>>();
assert_eq!(
failover_plans,
vec![
json!("provider-codex-search-1"),
json!("provider-codex-search-2")
]
);
let failover_candidates = request_candidates
.list_by_request_id("trace-search-failover-1")
.await
.expect("failover request candidates should read");
assert_eq!(failover_candidates.len(), 2);
assert_eq!(
failover_candidates[0].status,
RequestCandidateStatus::Failed
);
assert_eq!(failover_candidates[0].status_code, Some(500));
assert_eq!(
failover_candidates[1].status,
RequestCandidateStatus::Success
);
assert_eq!(failover_candidates[1].status_code, Some(200));
gateway_handle.abort();
execution_runtime_handle.abort();
}
@@ -559,6 +559,103 @@ async fn gateway_creates_admin_provider_endpoint_locally_with_trusted_admin_prin
upstream_handle.abort();
}
#[tokio::test]
async fn gateway_rejects_streaming_policy_for_search_endpoint_before_catalog_write() {
let mut create_provider = sample_provider("provider-search-create", "search-create", 10);
create_provider.provider_type = "custom".to_string();
let mut update_provider = sample_provider("provider-search-update", "search-update", 20);
update_provider.provider_type = "custom".to_string();
let mut existing_endpoint = sample_endpoint(
"endpoint-search-update",
"provider-search-update",
"openai:search",
"https://search.example/v1",
);
existing_endpoint.config = Some(json!({"marker": "kept"}));
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![create_provider, update_provider],
vec![existing_endpoint],
vec![],
));
let gateway = build_router_with_state(
AppState::new()
.expect("gateway should build")
.with_data_state_for_tests(
GatewayDataState::with_provider_catalog_repository_for_tests(
provider_catalog_repository.clone(),
),
),
);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let client = reqwest::Client::new();
let create_response = client
.post(format!(
"{gateway_url}/api/admin/endpoints/providers/provider-search-create/endpoints"
))
.header(GATEWAY_HEADER, "rust-phase3b")
.header(TRUSTED_ADMIN_USER_ID_HEADER, "admin-user-123")
.header(TRUSTED_ADMIN_USER_ROLE_HEADER, "admin")
.header(TRUSTED_ADMIN_SESSION_ID_HEADER, "session-123")
.json(&json!({
"provider_id": "provider-search-create",
"api_format": "openai:search",
"base_url": "https://search.example/v1",
"config": {"upstream_stream_policy": "force_stream"}
}))
.send()
.await
.expect("request should succeed");
assert_eq!(create_response.status(), StatusCode::BAD_REQUEST);
let create_payload: serde_json::Value = create_response
.json()
.await
.expect("json body should parse");
assert_eq!(
create_payload["detail"],
"OpenAI Search 端点仅支持非流式上游请求"
);
let update_response = client
.put(format!(
"{gateway_url}/api/admin/endpoints/endpoint-search-update"
))
.header(GATEWAY_HEADER, "rust-phase3b")
.header(TRUSTED_ADMIN_USER_ID_HEADER, "admin-user-123")
.header(TRUSTED_ADMIN_USER_ROLE_HEADER, "admin")
.header(TRUSTED_ADMIN_SESSION_ID_HEADER, "session-123")
.json(&json!({
"config": {"upstreamStreamPolicy": true}
}))
.send()
.await
.expect("request should succeed");
assert_eq!(update_response.status(), StatusCode::BAD_REQUEST);
let update_payload: serde_json::Value = update_response
.json()
.await
.expect("json body should parse");
assert_eq!(
update_payload["detail"],
"OpenAI Search 端点仅支持非流式上游请求"
);
let created = provider_catalog_repository
.list_endpoints_by_provider_ids(&["provider-search-create".to_string()])
.await
.expect("endpoints should read");
assert!(created.is_empty());
let unchanged = provider_catalog_repository
.list_endpoints_by_ids(&["endpoint-search-update".to_string()])
.await
.expect("endpoint should read");
assert_eq!(unchanged.len(), 1);
assert_eq!(unchanged[0].config, Some(json!({"marker": "kept"})));
gateway_handle.abort();
}
#[tokio::test]
async fn gateway_updates_admin_provider_endpoint_locally_with_trusted_admin_principal() {
let upstream_hits = Arc::new(Mutex::new(0usize));
@@ -6086,6 +6086,9 @@ async fn gateway_manual_codex_oauth_refresh_reconciles_missing_fixed_endpoint_im
.iter()
.find(|endpoint| endpoint.api_format == "openai:responses")
.expect("openai responses endpoint should be reconciled");
assert!(endpoints
.iter()
.any(|endpoint| endpoint.api_format == "openai:search"));
assert_eq!(
responses_endpoint.base_url,
"https://chatgpt.com/backend-api/codex"
@@ -2399,6 +2399,188 @@ async fn gateway_streams_codex_openai_responses_upstream_for_admin_pool_model_te
execution_runtime_handle.abort();
}
#[test]
fn gateway_executes_codex_search_admin_pool_model_test_with_search_contract() {
run_provider_query_test(
"gateway_executes_codex_search_admin_pool_model_test_with_search_contract",
gateway_executes_codex_search_admin_pool_model_test_with_search_contract_impl,
);
}
async fn gateway_executes_codex_search_admin_pool_model_test_with_search_contract_impl() {
let execution_runtime = Router::new().route(
"/v1/execute/sync",
any(move |Json(plan): Json<ExecutionPlan>| async move {
assert_eq!(plan.provider_id, "provider-codex-search");
assert_eq!(plan.endpoint_id, "endpoint-codex-search");
assert_eq!(plan.key_id, "key-codex-search");
assert_eq!(plan.client_api_format, "openai:search");
assert_eq!(plan.provider_api_format, "openai:search");
assert_eq!(
plan.url,
"https://chatgpt.com/backend-api/codex/alpha/search"
);
assert_eq!(plan.model_name.as_deref(), Some("gpt-5.6-sol"));
assert!(!plan.stream, "Codex Search is a synchronous JSON protocol");
assert_eq!(
plan.timeouts
.as_ref()
.and_then(|timeouts| timeouts.total_ms),
Some(900_000)
);
assert_eq!(
plan.headers.get("authorization").map(String::as_str),
Some("Bearer codex-search-access-token")
);
assert_eq!(
plan.headers.get("chatgpt-account-id").map(String::as_str),
Some("account-search-admin")
);
assert_eq!(
plan.headers.get("x-openai-fedramp").map(String::as_str),
Some("true")
);
assert_eq!(
plan.headers.get("originator").map(String::as_str),
Some("codex_cli_rs")
);
assert!(plan
.headers
.get("user-agent")
.is_some_and(|value| value.starts_with("codex_cli_rs/")));
assert!(!plan.headers.contains_key("openai-beta"));
assert!(!plan
.headers
.contains_key("x-openai-internal-codex-responses-lite"));
assert_ne!(
plan.headers.get("accept").map(String::as_str),
Some("text/event-stream")
);
let body = plan.body.json_body.as_ref().expect("search json body");
assert_eq!(
body["id"],
json!("aether-model-test-provider-query-search-trace")
);
assert_eq!(body["model"], json!("gpt-5.6-sol"));
assert_eq!(body["input"], json!("find current OpenAI documentation"));
assert_eq!(
body["commands"]["search_query"][0]["q"],
json!("OpenAI Codex Search")
);
assert!(body.get("stream").is_none());
assert!(body.get("store").is_none());
assert!(body.get("service_tier").is_none());
assert!(body.get("unknown_field").is_none());
Json(json!({
"request_id": plan.request_id,
"candidate_id": plan.candidate_id,
"status_code": 200,
"headers": {
"content-type": "application/json"
},
"body": {
"json_body": {
"output": "search result"
}
},
"telemetry": {
"elapsed_ms": 21
}
}))
}),
);
let (execution_runtime_url, execution_runtime_handle) = start_server(execution_runtime).await;
let mut provider = sample_provider("provider-codex-search", "Codex Search", 10);
provider.provider_type = "codex".to_string();
provider.request_timeout_secs = Some(900.0);
let mut endpoint = sample_endpoint(
"endpoint-codex-search",
"provider-codex-search",
"openai:search",
"https://chatgpt.com/backend-api/codex",
);
endpoint.config = Some(json!({"upstream_stream_policy": "force_stream"}));
let mut key = sample_key(
"key-codex-search",
"provider-codex-search",
"openai:search",
"codex-search-access-token",
);
key.auth_type = "oauth".to_string();
key.encrypted_auth_config = Some(
aether_crypto::encrypt_python_fernet_plaintext(
DEVELOPMENT_ENCRYPTION_KEY,
r#"{"provider_type":"codex","account_id":"account-search-admin","is_fedramp":true}"#,
)
.expect("auth config should encrypt"),
);
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![provider],
vec![endpoint],
vec![key],
));
let gateway = build_router_with_state(
build_state_with_execution_runtime_override(execution_runtime_url)
.with_data_state_for_tests(GatewayDataState::with_provider_transport_reader_for_tests(
provider_catalog_repository,
DEVELOPMENT_ENCRYPTION_KEY.to_string(),
)),
);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
.post(format!(
"{gateway_url}/api/admin/provider-query/test-model-failover"
))
.header(GATEWAY_HEADER, "rust-phase3b")
.header(TRUSTED_ADMIN_USER_ID_HEADER, "admin-user-123")
.header(TRUSTED_ADMIN_USER_ROLE_HEADER, "admin")
.header(TRUSTED_ADMIN_SESSION_ID_HEADER, "session-123")
.json(&json!({
"provider_id": "provider-codex-search",
"mode": "pool",
"model": "gpt-5.6-sol",
"failover_models": ["gpt-5.6-sol"],
"api_format": "openai:search",
"endpoint_id": "endpoint-codex-search",
"request_id": "provider-query-search-trace",
"request_body": {
"model": "gpt-5.6-sol",
"input": "find current OpenAI documentation",
"commands": {
"search_query": [{"q": "OpenAI Codex Search"}]
},
"max_output_tokens": 256,
"stream": true,
"store": false,
"service_tier": "priority",
"unknown_field": true
}
}))
.send()
.await
.expect("request should succeed");
assert_eq!(response.status(), StatusCode::OK);
let payload: serde_json::Value = response.json().await.expect("json body should parse");
assert_eq!(payload["success"], json!(true), "payload={payload}");
assert_eq!(
payload["attempts"][0]["request_body"]["id"],
json!("aether-model-test-provider-query-search-trace")
);
assert_eq!(
payload["attempts"][0]["response_body"]["output"],
json!("search result")
);
gateway_handle.abort();
execution_runtime_handle.abort();
}
#[test]
fn gateway_routes_grok_responses_admin_pool_model_test_through_grok_runtime() {
run_provider_query_test(
@@ -832,7 +832,7 @@ async fn gateway_updates_admin_provider_locally_with_trusted_admin_principal() {
"is_active": false,
"concurrent_limit": 8,
"max_retries": 6,
"request_timeout": 55.0,
"request_timeout": aether_contracts::MAX_EXECUTION_REQUEST_TIMEOUT_SECS,
"stream_first_byte_timeout": 11.0,
"enable_format_conversion": false,
"config": {
@@ -860,7 +860,10 @@ async fn gateway_updates_admin_provider_locally_with_trusted_admin_principal() {
assert_eq!(payload["enable_format_conversion"], false);
assert_eq!(payload["is_active"], false);
assert_eq!(payload["max_retries"], 6);
assert_eq!(payload["request_timeout"], 55.0);
assert_eq!(
payload["request_timeout"],
aether_contracts::MAX_EXECUTION_REQUEST_TIMEOUT_SECS
);
assert_eq!(payload["stream_first_byte_timeout"], 11.0);
assert_eq!(payload["proxy"], json!({"url": "https://proxy.example"}));
assert_eq!(payload["claude_code_advanced"], json!({"pool_size": 2}));
@@ -870,6 +873,21 @@ async fn gateway_updates_admin_provider_locally_with_trusted_admin_principal() {
assert_eq!(payload["ops_configured"], true);
assert_eq!(payload["ops_architecture_id"], "cubence");
let invalid_timeout_response = reqwest::Client::new()
.patch(format!("{gateway_url}/api/admin/providers/provider-openai"))
.header(crate::constants::GATEWAY_HEADER, "rust-phase3b")
.header(TRUSTED_ADMIN_USER_ID_HEADER, "admin-user-123")
.header(TRUSTED_ADMIN_USER_ROLE_HEADER, "admin")
.header(TRUSTED_ADMIN_SESSION_ID_HEADER, "session-123")
.json(&json!({
"request_timeout":
aether_contracts::MAX_EXECUTION_REQUEST_TIMEOUT_SECS + 1
}))
.send()
.await
.expect("request should succeed");
assert_eq!(invalid_timeout_response.status(), StatusCode::BAD_REQUEST);
let disable_response = reqwest::Client::new()
.patch(format!("{gateway_url}/api/admin/providers/provider-openai"))
.header(crate::constants::GATEWAY_HEADER, "rust-phase3b")
@@ -924,6 +942,10 @@ async fn gateway_updates_admin_provider_locally_with_trusted_admin_principal() {
.iter()
.find(|provider| provider.id == "provider-openai")
.expect("provider should exist");
assert_eq!(
updated_provider.request_timeout_secs,
Some(aether_contracts::MAX_EXECUTION_REQUEST_TIMEOUT_SECS as f64)
);
assert_eq!(
updated_provider
.config
@@ -1000,6 +1022,7 @@ async fn gateway_creates_admin_provider_locally_with_trusted_admin_principal() {
"website": "codex.example",
"keep_priority_on_conversion": true,
"max_retries": 7,
"request_timeout": aether_contracts::MAX_EXECUTION_REQUEST_TIMEOUT_SECS,
"config": {"chat_pii_redaction": {"enabled": true}},
"pool_advanced": {},
"failover_rules": {"strategy": "ordered"},
@@ -1034,6 +1057,10 @@ async fn gateway_creates_admin_provider_locally_with_trusted_admin_principal() {
assert_eq!(created.website.as_deref(), Some("https://codex.example"));
assert!(created.enable_format_conversion);
assert_eq!(created.max_retries, Some(7));
assert_eq!(
created.request_timeout_secs,
Some(aether_contracts::MAX_EXECUTION_REQUEST_TIMEOUT_SECS as f64)
);
assert_eq!(created.keep_priority_on_conversion, true);
assert_eq!(
created
@@ -1080,7 +1107,7 @@ async fn gateway_creates_admin_provider_locally_with_trusted_admin_principal() {
.list_endpoints_by_provider_ids(std::slice::from_ref(&created.id))
.await
.expect("endpoints should list");
assert_eq!(endpoints.len(), 3);
assert_eq!(endpoints.len(), 4);
let responses_endpoint = endpoints
.iter()
.find(|endpoint| endpoint.api_format == "openai:responses")
@@ -1089,6 +1116,10 @@ async fn gateway_creates_admin_provider_locally_with_trusted_admin_principal() {
.iter()
.find(|endpoint| endpoint.api_format == "openai:responses:compact")
.expect("compact endpoint should exist");
let search_endpoint = endpoints
.iter()
.find(|endpoint| endpoint.api_format == "openai:search")
.expect("search endpoint should exist");
let image_endpoint = endpoints
.iter()
.find(|endpoint| endpoint.api_format == "openai:image")
@@ -1101,12 +1132,17 @@ async fn gateway_creates_admin_provider_locally_with_trusted_admin_principal() {
compact_endpoint.base_url,
"https://chatgpt.com/backend-api/codex"
);
assert_eq!(
search_endpoint.base_url,
"https://chatgpt.com/backend-api/codex"
);
assert_eq!(
image_endpoint.base_url,
"https://chatgpt.com/backend-api/codex"
);
assert_eq!(responses_endpoint.max_retries, Some(7));
assert_eq!(compact_endpoint.max_retries, Some(7));
assert_eq!(search_endpoint.max_retries, Some(7));
assert_eq!(image_endpoint.max_retries, Some(7));
assert_eq!(
responses_endpoint
@@ -1116,6 +1152,14 @@ async fn gateway_creates_admin_provider_locally_with_trusted_admin_principal() {
.and_then(serde_json::Value::as_str),
Some("force_stream")
);
assert_eq!(
search_endpoint
.config
.as_ref()
.and_then(|value| value.get("upstream_stream_policy"))
.and_then(serde_json::Value::as_str),
None
);
assert_eq!(
image_endpoint
.config
@@ -1126,6 +1170,7 @@ async fn gateway_creates_admin_provider_locally_with_trusted_admin_principal() {
);
assert!(responses_endpoint.body_rules.is_none());
assert!(compact_endpoint.body_rules.is_none());
assert!(search_endpoint.body_rules.is_none());
assert!(image_endpoint.body_rules.is_none());
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
@@ -1216,7 +1261,7 @@ async fn gateway_updates_fixed_provider_and_reconciles_template_managed_endpoint
.list_endpoints_by_provider_ids(&["provider-codex".to_string()])
.await
.expect("endpoints should list");
assert_eq!(endpoints.len(), 3);
assert_eq!(endpoints.len(), 4);
let responses_endpoint = endpoints
.iter()
.find(|endpoint| endpoint.api_format == "openai:responses")
@@ -1225,6 +1270,10 @@ async fn gateway_updates_fixed_provider_and_reconciles_template_managed_endpoint
.iter()
.find(|endpoint| endpoint.api_format == "openai:responses:compact")
.expect("compact endpoint should exist");
let search_endpoint = endpoints
.iter()
.find(|endpoint| endpoint.api_format == "openai:search")
.expect("search endpoint should exist");
let image_endpoint = endpoints
.iter()
.find(|endpoint| endpoint.api_format == "openai:image")
@@ -1232,6 +1281,7 @@ async fn gateway_updates_fixed_provider_and_reconciles_template_managed_endpoint
assert_eq!(responses_endpoint.max_retries, Some(9));
assert_eq!(compact_endpoint.max_retries, Some(9));
assert_eq!(search_endpoint.max_retries, Some(9));
assert_eq!(image_endpoint.max_retries, Some(9));
assert_eq!(
responses_endpoint
@@ -1242,6 +1292,23 @@ async fn gateway_updates_fixed_provider_and_reconciles_template_managed_endpoint
.and_then(serde_json::Value::as_bool),
Some(true)
);
assert_eq!(
search_endpoint
.config
.as_ref()
.and_then(|value| value.get("_aether_fixed_provider_template"))
.and_then(|value| value.get("managed"))
.and_then(serde_json::Value::as_bool),
Some(true)
);
assert_eq!(
search_endpoint
.config
.as_ref()
.and_then(|value| value.get("upstream_stream_policy"))
.and_then(serde_json::Value::as_str),
None
);
assert_eq!(
image_endpoint
.config
@@ -1255,7 +1322,7 @@ async fn gateway_updates_fixed_provider_and_reconciles_template_managed_endpoint
.await
.expect("keys should list");
assert_eq!(keys.len(), 1);
assert!(keys[0].api_formats.is_none());
assert_eq!(keys[0].api_formats, Some(json!(["openai:responses"])));
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
gateway_handle.abort();
+125 -15
View File
@@ -10,6 +10,8 @@ use futures_util::stream;
use http::header::HeaderValue;
use http::StatusCode;
use serde_json::json;
use std::collections::HashMap;
use std::time::Duration;
use super::{
build_router_with_state, sample_proxy_node, start_server, AppState, GatewayDataState,
@@ -318,6 +320,88 @@ async fn gateway_forwards_tunnel_relay_to_attachment_owner() {
owner_handle.abort();
}
#[tokio::test]
async fn gateway_owner_relay_uses_non_stream_timeout_from_envelope() {
let owner = Router::new().route(
"/api/internal/tunnel/relay/node-123",
post(|body: Body| async move {
let body = axum::body::to_bytes(body, usize::MAX)
.await
.expect("body should read");
tokio::time::sleep(Duration::from_millis(40)).await;
(StatusCode::OK, Body::from(body))
}),
);
let (owner_url, owner_handle) = start_server(owner).await;
let data_state = GatewayDataState::disabled().with_system_config_values_for_tests(vec![(
"tunnel.attachments.node-123".to_string(),
json!({
"gateway_instance_id": "gateway-b",
"relay_base_url": owner_url,
"conn_count": 1,
"observed_at_unix_secs": 4_102_444_800u64,
}),
)]);
let mut state = AppState::new()
.expect("gateway should build")
.with_data_state_for_tests(data_state)
.with_tunnel_identity_for_tests("gateway-a", Some("http://gateway-a.internal"));
let short_timeout_client = reqwest::Client::builder()
.timeout(Duration::from_millis(10))
.build()
.expect("test client should build");
state.client = short_timeout_client.clone();
state.owner_forward_client = short_timeout_client;
let gateway = build_router_with_state(state);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let meta = aether_contracts::tunnel::RequestMeta {
provider_id: Some("provider-1".to_string()),
endpoint_id: Some("endpoint-1".to_string()),
key_id: Some("key-1".to_string()),
method: "POST".to_string(),
url: "https://example.com/responses".to_string(),
headers: HashMap::new(),
stream: false,
request_timeout_ms: Some(100),
stream_first_byte_timeout_ms: None,
timeout: 60,
follow_redirects: None,
http1_only: false,
transport_profile: None,
};
let encoded_meta = serde_json::to_vec(&meta).expect("metadata should encode");
let mut envelope = Vec::new();
envelope.extend_from_slice(&(encoded_meta.len() as u32).to_be_bytes());
envelope.extend_from_slice(&encoded_meta);
envelope.extend_from_slice(b"relay-body");
let split_at = 4 + encoded_meta.len() / 2;
let request_body = reqwest::Body::wrap_stream(stream::iter(vec![
Ok::<Bytes, io::Error>(Bytes::copy_from_slice(&envelope[..split_at])),
Ok::<Bytes, io::Error>(Bytes::copy_from_slice(&envelope[split_at..])),
]));
let response = reqwest::Client::builder()
.timeout(Duration::from_secs(1))
.build()
.expect("request client should build")
.post(format!("{gateway_url}/api/internal/tunnel/relay/node-123"))
.body(request_body)
.send()
.await
.expect("request should succeed");
assert_eq!(response.status(), StatusCode::OK);
assert_eq!(
response.bytes().await.expect("response body should read"),
Bytes::from(envelope)
);
gateway_handle.abort();
owner_handle.abort();
}
#[tokio::test]
async fn gateway_streams_tunnel_relay_body_to_attachment_owner() {
let owner_hits = Arc::new(Mutex::new(0usize));
@@ -331,8 +415,12 @@ async fn gateway_streams_tunnel_relay_body_to_attachment_owner() {
let body = axum::body::to_bytes(body, usize::MAX)
.await
.expect("body should read");
assert_eq!(body, Bytes::from_static(b"relay-stream-envelope"));
(StatusCode::OK, Body::from("stream-ok"))
let response_body = Body::from_stream(async_stream::stream! {
yield Ok::<_, io::Error>(Bytes::from_static(b"stream-"));
tokio::time::sleep(Duration::from_millis(40)).await;
yield Ok::<_, io::Error>(Bytes::from_static(b"ok"));
});
(StatusCode::OK, response_body)
}
}),
);
@@ -347,19 +435,41 @@ async fn gateway_streams_tunnel_relay_body_to_attachment_owner() {
"observed_at_unix_secs": 4_102_444_800u64,
}),
)]);
let gateway = build_router_with_state(
AppState::new()
.expect("gateway should build")
.with_data_state_for_tests(data_state)
.with_tunnel_identity_for_tests("gateway-a", Some("http://gateway-a.internal")),
);
let mut state = AppState::new()
.expect("gateway should build")
.with_data_state_for_tests(data_state)
.with_tunnel_identity_for_tests("gateway-a", Some("http://gateway-a.internal"));
state.client = reqwest::Client::builder()
.timeout(Duration::from_millis(10))
.build()
.expect("short shared client should build");
let gateway = build_router_with_state(state);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let request_body = reqwest::Body::wrap_stream(stream::iter(vec![
Ok::<Bytes, io::Error>(Bytes::from_static(b"relay-")),
Ok::<Bytes, io::Error>(Bytes::from_static(b"stream-")),
Ok::<Bytes, io::Error>(Bytes::from_static(b"envelope")),
]));
let meta = aether_contracts::tunnel::RequestMeta {
provider_id: Some("provider-1".to_string()),
endpoint_id: Some("endpoint-1".to_string()),
key_id: Some("key-1".to_string()),
method: "POST".to_string(),
url: "https://example.com/responses".to_string(),
headers: HashMap::new(),
stream: true,
request_timeout_ms: Some(900_000),
stream_first_byte_timeout_ms: Some(100),
timeout: 60,
follow_redirects: None,
http1_only: false,
transport_profile: None,
};
let encoded_meta = serde_json::to_vec(&meta).expect("metadata should encode");
let mut envelope = Vec::new();
envelope.extend_from_slice(&(encoded_meta.len() as u32).to_be_bytes());
envelope.extend_from_slice(&encoded_meta);
envelope.extend_from_slice(b"relay-stream-envelope");
let expected_envelope = Bytes::copy_from_slice(&envelope);
let request_body = reqwest::Body::wrap_stream(stream::iter(vec![Ok::<Bytes, io::Error>(
expected_envelope.clone(),
)]));
let response = reqwest::Client::new()
.post(format!("{gateway_url}/api/internal/tunnel/relay/node-123"))
.body(request_body)
@@ -369,8 +479,8 @@ async fn gateway_streams_tunnel_relay_body_to_attachment_owner() {
assert_eq!(response.status(), StatusCode::OK);
assert_eq!(
response.text().await.expect("body should read"),
"stream-ok"
response.bytes().await.expect("body should read"),
Bytes::from_static(b"stream-ok")
);
assert_eq!(*owner_hits.lock().expect("mutex should lock"), 1);
@@ -70,6 +70,7 @@ async fn gateway_exposes_frontdoor_manifest_without_proxying_upstream() {
assert!(owned_routes
.iter()
.any(|value| value == "/v1/responses/compact"));
assert!(owned_routes.iter().any(|value| value == "/v1/alpha/search"));
assert!(owned_routes.iter().any(|value| value == "/health"));
assert!(owned_routes.iter().any(|value| value == "/v1/health"));
assert!(owned_routes.iter().any(|value| value == "/v1/providers"));
+64 -26
View File
@@ -89,6 +89,13 @@ fn sample_cli_auth_snapshot(
}
fn sample_provider(provider_id: &str) -> StoredProviderCatalogProvider {
sample_provider_with_request_timeout(provider_id, None)
}
fn sample_provider_with_request_timeout(
provider_id: &str,
request_timeout_secs: Option<f64>,
) -> StoredProviderCatalogProvider {
StoredProviderCatalogProvider::new(
provider_id.to_string(),
provider_id.to_string(),
@@ -96,7 +103,17 @@ fn sample_provider(provider_id: &str) -> StoredProviderCatalogProvider {
"custom".to_string(),
)
.expect("provider should build")
.with_transport_fields(true, false, false, None, None, None, None, None, None)
.with_transport_fields(
true,
false,
false,
None,
None,
None,
request_timeout_secs,
None,
None,
)
}
fn sample_endpoint(endpoint_id: &str, provider_id: &str) -> StoredProviderCatalogEndpoint {
@@ -435,6 +452,7 @@ async fn gateway_forwards_public_request_to_remote_tunnel_owner_before_fallback_
async move {
let (parts, body) = request.into_parts();
let raw_body = to_bytes(body, usize::MAX).await.expect("body should read");
tokio::time::sleep(Duration::from_millis(40)).await;
*seen_owner_inner.lock().expect("mutex should lock") = Some(SeenOwnerRequest {
path: parts
.uri
@@ -511,7 +529,10 @@ async fn gateway_forwards_public_request_to_remote_tunnel_owner_before_fallback_
let (owner_url, owner_handle) = start_server(owner).await;
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![sample_provider("provider-owner")],
vec![sample_provider_with_request_timeout(
"provider-owner",
Some(0.1),
)],
vec![sample_endpoint("endpoint-owner", "provider-owner")],
vec![sample_key("key-owner", "provider-owner", "node-owner")],
));
@@ -540,6 +561,12 @@ async fn gateway_forwards_public_request_to_remote_tunnel_owner_before_fallback_
state = state
.with_data_state_for_tests(data_state)
.with_tunnel_identity_for_tests("gateway-a", Some("http://gateway-a:8080"));
let short_timeout_client = reqwest::Client::builder()
.timeout(Duration::from_millis(10))
.build()
.expect("test client should build");
state.client = short_timeout_client.clone();
state.owner_forward_client = short_timeout_client;
state.remember_scheduler_affinity_target(
"scheduler_affinity:api-key-affinity-1:openai:chat:gpt-4.1",
crate::cache::SchedulerAffinityTarget {
@@ -922,32 +949,39 @@ async fn gateway_streamifies_sync_json_from_remote_tunnel_owner_before_returning
.unwrap_or_default()
.to_string(),
});
let encoded_response = serde_json::to_vec(&json!({
"id": "resp-codex-affinity-stream-123",
"object": "response",
"model": "gpt-5.4",
"status": "completed",
"output": [{
"type": "message",
"id": "msg-codex-affinity-stream-123",
"role": "assistant",
"content": [{
"type": "output_text",
"text": "Hello from affinity sync json",
"annotations": []
}]
}],
"usage": {
"input_tokens": 1,
"output_tokens": 2,
"total_tokens": 3
}
}))
.expect("body should encode");
let split_at = encoded_response.len() / 2;
let first = axum::body::Bytes::copy_from_slice(&encoded_response[..split_at]);
let second = axum::body::Bytes::copy_from_slice(&encoded_response[split_at..]);
let response_body = Body::from_stream(async_stream::stream! {
yield Ok::<_, std::io::Error>(first);
tokio::time::sleep(Duration::from_millis(40)).await;
yield Ok::<_, std::io::Error>(second);
});
let mut response = Response::builder()
.status(StatusCode::OK)
.body(Body::from(
serde_json::to_vec(&json!({
"id": "resp-codex-affinity-stream-123",
"object": "response",
"model": "gpt-5.4",
"status": "completed",
"output": [{
"type": "message",
"id": "msg-codex-affinity-stream-123",
"role": "assistant",
"content": [{
"type": "output_text",
"text": "Hello from affinity sync json",
"annotations": []
}]
}],
"usage": {
"input_tokens": 1,
"output_tokens": 2,
"total_tokens": 3
}
}))
.expect("body should encode"),
))
.body(response_body)
.expect("response should build");
response.headers_mut().insert(
http::header::CONTENT_TYPE,
@@ -1001,6 +1035,10 @@ async fn gateway_streamifies_sync_json_from_remote_tunnel_owner_before_returning
state = state
.with_data_state_for_tests(data_state)
.with_tunnel_identity_for_tests("gateway-a", Some("http://gateway-a:8080"));
state.client = reqwest::Client::builder()
.timeout(Duration::from_millis(10))
.build()
.expect("short shared client should build");
state.remember_scheduler_affinity_target(
"scheduler_affinity:api-key-affinity-cli-1:openai:responses:gpt-5.4",
crate::cache::SchedulerAffinityTarget {
@@ -2,7 +2,10 @@ use std::io;
use std::net::SocketAddr;
use std::time::Duration;
use aether_contracts::tunnel::TUNNEL_RELAY_FORWARDED_BY_HEADER;
use aether_contracts::tunnel::{
resolve_tunnel_request_timeouts, try_decode_tunnel_relay_request_meta,
TUNNEL_RELAY_FORWARDED_BY_HEADER,
};
use aether_runtime::{maybe_hold_axum_response_permit, AdmissionPermit};
use async_stream::stream;
use axum::body::{Body, Bytes};
@@ -23,9 +26,6 @@ use super::protocol;
use super::AppState;
pub const TUNNEL_ERROR_HEADER: &str = "x-aether-tunnel-error";
const MAX_RELAY_META_LEN: usize = 256 * 1024;
const MIN_RELAY_TIMEOUT_MS: u64 = 1;
const MAX_RELAY_TIMEOUT_MS: u64 = 300_000;
struct StreamGuard {
hub: std::sync::Arc<super::hub::HubRouter>,
@@ -156,15 +156,7 @@ fn map_request_admission_error(error: super::RequestAdmissionError) -> String {
}
fn relay_header_timeout(meta: &protocol::RequestMeta) -> Duration {
let timeout_ms = if meta.stream {
meta.stream_first_byte_timeout_ms
.unwrap_or_else(|| meta.timeout.saturating_mul(1_000))
} else {
meta.request_timeout_ms
.or(meta.stream_first_byte_timeout_ms)
.unwrap_or_else(|| meta.timeout.saturating_mul(1_000))
};
Duration::from_millis(timeout_ms.clamp(MIN_RELAY_TIMEOUT_MS, MAX_RELAY_TIMEOUT_MS))
Duration::from_millis(resolve_tunnel_request_timeouts(meta).first_byte_ms)
}
pub async fn relay_request(
@@ -252,15 +244,17 @@ pub async fn relay_request(
if stream.is_none() {
envelope_buf.extend_from_slice(&chunk);
let Some((parsed_meta, body_offset)) = (match try_decode_envelope_meta(&envelope_buf) {
Ok(result) => result,
Err(error) => {
return release_permit_response(
tunnel_error_response(StatusCode::BAD_REQUEST, "bad_request", &error),
request_permit,
);
}
}) else {
let Some((parsed_meta, body_offset)) =
(match try_decode_tunnel_relay_request_meta(&envelope_buf) {
Ok(result) => result,
Err(error) => {
return release_permit_response(
tunnel_error_response(StatusCode::BAD_REQUEST, "bad_request", &error),
request_permit,
);
}
})
else {
continue;
};
@@ -427,27 +421,6 @@ fn release_permit_response(
response
}
fn try_decode_envelope_meta(
buffer: &BytesMut,
) -> Result<Option<(protocol::RequestMeta, usize)>, String> {
if buffer.len() < 4 {
return Ok(None);
}
let meta_len = u32::from_be_bytes([buffer[0], buffer[1], buffer[2], buffer[3]]) as usize;
if meta_len > MAX_RELAY_META_LEN {
return Err("relay metadata too large".to_string());
}
let meta_end = 4usize
.checked_add(meta_len)
.ok_or_else(|| "relay envelope length overflow".to_string())?;
if buffer.len() < meta_end {
return Ok(None);
}
let meta = serde_json::from_slice::<protocol::RequestMeta>(&buffer[4..meta_end])
.map_err(|e| format!("invalid relay metadata: {e}"))?;
Ok(Some((meta, meta_end)))
}
fn append_headers(target: &mut HeaderMap, headers: &[(String, String)]) {
for (name, value) in headers {
if should_skip_local_relay_response_header(name) {
@@ -542,6 +515,30 @@ mod tests {
assert_eq!(relay_header_timeout(&meta), Duration::from_secs(7));
}
#[test]
fn relay_header_timeout_keeps_the_protocol_maximum_for_non_stream_requests() {
let meta = protocol::RequestMeta {
provider_id: None,
endpoint_id: None,
key_id: None,
method: "POST".to_string(),
url: "https://example.com/responses".to_string(),
headers: HashMap::new(),
stream: false,
request_timeout_ms: Some(aether_contracts::MAX_EXECUTION_REQUEST_TIMEOUT_MS),
stream_first_byte_timeout_ms: None,
timeout: 60,
follow_redirects: None,
http1_only: false,
transport_profile: None,
};
assert_eq!(
relay_header_timeout(&meta),
Duration::from_millis(aether_contracts::MAX_EXECUTION_REQUEST_TIMEOUT_MS)
);
}
fn sample_connected_proxy_node(node_id: &str) -> StoredProxyNode {
StoredProxyNode::new(
node_id.to_string(),
+125 -16
View File
@@ -8,6 +8,7 @@ use std::sync::Arc;
use std::time::{Duration, SystemTime};
use aether_contracts::tunnel::{
resolve_tunnel_request_timeouts, try_decode_tunnel_relay_request_meta, RequestMeta,
TUNNEL_RELAY_FORWARDED_BY_HEADER, TUNNEL_RELAY_OWNER_INSTANCE_HEADER,
};
use aether_data::repository::proxy_nodes::{
@@ -21,6 +22,7 @@ use axum::extract::ws::WebSocketUpgrade;
use axum::extract::{ConnectInfo, Path, Request, State};
use axum::http::{HeaderMap, StatusCode};
use axum::response::IntoResponse;
use bytes::BytesMut;
use futures_util::StreamExt;
use serde::{Deserialize, Serialize};
use serde_json::json;
@@ -61,6 +63,22 @@ const TUNNEL_INSTANCE_ID_ENV: &str = "AETHER_GATEWAY_INSTANCE_ID";
const TUNNEL_RELAY_BASE_URL_ENV: &str = "AETHER_TUNNEL_RELAY_BASE_URL";
const TUNNEL_ATTACHMENT_TTL_ENV: &str = "AETHER_TUNNEL_ATTACHMENT_TTL_SECS";
pub(crate) async fn send_owner_forward_request(
request: reqwest::RequestBuilder,
first_byte_timeout: Option<Duration>,
) -> Result<reqwest::Response, String> {
match first_byte_timeout {
Some(timeout) => match tokio::time::timeout(timeout, request.send()).await {
Ok(result) => result.map_err(|error| error.to_string()),
Err(_) => Err(format!(
"owner gateway first byte timeout after {} ms",
timeout.as_millis()
)),
},
None => request.send().await.map_err(|error| error.to_string()),
}
}
#[derive(Debug, Deserialize)]
struct InternalTunnelHeartbeatRequest {
node_id: String,
@@ -776,8 +794,29 @@ async fn forward_relay_request_to_owner(
);
}
let limit_exceeded = Arc::new(AtomicBool::new(false));
let prepared_body =
match prepare_owner_relay_request_body(body, body_limit, Arc::clone(&limit_exceeded)).await
{
Ok(prepared_body) => prepared_body,
Err(_) if limit_exceeded.load(Ordering::SeqCst) => {
return build_local_http_error_response(
trace_id,
None,
StatusCode::PAYLOAD_TOO_LARGE,
&format!("tunnel relay body exceeds {body_limit} bytes"),
);
}
Err(error) => {
return build_local_http_error_response(
trace_id,
None,
StatusCode::BAD_REQUEST,
&error,
);
}
};
let mut upstream_request = state.client.post(owner_url);
let mut upstream_request = state.owner_forward_client.post(owner_url);
for (name, value) in &parts.headers {
if should_skip_request_header(name.as_str()) || name == http::header::HOST {
continue;
@@ -793,18 +832,23 @@ async fn forward_relay_request_to_owner(
TUNNEL_RELAY_OWNER_INSTANCE_HEADER,
owner.gateway_instance_id.as_str(),
);
let resolved_timeouts = resolve_tunnel_request_timeouts(&prepared_body.meta);
if let Some(timeout_ms) = resolved_timeouts.response_body_ms {
upstream_request = upstream_request.timeout(Duration::from_millis(timeout_ms));
}
if !parts.headers.contains_key(TRACE_ID_HEADER) {
upstream_request = upstream_request.header(TRACE_ID_HEADER, trace_id);
}
let upstream_response = match upstream_request
.body(build_owner_relay_request_body(
body,
body_limit,
Arc::clone(&limit_exceeded),
))
.send()
.await
let first_byte_timeout = prepared_body
.meta
.stream
.then_some(Duration::from_millis(resolved_timeouts.first_byte_ms));
let upstream_response = match send_owner_forward_request(
upstream_request.body(prepared_body.body),
first_byte_timeout,
)
.await
{
Ok(response) => response,
Err(err) if limit_exceeded.load(Ordering::SeqCst) => {
@@ -862,14 +906,53 @@ fn request_content_length_exceeds_limit(headers: &HeaderMap, body_limit: usize)
.is_some_and(|value| value > body_limit)
}
fn build_owner_relay_request_body(
struct PreparedOwnerRelayRequestBody {
body: reqwest::Body,
meta: RequestMeta,
}
async fn prepare_owner_relay_request_body(
body: Body,
body_limit: usize,
limit_exceeded: Arc<AtomicBool>,
) -> reqwest::Body {
) -> Result<PreparedOwnerRelayRequestBody, String> {
let mut body_stream = body.into_data_stream();
reqwest::Body::wrap_stream(stream! {
let mut forwarded = 0usize;
let mut buffered_chunks = Vec::new();
let mut meta_buffer = BytesMut::new();
let mut forwarded = 0usize;
let mut meta = None;
while meta.is_none() {
let Some(next_chunk) = body_stream.next().await else {
return Err("incomplete tunnel relay metadata".to_string());
};
match next_chunk {
Ok(chunk) => {
let next_forwarded = forwarded.saturating_add(chunk.len());
if next_forwarded > body_limit {
limit_exceeded.store(true, Ordering::SeqCst);
return Err(format!("tunnel relay body exceeds {body_limit} bytes"));
}
forwarded = next_forwarded;
meta_buffer.extend_from_slice(&chunk);
buffered_chunks.push(chunk);
match try_decode_tunnel_relay_request_meta(&meta_buffer) {
Ok(Some((parsed, _))) => meta = Some(parsed),
Ok(None) => {}
Err(error) => return Err(error),
}
}
Err(error) => {
return Err(format!("tunnel relay body read failed: {error}"));
}
}
}
let meta = meta.ok_or_else(|| "incomplete tunnel relay metadata".to_string())?;
let forwarded_body = reqwest::Body::wrap_stream(stream! {
for chunk in buffered_chunks {
yield Ok::<Bytes, io::Error>(chunk);
}
while let Some(next_chunk) = body_stream.next().await {
match next_chunk {
Ok(chunk) => {
@@ -884,12 +967,17 @@ fn build_owner_relay_request_body(
}
yield Ok::<Bytes, io::Error>(chunk);
}
Err(err) => {
yield Err::<Bytes, io::Error>(io::Error::other(err));
Err(error) => {
yield Err::<Bytes, io::Error>(io::Error::other(error));
break;
}
}
}
});
Ok(PreparedOwnerRelayRequestBody {
body: forwarded_body,
meta,
})
}
@@ -1047,12 +1135,15 @@ fn parse_embedded_tunnel_heartbeat_request(
mod tests {
use super::{
apply_embedded_tunnel_heartbeat, apply_embedded_tunnel_node_status, current_unix_secs,
tunnel_attachment_key, GatewayDataState, TunnelAttachmentDirectory, TunnelAttachmentRecord,
prepare_owner_relay_request_body, tunnel_attachment_key, GatewayDataState,
TunnelAttachmentDirectory, TunnelAttachmentRecord,
};
use aether_data::repository::proxy_nodes::{
InMemoryProxyNodeRepository, ProxyNodeReadRepository, StoredProxyNode,
};
use axum::body::Body;
use serde_json::json;
use std::sync::atomic::AtomicBool;
use std::sync::Arc;
fn sample_proxy_node(node_id: &str) -> StoredProxyNode {
@@ -1092,6 +1183,24 @@ mod tests {
)
}
#[tokio::test]
async fn owner_relay_body_preparation_rejects_invalid_metadata() {
let mut envelope = Vec::new();
envelope.extend_from_slice(&1u32.to_be_bytes());
envelope.push(b'{');
let error = prepare_owner_relay_request_body(
Body::from(envelope),
1024,
Arc::new(AtomicBool::new(false)),
)
.await
.err()
.expect("invalid metadata should fail");
assert!(error.contains("invalid relay metadata"));
}
#[tokio::test]
async fn embedded_tunnel_heartbeat_updates_proxy_node_repository() {
let repository = Arc::new(InMemoryProxyNodeRepository::seed(vec![sample_proxy_node(
+15 -36
View File
@@ -130,10 +130,6 @@ fn try_send_window_update(frame_tx: &FrameSender, stream_id: u32, bytes: usize)
}
}
/// Minimum allowed upstream request timeout (milliseconds).
const MIN_TIMEOUT_MS: u64 = 1;
/// Maximum allowed upstream request timeout (milliseconds).
const MAX_TIMEOUT_MS: u64 = 300_000;
/// Match reqwest's default redirect budget so direct execution and tunnel relay
/// fail at the same point instead of diverging after a different number of hops.
const MAX_REDIRECTS: usize = 10;
@@ -707,43 +703,14 @@ fn remaining_timeout(deadline: Instant) -> Option<Duration> {
}
fn resolve_request_timeouts(meta: &RequestMeta) -> RequestTimeouts {
let first_byte_timeout = if meta.stream {
meta.stream_first_byte_timeout_ms
.map(timeout_duration_from_ms)
.unwrap_or_else(|| timeout_duration_from_legacy_secs(meta.timeout))
} else {
meta.request_timeout_ms
.or(meta.stream_first_byte_timeout_ms)
.map(timeout_duration_from_ms)
.unwrap_or_else(|| timeout_duration_from_legacy_secs(meta.timeout))
};
let response_body_timeout = if meta.stream {
None
} else {
Some(
meta.request_timeout_ms
.or(meta.stream_first_byte_timeout_ms)
.map(timeout_duration_from_ms)
.unwrap_or_else(|| timeout_duration_from_legacy_secs(meta.timeout)),
)
};
let resolved = aether_contracts::tunnel::resolve_tunnel_request_timeouts(meta);
RequestTimeouts {
first_byte_timeout,
response_body_timeout,
first_byte_timeout: Duration::from_millis(resolved.first_byte_ms),
response_body_timeout: resolved.response_body_ms.map(Duration::from_millis),
}
}
fn timeout_duration_from_ms(ms: u64) -> Duration {
Duration::from_millis(ms.clamp(MIN_TIMEOUT_MS, MAX_TIMEOUT_MS))
}
fn timeout_duration_from_legacy_secs(secs: u64) -> Duration {
let ms = secs.saturating_mul(1_000);
timeout_duration_from_ms(ms)
}
async fn spool_request_body(
stream_id: u32,
mut body_rx: mpsc::Receiver<TunnelFrame>,
@@ -2176,6 +2143,18 @@ mod tests {
);
}
#[test]
fn non_stream_request_timeouts_keep_the_protocol_maximum() {
let mut meta = sample_request_meta();
meta.request_timeout_ms = Some(aether_contracts::MAX_EXECUTION_REQUEST_TIMEOUT_MS);
let timeouts = resolve_request_timeouts(&meta);
let expected = Duration::from_millis(aether_contracts::MAX_EXECUTION_REQUEST_TIMEOUT_MS);
assert_eq!(timeouts.first_byte_timeout, expected);
assert_eq!(timeouts.response_body_timeout, Some(expected));
}
#[test]
fn resolve_redirect_changes_post_to_get_for_302() {
let current_url = url::Url::parse("https://redirect.test/start").expect("url");
+28 -8
View File
@@ -951,14 +951,7 @@ fn admin_usage_api_format_defaults_to_non_stream(item: &StoredRequestUsageAudit)
let Some(value) = api_format else {
return false;
};
matches!(
aether_ai_formats::normalize_api_format_alias(value).as_str(),
"openai:chat"
| "openai:responses"
| "openai:responses:compact"
| "openai:image"
| "claude:messages"
)
aether_ai_formats::api_format_defaults_to_non_stream(value)
}
fn admin_usage_request_body_implies_default_non_stream(item: &StoredRequestUsageAudit) -> bool {
@@ -2809,6 +2802,33 @@ mod tests {
assert_eq!(record["client_is_stream"], false);
}
#[test]
fn client_requested_stream_defaults_to_non_stream_for_openai_search() {
let item = StoredRequestUsageAudit {
is_stream: true,
api_format: Some("openai:search".to_string()),
request_body: Some(json!({
"id": "session-search-1",
"model": "gpt-5.6-sol",
"input": "current documentation"
})),
..sample_usage("completed", Some(200), None)
};
assert!(!admin_usage_client_is_stream(&item));
let record = admin_usage_record_json(
&item,
&BTreeMap::new(),
&BTreeMap::new(),
false,
false,
None,
);
assert_eq!(record["client_requested_stream"], false);
assert_eq!(record["client_is_stream"], false);
}
#[test]
fn upstream_stream_prefers_request_metadata_flag() {
let item = StoredRequestUsageAudit {
+69 -13
View File
@@ -82,33 +82,89 @@ pub fn endpoint_key_counts_by_format(
) -> (BTreeMap<String, usize>, BTreeMap<String, usize>) {
let mut total = BTreeMap::new();
let mut active = BTreeMap::new();
let inherited_api_formats = active_endpoint_api_formats(endpoints);
let endpoint_api_formats = active_endpoint_api_formats(endpoints);
for key in keys {
if fixed_provider_key_inherits_api_formats(
let inherits_api_formats = fixed_provider_key_inherits_api_formats(
provider_type,
&key.auth_type,
key.encrypted_auth_config.as_deref(),
) {
for api_format in &inherited_api_formats {
*total.entry(api_format.clone()).or_insert(0) += 1;
if key.is_active {
*active.entry(api_format.clone()).or_insert(0) += 1;
}
}
continue;
}
);
let has_unrestricted_api_format_scope = key
.api_formats
.as_ref()
.is_none_or(serde_json::Value::is_null);
let configured_api_formats = configured_key_api_formats(key);
for api_format in configured_key_api_formats(key) {
for api_format in endpoint_api_formats.iter().filter(|api_format| {
inherits_api_formats
|| has_unrestricted_api_format_scope
|| configured_api_formats.iter().any(|allowed| {
aether_ai_formats::api_format_permission_covers(allowed, api_format)
})
}) {
*total.entry(api_format.clone()).or_insert(0) += 1;
if key.is_active {
*active.entry(api_format).or_insert(0) += 1;
*active.entry(api_format.clone()).or_insert(0) += 1;
}
}
}
(total, active)
}
#[cfg(test)]
mod endpoint_key_count_tests {
use super::*;
fn sample_endpoint(id: &str, api_format: &str) -> StoredProviderCatalogEndpoint {
StoredProviderCatalogEndpoint::new(
id.to_string(),
"provider-1".to_string(),
api_format.to_string(),
None,
None,
true,
)
.expect("endpoint should build")
}
fn sample_key(id: &str, api_format: Option<&str>) -> StoredProviderCatalogKey {
let mut key = StoredProviderCatalogKey::new(
id.to_string(),
"provider-1".to_string(),
id.to_string(),
"api_key".to_string(),
None,
true,
)
.expect("key should build");
key.api_formats = api_format.map(|api_format| json!([api_format]));
key
}
#[test]
fn endpoint_counts_follow_one_way_format_permissions() {
let endpoints = vec![
sample_endpoint("responses", "openai:responses"),
sample_endpoint("search", "openai:search"),
];
let mut empty_scope_key = sample_key("empty-scope-key", None);
empty_scope_key.api_formats = Some(json!([]));
let keys = vec![
sample_key("responses-key", Some("openai:responses")),
sample_key("search-key", Some("openai:search")),
sample_key("unrestricted-key", None),
empty_scope_key,
];
let (total, active) = endpoint_key_counts_by_format("custom", &endpoints, &keys);
assert_eq!(total.get("openai:responses"), Some(&2));
assert_eq!(total.get("openai:search"), Some(&3));
assert_eq!(active, total);
}
}
fn masked_proxy_value(proxy: Option<&serde_json::Value>) -> serde_json::Value {
let Some(proxy) = proxy.and_then(serde_json::Value::as_object) else {
return serde_json::Value::Null;
+6
View File
@@ -769,6 +769,12 @@ const ADMIN_API_FORMAT_DEFINITIONS: &[AdminApiFormatDefinition] = &[
default_path: "/v1/responses/compact",
aliases: &["responses_compact"],
},
AdminApiFormatDefinition {
value: "openai:search",
label: "OpenAI Search",
default_path: "/v1/alpha/search",
aliases: &["openai_search", "search"],
},
AdminApiFormatDefinition {
value: "openai:embedding",
label: "OpenAI Embedding",
+6 -3
View File
@@ -39,7 +39,8 @@ pub use crate::contracts::{
OPENAI_RESPONSES_COMPACT_SYNC_SUCCESS_REPORT_KIND, OPENAI_RESPONSES_STREAM_PLAN_KIND,
OPENAI_RESPONSES_STREAM_SUCCESS_REPORT_KIND, OPENAI_RESPONSES_SYNC_ERROR_REPORT_KIND,
OPENAI_RESPONSES_SYNC_FINALIZE_REPORT_KIND, OPENAI_RESPONSES_SYNC_PLAN_KIND,
OPENAI_RESPONSES_SYNC_SUCCESS_REPORT_KIND, OPENAI_VIDEO_CANCEL_SYNC_PLAN_KIND,
OPENAI_RESPONSES_SYNC_SUCCESS_REPORT_KIND, OPENAI_SEARCH_SYNC_PLAN_KIND,
OPENAI_SEARCH_SYNC_SUCCESS_REPORT_KIND, OPENAI_VIDEO_CANCEL_SYNC_PLAN_KIND,
OPENAI_VIDEO_CONTENT_PLAN_KIND, OPENAI_VIDEO_CREATE_SYNC_FINALIZE_REPORT_KIND,
OPENAI_VIDEO_CREATE_SYNC_PLAN_KIND, OPENAI_VIDEO_DELETE_SYNC_PLAN_KIND,
OPENAI_VIDEO_REMIX_SYNC_PLAN_KIND,
@@ -274,8 +275,10 @@ pub use aether_ai_formats::formats::conversion::response::{
convert_openai_responses_response_to_openai_chat, OpenAiResponsesResponseUsage,
};
pub use aether_ai_formats::{
api_format_alias_matches, api_format_storage_aliases, is_openai_responses_compact_format,
is_openai_responses_family_format, is_openai_responses_format, normalize_api_format_alias,
api_format_alias_matches, api_format_permission_covers, api_format_permission_storage_aliases,
api_format_storage_aliases, intersect_api_format_allowed_lists,
is_openai_responses_compact_format, is_openai_responses_family_format,
is_openai_responses_format, normalize_api_format_alias,
};
pub use aether_ai_formats::{
canonical_request_unknown_block_count, canonical_response_unknown_block_count,
@@ -23,9 +23,9 @@ pub use plan_kinds::{
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_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,
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,
};
pub use report_kinds::{
core_error_background_report_kind, core_error_default_client_api_format,
@@ -52,5 +52,6 @@ pub use report_kinds::{
OPENAI_RESPONSES_COMPACT_SYNC_FINALIZE_REPORT_KIND,
OPENAI_RESPONSES_COMPACT_SYNC_SUCCESS_REPORT_KIND, OPENAI_RESPONSES_STREAM_SUCCESS_REPORT_KIND,
OPENAI_RESPONSES_SYNC_ERROR_REPORT_KIND, OPENAI_RESPONSES_SYNC_FINALIZE_REPORT_KIND,
OPENAI_RESPONSES_SYNC_SUCCESS_REPORT_KIND, OPENAI_VIDEO_CREATE_SYNC_FINALIZE_REPORT_KIND,
OPENAI_RESPONSES_SYNC_SUCCESS_REPORT_KIND, OPENAI_SEARCH_SYNC_SUCCESS_REPORT_KIND,
OPENAI_VIDEO_CREATE_SYNC_FINALIZE_REPORT_KIND,
};
@@ -23,6 +23,7 @@ pub const OPENAI_VIDEO_CREATE_SYNC_PLAN_KIND: &str = "openai_video_create_sync";
pub const OPENAI_CHAT_SYNC_PLAN_KIND: &str = "openai_chat_sync";
pub const OPENAI_EMBEDDING_SYNC_PLAN_KIND: &str = "openai_embedding_sync";
pub const OPENAI_RERANK_SYNC_PLAN_KIND: &str = "openai_rerank_sync";
pub const OPENAI_SEARCH_SYNC_PLAN_KIND: &str = "openai_search_sync";
pub const GEMINI_EMBEDDING_SYNC_PLAN_KIND: &str = "gemini_embedding_sync";
pub const OPENAI_RESPONSES_SYNC_PLAN_KIND: &str = "openai_responses_sync";
pub const OPENAI_RESPONSES_COMPACT_SYNC_PLAN_KIND: &str = "openai_responses_compact_sync";
@@ -29,6 +29,7 @@ pub const OPENAI_RESPONSES_SYNC_SUCCESS_REPORT_KIND: &str = "openai_responses_sy
pub const OPENAI_RESPONSES_COMPACT_SYNC_SUCCESS_REPORT_KIND: &str =
"openai_responses_compact_sync_success";
pub const OPENAI_EMBEDDING_SYNC_SUCCESS_REPORT_KIND: &str = "openai_embedding_sync_success";
pub const OPENAI_SEARCH_SYNC_SUCCESS_REPORT_KIND: &str = "openai_search_sync_success";
pub const GEMINI_EMBEDDING_SYNC_SUCCESS_REPORT_KIND: &str = "gemini_embedding_sync_success";
pub const OPENAI_IMAGE_SYNC_SUCCESS_REPORT_KIND: &str = "openai_image_sync_success";
pub const CLAUDE_CLI_SYNC_SUCCESS_REPORT_KIND: &str = "claude_cli_sync_success";
+172 -1
View File
@@ -23,6 +23,7 @@ pub enum FormatId {
OpenAiChat,
OpenAiResponses,
OpenAiResponsesCompact,
OpenAiSearch,
OpenAiEmbedding,
OpenAiRerank,
ClaudeMessages,
@@ -49,6 +50,7 @@ impl FormatId {
Self::OpenAiChat
| Self::OpenAiResponses
| Self::OpenAiResponsesCompact
| Self::OpenAiSearch
| Self::OpenAiEmbedding
| Self::OpenAiRerank => FormatFamily::OpenAi,
Self::ClaudeMessages => FormatFamily::Claude,
@@ -73,6 +75,7 @@ impl FormatId {
Self::OpenAiChat => "openai:chat",
Self::OpenAiResponses => "openai:responses",
Self::OpenAiResponsesCompact => "openai:responses:compact",
Self::OpenAiSearch => "openai:search",
Self::OpenAiEmbedding => "openai:embedding",
Self::OpenAiRerank => "openai:rerank",
Self::ClaudeMessages => "claude:messages",
@@ -103,6 +106,9 @@ impl FromStr for FormatId {
"openai:responses:compact" | "/v1/responses/compact" => {
Ok(Self::OpenAiResponsesCompact)
}
"openai:search" | "openai_search" | "search" | "/v1/alpha/search" => {
Ok(Self::OpenAiSearch)
}
"openai:embedding" | "/v1/embeddings" => Ok(Self::OpenAiEmbedding),
"openai:rerank" | "/v1/rerank" => Ok(Self::OpenAiRerank),
"claude:messages" | "/v1/messages" => Ok(Self::ClaudeMessages),
@@ -139,6 +145,56 @@ pub fn api_format_alias_matches(left: &str, right: &str) -> bool {
normalize_api_format_alias(left) == normalize_api_format_alias(right)
}
pub fn api_format_defaults_to_non_stream(value: &str) -> bool {
matches!(
normalize_api_format_alias(value).as_str(),
"openai:chat"
| "openai:responses"
| "openai:responses:compact"
| "openai:search"
| "openai:image"
| "claude:messages"
)
}
pub fn api_format_defaults_to_client_error_failover(value: &str) -> bool {
!matches!(
FormatId::parse(value).map(FormatId::canonical),
Some(FormatId::OpenAiSearch)
)
}
pub fn api_format_permission_covers(allowed_value: &str, requested_api_format: &str) -> bool {
let allowed_value = normalize_api_format_alias(allowed_value);
let requested_api_format = normalize_api_format_alias(requested_api_format);
!allowed_value.is_empty()
&& !requested_api_format.is_empty()
&& (allowed_value == requested_api_format
|| allowed_value == "openai:responses" && requested_api_format == "openai:search")
}
pub fn intersect_api_format_allowed_lists(left: &[String], right: &[String]) -> Vec<String> {
let mut effective = Vec::new();
for left_value in left {
for right_value in right {
let intersection = if api_format_permission_covers(right_value, left_value) {
Some(left_value)
} else if api_format_permission_covers(left_value, right_value) {
Some(right_value)
} else {
None
};
if let Some(value) = intersection {
let normalized = normalize_api_format_alias(value);
if !effective.iter().any(|item| item == &normalized) {
effective.push(normalized);
}
}
}
}
effective
}
pub fn api_format_storage_aliases(value: &str) -> Vec<String> {
match FormatId::parse(value).map(FormatId::canonical) {
Some(FormatId::AliyunMultimodalEmbedding) => vec![
@@ -149,6 +205,22 @@ pub fn api_format_storage_aliases(value: &str) -> Vec<String> {
}
}
pub fn api_format_permission_storage_aliases(value: &str) -> Vec<String> {
let requested_api_format = normalize_api_format_alias(value);
let mut aliases = api_format_storage_aliases(&requested_api_format);
for allowed_api_format in [FormatId::OpenAiResponses.as_str()] {
if !api_format_permission_covers(allowed_api_format, &requested_api_format) {
continue;
}
for alias in api_format_storage_aliases(allowed_api_format) {
if !aliases.iter().any(|existing| existing == &alias) {
aliases.push(alias);
}
}
}
aliases
}
pub fn is_openai_responses_format(value: &str) -> bool {
normalize_api_format_alias(value) == "openai:responses"
}
@@ -179,7 +251,10 @@ pub fn api_format_uses_body_stream_field(value: &str) -> bool {
#[cfg(test)]
mod tests {
use super::{
api_format_alias_matches, api_format_storage_aliases, api_format_uses_body_stream_field,
api_format_alias_matches, api_format_defaults_to_client_error_failover,
api_format_defaults_to_non_stream, api_format_permission_covers,
api_format_permission_storage_aliases, api_format_storage_aliases,
api_format_uses_body_stream_field, intersect_api_format_allowed_lists,
normalize_api_format_alias, FormatId,
};
@@ -193,6 +268,102 @@ mod tests {
assert_eq!(FormatId::parse("gemini:cli"), None);
}
#[test]
fn responses_permission_covers_only_its_search_companion() {
assert!(api_format_permission_covers(
"OPENAI:RESPONSES",
"openai:search"
));
assert!(api_format_permission_covers(
"openai:search",
"openai:search"
));
assert!(!api_format_permission_covers(
"openai:search",
"openai:responses"
));
assert!(!api_format_permission_covers(
"openai:responses",
"openai:chat"
));
assert_eq!(
api_format_permission_storage_aliases("openai:search"),
vec!["openai:search".to_string(), "openai:responses".to_string()]
);
assert_eq!(
api_format_permission_storage_aliases("openai:responses"),
vec!["openai:responses".to_string()]
);
}
#[test]
fn normalizes_openai_search_aliases() {
for alias in [
"openai:search",
"OPENAI_SEARCH",
"search",
"/v1/alpha/search",
] {
assert_eq!(FormatId::parse(alias), Some(FormatId::OpenAiSearch));
assert_eq!(normalize_api_format_alias(alias), "openai:search");
}
assert!(!api_format_uses_body_stream_field("openai:search"));
}
#[test]
fn identifies_default_non_stream_formats_from_aliases() {
for format in [
"/v1/chat/completions",
"/v1/responses",
"/v1/responses/compact",
"/v1/alpha/search",
"openai:image",
"/v1/messages",
] {
assert!(api_format_defaults_to_non_stream(format), "{format}");
}
assert!(!api_format_defaults_to_non_stream("gemini:interactions"));
}
#[test]
fn search_defaults_to_passthrough_for_client_errors() {
for format in ["openai:search", "OPENAI_SEARCH", "/v1/alpha/search"] {
assert!(
!api_format_defaults_to_client_error_failover(format),
"{format}"
);
}
assert!(api_format_defaults_to_client_error_failover(
"openai:responses"
));
assert!(api_format_defaults_to_client_error_failover(
"custom:unknown"
));
}
#[test]
fn api_format_policy_intersection_keeps_the_narrowest_companion_scope() {
assert_eq!(
intersect_api_format_allowed_lists(
&["openai:responses".to_string()],
&["openai:search".to_string()],
),
vec!["openai:search".to_string()]
);
assert_eq!(
intersect_api_format_allowed_lists(
&["openai:search".to_string()],
&["OPENAI:RESPONSES".to_string()],
),
vec!["openai:search".to_string()]
);
assert!(intersect_api_format_allowed_lists(
&["openai:search".to_string()],
&["openai:chat".to_string()],
)
.is_empty());
}
#[test]
fn parses_embedding_api_formats() {
assert_eq!(
@@ -57,6 +57,9 @@ pub fn request_candidate_api_format_preference(
if client_api_format == "openai:responses:compact" {
return (provider_api_format == "openai:responses:compact").then_some((0, 0));
}
if client_api_format == "openai:search" {
return (provider_api_format == "openai:search").then_some((0, 0));
}
if is_gemini_interactions_api_format(client_api_format.as_str()) {
return (provider_api_format == "gemini:interactions").then_some((0, 0));
}
@@ -109,6 +112,9 @@ pub fn request_candidate_api_formats(
if client_api_format == "openai:responses:compact" {
return vec!["openai:responses:compact"];
}
if client_api_format == "openai:search" {
return vec!["openai:search"];
}
if is_gemini_interactions_api_format(client_api_format.as_str()) {
return GEMINI_INTERACTIONS_CANDIDATE_API_FORMATS.to_vec();
}
@@ -377,6 +383,10 @@ mod tests {
request_conversion_kind("openai:compact", "openai:responses:compact"),
None
);
assert_eq!(
request_conversion_kind("openai:search", "openai:responses"),
None
);
assert_eq!(
request_conversion_kind("gemini:generate_content", "claude:messages"),
Some(RequestConversionKind::ToClaudeStandard)
@@ -409,6 +419,26 @@ mod tests {
}
}
#[test]
fn search_candidate_registry_keeps_exact_protocol_identity() {
assert_eq!(
request_candidate_api_formats("openai:search", false),
vec!["openai:search"]
);
assert_eq!(
request_candidate_api_formats("/v1/alpha/search", true),
vec!["openai:search"]
);
assert_eq!(
request_candidate_api_format_preference("openai:search", "openai:search"),
Some((0, 0))
);
assert_eq!(
request_candidate_api_format_preference("openai:search", "openai:responses"),
None
);
}
#[test]
fn sync_response_conversion_registry_supports_bidirectional_standard_matrix() {
assert_eq!(
+4 -1
View File
@@ -13,7 +13,10 @@ pub mod shared;
pub use context::{FormatContext, FormatError};
pub use id::{
api_format_alias_matches, api_format_storage_aliases, is_openai_responses_compact_format,
api_format_alias_matches, api_format_defaults_to_client_error_failover,
api_format_defaults_to_non_stream, api_format_permission_covers,
api_format_permission_storage_aliases, api_format_storage_aliases,
intersect_api_format_allowed_lists, is_openai_responses_compact_format,
is_openai_responses_family_format, is_openai_responses_format, normalize_api_format_alias,
FormatFamily, FormatId, FormatProfile,
};
@@ -6,5 +6,6 @@ pub mod reasoning;
pub mod request_contract;
pub mod rerank;
pub mod responses;
pub mod search;
pub mod shared;
pub mod video;
@@ -68,19 +68,21 @@ pub(crate) fn validate_openai_reasoning_request_with_model_profile(
};
let source_api_format = crate::normalize_api_format_alias(source_api_format);
let reasoning = match source_api_format.as_str() {
"openai:responses" | "openai:responses:compact" => match object.get("reasoning") {
Some(Value::Object(reasoning)) => Some(reasoning),
Some(Value::Null) => None,
Some(value) => {
return Err(OpenAiReasoningContractViolation {
kind: OpenAiReasoningViolationKind::InvalidType,
field: "reasoning".to_string(),
value: Some(value.to_string()),
reason: "reasoning must be an object".to_string(),
});
"openai:responses" | "openai:responses:compact" | "openai:search" => {
match object.get("reasoning") {
Some(Value::Object(reasoning)) => Some(reasoning),
Some(Value::Null) => None,
Some(value) => {
return Err(OpenAiReasoningContractViolation {
kind: OpenAiReasoningViolationKind::InvalidType,
field: "reasoning".to_string(),
value: Some(value.to_string()),
reason: "reasoning must be an object".to_string(),
});
}
None => None,
}
None => None,
},
}
"openai:chat" => None,
_ => return Ok(()),
};
@@ -94,7 +96,7 @@ pub(crate) fn validate_openai_reasoning_request_with_model_profile(
let effort = match source_api_format.as_str() {
"openai:chat" => object.get("reasoning_effort"),
"openai:responses" | "openai:responses:compact" => {
"openai:responses" | "openai:responses:compact" | "openai:search" => {
reasoning.and_then(|reasoning| reasoning.get("effort"))
}
_ => None,
@@ -110,11 +112,13 @@ pub(crate) fn validate_openai_reasoning_request_with_model_profile(
)?;
}
if let Some(mode) = reasoning
.and_then(|reasoning| reasoning.get("mode"))
.filter(|value| !value.is_null())
{
validate_reasoning_mode(mode, provider_model, source_model, supports_reasoning_mode)?;
if source_api_format != "openai:search" {
if let Some(mode) = reasoning
.and_then(|reasoning| reasoning.get("mode"))
.filter(|value| !value.is_null())
{
validate_reasoning_mode(mode, provider_model, source_model, supports_reasoning_mode)?;
}
}
if let Some(context) = reasoning
.and_then(|reasoning| reasoning.get("context"))
@@ -35,20 +35,21 @@ pub fn finalize_openai_provider_request_with_codex_model_capabilities(
finalization: OpenAiProviderRequestFinalization<'_>,
model_capabilities: Option<&super::responses::codex::CodexResponsesModelCapabilities>,
) -> Result<(), OpenAiProviderRequestContractViolation> {
let is_codex_responses = finalization
let is_codex_reasoning_endpoint = finalization
.provider_type
.trim()
.eq_ignore_ascii_case("codex")
&& crate::is_openai_responses_family_format(finalization.provider_api_format);
let resolved_model_capabilities =
(is_codex_responses && model_capabilities.is_none()).then(|| {
&& (crate::is_openai_responses_family_format(finalization.provider_api_format)
|| crate::api_format_alias_matches(finalization.provider_api_format, "openai:search"));
let resolved_model_capabilities = (is_codex_reasoning_endpoint && model_capabilities.is_none())
.then(|| {
super::responses::codex::resolve_codex_responses_model_capabilities(
finalization.provider_model,
finalization.source_model,
None,
)
});
let model_capabilities = is_codex_responses
let model_capabilities = is_codex_reasoning_endpoint
.then(|| model_capabilities.or(resolved_model_capabilities.as_ref()))
.flatten();
match crate::normalize_api_format_alias(finalization.source_api_format).as_str() {
@@ -75,6 +76,11 @@ pub fn finalize_openai_provider_request_with_codex_model_capabilities(
)
}
}
super::responses::codex::normalize_codex_openai_reasoning_wire_effort(
body,
finalization.provider_type,
finalization.provider_api_format,
);
super::responses::codex::apply_openai_responses_compact_special_body_edits(
body,
finalization.provider_api_format,
@@ -85,6 +91,7 @@ pub fn finalize_openai_provider_request_with_codex_model_capabilities(
finalization.upstream_is_stream,
finalization.require_body_stream_field,
);
super::search::apply_openai_search_request_projection(body, finalization.provider_api_format);
let provider_model = body
.get("model")
.and_then(Value::as_str)
@@ -415,6 +422,81 @@ mod tests {
assert_eq!(ultra_only_body["reasoning"]["effort"], "max");
}
#[test]
fn codex_search_normalizes_reasoning_and_projects_the_typed_request() {
let finalization = OpenAiProviderRequestFinalization {
source_api_format: "openai:search",
provider_api_format: "openai:search",
provider_type: "codex",
provider_model: "gpt-5.6-sol",
source_model: "gpt-5.6-sol",
body_rules: None,
upstream_is_stream: false,
require_body_stream_field: false,
};
let mut body = json!({
"id": "session-1",
"model": "gpt-5.6-sol",
"reasoning": {
"effort": "ultra",
"summary": "auto",
"context": "current_turn",
"future_reasoning_field": true
},
"commands": {"search_query": [{"q": "Aether"}]},
"store": false,
"future_request_field": {"enabled": true},
"stream": true
});
finalize_openai_provider_request(&mut body, finalization)
.expect("Codex Search request should finalize");
assert_eq!(body["reasoning"]["effort"], "max");
assert_eq!(body["reasoning"]["summary"], "auto");
assert_eq!(body["reasoning"]["context"], "current_turn");
assert!(body["reasoning"].get("future_reasoning_field").is_none());
assert_eq!(body["commands"]["search_query"][0]["q"], "Aether");
assert!(body.get("store").is_none());
assert!(body.get("future_request_field").is_none());
assert!(body.get("stream").is_none());
assert!(body.get("tool_choice").is_none());
assert!(body.get("include").is_none());
}
#[test]
fn codex_search_validates_reasoning_effort_against_model_card() {
let finalization = OpenAiProviderRequestFinalization {
source_api_format: "openai:search",
provider_api_format: "openai:search",
provider_type: "codex",
provider_model: "gpt-5.6-sol",
source_model: "gpt-5.6-sol",
body_rules: None,
upstream_is_stream: false,
require_body_stream_field: false,
};
let mut supported = json!({
"id": "session-1",
"model": "gpt-5.6-sol",
"reasoning": {"effort": "high"}
});
finalize_openai_provider_request(&mut supported, finalization)
.expect("published Search effort should pass");
let mut unsupported = json!({
"id": "session-1",
"model": "gpt-5.6-sol",
"reasoning": {"effort": "none"}
});
let error = finalize_openai_provider_request(&mut unsupported, finalization)
.expect_err("unpublished Search effort should be rejected");
assert!(matches!(
error,
super::OpenAiProviderRequestContractViolation::Reasoning(_)
));
}
#[test]
fn dynamic_codex_card_preserves_custom_reasoning_effort_case() {
let finalization = OpenAiProviderRequestFinalization {
@@ -65,6 +65,7 @@ fn is_openai_image_request(provider_api_format: &str) -> bool {
enum CodexOpenAiEndpointKind {
Responses,
Compact,
Search,
Images,
}
@@ -601,6 +602,8 @@ fn codex_openai_endpoint_kind(
Some(CodexOpenAiEndpointKind::Compact)
} else if aether_ai_formats::is_openai_responses_format(provider_api_format) {
Some(CodexOpenAiEndpointKind::Responses)
} else if aether_ai_formats::api_format_alias_matches(provider_api_format, "openai:search") {
Some(CodexOpenAiEndpointKind::Search)
} else if is_openai_image_request(provider_api_format) {
Some(CodexOpenAiEndpointKind::Images)
} else {
@@ -786,6 +789,16 @@ fn remove_btree_header(headers: &mut BTreeMap<String, String>, header_name: &str
headers.retain(|name, _| !name.trim().eq_ignore_ascii_case(header_name));
}
fn header_value_contains_media_type(value: &str, media_type: &str) -> bool {
value.split(',').any(|media_range| {
media_range
.split(';')
.next()
.map(str::trim)
.is_some_and(|value| value.eq_ignore_ascii_case(media_type))
})
}
#[derive(Debug, Clone, Default, PartialEq, Eq)]
pub struct CodexAuthIdentity {
pub account_id: Option<String>,
@@ -1179,6 +1192,23 @@ fn normalize_codex_reasoning_effort(body_object: &mut serde_json::Map<String, Va
}
}
pub fn normalize_codex_openai_reasoning_wire_effort(
provider_request_body: &mut Value,
provider_type: &str,
provider_api_format: &str,
) {
if !provider_type.trim().eq_ignore_ascii_case("codex")
|| !(aether_ai_formats::is_openai_responses_family_format(provider_api_format)
|| aether_ai_formats::api_format_alias_matches(provider_api_format, "openai:search"))
{
return;
}
let Some(body_object) = provider_request_body.as_object_mut() else {
return;
};
normalize_codex_reasoning_effort(body_object);
}
fn is_codex_responses_lite_additional_tools_item(value: &Value) -> bool {
value
.get("type")
@@ -1719,6 +1749,17 @@ pub fn apply_codex_openai_special_headers(
"originator",
CODEX_CLIENT_ORIGINATOR,
);
if endpoint_kind == CodexOpenAiEndpointKind::Search {
remove_btree_header(provider_request_headers, CODEX_RESPONSES_LITE_HEADER);
remove_btree_header(provider_request_headers, "openai-beta");
if provider_request_headers.iter().any(|(name, value)| {
name.eq_ignore_ascii_case("accept")
&& header_value_contains_media_type(value, "text/event-stream")
}) {
remove_btree_header(provider_request_headers, "accept");
}
return;
}
if endpoint_kind == CodexOpenAiEndpointKind::Images {
return;
}
@@ -1753,8 +1794,9 @@ mod tests {
apply_codex_openai_special_headers, apply_openai_responses_compact_special_body_edits,
build_codex_model_catalog_metadata, bundled_codex_model_cards, effective_codex_model_cards,
resolve_codex_responses_model_capabilities,
validate_codex_openai_responses_compact_request_contract,
CODEX_OPENAI_IMAGE_INTERNAL_MODEL, CODEX_OPENAI_RESPONSES_UNSUPPORTED_BODY_FIELDS,
validate_codex_openai_responses_compact_request_contract, CODEX_CLIENT_ORIGINATOR,
CODEX_CLIENT_USER_AGENT, CODEX_OPENAI_IMAGE_INTERNAL_MODEL,
CODEX_OPENAI_RESPONSES_UNSUPPORTED_BODY_FIELDS, CODEX_RESPONSES_LITE_HEADER,
};
use serde_json::{json, Value};
@@ -2265,6 +2307,51 @@ mod tests {
);
}
#[test]
fn codex_search_uses_identity_headers_without_responses_protocol_headers() {
let mut headers = std::collections::BTreeMap::from([
(
"x-openai-internal-codex-responses-lite".to_string(),
"true".to_string(),
),
("openai-beta".to_string(), "responses=v1".to_string()),
(
"accept".to_string(),
"application/json, Text/Event-Stream; q=0.9".to_string(),
),
]);
apply_codex_openai_special_headers(
&mut headers,
&json!({"id": "session-1", "model": "gpt-5.6-luna"}),
&http::HeaderMap::new(),
"codex",
"openai:search",
Some("request-search"),
Some(r#"{"account_id":"account-1","is_fedramp":true}"#),
);
assert_eq!(
headers.get("chatgpt-account-id").map(String::as_str),
Some("account-1")
);
assert_eq!(
headers.get("x-openai-fedramp").map(String::as_str),
Some("true")
);
assert_eq!(
headers.get("user-agent").map(String::as_str),
Some(CODEX_CLIENT_USER_AGENT)
);
assert_eq!(
headers.get("originator").map(String::as_str),
Some(CODEX_CLIENT_ORIGINATOR)
);
assert!(!headers.contains_key(CODEX_RESPONSES_LITE_HEADER));
assert!(!headers.contains_key("openai-beta"));
assert!(!headers.contains_key("accept"));
}
#[test]
fn standard_codex_models_do_not_send_the_responses_lite_header() {
let mut headers = std::collections::BTreeMap::from([(
@@ -0,0 +1,219 @@
use serde_json::Value;
const REQUEST_FIELDS: &[&str] = &[
"id",
"model",
"reasoning",
"input",
"commands",
"settings",
"max_output_tokens",
];
const REASONING_FIELDS: &[&str] = &["effort", "summary", "context"];
const COMMAND_FIELDS: &[&str] = &[
"search_query",
"image_query",
"open",
"click",
"find",
"screenshot",
"finance",
"weather",
"sports",
"time",
"response_length",
];
const COMMAND_ITEM_FIELDS: &[(&str, &[&str])] = &[
("search_query", &["q", "recency", "domains"]),
("image_query", &["q", "recency", "domains"]),
("open", &["ref_id", "lineno"]),
("click", &["ref_id", "id"]),
("find", &["ref_id", "pattern"]),
("screenshot", &["ref_id", "pageno"]),
("finance", &["ticker", "type", "market"]),
("weather", &["location", "start", "duration"]),
(
"sports",
&[
"tool",
"fn",
"league",
"team",
"opponent",
"date_from",
"date_to",
"num_games",
"locale",
],
),
("time", &["utc_offset"]),
];
const SETTINGS_FIELDS: &[&str] = &[
"user_location",
"search_context_size",
"filters",
"image_settings",
"allowed_callers",
"external_web_access",
];
const USER_LOCATION_FIELDS: &[&str] = &["type", "country", "region", "city", "timezone"];
const FILTER_FIELDS: &[&str] = &["allowed_domains", "blocked_domains"];
const IMAGE_SETTINGS_FIELDS: &[&str] = &["max_results", "caption"];
fn retain_object_fields(value: &mut Value, fields: &[&str]) {
if let Some(object) = value.as_object_mut() {
object.retain(|field, _| fields.contains(&field.as_str()));
}
}
fn retain_array_object_fields(
object: &mut serde_json::Map<String, Value>,
key: &str,
fields: &[&str],
) {
if let Some(items) = object.get_mut(key).and_then(Value::as_array_mut) {
for item in items {
retain_object_fields(item, fields);
}
}
}
pub fn apply_openai_search_request_projection(body: &mut Value, provider_api_format: &str) {
if !crate::api_format_alias_matches(provider_api_format, "openai:search") {
return;
}
let Some(body_object) = body.as_object_mut() else {
return;
};
body_object.retain(|field, _| REQUEST_FIELDS.contains(&field.as_str()));
if let Some(reasoning) = body_object.get_mut("reasoning") {
retain_object_fields(reasoning, REASONING_FIELDS);
}
if let Some(commands) = body_object
.get_mut("commands")
.and_then(Value::as_object_mut)
{
commands.retain(|field, _| COMMAND_FIELDS.contains(&field.as_str()));
for (key, fields) in COMMAND_ITEM_FIELDS {
retain_array_object_fields(commands, key, fields);
}
}
if let Some(settings) = body_object
.get_mut("settings")
.and_then(Value::as_object_mut)
{
settings.retain(|field, _| SETTINGS_FIELDS.contains(&field.as_str()));
if let Some(user_location) = settings.get_mut("user_location") {
retain_object_fields(user_location, USER_LOCATION_FIELDS);
}
if let Some(filters) = settings.get_mut("filters") {
retain_object_fields(filters, FILTER_FIELDS);
}
if let Some(image_settings) = settings.get_mut("image_settings") {
retain_object_fields(image_settings, IMAGE_SETTINGS_FIELDS);
}
}
}
#[cfg(test)]
mod tests {
use serde_json::json;
use super::apply_openai_search_request_projection;
#[test]
fn projects_the_typed_search_request_contract() {
let mut body = json!({
"id": "session-1",
"model": "gpt-5.6-sol",
"reasoning": {
"effort": "max",
"summary": "auto",
"context": "current_turn",
"mode": "pro"
},
"input": "find documentation",
"commands": {
"search_query": [{"q": "Aether", "recency": 7, "unknown": true}],
"image_query": [{"q": "Aether UI", "domains": ["example.com"], "unknown": true}],
"open": [{"ref_id": "turn0search0", "lineno": 12, "unknown": true}],
"click": [{"ref_id": "turn0fetch0", "id": 3, "unknown": true}],
"find": [{"ref_id": "turn0fetch0", "pattern": "Aether", "unknown": true}],
"screenshot": [{"ref_id": "turn0fetch0", "pageno": 2, "unknown": true}],
"finance": [{"ticker": "OPENAI", "type": "equity", "market": "USA", "unknown": true}],
"weather": [{"location": "US, CA, San Francisco", "duration": 3, "unknown": true}],
"sports": [{"tool": "sports", "fn": "schedule", "league": "nba", "team": "GSW", "unknown": true}],
"time": [{"utc_offset": "+08:00", "unknown": true}],
"response_length": "short",
"unknown": true
},
"settings": {
"user_location": {"type": "approximate", "country": "US", "unknown": true},
"search_context_size": "high",
"filters": {"allowed_domains": ["openai.com"], "unknown": true},
"image_settings": {"max_results": 3, "unknown": true},
"allowed_callers": ["direct"],
"external_web_access": "live",
"unknown": true
},
"max_output_tokens": 1024,
"store": false,
"stream": true,
"service_tier": "priority",
"unknown": true
});
apply_openai_search_request_projection(&mut body, "/v1/alpha/search");
assert_eq!(
body["reasoning"],
json!({
"effort": "max",
"summary": "auto",
"context": "current_turn"
})
);
assert_eq!(
body["commands"],
json!({
"search_query": [{"q": "Aether", "recency": 7}],
"image_query": [{"q": "Aether UI", "domains": ["example.com"]}],
"open": [{"ref_id": "turn0search0", "lineno": 12}],
"click": [{"ref_id": "turn0fetch0", "id": 3}],
"find": [{"ref_id": "turn0fetch0", "pattern": "Aether"}],
"screenshot": [{"ref_id": "turn0fetch0", "pageno": 2}],
"finance": [{"ticker": "OPENAI", "type": "equity", "market": "USA"}],
"weather": [{"location": "US, CA, San Francisco", "duration": 3}],
"sports": [{"tool": "sports", "fn": "schedule", "league": "nba", "team": "GSW"}],
"time": [{"utc_offset": "+08:00"}],
"response_length": "short"
})
);
assert_eq!(
body["settings"],
json!({
"user_location": {"type": "approximate", "country": "US"},
"search_context_size": "high",
"filters": {"allowed_domains": ["openai.com"]},
"image_settings": {"max_results": 3},
"allowed_callers": ["direct"],
"external_web_access": "live"
})
);
assert!(body.get("store").is_none());
assert!(body.get("stream").is_none());
assert!(body.get("service_tier").is_none());
assert!(body.get("unknown").is_none());
}
#[test]
fn leaves_other_formats_unchanged() {
let mut body = json!({"model": "gpt-5.6-sol", "store": true});
let expected = body.clone();
apply_openai_search_request_projection(&mut body, "openai:responses");
assert_eq!(body, expected);
}
}
@@ -38,7 +38,7 @@ pub fn parse_request(
FormatId::GeminiEmbedding => gemini::embedding::request::from(body, ctx),
FormatId::DoubaoEmbedding => doubao::embedding::request::from(body, ctx),
FormatId::AliyunMultimodalEmbedding => aliyun::embedding::request::from(body, ctx),
FormatId::GeminiInteractions => None,
FormatId::OpenAiSearch | FormatId::GeminiInteractions => None,
}
.ok_or_else(|| FormatError::RequestParseFailed {
format: source.as_str().to_string(),
@@ -72,7 +72,7 @@ fn emit_request_inner(
FormatId::GeminiEmbedding => gemini::embedding::request::to(request, ctx),
FormatId::DoubaoEmbedding => doubao::embedding::request::to(request, ctx),
FormatId::AliyunMultimodalEmbedding => aliyun::embedding::request::to(request, ctx),
FormatId::GeminiInteractions => None,
FormatId::OpenAiSearch | FormatId::GeminiInteractions => None,
}
.ok_or_else(|| FormatError::RequestEmitFailed {
format: target.as_str().to_string(),
@@ -280,6 +280,7 @@ pub fn parse_response(
FormatId::ClaudeMessages => claude_messages::response::from(body, ctx),
FormatId::GeminiGenerateContent => gemini_generate_content::response::from(body, ctx),
FormatId::OpenAiEmbedding
| FormatId::OpenAiSearch
| FormatId::JinaEmbedding
| FormatId::OpenAiRerank
| FormatId::JinaRerank
@@ -318,6 +319,7 @@ fn emit_response_inner(
FormatId::ClaudeMessages => claude_messages::response::to(response, ctx),
FormatId::GeminiGenerateContent => gemini_generate_content::response::to(response, ctx),
FormatId::OpenAiEmbedding
| FormatId::OpenAiSearch
| FormatId::JinaEmbedding
| FormatId::OpenAiRerank
| FormatId::JinaRerank
@@ -1055,6 +1057,7 @@ fn standard_request_root_field_is_audited(source: FormatId, key: &str) -> bool {
| "tools"
),
FormatId::OpenAiEmbedding
| FormatId::OpenAiSearch
| FormatId::OpenAiRerank
| FormatId::GeminiEmbedding
| FormatId::JinaEmbedding
@@ -1535,6 +1538,7 @@ fn validate_source_response_stop_enums(
FormatId::GeminiGenerateContent => validate_gemini_response_finish_reasons(body, target),
FormatId::GeminiInteractions => Ok(()),
FormatId::OpenAiEmbedding
| FormatId::OpenAiSearch
| FormatId::OpenAiRerank
| FormatId::GeminiEmbedding
| FormatId::JinaEmbedding
@@ -38,7 +38,11 @@ pub fn build_core_error_body_for_client_format(
error_object.insert("message".to_string(), Value::String(message.to_string()));
match aether_ai_formats::normalize_api_format_alias(client_api_format).as_str() {
"openai:chat" | "openai:responses" | "openai:responses:compact" | "openai:embedding" => {
"openai:chat"
| "openai:responses"
| "openai:responses:compact"
| "openai:search"
| "openai:embedding" => {
error_object.insert(
"type".to_string(),
Value::String(map_local_sync_error_kind_to_openai_type(kind).to_string()),
@@ -157,6 +161,21 @@ mod tests {
assert_eq!(body["error"]["code"], "invalid_request");
}
#[test]
fn builds_openai_search_core_error_body() {
let body = build_core_error_body_for_client_format(
"openai:search",
"search unavailable",
Some("upstream_unavailable"),
LocalCoreSyncErrorKind::ServerError,
)
.expect("body should build");
assert_eq!(body["error"]["message"], "search unavailable");
assert_eq!(body["error"]["type"], "server_error");
assert_eq!(body["error"]["code"], "upstream_unavailable");
}
#[test]
fn recognizes_finalize_kind_and_success_mapping() {
assert!(is_core_error_finalize_kind("openai_chat_sync_finalize"));
@@ -1,9 +1,10 @@
use serde_json::{json, Value};
pub const MODEL_DIRECTIVE_API_FORMATS: [&str; 5] = [
pub const MODEL_DIRECTIVE_API_FORMATS: [&str; 6] = [
"openai:chat",
"openai:responses",
"openai:responses:compact",
"openai:search",
"claude:messages",
"gemini:generate_content",
];
@@ -433,7 +434,7 @@ pub fn default_model_directive_mapping_patch(
pub fn default_model_directive_suffixes(provider_api_format: &str) -> &'static [&'static str] {
match crate::normalize_api_format_alias(provider_api_format).as_str() {
"openai:chat" | "openai:responses" | "openai:responses:compact" => {
"openai:chat" | "openai:responses" | "openai:responses:compact" | "openai:search" => {
&OPENAI_MODEL_DIRECTIVE_SUFFIXES
}
"claude:messages" | "gemini:generate_content" => &CROSS_PROVIDER_MODEL_DIRECTIVE_SUFFIXES,
@@ -503,7 +504,7 @@ fn apply_reasoning_effort_override(
"reasoning_effort",
effort.as_openai_model_directive_value(),
),
"openai:responses" | "openai:responses:compact" => {
"openai:responses" | "openai:responses:compact" | "openai:search" => {
set_openai_responses_reasoning_effort(provider_request_body, effort)
}
"claude:messages" => {
@@ -525,7 +526,7 @@ fn apply_codex_reasoning_preset_override(
"openai:chat" => {
set_object_string(provider_request_body, "reasoning_effort", preset.as_str())
}
"openai:responses" | "openai:responses:compact" => {
"openai:responses" | "openai:responses:compact" | "openai:search" => {
let object = provider_request_body.as_object_mut()?;
let reasoning = object
.entry("reasoning".to_string())
@@ -550,6 +551,7 @@ fn apply_service_tier_override(
"service_tier",
tier.as_openai_value(),
),
"openai:search" => Some(()),
_ => None,
}
}
@@ -710,7 +712,7 @@ pub fn reasoning_effort_supported_for_model(
effort: ReasoningEffort,
) -> bool {
match crate::normalize_api_format_alias(provider_api_format).as_str() {
"openai:chat" | "openai:responses" | "openai:responses:compact" => {
"openai:chat" | "openai:responses" | "openai:responses:compact" | "openai:search" => {
match resolved_openai_model_identity(provider_model, source_model).0 {
OpenAiModelIdentity::Gpt56 => effort != ReasoningEffort::Minimal,
OpenAiModelIdentity::ConcreteOther => effort != ReasoningEffort::Max,
@@ -5,7 +5,8 @@ use crate::contracts::{
GEMINI_EMBEDDING_SYNC_SUCCESS_REPORT_KIND, GEMINI_INTERACTIONS_STREAM_PLAN_KIND,
GEMINI_INTERACTIONS_STREAM_SUCCESS_REPORT_KIND, GEMINI_INTERACTIONS_SYNC_PLAN_KIND,
GEMINI_INTERACTIONS_SYNC_SUCCESS_REPORT_KIND, OPENAI_EMBEDDING_SYNC_PLAN_KIND,
OPENAI_RERANK_SYNC_PLAN_KIND,
OPENAI_RERANK_SYNC_PLAN_KIND, OPENAI_SEARCH_SYNC_PLAN_KIND,
OPENAI_SEARCH_SYNC_SUCCESS_REPORT_KIND,
};
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
@@ -81,6 +82,13 @@ pub fn resolve_sync_spec(plan_kind: &str) -> Option<LocalSameFormatProviderSpec>
family: LocalSameFormatProviderFamily::Standard,
require_streaming: false,
}),
OPENAI_SEARCH_SYNC_PLAN_KIND => Some(LocalSameFormatProviderSpec {
api_format: "openai:search",
decision_kind: OPENAI_SEARCH_SYNC_PLAN_KIND,
report_kind: OPENAI_SEARCH_SYNC_SUCCESS_REPORT_KIND,
family: LocalSameFormatProviderFamily::Standard,
require_streaming: false,
}),
_ => None,
}
}
@@ -185,4 +193,12 @@ mod tests {
assert_eq!(spec.report_kind, "openai_rerank_sync_success");
assert!(!spec.require_streaming);
}
#[test]
fn resolves_openai_search_sync_spec() {
let spec = resolve_sync_spec("openai_search_sync").expect("spec");
assert_eq!(spec.api_format, "openai:search");
assert_eq!(spec.report_kind, "openai_search_sync_success");
assert!(!spec.require_streaming);
}
}
@@ -50,7 +50,8 @@ pub fn forbid_upstream_streaming_for_provider(
provider_type: &str,
provider_api_format: &str,
) -> bool {
aether_ai_formats::is_openai_responses_compact_format(provider_api_format)
aether_ai_formats::api_format_alias_matches(provider_api_format, "openai:search")
|| aether_ai_formats::is_openai_responses_compact_format(provider_api_format)
|| (provider_type.trim().eq_ignore_ascii_case("codex")
&& provider_api_format
.trim()
@@ -258,7 +259,15 @@ mod tests {
}
#[test]
fn forbids_streaming_for_compact_and_codex_images() {
fn forbids_streaming_for_sync_only_openai_formats() {
assert!(forbid_upstream_streaming_for_provider(
"codex",
"openai:search"
));
assert!(forbid_upstream_streaming_for_provider(
"custom",
"/v1/alpha/search"
));
assert!(forbid_upstream_streaming_for_provider(
"codex",
"openai:responses:compact"
@@ -461,6 +470,13 @@ mod tests {
#[test]
fn provider_policy_gives_non_stream_contracts_precedence() {
let force_stream = json!({"upstream_stream_policy": "force_stream"});
assert!(!resolve_upstream_is_stream_for_provider(
Some(&force_stream),
"codex",
"openai:search",
true,
true,
));
assert!(!resolve_upstream_is_stream_for_provider(
Some(&force_stream),
"codex",
@@ -13,9 +13,9 @@ use crate::contracts::{
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_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,
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::formats::openai::image::request::is_openai_image_stream_request;
@@ -216,6 +216,14 @@ pub fn resolve_execution_runtime_sync_plan_kind(
return Some(OPENAI_RERANK_SYNC_PLAN_KIND);
}
if route_family == Some("openai")
&& route_kind == Some("search")
&& *method == Method::POST
&& path == "/v1/alpha/search"
{
return Some(OPENAI_SEARCH_SYNC_PLAN_KIND);
}
if route_family == Some("openai")
&& route_kind == Some("image")
&& *method == Method::POST
@@ -433,6 +441,7 @@ pub fn supports_sync_execution_decision_kind(plan_kind: &str) -> bool {
OPENAI_CHAT_SYNC_PLAN_KIND
| OPENAI_EMBEDDING_SYNC_PLAN_KIND
| OPENAI_RERANK_SYNC_PLAN_KIND
| OPENAI_SEARCH_SYNC_PLAN_KIND
| OPENAI_IMAGE_SYNC_PLAN_KIND
| OPENAI_RESPONSES_SYNC_PLAN_KIND
| OPENAI_RESPONSES_COMPACT_SYNC_PLAN_KIND
@@ -492,6 +501,7 @@ mod tests {
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,
};
#[test]
@@ -606,6 +616,46 @@ mod tests {
));
}
#[test]
fn resolves_openai_search_as_sync_only() {
assert_eq!(
resolve_execution_runtime_sync_plan_kind(
Some("ai_public"),
Some("openai"),
Some("search"),
None,
&Method::POST,
"/v1/alpha/search",
),
Some(OPENAI_SEARCH_SYNC_PLAN_KIND)
);
assert_eq!(
resolve_execution_runtime_stream_plan_kind(
Some("ai_public"),
Some("openai"),
Some("search"),
None,
&Method::POST,
"/v1/alpha/search",
),
None
);
assert!(supports_sync_execution_decision_kind(
OPENAI_SEARCH_SYNC_PLAN_KIND
));
assert_eq!(
resolve_execution_runtime_sync_plan_kind(
Some("ai_public"),
Some("openai"),
Some("search"),
None,
&Method::POST,
"/backend-api/codex/alpha/search",
),
None
);
}
#[test]
fn resolves_claude_messages_plan_kinds_by_request_auth_channel() {
assert_eq!(
@@ -309,6 +309,7 @@ impl ProviderStreamParser {
FormatId::ClaudeMessages => Self::Claude(ClaudeProviderState::default()),
FormatId::GeminiGenerateContent => Self::Gemini(GeminiProviderState::default()),
FormatId::OpenAiEmbedding
| FormatId::OpenAiSearch
| FormatId::OpenAiRerank
| FormatId::GeminiEmbedding
| FormatId::GeminiInteractions
@@ -409,6 +410,7 @@ impl ClientStreamEmitter {
FormatId::ClaudeMessages => Self::Claude(ClaudeClientEmitter::default()),
FormatId::GeminiGenerateContent => Self::Gemini(GeminiClientEmitter::default()),
FormatId::OpenAiEmbedding
| FormatId::OpenAiSearch
| FormatId::OpenAiRerank
| FormatId::GeminiEmbedding
| FormatId::GeminiInteractions
@@ -531,6 +533,7 @@ fn parse_provider_error(
parse_gemini_error(payload)
}
FormatId::OpenAiEmbedding
| FormatId::OpenAiSearch
| FormatId::OpenAiRerank
| FormatId::GeminiEmbedding
| FormatId::JinaEmbedding
+4 -1
View File
@@ -11,7 +11,10 @@ pub use formats::context::{
FormatError,
};
pub use formats::id::{
api_format_alias_matches, api_format_storage_aliases, api_format_uses_body_stream_field,
api_format_alias_matches, api_format_defaults_to_client_error_failover,
api_format_defaults_to_non_stream, api_format_permission_covers,
api_format_permission_storage_aliases, api_format_storage_aliases,
api_format_uses_body_stream_field, intersect_api_format_allowed_lists,
is_openai_responses_compact_format, is_openai_responses_family_format,
is_openai_responses_format, normalize_api_format_alias, FormatFamily, FormatId, FormatProfile,
};
+3 -1
View File
@@ -11,7 +11,9 @@ pub use frame::{StreamFrame, StreamFramePayload, StreamFrameType};
pub use plan::{
ExecutionPlan, ExecutionTimeouts, ProxySnapshot, RequestBody, ResolvedTransportProfile,
EXECUTION_REQUEST_ACCEPT_INVALID_CERTS_HEADER, EXECUTION_REQUEST_FOLLOW_REDIRECTS_HEADER,
EXECUTION_REQUEST_HTTP1_ONLY_HEADER, TRANSPORT_BACKEND_BROWSER_WREQ,
EXECUTION_REQUEST_HTTP1_ONLY_HEADER, MAX_EXECUTION_REQUEST_TIMEOUT_MS,
MAX_EXECUTION_REQUEST_TIMEOUT_SECS, MAX_EXECUTION_STREAM_FIRST_BYTE_TIMEOUT_MS,
MAX_EXECUTION_STREAM_FIRST_BYTE_TIMEOUT_SECS, TRANSPORT_BACKEND_BROWSER_WREQ,
TRANSPORT_BACKEND_HYPER_RUSTLS, TRANSPORT_BACKEND_REQWEST_RUSTLS, TRANSPORT_HTTP_MODE_AUTO,
TRANSPORT_HTTP_MODE_H2C_PRIOR_KNOWLEDGE, TRANSPORT_HTTP_MODE_HTTP1_ONLY,
TRANSPORT_POOL_SCOPE_KEY,
+5
View File
@@ -7,6 +7,11 @@ pub const EXECUTION_REQUEST_FOLLOW_REDIRECTS_HEADER: &str = "x-aether-execution-
pub const EXECUTION_REQUEST_HTTP1_ONLY_HEADER: &str = "x-aether-execution-http1-only";
pub const EXECUTION_REQUEST_ACCEPT_INVALID_CERTS_HEADER: &str =
"x-aether-execution-accept-invalid-certs";
pub const MAX_EXECUTION_REQUEST_TIMEOUT_SECS: u64 = 1_200;
pub const MAX_EXECUTION_REQUEST_TIMEOUT_MS: u64 = MAX_EXECUTION_REQUEST_TIMEOUT_SECS * 1_000;
pub const MAX_EXECUTION_STREAM_FIRST_BYTE_TIMEOUT_SECS: u64 = 300;
pub const MAX_EXECUTION_STREAM_FIRST_BYTE_TIMEOUT_MS: u64 =
MAX_EXECUTION_STREAM_FIRST_BYTE_TIMEOUT_SECS * 1_000;
#[derive(Debug, Clone, Default, Serialize, Deserialize, PartialEq, Eq)]
#[serde(default)]
+163 -3
View File
@@ -12,6 +12,7 @@ pub const TUNNEL_PROTOCOL_VERSION_HEADER: &str = "x-aether-tunnel-protocol-versi
pub const TUNNEL_NODE_NAME_B64_HEADER: &str = "x-aether-tunnel-node-name-b64";
pub const CURRENT_TUNNEL_PROTOCOL_VERSION: u8 = 3;
pub const CURRENT_TUNNEL_PROTOCOL_VERSION_STR: &str = "3";
pub const MAX_TUNNEL_RELAY_META_LEN: usize = 256 * 1024;
pub mod flags {
pub const END_STREAM: u8 = 0x01;
@@ -220,6 +221,63 @@ pub struct RequestMeta {
pub transport_profile: Option<crate::ResolvedTransportProfile>,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct ResolvedTunnelRequestTimeouts {
pub first_byte_ms: u64,
pub response_body_ms: Option<u64>,
}
pub fn resolve_tunnel_request_timeouts(meta: &RequestMeta) -> ResolvedTunnelRequestTimeouts {
let legacy_timeout_ms = meta.timeout.saturating_mul(1_000);
let first_byte_ms = if meta.stream {
meta.stream_first_byte_timeout_ms
.unwrap_or(legacy_timeout_ms)
} else {
meta.request_timeout_ms
.or(meta.stream_first_byte_timeout_ms)
.unwrap_or(legacy_timeout_ms)
};
let response_body_ms = (!meta.stream).then_some(first_byte_ms);
ResolvedTunnelRequestTimeouts {
first_byte_ms: if meta.stream {
clamp_stream_first_byte_timeout_ms(first_byte_ms)
} else {
clamp_upstream_request_timeout_ms(first_byte_ms)
},
response_body_ms: response_body_ms.map(clamp_upstream_request_timeout_ms),
}
}
pub fn try_decode_tunnel_relay_request_meta(
buffer: &[u8],
) -> Result<Option<(RequestMeta, usize)>, String> {
if buffer.len() < 4 {
return Ok(None);
}
let meta_len = u32::from_be_bytes([buffer[0], buffer[1], buffer[2], buffer[3]]) as usize;
if meta_len > MAX_TUNNEL_RELAY_META_LEN {
return Err("relay metadata too large".to_string());
}
let meta_end = 4usize
.checked_add(meta_len)
.ok_or_else(|| "relay envelope length overflow".to_string())?;
if buffer.len() < meta_end {
return Ok(None);
}
let meta = serde_json::from_slice::<RequestMeta>(&buffer[4..meta_end])
.map_err(|error| format!("invalid relay metadata: {error}"))?;
Ok(Some((meta, meta_end)))
}
fn clamp_upstream_request_timeout_ms(timeout_ms: u64) -> u64 {
timeout_ms.clamp(1, crate::MAX_EXECUTION_REQUEST_TIMEOUT_MS)
}
fn clamp_stream_first_byte_timeout_ms(timeout_ms: u64) -> u64 {
timeout_ms.clamp(1, crate::MAX_EXECUTION_STREAM_FIRST_BYTE_TIMEOUT_MS)
}
fn default_timeout() -> u64 {
60
}
@@ -461,13 +519,115 @@ fn compress_gzip(data: &[u8]) -> Result<Bytes, std::io::Error> {
mod tests {
use super::{
compress_payload, decode_payload, encode_frame, encode_goaway_v3, encode_ping,
encode_reset_stream, encode_window_update, raw_payload, Frame, FrameHeader, GoAwayPayload,
MsgType, RequestMeta, ResetStreamPayload, WindowUpdatePayload,
CURRENT_TUNNEL_PROTOCOL_VERSION, CURRENT_TUNNEL_PROTOCOL_VERSION_STR, FLAG_GZIP_COMPRESSED,
encode_reset_stream, encode_window_update, raw_payload, resolve_tunnel_request_timeouts,
try_decode_tunnel_relay_request_meta, Frame, FrameHeader, GoAwayPayload, MsgType,
RequestMeta, ResetStreamPayload, WindowUpdatePayload, CURRENT_TUNNEL_PROTOCOL_VERSION,
CURRENT_TUNNEL_PROTOCOL_VERSION_STR, FLAG_GZIP_COMPRESSED, MAX_TUNNEL_RELAY_META_LEN,
REQUEST_HEADERS, TUNNEL_PROTOCOL_VERSION_HEADER,
};
use bytes::Bytes;
fn request_meta(stream: bool) -> RequestMeta {
RequestMeta {
provider_id: None,
endpoint_id: None,
key_id: None,
method: "POST".to_string(),
url: "https://example.com/responses".to_string(),
headers: std::collections::HashMap::new(),
stream,
request_timeout_ms: None,
stream_first_byte_timeout_ms: None,
timeout: 60,
follow_redirects: None,
http1_only: false,
transport_profile: None,
}
}
#[test]
fn tunnel_request_timeouts_preserve_non_stream_total_timeout() {
let mut meta = request_meta(false);
meta.request_timeout_ms = Some(900_000);
meta.stream_first_byte_timeout_ms = Some(12_000);
let resolved = resolve_tunnel_request_timeouts(&meta);
assert_eq!(resolved.first_byte_ms, 900_000);
assert_eq!(resolved.response_body_ms, Some(900_000));
}
#[test]
fn tunnel_request_timeouts_keep_stream_body_unbounded() {
let mut meta = request_meta(true);
meta.request_timeout_ms = Some(900_000);
meta.stream_first_byte_timeout_ms = Some(12_000);
let resolved = resolve_tunnel_request_timeouts(&meta);
assert_eq!(resolved.first_byte_ms, 12_000);
assert_eq!(resolved.response_body_ms, None);
}
#[test]
fn tunnel_request_timeouts_keep_stream_first_byte_protocol_limit() {
let mut meta = request_meta(true);
meta.stream_first_byte_timeout_ms =
Some(crate::MAX_EXECUTION_STREAM_FIRST_BYTE_TIMEOUT_MS + 1);
let resolved = resolve_tunnel_request_timeouts(&meta);
assert_eq!(
resolved.first_byte_ms,
crate::MAX_EXECUTION_STREAM_FIRST_BYTE_TIMEOUT_MS
);
assert_eq!(resolved.response_body_ms, None);
}
#[test]
fn tunnel_request_timeouts_clamp_only_out_of_range_protocol_values() {
let mut meta = request_meta(false);
meta.request_timeout_ms = Some(crate::MAX_EXECUTION_REQUEST_TIMEOUT_MS + 1);
let resolved = resolve_tunnel_request_timeouts(&meta);
assert_eq!(
resolved.first_byte_ms,
crate::MAX_EXECUTION_REQUEST_TIMEOUT_MS
);
assert_eq!(
resolved.response_body_ms,
Some(crate::MAX_EXECUTION_REQUEST_TIMEOUT_MS)
);
}
#[test]
fn tunnel_relay_request_meta_decodes_from_a_partial_prefix() {
let meta = request_meta(false);
let encoded_meta = serde_json::to_vec(&meta).expect("meta should encode");
let mut envelope = Vec::new();
envelope.extend_from_slice(&(encoded_meta.len() as u32).to_be_bytes());
envelope.extend_from_slice(&encoded_meta);
envelope.extend_from_slice(b"request-body");
assert!(try_decode_tunnel_relay_request_meta(&envelope[..3])
.expect("partial prefix should be valid")
.is_none());
let (decoded, body_offset) = try_decode_tunnel_relay_request_meta(&envelope)
.expect("envelope should be valid")
.expect("metadata should be complete");
assert_eq!(decoded.request_timeout_ms, meta.request_timeout_ms);
assert_eq!(&envelope[body_offset..], b"request-body");
}
#[test]
fn tunnel_relay_request_meta_rejects_oversized_prefix() {
let oversized = (MAX_TUNNEL_RELAY_META_LEN as u32 + 1).to_be_bytes();
assert!(try_decode_tunnel_relay_request_meta(&oversized).is_err());
}
#[test]
fn request_meta_accepts_integer_timeout() {
let raw = br#"{"method":"GET","url":"https://example.com","headers":{},"timeout":15}"#;
@@ -97,13 +97,13 @@ impl StoredMinimalCandidateSelectionRow {
None => true,
Some(formats) => formats
.iter()
.any(|value| api_format_matches(value, api_format)),
.any(|value| api_format_permission_covers(value, api_format)),
}
}
}
fn api_format_matches(left: &str, right: &str) -> bool {
aether_ai_formats::api_format_alias_matches(left, right)
fn api_format_permission_covers(allowed: &str, requested: &str) -> bool {
aether_ai_formats::api_format_permission_covers(allowed, requested)
}
#[async_trait]
@@ -247,7 +247,7 @@ impl ResolvedAuthApiKeySnapshot {
&mut self.user_allowed_providers,
&mut self.api_key_allowed_providers,
);
constrain_api_key_list_policy_to_user_policy(
constrain_api_key_api_format_policy_to_user_policy(
&mut self.user_allowed_api_formats,
&mut self.api_key_allowed_api_formats,
);
@@ -277,6 +277,22 @@ fn constrain_api_key_list_policy_to_user_policy(
*api_key_policy = Some(effective);
}
fn constrain_api_key_api_format_policy_to_user_policy(
user_policy: &mut Option<Vec<String>>,
api_key_policy: &mut Option<Vec<String>>,
) {
let Some(api_key_values) = api_key_policy.as_ref() else {
return;
};
let Some(user_values) = user_policy.as_ref() else {
return;
};
let effective =
aether_ai_formats::intersect_api_format_allowed_lists(api_key_values, user_values);
*user_policy = Some(effective.clone());
*api_key_policy = Some(effective);
}
fn intersect_allowed_lists(left: &[String], right: &[String]) -> Vec<String> {
let right_values = right.iter().collect::<std::collections::BTreeSet<_>>();
left.iter()
@@ -810,6 +826,29 @@ mod tests {
StoredAuthApiKeySnapshot,
};
#[test]
fn api_format_policy_intersection_preserves_search_companion_scope() {
assert_eq!(
aether_ai_formats::intersect_api_format_allowed_lists(
&["openai:search".to_string()],
&["openai:responses".to_string()],
),
vec!["openai:search".to_string()]
);
assert_eq!(
aether_ai_formats::intersect_api_format_allowed_lists(
&["openai:responses".to_string()],
&["openai:search".to_string()],
),
vec!["openai:search".to_string()]
);
assert!(aether_ai_formats::intersect_api_format_allowed_lists(
&["openai:search".to_string()],
&["openai:chat".to_string()],
)
.is_empty());
}
#[test]
fn rejects_non_array_allowed_providers() {
assert!(StoredAuthApiKeySnapshot::new(
@@ -230,7 +230,7 @@ fn row_matches_requested_model(
mapping.api_formats.as_ref().is_none_or(|formats| {
formats
.iter()
.any(|value| api_format_matches(value, api_format))
.any(|value| api_format_scope_covers(value, api_format))
}) && mapping.endpoint_ids.as_ref().is_none_or(|endpoint_ids| {
endpoint_ids
.iter()
@@ -286,7 +286,7 @@ fn mapping_scope_matches(
mapping.api_formats.as_ref().is_none_or(|formats| {
formats
.iter()
.any(|value| api_format_matches(value, api_format))
.any(|value| api_format_scope_covers(value, api_format))
}) && mapping.endpoint_ids.as_ref().is_none_or(|endpoint_ids| {
endpoint_ids
.iter()
@@ -294,6 +294,10 @@ fn mapping_scope_matches(
})
}
fn api_format_scope_covers(allowed: &str, requested: &str) -> bool {
aether_ai_formats::api_format_permission_covers(allowed, requested)
}
fn key_auth_channel_matches(row: &StoredMinimalCandidateSelectionRow, api_format: &str) -> bool {
let provider_type = row.provider_type.trim().to_ascii_lowercase();
let auth_type = row.key_auth_type.trim().to_ascii_lowercase();
@@ -303,7 +307,10 @@ fn key_auth_channel_matches(row: &StoredMinimalCandidateSelectionRow, api_format
auth_type == "oauth"
&& matches!(
api_format.as_str(),
"openai:responses" | "openai:responses:compact" | "openai:image"
"openai:responses"
| "openai:responses:compact"
| "openai:search"
| "openai:image"
)
}
"chatgpt_web" => {
@@ -434,6 +441,47 @@ mod tests {
assert_eq!(rows[0].provider_id, "provider-1");
}
#[tokio::test]
async fn search_uses_responses_key_and_model_permissions_with_exact_endpoint_identity() {
let mut search = sample_row(
"provider-search",
"openai:search",
"global-search-model",
10,
);
search.provider_type = "codex".to_string();
search.key_auth_type = "oauth".to_string();
search.key_api_formats = Some(vec!["openai:responses".to_string()]);
search.model_provider_model_name = "upstream-search-model".to_string();
search.model_provider_model_mappings = Some(vec![StoredProviderModelMapping {
name: "gpt-5.6-sol".to_string(),
priority: 0,
api_formats: Some(vec!["openai:responses".to_string()]),
endpoint_ids: None,
}]);
let mut responses = search.clone();
responses.provider_id = "provider-responses".to_string();
responses.endpoint_id = "endpoint-responses".to_string();
responses.endpoint_api_format = "openai:responses".to_string();
responses.key_id = "key-responses".to_string();
responses.model_id = "model-responses".to_string();
let repository =
InMemoryMinimalCandidateSelectionReadRepository::seed(vec![responses, search]);
let rows = repository
.list_for_exact_api_format_and_requested_model("openai:search", "gpt-5.6-sol")
.await
.expect("Search candidate should load");
assert_eq!(rows.len(), 1);
assert_eq!(rows[0].endpoint_api_format, "openai:search");
assert_eq!(
rows[0].key_api_formats,
Some(vec!["openai:responses".to_string()])
);
}
#[tokio::test]
async fn includes_grok_oauth_rows_for_chat_models() {
let mut row = sample_row(
@@ -409,7 +409,7 @@ fn mapping_scope_matches(
mapping.api_formats.as_ref().is_none_or(|formats| {
formats
.iter()
.any(|value| api_format_matches(value, api_format))
.any(|value| api_format_scope_covers(value, api_format))
}) && mapping.endpoint_ids.as_ref().is_none_or(|endpoint_ids| {
endpoint_ids
.iter()
@@ -426,7 +426,10 @@ fn key_auth_channel_matches(row: &CandidateSelectionRow, api_format: &str) -> bo
auth_type == "oauth"
&& matches!(
api_format.as_str(),
"openai:responses" | "openai:responses:compact" | "openai:image"
"openai:responses"
| "openai:responses:compact"
| "openai:search"
| "openai:image"
)
}
"chatgpt_web" => {
@@ -763,6 +766,10 @@ fn api_format_matches(left: &str, right: &str) -> bool {
aether_ai_formats::api_format_alias_matches(left, right)
}
fn api_format_scope_covers(allowed: &str, requested: &str) -> bool {
aether_ai_formats::api_format_permission_covers(allowed, requested)
}
fn sql_match_aliases(api_formats: &[String]) -> Vec<String> {
api_formats
.iter()
@@ -71,7 +71,7 @@ INNER JOIN LATERAL (
(
LOWER(BTRIM(p.provider_type)) = 'codex'
AND LOWER(BTRIM(pak.auth_type)) = 'oauth'
AND LOWER($3) IN ('openai:responses', 'openai:responses:compact', 'openai:image')
AND LOWER($3) IN ('openai:responses', 'openai:responses:compact', 'openai:search', 'openai:image')
)
OR (
LOWER(BTRIM(p.provider_type)) = 'chatgpt_web'
@@ -164,7 +164,7 @@ WHERE p.is_active = TRUE
(
LOWER(BTRIM(p.provider_type)) = 'codex'
AND LOWER(BTRIM(pak.auth_type)) = 'oauth'
AND LOWER($3) IN ('openai:responses', 'openai:responses:compact', 'openai:image')
AND LOWER($3) IN ('openai:responses', 'openai:responses:compact', 'openai:search', 'openai:image')
)
OR (
LOWER(BTRIM(p.provider_type)) = 'chatgpt_web'
@@ -350,7 +350,7 @@ INNER JOIN LATERAL (
(
LOWER(BTRIM(p.provider_type)) = 'codex'
AND LOWER(BTRIM(pak.auth_type)) = 'oauth'
AND LOWER($4) IN ('openai:responses', 'openai:responses:compact', 'openai:image')
AND LOWER($4) IN ('openai:responses', 'openai:responses:compact', 'openai:search', 'openai:image')
)
OR (
LOWER(BTRIM(p.provider_type)) = 'chatgpt_web'
@@ -444,7 +444,7 @@ WHERE p.is_active = TRUE
(
LOWER(BTRIM(p.provider_type)) = 'codex'
AND LOWER(BTRIM(pak.auth_type)) = 'oauth'
AND LOWER($4) IN ('openai:responses', 'openai:responses:compact', 'openai:image')
AND LOWER($4) IN ('openai:responses', 'openai:responses:compact', 'openai:search', 'openai:image')
)
OR (
LOWER(BTRIM(p.provider_type)) = 'chatgpt_web'
@@ -638,7 +638,7 @@ WHERE p.is_active = TRUE
(
LOWER(BTRIM(p.provider_type)) = 'codex'
AND LOWER(BTRIM(pak.auth_type)) = 'oauth'
AND LOWER($6) IN ('openai:responses', 'openai:responses:compact', 'openai:image')
AND LOWER($6) IN ('openai:responses', 'openai:responses:compact', 'openai:search', 'openai:image')
)
OR (
LOWER(BTRIM(p.provider_type)) = 'chatgpt_web'
@@ -784,7 +784,8 @@ impl SqlxMinimalCandidateSelectionReadRepository {
let mut rows = Vec::new();
let canonical_api_format = normalize_api_format(api_format);
let storage_aliases = api_format_aliases(&canonical_api_format);
let sql_match_aliases = sql_match_aliases(&storage_aliases);
let sql_match_aliases =
sql_match_aliases(&api_format_permission_aliases(&canonical_api_format));
for api_format in storage_aliases {
rows.extend(
Self::collect_query_rows(
@@ -809,7 +810,8 @@ impl SqlxMinimalCandidateSelectionReadRepository {
let mut rows = Vec::new();
let canonical_api_format = normalize_api_format(api_format);
let storage_aliases = api_format_aliases(&canonical_api_format);
let sql_match_aliases = sql_match_aliases(&storage_aliases);
let sql_match_aliases =
sql_match_aliases(&api_format_permission_aliases(&canonical_api_format));
for api_format in storage_aliases {
rows.extend(
Self::collect_query_rows(
@@ -835,7 +837,8 @@ impl SqlxMinimalCandidateSelectionReadRepository {
let mut rows = Vec::new();
let canonical_api_format = normalize_api_format(api_format);
let storage_aliases = api_format_aliases(&canonical_api_format);
let sql_match_aliases = sql_match_aliases(&storage_aliases);
let sql_match_aliases =
sql_match_aliases(&api_format_permission_aliases(&canonical_api_format));
let sql = requested_model_selection_sql();
for api_format in storage_aliases {
rows.extend(
@@ -861,7 +864,8 @@ impl SqlxMinimalCandidateSelectionReadRepository {
let mut rows = Vec::new();
let canonical_api_format = normalize_api_format(&query.api_format);
let storage_aliases = api_format_aliases(&canonical_api_format);
let sql_match_aliases = sql_match_aliases(&storage_aliases);
let sql_match_aliases =
sql_match_aliases(&api_format_permission_aliases(&canonical_api_format));
let limit = i64::from(query.limit.max(1));
let offset = i64::from(query.offset);
let sql = requested_model_selection_page_sql();
@@ -891,7 +895,8 @@ impl SqlxMinimalCandidateSelectionReadRepository {
let mut rows = Vec::new();
let canonical_api_format = normalize_api_format(&query.api_format);
let storage_aliases = api_format_aliases(&canonical_api_format);
let sql_match_aliases = sql_match_aliases(&storage_aliases);
let sql_match_aliases =
sql_match_aliases(&api_format_permission_aliases(&canonical_api_format));
let limit = i64::from(query.limit.max(1));
let offset = i64::from(query.offset);
let sql = pool_key_candidate_selection_sql(&query.order);
@@ -929,7 +934,8 @@ impl SqlxMinimalCandidateSelectionReadRepository {
let mut rows = Vec::new();
let canonical_api_format = normalize_api_format(&query.api_format);
let storage_aliases = api_format_aliases(&canonical_api_format);
let sql_match_aliases = sql_match_aliases(&storage_aliases);
let sql_match_aliases =
sql_match_aliases(&api_format_permission_aliases(&canonical_api_format));
let sql = pool_key_candidate_selection_by_key_ids_sql();
for api_format in storage_aliases {
rows.extend(
@@ -1111,6 +1117,10 @@ fn api_format_aliases(api_format: &str) -> Vec<String> {
aether_ai_formats::api_format_storage_aliases(api_format)
}
fn api_format_permission_aliases(api_format: &str) -> Vec<String> {
aether_ai_formats::api_format_permission_storage_aliases(api_format)
}
fn normalize_api_format(api_format: &str) -> String {
aether_ai_formats::normalize_api_format_alias(api_format)
}
@@ -125,7 +125,8 @@ impl SqliteMinimalCandidateSelectionReadRepository {
) -> Result<Vec<StoredMinimalCandidateSelectionRow>, DataLayerError> {
let canonical_api_format = normalize_api_format(api_format);
let storage_aliases = api_format_aliases(&canonical_api_format);
let match_aliases = sql_match_aliases(&storage_aliases);
let match_aliases =
sql_match_aliases(&api_format_permission_aliases(&canonical_api_format));
let mut rows = Vec::new();
for storage_api_format in storage_aliases {
@@ -270,7 +271,8 @@ impl MinimalCandidateSelectionReadRepository for SqliteMinimalCandidateSelection
) -> Result<Vec<StoredMinimalCandidateSelectionRow>, DataLayerError> {
let canonical_api_format = normalize_api_format(&query.api_format);
let storage_aliases = api_format_aliases(&canonical_api_format);
let match_aliases = sql_match_aliases(&storage_aliases);
let match_aliases =
sql_match_aliases(&api_format_permission_aliases(&canonical_api_format));
let mut rows = Vec::<CandidateSelectionRow>::new();
let page_in_sql = !matches!(query.order, StoredPoolKeyCandidateOrder::LoadBalance { .. });
@@ -337,7 +339,8 @@ impl MinimalCandidateSelectionReadRepository for SqliteMinimalCandidateSelection
.collect::<BTreeMap<_, _>>();
let canonical_api_format = normalize_api_format(&query.api_format);
let storage_aliases = api_format_aliases(&canonical_api_format);
let match_aliases = sql_match_aliases(&storage_aliases);
let match_aliases =
sql_match_aliases(&api_format_permission_aliases(&canonical_api_format));
let mut rows = Vec::new();
for storage_api_format in storage_aliases {
@@ -493,7 +496,7 @@ fn push_key_auth_channel_sql_filter(
);
builder.push_bind(api_format.clone());
builder.push(
r#" IN ('openai:responses', 'openai:responses:compact', 'openai:image')
r#" IN ('openai:responses', 'openai:responses:compact', 'openai:search', 'openai:image')
)
OR (
LOWER(TRIM(p.provider_type)) = 'chatgpt_web'
@@ -808,7 +811,7 @@ fn mapping_scope_matches(
mapping.api_formats.as_ref().is_none_or(|formats| {
formats
.iter()
.any(|value| api_format_matches(value, api_format))
.any(|value| api_format_scope_covers(value, api_format))
}) && mapping.endpoint_ids.as_ref().is_none_or(|endpoint_ids| {
endpoint_ids
.iter()
@@ -825,7 +828,10 @@ fn key_auth_channel_matches(row: &CandidateSelectionRow, api_format: &str) -> bo
auth_type == "oauth"
&& matches!(
api_format.as_str(),
"openai:responses" | "openai:responses:compact" | "openai:image"
"openai:responses"
| "openai:responses:compact"
| "openai:search"
| "openai:image"
)
}
"chatgpt_web" => {
@@ -1145,6 +1151,10 @@ fn api_format_aliases(api_format: &str) -> Vec<String> {
aether_ai_formats::api_format_storage_aliases(api_format)
}
fn api_format_permission_aliases(api_format: &str) -> Vec<String> {
aether_ai_formats::api_format_permission_storage_aliases(api_format)
}
fn normalize_api_format(api_format: &str) -> String {
aether_ai_formats::normalize_api_format_alias(api_format)
}
@@ -1153,6 +1163,10 @@ fn api_format_matches(left: &str, right: &str) -> bool {
aether_ai_formats::api_format_alias_matches(left, right)
}
fn api_format_scope_covers(allowed: &str, requested: &str) -> bool {
aether_ai_formats::api_format_permission_covers(allowed, requested)
}
fn sql_match_aliases(api_formats: &[String]) -> Vec<String> {
api_formats
.iter()
@@ -1264,6 +1278,21 @@ mod tests {
.collect::<Vec<_>>(),
vec!["key-chatgpt-web-oauth", "key-chatgpt-web-bearer"]
);
let search_rows = repository
.list_for_exact_api_format_and_requested_model_page(
&StoredRequestedModelCandidateRowsQuery {
api_format: "openai:search".to_string(),
requested_model_name: "gpt-5.6-sol".to_string(),
offset: 0,
limit: 10,
},
)
.await
.expect("Codex Search rows should load through Responses permissions");
assert_eq!(search_rows.len(), 1);
assert_eq!(search_rows[0].key_id, "key-codex-search");
assert_eq!(search_rows[0].endpoint_api_format, "openai:search");
}
async fn seed_candidate_selection(pool: &sqlx::SqlitePool) {
@@ -1296,6 +1325,11 @@ INSERT INTO providers (
)
VALUES ('provider-windsurf', 'Windsurf', 'windsurf', 15, 1, 1, 1);
INSERT INTO providers (
id, name, provider_type, provider_priority, is_active, created_at, updated_at
)
VALUES ('provider-codex-search', 'Codex Search', 'codex', 12, 1, 1, 1);
INSERT INTO provider_endpoints (
id, provider_id, name, base_url, api_format, is_active, created_at, updated_at
)
@@ -1312,6 +1346,14 @@ VALUES (
'https://server.codeium.com', 'openai:chat', 1, 1, 1
);
INSERT INTO provider_endpoints (
id, provider_id, name, base_url, api_format, is_active, created_at, updated_at
)
VALUES (
'endpoint-codex-search', 'provider-codex-search', 'Codex Search',
'https://chatgpt.com/backend-api/codex', 'openai:search', 1, 1, 1
);
INSERT INTO provider_api_keys (
id, provider_id, name, auth_type, api_formats, internal_priority, is_active, created_at, updated_at
)
@@ -1327,13 +1369,22 @@ VALUES (
'key-windsurf-oauth', 'provider-windsurf', 'OAuth', 'oauth', '["openai:chat"]', 10, 1, 1, 1
);
INSERT INTO provider_api_keys (
id, provider_id, name, auth_type, api_formats, internal_priority, is_active, created_at, updated_at
)
VALUES (
'key-codex-search', 'provider-codex-search', 'OAuth', 'oauth',
'["openai:responses"]', 10, 1, 1, 1
);
INSERT INTO global_models (
id, name, config, is_active, created_at, updated_at
)
VALUES
('global-1', 'gpt-5', '{"model_mappings":["alias-global"],"streaming":true}', 1, 1, 1),
('global-image-1', 'gpt-image-2', NULL, 1, 1, 1),
('global-windsurf-1', 'claude-opus-4-7', '{"streaming":true}', 1, 1, 1);
('global-windsurf-1', 'claude-opus-4-7', '{"streaming":true}', 1, 1, 1),
('global-codex-search-1', 'search-global', '{"streaming":false}', 1, 1, 1);
INSERT INTO models (
id, provider_id, global_model_id, provider_model_name, provider_model_mappings,
@@ -1351,6 +1402,11 @@ VALUES (
(
'model-windsurf-opus', 'provider-windsurf', 'global-windsurf-1', 'claude-opus-4-7',
NULL, NULL, 1, 1, 1, 1
),
(
'model-codex-search', 'provider-codex-search', 'global-codex-search-1', 'search-upstream',
'[{"name":"gpt-5.6-sol","api_formats":["openai:responses"],"priority":1}]',
0, 1, 1, 1, 1
);
"#,
)
@@ -354,9 +354,9 @@ fn transport_key_supports_api_format(
match transport.key.api_formats.as_deref() {
None => true,
Some(formats) => formats
.iter()
.any(|value| aether_ai_formats::api_format_alias_matches(value, endpoint_api_format)),
Some(formats) => formats.iter().any(|value| {
aether_ai_formats::api_format_permission_covers(value, endpoint_api_format)
}),
}
}
@@ -818,6 +818,29 @@ mod tests {
);
}
#[test]
fn responses_key_permission_covers_search_without_changing_endpoint_identity() {
let mut transport = transport_snapshot("custom", "openai:search", "bearer", true, None);
transport.key.api_formats = Some(vec!["openai:responses".to_string()]);
assert_eq!(
candidate_common_transport_skip_reason(
&transport,
candidate_facts("openai:search"),
None,
),
None
);
assert_eq!(
candidate_common_transport_skip_reason(
&transport,
candidate_facts("openai:responses"),
None,
),
Some("endpoint_api_format_changed")
);
}
#[test]
fn candidate_transport_pair_policy_reports_disabled_conversion_and_unsupported_pairs() {
let transport = transport_snapshot("custom", "openai:responses", "bearer", false, None);
@@ -454,6 +454,21 @@ mod tests {
assert_eq!(timeouts.first_byte_ms, Some(30_000));
}
#[test]
fn transport_execution_timeouts_preserve_the_configurable_maximum() {
let mut transport = sample_transport();
transport.provider.request_timeout_secs =
Some(aether_contracts::MAX_EXECUTION_REQUEST_TIMEOUT_SECS as f64);
let timeouts = resolve_transport_execution_timeouts(&transport)
.expect("provider timeouts should resolve");
assert_eq!(
timeouts.total_ms,
Some(aether_contracts::MAX_EXECUTION_REQUEST_TIMEOUT_MS)
);
}
#[test]
fn transport_execution_timeouts_preserve_configured_first_byte_value() {
let mut transport = sample_transport();
@@ -306,6 +306,12 @@ const CODEX_FIXED_PROVIDER_TEMPLATE: FixedProviderTemplate = FixedProviderTempla
custom_path: None,
config_defaults: EMPTY_ENDPOINT_CONFIG_DEFAULTS,
},
FixedProviderEndpointTemplate {
item_key: "openai:search",
api_format: "openai:search",
custom_path: None,
config_defaults: EMPTY_ENDPOINT_CONFIG_DEFAULTS,
},
FixedProviderEndpointTemplate {
item_key: "openai:image",
api_format: "openai:image",
@@ -639,7 +645,7 @@ mod tests {
};
#[test]
fn codex_fixed_provider_template_includes_openai_image() {
fn codex_fixed_provider_template_includes_codex_companion_endpoints() {
let template = fixed_provider_template("codex").expect("codex template should exist");
assert_eq!(template.base_url, "https://chatgpt.com/backend-api/codex");
assert_eq!(template.version, 1);
@@ -652,6 +658,7 @@ mod tests {
vec![
"openai:responses",
"openai:responses:compact",
"openai:search",
"openai:image"
]
);
@@ -660,6 +667,11 @@ mod tests {
fixed_provider_endpoint_template_by_api_format("codex", "openai:image")
.expect("codex image endpoint should exist");
assert!(image_template.config_defaults.is_empty());
let search_template =
fixed_provider_endpoint_template_by_api_format("codex", "openai:search")
.expect("codex search endpoint should exist");
assert!(search_template.config_defaults.is_empty());
}
#[test]
@@ -16,7 +16,8 @@ use crate::gemini_cli::{
use crate::snapshot::GatewayProviderTransportSnapshot;
use crate::url::{
build_claude_messages_url, build_gemini_content_url, build_openai_chat_url,
build_openai_responses_url, build_passthrough_path_url, normalize_gemini_content_action_path,
build_openai_responses_url, build_openai_search_url, build_passthrough_path_url,
normalize_gemini_content_action_path,
};
use crate::vertex::{
build_vertex_api_key_gemini_content_url, build_vertex_api_key_gemini_embedding_url,
@@ -124,6 +125,10 @@ fn build_transport_request_url_inner(
params.request_query,
true,
)),
"openai:search" => Some(build_openai_search_url(
&transport.endpoint.base_url,
params.request_query,
)),
"openai:embedding" | "jina:embedding" => {
build_provider_embedding_v1_url(&transport.endpoint.base_url, params.request_query)
}
@@ -868,6 +873,33 @@ mod tests {
assert_eq!(url, "https://api.openai.example/v1/responses?tenant=demo");
}
#[test]
fn builds_openai_search_url_for_codex_provider_root() {
let transport = sample_transport(
"codex",
"openai:search",
"https://chatgpt.com/backend-api/codex",
None,
);
let url = build_transport_request_url(
&transport,
TransportRequestUrlParams {
provider_api_format: "openai:search",
mapped_model: Some("gpt-5.6-luna"),
upstream_is_stream: false,
request_query: Some("tenant=demo"),
kiro_api_region: None,
},
)
.expect("openai search url");
assert_eq!(
url,
"https://chatgpt.com/backend-api/codex/alpha/search?tenant=demo"
);
}
#[test]
fn expands_custom_path_templates_when_hook_does_not_apply() {
let transport = sample_transport(
@@ -149,9 +149,11 @@ pub fn classify_same_format_provider_request_behavior(
params.require_streaming,
is_kiro || is_antigravity || gemini_cli_requires_upstream_streaming,
);
let force_body_stream_field = aether_ai_formats::endpoint_config_forces_upstream_stream_policy(
transport.endpoint.config.as_ref(),
);
let force_body_stream_field =
aether_ai_formats::api_format_uses_body_stream_field(params.provider_api_format)
&& 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 {
@@ -622,6 +624,7 @@ pub fn same_format_provider_transport_unsupported_reason_for_trace(
"openai:chat" => "openai:chat",
"openai:responses" => "openai:responses",
"openai:responses:compact" => "openai:responses:compact",
"openai:search" => "openai:search",
"claude:messages" => "claude:messages",
"gemini:generate_content" => "gemini:generate_content",
"gemini:interactions" => "gemini:interactions",
@@ -702,7 +705,9 @@ 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") {
if aether_ai_formats::api_format_alias_matches(provider_api_format, "openai:embedding")
|| aether_ai_formats::api_format_alias_matches(provider_api_format, "openai:search")
{
resolve_local_openai_bearer_auth(transport)
} else {
resolve_local_standard_auth(transport)
@@ -886,6 +891,21 @@ mod tests {
},
);
assert!(!compact_behavior.upstream_is_stream);
let mut search = sample_transport("codex");
search.endpoint.config = Some(json!({
"upstream_stream_policy": "force_stream"
}));
let search_behavior = classify_same_format_provider_request_behavior(
&search,
SameFormatProviderRequestBehaviorParams {
require_streaming: true,
provider_api_format: "openai:search",
report_kind: "openai_search_sync_success",
},
);
assert!(!search_behavior.upstream_is_stream);
assert!(!search_behavior.force_body_stream_field);
}
#[test]
@@ -1035,6 +1055,31 @@ mod tests {
);
}
#[test]
fn resolves_openai_search_direct_auth_with_bearer_header() {
let mut transport = sample_transport("custom");
transport.endpoint.api_format = "openai:search".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:search",
report_kind: "openai_search_sync_success",
},
);
assert_eq!(
resolve_same_format_provider_direct_auth(
&behavior,
&transport,
SameFormatProviderFamily::Standard,
"openai:search",
),
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");
@@ -1770,6 +1815,44 @@ mod tests {
assert_eq!(body["metadata"]["body_rule_seen"], true);
}
#[test]
fn search_body_projects_the_protocol_contract_and_keeps_fast_routing_only() {
let body = build_same_format_provider_request_body(SameFormatProviderRequestBodyInput {
body_json: &json!({
"id": "session-1",
"model": "gpt-5.6-luna-high-fast",
"commands": {"search_query": [{"q": "Aether"}]},
"settings": {"allowed_callers": ["direct"]},
"store": false,
"future_extension": {"enabled": true},
"stream": true
}),
mapped_model: "gpt-5.6-luna",
client_api_format: "openai:search",
provider_api_format: "openai:search",
source_model: Some("gpt-5.6-luna-high-fast"),
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: true,
})
.expect("search body should build");
assert_eq!(body["id"], "session-1");
assert_eq!(body["model"], "gpt-5.6-luna");
assert_eq!(body["commands"]["search_query"][0]["q"], "Aether");
assert_eq!(body["settings"]["allowed_callers"][0], "direct");
assert_eq!(body["reasoning"]["effort"], "high");
assert!(body.get("store").is_none());
assert!(body.get("future_extension").is_none());
assert!(body.get("service_tier").is_none());
assert!(body.get("stream").is_none());
}
#[test]
fn builds_same_format_headers_with_auth_and_stream_accept() {
let provider_request_body = json!({"model": "upstream-model"});
+24 -1
View File
@@ -28,6 +28,14 @@ pub fn build_openai_responses_url(
url
}
pub fn build_openai_search_url(upstream_base_url: &str, query: Option<&str>) -> String {
let (trimmed, base_query) = split_base_url_query(upstream_base_url);
let trimmed = trimmed.trim_end_matches('/');
let mut url = format!("{trimmed}/alpha/search");
append_merged_query(&mut url, base_query, None, query, &[]);
url
}
pub fn build_openai_image_url(
upstream_base_url: &str,
request_path: Option<&str>,
@@ -424,7 +432,7 @@ mod tests {
build_bigmodel_coding_models_url, build_claude_messages_url, build_gemini_content_url,
build_gemini_files_passthrough_url, build_gemini_video_predict_long_running_url,
build_openai_chat_url, build_openai_compatible_models_url, build_openai_image_url,
build_openai_responses_url, build_passthrough_path_url,
build_openai_responses_url, build_openai_search_url, build_passthrough_path_url,
normalize_gemini_content_action_path,
};
@@ -573,6 +581,21 @@ mod tests {
);
}
#[test]
fn openai_search_url_preserves_api_and_codex_roots() {
assert_eq!(
build_openai_search_url(
"https://api.openai.com/v1?tenant=base",
Some("trace=1&tenant=request")
),
"https://api.openai.com/v1/alpha/search?tenant=request&trace=1"
);
assert_eq!(
build_openai_search_url("https://chatgpt.com/backend-api/codex/", None),
"https://chatgpt.com/backend-api/codex/alpha/search"
);
}
#[test]
fn openai_image_url_uses_images_surface() {
assert_eq!(
+13 -1
View File
@@ -56,7 +56,7 @@ pub fn api_format_matches_allowed_value(allowed_value: &str, api_format: &str) -
if allowed_value.is_empty() || api_format.is_empty() {
return false;
}
crate::normalize_api_format(allowed_value) == crate::normalize_api_format(api_format)
aether_ai_formats::api_format_permission_covers(allowed_value, api_format)
}
pub fn auth_constraints_allow_model(
@@ -248,6 +248,18 @@ mod tests {
"openai:responses",
"openai:responses"
));
assert!(api_format_matches_allowed_value(
"openai:responses",
"openai:search"
));
assert!(api_format_matches_allowed_value(
"openai:search",
"openai:search"
));
assert!(!api_format_matches_allowed_value(
"openai:search",
"openai:responses"
));
assert!(!api_format_matches_allowed_value(
"openai:responses",
"claude:messages"
+40 -1
View File
@@ -279,7 +279,7 @@ fn mapping_scope_matches(
let api_format_matches_scope = mapping.api_formats.as_ref().is_none_or(|api_formats| {
api_formats
.iter()
.any(|value| api_format_matches(value, api_format))
.any(|value| api_format_scope_covers(value, api_format))
});
if !api_format_matches_scope {
return false;
@@ -419,6 +419,10 @@ fn api_format_matches(left: &str, right: &str) -> bool {
normalize_api_format(left) == normalize_api_format(right)
}
fn api_format_scope_covers(allowed: &str, requested: &str) -> bool {
aether_ai_formats::api_format_permission_covers(allowed, requested)
}
fn requested_model_name_candidates(
requested_model_name: &str,
enable_model_directives: bool,
@@ -568,6 +572,41 @@ mod tests {
);
}
#[test]
fn responses_model_mapping_scope_covers_search_in_one_direction() {
let mut row = sample_row("search-global", "search-default");
row.endpoint_api_format = "openai:search".to_string();
row.model_provider_model_mappings = Some(vec![StoredProviderModelMapping {
name: "gpt-5.6-sol".to_string(),
priority: 1,
api_formats: Some(vec!["openai:responses".to_string()]),
endpoint_ids: None,
}]);
assert!(row_supports_requested_model(
&row,
"gpt-5.6-sol",
"openai:search"
));
assert_eq!(
resolve_provider_model_name(&row, "gpt-5.6-sol", "openai:search")
.map(|resolved| resolved.0),
Some("gpt-5.6-sol".to_string())
);
row.model_provider_model_mappings = Some(vec![StoredProviderModelMapping {
name: "search-only".to_string(),
priority: 1,
api_formats: Some(vec!["openai:search".to_string()]),
endpoint_ids: None,
}]);
assert!(!row_supports_requested_model(
&row,
"search-only",
"openai:responses"
));
}
#[test]
fn model_directive_suffix_prefers_exact_model_before_base_fallback() {
let exact = sample_row("gpt-5.4-high", "gpt-5.4-high-upstream");
+25
View File
@@ -109,6 +109,9 @@ pub fn infer_internal_finalize_signature(payload: &GatewaySyncReportRequest) ->
if report_kind.starts_with("openai_image_") {
return Some("openai:image".to_string());
}
if report_kind.starts_with("openai_search_") {
return Some("openai:search".to_string());
}
if report_kind.starts_with("openai_cli_") {
return Some("openai:responses".to_string());
}
@@ -150,6 +153,11 @@ pub fn resolve_internal_finalize_route(signature: &str) -> Option<InternalFinali
route_family: "openai",
route_kind: "responses:compact",
}),
"openai:search" => Some(InternalFinalizeRoute {
public_path: "/v1/alpha/search",
route_family: "openai",
route_kind: "search",
}),
"openai:image" => Some(InternalFinalizeRoute {
public_path: "/v1/images/generations",
route_family: "openai",
@@ -246,6 +254,7 @@ pub fn is_local_ai_sync_report_kind(report_kind: &str) -> bool {
| "openai_responses_compact_sync_error"
| "openai_cli_sync_success"
| "openai_image_sync_success"
| "openai_search_sync_success"
| "openai_image_sync_error"
| "openai_embedding_sync_success"
| "openai_embedding_sync_error"
@@ -872,6 +881,7 @@ mod tests {
"openai_responses_compact_sync_error"
));
assert!(is_local_ai_sync_report_kind("openai_image_sync_success"));
assert!(is_local_ai_sync_report_kind("openai_search_sync_success"));
assert!(is_local_ai_sync_report_kind("openai_image_sync_error"));
assert!(is_local_ai_sync_report_kind(
"openai_embedding_sync_success"
@@ -1082,12 +1092,27 @@ mod tests {
Some("openai:responses:compact".to_string())
);
let from_search_report_kind =
sample_sync_report_with_context("openai_search_sync_finalize", json!({}));
assert_eq!(
infer_internal_finalize_signature(&from_search_report_kind),
Some("openai:search".to_string())
);
let unknown = sample_sync_report("unknown_sync_finalize", 200);
assert_eq!(infer_internal_finalize_signature(&unknown), None);
}
#[test]
fn resolves_internal_finalize_route_for_supported_signatures() {
assert_eq!(
resolve_internal_finalize_route("openai:search"),
Some(InternalFinalizeRoute {
public_path: "/v1/alpha/search",
route_family: "openai",
route_kind: "search",
})
);
assert_eq!(
resolve_internal_finalize_route("openai:responses:compact"),
Some(InternalFinalizeRoute {
@@ -1,6 +1,9 @@
import { describe, expect, it } from 'vitest'
import { buildModelsDevTieredPricing } from '@/api/models-dev-pricing'
import {
buildModelsDevTieredPricing,
resolveModelsDevTieredPricing,
} from '@/api/models-dev-pricing'
describe('buildModelsDevTieredPricing', () => {
it('maps context bands and cache prices without flattening them', () => {
@@ -99,3 +102,75 @@ describe('buildModelsDevTieredPricing', () => {
expect(buildModelsDevTieredPricing(cost)).toBeNull()
})
})
describe('resolveModelsDevTieredPricing', () => {
it.each([
{
modelId: 'gpt-5.6-sol',
standard: [5, 30, 6.25, 0.5],
longContext: [10, 45, 12.5, 1],
},
{
modelId: 'gpt-5.6-terra',
standard: [2.5, 15, 3.125, 0.25],
longContext: [5, 22.5, 6.25, 0.5],
},
{
modelId: 'gpt-5.6-luna',
standard: [1, 6, 1.25, 0.1],
longContext: [2, 9, 2.5, 0.2],
},
])('uses the complete OpenAI catalog for $modelId', ({ modelId, standard, longContext }) => {
const tier = (
upTo: number | null,
prices: number[],
multiplier: number,
) => ({
up_to: upTo,
input_price_per_1m: prices[0] * multiplier,
output_price_per_1m: prices[1] * multiplier,
cache_creation_price_per_1m: prices[2] * multiplier,
cache_read_price_per_1m: prices[3] * multiplier,
})
expect(resolveModelsDevTieredPricing('openai', modelId, { input: 999, output: 999 }))
.toEqual({
tiers: [
tier(272_000, standard, 1),
tier(null, longContext, 1),
],
processing_tiers: {
flex: {
tiers: [
tier(272_000, standard, 0.5),
tier(null, longContext, 0.5),
],
},
priority: {
tiers: [tier(272_000, standard, 2)],
},
},
})
})
it('keeps the models.dev lower-bound conversion for models outside the catalog', () => {
expect(resolveModelsDevTieredPricing('openai', 'other-model', {
input: 1,
output: 2,
tiers: [{ input: 3, output: 4, tier: { type: 'context', size: 272_000 } }],
})?.tiers.map(tier => tier.up_to)).toEqual([271_999, null])
})
it.each([
['other-provider', 'gpt-5.6-sol'],
['openai', 'GPT-5.6-SOL'],
['openai', 'gpt-5.6-sol-latest'],
['openai', '__proto__'],
['openai', 'constructor'],
])('matches provider and model identities exactly for %s/%s', (providerId, modelId) => {
expect(resolveModelsDevTieredPricing(providerId, modelId, { input: 1, output: 2 }))
.toEqual({
tiers: [{ up_to: null, input_price_per_1m: 1, output_price_per_1m: 2 }],
})
})
})
@@ -0,0 +1,71 @@
import type { PricingTier, TieredPricingConfig } from './endpoints/types'
interface TokenPrices {
input: number
output: number
cacheCreation: number
cacheRead: number
}
interface ContextPricing {
standard: TokenPrices
longContext: TokenPrices
}
const OPENAI_GPT_56_PRICING = new Map<string, ContextPricing>([
['gpt-5.6-sol', {
standard: { input: 5, output: 30, cacheCreation: 6.25, cacheRead: 0.5 },
longContext: { input: 10, output: 45, cacheCreation: 12.5, cacheRead: 1 },
}],
['gpt-5.6-terra', {
standard: { input: 2.5, output: 15, cacheCreation: 3.125, cacheRead: 0.25 },
longContext: { input: 5, output: 22.5, cacheCreation: 6.25, cacheRead: 0.5 },
}],
['gpt-5.6-luna', {
standard: { input: 1, output: 6, cacheCreation: 1.25, cacheRead: 0.1 },
longContext: { input: 2, output: 9, cacheCreation: 2.5, cacheRead: 0.2 },
}],
])
const STANDARD_CONTEXT_LIMIT = 272_000
function pricingTier(
upTo: number | null,
prices: TokenPrices,
multiplier = 1,
): PricingTier {
return {
up_to: upTo,
input_price_per_1m: prices.input * multiplier,
output_price_per_1m: prices.output * multiplier,
cache_creation_price_per_1m: prices.cacheCreation * multiplier,
cache_read_price_per_1m: prices.cacheRead * multiplier,
}
}
function contextPricingTiers(pricing: ContextPricing, multiplier: number): PricingTier[] {
return [
pricingTier(STANDARD_CONTEXT_LIMIT, pricing.standard, multiplier),
pricingTier(null, pricing.longContext, multiplier),
]
}
export function getAuthoritativeModelPricing(
providerId: string,
modelId: string,
): TieredPricingConfig | null {
if (providerId !== 'openai') return null
const pricing = OPENAI_GPT_56_PRICING.get(modelId)
if (!pricing) return null
return {
tiers: contextPricingTiers(pricing, 1),
processing_tiers: {
flex: { tiers: contextPricingTiers(pricing, 0.5) },
priority: {
tiers: [pricingTier(STANDARD_CONTEXT_LIMIT, pricing.standard, 2)],
},
},
}
}
@@ -2,6 +2,7 @@ import { describe, expect, it } from 'vitest'
import {
API_FORMATS,
apiFormatPermissionCovers,
formatApiFormat,
formatApiFormatShort,
groupApiFormats,
@@ -14,6 +15,8 @@ describe('api format display helpers', () => {
expect(normalizeApiFormatAlias('CLAUDE_MESSAGES')).toBe(API_FORMATS.CLAUDE_MESSAGES)
expect(normalizeApiFormatAlias('OPENAI_RESPONSES')).toBe(API_FORMATS.OPENAI_RESPONSES)
expect(normalizeApiFormatAlias('OPENAI_RESPONSES_COMPACT')).toBe(API_FORMATS.OPENAI_RESPONSES_COMPACT)
expect(normalizeApiFormatAlias('OPENAI_SEARCH')).toBe(API_FORMATS.OPENAI_SEARCH)
expect(normalizeApiFormatAlias('SEARCH')).toBe(API_FORMATS.OPENAI_SEARCH)
expect(normalizeApiFormatAlias('GEMINI_GENERATE_CONTENT')).toBe(API_FORMATS.GEMINI_GENERATE_CONTENT)
expect(normalizeApiFormatAlias('OPENAI_EMBEDDING')).toBe(API_FORMATS.OPENAI_EMBEDDING)
expect(normalizeApiFormatAlias('OPENAI_RERANK')).toBe(API_FORMATS.OPENAI_RERANK)
@@ -33,6 +36,25 @@ describe('api format display helpers', () => {
expect(formatApiFormatShort(API_FORMATS.JINA_RERANK)).toBe('JR')
})
it('formats OpenAI Search as a first-class api format', () => {
expect(formatApiFormat(API_FORMATS.OPENAI_SEARCH)).toBe('OpenAI Search')
expect(formatApiFormatShort(API_FORMATS.OPENAI_SEARCH)).toBe('OS')
expect(sortApiFormats([
API_FORMATS.OPENAI_EMBEDDING,
API_FORMATS.OPENAI_SEARCH,
API_FORMATS.OPENAI_RESPONSES,
])).toEqual([
API_FORMATS.OPENAI_RESPONSES,
API_FORMATS.OPENAI_SEARCH,
API_FORMATS.OPENAI_EMBEDDING,
])
})
it('applies Responses to Search permissions in one direction', () => {
expect(apiFormatPermissionCovers('OPENAI_RESPONSES', 'openai:search')).toBe(true)
expect(apiFormatPermissionCovers('openai:search', 'openai:responses')).toBe(false)
})
it('formats embedding api format ids distinctly from chat formats', () => {
expect(formatApiFormat(API_FORMATS.GEMINI_INTERACTIONS)).toBe('Gemini Interactions')
expect(formatApiFormatShort(API_FORMATS.GEMINI_INTERACTIONS)).toBe('GI')

Some files were not shown because too many files have changed in this diff Show More