From 2316df5c9a03dabcc83dea0ac98c97a071b15335 Mon Sep 17 00:00:00 2001 From: MMEXA Date: Sun, 12 Jul 2026 03:04:15 +0800 Subject: [PATCH] feat(codex): align Search and execution protocol --- .../src/ai_serving/planner/common.rs | 7 +- .../planner/decision/control_plan.rs | 11 +- .../src/ai_serving/planner/decision_input.rs | 23 +- .../passthrough/provider/family/candidates.rs | 8 +- .../passthrough/provider/family/request.rs | 1 + .../src/ai_serving/planner/standard/codex.rs | 29 +- .../planner/standard/codex/tests.rs | 44 +- .../aether-gateway/src/ai_serving/pure/mod.rs | 3 +- apps/aether-gateway/src/api/ai/openai.rs | 2 + apps/aether-gateway/src/api/ai/registry.rs | 2 + .../src/client_session_affinity.rs | 64 ++ apps/aether-gateway/src/constants.rs | 1 + .../src/control/auth/resolution.rs | 7 +- apps/aether-gateway/src/control/route/ai.rs | 8 + apps/aether-gateway/src/control/tests/ai.rs | 22 + apps/aether-gateway/src/data/state/auth.rs | 130 +++- .../src/execution_runtime/fallback.rs | 1 + .../src/execution_runtime/transport.rs | 43 +- .../src/frontdoor_loop_guard.rs | 1 + .../monitoring/cache_affinity.rs | 15 + .../handlers/admin/provider/crud/writes.rs | 12 +- .../admin/provider/query/models/model_test.rs | 169 ++++- .../query/models/model_test/adapter.rs | 5 + .../query/models/model_test/model_mapping.rs | 41 +- .../provider/query/models/model_test/tests.rs | 120 ++++ .../handlers/admin/provider/write/provider.rs | 70 +- .../admin/provider/write/provider/create.rs | 13 +- .../admin/provider/write/provider/template.rs | 252 ++++++-- .../admin/provider/write/provider/update.rs | 16 +- .../admin/request/provider/builders.rs | 66 ++ apps/aether-gateway/src/handlers/proxy/mod.rs | 191 +++++- .../handlers/public/support/user_me_usage.rs | 30 +- .../src/handlers/shared/catalog.rs | 21 +- apps/aether-gateway/src/maintenance/mod.rs | 4 +- .../aether-gateway/src/maintenance/runtime.rs | 5 + .../runtime/fixed_provider_reconciliation.rs | 301 +++++++++ .../src/maintenance/runtime/tests.rs | 11 +- .../src/orchestration/classifier.rs | 45 +- .../src/orchestration/effects.rs | 6 +- .../src/orchestration/policy.rs | 42 +- apps/aether-gateway/src/privacy/mod.rs | 181 +++++- apps/aether-gateway/src/state/app.rs | 1 + apps/aether-gateway/src/state/core.rs | 14 + apps/aether-gateway/src/task_runtime/mod.rs | 11 + .../src/tests/ai_execute/sync/mod.rs | 1 + .../src/tests/ai_execute/sync/search.rs | 597 ++++++++++++++++++ .../tests/control/admin/endpoints/routes.rs | 97 +++ .../src/tests/control/admin/oauth.rs | 3 + .../src/tests/control/admin/provider_query.rs | 182 ++++++ .../src/tests/control/admin/providers.rs | 77 ++- .../src/tests/control/internal.rs | 140 +++- .../aether-gateway/src/tests/frontdoor/ops.rs | 1 + apps/aether-gateway/src/tests/proxy.rs | 90 ++- .../src/tunnel/embedded/local_relay.rs | 83 ++- apps/aether-gateway/src/tunnel/mod.rs | 141 ++++- .../src/tunnel/stream_handler.rs | 51 +- .../aether-admin/src/observability/usage.rs | 36 +- crates/aether-admin/src/provider/endpoints.rs | 82 ++- crates/aether-admin/src/system.rs | 6 + crates/aether-ai-formats/src/api.rs | 9 +- crates/aether-ai-formats/src/contracts/mod.rs | 9 +- .../src/contracts/plan_kinds.rs | 1 + .../src/contracts/report_kinds.rs | 1 + crates/aether-ai-formats/src/formats/id.rs | 173 ++++- .../aether-ai-formats/src/formats/matrix.rs | 30 + crates/aether-ai-formats/src/formats/mod.rs | 5 +- .../src/formats/openai/mod.rs | 1 + .../src/formats/openai/reasoning.rs | 40 +- .../src/formats/openai/request_contract.rs | 92 ++- .../src/formats/openai/responses/codex.rs | 91 ++- .../src/formats/openai/search.rs | 219 +++++++ .../aether-ai-formats/src/formats/registry.rs | 8 +- .../src/formats/shared/error_body.rs | 21 +- .../src/formats/shared/model_directives.rs | 12 +- .../src/formats/shared/passthrough.rs | 18 +- .../src/formats/shared/request.rs | 20 +- .../src/formats/shared/routing.rs | 56 +- .../shared/stream_core/format_matrix.rs | 3 + crates/aether-ai-formats/src/lib.rs | 5 +- crates/aether-contracts/src/lib.rs | 4 +- crates/aether-contracts/src/plan.rs | 5 + crates/aether-contracts/src/tunnel.rs | 166 ++++- .../repository/candidate_selection/types.rs | 6 +- .../aether-data/src/repository/auth/types.rs | 41 +- .../repository/candidate_selection/memory.rs | 54 +- .../repository/candidate_selection/mysql.rs | 11 +- .../candidate_selection/postgres.rs | 32 +- .../repository/candidate_selection/sqlite.rs | 70 +- .../src/conversion.rs | 29 +- .../aether-provider-transport/src/network.rs | 15 + .../src/provider_types.rs | 14 +- .../src/request_url/mod.rs | 34 +- .../src/same_format_provider/mod.rs | 91 ++- crates/aether-provider-transport/src/url.rs | 25 +- crates/aether-scheduler-core/src/auth.rs | 14 +- crates/aether-scheduler-core/src/model.rs | 41 +- crates/aether-usage-runtime/src/report.rs | 25 + .../api/__tests__/models-dev-pricing.spec.ts | 77 ++- .../src/api/authoritative-model-pricing.ts | 71 +++ .../types/__tests__/api-format.spec.ts | 22 + .../src/api/endpoints/types/api-format.ts | 40 +- frontend/src/api/models-dev-pricing.ts | 10 + frontend/src/api/models-dev.ts | 9 +- .../components/EndpointFormDialog.vue | 21 +- .../components/PriorityManagementDialog.vue | 32 +- .../components/ProviderFormDialog.vue | 2 +- .../__tests__/endpoint-default-paths.spec.ts | 25 + .../endpoint-protocol-policy.spec.ts | 15 + .../components/endpoint-default-paths.ts | 5 + .../components/endpoint-protocol-policy.ts | 19 + .../provider-tabs/ModelMappingTab.vue | 17 +- .../__tests__/model-test-request.spec.ts | 75 ++- .../provider-tabs/model-test-capabilities.ts | 7 +- .../provider-tabs/model-test-request.ts | 54 +- .../composables/useEndpointStatus.ts | 24 +- .../composables/useProviderFilters.ts | 15 +- .../usage/utils/__tests__/status.spec.ts | 8 + frontend/src/features/usage/utils/status.ts | 1 + frontend/src/mocks/data.ts | 1 + .../__tests__/modelDirectivesConfig.spec.ts | 1 + .../modelDirectivesConfig.ts | 6 + 121 files changed, 5141 insertions(+), 543 deletions(-) create mode 100644 apps/aether-gateway/src/maintenance/runtime/fixed_provider_reconciliation.rs create mode 100644 apps/aether-gateway/src/tests/ai_execute/sync/search.rs create mode 100644 crates/aether-ai-formats/src/formats/openai/search.rs create mode 100644 frontend/src/api/authoritative-model-pricing.ts create mode 100644 frontend/src/features/providers/components/__tests__/endpoint-protocol-policy.spec.ts create mode 100644 frontend/src/features/providers/components/endpoint-protocol-policy.ts diff --git a/apps/aether-gateway/src/ai_serving/planner/common.rs b/apps/aether-gateway/src/ai_serving/planner/common.rs index 1fa228a2f..33565a583 100644 --- a/apps/aether-gateway/src/ai_serving/planner/common.rs +++ b/apps/aether-gateway/src/ai_serving/planner/common.rs @@ -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; diff --git a/apps/aether-gateway/src/ai_serving/planner/decision/control_plan.rs b/apps/aether-gateway/src/ai_serving/planner/decision/control_plan.rs index f681f8558..276df03c0 100644 --- a/apps/aether-gateway/src/ai_serving/planner/decision/control_plan.rs +++ b/apps/aether-gateway/src/ai_serving/planner/decision/control_plan.rs @@ -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)? } diff --git a/apps/aether-gateway/src/ai_serving/planner/decision_input.rs b/apps/aether-gateway/src/ai_serving/planner/decision_input.rs index 87ddf6862..5f5d1d7f8 100644 --- a/apps/aether-gateway/src/ai_serving/planner/decision_input.rs +++ b/apps/aether-gateway/src/ai_serving/planner/decision_input.rs @@ -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; diff --git a/apps/aether-gateway/src/ai_serving/planner/passthrough/provider/family/candidates.rs b/apps/aether-gateway/src/ai_serving/planner/passthrough/provider/family/candidates.rs index 33bfebb46..fb9d7a102 100644 --- a/apps/aether-gateway/src/ai_serving/planner/passthrough/provider/family/candidates.rs +++ b/apps/aether-gateway/src/ai_serving/planner/passthrough/provider/family/candidates.rs @@ -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, diff --git a/apps/aether-gateway/src/ai_serving/planner/passthrough/provider/family/request.rs b/apps/aether-gateway/src/ai_serving/planner/passthrough/provider/family/request.rs index eb09a76d3..a888f49fe 100644 --- a/apps/aether-gateway/src/ai_serving/planner/passthrough/provider/family/request.rs +++ b/apps/aether-gateway/src/ai_serving/planner/passthrough/provider/family/request.rs @@ -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", diff --git a/apps/aether-gateway/src/ai_serving/planner/standard/codex.rs b/apps/aether-gateway/src/ai_serving/planner/standard/codex.rs index c7aaccfbc..9bdc65adb 100644 --- a/apps/aether-gateway/src/ai_serving/planner/standard/codex.rs +++ b/apps/aether-gateway/src/ai_serving/planner/standard/codex.rs @@ -12,20 +12,33 @@ pub(crate) fn codex_model_capabilities_for_transport( provider_model: &str, source_model: &str, ) -> Option { - 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 { + 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, ), ) } diff --git a/apps/aether-gateway/src/ai_serving/planner/standard/codex/tests.rs b/apps/aether-gateway/src/ai_serving/planner/standard/codex/tests.rs index 75b9c0521..10bef0e6b 100644 --- a/apps/aether-gateway/src/ai_serving/planner/standard/codex/tests.rs +++ b/apps/aether-gateway/src/ai_serving/planner/standard/codex/tests.rs @@ -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!({ diff --git a/apps/aether-gateway/src/ai_serving/pure/mod.rs b/apps/aether-gateway/src/ai_serving/pure/mod.rs index a5434680e..985ab0c23 100644 --- a/apps/aether-gateway/src/ai_serving/pure/mod.rs +++ b/apps/aether-gateway/src/ai_serving/pure/mod.rs @@ -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, diff --git a/apps/aether-gateway/src/api/ai/openai.rs b/apps/aether-gateway/src/api/ai/openai.rs index 999c356c6..85d9cbe2e 100644 --- a/apps/aether-gateway/src/api/ai/openai.rs +++ b/apps/aether-gateway/src/api/ai/openai.rs @@ -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, diff --git a/apps/aether-gateway/src/api/ai/registry.rs b/apps/aether-gateway/src/api/ai/registry.rs index 5d92375c5..7ae1df4e6 100644 --- a/apps/aether-gateway/src/api/ai/registry.rs +++ b/apps/aether-gateway/src/api/ai/registry.rs @@ -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!( diff --git a/apps/aether-gateway/src/client_session_affinity.rs b/apps/aether-gateway/src/client_session_affinity.rs index 832554e16..66e009184 100644 --- a/apps/aether-gateway/src/client_session_affinity.rs +++ b/apps/aether-gateway/src/client_session_affinity.rs @@ -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 { + 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 { + 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 { + 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()); + } } diff --git a/apps/aether-gateway/src/constants.rs b/apps/aether-gateway/src/constants.rs index 8052f2001..44e2f4788 100644 --- a/apps/aether-gateway/src/constants.rs +++ b/apps/aether-gateway/src/constants.rs @@ -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", diff --git a/apps/aether-gateway/src/control/auth/resolution.rs b/apps/aether-gateway/src/control/auth/resolution.rs index 9e250124f..1515b09e4 100644 --- a/apps/aether-gateway/src/control/auth/resolution.rs +++ b/apps/aether-gateway/src/control/auth/resolution.rs @@ -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, } } diff --git a/apps/aether-gateway/src/control/route/ai.rs b/apps/aether-gateway/src/control/route/ai.rs index ec8ba8d96..65a4781ef 100644 --- a/apps/aether-gateway/src/control/route/ai.rs +++ b/apps/aether-gateway/src/control/route/ai.rs @@ -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, diff --git a/apps/aether-gateway/src/control/tests/ai.rs b/apps/aether-gateway/src/control/tests/ai.rs index 87c0097b5..c645cb54c 100644 --- a/apps/aether-gateway/src/control/tests/ai.rs +++ b/apps/aether-gateway/src/control/tests/ai.rs @@ -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")]); diff --git a/apps/aether-gateway/src/data/state/auth.rs b/apps/aether-gateway/src/data/state/auth.rs index b80148708..466b3b519 100644 --- a/apps/aether-gateway/src/data/state/auth.rs +++ b/apps/aether-gateway/src/data/state/auth.rs @@ -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>, + user_mode: &str, + groups: &[aether_data::repository::users::StoredUserGroup], + group_field: impl Fn( + &aether_data::repository::users::StoredUserGroup, + ) -> (&str, Option>), +) -> Option> { + 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>, + right: Option>, +) -> Option> { + 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, right: Option, @@ -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>, + api_key_policy: &mut Option>, +) { + 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( diff --git a/apps/aether-gateway/src/execution_runtime/fallback.rs b/apps/aether-gateway/src/execution_runtime/fallback.rs index a50e6ef30..bc4b811de 100644 --- a/apps/aether-gateway/src/execution_runtime/fallback.rs +++ b/apps/aether-gateway/src/execution_runtime/fallback.rs @@ -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, } ); } diff --git a/apps/aether-gateway/src/execution_runtime/transport.rs b/apps/aether-gateway/src/execution_runtime/transport.rs index c13df853c..231e33923 100644 --- a/apps/aether-gateway/src/execution_runtime/transport.rs +++ b/apps/aether-gateway/src/execution_runtime/transport.rs @@ -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 { }) } -fn resolve_non_stream_total_timeout(plan: &ExecutionPlan) -> Option { - 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 { + 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 { - if !plan.stream { +fn resolve_non_stream_total_timeout(plan: &ExecutionPlan) -> Option { + 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 { + 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 { + resolve_stream_first_byte_timeout_for_request(plan.stream, plan.timeouts.as_ref()) +} + pub(crate) async fn with_non_stream_total_timeout( 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 { diff --git a/apps/aether-gateway/src/frontdoor_loop_guard.rs b/apps/aether-gateway/src/frontdoor_loop_guard.rs index 2b8f0f6e2..827faf7f9 100644 --- a/apps/aether-gateway/src/frontdoor_loop_guard.rs +++ b/apps/aether-gateway/src/frontdoor_loop_guard.rs @@ -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" diff --git a/apps/aether-gateway/src/handlers/admin/observability/monitoring/cache_affinity.rs b/apps/aether-gateway/src/handlers/admin/observability/monitoring/cache_affinity.rs index a1a569ce8..47c1d19bf 100644 --- a/apps/aether-gateway/src/handlers/admin/observability/monitoring/cache_affinity.rs +++ b/apps/aether-gateway/src/handlers/admin/observability/monitoring/cache_affinity.rs @@ -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( diff --git a/apps/aether-gateway/src/handlers/admin/provider/crud/writes.rs b/apps/aether-gateway/src/handlers/admin/provider/crud/writes.rs index cbf44424d..fa4b96728 100644 --- a/apps/aether-gateway/src/handlers/admin/provider/crud/writes.rs +++ b/apps/aether-gateway/src/handlers/admin/provider/crud/writes.rs @@ -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 diff --git a/apps/aether-gateway/src/handlers/admin/provider/query/models/model_test.rs b/apps/aether-gateway/src/handlers/admin/provider/query/models/model_test.rs index cf3455edf..bf357cbe9 100644 --- a/apps/aether-gateway/src/handlers/admin/provider/query/models/model_test.rs +++ b/apps/aether-gateway/src/handlers/admin/provider/query/models/model_test.rs @@ -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, + 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, diff --git a/apps/aether-gateway/src/handlers/admin/provider/query/models/model_test/adapter.rs b/apps/aether-gateway/src/handlers/admin/provider/query/models/model_test/adapter.rs index 86333dc64..9bdd7802a 100644 --- a/apps/aether-gateway/src/handlers/admin/provider/query/models/model_test/adapter.rs +++ b/apps/aether-gateway/src/handlers/admin/provider/query/models/model_test/adapter.rs @@ -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" diff --git a/apps/aether-gateway/src/handlers/admin/provider/query/models/model_test/model_mapping.rs b/apps/aether-gateway/src/handlers/admin/provider/query/models/model_test/model_mapping.rs index 65bcdfa19..b7ac245e5 100644 --- a/apps/aether-gateway/src/handlers/admin/provider/query/models/model_test/model_mapping.rs +++ b/apps/aether-gateway/src/handlers/admin/provider/query/models/model_test/model_mapping.rs @@ -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"), + )); + } +} diff --git a/apps/aether-gateway/src/handlers/admin/provider/query/models/model_test/tests.rs b/apps/aether-gateway/src/handlers/admin/provider/query/models/model_test/tests.rs index 80f1f75de..dd7285a5f 100644 --- a/apps/aether-gateway/src/handlers/admin/provider/query/models/model_test/tests.rs +++ b/apps/aether-gateway/src/handlers/admin/provider/query/models/model_test/tests.rs @@ -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) diff --git a/apps/aether-gateway/src/handlers/admin/provider/write/provider.rs b/apps/aether-gateway/src/handlers/admin/provider/write/provider.rs index a9a727dbd..c8cea50e2 100644 --- a/apps/aether-gateway/src/handlers/admin/provider/write/provider.rs +++ b/apps/aether-gateway/src/handlers/admin/provider/write/provider.rs @@ -3,11 +3,79 @@ mod endpoint; mod template; mod update; +fn normalize_provider_request_timeout(value: Option) -> Result, 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) -> Result, 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()); + } +} diff --git a/apps/aether-gateway/src/handlers/admin/provider/write/provider/create.rs b/apps/aether-gateway/src/handlers/admin/provider/write/provider/create.rs index b83c7638d..25a63e891 100644 --- a/apps/aether-gateway/src/handlers/admin/provider/write/provider/create.rs +++ b/apps/aether-gateway/src/handlers/admin/provider/write/provider/create.rs @@ -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()) diff --git a/apps/aether-gateway/src/handlers/admin/provider/write/provider/template.rs b/apps/aether-gateway/src/handlers/admin/provider/write/provider/template.rs index 7fa260e11..dc432355d 100644 --- a/apps/aether-gateway/src/handlers/admin/provider/write/provider/template.rs +++ b/apps/aether-gateway/src/handlers/admin/provider/write/provider/template.rs @@ -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 { + 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 { 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 { - 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(overrides: &mut BTreeSet, 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); + } } diff --git a/apps/aether-gateway/src/handlers/admin/provider/write/provider/update.rs b/apps/aether-gateway/src/handlers/admin/provider/write/provider/update.rs index cd3ec295e..3c9bc1272 100644 --- a/apps/aether-gateway/src/handlers/admin/provider/write/provider/update.rs +++ b/apps/aether-gateway/src/handlers/admin/provider/write/provider/update.rs @@ -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") { diff --git a/apps/aether-gateway/src/handlers/admin/request/provider/builders.rs b/apps/aether-gateway/src/handlers/admin/request/provider/builders.rs index 0f31ca7bb..5690d1c28 100644 --- a/apps/aether-gateway/src/handlers/admin/request/provider/builders.rs +++ b/apps/aether-gateway/src/handlers/admin/request/provider/builders.rs @@ -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()); + } +} diff --git a/apps/aether-gateway/src/handlers/proxy/mod.rs b/apps/aether-gateway/src/handlers/proxy/mod.rs index 648991fd1..602d670e1 100644 --- a/apps/aether-gateway/src/handlers/proxy/mod.rs +++ b/apps/aether-gateway/src/handlers/proxy/mod.rs @@ -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::(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"); diff --git a/apps/aether-gateway/src/handlers/public/support/user_me_usage.rs b/apps/aether-gateway/src/handlers/public/support/user_me_usage.rs index 29cf23940..79eb24bd5 100644 --- a/apps/aether-gateway/src/handlers/public/support/user_me_usage.rs +++ b/apps/aether-gateway/src/handlers/public/support/user_me_usage.rs @@ -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 { diff --git a/apps/aether-gateway/src/handlers/shared/catalog.rs b/apps/aether-gateway/src/handlers/shared/catalog.rs index d34e14848..e558bc0e0 100644 --- a/apps/aether-gateway/src/handlers/shared/catalog.rs +++ b/apps/aether-gateway/src/handlers/shared/catalog.rs @@ -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"); diff --git a/apps/aether-gateway/src/maintenance/mod.rs b/apps/aether-gateway/src/maintenance/mod.rs index 98facca26..5cb50439b 100644 --- a/apps/aether-gateway/src/maintenance/mod.rs +++ b/apps/aether-gateway/src/maintenance/mod.rs @@ -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, diff --git a/apps/aether-gateway/src/maintenance/runtime.rs b/apps/aether-gateway/src/maintenance/runtime.rs index f2dfd9c0b..b513c2a82 100644 --- a/apps/aether-gateway/src/maintenance/runtime.rs +++ b/apps/aether-gateway/src/maintenance/runtime.rs @@ -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, }; diff --git a/apps/aether-gateway/src/maintenance/runtime/fixed_provider_reconciliation.rs b/apps/aether-gateway/src/maintenance/runtime/fixed_provider_reconciliation.rs new file mode 100644 index 000000000..9ebf13544 --- /dev/null +++ b/apps/aether-gateway/src/maintenance/runtime/fixed_provider_reconciliation.rs @@ -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 { + 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> { + 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); + } +} diff --git a/apps/aether-gateway/src/maintenance/runtime/tests.rs b/apps/aether-gateway/src/maintenance/runtime/tests.rs index 673a5d57d..70ee7e514 100644 --- a/apps/aether-gateway/src/maintenance/runtime/tests.rs +++ b/apps/aether-gateway/src/maintenance/runtime/tests.rs @@ -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![])); diff --git a/apps/aether-gateway/src/orchestration/classifier.rs b/apps/aether-gateway/src/orchestration/classifier.rs index ffce19805..cafffafc6 100644 --- a/apps/aether-gateway/src/orchestration/classifier.rs +++ b/apps/aether-gateway/src/orchestration/classifier.rs @@ -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!( diff --git a/apps/aether-gateway/src/orchestration/effects.rs b/apps/aether-gateway/src/orchestration/effects.rs index b1c81b7d2..241649344 100644 --- a/apps/aether-gateway/src/orchestration/effects.rs +++ b/apps/aether-gateway/src/orchestration/effects.rs @@ -243,7 +243,11 @@ fn local_client_session_affinity(report_context: Option<&Value>) -> Option) -> http::HeaderMap { diff --git a/apps/aether-gateway/src/orchestration/policy.rs b/apps/aether-gateway/src/orchestration/policy.rs index 6ae4028dd..6e0a05adc 100644 --- a/apps/aether-gateway/src/orchestration/policy.rs +++ b/apps/aether-gateway/src/orchestration/policy.rs @@ -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, pub(crate) stop_status_codes: BTreeSet, @@ -15,6 +15,21 @@ pub(crate) struct LocalFailoverPolicy { pub(crate) success_failover_patterns: Vec, pub(crate) error_stop_patterns: Vec, 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::>(), "error_stop_patterns": policy.error_stop_patterns.iter().map(local_failover_regex_rule_to_value).collect::>(), "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); diff --git a/apps/aether-gateway/src/privacy/mod.rs b/apps/aether-gateway/src/privacy/mod.rs index 42c2caad4..f2c91e03a 100644 --- a/apps/aether-gateway/src/privacy/mod.rs +++ b/apps/aether-gateway/src/privacy/mod.rs @@ -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 { 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 { + 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) { 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 { + 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 { + 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); diff --git a/apps/aether-gateway/src/state/app.rs b/apps/aether-gateway/src/state/app.rs index ca1682f80..ee91b9207 100644 --- a/apps/aether-gateway/src/state/app.rs +++ b/apps/aether-gateway/src/state/app.rs @@ -344,6 +344,7 @@ pub struct AppState { pub(crate) upstream_target_admission: Arc, pub(crate) distributed_request_gate: Option>, pub(crate) client: reqwest::Client, + pub(crate) owner_forward_client: reqwest::Client, pub(crate) auth_context_cache: Arc, pub(crate) auth_snapshot_cache: Arc, pub(crate) user_model_capability_settings_cache: Arc>, diff --git a/apps/aether-gateway/src/state/core.rs b/apps/aether-gateway/src/state/core.rs index c3bb82b36..5a7c5f9aa 100644 --- a/apps/aether-gateway/src/state/core.rs +++ b/apps/aether-gateway/src/state/core.rs @@ -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>| { if let Some(handle) = handle { diff --git a/apps/aether-gateway/src/task_runtime/mod.rs b/apps/aether-gateway/src/task_runtime/mod.rs index 78a2b76a6..61e8cdde8 100644 --- a/apps/aether-gateway/src/task_runtime/mod.rs +++ b/apps/aether-gateway/src/task_runtime/mod.rs @@ -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, diff --git a/apps/aether-gateway/src/tests/ai_execute/sync/mod.rs b/apps/aether-gateway/src/tests/ai_execute/sync/mod.rs index f3142119b..d93e4e742 100644 --- a/apps/aether-gateway/src/tests/ai_execute/sync/mod.rs +++ b/apps/aether-gateway/src/tests/ai_execute/sync/mod.rs @@ -29,3 +29,4 @@ mod cli; mod gemini; mod image; mod pii_redaction_formats; +mod search; diff --git a/apps/aether-gateway/src/tests/ai_execute/sync/search.rs b/apps/aether-gateway/src/tests/ai_execute/sync/search.rs new file mode 100644 index 000000000..eb1de4c43 --- /dev/null +++ b/apps/aether-gateway/src/tests/ai_execute/sync/search.rs @@ -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(test_name: &'static str, make_future: F) +where + F: FnOnce() -> Fut + Send + 'static, + Fut: std::future::Future + '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::::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::() + .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::() + .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::>(); + 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(); +} diff --git a/apps/aether-gateway/src/tests/control/admin/endpoints/routes.rs b/apps/aether-gateway/src/tests/control/admin/endpoints/routes.rs index 39e6e9efb..83e4c9517 100644 --- a/apps/aether-gateway/src/tests/control/admin/endpoints/routes.rs +++ b/apps/aether-gateway/src/tests/control/admin/endpoints/routes.rs @@ -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)); diff --git a/apps/aether-gateway/src/tests/control/admin/oauth.rs b/apps/aether-gateway/src/tests/control/admin/oauth.rs index cee2f189a..27afe8625 100644 --- a/apps/aether-gateway/src/tests/control/admin/oauth.rs +++ b/apps/aether-gateway/src/tests/control/admin/oauth.rs @@ -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" diff --git a/apps/aether-gateway/src/tests/control/admin/provider_query.rs b/apps/aether-gateway/src/tests/control/admin/provider_query.rs index 2060eeff7..51f435991 100644 --- a/apps/aether-gateway/src/tests/control/admin/provider_query.rs +++ b/apps/aether-gateway/src/tests/control/admin/provider_query.rs @@ -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| 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( diff --git a/apps/aether-gateway/src/tests/control/admin/providers.rs b/apps/aether-gateway/src/tests/control/admin/providers.rs index fc019cc43..185a66774 100644 --- a/apps/aether-gateway/src/tests/control/admin/providers.rs +++ b/apps/aether-gateway/src/tests/control/admin/providers.rs @@ -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(); diff --git a/apps/aether-gateway/src/tests/control/internal.rs b/apps/aether-gateway/src/tests/control/internal.rs index 9a74ec153..efa7a5934 100644 --- a/apps/aether-gateway/src/tests/control/internal.rs +++ b/apps/aether-gateway/src/tests/control/internal.rs @@ -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::copy_from_slice(&envelope[..split_at])), + Ok::(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::from_static(b"relay-")), - Ok::(Bytes::from_static(b"stream-")), - Ok::(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::( + 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); diff --git a/apps/aether-gateway/src/tests/frontdoor/ops.rs b/apps/aether-gateway/src/tests/frontdoor/ops.rs index cac36fd9b..e072e16f9 100644 --- a/apps/aether-gateway/src/tests/frontdoor/ops.rs +++ b/apps/aether-gateway/src/tests/frontdoor/ops.rs @@ -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")); diff --git a/apps/aether-gateway/src/tests/proxy.rs b/apps/aether-gateway/src/tests/proxy.rs index 965c242fe..e4c646bb0 100644 --- a/apps/aether-gateway/src/tests/proxy.rs +++ b/apps/aether-gateway/src/tests/proxy.rs @@ -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, +) -> 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 { diff --git a/apps/aether-gateway/src/tunnel/embedded/local_relay.rs b/apps/aether-gateway/src/tunnel/embedded/local_relay.rs index a7767fbfc..f599daa34 100644 --- a/apps/aether-gateway/src/tunnel/embedded/local_relay.rs +++ b/apps/aether-gateway/src/tunnel/embedded/local_relay.rs @@ -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, @@ -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, 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::(&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(), diff --git a/apps/aether-gateway/src/tunnel/mod.rs b/apps/aether-gateway/src/tunnel/mod.rs index 8845db8be..0b517e78a 100644 --- a/apps/aether-gateway/src/tunnel/mod.rs +++ b/apps/aether-gateway/src/tunnel/mod.rs @@ -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, +) -> Result { + 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, -) -> reqwest::Body { +) -> Result { 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::(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::(chunk); } - Err(err) => { - yield Err::(io::Error::other(err)); + Err(error) => { + yield Err::(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( diff --git a/apps/aether-tunnel/src/tunnel/stream_handler.rs b/apps/aether-tunnel/src/tunnel/stream_handler.rs index c626ef948..466cf6ab3 100644 --- a/apps/aether-tunnel/src/tunnel/stream_handler.rs +++ b/apps/aether-tunnel/src/tunnel/stream_handler.rs @@ -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 { } 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, @@ -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"); diff --git a/crates/aether-admin/src/observability/usage.rs b/crates/aether-admin/src/observability/usage.rs index 7ed5cd8c8..5ac2f2cb0 100644 --- a/crates/aether-admin/src/observability/usage.rs +++ b/crates/aether-admin/src/observability/usage.rs @@ -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 { diff --git a/crates/aether-admin/src/provider/endpoints.rs b/crates/aether-admin/src/provider/endpoints.rs index 856f4f499..35616dbf2 100644 --- a/crates/aether-admin/src/provider/endpoints.rs +++ b/crates/aether-admin/src/provider/endpoints.rs @@ -82,33 +82,89 @@ pub fn endpoint_key_counts_by_format( ) -> (BTreeMap, BTreeMap) { 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; diff --git a/crates/aether-admin/src/system.rs b/crates/aether-admin/src/system.rs index 0f1ae648d..020366a9d 100644 --- a/crates/aether-admin/src/system.rs +++ b/crates/aether-admin/src/system.rs @@ -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", diff --git a/crates/aether-ai-formats/src/api.rs b/crates/aether-ai-formats/src/api.rs index ef721721e..889b7ac45 100644 --- a/crates/aether-ai-formats/src/api.rs +++ b/crates/aether-ai-formats/src/api.rs @@ -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, diff --git a/crates/aether-ai-formats/src/contracts/mod.rs b/crates/aether-ai-formats/src/contracts/mod.rs index 3f7752f66..62f11b066 100644 --- a/crates/aether-ai-formats/src/contracts/mod.rs +++ b/crates/aether-ai-formats/src/contracts/mod.rs @@ -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, }; diff --git a/crates/aether-ai-formats/src/contracts/plan_kinds.rs b/crates/aether-ai-formats/src/contracts/plan_kinds.rs index d0d4eecb7..2a9e61329 100644 --- a/crates/aether-ai-formats/src/contracts/plan_kinds.rs +++ b/crates/aether-ai-formats/src/contracts/plan_kinds.rs @@ -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"; diff --git a/crates/aether-ai-formats/src/contracts/report_kinds.rs b/crates/aether-ai-formats/src/contracts/report_kinds.rs index bb2e769e0..ce19b1b4f 100644 --- a/crates/aether-ai-formats/src/contracts/report_kinds.rs +++ b/crates/aether-ai-formats/src/contracts/report_kinds.rs @@ -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"; diff --git a/crates/aether-ai-formats/src/formats/id.rs b/crates/aether-ai-formats/src/formats/id.rs index 1e064f928..fc83ffdc4 100644 --- a/crates/aether-ai-formats/src/formats/id.rs +++ b/crates/aether-ai-formats/src/formats/id.rs @@ -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 { + 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 { match FormatId::parse(value).map(FormatId::canonical) { Some(FormatId::AliyunMultimodalEmbedding) => vec![ @@ -149,6 +205,22 @@ pub fn api_format_storage_aliases(value: &str) -> Vec { } } +pub fn api_format_permission_storage_aliases(value: &str) -> Vec { + 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!( diff --git a/crates/aether-ai-formats/src/formats/matrix.rs b/crates/aether-ai-formats/src/formats/matrix.rs index be66d4833..0c2568e54 100644 --- a/crates/aether-ai-formats/src/formats/matrix.rs +++ b/crates/aether-ai-formats/src/formats/matrix.rs @@ -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!( diff --git a/crates/aether-ai-formats/src/formats/mod.rs b/crates/aether-ai-formats/src/formats/mod.rs index 6dd9a9bf1..a50a036cd 100644 --- a/crates/aether-ai-formats/src/formats/mod.rs +++ b/crates/aether-ai-formats/src/formats/mod.rs @@ -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, }; diff --git a/crates/aether-ai-formats/src/formats/openai/mod.rs b/crates/aether-ai-formats/src/formats/openai/mod.rs index a603b429e..ba6e87780 100644 --- a/crates/aether-ai-formats/src/formats/openai/mod.rs +++ b/crates/aether-ai-formats/src/formats/openai/mod.rs @@ -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; diff --git a/crates/aether-ai-formats/src/formats/openai/reasoning.rs b/crates/aether-ai-formats/src/formats/openai/reasoning.rs index 006f21604..e0632bdb5 100644 --- a/crates/aether-ai-formats/src/formats/openai/reasoning.rs +++ b/crates/aether-ai-formats/src/formats/openai/reasoning.rs @@ -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")) diff --git a/crates/aether-ai-formats/src/formats/openai/request_contract.rs b/crates/aether-ai-formats/src/formats/openai/request_contract.rs index 0eccc2fbf..722c4aa13 100644 --- a/crates/aether-ai-formats/src/formats/openai/request_contract.rs +++ b/crates/aether-ai-formats/src/formats/openai/request_contract.rs @@ -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 { diff --git a/crates/aether-ai-formats/src/formats/openai/responses/codex.rs b/crates/aether-ai-formats/src/formats/openai/responses/codex.rs index cd00f63dd..116b978a7 100644 --- a/crates/aether-ai-formats/src/formats/openai/responses/codex.rs +++ b/crates/aether-ai-formats/src/formats/openai/responses/codex.rs @@ -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, 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, @@ -1179,6 +1192,23 @@ fn normalize_codex_reasoning_effort(body_object: &mut serde_json::Map 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([( diff --git a/crates/aether-ai-formats/src/formats/openai/search.rs b/crates/aether-ai-formats/src/formats/openai/search.rs new file mode 100644 index 000000000..1553acbf0 --- /dev/null +++ b/crates/aether-ai-formats/src/formats/openai/search.rs @@ -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, + 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); + } +} diff --git a/crates/aether-ai-formats/src/formats/registry.rs b/crates/aether-ai-formats/src/formats/registry.rs index fcc052132..b56244536 100644 --- a/crates/aether-ai-formats/src/formats/registry.rs +++ b/crates/aether-ai-formats/src/formats/registry.rs @@ -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 diff --git a/crates/aether-ai-formats/src/formats/shared/error_body.rs b/crates/aether-ai-formats/src/formats/shared/error_body.rs index 53ace4d28..a00f053af 100644 --- a/crates/aether-ai-formats/src/formats/shared/error_body.rs +++ b/crates/aether-ai-formats/src/formats/shared/error_body.rs @@ -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")); diff --git a/crates/aether-ai-formats/src/formats/shared/model_directives.rs b/crates/aether-ai-formats/src/formats/shared/model_directives.rs index 215b32b28..b97501c62 100644 --- a/crates/aether-ai-formats/src/formats/shared/model_directives.rs +++ b/crates/aether-ai-formats/src/formats/shared/model_directives.rs @@ -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, diff --git a/crates/aether-ai-formats/src/formats/shared/passthrough.rs b/crates/aether-ai-formats/src/formats/shared/passthrough.rs index a0fef4ed8..4c1fd03ce 100644 --- a/crates/aether-ai-formats/src/formats/shared/passthrough.rs +++ b/crates/aether-ai-formats/src/formats/shared/passthrough.rs @@ -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 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); + } } diff --git a/crates/aether-ai-formats/src/formats/shared/request.rs b/crates/aether-ai-formats/src/formats/shared/request.rs index d7cafc9e6..e4c81f6a8 100644 --- a/crates/aether-ai-formats/src/formats/shared/request.rs +++ b/crates/aether-ai-formats/src/formats/shared/request.rs @@ -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", diff --git a/crates/aether-ai-formats/src/formats/shared/routing.rs b/crates/aether-ai-formats/src/formats/shared/routing.rs index 6fd256edf..fab37faa0 100644 --- a/crates/aether-ai-formats/src/formats/shared/routing.rs +++ b/crates/aether-ai-formats/src/formats/shared/routing.rs @@ -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!( diff --git a/crates/aether-ai-formats/src/formats/shared/stream_core/format_matrix.rs b/crates/aether-ai-formats/src/formats/shared/stream_core/format_matrix.rs index 5e441220b..dad57cf81 100644 --- a/crates/aether-ai-formats/src/formats/shared/stream_core/format_matrix.rs +++ b/crates/aether-ai-formats/src/formats/shared/stream_core/format_matrix.rs @@ -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 diff --git a/crates/aether-ai-formats/src/lib.rs b/crates/aether-ai-formats/src/lib.rs index f6e2d140f..b5ce22cf9 100644 --- a/crates/aether-ai-formats/src/lib.rs +++ b/crates/aether-ai-formats/src/lib.rs @@ -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, }; diff --git a/crates/aether-contracts/src/lib.rs b/crates/aether-contracts/src/lib.rs index a84b35e04..c450eac37 100644 --- a/crates/aether-contracts/src/lib.rs +++ b/crates/aether-contracts/src/lib.rs @@ -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, diff --git a/crates/aether-contracts/src/plan.rs b/crates/aether-contracts/src/plan.rs index 2652013f9..3609fb561 100644 --- a/crates/aether-contracts/src/plan.rs +++ b/crates/aether-contracts/src/plan.rs @@ -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)] diff --git a/crates/aether-contracts/src/tunnel.rs b/crates/aether-contracts/src/tunnel.rs index 0e8033432..97528adbc 100644 --- a/crates/aether-contracts/src/tunnel.rs +++ b/crates/aether-contracts/src/tunnel.rs @@ -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, } +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct ResolvedTunnelRequestTimeouts { + pub first_byte_ms: u64, + pub response_body_ms: Option, +} + +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, 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::(&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 { 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}"#; diff --git a/crates/aether-data-contracts/src/repository/candidate_selection/types.rs b/crates/aether-data-contracts/src/repository/candidate_selection/types.rs index cc153ffda..309773285 100644 --- a/crates/aether-data-contracts/src/repository/candidate_selection/types.rs +++ b/crates/aether-data-contracts/src/repository/candidate_selection/types.rs @@ -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] diff --git a/crates/aether-data/src/repository/auth/types.rs b/crates/aether-data/src/repository/auth/types.rs index c12e3a1e2..f3b7039b4 100644 --- a/crates/aether-data/src/repository/auth/types.rs +++ b/crates/aether-data/src/repository/auth/types.rs @@ -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>, + api_key_policy: &mut Option>, +) { + 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 { let right_values = right.iter().collect::>(); 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( diff --git a/crates/aether-data/src/repository/candidate_selection/memory.rs b/crates/aether-data/src/repository/candidate_selection/memory.rs index 044847f55..4b136bba4 100644 --- a/crates/aether-data/src/repository/candidate_selection/memory.rs +++ b/crates/aether-data/src/repository/candidate_selection/memory.rs @@ -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( diff --git a/crates/aether-data/src/repository/candidate_selection/mysql.rs b/crates/aether-data/src/repository/candidate_selection/mysql.rs index 691bc1ce4..9ae01cb74 100644 --- a/crates/aether-data/src/repository/candidate_selection/mysql.rs +++ b/crates/aether-data/src/repository/candidate_selection/mysql.rs @@ -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 { api_formats .iter() diff --git a/crates/aether-data/src/repository/candidate_selection/postgres.rs b/crates/aether-data/src/repository/candidate_selection/postgres.rs index 227214e16..ccf7d5835 100644 --- a/crates/aether-data/src/repository/candidate_selection/postgres.rs +++ b/crates/aether-data/src/repository/candidate_selection/postgres.rs @@ -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 { aether_ai_formats::api_format_storage_aliases(api_format) } +fn api_format_permission_aliases(api_format: &str) -> Vec { + 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) } diff --git a/crates/aether-data/src/repository/candidate_selection/sqlite.rs b/crates/aether-data/src/repository/candidate_selection/sqlite.rs index fd9f4f474..4190d56af 100644 --- a/crates/aether-data/src/repository/candidate_selection/sqlite.rs +++ b/crates/aether-data/src/repository/candidate_selection/sqlite.rs @@ -125,7 +125,8 @@ impl SqliteMinimalCandidateSelectionReadRepository { ) -> Result, 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, 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::::new(); let page_in_sql = !matches!(query.order, StoredPoolKeyCandidateOrder::LoadBalance { .. }); @@ -337,7 +339,8 @@ impl MinimalCandidateSelectionReadRepository for SqliteMinimalCandidateSelection .collect::>(); 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 { aether_ai_formats::api_format_storage_aliases(api_format) } +fn api_format_permission_aliases(api_format: &str) -> Vec { + 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 { api_formats .iter() @@ -1264,6 +1278,21 @@ mod tests { .collect::>(), 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 ); "#, ) diff --git a/crates/aether-provider-transport/src/conversion.rs b/crates/aether-provider-transport/src/conversion.rs index c4fd5ccbb..1e88d4da9 100644 --- a/crates/aether-provider-transport/src/conversion.rs +++ b/crates/aether-provider-transport/src/conversion.rs @@ -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); diff --git a/crates/aether-provider-transport/src/network.rs b/crates/aether-provider-transport/src/network.rs index c2678df85..725105706 100644 --- a/crates/aether-provider-transport/src/network.rs +++ b/crates/aether-provider-transport/src/network.rs @@ -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(); diff --git a/crates/aether-provider-transport/src/provider_types.rs b/crates/aether-provider-transport/src/provider_types.rs index ced7afa8f..f8efa4eaf 100644 --- a/crates/aether-provider-transport/src/provider_types.rs +++ b/crates/aether-provider-transport/src/provider_types.rs @@ -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] diff --git a/crates/aether-provider-transport/src/request_url/mod.rs b/crates/aether-provider-transport/src/request_url/mod.rs index 208aea3cd..ab77d7e4b 100644 --- a/crates/aether-provider-transport/src/request_url/mod.rs +++ b/crates/aether-provider-transport/src/request_url/mod.rs @@ -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( diff --git a/crates/aether-provider-transport/src/same_format_provider/mod.rs b/crates/aether-provider-transport/src/same_format_provider/mod.rs index 58b70a1ee..8dc7c2190 100644 --- a/crates/aether-provider-transport/src/same_format_provider/mod.rs +++ b/crates/aether-provider-transport/src/same_format_provider/mod.rs @@ -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"}); diff --git a/crates/aether-provider-transport/src/url.rs b/crates/aether-provider-transport/src/url.rs index 14661413b..07af0239b 100644 --- a/crates/aether-provider-transport/src/url.rs +++ b/crates/aether-provider-transport/src/url.rs @@ -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!( diff --git a/crates/aether-scheduler-core/src/auth.rs b/crates/aether-scheduler-core/src/auth.rs index 659d70f5b..086b3357d 100644 --- a/crates/aether-scheduler-core/src/auth.rs +++ b/crates/aether-scheduler-core/src/auth.rs @@ -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" diff --git a/crates/aether-scheduler-core/src/model.rs b/crates/aether-scheduler-core/src/model.rs index 044b15a90..f143645f2 100644 --- a/crates/aether-scheduler-core/src/model.rs +++ b/crates/aether-scheduler-core/src/model.rs @@ -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"); diff --git a/crates/aether-usage-runtime/src/report.rs b/crates/aether-usage-runtime/src/report.rs index 82997d223..9f7fca67b 100644 --- a/crates/aether-usage-runtime/src/report.rs +++ b/crates/aether-usage-runtime/src/report.rs @@ -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 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 { diff --git a/frontend/src/api/__tests__/models-dev-pricing.spec.ts b/frontend/src/api/__tests__/models-dev-pricing.spec.ts index b58f72c11..0cd7e1588 100644 --- a/frontend/src/api/__tests__/models-dev-pricing.spec.ts +++ b/frontend/src/api/__tests__/models-dev-pricing.spec.ts @@ -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 }], + }) + }) +}) diff --git a/frontend/src/api/authoritative-model-pricing.ts b/frontend/src/api/authoritative-model-pricing.ts new file mode 100644 index 000000000..448c82c76 --- /dev/null +++ b/frontend/src/api/authoritative-model-pricing.ts @@ -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([ + ['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)], + }, + }, + } +} diff --git a/frontend/src/api/endpoints/types/__tests__/api-format.spec.ts b/frontend/src/api/endpoints/types/__tests__/api-format.spec.ts index 44122b355..70b002c6c 100644 --- a/frontend/src/api/endpoints/types/__tests__/api-format.spec.ts +++ b/frontend/src/api/endpoints/types/__tests__/api-format.spec.ts @@ -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') diff --git a/frontend/src/api/endpoints/types/api-format.ts b/frontend/src/api/endpoints/types/api-format.ts index bae97e933..636355b0c 100644 --- a/frontend/src/api/endpoints/types/api-format.ts +++ b/frontend/src/api/endpoints/types/api-format.ts @@ -6,6 +6,7 @@ export const API_FORMATS = { OPENAI: 'openai:chat', OPENAI_RESPONSES: 'openai:responses', OPENAI_RESPONSES_COMPACT: 'openai:responses:compact', + OPENAI_SEARCH: 'openai:search', OPENAI_IMAGE: 'openai:image', OPENAI_VIDEO: 'openai:video', OPENAI_EMBEDDING: 'openai:embedding', @@ -30,6 +31,7 @@ export const API_FORMAT_LABELS: Record = { [API_FORMATS.OPENAI]: 'OpenAI Chat', [API_FORMATS.OPENAI_RESPONSES]: 'OpenAI Responses', [API_FORMATS.OPENAI_RESPONSES_COMPACT]: 'OpenAI Responses Compact', + [API_FORMATS.OPENAI_SEARCH]: 'OpenAI Search', [API_FORMATS.OPENAI_IMAGE]: 'OpenAI Image', [API_FORMATS.OPENAI_VIDEO]: 'OpenAI Video', [API_FORMATS.OPENAI_EMBEDDING]: 'OpenAI Embedding', @@ -48,6 +50,7 @@ export const API_FORMAT_LABELS: Record = { OPENAI: 'OpenAI Chat', OPENAI_RESPONSES: 'OpenAI Responses', OPENAI_RESPONSES_COMPACT: 'OpenAI Responses Compact', + OPENAI_SEARCH: 'OpenAI Search', OPENAI_IMAGE: 'OpenAI Image', OPENAI_VIDEO: 'OpenAI Video', OPENAI_EMBEDDING: 'OpenAI Embedding', @@ -69,6 +72,7 @@ export const API_FORMAT_SHORT: Record = { [API_FORMATS.OPENAI]: 'O', [API_FORMATS.OPENAI_RESPONSES]: 'OR', [API_FORMATS.OPENAI_RESPONSES_COMPACT]: 'ORC', + [API_FORMATS.OPENAI_SEARCH]: 'OS', [API_FORMATS.OPENAI_IMAGE]: 'OI', [API_FORMATS.OPENAI_VIDEO]: 'OV', [API_FORMATS.OPENAI_EMBEDDING]: 'OE', @@ -86,6 +90,7 @@ export const API_FORMAT_SHORT: Record = { OPENAI: 'O', OPENAI_RESPONSES: 'OR', OPENAI_RESPONSES_COMPACT: 'ORC', + OPENAI_SEARCH: 'OS', OPENAI_IMAGE: 'OI', OPENAI_VIDEO: 'OV', OPENAI_EMBEDDING: 'OE', @@ -109,6 +114,7 @@ export const API_FORMAT_ORDER: string[] = [ API_FORMATS.OPENAI, API_FORMATS.OPENAI_RESPONSES, API_FORMATS.OPENAI_RESPONSES_COMPACT, + API_FORMATS.OPENAI_SEARCH, API_FORMATS.OPENAI_EMBEDDING, API_FORMATS.OPENAI_RERANK, API_FORMATS.OPENAI_IMAGE, @@ -140,6 +146,7 @@ export const API_FORMAT_KIND_LABELS: Record = { chat: 'Chat', responses: 'Responses', 'responses:compact': 'Responses Compact', + search: 'Search', messages: 'Messages', generate_content: 'Generate Content', interactions: 'Interactions', @@ -174,6 +181,9 @@ export function normalizeApiFormatAlias(format: string | null | undefined): stri return API_FORMATS.OPENAI_RESPONSES case 'OPENAI_RESPONSES_COMPACT': return API_FORMATS.OPENAI_RESPONSES_COMPACT + case 'OPENAI_SEARCH': + case 'SEARCH': + return API_FORMATS.OPENAI_SEARCH case 'OPENAI_IMAGE': return API_FORMATS.OPENAI_IMAGE case 'OPENAI_VIDEO': @@ -216,6 +226,18 @@ export function normalizeApiFormatAlias(format: string | null | undefined): stri } } +export function apiFormatPermissionCovers( + allowedFormat: string | null | undefined, + requestedFormat: string | null | undefined, +): boolean { + const allowed = normalizeApiFormatAlias(allowedFormat) + const requested = normalizeApiFormatAlias(requestedFormat) + return Boolean(allowed) + && Boolean(requested) + && (allowed === requested + || (allowed === API_FORMATS.OPENAI_RESPONSES && requested === API_FORMATS.OPENAI_SEARCH)) +} + // 工具函数：按 family 分组并排序 API 格式数组 export interface ApiFormatGroup { family: string @@ -272,14 +294,16 @@ export function formatApiFormatShort(format: string | null | undefined): string // 工具函数：按标准顺序排序 API 格式数组 export function sortApiFormats(formats: string[]): string[] { - return [...formats].sort((a, b) => { - const aIdx = API_FORMAT_ORDER.indexOf(normalizeApiFormatAlias(a)) - const bIdx = API_FORMAT_ORDER.indexOf(normalizeApiFormatAlias(b)) - if (aIdx === -1 && bIdx === -1) return 0 - if (aIdx === -1) return 1 - if (bIdx === -1) return -1 - return aIdx - bIdx - }) + return [...formats].sort(compareApiFormats) +} + +export function compareApiFormats(a: string, b: string): number { + const aIdx = API_FORMAT_ORDER.indexOf(normalizeApiFormatAlias(a)) + const bIdx = API_FORMAT_ORDER.indexOf(normalizeApiFormatAlias(b)) + if (aIdx === -1 && bIdx === -1) return 0 + if (aIdx === -1) return 1 + if (bIdx === -1) return -1 + return aIdx - bIdx } // openai family 格式只支持 bearer（Authorization header），不允许覆盖认证方式 diff --git a/frontend/src/api/models-dev-pricing.ts b/frontend/src/api/models-dev-pricing.ts index 09cccb76f..d875f6b39 100644 --- a/frontend/src/api/models-dev-pricing.ts +++ b/frontend/src/api/models-dev-pricing.ts @@ -1,4 +1,5 @@ import type { PricingTier, TieredPricingConfig } from './endpoints/types' +import { getAuthoritativeModelPricing } from './authoritative-model-pricing' export interface ModelsDevTokenCost { input: number @@ -90,3 +91,12 @@ export function buildModelsDevTieredPricing(cost: unknown): TieredPricingConfig return { tiers } } + +export function resolveModelsDevTieredPricing( + providerId: string, + modelId: string, + cost: unknown, +): TieredPricingConfig | null { + return getAuthoritativeModelPricing(providerId, modelId) + ?? buildModelsDevTieredPricing(cost) +} diff --git a/frontend/src/api/models-dev.ts b/frontend/src/api/models-dev.ts index a4fd640b1..f73d22723 100644 --- a/frontend/src/api/models-dev.ts +++ b/frontend/src/api/models-dev.ts @@ -5,7 +5,7 @@ import api from './client' import { - buildModelsDevTieredPricing, + resolveModelsDevTieredPricing, type ModelsDevCost, } from './models-dev-pricing' import type { TieredPricingConfig } from './endpoints/types' @@ -166,15 +166,16 @@ export async function getModelsDevList(officialOnly: boolean = true): Promise = { - CLAUDE: 'claude:messages', - CLAUDE_MESSAGES: 'claude:messages', - OPENAI: 'openai:chat', - OPENAI_RESPONSES: 'openai:responses', - OPENAI_RESPONSES_COMPACT: 'openai:responses:compact', - OPENAI_VIDEO: 'openai:video', - GEMINI: 'gemini:generate_content', - GEMINI_GENERATE_CONTENT: 'gemini:generate_content', - GEMINI_VIDEO: 'gemini:video', - GEMINI_FILES: 'gemini:files', -} - function normalizeApiFormatKey(value: string | null | undefined): string { - const raw = String(value || '').trim() - if (!raw) return '' + const normalized = normalizeApiFormatAlias(value) + if (!normalized) return '' - if (raw.includes(':')) { - const [family, kind] = raw.split(':', 2) - const familyNorm = family?.trim().toLowerCase() - const kindNorm = kind?.trim().toLowerCase() - if (familyNorm === 'claude' && ['chat', 'cli', 'messages'].includes(kindNorm)) { + if (normalized.includes(':')) { + const { family, kind } = parseApiFormat(normalized) + if (family === 'claude' && ['chat', 'cli', 'messages'].includes(kind)) { return 'claude:messages' } - if (familyNorm === 'gemini' && ['chat', 'cli', 'generate_content'].includes(kindNorm)) { + if (family === 'gemini' && ['chat', 'cli', 'generate_content'].includes(kind)) { return 'gemini:generate_content' } - if (familyNorm && kindNorm) return `${familyNorm}:${kindNorm}` } - - const legacy = raw.toUpperCase().replace(/-/g, '_') - return LEGACY_API_FORMAT_MAP[legacy] || raw.toLowerCase() + return normalized } function normalizePriorityMap( diff --git a/frontend/src/features/providers/components/ProviderFormDialog.vue b/frontend/src/features/providers/components/ProviderFormDialog.vue index 790e58ee4..8db488c0c 100644 --- a/frontend/src/features/providers/components/ProviderFormDialog.vue +++ b/frontend/src/features/providers/components/ProviderFormDialog.vue @@ -197,7 +197,7 @@ :model-value="form.request_timeout ?? ''" type="number" min="1" - max="600" + max="1200" step="1" placeholder="300" @update:model-value="(v) => form.request_timeout = parseNumberInput(v)" diff --git a/frontend/src/features/providers/components/__tests__/endpoint-default-paths.spec.ts b/frontend/src/features/providers/components/__tests__/endpoint-default-paths.spec.ts index 183ebb5bd..5119b9a21 100644 --- a/frontend/src/features/providers/components/__tests__/endpoint-default-paths.spec.ts +++ b/frontend/src/features/providers/components/__tests__/endpoint-default-paths.spec.ts @@ -9,6 +9,7 @@ const apiFormats = [ { value: 'gemini:embedding', default_path: '/v1beta/models/{model}:embedContent' }, { value: 'gemini:video', default_path: '/v1beta/models/{model}:predictLongRunning' }, { value: 'openai:responses', default_path: '/v1/responses' }, + { value: 'openai:search', default_path: '/v1/alpha/search' }, { value: 'openai:embedding', default_path: '/v1/embeddings' }, { value: 'openai:rerank', default_path: '/v1/rerank' }, { value: 'openai:image', default_path: '/v1/images/generations' }, @@ -75,6 +76,30 @@ describe('endpoint default paths', () => { })).toBe('/responses') }) + it('uses the Search path relative to OpenAI and Codex api roots', () => { + expect(getDefaultEndpointBaseUrl({ + apiFormat: 'openai:search', + baseUrl: 'https://api.openai.com', + })).toBe('https://api.openai.com/v1') + expect(getDefaultEndpointPath({ + apiFormat: 'openai:search', + providerType: 'openai', + baseUrl: 'https://api.openai.com/v1', + apiFormats, + })).toBe('/alpha/search') + + expect(getDefaultEndpointBaseUrl({ + apiFormat: 'openai:search', + baseUrl: 'https://chatgpt.com/backend-api/codex', + })).toBe('https://chatgpt.com/backend-api/codex') + expect(getDefaultEndpointPath({ + apiFormat: 'openai:search', + providerType: 'codex', + baseUrl: 'https://chatgpt.com/backend-api/codex', + apiFormats, + })).toBe('/alpha/search') + }) + it('drops /v1 from API-root defaults because base URL is the API root', () => { expect(getDefaultEndpointPath({ apiFormat: 'openai:chat', diff --git a/frontend/src/features/providers/components/__tests__/endpoint-protocol-policy.spec.ts b/frontend/src/features/providers/components/__tests__/endpoint-protocol-policy.spec.ts new file mode 100644 index 000000000..a879b7f07 --- /dev/null +++ b/frontend/src/features/providers/components/__tests__/endpoint-protocol-policy.spec.ts @@ -0,0 +1,15 @@ +import { describe, expect, it } from 'vitest' + +import { fixedEndpointUpstreamStreamPolicy } from '../endpoint-protocol-policy' + +describe('fixedEndpointUpstreamStreamPolicy', () => { + it('uses the fixed Codex transport mode for each protocol', () => { + expect(fixedEndpointUpstreamStreamPolicy('codex', 'openai:responses')).toBe('force_stream') + expect(fixedEndpointUpstreamStreamPolicy('codex', 'openai:search')).toBe('force_non_stream') + }) + + it('applies Search synchronization by format and keeps other formats configurable', () => { + expect(fixedEndpointUpstreamStreamPolicy('custom', 'openai:search')).toBe('force_non_stream') + expect(fixedEndpointUpstreamStreamPolicy('codex', 'openai:responses:compact')).toBeNull() + }) +}) diff --git a/frontend/src/features/providers/components/endpoint-default-paths.ts b/frontend/src/features/providers/components/endpoint-default-paths.ts index 0fcd451d0..f49f3ef9f 100644 --- a/frontend/src/features/providers/components/endpoint-default-paths.ts +++ b/frontend/src/features/providers/components/endpoint-default-paths.ts @@ -90,6 +90,7 @@ function usesVersionedApiRootByDefault(apiFormat: string): boolean { return apiFormat === 'openai:chat' || apiFormat === 'openai:responses' || apiFormat === 'openai:responses:compact' + || apiFormat === 'openai:search' || apiFormat === 'openai:embedding' || apiFormat === 'openai:rerank' || apiFormat === 'openai:image' @@ -129,6 +130,7 @@ function skipsVersionedApiRootDefault(apiFormat: string, baseUrl: string): boole return false } return isDeepSeekApiRoot(baseUrl) + || isCodexUrl(baseUrl) || isBigModelCodingApiRoot(baseUrl) || isGoogleOpenAiCompatApiRoot(baseUrl) || isVertexOpenAiCompatApiRoot(baseUrl) @@ -198,6 +200,9 @@ export function getDefaultEndpointPath(params: { if (normalizedApiFormat === 'openai:responses' && isCodex) { return '/responses' } + if (normalizedApiFormat === 'openai:search' && isCodex) { + return '/alpha/search' + } if (usesVersionedApiRootByDefault(normalizedApiFormat)) { return stripVersionPrefixForApiRoot(defaultPath) } diff --git a/frontend/src/features/providers/components/endpoint-protocol-policy.ts b/frontend/src/features/providers/components/endpoint-protocol-policy.ts new file mode 100644 index 000000000..90710989f --- /dev/null +++ b/frontend/src/features/providers/components/endpoint-protocol-policy.ts @@ -0,0 +1,19 @@ +import { normalizeEndpointApiFormat } from './endpoint-default-paths' + +export type FixedUpstreamStreamPolicy = 'force_stream' | 'force_non_stream' + +export function fixedEndpointUpstreamStreamPolicy( + providerType: string | null | undefined, + apiFormat: string, +): FixedUpstreamStreamPolicy | null { + const normalizedApiFormat = normalizeEndpointApiFormat(apiFormat) + if (normalizedApiFormat === 'openai:search') return 'force_non_stream' + if (providerType?.trim().toLowerCase() !== 'codex') return null + + switch (normalizedApiFormat) { + case 'openai:responses': + return 'force_stream' + default: + return null + } +} diff --git a/frontend/src/features/providers/components/provider-tabs/ModelMappingTab.vue b/frontend/src/features/providers/components/provider-tabs/ModelMappingTab.vue index a217a4b9d..7a0b99337 100644 --- a/frontend/src/features/providers/components/provider-tabs/ModelMappingTab.vue +++ b/frontend/src/features/providers/components/provider-tabs/ModelMappingTab.vue @@ -366,12 +366,12 @@ import { type EndpointAPIKey } from '@/api/endpoints/keys' import { updateModel } from '@/api/endpoints/models' import { parseApiError } from '@/utils/errorParser' import type { ProviderWithEndpointsSummary } from '@/api/endpoints' -import { normalizeApiFormatAlias } from '@/api/endpoints/types/api-format' import { buildDefaultModelTestRequestHeaders, buildDefaultModelTestRequestBody, isModelTestableApiFormat, isModelTestableEndpoint, + modelTestMappingScopeMatchesEndpoint, parseModelTestRequestHeadersDraft, parseModelTestRequestBodyDraft, selectPreferredModelTestEndpoint, @@ -852,16 +852,11 @@ function scopedMappingEndpoints(item: CombinedMapping): ProviderEndpoint[] { const group = item.group if (!group) return activeEndpoints.value - const apiFormats = new Set(normalizeStringList(group.apiFormats).map(normalizeApiFormatAlias)) - const endpointIds = new Set(normalizeStringList(group.endpointIds)) - const matched = activeEndpoints.value.filter((endpoint) => { - const apiFormatMatched = apiFormats.size === 0 - || apiFormats.has(normalizeApiFormatAlias(endpoint.api_format)) - const endpointMatched = endpointIds.size === 0 || endpointIds.has(endpoint.id) - return apiFormatMatched && endpointMatched - }) - - return matched.length > 0 ? matched : activeEndpoints.value + return activeEndpoints.value.filter(endpoint => modelTestMappingScopeMatchesEndpoint( + group.apiFormats, + group.endpointIds, + endpoint, + )) } // 测试精确映射 diff --git a/frontend/src/features/providers/components/provider-tabs/__tests__/model-test-request.spec.ts b/frontend/src/features/providers/components/provider-tabs/__tests__/model-test-request.spec.ts index afab2e00e..ce3c6792e 100644 --- a/frontend/src/features/providers/components/provider-tabs/__tests__/model-test-request.spec.ts +++ b/frontend/src/features/providers/components/provider-tabs/__tests__/model-test-request.spec.ts @@ -10,6 +10,7 @@ import { isModelTestableEndpoint, isModelTestableApiFormat, listModelTestMappedModelOptions, + modelTestMappingScopeMatchesEndpoint, normalizeModelTestMappedModelSelection, selectPreferredModelTestEndpoint, setModelTestRequestBodyModel, @@ -97,6 +98,20 @@ describe('buildDefaultModelTestRequestBody', () => { expect(body.messages).toBeUndefined() }) + it('uses the synchronous OpenAI Search request contract', () => { + const body = JSON.parse(buildDefaultModelTestRequestBody('gpt-5.6-sol', 'openai:search')) + + expect(body).toEqual({ + id: expect.stringMatching(/^aether-model-test-/), + model: 'gpt-5.6-sol', + input: 'Hello! This is a test message.', + commands: { + search_query: [{ q: 'Hello! This is a test message.' }], + }, + max_output_tokens: 256, + }) + }) + it('uses image generation tools for image models on OpenAI Responses endpoints', () => { const body = JSON.parse(buildDefaultModelTestRequestBody( 'gpt-image-2', @@ -177,6 +192,52 @@ describe('buildDefaultModelTestRequestBody', () => { expect(options).toEqual([]) }) + it('uses one-way format permissions for Search mapping options', () => { + const responsesOptions = listModelTestMappedModelOptions({ + provider_model_name: 'gpt-5.6-luna', + provider_model_mappings: [{ + name: 'gpt-5.6-luna', + priority: 1, + api_formats: ['openai:responses'], + }], + }, { + id: 'search-endpoint', + api_format: 'openai:search', + }) + const searchOptions = listModelTestMappedModelOptions({ + provider_model_name: 'gpt-5.6-luna', + provider_model_mappings: [{ + name: 'gpt-5.6-luna', + priority: 1, + api_formats: ['openai:search'], + }], + }, { + id: 'responses-endpoint', + api_format: 'openai:responses', + }) + + expect(responsesOptions).toEqual([{ name: 'gpt-5.6-luna', priority: 1 }]) + expect(searchOptions).toEqual([]) + }) + + it('keeps explicit mapping scopes strict when no endpoint matches', () => { + expect(modelTestMappingScopeMatchesEndpoint( + ['openai:responses'], + ['search-endpoint'], + { id: 'search-endpoint', api_format: 'openai:search' }, + )).toBe(true) + expect(modelTestMappingScopeMatchesEndpoint( + ['openai:responses'], + ['another-endpoint'], + { id: 'search-endpoint', api_format: 'openai:search' }, + )).toBe(false) + expect(modelTestMappingScopeMatchesEndpoint( + ['openai:search'], + undefined, + { id: 'responses-endpoint', api_format: 'openai:responses' }, + )).toBe(false) + }) + it('keeps the current model selected by default until a mapped model is chosen', () => { const options = [ { name: 'MiniMax-M2.7-highspeed', priority: 1 }, @@ -259,7 +320,7 @@ describe('isModelTestableApiFormat', () => { 'gemini:video', 'gemini:files', ' OPENAI:VIDEO ', - ])('excludes task and file endpoint formats from model tests: %s', (apiFormat) => { + ])('excludes endpoint formats without a model-test request contract: %s', (apiFormat) => { expect(isModelTestableApiFormat(apiFormat)).toBe(false) }) @@ -267,6 +328,7 @@ describe('isModelTestableApiFormat', () => { 'openai:chat', 'openai:responses', 'openai:responses:compact', + 'openai:search', 'openai:image', 'claude:messages', 'gemini:generate_content', @@ -379,6 +441,17 @@ describe('isModelTestableEndpoint', () => { }, keys)).toBe(true) }) + it('uses one-way format permissions for Search endpoint keys', () => { + expect(isModelTestableEndpoint({ + api_format: 'openai:search', + is_active: true, + }, [{ api_formats: ['openai:responses'], is_active: true }])).toBe(true) + expect(isModelTestableEndpoint({ + api_format: 'openai:responses', + is_active: true, + }, [{ api_formats: ['openai:search'], is_active: true }])).toBe(false) + }) + it('lets fixed provider OAuth keys inherit testable endpoint formats', () => { const keys = [{ api_formats: ['legacy:mismatch'], diff --git a/frontend/src/features/providers/components/provider-tabs/model-test-capabilities.ts b/frontend/src/features/providers/components/provider-tabs/model-test-capabilities.ts index b57dcb6ea..ecb82bb46 100644 --- a/frontend/src/features/providers/components/provider-tabs/model-test-capabilities.ts +++ b/frontend/src/features/providers/components/provider-tabs/model-test-capabilities.ts @@ -1,4 +1,7 @@ -import { normalizeApiFormatAlias } from '@/api/endpoints/types/api-format' +import { + apiFormatPermissionCovers, + normalizeApiFormatAlias, +} from '@/api/endpoints/types/api-format' import type { ModelTestCapabilities, OpenAiImageModelTestCapability } from '@/api/endpoints/types' export type ModelTestEndpointSource = { @@ -71,7 +74,7 @@ export function modelTestKeySupportsEndpoint( const keyFormats = normalizeModelTestStringList(key.api_formats) if (keyFormats.length === 0) return true - return keyFormats.some(format => normalizeApiFormatAlias(format) === endpointFormat) + return keyFormats.some(format => apiFormatPermissionCovers(format, endpointFormat)) } export function isModelTestableEndpoint( diff --git a/frontend/src/features/providers/components/provider-tabs/model-test-request.ts b/frontend/src/features/providers/components/provider-tabs/model-test-request.ts index 9b5fcb65f..f0c2cde74 100644 --- a/frontend/src/features/providers/components/provider-tabs/model-test-request.ts +++ b/frontend/src/features/providers/components/provider-tabs/model-test-request.ts @@ -1,4 +1,7 @@ -import { normalizeApiFormatAlias } from '@/api/endpoints/types/api-format' +import { + apiFormatPermissionCovers, + normalizeApiFormatAlias, +} from '@/api/endpoints/types/api-format' import type { ProviderModelMapping } from '@/api/endpoints/types' import type { TestModelRequest } from '@/api/endpoints/providers' import { @@ -29,12 +32,18 @@ export type { ModelTestImagePreview } from './model-test-preview' const DEFAULT_MODEL_TEST_MESSAGE = 'Hello! This is a test message.' +function buildModelTestSearchSessionId(): string { + const id = globalThis.crypto?.randomUUID?.() + ?? `${Date.now().toString(36)}-${Math.random().toString(36).slice(2)}` + return `aether-model-test-${id}` +} + type ModelTestMappingSource = { provider_model_name: string provider_model_mappings?: ProviderModelMapping[] | null } -type ModelTestMappingEndpoint = { +export type ModelTestMappingEndpoint = { id: string api_format: string } @@ -44,17 +53,19 @@ export type ModelTestMappedModelOption = { priority: number } -function mappingApiFormatMatches(mapping: ProviderModelMapping, endpoint: ModelTestMappingEndpoint): boolean { - const apiFormats = normalizeModelTestStringList(mapping.api_formats) - if (apiFormats.length === 0) return true +export function modelTestMappingScopeMatchesEndpoint( + apiFormats: string[] | null | undefined, + endpointIds: string[] | null | undefined, + endpoint: ModelTestMappingEndpoint, +): boolean { + const normalizedApiFormats = normalizeModelTestStringList(apiFormats) const endpointFormat = normalizeApiFormatAlias(endpoint.api_format) - return apiFormats.some(format => normalizeApiFormatAlias(format) === endpointFormat) -} - -function mappingEndpointMatches(mapping: ProviderModelMapping, endpoint: ModelTestMappingEndpoint): boolean { - const endpointIds = normalizeModelTestStringList(mapping.endpoint_ids) - if (endpointIds.length === 0) return true - return endpointIds.includes(endpoint.id) + const apiFormatMatches = normalizedApiFormats.length === 0 + || normalizedApiFormats.some(format => apiFormatPermissionCovers(format, endpointFormat)) + const normalizedEndpointIds = normalizeModelTestStringList(endpointIds) + const endpointMatches = normalizedEndpointIds.length === 0 + || normalizedEndpointIds.includes(endpoint.id) + return apiFormatMatches && endpointMatches } export function listModelTestMappedModelOptions( @@ -65,8 +76,11 @@ export function listModelTestMappedModelOptions( const matchedMappings = model.provider_model_mappings .filter(mapping => mapping.name.trim()) - .filter(mapping => mappingApiFormatMatches(mapping, endpoint)) - .filter(mapping => mappingEndpointMatches(mapping, endpoint)) + .filter(mapping => modelTestMappingScopeMatchesEndpoint( + mapping.api_formats, + mapping.endpoint_ids, + endpoint, + )) .sort((left, right) => { const leftPriority = Number.isFinite(left.priority) ? left.priority : 1 const rightPriority = Number.isFinite(right.priority) ? right.priority : 1 @@ -204,6 +218,18 @@ export function buildDefaultModelTestRequestBody( }, null, 2) } + if (normalizedApiFormat === 'openai:search') { + return JSON.stringify({ + id: buildModelTestSearchSessionId(), + model: modelName, + input: DEFAULT_MODEL_TEST_MESSAGE, + commands: { + search_query: [{ q: DEFAULT_MODEL_TEST_MESSAGE }], + }, + max_output_tokens: 256, + }, null, 2) + } + if (normalizedApiFormat === 'openai:responses' && modelSupportsImageGeneration(model)) { return JSON.stringify({ model: modelName, diff --git a/frontend/src/features/providers/composables/useEndpointStatus.ts b/frontend/src/features/providers/composables/useEndpointStatus.ts index 0dd65645e..75feb17f1 100644 --- a/frontend/src/features/providers/composables/useEndpointStatus.ts +++ b/frontend/src/features/providers/composables/useEndpointStatus.ts @@ -1,35 +1,15 @@ import type { EndpointHealthDetail } from '@/api/endpoints' +import { compareApiFormats } from '@/api/endpoints/types/api-format' import { defaultLocale, translateLegacyText, type Locale } from '@/i18n/messages' // 端点状态枚举 export type EndpointStatus = 'disabled' | 'no_keys' | 'keys_disabled' | 'available' -const ENDPOINT_SORT_ORDER = [ - 'claude:messages', - 'openai:chat', - 'openai:responses', - 'openai:responses:compact', - 'openai:embedding', - 'openai:rerank', - 'gemini:generate_content', - 'gemini:interactions', - 'gemini:embedding', - 'openai:video', - 'gemini:video', - 'gemini:files', - 'jina:embedding', - 'jina:rerank', - 'doubao:embedding', - 'aliyun:multimodal_embedding', -] - /** * 端点排序 */ export function sortEndpoints(endpoints: T[]): T[] { - return [...endpoints].sort((a, b) => { - return ENDPOINT_SORT_ORDER.indexOf(a.api_format) - ENDPOINT_SORT_ORDER.indexOf(b.api_format) - }) + return [...endpoints].sort((a, b) => compareApiFormats(a.api_format, b.api_format)) } /** diff --git a/frontend/src/features/providers/composables/useProviderFilters.ts b/frontend/src/features/providers/composables/useProviderFilters.ts index 1c5bf711e..7acba18b3 100644 --- a/frontend/src/features/providers/composables/useProviderFilters.ts +++ b/frontend/src/features/providers/composables/useProviderFilters.ts @@ -1,5 +1,6 @@ import { ref, computed, watch } from 'vue' import type { ProviderSummaryQuery } from '@/api/endpoints' +import { API_FORMAT_ORDER, formatApiFormat } from '@/api/endpoints/types/api-format' import { useI18n } from '@/i18n' export interface FilterOption { @@ -25,19 +26,7 @@ export function useProviderFilters( const apiFormatFilters = computed(() => [ { value: 'all', label: legacyT('全部格式') }, - { value: 'claude:messages', label: 'Claude Messages' }, - { value: 'openai:chat', label: 'OpenAI Chat' }, - { value: 'openai:responses', label: 'OpenAI Responses' }, - { value: 'openai:responses:compact', label: 'OpenAI Responses Compact' }, - { value: 'openai:embedding', label: 'OpenAI Embedding' }, - { value: 'openai:rerank', label: 'OpenAI Rerank' }, - { value: 'gemini:generate_content', label: 'Gemini Generate Content' }, - { value: 'gemini:interactions', label: 'Gemini Interactions' }, - { value: 'gemini:embedding', label: 'Gemini Embedding' }, - { value: 'jina:embedding', label: 'Jina Embedding' }, - { value: 'jina:rerank', label: 'Jina Rerank' }, - { value: 'doubao:embedding', label: 'Doubao Embedding' }, - { value: 'aliyun:multimodal_embedding', label: 'Aliyun Multimodal Embedding' }, + ...API_FORMAT_ORDER.map(value => ({ value, label: formatApiFormat(value) })), ]) const modelFilters = computed(() => { diff --git a/frontend/src/features/usage/utils/__tests__/status.spec.ts b/frontend/src/features/usage/utils/__tests__/status.spec.ts index 83f5a6daf..5b809b1d4 100644 --- a/frontend/src/features/usage/utils/__tests__/status.spec.ts +++ b/frontend/src/features/usage/utils/__tests__/status.spec.ts @@ -197,6 +197,14 @@ describe('usage status helpers', () => { client_is_stream: undefined, }))).toBe('标准->流式') + expect(formatUsageStreamLabel(buildUsageRecord({ + api_format: 'openai:search', + is_stream: false, + upstream_is_stream: false, + client_requested_stream: undefined, + client_is_stream: undefined, + }))).toBe('标准') + expect(formatUsageStreamLabel(buildUsageRecord({ api_format: 'claude:messages', is_stream: false, diff --git a/frontend/src/features/usage/utils/status.ts b/frontend/src/features/usage/utils/status.ts index c4355c4c1..9ed987fc2 100644 --- a/frontend/src/features/usage/utils/status.ts +++ b/frontend/src/features/usage/utils/status.ts @@ -160,6 +160,7 @@ function usageApiFormatDefaultsToNonStream(apiFormat: string): boolean { case 'openai:chat': case 'openai:responses': case 'openai:responses:compact': + case 'openai:search': case 'openai:image': case 'claude:messages': return true diff --git a/frontend/src/mocks/data.ts b/frontend/src/mocks/data.ts index 9d9a74c2a..ca0c7d1d2 100644 --- a/frontend/src/mocks/data.ts +++ b/frontend/src/mocks/data.ts @@ -1187,6 +1187,7 @@ export const MOCK_API_FORMATS = { { value: 'openai:chat', label: 'OpenAI Chat', default_path: '/v1/chat/completions', aliases: [] }, { value: 'openai:responses', label: 'OpenAI Responses', default_path: '/v1/responses', aliases: [] }, { value: 'openai:responses:compact', label: 'OpenAI Responses Compact', default_path: '/v1/responses/compact', aliases: [] }, + { value: 'openai:search', label: 'OpenAI Search', default_path: '/v1/alpha/search', aliases: ['openai_search', 'search'] }, { value: 'openai:embedding', label: 'OpenAI Embedding', default_path: '/v1/embeddings', aliases: [] }, { value: 'openai:rerank', label: 'OpenAI Rerank', default_path: '/v1/rerank', aliases: [] }, { value: 'openai:image', label: 'OpenAI Image', default_path: '/v1/images/generations', aliases: [] }, diff --git a/frontend/src/views/admin/module-management/__tests__/modelDirectivesConfig.spec.ts b/frontend/src/views/admin/module-management/__tests__/modelDirectivesConfig.spec.ts index 855489d66..6c1e617cf 100644 --- a/frontend/src/views/admin/module-management/__tests__/modelDirectivesConfig.spec.ts +++ b/frontend/src/views/admin/module-management/__tests__/modelDirectivesConfig.spec.ts @@ -28,6 +28,7 @@ describe('modelDirectivesConfig', () => { 'fast', ]) expect(defaultModelDirectiveSuffixesForApiFormat('openai:responses')).toContain('ultra') + expect(defaultModelDirectiveSuffixesForApiFormat('openai:search')).toEqual(MODEL_DIRECTIVE_SUFFIXES) expect(defaultModelDirectiveSuffixesForApiFormat('claude:messages')).not.toContain('ultra') expect(defaultModelDirectiveSuffixesForApiFormat('gemini:generate_content')).not.toContain('ultra') }) diff --git a/frontend/src/views/admin/module-management/modelDirectivesConfig.ts b/frontend/src/views/admin/module-management/modelDirectivesConfig.ts index 19dff3e6b..9e307e025 100644 --- a/frontend/src/views/admin/module-management/modelDirectivesConfig.ts +++ b/frontend/src/views/admin/module-management/modelDirectivesConfig.ts @@ -63,6 +63,11 @@ export const MODEL_DIRECTIVE_API_FORMATS = [ label: 'OpenAI Responses Compact', parameter: 'reasoning.effort', }, + { + key: 'openai:search', + label: 'OpenAI Search', + parameter: 'reasoning.effort', + }, { key: 'claude:messages', label: 'Claude Messages', @@ -90,6 +95,7 @@ export function defaultModelDirectiveSuffixesForApiFormat( case 'openai:chat': case 'openai:responses': case 'openai:responses:compact': + case 'openai:search': return MODEL_DIRECTIVE_SUFFIXES case 'claude:messages': case 'gemini:generate_content':