mirror of
https://github.com/fawney19/Aether.git
synced 2026-09-10 05:00: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(
|
||||
|
||||
@@ -130,10 +130,6 @@ fn try_send_window_update(frame_tx: &FrameSender, stream_id: u32, bytes: usize)
|
||||
}
|
||||
}
|
||||
|
||||
/// Minimum allowed upstream request timeout (milliseconds).
|
||||
const MIN_TIMEOUT_MS: u64 = 1;
|
||||
/// Maximum allowed upstream request timeout (milliseconds).
|
||||
const MAX_TIMEOUT_MS: u64 = 300_000;
|
||||
/// Match reqwest's default redirect budget so direct execution and tunnel relay
|
||||
/// fail at the same point instead of diverging after a different number of hops.
|
||||
const MAX_REDIRECTS: usize = 10;
|
||||
@@ -707,43 +703,14 @@ fn remaining_timeout(deadline: Instant) -> Option<Duration> {
|
||||
}
|
||||
|
||||
fn resolve_request_timeouts(meta: &RequestMeta) -> RequestTimeouts {
|
||||
let first_byte_timeout = if meta.stream {
|
||||
meta.stream_first_byte_timeout_ms
|
||||
.map(timeout_duration_from_ms)
|
||||
.unwrap_or_else(|| timeout_duration_from_legacy_secs(meta.timeout))
|
||||
} else {
|
||||
meta.request_timeout_ms
|
||||
.or(meta.stream_first_byte_timeout_ms)
|
||||
.map(timeout_duration_from_ms)
|
||||
.unwrap_or_else(|| timeout_duration_from_legacy_secs(meta.timeout))
|
||||
};
|
||||
|
||||
let response_body_timeout = if meta.stream {
|
||||
None
|
||||
} else {
|
||||
Some(
|
||||
meta.request_timeout_ms
|
||||
.or(meta.stream_first_byte_timeout_ms)
|
||||
.map(timeout_duration_from_ms)
|
||||
.unwrap_or_else(|| timeout_duration_from_legacy_secs(meta.timeout)),
|
||||
)
|
||||
};
|
||||
let resolved = aether_contracts::tunnel::resolve_tunnel_request_timeouts(meta);
|
||||
|
||||
RequestTimeouts {
|
||||
first_byte_timeout,
|
||||
response_body_timeout,
|
||||
first_byte_timeout: Duration::from_millis(resolved.first_byte_ms),
|
||||
response_body_timeout: resolved.response_body_ms.map(Duration::from_millis),
|
||||
}
|
||||
}
|
||||
|
||||
fn timeout_duration_from_ms(ms: u64) -> Duration {
|
||||
Duration::from_millis(ms.clamp(MIN_TIMEOUT_MS, MAX_TIMEOUT_MS))
|
||||
}
|
||||
|
||||
fn timeout_duration_from_legacy_secs(secs: u64) -> Duration {
|
||||
let ms = secs.saturating_mul(1_000);
|
||||
timeout_duration_from_ms(ms)
|
||||
}
|
||||
|
||||
async fn spool_request_body(
|
||||
stream_id: u32,
|
||||
mut body_rx: mpsc::Receiver<TunnelFrame>,
|
||||
@@ -2176,6 +2143,18 @@ mod tests {
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn non_stream_request_timeouts_keep_the_protocol_maximum() {
|
||||
let mut meta = sample_request_meta();
|
||||
meta.request_timeout_ms = Some(aether_contracts::MAX_EXECUTION_REQUEST_TIMEOUT_MS);
|
||||
|
||||
let timeouts = resolve_request_timeouts(&meta);
|
||||
|
||||
let expected = Duration::from_millis(aether_contracts::MAX_EXECUTION_REQUEST_TIMEOUT_MS);
|
||||
assert_eq!(timeouts.first_byte_timeout, expected);
|
||||
assert_eq!(timeouts.response_body_timeout, Some(expected));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn resolve_redirect_changes_post_to_get_for_302() {
|
||||
let current_url = url::Url::parse("https://redirect.test/start").expect("url");
|
||||
|
||||
@@ -951,14 +951,7 @@ fn admin_usage_api_format_defaults_to_non_stream(item: &StoredRequestUsageAudit)
|
||||
let Some(value) = api_format else {
|
||||
return false;
|
||||
};
|
||||
matches!(
|
||||
aether_ai_formats::normalize_api_format_alias(value).as_str(),
|
||||
"openai:chat"
|
||||
| "openai:responses"
|
||||
| "openai:responses:compact"
|
||||
| "openai:image"
|
||||
| "claude:messages"
|
||||
)
|
||||
aether_ai_formats::api_format_defaults_to_non_stream(value)
|
||||
}
|
||||
|
||||
fn admin_usage_request_body_implies_default_non_stream(item: &StoredRequestUsageAudit) -> bool {
|
||||
@@ -2809,6 +2802,33 @@ mod tests {
|
||||
assert_eq!(record["client_is_stream"], false);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn client_requested_stream_defaults_to_non_stream_for_openai_search() {
|
||||
let item = StoredRequestUsageAudit {
|
||||
is_stream: true,
|
||||
api_format: Some("openai:search".to_string()),
|
||||
request_body: Some(json!({
|
||||
"id": "session-search-1",
|
||||
"model": "gpt-5.6-sol",
|
||||
"input": "current documentation"
|
||||
})),
|
||||
..sample_usage("completed", Some(200), None)
|
||||
};
|
||||
|
||||
assert!(!admin_usage_client_is_stream(&item));
|
||||
|
||||
let record = admin_usage_record_json(
|
||||
&item,
|
||||
&BTreeMap::new(),
|
||||
&BTreeMap::new(),
|
||||
false,
|
||||
false,
|
||||
None,
|
||||
);
|
||||
assert_eq!(record["client_requested_stream"], false);
|
||||
assert_eq!(record["client_is_stream"], false);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn upstream_stream_prefers_request_metadata_flag() {
|
||||
let item = StoredRequestUsageAudit {
|
||||
|
||||
@@ -82,33 +82,89 @@ pub fn endpoint_key_counts_by_format(
|
||||
) -> (BTreeMap<String, usize>, BTreeMap<String, usize>) {
|
||||
let mut total = BTreeMap::new();
|
||||
let mut active = BTreeMap::new();
|
||||
let inherited_api_formats = active_endpoint_api_formats(endpoints);
|
||||
let endpoint_api_formats = active_endpoint_api_formats(endpoints);
|
||||
|
||||
for key in keys {
|
||||
if fixed_provider_key_inherits_api_formats(
|
||||
let inherits_api_formats = fixed_provider_key_inherits_api_formats(
|
||||
provider_type,
|
||||
&key.auth_type,
|
||||
key.encrypted_auth_config.as_deref(),
|
||||
) {
|
||||
for api_format in &inherited_api_formats {
|
||||
*total.entry(api_format.clone()).or_insert(0) += 1;
|
||||
if key.is_active {
|
||||
*active.entry(api_format.clone()).or_insert(0) += 1;
|
||||
}
|
||||
}
|
||||
continue;
|
||||
}
|
||||
);
|
||||
let has_unrestricted_api_format_scope = key
|
||||
.api_formats
|
||||
.as_ref()
|
||||
.is_none_or(serde_json::Value::is_null);
|
||||
let configured_api_formats = configured_key_api_formats(key);
|
||||
|
||||
for api_format in configured_key_api_formats(key) {
|
||||
for api_format in endpoint_api_formats.iter().filter(|api_format| {
|
||||
inherits_api_formats
|
||||
|| has_unrestricted_api_format_scope
|
||||
|| configured_api_formats.iter().any(|allowed| {
|
||||
aether_ai_formats::api_format_permission_covers(allowed, api_format)
|
||||
})
|
||||
}) {
|
||||
*total.entry(api_format.clone()).or_insert(0) += 1;
|
||||
if key.is_active {
|
||||
*active.entry(api_format).or_insert(0) += 1;
|
||||
*active.entry(api_format.clone()).or_insert(0) += 1;
|
||||
}
|
||||
}
|
||||
}
|
||||
(total, active)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod endpoint_key_count_tests {
|
||||
use super::*;
|
||||
|
||||
fn sample_endpoint(id: &str, api_format: &str) -> StoredProviderCatalogEndpoint {
|
||||
StoredProviderCatalogEndpoint::new(
|
||||
id.to_string(),
|
||||
"provider-1".to_string(),
|
||||
api_format.to_string(),
|
||||
None,
|
||||
None,
|
||||
true,
|
||||
)
|
||||
.expect("endpoint should build")
|
||||
}
|
||||
|
||||
fn sample_key(id: &str, api_format: Option<&str>) -> StoredProviderCatalogKey {
|
||||
let mut key = StoredProviderCatalogKey::new(
|
||||
id.to_string(),
|
||||
"provider-1".to_string(),
|
||||
id.to_string(),
|
||||
"api_key".to_string(),
|
||||
None,
|
||||
true,
|
||||
)
|
||||
.expect("key should build");
|
||||
key.api_formats = api_format.map(|api_format| json!([api_format]));
|
||||
key
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn endpoint_counts_follow_one_way_format_permissions() {
|
||||
let endpoints = vec![
|
||||
sample_endpoint("responses", "openai:responses"),
|
||||
sample_endpoint("search", "openai:search"),
|
||||
];
|
||||
let mut empty_scope_key = sample_key("empty-scope-key", None);
|
||||
empty_scope_key.api_formats = Some(json!([]));
|
||||
let keys = vec![
|
||||
sample_key("responses-key", Some("openai:responses")),
|
||||
sample_key("search-key", Some("openai:search")),
|
||||
sample_key("unrestricted-key", None),
|
||||
empty_scope_key,
|
||||
];
|
||||
|
||||
let (total, active) = endpoint_key_counts_by_format("custom", &endpoints, &keys);
|
||||
|
||||
assert_eq!(total.get("openai:responses"), Some(&2));
|
||||
assert_eq!(total.get("openai:search"), Some(&3));
|
||||
assert_eq!(active, total);
|
||||
}
|
||||
}
|
||||
|
||||
fn masked_proxy_value(proxy: Option<&serde_json::Value>) -> serde_json::Value {
|
||||
let Some(proxy) = proxy.and_then(serde_json::Value::as_object) else {
|
||||
return serde_json::Value::Null;
|
||||
|
||||
@@ -769,6 +769,12 @@ const ADMIN_API_FORMAT_DEFINITIONS: &[AdminApiFormatDefinition] = &[
|
||||
default_path: "/v1/responses/compact",
|
||||
aliases: &["responses_compact"],
|
||||
},
|
||||
AdminApiFormatDefinition {
|
||||
value: "openai:search",
|
||||
label: "OpenAI Search",
|
||||
default_path: "/v1/alpha/search",
|
||||
aliases: &["openai_search", "search"],
|
||||
},
|
||||
AdminApiFormatDefinition {
|
||||
value: "openai:embedding",
|
||||
label: "OpenAI Embedding",
|
||||
|
||||
@@ -39,7 +39,8 @@ pub use crate::contracts::{
|
||||
OPENAI_RESPONSES_COMPACT_SYNC_SUCCESS_REPORT_KIND, OPENAI_RESPONSES_STREAM_PLAN_KIND,
|
||||
OPENAI_RESPONSES_STREAM_SUCCESS_REPORT_KIND, OPENAI_RESPONSES_SYNC_ERROR_REPORT_KIND,
|
||||
OPENAI_RESPONSES_SYNC_FINALIZE_REPORT_KIND, OPENAI_RESPONSES_SYNC_PLAN_KIND,
|
||||
OPENAI_RESPONSES_SYNC_SUCCESS_REPORT_KIND, OPENAI_VIDEO_CANCEL_SYNC_PLAN_KIND,
|
||||
OPENAI_RESPONSES_SYNC_SUCCESS_REPORT_KIND, OPENAI_SEARCH_SYNC_PLAN_KIND,
|
||||
OPENAI_SEARCH_SYNC_SUCCESS_REPORT_KIND, OPENAI_VIDEO_CANCEL_SYNC_PLAN_KIND,
|
||||
OPENAI_VIDEO_CONTENT_PLAN_KIND, OPENAI_VIDEO_CREATE_SYNC_FINALIZE_REPORT_KIND,
|
||||
OPENAI_VIDEO_CREATE_SYNC_PLAN_KIND, OPENAI_VIDEO_DELETE_SYNC_PLAN_KIND,
|
||||
OPENAI_VIDEO_REMIX_SYNC_PLAN_KIND,
|
||||
@@ -274,8 +275,10 @@ pub use aether_ai_formats::formats::conversion::response::{
|
||||
convert_openai_responses_response_to_openai_chat, OpenAiResponsesResponseUsage,
|
||||
};
|
||||
pub use aether_ai_formats::{
|
||||
api_format_alias_matches, api_format_storage_aliases, is_openai_responses_compact_format,
|
||||
is_openai_responses_family_format, is_openai_responses_format, normalize_api_format_alias,
|
||||
api_format_alias_matches, api_format_permission_covers, api_format_permission_storage_aliases,
|
||||
api_format_storage_aliases, intersect_api_format_allowed_lists,
|
||||
is_openai_responses_compact_format, is_openai_responses_family_format,
|
||||
is_openai_responses_format, normalize_api_format_alias,
|
||||
};
|
||||
pub use aether_ai_formats::{
|
||||
canonical_request_unknown_block_count, canonical_response_unknown_block_count,
|
||||
|
||||
@@ -23,9 +23,9 @@ pub use plan_kinds::{
|
||||
OPENAI_IMAGE_STREAM_PLAN_KIND, OPENAI_IMAGE_SYNC_PLAN_KIND, OPENAI_RERANK_SYNC_PLAN_KIND,
|
||||
OPENAI_RESPONSES_COMPACT_STREAM_PLAN_KIND, OPENAI_RESPONSES_COMPACT_SYNC_PLAN_KIND,
|
||||
OPENAI_RESPONSES_STREAM_PLAN_KIND, OPENAI_RESPONSES_SYNC_PLAN_KIND,
|
||||
OPENAI_VIDEO_CANCEL_SYNC_PLAN_KIND, OPENAI_VIDEO_CONTENT_PLAN_KIND,
|
||||
OPENAI_VIDEO_CREATE_SYNC_PLAN_KIND, OPENAI_VIDEO_DELETE_SYNC_PLAN_KIND,
|
||||
OPENAI_VIDEO_REMIX_SYNC_PLAN_KIND,
|
||||
OPENAI_SEARCH_SYNC_PLAN_KIND, OPENAI_VIDEO_CANCEL_SYNC_PLAN_KIND,
|
||||
OPENAI_VIDEO_CONTENT_PLAN_KIND, OPENAI_VIDEO_CREATE_SYNC_PLAN_KIND,
|
||||
OPENAI_VIDEO_DELETE_SYNC_PLAN_KIND, OPENAI_VIDEO_REMIX_SYNC_PLAN_KIND,
|
||||
};
|
||||
pub use report_kinds::{
|
||||
core_error_background_report_kind, core_error_default_client_api_format,
|
||||
@@ -52,5 +52,6 @@ pub use report_kinds::{
|
||||
OPENAI_RESPONSES_COMPACT_SYNC_FINALIZE_REPORT_KIND,
|
||||
OPENAI_RESPONSES_COMPACT_SYNC_SUCCESS_REPORT_KIND, OPENAI_RESPONSES_STREAM_SUCCESS_REPORT_KIND,
|
||||
OPENAI_RESPONSES_SYNC_ERROR_REPORT_KIND, OPENAI_RESPONSES_SYNC_FINALIZE_REPORT_KIND,
|
||||
OPENAI_RESPONSES_SYNC_SUCCESS_REPORT_KIND, OPENAI_VIDEO_CREATE_SYNC_FINALIZE_REPORT_KIND,
|
||||
OPENAI_RESPONSES_SYNC_SUCCESS_REPORT_KIND, OPENAI_SEARCH_SYNC_SUCCESS_REPORT_KIND,
|
||||
OPENAI_VIDEO_CREATE_SYNC_FINALIZE_REPORT_KIND,
|
||||
};
|
||||
|
||||
@@ -23,6 +23,7 @@ pub const OPENAI_VIDEO_CREATE_SYNC_PLAN_KIND: &str = "openai_video_create_sync";
|
||||
pub const OPENAI_CHAT_SYNC_PLAN_KIND: &str = "openai_chat_sync";
|
||||
pub const OPENAI_EMBEDDING_SYNC_PLAN_KIND: &str = "openai_embedding_sync";
|
||||
pub const OPENAI_RERANK_SYNC_PLAN_KIND: &str = "openai_rerank_sync";
|
||||
pub const OPENAI_SEARCH_SYNC_PLAN_KIND: &str = "openai_search_sync";
|
||||
pub const GEMINI_EMBEDDING_SYNC_PLAN_KIND: &str = "gemini_embedding_sync";
|
||||
pub const OPENAI_RESPONSES_SYNC_PLAN_KIND: &str = "openai_responses_sync";
|
||||
pub const OPENAI_RESPONSES_COMPACT_SYNC_PLAN_KIND: &str = "openai_responses_compact_sync";
|
||||
|
||||
@@ -29,6 +29,7 @@ pub const OPENAI_RESPONSES_SYNC_SUCCESS_REPORT_KIND: &str = "openai_responses_sy
|
||||
pub const OPENAI_RESPONSES_COMPACT_SYNC_SUCCESS_REPORT_KIND: &str =
|
||||
"openai_responses_compact_sync_success";
|
||||
pub const OPENAI_EMBEDDING_SYNC_SUCCESS_REPORT_KIND: &str = "openai_embedding_sync_success";
|
||||
pub const OPENAI_SEARCH_SYNC_SUCCESS_REPORT_KIND: &str = "openai_search_sync_success";
|
||||
pub const GEMINI_EMBEDDING_SYNC_SUCCESS_REPORT_KIND: &str = "gemini_embedding_sync_success";
|
||||
pub const OPENAI_IMAGE_SYNC_SUCCESS_REPORT_KIND: &str = "openai_image_sync_success";
|
||||
pub const CLAUDE_CLI_SYNC_SUCCESS_REPORT_KIND: &str = "claude_cli_sync_success";
|
||||
|
||||
@@ -23,6 +23,7 @@ pub enum FormatId {
|
||||
OpenAiChat,
|
||||
OpenAiResponses,
|
||||
OpenAiResponsesCompact,
|
||||
OpenAiSearch,
|
||||
OpenAiEmbedding,
|
||||
OpenAiRerank,
|
||||
ClaudeMessages,
|
||||
@@ -49,6 +50,7 @@ impl FormatId {
|
||||
Self::OpenAiChat
|
||||
| Self::OpenAiResponses
|
||||
| Self::OpenAiResponsesCompact
|
||||
| Self::OpenAiSearch
|
||||
| Self::OpenAiEmbedding
|
||||
| Self::OpenAiRerank => FormatFamily::OpenAi,
|
||||
Self::ClaudeMessages => FormatFamily::Claude,
|
||||
@@ -73,6 +75,7 @@ impl FormatId {
|
||||
Self::OpenAiChat => "openai:chat",
|
||||
Self::OpenAiResponses => "openai:responses",
|
||||
Self::OpenAiResponsesCompact => "openai:responses:compact",
|
||||
Self::OpenAiSearch => "openai:search",
|
||||
Self::OpenAiEmbedding => "openai:embedding",
|
||||
Self::OpenAiRerank => "openai:rerank",
|
||||
Self::ClaudeMessages => "claude:messages",
|
||||
@@ -103,6 +106,9 @@ impl FromStr for FormatId {
|
||||
"openai:responses:compact" | "/v1/responses/compact" => {
|
||||
Ok(Self::OpenAiResponsesCompact)
|
||||
}
|
||||
"openai:search" | "openai_search" | "search" | "/v1/alpha/search" => {
|
||||
Ok(Self::OpenAiSearch)
|
||||
}
|
||||
"openai:embedding" | "/v1/embeddings" => Ok(Self::OpenAiEmbedding),
|
||||
"openai:rerank" | "/v1/rerank" => Ok(Self::OpenAiRerank),
|
||||
"claude:messages" | "/v1/messages" => Ok(Self::ClaudeMessages),
|
||||
@@ -139,6 +145,56 @@ pub fn api_format_alias_matches(left: &str, right: &str) -> bool {
|
||||
normalize_api_format_alias(left) == normalize_api_format_alias(right)
|
||||
}
|
||||
|
||||
pub fn api_format_defaults_to_non_stream(value: &str) -> bool {
|
||||
matches!(
|
||||
normalize_api_format_alias(value).as_str(),
|
||||
"openai:chat"
|
||||
| "openai:responses"
|
||||
| "openai:responses:compact"
|
||||
| "openai:search"
|
||||
| "openai:image"
|
||||
| "claude:messages"
|
||||
)
|
||||
}
|
||||
|
||||
pub fn api_format_defaults_to_client_error_failover(value: &str) -> bool {
|
||||
!matches!(
|
||||
FormatId::parse(value).map(FormatId::canonical),
|
||||
Some(FormatId::OpenAiSearch)
|
||||
)
|
||||
}
|
||||
|
||||
pub fn api_format_permission_covers(allowed_value: &str, requested_api_format: &str) -> bool {
|
||||
let allowed_value = normalize_api_format_alias(allowed_value);
|
||||
let requested_api_format = normalize_api_format_alias(requested_api_format);
|
||||
!allowed_value.is_empty()
|
||||
&& !requested_api_format.is_empty()
|
||||
&& (allowed_value == requested_api_format
|
||||
|| allowed_value == "openai:responses" && requested_api_format == "openai:search")
|
||||
}
|
||||
|
||||
pub fn intersect_api_format_allowed_lists(left: &[String], right: &[String]) -> Vec<String> {
|
||||
let mut effective = Vec::new();
|
||||
for left_value in left {
|
||||
for right_value in right {
|
||||
let intersection = if api_format_permission_covers(right_value, left_value) {
|
||||
Some(left_value)
|
||||
} else if api_format_permission_covers(left_value, right_value) {
|
||||
Some(right_value)
|
||||
} else {
|
||||
None
|
||||
};
|
||||
if let Some(value) = intersection {
|
||||
let normalized = normalize_api_format_alias(value);
|
||||
if !effective.iter().any(|item| item == &normalized) {
|
||||
effective.push(normalized);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
effective
|
||||
}
|
||||
|
||||
pub fn api_format_storage_aliases(value: &str) -> Vec<String> {
|
||||
match FormatId::parse(value).map(FormatId::canonical) {
|
||||
Some(FormatId::AliyunMultimodalEmbedding) => vec![
|
||||
@@ -149,6 +205,22 @@ pub fn api_format_storage_aliases(value: &str) -> Vec<String> {
|
||||
}
|
||||
}
|
||||
|
||||
pub fn api_format_permission_storage_aliases(value: &str) -> Vec<String> {
|
||||
let requested_api_format = normalize_api_format_alias(value);
|
||||
let mut aliases = api_format_storage_aliases(&requested_api_format);
|
||||
for allowed_api_format in [FormatId::OpenAiResponses.as_str()] {
|
||||
if !api_format_permission_covers(allowed_api_format, &requested_api_format) {
|
||||
continue;
|
||||
}
|
||||
for alias in api_format_storage_aliases(allowed_api_format) {
|
||||
if !aliases.iter().any(|existing| existing == &alias) {
|
||||
aliases.push(alias);
|
||||
}
|
||||
}
|
||||
}
|
||||
aliases
|
||||
}
|
||||
|
||||
pub fn is_openai_responses_format(value: &str) -> bool {
|
||||
normalize_api_format_alias(value) == "openai:responses"
|
||||
}
|
||||
@@ -179,7 +251,10 @@ pub fn api_format_uses_body_stream_field(value: &str) -> bool {
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::{
|
||||
api_format_alias_matches, api_format_storage_aliases, api_format_uses_body_stream_field,
|
||||
api_format_alias_matches, api_format_defaults_to_client_error_failover,
|
||||
api_format_defaults_to_non_stream, api_format_permission_covers,
|
||||
api_format_permission_storage_aliases, api_format_storage_aliases,
|
||||
api_format_uses_body_stream_field, intersect_api_format_allowed_lists,
|
||||
normalize_api_format_alias, FormatId,
|
||||
};
|
||||
|
||||
@@ -193,6 +268,102 @@ mod tests {
|
||||
assert_eq!(FormatId::parse("gemini:cli"), None);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn responses_permission_covers_only_its_search_companion() {
|
||||
assert!(api_format_permission_covers(
|
||||
"OPENAI:RESPONSES",
|
||||
"openai:search"
|
||||
));
|
||||
assert!(api_format_permission_covers(
|
||||
"openai:search",
|
||||
"openai:search"
|
||||
));
|
||||
assert!(!api_format_permission_covers(
|
||||
"openai:search",
|
||||
"openai:responses"
|
||||
));
|
||||
assert!(!api_format_permission_covers(
|
||||
"openai:responses",
|
||||
"openai:chat"
|
||||
));
|
||||
assert_eq!(
|
||||
api_format_permission_storage_aliases("openai:search"),
|
||||
vec!["openai:search".to_string(), "openai:responses".to_string()]
|
||||
);
|
||||
assert_eq!(
|
||||
api_format_permission_storage_aliases("openai:responses"),
|
||||
vec!["openai:responses".to_string()]
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn normalizes_openai_search_aliases() {
|
||||
for alias in [
|
||||
"openai:search",
|
||||
"OPENAI_SEARCH",
|
||||
"search",
|
||||
"/v1/alpha/search",
|
||||
] {
|
||||
assert_eq!(FormatId::parse(alias), Some(FormatId::OpenAiSearch));
|
||||
assert_eq!(normalize_api_format_alias(alias), "openai:search");
|
||||
}
|
||||
assert!(!api_format_uses_body_stream_field("openai:search"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn identifies_default_non_stream_formats_from_aliases() {
|
||||
for format in [
|
||||
"/v1/chat/completions",
|
||||
"/v1/responses",
|
||||
"/v1/responses/compact",
|
||||
"/v1/alpha/search",
|
||||
"openai:image",
|
||||
"/v1/messages",
|
||||
] {
|
||||
assert!(api_format_defaults_to_non_stream(format), "{format}");
|
||||
}
|
||||
assert!(!api_format_defaults_to_non_stream("gemini:interactions"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn search_defaults_to_passthrough_for_client_errors() {
|
||||
for format in ["openai:search", "OPENAI_SEARCH", "/v1/alpha/search"] {
|
||||
assert!(
|
||||
!api_format_defaults_to_client_error_failover(format),
|
||||
"{format}"
|
||||
);
|
||||
}
|
||||
assert!(api_format_defaults_to_client_error_failover(
|
||||
"openai:responses"
|
||||
));
|
||||
assert!(api_format_defaults_to_client_error_failover(
|
||||
"custom:unknown"
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn api_format_policy_intersection_keeps_the_narrowest_companion_scope() {
|
||||
assert_eq!(
|
||||
intersect_api_format_allowed_lists(
|
||||
&["openai:responses".to_string()],
|
||||
&["openai:search".to_string()],
|
||||
),
|
||||
vec!["openai:search".to_string()]
|
||||
);
|
||||
assert_eq!(
|
||||
intersect_api_format_allowed_lists(
|
||||
&["openai:search".to_string()],
|
||||
&["OPENAI:RESPONSES".to_string()],
|
||||
),
|
||||
vec!["openai:search".to_string()]
|
||||
);
|
||||
assert!(intersect_api_format_allowed_lists(
|
||||
&["openai:search".to_string()],
|
||||
&["openai:chat".to_string()],
|
||||
)
|
||||
.is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn parses_embedding_api_formats() {
|
||||
assert_eq!(
|
||||
|
||||
@@ -57,6 +57,9 @@ pub fn request_candidate_api_format_preference(
|
||||
if client_api_format == "openai:responses:compact" {
|
||||
return (provider_api_format == "openai:responses:compact").then_some((0, 0));
|
||||
}
|
||||
if client_api_format == "openai:search" {
|
||||
return (provider_api_format == "openai:search").then_some((0, 0));
|
||||
}
|
||||
if is_gemini_interactions_api_format(client_api_format.as_str()) {
|
||||
return (provider_api_format == "gemini:interactions").then_some((0, 0));
|
||||
}
|
||||
@@ -109,6 +112,9 @@ pub fn request_candidate_api_formats(
|
||||
if client_api_format == "openai:responses:compact" {
|
||||
return vec!["openai:responses:compact"];
|
||||
}
|
||||
if client_api_format == "openai:search" {
|
||||
return vec!["openai:search"];
|
||||
}
|
||||
if is_gemini_interactions_api_format(client_api_format.as_str()) {
|
||||
return GEMINI_INTERACTIONS_CANDIDATE_API_FORMATS.to_vec();
|
||||
}
|
||||
@@ -377,6 +383,10 @@ mod tests {
|
||||
request_conversion_kind("openai:compact", "openai:responses:compact"),
|
||||
None
|
||||
);
|
||||
assert_eq!(
|
||||
request_conversion_kind("openai:search", "openai:responses"),
|
||||
None
|
||||
);
|
||||
assert_eq!(
|
||||
request_conversion_kind("gemini:generate_content", "claude:messages"),
|
||||
Some(RequestConversionKind::ToClaudeStandard)
|
||||
@@ -409,6 +419,26 @@ mod tests {
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn search_candidate_registry_keeps_exact_protocol_identity() {
|
||||
assert_eq!(
|
||||
request_candidate_api_formats("openai:search", false),
|
||||
vec!["openai:search"]
|
||||
);
|
||||
assert_eq!(
|
||||
request_candidate_api_formats("/v1/alpha/search", true),
|
||||
vec!["openai:search"]
|
||||
);
|
||||
assert_eq!(
|
||||
request_candidate_api_format_preference("openai:search", "openai:search"),
|
||||
Some((0, 0))
|
||||
);
|
||||
assert_eq!(
|
||||
request_candidate_api_format_preference("openai:search", "openai:responses"),
|
||||
None
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn sync_response_conversion_registry_supports_bidirectional_standard_matrix() {
|
||||
assert_eq!(
|
||||
|
||||
@@ -13,7 +13,10 @@ pub mod shared;
|
||||
|
||||
pub use context::{FormatContext, FormatError};
|
||||
pub use id::{
|
||||
api_format_alias_matches, api_format_storage_aliases, is_openai_responses_compact_format,
|
||||
api_format_alias_matches, api_format_defaults_to_client_error_failover,
|
||||
api_format_defaults_to_non_stream, api_format_permission_covers,
|
||||
api_format_permission_storage_aliases, api_format_storage_aliases,
|
||||
intersect_api_format_allowed_lists, is_openai_responses_compact_format,
|
||||
is_openai_responses_family_format, is_openai_responses_format, normalize_api_format_alias,
|
||||
FormatFamily, FormatId, FormatProfile,
|
||||
};
|
||||
|
||||
@@ -6,5 +6,6 @@ pub mod reasoning;
|
||||
pub mod request_contract;
|
||||
pub mod rerank;
|
||||
pub mod responses;
|
||||
pub mod search;
|
||||
pub mod shared;
|
||||
pub mod video;
|
||||
|
||||
@@ -68,19 +68,21 @@ pub(crate) fn validate_openai_reasoning_request_with_model_profile(
|
||||
};
|
||||
let source_api_format = crate::normalize_api_format_alias(source_api_format);
|
||||
let reasoning = match source_api_format.as_str() {
|
||||
"openai:responses" | "openai:responses:compact" => match object.get("reasoning") {
|
||||
Some(Value::Object(reasoning)) => Some(reasoning),
|
||||
Some(Value::Null) => None,
|
||||
Some(value) => {
|
||||
return Err(OpenAiReasoningContractViolation {
|
||||
kind: OpenAiReasoningViolationKind::InvalidType,
|
||||
field: "reasoning".to_string(),
|
||||
value: Some(value.to_string()),
|
||||
reason: "reasoning must be an object".to_string(),
|
||||
});
|
||||
"openai:responses" | "openai:responses:compact" | "openai:search" => {
|
||||
match object.get("reasoning") {
|
||||
Some(Value::Object(reasoning)) => Some(reasoning),
|
||||
Some(Value::Null) => None,
|
||||
Some(value) => {
|
||||
return Err(OpenAiReasoningContractViolation {
|
||||
kind: OpenAiReasoningViolationKind::InvalidType,
|
||||
field: "reasoning".to_string(),
|
||||
value: Some(value.to_string()),
|
||||
reason: "reasoning must be an object".to_string(),
|
||||
});
|
||||
}
|
||||
None => None,
|
||||
}
|
||||
None => None,
|
||||
},
|
||||
}
|
||||
"openai:chat" => None,
|
||||
_ => return Ok(()),
|
||||
};
|
||||
@@ -94,7 +96,7 @@ pub(crate) fn validate_openai_reasoning_request_with_model_profile(
|
||||
|
||||
let effort = match source_api_format.as_str() {
|
||||
"openai:chat" => object.get("reasoning_effort"),
|
||||
"openai:responses" | "openai:responses:compact" => {
|
||||
"openai:responses" | "openai:responses:compact" | "openai:search" => {
|
||||
reasoning.and_then(|reasoning| reasoning.get("effort"))
|
||||
}
|
||||
_ => None,
|
||||
@@ -110,11 +112,13 @@ pub(crate) fn validate_openai_reasoning_request_with_model_profile(
|
||||
)?;
|
||||
}
|
||||
|
||||
if let Some(mode) = reasoning
|
||||
.and_then(|reasoning| reasoning.get("mode"))
|
||||
.filter(|value| !value.is_null())
|
||||
{
|
||||
validate_reasoning_mode(mode, provider_model, source_model, supports_reasoning_mode)?;
|
||||
if source_api_format != "openai:search" {
|
||||
if let Some(mode) = reasoning
|
||||
.and_then(|reasoning| reasoning.get("mode"))
|
||||
.filter(|value| !value.is_null())
|
||||
{
|
||||
validate_reasoning_mode(mode, provider_model, source_model, supports_reasoning_mode)?;
|
||||
}
|
||||
}
|
||||
if let Some(context) = reasoning
|
||||
.and_then(|reasoning| reasoning.get("context"))
|
||||
|
||||
@@ -35,20 +35,21 @@ pub fn finalize_openai_provider_request_with_codex_model_capabilities(
|
||||
finalization: OpenAiProviderRequestFinalization<'_>,
|
||||
model_capabilities: Option<&super::responses::codex::CodexResponsesModelCapabilities>,
|
||||
) -> Result<(), OpenAiProviderRequestContractViolation> {
|
||||
let is_codex_responses = finalization
|
||||
let is_codex_reasoning_endpoint = finalization
|
||||
.provider_type
|
||||
.trim()
|
||||
.eq_ignore_ascii_case("codex")
|
||||
&& crate::is_openai_responses_family_format(finalization.provider_api_format);
|
||||
let resolved_model_capabilities =
|
||||
(is_codex_responses && model_capabilities.is_none()).then(|| {
|
||||
&& (crate::is_openai_responses_family_format(finalization.provider_api_format)
|
||||
|| crate::api_format_alias_matches(finalization.provider_api_format, "openai:search"));
|
||||
let resolved_model_capabilities = (is_codex_reasoning_endpoint && model_capabilities.is_none())
|
||||
.then(|| {
|
||||
super::responses::codex::resolve_codex_responses_model_capabilities(
|
||||
finalization.provider_model,
|
||||
finalization.source_model,
|
||||
None,
|
||||
)
|
||||
});
|
||||
let model_capabilities = is_codex_responses
|
||||
let model_capabilities = is_codex_reasoning_endpoint
|
||||
.then(|| model_capabilities.or(resolved_model_capabilities.as_ref()))
|
||||
.flatten();
|
||||
match crate::normalize_api_format_alias(finalization.source_api_format).as_str() {
|
||||
@@ -75,6 +76,11 @@ pub fn finalize_openai_provider_request_with_codex_model_capabilities(
|
||||
)
|
||||
}
|
||||
}
|
||||
super::responses::codex::normalize_codex_openai_reasoning_wire_effort(
|
||||
body,
|
||||
finalization.provider_type,
|
||||
finalization.provider_api_format,
|
||||
);
|
||||
super::responses::codex::apply_openai_responses_compact_special_body_edits(
|
||||
body,
|
||||
finalization.provider_api_format,
|
||||
@@ -85,6 +91,7 @@ pub fn finalize_openai_provider_request_with_codex_model_capabilities(
|
||||
finalization.upstream_is_stream,
|
||||
finalization.require_body_stream_field,
|
||||
);
|
||||
super::search::apply_openai_search_request_projection(body, finalization.provider_api_format);
|
||||
let provider_model = body
|
||||
.get("model")
|
||||
.and_then(Value::as_str)
|
||||
@@ -415,6 +422,81 @@ mod tests {
|
||||
assert_eq!(ultra_only_body["reasoning"]["effort"], "max");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn codex_search_normalizes_reasoning_and_projects_the_typed_request() {
|
||||
let finalization = OpenAiProviderRequestFinalization {
|
||||
source_api_format: "openai:search",
|
||||
provider_api_format: "openai:search",
|
||||
provider_type: "codex",
|
||||
provider_model: "gpt-5.6-sol",
|
||||
source_model: "gpt-5.6-sol",
|
||||
body_rules: None,
|
||||
upstream_is_stream: false,
|
||||
require_body_stream_field: false,
|
||||
};
|
||||
let mut body = json!({
|
||||
"id": "session-1",
|
||||
"model": "gpt-5.6-sol",
|
||||
"reasoning": {
|
||||
"effort": "ultra",
|
||||
"summary": "auto",
|
||||
"context": "current_turn",
|
||||
"future_reasoning_field": true
|
||||
},
|
||||
"commands": {"search_query": [{"q": "Aether"}]},
|
||||
"store": false,
|
||||
"future_request_field": {"enabled": true},
|
||||
"stream": true
|
||||
});
|
||||
|
||||
finalize_openai_provider_request(&mut body, finalization)
|
||||
.expect("Codex Search request should finalize");
|
||||
|
||||
assert_eq!(body["reasoning"]["effort"], "max");
|
||||
assert_eq!(body["reasoning"]["summary"], "auto");
|
||||
assert_eq!(body["reasoning"]["context"], "current_turn");
|
||||
assert!(body["reasoning"].get("future_reasoning_field").is_none());
|
||||
assert_eq!(body["commands"]["search_query"][0]["q"], "Aether");
|
||||
assert!(body.get("store").is_none());
|
||||
assert!(body.get("future_request_field").is_none());
|
||||
assert!(body.get("stream").is_none());
|
||||
assert!(body.get("tool_choice").is_none());
|
||||
assert!(body.get("include").is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn codex_search_validates_reasoning_effort_against_model_card() {
|
||||
let finalization = OpenAiProviderRequestFinalization {
|
||||
source_api_format: "openai:search",
|
||||
provider_api_format: "openai:search",
|
||||
provider_type: "codex",
|
||||
provider_model: "gpt-5.6-sol",
|
||||
source_model: "gpt-5.6-sol",
|
||||
body_rules: None,
|
||||
upstream_is_stream: false,
|
||||
require_body_stream_field: false,
|
||||
};
|
||||
let mut supported = json!({
|
||||
"id": "session-1",
|
||||
"model": "gpt-5.6-sol",
|
||||
"reasoning": {"effort": "high"}
|
||||
});
|
||||
finalize_openai_provider_request(&mut supported, finalization)
|
||||
.expect("published Search effort should pass");
|
||||
|
||||
let mut unsupported = json!({
|
||||
"id": "session-1",
|
||||
"model": "gpt-5.6-sol",
|
||||
"reasoning": {"effort": "none"}
|
||||
});
|
||||
let error = finalize_openai_provider_request(&mut unsupported, finalization)
|
||||
.expect_err("unpublished Search effort should be rejected");
|
||||
assert!(matches!(
|
||||
error,
|
||||
super::OpenAiProviderRequestContractViolation::Reasoning(_)
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn dynamic_codex_card_preserves_custom_reasoning_effort_case() {
|
||||
let finalization = OpenAiProviderRequestFinalization {
|
||||
|
||||
@@ -65,6 +65,7 @@ fn is_openai_image_request(provider_api_format: &str) -> bool {
|
||||
enum CodexOpenAiEndpointKind {
|
||||
Responses,
|
||||
Compact,
|
||||
Search,
|
||||
Images,
|
||||
}
|
||||
|
||||
@@ -601,6 +602,8 @@ fn codex_openai_endpoint_kind(
|
||||
Some(CodexOpenAiEndpointKind::Compact)
|
||||
} else if aether_ai_formats::is_openai_responses_format(provider_api_format) {
|
||||
Some(CodexOpenAiEndpointKind::Responses)
|
||||
} else if aether_ai_formats::api_format_alias_matches(provider_api_format, "openai:search") {
|
||||
Some(CodexOpenAiEndpointKind::Search)
|
||||
} else if is_openai_image_request(provider_api_format) {
|
||||
Some(CodexOpenAiEndpointKind::Images)
|
||||
} else {
|
||||
@@ -786,6 +789,16 @@ fn remove_btree_header(headers: &mut BTreeMap<String, String>, header_name: &str
|
||||
headers.retain(|name, _| !name.trim().eq_ignore_ascii_case(header_name));
|
||||
}
|
||||
|
||||
fn header_value_contains_media_type(value: &str, media_type: &str) -> bool {
|
||||
value.split(',').any(|media_range| {
|
||||
media_range
|
||||
.split(';')
|
||||
.next()
|
||||
.map(str::trim)
|
||||
.is_some_and(|value| value.eq_ignore_ascii_case(media_type))
|
||||
})
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Default, PartialEq, Eq)]
|
||||
pub struct CodexAuthIdentity {
|
||||
pub account_id: Option<String>,
|
||||
@@ -1179,6 +1192,23 @@ fn normalize_codex_reasoning_effort(body_object: &mut serde_json::Map<String, Va
|
||||
}
|
||||
}
|
||||
|
||||
pub fn normalize_codex_openai_reasoning_wire_effort(
|
||||
provider_request_body: &mut Value,
|
||||
provider_type: &str,
|
||||
provider_api_format: &str,
|
||||
) {
|
||||
if !provider_type.trim().eq_ignore_ascii_case("codex")
|
||||
|| !(aether_ai_formats::is_openai_responses_family_format(provider_api_format)
|
||||
|| aether_ai_formats::api_format_alias_matches(provider_api_format, "openai:search"))
|
||||
{
|
||||
return;
|
||||
}
|
||||
let Some(body_object) = provider_request_body.as_object_mut() else {
|
||||
return;
|
||||
};
|
||||
normalize_codex_reasoning_effort(body_object);
|
||||
}
|
||||
|
||||
fn is_codex_responses_lite_additional_tools_item(value: &Value) -> bool {
|
||||
value
|
||||
.get("type")
|
||||
@@ -1719,6 +1749,17 @@ pub fn apply_codex_openai_special_headers(
|
||||
"originator",
|
||||
CODEX_CLIENT_ORIGINATOR,
|
||||
);
|
||||
if endpoint_kind == CodexOpenAiEndpointKind::Search {
|
||||
remove_btree_header(provider_request_headers, CODEX_RESPONSES_LITE_HEADER);
|
||||
remove_btree_header(provider_request_headers, "openai-beta");
|
||||
if provider_request_headers.iter().any(|(name, value)| {
|
||||
name.eq_ignore_ascii_case("accept")
|
||||
&& header_value_contains_media_type(value, "text/event-stream")
|
||||
}) {
|
||||
remove_btree_header(provider_request_headers, "accept");
|
||||
}
|
||||
return;
|
||||
}
|
||||
if endpoint_kind == CodexOpenAiEndpointKind::Images {
|
||||
return;
|
||||
}
|
||||
@@ -1753,8 +1794,9 @@ mod tests {
|
||||
apply_codex_openai_special_headers, apply_openai_responses_compact_special_body_edits,
|
||||
build_codex_model_catalog_metadata, bundled_codex_model_cards, effective_codex_model_cards,
|
||||
resolve_codex_responses_model_capabilities,
|
||||
validate_codex_openai_responses_compact_request_contract,
|
||||
CODEX_OPENAI_IMAGE_INTERNAL_MODEL, CODEX_OPENAI_RESPONSES_UNSUPPORTED_BODY_FIELDS,
|
||||
validate_codex_openai_responses_compact_request_contract, CODEX_CLIENT_ORIGINATOR,
|
||||
CODEX_CLIENT_USER_AGENT, CODEX_OPENAI_IMAGE_INTERNAL_MODEL,
|
||||
CODEX_OPENAI_RESPONSES_UNSUPPORTED_BODY_FIELDS, CODEX_RESPONSES_LITE_HEADER,
|
||||
};
|
||||
use serde_json::{json, Value};
|
||||
|
||||
@@ -2265,6 +2307,51 @@ mod tests {
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn codex_search_uses_identity_headers_without_responses_protocol_headers() {
|
||||
let mut headers = std::collections::BTreeMap::from([
|
||||
(
|
||||
"x-openai-internal-codex-responses-lite".to_string(),
|
||||
"true".to_string(),
|
||||
),
|
||||
("openai-beta".to_string(), "responses=v1".to_string()),
|
||||
(
|
||||
"accept".to_string(),
|
||||
"application/json, Text/Event-Stream; q=0.9".to_string(),
|
||||
),
|
||||
]);
|
||||
|
||||
apply_codex_openai_special_headers(
|
||||
&mut headers,
|
||||
&json!({"id": "session-1", "model": "gpt-5.6-luna"}),
|
||||
&http::HeaderMap::new(),
|
||||
"codex",
|
||||
"openai:search",
|
||||
Some("request-search"),
|
||||
Some(r#"{"account_id":"account-1","is_fedramp":true}"#),
|
||||
);
|
||||
|
||||
assert_eq!(
|
||||
headers.get("chatgpt-account-id").map(String::as_str),
|
||||
Some("account-1")
|
||||
);
|
||||
assert_eq!(
|
||||
headers.get("x-openai-fedramp").map(String::as_str),
|
||||
Some("true")
|
||||
);
|
||||
assert_eq!(
|
||||
headers.get("user-agent").map(String::as_str),
|
||||
Some(CODEX_CLIENT_USER_AGENT)
|
||||
);
|
||||
assert_eq!(
|
||||
headers.get("originator").map(String::as_str),
|
||||
Some(CODEX_CLIENT_ORIGINATOR)
|
||||
);
|
||||
assert!(!headers.contains_key(CODEX_RESPONSES_LITE_HEADER));
|
||||
assert!(!headers.contains_key("openai-beta"));
|
||||
assert!(!headers.contains_key("accept"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn standard_codex_models_do_not_send_the_responses_lite_header() {
|
||||
let mut headers = std::collections::BTreeMap::from([(
|
||||
|
||||
@@ -0,0 +1,219 @@
|
||||
use serde_json::Value;
|
||||
|
||||
const REQUEST_FIELDS: &[&str] = &[
|
||||
"id",
|
||||
"model",
|
||||
"reasoning",
|
||||
"input",
|
||||
"commands",
|
||||
"settings",
|
||||
"max_output_tokens",
|
||||
];
|
||||
const REASONING_FIELDS: &[&str] = &["effort", "summary", "context"];
|
||||
const COMMAND_FIELDS: &[&str] = &[
|
||||
"search_query",
|
||||
"image_query",
|
||||
"open",
|
||||
"click",
|
||||
"find",
|
||||
"screenshot",
|
||||
"finance",
|
||||
"weather",
|
||||
"sports",
|
||||
"time",
|
||||
"response_length",
|
||||
];
|
||||
const COMMAND_ITEM_FIELDS: &[(&str, &[&str])] = &[
|
||||
("search_query", &["q", "recency", "domains"]),
|
||||
("image_query", &["q", "recency", "domains"]),
|
||||
("open", &["ref_id", "lineno"]),
|
||||
("click", &["ref_id", "id"]),
|
||||
("find", &["ref_id", "pattern"]),
|
||||
("screenshot", &["ref_id", "pageno"]),
|
||||
("finance", &["ticker", "type", "market"]),
|
||||
("weather", &["location", "start", "duration"]),
|
||||
(
|
||||
"sports",
|
||||
&[
|
||||
"tool",
|
||||
"fn",
|
||||
"league",
|
||||
"team",
|
||||
"opponent",
|
||||
"date_from",
|
||||
"date_to",
|
||||
"num_games",
|
||||
"locale",
|
||||
],
|
||||
),
|
||||
("time", &["utc_offset"]),
|
||||
];
|
||||
const SETTINGS_FIELDS: &[&str] = &[
|
||||
"user_location",
|
||||
"search_context_size",
|
||||
"filters",
|
||||
"image_settings",
|
||||
"allowed_callers",
|
||||
"external_web_access",
|
||||
];
|
||||
const USER_LOCATION_FIELDS: &[&str] = &["type", "country", "region", "city", "timezone"];
|
||||
const FILTER_FIELDS: &[&str] = &["allowed_domains", "blocked_domains"];
|
||||
const IMAGE_SETTINGS_FIELDS: &[&str] = &["max_results", "caption"];
|
||||
|
||||
fn retain_object_fields(value: &mut Value, fields: &[&str]) {
|
||||
if let Some(object) = value.as_object_mut() {
|
||||
object.retain(|field, _| fields.contains(&field.as_str()));
|
||||
}
|
||||
}
|
||||
|
||||
fn retain_array_object_fields(
|
||||
object: &mut serde_json::Map<String, Value>,
|
||||
key: &str,
|
||||
fields: &[&str],
|
||||
) {
|
||||
if let Some(items) = object.get_mut(key).and_then(Value::as_array_mut) {
|
||||
for item in items {
|
||||
retain_object_fields(item, fields);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub fn apply_openai_search_request_projection(body: &mut Value, provider_api_format: &str) {
|
||||
if !crate::api_format_alias_matches(provider_api_format, "openai:search") {
|
||||
return;
|
||||
}
|
||||
let Some(body_object) = body.as_object_mut() else {
|
||||
return;
|
||||
};
|
||||
body_object.retain(|field, _| REQUEST_FIELDS.contains(&field.as_str()));
|
||||
|
||||
if let Some(reasoning) = body_object.get_mut("reasoning") {
|
||||
retain_object_fields(reasoning, REASONING_FIELDS);
|
||||
}
|
||||
if let Some(commands) = body_object
|
||||
.get_mut("commands")
|
||||
.and_then(Value::as_object_mut)
|
||||
{
|
||||
commands.retain(|field, _| COMMAND_FIELDS.contains(&field.as_str()));
|
||||
for (key, fields) in COMMAND_ITEM_FIELDS {
|
||||
retain_array_object_fields(commands, key, fields);
|
||||
}
|
||||
}
|
||||
if let Some(settings) = body_object
|
||||
.get_mut("settings")
|
||||
.and_then(Value::as_object_mut)
|
||||
{
|
||||
settings.retain(|field, _| SETTINGS_FIELDS.contains(&field.as_str()));
|
||||
if let Some(user_location) = settings.get_mut("user_location") {
|
||||
retain_object_fields(user_location, USER_LOCATION_FIELDS);
|
||||
}
|
||||
if let Some(filters) = settings.get_mut("filters") {
|
||||
retain_object_fields(filters, FILTER_FIELDS);
|
||||
}
|
||||
if let Some(image_settings) = settings.get_mut("image_settings") {
|
||||
retain_object_fields(image_settings, IMAGE_SETTINGS_FIELDS);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use serde_json::json;
|
||||
|
||||
use super::apply_openai_search_request_projection;
|
||||
|
||||
#[test]
|
||||
fn projects_the_typed_search_request_contract() {
|
||||
let mut body = json!({
|
||||
"id": "session-1",
|
||||
"model": "gpt-5.6-sol",
|
||||
"reasoning": {
|
||||
"effort": "max",
|
||||
"summary": "auto",
|
||||
"context": "current_turn",
|
||||
"mode": "pro"
|
||||
},
|
||||
"input": "find documentation",
|
||||
"commands": {
|
||||
"search_query": [{"q": "Aether", "recency": 7, "unknown": true}],
|
||||
"image_query": [{"q": "Aether UI", "domains": ["example.com"], "unknown": true}],
|
||||
"open": [{"ref_id": "turn0search0", "lineno": 12, "unknown": true}],
|
||||
"click": [{"ref_id": "turn0fetch0", "id": 3, "unknown": true}],
|
||||
"find": [{"ref_id": "turn0fetch0", "pattern": "Aether", "unknown": true}],
|
||||
"screenshot": [{"ref_id": "turn0fetch0", "pageno": 2, "unknown": true}],
|
||||
"finance": [{"ticker": "OPENAI", "type": "equity", "market": "USA", "unknown": true}],
|
||||
"weather": [{"location": "US, CA, San Francisco", "duration": 3, "unknown": true}],
|
||||
"sports": [{"tool": "sports", "fn": "schedule", "league": "nba", "team": "GSW", "unknown": true}],
|
||||
"time": [{"utc_offset": "+08:00", "unknown": true}],
|
||||
"response_length": "short",
|
||||
"unknown": true
|
||||
},
|
||||
"settings": {
|
||||
"user_location": {"type": "approximate", "country": "US", "unknown": true},
|
||||
"search_context_size": "high",
|
||||
"filters": {"allowed_domains": ["openai.com"], "unknown": true},
|
||||
"image_settings": {"max_results": 3, "unknown": true},
|
||||
"allowed_callers": ["direct"],
|
||||
"external_web_access": "live",
|
||||
"unknown": true
|
||||
},
|
||||
"max_output_tokens": 1024,
|
||||
"store": false,
|
||||
"stream": true,
|
||||
"service_tier": "priority",
|
||||
"unknown": true
|
||||
});
|
||||
|
||||
apply_openai_search_request_projection(&mut body, "/v1/alpha/search");
|
||||
|
||||
assert_eq!(
|
||||
body["reasoning"],
|
||||
json!({
|
||||
"effort": "max",
|
||||
"summary": "auto",
|
||||
"context": "current_turn"
|
||||
})
|
||||
);
|
||||
assert_eq!(
|
||||
body["commands"],
|
||||
json!({
|
||||
"search_query": [{"q": "Aether", "recency": 7}],
|
||||
"image_query": [{"q": "Aether UI", "domains": ["example.com"]}],
|
||||
"open": [{"ref_id": "turn0search0", "lineno": 12}],
|
||||
"click": [{"ref_id": "turn0fetch0", "id": 3}],
|
||||
"find": [{"ref_id": "turn0fetch0", "pattern": "Aether"}],
|
||||
"screenshot": [{"ref_id": "turn0fetch0", "pageno": 2}],
|
||||
"finance": [{"ticker": "OPENAI", "type": "equity", "market": "USA"}],
|
||||
"weather": [{"location": "US, CA, San Francisco", "duration": 3}],
|
||||
"sports": [{"tool": "sports", "fn": "schedule", "league": "nba", "team": "GSW"}],
|
||||
"time": [{"utc_offset": "+08:00"}],
|
||||
"response_length": "short"
|
||||
})
|
||||
);
|
||||
assert_eq!(
|
||||
body["settings"],
|
||||
json!({
|
||||
"user_location": {"type": "approximate", "country": "US"},
|
||||
"search_context_size": "high",
|
||||
"filters": {"allowed_domains": ["openai.com"]},
|
||||
"image_settings": {"max_results": 3},
|
||||
"allowed_callers": ["direct"],
|
||||
"external_web_access": "live"
|
||||
})
|
||||
);
|
||||
assert!(body.get("store").is_none());
|
||||
assert!(body.get("stream").is_none());
|
||||
assert!(body.get("service_tier").is_none());
|
||||
assert!(body.get("unknown").is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn leaves_other_formats_unchanged() {
|
||||
let mut body = json!({"model": "gpt-5.6-sol", "store": true});
|
||||
let expected = body.clone();
|
||||
|
||||
apply_openai_search_request_projection(&mut body, "openai:responses");
|
||||
|
||||
assert_eq!(body, expected);
|
||||
}
|
||||
}
|
||||
@@ -38,7 +38,7 @@ pub fn parse_request(
|
||||
FormatId::GeminiEmbedding => gemini::embedding::request::from(body, ctx),
|
||||
FormatId::DoubaoEmbedding => doubao::embedding::request::from(body, ctx),
|
||||
FormatId::AliyunMultimodalEmbedding => aliyun::embedding::request::from(body, ctx),
|
||||
FormatId::GeminiInteractions => None,
|
||||
FormatId::OpenAiSearch | FormatId::GeminiInteractions => None,
|
||||
}
|
||||
.ok_or_else(|| FormatError::RequestParseFailed {
|
||||
format: source.as_str().to_string(),
|
||||
@@ -72,7 +72,7 @@ fn emit_request_inner(
|
||||
FormatId::GeminiEmbedding => gemini::embedding::request::to(request, ctx),
|
||||
FormatId::DoubaoEmbedding => doubao::embedding::request::to(request, ctx),
|
||||
FormatId::AliyunMultimodalEmbedding => aliyun::embedding::request::to(request, ctx),
|
||||
FormatId::GeminiInteractions => None,
|
||||
FormatId::OpenAiSearch | FormatId::GeminiInteractions => None,
|
||||
}
|
||||
.ok_or_else(|| FormatError::RequestEmitFailed {
|
||||
format: target.as_str().to_string(),
|
||||
@@ -280,6 +280,7 @@ pub fn parse_response(
|
||||
FormatId::ClaudeMessages => claude_messages::response::from(body, ctx),
|
||||
FormatId::GeminiGenerateContent => gemini_generate_content::response::from(body, ctx),
|
||||
FormatId::OpenAiEmbedding
|
||||
| FormatId::OpenAiSearch
|
||||
| FormatId::JinaEmbedding
|
||||
| FormatId::OpenAiRerank
|
||||
| FormatId::JinaRerank
|
||||
@@ -318,6 +319,7 @@ fn emit_response_inner(
|
||||
FormatId::ClaudeMessages => claude_messages::response::to(response, ctx),
|
||||
FormatId::GeminiGenerateContent => gemini_generate_content::response::to(response, ctx),
|
||||
FormatId::OpenAiEmbedding
|
||||
| FormatId::OpenAiSearch
|
||||
| FormatId::JinaEmbedding
|
||||
| FormatId::OpenAiRerank
|
||||
| FormatId::JinaRerank
|
||||
@@ -1055,6 +1057,7 @@ fn standard_request_root_field_is_audited(source: FormatId, key: &str) -> bool {
|
||||
| "tools"
|
||||
),
|
||||
FormatId::OpenAiEmbedding
|
||||
| FormatId::OpenAiSearch
|
||||
| FormatId::OpenAiRerank
|
||||
| FormatId::GeminiEmbedding
|
||||
| FormatId::JinaEmbedding
|
||||
@@ -1535,6 +1538,7 @@ fn validate_source_response_stop_enums(
|
||||
FormatId::GeminiGenerateContent => validate_gemini_response_finish_reasons(body, target),
|
||||
FormatId::GeminiInteractions => Ok(()),
|
||||
FormatId::OpenAiEmbedding
|
||||
| FormatId::OpenAiSearch
|
||||
| FormatId::OpenAiRerank
|
||||
| FormatId::GeminiEmbedding
|
||||
| FormatId::JinaEmbedding
|
||||
|
||||
@@ -38,7 +38,11 @@ pub fn build_core_error_body_for_client_format(
|
||||
error_object.insert("message".to_string(), Value::String(message.to_string()));
|
||||
|
||||
match aether_ai_formats::normalize_api_format_alias(client_api_format).as_str() {
|
||||
"openai:chat" | "openai:responses" | "openai:responses:compact" | "openai:embedding" => {
|
||||
"openai:chat"
|
||||
| "openai:responses"
|
||||
| "openai:responses:compact"
|
||||
| "openai:search"
|
||||
| "openai:embedding" => {
|
||||
error_object.insert(
|
||||
"type".to_string(),
|
||||
Value::String(map_local_sync_error_kind_to_openai_type(kind).to_string()),
|
||||
@@ -157,6 +161,21 @@ mod tests {
|
||||
assert_eq!(body["error"]["code"], "invalid_request");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn builds_openai_search_core_error_body() {
|
||||
let body = build_core_error_body_for_client_format(
|
||||
"openai:search",
|
||||
"search unavailable",
|
||||
Some("upstream_unavailable"),
|
||||
LocalCoreSyncErrorKind::ServerError,
|
||||
)
|
||||
.expect("body should build");
|
||||
|
||||
assert_eq!(body["error"]["message"], "search unavailable");
|
||||
assert_eq!(body["error"]["type"], "server_error");
|
||||
assert_eq!(body["error"]["code"], "upstream_unavailable");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn recognizes_finalize_kind_and_success_mapping() {
|
||||
assert!(is_core_error_finalize_kind("openai_chat_sync_finalize"));
|
||||
|
||||
@@ -1,9 +1,10 @@
|
||||
use serde_json::{json, Value};
|
||||
|
||||
pub const MODEL_DIRECTIVE_API_FORMATS: [&str; 5] = [
|
||||
pub const MODEL_DIRECTIVE_API_FORMATS: [&str; 6] = [
|
||||
"openai:chat",
|
||||
"openai:responses",
|
||||
"openai:responses:compact",
|
||||
"openai:search",
|
||||
"claude:messages",
|
||||
"gemini:generate_content",
|
||||
];
|
||||
@@ -433,7 +434,7 @@ pub fn default_model_directive_mapping_patch(
|
||||
|
||||
pub fn default_model_directive_suffixes(provider_api_format: &str) -> &'static [&'static str] {
|
||||
match crate::normalize_api_format_alias(provider_api_format).as_str() {
|
||||
"openai:chat" | "openai:responses" | "openai:responses:compact" => {
|
||||
"openai:chat" | "openai:responses" | "openai:responses:compact" | "openai:search" => {
|
||||
&OPENAI_MODEL_DIRECTIVE_SUFFIXES
|
||||
}
|
||||
"claude:messages" | "gemini:generate_content" => &CROSS_PROVIDER_MODEL_DIRECTIVE_SUFFIXES,
|
||||
@@ -503,7 +504,7 @@ fn apply_reasoning_effort_override(
|
||||
"reasoning_effort",
|
||||
effort.as_openai_model_directive_value(),
|
||||
),
|
||||
"openai:responses" | "openai:responses:compact" => {
|
||||
"openai:responses" | "openai:responses:compact" | "openai:search" => {
|
||||
set_openai_responses_reasoning_effort(provider_request_body, effort)
|
||||
}
|
||||
"claude:messages" => {
|
||||
@@ -525,7 +526,7 @@ fn apply_codex_reasoning_preset_override(
|
||||
"openai:chat" => {
|
||||
set_object_string(provider_request_body, "reasoning_effort", preset.as_str())
|
||||
}
|
||||
"openai:responses" | "openai:responses:compact" => {
|
||||
"openai:responses" | "openai:responses:compact" | "openai:search" => {
|
||||
let object = provider_request_body.as_object_mut()?;
|
||||
let reasoning = object
|
||||
.entry("reasoning".to_string())
|
||||
@@ -550,6 +551,7 @@ fn apply_service_tier_override(
|
||||
"service_tier",
|
||||
tier.as_openai_value(),
|
||||
),
|
||||
"openai:search" => Some(()),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
@@ -710,7 +712,7 @@ pub fn reasoning_effort_supported_for_model(
|
||||
effort: ReasoningEffort,
|
||||
) -> bool {
|
||||
match crate::normalize_api_format_alias(provider_api_format).as_str() {
|
||||
"openai:chat" | "openai:responses" | "openai:responses:compact" => {
|
||||
"openai:chat" | "openai:responses" | "openai:responses:compact" | "openai:search" => {
|
||||
match resolved_openai_model_identity(provider_model, source_model).0 {
|
||||
OpenAiModelIdentity::Gpt56 => effort != ReasoningEffort::Minimal,
|
||||
OpenAiModelIdentity::ConcreteOther => effort != ReasoningEffort::Max,
|
||||
|
||||
@@ -5,7 +5,8 @@ use crate::contracts::{
|
||||
GEMINI_EMBEDDING_SYNC_SUCCESS_REPORT_KIND, GEMINI_INTERACTIONS_STREAM_PLAN_KIND,
|
||||
GEMINI_INTERACTIONS_STREAM_SUCCESS_REPORT_KIND, GEMINI_INTERACTIONS_SYNC_PLAN_KIND,
|
||||
GEMINI_INTERACTIONS_SYNC_SUCCESS_REPORT_KIND, OPENAI_EMBEDDING_SYNC_PLAN_KIND,
|
||||
OPENAI_RERANK_SYNC_PLAN_KIND,
|
||||
OPENAI_RERANK_SYNC_PLAN_KIND, OPENAI_SEARCH_SYNC_PLAN_KIND,
|
||||
OPENAI_SEARCH_SYNC_SUCCESS_REPORT_KIND,
|
||||
};
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
@@ -81,6 +82,13 @@ pub fn resolve_sync_spec(plan_kind: &str) -> Option<LocalSameFormatProviderSpec>
|
||||
family: LocalSameFormatProviderFamily::Standard,
|
||||
require_streaming: false,
|
||||
}),
|
||||
OPENAI_SEARCH_SYNC_PLAN_KIND => Some(LocalSameFormatProviderSpec {
|
||||
api_format: "openai:search",
|
||||
decision_kind: OPENAI_SEARCH_SYNC_PLAN_KIND,
|
||||
report_kind: OPENAI_SEARCH_SYNC_SUCCESS_REPORT_KIND,
|
||||
family: LocalSameFormatProviderFamily::Standard,
|
||||
require_streaming: false,
|
||||
}),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
@@ -185,4 +193,12 @@ mod tests {
|
||||
assert_eq!(spec.report_kind, "openai_rerank_sync_success");
|
||||
assert!(!spec.require_streaming);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn resolves_openai_search_sync_spec() {
|
||||
let spec = resolve_sync_spec("openai_search_sync").expect("spec");
|
||||
assert_eq!(spec.api_format, "openai:search");
|
||||
assert_eq!(spec.report_kind, "openai_search_sync_success");
|
||||
assert!(!spec.require_streaming);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -50,7 +50,8 @@ pub fn forbid_upstream_streaming_for_provider(
|
||||
provider_type: &str,
|
||||
provider_api_format: &str,
|
||||
) -> bool {
|
||||
aether_ai_formats::is_openai_responses_compact_format(provider_api_format)
|
||||
aether_ai_formats::api_format_alias_matches(provider_api_format, "openai:search")
|
||||
|| aether_ai_formats::is_openai_responses_compact_format(provider_api_format)
|
||||
|| (provider_type.trim().eq_ignore_ascii_case("codex")
|
||||
&& provider_api_format
|
||||
.trim()
|
||||
@@ -258,7 +259,15 @@ mod tests {
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn forbids_streaming_for_compact_and_codex_images() {
|
||||
fn forbids_streaming_for_sync_only_openai_formats() {
|
||||
assert!(forbid_upstream_streaming_for_provider(
|
||||
"codex",
|
||||
"openai:search"
|
||||
));
|
||||
assert!(forbid_upstream_streaming_for_provider(
|
||||
"custom",
|
||||
"/v1/alpha/search"
|
||||
));
|
||||
assert!(forbid_upstream_streaming_for_provider(
|
||||
"codex",
|
||||
"openai:responses:compact"
|
||||
@@ -461,6 +470,13 @@ mod tests {
|
||||
#[test]
|
||||
fn provider_policy_gives_non_stream_contracts_precedence() {
|
||||
let force_stream = json!({"upstream_stream_policy": "force_stream"});
|
||||
assert!(!resolve_upstream_is_stream_for_provider(
|
||||
Some(&force_stream),
|
||||
"codex",
|
||||
"openai:search",
|
||||
true,
|
||||
true,
|
||||
));
|
||||
assert!(!resolve_upstream_is_stream_for_provider(
|
||||
Some(&force_stream),
|
||||
"codex",
|
||||
|
||||
@@ -13,9 +13,9 @@ use crate::contracts::{
|
||||
OPENAI_IMAGE_STREAM_PLAN_KIND, OPENAI_IMAGE_SYNC_PLAN_KIND, OPENAI_RERANK_SYNC_PLAN_KIND,
|
||||
OPENAI_RESPONSES_COMPACT_STREAM_PLAN_KIND, OPENAI_RESPONSES_COMPACT_SYNC_PLAN_KIND,
|
||||
OPENAI_RESPONSES_STREAM_PLAN_KIND, OPENAI_RESPONSES_SYNC_PLAN_KIND,
|
||||
OPENAI_VIDEO_CANCEL_SYNC_PLAN_KIND, OPENAI_VIDEO_CONTENT_PLAN_KIND,
|
||||
OPENAI_VIDEO_CREATE_SYNC_PLAN_KIND, OPENAI_VIDEO_DELETE_SYNC_PLAN_KIND,
|
||||
OPENAI_VIDEO_REMIX_SYNC_PLAN_KIND,
|
||||
OPENAI_SEARCH_SYNC_PLAN_KIND, OPENAI_VIDEO_CANCEL_SYNC_PLAN_KIND,
|
||||
OPENAI_VIDEO_CONTENT_PLAN_KIND, OPENAI_VIDEO_CREATE_SYNC_PLAN_KIND,
|
||||
OPENAI_VIDEO_DELETE_SYNC_PLAN_KIND, OPENAI_VIDEO_REMIX_SYNC_PLAN_KIND,
|
||||
};
|
||||
use crate::formats::openai::image::request::is_openai_image_stream_request;
|
||||
|
||||
@@ -216,6 +216,14 @@ pub fn resolve_execution_runtime_sync_plan_kind(
|
||||
return Some(OPENAI_RERANK_SYNC_PLAN_KIND);
|
||||
}
|
||||
|
||||
if route_family == Some("openai")
|
||||
&& route_kind == Some("search")
|
||||
&& *method == Method::POST
|
||||
&& path == "/v1/alpha/search"
|
||||
{
|
||||
return Some(OPENAI_SEARCH_SYNC_PLAN_KIND);
|
||||
}
|
||||
|
||||
if route_family == Some("openai")
|
||||
&& route_kind == Some("image")
|
||||
&& *method == Method::POST
|
||||
@@ -433,6 +441,7 @@ pub fn supports_sync_execution_decision_kind(plan_kind: &str) -> bool {
|
||||
OPENAI_CHAT_SYNC_PLAN_KIND
|
||||
| OPENAI_EMBEDDING_SYNC_PLAN_KIND
|
||||
| OPENAI_RERANK_SYNC_PLAN_KIND
|
||||
| OPENAI_SEARCH_SYNC_PLAN_KIND
|
||||
| OPENAI_IMAGE_SYNC_PLAN_KIND
|
||||
| OPENAI_RESPONSES_SYNC_PLAN_KIND
|
||||
| OPENAI_RESPONSES_COMPACT_SYNC_PLAN_KIND
|
||||
@@ -492,6 +501,7 @@ mod tests {
|
||||
OPENAI_IMAGE_STREAM_PLAN_KIND, OPENAI_IMAGE_SYNC_PLAN_KIND, OPENAI_RERANK_SYNC_PLAN_KIND,
|
||||
OPENAI_RESPONSES_COMPACT_STREAM_PLAN_KIND, OPENAI_RESPONSES_COMPACT_SYNC_PLAN_KIND,
|
||||
OPENAI_RESPONSES_STREAM_PLAN_KIND, OPENAI_RESPONSES_SYNC_PLAN_KIND,
|
||||
OPENAI_SEARCH_SYNC_PLAN_KIND,
|
||||
};
|
||||
|
||||
#[test]
|
||||
@@ -606,6 +616,46 @@ mod tests {
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn resolves_openai_search_as_sync_only() {
|
||||
assert_eq!(
|
||||
resolve_execution_runtime_sync_plan_kind(
|
||||
Some("ai_public"),
|
||||
Some("openai"),
|
||||
Some("search"),
|
||||
None,
|
||||
&Method::POST,
|
||||
"/v1/alpha/search",
|
||||
),
|
||||
Some(OPENAI_SEARCH_SYNC_PLAN_KIND)
|
||||
);
|
||||
assert_eq!(
|
||||
resolve_execution_runtime_stream_plan_kind(
|
||||
Some("ai_public"),
|
||||
Some("openai"),
|
||||
Some("search"),
|
||||
None,
|
||||
&Method::POST,
|
||||
"/v1/alpha/search",
|
||||
),
|
||||
None
|
||||
);
|
||||
assert!(supports_sync_execution_decision_kind(
|
||||
OPENAI_SEARCH_SYNC_PLAN_KIND
|
||||
));
|
||||
assert_eq!(
|
||||
resolve_execution_runtime_sync_plan_kind(
|
||||
Some("ai_public"),
|
||||
Some("openai"),
|
||||
Some("search"),
|
||||
None,
|
||||
&Method::POST,
|
||||
"/backend-api/codex/alpha/search",
|
||||
),
|
||||
None
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn resolves_claude_messages_plan_kinds_by_request_auth_channel() {
|
||||
assert_eq!(
|
||||
|
||||
@@ -309,6 +309,7 @@ impl ProviderStreamParser {
|
||||
FormatId::ClaudeMessages => Self::Claude(ClaudeProviderState::default()),
|
||||
FormatId::GeminiGenerateContent => Self::Gemini(GeminiProviderState::default()),
|
||||
FormatId::OpenAiEmbedding
|
||||
| FormatId::OpenAiSearch
|
||||
| FormatId::OpenAiRerank
|
||||
| FormatId::GeminiEmbedding
|
||||
| FormatId::GeminiInteractions
|
||||
@@ -409,6 +410,7 @@ impl ClientStreamEmitter {
|
||||
FormatId::ClaudeMessages => Self::Claude(ClaudeClientEmitter::default()),
|
||||
FormatId::GeminiGenerateContent => Self::Gemini(GeminiClientEmitter::default()),
|
||||
FormatId::OpenAiEmbedding
|
||||
| FormatId::OpenAiSearch
|
||||
| FormatId::OpenAiRerank
|
||||
| FormatId::GeminiEmbedding
|
||||
| FormatId::GeminiInteractions
|
||||
@@ -531,6 +533,7 @@ fn parse_provider_error(
|
||||
parse_gemini_error(payload)
|
||||
}
|
||||
FormatId::OpenAiEmbedding
|
||||
| FormatId::OpenAiSearch
|
||||
| FormatId::OpenAiRerank
|
||||
| FormatId::GeminiEmbedding
|
||||
| FormatId::JinaEmbedding
|
||||
|
||||
@@ -11,7 +11,10 @@ pub use formats::context::{
|
||||
FormatError,
|
||||
};
|
||||
pub use formats::id::{
|
||||
api_format_alias_matches, api_format_storage_aliases, api_format_uses_body_stream_field,
|
||||
api_format_alias_matches, api_format_defaults_to_client_error_failover,
|
||||
api_format_defaults_to_non_stream, api_format_permission_covers,
|
||||
api_format_permission_storage_aliases, api_format_storage_aliases,
|
||||
api_format_uses_body_stream_field, intersect_api_format_allowed_lists,
|
||||
is_openai_responses_compact_format, is_openai_responses_family_format,
|
||||
is_openai_responses_format, normalize_api_format_alias, FormatFamily, FormatId, FormatProfile,
|
||||
};
|
||||
|
||||
@@ -11,7 +11,9 @@ pub use frame::{StreamFrame, StreamFramePayload, StreamFrameType};
|
||||
pub use plan::{
|
||||
ExecutionPlan, ExecutionTimeouts, ProxySnapshot, RequestBody, ResolvedTransportProfile,
|
||||
EXECUTION_REQUEST_ACCEPT_INVALID_CERTS_HEADER, EXECUTION_REQUEST_FOLLOW_REDIRECTS_HEADER,
|
||||
EXECUTION_REQUEST_HTTP1_ONLY_HEADER, TRANSPORT_BACKEND_BROWSER_WREQ,
|
||||
EXECUTION_REQUEST_HTTP1_ONLY_HEADER, MAX_EXECUTION_REQUEST_TIMEOUT_MS,
|
||||
MAX_EXECUTION_REQUEST_TIMEOUT_SECS, MAX_EXECUTION_STREAM_FIRST_BYTE_TIMEOUT_MS,
|
||||
MAX_EXECUTION_STREAM_FIRST_BYTE_TIMEOUT_SECS, TRANSPORT_BACKEND_BROWSER_WREQ,
|
||||
TRANSPORT_BACKEND_HYPER_RUSTLS, TRANSPORT_BACKEND_REQWEST_RUSTLS, TRANSPORT_HTTP_MODE_AUTO,
|
||||
TRANSPORT_HTTP_MODE_H2C_PRIOR_KNOWLEDGE, TRANSPORT_HTTP_MODE_HTTP1_ONLY,
|
||||
TRANSPORT_POOL_SCOPE_KEY,
|
||||
|
||||
@@ -7,6 +7,11 @@ pub const EXECUTION_REQUEST_FOLLOW_REDIRECTS_HEADER: &str = "x-aether-execution-
|
||||
pub const EXECUTION_REQUEST_HTTP1_ONLY_HEADER: &str = "x-aether-execution-http1-only";
|
||||
pub const EXECUTION_REQUEST_ACCEPT_INVALID_CERTS_HEADER: &str =
|
||||
"x-aether-execution-accept-invalid-certs";
|
||||
pub const MAX_EXECUTION_REQUEST_TIMEOUT_SECS: u64 = 1_200;
|
||||
pub const MAX_EXECUTION_REQUEST_TIMEOUT_MS: u64 = MAX_EXECUTION_REQUEST_TIMEOUT_SECS * 1_000;
|
||||
pub const MAX_EXECUTION_STREAM_FIRST_BYTE_TIMEOUT_SECS: u64 = 300;
|
||||
pub const MAX_EXECUTION_STREAM_FIRST_BYTE_TIMEOUT_MS: u64 =
|
||||
MAX_EXECUTION_STREAM_FIRST_BYTE_TIMEOUT_SECS * 1_000;
|
||||
|
||||
#[derive(Debug, Clone, Default, Serialize, Deserialize, PartialEq, Eq)]
|
||||
#[serde(default)]
|
||||
|
||||
@@ -12,6 +12,7 @@ pub const TUNNEL_PROTOCOL_VERSION_HEADER: &str = "x-aether-tunnel-protocol-versi
|
||||
pub const TUNNEL_NODE_NAME_B64_HEADER: &str = "x-aether-tunnel-node-name-b64";
|
||||
pub const CURRENT_TUNNEL_PROTOCOL_VERSION: u8 = 3;
|
||||
pub const CURRENT_TUNNEL_PROTOCOL_VERSION_STR: &str = "3";
|
||||
pub const MAX_TUNNEL_RELAY_META_LEN: usize = 256 * 1024;
|
||||
|
||||
pub mod flags {
|
||||
pub const END_STREAM: u8 = 0x01;
|
||||
@@ -220,6 +221,63 @@ pub struct RequestMeta {
|
||||
pub transport_profile: Option<crate::ResolvedTransportProfile>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub struct ResolvedTunnelRequestTimeouts {
|
||||
pub first_byte_ms: u64,
|
||||
pub response_body_ms: Option<u64>,
|
||||
}
|
||||
|
||||
pub fn resolve_tunnel_request_timeouts(meta: &RequestMeta) -> ResolvedTunnelRequestTimeouts {
|
||||
let legacy_timeout_ms = meta.timeout.saturating_mul(1_000);
|
||||
let first_byte_ms = if meta.stream {
|
||||
meta.stream_first_byte_timeout_ms
|
||||
.unwrap_or(legacy_timeout_ms)
|
||||
} else {
|
||||
meta.request_timeout_ms
|
||||
.or(meta.stream_first_byte_timeout_ms)
|
||||
.unwrap_or(legacy_timeout_ms)
|
||||
};
|
||||
let response_body_ms = (!meta.stream).then_some(first_byte_ms);
|
||||
|
||||
ResolvedTunnelRequestTimeouts {
|
||||
first_byte_ms: if meta.stream {
|
||||
clamp_stream_first_byte_timeout_ms(first_byte_ms)
|
||||
} else {
|
||||
clamp_upstream_request_timeout_ms(first_byte_ms)
|
||||
},
|
||||
response_body_ms: response_body_ms.map(clamp_upstream_request_timeout_ms),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn try_decode_tunnel_relay_request_meta(
|
||||
buffer: &[u8],
|
||||
) -> Result<Option<(RequestMeta, usize)>, String> {
|
||||
if buffer.len() < 4 {
|
||||
return Ok(None);
|
||||
}
|
||||
let meta_len = u32::from_be_bytes([buffer[0], buffer[1], buffer[2], buffer[3]]) as usize;
|
||||
if meta_len > MAX_TUNNEL_RELAY_META_LEN {
|
||||
return Err("relay metadata too large".to_string());
|
||||
}
|
||||
let meta_end = 4usize
|
||||
.checked_add(meta_len)
|
||||
.ok_or_else(|| "relay envelope length overflow".to_string())?;
|
||||
if buffer.len() < meta_end {
|
||||
return Ok(None);
|
||||
}
|
||||
let meta = serde_json::from_slice::<RequestMeta>(&buffer[4..meta_end])
|
||||
.map_err(|error| format!("invalid relay metadata: {error}"))?;
|
||||
Ok(Some((meta, meta_end)))
|
||||
}
|
||||
|
||||
fn clamp_upstream_request_timeout_ms(timeout_ms: u64) -> u64 {
|
||||
timeout_ms.clamp(1, crate::MAX_EXECUTION_REQUEST_TIMEOUT_MS)
|
||||
}
|
||||
|
||||
fn clamp_stream_first_byte_timeout_ms(timeout_ms: u64) -> u64 {
|
||||
timeout_ms.clamp(1, crate::MAX_EXECUTION_STREAM_FIRST_BYTE_TIMEOUT_MS)
|
||||
}
|
||||
|
||||
fn default_timeout() -> u64 {
|
||||
60
|
||||
}
|
||||
@@ -461,13 +519,115 @@ fn compress_gzip(data: &[u8]) -> Result<Bytes, std::io::Error> {
|
||||
mod tests {
|
||||
use super::{
|
||||
compress_payload, decode_payload, encode_frame, encode_goaway_v3, encode_ping,
|
||||
encode_reset_stream, encode_window_update, raw_payload, Frame, FrameHeader, GoAwayPayload,
|
||||
MsgType, RequestMeta, ResetStreamPayload, WindowUpdatePayload,
|
||||
CURRENT_TUNNEL_PROTOCOL_VERSION, CURRENT_TUNNEL_PROTOCOL_VERSION_STR, FLAG_GZIP_COMPRESSED,
|
||||
encode_reset_stream, encode_window_update, raw_payload, resolve_tunnel_request_timeouts,
|
||||
try_decode_tunnel_relay_request_meta, Frame, FrameHeader, GoAwayPayload, MsgType,
|
||||
RequestMeta, ResetStreamPayload, WindowUpdatePayload, CURRENT_TUNNEL_PROTOCOL_VERSION,
|
||||
CURRENT_TUNNEL_PROTOCOL_VERSION_STR, FLAG_GZIP_COMPRESSED, MAX_TUNNEL_RELAY_META_LEN,
|
||||
REQUEST_HEADERS, TUNNEL_PROTOCOL_VERSION_HEADER,
|
||||
};
|
||||
use bytes::Bytes;
|
||||
|
||||
fn request_meta(stream: bool) -> RequestMeta {
|
||||
RequestMeta {
|
||||
provider_id: None,
|
||||
endpoint_id: None,
|
||||
key_id: None,
|
||||
method: "POST".to_string(),
|
||||
url: "https://example.com/responses".to_string(),
|
||||
headers: std::collections::HashMap::new(),
|
||||
stream,
|
||||
request_timeout_ms: None,
|
||||
stream_first_byte_timeout_ms: None,
|
||||
timeout: 60,
|
||||
follow_redirects: None,
|
||||
http1_only: false,
|
||||
transport_profile: None,
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn tunnel_request_timeouts_preserve_non_stream_total_timeout() {
|
||||
let mut meta = request_meta(false);
|
||||
meta.request_timeout_ms = Some(900_000);
|
||||
meta.stream_first_byte_timeout_ms = Some(12_000);
|
||||
|
||||
let resolved = resolve_tunnel_request_timeouts(&meta);
|
||||
|
||||
assert_eq!(resolved.first_byte_ms, 900_000);
|
||||
assert_eq!(resolved.response_body_ms, Some(900_000));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn tunnel_request_timeouts_keep_stream_body_unbounded() {
|
||||
let mut meta = request_meta(true);
|
||||
meta.request_timeout_ms = Some(900_000);
|
||||
meta.stream_first_byte_timeout_ms = Some(12_000);
|
||||
|
||||
let resolved = resolve_tunnel_request_timeouts(&meta);
|
||||
|
||||
assert_eq!(resolved.first_byte_ms, 12_000);
|
||||
assert_eq!(resolved.response_body_ms, None);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn tunnel_request_timeouts_keep_stream_first_byte_protocol_limit() {
|
||||
let mut meta = request_meta(true);
|
||||
meta.stream_first_byte_timeout_ms =
|
||||
Some(crate::MAX_EXECUTION_STREAM_FIRST_BYTE_TIMEOUT_MS + 1);
|
||||
|
||||
let resolved = resolve_tunnel_request_timeouts(&meta);
|
||||
|
||||
assert_eq!(
|
||||
resolved.first_byte_ms,
|
||||
crate::MAX_EXECUTION_STREAM_FIRST_BYTE_TIMEOUT_MS
|
||||
);
|
||||
assert_eq!(resolved.response_body_ms, None);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn tunnel_request_timeouts_clamp_only_out_of_range_protocol_values() {
|
||||
let mut meta = request_meta(false);
|
||||
meta.request_timeout_ms = Some(crate::MAX_EXECUTION_REQUEST_TIMEOUT_MS + 1);
|
||||
|
||||
let resolved = resolve_tunnel_request_timeouts(&meta);
|
||||
|
||||
assert_eq!(
|
||||
resolved.first_byte_ms,
|
||||
crate::MAX_EXECUTION_REQUEST_TIMEOUT_MS
|
||||
);
|
||||
assert_eq!(
|
||||
resolved.response_body_ms,
|
||||
Some(crate::MAX_EXECUTION_REQUEST_TIMEOUT_MS)
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn tunnel_relay_request_meta_decodes_from_a_partial_prefix() {
|
||||
let meta = request_meta(false);
|
||||
let encoded_meta = serde_json::to_vec(&meta).expect("meta should encode");
|
||||
let mut envelope = Vec::new();
|
||||
envelope.extend_from_slice(&(encoded_meta.len() as u32).to_be_bytes());
|
||||
envelope.extend_from_slice(&encoded_meta);
|
||||
envelope.extend_from_slice(b"request-body");
|
||||
|
||||
assert!(try_decode_tunnel_relay_request_meta(&envelope[..3])
|
||||
.expect("partial prefix should be valid")
|
||||
.is_none());
|
||||
let (decoded, body_offset) = try_decode_tunnel_relay_request_meta(&envelope)
|
||||
.expect("envelope should be valid")
|
||||
.expect("metadata should be complete");
|
||||
|
||||
assert_eq!(decoded.request_timeout_ms, meta.request_timeout_ms);
|
||||
assert_eq!(&envelope[body_offset..], b"request-body");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn tunnel_relay_request_meta_rejects_oversized_prefix() {
|
||||
let oversized = (MAX_TUNNEL_RELAY_META_LEN as u32 + 1).to_be_bytes();
|
||||
|
||||
assert!(try_decode_tunnel_relay_request_meta(&oversized).is_err());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn request_meta_accepts_integer_timeout() {
|
||||
let raw = br#"{"method":"GET","url":"https://example.com","headers":{},"timeout":15}"#;
|
||||
|
||||
@@ -97,13 +97,13 @@ impl StoredMinimalCandidateSelectionRow {
|
||||
None => true,
|
||||
Some(formats) => formats
|
||||
.iter()
|
||||
.any(|value| api_format_matches(value, api_format)),
|
||||
.any(|value| api_format_permission_covers(value, api_format)),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn api_format_matches(left: &str, right: &str) -> bool {
|
||||
aether_ai_formats::api_format_alias_matches(left, right)
|
||||
fn api_format_permission_covers(allowed: &str, requested: &str) -> bool {
|
||||
aether_ai_formats::api_format_permission_covers(allowed, requested)
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
|
||||
@@ -247,7 +247,7 @@ impl ResolvedAuthApiKeySnapshot {
|
||||
&mut self.user_allowed_providers,
|
||||
&mut self.api_key_allowed_providers,
|
||||
);
|
||||
constrain_api_key_list_policy_to_user_policy(
|
||||
constrain_api_key_api_format_policy_to_user_policy(
|
||||
&mut self.user_allowed_api_formats,
|
||||
&mut self.api_key_allowed_api_formats,
|
||||
);
|
||||
@@ -277,6 +277,22 @@ fn constrain_api_key_list_policy_to_user_policy(
|
||||
*api_key_policy = Some(effective);
|
||||
}
|
||||
|
||||
fn constrain_api_key_api_format_policy_to_user_policy(
|
||||
user_policy: &mut Option<Vec<String>>,
|
||||
api_key_policy: &mut Option<Vec<String>>,
|
||||
) {
|
||||
let Some(api_key_values) = api_key_policy.as_ref() else {
|
||||
return;
|
||||
};
|
||||
let Some(user_values) = user_policy.as_ref() else {
|
||||
return;
|
||||
};
|
||||
let effective =
|
||||
aether_ai_formats::intersect_api_format_allowed_lists(api_key_values, user_values);
|
||||
*user_policy = Some(effective.clone());
|
||||
*api_key_policy = Some(effective);
|
||||
}
|
||||
|
||||
fn intersect_allowed_lists(left: &[String], right: &[String]) -> Vec<String> {
|
||||
let right_values = right.iter().collect::<std::collections::BTreeSet<_>>();
|
||||
left.iter()
|
||||
@@ -810,6 +826,29 @@ mod tests {
|
||||
StoredAuthApiKeySnapshot,
|
||||
};
|
||||
|
||||
#[test]
|
||||
fn api_format_policy_intersection_preserves_search_companion_scope() {
|
||||
assert_eq!(
|
||||
aether_ai_formats::intersect_api_format_allowed_lists(
|
||||
&["openai:search".to_string()],
|
||||
&["openai:responses".to_string()],
|
||||
),
|
||||
vec!["openai:search".to_string()]
|
||||
);
|
||||
assert_eq!(
|
||||
aether_ai_formats::intersect_api_format_allowed_lists(
|
||||
&["openai:responses".to_string()],
|
||||
&["openai:search".to_string()],
|
||||
),
|
||||
vec!["openai:search".to_string()]
|
||||
);
|
||||
assert!(aether_ai_formats::intersect_api_format_allowed_lists(
|
||||
&["openai:search".to_string()],
|
||||
&["openai:chat".to_string()],
|
||||
)
|
||||
.is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn rejects_non_array_allowed_providers() {
|
||||
assert!(StoredAuthApiKeySnapshot::new(
|
||||
|
||||
@@ -230,7 +230,7 @@ fn row_matches_requested_model(
|
||||
mapping.api_formats.as_ref().is_none_or(|formats| {
|
||||
formats
|
||||
.iter()
|
||||
.any(|value| api_format_matches(value, api_format))
|
||||
.any(|value| api_format_scope_covers(value, api_format))
|
||||
}) && mapping.endpoint_ids.as_ref().is_none_or(|endpoint_ids| {
|
||||
endpoint_ids
|
||||
.iter()
|
||||
@@ -286,7 +286,7 @@ fn mapping_scope_matches(
|
||||
mapping.api_formats.as_ref().is_none_or(|formats| {
|
||||
formats
|
||||
.iter()
|
||||
.any(|value| api_format_matches(value, api_format))
|
||||
.any(|value| api_format_scope_covers(value, api_format))
|
||||
}) && mapping.endpoint_ids.as_ref().is_none_or(|endpoint_ids| {
|
||||
endpoint_ids
|
||||
.iter()
|
||||
@@ -294,6 +294,10 @@ fn mapping_scope_matches(
|
||||
})
|
||||
}
|
||||
|
||||
fn api_format_scope_covers(allowed: &str, requested: &str) -> bool {
|
||||
aether_ai_formats::api_format_permission_covers(allowed, requested)
|
||||
}
|
||||
|
||||
fn key_auth_channel_matches(row: &StoredMinimalCandidateSelectionRow, api_format: &str) -> bool {
|
||||
let provider_type = row.provider_type.trim().to_ascii_lowercase();
|
||||
let auth_type = row.key_auth_type.trim().to_ascii_lowercase();
|
||||
@@ -303,7 +307,10 @@ fn key_auth_channel_matches(row: &StoredMinimalCandidateSelectionRow, api_format
|
||||
auth_type == "oauth"
|
||||
&& matches!(
|
||||
api_format.as_str(),
|
||||
"openai:responses" | "openai:responses:compact" | "openai:image"
|
||||
"openai:responses"
|
||||
| "openai:responses:compact"
|
||||
| "openai:search"
|
||||
| "openai:image"
|
||||
)
|
||||
}
|
||||
"chatgpt_web" => {
|
||||
@@ -434,6 +441,47 @@ mod tests {
|
||||
assert_eq!(rows[0].provider_id, "provider-1");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn search_uses_responses_key_and_model_permissions_with_exact_endpoint_identity() {
|
||||
let mut search = sample_row(
|
||||
"provider-search",
|
||||
"openai:search",
|
||||
"global-search-model",
|
||||
10,
|
||||
);
|
||||
search.provider_type = "codex".to_string();
|
||||
search.key_auth_type = "oauth".to_string();
|
||||
search.key_api_formats = Some(vec!["openai:responses".to_string()]);
|
||||
search.model_provider_model_name = "upstream-search-model".to_string();
|
||||
search.model_provider_model_mappings = Some(vec![StoredProviderModelMapping {
|
||||
name: "gpt-5.6-sol".to_string(),
|
||||
priority: 0,
|
||||
api_formats: Some(vec!["openai:responses".to_string()]),
|
||||
endpoint_ids: None,
|
||||
}]);
|
||||
|
||||
let mut responses = search.clone();
|
||||
responses.provider_id = "provider-responses".to_string();
|
||||
responses.endpoint_id = "endpoint-responses".to_string();
|
||||
responses.endpoint_api_format = "openai:responses".to_string();
|
||||
responses.key_id = "key-responses".to_string();
|
||||
responses.model_id = "model-responses".to_string();
|
||||
|
||||
let repository =
|
||||
InMemoryMinimalCandidateSelectionReadRepository::seed(vec![responses, search]);
|
||||
let rows = repository
|
||||
.list_for_exact_api_format_and_requested_model("openai:search", "gpt-5.6-sol")
|
||||
.await
|
||||
.expect("Search candidate should load");
|
||||
|
||||
assert_eq!(rows.len(), 1);
|
||||
assert_eq!(rows[0].endpoint_api_format, "openai:search");
|
||||
assert_eq!(
|
||||
rows[0].key_api_formats,
|
||||
Some(vec!["openai:responses".to_string()])
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn includes_grok_oauth_rows_for_chat_models() {
|
||||
let mut row = sample_row(
|
||||
|
||||
@@ -409,7 +409,7 @@ fn mapping_scope_matches(
|
||||
mapping.api_formats.as_ref().is_none_or(|formats| {
|
||||
formats
|
||||
.iter()
|
||||
.any(|value| api_format_matches(value, api_format))
|
||||
.any(|value| api_format_scope_covers(value, api_format))
|
||||
}) && mapping.endpoint_ids.as_ref().is_none_or(|endpoint_ids| {
|
||||
endpoint_ids
|
||||
.iter()
|
||||
@@ -426,7 +426,10 @@ fn key_auth_channel_matches(row: &CandidateSelectionRow, api_format: &str) -> bo
|
||||
auth_type == "oauth"
|
||||
&& matches!(
|
||||
api_format.as_str(),
|
||||
"openai:responses" | "openai:responses:compact" | "openai:image"
|
||||
"openai:responses"
|
||||
| "openai:responses:compact"
|
||||
| "openai:search"
|
||||
| "openai:image"
|
||||
)
|
||||
}
|
||||
"chatgpt_web" => {
|
||||
@@ -763,6 +766,10 @@ fn api_format_matches(left: &str, right: &str) -> bool {
|
||||
aether_ai_formats::api_format_alias_matches(left, right)
|
||||
}
|
||||
|
||||
fn api_format_scope_covers(allowed: &str, requested: &str) -> bool {
|
||||
aether_ai_formats::api_format_permission_covers(allowed, requested)
|
||||
}
|
||||
|
||||
fn sql_match_aliases(api_formats: &[String]) -> Vec<String> {
|
||||
api_formats
|
||||
.iter()
|
||||
|
||||
@@ -71,7 +71,7 @@ INNER JOIN LATERAL (
|
||||
(
|
||||
LOWER(BTRIM(p.provider_type)) = 'codex'
|
||||
AND LOWER(BTRIM(pak.auth_type)) = 'oauth'
|
||||
AND LOWER($3) IN ('openai:responses', 'openai:responses:compact', 'openai:image')
|
||||
AND LOWER($3) IN ('openai:responses', 'openai:responses:compact', 'openai:search', 'openai:image')
|
||||
)
|
||||
OR (
|
||||
LOWER(BTRIM(p.provider_type)) = 'chatgpt_web'
|
||||
@@ -164,7 +164,7 @@ WHERE p.is_active = TRUE
|
||||
(
|
||||
LOWER(BTRIM(p.provider_type)) = 'codex'
|
||||
AND LOWER(BTRIM(pak.auth_type)) = 'oauth'
|
||||
AND LOWER($3) IN ('openai:responses', 'openai:responses:compact', 'openai:image')
|
||||
AND LOWER($3) IN ('openai:responses', 'openai:responses:compact', 'openai:search', 'openai:image')
|
||||
)
|
||||
OR (
|
||||
LOWER(BTRIM(p.provider_type)) = 'chatgpt_web'
|
||||
@@ -350,7 +350,7 @@ INNER JOIN LATERAL (
|
||||
(
|
||||
LOWER(BTRIM(p.provider_type)) = 'codex'
|
||||
AND LOWER(BTRIM(pak.auth_type)) = 'oauth'
|
||||
AND LOWER($4) IN ('openai:responses', 'openai:responses:compact', 'openai:image')
|
||||
AND LOWER($4) IN ('openai:responses', 'openai:responses:compact', 'openai:search', 'openai:image')
|
||||
)
|
||||
OR (
|
||||
LOWER(BTRIM(p.provider_type)) = 'chatgpt_web'
|
||||
@@ -444,7 +444,7 @@ WHERE p.is_active = TRUE
|
||||
(
|
||||
LOWER(BTRIM(p.provider_type)) = 'codex'
|
||||
AND LOWER(BTRIM(pak.auth_type)) = 'oauth'
|
||||
AND LOWER($4) IN ('openai:responses', 'openai:responses:compact', 'openai:image')
|
||||
AND LOWER($4) IN ('openai:responses', 'openai:responses:compact', 'openai:search', 'openai:image')
|
||||
)
|
||||
OR (
|
||||
LOWER(BTRIM(p.provider_type)) = 'chatgpt_web'
|
||||
@@ -638,7 +638,7 @@ WHERE p.is_active = TRUE
|
||||
(
|
||||
LOWER(BTRIM(p.provider_type)) = 'codex'
|
||||
AND LOWER(BTRIM(pak.auth_type)) = 'oauth'
|
||||
AND LOWER($6) IN ('openai:responses', 'openai:responses:compact', 'openai:image')
|
||||
AND LOWER($6) IN ('openai:responses', 'openai:responses:compact', 'openai:search', 'openai:image')
|
||||
)
|
||||
OR (
|
||||
LOWER(BTRIM(p.provider_type)) = 'chatgpt_web'
|
||||
@@ -784,7 +784,8 @@ impl SqlxMinimalCandidateSelectionReadRepository {
|
||||
let mut rows = Vec::new();
|
||||
let canonical_api_format = normalize_api_format(api_format);
|
||||
let storage_aliases = api_format_aliases(&canonical_api_format);
|
||||
let sql_match_aliases = sql_match_aliases(&storage_aliases);
|
||||
let sql_match_aliases =
|
||||
sql_match_aliases(&api_format_permission_aliases(&canonical_api_format));
|
||||
for api_format in storage_aliases {
|
||||
rows.extend(
|
||||
Self::collect_query_rows(
|
||||
@@ -809,7 +810,8 @@ impl SqlxMinimalCandidateSelectionReadRepository {
|
||||
let mut rows = Vec::new();
|
||||
let canonical_api_format = normalize_api_format(api_format);
|
||||
let storage_aliases = api_format_aliases(&canonical_api_format);
|
||||
let sql_match_aliases = sql_match_aliases(&storage_aliases);
|
||||
let sql_match_aliases =
|
||||
sql_match_aliases(&api_format_permission_aliases(&canonical_api_format));
|
||||
for api_format in storage_aliases {
|
||||
rows.extend(
|
||||
Self::collect_query_rows(
|
||||
@@ -835,7 +837,8 @@ impl SqlxMinimalCandidateSelectionReadRepository {
|
||||
let mut rows = Vec::new();
|
||||
let canonical_api_format = normalize_api_format(api_format);
|
||||
let storage_aliases = api_format_aliases(&canonical_api_format);
|
||||
let sql_match_aliases = sql_match_aliases(&storage_aliases);
|
||||
let sql_match_aliases =
|
||||
sql_match_aliases(&api_format_permission_aliases(&canonical_api_format));
|
||||
let sql = requested_model_selection_sql();
|
||||
for api_format in storage_aliases {
|
||||
rows.extend(
|
||||
@@ -861,7 +864,8 @@ impl SqlxMinimalCandidateSelectionReadRepository {
|
||||
let mut rows = Vec::new();
|
||||
let canonical_api_format = normalize_api_format(&query.api_format);
|
||||
let storage_aliases = api_format_aliases(&canonical_api_format);
|
||||
let sql_match_aliases = sql_match_aliases(&storage_aliases);
|
||||
let sql_match_aliases =
|
||||
sql_match_aliases(&api_format_permission_aliases(&canonical_api_format));
|
||||
let limit = i64::from(query.limit.max(1));
|
||||
let offset = i64::from(query.offset);
|
||||
let sql = requested_model_selection_page_sql();
|
||||
@@ -891,7 +895,8 @@ impl SqlxMinimalCandidateSelectionReadRepository {
|
||||
let mut rows = Vec::new();
|
||||
let canonical_api_format = normalize_api_format(&query.api_format);
|
||||
let storage_aliases = api_format_aliases(&canonical_api_format);
|
||||
let sql_match_aliases = sql_match_aliases(&storage_aliases);
|
||||
let sql_match_aliases =
|
||||
sql_match_aliases(&api_format_permission_aliases(&canonical_api_format));
|
||||
let limit = i64::from(query.limit.max(1));
|
||||
let offset = i64::from(query.offset);
|
||||
let sql = pool_key_candidate_selection_sql(&query.order);
|
||||
@@ -929,7 +934,8 @@ impl SqlxMinimalCandidateSelectionReadRepository {
|
||||
let mut rows = Vec::new();
|
||||
let canonical_api_format = normalize_api_format(&query.api_format);
|
||||
let storage_aliases = api_format_aliases(&canonical_api_format);
|
||||
let sql_match_aliases = sql_match_aliases(&storage_aliases);
|
||||
let sql_match_aliases =
|
||||
sql_match_aliases(&api_format_permission_aliases(&canonical_api_format));
|
||||
let sql = pool_key_candidate_selection_by_key_ids_sql();
|
||||
for api_format in storage_aliases {
|
||||
rows.extend(
|
||||
@@ -1111,6 +1117,10 @@ fn api_format_aliases(api_format: &str) -> Vec<String> {
|
||||
aether_ai_formats::api_format_storage_aliases(api_format)
|
||||
}
|
||||
|
||||
fn api_format_permission_aliases(api_format: &str) -> Vec<String> {
|
||||
aether_ai_formats::api_format_permission_storage_aliases(api_format)
|
||||
}
|
||||
|
||||
fn normalize_api_format(api_format: &str) -> String {
|
||||
aether_ai_formats::normalize_api_format_alias(api_format)
|
||||
}
|
||||
|
||||
@@ -125,7 +125,8 @@ impl SqliteMinimalCandidateSelectionReadRepository {
|
||||
) -> Result<Vec<StoredMinimalCandidateSelectionRow>, DataLayerError> {
|
||||
let canonical_api_format = normalize_api_format(api_format);
|
||||
let storage_aliases = api_format_aliases(&canonical_api_format);
|
||||
let match_aliases = sql_match_aliases(&storage_aliases);
|
||||
let match_aliases =
|
||||
sql_match_aliases(&api_format_permission_aliases(&canonical_api_format));
|
||||
let mut rows = Vec::new();
|
||||
|
||||
for storage_api_format in storage_aliases {
|
||||
@@ -270,7 +271,8 @@ impl MinimalCandidateSelectionReadRepository for SqliteMinimalCandidateSelection
|
||||
) -> Result<Vec<StoredMinimalCandidateSelectionRow>, DataLayerError> {
|
||||
let canonical_api_format = normalize_api_format(&query.api_format);
|
||||
let storage_aliases = api_format_aliases(&canonical_api_format);
|
||||
let match_aliases = sql_match_aliases(&storage_aliases);
|
||||
let match_aliases =
|
||||
sql_match_aliases(&api_format_permission_aliases(&canonical_api_format));
|
||||
let mut rows = Vec::<CandidateSelectionRow>::new();
|
||||
let page_in_sql = !matches!(query.order, StoredPoolKeyCandidateOrder::LoadBalance { .. });
|
||||
|
||||
@@ -337,7 +339,8 @@ impl MinimalCandidateSelectionReadRepository for SqliteMinimalCandidateSelection
|
||||
.collect::<BTreeMap<_, _>>();
|
||||
let canonical_api_format = normalize_api_format(&query.api_format);
|
||||
let storage_aliases = api_format_aliases(&canonical_api_format);
|
||||
let match_aliases = sql_match_aliases(&storage_aliases);
|
||||
let match_aliases =
|
||||
sql_match_aliases(&api_format_permission_aliases(&canonical_api_format));
|
||||
let mut rows = Vec::new();
|
||||
|
||||
for storage_api_format in storage_aliases {
|
||||
@@ -493,7 +496,7 @@ fn push_key_auth_channel_sql_filter(
|
||||
);
|
||||
builder.push_bind(api_format.clone());
|
||||
builder.push(
|
||||
r#" IN ('openai:responses', 'openai:responses:compact', 'openai:image')
|
||||
r#" IN ('openai:responses', 'openai:responses:compact', 'openai:search', 'openai:image')
|
||||
)
|
||||
OR (
|
||||
LOWER(TRIM(p.provider_type)) = 'chatgpt_web'
|
||||
@@ -808,7 +811,7 @@ fn mapping_scope_matches(
|
||||
mapping.api_formats.as_ref().is_none_or(|formats| {
|
||||
formats
|
||||
.iter()
|
||||
.any(|value| api_format_matches(value, api_format))
|
||||
.any(|value| api_format_scope_covers(value, api_format))
|
||||
}) && mapping.endpoint_ids.as_ref().is_none_or(|endpoint_ids| {
|
||||
endpoint_ids
|
||||
.iter()
|
||||
@@ -825,7 +828,10 @@ fn key_auth_channel_matches(row: &CandidateSelectionRow, api_format: &str) -> bo
|
||||
auth_type == "oauth"
|
||||
&& matches!(
|
||||
api_format.as_str(),
|
||||
"openai:responses" | "openai:responses:compact" | "openai:image"
|
||||
"openai:responses"
|
||||
| "openai:responses:compact"
|
||||
| "openai:search"
|
||||
| "openai:image"
|
||||
)
|
||||
}
|
||||
"chatgpt_web" => {
|
||||
@@ -1145,6 +1151,10 @@ fn api_format_aliases(api_format: &str) -> Vec<String> {
|
||||
aether_ai_formats::api_format_storage_aliases(api_format)
|
||||
}
|
||||
|
||||
fn api_format_permission_aliases(api_format: &str) -> Vec<String> {
|
||||
aether_ai_formats::api_format_permission_storage_aliases(api_format)
|
||||
}
|
||||
|
||||
fn normalize_api_format(api_format: &str) -> String {
|
||||
aether_ai_formats::normalize_api_format_alias(api_format)
|
||||
}
|
||||
@@ -1153,6 +1163,10 @@ fn api_format_matches(left: &str, right: &str) -> bool {
|
||||
aether_ai_formats::api_format_alias_matches(left, right)
|
||||
}
|
||||
|
||||
fn api_format_scope_covers(allowed: &str, requested: &str) -> bool {
|
||||
aether_ai_formats::api_format_permission_covers(allowed, requested)
|
||||
}
|
||||
|
||||
fn sql_match_aliases(api_formats: &[String]) -> Vec<String> {
|
||||
api_formats
|
||||
.iter()
|
||||
@@ -1264,6 +1278,21 @@ mod tests {
|
||||
.collect::<Vec<_>>(),
|
||||
vec!["key-chatgpt-web-oauth", "key-chatgpt-web-bearer"]
|
||||
);
|
||||
|
||||
let search_rows = repository
|
||||
.list_for_exact_api_format_and_requested_model_page(
|
||||
&StoredRequestedModelCandidateRowsQuery {
|
||||
api_format: "openai:search".to_string(),
|
||||
requested_model_name: "gpt-5.6-sol".to_string(),
|
||||
offset: 0,
|
||||
limit: 10,
|
||||
},
|
||||
)
|
||||
.await
|
||||
.expect("Codex Search rows should load through Responses permissions");
|
||||
assert_eq!(search_rows.len(), 1);
|
||||
assert_eq!(search_rows[0].key_id, "key-codex-search");
|
||||
assert_eq!(search_rows[0].endpoint_api_format, "openai:search");
|
||||
}
|
||||
|
||||
async fn seed_candidate_selection(pool: &sqlx::SqlitePool) {
|
||||
@@ -1296,6 +1325,11 @@ INSERT INTO providers (
|
||||
)
|
||||
VALUES ('provider-windsurf', 'Windsurf', 'windsurf', 15, 1, 1, 1);
|
||||
|
||||
INSERT INTO providers (
|
||||
id, name, provider_type, provider_priority, is_active, created_at, updated_at
|
||||
)
|
||||
VALUES ('provider-codex-search', 'Codex Search', 'codex', 12, 1, 1, 1);
|
||||
|
||||
INSERT INTO provider_endpoints (
|
||||
id, provider_id, name, base_url, api_format, is_active, created_at, updated_at
|
||||
)
|
||||
@@ -1312,6 +1346,14 @@ VALUES (
|
||||
'https://server.codeium.com', 'openai:chat', 1, 1, 1
|
||||
);
|
||||
|
||||
INSERT INTO provider_endpoints (
|
||||
id, provider_id, name, base_url, api_format, is_active, created_at, updated_at
|
||||
)
|
||||
VALUES (
|
||||
'endpoint-codex-search', 'provider-codex-search', 'Codex Search',
|
||||
'https://chatgpt.com/backend-api/codex', 'openai:search', 1, 1, 1
|
||||
);
|
||||
|
||||
INSERT INTO provider_api_keys (
|
||||
id, provider_id, name, auth_type, api_formats, internal_priority, is_active, created_at, updated_at
|
||||
)
|
||||
@@ -1327,13 +1369,22 @@ VALUES (
|
||||
'key-windsurf-oauth', 'provider-windsurf', 'OAuth', 'oauth', '["openai:chat"]', 10, 1, 1, 1
|
||||
);
|
||||
|
||||
INSERT INTO provider_api_keys (
|
||||
id, provider_id, name, auth_type, api_formats, internal_priority, is_active, created_at, updated_at
|
||||
)
|
||||
VALUES (
|
||||
'key-codex-search', 'provider-codex-search', 'OAuth', 'oauth',
|
||||
'["openai:responses"]', 10, 1, 1, 1
|
||||
);
|
||||
|
||||
INSERT INTO global_models (
|
||||
id, name, config, is_active, created_at, updated_at
|
||||
)
|
||||
VALUES
|
||||
('global-1', 'gpt-5', '{"model_mappings":["alias-global"],"streaming":true}', 1, 1, 1),
|
||||
('global-image-1', 'gpt-image-2', NULL, 1, 1, 1),
|
||||
('global-windsurf-1', 'claude-opus-4-7', '{"streaming":true}', 1, 1, 1);
|
||||
('global-windsurf-1', 'claude-opus-4-7', '{"streaming":true}', 1, 1, 1),
|
||||
('global-codex-search-1', 'search-global', '{"streaming":false}', 1, 1, 1);
|
||||
|
||||
INSERT INTO models (
|
||||
id, provider_id, global_model_id, provider_model_name, provider_model_mappings,
|
||||
@@ -1351,6 +1402,11 @@ VALUES (
|
||||
(
|
||||
'model-windsurf-opus', 'provider-windsurf', 'global-windsurf-1', 'claude-opus-4-7',
|
||||
NULL, NULL, 1, 1, 1, 1
|
||||
),
|
||||
(
|
||||
'model-codex-search', 'provider-codex-search', 'global-codex-search-1', 'search-upstream',
|
||||
'[{"name":"gpt-5.6-sol","api_formats":["openai:responses"],"priority":1}]',
|
||||
0, 1, 1, 1, 1
|
||||
);
|
||||
"#,
|
||||
)
|
||||
|
||||
@@ -354,9 +354,9 @@ fn transport_key_supports_api_format(
|
||||
|
||||
match transport.key.api_formats.as_deref() {
|
||||
None => true,
|
||||
Some(formats) => formats
|
||||
.iter()
|
||||
.any(|value| aether_ai_formats::api_format_alias_matches(value, endpoint_api_format)),
|
||||
Some(formats) => formats.iter().any(|value| {
|
||||
aether_ai_formats::api_format_permission_covers(value, endpoint_api_format)
|
||||
}),
|
||||
}
|
||||
}
|
||||
|
||||
@@ -818,6 +818,29 @@ mod tests {
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn responses_key_permission_covers_search_without_changing_endpoint_identity() {
|
||||
let mut transport = transport_snapshot("custom", "openai:search", "bearer", true, None);
|
||||
transport.key.api_formats = Some(vec!["openai:responses".to_string()]);
|
||||
|
||||
assert_eq!(
|
||||
candidate_common_transport_skip_reason(
|
||||
&transport,
|
||||
candidate_facts("openai:search"),
|
||||
None,
|
||||
),
|
||||
None
|
||||
);
|
||||
assert_eq!(
|
||||
candidate_common_transport_skip_reason(
|
||||
&transport,
|
||||
candidate_facts("openai:responses"),
|
||||
None,
|
||||
),
|
||||
Some("endpoint_api_format_changed")
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn candidate_transport_pair_policy_reports_disabled_conversion_and_unsupported_pairs() {
|
||||
let transport = transport_snapshot("custom", "openai:responses", "bearer", false, None);
|
||||
|
||||
@@ -454,6 +454,21 @@ mod tests {
|
||||
assert_eq!(timeouts.first_byte_ms, Some(30_000));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn transport_execution_timeouts_preserve_the_configurable_maximum() {
|
||||
let mut transport = sample_transport();
|
||||
transport.provider.request_timeout_secs =
|
||||
Some(aether_contracts::MAX_EXECUTION_REQUEST_TIMEOUT_SECS as f64);
|
||||
|
||||
let timeouts = resolve_transport_execution_timeouts(&transport)
|
||||
.expect("provider timeouts should resolve");
|
||||
|
||||
assert_eq!(
|
||||
timeouts.total_ms,
|
||||
Some(aether_contracts::MAX_EXECUTION_REQUEST_TIMEOUT_MS)
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn transport_execution_timeouts_preserve_configured_first_byte_value() {
|
||||
let mut transport = sample_transport();
|
||||
|
||||
@@ -306,6 +306,12 @@ const CODEX_FIXED_PROVIDER_TEMPLATE: FixedProviderTemplate = FixedProviderTempla
|
||||
custom_path: None,
|
||||
config_defaults: EMPTY_ENDPOINT_CONFIG_DEFAULTS,
|
||||
},
|
||||
FixedProviderEndpointTemplate {
|
||||
item_key: "openai:search",
|
||||
api_format: "openai:search",
|
||||
custom_path: None,
|
||||
config_defaults: EMPTY_ENDPOINT_CONFIG_DEFAULTS,
|
||||
},
|
||||
FixedProviderEndpointTemplate {
|
||||
item_key: "openai:image",
|
||||
api_format: "openai:image",
|
||||
@@ -639,7 +645,7 @@ mod tests {
|
||||
};
|
||||
|
||||
#[test]
|
||||
fn codex_fixed_provider_template_includes_openai_image() {
|
||||
fn codex_fixed_provider_template_includes_codex_companion_endpoints() {
|
||||
let template = fixed_provider_template("codex").expect("codex template should exist");
|
||||
assert_eq!(template.base_url, "https://chatgpt.com/backend-api/codex");
|
||||
assert_eq!(template.version, 1);
|
||||
@@ -652,6 +658,7 @@ mod tests {
|
||||
vec![
|
||||
"openai:responses",
|
||||
"openai:responses:compact",
|
||||
"openai:search",
|
||||
"openai:image"
|
||||
]
|
||||
);
|
||||
@@ -660,6 +667,11 @@ mod tests {
|
||||
fixed_provider_endpoint_template_by_api_format("codex", "openai:image")
|
||||
.expect("codex image endpoint should exist");
|
||||
assert!(image_template.config_defaults.is_empty());
|
||||
|
||||
let search_template =
|
||||
fixed_provider_endpoint_template_by_api_format("codex", "openai:search")
|
||||
.expect("codex search endpoint should exist");
|
||||
assert!(search_template.config_defaults.is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
|
||||
@@ -16,7 +16,8 @@ use crate::gemini_cli::{
|
||||
use crate::snapshot::GatewayProviderTransportSnapshot;
|
||||
use crate::url::{
|
||||
build_claude_messages_url, build_gemini_content_url, build_openai_chat_url,
|
||||
build_openai_responses_url, build_passthrough_path_url, normalize_gemini_content_action_path,
|
||||
build_openai_responses_url, build_openai_search_url, build_passthrough_path_url,
|
||||
normalize_gemini_content_action_path,
|
||||
};
|
||||
use crate::vertex::{
|
||||
build_vertex_api_key_gemini_content_url, build_vertex_api_key_gemini_embedding_url,
|
||||
@@ -124,6 +125,10 @@ fn build_transport_request_url_inner(
|
||||
params.request_query,
|
||||
true,
|
||||
)),
|
||||
"openai:search" => Some(build_openai_search_url(
|
||||
&transport.endpoint.base_url,
|
||||
params.request_query,
|
||||
)),
|
||||
"openai:embedding" | "jina:embedding" => {
|
||||
build_provider_embedding_v1_url(&transport.endpoint.base_url, params.request_query)
|
||||
}
|
||||
@@ -868,6 +873,33 @@ mod tests {
|
||||
assert_eq!(url, "https://api.openai.example/v1/responses?tenant=demo");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn builds_openai_search_url_for_codex_provider_root() {
|
||||
let transport = sample_transport(
|
||||
"codex",
|
||||
"openai:search",
|
||||
"https://chatgpt.com/backend-api/codex",
|
||||
None,
|
||||
);
|
||||
|
||||
let url = build_transport_request_url(
|
||||
&transport,
|
||||
TransportRequestUrlParams {
|
||||
provider_api_format: "openai:search",
|
||||
mapped_model: Some("gpt-5.6-luna"),
|
||||
upstream_is_stream: false,
|
||||
request_query: Some("tenant=demo"),
|
||||
kiro_api_region: None,
|
||||
},
|
||||
)
|
||||
.expect("openai search url");
|
||||
|
||||
assert_eq!(
|
||||
url,
|
||||
"https://chatgpt.com/backend-api/codex/alpha/search?tenant=demo"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn expands_custom_path_templates_when_hook_does_not_apply() {
|
||||
let transport = sample_transport(
|
||||
|
||||
@@ -149,9 +149,11 @@ pub fn classify_same_format_provider_request_behavior(
|
||||
params.require_streaming,
|
||||
is_kiro || is_antigravity || gemini_cli_requires_upstream_streaming,
|
||||
);
|
||||
let force_body_stream_field = aether_ai_formats::endpoint_config_forces_upstream_stream_policy(
|
||||
transport.endpoint.config.as_ref(),
|
||||
);
|
||||
let force_body_stream_field =
|
||||
aether_ai_formats::api_format_uses_body_stream_field(params.provider_api_format)
|
||||
&& aether_ai_formats::endpoint_config_forces_upstream_stream_policy(
|
||||
transport.endpoint.config.as_ref(),
|
||||
);
|
||||
let report_kind = if is_kiro && !params.require_streaming {
|
||||
"claude_cli_sync_finalize"
|
||||
} else if (is_gemini_cli || is_antigravity) && !params.require_streaming {
|
||||
@@ -622,6 +624,7 @@ pub fn same_format_provider_transport_unsupported_reason_for_trace(
|
||||
"openai:chat" => "openai:chat",
|
||||
"openai:responses" => "openai:responses",
|
||||
"openai:responses:compact" => "openai:responses:compact",
|
||||
"openai:search" => "openai:search",
|
||||
"claude:messages" => "claude:messages",
|
||||
"gemini:generate_content" => "gemini:generate_content",
|
||||
"gemini:interactions" => "gemini:interactions",
|
||||
@@ -702,7 +705,9 @@ fn resolve_same_format_standard_direct_auth(
|
||||
transport: &GatewayProviderTransportSnapshot,
|
||||
provider_api_format: &str,
|
||||
) -> Option<(String, String)> {
|
||||
if aether_ai_formats::api_format_alias_matches(provider_api_format, "openai:embedding") {
|
||||
if aether_ai_formats::api_format_alias_matches(provider_api_format, "openai:embedding")
|
||||
|| aether_ai_formats::api_format_alias_matches(provider_api_format, "openai:search")
|
||||
{
|
||||
resolve_local_openai_bearer_auth(transport)
|
||||
} else {
|
||||
resolve_local_standard_auth(transport)
|
||||
@@ -886,6 +891,21 @@ mod tests {
|
||||
},
|
||||
);
|
||||
assert!(!compact_behavior.upstream_is_stream);
|
||||
|
||||
let mut search = sample_transport("codex");
|
||||
search.endpoint.config = Some(json!({
|
||||
"upstream_stream_policy": "force_stream"
|
||||
}));
|
||||
let search_behavior = classify_same_format_provider_request_behavior(
|
||||
&search,
|
||||
SameFormatProviderRequestBehaviorParams {
|
||||
require_streaming: true,
|
||||
provider_api_format: "openai:search",
|
||||
report_kind: "openai_search_sync_success",
|
||||
},
|
||||
);
|
||||
assert!(!search_behavior.upstream_is_stream);
|
||||
assert!(!search_behavior.force_body_stream_field);
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -1035,6 +1055,31 @@ mod tests {
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn resolves_openai_search_direct_auth_with_bearer_header() {
|
||||
let mut transport = sample_transport("custom");
|
||||
transport.endpoint.api_format = "openai:search".to_string();
|
||||
transport.key.auth_type = "api_key".to_string();
|
||||
let behavior = classify_same_format_provider_request_behavior(
|
||||
&transport,
|
||||
SameFormatProviderRequestBehaviorParams {
|
||||
require_streaming: false,
|
||||
provider_api_format: "openai:search",
|
||||
report_kind: "openai_search_sync_success",
|
||||
},
|
||||
);
|
||||
|
||||
assert_eq!(
|
||||
resolve_same_format_provider_direct_auth(
|
||||
&behavior,
|
||||
&transport,
|
||||
SameFormatProviderFamily::Standard,
|
||||
"openai:search",
|
||||
),
|
||||
Some(("authorization".to_string(), "Bearer secret".to_string()))
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn keeps_claude_same_format_api_key_on_x_api_key_header() {
|
||||
let mut transport = sample_transport("custom");
|
||||
@@ -1770,6 +1815,44 @@ mod tests {
|
||||
assert_eq!(body["metadata"]["body_rule_seen"], true);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn search_body_projects_the_protocol_contract_and_keeps_fast_routing_only() {
|
||||
let body = build_same_format_provider_request_body(SameFormatProviderRequestBodyInput {
|
||||
body_json: &json!({
|
||||
"id": "session-1",
|
||||
"model": "gpt-5.6-luna-high-fast",
|
||||
"commands": {"search_query": [{"q": "Aether"}]},
|
||||
"settings": {"allowed_callers": ["direct"]},
|
||||
"store": false,
|
||||
"future_extension": {"enabled": true},
|
||||
"stream": true
|
||||
}),
|
||||
mapped_model: "gpt-5.6-luna",
|
||||
client_api_format: "openai:search",
|
||||
provider_api_format: "openai:search",
|
||||
source_model: Some("gpt-5.6-luna-high-fast"),
|
||||
family: SameFormatProviderFamily::Standard,
|
||||
body_rules: None,
|
||||
request_headers: None,
|
||||
upstream_is_stream: false,
|
||||
force_body_stream_field: false,
|
||||
kiro_auth_config: None,
|
||||
is_claude_code: false,
|
||||
enable_model_directives: true,
|
||||
})
|
||||
.expect("search body should build");
|
||||
|
||||
assert_eq!(body["id"], "session-1");
|
||||
assert_eq!(body["model"], "gpt-5.6-luna");
|
||||
assert_eq!(body["commands"]["search_query"][0]["q"], "Aether");
|
||||
assert_eq!(body["settings"]["allowed_callers"][0], "direct");
|
||||
assert_eq!(body["reasoning"]["effort"], "high");
|
||||
assert!(body.get("store").is_none());
|
||||
assert!(body.get("future_extension").is_none());
|
||||
assert!(body.get("service_tier").is_none());
|
||||
assert!(body.get("stream").is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn builds_same_format_headers_with_auth_and_stream_accept() {
|
||||
let provider_request_body = json!({"model": "upstream-model"});
|
||||
|
||||
@@ -28,6 +28,14 @@ pub fn build_openai_responses_url(
|
||||
url
|
||||
}
|
||||
|
||||
pub fn build_openai_search_url(upstream_base_url: &str, query: Option<&str>) -> String {
|
||||
let (trimmed, base_query) = split_base_url_query(upstream_base_url);
|
||||
let trimmed = trimmed.trim_end_matches('/');
|
||||
let mut url = format!("{trimmed}/alpha/search");
|
||||
append_merged_query(&mut url, base_query, None, query, &[]);
|
||||
url
|
||||
}
|
||||
|
||||
pub fn build_openai_image_url(
|
||||
upstream_base_url: &str,
|
||||
request_path: Option<&str>,
|
||||
@@ -424,7 +432,7 @@ mod tests {
|
||||
build_bigmodel_coding_models_url, build_claude_messages_url, build_gemini_content_url,
|
||||
build_gemini_files_passthrough_url, build_gemini_video_predict_long_running_url,
|
||||
build_openai_chat_url, build_openai_compatible_models_url, build_openai_image_url,
|
||||
build_openai_responses_url, build_passthrough_path_url,
|
||||
build_openai_responses_url, build_openai_search_url, build_passthrough_path_url,
|
||||
normalize_gemini_content_action_path,
|
||||
};
|
||||
|
||||
@@ -573,6 +581,21 @@ mod tests {
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn openai_search_url_preserves_api_and_codex_roots() {
|
||||
assert_eq!(
|
||||
build_openai_search_url(
|
||||
"https://api.openai.com/v1?tenant=base",
|
||||
Some("trace=1&tenant=request")
|
||||
),
|
||||
"https://api.openai.com/v1/alpha/search?tenant=request&trace=1"
|
||||
);
|
||||
assert_eq!(
|
||||
build_openai_search_url("https://chatgpt.com/backend-api/codex/", None),
|
||||
"https://chatgpt.com/backend-api/codex/alpha/search"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn openai_image_url_uses_images_surface() {
|
||||
assert_eq!(
|
||||
|
||||
@@ -56,7 +56,7 @@ pub fn api_format_matches_allowed_value(allowed_value: &str, api_format: &str) -
|
||||
if allowed_value.is_empty() || api_format.is_empty() {
|
||||
return false;
|
||||
}
|
||||
crate::normalize_api_format(allowed_value) == crate::normalize_api_format(api_format)
|
||||
aether_ai_formats::api_format_permission_covers(allowed_value, api_format)
|
||||
}
|
||||
|
||||
pub fn auth_constraints_allow_model(
|
||||
@@ -248,6 +248,18 @@ mod tests {
|
||||
"openai:responses",
|
||||
"openai:responses"
|
||||
));
|
||||
assert!(api_format_matches_allowed_value(
|
||||
"openai:responses",
|
||||
"openai:search"
|
||||
));
|
||||
assert!(api_format_matches_allowed_value(
|
||||
"openai:search",
|
||||
"openai:search"
|
||||
));
|
||||
assert!(!api_format_matches_allowed_value(
|
||||
"openai:search",
|
||||
"openai:responses"
|
||||
));
|
||||
assert!(!api_format_matches_allowed_value(
|
||||
"openai:responses",
|
||||
"claude:messages"
|
||||
|
||||
@@ -279,7 +279,7 @@ fn mapping_scope_matches(
|
||||
let api_format_matches_scope = mapping.api_formats.as_ref().is_none_or(|api_formats| {
|
||||
api_formats
|
||||
.iter()
|
||||
.any(|value| api_format_matches(value, api_format))
|
||||
.any(|value| api_format_scope_covers(value, api_format))
|
||||
});
|
||||
if !api_format_matches_scope {
|
||||
return false;
|
||||
@@ -419,6 +419,10 @@ fn api_format_matches(left: &str, right: &str) -> bool {
|
||||
normalize_api_format(left) == normalize_api_format(right)
|
||||
}
|
||||
|
||||
fn api_format_scope_covers(allowed: &str, requested: &str) -> bool {
|
||||
aether_ai_formats::api_format_permission_covers(allowed, requested)
|
||||
}
|
||||
|
||||
fn requested_model_name_candidates(
|
||||
requested_model_name: &str,
|
||||
enable_model_directives: bool,
|
||||
@@ -568,6 +572,41 @@ mod tests {
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn responses_model_mapping_scope_covers_search_in_one_direction() {
|
||||
let mut row = sample_row("search-global", "search-default");
|
||||
row.endpoint_api_format = "openai:search".to_string();
|
||||
row.model_provider_model_mappings = Some(vec![StoredProviderModelMapping {
|
||||
name: "gpt-5.6-sol".to_string(),
|
||||
priority: 1,
|
||||
api_formats: Some(vec!["openai:responses".to_string()]),
|
||||
endpoint_ids: None,
|
||||
}]);
|
||||
|
||||
assert!(row_supports_requested_model(
|
||||
&row,
|
||||
"gpt-5.6-sol",
|
||||
"openai:search"
|
||||
));
|
||||
assert_eq!(
|
||||
resolve_provider_model_name(&row, "gpt-5.6-sol", "openai:search")
|
||||
.map(|resolved| resolved.0),
|
||||
Some("gpt-5.6-sol".to_string())
|
||||
);
|
||||
|
||||
row.model_provider_model_mappings = Some(vec![StoredProviderModelMapping {
|
||||
name: "search-only".to_string(),
|
||||
priority: 1,
|
||||
api_formats: Some(vec!["openai:search".to_string()]),
|
||||
endpoint_ids: None,
|
||||
}]);
|
||||
assert!(!row_supports_requested_model(
|
||||
&row,
|
||||
"search-only",
|
||||
"openai:responses"
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn model_directive_suffix_prefers_exact_model_before_base_fallback() {
|
||||
let exact = sample_row("gpt-5.4-high", "gpt-5.4-high-upstream");
|
||||
|
||||
@@ -109,6 +109,9 @@ pub fn infer_internal_finalize_signature(payload: &GatewaySyncReportRequest) ->
|
||||
if report_kind.starts_with("openai_image_") {
|
||||
return Some("openai:image".to_string());
|
||||
}
|
||||
if report_kind.starts_with("openai_search_") {
|
||||
return Some("openai:search".to_string());
|
||||
}
|
||||
if report_kind.starts_with("openai_cli_") {
|
||||
return Some("openai:responses".to_string());
|
||||
}
|
||||
@@ -150,6 +153,11 @@ pub fn resolve_internal_finalize_route(signature: &str) -> Option<InternalFinali
|
||||
route_family: "openai",
|
||||
route_kind: "responses:compact",
|
||||
}),
|
||||
"openai:search" => Some(InternalFinalizeRoute {
|
||||
public_path: "/v1/alpha/search",
|
||||
route_family: "openai",
|
||||
route_kind: "search",
|
||||
}),
|
||||
"openai:image" => Some(InternalFinalizeRoute {
|
||||
public_path: "/v1/images/generations",
|
||||
route_family: "openai",
|
||||
@@ -246,6 +254,7 @@ pub fn is_local_ai_sync_report_kind(report_kind: &str) -> bool {
|
||||
| "openai_responses_compact_sync_error"
|
||||
| "openai_cli_sync_success"
|
||||
| "openai_image_sync_success"
|
||||
| "openai_search_sync_success"
|
||||
| "openai_image_sync_error"
|
||||
| "openai_embedding_sync_success"
|
||||
| "openai_embedding_sync_error"
|
||||
@@ -872,6 +881,7 @@ mod tests {
|
||||
"openai_responses_compact_sync_error"
|
||||
));
|
||||
assert!(is_local_ai_sync_report_kind("openai_image_sync_success"));
|
||||
assert!(is_local_ai_sync_report_kind("openai_search_sync_success"));
|
||||
assert!(is_local_ai_sync_report_kind("openai_image_sync_error"));
|
||||
assert!(is_local_ai_sync_report_kind(
|
||||
"openai_embedding_sync_success"
|
||||
@@ -1082,12 +1092,27 @@ mod tests {
|
||||
Some("openai:responses:compact".to_string())
|
||||
);
|
||||
|
||||
let from_search_report_kind =
|
||||
sample_sync_report_with_context("openai_search_sync_finalize", json!({}));
|
||||
assert_eq!(
|
||||
infer_internal_finalize_signature(&from_search_report_kind),
|
||||
Some("openai:search".to_string())
|
||||
);
|
||||
|
||||
let unknown = sample_sync_report("unknown_sync_finalize", 200);
|
||||
assert_eq!(infer_internal_finalize_signature(&unknown), None);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn resolves_internal_finalize_route_for_supported_signatures() {
|
||||
assert_eq!(
|
||||
resolve_internal_finalize_route("openai:search"),
|
||||
Some(InternalFinalizeRoute {
|
||||
public_path: "/v1/alpha/search",
|
||||
route_family: "openai",
|
||||
route_kind: "search",
|
||||
})
|
||||
);
|
||||
assert_eq!(
|
||||
resolve_internal_finalize_route("openai:responses:compact"),
|
||||
Some(InternalFinalizeRoute {
|
||||
|
||||
@@ -1,6 +1,9 @@
|
||||
import { describe, expect, it } from 'vitest'
|
||||
|
||||
import { buildModelsDevTieredPricing } from '@/api/models-dev-pricing'
|
||||
import {
|
||||
buildModelsDevTieredPricing,
|
||||
resolveModelsDevTieredPricing,
|
||||
} from '@/api/models-dev-pricing'
|
||||
|
||||
describe('buildModelsDevTieredPricing', () => {
|
||||
it('maps context bands and cache prices without flattening them', () => {
|
||||
@@ -99,3 +102,75 @@ describe('buildModelsDevTieredPricing', () => {
|
||||
expect(buildModelsDevTieredPricing(cost)).toBeNull()
|
||||
})
|
||||
})
|
||||
|
||||
describe('resolveModelsDevTieredPricing', () => {
|
||||
it.each([
|
||||
{
|
||||
modelId: 'gpt-5.6-sol',
|
||||
standard: [5, 30, 6.25, 0.5],
|
||||
longContext: [10, 45, 12.5, 1],
|
||||
},
|
||||
{
|
||||
modelId: 'gpt-5.6-terra',
|
||||
standard: [2.5, 15, 3.125, 0.25],
|
||||
longContext: [5, 22.5, 6.25, 0.5],
|
||||
},
|
||||
{
|
||||
modelId: 'gpt-5.6-luna',
|
||||
standard: [1, 6, 1.25, 0.1],
|
||||
longContext: [2, 9, 2.5, 0.2],
|
||||
},
|
||||
])('uses the complete OpenAI catalog for $modelId', ({ modelId, standard, longContext }) => {
|
||||
const tier = (
|
||||
upTo: number | null,
|
||||
prices: number[],
|
||||
multiplier: number,
|
||||
) => ({
|
||||
up_to: upTo,
|
||||
input_price_per_1m: prices[0] * multiplier,
|
||||
output_price_per_1m: prices[1] * multiplier,
|
||||
cache_creation_price_per_1m: prices[2] * multiplier,
|
||||
cache_read_price_per_1m: prices[3] * multiplier,
|
||||
})
|
||||
|
||||
expect(resolveModelsDevTieredPricing('openai', modelId, { input: 999, output: 999 }))
|
||||
.toEqual({
|
||||
tiers: [
|
||||
tier(272_000, standard, 1),
|
||||
tier(null, longContext, 1),
|
||||
],
|
||||
processing_tiers: {
|
||||
flex: {
|
||||
tiers: [
|
||||
tier(272_000, standard, 0.5),
|
||||
tier(null, longContext, 0.5),
|
||||
],
|
||||
},
|
||||
priority: {
|
||||
tiers: [tier(272_000, standard, 2)],
|
||||
},
|
||||
},
|
||||
})
|
||||
})
|
||||
|
||||
it('keeps the models.dev lower-bound conversion for models outside the catalog', () => {
|
||||
expect(resolveModelsDevTieredPricing('openai', 'other-model', {
|
||||
input: 1,
|
||||
output: 2,
|
||||
tiers: [{ input: 3, output: 4, tier: { type: 'context', size: 272_000 } }],
|
||||
})?.tiers.map(tier => tier.up_to)).toEqual([271_999, null])
|
||||
})
|
||||
|
||||
it.each([
|
||||
['other-provider', 'gpt-5.6-sol'],
|
||||
['openai', 'GPT-5.6-SOL'],
|
||||
['openai', 'gpt-5.6-sol-latest'],
|
||||
['openai', '__proto__'],
|
||||
['openai', 'constructor'],
|
||||
])('matches provider and model identities exactly for %s/%s', (providerId, modelId) => {
|
||||
expect(resolveModelsDevTieredPricing(providerId, modelId, { input: 1, output: 2 }))
|
||||
.toEqual({
|
||||
tiers: [{ up_to: null, input_price_per_1m: 1, output_price_per_1m: 2 }],
|
||||
})
|
||||
})
|
||||
})
|
||||
|
||||
@@ -0,0 +1,71 @@
|
||||
import type { PricingTier, TieredPricingConfig } from './endpoints/types'
|
||||
|
||||
interface TokenPrices {
|
||||
input: number
|
||||
output: number
|
||||
cacheCreation: number
|
||||
cacheRead: number
|
||||
}
|
||||
|
||||
interface ContextPricing {
|
||||
standard: TokenPrices
|
||||
longContext: TokenPrices
|
||||
}
|
||||
|
||||
const OPENAI_GPT_56_PRICING = new Map<string, ContextPricing>([
|
||||
['gpt-5.6-sol', {
|
||||
standard: { input: 5, output: 30, cacheCreation: 6.25, cacheRead: 0.5 },
|
||||
longContext: { input: 10, output: 45, cacheCreation: 12.5, cacheRead: 1 },
|
||||
}],
|
||||
['gpt-5.6-terra', {
|
||||
standard: { input: 2.5, output: 15, cacheCreation: 3.125, cacheRead: 0.25 },
|
||||
longContext: { input: 5, output: 22.5, cacheCreation: 6.25, cacheRead: 0.5 },
|
||||
}],
|
||||
['gpt-5.6-luna', {
|
||||
standard: { input: 1, output: 6, cacheCreation: 1.25, cacheRead: 0.1 },
|
||||
longContext: { input: 2, output: 9, cacheCreation: 2.5, cacheRead: 0.2 },
|
||||
}],
|
||||
])
|
||||
|
||||
const STANDARD_CONTEXT_LIMIT = 272_000
|
||||
|
||||
function pricingTier(
|
||||
upTo: number | null,
|
||||
prices: TokenPrices,
|
||||
multiplier = 1,
|
||||
): PricingTier {
|
||||
return {
|
||||
up_to: upTo,
|
||||
input_price_per_1m: prices.input * multiplier,
|
||||
output_price_per_1m: prices.output * multiplier,
|
||||
cache_creation_price_per_1m: prices.cacheCreation * multiplier,
|
||||
cache_read_price_per_1m: prices.cacheRead * multiplier,
|
||||
}
|
||||
}
|
||||
|
||||
function contextPricingTiers(pricing: ContextPricing, multiplier: number): PricingTier[] {
|
||||
return [
|
||||
pricingTier(STANDARD_CONTEXT_LIMIT, pricing.standard, multiplier),
|
||||
pricingTier(null, pricing.longContext, multiplier),
|
||||
]
|
||||
}
|
||||
|
||||
export function getAuthoritativeModelPricing(
|
||||
providerId: string,
|
||||
modelId: string,
|
||||
): TieredPricingConfig | null {
|
||||
if (providerId !== 'openai') return null
|
||||
|
||||
const pricing = OPENAI_GPT_56_PRICING.get(modelId)
|
||||
if (!pricing) return null
|
||||
|
||||
return {
|
||||
tiers: contextPricingTiers(pricing, 1),
|
||||
processing_tiers: {
|
||||
flex: { tiers: contextPricingTiers(pricing, 0.5) },
|
||||
priority: {
|
||||
tiers: [pricingTier(STANDARD_CONTEXT_LIMIT, pricing.standard, 2)],
|
||||
},
|
||||
},
|
||||
}
|
||||
}
|
||||
@@ -2,6 +2,7 @@ import { describe, expect, it } from 'vitest'
|
||||
|
||||
import {
|
||||
API_FORMATS,
|
||||
apiFormatPermissionCovers,
|
||||
formatApiFormat,
|
||||
formatApiFormatShort,
|
||||
groupApiFormats,
|
||||
@@ -14,6 +15,8 @@ describe('api format display helpers', () => {
|
||||
expect(normalizeApiFormatAlias('CLAUDE_MESSAGES')).toBe(API_FORMATS.CLAUDE_MESSAGES)
|
||||
expect(normalizeApiFormatAlias('OPENAI_RESPONSES')).toBe(API_FORMATS.OPENAI_RESPONSES)
|
||||
expect(normalizeApiFormatAlias('OPENAI_RESPONSES_COMPACT')).toBe(API_FORMATS.OPENAI_RESPONSES_COMPACT)
|
||||
expect(normalizeApiFormatAlias('OPENAI_SEARCH')).toBe(API_FORMATS.OPENAI_SEARCH)
|
||||
expect(normalizeApiFormatAlias('SEARCH')).toBe(API_FORMATS.OPENAI_SEARCH)
|
||||
expect(normalizeApiFormatAlias('GEMINI_GENERATE_CONTENT')).toBe(API_FORMATS.GEMINI_GENERATE_CONTENT)
|
||||
expect(normalizeApiFormatAlias('OPENAI_EMBEDDING')).toBe(API_FORMATS.OPENAI_EMBEDDING)
|
||||
expect(normalizeApiFormatAlias('OPENAI_RERANK')).toBe(API_FORMATS.OPENAI_RERANK)
|
||||
@@ -33,6 +36,25 @@ describe('api format display helpers', () => {
|
||||
expect(formatApiFormatShort(API_FORMATS.JINA_RERANK)).toBe('JR')
|
||||
})
|
||||
|
||||
it('formats OpenAI Search as a first-class api format', () => {
|
||||
expect(formatApiFormat(API_FORMATS.OPENAI_SEARCH)).toBe('OpenAI Search')
|
||||
expect(formatApiFormatShort(API_FORMATS.OPENAI_SEARCH)).toBe('OS')
|
||||
expect(sortApiFormats([
|
||||
API_FORMATS.OPENAI_EMBEDDING,
|
||||
API_FORMATS.OPENAI_SEARCH,
|
||||
API_FORMATS.OPENAI_RESPONSES,
|
||||
])).toEqual([
|
||||
API_FORMATS.OPENAI_RESPONSES,
|
||||
API_FORMATS.OPENAI_SEARCH,
|
||||
API_FORMATS.OPENAI_EMBEDDING,
|
||||
])
|
||||
})
|
||||
|
||||
it('applies Responses to Search permissions in one direction', () => {
|
||||
expect(apiFormatPermissionCovers('OPENAI_RESPONSES', 'openai:search')).toBe(true)
|
||||
expect(apiFormatPermissionCovers('openai:search', 'openai:responses')).toBe(false)
|
||||
})
|
||||
|
||||
it('formats embedding api format ids distinctly from chat formats', () => {
|
||||
expect(formatApiFormat(API_FORMATS.GEMINI_INTERACTIONS)).toBe('Gemini Interactions')
|
||||
expect(formatApiFormatShort(API_FORMATS.GEMINI_INTERACTIONS)).toBe('GI')
|
||||
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user