mirror of
https://github.com/zyycn/codex-proxy-rs.git
synced 2026-11-06 08:06:19 +00:00
fix: contain connection failures and retry unsent requests (#276)
Integrate PR #276 on the current main baseline. Keep local connection
capacity failures out of upstream retry and WebSocket breaker accounting.
Keep existing settings assertions valid when the #192 migration adds columns.
Source-commit: f9828a588f
Co-authored-by: Zyy <zyycn@users.noreply.github.com>
Co-authored-by: Codex <noreply@openai.com>
This commit is contained in:
committed by
Zyy
co-authored by
Zyy
Codex
parent
1e14250b82
commit
24a9f12a78
Generated
+1
@@ -3255,6 +3255,7 @@ dependencies = [
|
||||
"tokio-rustls",
|
||||
"tokio-tungstenite 0.28.0",
|
||||
"tokio-util",
|
||||
"tower-layer",
|
||||
"tower-service",
|
||||
"tracing",
|
||||
"tracing-subscriber",
|
||||
|
||||
@@ -639,18 +639,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;
|
||||
@@ -659,53 +655,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),
|
||||
);
|
||||
@@ -746,7 +705,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,
|
||||
@@ -763,25 +722,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]
|
||||
@@ -813,3 +753,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,
|
||||
@@ -801,7 +805,7 @@ where
|
||||
PollBoundary::Deadline => {
|
||||
// 会话 deadline 是网关自身的请求预算,不是上游超时;
|
||||
// 真正的上游超时会作为流错误进入 `handle_stream_error` 记账。
|
||||
// 这里不写 provider 失败,避免长流集中到期误触 provider 熔断。
|
||||
// 这里不写 provider 失败,避免把本地预算到期归因为上游故障。
|
||||
self.finish_interruption(&EngineError::Deadline).await?;
|
||||
return Err(EngineError::Deadline);
|
||||
}
|
||||
@@ -885,7 +889,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?;
|
||||
@@ -991,6 +1001,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))
|
||||
@@ -1043,9 +1054,7 @@ where
|
||||
return Err(EngineError::Cancelled);
|
||||
}
|
||||
ProviderBoundary::Deadline => {
|
||||
// 网关预算到期同样不是候选 Provider 的上游超时,不计入熔断;
|
||||
// Provider 自身的握手/传输超时会以 `ProviderErrorKind::Timeout`
|
||||
// 错误返回并在下方 `Result` 分支记账。
|
||||
// 网关预算到期是本请求的终态,不推定候选 Provider 不可用。
|
||||
self.finish_interruption(&EngineError::Deadline).await?;
|
||||
return Err(EngineError::Deadline);
|
||||
}
|
||||
@@ -1375,6 +1384,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 {
|
||||
@@ -1406,11 +1419,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 =>
|
||||
@@ -1419,6 +1455,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 =>
|
||||
@@ -1427,7 +1464,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()
|
||||
@@ -1477,6 +1516,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);
|
||||
}
|
||||
@@ -1430,7 +1402,6 @@ impl DefaultExecutionService {
|
||||
core,
|
||||
admission,
|
||||
active_request,
|
||||
Arc::clone(&self.circuits),
|
||||
Arc::clone(&self.continuation),
|
||||
self.budget.clone(),
|
||||
)),
|
||||
@@ -1856,48 +1827,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,
|
||||
@@ -1988,11 +1917,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) => {
|
||||
@@ -2510,9 +2434,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>>,
|
||||
}
|
||||
@@ -2522,7 +2444,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 {
|
||||
@@ -2531,9 +2452,7 @@ impl DefaultExecutionSession {
|
||||
admission: Some(admission),
|
||||
active_request,
|
||||
cleanup: None,
|
||||
circuits,
|
||||
continuation,
|
||||
observed_provider_outcomes: 0,
|
||||
continuation_recorded: false,
|
||||
budget,
|
||||
}
|
||||
@@ -2563,16 +2482,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;
|
||||
@@ -2591,7 +2500,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;
|
||||
}
|
||||
}
|
||||
@@ -2613,7 +2521,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
|
||||
})
|
||||
@@ -2626,7 +2533,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
|
||||
})
|
||||
@@ -2650,7 +2556,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
|
||||
})
|
||||
@@ -2662,7 +2567,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
|
||||
})
|
||||
@@ -2684,39 +2588,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,
|
||||
|
||||
@@ -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(())
|
||||
|
||||
@@ -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 {
|
||||
|
||||
@@ -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),
|
||||
);
|
||||
|
||||
@@ -1046,7 +1046,7 @@ impl RuntimeSnapshot {
|
||||
self.client_policies.get(id)
|
||||
}
|
||||
|
||||
/// 返回当前 Key 范围、本代次注册表与 circuit 条件共同允许的 Provider。
|
||||
/// 返回当前 Key 范围、本代次注册表与请求路由限制共同允许的 Provider。
|
||||
#[must_use]
|
||||
pub fn available_providers(
|
||||
&self,
|
||||
@@ -1167,7 +1167,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());
|
||||
}
|
||||
@@ -4738,3 +4738,199 @@ fn provider_default_profile_is_applied_once_and_frozen_across_account_retries()
|
||||
);
|
||||
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()),
|
||||
)
|
||||
@@ -972,7 +967,6 @@ fn completed_parent_rejects_new_nested_execution_and_cancels_an_active_child() {
|
||||
Arc::new(TrackingExecutionStore::default()),
|
||||
provider_index.clone(),
|
||||
Arc::new(Admissions::default()),
|
||||
Arc::new(UnusedCircuits),
|
||||
Arc::new(UnusedContinuation),
|
||||
Arc::new(RecordingClientApiKeyUsage::default()),
|
||||
)
|
||||
@@ -1164,7 +1158,6 @@ fn native_response_processing_uses_the_real_translation_boundary() {
|
||||
Arc::new(TrackingExecutionStore::default()),
|
||||
provider_index.clone(),
|
||||
Arc::new(Admissions::default()),
|
||||
Arc::new(UnusedCircuits),
|
||||
Arc::new(UnusedContinuation),
|
||||
Arc::new(RecordingClientApiKeyUsage::default()),
|
||||
)
|
||||
@@ -1315,7 +1308,6 @@ fn response_translation_zero_output_preserves_canonical_usage_and_finalization()
|
||||
store.clone(),
|
||||
provider_index.clone(),
|
||||
Arc::new(Admissions::default()),
|
||||
Arc::new(UnusedCircuits),
|
||||
Arc::new(UnusedContinuation),
|
||||
Arc::new(RecordingClientApiKeyUsage::default()),
|
||||
);
|
||||
@@ -1469,7 +1461,6 @@ fn native_response_failure_after_downstream_commit_is_not_replayed() {
|
||||
Arc::new(TrackingExecutionStore::default()),
|
||||
provider_index.clone(),
|
||||
Arc::new(Admissions::default()),
|
||||
Arc::new(UnusedCircuits),
|
||||
Arc::new(UnusedContinuation),
|
||||
Arc::new(RecordingClientApiKeyUsage::default()),
|
||||
);
|
||||
@@ -1542,7 +1533,6 @@ fn model_routing_policy_selects_native_provider() {
|
||||
Arc::new(TrackingExecutionStore::default()),
|
||||
provider_index.clone(),
|
||||
Arc::new(Admissions::default()),
|
||||
Arc::new(UnusedCircuits),
|
||||
Arc::new(UnusedContinuation),
|
||||
Arc::new(RecordingClientApiKeyUsage::default()),
|
||||
)
|
||||
@@ -1596,7 +1586,6 @@ fn model_routing_reject_stops_before_provider_execution() {
|
||||
Arc::new(TrackingExecutionStore::default()),
|
||||
provider_index.clone(),
|
||||
Arc::new(Admissions::default()),
|
||||
Arc::new(UnusedCircuits),
|
||||
Arc::new(UnusedContinuation),
|
||||
Arc::new(RecordingClientApiKeyUsage::default()),
|
||||
)
|
||||
@@ -1681,7 +1670,6 @@ fn model_routing_cannot_expand_the_frozen_key_model_scope() {
|
||||
Arc::new(TrackingExecutionStore::default()),
|
||||
provider_index.clone(),
|
||||
Arc::new(Admissions::default()),
|
||||
Arc::new(UnusedCircuits),
|
||||
Arc::new(UnusedContinuation),
|
||||
Arc::new(RecordingClientApiKeyUsage::default()),
|
||||
)
|
||||
@@ -1828,7 +1816,6 @@ fn scheduler_reject_is_a_terminal_policy_rejection_before_upstream_send() {
|
||||
Arc::new(TrackingExecutionStore::default()),
|
||||
provider_index.clone(),
|
||||
Arc::new(Admissions::default()),
|
||||
Arc::new(UnusedCircuits),
|
||||
Arc::new(UnusedContinuation),
|
||||
Arc::new(RecordingClientApiKeyUsage::default()),
|
||||
)
|
||||
@@ -2103,7 +2090,6 @@ fn observation_service(
|
||||
Arc::new(TrackingExecutionStore::default()),
|
||||
providers,
|
||||
admissions,
|
||||
Arc::new(UnusedCircuits),
|
||||
Arc::new(UnusedContinuation),
|
||||
Arc::new(RecordingClientApiKeyUsage::default()),
|
||||
)
|
||||
@@ -2141,7 +2127,6 @@ fn fallback_exhaustion_observation_uses_the_last_actual_provider_not_the_candida
|
||||
Arc::new(TrackingExecutionStore::default()),
|
||||
providers,
|
||||
Arc::new(Admissions::default()),
|
||||
Arc::new(UnusedCircuits),
|
||||
Arc::new(UnusedContinuation),
|
||||
Arc::new(RecordingClientApiKeyUsage::default()),
|
||||
)
|
||||
@@ -2206,7 +2191,6 @@ fn routing_external_effect_stops_a_not_sent_provider_retry() {
|
||||
Arc::new(TrackingExecutionStore::default()),
|
||||
providers,
|
||||
Arc::new(Admissions::default()),
|
||||
Arc::new(UnusedCircuits),
|
||||
Arc::new(UnusedContinuation),
|
||||
Arc::new(RecordingClientApiKeyUsage::default()),
|
||||
)
|
||||
@@ -2408,7 +2392,6 @@ fn charged_service(
|
||||
}) as Arc<dyn Provider>])
|
||||
.unwrap(),
|
||||
admissions,
|
||||
Arc::new(UnusedCircuits),
|
||||
Arc::new(UnusedContinuation),
|
||||
Arc::new(RecordingClientApiKeyUsage::default()),
|
||||
);
|
||||
@@ -2686,9 +2669,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::nested::NestedModelExecutionRequest;
|
||||
use gateway_core::engine::observation::{
|
||||
@@ -2735,30 +2717,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());
|
||||
@@ -2767,7 +2725,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()),
|
||||
);
|
||||
@@ -2799,7 +2756,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()),
|
||||
);
|
||||
@@ -2845,7 +2801,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()),
|
||||
);
|
||||
@@ -2888,7 +2843,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()),
|
||||
);
|
||||
@@ -2920,7 +2874,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()),
|
||||
);
|
||||
@@ -3015,7 +2968,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()),
|
||||
);
|
||||
@@ -3167,7 +3119,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(),
|
||||
);
|
||||
@@ -3187,7 +3138,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(),
|
||||
);
|
||||
@@ -3209,7 +3159,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(),
|
||||
);
|
||||
@@ -3243,7 +3192,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()),
|
||||
);
|
||||
@@ -3289,84 +3237,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()),
|
||||
);
|
||||
@@ -3400,6 +3278,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]
|
||||
@@ -3409,7 +3291,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()),
|
||||
);
|
||||
@@ -3437,7 +3318,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()),
|
||||
);
|
||||
@@ -3494,6 +3374,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>>,
|
||||
@@ -3556,6 +3437,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);
|
||||
@@ -3673,81 +3562,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 {
|
||||
@@ -3993,7 +3807,6 @@ fn account_wait_inherits_the_budget_spent_during_client_admission() {
|
||||
)
|
||||
.unwrap(),
|
||||
admissions.clone(),
|
||||
Arc::new(UnusedCircuits),
|
||||
Arc::new(UnusedContinuation),
|
||||
Arc::new(RecordingClientApiKeyUsage::default()),
|
||||
);
|
||||
@@ -4101,7 +3914,6 @@ fn queue_service(
|
||||
Arc::new(TrackingExecutionStore::default()),
|
||||
ProviderRegistry::default(),
|
||||
admissions.clone(),
|
||||
Arc::new(UnusedCircuits),
|
||||
Arc::new(UnusedContinuation),
|
||||
Arc::new(RecordingClientApiKeyUsage::default()),
|
||||
);
|
||||
@@ -4226,7 +4038,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()),
|
||||
);
|
||||
@@ -4246,3 +4057,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;
|
||||
|
||||
@@ -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(),
|
||||
}
|
||||
}
|
||||
@@ -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;
|
||||
}
|
||||
|
||||
@@ -176,7 +176,9 @@ async fn plugin_migration_preserves_main_settings_and_accepts_new_installations(
|
||||
.fetch_one(&database.pool)
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(after, before);
|
||||
for (key, value) in before.as_object().unwrap() {
|
||||
assert_eq!(after.get(key), Some(value), "existing setting {key}");
|
||||
}
|
||||
let store = PgPluginStore::new(database.pool.clone());
|
||||
store
|
||||
.install_artifact(
|
||||
|
||||
@@ -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 = !local_connection_capacity
|
||||
&& 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)
|
||||
@@ -790,6 +801,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
|
||||
@@ -564,6 +583,7 @@ pub enum CodexTransportDecision {
|
||||
ExactWebSocket,
|
||||
RequiredWebSocket,
|
||||
Http2WebSocketBudgetExhausted,
|
||||
Http2LocalConnectionCapacity,
|
||||
Http2BreakerOpen,
|
||||
Http2PoolUnavailable,
|
||||
}
|
||||
@@ -577,6 +597,7 @@ impl CodexTransportDecision {
|
||||
Self::ExactWebSocket => "ws_exact_required",
|
||||
Self::RequiredWebSocket => "ws_required",
|
||||
Self::Http2WebSocketBudgetExhausted => "http2_ws_budget_exhausted",
|
||||
Self::Http2LocalConnectionCapacity => "http2_local_connection_capacity",
|
||||
Self::Http2BreakerOpen => "http2_breaker_open",
|
||||
Self::Http2PoolUnavailable => "http2_pool_unavailable",
|
||||
}
|
||||
@@ -668,6 +689,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 +705,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 {
|
||||
@@ -1015,6 +1072,9 @@ pub(super) fn websocket_success_decision(
|
||||
pub(super) fn local_http_fallback_decision(
|
||||
error: &CodexWebSocketExchangeError,
|
||||
) -> Option<CodexTransportDecision> {
|
||||
if crate::transport::connection::is_admission_failure(error) {
|
||||
return Some(CodexTransportDecision::Http2LocalConnectionCapacity);
|
||||
}
|
||||
match error.classified() {
|
||||
CodexWebSocketExchangeError::OriginCircuitOpen
|
||||
| CodexWebSocketExchangeError::OriginHalfOpenBusy => {
|
||||
|
||||
@@ -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,98 @@
|
||||
//! 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;
|
||||
static ACTIVE: OnceLock<Semaphore> = OnceLock::new();
|
||||
static WAITING: OnceLock<Semaphore> = OnceLock::new();
|
||||
|
||||
#[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
|
||||
}
|
||||
|
||||
fn admission_rejected() -> std::io::Error {
|
||||
std::io::Error::new(std::io::ErrorKind::WouldBlock, ConnectionAdmissionRejected)
|
||||
}
|
||||
|
||||
/// 可回退 HTTP 的 WS 快路径不排队,避免本地等待被当作上游 opening 变慢。
|
||||
pub(super) fn try_acquire() -> Result<SemaphorePermit<'static>, std::io::Error> {
|
||||
ACTIVE
|
||||
.get_or_init(|| Semaphore::new(ACTIVE_CONNECTIONS))
|
||||
.try_acquire()
|
||||
.map_err(|_| admission_rejected())
|
||||
}
|
||||
|
||||
pub(super) async fn acquire() -> Result<SemaphorePermit<'static>, std::io::Error> {
|
||||
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(|_| admission_rejected())?;
|
||||
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;
|
||||
|
||||
@@ -268,6 +268,7 @@ async fn prepare_pooled_websocket(
|
||||
pool.clone(),
|
||||
connect_lease,
|
||||
permit,
|
||||
fast_path_budget.is_some(),
|
||||
);
|
||||
match waiter.wait(fast_path_budget).await? {
|
||||
WebSocketFastPath::Ready(handoff) => Ok(WebSocketFastPath::Ready(PreparedWebSocket {
|
||||
@@ -323,6 +324,7 @@ fn start_pooled_websocket_connect(
|
||||
pool: CodexWebSocketPool,
|
||||
connect_lease: WebSocketPoolConnectLease,
|
||||
permit: Option<WebSocketOriginBreakerPermit>,
|
||||
fast_path: bool,
|
||||
) -> PooledWebSocketConnectWaiter {
|
||||
let task_key = connect_lease.key().clone();
|
||||
let started_at = connect_lease.started_at();
|
||||
@@ -352,7 +354,7 @@ fn start_pooled_websocket_connect(
|
||||
);
|
||||
return;
|
||||
}
|
||||
result = connect_websocket_connection(&connection, keepalive, context) => result,
|
||||
result = connect_websocket_connection(&connection, keepalive, context, fast_path) => result,
|
||||
};
|
||||
match finish_breaker_attempt(permit, connected) {
|
||||
Ok((connection, connect_elapsed)) => {
|
||||
@@ -460,7 +462,7 @@ async fn connect_with_budget(
|
||||
match wait_for_fast_path(
|
||||
TokioInstant::now(),
|
||||
fast_path_budget,
|
||||
connect_websocket_connection(connection, keepalive, context),
|
||||
connect_websocket_connection(connection, keepalive, context, fast_path_budget.is_some()),
|
||||
)
|
||||
.await
|
||||
{
|
||||
@@ -473,9 +475,11 @@ async fn connect_websocket_connection(
|
||||
connection: &CodexWebSocketConnection,
|
||||
keepalive: PumpKeepalive,
|
||||
context: PumpLogContext,
|
||||
fast_path: bool,
|
||||
) -> Result<(PooledWebSocketConnection, Duration), CodexWebSocketExchangeError> {
|
||||
let started_at = Instant::now();
|
||||
let (websocket, response) = connect_pumped_websocket(connection, keepalive, context).await?;
|
||||
let (websocket, response) =
|
||||
connect_pumped_websocket(connection, keepalive, context, fast_path).await?;
|
||||
Ok((
|
||||
PooledWebSocketConnection {
|
||||
websocket,
|
||||
@@ -516,6 +520,10 @@ fn finish_breaker_attempt(
|
||||
permit.succeed();
|
||||
Ok(connection)
|
||||
}
|
||||
Err(error) if crate::transport::connection::is_admission_failure(&error) => {
|
||||
permit.cancel();
|
||||
Err(error)
|
||||
}
|
||||
Err(CodexWebSocketExchangeError::Upstream(upstream)) if upstream.status_code < 500 => {
|
||||
// 账号或请求级 opening 响应证明 origin 可达,不得污染 transport 熔断器。
|
||||
permit.succeed();
|
||||
|
||||
@@ -66,6 +66,7 @@ impl CodexWebSocketConnection {
|
||||
endpoint,
|
||||
headers,
|
||||
outbound_proxy: None,
|
||||
connection_budget: None,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -121,8 +122,9 @@ pub(super) async fn connect_pumped_websocket(
|
||||
connection: &CodexWebSocketConnection,
|
||||
keepalive: PumpKeepalive,
|
||||
context: PumpLogContext,
|
||||
fast_path: bool,
|
||||
) -> Result<(PumpedWebSocket, WsResponse<Option<Vec<u8>>>), CodexWebSocketExchangeError> {
|
||||
let (raw, response) = connect_websocket(connection).await?;
|
||||
let (raw, response) = connect_websocket(connection, fast_path).await?;
|
||||
Ok((PumpedWebSocket::new(raw, keepalive, context), response))
|
||||
}
|
||||
|
||||
@@ -158,6 +160,7 @@ pub(super) fn websocket_connection_metadata(
|
||||
|
||||
async fn connect_websocket(
|
||||
connection: &CodexWebSocketConnection,
|
||||
fast_path: bool,
|
||||
) -> Result<(RawWsStream, WsResponse<Option<Vec<u8>>>), CodexWebSocketExchangeError> {
|
||||
let request = websocket_handshake_request(connection)?;
|
||||
let connector = tls::maybe_build_rustls_client_config_with_custom_ca()
|
||||
@@ -167,7 +170,27 @@ 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 = if fast_path {
|
||||
crate::transport::connection::try_acquire()
|
||||
} else {
|
||||
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 +225,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 +296,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,
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -11145,3 +11145,198 @@ async fn invalid_oauth_candidate_does_not_block_another_websocket_account() {
|
||||
"acct_provider_contract"
|
||||
);
|
||||
}
|
||||
|
||||
#[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")
|
||||
);
|
||||
assert!(failure.pre_delivery_retry().is_none());
|
||||
}
|
||||
|
||||
@@ -1,2 +1,4 @@
|
||||
mod contract;
|
||||
mod failure;
|
||||
|
||||
pub(crate) use contract::assert_local_connection_capacity_is_not_an_upstream_failure;
|
||||
|
||||
@@ -0,0 +1,152 @@
|
||||
//! 在独立测试进程中验证共享建连准入,避免其他网络测试占用全局名额。
|
||||
|
||||
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;
|
||||
|
||||
let upstream = TcpListener::bind("127.0.0.1:0").await.unwrap();
|
||||
let base_url = format!("http://{}", upstream.local_addr().unwrap());
|
||||
let server = tokio::spawn(async move {
|
||||
for _ in 0..3 {
|
||||
let (mut http, _) = upstream.accept().await.unwrap();
|
||||
assert!(
|
||||
read_http_request(&mut http)
|
||||
.await
|
||||
.starts_with("POST /codex/responses")
|
||||
);
|
||||
write_completed_sse_response(&mut http).await;
|
||||
}
|
||||
let (stream, _) = upstream.accept().await.unwrap();
|
||||
let mut websocket = accept_codex_test_websocket(stream).await;
|
||||
websocket.next().await.unwrap().unwrap();
|
||||
websocket
|
||||
.send(tokio_tungstenite::tungstenite::Message::Text(
|
||||
completed_websocket_response("resp_after_local_capacity", 2, 1).into(),
|
||||
))
|
||||
.await
|
||||
.unwrap();
|
||||
});
|
||||
let backend = CodexBackendClient::new(
|
||||
reqwest::Client::builder().no_proxy().build().unwrap(),
|
||||
base_url,
|
||||
test_wire_profile(),
|
||||
);
|
||||
for attempt in 0..3 {
|
||||
let mut request = codex_request("gpt-5.5", "be brief", Vec::new());
|
||||
request.use_websocket = true;
|
||||
request.local_conversation_id = Some(format!("capacity-{attempt}"));
|
||||
let response = backend
|
||||
.create_response(
|
||||
&request,
|
||||
request_context("req_local_ws_capacity", Some("chatgpt-account")),
|
||||
)
|
||||
.await
|
||||
.expect("local WS capacity should fall back to HTTP");
|
||||
assert_eq!(response.transport, CodexBackendTransport::HttpSse);
|
||||
assert_eq!(
|
||||
response.transport_metrics.decision,
|
||||
Some(CodexTransportDecision::Http2LocalConnectionCapacity)
|
||||
);
|
||||
}
|
||||
|
||||
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;
|
||||
drop(sockets);
|
||||
|
||||
let mut request = codex_request("gpt-5.5", "be brief", Vec::new());
|
||||
request.use_websocket = true;
|
||||
request.local_conversation_id = Some("capacity-recovered".to_owned());
|
||||
let response = timeout(
|
||||
Duration::from_secs(3),
|
||||
backend.create_response(
|
||||
&request,
|
||||
request_context("req_ws_after_capacity", Some("chatgpt-account")),
|
||||
),
|
||||
)
|
||||
.await
|
||||
.unwrap()
|
||||
.expect("local capacity must not open the WS origin breaker");
|
||||
assert_eq!(response.transport, CodexBackendTransport::WebSocket);
|
||||
timeout(Duration::from_secs(3), server)
|
||||
.await
|
||||
.unwrap()
|
||||
.unwrap();
|
||||
}
|
||||
@@ -49,6 +49,7 @@ mod account_proxy;
|
||||
mod canonical;
|
||||
mod catalog;
|
||||
mod client;
|
||||
mod connection;
|
||||
mod diagnostics;
|
||||
mod downstream;
|
||||
mod endpoints;
|
||||
|
||||
Reference in New Issue
Block a user