mirror of
https://github.com/zyycn/codex-proxy-rs.git
synced 2026-11-11 16:06:13 +00:00
Compare commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
bd69329a37 | ||
|
|
151a3c6bfd | ||
|
|
e2f194875f | ||
|
|
07b61bdc83 | ||
|
|
c0fdc1bbff | ||
|
|
343c579017 | ||
|
|
914c00785c | ||
|
|
b3de0d637d | ||
|
|
1271fa0186 | ||
|
|
055da40476 | ||
|
|
36ba79ca5e | ||
|
|
1b24029a97 | ||
|
|
431992a7a4 | ||
|
|
d3143406dc | ||
|
|
478f2d1f21 |
@@ -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(), ®istry, 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",
|
||||
|
||||
@@ -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 {
|
||||
|
||||
@@ -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";
|
||||
|
||||
@@ -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()
|
||||
);
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
|
||||
/// 只按账号整体的当前事实迁移主绑定,临时容量和单模型限制仍允许恢复热缓存
|
||||
|
||||
@@ -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("a),
|
||||
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("a),
|
||||
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]
|
||||
|
||||
+279
-637
File diff suppressed because it is too large
Load Diff
@@ -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("a),
|
||||
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()
|
||||
|
||||
+42
-5
@@ -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;
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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
Reference in New Issue
Block a user