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 <noreply@openai.com>
This commit is contained in:
zyycn
2026-07-25 00:33:12 +08:00
co-authored by Codex
parent bcf57351d6
commit 98250268b0
48 changed files with 2260 additions and 565 deletions
@@ -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<ApiState>, 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::<Vec<_>>();
(StatusCode::OK, Json(models)).into_response()
}
/// `GET /v1/models/{model_id}/info`。
pub(crate) async fn model_info(
State(state): State<ApiState>,
headers: HeaderMap,
Path(model_id): Path<String>,
) -> 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",
})
}
@@ -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<ApiState> {
.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))
@@ -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) => {
@@ -253,9 +253,15 @@ struct UnusedContinuation;
impl NativeContinuationPort for UnusedContinuation {
fn resolve<'a>(
&'a self,
_: &'a ClientApiKeyId,
_: &'a PreviousResponseId,
) -> BoxFuture<'a, Result<Option<NativeContinuationPin>, 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") })
}
}
@@ -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())
@@ -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<ProviderSessionState>,
}
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(|_| "<redacted>"),
)
.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<Option<NativeContinuationPin>, NativeContinuationStoreError>>;
fn record<'a>(
&'a self,
pin: NativeContinuationPin,
) -> BoxFuture<'a, Result<(), NativeContinuationStoreError>>;
}
@@ -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<NativeContinuationPin> {
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();
@@ -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<StartedExecution, GatewayError> {
async fn start_inner(
&self,
mut request: StartExecution,
) -> Result<StartedExecution, GatewayError> {
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<ResponseExecutionSession<dyn ExecutionStore>>,
admission: Option<AdmissionLease>,
circuits: Arc<dyn ProviderCircuitPort>,
continuation: Arc<dyn NativeContinuationPort>,
observed_provider_outcomes: usize,
continuation_recorded: bool,
}
impl DefaultExecutionSession {
@@ -600,12 +620,15 @@ impl DefaultExecutionSession {
core: ResponseExecutionSession<dyn ExecutionStore>,
admission: AdmissionLease,
circuits: Arc<dyn ProviderCircuitPort>,
continuation: Arc<dyn NativeContinuationPort>,
) -> 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<Option<CoordinatedEvent>, 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<Vec<ProviderEvent>, 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 => {
@@ -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<ProviderEvent> {
self.events
}
/// 返回本批次中 Provider 交付的私有会话更新。
#[must_use]
pub fn session_update(&self) -> Option<&ProviderSessionState> {
self.events.iter().find_map(ProviderEvent::session_update)
}
}
#[derive(Debug, Error)]
@@ -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<String, Value>,
@@ -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<bool, ProviderStoreError>>;
}
/// Provider 会话内已失败账号的可丢失排除集。
///
/// Provider 自行派生会话键并决定何时写入或清理;Core 只承载调度所需的账号 ID
/// 与 compare-and-swap revision,不解释任一 Provider 协议字段。
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ProviderSessionExclusions {
excluded_accounts: BTreeSet<ProviderAccountId>,
revision: String,
}
impl ProviderSessionExclusions {
#[must_use]
pub const fn new(excluded_accounts: BTreeSet<ProviderAccountId>, revision: String) -> Self {
Self {
excluded_accounts,
revision,
}
}
#[must_use]
pub const fn excluded_accounts(&self) -> &BTreeSet<ProviderAccountId> {
&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<Option<ProviderSessionExclusions>, ProviderStoreError>>;
fn record_failure<'a>(
&'a self,
provider_kind: &'a ProviderKind,
key: &'a ProviderSessionAffinityKey,
account_id: &'a ProviderAccountId,
ttl: Duration,
) -> BoxFuture<'a, Result<ProviderSessionExclusions, ProviderStoreError>>;
fn clear<'a>(
&'a self,
provider_kind: &'a ProviderKind,
key: &'a ProviderSessionAffinityKey,
expected_revision: &'a str,
) -> BoxFuture<'a, Result<bool, ProviderStoreError>>;
}
/// 所有 Provider 共享的 OAuth refresh 并发容量。
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct ProviderRefreshCapacityRequest {
@@ -704,6 +760,7 @@ pub struct ProviderStorePorts {
accounts: Arc<dyn ProviderAccountStore>,
leases: Arc<dyn ProviderLeasePort>,
session_affinity: Arc<dyn ProviderSessionAffinityPort>,
session_exclusions: Arc<dyn ProviderSessionExclusionPort>,
account_feedback: Arc<AccountFeedbackStats>,
catalog_cache: Arc<dyn ProviderCatalogCachePort>,
credential_state: Arc<dyn ProviderCredentialStatePort>,
@@ -720,6 +777,7 @@ impl ProviderStorePorts {
accounts: Arc<dyn ProviderAccountStore>,
leases: Arc<dyn ProviderLeasePort>,
session_affinity: Arc<dyn ProviderSessionAffinityPort>,
session_exclusions: Arc<dyn ProviderSessionExclusionPort>,
catalog_cache: Arc<dyn ProviderCatalogCachePort>,
credential_state: Arc<dyn ProviderCredentialStatePort>,
cooldowns: Arc<dyn ProviderCooldownPort>,
@@ -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<dyn ProviderSessionExclusionPort> {
Arc::clone(&self.session_exclusions)
}
#[must_use]
pub fn account_feedback(&self) -> Arc<AccountFeedbackStats> {
Arc::clone(&self.account_feedback)
@@ -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(&current)
.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<Item = &ClientPolicy> {
@@ -238,11 +238,17 @@ struct UnusedContinuation;
impl NativeContinuationPort for UnusedContinuation {
fn resolve<'a>(
&'a self,
_: &'a ClientApiKeyId,
_: &'a PreviousResponseId,
) -> BoxFuture<'a, Result<Option<NativeContinuationPin>, NativeContinuationStoreError>> {
Box::pin(async { Ok(None) })
}
fn record<'a>(
&'a self,
_: NativeContinuationPin,
) -> BoxFuture<'a, Result<(), NativeContinuationStoreError>> {
Box::pin(async { Ok(()) })
}
}
fn probe_snapshot() -> RuntimeSnapshot {
@@ -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"),
));
+8 -1
View File
@@ -304,7 +304,10 @@ pub async fn initialize(mut config: StoreConfig) -> StoreResult<StoreBundle> {
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<StoreBundle> {
account_store,
provider_leases,
provider_session_affinity,
Arc::new(redis::RedisProviderSessionExclusionRepository::new(
redis_connection.clone(),
REDIS_NAMESPACE,
)?),
credential_state.clone(),
credential_state,
cooldowns,
@@ -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<Option<ModelRequestHistoryRecord>>;
async fn find_model_request_by_client_response_id(
&self,
client_response_id: &str,
caller_client_api_key_ref: &str,
) -> StoreResult<Option<ModelRequestHistoryRecord>>;
async fn resolve_native_continuation_pin(
&self,
client_response_id: &str,
caller_client_api_key_ref: &str,
) -> StoreResult<Option<NativeContinuationPin>>;
}
#[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<Option<ModelRequestHistoryRecord>> {
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<Option<NativeContinuationPin>> {
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<Option<NativeContinuationPin>, 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<Option<ModelRequestHistoryRecord>> {
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<NativeContinuationPin> {
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(),
}
}
@@ -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};
@@ -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<ProviderSessionState>,
}
impl RedisNativeContinuationRepository {
pub fn new(connection: ConnectionManager, key_namespace: &str) -> StoreResult<Self> {
Ok(Self {
connection,
namespace: format!("{}:response-affinity:v1", namespace(key_namespace)?),
})
}
fn entry_fingerprint(&self, response_id: &str) -> Result<String, NativeContinuationStoreError> {
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<Option<NativeContinuationPin>, 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::<Option<String>>(&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::<i64>(&mut connection)
.await
.map_err(|_| NativeContinuationStoreError)?;
Ok(())
})
}
}
fn encode_pin(pin: &NativeContinuationPin) -> Result<String, NativeContinuationStoreError> {
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<NativeContinuationPin, NativeContinuationStoreError> {
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<NativeContinuationScope, NativeContinuationStoreError> {
match scope {
"persisted" => Ok(NativeContinuationScope::Persisted),
"connection_local" => Ok(NativeContinuationScope::ConnectionLocal),
_ => Err(NativeContinuationStoreError),
}
}
@@ -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<String>,
#[serde(default)]
revision: String,
}
impl RedisProviderSessionExclusionRepository {
pub fn new(connection: ConnectionManager, key_namespace: &str) -> StoreResult<Self> {
Ok(Self {
connection,
namespace: namespace(key_namespace)?,
})
}
fn key(
&self,
provider_kind: &ProviderKind,
affinity_key: &ProviderSessionAffinityKey,
) -> Result<String, ProviderStoreError> {
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<Option<ProviderSessionExclusions>, ProviderStoreError>> {
Box::pin(async move {
let mut connection = self.connection.clone();
let payload = redis::cmd("GET")
.arg(self.key(provider_kind, key)?)
.query_async::<Option<String>>(&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<ProviderSessionExclusions, ProviderStoreError>> {
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::<String>(&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<bool, ProviderStoreError>> {
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::<u64>(&mut connection)
.await
.map_err(|_| provider_unavailable("clear provider session exclusion"))?;
Ok(removed > 0)
})
}
}
fn decode_state(payload: impl AsRef<str>) -> Result<ProviderSessionExclusions, ProviderStoreError> {
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::<Result<BTreeSet<_>, _>>()?;
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)
}
@@ -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<String> = 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(())
}
@@ -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;
@@ -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::<i64>(&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<String> {
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::<i64>(connection)
.await
.expect("delete isolated native continuation keys");
}
}
@@ -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::<i64>(&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<String> {
redis::cmd("KEYS")
.arg(format!("{namespace}:*"))
.query_async(connection)
.await
.expect("list isolated session exclusion keys")
}
+252 -4
View File
@@ -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, CodexSessionIdentityError> {
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<Utc>,
}
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 [
@@ -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<ProviderSessionAffinityKey> {
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<ProviderSessionAffinityKey> {
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())
}
@@ -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<CodexAgentIdentityTaskService>,
cache: Arc<RwLock<CatalogCacheState>>,
etags: Arc<Mutex<CatalogEtagState>>,
@@ -135,12 +135,14 @@ impl CodexCredentialCatalogService {
repository: CodexCredentialRepository,
profile: CodexWireProfileState,
http: reqwest::Client,
base_url: String,
agent_identity: Arc<CodexAgentIdentityTaskService>,
) -> 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::<String, CodexCatalogModel>::new();
@@ -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<dyn CodexJwksSource>,
cache: Mutex<JwksCache>,
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<String>) -> 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::<AccessTokenClaims>(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::<IdTokenClaims>(id_token, &decoding_key, &validation)
@@ -449,7 +458,7 @@ impl CodexSignedIdentityVerifier for CodexJwtIdentityVerifier {
) -> Result<CodexSignedIdentity, CodexIdentityVerificationError> {
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<CodexSignedIdentity, CodexIdentityVerificationError> {
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<Self, CodexIdentityVerificationError> {
pub fn new(
profile: CodexWireProfileState,
base_url: String,
) -> Result<Self, CodexIdentityVerificationError> {
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<CodexSignedIdentity, CodexIdentityVerificationError> {
fn into_signed_identity(
self,
api_audience: &str,
) -> Result<CodexSignedIdentity, CodexIdentityVerificationError> {
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
@@ -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,
@@ -333,11 +333,12 @@ pub struct CodexOAuthAdminService {
store: Arc<dyn ProviderAccountStore>,
runtime_policy: Arc<dyn ProviderRuntimePolicyPort>,
credentials: CodexCredentialAdmin,
oauth_client_id: String,
}
impl CodexOAuthAdminService {
#[must_use]
pub const fn new(
pub fn new(
pending: Arc<dyn CodexOAuthPendingStore>,
exchanger: Arc<dyn AuthorizationCodeExchanger>,
verifier: Arc<dyn CodexAccountIdentityVerifier>,
@@ -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<String>) -> 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<String, CodexOAuthAdminError> {
fn authorization_url(
pending: &CodexPendingAuthorization,
oauth_client_id: &str,
) -> Result<String, CodexOAuthAdminError> {
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<String, Code
));
url.query_pairs_mut()
.append_pair("response_type", "code")
.append_pair("client_id", OFFICIAL_CODEX_OAUTH_CLIENT_ID)
.append_pair("client_id", oauth_client_id)
.append_pair("redirect_uri", OFFICIAL_CODEX_REDIRECT_URI)
.append_pair("scope", AUTHORIZATION_SCOPE)
.append_pair("state", pending.state.expose_secret())
@@ -19,7 +19,6 @@ use thiserror::Error;
use tokio::sync::Mutex;
use uuid::Uuid;
use crate::provider::OFFICIAL_CODEX_BASE_URL;
use crate::transport::profile::CodexWireProfileState;
use crate::transport::{CodexBackendClient, CodexClientError, CodexRequestContext};
@@ -32,6 +31,27 @@ const QUOTA_SCHEDULING_TTL: Duration = Duration::from_secs(10 * 60);
const QUOTA_HYDRATION_FAILURE_TTL: Duration = Duration::from_secs(5);
const QUOTA_PERIODIC_REFRESH_MIN_INTERVAL: Duration = Duration::from_secs(30 * 60);
/// OpenAI Provider 主动额度刷新的调度策略。
///
/// 账号筛选仍由 core 的可调度状态统一处理;该策略只控制 Provider 自己的
/// quota worker 频率,避免把 OpenAI 额度语义泄漏到公共调度层。
#[derive(Debug, Clone, Copy)]
pub struct CodexQuotaRefreshPolicy {
interval: Duration,
}
impl CodexQuotaRefreshPolicy {
#[must_use]
pub const fn new(interval: Duration) -> 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<u8>,
@@ -199,6 +219,7 @@ pub struct CodexCredentialQuotaService {
store: Arc<dyn ProviderAccountStore>,
profile: CodexWireProfileState,
http: Client,
base_url: String,
agent_identity: Arc<CodexAgentIdentityTaskService>,
scheduling: CodexQuotaSchedulingProjection,
}
@@ -436,6 +457,7 @@ impl CodexCredentialQuotaService {
repository: CodexCredentialRepository,
profile: CodexWireProfileState,
http: Client,
base_url: String,
agent_identity: Arc<CodexAgentIdentityTaskService>,
) -> 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();
@@ -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<ProviderSessionExclusions>,
}
pub struct CodexCredentialSelector {
provider_kind: ProviderKind,
repository: CodexCredentialRepository,
leases: Arc<dyn ProviderLeasePort>,
session_affinity: Arc<dyn ProviderSessionAffinityPort>,
session_exclusions: Arc<dyn ProviderSessionExclusionPort>,
catalog: Arc<CodexCredentialCatalogService>,
quota: Arc<CodexCredentialQuotaService>,
agent_identity: Arc<CodexAgentIdentityTaskService>,
cookie_policy: CodexCookiePolicy,
risk_recovery: Mutex<HashMap<String, RiskRecoveryState>>,
account_feedback: Arc<AccountFeedbackStats>,
skip_exhausted: bool,
}
enum SessionAffinityLookup {
@@ -117,29 +126,41 @@ impl CodexCredentialSelector {
repository: CodexCredentialRepository,
leases: Arc<dyn ProviderLeasePort>,
session_affinity: Arc<dyn ProviderSessionAffinityPort>,
session_exclusions: Arc<dyn ProviderSessionExclusionPort>,
catalog: Arc<CodexCredentialCatalogService>,
quota: Arc<CodexCredentialQuotaService>,
agent_identity: Arc<CodexAgentIdentityTaskService>,
account_feedback: Arc<AccountFeedbackStats>,
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<CodexCredentialLease, CredentialSelectionError> {
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<CodexCredentialLease, CredentialSelectionError> {
let accounts = self.repository.list_for_provider().await?;
let accounts = accounts
@@ -189,10 +210,23 @@ impl CodexCredentialSelector {
AccountCandidate { account, signals }
})
.collect::<Vec<_>>();
// 旧版开关只影响“候选资格”。保留原始账号供 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::<Vec<_>>()
});
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<CodexCyberPolicyScope> {
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<String, RiskRecoveryState>,
account_id: &str,
@@ -651,6 +800,7 @@ pub struct CodexCredentialLease {
authentication: CodexRuntimeAuthentication,
cookies: Vec<RuntimeCodexCookie>,
installation_id: String,
cyber_policy_scope: Option<CodexCyberPolicyScope>,
_guard: Box<dyn ProviderLeaseGuard>,
}
@@ -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 {
@@ -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<OpenAiTokenClient, TokenClientBuildError> {
pub fn openai_token_client(
config: TokenClientConfig,
) -> Result<OpenAiTokenClient, TokenClientBuildError> {
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<OpenAiTokenClient, TokenClientBuildError> {
openai_token_client(TokenClientConfig {
client_id: OFFICIAL_CODEX_OAUTH_CLIENT_ID.to_owned(),
token_endpoint: OFFICIAL_CODEX_TOKEN_ENDPOINT.to_owned(),
})
}
#[derive(Deserialize)]
+45 -22
View File
@@ -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<ProviderBundle, OpenAiInitializeError> {
config
.resolve_and_validate(Path::new("."))
.map_err(OpenAiInitializeError::Config)?;
let provider_kind =
ProviderKind::new("openai").map_err(|_| OpenAiInitializeError::InvalidProviderKind)?;
let accounts: Arc<dyn ProviderAccountStore> = 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(&quota),
Arc::clone(&agent_identity),
Arc::clone(&account_feedback),
CodexCookiePolicy::official().map_err(|_| OpenAiInitializeError::CookiePolicy)?,
quota_skip_exhausted,
));
let core_provider: Arc<dyn Provider> = 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<dyn CodexSignedIdentityVerifier> =
Arc::new(CodexJwtIdentityVerifier::new(Box::new(jwks)));
let signed: Arc<dyn CodexSignedIdentityVerifier> = Arc::new(
CodexJwtIdentityVerifier::new(Box::new(jwks))
.with_oauth_client_id(config.oauth_client_id()),
);
let account_source: Arc<dyn CodexAuthenticatedAccountSource> = Arc::new(
ReqwestCodexAuthenticatedAccountSource::new(profile.clone())
ReqwestCodexAuthenticatedAccountSource::new(profile.clone(), config.base_url().to_owned())
.map_err(|_| OpenAiInitializeError::Identity)?,
);
let identity: Arc<dyn CodexAccountIdentityVerifier> =
@@ -164,14 +175,17 @@ pub async fn initialize(
ports.oauth_pending(),
provider_kind.clone(),
));
let oauth_admin: Arc<dyn CodexOAuthAdmin> = Arc::new(CodexOAuthAdminService::new(
pending,
exchanger,
Arc::clone(&identity),
Arc::clone(&accounts),
Arc::clone(&runtime_policy),
CodexCredentialAdmin,
));
let oauth_admin: Arc<dyn CodexOAuthAdmin> = 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<dyn ProviderAdmin> = 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)]
+109 -27
View File
@@ -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<AccountFeedbackStats>,
client: CodexBackendClient,
responses_url: Url,
session_identity: Option<CodexSessionIdentity>,
}
impl CodexProvider {
@@ -108,13 +111,13 @@ impl CodexProvider {
account_feedback: Arc<AccountFeedbackStats>,
http: Client,
profile: CodexWireProfileState,
base_url: String,
websocket_pool: Arc<CodexWebSocketPool>,
) -> Result<Self, CodexProviderConfigError> {
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<String>,
#[serde(default)]
turn_state: Option<String>,
continuation_scope: OpenAiContinuationScope,
transcript: Vec<OpenAiReplayItem>,
}
@@ -424,6 +453,7 @@ struct OpenAiSessionCapture {
request_input: Vec<Value>,
account_id: String,
conversation_id: Option<String>,
turn_state: Option<String>,
response_store: bool,
continuation_scope: Option<OpenAiContinuationScope>,
}
@@ -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 {
&quota,
&active_account,
&response_origin,
cyber_policy_scope.as_ref(),
&failure,
)
.await;
@@ -727,6 +760,7 @@ fn cold_response_stream(response: ColdResponse) -> EventStream {
&quota,
&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 {
&quota,
&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 {
&quota,
&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(&quota, &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 {
&quota,
&active_account,
&response_origin,
cyber_policy_scope.as_ref(),
failure,
)
.await;
@@ -876,6 +921,7 @@ fn cold_response_stream(response: ColdResponse) -> EventStream {
&quota,
&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 {
&quota,
&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 {
&quota,
&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<CodexAccountFailure>,
cyber_policy_failure: bool,
set_cookie_headers: Vec<String>,
rate_limit_headers: Vec<(String, String)>,
observation: Option<ProviderResponseObservation>,
@@ -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<CodexCredentialRefreshService>,
quota: Arc<CodexCredentialQuotaService>,
catalog: Arc<CodexCredentialCatalogService>,
quota_refresh_policy: CodexQuotaRefreshPolicy,
oauth_refresh_enabled: bool,
cli_release: Arc<CodexCliReleaseService>,
desktop_release: Arc<CodexDesktopReleaseService>,
) -> Result<Vec<WorkerContribution>, 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(
@@ -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;
@@ -372,38 +372,17 @@ fn non_empty_string(value: Option<&Value>) -> Option<String> {
.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) {
@@ -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<Self, CodexSessionIdentityError> {
match read_secret(path)? {
Some(secret) => Ok(Self { secret }),
None => Self::create(path),
}
}
fn create(path: &Path) -> Result<Self, CodexSessionIdentityError> {
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<Option<[u8; 32]>, 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()
}
+48 -25
View File
@@ -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,
}
}
+60 -16
View File
@@ -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
}
@@ -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));
@@ -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<MemoryAccountStore>) -> CodexCredentialCatalogService {
.no_proxy()
.build()
.expect("client"),
OFFICIAL_CODEX_BASE_URL.to_owned(),
agent_identity_service(store),
)
}
@@ -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<MemoryAccountStore>, 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<TestLeaseCoordinator>,
session_affinity: Arc<MemorySessionAffinity>,
account_feedback: Arc<AccountFeedbackStats>,
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());
@@ -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),
)
}
@@ -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(&quota),
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")
@@ -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<BTreeMap<(String, String), ProviderSessionExclusions>>,
revision: AtomicUsize,
}
impl ProviderSessionExclusionPort for MemorySessionExclusions {
fn load<'a>(
&'a self,
provider_kind: &'a ProviderKind,
key: &'a ProviderSessionAffinityKey,
) -> BoxFuture<'a, Result<Option<ProviderSessionExclusions>, 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<ProviderSessionExclusions, ProviderStoreError>> {
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<bool, ProviderStoreError>> {
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 {
@@ -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<VecDeque<Result<String, CodexCliReleaseError>>>,
}
impl CliReleaseTransport {
fn new(outcomes: impl IntoIterator<Item = Result<String, CodexCliReleaseError>>) -> Self {
Self {
outcomes: Mutex::new(outcomes.into_iter().collect()),
}
}
}
impl CodexCliReleaseTransport for CliReleaseTransport {
fn fetch(&self) -> BoxFuture<'_, Result<String, CodexCliReleaseError>> {
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());
+3 -1
View File
@@ -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),
+40 -2
View File
@@ -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<Option<ProviderSessionExclusions>, 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<ProviderSessionExclusions, ProviderStoreError>> {
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<bool, ProviderStoreError>> {
Box::pin(async { Ok(false) })
}
}
impl ProviderRuntimePolicyPort for StaticRuntimePolicy {
fn load_refresh_policy(
&self,