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(