diff --git a/backend/crates/providers/openai/src/credential/catalog.rs b/backend/crates/providers/openai/src/credential/catalog.rs index 1055617ca..f9ac83dae 100644 --- a/backend/crates/providers/openai/src/credential/catalog.rs +++ b/backend/crates/providers/openai/src/credential/catalog.rs @@ -32,6 +32,8 @@ const PLAN_CATALOG_CACHE_TTL: Duration = Duration::from_secs(5 * 60); const MAX_CATALOG_FETCH_ATTEMPTS: usize = 3; const MAX_PLAN_CATALOG_MODELS: usize = 2_048; const MAX_CLIENT_CATALOGS: usize = 32; +// 保持原先 32 份、每份 1 MiB 的原生正文缓存预算,不把缓存容量变成上游目录合同 +const MAX_CLIENT_CATALOG_CACHE_BYTES: usize = 32 * 1024 * 1024; const MODEL_CATALOG_TIMEOUT: Duration = Duration::from_secs(15); const MAX_CONCURRENT_API_CATALOGS: usize = 4; const CLIENT_CATALOG_ERROR_TTL: Duration = Duration::from_secs(5); @@ -218,6 +220,20 @@ struct CachedClientCatalog { result: ClientCatalogResult, } +impl CachedClientCatalog { + fn native_bytes(&self) -> usize { + self.result.as_ref().map_or(0, |models| { + models + .iter() + .map(|model| match &model.content { + ProviderModelContent::Native(payload) => payload.body().len(), + ProviderModelContent::Adapted(_) => 0, + }) + .sum() + }) + } +} + struct ClientCatalogEntry { key: ClientCatalogKey, created_at: Instant, @@ -460,6 +476,30 @@ impl CodexCredentialCatalogService { } }) .await; + // 完成中的 OnceCell 仍由并发读者持有;移出缓存只影响后续查询,不丢弃本次成功结果 + let mut cache = self + .client_catalogs + .lock() + .map_err(|_| CodexCredentialCatalogError::Cache)?; + if result.native_bytes() > MAX_CLIENT_CATALOG_CACHE_BYTES { + cache.retain(|entry| !Arc::ptr_eq(&entry.value, &value)); + } else { + let mut bytes: usize = cache + .iter() + .filter_map(|entry| entry.value.get()) + .map(CachedClientCatalog::native_bytes) + .sum(); + while bytes > MAX_CLIENT_CATALOG_CACHE_BYTES { + let Some(index) = cache.iter().position(|entry| entry.value.initialized()) else { + break; + }; + let entry = cache.remove(index); + bytes -= entry + .value + .get() + .map_or(0, CachedClientCatalog::native_bytes); + } + } result.result.clone() } diff --git a/backend/crates/providers/openai/src/provider/execution.rs b/backend/crates/providers/openai/src/provider/execution.rs index cf21b514a..83fc9b375 100644 --- a/backend/crates/providers/openai/src/provider/execution.rs +++ b/backend/crates/providers/openai/src/provider/execution.rs @@ -396,6 +396,7 @@ pub(super) struct ColdResponse { pub(super) session_transport_key: Option, pub(super) session_transport_key_hash: Option, pub(super) session_transport_recovery: CodexSessionTransportRecovery, + pub(super) guardian_routing: guardian::GuardianRouting, pub(super) websocket_retry_count: u32, pub(super) stream_max_retries: u32, pub(super) session_capture: Option, @@ -781,6 +782,7 @@ pub(super) fn cold_response_stream(response: ColdResponse) -> EventStream { session_transport_key, session_transport_key_hash, session_transport_recovery, + guardian_routing, websocket_retry_count, stream_max_retries, mut session_capture, @@ -951,6 +953,7 @@ pub(super) fn cold_response_stream(response: ColdResponse) -> EventStream { } } let response_transport = response.transport; + let routing_affinity = response.routing_affinity; let websocket_connection_id = response.websocket_connection_id; let mut body = response.body; let mut failure_diagnostics = response.diagnostics.clone(); @@ -1132,6 +1135,13 @@ pub(super) fn cold_response_stream(response: ColdResponse) -> EventStream { ) }); let timing_signals = decoder.take_timing_signals(); + if let Some(affinity) = &routing_affinity { + for fact in events.iter().flat_map(ProviderEvent::canonical_facts) { + if let GatewayEvent::Started(meta) = fact { + guardian_routing.remember(context.client_api_key_ref(), active_account.id(), meta.response_id(), affinity); + } + } + } let output_timing_changed = observation_state .observe_timing_signals(timing_signals, context.timing_started_at()); let upstream_timing_changed = observation_state diff --git a/backend/crates/providers/openai/src/provider/guardian.rs b/backend/crates/providers/openai/src/provider/guardian.rs new file mode 100644 index 000000000..8c672063c --- /dev/null +++ b/backend/crates/providers/openai/src/provider/guardian.rs @@ -0,0 +1,95 @@ +//! 按 Client Key 和账号关联父响应的连接路由快照,不持久化 Cookie 或改变选号 + +use std::{ + collections::HashMap, + sync::{Arc, Mutex}, + time::{Duration, Instant}, +}; + +use gateway_core::{account::ProviderAccountId, policy::ClientApiKeyId}; +use sha2::{Digest, Sha256}; + +use crate::transport::routing_affinity::ChatGptRoutingAffinity; + +// 与连接续接 tombstone 使用相同保留窗口;最多 4096 份、每份两个 8 KiB Cookie +const RETENTION: Duration = Duration::from_mins(30); +const MAX_PARENT_RESPONSES: usize = 4096; + +#[derive(Clone, Default)] +pub(super) struct GuardianRouting { + parents: Arc>>, +} + +struct ParentRouting { + captured_at: Instant, + affinity: ChatGptRoutingAffinity, +} + +impl GuardianRouting { + pub(super) fn get( + &self, + client_key: &ClientApiKeyId, + account: &ProviderAccountId, + response_id: &str, + ) -> Option { + let key = parent_key(client_key, account, response_id); + let mut parents = self + .parents + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + let parent = parents.get(&key)?; + if parent.captured_at.elapsed() >= RETENTION { + parents.remove(&key); + return None; + } + Some(parent.affinity.clone()) + } + + pub(super) fn remember( + &self, + client_key: &ClientApiKeyId, + account: &ProviderAccountId, + response_id: &str, + affinity: &ChatGptRoutingAffinity, + ) { + let key = parent_key(client_key, account, response_id); + let mut parents = self + .parents + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + if parents.contains_key(&key) { + return; + } + if parents.len() >= MAX_PARENT_RESPONSES { + parents.retain(|_, parent| parent.captured_at.elapsed() < RETENTION); + if parents.len() >= MAX_PARENT_RESPONSES + && let Some(oldest) = parents + .iter() + .min_by_key(|(_, parent)| parent.captured_at) + .map(|(key, _)| *key) + { + parents.remove(&oldest); + } + } + parents.insert( + key, + ParentRouting { + captured_at: Instant::now(), + affinity: affinity.clone(), + }, + ); + } +} + +fn parent_key( + client_key: &ClientApiKeyId, + account: &ProviderAccountId, + response_id: &str, +) -> [u8; 32] { + let mut hash = Sha256::new(); + for part in [client_key.as_str(), account.as_str(), response_id] { + hash.update(part.len().to_le_bytes()); + hash.update(part.as_bytes()); + } + hash.finalize().into() +} diff --git a/backend/crates/providers/openai/src/provider/mod.rs b/backend/crates/providers/openai/src/provider/mod.rs index 0f5eab5e5..1965ffdb6 100644 --- a/backend/crates/providers/openai/src/provider/mod.rs +++ b/backend/crates/providers/openai/src/provider/mod.rs @@ -104,6 +104,7 @@ use crate::transport::{ mod execution; mod failure; +mod guardian; mod live; mod observation; mod upstream_adapter; @@ -156,6 +157,7 @@ pub struct CodexProvider { live_calls_url: Url, session_identity: Option, session_transport_recovery: CodexSessionTransportRecovery, + guardian_routing: guardian::GuardianRouting, stream_max_retries: u32, live_registry: Arc, live_gateway: Option>, @@ -360,6 +362,7 @@ impl CodexProvider { live_calls_url, session_identity: None, session_transport_recovery: CodexSessionTransportRecovery::default(), + guardian_routing: guardian::GuardianRouting::default(), stream_max_retries, live_registry: Arc::new(CodexLiveRegistry::default()), live_gateway: None, @@ -936,11 +939,43 @@ impl CodexProvider { let state_before = trace .is_enabled() .then(|| crate::transport::diagnostics::evidence::request_state(&upstream_request)); + let guardian_parent = (!state_owner_cross_account + && upstream_request.is_guardian() + && matches!( + lease.authentication(), + crate::credential::CodexRuntimeAuthentication::OAuth(_) + )) + .then(|| { + upstream_request + .client_metadata()? + .get("parent_response_id")? + .as_str() + }) + .flatten() + .and_then(|response_id| { + self.guardian_routing + .get( + context.client_api_key_ref(), + lease.account_id(), + response_id, + ) + .map(|affinity| (response_id.to_owned(), affinity)) + }); scope_request_to_account( &mut upstream_request, lease.installation_id(), account_scope, ); + // 父响应的可信关联只恢复此字段,不把首次 Guardian 请求的其他账号状态视为可信 + if let Some((response_id, _)) = &guardian_parent + && let Some(Value::Object(metadata)) = + upstream_request.body_mut().get_mut("client_metadata") + { + metadata.insert( + "parent_response_id".to_owned(), + Value::String(response_id.clone()), + ); + } if let Some(before) = state_before { trace.record_evidence("provider.request.account_scope", serde_json::json!({ "scope": match account_scope { @@ -1107,6 +1142,7 @@ impl CodexProvider { .for_account(lease.account()) .map_err(|error| map_client_error(error, UpstreamSendState::NotSent, false).error)? .with_responses_api_base_url(lease.authentication().responses_api_base_url()) + .with_routing_affinity(guardian_parent.map(|(_, affinity)| affinity)) .with_connection_budget(context.connection_budget().clone()) .with_middleware_headers(middleware_headers) .with_outbound_middleware(outbound_middleware) @@ -1123,6 +1159,7 @@ impl CodexProvider { session_transport_key, session_transport_key_hash, session_transport_recovery: self.session_transport_recovery.clone(), + guardian_routing: self.guardian_routing.clone(), websocket_retry_count, stream_max_retries: self.stream_max_retries, session_capture, diff --git a/backend/crates/providers/openai/src/transport/catalog.rs b/backend/crates/providers/openai/src/transport/catalog.rs index bb1e50b64..ff46ea0a4 100644 --- a/backend/crates/providers/openai/src/transport/catalog.rs +++ b/backend/crates/providers/openai/src/transport/catalog.rs @@ -11,9 +11,6 @@ use serde::Deserialize; use super::client::OpenAiUpstreamProtocol; -/// 单次 Codex 模型目录响应允许的最大字节数 -pub const MAX_CODEX_MODEL_CATALOG_BYTES: usize = 1024 * 1024; - const MAX_CATALOG_MODELS: usize = 2_048; const MAX_DISPLAY_NAME_BYTES: usize = 256; const MAX_DESCRIPTION_BYTES: usize = 4 * 1024; @@ -256,9 +253,6 @@ impl CodexModelCatalogSnapshot { /// Codex 模型目录不满足完整快照约束 #[derive(Debug, Clone, Copy, PartialEq, Eq, thiserror::Error)] pub enum CodexModelCatalogError { - /// 响应超过硬上限 - #[error("Codex model catalog response exceeds the byte limit")] - ResponseTooLarge, /// 顶层不是唯一受支持的 `{models:[...]}` wire #[error("Codex model catalog response violates the official wire contract")] InvalidWire, @@ -351,9 +345,6 @@ fn parse_native_model_catalog( etag: Option<&str>, protocol: OpenAiUpstreamProtocol, ) -> Result { - if body.len() > MAX_CODEX_MODEL_CATALOG_BYTES { - return Err(CodexModelCatalogError::ResponseTooLarge); - } let wire: CodexModelsWire = serde_json::from_slice(body).map_err(|_| CodexModelCatalogError::InvalidWire)?; if wire.models.is_empty() { @@ -615,9 +606,6 @@ pub(crate) fn parse_api_model_catalog( id: String, display_name: Option, } - if body.len() > MAX_CODEX_MODEL_CATALOG_BYTES { - return Err(CodexModelCatalogError::ResponseTooLarge); - } let value: serde_json::Value = serde_json::from_slice(body).map_err(|_| CodexModelCatalogError::InvalidWire)?; // 部分 Responses 上游按客户端版本返回完整 Codex 目录;不能降格为只有 ID 的画像 diff --git a/backend/crates/providers/openai/src/transport/client.rs b/backend/crates/providers/openai/src/transport/client.rs index 280caeb66..eceb14b68 100644 --- a/backend/crates/providers/openai/src/transport/client.rs +++ b/backend/crates/providers/openai/src/transport/client.rs @@ -630,6 +630,7 @@ pub type CodexResponseMetadataUpdates = CodexWebSocketResponseMetadataUpdates; /// Codex Responses 上游流式响应 pub struct CodexBackendStreamingResponse { + pub(crate) routing_affinity: Option, /// HTTP 为 SSE 字节块,WebSocket 为完整 UTF-8 文本消息 pub body: CodexBackendSseStream, /// 实际使用的上游传输 @@ -696,6 +697,7 @@ impl OpenAiUpstreamProtocol { /// Codex HTTP/SSE 上游客户端 #[derive(Clone)] pub struct CodexBackendClient { + pub(super) routing_affinity: Option, pub(super) privacy: Option<( Arc, gateway_core::lifecycle::CancellationToken, @@ -718,6 +720,14 @@ pub struct CodexBackendClient { } impl CodexBackendClient { + pub(crate) fn with_routing_affinity( + mut self, + affinity: Option, + ) -> Self { + self.routing_affinity = affinity; + self + } + #[must_use] pub(crate) fn with_timezone( mut self, diff --git a/backend/crates/providers/openai/src/transport/client_sse.rs b/backend/crates/providers/openai/src/transport/client_sse.rs index 34229240b..7f618dc83 100644 --- a/backend/crates/providers/openai/src/transport/client_sse.rs +++ b/backend/crates/providers/openai/src/transport/client_sse.rs @@ -17,10 +17,7 @@ use reqwest::{ use tokio_tungstenite::tungstenite::handshake::client::generate_key; use crate::transport::{ - catalog::{ - CodexModelCatalogError, CodexModelCatalogSnapshot, MAX_CODEX_MODEL_CATALOG_BYTES, - catalog_etag, parse_codex_model_catalog, - }, + catalog::{CodexModelCatalogSnapshot, catalog_etag, parse_codex_model_catalog}, diagnostics::CodexUpstreamSendPhase, endpoints::endpoint_url, headers::websocket_header_pairs, @@ -66,6 +63,7 @@ impl CodexBackendClient { ) -> Self { let base_url = base_url.into().trim_end_matches('/').to_string(); Self { + routing_affinity: None, privacy: None, timezone: Default::default(), connection_budget: None, @@ -125,6 +123,10 @@ impl CodexBackendClient { let profile = self.profile.snapshot(); let mut headers = self.request_headers_for_http_response(&profile, upstream_request, context)?; + let endpoint = endpoint_url(&self.base_url, self.protocol.responses_path()); + if let Some(affinity) = &self.routing_affinity { + affinity.apply(&endpoint, &mut headers); + } // OAuth 请求遵循 Codex 压缩合同;API Key 上游使用普通 JSON // Codex 上游只交付 SSE;即使下游请求 `stream: false`,也要上游流式执行, // 再由 API 层收集 canonical events 并返回完整 JSON @@ -145,7 +147,6 @@ impl CodexBackendClient { self.apply_privacy(&mut upstream_body, &mut headers)?; let body = serde_json::to_vec(&upstream_body).map_err(CodexClientError::RequestBodyEncode)?; - let endpoint = endpoint_url(&self.base_url, self.protocol.responses_path()); let trace = context .trace .cloned() @@ -261,6 +262,7 @@ impl CodexBackendClient { let rate_limit_updates = Arc::new(tokio::sync::Mutex::new(Vec::new())); Ok(CodexBackendStreamingResponse { + routing_affinity: None, body: http_sse_stream(response, Arc::clone(&rate_limit_updates), trace), transport: CodexBackendTransport::HttpSse, websocket_connection_id: None, @@ -324,6 +326,12 @@ impl CodexBackendClient { let profile = self.profile.snapshot(); let mut headers = self.request_headers_for_websocket_response(&profile, &websocket_request, context)?; + if let Some(affinity) = &self.routing_affinity { + affinity.apply( + &endpoint_url(&self.base_url, self.protocol.responses_path()), + &mut headers, + ); + } let mut body = serde_json::Value::Object(std::mem::take(websocket_request.body_mut())); profile.apply_context_versions(&mut body); super::correlation::scope_request(&mut body, &mut headers, context.installation_id); @@ -600,6 +608,9 @@ impl CodexBackendClient { })); } Ok(CodexBackendStreamingResponse { + routing_affinity: (self.protocol == OpenAiUpstreamProtocol::Codex) + .then_some(exchange.routing_affinity) + .flatten(), body: Box::pin( exchange .body @@ -644,6 +655,7 @@ impl CodexBackendClient { .with_egress_key(&self.egress_key) .with_outbound_proxy(self.outbound_proxy.as_ref()) .with_connection_profile(connection_profile); + key = key.with_routing_affinity(self.routing_affinity.as_ref().map(|pin| pin.id())); if let Some(connection_id) = request.downstream_websocket_connection_id.as_deref() { key = key.with_downstream_connection_id(connection_id); } @@ -684,9 +696,9 @@ impl CodexBackendClient { .then(|| catalog_etag(response.headers())) .transpose()? .flatten(); - let body = read_model_catalog_body(response).await?; + let body = response.bytes().await?; if !status.is_success() { - let body = String::from_utf8_lossy_owned(body); + let body = String::from_utf8_lossy(&body).into_owned(); return Err(CodexClientError::Upstream { status, retry_after_seconds: retry_after_seconds @@ -721,28 +733,6 @@ fn is_websocket_replayable_prelude(frame: &[u8]) -> bool { }) } -async fn read_model_catalog_body(response: ReqwestResponse) -> CodexClientResult> { - if response - .content_length() - .is_some_and(|length| length > MAX_CODEX_MODEL_CATALOG_BYTES as u64) - { - return Err(CodexModelCatalogError::ResponseTooLarge.into()); - } - let mut body = Vec::new(); - let mut stream = response.bytes_stream(); - while let Some(chunk) = stream.next().await { - let chunk = chunk?; - let Some(next_len) = body.len().checked_add(chunk.len()) else { - return Err(CodexModelCatalogError::ResponseTooLarge.into()); - }; - if next_len > MAX_CODEX_MODEL_CATALOG_BYTES { - return Err(CodexModelCatalogError::ResponseTooLarge.into()); - } - body.extend_from_slice(&chunk); - } - Ok(body) -} - fn websocket_connection_profile( headers: &HeaderMap, middleware_headers: &[gateway_core::engine::middleware::MiddlewareHeader], diff --git a/backend/crates/providers/openai/src/transport/mod.rs b/backend/crates/providers/openai/src/transport/mod.rs index 58aa1fcf5..febd5520c 100644 --- a/backend/crates/providers/openai/src/transport/mod.rs +++ b/backend/crates/providers/openai/src/transport/mod.rs @@ -20,6 +20,7 @@ pub mod protocol; pub mod request; pub mod reset_credits; mod response_meta; +pub(crate) mod routing_affinity; pub(crate) mod session; pub mod subscription; pub(crate) use endpoints::parse_upstream_base_url; @@ -32,7 +33,7 @@ pub use self::{ catalog::{ CodexCatalogCapabilities, CodexCatalogCapabilityEvidence, CodexCatalogLimits, CodexCatalogMetadata, CodexCatalogModel, CodexCatalogVisibility, CodexModelCatalogError, - CodexModelCatalogSnapshot, MAX_CODEX_MODEL_CATALOG_BYTES, parse_codex_model_catalog, + CodexModelCatalogSnapshot, parse_codex_model_catalog, }, client::{ CodexAccountSelectionTelemetry, CodexBackendClient, CodexBackendJsonResponse, diff --git a/backend/crates/providers/openai/src/transport/routing_affinity.rs b/backend/crates/providers/openai/src/transport/routing_affinity.rs new file mode 100644 index 000000000..3facc949b --- /dev/null +++ b/backend/crates/providers/openai/src/transport/routing_affinity.rs @@ -0,0 +1,105 @@ +//! Responses 连接握手的路由 Cookie 快照,仅用于同一上游端点的 Guardian 请求 + +use std::{fmt, sync::Arc}; + +use reqwest::{ + cookie::{CookieStore, Jar}, + header::{COOKIE, HeaderMap, HeaderValue, SET_COOKIE}, +}; +use url::Url; +use uuid::Uuid; + +const ROUTING_COOKIE_NAMES: [&str; 2] = ["__oailb", "__cflb"]; + +#[derive(Clone)] +pub(crate) struct ChatGptRoutingAffinity { + id: Uuid, + url: Url, + jar: Arc, +} + +impl fmt::Debug for ChatGptRoutingAffinity { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + formatter + .debug_struct("ChatGptRoutingAffinity") + .finish_non_exhaustive() + } +} + +impl ChatGptRoutingAffinity { + pub(crate) fn from_response(endpoint: &str, headers: &HeaderMap) -> Option { + let url = cookie_url(endpoint)?; + let mut cookies = Vec::with_capacity(2); + for name in ROUTING_COOKIE_NAMES { + let mut matching = headers.get_all(SET_COOKIE).iter().filter(|header| { + header + .to_str() + .ok() + .and_then(|value| value.split_once('=')) + .is_some_and(|(candidate, _)| candidate.trim() == name) + }); + if let Some(cookie) = matching.next() { + // 与 Codex 0.166 的连接快照一致,歧义或超长的 pin 不用于亲和 + if cookie.as_bytes().len() > 8192 || matching.next().is_some() { + return None; + } + cookies.push(cookie); + } + } + let jar = Jar::default(); + jar.set_cookies(&mut cookies.into_iter(), &url); + jar.cookies(&url)?; + Some(Self { + id: Uuid::new_v4(), + url, + jar: Arc::new(jar), + }) + } + + pub(crate) const fn id(&self) -> Uuid { + self.id + } + + pub(crate) fn apply(&self, endpoint: &str, headers: &mut HeaderMap) { + let Some(url) = cookie_url(endpoint) else { + return; + }; + // 网关可配置 OAuth backend;快照只信任实际握手的 origin 与路径 + if url.origin() != self.url.origin() || url.path() != self.url.path() { + return; + } + let Some(pin) = self.jar.cookies(&url) else { + return; + }; + let Ok(pin) = pin.to_str() else { return }; + let mut parts: Vec<_> = headers + .get_all(COOKIE) + .iter() + .filter_map(|value| value.to_str().ok()) + .flat_map(|value| value.split(';')) + .map(str::trim) + .filter(|part| { + !part.is_empty() + && !part + .split_once('=') + .is_some_and(|(name, _)| ROUTING_COOKIE_NAMES.contains(&name.trim())) + }) + .collect(); + parts.push(pin); + if let Ok(mut value) = HeaderValue::from_str(&parts.join("; ")) { + value.set_sensitive(true); + headers.insert(COOKIE, value); + } + } +} + +fn cookie_url(endpoint: &str) -> Option { + let mut url = Url::parse(endpoint).ok()?; + let scheme = match url.scheme() { + "https" | "wss" => "https", + "http" | "ws" => "http", + _ => return None, + }; + url.set_scheme(scheme).ok()?; + Some(url) +} diff --git a/backend/crates/providers/openai/src/transport/websocket/exchange/mod.rs b/backend/crates/providers/openai/src/transport/websocket/exchange/mod.rs index 50ce7b8b7..5518b63be 100644 --- a/backend/crates/providers/openai/src/transport/websocket/exchange/mod.rs +++ b/backend/crates/providers/openai/src/transport/websocket/exchange/mod.rs @@ -25,6 +25,7 @@ pub(super) use self::stream::{WebSocketStreamPoolReturn, stream_websocket_respon /// Responses WebSocket 文本流交互结果 pub struct CodexWebSocketStreamingExchange { + pub(crate) routing_affinity: Option, /// 关联请求日志与底层 pump 生命周期日志的连接标识 pub(crate) websocket_connection_id: Uuid, /// 上游完整 WebSocket 文本消息流 diff --git a/backend/crates/providers/openai/src/transport/websocket/exchange/stream.rs b/backend/crates/providers/openai/src/transport/websocket/exchange/stream.rs index 486a5fe81..4133a6969 100644 --- a/backend/crates/providers/openai/src/transport/websocket/exchange/stream.rs +++ b/backend/crates/providers/openai/src/transport/websocket/exchange/stream.rs @@ -116,6 +116,7 @@ pub(in crate::transport::websocket) fn stream_websocket_response( }); CodexWebSocketStreamingExchange { + routing_affinity: response_metadata.routing_affinity, websocket_connection_id, body: Box::pin(body), turn_state: response_metadata.turn_state, diff --git a/backend/crates/providers/openai/src/transport/websocket/handshake.rs b/backend/crates/providers/openai/src/transport/websocket/handshake.rs index a2084dbce..43e39774e 100644 --- a/backend/crates/providers/openai/src/transport/websocket/handshake.rs +++ b/backend/crates/providers/openai/src/transport/websocket/handshake.rs @@ -181,6 +181,10 @@ pub(super) fn websocket_connection_metadata( connection: &CodexWebSocketConnection, ) -> CodexWebSocketConnectionMetadata { CodexWebSocketConnectionMetadata { + routing_affinity: crate::transport::routing_affinity::ChatGptRoutingAffinity::from_response( + &connection.endpoint, + response.headers(), + ), turn_state: response_meta::turn_state(response.headers()), set_cookie_headers: response_meta::set_cookie_headers(response.headers()), rate_limit_headers: response_meta::rate_limit_headers(response.headers()), diff --git a/backend/crates/providers/openai/src/transport/websocket/pool/state.rs b/backend/crates/providers/openai/src/transport/websocket/pool/state.rs index 981aeeffc..09e4366a1 100644 --- a/backend/crates/providers/openai/src/transport/websocket/pool/state.rs +++ b/backend/crates/providers/openai/src/transport/websocket/pool/state.rs @@ -24,6 +24,7 @@ pub struct CodexWebSocketPoolKey { account_id: String, conversation_id: String, connection_profile: String, + routing_affinity_id: Option, pub(super) downstream_connection_id: String, egress_key: String, pub(super) proxy_capacity_key: String, @@ -41,6 +42,7 @@ impl CodexWebSocketPoolKey { account_id: account_id.into(), conversation_id: conversation_id.into(), connection_profile: String::new(), + routing_affinity_id: None, downstream_connection_id: String::new(), egress_key: String::new(), proxy_capacity_key: String::new(), @@ -66,6 +68,11 @@ impl CodexWebSocketPoolKey { &self.account_id } + pub(crate) fn with_routing_affinity(mut self, id: Option) -> Self { + self.routing_affinity_id = id; + self + } + pub(crate) fn with_egress_key(mut self, key: &str) -> Self { self.egress_key = key.to_owned(); self @@ -84,6 +91,10 @@ impl CodexWebSocketPoolKey { } pub(crate) fn stable_hash(&self) -> String { + let routing_affinity = self + .routing_affinity_id + .map(|id| id.to_string()) + .unwrap_or_default(); short_sha256([ self.base_url.as_str(), self.account_id.as_str(), @@ -91,6 +102,7 @@ impl CodexWebSocketPoolKey { self.connection_profile.as_str(), self.downstream_connection_id.as_str(), self.egress_key.as_str(), + routing_affinity.as_str(), ]) } @@ -99,6 +111,7 @@ impl CodexWebSocketPoolKey { && self.account_id == other.account_id && self.conversation_id == other.conversation_id && self.egress_key == other.egress_key + && self.routing_affinity_id == other.routing_affinity_id } } @@ -171,6 +184,7 @@ impl WebSocketPoolState { #[derive(Clone)] pub(crate) struct CodexWebSocketConnectionMetadata { + pub(crate) routing_affinity: Option, pub(crate) turn_state: Option, pub(crate) set_cookie_headers: Vec, pub(crate) rate_limit_headers: Vec<(String, String)>, diff --git a/backend/crates/providers/openai/tests/credential/catalog.rs b/backend/crates/providers/openai/tests/credential/catalog.rs index 38b0d50f6..c76261a95 100644 --- a/backend/crates/providers/openai/tests/credential/catalog.rs +++ b/backend/crates/providers/openai/tests/credential/catalog.rs @@ -1573,3 +1573,41 @@ async fn catalog_account_membership_changes_advance_generation_without_model_cha BTreeSet::from([first.id().clone(), second.id().clone()]) ); } + +#[tokio::test] +async fn client_catalog_byte_budget_evicts_completed_entries_without_rejecting_large_catalogs() { + let store = Arc::new(MemoryAccountStore::default()); + let account = seed_account(&store, "acct_catalog_bytes").await; + let scope = client_scope(&[account]); + let server = MockServer::start().await; + for (version, bytes, calls) in [ + ("small", 1024, 1), + ("large", 17 * 1024 * 1024, 2), + ("other", 17 * 1024 * 1024, 1), + ("oversized", 33 * 1024 * 1024, 2), + ] { + Mock::given(method("GET")).and(path("/codex/models")) + .and(query_param("client_version", version)) + .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({"models":[{ + "slug":"gpt-large", "display_name":"Large", "base_instructions":"x".repeat(bytes) + }]}))).expect(calls).mount(&server).await; + } + let service = service_with_catalog_cache(&store, server.uri(), catalog_cache()); + for version in [ + "large", + "other", + "other", + "large", + "small", + "oversized", + "oversized", + "small", + ] { + let models = service + .client_model_catalog(&scope, version) + .await + .expect("large catalog delivered"); + assert_eq!(models.len(), 1); + } + server.verify().await; +} diff --git a/backend/crates/providers/openai/tests/provider/contract/account_isolation.rs b/backend/crates/providers/openai/tests/provider/contract/account_isolation.rs index 7d92dc4bc..4ee7cc1c8 100644 --- a/backend/crates/providers/openai/tests/provider/contract/account_isolation.rs +++ b/backend/crates/providers/openai/tests/provider/contract/account_isolation.rs @@ -281,3 +281,293 @@ async fn middleware_can_override_account_headers_on_http_and_websocket() { } } } + +struct RoutingUpstream { + url: String, + requests: Arc>>, + task: tokio::task::JoinHandle<()>, +} + +impl Drop for RoutingUpstream { + fn drop(&mut self) { + self.task.abort(); + } +} + +impl RoutingUpstream { + async fn start(cookies: Option>) -> Self { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let url = format!("http://{}", listener.local_addr().unwrap()); + let requests = Arc::new(Mutex::new(Vec::new())); + let capture = Arc::clone(&requests); + let task = tokio::spawn(async move { + let mut connections = tokio::task::JoinSet::new(); + let mut index = 0; + loop { + let (socket, _) = listener.accept().await.unwrap(); + let capture = Arc::clone(&capture); + let cookies = cookies.clone().unwrap_or_else(|| { + vec![ + format!("__oailb=pin-{index}; Path=/codex/responses; Max-Age=300"), + format!("__cflb=edge-{index}; Path=/; Max-Age=300"), + ] + }); + connections.spawn(async move { + let mut headers = reqwest::header::HeaderMap::new(); + let mut socket = crate::transport::accept_codex_test_websocket_with(socket, |request, response| { + headers = request.headers().clone(); + for cookie in cookies { + response.headers_mut().append("set-cookie", cookie.parse().unwrap()); + } + }).await; + while let Some(Ok(message)) = socket.next().await { + let Message::Text(text) = message else { continue }; + let body: Value = serde_json::from_str(&text).unwrap(); + let id = body["client_metadata"]["fixture_response_id"].as_str().unwrap().to_owned(); + capture.lock().unwrap().push((index, headers.clone(), body)); + for event in ["response.created", "response.completed"] { + socket.send(Message::Text(json!({ + "type":event, "response":{"id":id,"model":"gpt-5.4","status":"completed","output":[]} + }).to_string().into())).await.unwrap(); + } + } + }); + index += 1; + } + }); + Self { + url, + requests, + task, + } + } +} + +fn routing_operation(id: &str, parent: Option<&str>, conversation: Option<&str>) -> Operation { + let mut metadata = json!({"fixture_response_id":id}); + if let Some(parent) = parent { + metadata["parent_response_id"] = json!(parent); + metadata["x-openai-subagent"] = json!("guardian"); + } + let mut protocol_context = Map::from_iter([("use_websocket".to_owned(), json!(true))]); + if parent.is_some() { + protocol_context.insert( + "opaque_request_headers".to_owned(), + json!([["x-codex-guardian", "cmV2aWV3ZXI="]]), + ); + } + if let Some(conversation) = conversation { + protocol_context.insert( + "downstream_websocket_connection_id".to_owned(), + json!(conversation), + ); + } + let payload = ProtocolPayload::json_object( + "openai", + json!({ + "model":"gpt-5.4", "input":"hello", "client_metadata":metadata + }) + .as_object() + .unwrap() + .clone(), + ) + .unwrap() + .with_context(protocol_context); + let mut generate = GenerateRequest::from_protocol_payload(payload); + if let Some(conversation) = conversation { + generate = generate.with_provider_session_state(ProviderSessionState::new("openai", json!({ + "account_id":SELECTED_ACCOUNT, "conversation_id":conversation, "continuation_scope":"persisted" + }).as_object().unwrap().clone()).unwrap()); + } + Operation::Generate(generate) +} + +fn routing_context(id: &str, key: &str, account: &str) -> AttemptContext { + AttemptContext::new( + RequestAttemptContext::new( + ModelRequestId::new(format!("req_{id}")).unwrap(), + ClientApiKeyId::new(key).unwrap(), + ), + NonZeroU32::new(1).unwrap(), + SystemTime::now() + Duration::from_secs(30), + account_policy(), + AccountAttemptContext::new( + BTreeSet::new(), + Some(ProviderAccountId::new(account).unwrap()), + None, + ) + .with_account_scope(contract_account_scope()), + None, + CancellationToken::new(), + ) +} + +async fn execute_routing( + provider: &Arc, + operation: Operation, + context: AttemptContext, +) { + let mut stream = Arc::clone(provider) + .execute(planned_request("openai", operation), context) + .await + .unwrap(); + while let Some(event) = stream.next().await { + event.unwrap(); + } +} + +#[tokio::test] +async fn guardian_pins_follow_parent_connection_and_separate_reviewer_pool_reuse() { + let store = Arc::new(MemoryAccountStore::default()); + create_account(&store, SELECTED_ACCOUNT).await; + let server = RoutingUpstream::start(None).await; + let provider = provider_with_base_url(&store, server.url.clone()); + for (id, parent, conversation) in [ + ("parent_a", None, "root_a"), + ("parent_b", None, "root_b"), + ("review_a", Some("parent_a"), "reviewer"), + ("review_a_again", Some("parent_a"), "reviewer"), + ("review_b", Some("parent_b"), "reviewer"), + ] { + execute_routing( + &provider, + routing_operation(id, parent, Some(conversation)), + routing_context(id, "key_routing", SELECTED_ACCOUNT), + ) + .await; + } + let requests = server.requests.lock().unwrap(); + assert_eq!(requests.len(), 5); + assert_eq!( + requests[2].0, requests[3].0, + "same parent keeps the reviewer socket" + ); + assert_ne!( + requests[3].0, requests[4].0, + "another parent needs its own routed socket" + ); + for (review, parent) in [(2, 0), (3, 0), (4, 1)] { + let cookie = requests[review].1["cookie"].to_str().unwrap(); + assert!(cookie.contains(&format!("__oailb=pin-{parent}"))); + assert!(cookie.contains(&format!("__cflb=edge-{parent}"))); + } + assert!( + !requests[1].1.contains_key("cookie"), + "ordinary traffic cannot inherit another connection's pins" + ); +} + +#[tokio::test] +async fn guardian_parent_pins_require_matching_client_key_and_selected_account() { + let store = Arc::new(MemoryAccountStore::default()); + create_account(&store, SELECTED_ACCOUNT).await; + create_account(&store, "acct_scope_old").await; + let server = RoutingUpstream::start(None).await; + let provider = provider_with_base_url(&store, server.url.clone()); + execute_routing( + &provider, + routing_operation("parent", None, None), + routing_context("parent", "key_routing", SELECTED_ACCOUNT), + ) + .await; + for (id, key, account, parent) in [ + ("same", "key_routing", SELECTED_ACCOUNT, "parent"), + ("other_key", "key_other", SELECTED_ACCOUNT, "parent"), + ("other_account", "key_routing", "acct_scope_old", "parent"), + ("unknown", "key_routing", SELECTED_ACCOUNT, "unknown_parent"), + ] { + execute_routing( + &provider, + routing_operation(id, Some(parent), None), + routing_context(id, key, account), + ) + .await; + } + let requests = server.requests.lock().unwrap(); + assert_eq!(requests.len(), 5); + for (index, (_, headers, body)) in requests.iter().enumerate().skip(1) { + assert_eq!(headers.contains_key("cookie"), index == 1); + assert_eq!( + body.pointer("/client_metadata/parent_response_id") + .and_then(Value::as_str), + (index == 1).then_some("parent") + ); + } +} + +#[tokio::test] +async fn guardian_cookie_expiry_keeps_existing_socket_but_does_not_refresh_parent_pin() { + let store = Arc::new(MemoryAccountStore::default()); + create_account(&store, SELECTED_ACCOUNT).await; + let server = RoutingUpstream::start(Some(vec![ + "__oailb=short_lived; Path=/; Max-Age=1".to_owned(), + ])) + .await; + let provider = provider_with_base_url(&store, server.url.clone()); + for (id, parent, conversation) in [ + ("parent_a", None, "root"), + ("review_a", Some("parent_a"), "reviewer"), + ] { + execute_routing( + &provider, + routing_operation(id, parent, Some(conversation)), + routing_context(id, "key_routing", SELECTED_ACCOUNT), + ) + .await; + } + tokio::time::sleep(Duration::from_millis(1100)).await; + for (id, parent, conversation) in [ + ("parent_b", None, "root"), + ("review_b", Some("parent_b"), "reviewer"), + ("review_new", Some("parent_b"), "another_reviewer"), + ] { + execute_routing( + &provider, + routing_operation(id, parent, Some(conversation)), + routing_context(id, "key_routing", SELECTED_ACCOUNT), + ) + .await; + } + let requests = server.requests.lock().unwrap(); + assert_eq!(requests.len(), 5); + assert_eq!(requests[0].0, requests[2].0); + assert_eq!(requests[1].0, requests[3].0); + assert_ne!(requests[3].0, requests[4].0); + assert!(requests[1].1.contains_key("cookie")); + assert!(!requests[4].1.contains_key("cookie")); + assert_eq!( + requests[4].2["client_metadata"]["parent_response_id"], + "parent_b" + ); +} + +#[tokio::test] +async fn guardian_ignores_expired_ambiguous_and_wrong_path_routing_cookies() { + for cookies in [ + vec!["__oailb=expired; Path=/; Max-Age=0"], + vec!["__oailb=first; Path=/", "__oailb=second; Path=/"], + vec!["__oailb=other; Path=/other"], + ] { + let store = Arc::new(MemoryAccountStore::default()); + create_account(&store, SELECTED_ACCOUNT).await; + let server = + RoutingUpstream::start(Some(cookies.into_iter().map(str::to_owned).collect())).await; + let provider = provider_with_base_url(&store, server.url.clone()); + for (id, parent) in [("parent", None), ("review", Some("parent"))] { + execute_routing( + &provider, + routing_operation(id, parent, None), + routing_context(id, "key_routing", SELECTED_ACCOUNT), + ) + .await; + } + let requests = server.requests.lock().unwrap(); + assert!(!requests[1].1.contains_key("cookie")); + assert!( + requests[1] + .2 + .pointer("/client_metadata/parent_response_id") + .is_none() + ); + } +} diff --git a/backend/crates/providers/openai/tests/transport/catalog.rs b/backend/crates/providers/openai/tests/transport/catalog.rs index 94446e5a3..ea4ea1f09 100644 --- a/backend/crates/providers/openai/tests/transport/catalog.rs +++ b/backend/crates/providers/openai/tests/transport/catalog.rs @@ -4,9 +4,8 @@ use chrono::{TimeZone, Utc}; use gateway_core::routing::ModelServiceTier; use provider_openai::transport::profile::{CodexWireProfile, CodexWireProfileState}; use provider_openai::transport::{ - CodexBackendClient, CodexCatalogCapabilityEvidence, CodexCatalogVisibility, CodexClientError, - CodexModelCatalogError, CodexRequestContext, MAX_CODEX_MODEL_CATALOG_BYTES, - build_reqwest_client, parse_codex_model_catalog, + CodexBackendClient, CodexCatalogCapabilityEvidence, CodexCatalogVisibility, + CodexModelCatalogError, CodexRequestContext, build_reqwest_client, parse_codex_model_catalog, }; use serde_json::{Value, json}; use wiremock::matchers::{header, method, path, query_param}; @@ -318,17 +317,6 @@ fn invalid_request_slug_should_fail_the_entire_snapshot() { )); } -#[test] -fn body_over_hard_limit_should_fail_before_json_parsing() { - let body = vec![b' '; MAX_CODEX_MODEL_CATALOG_BYTES + 1]; - let result = parse_codex_model_catalog(&body, None); - - assert!(matches!( - result, - Err(CodexModelCatalogError::ResponseTooLarge) - )); -} - #[tokio::test] async fn fetch_should_send_official_catalog_headers_and_query() { let server = MockServer::start().await; @@ -386,15 +374,16 @@ async fn fetch_should_send_official_catalog_headers_and_query() { } #[tokio::test] -async fn fetch_should_reject_streamed_body_over_hard_limit() { +async fn fetch_preserves_large_model_instructions() { + let instructions = "model instructions\n".repeat(100_000); + let body = json!({"models":[{ + "slug":"gpt-native", "display_name":"Native", "base_instructions":instructions, + "model_messages":{"future":{"template":"preserved"}} + }]}); let server = MockServer::start().await; Mock::given(method("GET")) .and(path("/codex/models")) - .respond_with(ResponseTemplate::new(200).set_body_bytes(vec![ - b' '; - MAX_CODEX_MODEL_CATALOG_BYTES - + 1 - ])) + .respond_with(ResponseTemplate::new(200).set_body_json(&body)) .expect(1) .mount(&server) .await; @@ -404,14 +393,12 @@ async fn fetch_should_reject_streamed_body_over_hard_limit() { profile(), ); - let result = client.fetch_models_with_context(context(), None).await; - - assert!(matches!( - result, - Err(CodexClientError::ModelCatalog( - CodexModelCatalogError::ResponseTooLarge - )) - )); + let result = client + .fetch_models_with_context(context(), Some("0.166.0")) + .await + .unwrap(); + let preserved: Value = serde_json::from_slice(result.models()[0].document().body()).unwrap(); + assert!(preserved == body["models"][0]); } fn profile() -> CodexWireProfileState { diff --git a/backend/crates/providers/openai/tests/transport/mod.rs b/backend/crates/providers/openai/tests/transport/mod.rs index 328d10df7..108e9d577 100644 --- a/backend/crates/providers/openai/tests/transport/mod.rs +++ b/backend/crates/providers/openai/tests/transport/mod.rs @@ -261,6 +261,7 @@ async fn collect_backend_response( response_metadata, mut transport_metrics, connection_local_continuation, + .. } = response; let mut body_bytes = Vec::new(); while let Some(chunk) = body.next().await { diff --git a/docs/api.md b/docs/api.md index a4b3c0936..86c8a14ad 100644 --- a/docs/api.md +++ b/docs/api.md @@ -452,6 +452,10 @@ turn metadata 的 `workspaces` 也会保留仓库绝对路径、Git 远端地址 切号或归属未知时移除。`x-codex-guardian`、`guardian_credits_requested` 和序列化 `x-codex-turn-metadata` 内普通扩展的同名 `parent_response_id` 保持原样 +同一 Client Key、同一 OAuth 账号的 Guardian 请求可利用网关记录的父响应 WS 握手信息, +尽力复用父请求的上游路由。路由 Cookie 仅在网关与上游之间使用,不传给客户端; +网关重启、记录淘汰或 Cookie 到期后仍可按普通路由请求,不保证上游路由亲和 + #### WebSocket 与上游传输 公开 HTTP 请求在路由匹配前经过已绑定的 `http` 中间件;插件可以改写握手路径和请求参数。 diff --git a/docs/architecture.md b/docs/architecture.md index 881790164..90468b9f3 100644 --- a/docs/architecture.md +++ b/docs/architecture.md @@ -330,7 +330,8 @@ Provider 可为模型提供来源账号集合;Core 在公开目录查询时结 不解析上游字段;API 负责协议输出,不直接依赖具体 Provider。原生目录失败不会退回通用画像,只有不提供 该协议原生目录的 Provider(如 xAI 的 Codex 适配)才走画像转换。目录读取不创建计量请求或推理 attempt。 OpenAI 的账号/凭据 revision/客户端版本隔离、并发合并、TTL、容量与失效均归属 credential catalog -service;客户端原生对象保存在有界进程缓存中,与套餐 Redis 模型 ID 缓存分开管理 +service;客户端原生对象保存在有界进程缓存中,与套餐 Redis 模型 ID 缓存分开管理。 +该缓存最多保留 32 项,原生正文总预算为 32 MiB;单次结果超过预算时仍正常交付,但不驻留缓存 全局请求位置及开关由运行设置持久化;快照仅在开关开启时生成全局覆盖,并随 `RuntimeSnapshot → RoutingPlan → AttemptContext` 冻结传递,关闭时保留客户端原有字段; @@ -493,6 +494,11 @@ OpenAI 的连接池监督该取消信号,按当前归属保护绑定连接; 回收先移出可复用集合,再在锁外有界关闭并等待 pump/socket 释放,物理容量随 socket 归还;批量关闭并行执行。 协议合同与运行参数分别见 [Responses 传输](api.md#websocket-与上游传输)和[部署与运维](../deploy/README.md#公网访问) +OpenAI Provider 按 Client Key、所选 OAuth 账号和父响应 ID 关联 WS 握手的路由 Cookie 快照, +只供 Guardian 使用,不参与选号或写入账号凭据。进程内最多保存 4096 个父响应,保留 30 分钟; +缺失或失效时按普通上游路由执行。transport 负责 Cookie 有效期和同 origin、同端点路径约束, +连接池按快照身份隔离审查连接;复用父连接不会重新计算 Cookie 有效期,Cookie 到期不强制关闭已建立的审查连接 + OpenAI 的 OAuth 与 API Key 共用现有账号和事务。API Key 的 Base URL、密钥和传输策略属于 Provider 凭据 JSON, 随 credential revision 更新;普通详情只投影非敏感连接设置。API Key 客户端目录按账号、凭据版本和客户端版本隔离, 标准 `data` 模型列表通过通用画像输出,完整 Codex `models` 目录保留原生对象。通用账号层按 Provider 提交的