From dbf2809bd6d3dac34e17b71e90afa5e67e2baff2 Mon Sep 17 00:00:00 2001 From: AAEE86 Date: Sat, 1 Aug 2026 02:40:18 +0800 Subject: [PATCH] fix(ws): settle the previous attempt before transparent retry replanning MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 评审第 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 凭证。 --- .../proxy/websocket/responses/connection.rs | 45 +-- .../proxy/websocket/responses/lifecycle.rs | 96 +++++++ .../proxy/websocket/responses/quota.rs | 9 +- .../proxy/websocket/responses/turn_state.rs | 20 ++ .../tests/responses_websocket_e2e.rs | 260 ++++++++++++++++-- 5 files changed, 382 insertions(+), 48 deletions(-) diff --git a/apps/aether-gateway/src/handlers/proxy/websocket/responses/connection.rs b/apps/aether-gateway/src/handlers/proxy/websocket/responses/connection.rs index a224d33a3..ecc6937b5 100644 --- a/apps/aether-gateway/src/handlers/proxy/websocket/responses/connection.rs +++ b/apps/aether-gateway/src/handlers/proxy/websocket/responses/connection.rs @@ -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, diff --git a/apps/aether-gateway/src/handlers/proxy/websocket/responses/lifecycle.rs b/apps/aether-gateway/src/handlers/proxy/websocket/responses/lifecycle.rs index ba8814382..ba14076ef 100644 --- a/apps/aether-gateway/src/handlers/proxy/websocket/responses/lifecycle.rs +++ b/apps/aether-gateway/src/handlers/proxy/websocket/responses/lifecycle.rs @@ -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; + } +} diff --git a/apps/aether-gateway/src/handlers/proxy/websocket/responses/quota.rs b/apps/aether-gateway/src/handlers/proxy/websocket/responses/quota.rs index c3d1ca449..128c30069 100644 --- a/apps/aether-gateway/src/handlers/proxy/websocket/responses/quota.rs +++ b/apps/aether-gateway/src/handlers/proxy/websocket/responses/quota.rs @@ -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; diff --git a/apps/aether-gateway/src/handlers/proxy/websocket/responses/turn_state.rs b/apps/aether-gateway/src/handlers/proxy/websocket/responses/turn_state.rs index 2b03acf57..ad3c79889 100644 --- a/apps/aether-gateway/src/handlers/proxy/websocket/responses/turn_state.rs +++ b/apps/aether-gateway/src/handlers/proxy/websocket/responses/turn_state.rs @@ -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::::Idle; diff --git a/crates/aether-testing/integration/tests/responses_websocket_e2e.rs b/crates/aether-testing/integration/tests/responses_websocket_e2e.rs index da592caa6..cf73f917d 100644 --- a/crates/aether-testing/integration/tests/responses_websocket_e2e.rs +++ b/crates/aether-testing/integration/tests/responses_websocket_e2e.rs @@ -52,12 +52,18 @@ const PROVIDER_API_KEY: &str = "sk-upstream-responses-ws-e2e"; const PROVIDER_ID: &str = "provider-responses-ws-e2e"; const ENDPOINT_ID: &str = "endpoint-responses-ws-e2e"; const PROVIDER_KEY_ID: &str = "provider-key-responses-ws-e2e"; +/// 透明重试的替代 key。只有配额重试用例会 seed 它。 +const ALTERNATE_PROVIDER_KEY_ID: &str = "provider-key-responses-ws-e2e-alt"; +const ALTERNATE_PROVIDER_API_KEY: &str = "sk-upstream-responses-ws-e2e-alt"; const GLOBAL_MODEL_ID: &str = "global-model-responses-ws-e2e"; const PROVIDER_MODEL_ID: &str = "provider-model-responses-ws-e2e"; const API_KEY_ID: &str = "api-key-responses-ws-e2e"; const PUBLIC_MODEL: &str = "gpt-responses-ws-e2e"; const UPSTREAM_MODEL: &str = "gpt-responses-ws-upstream"; +/// 2100-01-01，保证 oauth 凭证在测试期间不会被判为过期。 +const FAR_FUTURE_UNIX_SECS: u64 = 4_102_444_800; + const INPUT_TOKENS: u64 = 4; const OUTPUT_TOKENS: u64 = 2; @@ -235,6 +241,101 @@ async fn upstream_drop_mid_turn_reports_an_error_and_settles_the_usage_row() -> Ok(()) } +/// 供应商配额耗尽后的透明重试：客户端不该看到 429，两个 attempt 都要结算。 +/// +/// 第一个 attempt 拿到 Codex 的 `usage_limit_reached`，网关换到第二把 key 重开一条 +/// 上游连接重放同一个 `response.create`。C6 之前，重试的规划发生在旧 attempt 结算 +/// 之前：规划读到的是旧 attempt 还没投射的 health / adaptive / pool 状态，而且旧 +/// attempt 的 pool key lease 还被它自己占着。 +/// +/// 顺序本身在这里无法确定性断言（结算与规划都在同一个任务里、DB 里看不到先后）， +/// 由 lifecycle 的单测确定性覆盖；这个用例保证整条路径真的能跑通，并且两个 +/// attempt 都留下了终态记账行。 +#[tokio::test] +async fn provider_quota_exhaustion_transparently_retries_onto_another_key() -> Result<(), BoxError> +{ + let harness = Harness::start_with_fixture( + UpstreamBehavior::QuotaExhaustedThenComplete, + ProviderFixture::CodexKeyPair, + ) + .await?; + let mut client = harness.connect().await?; + + client + .send(response_create(json!({"input": "retry after quota"}))) + .await?; + + // 客户端只应该看到重试之后那次成功的响应，看不到 429。 + let completed = receive_event(&mut client, "response.completed").await?; + assert_eq!( + completed.pointer("/response/status").and_then(Value::as_str), + Some("completed") + ); + + // 上游被连了两次：配额耗尽的那条 + 重试用的那条。 + assert_eq!( + harness.upstream.connections(), + 2, + "the transparent retry must open a second upstream connection" + ); + let observed = harness.upstream.observed_events().await; + assert_eq!( + observed.len(), + 2, + "the same response.create must be replayed once" + ); + + // 两把不同的 key 被用过：重试不能落回那把已经耗尽的 key。 + let authorizations = harness.upstream.authorization_headers().await; + assert_eq!(authorizations.len(), 2); + assert_ne!( + authorizations[0], authorizations[1], + "the retry must not reuse the exhausted key: {authorizations:?}" + ); + + // 两个 attempt 各自留下一条终态行：配额失败的那条 + 成功计费的那条。 + let audits = harness + .usage_audits_where(2, "settled attempts", |audit| !is_pending(audit)) + .await?; + let settled = audits + .iter() + .filter(|audit| !is_pending(audit)) + .collect::>(); + assert_eq!( + settled.len(), + 2, + "both attempts must reach a terminal accounting row: {:?}", + audits + .iter() + .map(|audit| (audit.status.clone(), audit.status_code, audit.total_tokens)) + .collect::>() + ); + assert!( + settled.iter().any(|audit| audit.status_code == Some(429)), + "the exhausted attempt keeps its own 429 row: {:?}", + settled + .iter() + .map(|audit| (audit.status.clone(), audit.status_code)) + .collect::>() + ); + + let billed = harness + .usage_audits_where(1, "the billed retry attempt", is_billed) + .await?; + let retry = billed + .iter() + .find(|audit| is_billed(audit)) + .ok_or("the successful retry attempt must be billed")?; + assert_eq!( + retry.total_tokens, + INPUT_TOKENS + OUTPUT_TOKENS, + "the retry attempt is billed for what it actually consumed" + ); + + client.close(None).await?; + Ok(()) +} + fn is_pending(audit: &StoredRequestUsageAudit) -> bool { audit.status.eq_ignore_ascii_case("pending") } @@ -258,14 +359,46 @@ struct Harness { _gateway_server: SpawnedServer, } +/// 供应商夹具形态。 +/// +/// 透明配额重试只有 Codex adapter 会开启（`retry_current_turn: true` 只从 +/// codex.rs 出），而且重试要有第二把 key 可挑，否则规划直接判无可用供应商。 +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum ProviderFixture { + /// 单个 openai 类型供应商、单把 key。 + SingleOpenAiKey, + /// codex 类型供应商 + 两把 key：第一把配额耗尽后重试落到第二把。 + CodexKeyPair, +} + +impl ProviderFixture { + const fn provider_type(self) -> &'static str { + match self { + Self::SingleOpenAiKey => "openai", + Self::CodexKeyPair => "codex", + } + } + + const fn has_alternate_key(self) -> bool { + matches!(self, Self::CodexKeyPair) + } +} + impl Harness { async fn start(behavior: UpstreamBehavior) -> Result { + Self::start_with_fixture(behavior, ProviderFixture::SingleOpenAiKey).await + } + + async fn start_with_fixture( + behavior: UpstreamBehavior, + fixture: ProviderFixture, + ) -> Result { let upstream = Arc::new(MockUpstreamState::new(behavior)); let upstream_server = SpawnedServer::start(mock_upstream_router(Arc::clone(&upstream))).await?; let database = TemporarySqlite::new(); - prepare_and_seed_database(&database.config, upstream_server.base_url()).await?; + prepare_and_seed_database(&database.config, upstream_server.base_url(), fixture).await?; let data_config = GatewayDataConfig::from_database_config(database.config.clone()) .with_encryption_key(DEVELOPMENT_ENCRYPTION_KEY); @@ -480,6 +613,11 @@ enum UpstreamBehavior { StallAfterCreated, /// Announce the response and then hang up mid-turn. CloseAfterCreated, + /// 第一轮只回一个 Codex 配额耗尽错误，之后的每一轮正常完成。 + /// + /// 第一轮刻意不发 `response.created`：任何标准 `response.*` 事件都会让 + /// codex adapter 把这一轮判成 replay-unsafe，透明重试就不会发生。 + QuotaExhaustedThenComplete, } #[derive(Debug)] @@ -571,6 +709,16 @@ async fn run_mock_upstream( let _ = send_mock_created(&mut socket, &response_id).await; break; } + UpstreamBehavior::QuotaExhaustedThenComplete => { + if turn == 1 { + let _ = send_mock_event(&mut socket, codex_quota_exhausted_error()) + .await; + break; + } + if send_mock_turn(&mut socket, &response_id).await.is_err() { + break; + } + } } } AxumWsMessage::Ping(payload) => { @@ -584,6 +732,24 @@ async fn run_mock_upstream( } } +/// Codex 的账户级配额耗尽信号。 +/// +/// `status_code: 429` + `error.type: usage_limit_reached` 是 adapter 识别 +/// 「配额耗尽、可透明重试」的最小载荷：解析出的元数据被强制标上 +/// `limit_reached: true`，于是 drain 指令带着 `retry_current_turn: true` 下来。 +fn codex_quota_exhausted_error() -> Value { + json!({ + "type": "error", + "status_code": 429, + "error": { + "type": "usage_limit_reached", + "message": "You have hit your usage limit", + "plan_type": "plus", + "resets_in_seconds": 3_600 + } + }) +} + async fn send_mock_created(socket: &mut WebSocket, response_id: &str) -> Result<(), axum::Error> { send_mock_event( socket, @@ -681,6 +847,7 @@ impl Drop for TemporarySqlite { async fn prepare_and_seed_database( database: &SqlDatabaseConfig, upstream_base_url: &str, + fixture: ProviderFixture, ) -> Result<(), BoxError> { let backends = DataBackends::from_config(DataLayerConfig::from_database(database.clone()))?; let pending = backends @@ -691,7 +858,7 @@ async fn prepare_and_seed_database( backends.run_database_migrations().await?; } - seed_provider_catalog(&backends, upstream_base_url).await?; + seed_provider_catalog(&backends, upstream_base_url, fixture).await?; seed_models(&backends).await?; let user_id = seed_user(&backends).await?; seed_client_api_key(&backends, &user_id).await?; @@ -716,6 +883,7 @@ async fn prepare_and_seed_database( async fn seed_provider_catalog( backends: &DataBackends, upstream_base_url: &str, + fixture: ProviderFixture, ) -> Result<(), BoxError> { let writer = backends .write() @@ -727,7 +895,7 @@ async fn seed_provider_catalog( PROVIDER_ID.to_string(), "Responses WebSocket E2E".to_string(), None, - "openai".to_string(), + fixture.provider_type().to_string(), )? .with_transport_fields( true, @@ -766,35 +934,71 @@ async fn seed_provider_catalog( ) .await?; writer - .create_key( - &StoredProviderCatalogKey::new( - PROVIDER_KEY_ID.to_string(), - PROVIDER_ID.to_string(), - "Responses WebSocket E2E".to_string(), - "api_key".to_string(), - Some(json!({"streaming": true})), - true, - )? - .with_transport_fields( - Some(json!(["openai:responses"])), - encrypt_python_fernet_plaintext(DEVELOPMENT_ENCRYPTION_KEY, PROVIDER_API_KEY)?, - None, - None, - Some(json!({"openai:responses": 1})), - Some(json!([PUBLIC_MODEL, UPSTREAM_MODEL])), - None, - None, - None, - )? - .with_health_fields( - Some(json!({"openai:responses": {"status": "healthy"}})), - Some(json!({"openai:responses": {"state": "closed"}})), - ), - ) + .create_key(&catalog_key(PROVIDER_KEY_ID, PROVIDER_API_KEY, fixture)?) .await?; + if fixture.has_alternate_key() { + writer + .create_key(&catalog_key( + ALTERNATE_PROVIDER_KEY_ID, + ALTERNATE_PROVIDER_API_KEY, + fixture, + )?) + .await?; + } Ok(()) } +/// 一把健康、可服务本用例模型的 key。 +/// +/// codex 类型的候选要求 `auth_type = oauth`（见 candidate_selection 的 +/// provider_type 约束），所以配额重试夹具走 oauth，凭证是一份未过期的 +/// access_token。 +fn catalog_key( + key_id: &str, + secret: &str, + fixture: ProviderFixture, +) -> Result { + let oauth = fixture.has_alternate_key(); + let auth_type = if oauth { "oauth" } else { "api_key" }; + let auth_config = if oauth { + Some(encrypt_python_fernet_plaintext( + DEVELOPMENT_ENCRYPTION_KEY, + &json!({ + "access_token": secret, + "refresh_token": format!("{secret}-refresh"), + "account_id": format!("{key_id}-account"), + "expires_at": FAR_FUTURE_UNIX_SECS, + }) + .to_string(), + )?) + } else { + None + }; + Ok(StoredProviderCatalogKey::new( + key_id.to_string(), + PROVIDER_ID.to_string(), + "Responses WebSocket E2E".to_string(), + auth_type.to_string(), + Some(json!({"streaming": true})), + true, + )? + .with_transport_fields( + Some(json!(["openai:responses"])), + encrypt_python_fernet_plaintext(DEVELOPMENT_ENCRYPTION_KEY, secret)?, + auth_config, + None, + Some(json!({"openai:responses": 1})), + Some(json!([PUBLIC_MODEL, UPSTREAM_MODEL])), + None, + None, + None, + )? + .with_health_fields( + Some(json!({"openai:responses": {"status": "healthy"}})), + Some(json!({"openai:responses": {"state": "closed"}})), + )) +} + async fn seed_models(backends: &DataBackends) -> Result<(), BoxError> { let writer = backends .write()