fix(ws): settle the previous attempt before transparent retry replanning

评审第 2 条。配额透明重试原来的顺序是「detach 旧 attempt → 规划并绑定新
attempt → 把旧 attempt 的结算排进队列」。规划因此读到的是旧 attempt 还没投射的
health / adaptive / pool 状态,而且旧 attempt 仍占着自己的 pool key lease——替代
key 的挑选看到的是一把仍被占用的 key,最坏情况下判成「无可用供应商」而放弃一次
本可以成功的重试。

普通的新 turn 早就挡住了这件事:client.rs 在处理 response.create 前调用
await_pending_turn_finalization,注释写的正是「不要让新 turn 基于陈旧的 health /
adaptive / pool 状态规划」。透明重试是同一个问题的另一条入口,漏了这一步。

现在顺序是:detach → 释放准入 → 结算旧 attempt 并等它落地 → 规划/绑定新 attempt。

新增 lifecycle::settle_turn_finalization:与 queue_turn_finalization 的区别只在于
「等」。后者把 handle 挂在连接上让 relay loop 继续跑,用在结算之后不再读取共享
状态的出口;前者用在必须先看到结算结果才能继续的路径上。

顺序用类型固定,而不是靠注释:settle_turn_finalization 返回
PreviousAttemptSettled,retry_active_turn_after_quota_exhaustion 要求这个参数。
凭证只能由 lifecycle 颁发(结算完成,或明确「没有 attempt 要结算」),所以把顺序
写反连编译都过不了。

重试失败路径随之变化:旧 attempt 已经结算,不再 resume 回去。logical turn 仍停在
Replanning,后续分支的 end() / finalize_active_turn 只清 logical turn、不交出
attempt,因此不存在重复结算。结算 outcome 取值不变(两条路径用的都是
terminal_outcome.unwrap_or_else(upstream_closed),而这条分支里 terminal_outcome
必为 Some——usage_limit_error 成立意味着有一个已解析的 error 终态帧)。

代价(都落在「重试失败」这一侧,且只影响已终态 attempt 的报告注解,不影响计费):
- 那条最终转发给客户端的 429 事件不再进旧 attempt 的 client capture;
  provider 侧 capture 早在 observe_upstream_frame 里就记下了。
- 如果转发 429 给客户端也失败,record_client_delivery_aborted 落在一个已经结算的
  attempt 上,成为 no-op。

测试:
- lifecycle:await_turn_finalization_handle 必须「等到落地」而不是「排进队列」
  (C6 依赖的性质);结算完成后规划才读状态的顺序型断言(计数器替身);结算任务
  panic 也必须放行调用方,不能卡死 relay loop。
- turn_state:Replanning 状态下 end() 不再交出第二个 attempt(无重复结算)。
- e2e 新增 provider_quota_exhaustion_transparently_retries_onto_another_key:
  mock 上游首轮只回 Codex 的 429 usage_limit_reached,网关换到第二把 key 重放同一个
  response.create;断言客户端看不到 429、上游被连两次、两次用的不是同一把 key、两个
  attempt 各留一条终态行(429 的那条 + 计费的那条)。已验证它在改动前后都通过——
  它覆盖的是整条路径可用,顺序由上面的单测确定性覆盖。
  夹具随之参数化出 ProviderFixture::CodexKeyPair:透明重试只有 Codex adapter 会
  开启,而 codex 候选要求 auth_type = oauth,所以这个夹具用未过期的 oauth 凭证。
This commit is contained in:
AAEE86
2026-08-17 14:53:40 +08:00
committed by ZheFox
parent 1d3051cb89
commit dbf2809bd6
5 changed files with 382 additions and 48 deletions
@@ -12,7 +12,7 @@ use super::client::{adapter_drain_ready, forward_client_message, RelayDispositio
use super::frame::ParsedResponsesWebSocketFrame;
use super::lifecycle::{
await_pending_adapter_observation, finalize_active_turn, queue_turn_finalization,
ActiveProviderAttempt,
settle_turn_finalization, ActiveProviderAttempt, PreviousAttemptSettled,
};
use super::quota::{
active_continuation_can_retry_from_full_input, detach_exhausted_upstream,
@@ -389,13 +389,23 @@ pub(super) async fn relay_bound_connection(
let mut quota_relay_action = classify_quota_relay(quota_facts);
if matches!(quota_relay_action, QuotaRelayAction::AttemptTransparentRetry) {
// detach_attempt 保留 logical turn:重试是同一轮请求的下一个 attempt。
let mut retry_turn = bound.turn_state.detach_attempt().map(ActiveProviderAttempt::disarm);
if let Some(turn) = retry_turn.as_mut() {
turn.release_admission().await;
}
if retry_active_turn_after_quota_exhaustion(bound, state, context).await {
if let Some(turn) = retry_turn {
queue_turn_finalization(
let retry_turn = bound
.turn_state
.detach_attempt()
.map(ActiveProviderAttempt::disarm);
// 先结算旧 attempt 并等它落地,再规划下一个 attempt。两个理由:
//
// 1. 规划要读 health / adaptive / pool 状态,而这些正是旧
// attempt 结算时才投射的。普通的新 turn 早就在 client.rs 里
// 用 await_pending_turn_finalization 挡住了「基于陈旧状态
// 规划」,透明重试这条路径原先漏了这一步。
// 2. 旧 attempt 还占着自己的 pool key lease。不先释放,重试就
// 可能因为「这把 key 仍被占用」而挑不到本该可用的替代 key,
// 或者干脆判成无可用供应商。
let settled = match retry_turn {
Some(mut turn) => {
turn.release_admission().await;
settle_turn_finalization(
bound,
state,
turn,
@@ -403,20 +413,17 @@ pub(super) async fn relay_bound_connection(
ResponsesWebSocketTurnOutcome::upstream_closed,
),
)
.await;
.await
}
None => PreviousAttemptSettled::nothing_to_settle(),
};
if retry_active_turn_after_quota_exhaustion(bound, state, context, settled).await
{
continue;
}
if let Some(turn) = retry_turn {
let restored = bound
.turn_state
.resume(ActiveProviderAttempt::new(state, turn));
debug_assert!(
restored.is_ok(),
"a failed transparent retry must be able to restore its attempt"
);
drop(restored);
}
// 重试失败。旧 attempt 已经结算,logical turn 仍停在
// Replanning,所以后面分支里的 end() / finalize_active_turn
// 只会清掉 logical turn 而不会交出 attempt——不存在重复结算。
quota_relay_action = classify_quota_relay(QuotaRelayFacts {
retry_current_turn: false,
transparent_retry_failed: true,
@@ -129,6 +129,38 @@ pub(super) async fn queue_turn_finalization(
Some(spawn_responses_websocket_turn_finalization(state.clone(), turn, outcome).await);
}
/// 「上一个 attempt 已经结算完毕」的凭证。
///
/// 只能由本模块颁发,且只有在结算真正落地之后。规划下一个 attempt 的入口
/// ([`super::quota::retry_active_turn_after_quota_exhaustion`]) 要求这个参数,
/// 于是「先结算、再规划」成为签名的一部分,而不是一句注释——顺序写反连编译都
/// 过不了。
pub(super) struct PreviousAttemptSettled(());
impl PreviousAttemptSettled {
/// 没有 attempt 要结算(连接此刻不在 `Responding`)。
pub(super) const fn nothing_to_settle() -> Self {
Self(())
}
}
/// 结算一个 attempt 并等它落地。
///
/// 与 [`queue_turn_finalization`] 的区别只在于「等」:后者把 handle 挂在连接上
/// 让 relay loop 继续跑,适用于结算之后不再需要读取共享状态的出口;这个用在
/// 必须先看到结算结果才能继续的路径上——典型的就是透明重试,它紧接着要按
/// health / adaptive / pool 状态规划下一个 attempt。
pub(super) async fn settle_turn_finalization(
bound: &mut BoundResponsesConnection,
state: &AppState,
turn: ResponsesProviderAttempt,
outcome: ResponsesWebSocketTurnOutcome,
) -> PreviousAttemptSettled {
queue_turn_finalization(bound, state, turn, outcome).await;
await_pending_turn_finalization(bound).await;
PreviousAttemptSettled(())
}
pub(super) async fn await_pending_adapter_observation(bound: &mut BoundResponsesConnection) {
if let Some(mut handle) = bound.pending_adapter_observation.take() {
match timeout(RESPONSES_WEBSOCKET_ADAPTER_OBSERVATION_TIMEOUT, &mut handle).await {
@@ -253,3 +285,67 @@ pub(super) fn responses_websocket_turn_start_close(error: &GatewayError) -> (u16
_ => (CLOSE_INTERNAL_ERROR, "turn_start_failed"),
}
}
#[cfg(test)]
mod tests {
use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
use std::sync::Arc;
use std::time::Duration;
use super::await_turn_finalization_handle;
/// C6 依赖的性质:结算是「等到落地」而不是「排进队列」。
///
/// 透明重试在这之后立刻按 health / adaptive / pool 状态规划下一个 attempt
/// 所以结算任务必须已经跑完——只把 handle 挂起来是不够的。
#[tokio::test]
async fn awaiting_a_finalization_handle_runs_the_settlement_to_completion() {
let settled = Arc::new(AtomicBool::new(false));
let flag = Arc::clone(&settled);
let handle = tokio::spawn(async move {
tokio::time::sleep(Duration::from_millis(60)).await;
flag.store(true, Ordering::SeqCst);
});
assert!(
!settled.load(Ordering::SeqCst),
"the settlement has not finished yet"
);
await_turn_finalization_handle(handle).await;
assert!(
settled.load(Ordering::SeqCst),
"the settlement must be complete before the caller proceeds"
);
}
/// 顺序型:结算的每一步都要排在规划之前。
///
/// 用计数器替身重放透明重试的两步——旧 attempt 结算完成写入 1,规划开始时
/// 读到的必须已经是 1。旧实现在这里先规划、再把结算排进队列,规划读到的是 0。
#[tokio::test]
async fn transparent_retry_replans_only_after_the_previous_attempt_is_settled() {
let steps = Arc::new(AtomicUsize::new(0));
// 第一步:结算旧 attempt(等到落地)。
let recorder = Arc::clone(&steps);
let settlement = tokio::spawn(async move {
tokio::time::sleep(Duration::from_millis(40)).await;
recorder.store(1, Ordering::SeqCst);
});
await_turn_finalization_handle(settlement).await;
// 第二步:规划下一个 attempt,它读到的状态必须是结算之后的。
let observed_at_planning = steps.load(Ordering::SeqCst);
assert_eq!(
observed_at_planning, 1,
"planning must observe the state projected by the settled attempt"
);
}
/// 结算任务失败(panic / cancel)也必须让调用方继续,不能把 relay loop 卡死。
#[tokio::test]
async fn a_failed_finalization_task_still_releases_the_caller() {
let handle = tokio::spawn(async { panic!("settlement task exploded") });
await_turn_finalization_handle(handle).await;
}
}
@@ -10,7 +10,7 @@ use super::adapter::{
resolve_responses_websocket_adapter, ResponsesWebSocketDrainDirective,
ResponsesWebSocketRebindSafety,
};
use super::lifecycle::{queue_turn_finalization, ActiveProviderAttempt};
use super::lifecycle::{queue_turn_finalization, ActiveProviderAttempt, PreviousAttemptSettled};
use super::request::{
build_planning_parts, planned_response_create_event, response_create_has_previous_response_id,
};
@@ -99,10 +99,17 @@ pub(super) fn record_exhausted_bound_key(
Some((key_id, exclusion_until))
}
/// 为同一个 logical turn 规划并绑定下一个 attempt。
///
/// `_previous_settled` 不被使用,它只是把「上一个 attempt 已经结算完毕」这个
/// 前置条件写进签名:规划要读 health / adaptive / pool 状态,而这些是上一个
/// attempt 结算时才投射的;它的 pool key lease 也要先释放,否则替代 key 的挑选
/// 会看到一把仍被占用的 key。
pub(super) async fn retry_active_turn_after_quota_exhaustion(
bound: &mut BoundResponsesConnection,
state: &AppState,
context: &WebSocketRequestContext,
_previous_settled: PreviousAttemptSettled,
) -> bool {
let Some(active) = bound.turn_state.logical_mut() else {
return false;
@@ -199,6 +199,26 @@ mod tests {
)
}
/// 透明重试失败之后:旧 attempt 已经被 detach 并结算过,logical turn 仍停在
/// `Replanning`。此时 `end()` 不能再交出 attempt,否则同一个 attempt 会被
/// 结算两次(两条 usage terminal、两次 pool lease 释放)。
#[test]
fn ending_a_replanning_turn_does_not_hand_out_a_second_attempt() {
let mut state = ResponsesTurnState::Idle;
state.begin(logical(), FakeAttempt(1));
let detached = state.detach_attempt();
assert_eq!(detached, Some(FakeAttempt(1)), "the attempt is settled once");
assert!(matches!(state, ResponsesTurnState::Replanning { .. }));
// 结算已经发生,没有第二个 attempt 可交。
assert!(
state.end().is_none(),
"a replanning turn must not yield a second attempt to settle"
);
assert!(state.accepts_new_response_create());
}
#[test]
fn idle_has_no_turn_and_accepts_a_new_response_create() {
let state = ResponsesTurnState::<FakeAttempt>::Idle;