fix: contain connection failures and retry unsent requests

Co-authored-by: Codex <noreply@openai.com>
This commit is contained in:
Yanshan Chen
2026-09-24 21:25:09 +08:00
co-authored by Codex
parent 1ef9f3b968
commit f9828a588f
42 changed files with 1233 additions and 1175 deletions
+1
View File
@@ -3255,6 +3255,7 @@ dependencies = [
"tokio-rustls",
"tokio-tungstenite 0.28.0",
"tokio-util",
"tower-layer",
"tower-service",
"tracing",
"tracing-subscriber",
@@ -609,18 +609,14 @@ async fn continuation_recovery_should_preserve_the_official_retry_signal() {
mod model_routing {
use std::collections::BTreeMap;
use std::sync::Arc;
use std::time::{Duration, SystemTime};
use axum::{
body::{Body, to_bytes},
http::{Request, StatusCode},
};
use futures::future::BoxFuture;
use gateway_core::engine::execution::{
DefaultExecutionService, ProviderCircuitDecision, ProviderCircuitError, ProviderCircuitPort,
};
use gateway_core::engine::execution::DefaultExecutionService;
use gateway_core::engine::provider::ProviderRegistry;
use gateway_core::routing::{ProviderKind, RuntimeSnapshot};
use gateway_core::routing::RuntimeSnapshot;
use gateway_core::runtime::RuntimeSnapshotHandle;
use serde_json::{Value, json};
use tower::ServiceExt;
@@ -629,53 +625,16 @@ mod model_routing {
IgnoredClientApiKeyUsage, UnusedAdmissions, UnusedContinuation, UnusedExecutionStore,
};
struct TestCircuits {
blocked: bool,
}
impl ProviderCircuitPort for TestCircuits {
fn decision<'a>(
&'a self,
_: &'a ProviderKind,
) -> BoxFuture<'a, Result<ProviderCircuitDecision, ProviderCircuitError>> {
Box::pin(async move {
Ok(if self.blocked {
ProviderCircuitDecision::BlockedUntil(
SystemTime::now() + Duration::from_secs(30),
)
} else {
ProviderCircuitDecision::Allow
})
})
}
fn observe_failure<'a>(
&'a self,
_: &'a ProviderKind,
) -> BoxFuture<'a, Result<(), ProviderCircuitError>> {
Box::pin(async { unreachable!("routing rejection must not execute a provider") })
}
fn observe_success<'a>(
&'a self,
_: &'a ProviderKind,
) -> BoxFuture<'a, Result<(), ProviderCircuitError>> {
Box::pin(async { unreachable!("routing rejection must not execute a provider") })
}
}
async fn request_model(
snapshot: RuntimeSnapshot,
model: &str,
stream: bool,
blocked: bool,
) -> (StatusCode, Value) {
let execution = DefaultExecutionService::new(
RuntimeSnapshotHandle::new(snapshot),
Arc::new(UnusedExecutionStore),
ProviderRegistry::default(),
Arc::new(UnusedAdmissions),
Arc::new(TestCircuits { blocked }),
Arc::new(UnusedContinuation),
Arc::new(IgnoredClientApiKeyUsage),
);
@@ -716,7 +675,7 @@ mod model_routing {
for stream in [false, true] {
let snapshot = crate::openai::snapshot("sk_model_routing", "openai")
.with_model_mappings(mappings.clone());
let response = request_model(snapshot, model, stream, false).await;
let response = request_model(snapshot, model, stream).await;
assert_eq!(
response,
@@ -733,25 +692,6 @@ mod model_routing {
}
}
}
#[tokio::test]
async fn an_existing_model_with_a_blocked_provider_should_remain_service_unavailable() {
let response = request_model(
crate::openai::snapshot("sk_model_routing", "openai"),
"model-a",
true,
true,
)
.await;
assert_eq!(
(response.0, response.1["error"]["code"].as_str()),
(
StatusCode::SERVICE_UNAVAILABLE,
Some("no_available_provider")
),
);
}
}
#[tokio::test]
@@ -783,3 +723,27 @@ async fn quota_recovery_http_response_uses_projection_without_replacing_upstream
"usage_limit_reached"
);
}
#[test]
fn final_connection_failures_and_local_capacity_have_distinct_http_statuses() {
for (kind, expected) in [
(ProviderErrorKind::Transport, StatusCode::BAD_GATEWAY),
(ProviderErrorKind::Timeout, StatusCode::GATEWAY_TIMEOUT),
(
ProviderErrorKind::ProviderInfrastructureUnavailable,
StatusCode::SERVICE_UNAVAILABLE,
),
] {
let error = EngineError::Provider(ProviderError::new(kind, UpstreamSendState::NotSent));
assert_eq!(engine_error_response(&error).status(), expected);
}
for kind in [
GatewayErrorKind::NoAvailableProvider,
GatewayErrorKind::AccountCapacityUnavailable,
] {
assert_eq!(
gateway_error_response(&GatewayError::new(kind, "temporarily unavailable")).status(),
StatusCode::SERVICE_UNAVAILABLE
);
}
}
@@ -26,7 +26,6 @@ use gateway_core::engine::continuation::{
};
use gateway_core::engine::execution::{
AuthenticatedClient, ClientApiKeyUsageSink, DefaultExecutionService, ExecutionService,
ProviderCircuitDecision, ProviderCircuitError, ProviderCircuitPort,
};
use gateway_core::engine::provider::ProviderRegistry;
use gateway_core::engine::{
@@ -89,7 +88,6 @@ pub(super) fn api_router_with_admin_and_client(
Arc::new(UnusedExecutionStore),
ProviderRegistry::default(),
Arc::new(UnusedAdmissions),
Arc::new(UnusedCircuits),
Arc::new(UnusedContinuation),
Arc::new(IgnoredClientApiKeyUsage),
));
@@ -105,7 +103,6 @@ pub(super) fn api_router_with_config(
Arc::new(UnusedExecutionStore),
ProviderRegistry::default(),
Arc::new(UnusedAdmissions),
Arc::new(UnusedCircuits),
Arc::new(UnusedContinuation),
Arc::new(IgnoredClientApiKeyUsage),
));
@@ -194,7 +191,6 @@ pub(super) fn authenticated_client_for_provider_with_limit(
Arc::new(UnusedExecutionStore),
ProviderRegistry::default(),
Arc::new(UnusedAdmissions),
Arc::new(UnusedCircuits),
Arc::new(UnusedContinuation),
Arc::new(IgnoredClientApiKeyUsage),
);
@@ -218,7 +214,6 @@ pub(super) fn authenticated_client_with_min_versions(
Arc::new(UnusedExecutionStore),
ProviderRegistry::default(),
Arc::new(UnusedAdmissions),
Arc::new(UnusedCircuits),
Arc::new(UnusedContinuation),
Arc::new(IgnoredClientApiKeyUsage),
);
@@ -395,31 +390,6 @@ impl ClientAdmissionPort for UnusedAdmissions {
}
}
struct UnusedCircuits;
impl ProviderCircuitPort for UnusedCircuits {
fn decision<'a>(
&'a self,
_: &'a ProviderKind,
) -> BoxFuture<'a, Result<ProviderCircuitDecision, ProviderCircuitError>> {
Box::pin(async { unreachable!("authentication fixture does not execute") })
}
fn observe_failure<'a>(
&'a self,
_: &'a ProviderKind,
) -> BoxFuture<'a, Result<(), ProviderCircuitError>> {
Box::pin(async { unreachable!("authentication fixture does not execute") })
}
fn observe_success<'a>(
&'a self,
_: &'a ProviderKind,
) -> BoxFuture<'a, Result<(), ProviderCircuitError>> {
Box::pin(async { unreachable!("authentication fixture does not execute") })
}
}
struct UnusedContinuation;
impl NativeContinuationPort for UnusedContinuation {
@@ -27,8 +27,7 @@ use gateway_core::engine::admission::{
use gateway_core::engine::budget::{ClientBudgetCharge, ClientBudgetError, ClientBudgetPort};
use gateway_core::engine::execution::{
AuthenticatedClient, ClientAuthenticationError, DefaultExecutionService, ExecutionService,
ExecutionSession, ProviderCircuitDecision, ProviderCircuitError, ProviderCircuitPort,
StartExecution, StartProviderExecution, StartedExecution,
ExecutionSession, StartExecution, StartProviderExecution, StartedExecution,
};
use gateway_core::engine::provider::{
Provider, ProviderCallMetadata, ProviderRegistry, ProviderRequest, ProviderStream,
@@ -48,9 +47,7 @@ use gateway_core::lifecycle::{
use gateway_core::metering::{Decimal, ProviderReportedCost};
use gateway_core::operation::Operation;
use gateway_core::policy::ClientApiKeyId;
use gateway_core::routing::{
ProviderCatalogGeneration, ProviderKind, ProviderModelCapabilities, PublicModelId,
};
use gateway_core::routing::{ProviderCatalogGeneration, ProviderModelCapabilities, PublicModelId};
use gateway_core::runtime::RuntimeSnapshotHandle;
use gateway_core::upstream::{UpstreamSendState, UpstreamTransport};
use gateway_protocol::openai::codex_responses_request_semantics;
@@ -984,29 +981,6 @@ impl ClientBudgetPort for SettlementPorts {
}
}
impl ProviderCircuitPort for SettlementPorts {
fn decision<'a>(
&'a self,
_: &'a ProviderKind,
) -> BoxFuture<'a, Result<ProviderCircuitDecision, ProviderCircuitError>> {
Box::pin(async { Ok(ProviderCircuitDecision::Allow) })
}
fn observe_failure<'a>(
&'a self,
_: &'a ProviderKind,
) -> BoxFuture<'a, Result<(), ProviderCircuitError>> {
Box::pin(async { Ok(()) })
}
fn observe_success<'a>(
&'a self,
_: &'a ProviderKind,
) -> BoxFuture<'a, Result<(), ProviderCircuitError>> {
Box::pin(async { Ok(()) })
}
}
struct ChargedWebSocketProvider;
fn websocket_charge() -> Decimal {
@@ -1114,7 +1088,6 @@ async fn websocket_disconnect_during_core_settlement_finishes_charge_before_rele
ProviderRegistry::new([Arc::new(ChargedWebSocketProvider) as Arc<dyn Provider>])
.unwrap(),
ports.clone(),
ports.clone(),
Arc::new(UnusedContinuation),
Arc::new(IgnoredClientApiKeyUsage),
)
@@ -0,0 +1,108 @@
//! 请求内共享的发送前连接预算;不投影账号或 Provider 健康状态。
use std::sync::{Arc, Mutex};
use std::time::{Duration, Instant};
const MAX_OPENINGS: u32 = 4;
const STARTUP_TIMEOUT: Duration = Duration::from_secs(30);
#[derive(Debug, Default)]
struct State {
started: Option<Instant>,
openings: u32,
completed: bool,
}
/// 换账号、换传输和内部重试共享同一预算;成功发送后的协议恢复不重置它。
#[derive(Debug, Clone, Default)]
pub struct ConnectionBudget(Arc<Mutex<State>>);
#[derive(Debug, Clone, Copy, PartialEq, Eq, thiserror::Error)]
#[error("connection recovery budget exhausted")]
pub struct ConnectionBudgetExhausted;
impl ConnectionBudget {
/// 在可能建立连接前预约一次 opening;返回本次连接允许使用的剩余时间。
/// HTTP 连接池隐藏了是否冷建连,因此一次 HTTP opening 也计入次数上限。
pub fn begin(&self) -> Result<Option<Duration>, ConnectionBudgetExhausted> {
let mut state = self
.0
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner);
if state.completed {
return Ok(None);
}
let started = *state.started.get_or_insert_with(Instant::now);
let remaining = STARTUP_TIMEOUT.saturating_sub(started.elapsed());
if state.openings >= MAX_OPENINGS || remaining.is_zero() {
return Err(ConnectionBudgetExhausted);
}
state.openings += 1;
Ok(Some(remaining))
}
/// 仅连接成功不足以清零;确认业务发送后退出新增的 NotSent 恢复策略。
pub fn complete(&self) {
self.0
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner)
.completed = true;
}
#[must_use]
pub fn remaining(&self) -> Option<Duration> {
let state = self
.0
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner);
if state.completed || state.openings >= MAX_OPENINGS {
return None;
}
let remaining = state.started.map_or(STARTUP_TIMEOUT, |started| {
STARTUP_TIMEOUT.saturating_sub(started.elapsed())
});
(!remaining.is_zero()).then_some(remaining)
}
#[must_use]
pub fn exhausted(&self) -> bool {
let state = self
.0
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner);
!state.completed
&& (state.openings >= MAX_OPENINGS
|| state
.started
.is_some_and(|started| started.elapsed() >= STARTUP_TIMEOUT))
}
/// 首次建连后的等待也计入恢复窗口;不缩短已发送请求的生成时间。
#[must_use]
pub fn startup_remaining(&self) -> Option<Duration> {
let state = self
.0
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner);
if state.completed {
return None;
}
state
.started
.map(|started| STARTUP_TIMEOUT.saturating_sub(started.elapsed()))
}
/// 额外 attempt 最多三个;抖动由请求 ID 派生,重试不依赖运行时随机源。
#[must_use]
pub fn retry_delay(&self, retries: u32, request_id: &str) -> Option<Duration> {
if retries >= MAX_OPENINGS - 1 {
return None;
}
let remaining = self.remaining()?;
let hash = request_id.bytes().fold(u64::from(retries), |hash, byte| {
hash.wrapping_mul(31).wrapping_add(u64::from(byte))
});
let delay = Duration::from_millis((500_u64 << retries) * (80 + hash % 41) / 100);
(delay < remaining).then_some(delay)
}
}
@@ -279,6 +279,8 @@ where
request_id,
client_api_key_ref,
concurrency_wait_budget: ConcurrencyWaitBudget::default(),
connection_budget: super::connection::ConnectionBudget::default(),
connection_retries: 0,
observation: ResponseObservation::new(timing_started_at),
request_observation,
budget_prior_attempts_usd: Decimal::ZERO,
@@ -395,6 +397,8 @@ pub struct ResponseExecutionSession<S: ?Sized> {
request_id: ModelRequestId,
client_api_key_ref: crate::policy::ClientApiKeyId,
concurrency_wait_budget: ConcurrencyWaitBudget,
connection_budget: super::connection::ConnectionBudget,
connection_retries: u32,
observation: ResponseObservation,
request_observation: Option<RequestObservationDispatch>,
budget_prior_attempts_usd: Decimal,
@@ -810,7 +814,7 @@ where
PollBoundary::Deadline => {
// 会话 deadline 是网关自身的请求预算,不是上游超时;
// 真正的上游超时会作为流错误进入 `handle_stream_error` 记账。
// 这里不写 provider 失败,避免长流集中到期误触 provider 熔断。
// 这里不写 provider 失败,避免把本地预算到期归因为上游故障。
self.finish_interruption(&EngineError::Deadline).await?;
return Err(EngineError::Deadline);
}
@@ -894,7 +898,13 @@ where
if let Some(recovery) = pending_retry.as_ref()
&& !recovery.delay.is_zero()
{
match poll_retry_delay(recovery.delay, self.cancellation.clone(), self.deadline).await {
let deadline = self
.connection_budget
.startup_remaining()
.map_or(self.deadline, |remaining| {
self.deadline.min(SystemTime::now() + remaining)
});
match poll_retry_delay(recovery.delay, self.cancellation.clone(), deadline).await {
RetryDelayBoundary::Elapsed => {}
RetryDelayBoundary::Cancelled => {
self.finish_interruption(&EngineError::Cancelled).await?;
@@ -1000,6 +1010,7 @@ where
.with_pricing(self.plan.pricing())
.with_request_location(self.plan.request_location().cloned())
.with_concurrency_wait_budget(self.concurrency_wait_budget.clone())
.with_connection_budget(self.connection_budget.clone())
.with_timing_started_at(self.observation.timing_started_at)
.with_request_policy(self.request_policy.clone())
.with_execution_effects(self.execution_effects.as_ref().map(Arc::clone))
@@ -1052,9 +1063,7 @@ where
return Err(EngineError::Cancelled);
}
ProviderBoundary::Deadline => {
// 网关预算到期同样不是候选 Provider 的上游超时,不计入熔断;
// Provider 自身的握手/传输超时会以 `ProviderErrorKind::Timeout`
// 错误返回并在下方 `Result` 分支记账。
// 网关预算到期是本请求的终态,不推定候选 Provider 不可用。
self.finish_interruption(&EngineError::Deadline).await?;
return Err(EngineError::Deadline);
}
@@ -1384,6 +1393,10 @@ where
} else {
error.send_state()
};
if attempt_send_state != UpstreamSendState::NotSent {
// 已发送的容量拒绝、凭据恢复沿用原策略,不再受首次建连窗口限制。
self.connection_budget.complete();
}
let execution_effect_observed = self.execution_effect_observed();
let send_state = self.raise_send_watermark(attempt_send_state);
if self.request_persisted {
@@ -1415,11 +1428,34 @@ where
Some(crate::error::PreDeliveryRetry::AccountRotation)
)
&& self.routing_attempts < self.plan.max_attempts().get();
// 只有尚未发送的逻辑请求进入新增策略;已有发送后安全恢复保持原有路由规则。
let connection_retry_requested = self.current_send_state() == UpstreamSendState::NotSent
&& matches!(
error.pre_delivery_retry(),
Some(crate::error::PreDeliveryRetry::SameAccountConnectionRetry { .. })
);
let transport_recovery = match error.pre_delivery_retry() {
Some(crate::error::PreDeliveryRetry::SameAccountConnectionRetry { transport })
if !execution_effect_observed
&& !self.connection_budget.exhausted()
&& self.downstream_committed_at.is_none()
&& !self.delivery_pending
&& attempt_send_state == UpstreamSendState::NotSent
&& self.current_send_state() == UpstreamSendState::NotSent
&& self.continuation_attempt == ContinuationAttempt::None =>
{
self.connection_budget
.retry_delay(self.connection_retries, self.request_id.as_str())
.map(|delay| {
self.connection_retries += 1;
(transport, delay)
})
}
Some(crate::error::PreDeliveryRetry::SameAccountTransportRetry {
retry_index,
delay,
}) if !execution_effect_observed
&& !self.connection_budget.exhausted()
&& self.downstream_committed_at.is_none()
&& !self.delivery_pending
&& attempt_send_state != UpstreamSendState::Ambiguous =>
@@ -1428,6 +1464,7 @@ where
}
Some(crate::error::PreDeliveryRetry::SameAccountTransportFallback)
if !execution_effect_observed
&& !self.connection_budget.exhausted()
&& self.downstream_committed_at.is_none()
&& !self.delivery_pending
&& attempt_send_state != UpstreamSendState::Ambiguous =>
@@ -1436,7 +1473,9 @@ where
}
_ => None,
};
let ordinary_retry = !execution_effect_observed
let ordinary_retry = !self.connection_budget.exhausted()
&& !connection_retry_requested
&& !execution_effect_observed
&& self.account_selection.required_account().is_none()
&& self.continuation_attempt == ContinuationAttempt::None
&& self.downstream_committed_at.is_none()
@@ -1486,6 +1525,9 @@ where
"sameAccountRetry": same_account_retry, "accountRotationRetry": account_rotation_retry,
"ordinaryRetry": ordinary_retry, "transportRecovery": transport_recovery.is_some(),
"transientRetry": transient_retry.is_some(),
"connectionRetry": connection_retry_requested,
"connectionRetries": self.connection_retries,
"connectionBudgetRemainingMs": self.connection_budget.remaining().map(duration_ms),
"executionEffectObserved": execution_effect_observed,
"delayMs": transient_retry.or(transport_recovery.map(|(_, delay)| delay)).map(duration_ms),
"downstreamCommitted": self.downstream_committed_at.is_some(),
@@ -1,9 +1,8 @@
//! 数据面执行用例:认证、准入、路由、continuation、circuit 与会话生命周期。
//! 数据面执行用例:认证、准入、路由、continuation 与会话生命周期。
use std::collections::{BTreeMap, BTreeSet};
use std::fmt;
use std::net::IpAddr;
use std::num::NonZeroU32;
use std::sync::{
Arc, Mutex, Weak,
atomic::{AtomicUsize, Ordering},
@@ -52,9 +51,9 @@ use crate::engine::provider::ProviderRegistry;
use crate::engine::{
AttemptCoordinator, AttemptRecord, CoordinatedEvent, EngineError, ExecutionStore,
GatewayEngine, IntermediateFailure, ModelRequestFinalization, ModelRequestId, NewModelRequest,
ProbeFailure, ProviderAccountId, ProviderAttemptOutcome, RecoveryReport, UpstreamSendState,
ProbeFailure, ProviderAccountId, RecoveryReport, UpstreamSendState,
};
use crate::error::{GatewayError, GatewayErrorKind, ProviderErrorKind, StoreError};
use crate::error::{GatewayError, GatewayErrorKind, StoreError};
use crate::event::{GatewayEvent, ProviderEvent, ProviderResponseHeader};
use crate::identity::ProviderKind;
use crate::lifecycle::CancellationToken;
@@ -649,47 +648,6 @@ pub trait ClientApiKeyUsageSink: Send + Sync {
fn record_used(&self, key_id: &ClientApiKeyId);
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ProviderCircuitDecision {
Allow,
BlockedUntil(SystemTime),
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, thiserror::Error)]
#[error("provider circuit store is unavailable")]
pub struct ProviderCircuitError;
/// Provider circuit 的可重建协调策略;由 Core 拥有并交给 Store adapter 执行。
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct ProviderCircuitPolicy {
pub failure_threshold: NonZeroU32,
pub open_duration: Duration,
}
impl Default for ProviderCircuitPolicy {
fn default() -> Self {
Self {
failure_threshold: NonZeroU32::new(3).unwrap_or(NonZeroU32::MIN),
open_duration: Duration::from_secs(30),
}
}
}
pub trait ProviderCircuitPort: Send + Sync {
fn decision<'a>(
&'a self,
provider_kind: &'a ProviderKind,
) -> BoxFuture<'a, Result<ProviderCircuitDecision, ProviderCircuitError>>;
fn observe_failure<'a>(
&'a self,
provider_kind: &'a ProviderKind,
) -> BoxFuture<'a, Result<(), ProviderCircuitError>>;
fn observe_success<'a>(
&'a self,
provider_kind: &'a ProviderKind,
) -> BoxFuture<'a, Result<(), ProviderCircuitError>>;
}
pub struct DefaultExecutionService {
snapshots: RuntimeSnapshotHandle,
/// probe 自身走 transient store,探测失败仍写入持久 store 的 ops_events。
@@ -697,7 +655,6 @@ pub struct DefaultExecutionService {
providers: ProviderRegistry,
admissions: Arc<dyn ClientAdmissionPort>,
admission_waiting: ConcurrencyWaitQueue<ClientApiKeyId>,
circuits: Arc<dyn ProviderCircuitPort>,
continuation: Arc<dyn NativeContinuationPort>,
client_api_key_usage: Arc<dyn ClientApiKeyUsageSink>,
budget: Option<Arc<dyn ClientBudgetPort>>,
@@ -717,7 +674,6 @@ impl DefaultExecutionService {
execution: Arc<dyn ExecutionStore>,
providers: ProviderRegistry,
admissions: Arc<dyn ClientAdmissionPort>,
circuits: Arc<dyn ProviderCircuitPort>,
continuation: Arc<dyn NativeContinuationPort>,
client_api_key_usage: Arc<dyn ClientApiKeyUsageSink>,
) -> Self {
@@ -727,7 +683,6 @@ impl DefaultExecutionService {
providers,
admissions,
admission_waiting: ConcurrencyWaitQueue::default(),
circuits,
continuation,
client_api_key_usage,
budget: None,
@@ -1029,14 +984,7 @@ impl DefaultExecutionService {
});
// 路由插件可以调用 host.model/host.affinity;仅这条扩展路径必须在首次 RPC
// 前取得父 Key 准入。无策略的原生路由仍保持“先路由、后准入”的既有顺序。
let mut start_guard = if request_policy.is_some() {
Some(
self.prepare_execution_start(&request, &request_id, &mut authorization)
.await?,
)
} else {
None
};
let mut start_guard = None;
let request_observation = self.request_observers.as_ref().and_then(|observers| {
let generation = request.client.snapshot.extensions()?.clone();
let plan = observers.resolve(&generation)?;
@@ -1055,10 +1003,12 @@ impl DefaultExecutionService {
))
});
let budget_key_id = request.client.policy.key_id().clone();
let mut entered_execution = false;
let result = async {
let mut routing_context = self
.route_context(authorization.account_scope.provider_kinds())
.await?;
if request_policy.is_some() {
start_guard = Some(self.prepare_execution_start(&request, &request_id, &mut authorization).await?);
}
let mut routing_context = RoutingContext::default();
if let Some(provider) = authorization.required_provider.clone() {
if !authorization.account_scope.provider_kinds().contains(&provider) {
return Err(GatewayError::new(
@@ -1217,6 +1167,7 @@ impl DefaultExecutionService {
request.metadata.transport,
)
.with_extension_scope(authorization.extension_scope.clone());
entered_execution = true;
self.start_without_continuation(
request,
PreparedExecutionStart {
@@ -1233,6 +1184,27 @@ impl DefaultExecutionService {
.await
}
.await;
if !entered_execution && let Err(error) = &result {
tracing::warn!(
request_id = request_id.as_str(),
key_id = budget_key_id.as_str(),
failure_kind = error.kind().as_str(),
"请求在路由或准入阶段被拒绝"
);
let rejection = super::EntryRejection {
request_id: request_id.clone(),
client_key_id: budget_key_id.clone(),
error: error.clone(),
latency: started_at.elapsed().unwrap_or_default(),
};
let write = self.observations.record_entry_rejection(rejection).fuse();
let timeout = Delay::new(COORDINATION_TIMEOUT).fuse();
pin_mut!(write, timeout);
select_biased! {
result = write => { if result.is_err() { tracing::warn!("入口拒绝观测写入失败"); } },
_ = timeout => tracing::warn!("入口拒绝观测写入超时"),
}
}
if let (Some(observation), Err(error)) = (&request_observation, &result) {
observation.reject(error);
}
@@ -1438,7 +1410,6 @@ impl DefaultExecutionService {
core,
admission,
active_request,
Arc::clone(&self.circuits),
Arc::clone(&self.continuation),
self.budget.clone(),
)),
@@ -1864,48 +1835,6 @@ impl DefaultExecutionService {
}
}
async fn route_context(
&self,
provider_kinds: &BTreeSet<ProviderKind>,
) -> Result<RoutingContext, GatewayError> {
let decisions = futures::future::join_all(provider_kinds.iter().map(|provider_kind| {
let circuits = Arc::clone(&self.circuits);
async move {
let decision = circuits.decision(provider_kind).fuse();
let timeout = Delay::new(COORDINATION_TIMEOUT).fuse();
pin_mut!(decision, timeout);
let decision = select_biased! {
result = decision => Some(result),
_ = timeout => None,
};
(provider_kind, decision)
}
}))
.await;
let mut blocked_providers = BTreeSet::new();
for (provider_kind, decision) in decisions {
match decision {
Some(Ok(ProviderCircuitDecision::BlockedUntil(_))) => {
blocked_providers.insert(provider_kind.clone());
}
Some(Err(error)) => tracing::warn!(
provider = provider_kind.as_str(),
%error,
"Provider circuit 读取失败,按可重建协调状态 fail-open"
),
None => tracing::warn!(
provider = provider_kind.as_str(),
"Provider circuit 读取超时,按可重建协调状态 fail-open"
),
Some(Ok(ProviderCircuitDecision::Allow)) => {}
}
}
Ok(RoutingContext {
required_provider: None,
blocked_providers,
})
}
async fn probe_inner(
&self,
request: AccountProbeRequest,
@@ -2004,11 +1933,6 @@ impl DefaultExecutionService {
}
};
let events = session.collect_uncommitted().await;
publish_provider_attempt_outcomes(
self.circuits.as_ref(),
session.provider_attempt_outcomes(),
)
.await;
let events = match events {
Ok(events) => events,
Err(error) => {
@@ -2531,9 +2455,7 @@ struct DefaultExecutionSession {
admission: Option<ExecutionAdmission>,
active_request: Option<ActiveRequestLease>,
cleanup: Option<BoxFuture<'static, ()>>,
circuits: Arc<dyn ProviderCircuitPort>,
continuation: Arc<dyn NativeContinuationPort>,
observed_provider_outcomes: usize,
continuation_recorded: bool,
budget: Option<Arc<dyn ClientBudgetPort>>,
}
@@ -2543,7 +2465,6 @@ impl DefaultExecutionSession {
core: ResponseExecutionSession<dyn ExecutionStore>,
admission: ExecutionAdmission,
active_request: Option<ActiveRequestLease>,
circuits: Arc<dyn ProviderCircuitPort>,
continuation: Arc<dyn NativeContinuationPort>,
budget: Option<Arc<dyn ClientBudgetPort>>,
) -> Self {
@@ -2552,9 +2473,7 @@ impl DefaultExecutionSession {
admission: Some(admission),
active_request,
cleanup: None,
circuits,
continuation,
observed_provider_outcomes: 0,
continuation_recorded: false,
budget,
}
@@ -2584,16 +2503,6 @@ impl DefaultExecutionSession {
}
}
async fn observe_provider_outcomes(&mut self) {
let outcomes = self.core.provider_attempt_outcomes();
let new_outcomes = outcomes
.get(self.observed_provider_outcomes..)
.unwrap_or_default()
.to_vec();
self.observed_provider_outcomes = outcomes.len();
publish_provider_attempt_outcomes(self.circuits.as_ref(), &new_outcomes).await;
}
async fn record_continuation(&mut self, state: Option<&ProviderSessionState>) {
if self.continuation_recorded {
return;
@@ -2612,7 +2521,6 @@ impl DefaultExecutionSession {
if let Err(error) = self.core.cancel_and_finalize().await {
tracing::warn!(%error, "Detached execution 终态收敛失败");
}
self.observe_provider_outcomes().await;
self.settle_if_finalized().await;
}
}
@@ -2634,7 +2542,6 @@ impl ExecutionSession for DefaultExecutionSession {
if let Ok(Some(event)) = result.as_ref() {
self.record_continuation(event.session_update()).await;
}
self.observe_provider_outcomes().await;
self.settle_if_finalized().await;
result
})
@@ -2647,7 +2554,6 @@ impl ExecutionSession for DefaultExecutionSession {
let state = events.iter().find_map(ProviderEvent::session_update);
self.record_continuation(state).await;
}
self.observe_provider_outcomes().await;
self.settle_if_finalized().await;
result
})
@@ -2671,7 +2577,6 @@ impl ExecutionSession for DefaultExecutionSession {
) -> BoxFuture<'_, Result<(), EngineError>> {
Box::pin(async move {
let result = self.core.commit_downstream(client_status_code).await;
self.observe_provider_outcomes().await;
self.settle_if_finalized().await;
result
})
@@ -2683,7 +2588,6 @@ impl ExecutionSession for DefaultExecutionSession {
) -> BoxFuture<'_, Result<(), EngineError>> {
Box::pin(async move {
let result = self.core.record_client_status(client_status_code).await;
self.observe_provider_outcomes().await;
self.settle_if_finalized().await;
result
})
@@ -2705,39 +2609,6 @@ impl ExecutionSession for DefaultExecutionSession {
}
}
#[must_use]
pub const fn provider_failure_affects_circuit(error_kind: ProviderErrorKind) -> bool {
matches!(
error_kind,
ProviderErrorKind::Timeout
| ProviderErrorKind::Transport
| ProviderErrorKind::Protocol
| ProviderErrorKind::Unavailable
)
}
async fn publish_provider_attempt_outcomes(
circuits: &dyn ProviderCircuitPort,
outcomes: &[ProviderAttemptOutcome],
) {
for outcome in outcomes {
let result = match outcome.error_kind() {
None => circuits.observe_success(outcome.provider_kind()).await,
Some(kind) if provider_failure_affects_circuit(kind) => {
circuits.observe_failure(outcome.provider_kind()).await
}
Some(_) => continue,
};
if let Err(error) = result {
tracing::warn!(
provider = outcome.provider_kind().as_str(),
%error,
"Provider circuit feedback 写入失败,数据面不受影响"
);
}
}
}
async fn record_native_continuation(
continuation: &dyn NativeContinuationPort,
pin: NativeContinuationPin,
+31 -2
View File
@@ -3,6 +3,7 @@
pub mod admission;
pub mod authentication;
pub mod budget;
pub mod connection;
pub mod continuation;
pub mod coordinator;
pub mod execution;
@@ -120,9 +121,9 @@ impl AttemptTrigger {
}
}
/// 一次实际上游调用对 Provider 健康度产生的事实。
/// 一次实际上游调用的诊断结果。
///
/// 该事实只描述调用结果,不在 Core 中定义 circuit 策略或持久化方式。
/// 该事实只描述调用结果,不参与跨请求的路由屏蔽。
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum ProviderAttemptOutcome {
/// 上游流自然完成且通过 canonical event 序列校验。
@@ -331,6 +332,7 @@ pub struct RequestAttemptContext {
timing_started_at: Instant,
trace: crate::diagnostics::TraceContext,
concurrency_wait_budget: crate::concurrency::ConcurrencyWaitBudget,
connection_budget: connection::ConnectionBudget,
request_policy: Option<policy::RequestPolicyContext>,
execution_effects: Option<Arc<nested::ExecutionEffects>>,
middleware: Option<middleware::FrozenMiddlewarePlan>,
@@ -383,6 +385,7 @@ impl RequestAttemptContext {
timing_started_at: Instant::now(),
trace: crate::diagnostics::TraceContext::default(),
concurrency_wait_budget: crate::concurrency::ConcurrencyWaitBudget::default(),
connection_budget: connection::ConnectionBudget::default(),
request_policy: None,
execution_effects: None,
middleware: None,
@@ -409,6 +412,13 @@ impl RequestAttemptContext {
self
}
/// 传递请求内共享的连接恢复预算。
#[must_use]
pub fn with_connection_budget(mut self, budget: connection::ConnectionBudget) -> Self {
self.connection_budget = budget;
self
}
/// 覆盖本次请求的单调计时原点。
#[must_use]
pub fn with_timing_started_at(mut self, timing_started_at: Instant) -> Self {
@@ -725,6 +735,11 @@ impl AttemptContext {
self.continuation_attempt
}
#[must_use]
pub const fn connection_budget(&self) -> &connection::ConnectionBudget {
&self.request.connection_budget
}
/// 返回本次 attempt 的 Provider 传输档位。
#[must_use]
pub const fn transport(&self) -> AttemptTransport {
@@ -817,6 +832,15 @@ pub struct IntermediateFailure {
pub latency: Duration,
}
/// 已认证有效请求在执行会话建立前的拒绝;不伪造模型执行或上游 attempt。
#[derive(Debug)]
pub struct EntryRejection {
pub request_id: ModelRequestId,
pub client_key_id: ClientApiKeyId,
pub error: GatewayError,
pub latency: Duration,
}
/// 不属于任何 `model_requests` 行的管理端账号探测失败。
#[derive(Debug)]
pub struct ProbeFailure {
@@ -927,6 +951,11 @@ pub trait ExecutionStore: Send + Sync {
&self,
failure: IntermediateFailure,
) -> Result<(), StoreError>;
/// 记录路由或准入拒绝;可恢复观测失败不改变客户端结果。
async fn record_entry_rejection(&self, _rejection: EntryRejection) -> Result<(), StoreError> {
Ok(())
}
/// 记录不挂在 `model_requests` 上的账号探测失败;默认丢弃。
async fn record_probe_failure(&self, _failure: ProbeFailure) -> Result<(), StoreError> {
Ok(())
+13
View File
@@ -135,6 +135,11 @@ pub enum PreDeliveryRetry {
/// 发起下一次传输尝试前的退避时长。
delay: Duration,
},
/// 可靠 NotSent 的建连恢复;次数与时间由 Core 的请求级预算裁决。
SameAccountConnectionRetry {
/// Provider 明确指定实际失败的传输,避免 HTTP 恢复重新进入 WS。
transport: crate::engine::AttemptTransport,
},
/// 固定本次账号,并要求 Provider 使用备用传输。
SameAccountTransportFallback,
}
@@ -625,6 +630,14 @@ impl ProviderError {
}));
}
#[must_use]
pub fn with_connection_retry(mut self, transport: crate::engine::AttemptTransport) -> Self {
self.pre_delivery_retry = Some(Box::new(PreDeliveryRetry::SameAccountConnectionRetry {
transport,
}));
self
}
/// 要求 Core 在账号凭据已恢复后仅对原账号重放一次。
#[must_use]
pub const fn with_same_account_retry(mut self) -> Self {
-5
View File
@@ -32,7 +32,6 @@ use engine::admission::{
use engine::continuation::NativeContinuationPort;
use engine::execution::{
ClientApiKeyUsageSink, ClientKeyVerifier, DefaultExecutionService, ExecutionService,
ProviderCircuitPort,
};
use engine::nested::{AffinityLookupPort, NestedModelExecutionPort};
use engine::probe::AccountProbe;
@@ -51,7 +50,6 @@ pub struct CoreStorePorts {
execution: Arc<dyn ExecutionStore>,
admissions: Arc<dyn ClientAdmissionPort>,
admission_recovery: Arc<dyn ClientAdmissionRecoveryPort>,
circuits: Arc<dyn ProviderCircuitPort>,
continuation: Arc<dyn NativeContinuationPort>,
snapshots: Arc<dyn SnapshotStorePort>,
snapshot_subscriptions: Arc<dyn SnapshotSubscriptionPort>,
@@ -68,7 +66,6 @@ impl CoreStorePorts {
Arc<dyn ClientAdmissionPort>,
Arc<dyn ClientAdmissionRecoveryPort>,
),
circuits: Arc<dyn ProviderCircuitPort>,
continuation: Arc<dyn NativeContinuationPort>,
(snapshots, snapshot_subscriptions): (
Arc<dyn SnapshotStorePort>,
@@ -80,7 +77,6 @@ impl CoreStorePorts {
execution,
admissions,
admission_recovery,
circuits,
continuation,
snapshots,
snapshot_subscriptions,
@@ -198,7 +194,6 @@ fn build_execution_service(
Arc::clone(&ports.execution),
providers,
Arc::clone(&ports.admissions),
Arc::clone(&ports.circuits),
Arc::clone(&ports.continuation),
Arc::clone(&ports.client_api_key_usage),
);
@@ -1062,7 +1062,7 @@ impl RuntimeSnapshot {
self.client_policies.get(id)
}
/// 返回当前 Key 范围、本代次注册表与 circuit 条件共同允许的 Provider。
/// 返回当前 Key 范围、本代次注册表与请求路由限制共同允许的 Provider。
#[must_use]
pub fn available_providers(
&self,
@@ -1183,7 +1183,7 @@ impl RuntimeSnapshot {
}
if candidates.is_empty() {
// 模型存在性必须包含被熔断的 Provider;目录未知时不能断言模型不存在。
// 模型存在性必须包含被请求路由限制排除的 Provider;目录未知时不能断言模型不存在。
let mut scoped_providers = providers.intersection(&self.providers).peekable();
if scoped_providers.peek().is_some()
&& scoped_providers.all(|provider| {
@@ -0,0 +1,26 @@
use gateway_core::engine::connection::ConnectionBudget;
#[test]
fn all_transports_share_one_opening_budget_but_other_requests_do_not() {
let request = ConnectionBudget::default();
let http = request.clone();
let websocket = request.clone();
for transport in [&http, &websocket, &http, &websocket] {
assert!(transport.begin().unwrap().is_some());
}
assert!(request.exhausted());
assert!(http.begin().is_err());
assert!(websocket.retry_delay(0, "req_shared").is_none());
assert!(ConnectionBudget::default().begin().is_ok());
}
#[test]
fn successful_opening_does_not_limit_later_generation_or_protocol_recovery() {
let request = ConnectionBudget::default();
request.begin().unwrap();
request.complete();
assert!(request.startup_remaining().is_none());
assert!(request.retry_delay(0, "req_sent").is_none());
assert!(!request.exhausted());
assert!(request.begin().unwrap().is_none());
}
@@ -4735,3 +4735,199 @@ fn provider_default_profile_is_applied_once_and_frozen_across_account_retries()
assert_eq!(first.expose_to_provider()["selection"], "plugin-default");
assert_eq!(contexts[1].request_profile(), Some(first));
}
#[test]
fn connection_recovery_keeps_one_account_and_http_transport_until_success() {
let operation = generate_operation();
let route_plan = plan(&operation);
let failure = || Script::Stream {
account_id: "acct_only",
items: vec![Err(ProviderError::new(
ProviderErrorKind::Transport,
UpstreamSendState::NotSent,
)
.with_connection_retry(AttemptTransport::Fallback))],
};
let (coordinator, store, provider) = coordinator(vec![
failure(),
failure(),
failure(),
Script::Stream {
account_id: "acct_only",
items: complete_stream(None),
},
]);
let mut session = block_on(coordinator.start(
model_request(&operation, SystemTime::now() + Duration::from_secs(30)),
operation,
route_plan,
None,
None,
CancellationToken::new(),
))
.unwrap();
block_on(session.collect_uncommitted()).unwrap();
block_on(session.commit_downstream(Some(200))).unwrap();
let contexts = provider.contexts.lock().unwrap();
assert_eq!(contexts.len(), 4);
let account = ProviderAccountId::new("acct_only").unwrap();
for context in contexts.iter().skip(1) {
assert_eq!(context.required_account(), Some(&account));
assert_eq!(context.transport(), AttemptTransport::Fallback);
assert!(!context.excluded_accounts().contains(&account));
}
let state = store.state.lock().unwrap();
assert_eq!(state.intermediate_failures, 3);
assert_eq!(state.finalizations[0].outcome, ExecutionOutcome::Succeeded);
assert_eq!(provider.released_leases.load(Ordering::SeqCst), 3);
drop(session);
assert_eq!(provider.released_leases.load(Ordering::SeqCst), 4);
}
#[test]
fn connection_recovery_exhaustion_preserves_transport_failure() {
let operation = generate_operation();
let route_plan = plan(&operation);
let (coordinator, store, provider) = coordinator(
(0..4)
.map(|_| Script::Stream {
account_id: "acct_only",
items: vec![Err(ProviderError::new(
ProviderErrorKind::Transport,
UpstreamSendState::NotSent,
)
.with_connection_retry(AttemptTransport::Fallback))],
})
.collect(),
);
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 error = block_on(session.collect_uncommitted()).unwrap_err();
assert!(
matches!(error, EngineError::Provider(ref error) if error.kind() == ProviderErrorKind::Transport)
);
assert_eq!(provider.contexts.lock().unwrap().len(), 4);
assert_eq!(
store.state.lock().unwrap().finalizations[0].attempt_count,
4
);
assert_eq!(provider.released_leases.load(Ordering::SeqCst), 3);
drop(session);
assert_eq!(provider.released_leases.load(Ordering::SeqCst), 4);
}
#[test]
fn connection_recovery_never_replays_sent_or_ambiguous_payloads() {
for send_state in [UpstreamSendState::Sent, UpstreamSendState::Ambiguous] {
let operation = generate_operation();
let route_plan = plan(&operation);
let (coordinator, store, provider) = coordinator(vec![Script::Stream {
account_id: "acct_only",
items: vec![Err(ProviderError::new(
ProviderErrorKind::Transport,
send_state,
)
.with_connection_retry(AttemptTransport::Fallback))],
}]);
let mut session = block_on(coordinator.start(
model_request(&operation, SystemTime::now() + Duration::from_secs(30)),
operation,
route_plan,
None,
None,
CancellationToken::new(),
))
.unwrap();
assert!(block_on(session.collect_uncommitted()).is_err());
assert_eq!(provider.contexts.lock().unwrap().len(), 1);
assert_eq!(store.state.lock().unwrap().intermediate_failures, 0);
}
}
#[test]
fn connection_backoff_cancellation_releases_account_without_another_attempt() {
let operation = generate_operation();
let route_plan = plan(&operation);
let (coordinator, _, provider) = coordinator(vec![Script::Stream {
account_id: "acct_first",
items: vec![Err(ProviderError::new(
ProviderErrorKind::Transport,
UpstreamSendState::NotSent,
)
.with_connection_retry(AttemptTransport::Fallback))],
}]);
let cancellation = CancellationToken::new();
let mut session = block_on(coordinator.start(
model_request(&operation, SystemTime::now() + Duration::from_secs(30)),
operation,
route_plan,
None,
None,
cancellation.clone(),
))
.expect("start execution");
let cancel = std::thread::spawn(move || {
std::thread::sleep(Duration::from_millis(50));
cancellation.cancel();
});
assert!(matches!(
block_on(session.collect_uncommitted()),
Err(EngineError::Cancelled)
));
cancel.join().expect("cancel task");
assert_eq!(provider.contexts.lock().expect("contexts").len(), 1);
assert_eq!(provider.released_leases.load(Ordering::SeqCst), 1);
}
#[test]
fn connection_failure_after_a_replay_safe_rejection_keeps_existing_account_rotation() {
let operation = generate_operation();
let route_plan = plan(&operation);
let (coordinator, _, provider) = coordinator(vec![
Script::Stream {
account_id: "acct_first",
items: vec![Err(ProviderError::new(
ProviderErrorKind::RateLimited,
UpstreamSendState::Sent,
)
.with_replay_safe())],
},
Script::Stream {
account_id: "acct_second",
items: vec![Err(ProviderError::new(
ProviderErrorKind::Transport,
UpstreamSendState::NotSent,
)
.with_connection_retry(AttemptTransport::Fallback))],
},
Script::Stream {
account_id: "acct_other",
items: complete_stream(None),
},
]);
let mut session = block_on(coordinator.start(
model_request(&operation, SystemTime::now() + Duration::from_secs(30)),
operation,
route_plan,
None,
None,
CancellationToken::new(),
))
.unwrap();
block_on(session.collect_uncommitted()).unwrap();
let contexts = provider.contexts.lock().unwrap();
assert_eq!(contexts.len(), 3);
assert!(contexts[2].required_account().is_none());
assert!(
contexts[2]
.excluded_accounts()
.contains(&ProviderAccountId::new("acct_second").unwrap())
);
}
@@ -111,7 +111,6 @@ fn service(admissions: Arc<Admissions>, budget: Arc<Budget>) -> DefaultExecution
Arc::new(TrackingExecutionStore::default()),
ProviderRegistry::default(),
admissions,
Arc::new(UnusedCircuits),
Arc::new(UnusedContinuation),
Arc::new(RecordingClientApiKeyUsage::default()),
)
@@ -280,7 +279,6 @@ fn reused_client_uses_updated_limits_for_each_execution() {
Arc::new(TrackingExecutionStore::default()),
ProviderRegistry::default(),
admissions.clone(),
Arc::new(UnusedCircuits),
Arc::new(UnusedContinuation),
Arc::new(RecordingClientApiKeyUsage::default()),
);
@@ -312,7 +310,6 @@ fn reused_client_cannot_start_after_key_disable_or_snapshot_suspension() {
Arc::new(TrackingExecutionStore::default()),
ProviderRegistry::default(),
admissions.clone(),
Arc::new(UnusedCircuits),
Arc::new(UnusedContinuation),
Arc::new(RecordingClientApiKeyUsage::default()),
);
@@ -416,7 +413,6 @@ fn reused_client_is_reauthenticated_by_the_current_frontend_plan_without_identit
Arc::new(TrackingExecutionStore::default()),
ProviderRegistry::default(),
admissions.clone(),
Arc::new(UnusedCircuits),
Arc::new(UnusedContinuation),
Arc::new(RecordingClientApiKeyUsage::default()),
)
@@ -508,7 +504,6 @@ fn early_failure_service(
store,
ProviderRegistry::new([Arc::new(LocalFailingProvider) as Arc<dyn Provider>]).unwrap(),
admissions,
Arc::new(UnusedCircuits),
Arc::new(UnusedContinuation),
Arc::new(RecordingClientApiKeyUsage::default()),
)
@@ -738,7 +733,6 @@ fn published_provider_update_does_not_replace_an_already_started_request() {
Arc::new(TrackingExecutionStore::default()),
index.registry(),
Arc::new(Admissions::default()),
Arc::new(UnusedCircuits),
Arc::new(UnusedContinuation),
Arc::new(RecordingClientApiKeyUsage::default()),
);
@@ -1020,7 +1014,6 @@ fn completed_parent_rejects_new_nested_execution_and_cancels_an_active_child() {
Arc::new(TrackingExecutionStore::default()),
provider_index.registry(),
Arc::new(Admissions::default()),
Arc::new(UnusedCircuits),
Arc::new(UnusedContinuation),
Arc::new(RecordingClientApiKeyUsage::default()),
)
@@ -1213,7 +1206,6 @@ fn native_response_processing_uses_the_real_translation_boundary() {
Arc::new(TrackingExecutionStore::default()),
provider_index.registry(),
Arc::new(Admissions::default()),
Arc::new(UnusedCircuits),
Arc::new(UnusedContinuation),
Arc::new(RecordingClientApiKeyUsage::default()),
)
@@ -1365,7 +1357,6 @@ fn response_translation_zero_output_preserves_canonical_usage_and_finalization()
store.clone(),
provider_index.registry(),
Arc::new(Admissions::default()),
Arc::new(UnusedCircuits),
Arc::new(UnusedContinuation),
Arc::new(RecordingClientApiKeyUsage::default()),
);
@@ -1522,7 +1513,6 @@ fn native_response_failure_after_downstream_commit_is_not_replayed() {
Arc::new(TrackingExecutionStore::default()),
provider_index.registry(),
Arc::new(Admissions::default()),
Arc::new(UnusedCircuits),
Arc::new(UnusedContinuation),
Arc::new(RecordingClientApiKeyUsage::default()),
);
@@ -1609,7 +1599,6 @@ fn model_routing_policy_selects_native_and_generation_plugin_providers() {
Arc::new(TrackingExecutionStore::default()),
provider_index.registry(),
Arc::new(Admissions::default()),
Arc::new(UnusedCircuits),
Arc::new(UnusedContinuation),
Arc::new(RecordingClientApiKeyUsage::default()),
)
@@ -1665,7 +1654,6 @@ fn model_routing_reject_stops_before_provider_execution() {
Arc::new(TrackingExecutionStore::default()),
provider_index.registry(),
Arc::new(Admissions::default()),
Arc::new(UnusedCircuits),
Arc::new(UnusedContinuation),
Arc::new(RecordingClientApiKeyUsage::default()),
)
@@ -1751,7 +1739,6 @@ fn model_routing_cannot_expand_the_frozen_key_model_scope() {
Arc::new(TrackingExecutionStore::default()),
provider_index.registry(),
Arc::new(Admissions::default()),
Arc::new(UnusedCircuits),
Arc::new(UnusedContinuation),
Arc::new(RecordingClientApiKeyUsage::default()),
)
@@ -1899,7 +1886,6 @@ fn scheduler_reject_is_a_terminal_policy_rejection_before_upstream_send() {
Arc::new(TrackingExecutionStore::default()),
provider_index.registry(),
Arc::new(Admissions::default()),
Arc::new(UnusedCircuits),
Arc::new(UnusedContinuation),
Arc::new(RecordingClientApiKeyUsage::default()),
)
@@ -2175,7 +2161,6 @@ fn observation_service(
Arc::new(TrackingExecutionStore::default()),
provider_extensions.registry(),
admissions,
Arc::new(UnusedCircuits),
Arc::new(UnusedContinuation),
Arc::new(RecordingClientApiKeyUsage::default()),
)
@@ -2216,7 +2201,6 @@ fn fallback_exhaustion_observation_uses_the_last_actual_provider_not_the_candida
Arc::new(TrackingExecutionStore::default()),
provider_extensions.registry(),
Arc::new(Admissions::default()),
Arc::new(UnusedCircuits),
Arc::new(UnusedContinuation),
Arc::new(RecordingClientApiKeyUsage::default()),
)
@@ -2284,7 +2268,6 @@ fn routing_external_effect_stops_a_not_sent_provider_retry() {
Arc::new(TrackingExecutionStore::default()),
provider_extensions.registry(),
Arc::new(Admissions::default()),
Arc::new(UnusedCircuits),
Arc::new(UnusedContinuation),
Arc::new(RecordingClientApiKeyUsage::default()),
)
@@ -2486,7 +2469,6 @@ fn charged_service(
}) as Arc<dyn Provider>])
.unwrap(),
admissions,
Arc::new(UnusedCircuits),
Arc::new(UnusedContinuation),
Arc::new(RecordingClientApiKeyUsage::default()),
);
@@ -2764,9 +2746,8 @@ use gateway_core::engine::continuation::{
};
use gateway_core::engine::execution::{
ClientApiKeyUsageSink, ClientKeyVerifier, ClientTransport, DefaultExecutionService,
ExecutionRequestMetadata, ExecutionService, ExecutionSession, ProviderCircuitDecision,
ProviderCircuitError, ProviderCircuitPort, StartExecution, StartProviderExecution,
provider_failure_affects_circuit,
ExecutionRequestMetadata, ExecutionService, ExecutionSession, StartExecution,
StartProviderExecution,
};
use gateway_core::engine::extensions::ProviderExtensionIndex;
use gateway_core::engine::nested::NestedModelExecutionRequest;
@@ -2814,30 +2795,6 @@ use gateway_core::runtime::{
use gateway_core::upstream::{UpstreamSendState, UpstreamTransport};
use serde_json::{Value, json};
#[test]
fn only_provider_attributable_failures_should_affect_circuit() {
assert!(provider_failure_affects_circuit(ProviderErrorKind::Timeout));
assert!(provider_failure_affects_circuit(
ProviderErrorKind::Transport
));
assert!(!provider_failure_affects_circuit(
ProviderErrorKind::RateLimited
));
assert!(!provider_failure_affects_circuit(
ProviderErrorKind::InvalidRequest
));
assert!(!provider_failure_affects_circuit(
ProviderErrorKind::ContinuationRecoveryRequired
));
assert!(!provider_failure_affects_circuit(
ProviderErrorKind::UpstreamCapacityUnavailable
));
// 上游 close 1009(message too big)是请求自身的问题:不得熔断 provider。
assert!(!provider_failure_affects_circuit(
ProviderErrorKind::MessageTooBig
));
}
#[test]
fn account_probe_should_not_write_to_the_persistent_execution_store() {
let store = Arc::new(TrackingExecutionStore::default());
@@ -2846,7 +2803,6 @@ fn account_probe_should_not_write_to_the_persistent_execution_store() {
store.clone(),
ProviderRegistry::default(),
Arc::new(UnusedAdmissions),
Arc::new(UnusedCircuits),
Arc::new(UnusedContinuation),
Arc::new(RecordingClientApiKeyUsage::default()),
);
@@ -2878,7 +2834,6 @@ fn probe_failures_should_be_observable_without_a_model_request_row() {
store.clone(),
providers,
Arc::new(UnusedAdmissions),
Arc::new(UnusedCircuits),
Arc::new(UnusedContinuation),
Arc::new(RecordingClientApiKeyUsage::default()),
);
@@ -2924,7 +2879,6 @@ fn provider_local_probe_failure_should_remain_distinct_from_upstream() {
store,
providers,
Arc::new(UnusedAdmissions),
Arc::new(UnusedCircuits),
Arc::new(UnusedContinuation),
Arc::new(RecordingClientApiKeyUsage::default()),
);
@@ -2967,7 +2921,6 @@ fn diagnostic_probe_does_not_apply_data_plane_account_model_policy() {
ProviderRegistry::new([Arc::new(LocalFailingProvider) as Arc<dyn Provider>])
.expect("providers"),
Arc::new(UnusedAdmissions),
Arc::new(UnusedCircuits),
Arc::new(UnusedContinuation),
Arc::new(RecordingClientApiKeyUsage::default()),
);
@@ -2999,7 +2952,6 @@ fn probe_observation_store_failure_preserves_the_provider_error() {
store.clone(),
providers,
Arc::new(UnusedAdmissions),
Arc::new(UnusedCircuits),
Arc::new(UnusedContinuation),
Arc::new(RecordingClientApiKeyUsage::default()),
);
@@ -3094,7 +3046,6 @@ fn client_catalog_forwards_scope_maps_whole_objects_and_omits_unroutable_models(
ProviderRegistry::new([Arc::new(NativeCatalogProvider { fail }) as Arc<dyn Provider>])
.expect("registry"),
Arc::new(UnusedAdmissions),
Arc::new(UnusedCircuits),
Arc::new(UnusedContinuation),
Arc::new(RecordingClientApiKeyUsage::default()),
);
@@ -3246,7 +3197,6 @@ fn successful_authentication_should_record_client_key_usage() {
Arc::new(TrackingExecutionStore::default()),
ProviderRegistry::default(),
Arc::new(UnusedAdmissions),
Arc::new(UnusedCircuits),
Arc::new(UnusedContinuation),
usage.clone(),
);
@@ -3266,7 +3216,6 @@ fn client_key_verification_should_not_record_client_key_usage() {
Arc::new(TrackingExecutionStore::default()),
ProviderRegistry::default(),
Arc::new(UnusedAdmissions),
Arc::new(UnusedCircuits),
Arc::new(UnusedContinuation),
usage.clone(),
);
@@ -3288,7 +3237,6 @@ fn request_verification_should_apply_entry_authentication_without_recording_key_
Arc::new(TrackingExecutionStore::default()),
ProviderRegistry::default(),
Arc::new(UnusedAdmissions),
Arc::new(UnusedCircuits),
Arc::new(UnusedContinuation),
usage.clone(),
);
@@ -3322,7 +3270,6 @@ fn assert_provider_endpoint_observation(model: Option<&str>) {
}) as Arc<dyn Provider>])
.expect("provider registry"),
Arc::new(UnusedAdmissions),
Arc::new(UnusedCircuits),
Arc::new(UnusedContinuation),
Arc::new(RecordingClientApiKeyUsage::default()),
);
@@ -3368,84 +3315,14 @@ fn assert_provider_endpoint_observation(model: Option<&str>) {
assert!(upstream_models.iter().all(Option::is_none));
}
#[test]
fn circuit_store_failure_should_fail_open_during_request_start() {
let service = DefaultExecutionService::new(
RuntimeSnapshotHandle::new(start_snapshot()),
Arc::new(TrackingExecutionStore::default()),
ProviderRegistry::default(),
Arc::new(UnusedAdmissions),
Arc::new(FailingDecisionCircuits),
Arc::new(UnusedContinuation),
Arc::new(RecordingClientApiKeyUsage::default()),
);
let client = service
.authenticate("sk_start_test")
.expect("authenticated client");
let started = block_on(service.start(StartExecution {
client,
public_model: PublicModelId::new("gpt-start").expect("public model"),
operation: start_operation(),
metadata: ExecutionRequestMetadata {
protocol: "openai".to_owned(),
endpoint: "/v1/responses".to_owned(),
transport: ClientTransport::HttpJson,
stream: false,
client_ip: None,
user_agent: None,
previous_response_id: None,
},
}))
.expect("recoverable circuit state must not reject the request");
assert!(!started.session.is_finalized());
}
#[test]
fn slow_circuit_store_should_time_out_and_fail_open_during_request_start() {
let service = DefaultExecutionService::new(
RuntimeSnapshotHandle::new(start_snapshot()),
Arc::new(TrackingExecutionStore::default()),
ProviderRegistry::default(),
Arc::new(UnusedAdmissions),
Arc::new(PendingDecisionCircuits),
Arc::new(UnusedContinuation),
Arc::new(RecordingClientApiKeyUsage::default()),
);
let client = service
.authenticate("sk_start_test")
.expect("authenticated client");
let started_at = Instant::now();
let started = block_on(service.start(StartExecution {
client,
public_model: PublicModelId::new("gpt-start").expect("public model"),
operation: start_operation(),
metadata: ExecutionRequestMetadata {
protocol: "openai".to_owned(),
endpoint: "/v1/responses".to_owned(),
transport: ClientTransport::HttpJson,
stream: false,
client_ip: None,
user_agent: None,
previous_response_id: None,
},
}))
.expect("slow recoverable circuit state must not reject the request");
assert!(started_at.elapsed() < Duration::from_secs(2));
assert!(!started.session.is_finalized());
}
#[test]
fn known_catalog_should_reject_a_model_that_the_provider_did_not_publish() {
let store = Arc::new(TrackingExecutionStore::default());
let service = DefaultExecutionService::new(
RuntimeSnapshotHandle::new(start_snapshot()),
Arc::new(TrackingExecutionStore::default()),
store.clone(),
ProviderRegistry::default(),
Arc::new(UnusedAdmissions),
Arc::new(UnusedCircuits),
Arc::new(UnusedContinuation),
Arc::new(RecordingClientApiKeyUsage::default()),
);
@@ -3479,6 +3356,10 @@ fn known_catalog_should_reject_a_model_that_the_provider_did_not_publish() {
"the requested model was not found in the provider catalogs available to this API key; check the model name",
)
);
let rejections = store.entry_rejections.lock().unwrap();
assert_eq!(rejections.len(), 1);
assert_eq!(rejections[0].error.kind(), GatewayErrorKind::ModelNotFound);
assert!(store.requests.lock().unwrap().is_empty());
}
#[test]
@@ -3488,7 +3369,6 @@ fn continuation_owned_by_another_client_api_key_should_fail_closed() {
Arc::new(TrackingExecutionStore::default()),
ProviderRegistry::default(),
Arc::new(UnusedAdmissions),
Arc::new(UnusedCircuits),
Arc::new(RejectedContinuation::OwnershipMismatch),
Arc::new(RecordingClientApiKeyUsage::default()),
);
@@ -3516,7 +3396,6 @@ fn invalid_continuation_record_should_not_be_forwarded_as_an_external_handle() {
Arc::new(TrackingExecutionStore::default()),
ProviderRegistry::default(),
Arc::new(UnusedAdmissions),
Arc::new(UnusedCircuits),
Arc::new(RejectedContinuation::InvalidData),
Arc::new(RecordingClientApiKeyUsage::default()),
);
@@ -3573,6 +3452,7 @@ impl ClientApiKeyUsageSink for RecordingClientApiKeyUsage {
struct TrackingExecutionStore {
touched: AtomicBool,
probe_failures: Mutex<Vec<String>>,
entry_rejections: Mutex<Vec<gateway_core::engine::EntryRejection>>,
requests: Mutex<Vec<NewModelRequest>>,
attempts: Mutex<Vec<AttemptRecord>>,
finalizations: Mutex<Vec<ModelRequestFinalization>>,
@@ -3635,6 +3515,14 @@ impl TrackingExecutionStore {
#[async_trait]
impl ExecutionStore for TrackingExecutionStore {
async fn record_entry_rejection(
&self,
rejection: gateway_core::engine::EntryRejection,
) -> Result<(), StoreError> {
self.entry_rejections.lock().unwrap().push(rejection);
Ok(())
}
async fn create_model_request(&self, request: NewModelRequest) -> Result<(), StoreError> {
self.touch();
self.creates.fetch_add(1, Ordering::SeqCst);
@@ -3752,81 +3640,6 @@ impl ClientAdmissionPort for UnusedAdmissions {
}
}
struct UnusedCircuits;
impl ProviderCircuitPort for UnusedCircuits {
fn decision<'a>(
&'a self,
_: &'a ProviderKind,
) -> BoxFuture<'a, Result<ProviderCircuitDecision, ProviderCircuitError>> {
Box::pin(async { Ok(ProviderCircuitDecision::Allow) })
}
fn observe_failure<'a>(
&'a self,
_: &'a ProviderKind,
) -> BoxFuture<'a, Result<(), ProviderCircuitError>> {
Box::pin(async { Ok(()) })
}
fn observe_success<'a>(
&'a self,
_: &'a ProviderKind,
) -> BoxFuture<'a, Result<(), ProviderCircuitError>> {
Box::pin(async { Ok(()) })
}
}
struct FailingDecisionCircuits;
impl ProviderCircuitPort for FailingDecisionCircuits {
fn decision<'a>(
&'a self,
_: &'a ProviderKind,
) -> BoxFuture<'a, Result<ProviderCircuitDecision, ProviderCircuitError>> {
Box::pin(async { Err(ProviderCircuitError) })
}
fn observe_failure<'a>(
&'a self,
_: &'a ProviderKind,
) -> BoxFuture<'a, Result<(), ProviderCircuitError>> {
Box::pin(async { Ok(()) })
}
fn observe_success<'a>(
&'a self,
_: &'a ProviderKind,
) -> BoxFuture<'a, Result<(), ProviderCircuitError>> {
Box::pin(async { Ok(()) })
}
}
struct PendingDecisionCircuits;
impl ProviderCircuitPort for PendingDecisionCircuits {
fn decision<'a>(
&'a self,
_: &'a ProviderKind,
) -> BoxFuture<'a, Result<ProviderCircuitDecision, ProviderCircuitError>> {
Box::pin(futures::future::pending())
}
fn observe_failure<'a>(
&'a self,
_: &'a ProviderKind,
) -> BoxFuture<'a, Result<(), ProviderCircuitError>> {
Box::pin(async { Ok(()) })
}
fn observe_success<'a>(
&'a self,
_: &'a ProviderKind,
) -> BoxFuture<'a, Result<(), ProviderCircuitError>> {
Box::pin(async { Ok(()) })
}
}
struct UnusedContinuation;
impl NativeContinuationPort for UnusedContinuation {
@@ -4072,7 +3885,6 @@ fn account_wait_inherits_the_budget_spent_during_client_admission() {
)
.unwrap(),
admissions.clone(),
Arc::new(UnusedCircuits),
Arc::new(UnusedContinuation),
Arc::new(RecordingClientApiKeyUsage::default()),
);
@@ -4180,7 +3992,6 @@ fn queue_service(
Arc::new(TrackingExecutionStore::default()),
ProviderRegistry::default(),
admissions.clone(),
Arc::new(UnusedCircuits),
Arc::new(UnusedContinuation),
Arc::new(RecordingClientApiKeyUsage::default()),
);
@@ -4305,7 +4116,6 @@ fn reused_websocket_client_gets_group_fast_policy_from_each_new_request_snapshot
Arc::new(TrackingExecutionStore::default()),
ProviderRegistry::new([provider.clone() as Arc<dyn Provider>]).unwrap(),
admissions,
Arc::new(UnusedCircuits),
Arc::new(UnusedContinuation),
Arc::new(RecordingClientApiKeyUsage::default()),
);
@@ -4325,3 +4135,54 @@ fn reused_websocket_client_gets_group_fast_policy_from_each_new_request_snapshot
}
assert_eq!(*provider.policies.lock().unwrap(), vec![true, false]);
}
#[test]
fn repeated_connection_failures_never_block_later_requests_for_the_provider() {
let store = Arc::new(TrackingExecutionStore::default());
let service = DefaultExecutionService::new(
RuntimeSnapshotHandle::new(client_snapshot()),
store.clone(),
ProviderRegistry::new([Arc::new(ColdFailingProvider {
requested_model: None,
}) as Arc<dyn Provider>])
.expect("provider registry"),
Arc::new(UnusedAdmissions),
Arc::new(UnusedContinuation),
Arc::new(RecordingClientApiKeyUsage::default()),
);
for _ in 0..5 {
let client = service
.authenticate("sk_usage_test")
.expect("authenticated client");
let operation = Operation::GenerateImage(ImageRequest::from_raw_json(
ImageRequestKind::Generation,
RawJsonPayload::new(
"openai",
Bytes::from_static(br#"{"model":"gpt-image-2","prompt":"hello"}"#),
)
.expect("image payload"),
));
let mut started = block_on(service.start_provider_endpoint(StartProviderExecution {
client,
provider: ProviderKind::new("openai").expect("provider"),
upstream_model: None,
operation,
metadata: ExecutionRequestMetadata {
protocol: "openai".to_owned(),
endpoint: "/v1/images/generations".to_owned(),
transport: ClientTransport::HttpJson,
stream: false,
client_ip: None,
user_agent: None,
previous_response_id: None,
},
}))
.expect("provider endpoint request should start without a text catalog entry");
let _error = block_on(started.session.collect_uncommitted())
.expect_err("the cold provider stops execution after persistence");
}
assert_eq!(store.requests.lock().unwrap().len(), 5);
assert!(store.entry_rejections.lock().unwrap().is_empty());
}
@@ -1,4 +1,5 @@
mod admission;
mod connection;
mod continuation;
mod coordinator;
mod execution;
+1 -11
View File
@@ -57,7 +57,7 @@ impl StoreBundle {
Ok(())
}
/// 命令成功、失败或取消后有界排空账本、Key 使用、准入释放与 circuit 反馈。
/// 命令成功、失败或取消后有界排空账本、Key 使用与准入释放。
pub async fn shutdown_command_line_writes(&mut self) -> Result<(), CommandStoreDrainError> {
match self.command_drain.take() {
Some(drain) => drain.shutdown().await,
@@ -192,12 +192,6 @@ async fn connect(
let admissions: Arc<dyn gateway_core::engine::admission::ClientAdmissionPort> = Arc::new(
redis::RedisClientAdmissionRepository::new(redis_connection.clone(), REDIS_NAMESPACE)?,
);
let circuits: Arc<dyn gateway_core::engine::execution::ProviderCircuitPort> =
Arc::new(redis::RedisProviderCircuitRepository::new(
redis_connection.clone(),
REDIS_NAMESPACE,
gateway_core::engine::execution::ProviderCircuitPolicy::default(),
)?);
// Continuation affinity 是下一轮请求的路由事实,Core 必须直接等待 Redis 确认。
let continuation: Arc<dyn gateway_core::engine::continuation::NativeContinuationPort> =
Arc::new(redis::RedisNativeContinuationRepository::new(
@@ -206,7 +200,6 @@ async fn connect(
)?);
let (admissions, admission_release_writer) =
redis::BufferedClientAdmissionPort::new(admissions);
let (circuits, circuit_feedback_writer) = redis::BufferedProviderCircuitPort::new(circuits);
let core_ports = CoreStorePorts::new(
execution,
(
@@ -215,7 +208,6 @@ async fn connect(
pool.clone(),
)),
),
Arc::new(circuits),
continuation,
(
Arc::new(postgres::PgRuntimeSnapshotRepository::new(pool.clone())),
@@ -263,7 +255,6 @@ async fn connect(
execution_writer,
client_key_usage_writer,
admission_release_writer,
circuit_feedback_writer,
retention,
)?,
None,
@@ -274,7 +265,6 @@ async fn connect(
execution: execution_writer,
client_key_usage: client_key_usage_writer,
admission_release: admission_release_writer,
circuit_feedback: circuit_feedback_writer,
}),
),
};
@@ -1108,6 +1108,48 @@ impl ExecutionStore for PgExecutionStore {
.map_err(core_store_error)
}
async fn record_entry_rejection(
&self,
rejection: gateway_core::engine::EntryRejection,
) -> Result<(), CoreStoreError> {
let error = rejection.error;
super::OpsEventRepository::append_ops_event(
&super::PgOpsEventRepository::new(self.pool.clone()),
super::OpsEvent {
id: Uuid::now_v7().to_string(),
model_request_id: None,
attempt_index: None,
level: super::OpsEventLevel::Warning,
component: "request_entry".to_owned(),
operation: "reject".to_owned(),
provider_kind: None,
provider_account_id: None,
provider_account_ref: None,
upstream_model_id: None,
failure_kind: error.kind().as_str().to_owned(),
upstream_send_state: None,
raw_upstream_error: None,
status_code: None,
provider_error_code: error.client_error_code().map(str::to_owned),
retry_after_ms: error
.retry_after()
.and_then(|delay| u64::try_from(delay.as_millis()).ok()),
upstream_request_id: None,
latency_ms: u64::try_from(rejection.latency.as_millis()).ok(),
// 入口尚无 model_requests 行,关联 ID 放在安全消息中,不能伪造外键。
message: serde_json::json!({
"requestId": rejection.request_id.as_str(),
"clientKeyId": rejection.client_key_id.as_str(),
"message": error.client_message(),
})
.to_string(),
created_at: Utc::now(),
},
)
.await
.map_err(core_store_error)
}
async fn record_probe_failure(&self, failure: ProbeFailure) -> Result<(), CoreStoreError> {
let error = failure.error;
let retry_after_ms = error
@@ -9,8 +9,8 @@ use std::time::{Duration, Instant, SystemTime};
use async_trait::async_trait;
use futures::future::try_join_all;
use gateway_core::engine::{
AttemptRecord, ExecutionStore, IntermediateFailure, ModelRequestFinalization, ModelRequestId,
NewModelRequest, ProbeFailure, RecoveryReport,
AttemptRecord, EntryRejection, ExecutionStore, IntermediateFailure, ModelRequestFinalization,
ModelRequestId, NewModelRequest, ProbeFailure, RecoveryReport,
};
use gateway_core::error::{ProviderError, StoreError};
use gateway_core::lifecycle::CancellationToken;
@@ -358,6 +358,13 @@ where
Ok(())
}
async fn record_entry_rejection(&self, rejection: EntryRejection) -> Result<(), StoreError> {
self.enqueue(ExecutionObservationWrite::EntryRejection(Box::new(
rejection,
)));
Ok(())
}
async fn record_probe_failure(&self, failure: ProbeFailure) -> Result<(), StoreError> {
self.enqueue(ExecutionObservationWrite::ProbeFailure(Box::new(failure)));
Ok(())
@@ -677,6 +684,7 @@ enum ExecutionObservationWrite {
},
IntermediateFailure(Box<IntermediateFailure>),
ProbeFailure(Box<ProbeFailure>),
EntryRejection(Box<EntryRejection>),
Finalize(Box<ModelRequestFinalization>),
}
@@ -691,6 +699,7 @@ impl ExecutionObservationWrite {
Self::RecordClientStatus { .. } => "record_client_status",
Self::IntermediateFailure(_) => "record_intermediate_failure",
Self::ProbeFailure(_) => "record_probe_failure",
Self::EntryRejection(_) => "record_entry_rejection",
Self::Finalize(_) => "finalize_model_request",
}
}
@@ -704,7 +713,7 @@ impl ExecutionObservationWrite {
| Self::MarkDownstreamCommitted { request_id, .. }
| Self::RecordClientStatus { request_id, .. } => Some(request_id.as_str()),
Self::IntermediateFailure(failure) => Some(failure.request_id.as_str()),
Self::ProbeFailure(_) => None,
Self::ProbeFailure(_) | Self::EntryRejection(_) => None,
Self::Finalize(finalization) => Some(finalization.request_id.as_str()),
}
}
@@ -730,6 +739,14 @@ impl ExecutionObservationWrite {
failure.upstream_request_id.as_deref(),
])
.saturating_add(provider_error_bytes(&failure.error)),
Self::EntryRejection(rejection) => text_bytes([
Some(rejection.request_id.as_str()),
Some(rejection.client_key_id.as_str()),
Some(rejection.error.client_message()),
rejection.error.client_error_code(),
rejection.error.client_error_type(),
])
.saturating_add(size_of::<EntryRejection>()),
Self::ProbeFailure(failure) => text_bytes([
Some(failure.provider_kind.as_str()),
Some(failure.account_id.as_str()),
@@ -799,6 +816,7 @@ impl ExecutionObservationWrite {
}
Self::IntermediateFailure(failure) => store.record_intermediate_failure(*failure).await,
Self::ProbeFailure(failure) => store.record_probe_failure(*failure).await,
Self::EntryRejection(rejection) => store.record_entry_rejection(*rejection).await,
Self::Finalize(finalization) => store.finalize_model_request(*finalization).await,
}
}
@@ -15,12 +15,8 @@ use gateway_core::engine::admission::{
ClientAdmissionDecision, ClientAdmissionError, ClientAdmissionPort, ClientAdmissionRecovery,
ClientAdmissionRequest, ClientAdmissionRestoreResult,
};
use gateway_core::engine::execution::{
ProviderCircuitDecision, ProviderCircuitError, ProviderCircuitPort,
};
use gateway_core::lifecycle::CancellationToken;
use gateway_core::policy::ClientApiKeyId;
use gateway_core::routing::ProviderKind;
use gateway_core::task::{DaemonTask, WorkerTaskError};
use tokio::sync::{Mutex, mpsc};
@@ -154,123 +150,6 @@ impl DaemonTask for ClientAdmissionReleaseWriter {
}
}
/// Circuit decision 仍直读 Redis;success/failure feedback 只做有界入队。
#[derive(Clone)]
pub struct BufferedProviderCircuitPort {
inner: Arc<dyn ProviderCircuitPort>,
sender: mpsc::Sender<CircuitFeedback>,
}
impl BufferedProviderCircuitPort {
#[must_use]
pub fn new(inner: Arc<dyn ProviderCircuitPort>) -> (Self, ProviderCircuitFeedbackWriter) {
Self::with_capacity(
inner,
NonZeroUsize::new(DEFAULT_QUEUE_CAPACITY).expect("queue capacity is non-zero"),
)
}
#[must_use]
pub fn with_capacity(
inner: Arc<dyn ProviderCircuitPort>,
capacity: NonZeroUsize,
) -> (Self, ProviderCircuitFeedbackWriter) {
let (sender, receiver) = mpsc::channel(capacity.get());
(
Self {
inner: Arc::clone(&inner),
sender,
},
ProviderCircuitFeedbackWriter {
inner,
receiver: Mutex::new(receiver),
},
)
}
fn enqueue(&self, feedback: CircuitFeedback) {
if let Err(error) = self.sender.try_send(feedback) {
let reason = match error {
mpsc::error::TrySendError::Full(_) => "full",
mpsc::error::TrySendError::Closed(_) => "closed",
};
tracing::warn!(
operation = "record_provider_circuit_feedback",
reason,
"Redis 协调队列不可用,已丢弃本次 circuit feedback"
);
}
}
}
impl ProviderCircuitPort for BufferedProviderCircuitPort {
fn decision<'a>(
&'a self,
provider_kind: &'a ProviderKind,
) -> BoxFuture<'a, Result<ProviderCircuitDecision, ProviderCircuitError>> {
self.inner.decision(provider_kind)
}
fn observe_failure<'a>(
&'a self,
provider_kind: &'a ProviderKind,
) -> BoxFuture<'a, Result<(), ProviderCircuitError>> {
self.enqueue(CircuitFeedback::Failure(provider_kind.clone()));
Box::pin(ready(Ok(())))
}
fn observe_success<'a>(
&'a self,
provider_kind: &'a ProviderKind,
) -> BoxFuture<'a, Result<(), ProviderCircuitError>> {
self.enqueue(CircuitFeedback::Success(provider_kind.clone()));
Box::pin(ready(Ok(())))
}
}
enum CircuitFeedback {
Failure(ProviderKind),
Success(ProviderKind),
}
pub struct ProviderCircuitFeedbackWriter {
inner: Arc<dyn ProviderCircuitPort>,
receiver: Mutex<mpsc::Receiver<CircuitFeedback>>,
}
impl DaemonTask for ProviderCircuitFeedbackWriter {
fn run(&self, cancellation: CancellationToken) -> BoxFuture<'_, Result<(), WorkerTaskError>> {
Box::pin(async move {
let mut receiver = self.receiver.lock().await;
loop {
let feedback = tokio::select! {
() = cancellation.cancelled() => {
drain_circuit_feedback(&mut receiver, self.inner.as_ref()).await;
return Ok(());
},
feedback = receiver.recv() => feedback,
};
let Some(feedback) = feedback else {
return Err(WorkerTaskError::safe(
"provider circuit feedback queue closed",
));
};
let result = match feedback {
CircuitFeedback::Failure(provider) => {
self.inner.observe_failure(&provider).await
}
CircuitFeedback::Success(provider) => {
self.inner.observe_success(&provider).await
}
};
if let Err(error) = result {
tracing::warn!(%error, "Provider circuit feedback 后台写入失败");
}
}
})
}
}
async fn drain_admission_releases(
receiver: &mut mpsc::Receiver<AdmissionRelease>,
port: &dyn ClientAdmissionPort,
@@ -292,30 +171,3 @@ async fn drain_admission_releases(
}
}
}
async fn drain_circuit_feedback(
receiver: &mut mpsc::Receiver<CircuitFeedback>,
port: &dyn ProviderCircuitPort,
) {
receiver.close();
let started_at = Instant::now();
while let Some(feedback) = receiver.recv().await {
let remaining = SHUTDOWN_DRAIN_TIMEOUT.saturating_sub(started_at.elapsed());
if remaining.is_zero() {
tracing::warn!("Provider circuit 关闭排空超时,剩余反馈已丢弃");
break;
}
let result = match feedback {
CircuitFeedback::Failure(provider) => {
tokio::time::timeout(remaining, port.observe_failure(&provider)).await
}
CircuitFeedback::Success(provider) => {
tokio::time::timeout(remaining, port.observe_success(&provider)).await
}
};
if result.is_err() {
tracing::warn!("Provider circuit 关闭排空超时,剩余反馈已丢弃");
break;
}
}
}
@@ -12,7 +12,6 @@ mod credential_leases;
mod credential_state;
mod native_continuation;
mod oauth_pending;
mod provider_circuit;
mod provider_session_affinity;
mod provider_session_exclusion;
mod runtime_change;
@@ -28,7 +27,6 @@ pub use credential_leases::*;
pub use credential_state::*;
pub use native_continuation::*;
pub use oauth_pending::*;
pub use provider_circuit::*;
pub use provider_session_affinity::*;
pub use provider_session_exclusion::*;
pub use runtime_change::*;
@@ -1,208 +0,0 @@
//! Provider 的可重建 Redis circuit。
use async_trait::async_trait;
use chrono::{DateTime, Utc};
use gateway_core::engine::execution::{
ProviderCircuitDecision as CoreCircuitDecision, ProviderCircuitError, ProviderCircuitPolicy,
ProviderCircuitPort,
};
use gateway_core::routing::ProviderKind;
use redis::{Script, aio::ConnectionManager};
use crate::{StoreError, StoreResult, redis_unavailable, require_nonempty};
use super::{namespace, resource_fingerprint};
const FAILURE_SCRIPT: &str = r#"
local clock = redis.call('TIME')
local now_ms = (tonumber(clock[1]) * 1000) + math.floor(tonumber(clock[2]) / 1000)
local failures = redis.call('HINCRBY', KEYS[1], 'failures', 1)
local open_until = tonumber(redis.call('HGET', KEYS[1], 'open_until_ms') or '0')
if failures >= tonumber(ARGV[1]) then
open_until = now_ms + tonumber(ARGV[2])
redis.call('HSET', KEYS[1], 'open_until_ms', tostring(open_until))
end
redis.call('HSET', KEYS[1], 'observed_at_ms', tostring(now_ms))
redis.call('PEXPIRE', KEYS[1], math.max(86400000, tonumber(ARGV[2]) * 2))
return {tostring(failures), tostring(open_until)}
"#;
const DECISION_SCRIPT: &str = r#"
local clock = redis.call('TIME')
local now_ms = (tonumber(clock[1]) * 1000) + math.floor(tonumber(clock[2]) / 1000)
local open_until = tonumber(redis.call('HGET', KEYS[1], 'open_until_ms') or '0')
if open_until > now_ms then return {0, tostring(open_until)} end
return {1, '0'}
"#;
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ProviderCircuitDecision {
Allow,
BlockedUntil(DateTime<Utc>),
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct ProviderCircuitObservation {
pub failure_count: u32,
pub open_until: Option<DateTime<Utc>>,
}
#[async_trait]
pub trait ProviderCircuitRepository: Send + Sync {
async fn provider_circuit_decision(
&self,
provider_kind: &str,
) -> StoreResult<ProviderCircuitDecision>;
async fn observe_provider_failure(
&self,
provider_kind: &str,
) -> StoreResult<ProviderCircuitObservation>;
async fn observe_provider_success(&self, provider_kind: &str) -> StoreResult<()>;
}
#[derive(Clone)]
pub struct RedisProviderCircuitRepository {
connection: ConnectionManager,
namespace: String,
policy: ProviderCircuitPolicy,
}
impl RedisProviderCircuitRepository {
pub fn new(
connection: ConnectionManager,
key_namespace: &str,
policy: ProviderCircuitPolicy,
) -> StoreResult<Self> {
if policy.open_duration.is_zero() {
return Err(invalid("open duration must be positive"));
}
Ok(Self {
connection,
namespace: namespace(key_namespace)?,
policy,
})
}
fn key(&self, provider_kind: &str) -> StoreResult<String> {
let fingerprint = resource_fingerprint("provider circuit", provider_kind)?;
Ok(format!("{}:provider:{fingerprint}:circuit", self.namespace))
}
}
#[async_trait]
impl ProviderCircuitRepository for RedisProviderCircuitRepository {
async fn provider_circuit_decision(
&self,
provider_kind: &str,
) -> StoreResult<ProviderCircuitDecision> {
require_nonempty("provider circuit", "provider_kind", provider_kind)?;
let mut connection = self.connection.clone();
let (allow, until): (i64, String) = Script::new(DECISION_SCRIPT)
.key(self.key(provider_kind)?)
.invoke_async(&mut connection)
.await
.map_err(|_| redis_unavailable("read provider circuit"))?;
if allow == 1 {
Ok(ProviderCircuitDecision::Allow)
} else {
Ok(ProviderCircuitDecision::BlockedUntil(timestamp(&until)?))
}
}
async fn observe_provider_failure(
&self,
provider_kind: &str,
) -> StoreResult<ProviderCircuitObservation> {
require_nonempty("provider circuit", "provider_kind", provider_kind)?;
let duration_ms = u64::try_from(self.policy.open_duration.as_millis())
.map_err(|_| invalid("open duration is too large"))?;
let mut connection = self.connection.clone();
let (failures, until): (String, String) = Script::new(FAILURE_SCRIPT)
.key(self.key(provider_kind)?)
.arg(self.policy.failure_threshold.get())
.arg(duration_ms)
.invoke_async(&mut connection)
.await
.map_err(|_| redis_unavailable("observe provider failure"))?;
let failure_count = failures
.parse()
.map_err(|_| invalid("Redis returned an invalid failure count"))?;
let open_until = if until == "0" {
None
} else {
Some(timestamp(&until)?)
};
Ok(ProviderCircuitObservation {
failure_count,
open_until,
})
}
async fn observe_provider_success(&self, provider_kind: &str) -> StoreResult<()> {
require_nonempty("provider circuit", "provider_kind", provider_kind)?;
let mut connection = self.connection.clone();
redis::cmd("DEL")
.arg(self.key(provider_kind)?)
.query_async::<i64>(&mut connection)
.await
.map_err(|_| redis_unavailable("reset provider circuit"))?;
Ok(())
}
}
impl ProviderCircuitPort for RedisProviderCircuitRepository {
fn decision<'a>(
&'a self,
provider_kind: &'a ProviderKind,
) -> futures::future::BoxFuture<'a, Result<CoreCircuitDecision, ProviderCircuitError>> {
Box::pin(async move {
self.provider_circuit_decision(provider_kind.as_str())
.await
.map(|decision| match decision {
ProviderCircuitDecision::Allow => CoreCircuitDecision::Allow,
ProviderCircuitDecision::BlockedUntil(until) => {
CoreCircuitDecision::BlockedUntil(until.into())
}
})
.map_err(|_| ProviderCircuitError)
})
}
fn observe_failure<'a>(
&'a self,
provider_kind: &'a ProviderKind,
) -> futures::future::BoxFuture<'a, Result<(), ProviderCircuitError>> {
Box::pin(async move {
self.observe_provider_failure(provider_kind.as_str())
.await
.map(|_| ())
.map_err(|_| ProviderCircuitError)
})
}
fn observe_success<'a>(
&'a self,
provider_kind: &'a ProviderKind,
) -> futures::future::BoxFuture<'a, Result<(), ProviderCircuitError>> {
Box::pin(async move {
self.observe_provider_success(provider_kind.as_str())
.await
.map_err(|_| ProviderCircuitError)
})
}
}
fn timestamp(value: &str) -> StoreResult<DateTime<Utc>> {
let value = value
.parse::<i64>()
.map_err(|_| invalid("Redis returned an invalid circuit timestamp"))?;
DateTime::from_timestamp_millis(value)
.ok_or_else(|| invalid("Redis returned an out-of-range circuit timestamp"))
}
fn invalid(message: &str) -> StoreError {
StoreError::InvalidData {
entity: "provider circuit",
message: message.to_owned(),
}
}
+1 -17
View File
@@ -9,10 +9,9 @@ pub(crate) struct CommandStoreWriters {
pub(crate) execution: postgres::ExecutionObservationWriter<postgres::PgExecutionStore>,
pub(crate) client_key_usage: postgres::PgClientApiKeyUsageWriter,
pub(crate) admission_release: redis::ClientAdmissionReleaseWriter,
pub(crate) circuit_feedback: redis::ProviderCircuitFeedbackWriter,
}
/// 短生命周期 CLI 只运行数据面必需的四个写泵,不注册恢复、保留或维护 Worker。
/// 短生命周期 CLI 只运行数据面必需的三个写泵,不注册恢复、保留或维护 Worker。
pub struct CommandStoreDrain {
cancellation: gateway_core::lifecycle::CancellationToken,
tasks: Vec<tokio::task::JoinHandle<Result<(), WorkerTaskError>>>,
@@ -30,14 +29,12 @@ impl CommandStoreWriters {
execution,
client_key_usage,
admission_release,
circuit_feedback,
} = self;
let execution_idle = execution.idle();
let tasks = vec![
spawn_command_writer(execution, cancellation.child_token()),
spawn_command_writer(client_key_usage, cancellation.child_token()),
spawn_command_writer(admission_release, cancellation.child_token()),
spawn_command_writer(circuit_feedback, cancellation.child_token()),
];
CommandStoreDrain {
cancellation,
@@ -102,7 +99,6 @@ pub(crate) fn store_worker_contributions(
execution_writer: postgres::ExecutionObservationWriter<postgres::PgExecutionStore>,
client_key_usage_writer: postgres::PgClientApiKeyUsageWriter,
admission_release_writer: redis::ClientAdmissionReleaseWriter,
circuit_feedback_writer: redis::ProviderCircuitFeedbackWriter,
retention: Arc<postgres::PgRetentionRepository>,
) -> StoreResult<Vec<WorkerContribution>> {
let stale_id = WorkerId::try_new(WorkerKind::StaleModelRequestRecovery, "postgres")
@@ -116,8 +112,6 @@ pub(crate) fn store_worker_contributions(
.map_err(worker_definition_error)?;
let admission_flush_id = WorkerId::try_new(WorkerKind::OpsFlush, "redis_admission")
.map_err(worker_definition_error)?;
let circuit_flush_id = WorkerId::try_new(WorkerKind::OpsFlush, "redis_circuit")
.map_err(worker_definition_error)?;
let ops_flush_restart =
DaemonRestartPolicy::try_new(Duration::from_secs(1), Duration::from_secs(60))
.map_err(worker_definition_error)?;
@@ -162,16 +156,6 @@ pub(crate) fn store_worker_contributions(
)
.map_err(worker_definition_error)?,
),
WorkerContribution::Registration(
WorkerRegistration::try_new(
circuit_flush_id,
WorkerRunnable::Daemon {
restart: ops_flush_restart,
task: Box::new(circuit_feedback_writer),
},
)
.map_err(worker_definition_error)?,
),
])
}
@@ -1995,3 +1995,36 @@ async fn zero_attempt_failure_does_not_enter_successful_usage_or_cost_aggregates
assert_eq!(charges, 0);
database.close().await;
}
#[tokio::test]
async fn entry_rejection_is_visible_without_a_fictitious_model_execution() {
let Some(database) = TestDatabase::create("entry_rejection").await else {
return;
};
let repository = PgExecutionStore::new(database.pool.clone());
repository
.record_entry_rejection(gateway_core::engine::EntryRejection {
request_id: ModelRequestId::new("req_entry_rejection").unwrap(),
client_key_id: ClientApiKeyId::new("key_entry").unwrap(),
error: GatewayError::new(GatewayErrorKind::NoAvailableProvider, "no route"),
latency: StdDuration::from_millis(12),
})
.await
.unwrap();
let row: (Option<String>, String, String, i64) = sqlx::query_as(
"select model_request_id, failure_kind, message, latency_ms from ops_events where component = 'request_entry' and operation = 'reject'"
).fetch_one(&database.pool).await.unwrap();
assert!(row.0.is_none());
assert_eq!(row.1, GatewayErrorKind::NoAvailableProvider.as_str());
assert_eq!(
serde_json::from_str::<Value>(&row.2).unwrap()["requestId"],
"req_entry_rejection"
);
assert_eq!(row.3, 12);
let count: i64 = sqlx::query_scalar("select count(*) from model_requests")
.fetch_one(&database.pool)
.await
.unwrap();
assert_eq!(count, 0);
database.close().await;
}
@@ -9,13 +9,9 @@ use gateway_core::engine::admission::{
ClientAdmissionDecision, ClientAdmissionError, ClientAdmissionPort, ClientAdmissionRecovery,
ClientAdmissionRejection, ClientAdmissionRequest, ClientAdmissionRestoreResult,
};
use gateway_core::engine::execution::{
ProviderCircuitDecision, ProviderCircuitError, ProviderCircuitPort,
};
use gateway_core::lifecycle::CancellationToken;
use gateway_core::policy::{ClientApiKeyId, RateLimits};
use gateway_core::routing::ProviderKind;
use gateway_store::redis::{BufferedClientAdmissionPort, BufferedProviderCircuitPort};
use gateway_store::redis::BufferedClientAdmissionPort;
#[derive(Default)]
struct RecordingCoordination {
@@ -73,35 +69,6 @@ impl ClientAdmissionPort for RecordingCoordination {
}
}
impl ProviderCircuitPort for RecordingCoordination {
fn decision<'a>(
&'a self,
_: &'a ProviderKind,
) -> BoxFuture<'a, Result<ProviderCircuitDecision, ProviderCircuitError>> {
Box::pin(async { Ok(ProviderCircuitDecision::Allow) })
}
fn observe_failure<'a>(
&'a self,
_: &'a ProviderKind,
) -> BoxFuture<'a, Result<(), ProviderCircuitError>> {
Box::pin(async move {
self.record("circuit_failure");
Ok(())
})
}
fn observe_success<'a>(
&'a self,
_: &'a ProviderKind,
) -> BoxFuture<'a, Result<(), ProviderCircuitError>> {
Box::pin(async move {
self.record("circuit_success");
Ok(())
})
}
}
struct CapacityCoordination {
active: Mutex<BTreeSet<String>>,
maximum: usize,
@@ -173,13 +140,8 @@ async fn full_recoverable_coordination_queues_should_drop_writes_without_waiting
inner.clone(),
NonZeroUsize::new(1).expect("capacity"),
);
let (circuits, _circuit_writer) = BufferedProviderCircuitPort::with_capacity(
inner.clone(),
NonZeroUsize::new(1).expect("capacity"),
);
let client = ClientApiKeyId::new("key_buffer_test").expect("client key");
let request = ModelRequestId::new("req_buffer_test").expect("request ID");
let provider = ProviderKind::new("openai").expect("provider");
tokio::time::timeout(Duration::from_millis(50), async {
admissions
@@ -190,14 +152,6 @@ async fn full_recoverable_coordination_queues_should_drop_writes_without_waiting
.release(&client, &request)
.await
.expect("full admission queue remains fail-open");
circuits
.observe_success(&provider)
.await
.expect("first circuit enqueue");
circuits
.observe_failure(&provider)
.await
.expect("full circuit queue remains fail-open");
})
.await
.expect("coordination enqueue must never wait for Redis");
@@ -212,29 +166,20 @@ async fn redis_coordination_writers_should_flush_each_side_effect() {
inner.clone(),
NonZeroUsize::new(8).expect("capacity"),
);
let (circuits, circuit_writer) = BufferedProviderCircuitPort::with_capacity(
inner.clone(),
NonZeroUsize::new(8).expect("capacity"),
);
let client = ClientApiKeyId::new("key_writer_test").expect("client key");
let request = ModelRequestId::new("req_writer_test").expect("request ID");
let provider = ProviderKind::new("openai").expect("provider");
admissions
.release(&client, &request)
.await
.expect("enqueue admission release");
circuits
.observe_success(&provider)
.await
.expect("enqueue circuit feedback");
let cancellation = CancellationToken::new();
let tasks = [
spawn_writer(Arc::new(admission_writer), cancellation.clone()),
spawn_writer(Arc::new(circuit_writer), cancellation.clone()),
];
let tasks = [spawn_writer(
Arc::new(admission_writer),
cancellation.clone(),
)];
tokio::time::timeout(Duration::from_secs(1), async {
loop {
if inner.operations().len() == 2 {
if inner.operations().len() == 1 {
break;
}
tokio::task::yield_now().await;
@@ -251,7 +196,7 @@ async fn redis_coordination_writers_should_flush_each_side_effect() {
let mut operations = inner.operations();
operations.sort_unstable();
assert_eq!(operations, ["admission", "circuit_success"]);
assert_eq!(operations, ["admission"]);
}
#[tokio::test]
@@ -7,7 +7,6 @@ mod credential_leases;
mod credential_state;
mod native_continuation;
mod oauth_pending;
mod provider_circuit;
mod provider_session_affinity;
mod provider_session_exclusion;
mod worker_lease;
@@ -1,191 +0,0 @@
use std::{num::NonZeroU32, time::Duration};
use gateway_core::{
engine::execution::{ProviderCircuitPolicy, provider_failure_affects_circuit},
error::ProviderErrorKind,
};
use gateway_store::redis::{
ProviderCircuitDecision, ProviderCircuitRepository, RedisProviderCircuitRepository,
};
use redis::aio::ConnectionManager;
use uuid::Uuid;
#[test]
fn provider_circuit_default_has_positive_threshold() {
assert!(ProviderCircuitPolicy::default().failure_threshold.get() > 0);
}
#[test]
fn provider_circuit_should_only_count_instance_attributable_failures() {
for error_kind in [
ProviderErrorKind::Timeout,
ProviderErrorKind::Transport,
ProviderErrorKind::Protocol,
ProviderErrorKind::Unavailable,
] {
assert!(provider_failure_affects_circuit(error_kind));
}
for error_kind in [
ProviderErrorKind::InvalidRequest,
ProviderErrorKind::Unsupported,
ProviderErrorKind::Unauthorized,
ProviderErrorKind::PermissionDenied,
ProviderErrorKind::RateLimited,
ProviderErrorKind::QuotaExhausted,
ProviderErrorKind::Cancelled,
ProviderErrorKind::ProcessTerminated,
] {
assert!(!provider_failure_affects_circuit(error_kind));
}
}
#[tokio::test]
async fn provider_circuit_opens_at_threshold_and_success_resets_it() {
let Some((repository, mut connection, namespace)) = repository(2).await else {
return;
};
let instance = "instance-circuit-primary";
assert_eq!(
repository
.provider_circuit_decision(instance)
.await
.expect("read empty provider circuit"),
ProviderCircuitDecision::Allow
);
let first = repository
.observe_provider_failure(instance)
.await
.expect("record first provider failure");
assert_eq!(first.failure_count, 1);
assert!(first.open_until.is_none());
let second = repository
.observe_provider_failure(instance)
.await
.expect("record threshold provider failure");
assert_eq!(second.failure_count, 2);
assert!(second.open_until.is_some());
assert!(matches!(
repository
.provider_circuit_decision(instance)
.await
.expect("read open provider circuit"),
ProviderCircuitDecision::BlockedUntil(_)
));
repository
.observe_provider_success(instance)
.await
.expect("reset provider circuit after success");
assert_eq!(
repository
.provider_circuit_decision(instance)
.await
.expect("read reset provider circuit"),
ProviderCircuitDecision::Allow
);
delete_namespace_keys(&mut connection, &namespace).await;
}
#[tokio::test]
async fn provider_circuit_keeps_instances_isolated() {
let Some((repository, mut connection, namespace)) = repository(1).await else {
return;
};
repository
.observe_provider_failure("instance-circuit-failed")
.await
.expect("open failed instance circuit");
assert!(matches!(
repository
.provider_circuit_decision("instance-circuit-failed")
.await
.expect("read failed instance circuit"),
ProviderCircuitDecision::BlockedUntil(_)
));
assert_eq!(
repository
.provider_circuit_decision("instance-circuit-healthy")
.await
.expect("read isolated healthy circuit"),
ProviderCircuitDecision::Allow
);
delete_namespace_keys(&mut connection, &namespace).await;
}
#[tokio::test]
async fn provider_circuit_reopens_after_redis_time_deadline() {
let Some((repository, mut connection, namespace)) =
repository_with_policy(1, Duration::from_millis(40)).await
else {
return;
};
let instance = "instance-circuit-deadline";
repository
.observe_provider_failure(instance)
.await
.expect("open short provider circuit");
assert!(matches!(
repository
.provider_circuit_decision(instance)
.await
.expect("read short open circuit"),
ProviderCircuitDecision::BlockedUntil(_)
));
tokio::time::sleep(Duration::from_millis(80)).await;
assert_eq!(
repository
.provider_circuit_decision(instance)
.await
.expect("read circuit after Redis deadline"),
ProviderCircuitDecision::Allow
);
delete_namespace_keys(&mut connection, &namespace).await;
}
async fn repository(
threshold: u32,
) -> Option<(RedisProviderCircuitRepository, ConnectionManager, String)> {
repository_with_policy(threshold, Duration::from_secs(30)).await
}
async fn repository_with_policy(
threshold: u32,
open_duration: Duration,
) -> Option<(RedisProviderCircuitRepository, ConnectionManager, String)> {
let redis_url = crate::support::test_env("CPR_TEST_REDIS_URL")?;
let client = redis::Client::open(redis_url).expect("valid CPR_TEST_REDIS_URL");
let connection = client
.get_connection_manager()
.await
.expect("connect test Redis");
let namespace = format!("gateway-store-circuit-test-{}", Uuid::new_v4());
let policy = ProviderCircuitPolicy {
failure_threshold: NonZeroU32::new(threshold).expect("positive threshold"),
open_duration,
};
let repository = RedisProviderCircuitRepository::new(connection.clone(), &namespace, policy)
.expect("valid provider circuit repository");
Some((repository, connection, namespace))
}
async fn delete_namespace_keys(connection: &mut ConnectionManager, namespace: &str) {
let keys = redis::cmd("KEYS")
.arg(format!("{namespace}:*"))
.query_async::<Vec<String>>(connection)
.await
.expect("list isolated provider circuit keys");
if !keys.is_empty() {
redis::cmd("DEL")
.arg(keys)
.query_async::<i64>(connection)
.await
.expect("delete isolated provider circuit keys");
}
}
@@ -52,6 +52,7 @@ tokio = { version = "=1.52.3", default-features = false, features = ["fs", "macr
tokio-rustls = { version = "=0.26.4", default-features = false }
tokio-util = { version = "0.7.18", features = ["rt"] }
tower-service = "0.3.3"
tower-layer = "0.3.3"
# 官方 Codex WebSocket 需要上游尚未发布的 proxy 与 permessage-deflate 实现,固定到已审计提交。
tokio-tungstenite = { git = "https://github.com/openai-oss-forks/tokio-tungstenite", rev = "0e5b2d73aa18dd9f0a50ee9ff199d5aef7594186", features = ["proxy", "rustls-tls-native-roots"] }
tracing = "0.1.44"
@@ -760,6 +760,7 @@ pub(super) fn cold_response_stream(response: ColdResponse) -> EventStream {
return;
}
};
context.connection_budget().complete();
if !accepts_backend_transport(transport_policy, response.transport) {
let failure = MappedProviderFailure::plain(provider_error(
ProviderErrorKind::Protocol,
@@ -810,7 +810,15 @@ pub(super) fn map_client_error(
uncertain_state: UpstreamSendState,
observe_transport: bool,
) -> MappedProviderFailure {
let diagnostic = client_diagnostic(&error);
let local_connection_capacity = crate::transport::connection::is_admission_failure(&error);
let diagnostic = if local_connection_capacity {
Some(
ProviderDiagnostic::new("OpenAI local connection capacity unavailable")
.with_classification("admission", "local_connection_capacity"),
)
} else {
client_diagnostic(&error)
};
let raw_upstream_error = match &error {
CodexClientError::WebSocket(error) => websocket_raw_error(error),
_ => None,
@@ -850,7 +858,13 @@ pub(super) fn map_client_error(
if let Some(failure) = error.upstream_failure() {
return map_upstream_failure(failure, observation, ReplayBoundary::BeforeSemanticOutput);
}
let connect_retry =
matches!(&error, CodexClientError::Http(error) if transient_http_connect(error));
let mut failure = match error {
CodexClientError::ConnectionBudgetExhausted => MappedProviderFailure::plain(
provider_error(ProviderErrorKind::Timeout, UpstreamSendState::NotSent)
.with_connection_retry(gateway_core::engine::AttemptTransport::Fallback),
),
CodexClientError::Upstream { .. } => MappedProviderFailure::plain(provider_error(
ProviderErrorKind::Protocol,
UpstreamSendState::Sent,
@@ -963,6 +977,17 @@ pub(super) fn map_client_error(
failure
}
};
if local_connection_capacity {
failure.error = provider_error(
ProviderErrorKind::ProviderInfrastructureUnavailable,
UpstreamSendState::NotSent,
);
}
if connect_retry {
failure.error = failure
.error
.with_connection_retry(gateway_core::engine::AttemptTransport::Fallback);
}
if let Some(continuation_failure) = continuation_failure {
failure.error = failure
.error
@@ -989,6 +1014,11 @@ pub(super) fn map_client_error(
/// 在消费 transport 错误前提取可持久化事实,不能使用可能携带 URL/凭据的 Display。
fn client_diagnostic(error: &CodexClientError) -> Option<ProviderDiagnostic> {
let (stage, code, message) = match error {
CodexClientError::ConnectionBudgetExhausted => (
"connect",
"connection_budget_exhausted",
"OpenAI connection recovery budget exhausted".to_owned(),
),
CodexClientError::WebSocket(error) => return Some(websocket_diagnostic(error)),
CodexClientError::ErrorBodyRead { status, source, .. } => {
return Some(
@@ -1561,3 +1591,28 @@ pub(super) fn remaining(deadline: SystemTime) -> Option<Duration> {
.ok()
.filter(|remaining| !remaining.is_zero())
}
/// 只恢复明确未发送的 HTTP 建连失败;已知 TLS/配置错误不反复尝试同一路径。
fn transient_http_connect(error: &reqwest::Error) -> bool {
if !error.is_connect() || error.is_builder() {
return false;
}
let mut source = std::error::Error::source(error);
while let Some(cause) = source {
if cause.is::<native_tls::Error>() || cause.is::<rustls::Error>() {
return false;
}
if let Some(io) = cause.downcast_ref::<std::io::Error>()
&& matches!(
io.kind(),
std::io::ErrorKind::PermissionDenied
| std::io::ErrorKind::InvalidInput
| std::io::ErrorKind::InvalidData
)
{
return false;
}
source = cause.source();
}
true
}
@@ -482,23 +482,34 @@ impl Provider for CodexProvider {
)
});
let selection_started_at = Instant::now();
let lease = self
.selector
.select_with_cyber_policy(
&SelectCodexCredential {
upstream_model: upstream_model.as_str(),
request_url: &self.responses_url,
attempt: &context,
session_affinity_key: selection_session_affinity
.as_ref()
.map(|affinity| affinity.key()),
},
selection_cyber_policy_key.as_ref(),
selection_session_affinity.as_ref(),
requires_websocket,
)
.await
.map_err(map_selection_error)?;
let selection = async {
self.selector
.select_with_cyber_policy(
&SelectCodexCredential {
upstream_model: upstream_model.as_str(),
request_url: &self.responses_url,
attempt: &context,
session_affinity_key: selection_session_affinity
.as_ref()
.map(|affinity| affinity.key()),
},
selection_cyber_policy_key.as_ref(),
selection_session_affinity.as_ref(),
requires_websocket,
)
.await
.map_err(map_selection_error)
};
// 恢复期间排队也消耗启动窗口;只包住选账号,不限制业务响应时长。
let lease = if let Some(remaining) = context.connection_budget().startup_remaining() {
tokio::time::timeout(remaining, selection)
.await
.map_err(|_| {
provider_error(ProviderErrorKind::Timeout, UpstreamSendState::NotSent)
})??
} else {
selection.await?
};
let account_selection_wait_ms =
u64::try_from(selection_started_at.elapsed().as_millis()).unwrap_or(u64::MAX);
let provider_kind = ProviderKind::new(PROVIDER_NAME)
@@ -787,6 +798,7 @@ impl CodexProvider {
provider_error(ProviderErrorKind::Unavailable, UpstreamSendState::NotSent)
})?
.with_authentication(lease.authentication())
.with_connection_budget(context.connection_budget().clone())
.with_middleware_headers(middleware_headers),
response_origin: self.responses_url.clone(),
request: upstream_request,
@@ -46,7 +46,7 @@ const UPSTREAM_CONNECT_TIMEOUT: Duration = Duration::from_secs(15);
pub(super) const UPSTREAM_STREAM_IDLE_TIMEOUT: Duration = Duration::from_secs(5 * 60);
const X_CODEX_WS_STREAM_REQUEST_START_MS_CLIENT_METADATA_KEY: &str =
"x-codex-ws-stream-request-start-ms";
type ReqwestClientCacheKey = (Option<String>, String);
type ReqwestClientCacheKey = (Option<String>, String, Duration);
type ReqwestClientCache = Mutex<HashMap<ReqwestClientCacheKey, Client>>;
/// 构建带缓存、自动协商 HTTP/2 的 reqwest Client。
@@ -57,9 +57,21 @@ pub fn build_reqwest_client() -> Result<Client, CustomCaError> {
pub fn build_account_http_client(
account_id: &str,
proxy: Option<&gateway_core::account::OutboundProxy>,
) -> Result<Client, CustomCaError> {
build_account_http_client_with_timeout(account_id, proxy, UPSTREAM_CONNECT_TIMEOUT)
}
pub(super) fn build_account_http_client_with_timeout(
account_id: &str,
proxy: Option<&gateway_core::account::OutboundProxy>,
timeout: Duration,
) -> Result<Client, CustomCaError> {
super::tls::ensure_rustls_provider();
let cache_key = (custom_ca_env_cache_key(), egress_key(account_id, proxy));
let cache_key = (
custom_ca_env_cache_key(),
egress_key(account_id, proxy),
timeout,
);
static CLIENTS: OnceLock<ReqwestClientCache> = OnceLock::new();
let cache = CLIENTS.get_or_init(|| Mutex::new(HashMap::new()));
if let Some(client) = cache
@@ -75,7 +87,8 @@ pub fn build_account_http_client(
.redirect(reqwest::redirect::Policy::none())
.pool_max_idle_per_host(4)
.pool_idle_timeout(None::<Duration>)
.connect_timeout(UPSTREAM_CONNECT_TIMEOUT)
.connect_timeout(timeout)
.connector_layer(super::connection::ConnectionLayer)
.tcp_keepalive(Duration::from_secs(30))
.http2_keep_alive_interval(Duration::from_secs(30))
.http2_keep_alive_timeout(Duration::from_secs(5))
@@ -188,6 +201,8 @@ impl fmt::Debug for CodexClientVisibleUpstreamResponse {
/// Codex 上游 HTTP 客户端错误。
#[derive(Error)]
pub enum CodexClientError {
#[error("connection recovery budget exhausted")]
ConnectionBudgetExhausted,
/// Reqwest 传输失败。
#[error("http transport error: {0}")]
Http(#[from] reqwest::Error),
@@ -269,6 +284,9 @@ pub enum CodexClientError {
impl fmt::Debug for CodexClientError {
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
match self {
Self::ConnectionBudgetExhausted => {
formatter.write_str("CodexClientError::ConnectionBudgetExhausted")
}
Self::Http(_) => formatter.write_str("CodexClientError::Http([REDACTED])"),
Self::HttpJson(_) => formatter.write_str("CodexClientError::HttpJson([REDACTED])"),
Self::ErrorBodyRead {
@@ -338,7 +356,8 @@ impl CodexClientError {
Self::Upstream { transport, .. } | Self::ErrorBodyRead { transport, .. } => {
Some(*transport)
}
Self::CustomCa(_)
Self::ConnectionBudgetExhausted
| Self::CustomCa(_)
| Self::InvalidHeaderName(_)
| Self::InvalidHeaderValue(_)
| Self::MiddlewareHeaderConflict
@@ -668,6 +687,7 @@ impl OpenAiUpstreamProtocol {
/// Codex HTTP/SSE 上游客户端。
#[derive(Clone)]
pub struct CodexBackendClient {
pub(super) connection_budget: Option<gateway_core::engine::connection::ConnectionBudget>,
pub(super) client: Client,
pub(super) direct_client: Client,
pub(super) base_url: String,
@@ -683,6 +703,41 @@ pub struct CodexBackendClient {
}
impl CodexBackendClient {
pub(crate) fn with_connection_budget(
mut self,
budget: gateway_core::engine::connection::ConnectionBudget,
) -> Self {
self.connection_budget = Some(budget);
self
}
pub(super) fn connection_opening(&self) -> Result<Option<Duration>, CodexClientError> {
self.connection_budget.as_ref().map_or(Ok(None), |budget| {
budget
.begin()
.map_err(|_| CodexClientError::ConnectionBudgetExhausted)
})
}
pub(super) fn http_opening_client(&self) -> Result<Client, CodexClientError> {
let Some(remaining) = self.connection_opening()? else {
return Ok(self.client.clone());
};
if remaining >= UPSTREAM_CONNECT_TIMEOUT {
return Ok(self.client.clone());
}
// 按整秒向下取整限制缓存种类;不足一秒不缓存新的 client 配置。
let timeout = Duration::from_secs(remaining.as_secs());
if timeout.is_zero() {
return Err(CodexClientError::ConnectionBudgetExhausted);
}
Ok(build_account_http_client_with_timeout(
&self.egress_key,
self.outbound_proxy.as_ref(),
timeout,
)?)
}
/// 覆盖官方账号接口基址,用于隔离上游联调与协议测试。
#[must_use]
pub fn with_official_base_url(mut self, official_base_url: impl Into<String>) -> Self {
@@ -64,6 +64,7 @@ impl CodexBackendClient {
) -> Self {
let base_url = base_url.into().trim_end_matches('/').to_string();
Self {
connection_budget: None,
direct_client: client.clone(),
client,
websocket_origin_key: websocket_origin_key(&base_url),
@@ -134,7 +135,8 @@ impl CodexBackendClient {
.map(|(name, value)| (name.as_str(), value.as_bytes())),
);
trace.capture("upstream.request.body", &body);
let mut outbound = self.client.post(endpoint).headers(headers);
let client = self.http_opening_client()?;
let mut outbound = client.post(endpoint).headers(headers);
let body = if self.protocol == OpenAiUpstreamProtocol::Codex {
outbound = outbound.header(CONTENT_ENCODING, HeaderValue::from_static("zstd"));
zstd::stream::encode_all(std::io::Cursor::new(body), 3)
@@ -282,6 +284,7 @@ impl CodexBackendClient {
)
.map_err(CodexClientError::WebSocketEncode)?;
websocket_create.connection.outbound_proxy = self.outbound_proxy.clone();
websocket_create.connection.connection_budget = self.connection_budget.clone();
context.trace.cloned().unwrap_or_default().headers(
"upstream.request.headers",
serde_json::json!({"transport": "websocket", "phase": "prepared_opening"}),
@@ -0,0 +1,88 @@
//! HTTP 与 WS 共用的建连资源上限,不保留出口健康或冷却状态。
use std::future::Future;
use std::pin::Pin;
use std::sync::OnceLock;
use std::task::{Context, Poll};
use tokio::sync::{Semaphore, SemaphorePermit};
use tower_layer::Layer;
use tower_service::Service;
// 为单副本的文件描述符和 TLS 握手保留硬上限;不限制已经建立的长流。
const ACTIVE_CONNECTIONS: usize = 128;
const WAITING_CONNECTIONS: usize = 1024;
#[derive(Debug, thiserror::Error)]
#[error("local connection capacity unavailable")]
struct ConnectionAdmissionRejected;
/// 只根据自有类型识别本地容量拒绝,不将任意 IO 文案当作上游故障。
pub(crate) fn is_admission_failure(error: &(dyn std::error::Error + 'static)) -> bool {
let mut source = Some(error);
while let Some(cause) = source {
if cause.is::<ConnectionAdmissionRejected>()
|| cause
.downcast_ref::<std::io::Error>()
.and_then(std::io::Error::get_ref)
.is_some_and(|inner| inner.is::<ConnectionAdmissionRejected>())
{
return true;
}
source = cause.source();
}
false
}
pub(super) async fn acquire() -> Result<SemaphorePermit<'static>, std::io::Error> {
static ACTIVE: OnceLock<Semaphore> = OnceLock::new();
static WAITING: OnceLock<Semaphore> = OnceLock::new();
let active = ACTIVE.get_or_init(|| Semaphore::new(ACTIVE_CONNECTIONS));
let waiting = WAITING.get_or_init(|| Semaphore::new(WAITING_CONNECTIONS));
let queued = waiting.try_acquire().map_err(|_| {
std::io::Error::new(std::io::ErrorKind::WouldBlock, ConnectionAdmissionRejected)
})?;
let permit = active
.acquire()
.await
.map_err(|_| std::io::Error::other(ConnectionAdmissionRejected))?;
drop(queued);
Ok(permit)
}
#[derive(Clone)]
pub(super) struct ConnectionLayer;
impl<S> Layer<S> for ConnectionLayer {
type Service = ConnectionService<S>;
fn layer(&self, inner: S) -> Self::Service {
ConnectionService(inner)
}
}
#[derive(Clone)]
pub(super) struct ConnectionService<S>(S);
type BoxError = Box<dyn std::error::Error + Send + Sync>;
impl<S, R> Service<R> for ConnectionService<S>
where
S: Service<R, Error = BoxError>,
S::Future: Send + 'static,
S::Response: 'static,
{
type Response = S::Response;
type Error = BoxError;
type Future = Pin<Box<dyn Future<Output = Result<Self::Response, Self::Error>> + Send>>;
fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
self.0.poll_ready(cx)
}
fn call(&mut self, request: R) -> Self::Future {
let connection = self.0.call(request);
Box::pin(async move {
let _permit = acquire().await?;
connection.await
})
}
}
@@ -5,6 +5,7 @@ pub mod catalog;
pub mod client;
mod client_json;
mod client_sse;
pub(crate) mod connection;
pub mod diagnostics;
mod downstream;
pub mod endpoints;
@@ -66,6 +66,7 @@ impl CodexWebSocketConnection {
endpoint,
headers,
outbound_proxy: None,
connection_budget: None,
}
}
@@ -167,7 +168,24 @@ async fn connect_websocket(
)))
})?
.map(Connector::Rustls);
let result = timeout(WEBSOCKET_CONNECT_TIMEOUT, async {
let remaining = connection
.connection_budget
.as_ref()
.map(|budget| budget.begin())
.transpose()
.map_err(|error| {
CodexWebSocketExchangeError::Connect(tungstenite::Error::Io(std::io::Error::other(
error,
)))
})?
.flatten();
let connect_timeout = remaining.map_or(WEBSOCKET_CONNECT_TIMEOUT, |remaining| {
remaining.min(WEBSOCKET_CONNECT_TIMEOUT)
});
let result = timeout(connect_timeout, async {
let _permit = crate::transport::connection::acquire()
.await
.map_err(tungstenite::Error::Io)?;
// Preserve the native direct handshake; explicit egress never inherits a global proxy.
if connection.outbound_proxy.is_none()
&& matches!(
@@ -202,7 +220,7 @@ async fn connect_websocket(
})
.await
.map_err(|_| CodexWebSocketExchangeError::ConnectTimeout {
timeout: WEBSOCKET_CONNECT_TIMEOUT,
timeout: connect_timeout,
})?;
match result {
Ok((websocket, response)) => Ok((websocket, response)),
@@ -273,7 +291,6 @@ async fn dial_account(
tokio_tungstenite::proxy::connect_via_proxy(BufWriter::new(stream), &config, &target, port)
.await
.map(BufWriter::into_inner)
.map_err(|_| invalid())
}
async fn connect_tcp(host: &str, port: u16) -> Result<tokio::net::TcpStream, tungstenite::Error> {
@@ -7,13 +7,24 @@ use crate::transport::protocol::responses::{
};
/// WebSocket opening 描述。
#[derive(Debug, Clone, PartialEq, Eq)]
#[derive(Debug, Clone)]
pub struct CodexWebSocketConnection {
pub(crate) connection_budget: Option<gateway_core::engine::connection::ConnectionBudget>,
pub(crate) outbound_proxy: Option<gateway_core::account::OutboundProxy>,
pub(super) endpoint: String,
pub(super) headers: Vec<(String, String)>,
}
// 请求预算不是连接身份;连接池复用不能由请求局部状态决定。
impl PartialEq for CodexWebSocketConnection {
fn eq(&self, other: &Self) -> bool {
self.endpoint == other.endpoint
&& self.headers == other.headers
&& self.outbound_proxy == other.outbound_proxy
}
}
impl Eq for CodexWebSocketConnection {}
/// 已构造完成的 Responses WebSocket 请求描述。
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct CodexWebSocketRequest {
@@ -115,6 +126,7 @@ impl CodexWebSocketConnection {
endpoint: endpoint.into(),
headers,
outbound_proxy: None,
connection_budget: None,
}
}
@@ -11166,3 +11166,197 @@ async fn websocket_new_chain_skips_http_only_oauth_credential_before_lease() {
assert_eq!(store.credential_loads(), 2);
assert_eq!(leases.requests.lock().expect("leases").len(), 1);
}
#[tokio::test]
async fn transient_socks_openings_offer_same_account_http_retry_and_send_one_payload() {
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
let proxy = format!("socks5h://{}", listener.local_addr().unwrap());
let server = tokio::spawn(async move {
for opening in 0..3 {
let (mut socket, _) = listener.accept().await.unwrap();
let mut greeting = [0_u8; 2];
socket.read_exact(&mut greeting).await.unwrap();
assert_eq!(greeting[0], 5);
let mut methods = vec![0; usize::from(greeting[1])];
socket.read_exact(&mut methods).await.unwrap();
if opening < 2 {
continue;
}
socket.write_all(&[5, 0]).await.unwrap();
let mut connect = [0_u8; 4];
socket.read_exact(&mut connect).await.unwrap();
assert_eq!(connect, [5, 1, 0, 3]);
let length = socket.read_u8().await.unwrap();
let mut target = vec![0; usize::from(length) + 2];
socket.read_exact(&mut target).await.unwrap();
socket
.write_all(&[5, 0, 0, 1, 127, 0, 0, 1, 0, 80])
.await
.unwrap();
read_http_request(&mut socket).await;
let body = CAPTURE_COMPLETED_SSE;
socket.write_all(format!("HTTP/1.1 200 OK\r\nContent-Type: text/event-stream\r\nContent-Length: {}\r\nConnection: close\r\n\r\n{body}", body.len()).as_bytes()).await.unwrap();
}
});
let store = Arc::new(MemoryAccountStore::default());
create_account(&store, "acct_provider_contract").await;
store.set_egress(
"acct_provider_contract",
Some(gateway_core::account::OutboundProxy::parse(&proxy).unwrap()),
None,
);
let provider = provider_with_base_url(&store, "http://upstream.invalid".to_owned());
let context = fallback_transport_context("req_socks_recovery");
for opening in 0..3 {
let mut stream = provider
.clone()
.execute(
planned_request("openai", generate_operation()),
context.clone(),
)
.await
.unwrap();
let mut failure = None;
while let Some(event) = stream.next().await {
if let Err(error) = event {
failure = Some(error);
break;
}
}
if opening < 2 {
let failure = failure.expect("SOCKS negotiation must fail before payload");
assert_eq!(failure.send_state(), UpstreamSendState::NotSent);
assert!(matches!(
failure.pre_delivery_retry(),
Some(PreDeliveryRetry::SameAccountConnectionRetry {
transport: AttemptTransport::Fallback
})
));
} else {
assert!(failure.is_none());
}
}
timeout(Duration::from_secs(5), server)
.await
.unwrap()
.unwrap();
}
#[tokio::test]
async fn lost_http_response_after_payload_does_not_request_connection_replay() {
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
let base = format!("http://{}", listener.local_addr().unwrap());
let server = tokio::spawn(async move {
let (mut socket, _) = listener.accept().await.unwrap();
read_http_request(&mut socket).await;
});
let store = Arc::new(MemoryAccountStore::default());
create_account(&store, "acct_provider_contract").await;
let mut stream = provider_with_base_url(&store, base)
.execute(
planned_request("openai", generate_operation()),
fallback_transport_context("req_response_lost"),
)
.await
.unwrap();
let mut failure = None;
while let Some(event) = stream.next().await {
if let Err(error) = event {
failure = Some(error);
break;
}
}
let failure = failure.unwrap();
assert_ne!(failure.send_state(), UpstreamSendState::NotSent);
assert!(!matches!(
failure.pre_delivery_retry(),
Some(PreDeliveryRetry::SameAccountConnectionRetry { .. })
));
server.await.unwrap();
}
#[tokio::test]
async fn stalled_proxy_connections_share_thirty_seconds_instead_of_resetting_timeout() {
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
let proxy = format!("socks5h://{}", listener.local_addr().unwrap());
let server = tokio::spawn(async move {
let mut sockets = Vec::new();
loop {
sockets.push(listener.accept().await.unwrap().0);
}
});
let store = Arc::new(MemoryAccountStore::default());
create_account(&store, "acct_provider_contract").await;
store.set_egress(
"acct_provider_contract",
Some(gateway_core::account::OutboundProxy::parse(&proxy).unwrap()),
None,
);
let provider = provider_with_base_url(&store, "http://upstream.invalid".to_owned());
let context = fallback_transport_context("req_socks_timeout");
let started = std::time::Instant::now();
timeout(Duration::from_secs(32), async {
for _ in 0..4 {
let mut stream = provider
.clone()
.execute(
planned_request("openai", generate_operation()),
context.clone(),
)
.await
.unwrap();
let mut failure = None;
while let Some(event) = stream.next().await {
if let Err(error) = event {
failure = Some(error);
break;
}
}
let failure = failure.unwrap();
assert_eq!(failure.kind(), ProviderErrorKind::Timeout);
assert_eq!(failure.send_state(), UpstreamSendState::NotSent);
}
})
.await
.expect("shared recovery window must not become four 15-second timeouts");
assert!(started.elapsed() >= Duration::from_secs(25));
assert!(context.connection_budget().exhausted());
server.abort();
let _ = server.await;
}
pub(crate) async fn assert_local_connection_capacity_is_not_an_upstream_failure() {
let store = Arc::new(MemoryAccountStore::default());
create_account(&store, "acct_provider_contract").await;
// 显式出口使 Provider 使用生产建连层,而不是测试注入的裸 reqwest client。
store.set_egress(
"acct_provider_contract",
Some(gateway_core::account::OutboundProxy::parse("http://127.0.0.1:9").unwrap()),
None,
);
let mut stream = provider_with_base_url(&store, "http://127.0.0.1:9".to_owned())
.execute(
planned_request("openai", generate_operation()),
fallback_transport_context("req_local_capacity"),
)
.await
.unwrap();
let mut failure = None;
while let Some(event) = stream.next().await {
if let Err(error) = event {
failure = Some(error);
break;
}
}
let failure = failure.unwrap();
assert_eq!(
failure.kind(),
ProviderErrorKind::ProviderInfrastructureUnavailable
);
assert_eq!(failure.send_state(), UpstreamSendState::NotSent);
assert_eq!(failure.diagnostic().unwrap().stage(), Some("admission"));
assert_eq!(
failure.diagnostic().unwrap().code(),
Some("local_connection_capacity")
);
}
@@ -1,2 +1,4 @@
mod contract;
mod failure;
pub(crate) use contract::assert_local_connection_capacity_is_not_an_upstream_failure;
@@ -0,0 +1,86 @@
//! 在独立测试进程中验证共享建连准入,避免其他网络测试占用全局名额。
use super::*;
use gateway_core::account::OutboundProxy;
use provider_openai::transport::client::build_account_http_client;
#[tokio::test]
async fn cold_connection_admission_bounds_queue_and_releases_cancelled_work() {
const CHILD: &str = "CPR_CONNECTION_ADMISSION_TEST_CHILD";
if std::env::var_os(CHILD).is_none() {
let status = Command::new(std::env::current_exe().unwrap())
.args(["--exact", "transport::connection::cold_connection_admission_bounds_queue_and_releases_cancelled_work", "--nocapture"])
.env(CHILD, "1")
.status().unwrap();
assert!(status.success());
return;
}
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
let proxy =
OutboundProxy::parse(&format!("http://{}", listener.local_addr().unwrap())).unwrap();
let client = build_account_http_client("acct_connection_admission", Some(&proxy)).unwrap();
let (completed, mut results) = tokio::sync::mpsc::unbounded_channel();
let mut tasks = Vec::new();
let spawn = |client: reqwest::Client, completed: tokio::sync::mpsc::UnboundedSender<bool>| {
tokio::spawn(async move {
let result = client.get("https://upstream.invalid/").send().await;
let _ = completed.send(result.is_err_and(|error| error.is_connect()));
})
};
for _ in 0..128 {
tasks.push(spawn(client.clone(), completed.clone()));
}
let mut sockets = Vec::new();
for _ in 0..128 {
sockets.push(
timeout(Duration::from_secs(10), listener.accept())
.await
.unwrap()
.unwrap()
.0,
);
}
// 128 条 CONNECT 握手保持挂起;1024 个等待者之外的一次请求应立即拒绝。
for _ in 0..1025 {
tasks.push(spawn(client.clone(), completed.clone()));
}
assert_eq!(
timeout(Duration::from_secs(3), results.recv())
.await
.unwrap(),
Some(true)
);
assert!(
timeout(Duration::from_millis(100), results.recv())
.await
.is_err()
);
assert!(
timeout(Duration::from_millis(100), listener.accept())
.await
.is_err()
);
crate::provider::assert_local_connection_capacity_is_not_an_upstream_failure().await;
tasks[0].abort();
sockets.push(
timeout(Duration::from_secs(3), listener.accept())
.await
.unwrap()
.unwrap()
.0,
);
for task in &tasks {
task.abort();
}
for task in tasks {
let _ = task.await;
}
// 同时取消活动与排队请求后,新请求仍能取得名额。
let fresh = spawn(client, completed);
let _socket = timeout(Duration::from_secs(3), listener.accept())
.await
.unwrap()
.unwrap();
fresh.abort();
let _ = fresh.await;
}
@@ -49,6 +49,7 @@ mod account_proxy;
mod canonical;
mod catalog;
mod client;
mod connection;
mod diagnostics;
mod downstream;
mod endpoints;
+6
View File
@@ -334,6 +334,12 @@ Responses wire 之间的协议转换层,转换只在 xAI Provider 内完成。
上游结构化错误的 message/code/type 按上述边界交付客户端,其中内嵌的账号指纹 UUID 已脱敏。模型映射是
全局精确映射,未命中时模型名原样交给候选 Provider;分组只限定账号集合,不参与模型改名。
未收到上游 HTTP 响应的连接失败,在安全的请求内恢复耗尽后返回 `502`;上游超时返回 `504`。
没有可用路由、账号容量或本地建连资源等暂时不可服务的情况返回 `503`;调用方自身限流继续返回 `429`。
这些状态区分故障原因,不能保证任意客户端都会重试。已提交 SSE/WS 输出后不能更改 HTTP 状态,
只能通过既有流内错误合同结束。仅在存在可信等待时长时提供 `Retry-After`,不推测出口恢复时间。
发送前连接恢复的次数、时间和重放边界见 [核心执行不变量](architecture.md)。
OpenAI 明确返回 `server_is_overloaded`、`slow_down` 或模型容量不足错误时,代理在允许安全重放且
尚未交付输出的前提下,先做最多 3 次同账号指数退避,再通过现有调度换号。默认间隔从 500ms 开始,
上游 `Retry-After` 参与退避计算,单次等待不超过 8 秒;重试同时受请求总尝试次数和截止时间约束。
+13 -1
View File
@@ -271,7 +271,9 @@ Client Key 鉴权完成后,API adapter 从有界请求头识别 Codex Desktop/
核心不变量:
- 一个客户端请求对应一条 `model_requests`;attempt 是请求内事实,不建立第二张权威表。
- 进入执行的客户端请求对应一条 `model_requests`;attempt 是请求内事实,不建立第二张权威表。
路由或准入阶段拒绝写入 `ops_events` 的 `request_entry/reject`,消息携带 request ID 和 Key ID,
不伪造执行行;这些拒绝尚不计入基于 `model_requests` 的请求统计。
- Provider 的一次 `execute` 只选择一个 credential 并返回一个冷流;换号、重试和 fallback 由 Core 决定。
- `not_sent`、`sent`、`ambiguous` 是单调的上游发送边界;结果不明确时不能假定上游未收到请求。
- downstream commit 是不可撤回的交付承诺。commit 后禁止换号、重试和 fallback。
@@ -279,6 +281,16 @@ Client Key 鉴权完成后,API adapter 从有界请求头识别 Codex Desktop/
重试与提交边界。具体错误合同见 [数据面接口](api.md#3-openai-数据面与模型目录)。
- Provider 可将明确容量拒绝标记为有界同账号退避,Core 在既有安全重放边界内执行,按账号维护请求内
预算,耗尽后复用普通换号路径。该退避消耗总路由预算,与 WS 传输恢复、OAuth 刷新及账号额度冷却分开。
- 普通请求或探测的 transport、timeout、protocol 错误不改变 Provider 类型的全局可路由性。
账号凭据/额度事实、可选账号容量冻结与 WS origin+出口 opening breaker 各自保持原有范围。
- HTTP Responses 的可靠 `not_sent` 建连失败由 Core 决定同账号重试,保留 HTTP 传输。
每请求最多 4 次 opening,共享首次建连起 30 秒窗口;HTTP opening 保守计数(包含池复用),
WS 只在物理建连时计数,换传输或账号不重置。退避为 500/1000/2000ms 加减 20% 抖动。
等待重试与恢复选账号计入窗口;该窗口只约束建立连接,不覆盖业务响应头等待和长流生成。
已发送、发送不确定、产生执行副作用、下游已提交或存在续写绑定时不使用这条新增重试路径。
已识别的 TLS/配置错误不申请这类重试;底层库未提供结构化原因的代理失败只做有界恢复。
- OpenAI HTTP/WS 的实际建连共用每进程 128 个活动名额、1024 个等待名额;等待包含在建连超时内,
取消或失败释放名额,已建立长流不占建连名额。此上限是资源保护配置,不代表上游容量承诺。
- 跨 Provider 只在账号范围和能力都允许,且请求尚未到达上游或已被证明可安全重放时发生。
- 可恢复观测写入失败不能替换已经确定的客户端协议结果。