From 98250268b0526d3f2cb07b50a8e9568cf741fb28 Mon Sep 17 00:00:00 2001 From: zyycn Date: Fri, 24 Jul 2026 16:19:30 +0800 Subject: [PATCH] fix(gateway): restore legacy session and provider semantics Recover OpenAI continuation affinity, session-level account exclusion, catalog/configuration behavior, and provider-owned runtime defaults without changing the transparent OpenAI wire path. Co-authored-by: Codex --- .../crates/gateway-api/src/openai/models.rs | 60 ++++ .../crates/gateway-api/src/openai/router.rs | 4 +- .../crates/gateway-api/src/openai/service.rs | 6 + .../crates/gateway-api/tests/openai/mod.rs | 8 +- .../crates/gateway-api/tests/openai/models.rs | 40 +++ .../gateway-core/src/engine/continuation.rs | 46 +++- .../gateway-core/src/engine/coordinator.rs | 36 ++- .../gateway-core/src/engine/execution.rs | 100 +++++-- backend/crates/gateway-core/src/engine/mod.rs | 7 + .../crates/gateway-core/src/operation/mod.rs | 3 +- .../gateway-core/src/provider_ports/mod.rs | 66 ++++- .../gateway-core/src/routing/snapshot.rs | 39 ++- .../gateway-core/tests/engine/execution.rs | 8 +- .../crates/gateway-core/tests/routing/mod.rs | 24 +- backend/crates/gateway-store/src/lib.rs | 9 +- .../gateway-store/src/postgres/history.rs | 148 +--------- backend/crates/gateway-store/src/redis/mod.rs | 4 + .../src/redis/native_continuation.rs | 221 +++++++++++++++ .../src/redis/provider_session_exclusion.rs | 194 +++++++++++++ .../gateway-store/tests/postgres/history.rs | 162 +---------- .../crates/gateway-store/tests/redis/mod.rs | 2 + .../tests/redis/native_continuation.rs | 123 +++++++++ .../tests/redis/provider_session_exclusion.rs | 113 ++++++++ backend/crates/providers/openai/src/config.rs | 256 +++++++++++++++++- .../openai/src/credential/affinity.rs | 41 +++ .../openai/src/credential/catalog.rs | 8 +- .../openai/src/credential/identity.rs | 45 ++- .../providers/openai/src/credential/mod.rs | 9 +- .../providers/openai/src/credential/oauth.rs | 19 +- .../providers/openai/src/credential/quota.rs | 29 +- .../openai/src/credential/selector.rs | 177 +++++++++++- .../openai/src/credential/token_client.rs | 22 +- backend/crates/providers/openai/src/lib.rs | 67 +++-- .../crates/providers/openai/src/provider.rs | 136 ++++++++-- .../providers/openai/src/transport/mod.rs | 1 + .../providers/openai/src/transport/request.rs | 66 +---- .../providers/openai/src/transport/session.rs | 125 +++++++++ .../crates/providers/openai/tests/admin.rs | 73 +++-- .../crates/providers/openai/tests/config.rs | 76 ++++-- .../openai/tests/credential/admin.rs | 3 +- .../openai/tests/credential/catalog.rs | 2 + .../openai/tests/credential/contract.rs | 56 +++- .../openai/tests/credential/quota.rs | 2 + .../openai/tests/provider/contract.rs | 11 +- .../crates/providers/openai/tests/support.rs | 80 +++++- .../openai/tests/transport/profile.rs | 52 +++- backend/crates/providers/xai/tests/admin.rs | 4 +- backend/crates/providers/xai/tests/support.rs | 42 ++- 48 files changed, 2260 insertions(+), 565 deletions(-) create mode 100644 backend/crates/gateway-store/src/redis/native_continuation.rs create mode 100644 backend/crates/gateway-store/src/redis/provider_session_exclusion.rs create mode 100644 backend/crates/gateway-store/tests/redis/native_continuation.rs create mode 100644 backend/crates/gateway-store/tests/redis/provider_session_exclusion.rs create mode 100644 backend/crates/providers/openai/src/transport/session.rs diff --git a/backend/crates/gateway-api/src/openai/models.rs b/backend/crates/gateway-api/src/openai/models.rs index 4460e9890..818fca48f 100644 --- a/backend/crates/gateway-api/src/openai/models.rs +++ b/backend/crates/gateway-api/src/openai/models.rs @@ -66,6 +66,50 @@ pub(crate) async fn model_detail( .into_response() } +/// `GET /v1/models/catalog`。 +/// +/// 新架构的目录由 core 按客户端 Provider 范围裁剪；这里不读取或重建任一 +/// Provider 的 wire catalog，只恢复 Codex 客户端依赖的稳定展示合同。 +pub(crate) async fn model_catalog(State(state): State, headers: HeaderMap) -> Response { + let service = state.openai(); + let client = match authenticate_client(service, &headers) { + Ok(client) => client, + Err(error) => return authentication_error_response(error), + }; + let models = service + .public_models(&client) + .into_iter() + .map(|model| catalog_model_json(model.as_str())) + .collect::>(); + + (StatusCode::OK, Json(models)).into_response() +} + +/// `GET /v1/models/{model_id}/info`。 +pub(crate) async fn model_info( + State(state): State, + headers: HeaderMap, + Path(model_id): Path, +) -> Response { + let service = state.openai(); + let client = match authenticate_client(service, &headers) { + Ok(client) => client, + Err(error) => return authentication_error_response(error), + }; + let Ok(public_model) = PublicModelId::new(model_id) else { + return model_not_found_response().into_response(); + }; + if !service.contains_public_model(&client, &public_model) { + return model_not_found_response().into_response(); + } + + ( + StatusCode::OK, + Json(catalog_model_json(public_model.as_str())), + ) + .into_response() +} + fn openai_model_json(id: &str) -> Value { json!({ "id": id, @@ -74,3 +118,19 @@ fn openai_model_json(id: &str) -> Value { "owned_by": "gateway", }) } + +fn catalog_model_json(id: &str) -> Value { + json!({ + "id": id, + "displayName": id, + "description": "", + "isDefault": false, + "supportedReasoningEfforts": [], + "defaultReasoningEffort": "", + "inputModalities": ["text"], + "outputModalities": ["text"], + "supportsPersonality": false, + "upgrade": Value::Null, + "source": "gateway", + }) +} diff --git a/backend/crates/gateway-api/src/openai/router.rs b/backend/crates/gateway-api/src/openai/router.rs index 4840f6c97..9a8b53f60 100644 --- a/backend/crates/gateway-api/src/openai/router.rs +++ b/backend/crates/gateway-api/src/openai/router.rs @@ -7,7 +7,7 @@ use axum::{ }; use super::{ - models::{model_detail, models}, + models::{model_catalog, model_detail, model_info, models}, responses::{responses, responses_websocket, review_responses}, }; @@ -22,6 +22,8 @@ pub(crate) fn router() -> Router { .route("/v1/responses", get(responses_websocket).post(responses)) .route("/v1/responses/review", post(review_responses)) .route("/v1/models", get(models)) + .route("/v1/models/catalog", get(model_catalog)) + .route("/v1/models/{model_id}/info", get(model_info)) // 官方 OpenAI 模型详情合同使用 path ID；它不属于 Admin API 约束。 .route("/v1/models/{model_id}", get(model_detail)) .layer(DefaultBodyLimit::max(MAX_CLIENT_REQUEST_BODY_BYTES)) diff --git a/backend/crates/gateway-api/src/openai/service.rs b/backend/crates/gateway-api/src/openai/service.rs index 07fee5277..3f9496de0 100644 --- a/backend/crates/gateway-api/src/openai/service.rs +++ b/backend/crates/gateway-api/src/openai/service.rs @@ -75,6 +75,12 @@ impl OpenAiService { "requested model was not found", ) })?; + if !self.execution.contains_public_model(&client, &public_model) { + return Err(GatewayError::new( + GatewayErrorKind::ModelNotFound, + "requested model was not found", + )); + } let previous_response_id = match metadata.continuation() { ContinuationIntent::None => None, ContinuationIntent::PreviousResponseId(value) => { diff --git a/backend/crates/gateway-api/tests/openai/mod.rs b/backend/crates/gateway-api/tests/openai/mod.rs index dd9066bf5..0889f6367 100644 --- a/backend/crates/gateway-api/tests/openai/mod.rs +++ b/backend/crates/gateway-api/tests/openai/mod.rs @@ -253,9 +253,15 @@ struct UnusedContinuation; impl NativeContinuationPort for UnusedContinuation { fn resolve<'a>( &'a self, - _: &'a ClientApiKeyId, _: &'a PreviousResponseId, ) -> BoxFuture<'a, Result, NativeContinuationStoreError>> { Box::pin(async { unreachable!("authentication fixture does not execute") }) } + + fn record<'a>( + &'a self, + _: NativeContinuationPin, + ) -> BoxFuture<'a, Result<(), NativeContinuationStoreError>> { + Box::pin(async { unreachable!("authentication fixture does not execute") }) + } } diff --git a/backend/crates/gateway-api/tests/openai/models.rs b/backend/crates/gateway-api/tests/openai/models.rs index b2f64fea1..511eb5e32 100644 --- a/backend/crates/gateway-api/tests/openai/models.rs +++ b/backend/crates/gateway-api/tests/openai/models.rs @@ -113,6 +113,46 @@ async fn model_detail_should_hide_unknown_models() { assert_eq!(response.status(), StatusCode::NOT_FOUND); } +#[tokio::test] +async fn model_catalog_should_keep_the_codex_catalog_contract() { + let response = api_router(ModelsExecution::new()) + .await + .oneshot(authorized_request("/v1/models/catalog")) + .await + .expect("catalog response"); + + assert_eq!(response.status(), StatusCode::OK); + let body = to_bytes(response.into_body(), usize::MAX) + .await + .expect("read catalog body"); + let value: serde_json::Value = serde_json::from_slice(&body).expect("catalog JSON"); + assert_eq!(value[0]["id"], "model-a"); + assert_eq!(value[0]["displayName"], "model-a"); + assert_eq!(value[0]["source"], "gateway"); +} + +#[tokio::test] +async fn model_info_should_return_a_single_visible_catalog_entry() { + let response = api_router(ModelsExecution::new()) + .await + .oneshot(authorized_request("/v1/models/model-a/info")) + .await + .expect("model info response"); + + assert_eq!(response.status(), StatusCode::OK); +} + +#[tokio::test] +async fn model_info_should_hide_unknown_models() { + let response = api_router(ModelsExecution::new()) + .await + .oneshot(authorized_request("/v1/models/model-private/info")) + .await + .expect("missing model info response"); + + assert_eq!(response.status(), StatusCode::NOT_FOUND); +} + #[tokio::test] async fn models_should_require_a_valid_bearer_client_key() { let response = api_router(ModelsExecution::new()) diff --git a/backend/crates/gateway-core/src/engine/continuation.rs b/backend/crates/gateway-core/src/engine/continuation.rs index 861fa34a6..94d180758 100644 --- a/backend/crates/gateway-core/src/engine/continuation.rs +++ b/backend/crates/gateway-core/src/engine/continuation.rs @@ -9,7 +9,7 @@ use futures::future::BoxFuture; use crate::engine::credential::ProviderAccountId; use crate::error::{IdentifierError, SafeUpstreamValue, validate_text}; -use crate::policy::ClientApiKeyId; +use crate::operation::ProviderSessionState; use crate::routing::ProviderKind; /// 客户端传入的 previous response ID。 @@ -43,19 +43,20 @@ pub enum NativeContinuationScope { ConnectionLocal, } -/// 外层已解析并认证的 native previous-response pin。 +/// 从可丢失会话亲和存储恢复的 native previous-response pin。 /// -/// 该值不代表数据库 transcript；它只阻止 native handle 被发送到错误的 -/// Provider 或账号。 -#[derive(Clone, PartialEq, Eq)] +/// 该值不依赖调用方 API Key；它只阻止 native handle 与不透明 Provider state +/// 被发送到错误的 Provider 或账号。 +#[derive(Clone, PartialEq)] pub struct NativeContinuationPin { - /// 客户端提交、仅供 Store 在调用方隔离下查找的网关 response ID。 + /// 客户端提交、用于查找可丢失会话亲和记录的 response ID。 previous_response_id: PreviousResponseId, - /// Store 从已提交成功请求解析出的 Provider 原生 response handle。 + /// Provider 原生 response handle。 upstream_response_id: SafeUpstreamValue, provider: ProviderKind, account: ProviderAccountId, scope: NativeContinuationScope, + session_state: Option, } impl NativeContinuationPin { @@ -72,6 +73,7 @@ impl NativeContinuationPin { provider, account, scope: NativeContinuationScope::ConnectionLocal, + session_state: None, } } @@ -82,6 +84,13 @@ impl NativeContinuationPin { self } + /// 附着仅由对应 Provider 解释的不透明会话状态。 + #[must_use] + pub fn with_session_state(mut self, state: ProviderSessionState) -> Self { + self.session_state = Some(state); + self + } + #[must_use] pub const fn previous_response_id(&self) -> &PreviousResponseId { &self.previous_response_id @@ -108,6 +117,12 @@ impl NativeContinuationPin { self.scope } + /// 返回与本 pin 同账号绑定的 Provider 私有会话状态。 + #[must_use] + pub const fn session_state(&self) -> Option<&ProviderSessionState> { + self.session_state.as_ref() + } + /// 校验本次 route/account 选择没有破坏 native pin。 #[must_use] pub fn matches(&self, provider: &ProviderKind, account: &ProviderAccountId) -> bool { @@ -124,6 +139,10 @@ impl fmt::Debug for NativeContinuationPin { .field("provider", &self.provider) .field("account", &self.account) .field("scope", &self.scope) + .field( + "session_state", + &self.session_state.as_ref().map(|_| ""), + ) .finish() } } @@ -132,7 +151,7 @@ impl fmt::Debug for NativeContinuationPin { /// /// 已命中网关历史的 handle 携带完整账号 pin；未命中历史的外部 handle 只保留 /// 客户端提交的 opaque ID，由目标 Provider 在首次且唯一一次 attempt 中解释。 -#[derive(Clone, PartialEq, Eq)] +#[derive(Clone, PartialEq)] pub enum ContinuationBinding { Pinned(NativeContinuationPin), External(PreviousResponseId), @@ -172,11 +191,18 @@ impl fmt::Debug for ContinuationBinding { #[error("native continuation store is unavailable")] pub struct NativeContinuationStoreError; -/// 调用方隔离后的 previous-response 解析端口。 +/// 可丢失的 previous-response 亲和存储端口。 +/// +/// 亲和记录不按调用方 API Key 隔离：与旧版一样，response ID 是会话级 opaque +/// handle；Redis 不可用或超时必须由调用方 fail-open。 pub trait NativeContinuationPort: Send + Sync { fn resolve<'a>( &'a self, - client_api_key_id: &'a ClientApiKeyId, previous_response_id: &'a PreviousResponseId, ) -> BoxFuture<'a, Result, NativeContinuationStoreError>>; + + fn record<'a>( + &'a self, + pin: NativeContinuationPin, + ) -> BoxFuture<'a, Result<(), NativeContinuationStoreError>>; } diff --git a/backend/crates/gateway-core/src/engine/coordinator.rs b/backend/crates/gateway-core/src/engine/coordinator.rs index 6ca107a35..07d0f69d2 100644 --- a/backend/crates/gateway-core/src/engine/coordinator.rs +++ b/backend/crates/gateway-core/src/engine/coordinator.rs @@ -6,7 +6,9 @@ use std::sync::Arc; use std::time::{Duration, SystemTime}; use crate::accounting::{CostEstimate, CostSource, Usage}; -use crate::engine::continuation::ContinuationBinding; +use crate::engine::continuation::{ + ContinuationBinding, NativeContinuationPin, NativeContinuationScope, PreviousResponseId, +}; use crate::engine::provider::{Provider, ProviderCallMetadata, ProviderRequest, ProviderStream}; use crate::engine::{ AccountAttemptContext, AttemptContext, AttemptRecord, AttemptTrigger, CancellationToken, @@ -19,7 +21,7 @@ use crate::error::{GatewayError, GatewayErrorKind, ProviderError, ProviderErrorK use crate::event::{ GatewayEvent, ProviderEvent, ProviderResponseHeader, ProviderResponseObservation, }; -use crate::operation::{Operation, RetrySafety}; +use crate::operation::{Operation, ProviderSessionState, RetrySafety}; use crate::routing::RoutingPlan; use futures::{FutureExt, StreamExt, pin_mut, select_biased}; use futures_timer::Delay; @@ -348,6 +350,36 @@ where .unwrap_or_default() } + /// 将已完成响应的账号事实与 Provider 私有状态封装为可丢失的亲和记录。 + /// + /// Core 不读取 `state` 内容；Provider 将在后续同账号 continuation 时自行解释。 + #[must_use] + pub fn native_continuation_pin( + &self, + state: &ProviderSessionState, + ) -> Option { + let current = self.current.as_ref()?; + let provider = current.metadata.provider().clone(); + if state.provider() != provider.as_str() { + return None; + } + let account = current.metadata.provider_account_id()?.clone(); + let response_id = self.upstream_response_id.as_deref()?; + let previous_response_id = PreviousResponseId::new(response_id.to_owned()).ok()?; + let upstream_response_id = + crate::error::SafeUpstreamValue::new(response_id.to_owned()).ok()?; + Some( + NativeContinuationPin::new( + previous_response_id, + upstream_response_id, + provider, + account, + ) + .with_scope(NativeContinuationScope::Persisted) + .with_session_state(state.clone()), + ) + } + /// 请求取消；实际终态在下一次会话 poll 时由 Core 持久化。 pub fn cancel(&self) { self.cancellation.cancel(); diff --git a/backend/crates/gateway-core/src/engine/execution.rs b/backend/crates/gateway-core/src/engine/execution.rs index bf6150f2b..6ba6b88f6 100644 --- a/backend/crates/gateway-core/src/engine/execution.rs +++ b/backend/crates/gateway-core/src/engine/execution.rs @@ -6,14 +6,15 @@ use std::net::IpAddr; use std::sync::Arc; use std::time::{Duration, SystemTime}; -use futures::future::BoxFuture; +use futures::{FutureExt, future::BoxFuture, pin_mut, select_biased}; +use futures_timer::Delay; use uuid::Uuid; use crate::engine::admission::{ ClientAdmissionDecision, ClientAdmissionPort, ClientAdmissionRejection, ClientAdmissionRequest, }; use crate::engine::continuation::{ - ContinuationBinding, NativeContinuationPort, PreviousResponseId, + ContinuationBinding, NativeContinuationPin, NativeContinuationPort, PreviousResponseId, }; use crate::engine::coordinator::ResponseExecutionSession; use crate::engine::probe::{AccountProbe, AccountProbeRequest, AccountProbeResult}; @@ -25,12 +26,13 @@ use crate::engine::{ }; use crate::error::{GatewayError, GatewayErrorKind, ProviderErrorKind, StoreError}; use crate::event::{GatewayEvent, ProviderEvent, ProviderResponseHeader}; -use crate::operation::Operation; +use crate::operation::{Operation, ProviderSessionState}; use crate::policy::{ClientApiKeyId, ClientPolicy}; use crate::routing::snapshot::RuntimeSnapshotHandle; use crate::routing::{ProviderKind, PublicModelId, RoutingContext, RuntimeSnapshot}; const MODEL_REQUEST_DEADLINE: Duration = Duration::from_secs(10 * 60); +const CONTINUATION_AFFINITY_TIMEOUT: Duration = Duration::from_millis(100); #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub enum ClientTransport { @@ -213,10 +215,22 @@ impl DefaultExecutionService { } } - async fn start_inner(&self, request: StartExecution) -> Result { + async fn start_inner( + &self, + mut request: StartExecution, + ) -> Result { request.client.policy.authorize().map_err(|_| { GatewayError::new(GatewayErrorKind::PolicyDenied, "client API key is disabled") })?; + if !request.client.snapshot.contains_public_model_for_provider( + &request.public_model, + request.client.policy.provider_kind(), + ) { + return Err(GatewayError::new( + GatewayErrorKind::ModelNotFound, + "requested model was not found", + )); + } let started_at = SystemTime::now(); let deadline_at = started_at .checked_add(MODEL_REQUEST_DEADLINE) @@ -233,19 +247,22 @@ impl DefaultExecutionService { .plan(&request.public_model, &request.operation, &routing_context) .map_err(map_routing_error)?; let continuation = match request.metadata.previous_response_id.as_ref() { - Some(previous) => match self - .continuation - .resolve(request.client.policy.key_id(), previous) - .await - .map_err(|_| { - GatewayError::new( - GatewayErrorKind::Internal, - "previous response history is temporarily unavailable", - ) - })? { - Some(pin) => ContinuationBinding::Pinned(pin), - None => ContinuationBinding::External(previous.clone()), - }, + Some(previous) => { + let resolve = self.continuation.resolve(previous).fuse(); + let timeout = Delay::new(CONTINUATION_AFFINITY_TIMEOUT).fuse(); + pin_mut!(resolve, timeout); + let pin = select_biased! { + result = resolve => result.ok().flatten(), + _ = timeout => None, + }; + match pin { + Some(pin) => { + attach_continuation_session_state(&mut request.operation, &pin); + ContinuationBinding::Pinned(pin) + } + None => ContinuationBinding::External(previous.clone()), + } + } None => { return self .start_without_continuation( @@ -361,6 +378,7 @@ impl DefaultExecutionService { core, admission, Arc::clone(&self.circuits), + Arc::clone(&self.continuation), )), }) } @@ -592,7 +610,9 @@ struct DefaultExecutionSession { core: Option>, admission: Option, circuits: Arc, + continuation: Arc, observed_provider_outcomes: usize, + continuation_recorded: bool, } impl DefaultExecutionSession { @@ -600,12 +620,15 @@ impl DefaultExecutionSession { core: ResponseExecutionSession, admission: AdmissionLease, circuits: Arc, + continuation: Arc, ) -> Self { Self { core: Some(core), admission: Some(admission), circuits, + continuation, observed_provider_outcomes: 0, + continuation_recorded: false, } } @@ -639,6 +662,30 @@ impl DefaultExecutionSession { publish_provider_attempt_outcomes(self.circuits.as_ref(), &new_outcomes).await; } + async fn record_continuation(&mut self, state: Option<&ProviderSessionState>) { + if self.continuation_recorded { + return; + } + let Some(state) = state else { + return; + }; + let Some(pin) = self + .core + .as_ref() + .and_then(|core| core.native_continuation_pin(state)) + else { + return; + }; + self.continuation_recorded = true; + let record = self.continuation.record(pin).fuse(); + let timeout = Delay::new(CONTINUATION_AFFINITY_TIMEOUT).fuse(); + pin_mut!(record, timeout); + select_biased! { + _ = record => {}, + _ = timeout => {}, + } + } + async fn finalize_detached(&mut self) { let Some(mut core) = self.core.take() else { return; @@ -670,6 +717,9 @@ impl ExecutionSession for DefaultExecutionSession { fn next_event(&mut self) -> BoxFuture<'_, Result, EngineError>> { Box::pin(async move { let result = self.core_mut()?.next_event().await; + if let Ok(Some(event)) = result.as_ref() { + self.record_continuation(event.session_update()).await; + } self.observe_provider_outcomes().await; self.settle_if_finalized().await; result @@ -679,6 +729,10 @@ impl ExecutionSession for DefaultExecutionSession { fn collect_uncommitted(&mut self) -> BoxFuture<'_, Result, EngineError>> { Box::pin(async move { let result = self.core_mut()?.collect_uncommitted().await; + if let Ok(events) = result.as_ref() { + let state = events.iter().find_map(ProviderEvent::session_update); + self.record_continuation(state).await; + } self.observe_provider_outcomes().await; self.settle_if_finalized().await; result @@ -794,6 +848,18 @@ fn map_routing_error(error: crate::error::RoutingError) -> GatewayError { } } +fn attach_continuation_session_state(operation: &mut Operation, pin: &NativeContinuationPin) { + let Some(state) = pin.session_state() else { + return; + }; + if state.provider() != pin.provider().as_str() + || operation.provider_session_state(state.provider()).is_some() + { + return; + } + *operation = operation.clone().with_provider_session_state(state.clone()); +} + pub fn gateway_error_from_engine(error: &EngineError) -> GatewayError { match error { EngineError::Cancelled => { diff --git a/backend/crates/gateway-core/src/engine/mod.rs b/backend/crates/gateway-core/src/engine/mod.rs index 1aef698e6..1d0553729 100644 --- a/backend/crates/gateway-core/src/engine/mod.rs +++ b/backend/crates/gateway-core/src/engine/mod.rs @@ -30,6 +30,7 @@ use crate::error::{ }; use crate::event::ProviderEvent; use crate::operation::OperationKind; +use crate::operation::ProviderSessionState; use crate::policy::ClientApiKeyId; use crate::routing::{ConfigRevision, ProviderKind, PublicModelId, UpstreamModelId}; @@ -649,6 +650,12 @@ impl CoordinatedEvent { pub fn into_provider_events(self) -> Vec { self.events } + + /// 返回本批次中 Provider 交付的私有会话更新。 + #[must_use] + pub fn session_update(&self) -> Option<&ProviderSessionState> { + self.events.iter().find_map(ProviderEvent::session_update) + } } #[derive(Debug, Error)] diff --git a/backend/crates/gateway-core/src/operation/mod.rs b/backend/crates/gateway-core/src/operation/mod.rs index 321390e8b..dc839db5a 100644 --- a/backend/crates/gateway-core/src/operation/mod.rs +++ b/backend/crates/gateway-core/src/operation/mod.rs @@ -8,6 +8,7 @@ use std::collections::BTreeSet; use std::fmt; use std::sync::Arc; +use serde::{Deserialize, Serialize}; use serde_json::{Map, Value}; use crate::error::{OperationError, validate_text}; @@ -135,7 +136,7 @@ impl CapabilityRequirements { /// /// Core 只在重试和路由过程中保持该值；协议层只能把 Provider 返回的状态原样带入 /// 下一轮，不能读取或改写正文。 -#[derive(Clone, PartialEq)] +#[derive(Clone, PartialEq, Serialize, Deserialize)] pub struct ProviderSessionState { provider: String, payload: Map, diff --git a/backend/crates/gateway-core/src/provider_ports/mod.rs b/backend/crates/gateway-core/src/provider_ports/mod.rs index e29f4c063..1ec8f2e85 100644 --- a/backend/crates/gateway-core/src/provider_ports/mod.rs +++ b/backend/crates/gateway-core/src/provider_ports/mod.rs @@ -1,6 +1,6 @@ //! Provider 运行时所需的中立存储能力。 -use std::collections::BTreeMap; +use std::collections::{BTreeMap, BTreeSet}; use std::fmt; use std::num::NonZeroU32; use std::sync::Arc; @@ -235,6 +235,62 @@ pub trait ProviderSessionAffinityPort: Send + Sync { ) -> BoxFuture<'a, Result>; } +/// Provider 会话内已失败账号的可丢失排除集。 +/// +/// Provider 自行派生会话键并决定何时写入或清理；Core 只承载调度所需的账号 ID +/// 与 compare-and-swap revision，不解释任一 Provider 协议字段。 +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct ProviderSessionExclusions { + excluded_accounts: BTreeSet, + revision: String, +} + +impl ProviderSessionExclusions { + #[must_use] + pub const fn new(excluded_accounts: BTreeSet, revision: String) -> Self { + Self { + excluded_accounts, + revision, + } + } + + #[must_use] + pub const fn excluded_accounts(&self) -> &BTreeSet { + &self.excluded_accounts + } + + #[must_use] + pub fn revision(&self) -> &str { + &self.revision + } +} + +/// 可丢失的 Provider 会话级账号排除状态。 +/// +/// 该端口不接收协议正文；Provider 只能以不可逆会话键、账号 ID 和固定 TTL 操作。 +pub trait ProviderSessionExclusionPort: Send + Sync { + fn load<'a>( + &'a self, + provider_kind: &'a ProviderKind, + key: &'a ProviderSessionAffinityKey, + ) -> BoxFuture<'a, Result, ProviderStoreError>>; + + fn record_failure<'a>( + &'a self, + provider_kind: &'a ProviderKind, + key: &'a ProviderSessionAffinityKey, + account_id: &'a ProviderAccountId, + ttl: Duration, + ) -> BoxFuture<'a, Result>; + + fn clear<'a>( + &'a self, + provider_kind: &'a ProviderKind, + key: &'a ProviderSessionAffinityKey, + expected_revision: &'a str, + ) -> BoxFuture<'a, Result>; +} + /// 所有 Provider 共享的 OAuth refresh 并发容量。 #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub struct ProviderRefreshCapacityRequest { @@ -704,6 +760,7 @@ pub struct ProviderStorePorts { accounts: Arc, leases: Arc, session_affinity: Arc, + session_exclusions: Arc, account_feedback: Arc, catalog_cache: Arc, credential_state: Arc, @@ -720,6 +777,7 @@ impl ProviderStorePorts { accounts: Arc, leases: Arc, session_affinity: Arc, + session_exclusions: Arc, catalog_cache: Arc, credential_state: Arc, cooldowns: Arc, @@ -730,6 +788,7 @@ impl ProviderStorePorts { accounts, leases, session_affinity, + session_exclusions, account_feedback: Arc::new(AccountFeedbackStats::default()), catalog_cache, credential_state, @@ -754,6 +813,11 @@ impl ProviderStorePorts { Arc::clone(&self.session_affinity) } + #[must_use] + pub fn session_exclusions(&self) -> Arc { + Arc::clone(&self.session_exclusions) + } + #[must_use] pub fn account_feedback(&self) -> Arc { Arc::clone(&self.account_feedback) diff --git a/backend/crates/gateway-core/src/routing/snapshot.rs b/backend/crates/gateway-core/src/routing/snapshot.rs index ca6024380..a6f249ba7 100644 --- a/backend/crates/gateway-core/src/routing/snapshot.rs +++ b/backend/crates/gateway-core/src/routing/snapshot.rs @@ -412,22 +412,47 @@ impl RuntimeSnapshot { .collect() } - /// 透明代理不以目录白名单拒绝模型；只要求该 Provider 已注册。 + /// 已取得目录时按目录与全局别名判断模型可用性；目录暂不可用时沿用旧版 + /// fail-open 语义，把具体 wire 可接受性留给上游。 #[must_use] pub fn contains_public_model_for_provider( &self, - _public_model: &PublicModelId, + public_model: &PublicModelId, provider: &ProviderKind, ) -> bool { - self.providers.contains(provider) + if !self.providers.contains(provider) { + return false; + } + if self.model_mappings.contains_key(public_model.as_str()) { + return true; + } + match self.provider_models.get(provider) { + Some(models) if !models.is_empty() => models + .keys() + .any(|model| model.as_str() == public_model.as_str()), + Some(_) | None => true, + } } #[must_use] pub fn mapped_model(&self, requested: &str) -> String { - self.model_mappings - .get(requested) - .cloned() - .unwrap_or_else(|| requested.to_owned()) + let original = requested.trim(); + let mut current = original.to_owned(); + let mut seen = BTreeSet::new(); + for _ in 0..20 { + let Some(target) = self + .model_mappings + .get(¤t) + .map(|target| target.trim()) + else { + return current; + }; + if !seen.insert(current.clone()) || seen.contains(target) { + return original.to_owned(); + } + current = target.to_owned(); + } + original.to_owned() } pub fn client_policies(&self) -> impl Iterator { diff --git a/backend/crates/gateway-core/tests/engine/execution.rs b/backend/crates/gateway-core/tests/engine/execution.rs index 11a8ea8ee..050d713eb 100644 --- a/backend/crates/gateway-core/tests/engine/execution.rs +++ b/backend/crates/gateway-core/tests/engine/execution.rs @@ -238,11 +238,17 @@ struct UnusedContinuation; impl NativeContinuationPort for UnusedContinuation { fn resolve<'a>( &'a self, - _: &'a ClientApiKeyId, _: &'a PreviousResponseId, ) -> BoxFuture<'a, Result, NativeContinuationStoreError>> { Box::pin(async { Ok(None) }) } + + fn record<'a>( + &'a self, + _: NativeContinuationPin, + ) -> BoxFuture<'a, Result<(), NativeContinuationStoreError>> { + Box::pin(async { Ok(()) }) + } } fn probe_snapshot() -> RuntimeSnapshot { diff --git a/backend/crates/gateway-core/tests/routing/mod.rs b/backend/crates/gateway-core/tests/routing/mod.rs index aaac74138..9c41c6e3f 100644 --- a/backend/crates/gateway-core/tests/routing/mod.rs +++ b/backend/crates/gateway-core/tests/routing/mod.rs @@ -148,6 +148,26 @@ fn selected_provider_should_use_global_model_mapping() { assert_eq!(plan.candidates()[0].upstream_model().as_str(), "gpt-5.5"); } +#[test] +fn model_mapping_should_follow_a_bounded_alias_chain() { + let snapshot = snapshot().with_model_mappings(BTreeMap::from([ + ("public-model".to_owned(), "compat-model".to_owned()), + ("compat-model".to_owned(), "gpt-5.5".to_owned()), + ])); + + assert_eq!(snapshot.mapped_model("public-model"), "gpt-5.5"); +} + +#[test] +fn cyclic_model_mapping_should_fall_back_to_the_original_name() { + let snapshot = snapshot().with_model_mappings(BTreeMap::from([ + ("first".to_owned(), "second".to_owned()), + ("second".to_owned(), "first".to_owned()), + ])); + + assert_eq!(snapshot.mapped_model("first"), "first"); +} + #[test] fn global_mapping_should_apply_to_every_provider() { let snapshot = snapshot(); @@ -274,8 +294,8 @@ fn public_catalog_should_include_discovered_models_and_aliases() { } #[test] -fn unknown_model_should_be_accepted_when_provider_is_registered() { - assert!(snapshot().contains_public_model_for_provider( +fn unknown_model_should_be_hidden_when_provider_catalog_is_available() { + assert!(!snapshot().contains_public_model_for_provider( &PublicModelId::new("future-model").expect("model"), &ProviderKind::new("openai").expect("provider"), )); diff --git a/backend/crates/gateway-store/src/lib.rs b/backend/crates/gateway-store/src/lib.rs index 59b1bdbaf..06bd5e2ea 100644 --- a/backend/crates/gateway-store/src/lib.rs +++ b/backend/crates/gateway-store/src/lib.rs @@ -304,7 +304,10 @@ pub async fn initialize(mut config: StoreConfig) -> StoreResult { REDIS_NAMESPACE, redis::ProviderCircuitPolicy::default(), )?), - Arc::new(postgres::PgHistoryRepository::new(pool.clone())), + Arc::new(redis::RedisNativeContinuationRepository::new( + redis_connection.clone(), + REDIS_NAMESPACE, + )?), ( Arc::new(postgres::PgRuntimeSnapshotRepository::new(pool.clone())), Arc::new(redis::RedisRuntimeChangeRepository::new( @@ -319,6 +322,10 @@ pub async fn initialize(mut config: StoreConfig) -> StoreResult { account_store, provider_leases, provider_session_affinity, + Arc::new(redis::RedisProviderSessionExclusionRepository::new( + redis_connection.clone(), + REDIS_NAMESPACE, + )?), credential_state.clone(), credential_state, cooldowns, diff --git a/backend/crates/gateway-store/src/postgres/history.rs b/backend/crates/gateway-store/src/postgres/history.rs index 441dca533..bb4e3ac62 100644 --- a/backend/crates/gateway-store/src/postgres/history.rs +++ b/backend/crates/gateway-store/src/postgres/history.rs @@ -2,19 +2,9 @@ use async_trait::async_trait; use chrono::{DateTime, Utc}; -use gateway_core::engine::continuation::{NativeContinuationPort, NativeContinuationStoreError}; -use gateway_core::policy::ClientApiKeyId; -use gateway_core::{ - engine::{ - continuation::{NativeContinuationPin, NativeContinuationScope, PreviousResponseId}, - credential::ProviderAccountId, - }, - error::SafeUpstreamValue, - routing::ProviderKind, -}; use sqlx::PgPool; -use crate::{StoreError, StoreResult, postgres_unavailable, require_nonempty}; +use crate::{StoreResult, postgres_unavailable, require_nonempty}; #[derive(Debug, Clone, PartialEq, Eq)] pub struct ModelRequestHistoryRecord { @@ -37,18 +27,6 @@ pub trait ModelRequestHistoryRepository: Send + Sync { &self, model_request_id: &str, ) -> StoreResult>; - - async fn find_model_request_by_client_response_id( - &self, - client_response_id: &str, - caller_client_api_key_ref: &str, - ) -> StoreResult>; - - async fn resolve_native_continuation_pin( - &self, - client_response_id: &str, - caller_client_api_key_ref: &str, - ) -> StoreResult>; } #[derive(Clone)] @@ -76,81 +54,6 @@ impl ModelRequestHistoryRepository for PgHistoryRepository { )?; fetch_history(&self.pool, "id", model_request_id).await } - - async fn find_model_request_by_client_response_id( - &self, - client_response_id: &str, - caller_client_api_key_ref: &str, - ) -> StoreResult> { - require_nonempty( - "model request history", - "client_response_id", - client_response_id, - )?; - require_nonempty( - "model request history", - "caller_client_api_key_ref", - caller_client_api_key_ref, - )?; - fetch_history_by_client_response(&self.pool, client_response_id, caller_client_api_key_ref) - .await - } - - async fn resolve_native_continuation_pin( - &self, - client_response_id: &str, - caller_client_api_key_ref: &str, - ) -> StoreResult> { - require_nonempty( - "native continuation", - "client_response_id", - client_response_id, - )?; - require_nonempty( - "native continuation", - "caller_client_api_key_ref", - caller_client_api_key_ref, - )?; - let row = sqlx::query_as::<_, (String, String, String)>( - "select mr.provider_kind, mr.provider_account_id, mr.upstream_response_id - from model_requests mr - join provider_accounts account - on account.id = mr.provider_account_id - and account.provider_kind = mr.provider_kind - where mr.client_response_id = $1 - and mr.client_api_key_ref = $2 - and mr.outcome = 'succeeded' - and mr.downstream_committed_at is not null - and mr.upstream_response_id is not null", - ) - .bind(client_response_id) - .bind(caller_client_api_key_ref) - .fetch_optional(&self.pool) - .await - .map_err(|_| postgres_unavailable("resolve native continuation pin"))?; - row.map(|row| native_pin_from_row(client_response_id, row)) - .transpose() - } -} - -impl NativeContinuationPort for PgHistoryRepository { - fn resolve<'a>( - &'a self, - client_api_key_id: &'a ClientApiKeyId, - previous_response_id: &'a PreviousResponseId, - ) -> futures::future::BoxFuture< - 'a, - Result, NativeContinuationStoreError>, - > { - Box::pin(async move { - self.resolve_native_continuation_pin( - previous_response_id.as_str(), - client_api_key_id.as_str(), - ) - .await - .map_err(|_| NativeContinuationStoreError) - }) - } } async fn fetch_history( @@ -177,27 +80,6 @@ async fn fetch_history( Ok(row.map(history_from_row)) } -async fn fetch_history_by_client_response( - pool: &PgPool, - client_response_id: &str, - caller_client_api_key_ref: &str, -) -> StoreResult> { - let row = sqlx::query_as::<_, HistoryRow>( - "select id, client_api_key_ref, requested_model_id, provider_kind, - provider_account_ref, upstream_model_id, outcome, - client_response_id, upstream_response_id, - started_at, completed_at - from model_requests - where client_response_id = $1 and client_api_key_ref = $2", - ) - .bind(client_response_id) - .bind(caller_client_api_key_ref) - .fetch_optional(pool) - .await - .map_err(|_| postgres_unavailable("read model request history by client response"))?; - Ok(row.map(history_from_row)) -} - type HistoryRow = ( String, String, @@ -227,31 +109,3 @@ fn history_from_row(row: HistoryRow) -> ModelRequestHistoryRecord { completed_at: row.10, } } - -fn native_pin_from_row( - client_response_id: &str, - row: (String, String, String), -) -> StoreResult { - let previous_response_id = PreviousResponseId::new(client_response_id.to_owned()) - .map_err(|_| invalid_native_pin("invalid client response ID"))?; - let provider = - ProviderKind::new(row.0).map_err(|_| invalid_native_pin("invalid provider kind"))?; - let upstream_response_id = SafeUpstreamValue::new(row.2) - .map_err(|_| invalid_native_pin("invalid upstream response ID"))?; - let account = ProviderAccountId::new(row.1) - .map_err(|_| invalid_native_pin("invalid provider account ID"))?; - Ok(NativeContinuationPin::new( - previous_response_id, - upstream_response_id, - provider, - account, - ) - .with_scope(NativeContinuationScope::Persisted)) -} - -fn invalid_native_pin(message: &str) -> StoreError { - StoreError::InvalidData { - entity: "native continuation", - message: message.to_owned(), - } -} diff --git a/backend/crates/gateway-store/src/redis/mod.rs b/backend/crates/gateway-store/src/redis/mod.rs index e0e6bb9de..c7798af4f 100644 --- a/backend/crates/gateway-store/src/redis/mod.rs +++ b/backend/crates/gateway-store/src/redis/mod.rs @@ -11,9 +11,11 @@ mod client_admission; mod credential_cooldown; mod credential_leases; mod credential_state; +mod native_continuation; mod oauth_pending; mod provider_circuit; mod provider_session_affinity; +mod provider_session_exclusion; mod runtime_change; pub(crate) mod worker_lease; @@ -21,9 +23,11 @@ pub use client_admission::*; pub use credential_cooldown::*; pub use credential_leases::*; pub use credential_state::*; +pub use native_continuation::*; pub use oauth_pending::*; pub use provider_circuit::*; pub use provider_session_affinity::*; +pub use provider_session_exclusion::*; pub use runtime_change::*; use crate::{StoreError, StoreResult, redis_unavailable, require_nonempty}; diff --git a/backend/crates/gateway-store/src/redis/native_continuation.rs b/backend/crates/gateway-store/src/redis/native_continuation.rs new file mode 100644 index 000000000..08e702e4f --- /dev/null +++ b/backend/crates/gateway-store/src/redis/native_continuation.rs @@ -0,0 +1,221 @@ +//! 旧版 Responses continuation 的 Redis 会话亲和记录。 +//! +//! 该存储只保存可丢失的 response → Provider/account pin；PostgreSQL 审计记录 +//! 不是 continuation 的可用性前提。键只含 response ID 的不可逆指纹，Provider +//! 私有状态作为不透明 JSON 保存，Store 不解释其内容。 + +use std::time::{Duration, SystemTime, UNIX_EPOCH}; + +use futures::future::BoxFuture; +use gateway_core::engine::continuation::{ + NativeContinuationPin, NativeContinuationPort, NativeContinuationScope, + NativeContinuationStoreError, PreviousResponseId, +}; +use gateway_core::engine::credential::ProviderAccountId; +use gateway_core::error::SafeUpstreamValue; +use gateway_core::operation::ProviderSessionState; +use gateway_core::routing::ProviderKind; +use redis::aio::ConnectionManager; +use serde::{Deserialize, Serialize}; + +use crate::StoreResult; + +use super::{namespace, resource_fingerprint}; + +const CONTINUATION_TTL: Duration = Duration::from_secs(4 * 60 * 60); +const MAX_CONTINUATIONS: usize = 65_536; +const MAX_RECORD_BYTES: usize = 64 * 1024; + +const RECORD_SCRIPT: &str = r#" +local entry_key = KEYS[1] +local index_key = KEYS[2] +local entry_id = ARGV[1] +local payload = ARGV[2] +local score = tonumber(ARGV[3]) +local ttl_seconds = tonumber(ARGV[4]) +local cutoff = tonumber(ARGV[5]) +local max_entries = tonumber(ARGV[6]) +local entry_prefix = ARGV[7] + +redis.call('SET', entry_key, payload, 'EX', ttl_seconds) +redis.call('ZADD', index_key, score, entry_id) +redis.call('EXPIRE', index_key, ttl_seconds) + +local stale = redis.call('ZRANGEBYSCORE', index_key, '-inf', cutoff) +for _, stale_id in ipairs(stale) do + redis.call('UNLINK', entry_prefix .. '{' .. stale_id .. '}') + redis.call('ZREM', index_key, stale_id) +end + +local count = redis.call('ZCARD', index_key) +local excess = count - max_entries +if excess > 0 then + local oldest = redis.call('ZRANGE', index_key, 0, excess - 1) + for _, old_id in ipairs(oldest) do + redis.call('UNLINK', entry_prefix .. '{' .. old_id .. '}') + redis.call('ZREM', index_key, old_id) + end +end + +return 1 +"#; + +/// Redis-backed, best-effort native response continuation affinity. +#[derive(Clone)] +pub struct RedisNativeContinuationRepository { + connection: ConnectionManager, + namespace: String, +} + +#[derive(Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +struct ContinuationWire { + upstream_response_id: String, + provider: String, + account: String, + scope: String, + #[serde(default)] + session_state: Option, +} + +impl RedisNativeContinuationRepository { + pub fn new(connection: ConnectionManager, key_namespace: &str) -> StoreResult { + Ok(Self { + connection, + namespace: format!("{}:response-affinity:v1", namespace(key_namespace)?), + }) + } + + fn entry_fingerprint(&self, response_id: &str) -> Result { + resource_fingerprint("native continuation", response_id) + .map_err(|_| NativeContinuationStoreError) + } + + fn entry_key(&self, fingerprint: &str) -> String { + format!("{}{{{fingerprint}}}", self.entry_prefix()) + } + + fn entry_prefix(&self) -> String { + format!("{}:entry:", self.namespace) + } + + fn index_key(&self) -> String { + format!("{}:global:index", self.namespace) + } +} + +impl NativeContinuationPort for RedisNativeContinuationRepository { + fn resolve<'a>( + &'a self, + previous_response_id: &'a PreviousResponseId, + ) -> BoxFuture<'a, Result, NativeContinuationStoreError>> { + Box::pin(async move { + let fingerprint = self.entry_fingerprint(previous_response_id.as_str())?; + let mut connection = self.connection.clone(); + let payload = redis::cmd("GET") + .arg(self.entry_key(&fingerprint)) + .query_async::>(&mut connection) + .await + .map_err(|_| NativeContinuationStoreError)?; + payload + .map(|payload| decode_pin(previous_response_id, &payload)) + .transpose() + }) + } + + fn record<'a>( + &'a self, + pin: NativeContinuationPin, + ) -> BoxFuture<'a, Result<(), NativeContinuationStoreError>> { + Box::pin(async move { + let fingerprint = self.entry_fingerprint(pin.previous_response_id().as_str())?; + let payload = encode_pin(&pin)?; + if payload.len() > MAX_RECORD_BYTES { + return Err(NativeContinuationStoreError); + } + let now_millis = SystemTime::now() + .duration_since(UNIX_EPOCH) + .map_err(|_| NativeContinuationStoreError)? + .as_millis(); + let now_millis = u64::try_from(now_millis).map_err(|_| NativeContinuationStoreError)?; + let ttl_millis = u64::try_from(CONTINUATION_TTL.as_millis()) + .map_err(|_| NativeContinuationStoreError)?; + let cutoff = now_millis.saturating_sub(ttl_millis); + let ttl_seconds = CONTINUATION_TTL.as_secs(); + let mut connection = self.connection.clone(); + redis::Script::new(RECORD_SCRIPT) + .key(self.entry_key(&fingerprint)) + .key(self.index_key()) + .arg(fingerprint) + .arg(payload) + .arg(now_millis) + .arg(ttl_seconds) + .arg(cutoff) + .arg(MAX_CONTINUATIONS) + .arg(self.entry_prefix()) + .invoke_async::(&mut connection) + .await + .map_err(|_| NativeContinuationStoreError)?; + Ok(()) + }) + } +} + +fn encode_pin(pin: &NativeContinuationPin) -> Result { + if pin + .session_state() + .is_some_and(|state| state.provider() != pin.provider().as_str()) + { + return Err(NativeContinuationStoreError); + } + serde_json::to_string(&ContinuationWire { + upstream_response_id: pin.upstream_response_id().as_str().to_owned(), + provider: pin.provider().as_str().to_owned(), + account: pin.account().as_str().to_owned(), + scope: scope_name(pin.scope()).to_owned(), + session_state: pin.session_state().cloned(), + }) + .map_err(|_| NativeContinuationStoreError) +} + +fn decode_pin( + previous_response_id: &PreviousResponseId, + payload: &str, +) -> Result { + let wire: ContinuationWire = + serde_json::from_str(payload).map_err(|_| NativeContinuationStoreError)?; + let provider = ProviderKind::new(wire.provider).map_err(|_| NativeContinuationStoreError)?; + let account = ProviderAccountId::new(wire.account).map_err(|_| NativeContinuationStoreError)?; + let upstream_response_id = SafeUpstreamValue::new(wire.upstream_response_id) + .map_err(|_| NativeContinuationStoreError)?; + let scope = parse_scope(&wire.scope)?; + let mut pin = NativeContinuationPin::new( + previous_response_id.clone(), + upstream_response_id, + provider.clone(), + account, + ) + .with_scope(scope); + if let Some(state) = wire.session_state { + if state.provider() != provider.as_str() { + return Err(NativeContinuationStoreError); + } + pin = pin.with_session_state(state); + } + Ok(pin) +} + +const fn scope_name(scope: NativeContinuationScope) -> &'static str { + match scope { + NativeContinuationScope::Persisted => "persisted", + NativeContinuationScope::ConnectionLocal => "connection_local", + } +} + +fn parse_scope(scope: &str) -> Result { + match scope { + "persisted" => Ok(NativeContinuationScope::Persisted), + "connection_local" => Ok(NativeContinuationScope::ConnectionLocal), + _ => Err(NativeContinuationStoreError), + } +} diff --git a/backend/crates/gateway-store/src/redis/provider_session_exclusion.rs b/backend/crates/gateway-store/src/redis/provider_session_exclusion.rs new file mode 100644 index 000000000..ba1dddac4 --- /dev/null +++ b/backend/crates/gateway-store/src/redis/provider_session_exclusion.rs @@ -0,0 +1,194 @@ +//! Provider 会话级账号排除的可丢失 Redis 状态。 +//! +//! Redis 只保存 Provider 派生的不可逆会话键对应的账号 ID 集合及 revision； +//! 不读取或保存任何 Provider 协议正文。 + +use std::collections::BTreeSet; +use std::time::Duration; + +use futures::future::BoxFuture; +use gateway_core::engine::credential::ProviderAccountId; +use gateway_core::provider_ports::{ + ProviderSessionAffinityKey, ProviderSessionExclusionPort, ProviderSessionExclusions, + ProviderStoreError, ProviderStoreErrorKind, +}; +use gateway_core::routing::ProviderKind; +use redis::aio::ConnectionManager; +use serde::{Deserialize, Serialize}; +use uuid::Uuid; + +use crate::StoreResult; + +use super::{namespace, resource_fingerprint}; + +const MAX_SESSION_EXCLUSION_TTL: Duration = Duration::from_secs(24 * 60 * 60); + +const RECORD_FAILURE_SCRIPT: &str = r#" +local raw = redis.call('GET', KEYS[1]) +local state = { failedAccountIds = {}, revision = '' } +if raw then + local ok, decoded = pcall(cjson.decode, raw) + if ok and type(decoded) == 'table' and type(decoded.failedAccountIds) == 'table' then + state = decoded + end +end + +local account_id = ARGV[1] +local known = false +for _, failed_account_id in ipairs(state.failedAccountIds) do + if failed_account_id == account_id then + known = true + break + end +end +if not known then + table.insert(state.failedAccountIds, account_id) +end + +state.revision = ARGV[2] +local encoded = cjson.encode(state) +redis.call('SET', KEYS[1], encoded, 'EX', tonumber(ARGV[3])) +return encoded +"#; + +const CLEAR_SCRIPT: &str = r#" +local raw = redis.call('GET', KEYS[1]) +if not raw then + return 0 +end +local ok, state = pcall(cjson.decode, raw) +if not ok or type(state) ~= 'table' then + return 0 +end +if tostring(state.revision or '') ~= ARGV[1] then + return 0 +end +return redis.call('DEL', KEYS[1]) +"#; + +#[derive(Clone)] +pub struct RedisProviderSessionExclusionRepository { + connection: ConnectionManager, + namespace: String, +} + +#[derive(Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +struct SessionExclusionWire { + failed_account_ids: Vec, + #[serde(default)] + revision: String, +} + +impl RedisProviderSessionExclusionRepository { + pub fn new(connection: ConnectionManager, key_namespace: &str) -> StoreResult { + Ok(Self { + connection, + namespace: namespace(key_namespace)?, + }) + } + + fn key( + &self, + provider_kind: &ProviderKind, + affinity_key: &ProviderSessionAffinityKey, + ) -> Result { + let scope = format!( + "{}\0{}", + provider_kind.as_str(), + affinity_key.expose_to_store() + ); + let fingerprint = resource_fingerprint("provider session exclusion", &scope) + .map_err(|_| provider_invalid("encode provider session exclusion key"))?; + Ok(format!( + "{}:scheduler:session-exclusion:{{{fingerprint}}}", + self.namespace + )) + } +} + +impl ProviderSessionExclusionPort for RedisProviderSessionExclusionRepository { + fn load<'a>( + &'a self, + provider_kind: &'a ProviderKind, + key: &'a ProviderSessionAffinityKey, + ) -> BoxFuture<'a, Result, ProviderStoreError>> { + Box::pin(async move { + let mut connection = self.connection.clone(); + let payload = redis::cmd("GET") + .arg(self.key(provider_kind, key)?) + .query_async::>(&mut connection) + .await + .map_err(|_| provider_unavailable("load provider session exclusion"))?; + payload.map(decode_state).transpose() + }) + } + + fn record_failure<'a>( + &'a self, + provider_kind: &'a ProviderKind, + key: &'a ProviderSessionAffinityKey, + account_id: &'a ProviderAccountId, + ttl: Duration, + ) -> BoxFuture<'a, Result> { + Box::pin(async move { + if ttl.is_zero() || ttl > MAX_SESSION_EXCLUSION_TTL { + return Err(provider_invalid("validate provider session exclusion TTL")); + } + let ttl_seconds = ttl.as_secs().max(1); + let mut connection = self.connection.clone(); + let payload = redis::Script::new(RECORD_FAILURE_SCRIPT) + .key(self.key(provider_kind, key)?) + .arg(account_id.as_str()) + .arg(Uuid::new_v4().to_string()) + .arg(ttl_seconds) + .invoke_async::(&mut connection) + .await + .map_err(|_| provider_unavailable("record provider session exclusion"))?; + decode_state(payload) + }) + } + + fn clear<'a>( + &'a self, + provider_kind: &'a ProviderKind, + key: &'a ProviderSessionAffinityKey, + expected_revision: &'a str, + ) -> BoxFuture<'a, Result> { + Box::pin(async move { + let mut connection = self.connection.clone(); + let removed = redis::Script::new(CLEAR_SCRIPT) + .key(self.key(provider_kind, key)?) + .arg(expected_revision) + .invoke_async::(&mut connection) + .await + .map_err(|_| provider_unavailable("clear provider session exclusion"))?; + Ok(removed > 0) + }) + } +} + +fn decode_state(payload: impl AsRef) -> Result { + let state: SessionExclusionWire = serde_json::from_str(payload.as_ref()) + .map_err(|_| provider_invalid("decode provider session exclusion"))?; + let excluded_accounts = state + .failed_account_ids + .into_iter() + .map(|account_id| { + ProviderAccountId::new(account_id) + .map_err(|_| provider_invalid("decode provider session exclusion")) + }) + .collect::, _>>()?; + Ok(ProviderSessionExclusions::new( + excluded_accounts, + state.revision, + )) +} + +fn provider_unavailable(operation: &'static str) -> ProviderStoreError { + ProviderStoreError::new(ProviderStoreErrorKind::Unavailable, operation) +} + +fn provider_invalid(operation: &'static str) -> ProviderStoreError { + ProviderStoreError::new(ProviderStoreErrorKind::InvalidData, operation) +} diff --git a/backend/crates/gateway-store/tests/postgres/history.rs b/backend/crates/gateway-store/tests/postgres/history.rs index 97079ea42..ed9338f9a 100644 --- a/backend/crates/gateway-store/tests/postgres/history.rs +++ b/backend/crates/gateway-store/tests/postgres/history.rs @@ -1,9 +1,5 @@ use chrono::Utc; -use gateway_store::postgres::{ - ModelRequestHistoryRecord, ModelRequestHistoryRepository, PgHistoryRepository, -}; - -use super::TestDatabase; +use gateway_store::postgres::ModelRequestHistoryRecord; #[test] fn history_contains_no_payload_field() { @@ -20,160 +16,6 @@ fn history_contains_no_payload_field() { started_at: Utc::now(), completed_at: None, }; + assert_eq!(history.id, "request-1"); } - -#[tokio::test] -async fn native_pin_is_caller_scoped_and_uses_live_account_and_upstream_handle() { - let Some(database) = TestDatabase::create("native_pin").await else { - return; - }; - seed_native_history(&database.pool) - .await - .expect("seed native history"); - let repository = PgHistoryRepository::new(database.pool.clone()); - - let pin = repository - .resolve_native_continuation_pin("resp_gateway_visible", "key_owner") - .await - .expect("resolve native pin") - .expect("eligible pin"); - assert_eq!(pin.previous_response_id().as_str(), "resp_gateway_visible"); - assert_eq!(pin.upstream_response_id().as_str(), "resp_upstream_native"); - assert_eq!(pin.account().as_str(), "acct_history"); - let debug = format!("{pin:?}"); - assert!(!debug.contains("resp_gateway_visible")); - assert!(!debug.contains("resp_upstream_native")); - - let other_client = repository - .resolve_native_continuation_pin("resp_gateway_visible", "key_other") - .await - .expect("resolve other caller"); - assert!(other_client.is_none()); - - sqlx::query("delete from provider_accounts where id = 'acct_history'") - .execute(&database.pool) - .await - .expect("delete live account while preserving historical ref"); - let deleted_account = repository - .resolve_native_continuation_pin("resp_gateway_visible", "key_owner") - .await - .expect("resolve deleted account"); - assert!(deleted_account.is_none()); - let historical_ref: Option = sqlx::query_scalar( - "select provider_account_ref from model_requests where id = 'req_history'", - ) - .fetch_one(&database.pool) - .await - .expect("load historical account ref"); - assert_eq!(historical_ref.as_deref(), Some("acct_history")); - - database.close().await; -} - -#[tokio::test] -async fn native_pin_requires_committed_success_but_not_a_specific_transport() { - let Some(database) = TestDatabase::create("native_pin_boundaries").await else { - return; - }; - seed_native_history(&database.pool) - .await - .expect("seed native history"); - let repository = PgHistoryRepository::new(database.pool.clone()); - - sqlx::query( - "update model_requests set downstream_committed_at = null where id = 'req_history'", - ) - .execute(&database.pool) - .await - .expect("clear downstream commit"); - assert!( - repository - .resolve_native_continuation_pin("resp_gateway_visible", "key_owner",) - .await - .expect("resolve uncommitted request") - .is_none() - ); - - sqlx::query( - "update model_requests - set downstream_committed_at = now(), outcome = 'failed' - where id = 'req_history'", - ) - .execute(&database.pool) - .await - .expect("mark request failed"); - assert!( - repository - .resolve_native_continuation_pin("resp_gateway_visible", "key_owner",) - .await - .expect("resolve failed request") - .is_none() - ); - - sqlx::query("update model_requests set outcome = 'succeeded' where id = 'req_history'") - .execute(&database.pool) - .await - .expect("restore successful request"); - assert!( - repository - .resolve_native_continuation_pin("resp_gateway_visible", "key_owner",) - .await - .expect("resolve restored request") - .is_some() - ); - - sqlx::query( - "update model_requests set upstream_transport = 'http_sse' where id = 'req_history'", - ) - .execute(&database.pool) - .await - .expect("mark response connection-independent"); - assert!( - repository - .resolve_native_continuation_pin("resp_gateway_visible", "key_owner",) - .await - .expect("resolve HTTP response") - .is_some() - ); - - database.close().await; -} - -async fn seed_native_history(pool: &sqlx::PgPool) -> Result<(), sqlx::Error> { - sqlx::query( - "insert into provider_accounts ( - id, provider_kind, name, upstream_user_id, - provider_credentials_json, credential_revision, has_refresh_token, - access_token_expires_at, enabled, availability, availability_observed_at, - created_at, updated_at - ) values ( - 'acct_history', 'openai', 'history', 'user-history', - '{}'::jsonb, 1, false, now() + interval '1 day', true, 'ready', - now(), now(), now() - )", - ) - .execute(pool) - .await?; - sqlx::query( - "insert into model_requests ( - id, client_api_key_ref, config_revision, protocol, operation, endpoint, - client_transport, requested_model_id, - provider_kind, provider_account_id, - provider_account_ref, upstream_model_id, upstream_transport, attempt_count, - upstream_send_state, downstream_committed_at, outcome, client_status_code, - upstream_status_code, client_response_id, upstream_response_id, - cost_source, started_at, deadline_at, completed_at - ) values ( - 'req_history', 'key_owner', 1, 'openai', 'responses', '/v1/responses', - 'http_sse', 'history-model', 'openai', 'acct_history', - 'acct_history', 'history-upstream', - 'websocket', 1, 'sent', now(), 'succeeded', 200, 200, - 'resp_gateway_visible', 'resp_upstream_native', 'unavailable', - now() - interval '1 minute', now() + interval '1 minute', now() - )", - ) - .execute(pool) - .await?; - Ok(()) -} diff --git a/backend/crates/gateway-store/tests/redis/mod.rs b/backend/crates/gateway-store/tests/redis/mod.rs index 433a3dbda..1fe2f8a13 100644 --- a/backend/crates/gateway-store/tests/redis/mod.rs +++ b/backend/crates/gateway-store/tests/redis/mod.rs @@ -2,9 +2,11 @@ mod client_admission; mod credential_cooldown; mod credential_leases; mod credential_state; +mod native_continuation; mod oauth_pending; mod provider_circuit; mod provider_session_affinity; +mod provider_session_exclusion; mod worker_lease; use chrono::Utc; diff --git a/backend/crates/gateway-store/tests/redis/native_continuation.rs b/backend/crates/gateway-store/tests/redis/native_continuation.rs new file mode 100644 index 000000000..6374a50e3 --- /dev/null +++ b/backend/crates/gateway-store/tests/redis/native_continuation.rs @@ -0,0 +1,123 @@ +use gateway_core::engine::continuation::{ + NativeContinuationPin, NativeContinuationPort, NativeContinuationScope, PreviousResponseId, +}; +use gateway_core::engine::credential::ProviderAccountId; +use gateway_core::error::SafeUpstreamValue; +use gateway_core::operation::ProviderSessionState; +use gateway_core::routing::ProviderKind; +use gateway_store::redis::RedisNativeContinuationRepository; +use redis::aio::ConnectionManager; +use serde_json::{Map, Value}; +use uuid::Uuid; + +#[tokio::test] +async fn native_continuation_round_trips_opaque_state_without_exposing_response_id_in_keys() { + let Some((repository, mut connection, namespace)) = repository().await else { + return; + }; + let previous_response_id = + PreviousResponseId::new("resp-client-visible-secret").expect("previous response ID"); + let state = ProviderSessionState::new( + "openai", + Map::from_iter([( + "turn_state".to_owned(), + Value::String("opaque-turn".to_owned()), + )]), + ) + .expect("provider session state"); + + repository + .record( + NativeContinuationPin::new( + previous_response_id.clone(), + SafeUpstreamValue::new("resp-upstream-handle").expect("upstream response ID"), + ProviderKind::new("openai").expect("provider"), + ProviderAccountId::new("acct_primary").expect("account"), + ) + .with_scope(NativeContinuationScope::Persisted) + .with_session_state(state.clone()), + ) + .await + .expect("record continuation"); + + let resolved = repository + .resolve(&previous_response_id) + .await + .expect("resolve continuation") + .expect("stored continuation"); + assert_eq!( + resolved.upstream_response_id().as_str(), + "resp-upstream-handle" + ); + assert_eq!(resolved.account().as_str(), "acct_primary"); + assert_eq!(resolved.session_state(), Some(&state)); + + let keys = namespace_keys(&mut connection, &namespace).await; + assert!( + keys.iter() + .all(|key| !key.contains("resp-client-visible-secret")) + ); + let entry_key = keys + .iter() + .find(|key| key.contains(":entry:{")) + .expect("continuation entry key"); + let ttl = redis::cmd("PTTL") + .arg(entry_key) + .query_async::(&mut connection) + .await + .expect("read continuation TTL"); + assert!((14_000_000..=14_400_000).contains(&ttl)); + + delete_namespace_keys(&mut connection, &namespace).await; +} + +#[tokio::test] +async fn native_continuation_rejects_provider_state_for_a_different_provider() { + let Some((repository, mut connection, namespace)) = repository().await else { + return; + }; + let pin = NativeContinuationPin::new( + PreviousResponseId::new("resp-mismatched-state").expect("previous response ID"), + SafeUpstreamValue::new("resp-upstream-handle").expect("upstream response ID"), + ProviderKind::new("openai").expect("provider"), + ProviderAccountId::new("acct_primary").expect("account"), + ) + .with_session_state( + ProviderSessionState::new("xai", Map::new()).expect("provider session state"), + ); + + assert!(repository.record(pin).await.is_err()); + assert!(namespace_keys(&mut connection, &namespace).await.is_empty()); +} + +async fn repository() -> Option<(RedisNativeContinuationRepository, ConnectionManager, String)> { + let redis_url = std::env::var("CPR_TEST_REDIS_URL").ok()?; + let client = redis::Client::open(redis_url).expect("valid CPR_TEST_REDIS_URL"); + let connection = client + .get_connection_manager() + .await + .expect("connect test Redis"); + let namespace = format!("gateway-store-native-continuation-test-{}", Uuid::new_v4()); + let repository = RedisNativeContinuationRepository::new(connection.clone(), &namespace) + .expect("valid test namespace"); + Some((repository, connection, namespace)) +} + +async fn namespace_keys(connection: &mut ConnectionManager, namespace: &str) -> Vec { + redis::cmd("KEYS") + .arg(format!("{namespace}:*")) + .query_async(connection) + .await + .expect("list isolated native continuation keys") +} + +async fn delete_namespace_keys(connection: &mut ConnectionManager, namespace: &str) { + let keys = namespace_keys(connection, namespace).await; + if !keys.is_empty() { + redis::cmd("DEL") + .arg(keys) + .query_async::(connection) + .await + .expect("delete isolated native continuation keys"); + } +} diff --git a/backend/crates/gateway-store/tests/redis/provider_session_exclusion.rs b/backend/crates/gateway-store/tests/redis/provider_session_exclusion.rs new file mode 100644 index 000000000..3546e2431 --- /dev/null +++ b/backend/crates/gateway-store/tests/redis/provider_session_exclusion.rs @@ -0,0 +1,113 @@ +use std::time::Duration; + +use gateway_core::engine::credential::ProviderAccountId; +use gateway_core::provider_ports::{ProviderSessionAffinityKey, ProviderSessionExclusionPort}; +use gateway_core::routing::ProviderKind; +use gateway_store::redis::RedisProviderSessionExclusionRepository; +use redis::aio::ConnectionManager; +use uuid::Uuid; + +#[tokio::test] +async fn session_exclusion_should_record_each_failed_account_for_one_hour() { + let Some((repository, mut connection, namespace)) = repository().await else { + return; + }; + let provider = ProviderKind::new("openai").expect("provider"); + let session = ProviderSessionAffinityKey::try_new("cyber-policy-session").expect("session"); + let first = ProviderAccountId::new("acct_first").expect("first account"); + let second = ProviderAccountId::new("acct_second").expect("second account"); + + repository + .record_failure(&provider, &session, &first, Duration::from_secs(60 * 60)) + .await + .expect("record first failure"); + let state = repository + .record_failure(&provider, &session, &second, Duration::from_secs(60 * 60)) + .await + .expect("record second failure"); + + assert_eq!(state.excluded_accounts().len(), 2); + assert!(state.excluded_accounts().contains(&first)); + assert!(state.excluded_accounts().contains(&second)); + let keys = namespace_keys(&mut connection, &namespace).await; + assert_eq!(keys.len(), 1); + assert!(!keys[0].contains("cyber-policy-session")); + let ttl = redis::cmd("PTTL") + .arg(&keys[0]) + .query_async::(&mut connection) + .await + .expect("read exclusion TTL"); + assert!((3_590_000..=3_600_000).contains(&ttl)); +} + +#[tokio::test] +async fn session_exclusion_should_clear_only_the_state_observed_by_the_successful_request() { + let Some((repository, _connection, _namespace)) = repository().await else { + return; + }; + let provider = ProviderKind::new("openai").expect("provider"); + let session = ProviderSessionAffinityKey::try_new("cyber-policy-cas").expect("session"); + let first = ProviderAccountId::new("acct_first").expect("first account"); + let second = ProviderAccountId::new("acct_second").expect("second account"); + + let observed = repository + .record_failure(&provider, &session, &first, Duration::from_secs(60 * 60)) + .await + .expect("record observed failure"); + let current = repository + .record_failure(&provider, &session, &second, Duration::from_secs(60 * 60)) + .await + .expect("record concurrent failure"); + + assert!( + !repository + .clear(&provider, &session, observed.revision()) + .await + .expect("compare-and-swap clear") + ); + assert_eq!( + repository + .load(&provider, &session) + .await + .expect("load retained state"), + Some(current.clone()) + ); + assert!( + repository + .clear(&provider, &session, current.revision()) + .await + .expect("clear current state") + ); + assert_eq!( + repository + .load(&provider, &session) + .await + .expect("load cleared state"), + None + ); +} + +async fn repository() -> Option<( + RedisProviderSessionExclusionRepository, + ConnectionManager, + String, +)> { + let redis_url = std::env::var("CPR_TEST_REDIS_URL").ok()?; + let client = redis::Client::open(redis_url).expect("valid CPR_TEST_REDIS_URL"); + let connection = client + .get_connection_manager() + .await + .expect("connect test Redis"); + let namespace = format!("gateway-store-session-exclusion-test-{}", Uuid::new_v4()); + let repository = RedisProviderSessionExclusionRepository::new(connection.clone(), &namespace) + .expect("valid test namespace"); + Some((repository, connection, namespace)) +} + +async fn namespace_keys(connection: &mut ConnectionManager, namespace: &str) -> Vec { + redis::cmd("KEYS") + .arg(format!("{namespace}:*")) + .query_async(connection) + .await + .expect("list isolated session exclusion keys") +} diff --git a/backend/crates/providers/openai/src/config.rs b/backend/crates/providers/openai/src/config.rs index d955c3f3f..9e43d53ca 100644 --- a/backend/crates/providers/openai/src/config.rs +++ b/backend/crates/providers/openai/src/config.rs @@ -1,29 +1,261 @@ //! OpenAI Provider 启动配置与 Codex Desktop 请求画像校验。 -use std::path::Path; +use std::path::{Path, PathBuf}; +use std::time::Duration; use chrono::{DateTime, Utc}; use serde::Deserialize; +use url::Url; +use crate::credential::CodexQuotaRefreshPolicy; use crate::transport::profile::{CodexWireProfile, CodexWireProfileState}; +use crate::transport::session::{CodexSessionIdentity, CodexSessionIdentityError}; +use crate::transport::websocket::CodexWebSocketPoolConfig; +use crate::{ + OFFICIAL_CODEX_BASE_URL, + credential::token_client::{ + OFFICIAL_CODEX_OAUTH_CLIENT_ID, OFFICIAL_CODEX_TOKEN_ENDPOINT, TokenClientConfig, + }, +}; /// OpenAI Provider 唯一启动配置。 #[derive(Debug, Clone, Deserialize, PartialEq, Eq)] #[serde(deny_unknown_fields)] pub struct OpenAiConfig { + #[serde(default)] + pub api: CodexApiConfig, + #[serde(default)] + pub ws_pool: CodexWebSocketPoolSettings, + #[serde(default)] + pub quota: CodexQuotaSettings, + #[serde(default)] + pub auth: CodexAuthSettings, pub wire_profile: CodexWireProfileConfig, + #[serde(skip)] + identity_secret_path: PathBuf, } impl OpenAiConfig { - /// 校验 Provider-owned 字段；当前配置不含文件系统路径。 - pub fn resolve_and_validate(&mut self, _source_dir: &Path) -> Result<(), OpenAiConfigError> { - self.wire_profile.validate() + /// 校验 Provider-owned 字段，并定位旧版会话锚点持续使用的身份密钥。 + pub fn resolve_and_validate(&mut self, source_dir: &Path) -> Result<(), OpenAiConfigError> { + self.api.validate()?; + self.ws_pool.validate()?; + self.quota.validate()?; + self.auth.validate()?; + self.wire_profile.validate()?; + self.identity_secret_path = source_dir + .parent() + .unwrap_or(source_dir) + .join(".runtime/data/identity_hmac_secret"); + Ok(()) } #[must_use] pub fn wire_profile_state(&self) -> CodexWireProfileState { CodexWireProfileState::new(self.wire_profile.clone().into()) } + + #[must_use] + pub fn base_url(&self) -> &str { + &self.api.base_url + } + + #[must_use] + pub fn websocket_pool_config(&self) -> CodexWebSocketPoolConfig { + self.ws_pool.pool_config() + } + + #[must_use] + pub fn quota_refresh_policy(&self) -> CodexQuotaRefreshPolicy { + self.quota.refresh_policy() + } + + #[must_use] + pub const fn quota_skip_exhausted(&self) -> bool { + self.quota.skip_exhausted + } + + #[must_use] + pub fn token_client_config(&self) -> TokenClientConfig { + TokenClientConfig { + client_id: self.auth.oauth_client_id.clone(), + token_endpoint: self.auth.oauth_token_endpoint.clone(), + } + } + + #[must_use] + pub fn oauth_client_id(&self) -> &str { + &self.auth.oauth_client_id + } + + #[must_use] + pub const fn oauth_refresh_enabled(&self) -> bool { + self.auth.refresh_enabled + } + + pub(crate) fn session_identity( + &self, + ) -> Result { + CodexSessionIdentity::load_or_create(&self.identity_secret_path) + } +} + +impl Default for OpenAiConfig { + fn default() -> Self { + Self { + api: CodexApiConfig::default(), + ws_pool: CodexWebSocketPoolSettings::default(), + quota: CodexQuotaSettings::default(), + auth: CodexAuthSettings::default(), + wire_profile: CodexWireProfileConfig::default(), + identity_secret_path: PathBuf::from(".runtime/data/identity_hmac_secret"), + } + } +} + +/// Codex 上游 API 的 Provider-owned 地址配置。 +#[derive(Debug, Clone, Deserialize, PartialEq, Eq)] +#[serde(deny_unknown_fields)] +pub struct CodexApiConfig { + pub base_url: String, +} + +impl Default for CodexApiConfig { + fn default() -> Self { + Self { + base_url: OFFICIAL_CODEX_BASE_URL.to_owned(), + } + } +} + +impl CodexApiConfig { + fn validate(&self) -> Result<(), OpenAiConfigError> { + let url = Url::parse(&self.base_url) + .map_err(|_| OpenAiConfigError::InvalidField("openai.api.base_url"))?; + if !matches!(url.scheme(), "http" | "https") || url.host_str().is_none() { + return Err(OpenAiConfigError::InvalidField("openai.api.base_url")); + } + Ok(()) + } +} + +/// Codex Responses WebSocket pool 的 Provider-owned 启动设置。 +#[derive(Debug, Clone, Deserialize, PartialEq, Eq)] +#[serde(deny_unknown_fields)] +pub struct CodexWebSocketPoolSettings { + pub enabled: bool, + pub max_age_ms: u64, + pub max_per_account: usize, + pub max_total: usize, + pub max_connecting: usize, + pub initial_event_timeout_ms: u64, +} + +impl Default for CodexWebSocketPoolSettings { + fn default() -> Self { + Self { + enabled: true, + max_age_ms: 55 * 60 * 1000, + max_per_account: 8, + max_total: 64, + max_connecting: 8, + initial_event_timeout_ms: 800, + } + } +} + +impl CodexWebSocketPoolSettings { + fn validate(&self) -> Result<(), OpenAiConfigError> { + if self.max_age_ms == 0 + || self.max_per_account == 0 + || self.max_total == 0 + || self.max_connecting == 0 + || self.max_connecting > self.max_total + { + return Err(OpenAiConfigError::InvalidField("openai.ws_pool")); + } + Ok(()) + } + + fn pool_config(&self) -> CodexWebSocketPoolConfig { + CodexWebSocketPoolConfig { + enabled: self.enabled, + max_age: Duration::from_millis(self.max_age_ms), + max_per_account: self.max_per_account, + max_total: self.max_total, + max_connecting: self.max_connecting, + initial_event_timeout: (self.initial_event_timeout_ms != 0) + .then(|| Duration::from_millis(self.initial_event_timeout_ms)), + ..CodexWebSocketPoolConfig::default() + } + } +} + +/// OpenAI Provider 的额度刷新策略。 +#[derive(Debug, Clone, Deserialize, PartialEq, Eq)] +#[serde(deny_unknown_fields)] +pub struct CodexQuotaSettings { + pub refresh_interval_minutes: u64, + pub skip_exhausted: bool, +} + +impl Default for CodexQuotaSettings { + fn default() -> Self { + Self { + refresh_interval_minutes: 5, + skip_exhausted: true, + } + } +} + +impl CodexQuotaSettings { + fn validate(&self) -> Result<(), OpenAiConfigError> { + if self.refresh_interval_minutes == 0 { + return Err(OpenAiConfigError::InvalidField( + "openai.quota.refresh_interval_minutes", + )); + } + Ok(()) + } + + fn refresh_policy(&self) -> CodexQuotaRefreshPolicy { + CodexQuotaRefreshPolicy::new(Duration::from_secs( + self.refresh_interval_minutes.saturating_mul(60), + )) + } +} + +/// OpenAI OAuth 的 Provider-owned 运行开关和端点。 +#[derive(Debug, Clone, Deserialize, PartialEq, Eq)] +#[serde(deny_unknown_fields)] +pub struct CodexAuthSettings { + pub refresh_enabled: bool, + pub oauth_client_id: String, + pub oauth_token_endpoint: String, +} + +impl Default for CodexAuthSettings { + fn default() -> Self { + Self { + refresh_enabled: true, + oauth_client_id: OFFICIAL_CODEX_OAUTH_CLIENT_ID.to_owned(), + oauth_token_endpoint: OFFICIAL_CODEX_TOKEN_ENDPOINT.to_owned(), + } + } +} + +impl CodexAuthSettings { + fn validate(&self) -> Result<(), OpenAiConfigError> { + if self.oauth_client_id.trim().is_empty() + || Url::parse(&self.oauth_token_endpoint) + .ok() + .filter(|url| matches!(url.scheme(), "http" | "https") && url.host_str().is_some()) + .is_none() + { + return Err(OpenAiConfigError::InvalidField("openai.auth")); + } + Ok(()) + } } /// 经审计固定的 Codex Desktop 上游请求画像。 @@ -41,6 +273,22 @@ pub struct CodexWireProfileConfig { pub verified_at: DateTime, } +impl Default for CodexWireProfileConfig { + fn default() -> Self { + Self { + originator: "Codex Desktop".to_owned(), + codex_version: "0.144.6".to_owned(), + desktop_version: "26.715.61943".to_owned(), + desktop_build: "5628".to_owned(), + os_type: "Mac OS".to_owned(), + os_version: "15.7.1".to_owned(), + arch: "arm64".to_owned(), + terminal: "unknown".to_owned(), + verified_at: Utc::now(), + } + } +} + impl CodexWireProfileConfig { fn validate(&self) -> Result<(), OpenAiConfigError> { for (field, value) in [ diff --git a/backend/crates/providers/openai/src/credential/affinity.rs b/backend/crates/providers/openai/src/credential/affinity.rs index aaabf4f8d..d94b89bf1 100644 --- a/backend/crates/providers/openai/src/credential/affinity.rs +++ b/backend/crates/providers/openai/src/credential/affinity.rs @@ -1,5 +1,6 @@ //! OpenAI 会话锚点到 Store 不透明亲和键的单向派生。 +use gateway_core::policy::ClientApiKeyId; use gateway_core::provider_ports::ProviderSessionAffinityKey; use sha2::{Digest, Sha256}; @@ -9,6 +10,14 @@ use crate::transport::request::derive_conversation_anchor; pub(crate) fn derive_codex_session_affinity_key( request: &CodexResponsesRequest, ) -> Option { + if let Some(conversation_id) = request + .local_conversation_id + .as_deref() + .map(str::trim) + .filter(|value| !value.is_empty()) + { + return ProviderSessionAffinityKey::try_new(conversation_id.to_owned()).ok(); + } let (domain, value) = derive_conversation_anchor(request)?; let mut hasher = Sha256::new(); hasher.update(b"codex-session-affinity-v1\0"); @@ -17,3 +26,35 @@ pub(crate) fn derive_codex_session_affinity_key( hasher.update(value.as_bytes()); ProviderSessionAffinityKey::try_new(hex::encode(hasher.finalize())).ok() } + +/// 恢复旧版 `cyber_policy` 的会话隔离键。 +/// +/// 它只接受显式 session/conversation 或客户端明确给出的 prompt cache key，避免将 +/// 请求内容哈希误当成长会话；`previous_response_id` 续写不参与该策略。 +pub(crate) fn derive_codex_cyber_policy_session_key( + request: &CodexResponsesRequest, + client_api_key_id: &ClientApiKeyId, +) -> Option { + if request.previous_response_id().is_some() { + return None; + } + let session_id = non_empty(request.client_session_id.as_deref()) + .or_else(|| non_empty(request.client_conversation_id.as_deref())) + .or_else(|| { + request + .explicit_prompt_cache_key + .then(|| request.prompt_cache_key()) + .flatten() + .and_then(|value| non_empty(Some(value))) + })?; + let mut hasher = Sha256::new(); + hasher.update(b"cyber-policy-session\0"); + hasher.update(client_api_key_id.as_str().as_bytes()); + hasher.update(b"\0"); + hasher.update(session_id.as_bytes()); + ProviderSessionAffinityKey::try_new(hex::encode(hasher.finalize())).ok() +} + +fn non_empty(value: Option<&str>) -> Option<&str> { + value.map(str::trim).filter(|value| !value.is_empty()) +} diff --git a/backend/crates/providers/openai/src/credential/catalog.rs b/backend/crates/providers/openai/src/credential/catalog.rs index 8e080d7af..8e2f5f319 100644 --- a/backend/crates/providers/openai/src/credential/catalog.rs +++ b/backend/crates/providers/openai/src/credential/catalog.rs @@ -16,7 +16,6 @@ use uuid::Uuid; use super::agent_identity::CodexAgentIdentityTaskService; use super::repository::{CodexCredentialRepository, CredentialRepositoryError}; -use crate::provider::OFFICIAL_CODEX_BASE_URL; use crate::transport::profile::CodexWireProfileState; use crate::transport::{CodexBackendClient, CodexCatalogModel, CodexRequestContext}; @@ -124,6 +123,7 @@ pub struct CodexCredentialCatalogService { repository: CodexCredentialRepository, profile: CodexWireProfileState, http: reqwest::Client, + base_url: String, agent_identity: Arc, cache: Arc>, etags: Arc>, @@ -135,12 +135,14 @@ impl CodexCredentialCatalogService { repository: CodexCredentialRepository, profile: CodexWireProfileState, http: reqwest::Client, + base_url: String, agent_identity: Arc, ) -> Self { Self { repository, profile, http, + base_url, agent_identity, cache: Arc::new(RwLock::new(CatalogCacheState::default())), etags: Arc::new(Mutex::new(CatalogEtagState::default())), @@ -210,7 +212,7 @@ impl CodexCredentialCatalogService { .ok_or(CodexCredentialCatalogError::NoEligibleCredential)?; let client = CodexBackendClient::new( self.http.clone(), - OFFICIAL_CODEX_BASE_URL, + self.base_url.clone(), self.profile.clone(), ); let cache_revision = self.cache_revision()?; @@ -254,7 +256,7 @@ impl CodexCredentialCatalogService { } let client = CodexBackendClient::new( self.http.clone(), - OFFICIAL_CODEX_BASE_URL, + self.base_url.clone(), self.profile.clone(), ); let mut union = BTreeMap::::new(); diff --git a/backend/crates/providers/openai/src/credential/identity.rs b/backend/crates/providers/openai/src/credential/identity.rs index 11e761f8f..d21f27e89 100644 --- a/backend/crates/providers/openai/src/credential/identity.rs +++ b/backend/crates/providers/openai/src/credential/identity.rs @@ -42,7 +42,6 @@ const MAX_TOKEN_LIFETIME_SECONDS: u64 = 366 * 24 * 60 * 60; const JWKS_CACHE_TTL: Duration = Duration::from_secs(60 * 60); const REQUEST_TIMEOUT: Duration = Duration::from_secs(30); const CONNECT_TIMEOUT: Duration = Duration::from_secs(10); -const OFFICIAL_CODEX_BASE_URL: &str = "https://chatgpt.com/backend-api"; /// OAuth identity verification 的稳定失败分类。 #[derive(Debug, Clone, Copy, PartialEq, Eq, thiserror::Error)] @@ -339,6 +338,8 @@ pub struct CodexJwtIdentityVerifier { jwks_source: Box, cache: Mutex, cache_ttl: Duration, + api_audience: String, + oauth_client_id: String, } impl CodexJwtIdentityVerifier { @@ -348,9 +349,17 @@ impl CodexJwtIdentityVerifier { jwks_source, cache: Mutex::new(JwksCache::default()), cache_ttl: JWKS_CACHE_TTL, + api_audience: OFFICIAL_OPENAI_API_AUDIENCE.to_owned(), + oauth_client_id: OFFICIAL_CODEX_OAUTH_CLIENT_ID.to_owned(), } } + #[must_use] + pub fn with_oauth_client_id(mut self, oauth_client_id: impl Into) -> Self { + self.oauth_client_id = oauth_client_id.into(); + self + } + async fn decoding_key( &self, key_id: &str, @@ -391,7 +400,7 @@ impl CodexJwtIdentityVerifier { .as_deref() .ok_or(CodexIdentityVerificationError::Rejected)?; let decoding_key = self.decoding_key(key_id).await?; - let mut validation = official_validation(OFFICIAL_OPENAI_API_AUDIENCE); + let mut validation = official_validation(&self.api_audience); validation.set_required_spec_claims(&["sub", "iss", "aud", "exp", "iat"]); Ok( decode::(access_token, &decoding_key, &validation) @@ -421,7 +430,7 @@ impl CodexJwtIdentityVerifier { .as_deref() .ok_or(CodexIdentityVerificationError::Rejected)?; let decoding_key = self.decoding_key(key_id).await?; - let mut validation = official_validation(OFFICIAL_CODEX_OAUTH_CLIENT_ID); + let mut validation = official_validation(&self.oauth_client_id); validation.set_required_spec_claims(&["sub", "iss", "aud", "exp", "iat", "nonce"]); Ok( decode::(id_token, &decoding_key, &validation) @@ -449,7 +458,7 @@ impl CodexSignedIdentityVerifier for CodexJwtIdentityVerifier { ) -> Result { self.decode_access_token(secret) .await? - .into_signed_identity() + .into_signed_identity(&self.api_audience) } async fn verify_authorization( @@ -460,24 +469,31 @@ impl CodexSignedIdentityVerifier for CodexJwtIdentityVerifier { ) -> Result { let access = self.decode_access_token(secret).await?; let id = self.decode_id_token(id_token).await?; - if access.sub != id.sub || !id.valid_for_nonce(expected_nonce.expose_secret()) { + if access.sub != id.sub + || !id.valid_for_nonce(expected_nonce.expose_secret(), &self.oauth_client_id) + { return Err(CodexIdentityVerificationError::Rejected); } - access.into_signed_identity() + access.into_signed_identity(&self.api_audience) } } /// 官方 `/wham/usage` 账号身份 source。 pub struct ReqwestCodexAuthenticatedAccountSource { client: CodexBackendClient, + base_url: String, } impl ReqwestCodexAuthenticatedAccountSource { - pub fn new(profile: CodexWireProfileState) -> Result { + pub fn new( + profile: CodexWireProfileState, + base_url: String, + ) -> Result { let http = build_reqwest_client().map_err(|_| CodexIdentityVerificationError::Unavailable)?; Ok(Self { - client: CodexBackendClient::new(http, OFFICIAL_CODEX_BASE_URL, profile), + client: CodexBackendClient::new(http, base_url.clone(), profile), + base_url, }) } } @@ -486,7 +502,7 @@ impl fmt::Debug for ReqwestCodexAuthenticatedAccountSource { fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { formatter .debug_struct("ReqwestCodexAuthenticatedAccountSource") - .field("endpoint", &OFFICIAL_CODEX_BASE_URL) + .field("endpoint", &self.base_url) .finish_non_exhaustive() } } @@ -648,10 +664,10 @@ struct IdTokenClaims { } impl IdTokenClaims { - fn valid_for_nonce(&self, expected_nonce: &str) -> bool { + fn valid_for_nonce(&self, expected_nonce: &str, oauth_client_id: &str) -> bool { let now = get_current_timestamp(); self.iss == OFFICIAL_OPENAI_ISSUER - && self.aud.is_exact(OFFICIAL_CODEX_OAUTH_CLIENT_ID) + && self.aud.is_exact(oauth_client_id) && self.exp > now && self.iat <= now && self.iat < self.exp @@ -666,10 +682,13 @@ impl IdTokenClaims { } impl AccessTokenClaims { - fn into_signed_identity(self) -> Result { + fn into_signed_identity( + self, + api_audience: &str, + ) -> Result { let now = get_current_timestamp(); if self.iss != OFFICIAL_OPENAI_ISSUER - || !self.aud.is_exact(OFFICIAL_OPENAI_API_AUDIENCE) + || !self.aud.is_exact(api_audience) || self.exp <= now || self.iat > now || self.iat >= self.exp diff --git a/backend/crates/providers/openai/src/credential/mod.rs b/backend/crates/providers/openai/src/credential/mod.rs index a6f1df1bb..c495e03da 100644 --- a/backend/crates/providers/openai/src/credential/mod.rs +++ b/backend/crates/providers/openai/src/credential/mod.rs @@ -16,7 +16,9 @@ pub mod token_client; mod types; pub(crate) use admin::refresh_time; -pub(crate) use affinity::derive_codex_session_affinity_key; +pub(crate) use affinity::{ + derive_codex_cyber_policy_session_key, derive_codex_session_affinity_key, +}; pub(crate) use oauth::oauth_owner_ref; pub use agent_identity::{ @@ -50,8 +52,8 @@ pub use oauth::{ }; pub use quota::{ CodexAccountQuotaSnapshot, CodexCredentialQuotaError, CodexCredentialQuotaService, - CodexQuotaFact, CodexQuotaSyncSummary, CodexQuotaWindow, CodexQuotaWindowKind, - CodexQuotaWindowRole, parse_codex_quota_usage, + CodexQuotaFact, CodexQuotaRefreshPolicy, CodexQuotaSyncSummary, CodexQuotaWindow, + CodexQuotaWindowKind, CodexQuotaWindowRole, parse_codex_quota_usage, }; pub use refresh::{ CodexCredentialRefreshError, CodexCredentialRefreshOutcome, CodexCredentialRefreshService, @@ -62,6 +64,7 @@ pub use security::{ CodexCredentialCodec, CodexCredentialDataError, CodexRuntimeAuthentication, CodexRuntimeCredential, }; +pub(crate) use selector::CodexCyberPolicyScope; pub use selector::{ CodexAccountFailure, CodexCredentialLease, CodexCredentialSelector, CredentialSelectionError, SelectCodexCredential, diff --git a/backend/crates/providers/openai/src/credential/oauth.rs b/backend/crates/providers/openai/src/credential/oauth.rs index 315b6e51e..abcbe9964 100644 --- a/backend/crates/providers/openai/src/credential/oauth.rs +++ b/backend/crates/providers/openai/src/credential/oauth.rs @@ -333,11 +333,12 @@ pub struct CodexOAuthAdminService { store: Arc, runtime_policy: Arc, credentials: CodexCredentialAdmin, + oauth_client_id: String, } impl CodexOAuthAdminService { #[must_use] - pub const fn new( + pub fn new( pending: Arc, exchanger: Arc, verifier: Arc, @@ -352,8 +353,15 @@ impl CodexOAuthAdminService { store, runtime_policy, credentials, + oauth_client_id: OFFICIAL_CODEX_OAUTH_CLIENT_ID.to_owned(), } } + + #[must_use] + pub fn with_oauth_client_id(mut self, oauth_client_id: impl Into) -> Self { + self.oauth_client_id = oauth_client_id.into(); + self + } } #[async_trait] @@ -520,7 +528,7 @@ impl CodexOAuthAdminService { .map_err(map_pending_error)?; Ok(CodexOAuthAuthorizationStarted { flow_id: pending.flow_id.clone(), - authorization_url: authorization_url(&pending)?, + authorization_url: authorization_url(&pending, &self.oauth_client_id)?, expires_at, }) } @@ -590,7 +598,10 @@ fn oauth_identity_expectation( .map_err(|_| CodexOAuthAdminError::Credential) } -fn authorization_url(pending: &CodexPendingAuthorization) -> Result { +fn authorization_url( + pending: &CodexPendingAuthorization, + oauth_client_id: &str, +) -> Result { let mut url = Url::parse(AUTHORIZATION_ENDPOINT).map_err(|_| CodexOAuthAdminError::InvalidInput)?; let challenge = URL_SAFE_NO_PAD.encode(Sha256::digest( @@ -598,7 +609,7 @@ fn authorization_url(pending: &CodexPendingAuthorization) -> Result Self { + Self { interval } + } + + #[must_use] + pub const fn interval(self) -> Duration { + self.interval + } +} + #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub struct CodexQuotaFact { remaining_percent: Option, @@ -199,6 +219,7 @@ pub struct CodexCredentialQuotaService { store: Arc, profile: CodexWireProfileState, http: Client, + base_url: String, agent_identity: Arc, scheduling: CodexQuotaSchedulingProjection, } @@ -436,6 +457,7 @@ impl CodexCredentialQuotaService { repository: CodexCredentialRepository, profile: CodexWireProfileState, http: Client, + base_url: String, agent_identity: Arc, ) -> Self { Self { @@ -443,6 +465,7 @@ impl CodexCredentialQuotaService { repository, profile, http, + base_url, agent_identity, scheduling: CodexQuotaSchedulingProjection::default(), } @@ -502,7 +525,7 @@ impl CodexCredentialQuotaService { } let client = CodexBackendClient::new( self.http.clone(), - OFFICIAL_CODEX_BASE_URL, + self.base_url.clone(), self.profile.clone(), ); for account in accounts { @@ -721,7 +744,7 @@ impl CodexCredentialQuotaService { .ok_or(CodexCredentialQuotaError::NotFound)?; let client = CodexBackendClient::new( self.http.clone(), - OFFICIAL_CODEX_BASE_URL, + self.base_url.clone(), self.profile.clone(), ); let observed_at = SystemTime::now(); diff --git a/backend/crates/providers/openai/src/credential/selector.rs b/backend/crates/providers/openai/src/credential/selector.rs index 8b65244dc..9fbde29cb 100644 --- a/backend/crates/providers/openai/src/credential/selector.rs +++ b/backend/crates/providers/openai/src/credential/selector.rs @@ -13,7 +13,7 @@ use gateway_core::engine::{AttemptContext, ContinuationAttempt}; use gateway_core::provider_ports::{ ProviderLeaseAcquisition, ProviderLeaseGuard, ProviderLeasePort, ProviderLeaseRequest, ProviderSchedulingLeaseRequest, ProviderSessionAffinityKey, ProviderSessionAffinityPort, - ProviderStoreError, + ProviderSessionExclusionPort, ProviderSessionExclusions, ProviderStoreError, }; use gateway_core::routing::ProviderKind; use secrecy::ExposeSecret; @@ -41,8 +41,9 @@ const CLOUDFLARE_CHALLENGE_BACKOFF: [Duration; 4] = [ Duration::from_secs(120), ]; const CLOUDFLARE_PATH_BLOCK_THRESHOLD: u32 = 3; -const SESSION_AFFINITY_TTL: Duration = Duration::from_secs(60 * 60); +const SESSION_AFFINITY_TTL: Duration = Duration::from_secs(4 * 60 * 60); const SESSION_AFFINITY_TIMEOUT: Duration = Duration::from_millis(100); +const CYBER_POLICY_SESSION_TTL: Duration = Duration::from_secs(60 * 60); #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub enum CodexAccountFailure { @@ -89,17 +90,25 @@ pub struct SelectCodexCredential<'a> { pub session_affinity_key: Option<&'a ProviderSessionAffinityKey>, } +#[derive(Clone)] +pub(crate) struct CodexCyberPolicyScope { + key: ProviderSessionAffinityKey, + state: Option, +} + pub struct CodexCredentialSelector { provider_kind: ProviderKind, repository: CodexCredentialRepository, leases: Arc, session_affinity: Arc, + session_exclusions: Arc, catalog: Arc, quota: Arc, agent_identity: Arc, cookie_policy: CodexCookiePolicy, risk_recovery: Mutex>, account_feedback: Arc, + skip_exhausted: bool, } enum SessionAffinityLookup { @@ -117,29 +126,41 @@ impl CodexCredentialSelector { repository: CodexCredentialRepository, leases: Arc, session_affinity: Arc, + session_exclusions: Arc, catalog: Arc, quota: Arc, agent_identity: Arc, account_feedback: Arc, cookie_policy: CodexCookiePolicy, + skip_exhausted: bool, ) -> Self { Self { provider_kind, repository, leases, session_affinity, + session_exclusions, catalog, quota, agent_identity, cookie_policy, risk_recovery: Mutex::new(HashMap::new()), account_feedback, + skip_exhausted, } } pub async fn select( &self, request: &SelectCodexCredential<'_>, + ) -> Result { + self.select_with_cyber_policy(request, None).await + } + + pub(crate) async fn select_with_cyber_policy( + &self, + request: &SelectCodexCredential<'_>, + cyber_policy_session_key: Option<&ProviderSessionAffinityKey>, ) -> Result { let accounts = self.repository.list_for_provider().await?; let accounts = accounts @@ -189,10 +210,23 @@ impl CodexCredentialSelector { AccountCandidate { account, signals } }) .collect::>(); + // 旧版开关只影响“候选资格”。保留原始账号供 revision-fenced 凭据加载， + // 用影子候选承载放宽后的资格，避免伪造已持久化账号的运行状态。 + let quota_exhausted_candidates = (!self.skip_exhausted).then(|| { + candidates + .iter() + .cloned() + .map(|mut candidate| { + candidate.account = quota_exhausted_scheduling_projection(&candidate.account); + candidate + }) + .collect::>() + }); + let scheduling_candidates = quota_exhausted_candidates.as_deref().unwrap_or(&candidates); let affinity_account = match request.session_affinity_key { Some(key) => match self.lookup_session_affinity(key).await { SessionAffinityLookup::Bound(account_id) - if candidates.iter().any(|candidate| { + if scheduling_candidates.iter().any(|candidate| { candidate.account.id() == &account_id && candidate.account.is_schedulable(SystemTime::now()) }) => @@ -207,7 +241,16 @@ impl CodexCredentialSelector { }, None => None, }; + let cyber_policy_scope = self + .prepare_cyber_policy_scope(cyber_policy_session_key) + .await; let mut excluded = request.attempt.excluded_accounts().clone(); + if let Some(state) = cyber_policy_scope + .as_ref() + .and_then(|scope| scope.state.as_ref()) + { + excluded.extend(state.excluded_accounts().iter().cloned()); + } let continuation_account = match request.attempt.continuation_attempt() { ContinuationAttempt::Native => request .attempt @@ -248,7 +291,7 @@ impl CodexCredentialSelector { preferred_account: preferred.clone(), round_robin_cursor, }; - let Some(selected) = AccountSelector.select(&candidates, &context) else { + let Some(selected) = AccountSelector.select(scheduling_candidates, &context) else { return match shortest_retry { Some(retry_after) => Err(CredentialSelectionError::CapacityUnavailable { retry_after: Some(retry_after), @@ -256,7 +299,11 @@ impl CodexCredentialSelector { None => Err(CredentialSelectionError::NoEligibleCredential), }; }; - let account = selected.account.clone(); + let account = candidates + .iter() + .find(|candidate| candidate.account.id() == selected.account.id()) + .map(|candidate| candidate.account.clone()) + .ok_or(CredentialSelectionError::InvalidCredential)?; let policy = request.attempt.account_selection_policy(); match self .leases @@ -311,6 +358,7 @@ impl CodexCredentialSelector { account, authentication: runtime.authentication, cookies, + cyber_policy_scope, _guard: guard, }); } @@ -364,6 +412,97 @@ impl CodexCredentialSelector { } } + async fn prepare_cyber_policy_scope( + &self, + key: Option<&ProviderSessionAffinityKey>, + ) -> Option { + let key = key?.clone(); + let state = match tokio::time::timeout( + SESSION_AFFINITY_TIMEOUT, + self.session_exclusions.load(&self.provider_kind, &key), + ) + .await + { + Ok(Ok(state)) => state, + Ok(Err(error)) => { + tracing::warn!(error = %error, "OpenAI cyber policy state read failed open"); + None + } + Err(_) => { + tracing::warn!( + timeout_ms = SESSION_AFFINITY_TIMEOUT.as_millis(), + "OpenAI cyber policy state read timed out" + ); + None + } + }; + Some(CodexCyberPolicyScope { key, state }) + } + + pub(crate) async fn record_cyber_policy_failure( + &self, + scope: Option<&CodexCyberPolicyScope>, + account: &ProviderAccount, + ) { + let Some(scope) = scope else { + return; + }; + match tokio::time::timeout( + SESSION_AFFINITY_TIMEOUT, + self.session_exclusions.record_failure( + &self.provider_kind, + &scope.key, + account.id(), + CYBER_POLICY_SESSION_TTL, + ), + ) + .await + { + Ok(Ok(_)) => {} + Ok(Err(error)) => { + tracing::warn!( + account_id = %account.id(), + error = %error, + "OpenAI cyber policy exclusion write failed open" + ); + } + Err(_) => { + tracing::warn!( + account_id = %account.id(), + timeout_ms = SESSION_AFFINITY_TIMEOUT.as_millis(), + "OpenAI cyber policy exclusion write timed out" + ); + } + } + } + + pub(crate) async fn observe_cyber_policy_success(&self, scope: Option<&CodexCyberPolicyScope>) { + let Some(scope) = scope.filter(|scope| scope.state.is_some()) else { + return; + }; + let Some(state) = scope.state.as_ref() else { + return; + }; + match tokio::time::timeout( + SESSION_AFFINITY_TIMEOUT, + self.session_exclusions + .clear(&self.provider_kind, &scope.key, state.revision()), + ) + .await + { + Ok(Ok(_)) => {} + Ok(Err(error)) => { + tracing::warn!(error = %error, "OpenAI cyber policy exclusion clear failed open"); + } + Err(_) => { + tracing::warn!( + timeout_ms = SESSION_AFFINITY_TIMEOUT.as_millis(), + "OpenAI cyber policy exclusion clear timed out" + ); + } + } + } + pub async fn record_failure( &self, account: &ProviderAccount, @@ -452,12 +591,11 @@ impl CodexCredentialSelector { let Some(key) = session_affinity_key else { return; }; - match self.lookup_session_affinity(key).await { - SessionAffinityLookup::Bound(bound_account) if &bound_account != account.id() => { - return; - } - SessionAffinityLookup::Unavailable => return, - SessionAffinityLookup::Missing | SessionAffinityLookup::Bound(_) => {} + if matches!( + self.lookup_session_affinity(key).await, + SessionAffinityLookup::Unavailable + ) { + return; } match tokio::time::timeout( SESSION_AFFINITY_TIMEOUT, @@ -615,6 +753,17 @@ impl CodexCredentialSelector { } } +fn quota_exhausted_scheduling_projection(account: &ProviderAccount) -> ProviderAccount { + if account.availability() == AccountAvailability::QuotaExhausted { + return account.clone().with_runtime_state( + account.enabled(), + AccountAvailability::Ready, + account.cooldown_until(), + ); + } + account.clone() +} + fn active_risk_recovery<'a>( recovery: &'a mut HashMap, account_id: &str, @@ -651,6 +800,7 @@ pub struct CodexCredentialLease { authentication: CodexRuntimeAuthentication, cookies: Vec, installation_id: String, + cyber_policy_scope: Option, _guard: Box, } @@ -679,6 +829,11 @@ impl CodexCredentialLease { pub fn installation_id(&self) -> &str { &self.installation_id } + + #[must_use] + pub(crate) const fn cyber_policy_scope(&self) -> Option<&CodexCyberPolicyScope> { + self.cyber_policy_scope.as_ref() + } } impl fmt::Debug for CodexCredentialLease { diff --git a/backend/crates/providers/openai/src/credential/token_client.rs b/backend/crates/providers/openai/src/credential/token_client.rs index 61660cbd4..e6ad9bd95 100644 --- a/backend/crates/providers/openai/src/credential/token_client.rs +++ b/backend/crates/providers/openai/src/credential/token_client.rs @@ -139,25 +139,29 @@ impl OpenAiTokenClient { } } -/// 构建固定官方 endpoint、禁止 redirect 且无自动重试的 Codex token client。 +/// 构建禁止 redirect 且无自动重试的 Codex token client。 /// /// # Errors /// /// 本地 TLS/HTTP client 初始化失败时返回脱敏错误。 -pub fn official_openai_token_client() -> Result { +pub fn openai_token_client( + config: TokenClientConfig, +) -> Result { let client = Client::builder() .redirect(Policy::none()) .connect_timeout(TOKEN_CONNECT_TIMEOUT) .timeout(TOKEN_REQUEST_TIMEOUT) .build() .map_err(|_| TokenClientBuildError)?; - Ok(OpenAiTokenClient::new( - client, - TokenClientConfig { - client_id: OFFICIAL_CODEX_OAUTH_CLIENT_ID.to_owned(), - token_endpoint: OFFICIAL_CODEX_TOKEN_ENDPOINT.to_owned(), - }, - )) + Ok(OpenAiTokenClient::new(client, config)) +} + +/// 构建固定官方 endpoint 的 Codex token client。 +pub fn official_openai_token_client() -> Result { + openai_token_client(TokenClientConfig { + client_id: OFFICIAL_CODEX_OAUTH_CLIENT_ID.to_owned(), + token_endpoint: OFFICIAL_CODEX_TOKEN_ENDPOINT.to_owned(), + }) } #[derive(Deserialize)] diff --git a/backend/crates/providers/openai/src/lib.rs b/backend/crates/providers/openai/src/lib.rs index 2c47e57e7..ec84c8bf1 100644 --- a/backend/crates/providers/openai/src/lib.rs +++ b/backend/crates/providers/openai/src/lib.rs @@ -4,7 +4,6 @@ mod admin; pub mod config; mod provider; -use std::path::Path; use std::sync::Arc; use gateway_admin::ports::provider::ProviderAdmin; @@ -54,20 +53,21 @@ pub struct ProviderBundle { /// 构造 OpenAI 数据面、Provider-owned 后台任务与 Redis OAuth pending owner。 pub async fn initialize( - mut config: OpenAiConfig, + config: OpenAiConfig, ports: ProviderStorePorts, ) -> Result { - config - .resolve_and_validate(Path::new(".")) - .map_err(OpenAiInitializeError::Config)?; let provider_kind = ProviderKind::new("openai").map_err(|_| OpenAiInitializeError::InvalidProviderKind)?; let accounts: Arc = ports.accounts(); let leases = ports.leases(); let session_affinity = ports.session_affinity(); + let session_exclusions = ports.session_exclusions(); let account_feedback = ports.account_feedback(); let runtime_policy = ports.runtime_policy(); let profile = config.wire_profile_state(); + let session_identity = config + .session_identity() + .map_err(|_| OpenAiInitializeError::SessionIdentity)?; let http = build_reqwest_client().map_err(|_| OpenAiInitializeError::Transport)?; let desktop_release = Arc::new(CodexDesktopReleaseService::new( profile.clone(), @@ -85,7 +85,9 @@ pub async fn initialize( ), )); let repository = CodexCredentialRepository::new(Arc::clone(&accounts)); - let websocket_pool = Arc::new(CodexWebSocketPool::default()); + let websocket_pool = Arc::new(CodexWebSocketPool::with_config( + config.websocket_pool_config(), + )); let agent_identity = Arc::new(CodexAgentIdentityTaskService::new( repository.clone(), Arc::new( @@ -98,24 +100,29 @@ pub async fn initialize( repository.clone(), profile.clone(), http.clone(), + config.base_url().to_owned(), Arc::clone(&agent_identity), )); let quota = Arc::new(CodexCredentialQuotaService::new( repository.clone(), profile.clone(), http.clone(), + config.base_url().to_owned(), Arc::clone(&agent_identity), )); + let quota_skip_exhausted = config.quota_skip_exhausted(); let selector = Arc::new(CodexCredentialSelector::new( provider_kind.clone(), repository.clone(), Arc::clone(&leases), session_affinity, + session_exclusions, Arc::clone(&catalog), Arc::clone("a), Arc::clone(&agent_identity), Arc::clone(&account_feedback), CodexCookiePolicy::official().map_err(|_| OpenAiInitializeError::CookiePolicy)?, + quota_skip_exhausted, )); let core_provider: Arc = Arc::new( CodexProvider::new( @@ -126,20 +133,24 @@ pub async fn initialize( account_feedback, http, profile.clone(), + config.base_url().to_owned(), Arc::clone(&websocket_pool), ) - .map_err(OpenAiInitializeError::Provider)?, + .map_err(OpenAiInitializeError::Provider)? + .with_session_identity(session_identity), ); let token_client = Arc::new( - credential::token_client::official_openai_token_client() + credential::token_client::openai_token_client(config.token_client_config()) .map_err(|_| OpenAiInitializeError::TokenClient)?, ); let jwks = ReqwestOpenAiJwksSource::new().map_err(|_| OpenAiInitializeError::Identity)?; - let signed: Arc = - Arc::new(CodexJwtIdentityVerifier::new(Box::new(jwks))); + let signed: Arc = Arc::new( + CodexJwtIdentityVerifier::new(Box::new(jwks)) + .with_oauth_client_id(config.oauth_client_id()), + ); let account_source: Arc = Arc::new( - ReqwestCodexAuthenticatedAccountSource::new(profile.clone()) + ReqwestCodexAuthenticatedAccountSource::new(profile.clone(), config.base_url().to_owned()) .map_err(|_| OpenAiInitializeError::Identity)?, ); let identity: Arc = @@ -164,14 +175,17 @@ pub async fn initialize( ports.oauth_pending(), provider_kind.clone(), )); - let oauth_admin: Arc = Arc::new(CodexOAuthAdminService::new( - pending, - exchanger, - Arc::clone(&identity), - Arc::clone(&accounts), - Arc::clone(&runtime_policy), - CodexCredentialAdmin, - )); + let oauth_admin: Arc = Arc::new( + CodexOAuthAdminService::new( + pending, + exchanger, + Arc::clone(&identity), + Arc::clone(&accounts), + Arc::clone(&runtime_policy), + CodexCredentialAdmin, + ) + .with_oauth_client_id(config.oauth_client_id()), + ); let admin_provider: Arc = Arc::new(OpenAiAdminProvider::new( provider_kind, profile, @@ -187,9 +201,16 @@ pub async fn initialize( websocket_pool, desktop_release_status, )); - let worker_contributions = - provider::worker_contributions(refresh, quota, catalog, cli_release, desktop_release) - .map_err(|_| OpenAiInitializeError::Worker)?; + let worker_contributions = provider::worker_contributions( + refresh, + quota, + catalog, + config.quota_refresh_policy(), + config.oauth_refresh_enabled(), + cli_release, + desktop_release, + ) + .map_err(|_| OpenAiInitializeError::Worker)?; Ok(ProviderBundle { core_provider, @@ -224,6 +245,8 @@ pub enum OpenAiInitializeError { RuntimePolicy, #[error("OpenAI Provider kind is invalid")] InvalidProviderKind, + #[error("OpenAI local session identity is unavailable")] + SessionIdentity, #[error("OpenAI transport could not initialize")] Transport, #[error(transparent)] diff --git a/backend/crates/providers/openai/src/provider.rs b/backend/crates/providers/openai/src/provider.rs index a7512255d..303b10446 100644 --- a/backend/crates/providers/openai/src/provider.rs +++ b/backend/crates/providers/openai/src/provider.rs @@ -42,8 +42,9 @@ use url::Url; use crate::credential::{ CodexAccountFailure, CodexAgentIdentityTaskService, CodexCredentialCatalogService, CodexCredentialLease, CodexCredentialQuotaService, CodexCredentialRefreshOutcome, - CodexCredentialRefreshService, CodexCredentialSelector, CredentialSelectionError, - RuntimeCodexCookie, SelectCodexCredential, derive_codex_session_affinity_key, + CodexCredentialRefreshService, CodexCredentialSelector, CodexCyberPolicyScope, + CodexQuotaRefreshPolicy, CredentialSelectionError, RuntimeCodexCookie, SelectCodexCredential, + derive_codex_cyber_policy_session_key, derive_codex_session_affinity_key, }; use crate::transport::canonical::{ CodexCanonicalDecoder, CodexCanonicalError, CodexCanonicalOutcome, @@ -61,6 +62,7 @@ use crate::transport::request::{ CodexRequestEncodeError, encode_generate_request, sanitize_cross_account_item, scope_request_to_account, }; +use crate::transport::session::CodexSessionIdentity; use crate::transport::websocket::{CodexWebSocketExchangeError, PreviousResponseUnavailableReason}; use crate::transport::{ CODEX_RESPONSES_PATH, CodexBackendClient, CodexBackendStreamingResponse, CodexBackendTransport, @@ -83,7 +85,7 @@ pub enum CodexProviderTransport { #[derive(Debug, Clone, PartialEq, Eq, thiserror::Error)] pub enum CodexProviderConfigError { - #[error("official Codex provider URL is invalid")] + #[error("Codex provider URL is invalid")] InvalidBaseUrl, } @@ -95,6 +97,7 @@ pub struct CodexProvider { account_feedback: Arc, client: CodexBackendClient, responses_url: Url, + session_identity: Option, } impl CodexProvider { @@ -108,13 +111,13 @@ impl CodexProvider { account_feedback: Arc, http: Client, profile: CodexWireProfileState, + base_url: String, websocket_pool: Arc, ) -> Result { - let responses_url = - Url::parse(&endpoint_url(OFFICIAL_CODEX_BASE_URL, CODEX_RESPONSES_PATH)) - .map_err(|_| CodexProviderConfigError::InvalidBaseUrl)?; - let client = CodexBackendClient::new(http, OFFICIAL_CODEX_BASE_URL, profile) - .with_websocket_pool(websocket_pool); + let responses_url = Url::parse(&endpoint_url(&base_url, CODEX_RESPONSES_PATH)) + .map_err(|_| CodexProviderConfigError::InvalidBaseUrl)?; + let client = + CodexBackendClient::new(http, base_url, profile).with_websocket_pool(websocket_pool); Ok(Self { selector, catalog, @@ -123,8 +126,14 @@ impl CodexProvider { account_feedback, client, responses_url, + session_identity: None, }) } + + pub(crate) fn with_session_identity(mut self, identity: CodexSessionIdentity) -> Self { + self.session_identity = Some(identity); + self + } } impl fmt::Debug for CodexProvider { @@ -230,19 +239,42 @@ impl Provider for CodexProvider { let mut upstream_request = encode_generate_request(generate, candidate.upstream_model().as_str()) .map_err(map_request_error)?; + if let Some(conversation_id) = previous_session + .as_ref() + .and_then(|state| state.conversation_id.as_ref()) + { + upstream_request.local_conversation_id = Some(conversation_id.clone()); + } + if let Some(identity) = &self.session_identity { + identity.prepare_local_conversation(&mut upstream_request); + } + if upstream_request + .turn_state + .as_deref() + .is_none_or(str::is_empty) + { + upstream_request.turn_state = previous_session + .as_ref() + .and_then(|state| state.turn_state.clone()); + } let session_affinity_key = derive_codex_session_affinity_key(&upstream_request); + let cyber_policy_session_key = + derive_codex_cyber_policy_session_key(&upstream_request, context.client_api_key_ref()); let request_input = upstream_request.input().to_vec(); let transport = selected_transport(&upstream_request); apply_transport(&mut upstream_request, transport); let lease = self .selector - .select(&SelectCodexCredential { - upstream_model: candidate.upstream_model().as_str(), - request_url: &self.responses_url, - attempt: &context, - session_affinity_key: session_affinity_key.as_ref(), - }) + .select_with_cyber_policy( + &SelectCodexCredential { + upstream_model: candidate.upstream_model().as_str(), + request_url: &self.responses_url, + attempt: &context, + session_affinity_key: session_affinity_key.as_ref(), + }, + cyber_policy_session_key.as_ref(), + ) .await .map_err(map_selection_error)?; let lease = Arc::new(lease); @@ -280,12 +312,6 @@ impl Provider for CodexProvider { lease.installation_id(), cross_account, ); - if let Some(conversation_id) = previous_session - .as_ref() - .and_then(|state| state.conversation_id.as_ref()) - { - upstream_request.local_conversation_id = Some(conversation_id.clone()); - } if context.continuation_attempt() == ContinuationAttempt::Native && !replay_previous_response && let Some(continuation) = context.continuation() @@ -346,6 +372,7 @@ impl Provider for CodexProvider { request_input, account_id: lease.account_id().as_str().to_owned(), conversation_id: upstream_request.local_conversation_id.clone(), + turn_state: upstream_request.turn_state.clone(), response_store, continuation_scope: None, }); @@ -389,6 +416,8 @@ struct ColdResponse { struct OpenAiSessionState { account_id: String, conversation_id: Option, + #[serde(default)] + turn_state: Option, continuation_scope: OpenAiContinuationScope, transcript: Vec, } @@ -424,6 +453,7 @@ struct OpenAiSessionCapture { request_input: Vec, account_id: String, conversation_id: Option, + turn_state: Option, response_store: bool, continuation_scope: Option, } @@ -490,6 +520,7 @@ fn attach_openai_session_update( let update = encode_openai_session_state(OpenAiSessionState { account_id: capture.account_id, conversation_id: capture.conversation_id, + turn_state: capture.turn_state, continuation_scope: capture .continuation_scope .ok_or_else(|| provider_error(ProviderErrorKind::Protocol, UpstreamSendState::Sent))?, @@ -625,6 +656,7 @@ fn cold_response_stream(response: ColdResponse) -> EventStream { mut session_capture, } = response; Box::pin(async_stream::try_stream! { + let cyber_policy_scope = lease.cyber_policy_scope().cloned(); let mut active_account = lease.account().clone(); let cookie_header = build_cookie_header(lease.cookies())?; let mut authorization = lease @@ -709,6 +741,7 @@ fn cold_response_stream(response: ColdResponse) -> EventStream { "a, &active_account, &response_origin, + cyber_policy_scope.as_ref(), &failure, ) .await; @@ -727,6 +760,7 @@ fn cold_response_stream(response: ColdResponse) -> EventStream { "a, &active_account, &response_origin, + cyber_policy_scope.as_ref(), &failure, ) .await; @@ -743,6 +777,7 @@ fn cold_response_stream(response: ColdResponse) -> EventStream { } else { OpenAiContinuationScope::ReplayRequired }); + capture.turn_state = response.turn_state.clone().or(capture.turn_state.clone()); } let observation = match codex_response_observation( response.transport, @@ -760,6 +795,7 @@ fn cold_response_stream(response: ColdResponse) -> EventStream { "a, &active_account, &response_origin, + cyber_policy_scope.as_ref(), &failure, ) .await; @@ -796,6 +832,7 @@ fn cold_response_stream(response: ColdResponse) -> EventStream { let failure_set_cookie_headers = response.set_cookie_headers.clone(); let failure_rate_limit_headers = response.rate_limit_headers.clone(); let rate_limit_updates = response.rate_limit_header_updates; + let turn_state_updates = response.turn_state_update; let mut decoder = CodexCanonicalDecoder::new(upstream_model.as_str()); if response.transport == CodexBackendTransport::HttpSse { decoder = decoder.with_raw_sse_passthrough(); @@ -830,6 +867,7 @@ fn cold_response_stream(response: ColdResponse) -> EventStream { "a, &active_account, &response_origin, + cyber_policy_scope.as_ref(), &failure, ) .await; @@ -842,6 +880,12 @@ fn cold_response_stream(response: ColdResponse) -> EventStream { let updates = std::mem::take(&mut *updates.lock().await); synchronize_passive_quota("a, &active_account, &updates).await; } + if let (Some(capture), Some(updates)) = + (session_capture.as_mut(), turn_state_updates.as_ref()) + && let Some(turn_state) = updates.lock().await.clone() + { + capture.turn_state = Some(turn_state); + } let (mut events, terminal_failure) = match decoder.push(&chunk) { CodexCanonicalOutcome::Events(events) => (events, None), CodexCanonicalOutcome::Failed(failure) => { @@ -863,6 +907,7 @@ fn cold_response_stream(response: ColdResponse) -> EventStream { "a, &active_account, &response_origin, + cyber_policy_scope.as_ref(), failure, ) .await; @@ -876,6 +921,7 @@ fn cold_response_stream(response: ColdResponse) -> EventStream { "a, &active_account, &response_origin, + cyber_policy_scope.as_ref(), &failure, ) .await; @@ -900,6 +946,9 @@ fn cold_response_stream(response: ColdResponse) -> EventStream { selector .record_success(&active_account, session_affinity_key.as_ref()) .await; + selector + .observe_cyber_policy_success(cyber_policy_scope.as_ref()) + .await; return; } } @@ -924,10 +973,17 @@ fn cold_response_stream(response: ColdResponse) -> EventStream { "a, &active_account, &response_origin, + cyber_policy_scope.as_ref(), failure, ) .await; } + if let (Some(capture), Some(updates)) = + (session_capture.as_mut(), turn_state_updates.as_ref()) + && let Some(turn_state) = updates.lock().await.clone() + { + capture.turn_state = Some(turn_state); + } if let Err(error) = attach_openai_session_update(&mut events, &mut session_capture) { if terminal_failure.is_none() { let failure = MappedProviderFailure::plain(error); @@ -937,6 +993,7 @@ fn cold_response_stream(response: ColdResponse) -> EventStream { "a, &active_account, &response_origin, + cyber_policy_scope.as_ref(), &failure, ) .await; @@ -961,6 +1018,9 @@ fn cold_response_stream(response: ColdResponse) -> EventStream { selector .record_success(&active_account, session_affinity_key.as_ref()) .await; + selector + .observe_cyber_policy_success(cyber_policy_scope.as_ref()) + .await; } }) } @@ -1254,6 +1314,7 @@ fn map_selection_error(error: CredentialSelectionError) -> ProviderError { struct MappedProviderFailure { error: ProviderError, account_failure: Option, + cyber_policy_failure: bool, set_cookie_headers: Vec, rate_limit_headers: Vec<(String, String)>, observation: Option, @@ -1265,6 +1326,7 @@ impl MappedProviderFailure { Self { error, account_failure: None, + cyber_policy_failure: false, set_cookie_headers: Vec::new(), rate_limit_headers: Vec::new(), observation: None, @@ -1299,9 +1361,15 @@ async fn apply_failure( quota: &CodexCredentialQuotaService, account: &ProviderAccount, response_origin: &Url, + cyber_policy_scope: Option<&CodexCyberPolicyScope>, failure: &MappedProviderFailure, ) { synchronize_passive_quota(quota, account, &failure.rate_limit_headers).await; + if failure.cyber_policy_failure { + selector + .record_cyber_policy_failure(cyber_policy_scope, account) + .await; + } if let Some(account_failure) = failure.account_failure { client.evict_websocket_account(account.id().as_str()).await; if let Err(error) = selector.record_failure(account, account_failure).await { @@ -1433,6 +1501,10 @@ fn map_upstream_failure( replay_boundary: ReplayBoundary, ) -> MappedProviderFailure { let category = failure.category(); + let cyber_policy_failure = failure + .status + .is_some_and(|status| status.is_client_error()) + && is_cyber_policy_code(failure.code.as_deref()); let continuation_failure = failure .persistable_code() .filter(|code| is_history_failure_code(code)) @@ -1452,7 +1524,7 @@ fn map_upstream_failure( error = error.with_status(status.as_u16()); } if replay_boundary.permits_provider_proof() - && (failure.replay_is_safe() || continuation_failure.is_some()) + && (failure.replay_is_safe() || continuation_failure.is_some() || cyber_policy_failure) { error = error.with_replay_safe(); } @@ -1478,6 +1550,7 @@ fn map_upstream_failure( MappedProviderFailure { error, account_failure: account_failure(category, failure.retry_after_seconds), + cyber_policy_failure, set_cookie_headers: failure.set_cookie_headers, rate_limit_headers: failure.rate_limit_headers, observation, @@ -1488,6 +1561,10 @@ fn map_upstream_failure( } } +fn is_cyber_policy_code(code: Option<&str>) -> bool { + code.is_some_and(|code| code.trim().eq_ignore_ascii_case("cyber_policy")) +} + fn is_history_failure_code(code: &str) -> bool { matches!( code, @@ -1622,7 +1699,6 @@ const WORKER_MAXIMUM_BACKOFF: Duration = Duration::from_secs(60); const WORKER_LEASE_TTL: Duration = Duration::from_secs(15 * 60); const WORKER_LEASE_RENEWAL: Duration = Duration::from_secs(5 * 60); const OAUTH_REFRESH_INTERVAL: Duration = Duration::from_secs(30); -const QUOTA_REFRESH_INTERVAL: Duration = Duration::from_secs(5 * 60); const DESKTOP_RELEASE_WORKER_OWNER: &str = "openai-desktop-release"; const CLI_RELEASE_WORKER_OWNER: &str = "openai-cli-release"; const MODEL_ETAG_WORKER_OWNER: &str = "openai-model-etag"; @@ -1631,6 +1707,8 @@ pub(crate) fn worker_contributions( refresh: Arc, quota: Arc, catalog: Arc, + quota_refresh_policy: CodexQuotaRefreshPolicy, + oauth_refresh_enabled: bool, cli_release: Arc, desktop_release: Arc, ) -> Result, WorkerDefinitionError> { @@ -1641,15 +1719,18 @@ pub(crate) fn worker_contributions( WorkerId::try_new(WorkerKind::QuotaCatalogHealth, DESKTOP_RELEASE_WORKER_OWNER)?; let cli_release_id = WorkerId::try_new(WorkerKind::QuotaCatalogHealth, CLI_RELEASE_WORKER_OWNER)?; - Ok(vec![ - WorkerContribution::Registration(scheduled_registration( + let mut contributions = Vec::new(); + if oauth_refresh_enabled { + contributions.push(WorkerContribution::Registration(scheduled_registration( refresh_id, OAUTH_REFRESH_INTERVAL, Box::new(OpenAiOAuthRefreshTask { service: refresh }), - )?), + )?)); + } + contributions.extend([ WorkerContribution::Registration(scheduled_registration( quota_id, - QUOTA_REFRESH_INTERVAL, + quota_refresh_policy.interval(), Box::new(OpenAiQuotaTask { quota }), )?), WorkerContribution::Registration(WorkerRegistration::try_new( @@ -1676,7 +1757,8 @@ pub(crate) fn worker_contributions( service: desktop_release, }), )?), - ]) + ]); + Ok(contributions) } fn scheduled_registration( diff --git a/backend/crates/providers/openai/src/transport/mod.rs b/backend/crates/providers/openai/src/transport/mod.rs index a3b54b624..63a79210b 100644 --- a/backend/crates/providers/openai/src/transport/mod.rs +++ b/backend/crates/providers/openai/src/transport/mod.rs @@ -11,6 +11,7 @@ pub mod profile; pub mod protocol; pub mod request; mod response_meta; +pub(crate) mod session; mod time; pub mod tls; pub mod usage; diff --git a/backend/crates/providers/openai/src/transport/request.rs b/backend/crates/providers/openai/src/transport/request.rs index 732a72c55..fb5e45b62 100644 --- a/backend/crates/providers/openai/src/transport/request.rs +++ b/backend/crates/providers/openai/src/transport/request.rs @@ -372,38 +372,17 @@ fn non_empty_string(value: Option<&Value>) -> Option { .map(ToOwned::to_owned) } -fn derive_account_scoped_local_conversation_id( - request: &mut CodexResponsesRequest, - installation_id: &str, -) { - let anchor = derive_conversation_anchor(request); - let Some((domain, value)) = anchor else { - return; - }; - let digest = hmac_sha256( - installation_id.as_bytes(), - &[ - b"codex-local-conversation-v2\0", - domain.as_bytes(), - b"\0", - value.as_bytes(), - ], - ); - request.local_conversation_id = Some(format!("lc_{}", hex::encode(digest))); -} - pub(crate) fn derive_conversation_anchor( request: &CodexResponsesRequest, ) -> Option<(&'static str, String)> { request - .client_session_id - .as_deref() - .map(|value| ("session", value.to_owned())) + .prompt_cache_key() + .map(|value| ("prompt-cache", value.to_owned())) .or_else(|| { request - .client_conversation_id + .client_session_id .as_deref() - .map(|value| ("conversation", value.to_owned())) + .map(|value| ("session", value.to_owned())) }) .or_else(|| { request @@ -413,8 +392,9 @@ pub(crate) fn derive_conversation_anchor( }) .or_else(|| { request - .prompt_cache_key() - .map(|value| ("prompt-cache", value.to_owned())) + .client_conversation_id + .as_deref() + .map(|value| ("conversation", value.to_owned())) }) .or_else(|| derive_stable_conversation_key(request).map(|value| ("request", value))) } @@ -494,37 +474,6 @@ fn normalize_conversation_anchor_text(text: &str) -> String { rest.to_owned() } -fn hmac_sha256(key: &[u8], message: &[&[u8]]) -> [u8; 32] { - const BLOCK_BYTES: usize = 64; - - let mut key_block = [0_u8; BLOCK_BYTES]; - if key.len() > BLOCK_BYTES { - key_block[..32].copy_from_slice(&Sha256::digest(key)); - } else { - key_block[..key.len()].copy_from_slice(key); - } - let mut inner_pad = key_block; - let mut outer_pad = key_block; - for byte in &mut inner_pad { - *byte ^= 0x36; - } - for byte in &mut outer_pad { - *byte ^= 0x5c; - } - - let mut inner = Sha256::new(); - inner.update(inner_pad); - for part in message { - inner.update(part); - } - let inner_digest = inner.finalize(); - - let mut outer = Sha256::new(); - outer.update(outer_pad); - outer.update(inner_digest); - outer.finalize().into() -} - /// 把客户端正文收敛到当前 lease 的账号身份边界。 /// /// 真实 account ID 与 installation ID 由随后构造的 `CodexRequestContext` 注入 @@ -639,7 +588,6 @@ pub(crate) fn scope_request_to_account( request.turn_state = turn_state; request.turn_metadata = turn_metadata; - derive_account_scoped_local_conversation_id(request, installation_id); } fn sanitize_cross_account_input(request: &mut CodexResponsesRequest) { diff --git a/backend/crates/providers/openai/src/transport/session.rs b/backend/crates/providers/openai/src/transport/session.rs new file mode 100644 index 000000000..da2a2dee8 --- /dev/null +++ b/backend/crates/providers/openai/src/transport/session.rs @@ -0,0 +1,125 @@ +//! Codex 本地会话锚点的稳定派生。 +//! +//! 旧版 Redis 会话亲和与 WebSocket 复用以 `.runtime/data/identity_hmac_secret` +//! 生成的 `lc_` 值为键。该密钥是持久运行数据，不是账号凭据；保留它可避免 +//! 架构迁移后同一客户端会话全部换键。 + +use std::fs; +use std::io::Write as _; +use std::path::Path; + +use base64::{Engine as _, engine::general_purpose::URL_SAFE_NO_PAD}; +use sha2::{Digest, Sha256}; + +use super::protocol::responses::CodexResponsesRequest; +use super::request::derive_conversation_anchor; + +/// OpenAI Provider 的持久会话锚点密钥。 +#[derive(Clone)] +pub(crate) struct CodexSessionIdentity { + secret: [u8; 32], +} + +impl CodexSessionIdentity { + pub(crate) fn load_or_create(path: &Path) -> Result { + match read_secret(path)? { + Some(secret) => Ok(Self { secret }), + None => Self::create(path), + } + } + + fn create(path: &Path) -> Result { + let parent = path + .parent() + .ok_or(CodexSessionIdentityError::InvalidPath)?; + fs::create_dir_all(parent).map_err(|_| CodexSessionIdentityError::Storage)?; + let mut secret = [0_u8; 32]; + getrandom::fill(&mut secret).map_err(|_| CodexSessionIdentityError::Storage)?; + let mut options = fs::OpenOptions::new(); + options.write(true).create_new(true); + #[cfg(unix)] + { + use std::os::unix::fs::OpenOptionsExt as _; + + options.mode(0o600); + } + match options.open(path) { + Ok(mut file) => { + file.write_all(hex::encode(secret).as_bytes()) + .map_err(|_| CodexSessionIdentityError::Storage)?; + file.sync_all() + .map_err(|_| CodexSessionIdentityError::Storage)?; + Ok(Self { secret }) + } + Err(error) if error.kind() == std::io::ErrorKind::AlreadyExists => read_secret(path)? + .map(|secret| Self { secret }) + .ok_or(CodexSessionIdentityError::InvalidSecret), + Err(_) => Err(CodexSessionIdentityError::Storage), + } + } + + /// 在账号选择前固定一次旧版 `lc_` 会话键。 + pub(crate) fn prepare_local_conversation(&self, request: &mut CodexResponsesRequest) { + if request.local_conversation_id.is_some() { + return; + } + let Some((_, anchor)) = derive_conversation_anchor(request) else { + return; + }; + request.local_conversation_id = Some(format!( + "lc_{}", + URL_SAFE_NO_PAD.encode(hmac_sha256( + &self.secret, + &[b"local-conversation", b"\0", b"\0", anchor.as_bytes()], + )) + )); + } +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, thiserror::Error)] +pub(crate) enum CodexSessionIdentityError { + #[error("Codex session identity path is invalid")] + InvalidPath, + #[error("Codex session identity secret is invalid")] + InvalidSecret, + #[error("Codex session identity storage is unavailable")] + Storage, +} + +fn read_secret(path: &Path) -> Result, CodexSessionIdentityError> { + let raw = match fs::read_to_string(path) { + Ok(raw) => raw, + Err(error) if error.kind() == std::io::ErrorKind::NotFound => return Ok(None), + Err(_) => return Err(CodexSessionIdentityError::Storage), + }; + let bytes = hex::decode(raw.trim()).map_err(|_| CodexSessionIdentityError::InvalidSecret)?; + bytes + .try_into() + .map(Some) + .map_err(|_| CodexSessionIdentityError::InvalidSecret) +} + +fn hmac_sha256(key: &[u8; 32], parts: &[&[u8]]) -> [u8; 32] { + const BLOCK_BYTES: usize = 64; + + let mut key_block = [0_u8; BLOCK_BYTES]; + key_block[..key.len()].copy_from_slice(key); + let mut inner_pad = key_block; + let mut outer_pad = key_block; + for byte in &mut inner_pad { + *byte ^= 0x36; + } + for byte in &mut outer_pad { + *byte ^= 0x5c; + } + let mut inner = Sha256::new(); + inner.update(inner_pad); + for part in parts { + inner.update(part); + } + let inner_digest = inner.finalize(); + let mut outer = Sha256::new(); + outer.update(outer_pad); + outer.update(inner_digest); + outer.finalize().into() +} diff --git a/backend/crates/providers/openai/tests/admin.rs b/backend/crates/providers/openai/tests/admin.rs index 064b91daa..85598bf93 100644 --- a/backend/crates/providers/openai/tests/admin.rs +++ b/backend/crates/providers/openai/tests/admin.rs @@ -48,9 +48,11 @@ use provider_openai::credential::ImportCodexOAuthCredential; use provider_openai::transport::profile::APPCAST_POLL_INTERVAL; use serde_json::{Map, Value, json}; use std::collections::BTreeSet; +use tempfile::TempDir; use crate::support::{ - MemoryAccountStore, MemorySessionAffinity, TestLeaseCoordinator, profile, secret, + MemoryAccountStore, MemorySessionAffinity, MemorySessionExclusions, TestLeaseCoordinator, + profile, secret, }; #[tokio::test] @@ -64,8 +66,9 @@ async fn real_openai_conversation_crosses_production_provider_boundaries() { panic!("real OpenAI fixture must be an object"); }; let store = Arc::new(MemoryAccountStore::default()); + let config = valid_config(); let bundle = provider_openai::initialize( - valid_config(), + config.config.clone(), provider_ports_with(Arc::clone(&store), Arc::new(TestOAuthPending::default())), ) .await @@ -190,7 +193,8 @@ fn real_attempt_context() -> AttemptContext { #[tokio::test] async fn openai_bundle_exposes_one_core_provider_and_drains_worker_contributions_once() { - let mut bundle = provider_openai::initialize(valid_config(), provider_ports()) + let config = valid_config(); + let mut bundle = provider_openai::initialize(config.config.clone(), provider_ports()) .await .expect("OpenAI bundle"); @@ -237,7 +241,8 @@ async fn openai_bundle_exposes_one_core_provider_and_drains_worker_contributions #[tokio::test] async fn openai_admin_provider_exposes_live_wire_profile_and_validated_billing() { - let bundle = provider_openai::initialize(valid_config(), provider_ports()) + let config = valid_config(); + let bundle = provider_openai::initialize(config.config.clone(), provider_ports()) .await .expect("OpenAI bundle"); let admin = bundle.admin_provider(); @@ -274,7 +279,8 @@ async fn openai_admin_provider_exposes_live_wire_profile_and_validated_billing() #[tokio::test] async fn openai_core_provider_projects_codex_request_observation_without_routing_side_effects() { - let bundle = provider_openai::initialize(valid_config(), provider_ports()) + let config = valid_config(); + let bundle = provider_openai::initialize(config.config.clone(), provider_ports()) .await .expect("OpenAI bundle"); let payload = ProtocolPayload::json_object( @@ -305,8 +311,9 @@ async fn openai_core_provider_projects_codex_request_observation_without_routing #[tokio::test] async fn openai_admin_provider_persists_the_full_pending_envelope_and_binds_owner() { let pending = Arc::new(TestOAuthPending::default()); + let config = valid_config(); let bundle = provider_openai::initialize( - valid_config(), + config.config.clone(), provider_ports_with( Arc::new(MemoryAccountStore::default()), Arc::clone(&pending), @@ -386,8 +393,9 @@ async fn openai_admin_provider_projects_cached_quota_models_and_canonical_export .account("acct_admin_projection") .expect("stored account"); let record = account_record(&account); + let config = valid_config(); let bundle = provider_openai::initialize( - valid_config(), + config.config.clone(), provider_ports_with(Arc::clone(&store), Arc::new(TestOAuthPending::default())), ) .await @@ -501,8 +509,9 @@ async fn openai_admin_provider_projects_codex_additional_limit_as_the_primary_qu }) .await .expect("persist quota"); + let config = valid_config(); let bundle = provider_openai::initialize( - valid_config(), + config.config.clone(), provider_ports_with(Arc::clone(&store), Arc::new(TestOAuthPending::default())), ) .await @@ -544,8 +553,9 @@ async fn openai_admin_provider_rejects_unprepared_mutations_before_store_commit( .await; let account = store.account("acct_admin_invalid").expect("stored account"); let record = account_record(&account); + let config = valid_config(); let bundle = provider_openai::initialize( - valid_config(), + config.config.clone(), provider_ports_with(store, Arc::new(TestOAuthPending::default())), ) .await @@ -591,6 +601,7 @@ fn provider_ports_with( accounts, Arc::new(TestLeaseCoordinator::default()), Arc::new(MemorySessionAffinity::default()), + Arc::new(MemorySessionExclusions::default()), Arc::new(TestCatalogCache::default()), Arc::new(TestCredentialState), Arc::new(TestCooldown), @@ -625,22 +636,34 @@ fn account_record(account: &ProviderAccount) -> AccountRecord { } } -fn valid_config() -> OpenAiConfig { - OpenAiConfig { - wire_profile: CodexWireProfileConfig { - originator: "Codex Desktop".to_owned(), - codex_version: "0.102.0".to_owned(), - desktop_version: "1.2026.190".to_owned(), - desktop_build: "19012345678".to_owned(), - os_type: "macOS".to_owned(), - os_version: "15.5.0".to_owned(), - arch: "arm64".to_owned(), - terminal: "xterm-256color".to_owned(), - verified_at: Utc - .with_ymd_and_hms(2026, 7, 19, 0, 0, 0) - .single() - .expect("valid test time"), - }, +struct TestOpenAiConfig { + config: OpenAiConfig, + _runtime: TempDir, +} + +fn valid_config() -> TestOpenAiConfig { + let mut config = OpenAiConfig::default(); + config.wire_profile = CodexWireProfileConfig { + originator: "Codex Desktop".to_owned(), + codex_version: "0.102.0".to_owned(), + desktop_version: "1.2026.190".to_owned(), + desktop_build: "19012345678".to_owned(), + os_type: "macOS".to_owned(), + os_version: "15.5.0".to_owned(), + arch: "arm64".to_owned(), + terminal: "xterm-256color".to_owned(), + verified_at: Utc + .with_ymd_and_hms(2026, 7, 19, 0, 0, 0) + .single() + .expect("valid test time"), + }; + let runtime = tempfile::tempdir().expect("test runtime directory"); + config + .resolve_and_validate(&runtime.path().join("deploy")) + .expect("valid OpenAI test configuration"); + TestOpenAiConfig { + config, + _runtime: runtime, } } diff --git a/backend/crates/providers/openai/tests/config.rs b/backend/crates/providers/openai/tests/config.rs index 5f82fbca6..f52c6b3e9 100644 --- a/backend/crates/providers/openai/tests/config.rs +++ b/backend/crates/providers/openai/tests/config.rs @@ -45,21 +45,65 @@ fn openai_config_rejects_noncanonical_versions_and_empty_fields() { ); } +#[test] +fn openai_config_keeps_the_quota_exhaustion_scheduling_switch() { + let mut config = valid_config(); + config.quota.skip_exhausted = false; + + assert!(!config.quota_skip_exhausted()); +} + +#[test] +fn openai_config_defaults_to_the_provider_owned_operating_values() { + let config = OpenAiConfig::default(); + + assert_eq!( + ( + config.api.base_url.as_str(), + config.ws_pool.enabled, + config.ws_pool.max_age_ms, + config.ws_pool.max_per_account, + config.ws_pool.max_total, + config.ws_pool.max_connecting, + config.ws_pool.initial_event_timeout_ms, + config.quota.refresh_interval_minutes, + config.quota.skip_exhausted, + config.auth.refresh_enabled, + config.auth.oauth_client_id.as_str(), + config.auth.oauth_token_endpoint.as_str(), + ), + ( + "https://chatgpt.com/backend-api", + true, + 3_300_000, + 8, + 64, + 8, + 800, + 5, + true, + true, + "app_EMoamEEZ73f0CkXaXp7hrann", + "https://auth.openai.com/oauth/token", + ) + ); +} + fn valid_config() -> OpenAiConfig { - OpenAiConfig { - wire_profile: CodexWireProfileConfig { - originator: "Codex Desktop".to_owned(), - codex_version: "0.102.0".to_owned(), - desktop_version: "1.2026.190".to_owned(), - desktop_build: "19012345678".to_owned(), - os_type: "macOS".to_owned(), - os_version: "15.5.0".to_owned(), - arch: "arm64".to_owned(), - terminal: "xterm-256color".to_owned(), - verified_at: Utc - .with_ymd_and_hms(2026, 7, 19, 0, 0, 0) - .single() - .expect("valid test time"), - }, - } + let mut config = OpenAiConfig::default(); + config.wire_profile = CodexWireProfileConfig { + originator: "Codex Desktop".to_owned(), + codex_version: "0.102.0".to_owned(), + desktop_version: "1.2026.190".to_owned(), + desktop_build: "19012345678".to_owned(), + os_type: "macOS".to_owned(), + os_version: "15.5.0".to_owned(), + arch: "arm64".to_owned(), + terminal: "xterm-256color".to_owned(), + verified_at: Utc + .with_ymd_and_hms(2026, 7, 19, 0, 0, 0) + .single() + .expect("valid test time"), + }; + config } diff --git a/backend/crates/providers/openai/tests/credential/admin.rs b/backend/crates/providers/openai/tests/credential/admin.rs index 575398579..b7307c4e6 100644 --- a/backend/crates/providers/openai/tests/credential/admin.rs +++ b/backend/crates/providers/openai/tests/credential/admin.rs @@ -12,6 +12,7 @@ use gateway_core::provider_ports::{ ProviderLeaseAcquisition, ProviderLeasePort, ProviderLeaseRequest, ProviderRefreshLeaseRequest, ProviderStoreError, }; +use provider_openai::OFFICIAL_CODEX_BASE_URL; use provider_openai::credential::token_client::{RefreshFailure, TokenPair, TokenRefresher}; use provider_openai::credential::{ CodexAccountIdentityService, CodexAccountIdentityVerifier, CodexCredentialAdmin, @@ -741,7 +742,7 @@ async fn real_cpr_fixture_import_contract() { .expect("fixture time"), }); let accounts = Arc::new( - ReqwestCodexAuthenticatedAccountSource::new(profile) + ReqwestCodexAuthenticatedAccountSource::new(profile, OFFICIAL_CODEX_BASE_URL.to_owned()) .expect("official authenticated account source"), ); let verifier = Arc::new(CodexAccountIdentityService::new(signed, accounts)); diff --git a/backend/crates/providers/openai/tests/credential/catalog.rs b/backend/crates/providers/openai/tests/credential/catalog.rs index c8e596d18..15aec4d03 100644 --- a/backend/crates/providers/openai/tests/credential/catalog.rs +++ b/backend/crates/providers/openai/tests/credential/catalog.rs @@ -3,6 +3,7 @@ use std::sync::Arc; use chrono::{TimeZone as _, Utc}; use futures::executor::block_on; use gateway_core::engine::credential::ProviderAccountId; +use provider_openai::OFFICIAL_CODEX_BASE_URL; use provider_openai::credential::{CodexCredentialCatalogError, CodexCredentialCatalogService}; use provider_openai::transport::profile::{CodexWireProfile, CodexWireProfileState}; @@ -33,6 +34,7 @@ fn service(store: &Arc) -> CodexCredentialCatalogService { .no_proxy() .build() .expect("client"), + OFFICIAL_CODEX_BASE_URL.to_owned(), agent_identity_service(store), ) } diff --git a/backend/crates/providers/openai/tests/credential/contract.rs b/backend/crates/providers/openai/tests/credential/contract.rs index 199daf58f..2d3127562 100644 --- a/backend/crates/providers/openai/tests/credential/contract.rs +++ b/backend/crates/providers/openai/tests/credential/contract.rs @@ -15,6 +15,7 @@ use gateway_core::engine::{ use gateway_core::policy::ClientApiKeyId; use gateway_core::provider_ports::{ProviderSessionAffinityKey, ProviderSessionAffinityPort}; use gateway_core::routing::ProviderKind; +use provider_openai::OFFICIAL_CODEX_BASE_URL; use provider_openai::credential::{ CodexAccountFailure, CodexCookiePolicy, CodexCredentialCatalogService, CodexCredentialCodec, CodexCredentialQuotaService, CodexCredentialSelector, CredentialSelectionError, @@ -25,8 +26,8 @@ use secrecy::ExposeSecret; use url::Url; use crate::support::{ - MemoryAccountStore, MemorySessionAffinity, TestLeaseCoordinator, account_policy, - agent_identity_service, profile, secret, + MemoryAccountStore, MemorySessionAffinity, MemorySessionExclusions, TestLeaseCoordinator, + account_policy, agent_identity_service, profile, secret, }; fn create_account(store: &Arc, id: &str, token: &str) { @@ -98,6 +99,7 @@ fn selector_with_affinity( leases, session_affinity, Arc::new(AccountFeedbackStats::default()), + true, ) } @@ -106,6 +108,7 @@ fn selector_with_runtime( leases: Arc, session_affinity: Arc, account_feedback: Arc, + skip_exhausted: bool, ) -> CodexCredentialSelector { let profile = CodexWireProfileState::new(CodexWireProfile { originator: "codex_cli_rs".to_owned(), @@ -124,12 +127,14 @@ fn selector_with_runtime( store.repository(), profile.clone(), http.clone(), + OFFICIAL_CODEX_BASE_URL.to_owned(), Arc::clone(&agent_identity), )); let quota = Arc::new(CodexCredentialQuotaService::new( store.repository(), profile, http, + OFFICIAL_CODEX_BASE_URL.to_owned(), Arc::clone(&agent_identity), )); CodexCredentialSelector::new( @@ -137,11 +142,13 @@ fn selector_with_runtime( store.repository(), leases, session_affinity, + Arc::new(MemorySessionExclusions::default()), catalog, quota, agent_identity, account_feedback, CodexCookiePolicy::official().expect("official cookie policy"), + skip_exhausted, ) } @@ -375,7 +382,7 @@ async fn selector_should_reuse_the_account_bound_to_the_same_session() { } #[tokio::test] -async fn selector_should_escape_a_busy_affinity_account_without_overwriting_the_binding() { +async fn selector_should_replace_a_busy_affinity_binding_after_the_fallback_succeeds() { let store = Arc::new(MemoryAccountStore::default()); create_account(&store, "acct_first", "at-first"); create_account(&store, "acct_second", "at-second"); @@ -416,8 +423,8 @@ async fn selector_should_escape_a_busy_affinity_account_without_overwriting_the_ affinity .load(&provider, &key) .await - .expect("load preserved affinity"), - Some(bound) + .expect("load replaced affinity"), + Some(ProviderAccountId::new("acct_second").expect("second account")) ); } @@ -440,6 +447,7 @@ async fn selector_should_escape_an_unhealthy_affinity_account() { Arc::new(TestLeaseCoordinator::default()), affinity, Arc::clone(&account_feedback), + true, ); for _ in 0..4 { account_feedback.report( @@ -628,6 +636,44 @@ fn rate_limited_failure_marks_account_quota_exhausted_with_a_cooldown() { ); } +#[test] +fn selector_keeps_a_quota_exhausted_account_eligible_when_configured_not_to_skip_it() { + let store = Arc::new(MemoryAccountStore::default()); + create_account(&store, "acct_primary", "at-primary"); + let strict_selector = selector(&store, Arc::new(TestLeaseCoordinator::default())); + let account = store.account("acct_primary").expect("account"); + block_on(strict_selector.record_failure(&account, CodexAccountFailure::QuotaExhausted)) + .expect("mark account exhausted"); + + let selector = selector_with_runtime( + &store, + Arc::new(TestLeaseCoordinator::default()), + Arc::new(MemorySessionAffinity::default()), + Arc::new(AccountFeedbackStats::default()), + false, + ); + let request_url = + Url::parse("https://chatgpt.com/backend-api/codex/responses").expect("request URL"); + let request_attempt = attempt(BTreeSet::new()); + + let lease = block_on(selector.select(&SelectCodexCredential { + upstream_model: "gpt-5.4", + request_url: &request_url, + attempt: &request_attempt, + session_affinity_key: None, + })) + .expect("select quota exhausted account when skipping is disabled"); + + assert_eq!(lease.account_id().as_str(), "acct_primary"); + assert_eq!( + store + .account("acct_primary") + .expect("persisted account") + .availability(), + AccountAvailability::QuotaExhausted + ); +} + #[test] fn identity_verification_failure_isolates_only_selected_account() { let store = Arc::new(MemoryAccountStore::default()); diff --git a/backend/crates/providers/openai/tests/credential/quota.rs b/backend/crates/providers/openai/tests/credential/quota.rs index ba1293984..944440fec 100644 --- a/backend/crates/providers/openai/tests/credential/quota.rs +++ b/backend/crates/providers/openai/tests/credential/quota.rs @@ -7,6 +7,7 @@ use gateway_core::engine::credential::{ AccountAvailability, AccountStateChange, OpaqueProviderData, ProviderAccountStore as _, QuotaObservation, QuotaWriteOutcome, }; +use provider_openai::OFFICIAL_CODEX_BASE_URL; use provider_openai::credential::{ CodexCredentialQuotaService, CodexQuotaSyncSummary, CodexQuotaWindowKind, ImportCodexOAuthCredential, parse_codex_quota_usage, @@ -153,6 +154,7 @@ fn quota_service_with_http( store.repository(), wire_profile(), http, + OFFICIAL_CODEX_BASE_URL.to_owned(), crate::support::agent_identity_service(store), ) } diff --git a/backend/crates/providers/openai/tests/provider/contract.rs b/backend/crates/providers/openai/tests/provider/contract.rs index 9d47ba5d6..685b5fb69 100644 --- a/backend/crates/providers/openai/tests/provider/contract.rs +++ b/backend/crates/providers/openai/tests/provider/contract.rs @@ -20,18 +20,18 @@ use gateway_core::routing::{ ConfigRevision, ModelCapabilities, ProviderKind, ProviderModel, PublicModelId, RoutingContext, RuntimeSnapshot, UpstreamModelId, }; -use provider_openai::CodexProvider; use provider_openai::credential::{ CodexCookiePolicy, CodexCredentialCatalogService, CodexCredentialQuotaService, CodexCredentialSelector, ImportCodexOAuthCredential, }; use provider_openai::transport::CodexWebSocketPool; use provider_openai::transport::profile::{CodexWireProfile, CodexWireProfileState}; +use provider_openai::{CodexProvider, OFFICIAL_CODEX_BASE_URL}; use serde_json::{Map, json}; use crate::support::{ - MemoryAccountStore, MemorySessionAffinity, TestLeaseCoordinator, account_policy, - agent_identity_service_with_pool, profile, secret, + MemoryAccountStore, MemorySessionAffinity, MemorySessionExclusions, TestLeaseCoordinator, + account_policy, agent_identity_service_with_pool, profile, secret, }; fn wire_profile() -> CodexWireProfileState { @@ -67,12 +67,14 @@ fn provider_with_affinity( store.repository(), profile.clone(), http.clone(), + OFFICIAL_CODEX_BASE_URL.to_owned(), Arc::clone(&agent_identity), )); let quota = Arc::new(CodexCredentialQuotaService::new( store.repository(), profile.clone(), http.clone(), + OFFICIAL_CODEX_BASE_URL.to_owned(), Arc::clone(&agent_identity), )); let account_feedback = Arc::new(AccountFeedbackStats::default()); @@ -81,11 +83,13 @@ fn provider_with_affinity( store.repository(), Arc::new(TestLeaseCoordinator::default()), session_affinity, + Arc::new(MemorySessionExclusions::default()), Arc::clone(&catalog), Arc::clone("a), Arc::clone(&agent_identity), Arc::clone(&account_feedback), CodexCookiePolicy::official().expect("cookie policy"), + true, )); CodexProvider::new( @@ -96,6 +100,7 @@ fn provider_with_affinity( account_feedback, http, profile, + OFFICIAL_CODEX_BASE_URL.to_owned(), websocket_pool, ) .expect("official OpenAI provider") diff --git a/backend/crates/providers/openai/tests/support.rs b/backend/crates/providers/openai/tests/support.rs index d299b371b..4ad52cdc3 100644 --- a/backend/crates/providers/openai/tests/support.rs +++ b/backend/crates/providers/openai/tests/support.rs @@ -18,7 +18,8 @@ use gateway_core::error::{StoreError, StoreErrorKind}; use gateway_core::provider_ports::{ ProviderLeaseAcquisition, ProviderLeasePort, ProviderLeaseRequest, ProviderRefreshPolicy, ProviderRuntimePolicyPort, ProviderSchedulingLeaseRequest, ProviderSchedulingState, - ProviderSessionAffinityKey, ProviderSessionAffinityPort, ProviderStoreError, + ProviderSessionAffinityKey, ProviderSessionAffinityPort, ProviderSessionExclusionPort, + ProviderSessionExclusions, ProviderStoreError, }; use gateway_core::routing::ProviderKind; use provider_openai::credential::{ @@ -501,6 +502,83 @@ impl ProviderSessionAffinityPort for MemorySessionAffinity { } } +#[derive(Default)] +pub(crate) struct MemorySessionExclusions { + states: Mutex>, + revision: AtomicUsize, +} + +impl ProviderSessionExclusionPort for MemorySessionExclusions { + fn load<'a>( + &'a self, + provider_kind: &'a ProviderKind, + key: &'a ProviderSessionAffinityKey, + ) -> BoxFuture<'a, Result, ProviderStoreError>> { + Box::pin(async move { + Ok(self + .states + .lock() + .expect("session exclusion lock") + .get(&( + provider_kind.as_str().to_owned(), + key.expose_to_store().to_owned(), + )) + .cloned()) + }) + } + + fn record_failure<'a>( + &'a self, + provider_kind: &'a ProviderKind, + key: &'a ProviderSessionAffinityKey, + account_id: &'a ProviderAccountId, + _ttl: Duration, + ) -> BoxFuture<'a, Result> { + Box::pin(async move { + let key = ( + provider_kind.as_str().to_owned(), + key.expose_to_store().to_owned(), + ); + let mut states = self.states.lock().expect("session exclusion lock"); + let mut excluded_accounts = states + .get(&key) + .map(|state| state.excluded_accounts().clone()) + .unwrap_or_default(); + excluded_accounts.insert(account_id.clone()); + let state = ProviderSessionExclusions::new( + excluded_accounts, + self.revision.fetch_add(1, Ordering::SeqCst).to_string(), + ); + states.insert(key, state.clone()); + Ok(state) + }) + } + + fn clear<'a>( + &'a self, + provider_kind: &'a ProviderKind, + key: &'a ProviderSessionAffinityKey, + expected_revision: &'a str, + ) -> BoxFuture<'a, Result> { + Box::pin(async move { + let key = ( + provider_kind.as_str().to_owned(), + key.expose_to_store().to_owned(), + ); + let mut states = self.states.lock().expect("session exclusion lock"); + if states + .get(&key) + .is_some_and(|state| state.revision() == expected_revision) + { + states.remove(&key); + Ok(true) + } else { + Ok(false) + } + }) + } +} + pub(crate) fn profile(account_id: &str) -> CodexAccountProfile { let now = chrono::Utc::now(); CodexAccountProfile { diff --git a/backend/crates/providers/openai/tests/transport/profile.rs b/backend/crates/providers/openai/tests/transport/profile.rs index 7cbcb871e..b304475f3 100644 --- a/backend/crates/providers/openai/tests/transport/profile.rs +++ b/backend/crates/providers/openai/tests/transport/profile.rs @@ -4,8 +4,9 @@ use std::sync::{Arc, Mutex}; use chrono::{TimeZone, Utc}; use futures::future::BoxFuture; use provider_openai::transport::profile::{ - CodexDesktopRelease, CodexDesktopReleaseError, CodexDesktopReleaseService, - CodexDesktopReleaseTransport, CodexWireProfile, CodexWireProfileState, parse_desktop_release, + CodexCliReleaseError, CodexCliReleaseService, CodexCliReleaseTransport, CodexDesktopRelease, + CodexDesktopReleaseError, CodexDesktopReleaseService, CodexDesktopReleaseTransport, + CodexWireProfile, CodexWireProfileState, parse_desktop_release, }; struct ReleaseTransport { @@ -34,6 +35,30 @@ impl CodexDesktopReleaseTransport for ReleaseTransport { } } +struct CliReleaseTransport { + outcomes: Mutex>>, +} + +impl CliReleaseTransport { + fn new(outcomes: impl IntoIterator>) -> Self { + Self { + outcomes: Mutex::new(outcomes.into_iter().collect()), + } + } +} + +impl CodexCliReleaseTransport for CliReleaseTransport { + fn fetch(&self) -> BoxFuture<'_, Result> { + Box::pin(async move { + self.outcomes + .lock() + .expect("CLI release outcomes") + .pop_front() + .expect("CLI release outcome") + }) + } +} + #[test] fn wire_profile_should_generate_codex_core_user_agent() { let profile = CodexWireProfile { @@ -136,6 +161,29 @@ async fn wire_profile_release_update_should_change_only_desktop_identity() { assert!(status.last_error.is_none()); } +#[tokio::test] +async fn cli_release_update_should_change_only_codex_core_identity() { + let original = wire_profile(); + let state = CodexWireProfileState::new(original.clone()); + let service = CodexCliReleaseService::new( + state.clone(), + Arc::new(CliReleaseTransport::new([Ok("0.145.0".to_owned())])), + ); + + service.refresh().await.expect("CLI release refresh"); + + let updated = state.snapshot(); + assert_eq!(updated.codex_version, "0.145.0"); + assert_eq!(updated.desktop_version, original.desktop_version); + assert_eq!(updated.desktop_build, original.desktop_build); + assert_eq!(updated.originator, original.originator); + assert_eq!(updated.os_type, original.os_type); + assert_eq!(updated.os_version, original.os_version); + assert_eq!(updated.arch, original.arch); + assert_eq!(updated.terminal, original.terminal); + assert_eq!(updated.verified_at, original.verified_at); +} + #[tokio::test] async fn failed_release_refresh_should_preserve_the_last_successful_profile_and_release() { let state = CodexWireProfileState::new(wire_profile()); diff --git a/backend/crates/providers/xai/tests/admin.rs b/backend/crates/providers/xai/tests/admin.rs index 39e3c50cd..fef00fd98 100644 --- a/backend/crates/providers/xai/tests/admin.rs +++ b/backend/crates/providers/xai/tests/admin.rs @@ -50,7 +50,8 @@ use serde_json::{Map, Value, json}; use sha2::{Digest as _, Sha256}; use crate::support::{ - MemoryProviderAccountStore, TestSessionAffinity, create_input, seed_input, xai_config, + MemoryProviderAccountStore, TestSessionAffinity, TestSessionExclusions, create_input, + seed_input, xai_config, }; #[tokio::test] @@ -582,6 +583,7 @@ fn provider_ports_with( accounts, Arc::new(TestLeases), Arc::new(TestSessionAffinity), + Arc::new(TestSessionExclusions), Arc::new(TestCatalogCache::default()), Arc::new(TestCredentialState), Arc::new(TestCooldown), diff --git a/backend/crates/providers/xai/tests/support.rs b/backend/crates/providers/xai/tests/support.rs index d0b4da251..9afd9cd01 100644 --- a/backend/crates/providers/xai/tests/support.rs +++ b/backend/crates/providers/xai/tests/support.rs @@ -1,4 +1,4 @@ -use std::collections::BTreeMap; +use std::collections::{BTreeMap, BTreeSet}; use std::net::IpAddr; use std::num::NonZeroU32; use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering}; @@ -16,7 +16,8 @@ use gateway_core::engine::credential::{ use gateway_core::error::{StoreError, StoreErrorKind}; use gateway_core::provider_ports::{ ProviderRefreshPolicy, ProviderRuntimePolicyPort, ProviderSessionAffinityKey, - ProviderSessionAffinityPort, ProviderStoreError, + ProviderSessionAffinityPort, ProviderSessionExclusionPort, ProviderSessionExclusions, + ProviderStoreError, }; use gateway_core::routing::ProviderKind; use provider_xai::{ @@ -551,6 +552,8 @@ pub struct StaticRuntimePolicy; pub struct TestSessionAffinity; +pub struct TestSessionExclusions; + impl ProviderSessionAffinityPort for TestSessionAffinity { fn load<'a>( &'a self, @@ -579,6 +582,41 @@ impl ProviderSessionAffinityPort for TestSessionAffinity { } } +impl ProviderSessionExclusionPort for TestSessionExclusions { + fn load<'a>( + &'a self, + _provider_kind: &'a ProviderKind, + _key: &'a ProviderSessionAffinityKey, + ) -> futures::future::BoxFuture<'a, Result, ProviderStoreError>> + { + Box::pin(async { Ok(None) }) + } + + fn record_failure<'a>( + &'a self, + _provider_kind: &'a ProviderKind, + _key: &'a ProviderSessionAffinityKey, + _account_id: &'a ProviderAccountId, + _ttl: Duration, + ) -> futures::future::BoxFuture<'a, Result> { + Box::pin(async { + Ok(ProviderSessionExclusions::new( + BTreeSet::new(), + String::new(), + )) + }) + } + + fn clear<'a>( + &'a self, + _provider_kind: &'a ProviderKind, + _key: &'a ProviderSessionAffinityKey, + _expected_revision: &'a str, + ) -> futures::future::BoxFuture<'a, Result> { + Box::pin(async { Ok(false) }) + } +} + impl ProviderRuntimePolicyPort for StaticRuntimePolicy { fn load_refresh_policy( &self,