Compare commits

...
15 Commits
Author SHA1 Message Date
ZyyandCodex bd69329a37 chore: bump version to 3.45.2
Co-authored-by: Codex <noreply@openai.com>
2026-11-11 17:04:39 +08:00
ZyyandCodex 151a3c6bfd fix(openai): negotiate Codex catalogs for Responses upstreams
Co-authored-by: Codex <noreply@openai.com>
2026-11-11 16:49:10 +08:00
ZyyandCodex e2f194875f refactor: clarify protocol ownership and request observations
Separate protocol identifiers from Provider identities and centralize shared
request facts and observations in gateway-protocol. Keep endpoint models
and private continuation state owned by Providers.

Reorganize the architecture documentation and clarify the shutdown guarantees
for buffered charges and observations.

Validated with workspace formatting, strict Clippy, 3013 affected-module
tests and 16 architecture checks. One reference-binary TLS test is ignored.

Co-authored-by: Codex <noreply@openai.com>
2026-11-11 16:04:16 +08:00
ZyyandCodex 07b61bdc83 fix(ui): keep help buttons aligned with form labels and clarify hints
Co-authored-by: Codex <noreply@openai.com>
2026-11-11 15:13:00 +08:00
c0fdc1bbff fix(openai): record explicit Anthropic reasoning effort (#680)
Observe top-level output_config.effort before the OpenAI encoder rejects
Messages payloads. Reuse the protocol parser without inferring an effort
from thinking budgets or changing upstream conversion behavior.

Cover absent and explicit effort values, native Messages forwarding, and
Messages-to-Responses conversion. Document the observation contract.

Co-authored-by: Zyy <zyycn@users.noreply.github.com>
Co-authored-by: Codex <noreply@openai.com>
2026-11-11 15:12:02 +08:00
343c579017 fix(openai): rotate new round-robin sessions without affinity hits (#678)
* fix(openai): rotate new round-robin sessions without affinity hits

Strict affinity keeps a session on one account, but the shared request
cursor advanced on every scheduling load, including affinity hits. A new
session then took whatever slot that request volume happened to leave.

Advance a separate session cursor only when round-robin assigns an
account, and keep the request cursor for other strategies.

* fix(openai): advance round robin when an account assignment succeeds

Commit the session binding and round-robin cursor together after acquiring the account lease. Release a conflicting lease before retrying so concurrent first requests share a binding without consuming extra turns. Count delegated selections and preserve explicit plugin picks.

Verified with OpenAI provider tests, store tests against Redis and PostgreSQL, architecture checks, and workspace Clippy.

Co-authored-by: Codex <noreply@openai.com>

---------

Co-authored-by: koeria <riennie@qq.com>
Co-authored-by: Zyy <zyycn@users.noreply.github.com>
Co-authored-by: Codex <noreply@openai.com>
2026-11-11 14:40:59 +08:00
koeriaandCursor Agent 914c00785c fix(openai): skip same-account retry on NO_MORE_RETRY (#675)
Treat x-retry-metadata: NO_MORE_RETRY as a hard signal to skip
same-account backoff and transport recovery, even when Retry-After
is within maxRetryAfterSeconds. Capacity failures still rotate under
the existing policy; server_error still returns the upstream error.

Co-authored-by: Cursor Agent <cursoragent@cursor.com>
2026-11-11 14:08:24 +08:00
ZyyandCodex b3de0d637d fix(peer): send routing headers only to discovered nodes
Reuse connection-scoped peer discovery for native HTTP and WebSocket forwarding. Strip the route header for ordinary upstreams and leave explicit plugin forwarding changes to plugins.

Co-authored-by: Codex <noreply@openai.com>
2026-11-11 13:14:25 +08:00
ZyyandCodex 1271fa0186 test(openai): verify platform TLS against official Codex
Compare native HTTP, custom CA HTTP, and WebSocket ClientHello captures with an official Codex executable on the same platform. Document the explicit reference test workflow.

Co-authored-by: Codex <noreply@openai.com>
2026-11-11 12:41:04 +08:00
ZyyandCodex 055da40476 feat(keys): enable async compaction in Codex configs
Co-authored-by: Codex <noreply@openai.com>
2026-11-11 12:37:12 +08:00
ZyyandCodex 36ba79ca5e fix(openai): preserve native Responses service tiers
Limit default tier omission to Messages-to-OAuth conversion so native Responses retain their requested tier and routing hint.

Co-authored-by: Codex <noreply@openai.com>
2026-11-11 11:33:01 +08:00
github-actions[bot] 1b24029a97 chore(release): update release channels to 3.45.1 2026-11-11 00:03:25 +00:00
zyycnandCodex 431992a7a4 chore: bump version to 3.45.1
Co-authored-by: Codex <noreply@openai.com>
2026-11-11 07:41:09 +08:00
ZyyandCodex d3143406dc fix(openai): restore account catalogs without blocking requests
Persist validated directories per credential revision. Restore them in bulk and refresh only missing, expired, or changed accounts with paced background requests. Filter large pools using current per-account evidence while preserving unknown and explicitly allowed candidates.

Co-authored-by: Codex <noreply@openai.com>
2026-11-11 03:46:25 +08:00
github-actions[bot] 478f2d1f21 chore(release): update release channels to 3.45.0 2026-11-10 18:15:24 +00:00
127 changed files with 5331 additions and 3086 deletions
@@ -52,7 +52,7 @@ impl PeerInfoService {
});
let limits = peer.and_then(|peer| peer.key_limits);
let info = ProviderUpstreamInfo {
discovered: peer.is_some() || stored_failure == Some(UpstreamInfoFailure::Unauthorized),
discovered: snapshot::is_discovered(peer, stored_failure),
observed_at,
stale,
connection_revision: connection.revision.clone(),
@@ -282,6 +282,21 @@ impl PeerInfoService {
}
impl UpstreamInfoSource for PeerInfoService {
fn is_peer(&self, account: &ProviderAccount) -> bool {
if !self.supported(account) {
return false;
}
let connection = UpstreamConnection::for_account(account);
self.observations
.read()
.unwrap_or_else(std::sync::PoisonError::into_inner)
.get(account.id())
.is_some_and(|cached| {
cached.connection == connection
&& snapshot::is_discovered(cached.summary.as_deref(), cached.stored_failure)
})
}
fn signal(
&self,
account: &ProviderAccount,
@@ -16,6 +16,10 @@ use serde::Deserialize as _;
pub(super) const MAX_AGE: Duration = Duration::from_secs(30 * 60);
pub(super) const FRESH_FOR: Duration = Duration::from_secs(5 * 60);
pub(super) fn is_discovered(peer: Option<&PeerLink>, failure: Option<UpstreamInfoFailure>) -> bool {
peer.is_some() || failure == Some(UpstreamInfoFailure::Unauthorized)
}
// 不支持协议是连接的识别结果,复用观测存储与版本保护,不伪造上游额度快照
pub(super) fn unsupported(connection: &UpstreamConnection) -> OpaqueProviderData {
OpaqueProviderData::new(serde_json::Map::from_iter([
@@ -140,6 +140,10 @@ impl PendingPeer {
}
impl UpstreamInfoSource for PendingPeer {
fn is_peer(&self, _: &ProviderAccount) -> bool {
false
}
fn signal(
&self,
_: &ProviderAccount,
@@ -239,6 +239,7 @@ async fn discovery_uses_provider_capabilities_and_publishes_neutral_signals() {
initialize_peer_upstream(accounts.clone(), &registry, transport.clone()).unwrap();
let initial_accounts = accounts.accounts.lock().unwrap().clone();
let reader = UpstreamInfoReader::default();
assert!(!reader.is_peer(&initial_accounts[0]));
assert!(
reader
.signal(&initial_accounts[0], Some("model"), SystemTime::now())
@@ -246,6 +247,7 @@ async fn discovery_uses_provider_capabilities_and_publishes_neutral_signals() {
);
let source = bundle.scheduling_source();
reader.bind(&source).unwrap();
assert!(!reader.is_peer(&initial_accounts[0]));
assert!(reader.bind(&source).is_err());
drop(source);
let WorkerContribution::Registration(worker) = bundle.take_worker_contributions().remove(0)
@@ -272,6 +274,10 @@ async fn discovery_uses_provider_capabilities_and_publishes_neutral_signals() {
}
);
for account in &initial_accounts[..2] {
assert_eq!(
reader.is_peer(account),
Some(account.provider().as_str()) != failed_provider
);
if Some(account.provider().as_str()) == failed_provider {
assert!(
reader
@@ -304,6 +310,7 @@ async fn discovery_uses_provider_capabilities_and_publishes_neutral_signals() {
);
drop(worker);
drop(bundle);
assert!(!reader.is_peer(&initial_accounts[0]));
assert!(
reader
.signal(&initial_accounts[0], Some("model"), SystemTime::now())
@@ -340,6 +347,7 @@ async fn account_edits_preserve_observation_and_refresh_budget_until_connection_
.is_some()
);
let after = service.refresh_quota(edited.id()).await.unwrap();
assert!(bundle.scheduling_source().is_peer(&edited));
assert_eq!(after.upstream, before.upstream);
assert_eq!(transport.0.lock().unwrap().len(), 1);
@@ -347,6 +355,7 @@ async fn account_edits_preserve_observation_and_refresh_budget_until_connection_
OutboundProxy::parse("http://proxy.invalid:8080").unwrap(),
));
*accounts.accounts.lock().unwrap() = vec![reconnected.clone()];
assert!(!bundle.scheduling_source().is_peer(&reconnected));
service.account_changed(reconnected.id()).await.unwrap();
assert!(
bundle
@@ -422,6 +431,21 @@ fn loaded(account: &ProviderAccount) -> LoadedCredential {
#[async_trait]
impl ProviderAccountStore for TestAccounts {
async fn load_account_catalogs(
&self,
_: &gateway_core::routing::ProviderKind,
) -> Result<Vec<gateway_core::account::StoredAccountCatalog>, gateway_core::error::StoreError>
{
Ok(Vec::new())
}
async fn save_account_catalog(
&self,
_: &gateway_core::account::AccountCatalogObservation,
_: &[u8],
) -> Result<bool, gateway_core::error::StoreError> {
unreachable!("fixture does not discover account catalogs")
}
async fn create_account(&self, _account: NewProviderAccount) -> Result<(), StoreError> {
panic!("unused account operation")
}
@@ -215,6 +215,10 @@ async fn known_peer_identity_changes_only_on_definitive_detection() {
quota.upstream
);
let account = accounts.accounts.lock().unwrap()[0].clone();
assert_eq!(
bundle.scheduling_source().is_peer(&account),
!clears_identity
);
assert_eq!(
bundle
.scheduling_source()
@@ -237,6 +241,7 @@ async fn known_peer_identity_changes_only_on_definitive_detection() {
drop(bundle);
let restarted = initialize(accounts, transport);
let restored = restarted.service().read_quota(&id).await.unwrap().unwrap();
assert!(!restarted.scheduling_source().is_peer(&account));
assert_eq!(
restored.upstream.as_ref().unwrap().discovered,
!clears_identity
@@ -258,6 +263,10 @@ async fn known_peer_identity_changes_only_on_definitive_detection() {
.upstream,
restored.upstream
);
assert_eq!(
restarted.scheduling_source().is_peer(&account),
!clears_identity
);
}
}
@@ -22,6 +22,7 @@ use gateway_core::{
error::GatewayError,
operation::Operation,
};
use gateway_protocol::openai::OPENAI_PROTOCOL;
use crate::client_error::{ProtocolError, ProtocolErrorBody};
use crate::middleware::headers::encode_headers;
@@ -65,7 +66,7 @@ where
Box::new(move |prepared, request| {
Box::pin(async move {
let (protocol, headers, body) = request_parts(request.clone())?;
if protocol != "openai" {
if protocol != OPENAI_PROTOCOL {
return Err(MiddlewareError::Rejected);
}
let operation = request.apply_capabilities(decode(body, headers).await?)?;
@@ -80,7 +81,8 @@ where
{
Ok(started) => started,
Err(error) => {
return buffered_response("openai", gateway_error_response(&error)).await;
return buffered_response(OPENAI_PROTOCOL, gateway_error_response(&error))
.await;
}
};
collect_raw_json_response(started).await
@@ -89,7 +91,9 @@ where
)
.await;
let response = match result {
Ok(response) => into_http_response(response, ExpectedBody::SingleJson, "openai").await,
Ok(response) => {
into_http_response(response, ExpectedBody::SingleJson, OPENAI_PROTOCOL).await
}
Err(error) => error_response(error),
};
if !execution_entered.load(Ordering::Relaxed) {
@@ -143,7 +147,7 @@ async fn collect_raw_json_response(
.await
.map_err(|_| MiddlewareError::Fault)?;
Ok(pending_execution_response(
"openai".to_owned(),
OPENAI_PROTOCOL.to_owned(),
parts.status.as_u16(),
encode_headers(&parts.headers),
MiddlewareFrame::new(bytes, framing, true).with_transformed(transformed),
@@ -155,7 +159,7 @@ fn raw_json_body(events: Vec<ProviderEvent>) -> Option<Bytes> {
let mut body = None;
for event in events {
let (_, wire) = event.into_parts();
let Some(wire) = wire.filter(|wire| wire.protocol() == "openai") else {
let Some(wire) = wire.filter(|wire| wire.protocol() == OPENAI_PROTOCOL) else {
continue;
};
let Some(raw) = wire.into_raw_json_body() else {
@@ -32,6 +32,7 @@ use gateway_core::engine::{
use gateway_core::error::{GatewayError, GatewayErrorKind};
use gateway_core::event::{ProviderEvent, ProviderResponseHeader};
use gateway_core::lifecycle::{CancellationToken, ConnectionGuard};
use gateway_protocol::openai::OPENAI_PROTOCOL;
use gateway_protocol::openai::sse::response_failed_sse_event_with_id;
use serde_json::Value;
use tokio::time::Instant;
@@ -491,7 +492,9 @@ pub async fn collect_execution_response(session: Box<dyn ExecutionSession>) -> R
)
.await
{
Ok(response) => into_http_response(response, ExpectedBody::SingleJson, "openai").await,
Ok(response) => {
into_http_response(response, ExpectedBody::SingleJson, OPENAI_PROTOCOL).await
}
Err(error) => error_response(error),
}
}
@@ -513,7 +516,9 @@ pub async fn stream_execution_response(
)
.await
{
Ok(response) => into_http_response(response, ExpectedBody::Sse, "openai").await,
Ok(response) => {
into_http_response(response, ExpectedBody::Sse, OPENAI_PROTOCOL).await
}
Err(error) => error_response(error),
}
},
@@ -15,6 +15,8 @@ use gateway_core::{
error::{GatewayError, GatewayErrorKind},
operation::OperationKind,
};
use gateway_protocol::anthropic::ANTHROPIC_PROTOCOL;
use gateway_protocol::openai::OPENAI_PROTOCOL;
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub(crate) enum GenerationProtocol {
@@ -27,7 +29,7 @@ impl GenerationProtocol {
pub(crate) fn from_request(request: &DecodedGenerationRequest) -> Self {
if request.operation().kind() == OperationKind::CountTokens {
Self::MessageTokenCount
} else if request.operation().protocol() == "anthropic" {
} else if request.operation().protocol() == ANTHROPIC_PROTOCOL {
Self::Messages
} else {
Self::Responses
@@ -36,8 +38,8 @@ impl GenerationProtocol {
pub(crate) const fn wire(self) -> &'static str {
match self {
Self::Responses => "openai",
Self::Messages | Self::MessageTokenCount => "anthropic",
Self::Responses => OPENAI_PROTOCOL,
Self::Messages | Self::MessageTokenCount => ANTHROPIC_PROTOCOL,
}
}
@@ -2,6 +2,7 @@
use axum::body::Bytes;
use gateway_core::{error::GatewayError, event::ProviderEvent};
use gateway_protocol::anthropic::ANTHROPIC_PROTOCOL;
use gateway_protocol::{
anthropic::MessagesAccumulator,
openai::sse::{
@@ -50,7 +51,7 @@ impl GenerationEncoder {
} => {
let Some(wire) = event
.wire_event()
.filter(|wire| wire.protocol() == "anthropic")
.filter(|wire| wire.protocol() == ANTHROPIC_PROTOCOL)
else {
return Vec::new();
};
@@ -1,6 +1,7 @@
//! 出口校验按协议维护边界事实,不在流式路径保存生成正文
use gateway_core::event::ProviderEvent;
use gateway_protocol::anthropic::ANTHROPIC_PROTOCOL;
use gateway_protocol::openai::sse::parse_sse_events;
use serde_json::Value;
use std::sync::{Arc, Mutex};
@@ -32,7 +33,7 @@ impl GenerationValidationFacts {
Self::Messages(id) => {
if let Some(value) = event
.wire_event()
.filter(|wire| wire.protocol() == "anthropic")
.filter(|wire| wire.protocol() == ANTHROPIC_PROTOCOL)
.map(|wire| wire.data())
&& let Some(value) = value
.pointer("/message/id")
@@ -13,6 +13,7 @@ use gateway_core::error::{GatewayError, GatewayErrorKind};
use gateway_core::operation::{
ImageRequest, ImageRequestKind, Operation, OperationKind, RawJsonPayload,
};
use gateway_protocol::openai::OPENAI_PROTOCOL;
use serde_json::Value;
use crate::ApiState;
@@ -25,7 +26,6 @@ use crate::openai::{
router::{IMAGE_EDITS_PATH, IMAGE_GENERATIONS_PATH},
};
const OPENAI_PROTOCOL: &str = "openai";
const IMAGE_TURN_ID_CONTEXT_KEY: &str = "image_turn_id";
/// `POST /v1/images/generations`
@@ -13,6 +13,7 @@ use gateway_core::engine::execution::StartedExecution;
use gateway_core::event::ProviderEvent;
use gateway_core::live::LiveHangupRequest;
use gateway_core::operation::Operation;
use gateway_protocol::openai::OPENAI_PROTOCOL;
use serde_json::Value;
use super::{
@@ -31,9 +32,6 @@ use crate::{
},
};
/// 引导 wire 事件使用的协议名;与 Provider 侧 `PROVIDER_NAME` 一致。
const LIVE_PROTOCOL: &str = "openai";
pub(crate) async fn live_call(
State(state): State<ApiState>,
connect_info: Option<Extension<ConnectInfo<SocketAddr>>>,
@@ -344,7 +342,7 @@ fn live_call_body(events: Vec<ProviderEvent>) -> Option<Bytes> {
let mut raw_events = 0usize;
for event in events {
let (_, wire) = event.into_parts();
let Some(wire) = wire.filter(|wire| wire.protocol() == LIVE_PROTOCOL) else {
let Some(wire) = wire.filter(|wire| wire.protocol() == OPENAI_PROTOCOL) else {
continue;
};
let raw = wire.into_raw_http_body()?;
@@ -20,6 +20,7 @@ use gateway_core::event::ProviderResponseHeader;
use gateway_core::operation::{
ProviderHttpHeader, ProviderHttpMethod, ProviderHttpRequest, RawHttpPayload,
};
use gateway_protocol::openai::OPENAI_PROTOCOL;
use serde_json::{Value, json};
use crate::openai::error::openai_error_response;
@@ -172,7 +173,7 @@ pub(crate) fn live_call_operation(
"invalid_request",
))
};
let payload = RawHttpPayload::new("openai", body).map_err(|_| invalid())?;
let payload = RawHttpPayload::new(OPENAI_PROTOCOL, body).map_err(|_| invalid())?;
ProviderHttpRequest::new(
LIVE_CALLS_ENDPOINT,
ProviderHttpMethod::Post,
@@ -20,6 +20,7 @@ use gateway_core::{
lifecycle::CancellationToken,
operation::OperationKind,
};
use gateway_protocol::openai::OPENAI_PROTOCOL;
use std::{future::Future, sync::Arc};
use super::http::{ExpectedBody, buffered_response, error_response, into_http_response};
@@ -38,7 +39,7 @@ impl RequestInput {
pub(crate) fn query(endpoint: String, model_hint: Option<String>, headers: &HeaderMap) -> Self {
Self {
endpoint,
protocol: "openai".to_owned(),
protocol: OPENAI_PROTOCOL.to_owned(),
operation: None,
transport: ClientTransport::HttpJson,
model_hint,
@@ -10,6 +10,7 @@ use gateway_core::{
engine::execution::AuthenticatedClient,
routing::{PublicModelDescriptor, PublicModelId, PublicModelProfile},
};
use gateway_protocol::anthropic::ANTHROPIC_PROTOCOL;
use serde::Deserialize;
use serde_json::{Value, json};
@@ -52,7 +53,7 @@ pub(crate) async fn models(
let mut input = RequestInput::query(super::router::MODELS_PATH.to_owned(), None, &headers);
if anthropic {
input.protocol = "anthropic".to_owned();
input.protocol = ANTHROPIC_PROTOCOL.to_owned();
}
let response = middleware::query_response(service.execution(), client, input, move |client| {
models_response(state, client, query, anthropic)
@@ -173,7 +174,7 @@ pub(crate) async fn model_detail(
&headers,
);
if anthropic {
input.protocol = "anthropic".to_owned();
input.protocol = ANTHROPIC_PROTOCOL.to_owned();
}
let response = middleware::query_response(
service.execution(),
@@ -6,6 +6,7 @@ use axum::http::{HeaderMap, HeaderName, HeaderValue};
use base64::{Engine as _, engine::general_purpose::STANDARD};
use bytes::Bytes;
use gateway_core::operation::{ClientTurn, GenerateRequest, Operation, ProtocolPayload};
use gateway_protocol::openai::OPENAI_PROTOCOL;
use gateway_protocol::openai::{
X_OPENAI_INTERNAL_CODEX_RESPONSES_LITE_HEADER, X_OPENAI_MEMGEN_REQUEST_HEADER,
is_transport_managed_request_header,
@@ -17,7 +18,6 @@ use crate::openai::generation::{
ContinuationIntent, DecodedGenerationRequest, GenerationRequestMetadata,
};
const OPENAI_PROTOCOL: &str = "openai";
const OPENAI_SUBAGENT_KEY: &str = "x-openai-subagent";
const CODEX_TURN_METADATA_KEY: &str = "x-codex-turn-metadata";
const PASSTHROUGH_HEADERS_CONTEXT_KEY: &str = "opaque_request_headers";
@@ -4,6 +4,7 @@ use std::collections::BTreeMap;
use bytes::Bytes;
use gateway_core::event::{GatewayEvent, ProtocolWireEvent, ProviderEvent};
use gateway_protocol::openai::OPENAI_PROTOCOL;
use gateway_protocol::openai::sse::{
encode_sse_event_with_metadata, response_failed_sse_data_from_error_event,
};
@@ -12,8 +13,6 @@ use serde_json::Value;
use super::error::ResponseEncodeError;
use crate::client_error::capacity_error_for_client;
const OPENAI_PROTOCOL: &str = "openai";
/// OpenAI Responses wire 的转发器
///
/// OpenAI Provider 与 xAI adapter 都必须交付 OpenAI wire
@@ -6,6 +6,7 @@ use std::{
};
use gateway_core::event::{GatewayEvent, ProtocolWireEvent, ProviderEvent};
use gateway_protocol::openai::OPENAI_PROTOCOL;
use gateway_protocol::openai::sse::{parse_sse_events, sse_frame_is_done};
use serde_json::{Map, Value};
@@ -44,7 +45,7 @@ impl ResponseValidationFacts {
if let Some(wire) = event
.middleware_origin_wire()
.or_else(|| event.wire_event())
.filter(|wire| wire.protocol() == "openai")
.filter(|wire| wire.protocol() == OPENAI_PROTOCOL)
{
facts.observe_wire(wire);
}
@@ -16,6 +16,7 @@ use gateway_core::engine::response_control::{ResponseControl, ResponseControlUna
use gateway_core::error::{GatewayError, GatewayErrorKind};
use gateway_core::event::ProviderResponseHeader;
use gateway_core::operation::ProviderSessionState;
use gateway_protocol::openai::OPENAI_PROTOCOL;
use crate::openai::error::{gateway_error_contract, gateway_error_from_engine};
use crate::openai::middleware::{PendingExecution, pending_execution_response};
@@ -118,7 +119,7 @@ pub(super) async fn execution_response(
.collect();
// 失败也要保留会话到洋葱链返回,否则释放会话会取消仍在等待 next 的插件
return Ok(pending_execution_response(
"openai".to_owned(),
OPENAI_PROTOCOL.to_owned(),
status,
response_headers,
MiddlewareFrame::new(
@@ -142,7 +143,7 @@ pub(super) async fn execution_response(
.map(|header| MiddlewareHeader::new(header.name(), header.value().clone()))
.collect();
Ok(MiddlewareResponse::new(
"openai".to_owned(),
OPENAI_PROTOCOL.to_owned(),
200,
headers,
Box::new(body),
@@ -159,7 +160,7 @@ pub(super) async fn forward_response(
response_control: &ResponseControl,
) -> ForwardOutcome {
let (protocol, status, headers, mut body, _) = response.into_parts();
if protocol != "openai" || StatusCode::from_u16(status).is_err() {
if protocol != OPENAI_PROTOCOL || StatusCode::from_u16(status).is_err() {
return fail_body(connection, body, MiddlewareError::InvalidState, &request_id).await;
}
let response_headers = headers
@@ -26,6 +26,7 @@ use gateway_core::{
lifecycle::{ConnectionGuard, ConnectionLifecycle},
operation::OperationKind,
};
use gateway_protocol::openai::OPENAI_PROTOCOL;
use crate::middleware::headers::encode_headers;
use crate::{
@@ -328,7 +329,7 @@ async fn serve_responses_websocket(socket: Socket, session: ResponsesWebSocketSe
let validation = ResponseValidationFacts::default();
let input = RequestInput {
endpoint: crate::openai::router::RESPONSES_PATH.to_owned(),
protocol: "openai".to_owned(),
protocol: OPENAI_PROTOCOL.to_owned(),
operation: Some(OperationKind::Generate),
transport: ClientTransport::WebSocket,
model_hint,
@@ -349,7 +350,7 @@ async fn serve_responses_websocket(socket: Socket, session: ResponsesWebSocketSe
Box::new(move |prepared, request| {
Box::pin(async move {
let (protocol, headers, body) = request_parts(request.clone())?;
if protocol != "openai" {
if protocol != OPENAI_PROTOCOL {
return Err(MiddlewareError::Rejected);
}
let payload =
@@ -11,6 +11,7 @@ use axum::{
use gateway_core::engine::execution::ClientTransport;
use gateway_core::error::{GatewayError, GatewayErrorKind};
use gateway_core::operation::{Operation, OperationKind, RawJsonPayload, StandaloneSearchRequest};
use gateway_protocol::openai::OPENAI_PROTOCOL;
use crate::ApiState;
use crate::middleware::headers::encode_headers;
@@ -22,8 +23,6 @@ use crate::openai::{
router::SEARCH_PATH,
};
const OPENAI_PROTOCOL: &str = "openai";
/// `POST /v1/alpha/search`
pub(crate) async fn standalone_search(
State(state): State<ApiState>,
@@ -15,6 +15,7 @@ use gateway_core::lifecycle::{ConnectionDraining, ConnectionGuard, ConnectionLif
use gateway_core::routing::{
ProviderCatalogUnavailable, ProviderKind, PublicModelDescriptor, PublicModelId, UpstreamModelId,
};
use gateway_protocol::openai::OPENAI_PROTOCOL;
use uuid::Uuid;
use super::auth::ClientApiKeyAuthError;
@@ -159,7 +160,7 @@ impl OpenAiService {
upstream_model,
operation,
ExecutionRequestMetadata {
protocol: "openai".to_owned(),
protocol: OPENAI_PROTOCOL.to_owned(),
endpoint,
transport: ClientTransport::HttpJson,
stream: false,
@@ -10,6 +10,14 @@ use std::time::SystemTime;
pub(super) struct TestPeerInfo(pub Arc<Mutex<Option<ProviderQuota>>>);
impl UpstreamInfoSource for TestPeerInfo {
fn is_peer(&self, _: &ProviderAccount) -> bool {
self.0
.lock()
.unwrap()
.as_ref()
.is_some_and(|quota| quota.upstream.as_ref().is_some_and(|info| info.discovered))
}
fn signal(
&self,
_: &ProviderAccount,
@@ -28,5 +28,7 @@ pub use model::*;
pub(crate) use selection::smart_score;
pub use selection::*;
pub use smart_scheduling::{SmartSchedulingConfig, SmartSchedulingConfigError};
pub use store::ProviderAccountStore;
pub use store::{
AccountCatalogObservation, AccountCatalogOmissions, ProviderAccountStore, StoredAccountCatalog,
};
pub use upstream::*;
@@ -12,9 +12,45 @@ use super::{
QuotaObservationTouch, QuotaObservationVersion, QuotaWriteOutcome,
};
/// 一个账号最后成功的目录观测;正文由 Provider 解释,归属不能按相同正文合并
#[derive(Clone)]
pub struct AccountCatalogObservation {
pub account_id: ProviderAccountId,
pub revision: CredentialRevision,
pub plan: Option<String>,
pub etag: Option<String>,
pub observed_at: std::time::SystemTime,
}
/// 批量恢复时每份不同的正文只传输、解析一次
pub struct StoredAccountCatalog {
pub accounts: Vec<AccountCatalogObservation>,
pub document: bytes::Bytes,
}
/// 批量选号的存储优化:只有账号自己的当前版本、有效时间和目录摘要都匹配才可跳过
pub struct AccountCatalogOmissions {
pub documents: Vec<[u8; 32]>,
pub exceptions: Vec<ProviderAccountId>,
pub observed_after: std::time::SystemTime,
pub as_of: std::time::SystemTime,
}
/// `provider_accounts` 的数据库中立端口
#[async_trait]
pub trait ProviderAccountStore: Send + Sync {
async fn load_account_catalogs(
&self,
provider: &ProviderKind,
) -> Result<Vec<StoredAccountCatalog>, StoreError>;
/// 账号删除或凭据已变化时返回 false,不发布迟到的观测
async fn save_account_catalog(
&self,
observation: &AccountCatalogObservation,
document: &[u8],
) -> Result<bool, StoreError>;
async fn create_account(&self, account: NewProviderAccount) -> Result<(), StoreError>;
async fn get_account(
@@ -36,6 +72,16 @@ pub trait ProviderAccountStore: Send + Sync {
account_ids: &[ProviderAccountId],
) -> Result<Vec<ProviderAccount>, StoreError>;
/// 存储可以下推目录筛选,也可以返回完整候选;调用方始终复核最终账号
async fn list_catalog_candidates(
&self,
provider: &ProviderKind,
accounts: &[ProviderAccountId],
_omissions: &AccountCatalogOmissions,
) -> Result<Vec<ProviderAccount>, StoreError> {
self.list_for_provider_accounts(provider, accounts).await
}
/// 一次有界查询返回账号事实和 revision-fenced 明文 credential
async fn list_refresh_candidates(
&self,
@@ -84,6 +84,9 @@ pub struct UpstreamSchedulingSignal {
}
pub trait UpstreamInfoSource: Send + Sync {
/// 当前连接是否已识别为互联节点,不以额度快照的新鲜度代替身份
fn is_peer(&self, account: &ProviderAccount) -> bool;
fn signal(
&self,
account: &ProviderAccount,
@@ -110,6 +113,13 @@ impl UpstreamInfoReader {
}
impl UpstreamInfoSource for UpstreamInfoReader {
fn is_peer(&self, account: &ProviderAccount) -> bool {
self.source
.get()
.and_then(Weak::upgrade)
.is_some_and(|source| source.is_peer(account))
}
fn signal(
&self,
account: &ProviderAccount,
@@ -158,6 +158,7 @@ impl TraceContext {
"recovery": {
"replaySafe": error.replay_is_safe(),
"retryProhibited": error.retry_is_prohibited(),
"sameAccountRetryBlocked": error.same_account_retry_blocked(),
"sameAccountCredentialRetry": error.retries_same_account(),
"intent": error.pre_delivery_retry().map(|intent| format!("{intent:?}")),
"continuationFailure": error.continuation_failure().map(|failure| format!("{failure:?}")),
@@ -346,6 +347,7 @@ impl TraceContext {
for key in [
"replaySafe",
"retryProhibited",
"sameAccountRetryBlocked",
"sameAccountCredentialRetry",
] {
kept[key] = json!(recovery.get(key).and_then(Value::as_bool));
@@ -1672,15 +1672,19 @@ where
error.pre_delivery_retry(),
Some(crate::error::PreDeliveryRetry::SameAccountConnectionRetry { .. })
);
// 服务器建议不可截短;超出请求冻结的上限后沿用安全换号或原错误出口
// 服务器建议不可截短
// 超出上限,或上游拒绝继续重试当前账号时,跳过同账号等待
// 较短的 Retry-After 也直接沿用安全换号或原错误出口
let retry_after_exceeds_limit = error
.retry_after()
.is_some_and(|delay| delay > self.plan.max_retry_after());
let same_account_retry_blocked = error.same_account_retry_blocked();
// Provider 的有界同账号传输恢复是尽力重试策略,不要求把 Sent/Ambiguous
// 改写成安全证明;外部副作用、取消、预算与下游提交仍由 Core 拦截
let transport_recovery = match error.pre_delivery_retry() {
Some(crate::error::PreDeliveryRetry::SameAccountConnectionRetry { transport })
if !execution_effect_observed
if !same_account_retry_blocked
&& !execution_effect_observed
&& !self.connection_budget.exhausted()
&& !self.replay_barrier_reached
&& !self.delivery_pending
@@ -1698,7 +1702,8 @@ where
Some(crate::error::PreDeliveryRetry::SameAccountTransportRetry {
retry_index,
delay,
}) if !retry_after_exceeds_limit
}) if !same_account_retry_blocked
&& !retry_after_exceeds_limit
&& !execution_effect_observed
&& !self.connection_budget.exhausted()
&& !self.replay_barrier_reached
@@ -1707,7 +1712,8 @@ where
Some((AttemptTransport::Retry(retry_index), delay))
}
Some(crate::error::PreDeliveryRetry::SameAccountTransportFallback)
if !retry_after_exceeds_limit
if !same_account_retry_blocked
&& !retry_after_exceeds_limit
&& !execution_effect_observed
&& !self.connection_budget.exhausted()
&& !self.replay_barrier_reached
@@ -1754,7 +1760,11 @@ where
max_delay,
jitter,
..
}) if ordinary_retry && !suppress_transient_retry && !retry_after_exceeds_limit => {
}) if !same_account_retry_blocked
&& ordinary_retry
&& !suppress_transient_retry
&& !retry_after_exceeds_limit =>
{
let retries = self
.transient_retry_counts
.entry(current.metadata.provider_account_id().clone())
@@ -1776,7 +1786,8 @@ where
}
_ => None,
};
let same_account_retry = !execution_effect_observed
let same_account_retry = !same_account_retry_blocked
&& !execution_effect_observed
&& error.retries_same_account()
&& provider_proved_replay_safe
&& !self.replay_barrier_reached
@@ -1834,6 +1845,7 @@ where
"reserveConnectionRotation": reserve_connection_rotation,
"transientRetry": transient_retry.is_some(),
"retryAfterExceeded": retry_after_exceeds_limit,
"sameAccountRetryBlocked": same_account_retry_blocked,
"maxRetryAfterSeconds": self.plan.max_retry_after().as_secs(),
"clientRetrySuppressed": suppress_transient_retry,
"rotationRetry": rotation_retry, "accountRotations": self.account_rotations,
@@ -801,16 +801,38 @@ impl DefaultExecutionService {
active_request,
admission_decision_ms,
} = start_guard;
let observation = plan
.candidates()
.first()
.map_or_else(Default::default, |candidate| {
providers.request_observation(
candidate.provider(),
&operation,
client.policy.key_id(),
)
let observation_provider = match &target {
ExecutionTarget::ProviderEndpoint { provider, .. } => Some(provider),
ExecutionTarget::Model(_) => plan
.candidates()
.iter()
.find(|candidate| {
operation
.provider_session_state(candidate.provider().as_str())
.is_some_and(|state| state.extension_owner().is_none())
})
.or_else(|| plan.candidates().first())
.map(|candidate| candidate.provider()),
};
let provider_observation =
observation_provider.map_or_else(Default::default, |candidate| {
providers.request_observation(candidate, &operation, client.policy.key_id())
});
// 纯协议事实不依赖候选顺序,Provider 仅补充所属会话的旧配置和恢复关联
let observation = match &operation {
Operation::Generate(request) => {
let payload = request.protocol_payload();
gateway_protocol::request::request_observation(
payload.protocol(),
payload.body(),
payload.context(),
provider_observation
.previous_reasoning_effort_override
.as_deref(),
)
}
_ => Default::default(),
};
let (request_kind, subagent_kind) = if let Some(nested) = authorization.nested.as_ref() {
(
Some("plugin_child_model".to_owned()),
@@ -843,7 +865,9 @@ impl DefaultExecutionService {
operation: operation.kind(),
endpoint: metadata.endpoint,
client_transport: metadata.transport.as_str().to_owned(),
requested_model: target.into_public_model().or(observation.requested_model),
requested_model: target
.into_public_model()
.or(provider_observation.requested_model),
client_ip: metadata.client_ip,
user_agent: metadata.user_agent,
reasoning_effort: observation.reasoning_effort,
@@ -851,7 +875,7 @@ impl DefaultExecutionService {
request_kind,
subagent_kind,
compact: observation.compact,
continuation: observation.continuation,
continuation: provider_observation.continuation,
image_generation_requested: operation.image_generation_requested(),
admission_decision_ms,
started_at,
@@ -1088,17 +1088,13 @@ pub struct ProviderRequest {
candidate: ProviderCandidate,
}
/// Provider 对公共观测表可解释的请求语义;未知字段保持空值
/// Provider 私有状态与独立端点的观测补充,通用请求语义由协议层提取
#[derive(Debug, Clone, Default, PartialEq, Eq)]
pub struct ProviderRequestObservation {
/// 独立 Provider 端点声明的请求模型,仅用于观测,不参与文本模型目录路由
pub requested_model: Option<PublicModelId>,
/// Provider 从客户端请求与续接配置中解析的推理强度
pub reasoning_effort: Option<String>,
pub reasoning_preset: Option<String>,
pub request_kind: Option<String>,
pub subagent_kind: Option<String>,
pub compact: bool,
/// 已核实归属的上一响应配置;当前请求是否继承由协议层判断
pub previous_reasoning_effort_override: Option<String>,
pub continuation: ContinuationRequestObservation,
}
@@ -516,6 +516,7 @@ pub struct ProviderError {
failure_observation: Option<Box<ProviderErrorFailureObservation>>,
replay_safe: bool,
retry_prohibited: bool,
same_account_retry_blocked: bool,
pre_delivery_retry: Option<Box<PreDeliveryRetry>>,
credential_recovery_required: bool,
retry_same_account: bool,
@@ -559,6 +560,7 @@ impl ProviderError {
failure_observation: None,
replay_safe: false,
retry_prohibited: false,
same_account_retry_blocked: false,
pre_delivery_retry: None,
credential_recovery_required: false,
retry_same_account: false,
@@ -685,6 +687,22 @@ impl ProviderError {
self.retry_prohibited
}
/// 保留已选择的耗尽策略,但禁止执行同账号退避和传输恢复
///
/// Provider 在上游明确拒绝继续使用当前账号时设置
/// Core 不根据错误码自行推断
#[must_use]
pub const fn with_same_account_retry_blocked(mut self) -> Self {
self.same_account_retry_blocked = true;
self
}
/// 返回上游是否拒绝继续重试当前账号
#[must_use]
pub const fn same_account_retry_blocked(&self) -> bool {
self.same_account_retry_blocked
}
/// 允许 Core 仅在客户端尚未收到任何事件时执行一次受预算约束的换号恢复
///
/// 该标记不证明上游未执行,也不等价于 [`Self::with_replay_safe`];它只表达
@@ -1008,6 +1026,7 @@ impl ProviderError {
failure_observation: self.failure_observation.clone(),
replay_safe: self.replay_safe,
retry_prohibited: self.retry_prohibited,
same_account_retry_blocked: self.same_account_retry_blocked,
pre_delivery_retry: self.pre_delivery_retry.clone(),
credential_recovery_required: self.credential_recovery_required,
retry_same_account: self.retry_same_account,
@@ -1036,6 +1055,10 @@ impl fmt::Debug for ProviderError {
)
.field("retry_after", &self.retry_after)
.field("retry_prohibited", &self.retry_prohibited)
.field(
"same_account_retry_blocked",
&self.same_account_retry_blocked,
)
.field("continuation_failure", &self.continuation_failure)
.field(
"continuation_recovery_disposition",
+4 -35
View File
@@ -3,7 +3,7 @@
//! 这里只保留网关需要解释、路由或结算的稳定语义
//! 协议 adapter 无法写入
//! wire body 与连接级事实随 [`ProtocolPayload`] 或 [`RawJsonPayload`] 不透明传递,
//! 由对应 Provider 自己解释
//! 纯协议事实由协议层解析,私有字段与会话由对应 Provider 解释
use std::collections::BTreeSet;
use std::fmt;
@@ -458,8 +458,8 @@ impl ClientTurn {
impl GenerateRequest {
/// 创建携带协议不透明正文的生成请求
///
/// 完整请求只以原始 wire object 保存;Core 仅局部只读已知字段以推导路由能力,
/// Provider 负责解释与转换正文
/// 完整请求以原始 wire object 保存,协议层提取路由与权限所需的纯事实
/// Provider 在选定账号后编排协议转换与私有字段处理
#[must_use]
pub fn from_protocol_payload(protocol_payload: ProtocolPayload) -> Self {
Self {
@@ -521,9 +521,7 @@ impl GenerateRequest {
/// 返回客户端是否请求了图片生成工具
#[must_use]
pub fn image_generation_requested(&self) -> bool {
self.payload.source_image_generation
|| (self.protocol_payload().protocol() == "openai"
&& requests_image_generation(self.body()))
self.payload.source_image_generation || self.payload.wire_facts.image_generation
}
/// 返回指定 Provider 的连接内会话状态
@@ -596,35 +594,6 @@ impl fmt::Debug for GenerateRequest {
}
}
/// 原生工具声明授权上游生成图片,Codex 的被动 namespace 与图片输入不属于该权限
fn requests_image_generation(body: &Map<String, Value>) -> bool {
contains_image_tool(body.get("tools"))
|| body
.get("input")
.and_then(Value::as_array)
.is_some_and(|items| {
items.iter().any(|item| {
item.get("type").and_then(Value::as_str) == Some("additional_tools")
&& contains_image_tool(item.get("tools"))
})
})
|| body.get("tool_choice").is_some_and(|choice| {
is_image_tool(choice)
|| (choice.get("type").and_then(Value::as_str) == Some("allowed_tools")
&& contains_image_tool(choice.get("tools")))
})
}
fn contains_image_tool(value: Option<&Value>) -> bool {
value
.and_then(Value::as_array)
.is_some_and(|tools| tools.iter().any(is_image_tool))
}
fn is_image_tool(value: &Value) -> bool {
value.get("type").and_then(Value::as_str) == Some("image_generation")
}
/// 图像 API 的稳定端点语义
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub enum ImageRequestKind {
@@ -1,7 +1,7 @@
//! Provider 会话亲和、绑定和排除事实
use super::{ProviderStoreError, ProviderStoreErrorKind};
use crate::{account::ProviderAccountId, identity::ProviderKind};
use crate::{account::ProviderAccountId, identity::ProviderKind, policy::ClientApiKeyId};
use futures::future::BoxFuture;
use std::{collections::BTreeSet, fmt, time::Duration};
@@ -81,6 +81,36 @@ pub struct ProviderSessionAlias {
pub follow_only: bool,
}
/// 发送前需要核对版本的会话绑定;首次认领的 expected 为空
#[derive(Clone, Copy)]
pub struct ProviderSessionClaim<'a> {
pub key: &'a ProviderSessionAffinityKey,
pub expected: Option<&'a ProviderSessionBinding>,
pub account: &'a ProviderAccountId,
pub ttl: Duration,
}
/// 仅成功的内置轮询推进游标;有会话的分配必须同时提交绑定
#[derive(Clone, Copy)]
pub enum ProviderAccountAdmission<'a> {
Session(ProviderSessionClaim<'a>),
RoundRobin {
client_key: &'a ClientApiKeyId,
cursor: u64,
session: Option<ProviderSessionClaim<'a>>,
},
}
impl<'a> ProviderAccountAdmission<'a> {
#[must_use]
pub const fn session(self) -> Option<ProviderSessionClaim<'a>> {
match self {
Self::Session(session) => Some(session),
Self::RoundRobin { session, .. } => session,
}
}
}
/// 客户端作用域内的会话账号绑定,原始会话身份由 Provider 哈希后传入
pub trait ProviderSessionAffinityPort: Send + Sync {
fn load<'a>(
@@ -89,16 +119,19 @@ pub trait ProviderSessionAffinityPort: Send + Sync {
key: &'a ProviderSessionAffinityKey,
) -> BoxFuture<'a, Result<Option<ProviderSessionBinding>, ProviderStoreError>>;
/// 在发送前原子认领、续期或迁移;None 表示快照已变化,调用方必须释放租约重新选择
/// expected 为 None 只允许首次认领,已有快照必须连同版本匹配
fn compare_and_bind<'a>(
/// 读取本 Client Key 和 Provider 的下一轮位置,不推进游标
fn round_robin_cursor<'a>(
&'a self,
provider_kind: &'a ProviderKind,
key: &'a ProviderSessionAffinityKey,
expected: Option<&'a ProviderSessionBinding>,
account_id: &'a ProviderAccountId,
ttl: Duration,
) -> BoxFuture<'a, Result<Option<ProviderSessionBinding>, ProviderStoreError>>;
provider: &'a ProviderKind,
client_key: &'a ClientApiKeyId,
) -> BoxFuture<'a, Result<u64, ProviderStoreError>>;
/// 原子核对并提交会话版本和轮询位置;冲突时均不写入,调用方释放租约重选
fn admit_account<'a>(
&'a self,
provider: &'a ProviderKind,
admission: ProviderAccountAdmission<'a>,
) -> BoxFuture<'a, Result<bool, ProviderStoreError>>;
/// 显式观测到的请求关联只指向会话键,不缓存账号,迁移后仍读取当前绑定
fn load_alias<'a>(
@@ -154,10 +154,16 @@ impl ProviderSchedulingLeaseRequest {
}
}
/// Lease 生命周期由具体 Store guard 管理,Provider 只能持有
pub trait ProviderLeaseGuard: Send + Sync + 'static {}
/// Drop 回收租约;同一请求需要重选时显式等待释放,避免被自己的租约阻塞
pub trait ProviderLeaseGuard: Send + Sync + 'static {
fn release(self: Box<Self>) -> BoxFuture<'static, Result<(), ProviderStoreError>>;
}
impl<T> ProviderLeaseGuard for T where T: Send + Sync + 'static {}
impl ProviderLeaseGuard for () {
fn release(self: Box<Self>) -> BoxFuture<'static, Result<(), ProviderStoreError>> {
Box::pin(async { Ok(()) })
}
}
pub enum ProviderLeaseAcquisition {
Acquired(Box<dyn ProviderLeaseGuard>),
@@ -10,8 +10,9 @@ pub use lease::{
};
mod affinity;
pub use affinity::{
ProviderSessionAffinityKey, ProviderSessionAffinityPort, ProviderSessionAlias,
ProviderSessionBinding, ProviderSessionExclusionPort, ProviderSessionExclusions,
ProviderAccountAdmission, ProviderSessionAffinityKey, ProviderSessionAffinityPort,
ProviderSessionAlias, ProviderSessionBinding, ProviderSessionClaim,
ProviderSessionExclusionPort, ProviderSessionExclusions,
};
mod catalog;
pub use catalog::{
@@ -4325,6 +4325,87 @@ fn retry_after_above_the_configured_limit_rotates_without_waiting() {
}
}
#[test]
fn no_more_retry_skips_same_account_wait_even_when_retry_after_is_short() {
let advised = Duration::from_secs(30);
for (rejection, rotates) in [
(
transient_rejection()
.with_same_account_retry_blocked()
.with_retry_after(advised),
true,
),
(
transient_rejection()
.with_transient_retry(
NonZeroU32::MIN,
Duration::ZERO,
Duration::ZERO,
Duration::ZERO,
TransientRetryPolicy::RetryThenReturnError,
)
.with_same_account_retry_blocked()
.with_retry_after(advised),
false,
),
] {
let operation = generate_operation();
let route_plan = retry_after_plan(&operation, 240);
let mut scripts = vec![Script::Stream {
account_id: "acct_first",
items: vec![Err(rejection)],
}];
if rotates {
scripts.push(Script::Stream {
account_id: "acct_second",
items: complete_stream(None),
});
}
let (coordinator, _, provider) = coordinator(scripts);
let mut session = block_on(coordinator.start(
model_request(&operation, SystemTime::now() + Duration::from_secs(30)),
operation,
route_plan,
None,
None,
CancellationToken::new(),
))
.unwrap();
let trace = session.trace().clone();
let mut collection = Box::pin(session.collect_uncommitted());
let polled = block_on(async { futures::poll!(&mut collection) });
let trace = trace.snapshot().unwrap();
let retry = trace["events"]
.as_array()
.unwrap()
.iter()
.find(|event| event["stage"] == "retry.decided")
.unwrap();
assert_eq!(retry["data"]["sameAccountRetryBlocked"], true);
assert_eq!(retry["data"]["transientRetry"], false);
assert_eq!(retry["data"]["retryAfterExceeded"], false);
if rotates {
assert!(matches!(polled, std::task::Poll::Ready(Ok(_))));
drop(collection);
let contexts = provider.contexts.lock().unwrap();
assert_eq!(contexts.len(), 2);
assert_eq!(contexts[1].required_account(), None);
assert!(
contexts[1]
.excluded_accounts()
.contains(&ProviderAccountId::new("acct_first").unwrap())
);
} else {
assert!(matches!(
polled,
std::task::Poll::Ready(Err(EngineError::Provider(error)))
if error.retry_after() == Some(advised)
));
assert_eq!(provider.contexts.lock().unwrap().len(), 1);
}
}
}
#[test]
fn excessive_retry_after_preserves_the_error_when_rotation_is_unavailable_or_disallowed() {
for (rejection, rotation_allowed) in [
@@ -148,6 +148,209 @@ fn request(service: &DefaultExecutionService, transport: ClientTransport) -> Sta
}
}
#[test]
fn request_observation_is_stable_across_candidate_order_and_respects_session_ownership() {
for order in [["alpha", "beta"], ["beta", "alpha"]] {
for (protocol, state_owner, extension_owned, body, expected) in [
(
"openai",
None,
false,
json!({"reasoning":{"effort":"medium"}}),
"medium",
),
(
"anthropic",
None,
false,
json!({"output_config":{"effort":" Future-Level "}}),
"Future-Level",
),
(
"openai",
Some("beta"),
false,
json!({"previous_response_id":"resp_previous","reasoning":{"effort":"medium"}}),
"high",
),
(
"openai",
Some("unavailable"),
false,
json!({"previous_response_id":"resp_previous","reasoning":{"effort":"medium"}}),
"medium",
),
(
"openai",
Some("beta"),
true,
json!({"previous_response_id":"resp_previous","reasoning":{"effort":"medium"}}),
"medium",
),
(
"openai",
Some("beta"),
false,
json!({"previous_response_id":"resp_previous","reasoning":{"effort":"medium"},"input":[{"type":"configuration_update","reasoning":{}}]}),
"medium",
),
] {
let kinds: Vec<_> = order
.iter()
.map(|name| ProviderKind::new(*name).unwrap())
.collect();
let directory = RuntimeAccountDirectory::new(
kinds
.iter()
.map(|kind| {
(
ProviderAccountId::new(format!("acct_{}", kind.as_str())).unwrap(),
RuntimeAccount::new(kind.clone(), BTreeSet::new()),
)
})
.collect(),
);
let snapshot = RuntimeSnapshot::new(
ConfigRevision::new(1).unwrap(),
gateway_core::settings::SettingsValues::new(
1,
"smart",
Default::default(),
None,
None,
),
kinds.clone(),
kinds
.into_iter()
.map(|kind| {
ProviderModel::new(
kind,
UpstreamModelId::new("gpt-start").unwrap(),
ModelCapabilities::new(BTreeSet::from([OperationKind::Generate]), None)
.with_upstream_feature_validation(),
)
})
.collect(),
vec![ClientPolicy::new(
ClientApiKeyId::new("key_start_test").unwrap(),
PlaintextClientApiKey::new("sk_start_test").unwrap(),
Arc::new(FrozenAccountScope::new(
Arc::new(directory),
ClientRoutingScope::all_accounts(),
)),
true,
RateLimits::unlimited(),
)],
)
.unwrap();
let store = Arc::new(TrackingExecutionStore::default());
let service =
DefaultExecutionService::new(
RuntimeSnapshotHandle::new(snapshot),
store.clone(),
ProviderRegistry::new(order.map(|name| {
Arc::new(SessionObservationProvider(name)) as Arc<dyn Provider>
}))
.unwrap(),
Arc::new(Admissions::default()),
Arc::new(UnusedContinuation),
Arc::new(RecordingClientApiKeyUsage::default()),
Arc::new(crate::RecordingDiagnostics::default()),
);
let mut generate = GenerateRequest::from_protocol_payload(
ProtocolPayload::json_object(protocol, body.as_object().unwrap().clone())
.unwrap()
.with_context(serde_json::Map::from_iter([(
"turn_metadata".to_owned(),
json!("{\"request_kind\":\"review\"}"),
)])),
);
if let Some(owner) = state_owner {
let mut state = gateway_core::operation::ProviderSessionState::new(
owner,
serde_json::Map::from_iter([("effort".to_owned(), json!("high"))]),
)
.unwrap();
if extension_owned {
state = state.with_extension_owner(
gateway_core::operation::ExtensionSessionOwner {
instance_id: "plugin".into(),
contribution_id: "adapter".into(),
adapter_id: "adapter".into(),
generation: 1,
incarnation: "incarnation".into(),
connection_local: true,
},
);
}
generate = generate.with_provider_session_state(state);
}
let mut start = request(&service, ClientTransport::HttpJson);
start.operation = Operation::Generate(generate);
start.metadata.protocol = protocol.to_owned();
let started = block_on(service.start(start)).expect("request observation is recorded");
block_on(started.session.detach_finalize());
let requests = store.requests.lock().unwrap();
let recorded = requests.first().expect("logical request");
assert_eq!(
recorded.reasoning_effort.as_deref(),
Some(expected),
"order={order:?}, state={state_owner:?}, extension={extension_owned}"
);
assert_eq!(
recorded.request_kind.as_deref(),
(protocol == "openai").then_some("review")
);
}
}
}
struct SessionObservationProvider(&'static str);
#[async_trait]
impl Provider for SessionObservationProvider {
fn name(&self) -> &'static str {
self.0
}
fn catalog_generation(&self) -> ProviderCatalogGeneration {
ProviderCatalogGeneration::default()
}
fn request_observation(
&self,
operation: &Operation,
_: &ClientApiKeyId,
) -> ProviderRequestObservation {
ProviderRequestObservation {
previous_reasoning_effort_override: operation
.provider_session_state(self.0)
.filter(|state| state.extension_owner().is_none())
.and_then(|state| state.payload().get("effort"))
.and_then(Value::as_str)
.map(str::to_owned),
..Default::default()
}
}
async fn query_model_capabilities(
&self,
) -> Result<Vec<ProviderModelCapabilities>, ProviderError> {
Ok(Vec::new())
}
async fn execute(
self: Arc<Self>,
_: ProviderRequest,
_: AttemptContext,
) -> Result<ProviderStream, ProviderError> {
Err(ProviderError::new(
ProviderErrorKind::Unavailable,
UpstreamSendState::NotSent,
))
}
}
#[test]
fn default_request_has_no_total_deadline_and_explicit_timeout_can_be_cleared() {
block_on(async {
+25 -33
View File
@@ -157,40 +157,28 @@ fn prompt_cache_and_image_generation_should_be_read_from_raw_body() {
}
#[test]
fn image_generation_detection_distinguishes_tools_from_image_inputs_and_passive_namespaces() {
for body in [
json!({"tools":[{"type":"image_generation"}]}),
json!({"tool_choice":{"type":"image_generation"}}),
json!({"tool_choice":{"type":"allowed_tools","tools":[{"type":"image_generation"}]}}),
json!({"input":[{"type":"additional_tools","tools":[{"type":"image_generation"}]}]}),
] {
assert!(
generate(body.clone()).image_generation_requested(),
"{body}"
);
}
for body in [
json!({"input":"Please generate an image"}),
json!({"input":[{"role":"user","content":[{"type":"input_image","image_url":"data:image/png;base64,AA=="}]}]}),
json!({"tools":[{"type":"namespace","name":"image_gen","tools":[{"type":"function","name":"imagegen"}]}]}),
json!({"tools":[{"type":"function","name":"image_generation"}]}),
] {
assert!(
!generate(body.clone()).image_generation_requested(),
"{body}"
);
}
let private = GenerateRequest::from_protocol_payload(
ProtocolPayload::json_object(
"private",
json!({"tools":[{"type":"image_generation"}]})
.as_object()
.unwrap()
.clone(),
fn image_generation_fact_tracks_body_replacement_and_ignores_context() {
let operation = Operation::Generate(generate(json!({"input":"hello"})));
let operation = operation
.replace_protocol_wire(
Some(Bytes::from_static(
br#"{"tools":[{"type":"image_generation"}]}"#,
)),
Map::new(),
)
.unwrap(),
);
assert!(!private.image_generation_requested());
.unwrap();
assert!(operation.image_generation_requested());
let operation = operation
.replace_protocol_wire(None, Map::from_iter([("tools".to_owned(), json!([]))]))
.unwrap();
assert!(operation.image_generation_requested());
let operation = operation
.replace_protocol_wire(
Some(Bytes::from_static(br#"{"input":"hello"}"#)),
Map::new(),
)
.unwrap();
assert!(!operation.image_generation_requested());
}
#[test]
@@ -203,6 +191,10 @@ fn image_generation_fact_survives_cross_protocol_translation() {
)
.unwrap();
assert!(translated.image_generation_requested());
let replaced = translated
.replace_protocol_wire(Some(Bytes::from_static(br#"{"messages":[]}"#)), Map::new())
.unwrap();
assert!(replaced.image_generation_requested());
}
#[test]
@@ -449,6 +449,9 @@ Ping 的自动 Pong 与关闭握手属于底层 WebSocket 协议事实;插件
或 `WebSocket { uri }`,不能在准备过程中切换传输类型。HTTP 正文是已经序列化、完成内容编码的实际字节,
插件自行处理压缩与 `Content-Encoding` 等头部的一致性;宿主不解码正文、补回字段或修复请求
内置转发在出站准备前,仅为已识别的互联上游组装 `X-CPR-Route`,向普通模型上游剥离。
插件改写目标或自行实现转发时,负责在互联节点间保留并维护路径,在最终模型上游前移除该头;宿主不在插件返回后重新判断或修复互联链路
`OutboundCall.metadata` 提供真实请求、Client Key、分组、Provider、模型、账号与 attempt 序号。
这些归属由宿主冻结,返回值只有请求;修改 URL、认证头或正文不改变账号租约、凭据版本与计费归属。
宿主不从最终字节重建模型、选号或续接事实;需要改变这些业务含义时,插件应在相应的 request / attempt 阶段保持一致。
@@ -24,6 +24,7 @@ pub use stream::{
MessagesAccumulator, MessagesToResponsesStream, ResponsesToMessagesStream,
message_to_stream_events,
};
pub use thinking::messages_request_effort;
pub use usage::MessageUsage;
/// Messages 载荷的协议标识
@@ -1,16 +1,22 @@
//! 按目标模型能力映射推理等级与预算;未知兼容上游沿用 CPA 手动预算
//! 提取 Messages 显式 effort,并按目标模型能力映射推理等级与预算
use serde_json::{Map, Value, json};
use super::MessagesModelCapabilities;
pub(super) fn to_responses(request: &Map<String, Value>) -> Value {
let configured = request
/// 返回客户端显式设置的 effort,预算与默认档位由转换器单独处理
#[must_use]
pub fn messages_request_effort(request: &Map<String, Value>) -> Option<&str> {
request
.get("output_config")
.and_then(|value| value.get("effort"))
.and_then(Value::as_str)
.map(|value| value.trim().to_ascii_lowercase())
.filter(|value| !value.is_empty());
.map(str::trim)
.filter(|value| !value.is_empty())
}
pub(super) fn to_responses(request: &Map<String, Value>) -> Value {
let configured = messages_request_effort(request).map(str::to_ascii_lowercase);
let thinking = request.get("thinking");
let effort = match thinking
.and_then(|value| value.get("type"))
@@ -127,7 +127,7 @@ pub fn codex_responses_request_semantics(
) -> CodexResponsesRequestSemantics {
codex_responses_request_semantics_with_turn_metadata(
body,
non_empty_string(context.get("turn_metadata")),
context.get("turn_metadata").and_then(Value::as_str),
None,
)
}
@@ -153,10 +153,12 @@ pub fn codex_responses_request_semantics_with_turn_metadata(
let reasoning_effort = reasoning_effort_override
.clone()
.or_else(|| reasoning_effort(body.get("reasoning")));
// 权威上下文与显式 turnMetadata 即使为空也不退回其他别名
let parsed_turn_metadata = turn_metadata
.or_else(|| body.get("turnMetadata").and_then(Value::as_str))
.or_else(|| request_turn_metadata(body))
.and_then(|value| serde_json::from_str::<Value>(value).ok());
let request_kind = parsed_turn_metadata
let mut request_kind = parsed_turn_metadata
.as_ref()
.and_then(|metadata| non_empty_string(metadata.get("request_kind")))
.map(ToOwned::to_owned);
@@ -174,6 +176,12 @@ pub fn codex_responses_request_semantics_with_turn_metadata(
&& proactive_multi_agent)
.then_some("ultra");
let compact = request_kind.as_deref() == Some("compaction") || compact_trigger;
// generate=false 才是预热,客户端分类提示不能把真实推理排除在用量统计外
if body.get("generate").and_then(Value::as_bool) == Some(false) {
request_kind = Some("prewarm".to_owned());
} else if request_kind.as_deref() == Some("prewarm") {
request_kind = None;
}
CodexResponsesRequestSemantics {
reasoning_effort,
@@ -217,7 +225,10 @@ fn latest_multi_agent_mode<'a>(
input: impl DoubleEndedIterator<Item = &'a Value>,
) -> Option<&'a str> {
input.rev().find_map(|item| {
if non_empty_string(item.get("role")) != Some("developer") {
// 显式 system 消息会由 Codex 编码器转为 developer,观测在编码前保持等价
let system_message = item.get("type").and_then(Value::as_str) == Some("message")
&& item.get("role").and_then(Value::as_str) == Some("system");
if non_empty_string(item.get("role")) != Some("developer") && !system_message {
return None;
}
item.get("content")?
@@ -1,4 +1,4 @@
//! 从 OpenAI 事件与响应头提取用量、额度和重试等待事实
//! 从 OpenAI 事件与响应头提取用量、额度和重试事实
use std::{collections::BTreeMap, time::Duration};
@@ -249,6 +249,37 @@ pub(super) fn retry_after_header(value: &Value) -> Option<(&Value, u64)> {
Some((value, seconds))
}
/// 判断上游错误是否携带 `x-retry-metadata: NO_MORE_RETRY`
///
/// 读取位置与 Retry-After 相同:优先错误对象的 headers,再看事件自身的 headers
/// 只承认这一明确信号,其他取值、数组和非字符串都不构成该事实
pub fn upstream_no_more_retry(value: &Value) -> bool {
let error = value
.pointer("/response/error")
.or_else(|| value.get("error"))
.unwrap_or(value);
no_more_retry_headers(error) || no_more_retry_headers(value)
}
/// 从上游错误正文判断是否声明不再重试当前账号
pub fn upstream_no_more_retry_from_body(body: &str) -> bool {
serde_json::from_str::<Value>(body).is_ok_and(|value| upstream_no_more_retry(&value))
}
fn no_more_retry_headers(value: &Value) -> bool {
value
.get("headers")
.and_then(Value::as_object)
.is_some_and(|headers| {
headers.iter().any(|(name, value)| {
name.eq_ignore_ascii_case("x-retry-metadata")
&& value
.as_str()
.is_some_and(|value| value.trim().eq_ignore_ascii_case("NO_MORE_RETRY"))
})
})
}
/// 标准化的单个限流窗口
#[derive(Debug, Clone, Copy, PartialEq)]
pub struct RateLimitWindow {
@@ -1,5 +1,8 @@
//! 共享的 OpenAI 线协议类型、解析与编码入口
/// OpenAI 线协议载荷与事件的协议标识
pub const OPENAI_PROTOCOL: &str = "openai";
/// Codex Responses Lite 的 HTTP 请求头
pub const X_OPENAI_INTERNAL_CODEX_RESPONSES_LITE_HEADER: &str =
"x-openai-internal-codex-responses-lite";
+84 -3
View File
@@ -1,7 +1,56 @@
//! 从协议正文提取调度需要的纯请求事实
//! 从协议正文提取能力要求与请求观测事实
use serde_json::{Map, Value};
use crate::anthropic::{ANTHROPIC_PROTOCOL, messages_request_effort};
use crate::openai::{OPENAI_PROTOCOL, codex_responses_request_semantics_with_turn_metadata};
/// 与执行 Provider 无关的逻辑请求观测
#[derive(Debug, Clone, Default, PartialEq, Eq)]
pub struct RequestObservation {
/// 当前请求与已核实续接配置共同确定的推理强度
pub reasoning_effort: Option<String>,
/// 协议声明派生的推理预设
pub reasoning_preset: Option<String>,
/// 协议声明的请求类型,可信内部分类由调用方覆盖
pub request_kind: Option<String>,
/// 协议声明的子代理类型
pub subagent_kind: Option<String>,
/// 当前请求是否声明压缩
pub compact: bool,
}
/// 旧配置只由状态所属 Provider 提供,是否继承仍由当前协议的续接语义决定
#[must_use]
pub fn request_observation(
protocol: &str,
body: &Map<String, Value>,
context: &Map<String, Value>,
previous_reasoning_effort_override: Option<&str>,
) -> RequestObservation {
match protocol {
OPENAI_PROTOCOL => {
let semantics = codex_responses_request_semantics_with_turn_metadata(
body,
context.get("turn_metadata").and_then(Value::as_str),
previous_reasoning_effort_override,
);
RequestObservation {
reasoning_effort: semantics.reasoning_effort,
reasoning_preset: semantics.reasoning_preset.map(str::to_owned),
request_kind: semantics.request_kind,
subagent_kind: semantics.subagent_kind,
compact: semantics.compact,
}
}
ANTHROPIC_PROTOCOL => RequestObservation {
reasoning_effort: messages_request_effort(body).map(str::to_owned),
..RequestObservation::default()
},
_ => RequestObservation::default(),
}
}
/// 请求正文表达的能力要求
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
pub struct RequestFacts {
@@ -11,6 +60,8 @@ pub struct RequestFacts {
pub tools: bool,
/// 图片输入
pub vision: bool,
/// 图片生成工具声明
pub image_generation: bool,
/// 推理配置
pub reasoning: bool,
/// JSON Schema 输出
@@ -30,12 +81,13 @@ pub fn request_facts(protocol: &str, body: &Map<String, Value>) -> RequestFacts
..RequestFacts::default()
};
match protocol {
"openai" => {
OPENAI_PROTOCOL => {
facts.max_output_tokens = body
.get("max_output_tokens")
.and_then(Value::as_u64)
.filter(|n| *n > 0);
facts.vision = contains_type(body.get("input"), "input_image");
facts.image_generation = requests_image_generation(body);
facts.reasoning = body.get("reasoning").is_some_and(Value::is_object);
facts.json_schema = body
.get("text")
@@ -48,7 +100,7 @@ pub fn request_facts(protocol: &str, body: &Map<String, Value>) -> RequestFacts
.and_then(Value::as_str)
.is_some();
}
"anthropic" => {
ANTHROPIC_PROTOCOL => {
facts.max_output_tokens = body.get("max_tokens").and_then(Value::as_u64);
facts.vision = contains_type(body.get("messages"), "image");
facts.reasoning = body
@@ -68,6 +120,35 @@ pub fn request_facts(protocol: &str, body: &Map<String, Value>) -> RequestFacts
facts
}
/// 原生工具声明授权上游生成图片,Codex 的被动 namespace 与图片输入不属于该权限
fn requests_image_generation(body: &Map<String, Value>) -> bool {
contains_image_tool(body.get("tools"))
|| body
.get("input")
.and_then(Value::as_array)
.is_some_and(|items| {
items.iter().any(|item| {
item.get("type").and_then(Value::as_str) == Some("additional_tools")
&& contains_image_tool(item.get("tools"))
})
})
|| body.get("tool_choice").is_some_and(|choice| {
is_image_tool(choice)
|| (choice.get("type").and_then(Value::as_str) == Some("allowed_tools")
&& contains_image_tool(choice.get("tools")))
})
}
fn contains_image_tool(value: Option<&Value>) -> bool {
value
.and_then(Value::as_array)
.is_some_and(|tools| tools.iter().any(is_image_tool))
}
fn is_image_tool(value: &Value) -> bool {
value.get("type").and_then(Value::as_str) == Some("image_generation")
}
fn contains_type(value: Option<&Value>, expected: &str) -> bool {
match value {
Some(Value::Array(values)) => values.iter().any(|v| contains_type(Some(v), expected)),
@@ -219,6 +219,27 @@ fn codex_semantics_should_detect_ultra_preset_from_proactive_multi_agent() {
assert_eq!(semantics.reasoning_preset, Some("ultra"));
assert!(!semantics.compact);
for (role, message_type, expected) in [
("system", Some("message"), Some("ultra")),
("system", None, None),
("system", Some("function_call"), None),
("user", Some("message"), None),
(" system ", Some("message"), None),
] {
let mut shaped = body.clone();
let item = shaped["input"][0].as_object_mut().unwrap();
item.insert("role".to_owned(), json!(role));
if let Some(message_type) = message_type {
item.insert("type".to_owned(), json!(message_type));
} else {
item.remove("type");
}
assert_eq!(
codex_responses_request_semantics(&shaped, &Map::new()).reasoning_preset,
expected
);
}
for (effort, expected_preset) in [("high", None), ("max", Some("ultra"))] {
body["input"]
.as_array_mut()
@@ -230,6 +251,27 @@ fn codex_semantics_should_detect_ultra_preset_from_proactive_multi_agent() {
}
}
#[test]
fn empty_authoritative_turn_metadata_does_not_fall_back_to_body_aliases() {
for (context, body) in [
(
json!({"turn_metadata":""}),
json!({"turnMetadata":"{\"request_kind\":\"review\"}"}),
),
(
json!({"turn_metadata":" "}),
json!({"turnMetadata":"{\"request_kind\":\"review\"}"}),
),
(
json!({}),
json!({"turnMetadata":"","turn_metadata":"{\"request_kind\":\"review\"}"}),
),
] {
let semantics = codex_responses_request_semantics(&object(body), &object(context));
assert_eq!(semantics.request_kind, None);
}
}
fn object(value: Value) -> Map<String, Value> {
let Value::Object(object) = value else {
panic!("test value must be an object");
@@ -267,3 +309,42 @@ fn account_identity_does_not_infer_a_session_from_turn_or_cache_keys() {
assert_eq!(codex_account_session_id(&object(body), &Map::new()), None);
}
}
#[test]
fn request_observation_classifies_prewarm_from_generate() {
for (generate, request_kind, expected_kind) in [
(Some(json!(false)), None, Some("prewarm")),
(Some(json!(false)), Some("prewarm"), Some("prewarm")),
(Some(json!(false)), Some("review"), Some("prewarm")),
(Some(json!(true)), Some("prewarm"), None),
(None, Some("prewarm"), None),
(Some(json!(null)), Some("prewarm"), None),
(Some(json!("false")), Some("prewarm"), None),
(None, Some("review"), Some("review")),
(None, None, None),
] {
let mut body = Map::from_iter([
("model".to_owned(), json!("gpt-test")),
("input".to_owned(), json!("hello")),
("store".to_owned(), json!(false)),
]);
if let Some(generate) = generate {
body.insert("generate".to_owned(), generate);
}
let context = Map::from_iter([(
"turn_metadata".to_owned(),
Value::String(
json!({"request_kind": request_kind, "subagent_kind": "worker"}).to_string(),
),
)]);
let observation = codex_responses_request_semantics(&body, &context);
assert_eq!(
observation.request_kind.as_deref(),
expected_kind,
"generate={:?}, request_kind={request_kind:?}",
body.get("generate"),
);
assert_eq!(observation.subagent_kind.as_deref(), Some("worker"));
}
}
@@ -4,7 +4,8 @@ use gateway_protocol::openai::events::{
RateLimitKeySource, RateLimitWindow, TokenUsage, billable_usage_is_complete, extract_sse_usage,
extract_usage, is_codex_quota_header_name, is_rate_limit_header_name,
parse_image_rate_limit_headers, parse_rate_limit_headers, parse_rate_limits_event,
rate_limits_to_header_pairs, retry_after_seconds_from_body,
rate_limits_to_header_pairs, retry_after_seconds_from_body, upstream_no_more_retry,
upstream_no_more_retry_from_body,
};
use serde_json::json;
@@ -429,6 +430,51 @@ fn retry_after_should_not_infer_retry_from_unrelated_error_message() {
assert_eq!(retry_after_seconds_from_body(&body), None);
}
#[test]
fn no_more_retry_should_read_the_error_header_without_replacing_retry_after() {
for (body, expected) in [
(
json!({"type":"error","error":{"type":"service_unavailable_error","code":"server_is_overloaded","headers":{"x-retry-metadata":"NO_MORE_RETRY","retry-after":"30","retry-after-ms":"30000"}}}),
true,
),
(
json!({"type":"response.failed","response":{"error":{"headers":{"X-Retry-Metadata":" no_more_retry "}}}}),
true,
),
(
json!({"type":"error","headers":{"x-retry-metadata":"NO_MORE_RETRY"},"error":{"code":"server_is_overloaded"}}),
true,
),
(
json!({"error":{"headers":{"x-retry-metadata":"RETRY","retry-after":"30"}}}),
false,
),
(
json!({"error":{"headers":{"x-retry-metadata":["NO_MORE_RETRY"]}}}),
false,
),
(
json!({"error":{"code":"server_is_overloaded","headers":{"retry-after":"30"}}}),
false,
),
] {
assert_eq!(upstream_no_more_retry(&body), expected, "{body}");
assert_eq!(
upstream_no_more_retry_from_body(&body.to_string()),
expected,
"{body}"
);
}
assert!(!upstream_no_more_retry_from_body("not-json"));
assert_eq!(
retry_after_seconds_from_body(
&json!({"error":{"headers":{"x-retry-metadata":"NO_MORE_RETRY","retry-after":"30"}}})
.to_string()
),
Some(30)
);
}
#[test]
fn retry_after_should_reject_malformed_json() {
assert_eq!(retry_after_seconds_from_body("not-json"), None);
@@ -1,7 +1,7 @@
//! 验证不同协议的共享请求事实
use gateway_protocol::request::request_facts;
use serde_json::json;
use gateway_protocol::request::{request_facts, request_observation};
use serde_json::{Map, json};
#[test]
fn messages_facts_include_images_inside_tool_results() {
@@ -20,3 +20,77 @@ fn responses_facts_keep_native_continuation_separate_from_history() {
assert!(facts.native_continuation);
assert!(!facts.tools && !facts.vision && !facts.reasoning && !facts.json_schema);
}
#[test]
fn image_generation_detection_distinguishes_tools_from_image_inputs_and_passive_namespaces() {
for body in [
json!({"tools":[{"type":"image_generation"}]}),
json!({"tool_choice":{"type":"image_generation"}}),
json!({"tool_choice":{"type":"allowed_tools","tools":[{"type":"image_generation"}]}}),
json!({"input":[{"type":"additional_tools","tools":[{"type":"image_generation"}]}]}),
] {
assert!(
request_facts("openai", body.as_object().expect("请求正文")).image_generation,
"{body}"
);
}
for body in [
json!({"input":"Please generate an image"}),
json!({"input":[{"role":"user","content":[{"type":"input_image","image_url":"data:image/png;base64,AA=="}]}]}),
json!({"tools":[{"type":"namespace","name":"image_gen","tools":[{"type":"function","name":"imagegen"}]}]}),
json!({"tools":[{"type":"function","name":"image_generation"}]}),
] {
assert!(
!request_facts("openai", body.as_object().expect("请求正文")).image_generation,
"{body}"
);
}
let body = json!({"tools":[{"type":"image_generation"}]});
for protocol in ["anthropic", "private"] {
assert!(!request_facts(protocol, body.as_object().expect("请求正文")).image_generation);
}
}
#[test]
fn messages_request_observation_preserves_explicit_effort_without_inventing_defaults() {
for (options, expected) in [
(json!({"output_config":{"effort":"xhigh"}}), Some("xhigh")),
(
json!({"thinking":{"type":"adaptive"},"output_config":{"effort":"low"}}),
Some("low"),
),
(
json!({"thinking":{"type":"disabled"},"output_config":{"effort":"high"}}),
Some("high"),
),
(
json!({"output_config":{"effort":" Future-Level "}}),
Some("Future-Level"),
),
(json!({"output_config":{"effort":" "}}), None),
(json!({"output_config":{"effort":null}}), None),
(json!({"output_config":{"effort":42}}), None),
(json!({"thinking":{"type":"adaptive"}}), None),
(
json!({"thinking":{"type":"enabled","budget_tokens":32768}}),
None,
),
(json!({}), None),
] {
let mut body = json!({"model":"gpt-5.4","max_tokens":128,"messages":[{"role":"user","content":"hello"}]});
body.as_object_mut()
.unwrap()
.extend(options.as_object().unwrap().clone());
let observation = request_observation(
"anthropic",
body.as_object().unwrap(),
&Map::new(),
Some("inapplicable-previous-effort"),
);
assert_eq!(
observation.reasoning_effort.as_deref(),
expected,
"{options}"
);
}
}
@@ -14,6 +14,21 @@ fn loaded_credential_from_record(
#[async_trait]
impl ProviderAccountStore for PgProviderAccountRepository {
async fn load_account_catalogs(
&self,
provider: &ProviderKind,
) -> Result<Vec<gateway_core::account::StoredAccountCatalog>, CoreStoreError> {
self.load_model_catalogs(provider).await
}
async fn save_account_catalog(
&self,
observation: &gateway_core::account::AccountCatalogObservation,
document: &[u8],
) -> Result<bool, CoreStoreError> {
self.save_model_catalog(observation, document).await
}
async fn create_account(&self, account: CoreNewProviderAccount) -> Result<(), CoreStoreError> {
if account.account.revision().get() != 1 {
return Err(CoreStoreError::new(CoreStoreErrorKind::InvalidData));
@@ -88,6 +103,29 @@ impl ProviderAccountStore for PgProviderAccountRepository {
.collect()
}
async fn list_catalog_candidates(
&self,
provider: &ProviderKind,
accounts: &[CoreProviderAccountId],
omissions: &gateway_core::account::AccountCatalogOmissions,
) -> Result<Vec<CoreProviderAccount>, CoreStoreError> {
let ids = accounts
.iter()
.map(CoreProviderAccountId::as_str)
.collect::<Vec<_>>();
self.list_provider_account_summaries(
Some(provider.as_str()),
false,
Some(&ids),
Some(omissions),
)
.await
.map_err(core_store_error)?
.into_iter()
.map(core_account_from_summary)
.collect()
}
async fn list_for_provider_accounts(
&self,
provider: &ProviderKind,
@@ -100,12 +138,17 @@ impl ProviderAccountStore for PgProviderAccountRepository {
.iter()
.map(CoreProviderAccountId::as_str)
.collect::<Vec<_>>();
self.list_provider_account_summaries(Some(provider.as_str()), false, Some(&account_ids))
.await
.map_err(core_store_error)?
.into_iter()
.map(core_account_from_summary)
.collect()
self.list_provider_account_summaries(
Some(provider.as_str()),
false,
Some(&account_ids),
None,
)
.await
.map_err(core_store_error)?
.into_iter()
.map(core_account_from_summary)
.collect()
}
async fn list_refresh_candidates(
@@ -58,6 +58,7 @@ mod admin_queries;
mod authorization;
mod core_adapter;
mod mapping;
mod model_catalogs;
mod quota;
mod quota_forecast;
mod repository;
@@ -0,0 +1,83 @@
//! 账号目录的版本约束写入和按相同正文合并的批量恢复
use super::{
CoreCredentialRevision, CoreProviderAccountId, CoreStoreError, CoreStoreErrorKind, DateTime,
PgProviderAccountRepository, ProviderKind, Utc,
};
use gateway_core::account::{AccountCatalogObservation, StoredAccountCatalog};
impl PgProviderAccountRepository {
pub(super) async fn load_model_catalogs(
&self,
provider: &ProviderKind,
) -> Result<Vec<StoredAccountCatalog>, CoreStoreError> {
// 先聚合轻量归属,再按主键读取代表正文,避免一万个相同账号展开一万份 TOAST
let rows = sqlx::query_as::<_, (Vec<String>, Vec<i64>, Vec<Option<String>>, Vec<Option<String>>, Vec<DateTime<Utc>>, Vec<u8>)>(
"select g.ids, g.revisions, g.plans, g.etags, g.times, d.document from (
select min(c.account_id) representative, array_agg(c.account_id order by c.account_id) ids,
array_agg(c.credential_revision order by c.account_id) revisions,
array_agg(c.plan_type order by c.account_id) plans, array_agg(c.etag order by c.account_id) etags, array_agg(c.observed_at order by c.account_id) times
from provider_account_model_catalogs c join provider_accounts a on a.id=c.account_id
where a.provider_kind=$1 and a.enabled and a.credential_state in ('unknown','ready')
and (a.access_token_expires_at is null or a.access_token_expires_at > now())
and a.credential_revision=c.credential_revision and a.plan_type is not distinct from c.plan_type
group by c.document_hash
) g join provider_account_model_catalogs d on d.account_id=g.representative")
.bind(provider.as_str()).fetch_all(&self.pool).await.map_err(unavailable)?;
rows.into_iter()
.map(|(ids, revisions, plans, etags, times, document)| {
let accounts = ids
.into_iter()
.zip(revisions)
.zip(plans)
.zip(etags)
.zip(times)
.map(|((((id, revision), plan), etag), time)| {
Ok(AccountCatalogObservation {
account_id: CoreProviderAccountId::new(id).map_err(|e| {
CoreStoreError::caused_by(CoreStoreErrorKind::InvalidData, e)
})?,
revision: CoreCredentialRevision::new(
u64::try_from(revision).map_err(|_| {
CoreStoreError::new(CoreStoreErrorKind::InvalidData)
})?,
)
.map_err(|e| {
CoreStoreError::caused_by(CoreStoreErrorKind::InvalidData, e)
})?,
plan,
etag,
observed_at: time.into(),
})
})
.collect::<Result<_, CoreStoreError>>()?;
Ok(StoredAccountCatalog {
accounts,
document: document.into(),
})
})
.collect()
}
pub(super) async fn save_model_catalog(
&self,
observation: &AccountCatalogObservation,
document: &[u8],
) -> Result<bool, CoreStoreError> {
let result = sqlx::query("insert into provider_account_model_catalogs (account_id, credential_revision, plan_type, observed_at, document, etag)
select a.id, a.credential_revision, a.plan_type, $4, $5, $6 from provider_accounts a
where a.id=$1 and a.credential_revision=$2 and a.plan_type is not distinct from $3 for share
on conflict (account_id) do update set credential_revision=excluded.credential_revision, plan_type=excluded.plan_type,
etag=excluded.etag, observed_at=excluded.observed_at, document=excluded.document
where provider_account_model_catalogs.credential_revision < excluded.credential_revision or
(provider_account_model_catalogs.credential_revision=excluded.credential_revision and provider_account_model_catalogs.observed_at <= excluded.observed_at)")
.bind(observation.account_id.as_str()).bind(i64::try_from(observation.revision.get()).map_err(|_|CoreStoreError::new(CoreStoreErrorKind::InvalidData))?)
.bind(observation.plan.as_deref()).bind(DateTime::<Utc>::from(observation.observed_at)).bind(document).bind(observation.etag.as_deref())
.execute(&self.pool).await.map_err(unavailable)?;
Ok(result.rows_affected() == 1)
}
}
fn unavailable(source: sqlx::Error) -> CoreStoreError {
CoreStoreError::caused_by(CoreStoreErrorKind::Unavailable, source)
}
@@ -104,6 +104,7 @@ impl PgProviderAccountRepository {
provider_kind: Option<&str>,
include_disabled: bool,
account_ids: Option<&[&str]>,
omissions: Option<&gateway_core::account::AccountCatalogOmissions>,
) -> StoreResult<Vec<ProviderAccountSummary>> {
let mut query = sqlx::QueryBuilder::<Postgres>::new(
"select coalesce(proxy_enabled, true) as outbound_proxy_enabled, coalesce(max_websocket_connections, 0) as outbound_proxy_max_websocket_connections, auto_location, detected_location_json, location_country, location_region, location_city, location_timezone, outbound_proxy_id, outbound_proxy_url, id, provider_kind, name, notes, email, upstream_user_id,
@@ -129,6 +130,24 @@ impl PgProviderAccountRepository {
.push_bind(account_ids)
.push("::text[])");
}
if let Some(omit) = omissions {
let hashes = omit
.documents
.iter()
.map(<[u8; 32]>::as_slice)
.collect::<Vec<_>>();
let exceptions = omit
.exceptions
.iter()
.map(|id| id.as_str())
.collect::<Vec<_>>();
query.push(" and not exists (select 1 from provider_account_model_catalogs c
where c.account_id=provider_accounts.id and c.credential_revision=provider_accounts.credential_revision
and c.plan_type is not distinct from provider_accounts.plan_type and c.document_hash=any(")
.push_bind(hashes).push("::bytea[]) and c.observed_at > ").push_bind(DateTime::<Utc>::from(omit.observed_after))
.push(" and c.observed_at <= ").push_bind(DateTime::<Utc>::from(omit.as_of))
.push(" and not (c.account_id=any(").push_bind(exceptions).push("::text[])))");
}
query.push(" order by provider_kind, name, id");
let rows = query
.build()
@@ -215,7 +234,7 @@ impl ProviderAccountRepository for PgProviderAccountRepository {
provider_kind: Option<&str>,
include_disabled: bool,
) -> StoreResult<Vec<ProviderAccountSummary>> {
self.list_provider_account_summaries(provider_kind, include_disabled, None)
self.list_provider_account_summaries(provider_kind, include_disabled, None, None)
.await
}
@@ -5,6 +5,9 @@ use std::time::Duration;
use futures::future::BoxFuture;
use gateway_core::lifecycle::CancellationToken;
use gateway_core::provider_ports::{
ProviderLeaseGuard, ProviderStoreError, ProviderStoreErrorKind,
};
use gateway_core::task::{DaemonTask, WorkerDaemonContext, WorkerTaskError};
use tokio::sync::{RwLock, RwLockReadGuard, oneshot};
use tokio::task::JoinSet;
@@ -22,11 +25,32 @@ struct State {
}
pub(crate) struct SchedulingLeaseGuard {
_release: oneshot::Sender<()>,
release: oneshot::Sender<oneshot::Sender<Result<(), ProviderStoreError>>>,
// 活跃请求可能比组合根活得久,不能因端口析构提前中止其续期与租约
_cleanup: SchedulingLeaseCleanup,
}
impl ProviderLeaseGuard for SchedulingLeaseGuard {
fn release(self: Box<Self>) -> BoxFuture<'static, Result<(), ProviderStoreError>> {
Box::pin(async move {
let (ack, released) = oneshot::channel();
self.release.send(ack).map_err(|_| {
ProviderStoreError::new(
ProviderStoreErrorKind::Unavailable,
"request scheduling lease release",
)
})?;
released.await.map_err(|source| {
ProviderStoreError::caused_by(
ProviderStoreErrorKind::Unavailable,
"await scheduling lease release",
source,
)
})?
})
}
}
impl SchedulingLeaseCleanup {
// 关闭必须等已发出的申请完成登记;不能漏掉 Redis 已创建、尚未返回的租约
pub(crate) async fn acquire(&self) -> Option<RwLockReadGuard<'_, ()>> {
@@ -39,7 +63,8 @@ impl SchedulingLeaseCleanup {
cancellation: CancellationToken,
release: impl Future<Output = crate::StoreResult<bool>> + Send + 'static,
) -> SchedulingLeaseGuard {
let (guard, dropped) = oneshot::channel::<()>();
let (guard, dropped) =
oneshot::channel::<oneshot::Sender<Result<(), ProviderStoreError>>>();
let stopping = self.0.stopping.clone();
let mut tasks = self
.0
@@ -50,17 +75,27 @@ impl SchedulingLeaseCleanup {
report_join(result);
}
tasks.spawn(async move {
tokio::select! {
let acknowledgement = tokio::select! {
biased;
() = stopping.cancelled() => cancellation.cancel(),
_ = dropped => {}
}
if let Err(error) = release.await {
() = stopping.cancelled() => { cancellation.cancel(); None },
ack = dropped => ack.ok(),
};
let result = release.await;
if let Err(error) = &result {
tracing::warn!(%error, "账号租约释放失败,依赖 TTL 回收");
}
if let Some(ack) = acknowledgement {
let _ = ack.send(result.map(|_| ()).map_err(|source| {
ProviderStoreError::caused_by(
ProviderStoreErrorKind::Unavailable,
"release scheduling lease",
source,
)
}));
}
});
SchedulingLeaseGuard {
_release: guard,
release: guard,
_cleanup: self.clone(),
}
}
@@ -249,6 +249,20 @@ pub struct CredentialLeaseGuard {
grant: Option<CredentialLeaseGrant>,
}
impl gateway_core::provider_ports::ProviderLeaseGuard for CredentialLeaseGuard {
fn release(
self: Box<Self>,
) -> futures::future::BoxFuture<'static, Result<(), ProviderStoreError>> {
Box::pin(async move {
(*self)
.release()
.await
.map(|_| ())
.map_err(|source| crate::provider_unavailable("release credential lease", source))
})
}
}
impl CredentialLeaseGuard {
/// 自动续期仅属于返回 guard;丢失租约会取消依赖该槽位的请求
pub fn maintain(
@@ -319,6 +333,16 @@ struct RenewingSchedulingLease {
_guard: CredentialLeaseGuard,
}
impl gateway_core::provider_ports::ProviderLeaseGuard for RenewingSchedulingLease {
fn release(
self: Box<Self>,
) -> futures::future::BoxFuture<'static, Result<(), ProviderStoreError>> {
let Self { _renewal, _guard } = *self;
drop(_renewal);
Box::new(_guard).release()
}
}
impl fmt::Debug for CredentialLeaseGuard {
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
formatter
@@ -430,7 +454,12 @@ impl RedisCredentialLeaseRepository {
client_api_key_id: &ClientApiKeyId,
provider_kind: &ProviderKind,
) -> StoreResult<u64> {
let key = self.scheduling_cursor_key(client_api_key_id, provider_kind)?;
let key = super::scheduling_cursor_key(
&self.namespace,
client_api_key_id,
provider_kind,
"cursor",
)?;
let mut connection = self.connection.clone();
let cursor = Script::new(ADVANCE_SCHEDULING_CURSOR_SCRIPT)
.key(key)
@@ -473,21 +502,6 @@ impl RedisCredentialLeaseRepository {
))
}
fn scheduling_cursor_key(
&self,
client_api_key_id: &ClientApiKeyId,
provider_kind: &ProviderKind,
) -> StoreResult<String> {
let fingerprint = resource_fingerprint(
"provider scheduling cursor",
&format!("{}\0{}", client_api_key_id.as_str(), provider_kind.as_str()),
)?;
Ok(format!(
"{}:scheduler:client-provider:{{{fingerprint}}}:cursor",
self.namespace
))
}
async fn acquire_with_limits(
&self,
request: &CredentialLeaseRequest,
@@ -43,6 +43,21 @@ pub(crate) fn resource_fingerprint(entity: &'static str, value: &str) -> StoreRe
Ok(hex::encode(Sha256::digest(value.as_bytes())))
}
fn scheduling_cursor_key(
namespace: &str,
client_key: &gateway_core::policy::ClientApiKeyId,
provider: &gateway_core::routing::ProviderKind,
suffix: &str,
) -> StoreResult<String> {
let fingerprint = resource_fingerprint(
"provider scheduling cursor",
&format!("{}\0{}", client_key.as_str(), provider.as_str()),
)?;
Ok(format!(
"{namespace}:scheduler:client-provider:{{{fingerprint}}}:{suffix}"
))
}
pub(crate) fn namespace(value: &str) -> StoreResult<String> {
require_nonempty("Redis namespace", "namespace", value)?;
if value.len() > 128
@@ -3,9 +3,10 @@
use std::time::Duration;
use gateway_core::account::ProviderAccountId;
use gateway_core::policy::ClientApiKeyId;
use gateway_core::provider_ports::{
ProviderSessionAffinityKey, ProviderSessionAffinityPort, ProviderSessionBinding,
ProviderStoreError, ProviderStoreErrorKind,
ProviderAccountAdmission, ProviderSessionAffinityKey, ProviderSessionAffinityPort,
ProviderSessionBinding, ProviderStoreError, ProviderStoreErrorKind,
};
use gateway_core::routing::ProviderKind;
use redis::aio::ConnectionManager;
@@ -17,14 +18,15 @@ use super::{namespace, resource_fingerprint};
const MAX_SESSION_AFFINITY_TTL: Duration =
Duration::from_secs(gateway_core::account::MAX_SESSION_AFFINITY_TTL_HOURS as u64 * 3600);
// 比较完整记录而非账号 ID;冲突时既不改绑定,也不续期
const COMPARE_AND_BIND_SCRIPT: &str = r#"
local current = redis.call('GET', KEYS[1])
if (not current and ARGV[1] == '') or current == ARGV[1] then
redis.call('PSETEX', KEYS[1], tonumber(ARGV[3]), ARGV[2])
return 1
end
return 0
const CURSOR_TTL_MILLIS: u64 = 24 * 60 * 60 * 1000;
// 先比较全部快照再写入;失去首绑或轮询位置的请求不能续期或消耗名额
const ADMIT_ACCOUNT_SCRIPT: &str = r#"
if ARGV[2] ~= '' and (redis.call('GET', KEYS[1]) or '') ~= ARGV[1] then return 0 end
if ARGV[4] ~= '' and (redis.call('GET', KEYS[#KEYS]) or '0') ~= ARGV[4] then return 0 end
if ARGV[2] ~= '' then redis.call('PSETEX', KEYS[1], ARGV[3], ARGV[2]) end
if ARGV[4] ~= '' then redis.call('PSETEX', KEYS[#KEYS], ARGV[6], ARGV[5]) end
return 1
"#;
#[derive(serde::Serialize, serde::Deserialize)]
@@ -75,6 +77,15 @@ impl RedisProviderSessionAffinityRepository {
})
}
fn cursor_key(
&self,
provider: &ProviderKind,
client_key: &ClientApiKeyId,
) -> Result<String, ProviderStoreError> {
super::scheduling_cursor_key(&self.namespace, client_key, provider, "session-cursor")
.map_err(|_| provider_invalid("encode provider round robin cursor key"))
}
fn key(
&self,
provider_kind: &ProviderKind,
@@ -113,42 +124,71 @@ impl ProviderSessionAffinityPort for RedisProviderSessionAffinityRepository {
})
}
fn compare_and_bind<'a>(
fn round_robin_cursor<'a>(
&'a self,
provider_kind: &'a ProviderKind,
key: &'a ProviderSessionAffinityKey,
expected: Option<&'a ProviderSessionBinding>,
account_id: &'a ProviderAccountId,
ttl: Duration,
) -> futures::future::BoxFuture<'a, Result<Option<ProviderSessionBinding>, ProviderStoreError>>
{
provider: &'a ProviderKind,
client_key: &'a ClientApiKeyId,
) -> futures::future::BoxFuture<'a, Result<u64, ProviderStoreError>> {
Box::pin(async move {
let ttl_millis = session_affinity_ttl_millis(ttl)?;
let expected_raw = expected
.map(BindingRecord::encode)
.transpose()?
.unwrap_or_default();
// 同账号命中仅续期;初次认领及换号生成新版本
let binding = match expected.filter(|binding| binding.account_id() == account_id) {
Some(binding) => binding.clone(),
None => ProviderSessionBinding::new(
account_id.clone(),
uuid::Uuid::new_v4().simple().to_string(),
)?,
};
let applied = redis::Script::new(COMPARE_AND_BIND_SCRIPT)
.key(self.key(provider_kind, key)?)
.arg(expected_raw)
.arg(BindingRecord::encode(&binding)?)
.arg(ttl_millis)
.invoke_async::<bool>(&mut self.connection.clone())
let key = self.cursor_key(provider, client_key)?;
redis::cmd("GET")
.arg(key)
.query_async::<Option<u64>>(&mut self.connection.clone())
.await
.map_err(|source| {
crate::provider_unavailable("admit provider session binding", source)
})?;
Ok(applied.then_some(binding))
.map(Option::unwrap_or_default)
.map_err(|source| crate::provider_unavailable("read round robin cursor", source))
})
}
fn admit_account<'a>(
&'a self,
provider: &'a ProviderKind,
admission: ProviderAccountAdmission<'a>,
) -> futures::future::BoxFuture<'a, Result<bool, ProviderStoreError>> {
Box::pin(async move {
let script = redis::Script::new(ADMIT_ACCOUNT_SCRIPT);
let mut invocation = script.prepare_invoke();
if let Some(claim) = admission.session() {
let binding = match claim.expected.filter(|b| b.account_id() == claim.account) {
Some(existing) => existing.clone(),
None => ProviderSessionBinding::new(
claim.account.clone(),
uuid::Uuid::new_v4().simple().to_string(),
)?,
};
invocation
.key(self.key(provider, claim.key)?)
.arg(
claim
.expected
.map(BindingRecord::encode)
.transpose()?
.unwrap_or_default(),
)
.arg(BindingRecord::encode(&binding)?)
.arg(session_affinity_ttl_millis(claim.ttl)?);
} else {
invocation.arg("").arg("").arg(0);
}
if let ProviderAccountAdmission::RoundRobin {
client_key, cursor, ..
} = admission
{
invocation
.key(self.cursor_key(provider, client_key)?)
.arg(cursor)
.arg(cursor.wrapping_add(1))
.arg(CURSOR_TTL_MILLIS);
} else {
invocation.arg("");
}
invocation
.invoke_async(&mut self.connection.clone())
.await
.map_err(|source| crate::provider_unavailable("admit provider account", source))
})
}
fn load_alias<'a>(
&'a self,
provider: &'a ProviderKind,
@@ -83,6 +83,14 @@ async fn shutdown_releases_owned_account_pools_and_preserves_other_instances() {
.unwrap(),
ProviderLeaseAcquisition::Busy { .. }
));
second_guard.release().await.unwrap();
let ProviderLeaseAcquisition::Acquired(second_guard) = second_port
.try_acquire(ProviderLeaseRequest::Scheduling(other.clone()))
.await
.unwrap()
else {
panic!("released slot must be available before reselection")
};
first_requests.push(request);
second_requests.push(other);
first_guards.push(first_guard);
@@ -103,3 +103,101 @@ async fn scoped_account_list_filters_ids_provider_and_live_enabled_state() {
);
database.close().await;
}
#[tokio::test]
async fn bulk_catalog_restore_keeps_ten_thousand_owners_and_fences_late_writes() {
use gateway_core::account::{AccountCatalogObservation, CredentialRevision};
let Some(database) = TestDatabase::create("account_catalog_bulk").await else {
return;
};
let repo = PgProviderAccountRepository::new(database.pool.clone());
sqlx::query("insert into provider_accounts (id,provider_kind,name,authentication_kind,provider_credentials_json,has_refresh_token,credential_state,credential_observed_at,created_at,updated_at,plan_type)
select 'acct_'||i,'openai','source '||i,'oauth','{}',false,'ready',now(),now(),now(),'pro' from generate_series(1,10000) i").execute(&database.pool).await.unwrap();
let observation = AccountCatalogObservation {
account_id: ProviderAccountId::new("acct_1").unwrap(),
revision: CredentialRevision::new(1).unwrap(),
observed_at: std::time::SystemTime::now(),
etag: Some("v1".into()),
plan: Some("pro".into()),
};
let document = vec![b'x'; 64 * 1024];
assert!(
repo.save_account_catalog(&observation, &document)
.await
.unwrap()
);
sqlx::query("insert into provider_account_model_catalogs (account_id,credential_revision,plan_type,observed_at,document)
select a.id,1,'pro',now(),c.document from provider_accounts a cross join provider_account_model_catalogs c where a.id<>'acct_1'").execute(&database.pool).await.unwrap();
let provider = ProviderKind::new("openai").unwrap();
let start = std::time::Instant::now();
let group = repo.load_account_catalogs(&provider).await.unwrap();
eprintln!(
"catalog bulk restore: 10000 accounts, {} documents, {} ms",
group.len(),
start.elapsed().as_millis()
);
assert_eq!(group.len(), 1);
assert_eq!(group[0].accounts.len(), 10000);
assert_eq!(&group[0].document[..], &document);
let older = AccountCatalogObservation {
observed_at: observation.observed_at - std::time::Duration::from_secs(1),
..observation.clone()
};
assert!(!repo.save_account_catalog(&older, b"late").await.unwrap());
sqlx::query("update provider_accounts set credential_revision=2 where id='acct_1'")
.execute(&database.pool)
.await
.unwrap();
assert!(
!repo
.save_account_catalog(&observation, b"old revision")
.await
.unwrap()
);
sqlx::query("update provider_accounts set enabled=false where id='acct_2'")
.execute(&database.pool)
.await
.unwrap();
sqlx::query("delete from provider_accounts where id='acct_3'")
.execute(&database.pool)
.await
.unwrap();
let group = repo.load_account_catalogs(&provider).await.unwrap();
assert_eq!(group[0].accounts.len(), 9997);
assert!(
group[0]
.accounts
.iter()
.all(|a| !matches!(a.account_id.as_str(), "acct_1" | "acct_2" | "acct_3"))
);
// 数据库只能排除有当前版本、未过期的账号证据;显式例外仍进入选号
let hash: Vec<u8> = sqlx::query_scalar("select sha256($1::bytea)")
.bind(&document)
.fetch_one(&database.pool)
.await
.unwrap();
let now = std::time::SystemTime::now();
sqlx::query("update provider_account_model_catalogs set observed_at=now()-interval '1 hour' where account_id='acct_4'").execute(&database.pool).await.unwrap();
let filter = gateway_core::account::AccountCatalogOmissions {
documents: vec![hash.try_into().unwrap()],
exceptions: vec![ProviderAccountId::new("acct_5").unwrap()],
observed_after: now - std::time::Duration::from_secs(1800),
as_of: now,
};
let ids = (1..=10000)
.map(|n| ProviderAccountId::new(format!("acct_{n}")).unwrap())
.collect::<Vec<_>>();
let candidates = repo
.list_catalog_candidates(&provider, &ids, &filter)
.await
.unwrap();
let ids = candidates
.iter()
.map(|a| a.id().as_str())
.collect::<std::collections::BTreeSet<_>>();
assert_eq!(
ids,
std::collections::BTreeSet::from(["acct_1", "acct_4", "acct_5"])
);
database.close().await;
}
@@ -3,12 +3,39 @@
use std::time::Duration;
use gateway_core::account::ProviderAccountId;
use gateway_core::provider_ports::{ProviderSessionAffinityKey, ProviderSessionAffinityPort};
use gateway_core::provider_ports::{
ProviderAccountAdmission, ProviderSessionAffinityKey, ProviderSessionAffinityPort,
ProviderSessionBinding, ProviderSessionClaim, ProviderStoreError,
};
use gateway_core::routing::ProviderKind;
use gateway_store::redis::RedisProviderSessionAffinityRepository;
use redis::aio::ConnectionManager;
use uuid::Uuid;
async fn bind(
repo: &RedisProviderSessionAffinityRepository,
provider: &ProviderKind,
key: &ProviderSessionAffinityKey,
expected: Option<&ProviderSessionBinding>,
account: &ProviderAccountId,
ttl: Duration,
) -> Result<Option<ProviderSessionBinding>, ProviderStoreError> {
let claim = ProviderSessionClaim {
key,
expected,
account,
ttl,
};
if repo
.admit_account(provider, ProviderAccountAdmission::Session(claim))
.await?
{
repo.load(provider, key).await
} else {
Ok(None)
}
}
#[tokio::test]
async fn concurrent_claims_admit_exactly_one_account() {
let Some((repo, mut connection, namespace)) = affinity_repository().await else {
@@ -19,8 +46,22 @@ async fn concurrent_claims_admit_exactly_one_account() {
let first = ProviderAccountId::new("acct_first").unwrap();
let second = ProviderAccountId::new("acct_second").unwrap();
let (a, b) = tokio::join!(
repo.compare_and_bind(&provider, &key, None, &first, Duration::from_secs(60)),
repo.compare_and_bind(&provider, &key, None, &second, Duration::from_secs(60))
bind(
&repo,
&provider,
&key,
None,
&first,
Duration::from_secs(60)
),
bind(
&repo,
&provider,
&key,
None,
&second,
Duration::from_secs(60)
)
);
let winners = [a.unwrap(), b.unwrap()]
.into_iter()
@@ -56,24 +97,21 @@ async fn stale_revision_cannot_overwrite_an_account_that_returned_to_the_session
let a = ProviderAccountId::new("acct_a").unwrap();
let b = ProviderAccountId::new("acct_b").unwrap();
let ttl = Duration::from_secs(60);
let first = repo
.compare_and_bind(&provider, &key, None, &a, ttl)
let first = bind(&repo, &provider, &key, None, &a, ttl)
.await
.unwrap()
.unwrap();
let second = repo
.compare_and_bind(&provider, &key, Some(&first), &b, ttl)
let second = bind(&repo, &provider, &key, Some(&first), &b, ttl)
.await
.unwrap()
.unwrap();
let third = repo
.compare_and_bind(&provider, &key, Some(&second), &a, ttl)
let third = bind(&repo, &provider, &key, Some(&second), &a, ttl)
.await
.unwrap()
.unwrap();
assert_ne!(first.revision(), third.revision());
assert!(
repo.compare_and_bind(&provider, &key, Some(&first), &b, ttl)
bind(&repo, &provider, &key, Some(&first), &b, ttl)
.await
.unwrap()
.is_none()
@@ -90,13 +128,11 @@ async fn renewal_preserves_revision_but_expiration_requires_fresh_claim() {
let key = ProviderSessionAffinityKey::try_new("expired-session").unwrap();
let a = ProviderAccountId::new("acct_a").unwrap();
let ttl = Duration::from_secs(60);
let first = repo
.compare_and_bind(&provider, &key, None, &a, ttl)
let first = bind(&repo, &provider, &key, None, &a, ttl)
.await
.unwrap()
.unwrap();
let renewed = repo
.compare_and_bind(&provider, &key, Some(&first), &a, ttl)
let renewed = bind(&repo, &provider, &key, Some(&first), &a, ttl)
.await
.unwrap()
.unwrap();
@@ -113,14 +149,13 @@ async fn renewal_preserves_revision_but_expiration_requires_fresh_claim() {
.await
.unwrap();
assert!(
repo.compare_and_bind(&provider, &key, Some(&first), &a, ttl)
bind(&repo, &provider, &key, Some(&first), &a, ttl)
.await
.unwrap()
.is_none()
);
assert!(repo.load(&provider, &key).await.unwrap().is_none());
let fresh = repo
.compare_and_bind(&provider, &key, None, &a, ttl)
let fresh = bind(&repo, &provider, &key, None, &a, ttl)
.await
.unwrap()
.unwrap();
@@ -136,7 +171,7 @@ async fn damaged_binding_is_an_error_and_cannot_be_claimed_as_absent() {
let key = ProviderSessionAffinityKey::try_new("damaged-session").unwrap();
let a = ProviderAccountId::new("acct_a").unwrap();
let ttl = Duration::from_secs(60);
repo.compare_and_bind(&provider, &key, None, &a, ttl)
bind(&repo, &provider, &key, None, &a, ttl)
.await
.unwrap()
.unwrap();
@@ -153,13 +188,140 @@ async fn damaged_binding_is_an_error_and_cannot_be_claimed_as_absent() {
.unwrap();
assert!(repo.load(&provider, &key).await.is_err());
assert!(
repo.compare_and_bind(&provider, &key, None, &a, ttl)
bind(&repo, &provider, &key, None, &a, ttl)
.await
.unwrap()
.is_none()
);
}
#[tokio::test]
async fn round_robin_commits_only_with_the_current_cursor_and_session_version() {
use gateway_core::policy::ClientApiKeyId;
let Some((repo, mut connection, namespace)) = affinity_repository().await else {
return;
};
let peer = RedisProviderSessionAffinityRepository::new(connection.clone(), &namespace).unwrap();
let provider = ProviderKind::new("openai").unwrap();
let client = ClientApiKeyId::new("client").unwrap();
let account = ProviderAccountId::new("acct_a").unwrap();
let key = ProviderSessionAffinityKey::try_new("rr-session").unwrap();
let other = ProviderSessionAffinityKey::try_new("rr-other").unwrap();
let claim = ProviderSessionClaim {
key: &key,
expected: None,
account: &account,
ttl: Duration::from_secs(60),
};
let step = |cursor, session| ProviderAccountAdmission::RoundRobin {
client_key: &client,
cursor,
session,
};
let first = step(0, Some(claim));
let (a, b) = tokio::join!(
repo.admit_account(&provider, first),
peer.admit_account(&provider, first)
);
assert_ne!(a.unwrap(), b.unwrap());
let winner = repo.load(&provider, &key).await.unwrap().unwrap();
assert_eq!(winner.account_id(), &account);
assert!(
!repo
.admit_account(&provider, step(1, Some(claim)))
.await
.unwrap()
);
assert_eq!(
peer.round_robin_cursor(&provider, &client).await.unwrap(),
1
);
let other_claim = ProviderSessionClaim {
key: &other,
..claim
};
assert!(
!peer
.admit_account(&provider, step(0, Some(other_claim)))
.await
.unwrap()
);
assert!(repo.load(&provider, &other).await.unwrap().is_none());
assert!(peer.admit_account(&provider, step(1, None)).await.unwrap());
assert_eq!(
repo.round_robin_cursor(&provider, &client).await.unwrap(),
2
);
for (p, c) in [
(&provider, ClientApiKeyId::new("other").unwrap()),
(&ProviderKind::new("xai").unwrap(), client.clone()),
] {
assert_eq!(repo.round_robin_cursor(p, &c).await.unwrap(), 0);
}
let leases =
gateway_store::redis::RedisCredentialLeaseRepository::new(connection.clone(), &namespace)
.unwrap();
assert_eq!(
leases
.advance_scheduling_cursor(&client, &provider)
.await
.unwrap(),
0
);
assert_eq!(
repo.round_robin_cursor(&provider, &client).await.unwrap(),
2
);
let cursor_keys: Vec<String> = redis::cmd("KEYS")
.arg(format!("{namespace}:*:session-cursor"))
.query_async(&mut connection)
.await
.unwrap();
let ttl: i64 = redis::cmd("PTTL")
.arg(&cursor_keys[0])
.query_async(&mut connection)
.await
.unwrap();
assert!((1..=86_400_000).contains(&ttl));
redis::cmd("DEL")
.arg(&cursor_keys)
.query_async::<i64>(&mut connection)
.await
.unwrap();
assert_eq!(
peer.round_robin_cursor(&provider, &client).await.unwrap(),
0
);
assert_eq!(repo.load(&provider, &key).await.unwrap(), Some(winner));
let fresh = ProviderSessionAffinityKey::try_new("rr-fresh").unwrap();
let claims = [
other_claim,
ProviderSessionClaim {
key: &fresh,
..claim
},
];
let (a, b) = tokio::join!(
repo.admit_account(&provider, step(0, Some(claims[0]))),
peer.admit_account(&provider, step(0, Some(claims[1])))
);
let won_first = a.unwrap();
assert_ne!(won_first, b.unwrap());
let loser = claims[usize::from(won_first)];
assert!(repo.load(&provider, loser.key).await.unwrap().is_none());
assert!(
repo.admit_account(&provider, step(1, Some(loser)))
.await
.unwrap()
);
assert_eq!(
peer.round_robin_cursor(&provider, &client).await.unwrap(),
2
);
}
async fn affinity_repository() -> Option<(
RedisProviderSessionAffinityRepository,
ConnectionManager,
@@ -221,12 +383,11 @@ async fn turn_alias_cannot_be_reassigned_and_follows_session_migration() {
.await
.unwrap()
);
let first = repo
.compare_and_bind(&provider, &session, None, &a, ttl)
let first = bind(&repo, &provider, &session, None, &a, ttl)
.await
.unwrap()
.unwrap();
repo.compare_and_bind(&provider, &session, Some(&first), &b, ttl)
bind(&repo, &provider, &session, Some(&first), &b, ttl)
.await
.unwrap()
.unwrap();
@@ -293,8 +454,7 @@ async fn binding_and_turn_alias_can_renew_beyond_one_day_and_shorten_without_sca
let mut binding = None;
for hours in [24, 168, 720, 1] {
let ttl = Duration::from_secs(hours * 3600);
binding = repo
.compare_and_bind(&provider, &session, binding.as_ref(), &account, ttl)
binding = bind(&repo, &provider, &session, binding.as_ref(), &account, ttl)
.await
.unwrap();
assert!(binding.is_some());
@@ -324,7 +484,7 @@ async fn binding_and_turn_alias_can_renew_beyond_one_day_and_shorten_without_sca
}
for ttl in [Duration::ZERO, Duration::from_secs(721 * 3600)] {
assert!(
repo.compare_and_bind(&provider, &session, binding.as_ref(), &account, ttl)
bind(&repo, &provider, &session, binding.as_ref(), &account, ttl)
.await
.is_err()
);
+1 -1
View File
@@ -79,7 +79,7 @@ openssl-sys = { version = "=0.9.116", features = ["vendored"] }
gateway-host = { path = "../../gateway-host" }
allocation-counter = "0.8.1"
tempfile = "3.27.0"
tokio = { version = "1.53.3", default-features = false, features = ["rt-multi-thread", "test-util"] }
tokio = { version = "1.53.3", default-features = false, features = ["process", "rt-multi-thread", "test-util"] }
tracing-subscriber = { version = "0.3.23", features = ["json"] }
wiremock = "0.6.5"
@@ -48,6 +48,7 @@ use gateway_core::provider_ports::{
ProviderStoreErrorKind,
};
use gateway_core::routing::{ProviderKind, UpstreamModelId};
use gateway_protocol::openai::OPENAI_PROTOCOL;
use secrecy::{ExposeSecret as _, SecretString};
use serde::Deserialize;
use serde_json::{Map, Number, Value};
@@ -297,7 +298,7 @@ impl ProviderAdmin for OpenAiAdminProvider {
return;
}
self.quota.invalidate_scheduling(account_ids);
if let Err(error) = self.catalog.invalidate() {
if let Err(error) = self.catalog.invalidate_accounts(account_ids).await {
tracing::warn!(
account_count = account_ids.len(),
error = %error,
@@ -1728,7 +1729,7 @@ fn build_connection_test_operation(
);
body.insert("stream".to_owned(), Value::Bool(true));
body.insert("store".to_owned(), Value::Bool(false));
let payload = ProtocolPayload::json_object("openai", body)
let payload = ProtocolPayload::json_object(OPENAI_PROTOCOL, body)
.map_err(|_| provider_admin_error(ProviderAdminErrorKind::Invalid))?;
Ok(Operation::Generate(GenerateRequest::from_protocol_payload(
payload,
@@ -1,8 +1,19 @@
//! 跨协议生成的 Messages fast 请求补充协议头,不改写原生 Messages 请求
//! Messages 跨协议转换的目标档位与 fast 协议头适配
use bytes::Bytes;
use gateway_core::engine::middleware::MiddlewareHeader;
use reqwest::header::HeaderMap;
use serde_json::{Map, Value};
pub(crate) fn normalize_translated_codex_service_tier(body: &mut Map<String, Value>) {
// 转换生成的默认档位沿用 Codex 客户端的省略表达,原生 Responses 保留调用方声明
if matches!(
body.get("service_tier").and_then(Value::as_str),
Some("auto" | "default")
) {
body.shift_remove("service_tier");
}
}
pub(crate) fn add_translated_fast_beta(headers: &mut Vec<MiddlewareHeader>, forwarded: &HeaderMap) {
const BETA: &[u8] = b"fast-mode-2026-02-01";
@@ -11,7 +11,7 @@ use gateway_core::engine::middleware::MiddlewareHeader;
use serde_json::{Map, Value};
pub(crate) use headers::is_non_codex_request_header;
pub(crate) use messages::add_translated_fast_beta;
pub(crate) use messages::{add_translated_fast_beta, normalize_translated_codex_service_tier};
pub(crate) use request::normalize_codex_request_body;
pub(crate) use search_tools::declares_hosted_search;
@@ -35,13 +35,6 @@ pub(super) fn normalize_non_codex_request_body(body: &mut Map<String, Value>) {
}
fn normalize_codex_oauth_options(body: &mut Map<String, Value>) {
// Codex Core 省略默认档位,OAuth 不接受公开 API 的 auto/default 字面值
if matches!(
body.get("service_tier").and_then(Value::as_str),
Some("auto" | "default")
) {
body.shift_remove("service_tier");
}
// Codex OAuth 拒绝这些客户端选项,即使采样参数取默认值
for field in [
"top_p",
File diff suppressed because it is too large Load Diff
@@ -0,0 +1,861 @@
//! OpenAI 账号模型目录的聚合、缓存与调度资格
mod refresh;
use std::collections::{BTreeMap, BTreeSet, btree_map::Entry};
use std::fmt;
use std::sync::{Arc, RwLock};
use std::time::{Duration, SystemTime};
use chrono::{DateTime, Utc};
use futures::{StreamExt, stream};
use gateway_core::account::AccountCatalogObservation;
use gateway_core::account::{
CredentialRevision, CredentialState, OpaqueProviderData, ProviderAccount, ProviderAccountId,
};
use gateway_core::provider_ports::{
ProviderCatalogCacheKey, ProviderCatalogCachePort, ProviderCatalogScope,
};
use gateway_core::routing::{
FrozenAccountScope, ModelPresentation, ProviderCatalogGeneration, ProviderKind,
ProviderModelContent, ProviderModelDescriptor, UpstreamModelId,
};
use secrecy::ExposeSecret;
use thiserror::Error;
use tokio::sync::Notify;
use uuid::Uuid;
use super::repository::{CodexCredentialRepository, CredentialRepositoryError};
use crate::transport::profile::CodexWireProfileState;
use crate::transport::{CodexBackendClient, CodexCatalogModel, CodexRequestContext};
const MAX_RESPONSE_ETAG_BYTES: usize = 256;
const ACCOUNT_CATALOG_CACHE_TTL: Duration = Duration::from_secs(5 * 60);
// 调度证据覆盖默认 15 分钟刷新周期及抖动,长期未刷新时恢复上游验证
const SCHEDULING_CATALOG_MAX_AGE: Duration = Duration::from_secs(30 * 60);
const MAX_ACCOUNT_CATALOG_MODELS: usize = 2_048;
const MODEL_CATALOG_TIMEOUT: Duration = Duration::from_secs(15);
const REFRESH_CONCURRENCY: usize = 32;
// 同时到期或导入大量账号时也限制起始速率,不能把连接并发上限当作请求速率上限
const REFRESH_SPACING: Duration = Duration::from_millis(10);
// 小账号池直接在内存复核,避免为少数账号增加数据库联表
const MIN_CATALOG_PUSHDOWN: usize = 128;
const PUBLICATION_BATCH: usize = 128;
/// 同一套餐也可能按账号开放模型,换凭据或套餐后必须重新取证
#[derive(Clone, Debug, PartialEq, Eq, PartialOrd, Ord)]
pub struct CodexCatalogScope {
account_id: ProviderAccountId,
revision: CredentialRevision,
plan: Option<String>,
}
impl CodexCatalogScope {
#[must_use]
pub fn for_account(account: &ProviderAccount) -> Self {
Self {
account_id: account.id().clone(),
revision: account.revision(),
plan: account
.plan_type()
.map(str::trim)
.filter(|plan| !plan.is_empty())
.map(str::to_ascii_lowercase),
}
}
/// Redis 作用域长度有上限,不直接拼接可变长度的账号与套餐文本
pub fn cache_scope(&self) -> Result<ProviderCatalogScope, CodexCredentialCatalogError> {
use sha2::{Digest, Sha256};
let mut digest = Sha256::new();
digest.update(self.account_id.as_str());
digest.update([0]);
digest.update(self.plan.as_deref().unwrap_or_default());
ProviderCatalogScope::new(format!(
"account:{}:{}",
hex::encode(digest.finalize()),
self.revision.get()
))
.map_err(|_| CodexCredentialCatalogError::Cache)
}
}
/// Redis TTL cache 中一条可重建的账号模型目录
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct CodexAccountCatalog {
scope: ProviderCatalogScope,
observed_at: DateTime<Utc>,
models: Vec<String>,
}
impl CodexAccountCatalog {
#[must_use]
pub const fn new(
scope: ProviderCatalogScope,
observed_at: DateTime<Utc>,
models: Vec<String>,
) -> Self {
Self {
scope,
observed_at,
models,
}
}
#[must_use]
pub const fn scope(&self) -> &ProviderCatalogScope {
&self.scope
}
#[must_use]
pub const fn observed_at(&self) -> DateTime<Utc> {
self.observed_at
}
#[must_use]
pub fn models(&self) -> &[String] {
&self.models
}
}
struct CatalogContent {
models: Vec<CodexCatalogModel>,
ids: Vec<String>,
digest: [u8; 32],
}
impl CatalogContent {
fn new(models: Vec<CodexCatalogModel>, document: &[u8]) -> Self {
use sha2::Digest;
Self {
digest: sha2::Sha256::digest(document).into(),
ids: model_ids(&models),
models,
}
}
}
#[derive(Clone)]
struct AccountEntry {
observed_at: SystemTime,
content: Arc<CatalogContent>,
invalidated: Arc<std::sync::atomic::AtomicBool>,
etag: Option<String>,
}
type Entries = BTreeMap<CodexCatalogScope, AccountEntry>;
pub struct CodexCredentialCatalogSnapshot {
models: Vec<CodexCatalogModel>,
entries: Entries,
sources: Vec<(Arc<CatalogContent>, Vec<ProviderAccountId>)>,
}
impl CodexCredentialCatalogSnapshot {
fn new(entries: Entries) -> Self {
let mut models: Vec<CodexCatalogModel> = Vec::new();
let mut indices = BTreeMap::<&str, usize>::new();
let mut sources = BTreeMap::<_, (Arc<CatalogContent>, Vec<ProviderAccountId>)>::new();
for (scope, entry) in &entries {
let (_, members) = sources
.entry(entry.content.digest)
.or_insert_with(|| (entry.content.clone(), Vec::new()));
members.push(scope.account_id.clone());
if members.len() > 1 {
continue;
}
for model in &entry.content.models {
match indices.entry(model.request_model().as_str()) {
Entry::Vacant(slot) => {
slot.insert(models.len());
models.push(model.clone());
}
Entry::Occupied(slot)
if model.document().protocol() == "codex"
&& models[*slot.get()].document().protocol() != "codex" =>
{
models[*slot.get()] = model.clone()
}
_ => {}
}
}
}
Self {
models,
entries,
sources: sources.into_values().collect(),
}
}
#[must_use]
pub fn models(&self) -> &[CodexCatalogModel] {
&self.models
}
pub(super) fn excludes(&self, account: &ProviderAccount, model: &str, now: SystemTime) -> bool {
self.entries
.get(&CodexCatalogScope::for_account(account))
.is_some_and(|e| {
!e.invalidated.load(std::sync::atomic::Ordering::Relaxed)
&& now
.duration_since(e.observed_at)
.is_ok_and(|age| age < SCHEDULING_CATALOG_MAX_AGE)
&& !e.content.ids.iter().any(|id| id == model)
})
}
pub fn account_models(&self, account: &ProviderAccount) -> Option<&[String]> {
self.entries
.get(&CodexCatalogScope::for_account(account))
.map(|e| e.content.ids.as_slice())
}
#[must_use]
pub fn model_catalog_accounts(&self) -> BTreeMap<&str, BTreeSet<ProviderAccountId>> {
let mut accounts = BTreeMap::<&str, BTreeSet<ProviderAccountId>>::new();
for (scope, entry) in &self.entries {
for model in &entry.content.ids {
accounts
.entry(model.as_str())
.or_default()
.insert(scope.account_id.clone());
}
}
accounts
}
}
impl fmt::Debug for CodexCredentialCatalogSnapshot {
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
formatter
.debug_struct("CodexCredentialCatalogSnapshot")
.field("model_count", &self.models.len())
.field("scope_count", &self.entries.len())
.finish()
}
}
#[derive(Debug, Clone, Error)]
pub enum CodexCredentialCatalogError {
#[error("Codex model catalog has no eligible account")]
NoEligibleCredential,
#[error("Codex model catalog account data is invalid")]
InvalidCredentialData,
#[error("Codex model catalog upstream query failed: {detail}")]
Upstream { detail: String },
#[error("Codex model catalog cache is unavailable")]
Cache,
#[error("Codex model catalog ETag is invalid")]
InvalidEtag,
#[error("Codex model catalog changed during refresh")]
ConcurrentUpdate,
}
#[derive(Default)]
struct CatalogCacheState {
restored: bool,
revision: u64,
generation: u64,
snapshot: Option<Arc<CodexCredentialCatalogSnapshot>>,
pending: BTreeMap<CodexCatalogScope, String>,
}
struct FetchedAccountModels {
models: Vec<CodexCatalogModel>,
document: bytes::Bytes,
etag: Option<String>,
}
impl From<CredentialRepositoryError> for CodexCredentialCatalogError {
fn from(_: CredentialRepositoryError) -> Self {
Self::InvalidCredentialData
}
}
#[derive(Clone)]
pub struct CodexCredentialCatalogService {
repository: CodexCredentialRepository,
profile: CodexWireProfileState,
http: reqwest::Client,
base_url: String,
account_catalog_cache: Arc<dyn ProviderCatalogCachePort>,
cache: Arc<RwLock<CatalogCacheState>>,
refresh_notification: Arc<Notify>,
refresh_lock: Arc<tokio::sync::Mutex<()>>,
publication: Arc<std::sync::Mutex<()>>,
}
impl CodexCredentialCatalogService {
pub fn new(
repository: CodexCredentialRepository,
profile: CodexWireProfileState,
http: reqwest::Client,
base_url: String,
account_catalog_cache: Arc<dyn ProviderCatalogCachePort>,
) -> Self {
Self {
repository,
profile,
http,
base_url,
account_catalog_cache,
cache: Arc::new(RwLock::new(CatalogCacheState::default())),
refresh_notification: Arc::new(Notify::new()),
refresh_lock: Arc::new(tokio::sync::Mutex::new(())),
publication: Default::default(),
}
}
/// 原生条目取自后台发现的账号目录;用当前账号版本和冻结的 Key 政策双重过滤
pub async fn client_model_catalog(
&self,
scope: &FrozenAccountScope,
) -> Result<Vec<ProviderModelDescriptor>, CodexCredentialCatalogError> {
let snapshot = self
.cached()?
.ok_or(CodexCredentialCatalogError::NoEligibleCredential)?;
// 正文相同的账号只需核对每个模型在当前政策下的首选来源
let mut candidates = BTreeSet::new();
for (content, members) in &snapshot.sources {
if let Some(id) = members.iter().find(|id| scope.allows(id)) {
candidates.insert(id.clone());
} else {
continue;
}
for id in &content.ids {
if let Some(account) = members
.iter()
.find(|account| scope.allows_model(account, id))
{
candidates.insert(account.clone());
}
}
}
let ids = candidates.into_iter().collect::<Vec<_>>();
let mut accounts = self.repository.list_for_accounts(&ids).await?;
let now = SystemTime::now();
// 代表账号刚被停用或轮换时,仅这一次扩大本地资格查询,仍不查询上游
if accounts.len() != ids.len()
|| accounts.iter().any(|a| {
!eligible_catalog_account(a, now)
|| !snapshot
.entries
.contains_key(&CodexCatalogScope::for_account(a))
})
{
let ids = snapshot
.entries
.keys()
.filter(|s| scope.allows(&s.account_id))
.map(|s| s.account_id.clone())
.collect::<Vec<_>>();
accounts = self.repository.list_for_accounts(&ids).await?;
}
let mut selected = BTreeMap::new();
let mut available = false;
for account in accounts.iter().filter(|a| eligible_catalog_account(a, now)) {
let Some(entry) = snapshot
.entries
.get(&CodexCatalogScope::for_account(account))
else {
continue;
};
available = true;
let api = account.authentication_kind() == super::CODEX_AUTHENTICATION_KIND_API_KEY;
for (position, model) in entry.content.models.iter().enumerate() {
if !scope.allows_model(account.id(), model.request_model().as_str()) {
continue;
}
let order = (
api,
model.document().protocol() != "codex",
account.id(),
position,
);
let slot = selected
.entry(model.request_model())
.or_insert((order, model));
if order < slot.0 {
*slot = (order, model);
}
}
}
if !available {
return Err(CodexCredentialCatalogError::NoEligibleCredential);
}
let mut models: Vec<_> = selected.into_values().collect();
models.sort_unstable_by_key(|(order, _)| *order);
Ok(models
.into_iter()
.map(|(_, model)| {
let id = model.request_model().as_str();
ProviderModelDescriptor {
model: model.request_model().clone(),
content: if model.document().protocol() == "codex" {
ProviderModelContent::Native(model.document().clone())
} else {
ProviderModelContent::Adapted(
ModelPresentation::new(Some(id.to_owned()), None)
.with_agent_tools(true, false)
.with_hidden(is_specialized_api_model(id)),
)
},
}
})
.collect())
}
pub(crate) fn messages_model_capabilities(
&self,
account: &ProviderAccount,
model: &UpstreamModelId,
) -> Option<gateway_protocol::anthropic::MessagesModelCapabilities> {
let snapshot = self.cached().ok()??;
snapshot
.entries
.get(&CodexCatalogScope::for_account(account))?
.content
.models
.iter()
.find(|entry| entry.request_model() == model)?
.messages_capabilities()
.cloned()
}
#[must_use]
pub fn catalog_generation(&self) -> ProviderCatalogGeneration {
let cache = self
.cache
.read()
.unwrap_or_else(std::sync::PoisonError::into_inner);
ProviderCatalogGeneration::new(cache.generation)
}
pub fn cached(
&self,
) -> Result<Option<Arc<CodexCredentialCatalogSnapshot>>, CodexCredentialCatalogError> {
Ok(self
.cache
.read()
.map_err(|_| CodexCredentialCatalogError::Cache)?
.snapshot
.as_ref()
.cloned())
}
/// 读取账号当前凭据对应的模型目录;不触发网络
pub fn cached_account_models(
&self,
account: &ProviderAccount,
) -> Result<Option<Vec<String>>, CodexCredentialCatalogError> {
let Some(snapshot) = self.cached()? else {
return Ok(None);
};
Ok(snapshot.account_models(account).map(<[String]>::to_vec))
}
/// 大池中相同正文只传递一个摘要;存储只跳过有当前账号证据的来源
pub(super) fn omissions(
&self,
model: &str,
ids: &[ProviderAccountId],
exempt: impl Fn(&ProviderAccountId) -> bool,
) -> Option<gateway_core::account::AccountCatalogOmissions> {
if ids.len() < MIN_CATALOG_PUSHDOWN {
return None;
}
let cache = self.cache.read().ok()?;
let snapshot = cache.snapshot.as_ref()?.clone();
let pending = cache
.pending
.keys()
.map(|s| s.account_id.clone())
.collect::<BTreeSet<_>>();
drop(cache);
let documents = snapshot
.sources
.iter()
.filter(|(c, _)| !c.ids.iter().any(|id| id == model))
.map(|(c, _)| c.digest)
.collect::<Vec<_>>();
if documents.is_empty() {
return None;
}
let as_of = SystemTime::now();
Some(gateway_core::account::AccountCatalogOmissions {
documents,
exceptions: ids
.iter()
.filter(|id| pending.contains(*id) || exempt(id))
.cloned()
.collect(),
observed_after: as_of.checked_sub(SCHEDULING_CATALOG_MAX_AGE)?,
as_of,
})
}
/// 优先读取账号目录缓存;缺失时只查询该账号
pub async fn cached_or_refresh_account_catalog(
&self,
account: &ProviderAccount,
) -> Result<CodexAccountCatalog, CodexCredentialCatalogError> {
if let Some(catalog) = self.read_account_catalog(account).await? {
return Ok(catalog);
}
self.refresh_account_catalog(account.id()).await
}
/// 实时刷新指定账号的模型集合,包括管理员显式查询的停用账号
pub async fn refresh_account_catalog(
&self,
account_id: &ProviderAccountId,
) -> Result<CodexAccountCatalog, CodexCredentialCatalogError> {
let account = self
.repository
.store()
.get_account(account_id)
.await
.map_err(|_| CodexCredentialCatalogError::InvalidCredentialData)?
.filter(|account| account.provider().as_str() == "openai")
.ok_or(CodexCredentialCatalogError::NoEligibleCredential)?;
let scope = CodexCatalogScope::for_account(&account);
let client = CodexBackendClient::new(
self.http.clone(),
self.base_url.clone(),
self.profile.clone(),
);
let revision = self.state_revision()?;
let (catalog, entry) = self
.fetch_observation(&client, &account, &scope, None)
.await?;
if eligible_catalog_account(&account, SystemTime::now())
&& let Err(error) = self.merge_observations(revision, vec![(scope, entry)])
{
if !matches!(error, CodexCredentialCatalogError::ConcurrentUpdate) {
return Err(error);
}
self.schedule_refresh();
}
Ok(catalog)
}
async fn fetch_observation(
&self,
client: &CodexBackendClient,
account: &ProviderAccount,
scope: &CodexCatalogScope,
etag: Option<String>,
) -> Result<(CodexAccountCatalog, AccountEntry), CodexCredentialCatalogError> {
let observed_at = SystemTime::now();
let fetched = self.fetch_account_models(client, account).await?;
let observed = AccountCatalogObservation {
account_id: account.id().clone(),
revision: account.revision(),
plan: account.plan_type().map(str::to_owned),
observed_at,
etag: etag.clone().or(fetched.etag),
};
if !self
.repository
.store()
.save_account_catalog(&observed, &fetched.document)
.await
.map_err(|_| CodexCredentialCatalogError::Cache)?
{
return Err(CodexCredentialCatalogError::ConcurrentUpdate);
}
let content = Arc::new(CatalogContent::new(fetched.models, &fetched.document));
// Redis 仅为管理端短期缓存,失败不能撤销已持久化的目录
let catalog = CodexAccountCatalog::new(
scope.cache_scope()?,
observed_at.into(),
content.ids.clone(),
);
if let Err(error) = self.replace_account_catalog(&catalog).await {
tracing::warn!(%error,"OpenAI account catalog cache write failed");
}
Ok((
catalog,
AccountEntry {
observed_at,
content,
etag: observed.etag,
invalidated: Default::default(),
},
))
}
/// 读取目标账号的原生目录条目;供管理端按账号导出目录文件使用
///
/// 按账号直接请求上游,避免其他账号的同名模型替换目标账号的原生对象;
/// 停用账号也不在常规快照候选中,仍允许管理员显式导出
/// 返回顺序沿用上游目录顺序,条目内容保持上游原样,不做别名改写或字段裁剪
pub async fn account_catalog_documents(
&self,
account: &ProviderAccount,
) -> Result<(Vec<CodexCatalogModel>, SystemTime), CodexCredentialCatalogError> {
let client = CodexBackendClient::new(
self.http.clone(),
self.base_url.clone(),
self.profile.clone(),
);
let fetched = self.fetch_account_models(&client, account).await?;
if fetched.models.is_empty() {
return Err(CodexCredentialCatalogError::NoEligibleCredential);
}
Ok((fetched.models, SystemTime::now()))
}
/// 读取当前账号的目录缓存,不触发上游请求
pub async fn read_account_catalog(
&self,
account: &ProviderAccount,
) -> Result<Option<CodexAccountCatalog>, CodexCredentialCatalogError> {
let scope = CodexCatalogScope::for_account(account);
let key = scope.cache_scope()?;
if let Some(snapshot) = self.cached()?
&& let Some(entry) = snapshot.entries.get(&scope)
{
return Ok(Some(CodexAccountCatalog::new(
key,
entry.observed_at.into(),
entry.content.ids.clone(),
)));
}
self.read_cached_catalog(&key).await
}
async fn replace_account_catalog(
&self,
catalog: &CodexAccountCatalog,
) -> Result<(), CodexCredentialCatalogError> {
let key = self.account_catalog_key(catalog.scope())?;
let document = encode_account_catalog(catalog);
self.account_catalog_cache
.replace(&key, &document, ACCOUNT_CATALOG_CACHE_TTL)
.await
.map_err(|_| CodexCredentialCatalogError::Cache)
}
async fn read_cached_catalog(
&self,
scope: &ProviderCatalogScope,
) -> Result<Option<CodexAccountCatalog>, CodexCredentialCatalogError> {
let key = self.account_catalog_key(scope)?;
self.account_catalog_cache
.read(&key)
.await
.map_err(|_| CodexCredentialCatalogError::Cache)?
.map(|document| decode_account_catalog(scope, document))
.transpose()
}
fn account_catalog_key(
&self,
scope: &ProviderCatalogScope,
) -> Result<ProviderCatalogCacheKey, CodexCredentialCatalogError> {
let provider_kind =
ProviderKind::new("openai").map_err(|_| CodexCredentialCatalogError::Cache)?;
Ok(ProviderCatalogCacheKey::new(provider_kind, scope.clone()))
}
async fn fetch_account_models(
&self,
client: &CodexBackendClient,
account: &ProviderAccount,
) -> Result<FetchedAccountModels, CodexCredentialCatalogError> {
tokio::time::timeout(MODEL_CATALOG_TIMEOUT, async {
let credential = self
.repository
.load_runtime_credential(account)
.await
.map_err(|_| CodexCredentialCatalogError::InvalidCredentialData)?;
let request_id = format!("catalog_{}", Uuid::now_v7().simple());
let authorization = credential
.authentication
.authorization_header()
.map_err(|_| CodexCredentialCatalogError::InvalidCredentialData)?;
let mut client = client
.for_account(account)
.map_err(|error| CodexCredentialCatalogError::Upstream {
detail: error.to_string(),
})?
.with_responses_api_base_url(credential.authentication.responses_api_base_url());
if let super::CodexRuntimeAuthentication::ApiKey(auth) = &credential.authentication
&& auth.configuration.protocol == super::ApiKeyProtocol::Messages
{
client = client.with_messages_api_base_url(&auth.configuration.base_url);
}
let result = client
.fetch_models_with_context(
CodexRequestContext::auxiliary(
authorization.expose_secret(),
account.upstream_account_id(),
&request_id,
None,
),
None,
)
.await;
let snapshot = result.map_err(|error| CodexCredentialCatalogError::Upstream {
detail: error.to_string(),
})?;
let (models, document, etag) = snapshot.into_parts();
Ok(FetchedAccountModels {
models,
document,
etag,
})
})
.await
.map_err(|_| CodexCredentialCatalogError::Upstream {
detail: "model catalog timed out".to_owned(),
})?
}
}
fn is_specialized_api_model(model: &str) -> bool {
// 普通目录不提供 Codex 可见性;仅按已知的专用模型 ID 隐藏,保留未知 ID 和直接调用
// 依据 https://developers.openai.com/api/docs/models 与 Codex 的语音通话模型
matches!(
model,
"dall-e-2"
| "dall-e-3"
| "chatgpt-image-latest"
| "gpt-image-1"
| "gpt-image-1-mini"
| "gpt-image-1.5"
| "gpt-image-1.5-2025-12-16"
| "gpt-image-2"
| "gpt-image-2-2026-04-21"
| "gpt-image-2.5-flare"
| "gpt-image-2.5-flare-2026-09-08"
| "gpt-image-2.5-sunburst"
| "gpt-image-2.5-sunburst-2026-09-08"
| "gpt-image-2.6-flare"
| "gpt-image-2.6-sunburst"
| "gpt-live-1"
| "gpt-live-1-codex"
| "gpt-live-transcribe"
| "gpt-realtime"
| "gpt-realtime-2025-08-28"
| "gpt-realtime-1.5"
| "gpt-realtime-2"
| "gpt-realtime-2.1"
| "gpt-realtime-2.1-mini"
| "gpt-realtime-mini"
| "gpt-realtime-mini-2025-10-06"
| "gpt-realtime-mini-2025-12-15"
| "gpt-realtime-translate"
| "gpt-realtime-whisper"
| "gpt-4o-realtime-preview"
| "gpt-4o-realtime-preview-2024-10-01"
| "gpt-4o-realtime-preview-2024-12-17"
| "gpt-4o-realtime-preview-2025-06-03"
| "gpt-4o-mini-realtime-preview"
| "gpt-4o-mini-realtime-preview-2024-12-17"
| "tts-1"
| "tts-1-hd"
| "gpt-4o-mini-tts"
| "gpt-4o-mini-tts-2025-03-20"
| "gpt-4o-mini-tts-2025-12-15"
)
}
fn model_ids(models: &[CodexCatalogModel]) -> Vec<String> {
models
.iter()
.map(|model| model.request_model().as_str().to_owned())
.collect()
}
fn encode_account_catalog(catalog: &CodexAccountCatalog) -> OpaqueProviderData {
let mut document = serde_json::Map::new();
document.insert("version".to_owned(), serde_json::Value::from(1));
document.insert(
"scope".to_owned(),
serde_json::Value::String(catalog.scope().as_str().to_owned()),
);
document.insert(
"observedAt".to_owned(),
serde_json::Value::String(catalog.observed_at().to_rfc3339()),
);
document.insert(
"models".to_owned(),
serde_json::Value::Array(
catalog
.models()
.iter()
.cloned()
.map(serde_json::Value::String)
.collect(),
),
);
OpaqueProviderData::new(document)
}
fn decode_account_catalog(
scope: &ProviderCatalogScope,
document: OpaqueProviderData,
) -> Result<CodexAccountCatalog, CodexCredentialCatalogError> {
let mut fields = document.into_inner();
if fields.remove("version").and_then(|value| value.as_u64()) != Some(1) {
return Err(CodexCredentialCatalogError::Cache);
}
if fields
.remove("scope")
.and_then(|value| value.as_str().map(ToOwned::to_owned))
.as_deref()
!= Some(scope.as_str())
{
return Err(CodexCredentialCatalogError::Cache);
}
let observed_at = fields
.remove("observedAt")
.and_then(|value| value.as_str().map(ToOwned::to_owned))
.and_then(|value| DateTime::parse_from_rfc3339(&value).ok())
.map(|value| value.with_timezone(&Utc))
.ok_or(CodexCredentialCatalogError::Cache)?;
let values = fields
.remove("models")
.and_then(|value| value.as_array().cloned())
.filter(|models| !models.is_empty() && models.len() <= MAX_ACCOUNT_CATALOG_MODELS)
.ok_or(CodexCredentialCatalogError::Cache)?;
if !fields.is_empty() {
return Err(CodexCredentialCatalogError::Cache);
}
let mut seen = BTreeSet::new();
let mut models = Vec::with_capacity(values.len());
for value in values {
let value = value
.as_str()
.map(ToOwned::to_owned)
.ok_or(CodexCredentialCatalogError::Cache)?;
let model = UpstreamModelId::new(value).map_err(|_| CodexCredentialCatalogError::Cache)?;
if !seen.insert(model.as_str().to_owned()) {
return Err(CodexCredentialCatalogError::Cache);
}
models.push(model.as_str().to_owned());
}
Ok(CodexAccountCatalog::new(scope.clone(), observed_at, models))
}
fn validate_response_etag(etag: &str) -> Result<(), CodexCredentialCatalogError> {
if etag.is_empty() || etag.len() > MAX_RESPONSE_ETAG_BYTES || etag.chars().any(char::is_control)
{
return Err(CodexCredentialCatalogError::InvalidEtag);
}
Ok(())
}
fn eligible_catalog_account(account: &ProviderAccount, now: SystemTime) -> bool {
account.enabled()
&& account
.access_token_expires_at()
.is_none_or(|expires_at| expires_at > now)
&& matches!(
account.credential_state(),
CredentialState::Unknown | CredentialState::Ready
)
}
@@ -0,0 +1,366 @@
//! 账号目录的批量恢复、变更通知和限速刷新
use super::*;
impl CodexCredentialCatalogService {
/// 启动只恢复已成功落库、仍匹配账号身份的目录
pub async fn restore(&self) -> Result<(), CodexCredentialCatalogError> {
let revision = self.state_revision()?;
let provider =
ProviderKind::new("openai").map_err(|_| CodexCredentialCatalogError::Cache)?;
let groups = self
.repository
.store()
.load_account_catalogs(&provider)
.await
.map_err(|_| CodexCredentialCatalogError::Cache)?;
let mut entries = Entries::new();
for group in groups {
let decoded =
match crate::transport::CodexModelCatalogSnapshot::restore(&group.document) {
Ok(decoded) => decoded,
Err(error) => {
tracing::warn!(%error,"OpenAI stored model catalog is invalid");
continue;
}
};
let (models, _, _) = decoded.into_parts();
let content = Arc::new(CatalogContent::new(models, &group.document));
for a in group.accounts {
let scope = CodexCatalogScope {
account_id: a.account_id,
revision: a.revision,
plan: a
.plan
.map(|p| p.trim().to_ascii_lowercase())
.filter(|p| !p.is_empty()),
};
entries.insert(
scope,
AccountEntry {
observed_at: a.observed_at,
content: content.clone(),
etag: a.etag,
invalidated: Default::default(),
},
);
}
}
let _writer = self
.publication
.lock()
.map_err(|_| CodexCredentialCatalogError::Cache)?;
self.publish(revision, entries)?;
self.cache
.write()
.map_err(|_| CodexCredentialCatalogError::Cache)?
.restored = true;
Ok(())
}
pub(super) fn state_revision(&self) -> Result<u64, CodexCredentialCatalogError> {
self.cache
.read()
.map(|c| c.revision)
.map_err(|_| CodexCredentialCatalogError::Cache)
}
pub async fn invalidate_accounts(
&self,
accounts: &[ProviderAccountId],
) -> Result<(), CodexCredentialCatalogError> {
let ids = accounts.iter().collect::<BTreeSet<_>>();
let current = self
.repository
.list_for_accounts(accounts)
.await
.inspect_err(|_| self.schedule_refresh())?;
let current = current
.iter()
.filter(|a| eligible_catalog_account(a, SystemTime::now()))
.map(CodexCatalogScope::for_account)
.collect::<BTreeSet<_>>();
let _writer = self
.publication
.lock()
.map_err(|_| CodexCredentialCatalogError::Cache)?;
let revision = {
let mut cache = self
.cache
.write()
.map_err(|_| CodexCredentialCatalogError::Cache)?;
cache.revision = cache.revision.wrapping_add(1);
cache
.pending
.retain(|scope, _| !ids.contains(&scope.account_id) || current.contains(scope));
cache.revision
};
if let Some(snapshot) = self.cached()? {
let keep = |scope: &CodexCatalogScope| {
!ids.contains(&scope.account_id) || current.contains(scope)
};
if snapshot.entries.keys().any(|s| !keep(s)) {
let mut entries = snapshot.entries.clone();
entries.retain(|s, _| keep(s));
self.publish(revision, entries)?;
}
}
self.schedule_refresh();
Ok(())
}
pub fn observe_response_etag(
&self,
account: &ProviderAccount,
etag: &str,
) -> Result<bool, CodexCredentialCatalogError> {
validate_response_etag(etag)?;
let scope = CodexCatalogScope::for_account(account);
let mut state = self
.cache
.write()
.map_err(|_| CodexCredentialCatalogError::Cache)?;
if state.pending.get(&scope).map(String::as_str) == Some(etag)
|| !state.pending.contains_key(&scope)
&& state
.snapshot
.as_ref()
.and_then(|s| s.entries.get(&scope))
.and_then(|e| e.etag.as_deref())
== Some(etag)
{
return Ok(false);
}
if let Some(entry) = state.snapshot.as_ref().and_then(|s| s.entries.get(&scope)) {
entry
.invalidated
.store(true, std::sync::atomic::Ordering::Relaxed);
}
state.pending.insert(scope, etag.to_owned());
drop(state);
self.refresh_notification.notify_one();
Ok(true)
}
pub(crate) fn schedule_refresh(&self) {
self.refresh_notification.notify_one();
}
pub async fn wait_for_refresh(&self) {
self.refresh_notification.notified().await;
}
/// 只合并成功的账号;手动刷新和后台任务同时完成时,保留较新的观测
pub(super) fn merge_observations(
&self,
revision: u64,
batch: Vec<(CodexCatalogScope, AccountEntry)>,
) -> Result<(), CodexCredentialCatalogError> {
let _writer = self
.publication
.lock()
.map_err(|_| CodexCredentialCatalogError::Cache)?;
if self.state_revision()? != revision {
return Err(CodexCredentialCatalogError::ConcurrentUpdate);
}
let mut entries = self
.cached()?
.map(|snapshot| snapshot.entries.clone())
.unwrap_or_default();
for (scope, entry) in batch {
if entries
.get(&scope)
.is_none_or(|old| old.observed_at < entry.observed_at)
{
entries.insert(scope, entry);
}
}
self.publish(revision, entries)?;
Ok(())
}
fn publish(
&self,
revision: u64,
entries: Entries,
) -> Result<Option<Arc<CodexCredentialCatalogSnapshot>>, CodexCredentialCatalogError> {
let next =
(!entries.is_empty()).then(|| Arc::new(CodexCredentialCatalogSnapshot::new(entries)));
let old = self.cached()?;
let changed = match (&old, &next) {
(Some(a), Some(b)) => {
a.entries.len() != b.entries.len()
|| a.entries.iter().zip(&b.entries).any(|((sa, a), (sb, b))| {
sa != sb
|| (!Arc::ptr_eq(&a.content, &b.content)
&& a.content.digest != b.content.digest)
})
}
_ => old.is_some() != next.is_some(),
};
let mut cache = self
.cache
.write()
.map_err(|_| CodexCredentialCatalogError::Cache)?;
if cache.revision != revision {
return Err(CodexCredentialCatalogError::ConcurrentUpdate);
}
if let Some(snapshot) = &next {
cache.pending.retain(|scope, tag| {
let Some(entry) = snapshot.entries.get(scope) else {
return true;
};
if entry.etag.as_deref() == Some(tag)
&& !entry.invalidated.load(std::sync::atomic::Ordering::Relaxed)
{
return false;
}
entry
.invalidated
.store(true, std::sync::atomic::Ordering::Relaxed);
true
});
}
cache.snapshot = next.clone();
if changed {
cache.generation = cache.generation.wrapping_add(1);
}
Ok(next)
}
/// 一次对账只探测缺失、过期或有新 ETag 的账号;成功逐批发布,失败不更新观测时间
pub async fn refresh_catalogs(
&self,
) -> Result<Arc<CodexCredentialCatalogSnapshot>, CodexCredentialCatalogError> {
let _writer = self.refresh_lock.lock().await;
// 启动时数据库暂不可用,应先重试恢复,不能把整池误判为没有历史
if !self
.cache
.read()
.map_err(|_| CodexCredentialCatalogError::Cache)?
.restored
{
self.restore().await?;
}
let revision = self.state_revision()?;
let now = SystemTime::now();
let accounts = self.repository.list_for_provider().await?;
let current: BTreeMap<_, _> = accounts
.into_iter()
.filter(|a| eligible_catalog_account(a, now))
.map(|a| (CodexCatalogScope::for_account(&a), a))
.collect();
let pending = {
let mut cache = self
.cache
.write()
.map_err(|_| CodexCredentialCatalogError::Cache)?;
cache.pending.retain(|scope, _| current.contains_key(scope));
cache.pending.clone()
};
let previous = self.cached()?;
let mut entries = previous
.as_ref()
.map(|s| s.entries.clone())
.unwrap_or_default();
entries.retain(|scope, _| current.contains_key(scope));
if previous
.as_ref()
.is_some_and(|s| s.entries.len() != entries.len())
{
let _writer = self
.publication
.lock()
.map_err(|_| CodexCredentialCatalogError::Cache)?;
let mut latest = self
.cached()?
.map(|s| s.entries.clone())
.unwrap_or_default();
latest.retain(|scope, _| current.contains_key(scope));
self.publish(revision, latest)?;
}
let mut contents: BTreeMap<_, _> = entries
.values()
.map(|e| (e.content.digest, e.content.clone()))
.collect();
let due = current
.into_iter()
.filter(|(scope, _)| {
pending.contains_key(scope)
|| !entries.get(scope).is_some_and(|e| {
now.duration_since(e.observed_at)
.is_ok_and(|age| age < SCHEDULING_CATALOG_MAX_AGE)
})
})
.collect::<Vec<_>>();
if due.is_empty() {
return self
.cached()?
.ok_or(CodexCredentialCatalogError::NoEligibleCredential);
}
let client = CodexBackendClient::new(
self.http.clone(),
self.base_url.clone(),
self.profile.clone(),
);
let mut pace = tokio::time::interval(REFRESH_SPACING);
pace.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Delay);
let starts = stream::unfold(pace, |mut clock| async move {
clock.tick().await;
Some(((), clock))
});
let fetches = stream::iter(due)
.zip(starts)
.map(|((scope, account), ())| {
let client = &client;
let etag = pending.get(&scope).cloned();
async move {
let (_, entry) = self
.fetch_observation(client, &account, &scope, etag.clone())
.await?;
Ok((scope, entry))
}
})
.buffer_unordered(REFRESH_CONCURRENCY);
tokio::pin!(fetches);
let mut batch = Vec::with_capacity(PUBLICATION_BATCH);
let mut completed = 0_usize;
let mut failures = 0_usize;
let mut last_error = None;
while let Some(result) = fetches.next().await {
match result {
Ok((scope, mut entry)) => {
entry.content = contents
.entry(entry.content.digest)
.or_insert_with(|| entry.content.clone())
.clone();
batch.push((scope, entry));
completed += 1;
if completed == 1 || completed.is_multiple_of(PUBLICATION_BATCH) {
self.merge_observations(revision, std::mem::take(&mut batch))?;
}
}
Err(error) => {
failures += 1;
last_error = Some(error);
}
}
}
if !batch.is_empty() {
self.merge_observations(revision, batch)?;
}
if let Some(error) = &last_error {
tracing::warn!(successes=completed,failures,%error,"OpenAI account catalog refresh incomplete");
}
if pending
.keys()
.any(|k| self.cache.read().is_ok_and(|c| c.pending.contains_key(k)))
&& let Some(error) = last_error.clone()
{
self.refresh_notification.notify_one();
return Err(error);
}
self.cached()?
.ok_or(last_error.unwrap_or(CodexCredentialCatalogError::NoEligibleCredential))
}
}
@@ -21,11 +21,11 @@ use gateway_core::concurrency::{
};
use gateway_core::engine::{AttemptContext, ContinuationAttempt, policy::AccountPolicyError};
use gateway_core::provider_ports::{
ProviderConcurrencyPool, ProviderCooldownPort, ProviderCooldownScope, ProviderLeaseAcquisition,
ProviderLeaseGuard, ProviderLeasePort, ProviderLeaseRequest, ProviderSchedulingLeaseRequest,
ProviderScopedCooldown, ProviderSessionAffinityKey, ProviderSessionAffinityPort,
ProviderSessionBinding, ProviderSessionExclusionPort, ProviderSessionExclusions,
ProviderStoreError,
ProviderAccountAdmission, ProviderConcurrencyPool, ProviderCooldownPort, ProviderCooldownScope,
ProviderLeaseAcquisition, ProviderLeaseGuard, ProviderLeasePort, ProviderLeaseRequest,
ProviderSchedulingLeaseRequest, ProviderScopedCooldown, ProviderSessionAffinityKey,
ProviderSessionAffinityPort, ProviderSessionBinding, ProviderSessionClaim,
ProviderSessionExclusionPort, ProviderSessionExclusions, ProviderStoreError,
};
use gateway_core::routing::ProviderKind;
use secrecy::ExposeSecret;
@@ -145,7 +145,7 @@ pub struct CodexCredentialSelector {
quota: Arc<CodexCredentialQuotaService>,
catalog: Arc<CodexCredentialCatalogService>,
cooldowns: Arc<dyn ProviderCooldownPort>,
upstream_info: Option<Arc<dyn gateway_core::account::UpstreamInfoSource>>,
upstream_info: Arc<dyn gateway_core::account::UpstreamInfoSource>,
cookie_policy: CodexCookiePolicy,
risk_recovery: Mutex<HashMap<String, RiskRecoveryState>>,
account_feedback: Arc<AccountFeedbackStats>,
@@ -296,6 +296,7 @@ impl CodexCredentialSelector {
cooldowns: Arc<dyn ProviderCooldownPort>,
account_feedback: Arc<AccountFeedbackStats>,
cookie_policy: CodexCookiePolicy,
upstream_info: Arc<dyn gateway_core::account::UpstreamInfoSource>,
) -> Self {
Self {
provider_kind,
@@ -306,7 +307,7 @@ impl CodexCredentialSelector {
quota,
catalog,
cooldowns,
upstream_info: None,
upstream_info,
cookie_policy,
risk_recovery: Mutex::new(HashMap::new()),
waiting: ConcurrencyWaitQueue::default(),
@@ -340,13 +341,8 @@ impl CodexCredentialSelector {
.await
}
#[must_use]
pub(crate) fn with_upstream_info(
mut self,
upstream_info: Arc<dyn gateway_core::account::UpstreamInfoSource>,
) -> Self {
self.upstream_info = Some(upstream_info);
self
pub(crate) fn is_peer(&self, account: &ProviderAccount) -> bool {
self.upstream_info.is_peer(account)
}
fn scheduling_signals(
@@ -363,8 +359,7 @@ impl CodexCredentialSelector {
}
let Some(signal) = self
.upstream_info
.as_ref()
.and_then(|source| source.signal(account, model.model(), SystemTime::now()))
.signal(account, model.model(), SystemTime::now())
else {
return (None, None);
};
@@ -426,12 +421,12 @@ impl CodexCredentialSelector {
.as_ref()
.is_some_and(|binding| binding.account_id() != &selected_account)
|| !self
.admit_session(
affinity.key(),
current.as_ref(),
&selected_account,
policy.openai_session_affinity_ttl(),
)
.admit_account(ProviderAccountAdmission::Session(ProviderSessionClaim {
key: affinity.key(),
expected: current.as_ref(),
account: &selected_account,
ttl: policy.openai_session_affinity_ttl(),
}))
.await?
{
return Err(CredentialSelectionError::SessionBound(Box::new(
@@ -752,7 +747,39 @@ impl CodexCredentialSelector {
.into_iter()
.flat_map(|scope| scope.allowed_account_ids().cloned())
.collect::<Vec<_>>();
self.repository.list_for_accounts(&account_ids).await?
let omissions = match (model, request.attempt.account_scope()) {
(CodexSelectionModel::Requested(model), Some(scope)) => {
self.catalog.omissions(model, &account_ids, |id| {
scope.directory().account(id).is_some_and(|a| {
a.model_access().mode() == AccountModelAccessMode::Allowlist
}) || adapter_model.as_ref().is_some_and(|m| {
!matches!(
request.attempt.upstream_adapter_for_account(
&self.provider_kind,
m,
id
),
Ok(None)
)
})
})
}
_ => None,
};
if let Some(filter) = omissions {
let candidates = self
.repository
.store()
.list_catalog_candidates(&self.provider_kind, &account_ids, &filter)
.await
.map_err(CredentialRepositoryError::from)?;
if request.attempt.trace().is_enabled() {
request.attempt.trace().record("account.catalog_candidates",serde_json::json!({"scopedCount":account_ids.len(),"candidateCount":candidates.len()}));
}
candidates
} else {
self.repository.list_for_accounts(&account_ids).await?
}
};
let mut model_access_rejected = 0_usize;
let accounts = accounts
@@ -778,20 +805,14 @@ impl CodexCredentialSelector {
})
.collect::<Vec<_>>();
let mut eligible = Vec::with_capacity(accounts.len());
let catalog_excluded = match (model, request.attempt.account_scope()) {
(CodexSelectionModel::Requested(model), Some(scope)) if !diagnostic => {
// 候选已通过冻结的模型政策,白名单精确命中的 ID 交给上游验证
let mut excluded =
self.catalog
.scheduling_exclusions(&accounts, model, SystemTime::now());
excluded.retain(|id| {
scope.directory().account(id).is_none_or(|account| {
account.model_access().mode() != AccountModelAccessMode::Allowlist
})
});
excluded
}
_ => Default::default(),
let catalog = match (model, request.attempt.account_scope()) {
(CodexSelectionModel::Requested(model), Some(scope)) if !diagnostic => self
.catalog
.cached()
.ok()
.flatten()
.map(|snapshot| (snapshot, model, scope, SystemTime::now())),
_ => None,
};
let mut catalog_rejected = 0_usize;
let mut native_accounts = std::collections::BTreeSet::new();
@@ -816,7 +837,15 @@ impl CodexCredentialSelector {
}
// 受管适配器拥有其上游模型合同,不套用原生目录与拒绝反馈
if adapter.is_none() {
if catalog_excluded.contains(account.id()) {
if catalog
.as_ref()
.is_some_and(|(snapshot, model, scope, now)| {
snapshot.excludes(&account, model, *now)
&& scope.directory().account(account.id()).is_none_or(|a| {
a.model_access().mode() != AccountModelAccessMode::Allowlist
})
})
{
catalog_rejected += 1;
continue;
}
@@ -966,15 +995,31 @@ impl CodexCredentialSelector {
}
}
}
let scheduling = self
.leases
.load_state(
request.attempt.client_api_key_ref(),
let policy = request.attempt.account_selection_policy();
let client_key = request.attempt.client_api_key_ref();
let use_assignment_cursor = policy.strategy() == RotationStrategy::RoundRobin
&& pinned_account.is_none()
&& !(binding.is_some()
&& (follow_only || affinity_mode == AccountAffinity::Preferred))
&& queued_account.is_none();
let (scheduling, assignment_cursor) = futures::try_join!(
self.leases.load_state(
client_key,
&self.provider_kind,
&account_ids,
concurrency_pool,
)
.await?;
concurrency_pool
),
async {
if use_assignment_cursor {
self.session_affinity
.round_robin_cursor(&self.provider_kind, client_key)
.await
.map(Some)
} else {
Ok(None)
}
},
)?;
let round_robin_cursor =
if affinity_mode == AccountAffinity::Preferred && binding.is_some() {
use sha2::{Digest, Sha256};
@@ -988,7 +1033,7 @@ impl CodexCredentialSelector {
},
)
} else {
scheduling.round_robin_cursor()
assignment_cursor.unwrap_or_else(|| scheduling.round_robin_cursor())
};
let mut candidates = accounts
.into_iter()
@@ -1101,8 +1146,6 @@ impl CodexCredentialSelector {
}
let mut shortest_retry = None;
let mut busy_accounts = std::collections::BTreeSet::new();
let policy = request.attempt.account_selection_policy();
loop {
let preferred = pinned_account
.clone()
@@ -1390,25 +1433,53 @@ impl CodexCredentialSelector {
{
binding_account = primary.account_id();
}
if let Some(key) = binding_key {
match self
.admit_session(
key,
binding.as_ref(),
binding_account,
policy.openai_session_affinity_ttl(),
)
.await
let claim = binding_key.map(|key| ProviderSessionClaim {
key,
expected: binding.as_ref(),
account: binding_account,
ttl: policy.openai_session_affinity_ttl(),
});
// 插件委托与普通内置选号共享轮询,显式选择、亲和命中及等待不消耗位置
let admission = match assignment_cursor {
Some(cursor)
if !selection.is_policy_choice()
&& selection.preferred() != PreferredAccountSelection::Hit =>
{
Ok(false) => {
drop(guard);
continue 'capacity;
}
Some(ProviderAccountAdmission::RoundRobin {
client_key,
cursor,
session: claim,
})
}
_ => claim.map(ProviderAccountAdmission::Session),
};
if let Some(admission) = admission {
let admitted = match self.admit_account(admission).await {
Ok(false) => false,
Err(error)
if cache_preference
&& affinity_mode == AccountAffinity::Preferred =>
&& affinity_mode == AccountAffinity::Preferred
&& admission.session().is_some() =>
{
tracing::warn!(%error, "OpenAI cache preference update unavailable");
// 缓存偏好允许丢失,实际的轮询分配仍须占用该位置
match admission {
ProviderAccountAdmission::RoundRobin {
client_key,
cursor,
..
} => {
self.admit_account(
ProviderAccountAdmission::RoundRobin {
client_key,
cursor,
session: None,
},
)
.await?
}
ProviderAccountAdmission::Session(_) => true,
}
}
Err(error) => return Err(error),
Ok(true) => {
@@ -1434,7 +1505,15 @@ impl CodexCredentialSelector {
}),
);
}
true
}
};
if !admitted {
guard.release().await?;
if request.attempt.deadline().is_elapsed() {
return Err(CredentialSelectionError::AccountSnapshotChanged);
}
continue 'capacity;
}
}
let affinity_telemetry = affinity.telemetry(account.id(), cache_preference);
@@ -1581,27 +1660,28 @@ impl CodexCredentialSelector {
.map_err(binding_store_error)
}
async fn admit_session(
async fn admit_account(
&self,
key: &ProviderSessionAffinityKey,
expected: Option<&ProviderSessionBinding>,
account: &ProviderAccountId,
session_ttl: Duration,
admission: ProviderAccountAdmission<'_>,
) -> Result<bool, CredentialSelectionError> {
tokio::time::timeout(
let result = tokio::time::timeout(
SESSION_AFFINITY_TIMEOUT,
self.session_affinity.compare_and_bind(
&self.provider_kind,
key,
expected,
account,
session_ttl,
),
self.session_affinity
.admit_account(&self.provider_kind, admission),
)
.await
.map_err(binding_store_error)?
.map(|binding| binding.is_some())
.map_err(binding_store_error)
.unwrap_or_else(|source| {
Err(ProviderStoreError::caused_by(
gateway_core::provider_ports::ProviderStoreErrorKind::Unavailable,
"admit provider account",
source,
))
});
if admission.session().is_some() {
result.map_err(binding_store_error)
} else {
result.map_err(Into::into)
}
}
/// 只按账号整体的当前事实迁移主绑定,临时容量和单模型限制仍允许恢复热缓存
+16 -15
View File
@@ -153,6 +153,9 @@ pub async fn initialize(
config.base_url().to_owned(),
ports.catalog_cache(),
));
if let Err(error) = catalog.restore().await {
tracing::warn!(%error,"OpenAI model catalog could not be restored");
}
let quota = Arc::new(CodexCredentialQuotaService::new(
repository.clone(),
profile.clone(),
@@ -167,21 +170,19 @@ pub async fn initialize(
http.clone(),
config.base_url().to_owned(),
));
let selector = Arc::new(
CodexCredentialSelector::new(
provider_kind.clone(),
repository.clone(),
Arc::clone(&leases),
session_affinity,
session_exclusions,
Arc::clone(&quota),
Arc::clone(&catalog),
ports.cooldowns(),
Arc::clone(&account_feedback),
CodexCookiePolicy::official().map_err(|_| OpenAiInitializeError::CookiePolicy)?,
)
.with_upstream_info(upstream_info),
);
let selector = Arc::new(CodexCredentialSelector::new(
provider_kind.clone(),
repository.clone(),
Arc::clone(&leases),
session_affinity,
session_exclusions,
Arc::clone(&quota),
Arc::clone(&catalog),
ports.cooldowns(),
Arc::clone(&account_feedback),
CodexCookiePolicy::official().map_err(|_| OpenAiInitializeError::CookiePolicy)?,
upstream_info,
));
let core_provider: Arc<dyn Provider> = Arc::new(
CodexProvider::new(
selector,
@@ -13,7 +13,7 @@ impl CodexProvider {
candidate: &ProviderCandidate,
context: AttemptContext,
) -> Result<ProviderStream, ProviderError> {
if image.payload().protocol() != PROVIDER_NAME || candidate.upstream_model().is_some() {
if image.payload().protocol() != OPENAI_PROTOCOL || candidate.upstream_model().is_some() {
return Err(provider_error(
ProviderErrorKind::InvalidRequest,
UpstreamSendState::NotSent,
@@ -112,7 +112,7 @@ impl CodexProvider {
candidate: &ProviderCandidate,
context: AttemptContext,
) -> Result<ProviderStream, ProviderError> {
if search.payload().protocol() != PROVIDER_NAME || candidate.upstream_model().is_some() {
if search.payload().protocol() != OPENAI_PROTOCOL || candidate.upstream_model().is_some() {
return Err(provider_error(
ProviderErrorKind::InvalidRequest,
UpstreamSendState::NotSent,
@@ -255,7 +255,8 @@ impl CodexProvider {
UpstreamSendState::NotSent,
));
};
if image.kind() != original.kind() || image.payload().protocol() != PROVIDER_NAME {
if image.kind() != original.kind() || image.payload().protocol() != OPENAI_PROTOCOL
{
return Err(provider_error(
ProviderErrorKind::InvalidRequest,
UpstreamSendState::NotSent,
@@ -271,7 +272,7 @@ impl CodexProvider {
}
Operation::Search(search) => {
if !matches!(&request.operation, Operation::Search(_))
|| search.payload().protocol() != PROVIDER_NAME
|| search.payload().protocol() != OPENAI_PROTOCOL
{
return Err(provider_error(
ProviderErrorKind::InvalidRequest,
@@ -330,7 +331,7 @@ impl CodexProvider {
));
// Standalone Provider 端点没有可证明的账号 owner;Images/Search metadata 必须按
// 跨账号输入收敛到当前 lease,不能沿用下游声明的账号或 installation identity
let client = self.client_for_request(&context)?;
let client = self.client_for_request(&context, lease.account())?;
let installation_id = client.profile_state().effective_installation_id(
context.client_api_key_ref(),
lease.account_id(),
@@ -538,7 +539,7 @@ pub(super) fn attach_openai_session_update(
pub(super) fn terminal_response_output(event: &ProviderEvent) -> Option<&[Value]> {
let wire = event.wire_event()?;
if wire.protocol() != PROVIDER_NAME {
if wire.protocol() != OPENAI_PROTOCOL {
return None;
}
let event_type = wire
@@ -750,7 +751,7 @@ pub(super) fn cold_json_response_stream(request: ColdJsonResponse) -> EventStrea
yield ProviderEvent::canonical(GatewayEvent::CalculatedCost(cost));
}
}
let wire = ProtocolWireEvent::raw_json(PROVIDER_NAME, response.body).map_err(|_| {
let wire = ProtocolWireEvent::raw_json(OPENAI_PROTOCOL, response.body).map_err(|_| {
provider_error(ProviderErrorKind::Protocol, UpstreamSendState::Sent)
})?;
yield ProviderEvent::wire(wire);
@@ -943,7 +944,7 @@ pub(super) fn cold_response_stream(response: ColdResponse) -> EventStream {
yield ProviderEvent::observation(observation);
}
if let Some(etag) = response.response_metadata.models_etag.as_deref()
&& let Err(error) = catalog.observe_response_etag(etag)
&& let Err(error) = catalog.observe_response_etag(&active_account,etag)
{
tracing::warn!(
error = %error,
@@ -488,8 +488,8 @@ pub(super) fn apply_websocket_recovery_policy(
failure: &mut MappedProviderFailure,
context: WebSocketRecoveryContext<'_>,
) {
// 未匹配网关恢复规则的上游拒绝直达客户端,不能转换成传输重试或污染会话回退状态
if failure.error.retry_is_prohibited() {
// 明确禁止重试,或上游拒绝继续使用当前账号时,不能改写成同账号传输恢复
if failure.error.retry_is_prohibited() || failure.error.same_account_retry_blocked() {
return;
}
// 明确账号拒绝走已有换号路径,业务瞬态失败走请求内退避,不消耗 WS 传输预算
@@ -1504,6 +1504,10 @@ pub(super) fn map_upstream_failure(
policy,
);
}
if failure.no_more_retry {
// 保留耗尽策略供 Core 决定换号或返回错误,但不再对当前账号退避
error = error.with_same_account_retry_blocked();
}
if let Some(code) = failure.upstream_code() {
error = error.with_upstream_code(OpaqueUpstreamValue::new(code.to_owned()));
}
@@ -365,7 +365,7 @@ impl CodexProvider {
.collect::<Vec<_>>();
let events = cold_live_call_stream(ColdLiveCall {
client: self
.client_for_request(&context)?
.client_for_request(&context, lease.account())?
.for_account(lease.account())
.map_err(|error| map_client_error(error, UpstreamSendState::NotSent, None).error)?
.with_responses_api_base_url(lease.authentication().responses_api_base_url())
@@ -526,7 +526,7 @@ fn cold_live_call_stream(request: ColdLiveCall) -> EventStream {
let response_meta = ResponseMeta::for_provider_endpoint(&request_id);
yield ProviderEvent::canonical(GatewayEvent::Started(response_meta.clone()));
let wire = ProtocolWireEvent::raw_http_body(PROVIDER_NAME, response.body).map_err(|_| {
let wire = ProtocolWireEvent::raw_http_body(OPENAI_PROTOCOL, response.body).map_err(|_| {
provider_error(ProviderErrorKind::Protocol, UpstreamSendState::Sent)
})?;
yield ProviderEvent::wire(wire);
@@ -27,13 +27,13 @@ impl CodexProvider {
return Err(invalid_request());
};
let source = generate.protocol_payload().protocol();
if !matches!(source, PROVIDER_NAME | ANTHROPIC_PROTOCOL) {
if !matches!(source, OPENAI_PROTOCOL | ANTHROPIC_PROTOCOL) {
return Err(invalid_request());
}
if selected.lease.authentication().api_protocol() == ApiKeyProtocol::Messages {
if context.continuation_attempt() == ContinuationAttempt::Native
|| generate.provider_session_state(PROVIDER_NAME).is_some()
|| (source == PROVIDER_NAME
|| (source == OPENAI_PROTOCOL
&& generate
.protocol_payload()
.body()
@@ -49,12 +49,10 @@ impl CodexProvider {
gateway_protocol::anthropic::strip_foreign_thinking_for_messages(&mut body);
body
} else {
let capabilities = tokio::select! {
biased;
_ = context.cancellation().cancelled() => return Err(provider_error(ProviderErrorKind::Cancelled, UpstreamSendState::NotSent)),
_ = context.deadline().wait() => return Err(provider_error(ProviderErrorKind::Timeout, UpstreamSendState::NotSent)),
result = self.catalog.messages_model_capabilities(selected.lease.account(), &model) => result.ok().flatten().unwrap_or_default(),
};
let capabilities = self
.catalog
.messages_model_capabilities(selected.lease.account(), &model)
.unwrap_or_default();
let mut body = responses_to_messages_request(original, &capabilities)
.map_err(|_| invalid_request())?;
// Responses 的 HTTP JSON 和 SSE 共用事件输出,跨协议上游统一请求事件流
@@ -63,7 +61,7 @@ impl CodexProvider {
};
body.insert("model".to_owned(), Value::String(model.as_str().to_owned()));
let mut source_context = generate.protocol_payload().context().clone();
if source == PROVIDER_NAME {
if source == OPENAI_PROTOCOL {
retain_target_headers(source, &mut source_context, &mut headers);
if body.get("speed").and_then(Value::as_str) == Some("fast") {
crate::compat::add_translated_fast_beta(
@@ -72,7 +70,7 @@ impl CodexProvider {
);
}
}
let translator = (source == PROVIDER_NAME)
let translator = (source == OPENAI_PROTOCOL)
.then(|| translator::MessagesResponseTranslator::new(original));
drop(operation);
let stream = self
@@ -89,13 +87,19 @@ impl CodexProvider {
let mut source_context = generate.protocol_payload().context().clone();
retain_target_headers(source, &mut source_context, &mut headers);
let translator = translator::ResponsesResponseTranslator::new(original);
let translated =
let mut translated =
messages_to_responses_request(original).map_err(|_| invalid_request())?;
if matches!(
selected.lease.authentication(),
CodexRuntimeAuthentication::OAuth(_)
) {
crate::compat::normalize_translated_codex_service_tier(&mut translated);
}
let body = serde_json::to_vec(&translated).map_err(|_| invalid_request())?;
let operation = operation
.replace_protocol_wire(None, source_context)
.map_err(|_| invalid_request())?
.translate_protocol_wire(PROVIDER_NAME, body.into())
.translate_protocol_wire(OPENAI_PROTOCOL, body.into())
.map_err(|_| invalid_request())?;
// 原始能力已在上面核对;Responses 终端只校验它实际消费的 wire 事实
if let Operation::Generate(generate) = &operation {
@@ -183,7 +187,7 @@ impl CodexProvider {
lease.capacity_snapshot(),
));
let client = self
.client_for_request(&context)?
.client_for_request(&context, lease.account())?
.for_account(lease.account())
.map_err(|error| map_client_error(error, UpstreamSendState::NotSent, None).error)?
.with_messages_api_base_url(&authentication.configuration.base_url)
@@ -21,7 +21,7 @@ impl NativeResponseTranslator for MessagesResponseTranslator {
ANTHROPIC_PROTOCOL
}
fn target_protocol(&self) -> &str {
PROVIDER_NAME
OPENAI_PROTOCOL
}
fn translate(
&mut self,
@@ -33,7 +33,7 @@ impl NativeResponseTranslator for MessagesResponseTranslator {
.stream
.complete_json_event(&message)
.map_err(|_| protocol_error())?;
return wire_events(PROVIDER_NAME, vec![event]);
return wire_events(OPENAI_PROTOCOL, vec![event]);
}
if !event.has_json_data() {
return Ok(Vec::new());
@@ -45,7 +45,7 @@ impl NativeResponseTranslator for MessagesResponseTranslator {
if event.data().get("type").and_then(Value::as_str) == Some("message_stop") {
self.stream.finish().map_err(|_| protocol_error())?;
}
wire_events(PROVIDER_NAME, events)
wire_events(OPENAI_PROTOCOL, events)
}
}
@@ -63,7 +63,7 @@ impl ResponsesResponseTranslator {
impl NativeResponseTranslator for ResponsesResponseTranslator {
fn source_protocol(&self) -> &str {
PROVIDER_NAME
OPENAI_PROTOCOL
}
fn target_protocol(&self) -> &str {
ANTHROPIC_PROTOCOL
@@ -44,6 +44,7 @@ use gateway_core::task::{
WorkerRunnable, WorkerSchedule, WorkerTaskError,
};
use gateway_core::upstream::{UpstreamSendState, UpstreamTransport};
use gateway_protocol::openai::OPENAI_PROTOCOL;
use gateway_protocol::openai::events::{
ParsedRateLimits, parse_rate_limit_headers, rate_limits_to_header_pairs,
};
@@ -301,9 +302,14 @@ impl CodexProvider {
fn client_for_request(
&self,
context: &AttemptContext,
account: &ProviderAccount,
) -> Result<CodexBackendClient, ProviderError> {
let mut client = self.client.clone();
client.peer_route = context.extension_scope().peer_route().cloned();
client.peer_route = context
.extension_scope()
.peer_route()
.filter(|_| self.selector.is_peer(account))
.cloned();
let Some(profile) = context.request_profile() else {
return Ok(client);
};
@@ -502,13 +508,11 @@ impl Provider for CodexProvider {
let Ok(encoded) = encode_generate_request(request, "observability", None) else {
return ProviderRequestObservation::default();
};
let previous_session = decode_openai_session_state(request);
let semantics = encoded.semantics(
previous_session
.as_ref()
.and_then(|state| state.reasoning_effort_override.as_deref()),
);
let reasoning_effort = semantics.reasoning_effort.clone();
let previous_reasoning_effort_override = request
.provider_session_state(PROVIDER_NAME)
.filter(|state| state.extension_owner().is_none())
.and_then(|_| decode_openai_session_state(request))
.and_then(|state| state.reasoning_effort_override);
let previous_response_id = encoded.previous_response_id();
let continuation = ContinuationRequestObservation {
affinity_hash: derive_codex_session_affinity(&encoded, client_api_key_id)
@@ -520,11 +524,7 @@ impl Provider for CodexProvider {
};
ProviderRequestObservation {
requested_model: None,
reasoning_effort,
reasoning_preset: semantics.reasoning_preset.map(str::to_owned),
request_kind: semantics.request_kind,
subagent_kind: semantics.subagent_kind,
compact: semantics.compact,
previous_reasoning_effort_override,
continuation,
}
}
@@ -536,9 +536,16 @@ impl Provider for CodexProvider {
async fn query_model_capabilities(
&self,
) -> Result<Vec<ProviderModelCapabilities>, ProviderError> {
let snapshot = self.catalog.synchronize().await.map_err(|_| {
provider_error(ProviderErrorKind::Unavailable, UpstreamSendState::NotSent)
})?;
// 能力发现读取已发布事实;冷启动和配置失效由 Provider 后台任务重建目录
let snapshot = self
.catalog
.cached()
.map_err(|_| {
provider_error(ProviderErrorKind::Unavailable, UpstreamSendState::NotSent)
})?
.ok_or_else(|| {
provider_error(ProviderErrorKind::Unavailable, UpstreamSendState::NotSent)
})?;
let mut accounts = snapshot.model_catalog_accounts();
Ok(snapshot
.models()
@@ -555,11 +562,11 @@ impl Provider for CodexProvider {
async fn query_client_model_catalog(
&self,
client_api_key_id: &gateway_core::policy::ClientApiKeyId,
_client_api_key_id: &gateway_core::policy::ClientApiKeyId,
scope: &gateway_core::account::scope::FrozenAccountScope,
protocol: &str,
client_version: &str,
peer_route: Option<&reqwest::header::HeaderValue>,
_client_version: &str,
_peer_route: Option<&reqwest::header::HeaderValue>,
) -> Result<
Option<Vec<gateway_core::routing::ProviderModelDescriptor>>,
gateway_core::routing::ProviderCatalogUnavailable,
@@ -568,7 +575,7 @@ impl Provider for CodexProvider {
return Ok(None);
}
self.catalog
.client_model_catalog(client_api_key_id, scope, client_version, peer_route)
.client_model_catalog(scope)
.await
.map(Some)
.map_err(|_| gateway_core::routing::ProviderCatalogUnavailable)
@@ -628,7 +635,7 @@ impl Provider for CodexProvider {
let mut operation = Operation::Generate(generate.clone());
if context.fast_mode() != gateway_core::account::FastMode::Default
&& let Operation::Generate(generate) = &operation
&& generate.protocol_payload().protocol() == PROVIDER_NAME
&& generate.protocol_payload().protocol() == OPENAI_PROTOCOL
{
let mut request =
CodexResponsesRequest::from_body(generate.protocol_payload().body().clone());
@@ -646,7 +653,7 @@ impl Provider for CodexProvider {
)
})?;
operation = operation
.replace_middleware_wire(PROVIDER_NAME, body.into())
.replace_middleware_wire(OPENAI_PROTOCOL, body.into())
.map_err(|_| {
provider_error(
ProviderErrorKind::InvalidRequest,
@@ -669,7 +676,7 @@ impl Provider for CodexProvider {
}
// 其他协议必须先取得真实账号,再按固定 attempt 阶段调用转换器;选号前不能
// 把未知正文当成 OpenAI wire 解释会话、亲和或传输字段
let upstream = (generate.protocol_payload().protocol() == PROVIDER_NAME)
let upstream = (generate.protocol_payload().protocol() == OPENAI_PROTOCOL)
.then(|| encode_generate_request(generate, upstream_model.as_str(), None))
.transpose()
.map_err(map_request_error)?;
@@ -836,7 +843,7 @@ impl CodexProvider {
UpstreamSendState::NotSent,
));
};
if generate.protocol_payload().protocol() != PROVIDER_NAME
if generate.protocol_payload().protocol() != OPENAI_PROTOCOL
|| generate.wire_capability_requirements() != frozen_requirements
{
return Err(provider_error(
@@ -993,7 +1000,7 @@ impl CodexProvider {
.get(context.client_api_key_ref(), lease.account(), response_id)
.map(|affinity| (response_id.to_owned(), affinity))
});
let client = self.client_for_request(&context)?;
let client = self.client_for_request(&context, lease.account())?;
let installation_id = client.profile_state().effective_installation_id(
context.client_api_key_ref(),
lease.account_id(),
@@ -5,7 +5,7 @@ use super::*;
pub(super) fn endpoint_requested_model(
payload: &gateway_core::operation::RawJsonPayload,
) -> Option<gateway_core::routing::PublicModelId> {
if payload.protocol() != PROVIDER_NAME {
if payload.protocol() != OPENAI_PROTOCOL {
return None;
}
// 选号与观测共用模型提取,跳过编辑请求中的图片且不改写原始正文
@@ -43,7 +43,7 @@ impl CodexProvider {
}
let requested_service_tier = match &operation {
Operation::Generate(generate)
if generate.protocol_payload().protocol() == PROVIDER_NAME =>
if generate.protocol_payload().protocol() == OPENAI_PROTOCOL =>
{
let request =
CodexResponsesRequest::from_body(generate.protocol_payload().body().clone());
@@ -46,7 +46,10 @@ pub(crate) fn worker_contributions(
contributions.push(WorkerContribution::Registration(scheduled_registration(
refresh_id,
OAUTH_REFRESH_INTERVAL,
Box::new(OpenAiOAuthRefreshTask { service: refresh }),
Box::new(OpenAiOAuthRefreshTask {
service: refresh,
catalog: Arc::clone(&catalog),
}),
)?));
}
contributions.extend([
@@ -132,6 +135,7 @@ pub(super) fn scheduled_registration(
}
pub(super) struct OpenAiOAuthRefreshTask {
catalog: Arc<CodexCredentialCatalogService>,
service: Arc<CodexCredentialRefreshService>,
}
@@ -144,6 +148,7 @@ impl ScheduledTask for OpenAiOAuthRefreshTask {
let outcomes = self.service.refresh_due().await.map_err(|error| {
WorkerTaskError::safe("OpenAI OAuth refresh failed").with_source(error)
})?;
let mut changed = Vec::new();
let mut refreshed = 0_u64;
let mut invalidated = 0_u64;
let mut banned = 0_u64;
@@ -155,9 +160,18 @@ impl ScheduledTask for OpenAiOAuthRefreshTask {
let mut failed_accounts = Vec::new();
for outcome in &outcomes {
match outcome {
CodexCredentialRefreshOutcome::Refreshed { .. } => refreshed += 1,
CodexCredentialRefreshOutcome::Invalidated { .. } => invalidated += 1,
CodexCredentialRefreshOutcome::Banned { .. } => banned += 1,
CodexCredentialRefreshOutcome::Refreshed { account_id, .. } => {
refreshed += 1;
changed.push(account_id);
}
CodexCredentialRefreshOutcome::Invalidated { account_id } => {
invalidated += 1;
changed.push(account_id);
}
CodexCredentialRefreshOutcome::Banned { account_id } => {
banned += 1;
changed.push(account_id);
}
CodexCredentialRefreshOutcome::Transient { account_id } => {
transient += 1;
transient_accounts.push(account_id);
@@ -172,6 +186,21 @@ impl ScheduledTask for OpenAiOAuthRefreshTask {
}
}
}
if !changed.is_empty() {
let ids = changed
.into_iter()
.map(|id| gateway_core::account::ProviderAccountId::new(id.as_str()))
.collect::<Result<Vec<_>, _>>()
.map_err(|error| {
WorkerTaskError::safe("OpenAI 账号 ID 无效").with_source(error)
})?;
self.catalog
.invalidate_accounts(&ids)
.await
.map_err(|error| {
WorkerTaskError::safe("OpenAI 凭据变更后目录失效失败").with_source(error)
})?;
}
if !outcomes.is_empty() {
tracing::info!(
refreshed,
@@ -306,12 +335,21 @@ impl DaemonTask for OpenAiCatalogEtagTask {
loop {
tokio::select! {
() = cancellation.cancelled() => return Ok(()),
() = self.catalog.wait_for_etag_refresh() => {},
() = self.catalog.wait_for_refresh() => {},
};
// 失败的 ETag 已重新入队,交回 Host 退避后再领取,避免立即重试
self.catalog.refresh().await.map_err(|error| {
WorkerTaskError::safe("OpenAI 模型目录 ETag 刷新失败").with_source(error)
})?;
let result = tokio::select! {
()=cancellation.cancelled()=>return Ok(()),
result=self.catalog.refresh_catalogs()=>result,
};
match result {
Ok(_) | Err(CodexCredentialCatalogError::NoEligibleCredential) => {}
Err(error) => {
self.catalog.schedule_refresh();
return Err(
WorkerTaskError::safe("OpenAI 模型目录刷新失败").with_source(error)
);
}
}
context.succeeded();
}
})
@@ -14,6 +14,7 @@ use gateway_core::event::{
};
use gateway_core::metering::Usage;
use gateway_core::upstream::UpstreamSendState;
use gateway_protocol::openai::OPENAI_PROTOCOL;
use gateway_protocol::openai::events::{TokenUsage, billable_usage_is_complete, extract_usage};
use gateway_protocol::openai::sse::{SseEvent, SseEventDecoder, SseFrame, sse_frame_is_done};
use serde_json::Value;
@@ -301,7 +302,7 @@ impl CodexCanonicalDecoder {
}
let raw = Bytes::from(raw);
if events.is_empty() {
if let Ok(wire) = ProtocolWireEvent::raw_sse("openai", raw) {
if let Ok(wire) = ProtocolWireEvent::raw_sse(OPENAI_PROTOCOL, raw) {
output.push(ProviderEvent::wire(wire));
}
continue;
@@ -346,10 +347,10 @@ impl CodexCanonicalDecoder {
);
let wire = match raw_sse_frame {
Some(RawResponseFrame::Sse(raw)) => {
ProtocolWireEvent::raw_sse("openai", raw).ok()
ProtocolWireEvent::raw_sse(OPENAI_PROTOCOL, raw).ok()
}
Some(RawResponseFrame::WebSocket(raw)) => {
ProtocolWireEvent::raw_websocket("openai", raw).ok()
ProtocolWireEvent::raw_websocket(OPENAI_PROTOCOL, raw).ok()
}
None => None,
};
@@ -474,12 +475,12 @@ impl CodexCanonicalDecoder {
let sse_id = event.id;
match raw_sse_frame {
Some(RawResponseFrame::WebSocket(raw)) => {
ProtocolWireEvent::json("openai", event_type, value)
ProtocolWireEvent::json(OPENAI_PROTOCOL, event_type, value)
.map(|wire| wire.with_raw_websocket_message(raw))
}
Some(RawResponseFrame::Sse(raw_sse_frame)) => {
ProtocolWireEvent::json_with_raw_sse_metadata(
"openai",
OPENAI_PROTOCOL,
event_type,
value,
raw_sse_frame,
@@ -488,7 +489,7 @@ impl CodexCanonicalDecoder {
)
}
None => ProtocolWireEvent::json_with_sse_metadata(
"openai",
OPENAI_PROTOCOL,
event_type,
value,
sse_id,
@@ -7,6 +7,7 @@ use bytes::Bytes;
use gateway_core::operation::RawJsonPayload;
use gateway_core::routing::{ModelServiceTier, UpstreamModelId};
use gateway_protocol::anthropic::MessagesModelCapabilities;
use gateway_protocol::openai::OPENAI_PROTOCOL;
use reqwest::header::{ETAG, HeaderMap};
use serde::Deserialize;
@@ -237,13 +238,40 @@ impl CodexCatalogModel {
}
/// 一次完整成功的 Codex 远端模型快照
#[derive(Debug, Clone, PartialEq, Eq)]
#[derive(Clone, PartialEq, Eq)]
pub struct CodexModelCatalogSnapshot {
document: Bytes,
models: Vec<CodexCatalogModel>,
etag: Option<String>,
}
impl std::fmt::Debug for CodexModelCatalogSnapshot {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("CodexModelCatalogSnapshot")
.field("models", &self.models)
.finish_non_exhaustive()
}
}
impl CodexModelCatalogSnapshot {
pub(crate) fn into_parts(self) -> (Vec<CodexCatalogModel>, Bytes, Option<String>) {
(self.models, self.document, self.etag)
}
/// 落库的目录沿用相同的上游校验器;协议标记区分普通 API 与 Messages 能力
pub(crate) fn restore(document: &[u8]) -> Result<Self, CodexModelCatalogError> {
match document.split_first() {
Some((0, body)) => parse_codex_model_catalog(body, None),
Some((1, body)) => {
parse_api_model_catalog(body, None, OpenAiUpstreamProtocol::ResponsesApi)
}
Some((2, body)) => {
parse_api_model_catalog(body, None, OpenAiUpstreamProtocol::MessagesApi)
}
_ => Err(CodexModelCatalogError::InvalidWire),
}
}
/// 返回本轮全部模型
#[must_use]
pub fn models(&self) -> &[CodexCatalogModel] {
@@ -289,7 +317,7 @@ pub enum CodexModelCatalogError {
InvalidEtag,
}
#[derive(Debug, Deserialize)]
#[derive(Debug, Deserialize, serde::Serialize)]
struct CodexModelsWire {
models: Vec<serde_json::Value>,
}
@@ -361,6 +389,7 @@ fn parse_native_model_catalog(
return Err(CodexModelCatalogError::TooManyModels);
}
let document = archive(protocol, &wire)?;
let etag = etag.map(validate_etag).transpose()?;
let mut seen = BTreeSet::new();
let mut models = Vec::with_capacity(wire.models.len());
@@ -380,7 +409,11 @@ fn parse_native_model_catalog(
}
models.push(model);
}
Ok(CodexModelCatalogSnapshot { models, etag })
Ok(CodexModelCatalogSnapshot {
document,
models,
etag,
})
}
pub(super) fn catalog_etag(headers: &HeaderMap) -> Result<Option<String>, CodexModelCatalogError> {
@@ -617,12 +650,30 @@ pub(super) fn parse_api_model_catalog(
#[serde(flatten)]
fields: serde_json::Map<String, serde_json::Value>,
}
let value: serde_json::Value =
let mut value: serde_json::Value =
serde_json::from_slice(body).map_err(|_| CodexModelCatalogError::InvalidWire)?;
// 部分 Responses 上游按客户端版本返回完整 Codex 目录;不能降格为只有 ID 的画像
if value.get("models").is_some() {
return parse_native_model_catalog(body, etag, OpenAiUpstreamProtocol::ResponsesApi);
}
if let Some(top) = value.as_object_mut() {
top.retain(|key, _| key == "data");
}
if let Some(items) = value
.get_mut("data")
.and_then(serde_json::Value::as_array_mut)
{
for item in items {
if let Some(fields) = item.as_object_mut() {
fields.retain(|key, _| {
matches!(key.as_str(), "id" | "display_name")
|| protocol == OpenAiUpstreamProtocol::MessagesApi
&& matches!(key.as_str(), "capabilities" | "max_tokens")
});
}
}
}
let document = archive(protocol, &value)?;
let wire: ApiModels =
serde_json::from_value(value).map_err(|_| CodexModelCatalogError::InvalidWire)?;
if wire.data.is_empty() {
@@ -647,7 +698,7 @@ pub(super) fn parse_api_model_catalog(
.unwrap_or_else(|| request_model.as_str().to_owned());
validate_public_text(&display_name, MAX_DISPLAY_NAME_BYTES, false)?;
let document = RawJsonPayload::new(
"openai",
OPENAI_PROTOCOL,
Bytes::from(
serde_json::to_vec(&serde_json::json!({"id": request_model.as_str()}))
.map_err(|_| CodexModelCatalogError::InvalidWire)?,
@@ -668,7 +719,21 @@ pub(super) fn parse_api_model_catalog(
});
}
Ok(CodexModelCatalogSnapshot {
document,
models,
etag: etag.map(validate_etag).transpose()?,
})
}
fn archive(
protocol: OpenAiUpstreamProtocol,
value: &impl serde::Serialize,
) -> Result<Bytes, CodexModelCatalogError> {
let mut bytes = vec![match protocol {
OpenAiUpstreamProtocol::Codex => 0,
OpenAiUpstreamProtocol::ResponsesApi => 1,
OpenAiUpstreamProtocol::MessagesApi => 2,
}];
serde_json::to_writer(&mut bytes, value).map_err(|_| CodexModelCatalogError::InvalidWire)?;
Ok(bytes.into())
}
@@ -10,6 +10,7 @@ use super::{
},
diagnostics::CodexUpstreamSendPhase,
endpoints::endpoint_url,
profile::CodexWireProfile,
response_meta,
};
@@ -20,9 +21,24 @@ impl CodexBackendClient {
context: CodexRequestContext<'_>,
client_version: Option<&str>,
) -> CodexClientResult<CodexModelCatalogSnapshot> {
let (body, etag) = self
.fetch_models_page(context, client_version, None)
.await?;
let profile = self.profile.snapshot();
let query = match self.protocol {
OpenAiUpstreamProtocol::Codex | OpenAiUpstreamProtocol::ResponsesApi => (
"client_version",
client_version.unwrap_or(&profile.codex_version),
),
OpenAiUpstreamProtocol::MessagesApi => ("limit", "1000"),
};
let (body, etag) = match self.fetch_models_page(context, &profile, &[query]).await {
Err(CodexClientError::Upstream {
status, ref body, ..
}) if self.protocol == OpenAiUpstreamProtocol::ResponsesApi
&& rejects_catalog_version(status, body) =>
{
self.fetch_models_page(context, &profile, &[]).await?
}
result => result?,
};
match self.protocol {
OpenAiUpstreamProtocol::Codex => Ok(parse_codex_model_catalog(&body, etag.as_deref())?),
OpenAiUpstreamProtocol::ResponsesApi => Ok(super::catalog::parse_api_model_catalog(
@@ -56,7 +72,9 @@ impl CodexBackendClient {
if !cursors.insert(cursor.to_owned()) {
return Err(super::CodexModelCatalogError::InvalidWire.into());
}
(page, _) = self.fetch_models_page(context, None, Some(cursor)).await?;
(page, _) = self
.fetch_models_page(context, &profile, &[query, ("after_id", cursor)])
.await?;
}
let body = serde_json::to_vec(&serde_json::json!({"data":models}))
.map_err(|_| super::CodexModelCatalogError::InvalidWire)?;
@@ -72,32 +90,20 @@ impl CodexBackendClient {
async fn fetch_models_page(
&self,
context: CodexRequestContext<'_>,
client_version: Option<&str>,
after_id: Option<&str>,
profile: &CodexWireProfile,
query: &[(&str, &str)],
) -> CodexClientResult<(bytes::Bytes, Option<String>)> {
let path = match self.protocol {
OpenAiUpstreamProtocol::Codex => "codex/models",
OpenAiUpstreamProtocol::ResponsesApi | OpenAiUpstreamProtocol::MessagesApi => "models",
};
let profile = self.profile.snapshot();
let headers = self.model_request_headers(&profile, context)?;
let headers = self.model_request_headers(profile, context)?;
let mut request = self
.client
.get(endpoint_url(&self.base_url, path))
.headers(headers);
let client_version = client_version.or_else(|| {
(self.protocol == OpenAiUpstreamProtocol::Codex)
.then_some(profile.codex_version.as_str())
});
if self.protocol == OpenAiUpstreamProtocol::MessagesApi {
request = request.query(&[("limit", "1000")]);
if let Some(after_id) = after_id {
request = request.query(&[("after_id", after_id)]);
}
} else if let Some(client_version) = client_version {
request = request.query(&[("client_version", client_version)]);
}
let mut request = request.build()?;
.headers(headers)
.query(query)
.build()?;
if let Some(route) = self.route_to(request.url()) {
request
.headers_mut()
@@ -134,3 +140,19 @@ impl CodexBackendClient {
Ok((body, etag))
}
}
/// 仅在上游明确不接受版本参数时读取普通目录,其他失败保留原错误
fn rejects_catalog_version(status: reqwest::StatusCode, body: &str) -> bool {
if !matches!(status.as_u16(), 400 | 422) {
return false;
}
let Ok(value) = serde_json::from_str::<serde_json::Value>(body) else {
return false;
};
let error = &value["error"];
error["param"].as_str() == Some("client_version")
&& matches!(
error["code"].as_str(),
Some("unknown_parameter" | "unsupported_parameter")
)
}
@@ -79,6 +79,8 @@ pub struct CodexUpstreamFailure {
/// 官方 `usage_limit_reached` 错误体携带的窗口重置时间(Unix 秒)
pub(crate) usage_limit_resets_at: Option<i64>,
pub(crate) retry_after_seconds: Option<u64>,
/// 上游错误头声明不再重试当前账号
pub(crate) no_more_retry: bool,
pub(crate) request_id: Option<String>,
pub(crate) response_origin: Option<url::Url>,
pub(crate) set_cookie_headers: Vec<String>,
@@ -127,6 +129,7 @@ impl CodexUpstreamFailure {
.then_some(fields.resets_at)
.flatten(),
retry_after_seconds,
no_more_retry: gateway_protocol::openai::events::upstream_no_more_retry_from_body(body),
request_id: diagnostics.request_id.clone(),
response_origin: diagnostics.response_origin.clone(),
set_cookie_headers: set_cookie_headers.to_vec(),
@@ -179,6 +182,7 @@ impl CodexUpstreamFailure {
identity_error_code: diagnostics.identity_error_code.clone(),
usage_limit_resets_at: None,
retry_after_seconds: failure.retry_after_seconds,
no_more_retry: failure.no_more_retry,
request_id: failure
.request_id
.clone()
@@ -26,7 +26,7 @@ impl CodexBackendClient {
self
}
// 官方端点不参与节点转发,按插件准备后的实际目标决定是否发送
// 已识别的互联账号才携带路径,官方端点不参与节点转发
pub(super) fn route_to(&self, target: &Url) -> Option<&http::HeaderValue> {
self.peer_route.as_ref().filter(|_| {
self.protocol != super::client::OpenAiUpstreamProtocol::Codex
@@ -40,20 +40,15 @@ impl CodexBackendClient {
pub(super) async fn prepare_http_request(
&self,
endpoint: String,
headers: HeaderMap,
mut headers: HeaderMap,
body: Bytes,
) -> CodexClientResult<reqwest::Request> {
let mut request = self.prepare_http_outbound(endpoint, headers, body).await?;
let path = self.route_to(request.url());
request
.headers_mut()
.remove(gateway_protocol::peer::ROUTE_HEADER);
if let Some(path) = path {
request
.headers_mut()
.insert(gateway_protocol::peer::ROUTE_HEADER, path.clone());
let target = Url::parse(&endpoint).map_err(|_| invalid_request())?;
headers.remove(gateway_protocol::peer::ROUTE_HEADER);
if let Some(path) = self.route_to(&target) {
headers.insert(gateway_protocol::peer::ROUTE_HEADER, path.clone());
}
Ok(request)
self.prepare_http_outbound(endpoint, headers, body).await
}
async fn prepare_http_outbound(
@@ -109,9 +104,8 @@ impl CodexBackendClient {
pub(super) async fn prepare_websocket_request(
&self,
request: CodexWebSocketRequest,
mut request: CodexWebSocketRequest,
) -> CodexClientResult<(CodexWebSocketRequest, Option<String>)> {
let (mut request, profile) = self.prepare_websocket_outbound(request).await?;
let path = if self.peer_route.is_some()
&& self.protocol != super::client::OpenAiUpstreamProtocol::Codex
{
@@ -125,7 +119,7 @@ impl CodexBackendClient {
None
};
request.connection.set_peer_route(path);
Ok((request, profile))
self.prepare_websocket_outbound(request).await
}
async fn prepare_websocket_outbound(
@@ -496,6 +496,8 @@ pub struct ResponsesSseFailure {
pub explicit_status_code: Option<u16>,
/// 上游显式重试间隔,或从官方限流消息中解析出的重试间隔
pub retry_after_seconds: Option<u64>,
/// 上游 `x-retry-metadata: NO_MORE_RETRY`,与重试间隔独立
pub no_more_retry: bool,
/// 当前错误事件自身携带的请求 ID;不是连接 opening ID
pub(crate) request_id: Option<String>,
/// 上游错误事件的原始 JSON data;只应在明确的失败审计边界读取
@@ -512,6 +514,7 @@ impl fmt::Debug for ResponsesSseFailure {
.field("has_upstream_type", &self.upstream_type.is_some())
.field("explicit_status_code", &self.explicit_status_code)
.field("retry_after_seconds", &self.retry_after_seconds)
.field("no_more_retry", &self.no_more_retry)
.finish()
}
}
@@ -531,6 +534,7 @@ impl ResponsesSseFailure {
upstream_type: failure_type(value),
explicit_status_code: failure_explicit_status_code(value),
retry_after_seconds: events::retry_after_seconds_from_value(value),
no_more_retry: events::upstream_no_more_retry(value),
request_id:
crate::transport::diagnostics::CodexUpstreamDiagnostics::error_event_request_id(
value,
@@ -768,19 +772,11 @@ impl CodexResponsesRequest {
&self,
previous_reasoning_effort_override: Option<&str>,
) -> CodexRequestSemantics {
let mut semantics = codex_responses_request_semantics_with_turn_metadata(
codex_responses_request_semantics_with_turn_metadata(
self.body(),
self.turn_metadata.as_deref(),
previous_reasoning_effort_override,
);
// 官方 generate=false 只准备连接与上下文,不是模型推理
// 预热分类参与用量筛选,不能仅凭客户端的 request_kind 提示排除真实推理
if !self.generate() {
semantics.request_kind = Some("prewarm".to_owned());
} else if semantics.request_kind.as_deref() == Some("prewarm") {
semantics.request_kind = None;
}
semantics
)
}
/// 返回请求语义与 child transport 隔离所用的子代理区分值
@@ -5,7 +5,8 @@ use std::io;
use base64::{Engine as _, engine::general_purpose::STANDARD};
use gateway_core::operation::GenerateRequest;
use gateway_protocol::openai::{
WS_REQUEST_HEADER_RESPONSES_LITE_CLIENT_METADATA_KEY, is_transport_managed_request_header,
OPENAI_PROTOCOL, WS_REQUEST_HEADER_RESPONSES_LITE_CLIENT_METADATA_KEY,
is_transport_managed_request_header,
};
use reqwest::header::{HeaderMap, HeaderName, HeaderValue};
use roxmltree::Document;
@@ -108,7 +109,7 @@ pub fn encode_generate_request(
location: Option<&CodexRequestLocation>,
) -> Result<CodexResponsesRequest, CodexRequestEncodeError> {
let payload = request.protocol_payload();
if payload.protocol() != "openai" {
if payload.protocol() != OPENAI_PROTOCOL {
return Err(CodexRequestEncodeError::InvalidProtocolPayload);
}
let mut body = payload.body().clone();
@@ -709,7 +709,7 @@ async fn reset_credit_explicit_http_rejection_should_preserve_the_raw_body() {
}
#[tokio::test]
async fn openai_core_provider_projects_codex_request_observation_without_routing_side_effects() {
async fn openai_core_provider_projects_private_request_observation_without_execution() {
let config = valid_config();
let bundle = provider_openai::initialize(
config.config.clone(),
@@ -723,14 +723,11 @@ async fn openai_core_provider_projects_codex_request_observation_without_routing
Map::from_iter([
("model".to_owned(), json!("gpt-5.4")),
("input".to_owned(), json!("summarize")),
("reasoning".to_owned(), json!({"effort": "high"})),
("session_id".to_owned(), json!("observation-session")),
("previous_response_id".to_owned(), json!("resp_previous")),
]),
)
.expect("OpenAI payload")
.with_context(Map::from_iter([(
"turn_metadata".to_owned(),
Value::String(r#"{"request_kind":"compaction","subagent_kind":"review"}"#.to_owned()),
)]));
.expect("OpenAI payload");
let operation = Operation::Generate(GenerateRequest::from_protocol_payload(payload));
let client_key_id = ClientApiKeyId::new("key_openai_admin_observation").expect("client key");
@@ -738,11 +735,9 @@ async fn openai_core_provider_projects_codex_request_observation_without_routing
.core_provider()
.request_observation(&operation, &client_key_id);
assert_eq!(observation.request_kind.as_deref(), Some("compaction"));
assert_eq!(observation.subagent_kind.as_deref(), Some("review"));
// Codex 当前只在特定多代理预设组合下给出 reasoning_preset;普通 high 保持空值
assert_eq!(observation.reasoning_preset, None);
assert!(observation.compact);
assert!(observation.continuation.requested);
assert!(observation.continuation.previous_response_id_hash.is_some());
assert!(observation.continuation.affinity_hash.is_some());
}
#[tokio::test]
@@ -617,3 +617,133 @@ async fn all_filtered_is_terminal_before_queue_or_lease_even_for_native_pins() {
assert!(leases.requests.lock().unwrap().is_empty());
}
}
#[tokio::test]
async fn round_robin_counts_concurrent_first_requests_as_one_session() {
let store = Arc::new(MemoryAccountStore::default());
create_account(&store, "acct_first", "test-first");
create_account(&store, "acct_second", "test-second");
let leases = Arc::new(TestLeaseCoordinator::default());
let affinity = Arc::new(MemorySessionAffinity::with_initial_claim_barrier(2));
let selector = selector_with_affinity(&store, leases, affinity.clone());
let url = Url::parse(OFFICIAL_CODEX_BASE_URL).unwrap();
let mut winners = Vec::new();
for key in ["session-1", "session-2"] {
let (a, b) = tokio::join!(biased;
select_strict_round_robin_session(&selector, &url, key),
select_strict_round_robin_session(&selector, &url, key),
);
assert_eq!(a, b);
winners.push(a);
}
assert_eq!(affinity.binding_count(), 2);
assert_ne!(winners[0], winners[1]);
}
#[tokio::test]
async fn busy_leases_do_not_consume_the_next_round_robin_assignment() {
let store = Arc::new(MemoryAccountStore::default());
create_account(&store, "acct_first", "test-first");
create_account(&store, "acct_second", "test-second");
let leases = Arc::new(TestLeaseCoordinator::default());
let selector = selector(&store, leases.clone());
let url = Url::parse(OFFICIAL_CODEX_BASE_URL).unwrap();
let attempt = round_robin_attempt();
let key = ProviderSessionAffinityKey::try_new("busy-session").unwrap();
*leases.busy.lock().unwrap() = true;
let result = selector
.select(&SelectCodexCredential {
upstream_model: "gpt-5.4",
request_url: &url,
attempt: &attempt,
session_affinity_key: Some(&key),
})
.await;
assert!(
matches!(
result,
Err(CredentialSelectionError::CapacityUnavailable { .. })
),
"{result:?}"
);
*leases.busy.lock().unwrap() = false;
assert_eq!(
select_strict_round_robin_session(&selector, &url, "available").await,
"acct_first"
);
}
#[tokio::test]
async fn round_robin_counts_delegated_assignments_and_skips_explicit_picks() {
for decision in [
AccountScheduleDecision::Delegate,
AccountScheduleDecision::DelegateFiltered(BTreeSet::from([ProviderAccountId::new(
"acct_second",
)
.unwrap()])),
] {
let store = Arc::new(MemoryAccountStore::default());
for id in ["acct_first", "acct_second", "acct_fallback"] {
create_account(&store, id, id);
}
let selector = selector(&store, Arc::new(TestLeaseCoordinator::default()));
let attempt_for = |decision| {
capacity_attempt(
AccountSelectionPolicy::new(
RotationStrategy::RoundRobin,
NonZeroU32::new(2).unwrap(),
),
Some(RequestPolicyContext::new(
Arc::new(Scheduler { decision }),
ExtensionSetReference::new(
ExtensionSetId::new("rr-policy".into()).unwrap(),
Arc::new(ExtensionLease),
),
ModelRequestId::new("req_capacity_contract").unwrap(),
ClientApiKeyId::new("key_codex_contract").unwrap(),
vec![],
)),
)
};
let url = Url::parse(OFFICIAL_CODEX_BASE_URL).unwrap();
let picked = attempt_for(AccountScheduleDecision::Pick(
ProviderAccountId::new("acct_second").unwrap(),
));
let delegated = attempt_for(decision.clone());
let mut actual = Vec::new();
for (session, attempt) in [
("pick", &picked),
("one", &delegated),
("one", &delegated),
("two", &delegated),
("three", &delegated),
] {
let key = ProviderSessionAffinityKey::try_new(session).unwrap();
let selected = selector
.select(&SelectCodexCredential {
upstream_model: "gpt-5.4",
request_url: &url,
attempt,
session_affinity_key: Some(&key),
})
.await
.unwrap();
actual.push(selected.account_id().as_str().to_owned());
}
let third = if matches!(decision, AccountScheduleDecision::Delegate) {
"acct_second"
} else {
"acct_fallback"
};
assert_eq!(
actual,
[
"acct_second",
"acct_fallback",
"acct_fallback",
"acct_first",
third
]
);
}
}
@@ -11,11 +11,11 @@ use std::time::{Duration, SystemTime};
use futures::executor::block_on;
use gateway_core::account::{
AccountAttemptFeedback, AccountConcurrencyLimit, AccountErrorReason, AccountFeedbackStats,
AccountSelectionPolicy, AccountStateChange, AccountWeight, CredentialState, OpaqueProviderData,
ProviderAccount, ProviderAccountId, ProviderAccountStore as _, QuotaAccessChange,
QuotaAccessState, QuotaEvidence, QuotaObservation, QuotaState, QuotaWriteOutcome,
RotationStrategy,
AccountAffinity, AccountAttemptFeedback, AccountConcurrencyLimit, AccountErrorReason,
AccountFeedbackStats, AccountSelectionPolicy, AccountStateChange, AccountWeight,
CredentialState, OpaqueProviderData, ProviderAccount, ProviderAccountId,
ProviderAccountStore as _, QuotaAccessChange, QuotaAccessState, QuotaEvidence,
QuotaObservation, QuotaState, QuotaWriteOutcome, RotationStrategy,
};
use gateway_core::engine::continuation::{
ContinuationBinding, NativeContinuationPin, PreviousResponseId,
@@ -263,6 +263,7 @@ fn selector_with_runtime(
cooldowns,
account_feedback,
CodexCookiePolicy::official().expect("official cookie policy"),
store.clone(),
)
}
@@ -795,8 +796,8 @@ fn selector_does_not_retry_invalid_credential_data() {
assert!(leases.requests.lock().expect("lease requests").is_empty());
}
#[test]
fn selector_round_robin_cursor_advances_across_requests() {
#[tokio::test]
async fn selector_round_robin_cursor_advances_across_requests() {
let store = Arc::new(MemoryAccountStore::default());
create_account(&store, "acct_first", "at-first");
create_account(&store, "acct_second", "at-second");
@@ -807,13 +808,15 @@ fn selector_round_robin_cursor_advances_across_requests() {
for _ in 0..4 {
let attempt = round_robin_attempt();
let lease = block_on(selector.select(&SelectCodexCredential {
upstream_model: "gpt-5.4",
request_url: &request_url,
attempt: &attempt,
session_affinity_key: None,
}))
.expect("select round robin account");
let lease = selector
.select(&SelectCodexCredential {
upstream_model: "gpt-5.4",
request_url: &request_url,
attempt: &attempt,
session_affinity_key: None,
})
.await
.expect("select round robin account");
selected.push(lease.account_id().as_str().to_owned());
}
@@ -823,6 +826,86 @@ fn selector_round_robin_cursor_advances_across_requests() {
);
}
#[tokio::test]
async fn strict_round_robin_rotates_each_new_session_across_eligible_accounts() {
let store = Arc::new(MemoryAccountStore::default());
create_account(&store, "acct_fallback", "at-fallback");
create_account(&store, "acct_first", "at-first");
create_account(&store, "acct_second", "at-second");
let leases = Arc::new(TestLeaseCoordinator::default());
let selector = selector_with_affinity(
&store,
Arc::clone(&leases),
Arc::new(MemorySessionAffinity::default()),
);
let request_url =
Url::parse("https://chatgpt.com/backend-api/codex/responses").expect("request URL");
assert_eq!(
select_strict_round_robin_session(&selector, &request_url, "session-1").await,
"acct_fallback"
);
for _ in 0..5 {
assert_eq!(
select_strict_round_robin_session(&selector, &request_url, "session-1").await,
"acct_fallback"
);
}
assert_eq!(
select_strict_round_robin_session(&selector, &request_url, "session-2").await,
"acct_first"
);
assert_eq!(
select_strict_round_robin_session(&selector, &request_url, "session-1").await,
"acct_fallback"
);
assert_eq!(
select_strict_round_robin_session(&selector, &request_url, "session-3").await,
"acct_second"
);
assert_eq!(
select_strict_round_robin_session(&selector, &request_url, "session-4").await,
"acct_fallback"
);
}
async fn select_strict_round_robin_session(
selector: &CodexCredentialSelector,
request_url: &Url,
key: &str,
) -> String {
let attempt = AttemptContext::new(
RequestAttemptContext::new(
ModelRequestId::new("req_codex_round_robin_session").expect("request id"),
ClientApiKeyId::new("key_codex_contract").expect("client key id"),
),
NonZeroU32::new(1).expect("attempt"),
SystemTime::now() + Duration::from_secs(30),
AccountSelectionPolicy::new(
RotationStrategy::RoundRobin,
NonZeroU32::new(2).expect("concurrency"),
)
.with_openai_account_affinity(AccountAffinity::Strict),
AccountAttemptContext::new(BTreeSet::new(), None, None)
.with_account_scope(contract_account_scope()),
None,
CancellationToken::new(),
);
let affinity_key = ProviderSessionAffinityKey::try_new(key).expect("affinity key");
selector
.select(&SelectCodexCredential {
upstream_model: "gpt-5.4",
request_url,
attempt: &attempt,
session_affinity_key: Some(&affinity_key),
})
.await
.expect("select round robin session")
.account_id()
.as_str()
.to_owned()
}
#[tokio::test]
async fn selector_should_claim_the_initial_session_account_before_upstream_send() {
let store = Arc::new(MemoryAccountStore::default());
@@ -149,6 +149,41 @@ async fn cache_preference_storage_timeouts_follow_the_affinity_mode() {
}
}
#[tokio::test]
async fn round_robin_retains_the_assignment_when_a_cache_preference_cannot_be_saved() {
use gateway_core::account::{AccountAffinity, AccountSelectionPolicy, RotationStrategy};
use gateway_core::provider_ports::ProviderSessionAffinityPort;
let store = Arc::new(MemoryAccountStore::default());
create_account(&store, "acct_affinity").await;
let affinity = Arc::new(MemorySessionAffinity::with_initial_claim_barrier(2));
let server = MockServer::start().await;
let provider = provider_with_affinity_and_base_url(&store, affinity.clone(), server.uri());
let attempt = session_binding::context_with_affinity_policy(
"req_cache_rr",
contract_account_scope(),
AccountSelectionPolicy::new(RotationStrategy::RoundRobin, NonZeroU32::MIN)
.with_openai_account_affinity(AccountAffinity::Preferred),
);
let key = attempt.client_api_key_ref().clone();
let result = provider
.execute(planned_request("openai", generate_operation()), attempt)
.await
.unwrap();
assert_eq!(
result.metadata().provider_account_id().as_str(),
"acct_affinity"
);
assert_eq!(affinity.binding_count(), 0);
assert_eq!(
affinity
.round_robin_cursor(&ProviderKind::new("openai").unwrap(), &key)
.await
.unwrap(),
1
);
}
#[tokio::test]
async fn strict_cache_preference_waits_and_preserves_binding_on_rejected_waits() {
use gateway_core::concurrency::ConcurrencyQueuePolicy;
@@ -0,0 +1,239 @@
//! 验证大账号池与慢目录上游不会阻塞启动时的能力读取
use super::*;
use gateway_core::engine::provider::{Provider, ProviderRegistry};
use gateway_core::routing::snapshot::{
RuntimeSnapshotCompiler, SnapshotFacts, SnapshotProviderAccountFacts, SnapshotStoreError,
SnapshotStorePort,
};
use std::sync::atomic::{AtomicUsize, Ordering};
use tokio::sync::Notify;
use wiremock::{Request, Respond};
struct StartupStore(SnapshotFacts);
impl SnapshotStorePort for StartupStore {
fn load_snapshot_facts(&self) -> BoxFuture<'_, Result<SnapshotFacts, SnapshotStoreError>> {
Box::pin(async { Ok(self.0.clone()) })
}
fn current_config_revision(&self) -> BoxFuture<'_, Result<ConfigRevision, SnapshotStoreError>> {
Box::pin(async { Ok(self.0.config_revision()) })
}
}
#[tokio::test]
async fn cold_catalog_with_ten_thousand_accounts_never_fetches_during_capability_reads() {
let store = Arc::new(MemoryAccountStore::default());
let mut accounts = Vec::new();
for index in 0..10_000 {
let id = format!("acct_startup_{index:05}");
create_account(&store, &id).await;
accounts.push(SnapshotProviderAccountFacts::new(
ProviderAccountId::new(id).unwrap(),
"openai",
));
}
let server = MockServer::start().await;
Mock::given(method("GET"))
.respond_with(ResponseTemplate::new(503).set_delay(Duration::from_secs(30)))
.expect(0)
.mount(&server)
.await;
let revision = ConfigRevision::new(1).unwrap();
let facts = Arc::new(StartupStore(SnapshotFacts::new(
revision,
revision,
SettingsValues::new(3, "smart", BTreeMap::new(), None, None),
vec![],
vec![],
accounts,
vec![],
)));
// 重建 Provider 模拟重启;Core 的读取只报告目录未知,不触发全量账号扫描
for _ in 0..3 {
let provider = provider_with_base_url(&store, server.uri());
let error =
tokio::time::timeout(Duration::from_secs(1), provider.query_model_capabilities())
.await
.expect("startup capability read must not wait for upstream")
.expect_err("a cold catalog is unknown");
assert_eq!(error.kind(), ProviderErrorKind::Unavailable);
let registry = Arc::new(ProviderRegistry::new([provider as Arc<dyn Provider>]).unwrap());
let compiler = RuntimeSnapshotCompiler::new(facts.clone(), registry);
let snapshot = tokio::time::timeout(Duration::from_secs(5), compiler.compile())
.await
.expect("startup snapshot must not wait for the account scan")
.expect("cold catalog must not prevent snapshot publication");
assert!(snapshot.public_models().is_empty());
}
server.verify().await;
}
struct SlowCatalog {
calls: Arc<AtomicUsize>,
first_batch: Arc<Notify>,
}
impl Respond for SlowCatalog {
fn respond(&self, _: &Request) -> ResponseTemplate {
if self.calls.fetch_add(1, Ordering::SeqCst) + 1 == 32 {
self.first_batch.notify_one();
}
ResponseTemplate::new(200)
.set_body_raw(OFFICIAL_FIXTURE.to_vec(), "application/json")
.set_delay(Duration::from_secs(30))
}
}
#[tokio::test]
async fn slow_background_catalog_is_bounded_and_does_not_block_capability_reads() {
let store = Arc::new(MemoryAccountStore::default());
for index in 0..64 {
create_account(&store, &format!("acct_background_{index:02}")).await;
}
let server = MockServer::start().await;
let calls = Arc::new(AtomicUsize::new(0));
let first_batch = Arc::new(Notify::new());
Mock::given(method("GET"))
.respond_with(SlowCatalog {
calls: Arc::clone(&calls),
first_batch: Arc::clone(&first_batch),
})
.mount(&server)
.await;
let (provider, _, _, catalog) = provider_and_quota_with_runtime_ports(
&store,
Arc::new(MemorySessionAffinity::default()),
server.uri(),
Arc::new(TestLeaseCoordinator::default()),
u32::try_from(DEFAULT_STREAM_MAX_RETRIES).unwrap(),
Arc::new(MemoryCooldownPort::new()),
crate::support::runtime_policy(),
);
let refresh = {
let catalog = Arc::clone(&catalog);
tokio::spawn(async move { catalog.refresh_catalogs().await })
};
tokio::time::timeout(Duration::from_secs(5), first_batch.notified())
.await
.expect("background scan must make progress beyond four accounts");
assert_eq!(calls.load(Ordering::SeqCst), 32);
assert!(
tokio::time::timeout(Duration::from_secs(1), provider.query_model_capabilities(),)
.await
.expect("capability read must not join a slow refresh")
.is_err()
);
// 取消整轮扫描后刷新锁必须释放,下一轮后台任务仍能完成
refresh.abort();
assert!(refresh.await.unwrap_err().is_cancelled());
server.reset().await;
Mock::given(method("GET"))
.respond_with(
ResponseTemplate::new(200).set_body_raw(OFFICIAL_FIXTURE.to_vec(), "application/json"),
)
.expect(64)
.mount(&server)
.await;
tokio::time::timeout(Duration::from_secs(5), catalog.refresh_catalogs())
.await
.expect("refresh cancellation must release the lock")
.expect("catalog refresh");
let generation = provider.catalog_generation();
for _ in 0..3 {
assert!(provider.query_model_capabilities().await.is_ok());
}
assert_eq!(provider.catalog_generation(), generation);
server.verify().await;
}
#[tokio::test]
async fn issue_666_same_plan_special_model_remains_visible_and_routes_only_to_its_source_after_restart()
{
let server = MockServer::start().await;
let store = Arc::new(MemoryAccountStore::default());
let source = "acct_provider_contract";
let other = "acct_scope_new";
for (id, models) in [
(source, vec!["common", "special-666"]),
(other, vec!["common"]),
] {
create_account(&store, id).await;
Mock::given(method("GET")).and(path("/codex/models")).and(header("chatgpt-account-id",format!("chatgpt-{id}")))
.respond_with(ResponseTemplate::new(200).set_body_json(json!({"models":models.iter().map(|id|json!({"slug":id,"display_name":id})).collect::<Vec<_>>()})))
.expect(1).mount(&server).await;
}
// 调度本来优先选 B,只有 A 自己的目录能授权特殊模型
store.set_scheduling(
other,
None,
gateway_core::account::AccountWeight::new(100).unwrap(),
);
Mock::given(method("POST"))
.and(path("/codex/responses"))
.and(header("chatgpt-account-id", format!("chatgpt-{source}")))
.respond_with(
ResponseTemplate::new(200).set_body_raw(CAPTURE_COMPLETED_SSE, "text/event-stream"),
)
.expect(2)
.mount(&server)
.await;
for restarted in [false, true] {
let (provider, _, _, catalog) = provider_and_quota_with_runtime_ports(
&store,
Arc::new(MemorySessionAffinity::default()),
server.uri(),
Arc::new(TestLeaseCoordinator::default()),
u32::try_from(DEFAULT_STREAM_MAX_RETRIES).unwrap(),
Arc::new(MemoryCooldownPort::new()),
crate::support::runtime_policy(),
);
catalog.restore().await.unwrap();
if !restarted {
assert!(provider.query_model_capabilities().await.is_err());
}
catalog.refresh_catalogs().await.unwrap();
assert!(
catalog
.client_model_catalog(&contract_account_scope())
.await
.unwrap()
.iter()
.any(|m| m.model.as_str() == "special-666")
);
let sources = catalog.cached().unwrap().unwrap();
assert_eq!(
sources.model_catalog_accounts()["special-666"],
std::collections::BTreeSet::from([gateway_core::account::ProviderAccountId::new(
source
)
.unwrap()])
);
let op = Operation::Generate(GenerateRequest::from_protocol_payload(
ProtocolPayload::json_object(
"openai",
json!({"model":"special-666","input":"hello"})
.as_object()
.unwrap()
.clone(),
)
.unwrap()
.with_context(Map::from_iter([("use_websocket".into(), json!(false))])),
));
let mut stream = provider
.execute(
planned_request_for_model("openai", op, "special-666"),
context("req_issue_666", CancellationToken::new()),
)
.await
.unwrap();
assert_eq!(stream.metadata().provider_account_id().as_str(), source);
while let Some(event) = stream.next().await {
event.unwrap();
}
}
server.verify().await;
}
@@ -35,6 +35,7 @@ fn cookie_provider(store: &Arc<MemoryAccountStore>, base_url: String) -> Arc<Cod
Arc::clone(&feedback),
// 测试仅允许本机模拟上游,生产 Cookie 白名单不放宽
CodexCookiePolicy::new(["cf_clearance"], ["127.0.0.1"]).unwrap(),
store.clone(),
));
Arc::new(
CodexProvider::new(
@@ -29,6 +29,7 @@ async fn seed(store: &Arc<MemoryAccountStore>, server: &MockServer, protocol: Ap
fn messages_operation(stream: bool) -> Operation {
let body = json!({
"model":"gpt-5.4", "max_tokens":128, "stream":stream,
"thinking":{"type":"adaptive"}, "output_config":{"effort":"xhigh"},
"messages":[
{"role":"assistant","content":[{"type":"thinking","thinking":"prior thought","signature":"real-native-signature"},{"type":"text","text":"before"}]},
{"role":"user","content":"hello"},
@@ -174,7 +175,8 @@ async fn native_messages_preserve_json_sse_headers_and_real_signatures() {
.await;
let operation = messages_operation(streaming);
let original = operation.middleware_body().unwrap();
let mut stream = provider(&store)
let provider = provider(&store);
let mut stream = provider
.execute(
planned_request("openai", operation),
context("req_messages_native", CancellationToken::new()),
@@ -290,6 +292,8 @@ async fn messages_and_responses_translate_with_target_header_boundaries() {
assert!(!headers.contains_key(name));
}
} else {
let body: Value = serde_json::from_slice(&requests[0].body).unwrap();
assert_eq!(body["reasoning"]["effort"], "xhigh");
for name in [
"anthropic-version",
"anthropic-beta",
@@ -302,6 +306,71 @@ async fn messages_and_responses_translate_with_target_header_boundaries() {
}
}
#[tokio::test]
async fn translated_messages_service_tiers_follow_the_selected_responses_upstream() {
for api_key in [false, true] {
for (options, oauth_tier, api_tier) in [
(json!({"service_tier":"auto"}), None, Some("auto")),
(
json!({"service_tier":"standard_only"}),
None,
Some("default"),
),
(json!({"speed":"fast"}), Some("priority"), Some("priority")),
] {
let server = MockServer::start().await;
let store = Arc::new(MemoryAccountStore::default());
if api_key {
seed(&store, &server, ApiKeyProtocol::Responses).await;
} else {
create_account(&store, "acct_provider_contract").await;
}
Mock::given(method("POST"))
.and(path(if api_key {
"/responses"
} else {
"/codex/responses"
}))
.respond_with(
ResponseTemplate::new(200)
.set_body_raw(CAPTURE_COMPLETED_SSE, "text/event-stream"),
)
.expect(1)
.mount(&server)
.await;
let mut body = json!({"model":"gpt-5.4","max_tokens":128,"stream":true,"messages":[{"role":"user","content":"hello"}]});
body.as_object_mut()
.unwrap()
.extend(options.as_object().unwrap().clone());
let operation = Operation::Generate(GenerateRequest::from_protocol_payload(
ProtocolPayload::json_object("anthropic", body.as_object().unwrap().clone())
.unwrap()
.with_context(Map::from_iter([("use_websocket".to_owned(), json!(false))])),
));
let mut stream = provider_with_base_url(&store, server.uri())
.execute(
planned_request("openai", operation),
context("req_translated_tier", CancellationToken::new()),
)
.await
.unwrap();
while let Some(event) = stream.next().await {
event.unwrap();
}
let requests = server.received_requests().await.unwrap();
let sent = captured_request_body(&requests[0]);
let expected = if api_key { api_tier } else { oauth_tier };
assert_eq!(sent.get("service_tier").and_then(Value::as_str), expected);
let hint = expected.map_or_else(
|| "model=gpt-5.4".to_owned(),
|tier| format!("model=gpt-5.4;tier={tier}"),
);
assert_eq!(requests[0].headers["x-codex-routing-hint"], hint);
server.verify().await;
}
}
}
#[tokio::test]
async fn custom_messages_and_responses_count_locally_without_billable_usage() {
for (target, source) in [
@@ -500,7 +569,7 @@ async fn translated_thinking_uses_selected_account_catalog_and_preserves_fast_be
.mount(server)
.await;
}
let provider = provider(&store);
let provider = provider_with_catalog_and_base_url(&store, adaptive.uri()).await;
for _ in 0..2 {
for account in ["acct_provider_contract", "acct_manual"] {
let body = json!({"model":"unrelated-client-alias","input":"hello","max_output_tokens":4096,"reasoning":{"effort":"xhigh"},"tools":[{"type":"function","name":"echo"}],"parallel_tool_calls":false,"service_tier":"priority"});
@@ -558,50 +627,20 @@ async fn translated_thinking_uses_selected_account_catalog_and_preserves_fast_be
}
#[tokio::test]
async fn cancelled_capability_lookup_never_sends_generation() {
async fn cancelled_messages_request_never_sends_generation_or_catalog_lookup() {
let server = MockServer::start().await;
let store = Arc::new(MemoryAccountStore::default());
seed(&store, &server, ApiKeyProtocol::Messages).await;
Mock::given(method("GET"))
.and(path("/models"))
.respond_with(
ResponseTemplate::new(200)
.set_delay(Duration::from_secs(10))
.set_body_json(json!({"data":[]})),
)
.mount(&server)
.await;
let cancellation = CancellationToken::new();
let cancelled = cancellation.clone();
let task = tokio::spawn(async move {
provider(&store)
.execute(
planned_request("openai", generate_operation()),
context("req_cancel_catalog", cancelled),
)
.await
});
tokio::time::timeout(Duration::from_secs(2), async {
while server.received_requests().await.unwrap().is_empty() {
tokio::task::yield_now().await;
}
})
.await
.unwrap();
cancellation.cancel();
let result = tokio::time::timeout(Duration::from_secs(1), task)
.await
.unwrap()
.unwrap();
let result = provider(&store)
.execute(
planned_request("openai", generate_operation()),
context("req_cancel_catalog", cancellation),
)
.await;
assert!(matches!(result,Err(error) if error.kind()==ProviderErrorKind::Cancelled));
assert!(
server
.received_requests()
.await
.unwrap()
.iter()
.all(|request| request.method == "GET")
);
assert!(server.received_requests().await.unwrap().is_empty());
}
#[tokio::test]
@@ -759,7 +798,7 @@ async fn oversized_messages_json_is_rejected_before_wire_delivery() {
}
#[tokio::test]
async fn slow_catalog_does_not_block_generation_and_populates_next_turn() {
async fn cold_directory_never_queries_models_during_generation() {
let server = MockServer::start().await;
let store = Arc::new(MemoryAccountStore::default());
seed(&store, &server, ApiKeyProtocol::Messages).await;
@@ -772,7 +811,7 @@ async fn slow_catalog_does_not_block_generation_and_populates_next_turn() {
json!({"data":[{"id":"gpt-5.4","max_tokens":2048}],"has_more":false}),
),
)
.expect(1)
.expect(0)
.mount(&server)
.await;
Mock::given(method("POST"))
@@ -782,7 +821,7 @@ async fn slow_catalog_does_not_block_generation_and_populates_next_turn() {
.mount(&server)
.await;
let provider = provider(&store);
for i in 0..2 {
for _ in 0..2 {
let mut stream = provider
.clone()
.execute(
@@ -794,9 +833,6 @@ async fn slow_catalog_does_not_block_generation_and_populates_next_turn() {
while let Some(event) = stream.next().await {
event.unwrap();
}
if i == 0 {
tokio::time::sleep(Duration::from_secs(2)).await;
}
}
let limits = server
.received_requests()
@@ -810,6 +846,6 @@ async fn slow_catalog_does_not_block_generation_and_populates_next_turn() {
.unwrap()
})
.collect::<Vec<_>>();
assert_eq!(limits, [32000, 2048]);
assert_eq!(limits, [32000, 32000]);
server.verify().await;
}
@@ -4,6 +4,7 @@ mod account_eligibility;
mod account_isolation;
mod cache_preference;
mod capacity;
mod catalog_startup;
mod cookies;
mod correlation;
mod error_details;
@@ -28,7 +29,7 @@ mod upstream_adapter;
use std::collections::{BTreeMap, BTreeSet};
use std::io::{self, Write};
use std::num::NonZeroU32;
use std::sync::{Arc, Mutex};
use std::sync::{Arc, Mutex, atomic::Ordering};
use std::time::{Duration, SystemTime};
use base64::{Engine as _, engine::general_purpose::STANDARD};
@@ -1207,7 +1208,7 @@ fn provider_and_quota_with_affinity_and_base_url_and_leases(
leases: Arc<TestLeaseCoordinator>,
stream_max_retries: u32,
) -> (Arc<CodexProvider>, Arc<CodexCredentialQuotaService>) {
let (provider, quota, _) = provider_and_quota_with_runtime_ports(
let (provider, quota, _, _) = provider_and_quota_with_runtime_ports(
store,
session_affinity,
base_url,
@@ -1219,6 +1220,44 @@ fn provider_and_quota_with_affinity_and_base_url_and_leases(
(provider, quota)
}
async fn provider_with_catalog_and_base_url(
store: &Arc<MemoryAccountStore>,
base_url: String,
) -> Arc<CodexProvider> {
provider_and_quota_with_catalog(
store,
Arc::new(MemorySessionAffinity::default()),
base_url,
Arc::new(TestLeaseCoordinator::default()),
u32::try_from(DEFAULT_STREAM_MAX_RETRIES).unwrap(),
)
.await
.0
}
async fn provider_and_quota_with_catalog(
store: &Arc<MemoryAccountStore>,
session_affinity: Arc<MemorySessionAffinity>,
base_url: String,
leases: Arc<TestLeaseCoordinator>,
stream_max_retries: u32,
) -> (Arc<CodexProvider>, Arc<CodexCredentialQuotaService>) {
let (provider, quota, _, catalog) = provider_and_quota_with_runtime_ports(
store,
session_affinity,
base_url,
leases,
stream_max_retries,
Arc::new(MemoryCooldownPort::new()),
crate::support::runtime_policy(),
);
catalog
.refresh_catalogs()
.await
.expect("background catalog");
(provider, quota)
}
fn provider_and_quota_with_runtime_ports(
store: &Arc<MemoryAccountStore>,
session_affinity: Arc<MemorySessionAffinity>,
@@ -1231,6 +1270,7 @@ fn provider_and_quota_with_runtime_ports(
Arc<CodexProvider>,
Arc<CodexCredentialQuotaService>,
Arc<CodexWebSocketPool>,
Arc<CodexCredentialCatalogService>,
) {
let profile = wire_profile();
let http = reqwest::Client::builder()
@@ -1265,11 +1305,12 @@ fn provider_and_quota_with_runtime_ports(
cooldowns,
Arc::clone(&account_feedback),
CodexCookiePolicy::official().expect("cookie policy"),
store.clone(),
));
let provider = CodexProvider::new(
selector,
catalog,
Arc::clone(&catalog),
Arc::clone(&quota),
account_feedback,
http,
@@ -1279,7 +1320,7 @@ fn provider_and_quota_with_runtime_ports(
stream_max_retries,
)
.expect("official OpenAI provider");
(Arc::new(provider), quota, websocket_pool)
(Arc::new(provider), quota, websocket_pool, catalog)
}
pub(super) async fn create_account(store: &Arc<MemoryAccountStore>, id: &str) {
@@ -1709,13 +1750,21 @@ fn diagnostic_context(request_id: &str, account_id: &str) -> AttemptContext {
)
}
fn request_with_peer_route(request_id: &str) -> RequestAttemptContext {
RequestAttemptContext::new(
ModelRequestId::new(request_id).unwrap(),
ClientApiKeyId::new("key_openai_contract").unwrap(),
)
.with_request_location(Some(global_request_location()))
.with_extension_scope(
gateway_core::engine::extensions::ExtensionCallScope::default()
.with_peer_route("0123456789abcdef0123456789abcdef".parse().unwrap()),
)
}
fn probe_context(request_id: &str, target: AccountProbeTarget) -> AttemptContext {
AttemptContext::new(
RequestAttemptContext::new(
ModelRequestId::new(request_id).expect("request id"),
ClientApiKeyId::new("key_openai_contract").expect("client key id"),
)
.with_request_location(Some(global_request_location())),
request_with_peer_route(request_id),
NonZeroU32::new(1).expect("attempt"),
SystemTime::now() + Duration::from_secs(30),
account_policy(),
@@ -2832,7 +2881,7 @@ async fn image_endpoints_bypass_only_the_text_catalog_and_preserve_the_current_c
.await;
}
let provider = provider_with_base_url(&store, server.uri());
let provider = provider_with_catalog_and_base_url(&store, server.uri()).await;
let catalog = provider
.query_model_capabilities()
.await
@@ -8286,7 +8335,7 @@ async fn responses_should_observe_and_bill_the_outbound_service_tier_on_both_tra
false,
Some(json!("default")),
Some("priority"),
None,
Some("default"),
Some(3_437_500),
),
(
@@ -8302,7 +8351,7 @@ async fn responses_should_observe_and_bill_the_outbound_service_tier_on_both_tra
false,
Some(json!("priority")),
Some("priority"),
None,
Some("default"),
Some(3_437_500),
),
(
@@ -8310,7 +8359,7 @@ async fn responses_should_observe_and_bill_the_outbound_service_tier_on_both_tra
false,
Some(json!("fast")),
None,
None,
Some("default"),
Some(3_437_500),
),
(
@@ -8318,7 +8367,7 @@ async fn responses_should_observe_and_bill_the_outbound_service_tier_on_both_tra
false,
Some(json!("default")),
None,
None,
Some("default"),
Some(3_437_500),
),
(FastMode::Disabled, false, None, None, None, Some(3_437_500)),
@@ -8391,8 +8440,8 @@ async fn responses_should_observe_and_bill_the_outbound_service_tier_on_both_tra
true,
Some(json!("auto")),
None,
Some("auto"),
None,
Some(3_437_500),
),
(FastMode::Enabled, false, None, None, None, Some(3_437_500)),
(
@@ -8400,7 +8449,7 @@ async fn responses_should_observe_and_bill_the_outbound_service_tier_on_both_tra
false,
Some(json!("default")),
None,
None,
Some("default"),
Some(3_437_500),
),
] {
@@ -9138,7 +9187,7 @@ fn provider_with_cooldowns(
base_url: String,
cooldowns: Arc<MemoryCooldownPort>,
) -> (Arc<CodexProvider>, Arc<CodexWebSocketPool>) {
let (provider, _, pool) = provider_and_quota_with_runtime_ports(
let (provider, _, pool, _) = provider_and_quota_with_runtime_ports(
store,
Arc::new(MemorySessionAffinity::default()),
base_url,
@@ -10458,80 +10507,6 @@ async fn successful_http_sse_rate_limit_event_persists_structured_exhaustion() {
assert!(store.has_quota(account_id));
}
#[test]
fn request_observation_reads_openai_metadata() {
let store = Arc::new(MemoryAccountStore::default());
let payload = ProtocolPayload::json_object(
"openai",
Map::from_iter([
("model".to_owned(), json!("gpt-5.4")),
("input".to_owned(), json!("hello")),
]),
)
.expect("OpenAI payload")
.with_context(Map::from_iter([(
"turn_metadata".to_owned(),
serde_json::Value::String(
r#"{"request_kind":"review","subagent_kind":"worker"}"#.to_owned(),
),
)]));
let generation = GenerateRequest::from_protocol_payload(payload);
let operation = Operation::Generate(generation);
let client_key_id = ClientApiKeyId::new("key_openai_observation").expect("client key");
let observation = provider(&store).request_observation(&operation, &client_key_id);
assert_eq!(observation.request_kind.as_deref(), Some("review"));
assert_eq!(observation.subagent_kind.as_deref(), Some("worker"));
assert!(!observation.compact);
}
#[test]
fn request_observation_classifies_prewarm_from_generate() {
let store = Arc::new(MemoryAccountStore::default());
let provider = provider(&store);
let client_key_id = ClientApiKeyId::new("key_prewarm_observation").expect("client key");
for (generate, request_kind, expected_kind) in [
(Some(json!(false)), None, Some("prewarm")),
(Some(json!(false)), Some("prewarm"), Some("prewarm")),
(Some(json!(false)), Some("review"), Some("prewarm")),
(Some(json!(true)), Some("prewarm"), None),
(None, Some("prewarm"), None),
(Some(json!(null)), Some("prewarm"), None),
(Some(json!("false")), Some("prewarm"), None),
(None, Some("review"), Some("review")),
(None, None, None),
] {
let mut body = Map::from_iter([
("model".to_owned(), json!("gpt-test")),
("input".to_owned(), json!("hello")),
("store".to_owned(), json!(false)),
]);
if let Some(generate) = generate {
body.insert("generate".to_owned(), generate);
}
let payload = ProtocolPayload::json_object("openai", body.clone())
.expect("OpenAI payload")
.with_context(Map::from_iter([(
"turn_metadata".to_owned(),
Value::String(
json!({"request_kind": request_kind, "subagent_kind": "worker"}).to_string(),
),
)]));
let operation = Operation::Generate(GenerateRequest::from_protocol_payload(payload));
let observation = provider.request_observation(&operation, &client_key_id);
assert_eq!(
observation.request_kind.as_deref(),
expected_kind,
"generate={:?}, request_kind={request_kind:?}",
body.get("generate"),
);
assert_eq!(observation.subagent_kind.as_deref(), Some("worker"));
}
}
#[test]
fn endpoint_observation_should_read_models_without_requiring_a_catalog() {
let store = Arc::new(MemoryAccountStore::default());
@@ -10571,26 +10546,6 @@ fn endpoint_observation_should_read_models_without_requiring_a_catalog() {
}
}
#[test]
fn request_observation_preserves_the_raw_reasoning_effort() {
let store = Arc::new(MemoryAccountStore::default());
for (effort, expected) in [
(json!("future-value"), "future-value"),
(json!(64), "64"),
(json!(0), "0"),
] {
let payload = ProtocolPayload::json_object(
"openai",
Map::from_iter([("reasoning".to_owned(), json!({"effort":effort}))]),
)
.unwrap();
let operation = Operation::Generate(GenerateRequest::from_protocol_payload(payload));
let client_key_id = ClientApiKeyId::new("key_openai_observation").unwrap();
let observation = provider(&store).request_observation(&operation, &client_key_id);
assert_eq!(observation.reasoning_effort.as_deref(), Some(expected));
}
}
#[tokio::test]
async fn reasoning_effort_observation_follows_configuration_updates_across_continuations() {
fn completed(index: usize) -> Value {
@@ -10696,7 +10651,19 @@ async fn reasoning_effort_observation_follows_configuration_updates_across_conti
generate = generate.with_provider_session_state(state);
}
let operation = Operation::Generate(generate);
let observation = provider.request_observation(&operation, &client_key);
let provider_observation = provider.request_observation(&operation, &client_key);
let Operation::Generate(generate) = &operation else {
panic!("generate operation");
};
let payload = generate.protocol_payload();
let observation = gateway_protocol::request::request_observation(
payload.protocol(),
payload.body(),
payload.context(),
provider_observation
.previous_reasoning_effort_override
.as_deref(),
);
assert_eq!(observation.reasoning_effort.as_deref(), Some(expected));
let binding = previous.as_ref().map(|previous| {
ContinuationBinding::Pinned(NativeContinuationPin::new(
@@ -10768,6 +10735,60 @@ async fn reasoning_effort_observation_follows_configuration_updates_across_conti
}
}
#[test]
fn request_observation_only_inherits_native_openai_session_configuration() {
let store = Arc::new(MemoryAccountStore::default());
let provider = provider(&store);
let client = ClientApiKeyId::new("key_session_observation").unwrap();
for (owner, extension_owned, expected) in [
("openai", false, Some("high")),
("xai", false, None),
("openai", true, None),
] {
let mut state = ProviderSessionState::new(
owner,
json!({
"account_id":"acct_observation", "conversation_id":"session",
"reasoning_effort_override":"high", "continuation_scope":"persisted"
})
.as_object()
.unwrap()
.clone(),
)
.unwrap();
if extension_owned {
state = state.with_extension_owner(gateway_core::operation::ExtensionSessionOwner {
instance_id: "plugin".into(),
contribution_id: "adapter".into(),
adapter_id: "adapter".into(),
generation: 1,
incarnation: "incarnation".into(),
connection_local: true,
});
}
let operation = Operation::Generate(
GenerateRequest::from_protocol_payload(
ProtocolPayload::json_object(
"openai",
json!({"previous_response_id":"resp_previous"})
.as_object()
.unwrap()
.clone(),
)
.unwrap(),
)
.with_provider_session_state(state),
);
assert_eq!(
provider
.request_observation(&operation, &client)
.previous_reasoning_effort_override
.as_deref(),
expected
);
}
}
#[tokio::test]
async fn provider_compiles_catalog_presentation_for_codex_models() {
let store = Arc::new(MemoryAccountStore::default());
@@ -10782,7 +10803,7 @@ async fn provider_compiles_catalog_presentation_for_codex_models() {
)
.mount(&server)
.await;
let provider = provider_with_base_url(&store, server.uri());
let provider = provider_with_catalog_and_base_url(&store, server.uri()).await;
let capabilities = provider
.query_model_capabilities()
@@ -10847,14 +10868,7 @@ async fn provider_compiles_catalog_presentation_for_codex_models() {
})
.expect("document");
assert_eq!(document, original["models"][0]);
assert!(
server
.received_requests()
.await
.expect("requests")
.iter()
.any(|request| request.url.query() == Some("client_version=0.154.0"))
);
assert_eq!(server.received_requests().await.unwrap().len(), 1);
}
#[tokio::test]
@@ -10881,7 +10895,7 @@ async fn provider_preserves_independent_catalog_context_windows() {
.expect(1)
.mount(&server)
.await;
let provider = provider_with_base_url(&store, server.uri());
let provider = provider_with_catalog_and_base_url(&store, server.uri()).await;
let capabilities = provider
.query_model_capabilities()
@@ -10916,7 +10930,7 @@ async fn provider_routes_catalog_model_when_supported_in_api_is_false() {
)
.mount(&server)
.await;
let provider = provider_with_base_url(&store, server.uri());
let provider = provider_with_catalog_and_base_url(&store, server.uri()).await;
let capabilities = provider
.query_model_capabilities()
@@ -11482,10 +11496,14 @@ async fn connection_limit_payload_survives_exhausted_retry_budget() {
#[tokio::test]
async fn api_key_native_endpoints_preserve_bodies_headers_and_own_base_url() {
for prefix in ["", "/v1", "/custom/v2"] {
for (prefix, discovered) in ["", "/v1", "/custom/v2"]
.into_iter()
.flat_map(|prefix| [false, true].map(|discovered| (prefix, discovered)))
{
let upstream = MockServer::start().await;
let oauth = MockServer::start().await;
let store = Arc::new(MemoryAccountStore::default());
store.peer_discovered.store(discovered, Ordering::SeqCst);
store
.seed_api_key(
"acct_provider_contract",
@@ -11503,10 +11521,6 @@ async fn api_key_native_endpoints_preserve_bodies_headers_and_own_base_url() {
Mock::given(method("POST"))
.and(path(format!("{prefix}{endpoint}")))
.and(header("authorization", "Bearer sk-api-test-only"))
.and(header(
gateway_protocol::peer::ROUTE_HEADER,
"0123456789abcdef0123456789abcdef",
))
.and(body_bytes(body.to_vec()))
.respond_with(
ResponseTemplate::new(200)
@@ -11575,6 +11589,13 @@ async fn api_key_native_endpoints_preserve_bodies_headers_and_own_base_url() {
}
assert!(oauth.received_requests().await.unwrap().is_empty());
for request in upstream.received_requests().await.unwrap() {
assert_eq!(
request
.headers
.get(gateway_protocol::peer::ROUTE_HEADER)
.map(|value| value.to_str().unwrap()),
discovered.then_some("0123456789abcdef0123456789abcdef")
);
for name in ["cookie", "chatgpt-account-id"] {
assert!(!request.headers.contains_key(name));
}
@@ -11616,7 +11637,8 @@ async fn api_key_responses_forward_lite_memgen_and_native_compaction_with_catalo
.expect(1)
.mount(&upstream)
.await;
let provider = provider(&store);
let provider =
provider_with_catalog_and_base_url(&store, OFFICIAL_CODEX_BASE_URL.to_owned()).await;
provider.query_model_capabilities().await.unwrap();
let body = json!({"model":"gpt-5.4","input":[{"type":"compaction_trigger"}],"reasoning":{"effort":"future-effort"},"tools":[{"type":"future-tool"}],"future_option":{"nested":true}});
let operation = Operation::Generate(GenerateRequest::from_protocol_payload(
@@ -11701,10 +11723,14 @@ fn generate_with_downstream_headers() -> Operation {
#[tokio::test]
async fn api_key_default_http_uses_own_prefix_plain_json_and_only_own_authentication() {
for prefix in ["", "/v1", "/custom/v2"] {
for (prefix, discovered) in ["", "/v1", "/custom/v2"]
.into_iter()
.flat_map(|prefix| [false, true].map(|discovered| (prefix, discovered)))
{
let upstream = MockServer::start().await;
let oauth = MockServer::start().await;
let store = Arc::new(MemoryAccountStore::default());
store.peer_discovered.store(discovered, Ordering::SeqCst);
store
.seed_api_key(
"acct_provider_contract",
@@ -11712,19 +11738,15 @@ async fn api_key_default_http_uses_own_prefix_plain_json_and_only_own_authentica
provider_openai::credential::ResponsesTransport::Http,
)
.await;
// 后台发现不协商客户端版本;客户端目录独立请求并按实际版本缓存
for query in [None, Some("client_version=1.0.0")] {
Mock::given(method("GET"))
.and(path(format!("{prefix}/models")))
.and(move |request: &wiremock::Request| request.url.query() == query)
.and(header("authorization", "Bearer sk-api-test-only"))
.respond_with(
ResponseTemplate::new(200).set_body_json(json!({"data":[{"id":"gpt-5.4"}]})),
)
.expect(1)
.mount(&upstream)
.await;
}
Mock::given(method("GET"))
.and(path(format!("{prefix}/models")))
.and(header("authorization", "Bearer sk-api-test-only"))
.respond_with(
ResponseTemplate::new(200).set_body_json(json!({"data":[{"id":"gpt-5.4"}]})),
)
.expect(1)
.mount(&upstream)
.await;
Mock::given(method("POST"))
.and(path(format!("{prefix}/responses")))
.and(header("authorization", "Bearer sk-api-test-only"))
@@ -11734,13 +11756,14 @@ async fn api_key_default_http_uses_own_prefix_plain_json_and_only_own_authentica
.expect(1)
.mount(&upstream)
.await;
let (provider, quota) = provider_and_quota_with_affinity_and_base_url_and_leases(
let (provider, quota) = provider_and_quota_with_catalog(
&store,
Arc::new(MemorySessionAffinity::default()),
oauth.uri(),
Arc::new(TestLeaseCoordinator::default()),
DEFAULT_STREAM_MAX_RETRIES as u32,
);
)
.await;
provider
.query_model_capabilities()
.await
@@ -11764,7 +11787,10 @@ async fn api_key_default_http_uses_own_prefix_plain_json_and_only_own_authentica
let mut stream = provider
.execute(
planned_request("openai", generate_with_downstream_headers()),
context("req_api_http", CancellationToken::new()),
context_for_request(
request_with_peer_route("req_api_http"),
CancellationToken::new(),
),
)
.await
.expect("scheduled API account");
@@ -11779,6 +11805,13 @@ async fn api_key_default_http_uses_own_prefix_plain_json_and_only_own_authentica
.find(|request| request.method == "POST")
.unwrap();
assert!(!request.headers.contains_key("content-encoding"));
assert_eq!(
request
.headers
.get(gateway_protocol::peer::ROUTE_HEADER)
.map(|value| value.to_str().unwrap()),
discovered.then_some("0123456789abcdef0123456789abcdef")
);
for name in API_KEY_DOWNSTREAM_HEADERS
.iter()
.filter(|name| **name != "authorization")
@@ -11806,12 +11839,10 @@ async fn api_key_default_http_uses_own_prefix_plain_json_and_only_own_authentica
.iter()
.filter(|request| request.method == "GET")
.collect::<Vec<_>>();
assert_eq!(model_requests.len(), 2);
for query in [None, Some("client_version=1.0.0")] {
let model_request = model_requests
.iter()
.find(|request| request.url.query() == query)
.expect("separate background and client catalog requests");
assert_eq!(model_requests.len(), 1);
{
let model_request = model_requests[0];
assert_eq!(model_request.url.query(), Some("client_version=0.144.0"));
assert_eq!(
model_request.headers["authorization"],
"Bearer sk-api-test-only"
@@ -11985,72 +12016,82 @@ async fn api_key_http_account_is_rejected_before_websocket_warmup_or_old_revisio
#[tokio::test]
async fn api_key_websocket_uses_api_path_and_bearer_without_oauth_identity() {
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
let base = format!("http://{}/custom/v2", listener.local_addr().unwrap());
let store = Arc::new(MemoryAccountStore::default());
store
.seed_api_key(
"acct_provider_contract",
base,
provider_openai::credential::ResponsesTransport::PreferWebsocket,
)
.await;
let server = tokio::spawn(async move {
let (stream, _) = listener.accept().await.unwrap();
let mut websocket =
crate::transport::accept_codex_test_websocket_with(stream, |request, _| {
assert_eq!(request.uri().path(), "/custom/v2/responses");
assert_eq!(
request.headers()["authorization"],
"Bearer sk-api-test-only"
);
assert!(!request.headers().contains_key("chatgpt-account-id"));
assert!(!request.headers().contains_key("cookie"));
assert_eq!(request.headers()["originator"], "codex_cli_rs");
assert_eq!(request.headers()["version"], "0.144.0");
assert_eq!(
request.headers()["user-agent"],
wire_profile().snapshot().user_agent()
);
for name in API_KEY_DOWNSTREAM_HEADERS
.iter()
.filter(|name| **name != "authorization")
{
assert!(!request.headers().contains_key(*name), "leaked {name}");
}
assert_eq!(request.headers()["x-business-extension"], "keep");
for name in API_KEY_BUSINESS_HEADERS {
if matches!(*name, "session-id" | "thread-id") {
assert_scoped_id(
request.headers()[*name].to_str().unwrap(),
"downstream-value",
);
} else {
assert_eq!(request.headers()[*name], "downstream-value", "lost {name}");
}
}
})
for discovered in [false, true] {
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
let base = format!("http://{}/custom/v2", listener.local_addr().unwrap());
let store = Arc::new(MemoryAccountStore::default());
store.peer_discovered.store(discovered, Ordering::SeqCst);
store
.seed_api_key(
"acct_provider_contract",
base,
provider_openai::credential::ResponsesTransport::PreferWebsocket,
)
.await;
let frame = websocket.next().await.unwrap().unwrap();
let payload: Value = serde_json::from_str(frame.to_text().unwrap()).unwrap();
assert_eq!(payload["type"], "response.create");
websocket.send(Message::Text(json!({"type":"response.completed","response":{"id":"resp_api_ws","model":"gpt-5.4","status":"completed","output":[],"usage":{"input_tokens":1,"output_tokens":1,"total_tokens":2}}}).to_string().into())).await.unwrap();
});
let provider = provider(&store);
let mut stream = provider
.execute(
planned_request("openai", generate_with_downstream_headers()),
diagnostic_context("req_api_ws", "acct_provider_contract"),
)
.await
.unwrap();
while let Some(event) = stream.next().await {
event.expect("API WebSocket completion");
let server = tokio::spawn(async move {
let (stream, _) = listener.accept().await.unwrap();
let mut websocket =
crate::transport::accept_codex_test_websocket_with(stream, |request, _| {
assert_eq!(request.uri().path(), "/custom/v2/responses");
assert_eq!(
request.headers()["authorization"],
"Bearer sk-api-test-only"
);
assert!(!request.headers().contains_key("chatgpt-account-id"));
assert!(!request.headers().contains_key("cookie"));
assert_eq!(
request
.headers()
.get(gateway_protocol::peer::ROUTE_HEADER)
.map(|value| value.to_str().unwrap()),
discovered.then_some("0123456789abcdef0123456789abcdef")
);
assert_eq!(request.headers()["originator"], "codex_cli_rs");
assert_eq!(request.headers()["version"], "0.144.0");
assert_eq!(
request.headers()["user-agent"],
wire_profile().snapshot().user_agent()
);
for name in API_KEY_DOWNSTREAM_HEADERS
.iter()
.filter(|name| **name != "authorization")
{
assert!(!request.headers().contains_key(*name), "leaked {name}");
}
assert_eq!(request.headers()["x-business-extension"], "keep");
for name in API_KEY_BUSINESS_HEADERS {
if matches!(*name, "session-id" | "thread-id") {
assert_scoped_id(
request.headers()[*name].to_str().unwrap(),
"downstream-value",
);
} else {
assert_eq!(request.headers()[*name], "downstream-value", "lost {name}");
}
}
})
.await;
let frame = websocket.next().await.unwrap().unwrap();
let payload: Value = serde_json::from_str(frame.to_text().unwrap()).unwrap();
assert_eq!(payload["type"], "response.create");
websocket.send(Message::Text(json!({"type":"response.completed","response":{"id":"resp_api_ws","model":"gpt-5.4","status":"completed","output":[],"usage":{"input_tokens":1,"output_tokens":1,"total_tokens":2}}}).to_string().into())).await.unwrap();
});
let provider = provider(&store);
let mut stream = provider
.execute(
planned_request("openai", generate_with_downstream_headers()),
diagnostic_context("req_api_ws", "acct_provider_contract"),
)
.await
.unwrap();
while let Some(event) = stream.next().await {
event.expect("API WebSocket completion");
}
timeout(Duration::from_secs(5), server)
.await
.unwrap()
.unwrap();
}
timeout(Duration::from_secs(5), server)
.await
.unwrap()
.unwrap();
}
#[tokio::test]
@@ -12592,18 +12633,27 @@ async fn fast_policy_changes_preserve_the_websocket_continuation_and_meter_each_
let mut websocket =
crate::transport::accept_codex_test_websocket_with(socket, |request, response| {
// 提示头只在首次握手发送;后续档位由各自 response.create 正文指定
assert_eq!(request.headers()["x-codex-routing-hint"], "model=gpt-5.4");
assert_eq!(
request.headers()["x-codex-routing-hint"],
"model=gpt-5.4;tier=default"
);
response.headers_mut().insert(
"sec-websocket-extensions",
"permessage-deflate".parse().unwrap(),
);
})
.await;
for (index, tier) in [None, Some("priority"), None, None].into_iter().enumerate() {
for (index, tier) in ["default", "priority", "default", "default"]
.into_iter()
.enumerate()
{
let message = websocket.next().await.unwrap().unwrap();
let frame: Value = serde_json::from_str(message.to_text().unwrap()).unwrap();
assert_eq!(frame["type"], "response.create");
assert_eq!(frame.get("service_tier").and_then(Value::as_str), tier);
assert_eq!(
frame.get("service_tier").and_then(Value::as_str),
Some(tier)
);
assert_eq!(
frame.get("previous_response_id").cloned(),
index
@@ -12760,7 +12810,11 @@ async fn fast_policy_changes_preserve_the_websocket_continuation_and_meter_each_
);
assert_eq!(
observed_tier.as_deref(),
(fast_mode == FastMode::Enabled).then_some("priority")
Some(if fast_mode == FastMode::Enabled {
"priority"
} else {
"default"
})
);
assert_eq!(
costs,
@@ -13313,7 +13367,7 @@ async fn public_catalog_filters_each_api_account_before_union_and_preserves_rout
.mount(server)
.await;
}
let provider = provider_with_base_url(&store, first.uri());
let provider = provider_with_catalog_and_base_url(&store, first.uri()).await;
let registry = Arc::new(ProviderRegistry::new([provider as Arc<dyn Provider>]).unwrap());
let policy = |mode, models: &[&str]| {
AccountModelAccess::new(mode, models.iter().map(|id| (*id).to_owned()).collect()).unwrap()
@@ -222,12 +222,14 @@ async fn account_allowlists_override_catalog_omissions_but_weight_affinity_and_p
store.set_scheduling("acct_scope_new", None, AccountWeight::new(100).unwrap());
let affinity = Arc::new(MemorySessionAffinity::default());
let leases = Arc::new(TestLeaseCoordinator::default());
let provider = provider_with_affinity_and_base_url_and_leases(
let (provider, _) = provider_and_quota_with_catalog(
&store,
affinity.clone(),
native.uri(),
leases.clone(),
);
u32::try_from(DEFAULT_STREAM_MAX_RETRIES).unwrap(),
)
.await;
provider.query_model_capabilities().await.unwrap();
for (model, expected) in [
("gpt-6-luna", "acct_scope_new"),
@@ -366,7 +368,7 @@ async fn only_a_current_account_catalog_can_exclude_a_model() {
use gateway_core::account::{CredentialCasOutcome, CredentialCasUpdate, ProviderAccountUpdate};
use provider_openai::credential::ResponsesTransport;
for scenario in ["missing", "unknown_plan", "rotated"] {
for scenario in ["missing", "unknown_plan", "rotated", "expired"] {
let native = MockServer::start().await;
let api = MockServer::start().await;
let store = Arc::new(MemoryAccountStore::default());
@@ -396,7 +398,7 @@ async fn only_a_current_account_catalog_can_exclude_a_model() {
store
.seed_api_key("acct_scope_new", api.uri(), ResponsesTransport::Http)
.await;
let response = if scenario == "missing" {
let response = if matches!(scenario, "missing" | "expired") {
ResponseTemplate::new(503)
} else {
ResponseTemplate::new(200).set_body_json(json!({"data": [{"id": "gpt-6-luna"}]}))
@@ -421,7 +423,42 @@ async fn only_a_current_account_catalog_can_exclude_a_model() {
.expect(if scenario == "unknown_plan" { 0 } else { 1 })
.mount(upstream)
.await;
let provider = provider_with_base_url(&store, native.uri());
let provider = if scenario == "expired" {
use gateway_core::account::AccountCatalogObservation;
let a = store.account("acct_scope_new").unwrap();
store
.save_account_catalog(
&AccountCatalogObservation {
account_id: a.id().clone(),
revision: a.revision(),
plan: a.plan_type().map(str::to_owned),
etag: None,
observed_at: SystemTime::now() - Duration::from_secs(3600),
},
b"\x01{\"data\":[{\"id\":\"gpt-6-luna\"}]}",
)
.await
.unwrap();
let (provider, _, _, catalog) = provider_and_quota_with_runtime_ports(
&store,
Arc::new(MemorySessionAffinity::default()),
native.uri(),
Arc::new(TestLeaseCoordinator::default()),
u32::try_from(DEFAULT_STREAM_MAX_RETRIES).unwrap(),
Arc::new(MemoryCooldownPort::new()),
crate::support::runtime_policy(),
);
catalog.restore().await.unwrap();
// B 的失败不能清空上次成功结果,也不能继续以旧的缺失条目拒绝请求
catalog.refresh_catalogs().await.unwrap();
assert_eq!(
catalog.cached_account_models(&a).unwrap().unwrap(),
["gpt-6-luna"]
);
provider
} else {
provider_with_catalog_and_base_url(&store, native.uri()).await
};
provider.query_model_capabilities().await.unwrap();
if scenario == "rotated" {
let account = store.account("acct_scope_new").unwrap();
@@ -83,6 +83,7 @@ async fn outbound_http_sends_final_encoded_bytes_and_target_without_restoring_de
request.target = Target::Http { method: reqwest::Method::PATCH, uri: uri.parse().unwrap() };
request.headers.retain(|header| !matches!(header.name(), "user-agent" | "content-encoding" | "content-type" | "authorization" | "accept"));
request.headers.extend([
MiddlewareHeader::new(gateway_protocol::peer::ROUTE_HEADER, Bytes::from_static(b"plugin-owned-route")),
MiddlewareHeader::new("x-plugin-values", Bytes::from_static(b"first")),
MiddlewareHeader::new("x-plugin-values", Bytes::from_static(b"second")),
MiddlewareHeader::new("x-plugin-bytes", Bytes::from_static(b"\x80")),
@@ -122,6 +123,10 @@ async fn outbound_http_sends_final_encoded_bytes_and_target_without_restoring_de
);
assert_eq!(request.url.query(), Some("source=outbound"));
assert_eq!(request.headers["x-plugin-bytes"].as_bytes(), b"\x80");
assert_eq!(
request.headers[gateway_protocol::peer::ROUTE_HEADER],
"plugin-owned-route"
);
}
#[tokio::test]
@@ -302,13 +302,28 @@ async fn websocket_messages_preserve_text_and_metadata_despite_unrecognized_json
}
#[tokio::test]
async fn default_service_tiers_are_omitted_only_for_codex_oauth() {
for tier in ["auto", "default", "priority", "flex"] {
async fn responses_service_tiers_reach_each_upstream_on_both_transports() {
for tier in [
None,
Some(Value::Null),
Some(json!("auto")),
Some(json!("default")),
Some(json!("priority")),
Some(json!("flex")),
] {
for api_key in [false, true] {
let body = json!({"model":"gpt-5.4","input":"hello","service_tier":tier});
let sent = capture_downstream_request(body, api_key, false, false).await;
let expected = (api_key || !matches!(tier, "auto" | "default")).then_some(tier);
assert_eq!(sent.get("service_tier").and_then(Value::as_str), expected);
for websocket in [false, true] {
let mut body = json!({"model":"gpt-5.4","input":"hello"});
if let Some(tier) = &tier {
body["service_tier"] = tier.clone();
}
let sent = capture_downstream_request(body, api_key, websocket, false).await;
assert_eq!(
sent.get("service_tier"),
tier.as_ref(),
"api_key={api_key}, websocket={websocket}"
);
}
}
}
}
@@ -483,3 +483,108 @@ mod native_tls {
server.await.unwrap();
}
}
fn overload_event(no_more_retry: bool) -> Value {
let mut headers = json!({"retry-after": "30", "retry-after-ms": "30000"});
if no_more_retry {
headers["x-retry-metadata"] = json!("NO_MORE_RETRY");
}
json!({
"type": "error",
"error": {
"type": "service_unavailable_error",
"code": "server_is_overloaded",
"message": "Our servers are currently overloaded. Please try again later.",
"headers": headers
}
})
}
#[tokio::test]
async fn upstream_no_more_retry_blocks_same_account_recovery_for_a_short_retry_after() {
for (name, websocket, status, body, content_type, blocked, policy, kind) in [
(
"sse-capacity",
false,
200,
format!("event: error\ndata: {}\n\n", overload_event(true)),
"text/event-stream",
true,
TransientRetryPolicy::RetryThenRotateAccount,
ProviderErrorKind::UpstreamCapacityUnavailable,
),
(
"sse-capacity-without-signal",
false,
200,
format!("event: error\ndata: {}\n\n", overload_event(false)),
"text/event-stream",
false,
TransientRetryPolicy::RetryThenRotateAccount,
ProviderErrorKind::UpstreamCapacityUnavailable,
),
(
"http-server-error",
false,
503,
json!({"error":{"code":"server_error","type":"server_error","message":"retry later","headers":{"x-retry-metadata":"NO_MORE_RETRY","retry-after":"30"}}}).to_string(),
"application/json",
true,
TransientRetryPolicy::RetryThenReturnError,
ProviderErrorKind::Unavailable,
),
(
"websocket-capacity",
true,
503,
json!({"error":{"code":"server_is_overloaded","type":"server_error","message":"Selected model is at capacity. Please try a different model.","headers":{"x-retry-metadata":"NO_MORE_RETRY","retry-after":"30"}}}).to_string(),
"application/json",
true,
TransientRetryPolicy::RetryThenRotateAccount,
ProviderErrorKind::UpstreamCapacityUnavailable,
),
] {
let store = Arc::new(MemoryAccountStore::default());
create_account(&store, "acct_provider_contract").await;
let server = MockServer::start().await;
let mut response = ResponseTemplate::new(status).insert_header("content-type", content_type);
if status != 200 {
response = response.insert_header("retry-after", "30");
}
Mock::given(method(if websocket { "GET" } else { "POST" }))
.and(path("/codex/responses"))
.respond_with(response.set_body_string(body))
.expect(1)
.mount(&server)
.await;
let operation = if websocket {
capacity_websocket_operation()
} else {
http_generate_operation()
};
let mut stream = provider_with_base_url(&store, server.uri())
.execute(
planned_request("openai", operation),
context("req_no_more_retry", CancellationToken::new()),
)
.await
.unwrap();
let (_, error) = stream_failure(&mut stream).await;
assert_eq!(error.kind(), kind, "{name}");
assert_eq!(error.same_account_retry_blocked(), blocked, "{name}");
assert!(!error.retry_is_prohibited(), "{name}");
assert_eq!(error.retry_after(), Some(Duration::from_secs(30)), "{name}");
assert!(
matches!(
error.pre_delivery_retry(),
Some(PreDeliveryRetry::SameAccountTransientRetry {
policy: actual,
..
}) if actual == policy
),
"{name} preserved the exhaustion policy: {:?}",
error.pre_delivery_retry()
);
server.verify().await;
}
}
+154 -43
View File
@@ -19,12 +19,12 @@ use gateway_core::account::{
use gateway_core::error::{StoreError, StoreErrorKind};
use gateway_core::policy::ClientApiKeyId;
use gateway_core::provider_ports::{
ProviderCatalogCacheKey, ProviderCatalogCachePort, ProviderCooldown, ProviderCooldownPort,
ProviderCooldownScope, ProviderLeaseAcquisition, ProviderLeasePort, ProviderLeaseRequest,
ProviderRefreshPolicy, ProviderRuntimePolicyPort, ProviderSchedulingLeaseRequest,
ProviderSchedulingState, ProviderScopedCooldown, ProviderSessionAffinityKey,
ProviderSessionAffinityPort, ProviderSessionExclusionPort, ProviderSessionExclusions,
ProviderStoreError,
ProviderAccountAdmission, ProviderCatalogCacheKey, ProviderCatalogCachePort, ProviderCooldown,
ProviderCooldownPort, ProviderCooldownScope, ProviderLeaseAcquisition, ProviderLeasePort,
ProviderLeaseRequest, ProviderRefreshPolicy, ProviderRuntimePolicyPort,
ProviderSchedulingLeaseRequest, ProviderSchedulingState, ProviderScopedCooldown,
ProviderSessionAffinityKey, ProviderSessionAffinityPort, ProviderSessionBinding,
ProviderSessionExclusionPort, ProviderSessionExclusions, ProviderStoreError,
};
use gateway_core::routing::ProviderKind;
use provider_openai::credential::{
@@ -52,7 +52,17 @@ type CredentialLoadHook = dyn for<'a> Fn(
#[derive(Default)]
pub(crate) struct MemoryAccountStore {
pub(crate) peer_discovered: AtomicBool,
accounts: Mutex<BTreeMap<ProviderAccountId, StoredAccount>>,
catalogs: Mutex<
BTreeMap<
ProviderAccountId,
(
gateway_core::account::AccountCatalogObservation,
bytes::Bytes,
),
>,
>,
account_reads: AtomicUsize,
provider_list_reads: AtomicUsize,
quota_reads: AtomicUsize,
@@ -64,6 +74,21 @@ pub(crate) struct MemoryAccountStore {
credential_loads: AtomicUsize,
}
impl gateway_core::account::UpstreamInfoSource for MemoryAccountStore {
fn is_peer(&self, _: &ProviderAccount) -> bool {
self.peer_discovered.load(Ordering::SeqCst)
}
fn signal(
&self,
_: &ProviderAccount,
_: Option<&str>,
_: SystemTime,
) -> Option<gateway_core::account::UpstreamSchedulingSignal> {
None
}
}
impl MemoryAccountStore {
pub(crate) fn with_pending_account_reads() -> Self {
Self {
@@ -296,6 +321,62 @@ impl MemoryAccountStore {
#[async_trait]
impl ProviderAccountStore for MemoryAccountStore {
async fn load_account_catalogs(
&self,
provider: &ProviderKind,
) -> Result<Vec<gateway_core::account::StoredAccountCatalog>, StoreError> {
let accounts = self.accounts.lock().unwrap();
let mut groups = BTreeMap::<bytes::Bytes, Vec<_>>::new();
for (id, (observation, document)) in self.catalogs.lock().unwrap().iter() {
if accounts.get(id).is_some_and(|row| {
row.account.provider() == provider
&& row.account.enabled()
&& matches!(
row.account.credential_state(),
CredentialState::Ready | CredentialState::Unknown
)
&& row
.account
.access_token_expires_at()
.is_none_or(|ts| ts > SystemTime::now())
&& row.account.revision() == observation.revision
&& row.account.plan_type() == observation.plan.as_deref()
}) {
groups
.entry(document.clone())
.or_default()
.push(observation.clone());
}
}
Ok(groups
.into_iter()
.map(
|(document, accounts)| gateway_core::account::StoredAccountCatalog {
accounts,
document,
},
)
.collect())
}
async fn save_account_catalog(
&self,
observation: &gateway_core::account::AccountCatalogObservation,
document: &[u8],
) -> Result<bool, StoreError> {
let rows = self.accounts.lock().unwrap();
if !rows.get(&observation.account_id).is_some_and(|a| {
a.account.revision() == observation.revision
&& a.account.plan_type() == observation.plan.as_deref()
}) {
return Ok(false);
}
self.catalogs.lock().unwrap().insert(
observation.account_id.clone(),
(observation.clone(), bytes::Bytes::copy_from_slice(document)),
);
Ok(true)
}
async fn create_account(&self, input: NewProviderAccount) -> Result<(), StoreError> {
let mut accounts = self.accounts.lock().expect("account store lock");
if accounts.contains_key(input.account.id()) {
@@ -383,13 +464,13 @@ impl ProviderAccountStore for MemoryAccountStore {
));
}
// 与 Postgres 实现的调度列表语义一致:停用账号不进入常规候选
Ok(self
.accounts
.lock()
.expect("account store lock")
.values()
let accounts = self.accounts.lock().unwrap();
Ok(account_ids
.iter()
.collect::<std::collections::BTreeSet<_>>()
.into_iter()
.filter_map(|id| accounts.get(id))
.filter(|stored| stored.account.provider() == provider && stored.account.enabled())
.filter(|stored| account_ids.contains(stored.account.id()))
.map(|stored| stored.account.clone())
.collect())
}
@@ -924,6 +1005,7 @@ impl ProviderLeasePort for TestLeaseCoordinator {
#[derive(Default)]
pub(crate) struct MemorySessionAffinity {
cursors: Mutex<BTreeMap<(ClientApiKeyId, ProviderKind), u64>>,
operation_delay: Duration,
initial_claim_barrier: Option<Arc<tokio::sync::Barrier>>,
aliases: Mutex<BTreeMap<(String, String), gateway_core::provider_ports::ProviderSessionAlias>>,
@@ -1035,48 +1117,77 @@ impl ProviderSessionAffinityPort for MemorySessionAffinity {
})
}
fn compare_and_bind<'a>(
fn round_robin_cursor<'a>(
&'a self,
provider: &'a ProviderKind,
key: &'a ProviderSessionAffinityKey,
expected: Option<&'a gateway_core::provider_ports::ProviderSessionBinding>,
account: &'a ProviderAccountId,
ttl: Duration,
) -> BoxFuture<
'a,
Result<Option<gateway_core::provider_ports::ProviderSessionBinding>, ProviderStoreError>,
> {
client: &'a ClientApiKeyId,
) -> BoxFuture<'a, Result<u64, ProviderStoreError>> {
Box::pin(async move {
Ok(self
.cursors
.lock()
.unwrap()
.get(&(client.clone(), provider.clone()))
.copied()
.unwrap_or_default())
})
}
fn admit_account<'a>(
&'a self,
provider: &'a ProviderKind,
admission: ProviderAccountAdmission<'a>,
) -> BoxFuture<'a, Result<bool, ProviderStoreError>> {
Box::pin(async move {
self.delay_operation().await;
// 让并发首请求都读到空绑定,再验证发送前的原子认领冲突
if expected.is_none()
if admission
.session()
.is_some_and(|claim| claim.expected.is_none())
&& let Some(barrier) = &self.initial_claim_barrier
{
barrier.wait().await;
}
let key = (
provider.as_str().to_owned(),
key.expose_to_store().to_owned(),
);
let mut bindings = self.bindings.lock().expect("bindings");
if bindings.get(&key) != expected {
return Ok(None);
}
let binding = expected
.filter(|binding| binding.account_id() == account)
.cloned()
.unwrap_or_else(|| {
gateway_core::provider_ports::ProviderSessionBinding::new(
account.clone(),
let mut cursors = self.cursors.lock().unwrap();
let step = if let ProviderAccountAdmission::RoundRobin {
client_key, cursor, ..
} = admission
{
let current = cursors
.entry((client_key.clone(), provider.clone()))
.or_default();
if *current != cursor {
return Ok(false);
}
Some(current)
} else {
None
};
if let Some(claim) = admission.session() {
let key = (
provider.as_str().to_owned(),
claim.key.expose_to_store().to_owned(),
);
let mut bindings = self.bindings.lock().unwrap();
if bindings.get(&key) != claim.expected {
return Ok(false);
}
let binding = match claim.expected.filter(|b| b.account_id() == claim.account) {
Some(binding) => binding.clone(),
None => ProviderSessionBinding::new(
claim.account.clone(),
uuid::Uuid::new_v4().simple().to_string(),
)
.expect("binding")
});
bindings.insert(key, binding.clone());
self.renewal_ttls.lock().expect("TTLs").push(ttl);
Ok(Some(binding))
)?,
};
bindings.insert(key, binding);
self.renewal_ttls.lock().unwrap().push(claim.ttl);
}
if let Some(cursor) = step {
*cursor = cursor.wrapping_add(1);
}
Ok(true)
})
}
fn load_alias<'a>(
&'a self,
provider: &'a ProviderKind,

Some files were not shown because too many files have changed in this diff Show More