fix(openai): restrict capacity cooldown to explicit overload

Count only confirmed capacity rejections toward account cooldown, excluding
generic server failures and unrecognized upstream errors. Classify streamed
invalid_prompt failures as invalid requests, matching official Codex.

Cover HTTP and WebSocket failures around the cooldown threshold while
preserving retry boundaries, upstream error payloads, and Smart scoring.
Update the capacity cooldown contract in the API and architecture docs.

Validation: Rustfmt, strict workspace Clippy, and all backend suites with
isolated PostgreSQL 18 / Redis 8 (--test-threads=4).

Reference: openai/codex@7498521d28,
codex-rs/codex-api/src/sse/responses.rs.

Refs #177

Co-authored-by: Codex <noreply@openai.com>
This commit is contained in:
Zyy
2026-09-20 14:01:29 +08:00
co-authored by Codex
parent 798d4f518e
commit c0ad2ae656
5 changed files with 121 additions and 34 deletions
@@ -1339,12 +1339,8 @@ pub(super) fn map_upstream_failure(
),
error_message: failure.client_message,
cyber_policy_failure,
// 仅上游明确返回的容量/服务不可用计数;本地连接保护也会映射为
// Unavailable,不能用通用 ProviderErrorKind 推断容量证据。
upstream_capacity_failure: matches!(
category,
CodexFailureCategory::CapacityUnavailable | CodexFailureCategory::Unavailable
),
// 普通 5xx 与未知流内错误不足以证明容量拒绝,不能用于冻结账号。
upstream_capacity_failure: capacity_unavailable,
set_cookie_headers: failure.set_cookie_headers,
rate_limit_headers: failure.rate_limit_headers,
observation,
@@ -453,6 +453,11 @@ fn classify_upstream_failure(
};
}
// 官方 Codex 将流内 invalid_prompt 归为请求错误,不依赖 HTTP 状态码。
if matches!(source, UpstreamFailureSource::SseFailure) && code == "invalid_prompt" {
return CodexFailureCategory::InvalidRequest;
}
if [code.as_str(), message.as_str(), body.as_str()]
.into_iter()
.any(is_model_unsupported)
@@ -6109,14 +6109,16 @@ fn provider_with_capacity_tracking(
}
#[tokio::test]
async fn capacity_feedback_counts_business_rejections_but_excludes_diagnostic_probes() {
async fn capacity_feedback_only_counts_overload_rejections_and_excludes_diagnostic_probes() {
use gateway_core::provider_ports::ProviderCooldownPort as _;
for websocket in [false, true] {
for (status, code) in [
(429, "slow_down"),
(503, "server_is_overloaded"),
(500, "server_error"),
for (status, code, capacity) in [
(429, "slow_down", true),
(503, "server_is_overloaded", true),
(500, "server_error", false),
(502, "server_error", false),
(503, "service_unavailable_error", false),
] {
for diagnostic in [false, true] {
let store = Arc::new(MemoryAccountStore::default());
@@ -6180,8 +6182,11 @@ async fn capacity_feedback_counts_business_rejections_but_excludes_diagnostic_pr
"unexpected upstream error: {error:?}"
);
let after = cooldowns.capacity_evidence(account.id());
if diagnostic {
assert_eq!(after, before, "probe must preserve capacity count and peak");
if diagnostic || !capacity {
assert_eq!(
after, before,
"only explicit overload from ordinary requests may change capacity evidence: {status}/{code}"
);
} else {
assert_eq!(
after.map(|(count, _)| count),
@@ -6280,23 +6285,66 @@ async fn websocket_usage_limit_rejection_preserves_quota_state_for_account_rotat
}
#[tokio::test]
async fn capacity_stream_rejection_only_retries_before_semantic_output_on_both_transports() {
for (use_websocket, code) in [
(false, "server_is_overloaded"),
(true, "server_is_overloaded"),
(false, "slow_down"),
(true, "slow_down"),
] {
async fn capacity_feedback_in_stream_only_counts_explicit_overload() {
use gateway_core::provider_ports::{ProviderCooldownKind, ProviderCooldownPort as _};
for (use_websocket, (code, message, expected_kind, scored)) in
[false, true].into_iter().flat_map(|websocket| {
[
(
"server_is_overloaded",
"Selected model is at capacity. Please try a different model.",
ProviderErrorKind::UpstreamCapacityUnavailable,
true,
),
(
"slow_down",
"slow_down",
ProviderErrorKind::UpstreamCapacityUnavailable,
true,
),
(
"invalid_prompt",
"Invalid prompt: we've limited access to this content for safety reasons.",
ProviderErrorKind::InvalidRequest,
false,
),
(
"server_error",
"An internal server error occurred.",
ProviderErrorKind::Unavailable,
true,
),
(
"unknown_error",
"Unrecognized upstream failure.",
ProviderErrorKind::Unavailable,
false,
),
]
.map(|case| (websocket, case))
})
{
for semantic_output in [false, true] {
let store = Arc::new(MemoryAccountStore::default());
create_account(&store, "acct_provider_contract").await;
let account = store.account("acct_provider_contract").expect("account");
let cooldowns = Arc::new(MemoryCooldownPort::new());
// 距离冻结阈值只差一次,验证非容量错误不会把可调度账号推入冷却。
for _ in 0..11 {
cooldowns
.record_capacity_failure(account.id(), Duration::from_secs(600), 20)
.await
.expect("seed capacity evidence");
}
let before = cooldowns.capacity_evidence(account.id());
let mut events = vec![
json!({"type": "response.created", "response": {"id": "resp_capacity", "model": "gpt-5.4"}}),
];
if semantic_output {
events.push(json!({"type": "response.output_text.delta", "output_index": 0, "content_index": 0, "delta": "hello"}));
}
let original = json!({"type": "response.failed", "response": {"id": "resp_capacity", "error": {"code": code, "message": "Selected model is at capacity. Please try a different model."}}});
let original = json!({"type": "response.failed", "response": {"id": "resp_capacity", "error": {"code": code, "message": message}}});
events.push(original.clone());
let (base_url, _http_server, websocket_server) = if use_websocket {
let listener = TcpListener::bind("127.0.0.1:0").await.expect("listener");
@@ -6339,7 +6387,9 @@ async fn capacity_stream_rejection_only_retries_before_semantic_output_on_both_t
} else {
http_generate_operation()
};
let mut stream = provider_with_base_url(&store, base_url)
let (provider, pool) =
provider_with_capacity_tracking(&store, base_url, Arc::clone(&cooldowns));
let mut stream = provider
.execute(
planned_request("openai", operation),
context("req_capacity_stream", CancellationToken::new()),
@@ -6355,15 +6405,48 @@ async fn capacity_stream_rejection_only_retries_before_semantic_output_on_both_t
}
}
Some(Err(error)) => break error,
None => panic!("expected capacity failure"),
None => panic!("expected upstream failure"),
}
};
assert_eq!(error.replay_is_safe(), !semantic_output);
assert_eq!(error.kind(), ProviderErrorKind::UpstreamCapacityUnavailable);
assert_eq!(error.pre_delivery_retry().is_some(), !semantic_output);
assert!(provider_openai::openai_failure_affects_account_score(
&error
));
let capacity = expected_kind == ProviderErrorKind::UpstreamCapacityUnavailable;
assert_eq!(
error.kind(),
expected_kind,
"unexpected error for {code}, websocket={use_websocket}, semantic_output={semantic_output}: {error:?}"
);
assert_eq!(error.replay_is_safe(), capacity && !semantic_output);
assert_eq!(error.upstream_status(), None);
assert_eq!(
error.pre_delivery_retry().is_some(),
capacity && !semantic_output
);
assert_eq!(
provider_openai::openai_failure_affects_account_score(&error),
scored
);
let cooldown = cooldowns.read(account.id()).await.expect("read cooldown");
if capacity {
assert_eq!(
cooldowns
.capacity_evidence(account.id())
.map(|(count, _)| count),
Some(12)
);
assert_eq!(
cooldown.expect("capacity cooldown").kind(),
ProviderCooldownKind::CapacityFreezeProbe
);
} else {
assert_eq!(
cooldowns.capacity_evidence(account.id()),
before,
"non-capacity error {code} changed capacity evidence"
);
assert!(
cooldown.is_none(),
"non-capacity error {code} froze the account"
);
}
assert_eq!(
serde_json::from_str::<Value>(
error.raw_upstream_error().expect("raw error").as_str()
@@ -6383,6 +6466,7 @@ async fn capacity_stream_rejection_only_retries_before_semantic_output_on_both_t
if let Some(server) = websocket_server {
server.await.expect("server");
}
pool.shutdown().await;
}
}
}
+3 -2
View File
@@ -1153,8 +1153,9 @@ User-Agent 使用 `grok-shell/<版本> (<系统>; <架构>)`,其中 `arm64`
### 账号自动冻结
账号自动冻结(`accountAutoFreezeEnabled`)默认关闭。启用后,在统计窗口内按尝试累计容量类上游错误(`server_is_overloaded`
等与 5xx 不可用),达到阈值后把该账号冻结为带恢复倒计时的 `rate_limited` 状态。`accountAutoFreezeThreshold`
账号自动冻结(`accountAutoFreezeEnabled`)默认关闭。启用后,在统计窗口内按尝试累计明确的上游容量拒绝
(`server_is_overloaded`、`slow_down` 或结构化错误中的明确过载提示),普通 5xx、`invalid_prompt` 和未识别的上游错误不计入。
达到阈值后把该账号冻结为带恢复倒计时的 `rate_limited` 状态。`accountAutoFreezeThreshold`
取值 2~1,000(默认 12,按普通请求的 attempt 计数,含请求内同账号重试,不含诊断探测与本地连接保护错误);`accountAutoFreezeWindowSeconds`
取值 60~3,600(默认 600,随每次失败滑动顺延);`accountAutoFreezeDurationSeconds` 取值 300~604,800
(默认 7,200,即 2 小时,探测失败后按该时长顺延)。`accountAutoFreezeProbeEnabled` 开启时恢复 worker
+4 -3
View File
@@ -555,9 +555,10 @@ Worker 由各 Bundle 贡献、由 Host 统一监督:
以及账号冻结恢复 worker(容量熔断的自适应并发下调与到期探测解冻);
- Provider:credential refresh、quota/catalog 健康和官方版本/etag 检查。
账号容量熔断默认关闭。启用后,仅普通请求收到的容量类上游错误(`server_is_overloaded` 等与
5xx 不可用)按滑动窗口计数,并把当时观测到的在途并发并入峰值证据;本地连接保护与诊断探测
不参与容量计数或峰值采样。这些错误不证明凭据或配额失效,不进入账号失败状态。达到阈值后写入
账号容量熔断默认关闭。启用后,仅普通请求收到的明确容量拒绝(`server_is_overloaded`、`slow_down`
或结构化错误中的明确过载提示)按滑动窗口计数,并把当时观测到的在途并发并入峰值证据;
普通 5xx、未识别的上游错误、本地连接保护与诊断探测不参与容量计数或峰值采样。
这些错误不证明凭据或配额失效,不进入账号失败状态。达到阈值后写入
账号级 Redis 容量冷却,调度立即跳过该账号;管理端沿用限流状态,通过原因区分容量冻结与上游限流。
冻结与计数由 Redis 保存,不改变 PostgreSQL 账号状态。需要探测的冻结在到期后仍阻止调度,恢复
worker 复用连接测试探针执行真实上游调用,并按冻结代次处理结果:成功清除,失败按配置时长顺延,