mirror of
https://github.com/fawney19/Aether.git
synced 2026-09-14 23:20:19 +08:00
feat(codex): align Search and execution protocol
This commit is contained in:
@@ -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;
|
||||
|
||||
+6
-2
@@ -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,
|
||||
|
||||
@@ -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());
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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")]);
|
||||
|
||||
@@ -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"
|
||||
|
||||
+39
-2
@@ -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());
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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");
|
||||
|
||||
@@ -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);
|
||||
|
||||
@@ -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);
|
||||
|
||||
@@ -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>>,
|
||||
|
||||
@@ -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();
|
||||
|
||||
@@ -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"));
|
||||
|
||||
@@ -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(),
|
||||
|
||||
@@ -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(
|
||||
|
||||
Reference in New Issue
Block a user