mirror of
https://github.com/zyycn/codex-proxy-rs.git
synced 2026-11-06 08:06:19 +00:00
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:
@@ -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(¤t)
|
||||
.map(|target| target.trim())
|
||||
else {
|
||||
return current;
|
||||
};
|
||||
if !seen.insert(current.clone()) || seen.contains(target) {
|
||||
return original.to_owned();
|
||||
}
|
||||
current = target.to_owned();
|
||||
}
|
||||
original.to_owned()
|
||||
}
|
||||
|
||||
pub fn client_policies(&self) -> impl Iterator<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"),
|
||||
));
|
||||
|
||||
@@ -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")
|
||||
}
|
||||
@@ -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)]
|
||||
|
||||
@@ -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("a),
|
||||
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)]
|
||||
|
||||
@@ -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 {
|
||||
"a,
|
||||
&active_account,
|
||||
&response_origin,
|
||||
cyber_policy_scope.as_ref(),
|
||||
&failure,
|
||||
)
|
||||
.await;
|
||||
@@ -727,6 +760,7 @@ fn cold_response_stream(response: ColdResponse) -> EventStream {
|
||||
"a,
|
||||
&active_account,
|
||||
&response_origin,
|
||||
cyber_policy_scope.as_ref(),
|
||||
&failure,
|
||||
)
|
||||
.await;
|
||||
@@ -743,6 +777,7 @@ fn cold_response_stream(response: ColdResponse) -> EventStream {
|
||||
} else {
|
||||
OpenAiContinuationScope::ReplayRequired
|
||||
});
|
||||
capture.turn_state = response.turn_state.clone().or(capture.turn_state.clone());
|
||||
}
|
||||
let observation = match codex_response_observation(
|
||||
response.transport,
|
||||
@@ -760,6 +795,7 @@ fn cold_response_stream(response: ColdResponse) -> EventStream {
|
||||
"a,
|
||||
&active_account,
|
||||
&response_origin,
|
||||
cyber_policy_scope.as_ref(),
|
||||
&failure,
|
||||
)
|
||||
.await;
|
||||
@@ -796,6 +832,7 @@ fn cold_response_stream(response: ColdResponse) -> EventStream {
|
||||
let failure_set_cookie_headers = response.set_cookie_headers.clone();
|
||||
let failure_rate_limit_headers = response.rate_limit_headers.clone();
|
||||
let rate_limit_updates = response.rate_limit_header_updates;
|
||||
let turn_state_updates = response.turn_state_update;
|
||||
let mut decoder = CodexCanonicalDecoder::new(upstream_model.as_str());
|
||||
if response.transport == CodexBackendTransport::HttpSse {
|
||||
decoder = decoder.with_raw_sse_passthrough();
|
||||
@@ -830,6 +867,7 @@ fn cold_response_stream(response: ColdResponse) -> EventStream {
|
||||
"a,
|
||||
&active_account,
|
||||
&response_origin,
|
||||
cyber_policy_scope.as_ref(),
|
||||
&failure,
|
||||
)
|
||||
.await;
|
||||
@@ -842,6 +880,12 @@ fn cold_response_stream(response: ColdResponse) -> EventStream {
|
||||
let updates = std::mem::take(&mut *updates.lock().await);
|
||||
synchronize_passive_quota("a, &active_account, &updates).await;
|
||||
}
|
||||
if let (Some(capture), Some(updates)) =
|
||||
(session_capture.as_mut(), turn_state_updates.as_ref())
|
||||
&& let Some(turn_state) = updates.lock().await.clone()
|
||||
{
|
||||
capture.turn_state = Some(turn_state);
|
||||
}
|
||||
let (mut events, terminal_failure) = match decoder.push(&chunk) {
|
||||
CodexCanonicalOutcome::Events(events) => (events, None),
|
||||
CodexCanonicalOutcome::Failed(failure) => {
|
||||
@@ -863,6 +907,7 @@ fn cold_response_stream(response: ColdResponse) -> EventStream {
|
||||
"a,
|
||||
&active_account,
|
||||
&response_origin,
|
||||
cyber_policy_scope.as_ref(),
|
||||
failure,
|
||||
)
|
||||
.await;
|
||||
@@ -876,6 +921,7 @@ fn cold_response_stream(response: ColdResponse) -> EventStream {
|
||||
"a,
|
||||
&active_account,
|
||||
&response_origin,
|
||||
cyber_policy_scope.as_ref(),
|
||||
&failure,
|
||||
)
|
||||
.await;
|
||||
@@ -900,6 +946,9 @@ fn cold_response_stream(response: ColdResponse) -> EventStream {
|
||||
selector
|
||||
.record_success(&active_account, session_affinity_key.as_ref())
|
||||
.await;
|
||||
selector
|
||||
.observe_cyber_policy_success(cyber_policy_scope.as_ref())
|
||||
.await;
|
||||
return;
|
||||
}
|
||||
}
|
||||
@@ -924,10 +973,17 @@ fn cold_response_stream(response: ColdResponse) -> EventStream {
|
||||
"a,
|
||||
&active_account,
|
||||
&response_origin,
|
||||
cyber_policy_scope.as_ref(),
|
||||
failure,
|
||||
)
|
||||
.await;
|
||||
}
|
||||
if let (Some(capture), Some(updates)) =
|
||||
(session_capture.as_mut(), turn_state_updates.as_ref())
|
||||
&& let Some(turn_state) = updates.lock().await.clone()
|
||||
{
|
||||
capture.turn_state = Some(turn_state);
|
||||
}
|
||||
if let Err(error) = attach_openai_session_update(&mut events, &mut session_capture) {
|
||||
if terminal_failure.is_none() {
|
||||
let failure = MappedProviderFailure::plain(error);
|
||||
@@ -937,6 +993,7 @@ fn cold_response_stream(response: ColdResponse) -> EventStream {
|
||||
"a,
|
||||
&active_account,
|
||||
&response_origin,
|
||||
cyber_policy_scope.as_ref(),
|
||||
&failure,
|
||||
)
|
||||
.await;
|
||||
@@ -961,6 +1018,9 @@ fn cold_response_stream(response: ColdResponse) -> EventStream {
|
||||
selector
|
||||
.record_success(&active_account, session_affinity_key.as_ref())
|
||||
.await;
|
||||
selector
|
||||
.observe_cyber_policy_success(cyber_policy_scope.as_ref())
|
||||
.await;
|
||||
}
|
||||
})
|
||||
}
|
||||
@@ -1254,6 +1314,7 @@ fn map_selection_error(error: CredentialSelectionError) -> ProviderError {
|
||||
struct MappedProviderFailure {
|
||||
error: ProviderError,
|
||||
account_failure: Option<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,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,
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -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("a),
|
||||
Arc::clone(&agent_identity),
|
||||
Arc::clone(&account_feedback),
|
||||
CodexCookiePolicy::official().expect("cookie policy"),
|
||||
true,
|
||||
));
|
||||
|
||||
CodexProvider::new(
|
||||
@@ -96,6 +100,7 @@ fn provider_with_affinity(
|
||||
account_feedback,
|
||||
http,
|
||||
profile,
|
||||
OFFICIAL_CODEX_BASE_URL.to_owned(),
|
||||
websocket_pool,
|
||||
)
|
||||
.expect("official OpenAI provider")
|
||||
|
||||
@@ -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());
|
||||
|
||||
@@ -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),
|
||||
|
||||
@@ -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,
|
||||
|
||||
Reference in New Issue
Block a user