From ab7fe4d1870269f9b1d5a96f9841c834f79e31d6 Mon Sep 17 00:00:00 2001 From: Mi Tom <6468993+MDX-Tom@users.noreply.github.com> Date: Thu, 20 Aug 2026 19:47:12 +0800 Subject: [PATCH] fix: WebSocket pre-response recovery and protocol alignment --- crates/service/src/agent_identity.rs | 48 ++- .../auth/tests/token_exchange_tests.rs | 42 +- .../src/gateway/auth/token_exchange.rs | 64 ++- crates/service/src/gateway/mod.rs | 1 + .../upstream/attempt_flow/transport.rs | 9 +- .../upstream/attempt_flow/transport_tests.rs | 9 +- .../service/src/http/responses_websocket.rs | 158 ++++--- .../src/http/responses_websocket_rebase.rs | 4 +- .../src/http/responses_websocket_tests.rs | 8 + .../src/http/tests/proxy_runtime_tests.rs | 401 ++++++++++++++++++ .../service/src/usage/usage_token_refresh.rs | 7 +- 11 files changed, 628 insertions(+), 123 deletions(-) diff --git a/crates/service/src/agent_identity.rs b/crates/service/src/agent_identity.rs index 4f73e6e0d..95974e37f 100644 --- a/crates/service/src/agent_identity.rs +++ b/crates/service/src/agent_identity.rs @@ -507,6 +507,12 @@ fn register_agent_identity( let mut request_builder = client .post(&url) .timeout(AGENT_REGISTRATION_TIMEOUT) + // Keep the auth request on the same first-party identity path as the + // official Codex client. The account bearer token alone is valid for + // normal Responses calls, but agent registration also relies on the + // Codex originator and User-Agent headers for request classification. + .header("originator", crate::gateway::current_wire_originator()) + .header("User-Agent", crate::gateway::current_codex_user_agent()) .bearer_auth(access_token) .json(&request); if is_fedramp { @@ -671,6 +677,10 @@ fn register_agent_identity_task( let response = client .post(&url) .timeout(AGENT_TASK_REGISTRATION_TIMEOUT) + // The official Codex auth client applies the same identity headers to + // task registration as to runtime registration. + .header("originator", crate::gateway::current_wire_originator()) + .header("User-Agent", crate::gateway::current_codex_user_agent()) .json(&request) .send() .map_err(|err| format!("agent task registration request failed: {err}"))?; @@ -1011,13 +1021,15 @@ mod tests { .expect("registration request"); let path = request.url().to_string(); let authorization = request_header(&request, "authorization"); + let originator = request_header(&request, "originator"); + let user_agent = request_header(&request, "user-agent"); let mut body = String::new(); request .as_reader() .read_to_string(&mut body) .expect("read registration request"); request_tx - .send((path, authorization, body)) + .send((path, authorization, originator, user_agent, body)) .expect("record request"); request .respond( @@ -1050,10 +1062,16 @@ mod tests { ); assert!(authorization.value.starts_with("AgentAssertion ")); - let (registration_path, registration_auth, registration_body) = request_rx + let ( + registration_path, + registration_auth, + registration_originator, + registration_user_agent, + registration_body, + ) = request_rx .recv_timeout(Duration::from_secs(5)) .expect("receive identity registration"); - let (task_path, task_auth, _task_body) = request_rx + let (task_path, task_auth, _task_originator, _task_user_agent, _task_body) = request_rx .recv_timeout(Duration::from_secs(5)) .expect("receive task registration"); server_handle.join().expect("join registration server"); @@ -1062,6 +1080,14 @@ mod tests { registration_auth.as_deref(), Some(format!("Bearer {access_token}").as_str()) ); + assert_eq!( + registration_originator.as_deref(), + Some(crate::gateway::current_wire_originator().as_str()) + ); + assert_eq!( + registration_user_agent.as_deref(), + Some(crate::gateway::current_codex_user_agent().as_str()) + ); let registration_body: serde_json::Value = serde_json::from_str(®istration_body).expect("parse registration body"); assert_eq!( @@ -1363,12 +1389,16 @@ mod tests { .expect("registration server timeout") .expect("registration request"); let path = request.url().to_string(); + let originator = request_header(&request, "originator"); + let user_agent = request_header(&request, "user-agent"); let mut body = String::new(); request .as_reader() .read_to_string(&mut body) .expect("read registration request"); - request_tx.send((path, body)).expect("record request"); + request_tx + .send((path, originator, user_agent, body)) + .expect("record request"); request .respond( Response::from_string(r#"{"taskId":"task-from-server"}"#) @@ -1385,11 +1415,19 @@ mod tests { register_agent_identity_task(&reqwest::blocking::Client::new(), &identity, &base_url) .expect("register task"); assert_eq!(task_id, "task-from-server"); - let (path, body) = request_rx + let (path, originator, user_agent, body) = request_rx .recv_timeout(Duration::from_secs(5)) .expect("receive request"); server_handle.join().expect("join server"); assert_eq!(path, "/v1/agent/agent-runtime-1/task/register"); + assert_eq!( + originator.as_deref(), + Some(crate::gateway::current_wire_originator().as_str()) + ); + assert_eq!( + user_agent.as_deref(), + Some(crate::gateway::current_codex_user_agent().as_str()) + ); let body: serde_json::Value = serde_json::from_str(&body).expect("parse request body"); let timestamp = body["timestamp"].as_str().expect("timestamp"); let signature = BASE64_STANDARD diff --git a/crates/service/src/gateway/auth/tests/token_exchange_tests.rs b/crates/service/src/gateway/auth/tests/token_exchange_tests.rs index 4d1fdd849..c88866b11 100644 --- a/crates/service/src/gateway/auth/tests/token_exchange_tests.rs +++ b/crates/service/src/gateway/auth/tests/token_exchange_tests.rs @@ -1,4 +1,5 @@ use super::*; +use base64::Engine as _; /// 函数 `same_account_reuses_exchange_lock` /// @@ -115,7 +116,7 @@ fn fallback_to_access_token_uses_runtime_access_token_when_exchange_fails() { assert_eq!(bearer, "runtime-access-token"); } -/// 函数 `api_key_exchange_subject_tokens_falls_back_to_imported_access_token` +/// 函数 `api_key_exchange_subject_token_omits_access_token_without_id_token` /// /// 作者: gaohongshun /// @@ -127,7 +128,7 @@ fn fallback_to_access_token_uses_runtime_access_token_when_exchange_fails() { /// # 返回 /// 无 #[test] -fn api_key_exchange_subject_tokens_falls_back_to_imported_access_token() { +fn api_key_exchange_subject_token_omits_access_token_without_id_token() { let token = Token { account_id: "acc-import-session".to_string(), id_token: String::new(), @@ -137,13 +138,10 @@ fn api_key_exchange_subject_tokens_falls_back_to_imported_access_token() { last_refresh: now_ts(), }; - assert_eq!( - api_key_exchange_subject_tokens(&token), - vec!["imported-session-access".to_string()] - ); + assert_eq!(api_key_exchange_subject_token(&token), None); } -/// 函数 `api_key_exchange_subject_tokens_prefers_access_token_before_id_token` +/// 函数 `api_key_exchange_subject_token_uses_id_token_only` /// /// 作者: gaohongshun /// @@ -155,7 +153,7 @@ fn api_key_exchange_subject_tokens_falls_back_to_imported_access_token() { /// # 返回 /// 无 #[test] -fn api_key_exchange_subject_tokens_prefers_access_token_before_id_token() { +fn api_key_exchange_subject_token_uses_id_token_only() { let token = Token { account_id: "acc-login".to_string(), id_token: "id-token".to_string(), @@ -166,8 +164,32 @@ fn api_key_exchange_subject_tokens_prefers_access_token_before_id_token() { }; assert_eq!( - api_key_exchange_subject_tokens(&token), - vec!["access-token".to_string(), "id-token".to_string()] + api_key_exchange_subject_token(&token), + Some("id-token".to_string()) + ); +} + +#[test] +fn api_key_exchange_client_id_prefers_id_token_claim() { + let jwt = |client_id: &str| { + let payload = serde_json::json!({"sub":"user-test","client_id": client_id}).to_string(); + format!( + "header.{}.signature", + base64::engine::general_purpose::URL_SAFE_NO_PAD.encode(payload) + ) + }; + let token = Token { + account_id: "acc-client-id".to_string(), + id_token: jwt("id-token-client"), + access_token: jwt("access-token-client"), + refresh_token: String::new(), + api_key_access_token: None, + last_refresh: now_ts(), + }; + + assert_eq!( + api_key_exchange_client_id(&token, "fallback-client"), + "id-token-client" ); } diff --git a/crates/service/src/gateway/auth/token_exchange.rs b/crates/service/src/gateway/auth/token_exchange.rs index 03688494a..3affefb7e 100644 --- a/crates/service/src/gateway/auth/token_exchange.rs +++ b/crates/service/src/gateway/auth/token_exchange.rs @@ -1,7 +1,7 @@ use std::collections::HashMap; use std::sync::{Arc, Mutex, OnceLock}; -use codexmanager_core::auth::extract_token_exp; +use codexmanager_core::auth::{extract_client_id_claim, extract_token_exp, DEFAULT_CLIENT_ID}; use codexmanager_core::storage::{now_ts, Account, Storage, Token}; use crate::account_status::mark_account_unavailable_for_auth_error; @@ -139,40 +139,39 @@ fn exchange_and_persist_api_key_access_token( issuer: &str, client_id: &str, ) -> Result { - let mut errors = Vec::new(); - for subject_token in api_key_exchange_subject_tokens(token) { - match auth_tokens::obtain_api_key(issuer, client_id, &subject_token) { - Ok(exchanged) => { - token.api_key_access_token = Some(exchanged.clone()); - let _ = storage.insert_token(token); - return Ok(exchanged); - } - Err(err) => errors.push(err), + let Some(subject_token) = api_key_exchange_subject_token(token) else { + return Err("id_token is unavailable for API key token exchange".to_string()); + }; + match auth_tokens::obtain_api_key(issuer, client_id, &subject_token) { + Ok(exchanged) => { + token.api_key_access_token = Some(exchanged.clone()); + let _ = storage.insert_token(token); + Ok(exchanged) } + Err(err) => Err(err), } - - let exchange_error = errors - .into_iter() - .next() - .unwrap_or_else(|| "api key exchange subject token is missing".to_string()); - Err(exchange_error) } -fn api_key_exchange_subject_tokens(token: &Token) -> Vec { - let mut subjects = Vec::new(); - // 中文注释：直接导入 api/auth/session JSON 时通常只有 accessToken； - // 登录授权路径已优先缓存 api_key_access_token，缓存缺失时也先按最新 access_token 兑换。 - push_unique_subject_token(&mut subjects, token.access_token.as_str()); - push_unique_subject_token(&mut subjects, token.id_token.as_str()); - subjects +fn api_key_exchange_subject_token(token: &Token) -> Option { + // `/oauth/token` uses the token-exchange grant and expects the OAuth ID + // token as its subject. An access token is the bearer fallback for the + // upstream request; sending it to this endpoint produces misleading + // "Invalid ID token" / audience errors even when the account is usable. + let id_token = token.id_token.trim(); + (!id_token.is_empty()).then(|| id_token.to_string()) } -fn push_unique_subject_token(subjects: &mut Vec, candidate: &str) { - let value = candidate.trim(); - if value.is_empty() || subjects.iter().any(|existing| existing == value) { - return; - } - subjects.push(value.to_string()); +pub(crate) fn api_key_exchange_client_id(token: &Token, fallback_client_id: &str) -> String { + // The exchange subject is the ID token, so prefer its client_id claim. A + // separately issued access token can carry a different audience/client + // claim and must not override the ID-token exchange client. + extract_client_id_claim(&token.id_token) + .or_else(|| extract_client_id_claim(&token.access_token)) + .or_else(|| { + let fallback = fallback_client_id.trim(); + (!fallback.is_empty()).then(|| fallback.to_string()) + }) + .unwrap_or_else(|| DEFAULT_CLIENT_ID.to_string()) } /// 函数 `fallback_to_access_token` @@ -256,7 +255,7 @@ pub(super) fn resolve_openai_bearer_token( } let fallback_client_id = super::runtime_config::token_exchange_client_id(); - let client_id = crate::usage_token_refresh::token_refresh_client_id(token, &fallback_client_id); + let client_id = api_key_exchange_client_id(token, &fallback_client_id); let issuer_env = super::runtime_config::token_exchange_default_issuer(); let issuer = if account.issuer.trim().is_empty() { issuer_env @@ -305,10 +304,7 @@ pub(super) fn resolve_openai_bearer_token( let _ = storage.insert_token(token); if !token.id_token.trim().is_empty() { - let refreshed_client_id = - crate::usage_token_refresh::token_refresh_client_id( - token, &client_id, - ); + let refreshed_client_id = api_key_exchange_client_id(token, &client_id); if let Ok(exchanged) = exchange_and_persist_api_key_access_token( storage, token, diff --git a/crates/service/src/gateway/mod.rs b/crates/service/src/gateway/mod.rs index 69edf7a45..d7507ec4f 100644 --- a/crates/service/src/gateway/mod.rs +++ b/crates/service/src/gateway/mod.rs @@ -436,6 +436,7 @@ pub(crate) use selection::{ }; #[cfg(test)] use token_exchange::account_token_exchange_lock; +pub(crate) use token_exchange::api_key_exchange_client_id; use token_exchange::resolve_openai_bearer_token; use upstream::proxy::proxy_validated_request; diff --git a/crates/service/src/gateway/upstream/attempt_flow/transport.rs b/crates/service/src/gateway/upstream/attempt_flow/transport.rs index af00dded5..68021ea04 100644 --- a/crates/service/src/gateway/upstream/attempt_flow/transport.rs +++ b/crates/service/src/gateway/upstream/attempt_flow/transport.rs @@ -1481,11 +1481,7 @@ fn is_websocket_upstream_terminal_text(text: &str) -> bool { .unwrap_or_default() .to_ascii_lowercase() .as_str(), - "response.completed" - | "response.done" - | "response.failed" - | "response.incomplete" - | "error" + "response.completed" | "response.failed" | "response.incomplete" | "error" ) } @@ -1521,7 +1517,7 @@ fn is_websocket_upstream_connection_limit_text(text: &str) -> bool { } fn is_websocket_upstream_transport_healthy_terminal_text(text: &str) -> bool { - is_websocket_upstream_terminal_text(text) && !is_websocket_upstream_connection_limit_text(text) + is_websocket_upstream_completed_text(text) && !is_websocket_upstream_connection_limit_text(text) } fn websocket_upstream_sse_event(text: &str) -> String { @@ -1533,7 +1529,6 @@ fn websocket_upstream_sse_event(text: &str) -> String { } } -#[cfg(test)] fn is_websocket_upstream_completed_text(text: &str) -> bool { serde_json::from_str::(text) .ok() diff --git a/crates/service/src/gateway/upstream/attempt_flow/transport_tests.rs b/crates/service/src/gateway/upstream/attempt_flow/transport_tests.rs index 240ffb800..6d8ffebfc 100644 --- a/crates/service/src/gateway/upstream/attempt_flow/transport_tests.rs +++ b/crates/service/src/gateway/upstream/attempt_flow/transport_tests.rs @@ -807,7 +807,7 @@ fn websocket_upstream_terminal_detection_parses_json_type() { assert!(super::is_websocket_upstream_terminal_text( r#"{"type":"response.completed"}"# )); - assert!(super::is_websocket_upstream_terminal_text( + assert!(!super::is_websocket_upstream_terminal_text( r#"{"type":"response.done"}"# )); assert!(super::is_websocket_upstream_terminal_text( @@ -991,7 +991,7 @@ fn send_websocket_upstream_request_builds_valid_handshake_and_stops_on_completed } #[test] -fn send_websocket_upstream_request_does_not_cooldown_after_application_failure() { +fn send_websocket_upstream_request_does_not_mark_recovery_completed_after_application_failure() { let _env_lock = crate::test_env_guard(); let _reload_guard = RuntimeConfigReloadGuard; let _proxy_guard = EnvGuard::set("CODEXMANAGER_UPSTREAM_PROXY_URL", ""); @@ -1016,6 +1016,11 @@ fn send_websocket_upstream_request_does_not_cooldown_after_application_failure() assert!( super::is_websocket_upstream_transport_healthy_terminal_text( + r#"{"type":"response.completed"}"# + ) + ); + assert!( + !super::is_websocket_upstream_transport_healthy_terminal_text( r#"{"type":"response.failed"}"# ) ); diff --git a/crates/service/src/http/responses_websocket.rs b/crates/service/src/http/responses_websocket.rs index f7e301700..74ca6bc15 100644 --- a/crates/service/src/http/responses_websocket.rs +++ b/crates/service/src/http/responses_websocket.rs @@ -51,6 +51,12 @@ const WEBSOCKET_CONNECTION_LIMIT_REACHED_CODE: &str = "websocket_connection_limi const RESPONSES_WS_REQUEST_IN_FLIGHT_CODE: &str = "response_in_flight"; const WEBSOCKET_CONNECTION_LIMIT_REACHED_MESSAGE: &str = "Responses websocket connection limit reached (60 minutes). Create a new websocket connection to continue."; +// A freshly handshaken socket can still be reset before its first client frame reaches the +// upstream. Allow one additional fresh socket, but keep recovery bounded and replay-free. +const RESPONSES_WS_MAX_PENDING_FRAME_SEND_ATTEMPTS: usize = 2; +// A response can emit only connection preamble events before a transport reset. Reconnect the +// lane at most twice in that state; once substantive output exists, replay is not safe. +const RESPONSES_WS_MAX_PRE_COMPLETION_RECOVERY_ATTEMPTS: u8 = 2; #[derive(Clone)] struct WsRequestContext { @@ -89,7 +95,7 @@ struct PendingWsRequestState { conversation_routing: Option, forwarded_upstream_event: bool, forwarded_non_preamble_event: bool, - replayed_after_upstream_disconnect: bool, + upstream_disconnect_recovery_attempts: u8, suppress_replayed_preamble: bool, buffered_upstream_preamble: Vec, buffer_retry_preamble: bool, @@ -379,7 +385,7 @@ async fn run_responses_websocket_session(mut socket: WebSocket, context: WsReque conversation_routing: upstream.conversation_routing.clone(), forwarded_upstream_event: false, forwarded_non_preamble_event: false, - replayed_after_upstream_disconnect: false, + upstream_disconnect_recovery_attempts: 0, suppress_replayed_preamble: false, buffered_upstream_preamble: Vec::new(), buffer_retry_preamble: should_buffer_ws_retry_preamble( @@ -589,7 +595,7 @@ async fn run_responses_websocket_session(mut socket: WebSocket, context: WsReque conversation_routing: upstream.conversation_routing.clone(), forwarded_upstream_event: false, forwarded_non_preamble_event: false, - replayed_after_upstream_disconnect: false, + upstream_disconnect_recovery_attempts: 0, suppress_replayed_preamble: false, buffered_upstream_preamble: Vec::new(), buffer_retry_preamble, @@ -1958,52 +1964,80 @@ async fn reconnect_upstream_for_pending_request( completed_responses: &CompletedWsResponseCache, completed_tool_calls: &CompletedWsToolCallCache, ) -> Result { - let mut replacement = connect_upstream_websocket_with_timeout( - context, - pending.prepared.model.as_deref(), - previous_account_id, - ) - .await?; - let account_changed = - previous_account_id.is_some_and(|account_id| account_id != replacement.account_id); - if let Err(err) = prepare_ws_request_for_new_connection( - pending, - completed_responses, - completed_tool_calls, - account_changed, - ) { - let _ = replacement.stream.close(None).await; - return Err(err); - } - - pending.attempted_account_ids.clear(); - pending - .attempted_account_ids - .insert(replacement.account_id.clone()); - pending.log.route_strategy = Some(replacement.route_strategy.to_string()); - pending.log.route_source = Some(replacement.route_source.to_string()); - pending.conversation_routing = replacement.conversation_routing.clone(); - pending.buffer_retry_preamble = should_buffer_ws_retry_preamble( - &replacement, - &pending.attempted_account_ids, - pending.prepared.text.as_str(), - pending.retried_missing_tool_call_context, - ); - if let Err(err) = replacement - .stream - .send(UpstreamMessage::Text(pending.prepared.text.clone().into())) - .await - { - let account_id = replacement.account_id.clone(); - let _ = replacement.stream.close(None).await; - return Err(WsSessionError::bad_gateway_bilingual( - "重连后发送上游 WebSocket 帧失败", - format!( - "send upstream websocket frame after reconnect failed for account {account_id}: {err}" - ), - )); + let mut previous_account_id = previous_account_id.map(str::to_owned); + let mut last_send_error = None; + + for attempt in 1..=RESPONSES_WS_MAX_PENDING_FRAME_SEND_ATTEMPTS { + let mut replacement = connect_upstream_websocket_with_timeout( + context, + pending.prepared.model.as_deref(), + previous_account_id.as_deref(), + ) + .await?; + let account_changed = previous_account_id + .as_deref() + .is_some_and(|account_id| account_id != replacement.account_id); + if let Err(err) = prepare_ws_request_for_new_connection( + pending, + completed_responses, + completed_tool_calls, + account_changed, + ) { + let _ = replacement.stream.close(None).await; + return Err(err); + } + + pending.attempted_account_ids.clear(); + pending + .attempted_account_ids + .insert(replacement.account_id.clone()); + pending.log.route_strategy = Some(replacement.route_strategy.to_string()); + pending.log.route_source = Some(replacement.route_source.to_string()); + pending.conversation_routing = replacement.conversation_routing.clone(); + pending.buffer_retry_preamble = should_buffer_ws_retry_preamble( + &replacement, + &pending.attempted_account_ids, + pending.prepared.text.as_str(), + pending.retried_missing_tool_call_context, + ); + match replacement + .stream + .send(UpstreamMessage::Text(pending.prepared.text.clone().into())) + .await + { + Ok(()) => { + if attempt > 1 { + log::info!( + "event=responses_ws_reconnect_send_recovered attempt={} account_id={}", + attempt, + replacement.account_id, + ); + } + return Ok(replacement); + } + Err(err) => { + let account_id = replacement.account_id.clone(); + log::warn!( + "event=responses_ws_reconnect_send_failed attempt={} max_attempts={} account_id={} err={err}", + attempt, + RESPONSES_WS_MAX_PENDING_FRAME_SEND_ATTEMPTS, + account_id, + ); + let _ = replacement.stream.close(None).await; + last_send_error = Some(format!( + "send upstream websocket frame after reconnect failed for account {account_id}: {err}" + )); + previous_account_id = Some(account_id); + } + } } - Ok(replacement) + + Err(WsSessionError::bad_gateway_bilingual( + "重连后发送上游 WebSocket 帧失败", + last_send_error.unwrap_or_else(|| { + "send upstream websocket frame after reconnect failed after bounded retries".to_string() + }), + )) } async fn retry_pending_request_after_upstream_disconnect( @@ -2014,25 +2048,25 @@ async fn retry_pending_request_after_upstream_disconnect( completed_tool_calls: &CompletedWsToolCallCache, reason: &str, ) -> Result { - if pending.forwarded_non_preamble_event || pending.replayed_after_upstream_disconnect { + if pending.forwarded_non_preamble_event + || pending.upstream_disconnect_recovery_attempts + >= RESPONSES_WS_MAX_PRE_COMPLETION_RECOVERY_ATTEMPTS + { return Ok(false); } - // Once a continuation has emitted a preamble, replaying the same request - // can duplicate an already accepted turn. The only exception retained for - // compatibility is a request without previous_response_id, where the - // bounded preamble replay remains deduplicated before reaching the client. - if pending.prepared.previous_response_id.is_some() && pending.forwarded_upstream_event { - return Err(WsSessionError::context_rebase_failed( - "上游 WebSocket 已接受增量请求，无法安全重放；请使用新的 response.create 继续", - )); - } - let suppress_replayed_preamble = pending.forwarded_upstream_event; - pending.replayed_after_upstream_disconnect = true; + pending.upstream_disconnect_recovery_attempts += 1; pending.suppress_replayed_preamble = suppress_replayed_preamble; pending.buffered_upstream_preamble.clear(); let previous_account_id = upstream.account_id.clone(); + log::info!( + "event=responses_ws_pre_completion_recovery_attempt account_id={} attempt={} max_attempts={} reason={}", + previous_account_id, + pending.upstream_disconnect_recovery_attempts, + RESPONSES_WS_MAX_PRE_COMPLETION_RECOVERY_ATTEMPTS, + reason, + ); let replacement = reconnect_upstream_for_pending_request( context, pending, @@ -2075,7 +2109,7 @@ async fn wait_for_client_request_and_reconnect_upstream( conversation_routing: None, forwarded_upstream_event: false, forwarded_non_preamble_event: false, - replayed_after_upstream_disconnect: false, + upstream_disconnect_recovery_attempts: 0, suppress_replayed_preamble: false, buffered_upstream_preamble: Vec::new(), buffer_retry_preamble: false, @@ -3470,7 +3504,7 @@ fn inspect_ws_terminal_event(text: &str) -> Option { let is_websocket_connection_limit = is_websocket_connection_limit_error(error_code.as_deref(), error.as_deref()); match event_type.as_str() { - "response.completed" | "response.done" => Some(WsTerminalEvent { + "response.completed" => Some(WsTerminalEvent { status_code: 200, usage: parse_ws_usage(&value), error: None, diff --git a/crates/service/src/http/responses_websocket_rebase.rs b/crates/service/src/http/responses_websocket_rebase.rs index af4496c46..ebad9d39b 100644 --- a/crates/service/src/http/responses_websocket_rebase.rs +++ b/crates/service/src/http/responses_websocket_rebase.rs @@ -95,7 +95,7 @@ impl CompletedWsToolCallCache { self.insert(item); } } - "response.completed" | "response.done" => { + "response.completed" => { if let Some(items) = value .get("response") .and_then(|response| response.get("output")) @@ -146,7 +146,7 @@ impl CompletedWsResponseCache { .unwrap_or_default() .trim() .to_ascii_lowercase(); - if !matches!(event_type.as_str(), "response.completed" | "response.done") { + if event_type != "response.completed" { return Ok(false); } let response = terminal diff --git a/crates/service/src/http/responses_websocket_tests.rs b/crates/service/src/http/responses_websocket_tests.rs index 0d3310cdc..9cd4549bd 100644 --- a/crates/service/src/http/responses_websocket_tests.rs +++ b/crates/service/src/http/responses_websocket_tests.rs @@ -682,6 +682,14 @@ fn inspect_ws_terminal_event_maps_incomplete_to_terminal_error() { assert_eq!(event.usage.total_tokens, Some(14)); } +#[test] +fn inspect_ws_terminal_event_requires_response_completed() { + assert!( + inspect_ws_terminal_event(r#"{"type":"response.done","response":{"id":"resp_done"}}"#,) + .is_none() + ); +} + #[test] fn websocket_frame_aligns_prompt_cache_key_with_native_conversation_anchor() { let _guard = crate::test_env_guard(); diff --git a/crates/service/src/http/tests/proxy_runtime_tests.rs b/crates/service/src/http/tests/proxy_runtime_tests.rs index d4ec75e42..02f53acd7 100644 --- a/crates/service/src/http/tests/proxy_runtime_tests.rs +++ b/crates/service/src/http/tests/proxy_runtime_tests.rs @@ -803,6 +803,93 @@ async fn start_mock_upstream_ws_resets_before_first_frame() -> ( (addr.to_string(), event_rx, handle) } +async fn start_mock_upstream_ws_resets_twice_before_first_frame() -> ( + String, + tokio::sync::mpsc::UnboundedReceiver<(usize, String)>, + tokio::task::JoinHandle<()>, +) { + let listener = tokio::net::TcpListener::bind("127.0.0.1:0") + .await + .expect("bind double-initial-send-reset mock upstream"); + let addr = listener + .local_addr() + .expect("double-initial-send-reset mock upstream addr"); + let (event_tx, event_rx) = tokio::sync::mpsc::unbounded_channel(); + let handle = tokio::spawn(async move { + let upstream_config = WebSocketConfig::default() + .max_message_size(Some(TEST_LARGE_RESPONSES_WS_FRAME_BYTES * 2)) + .max_frame_size(Some(TEST_LARGE_RESPONSES_WS_FRAME_BYTES * 2)); + + for round in 1..=2 { + let (stream, _) = listener + .accept() + .await + .expect("accept double initial-send-reset upstream"); + let websocket = accept_hdr_async_with_config( + stream, + |_: &Request, response: Response| Ok(response), + Some(upstream_config), + ) + .await + .expect("accept double initial-send-reset websocket handshake"); + let raw_stream = websocket + .into_inner() + .into_std() + .expect("convert double initial-send-reset stream"); + force_tcp_reset(&raw_stream); + drop(raw_stream); + event_tx + .send((round, String::new())) + .expect("record double initial-send-reset connection"); + } + + let (replacement_stream, _) = + match tokio::time::timeout(Duration::from_secs(3), listener.accept()).await { + Ok(result) => result.expect("accept final double initial-send-reset upstream"), + Err(_) => { + event_tx + .send(( + 0, + "final replacement connection was not attempted".to_string(), + )) + .expect("record missing final replacement connection"); + return; + } + }; + let mut replacement = accept_hdr_async_with_config( + replacement_stream, + |_: &Request, response: Response| Ok(response), + Some(upstream_config), + ) + .await + .expect("accept final double initial-send-reset websocket handshake"); + let text = match replacement.next().await { + Some(Ok(Message::Text(text))) => text.to_string(), + other => panic!("expected final replacement response.create, got {other:?}"), + }; + event_tx + .send((3, text)) + .expect("record final double initial-send-reset request"); + for payload in [ + serde_json::json!({ + "type": "response.created", + "response": { "id": "resp_ws_double_initial_send_recovery" } + }), + serde_json::json!({ + "type": "response.completed", + "response": { "id": "resp_ws_double_initial_send_recovery" } + }), + ] { + replacement + .send(Message::Text(payload.to_string().into())) + .await + .expect("send double initial-send-reset recovery response"); + } + let _ = replacement.next().await; + }); + (addr.to_string(), event_rx, handle) +} + async fn start_mock_upstream_ws_holds_first_response() -> ( String, tokio::sync::mpsc::UnboundedReceiver, @@ -1080,6 +1167,76 @@ async fn start_mock_upstream_ws_resets_after_preamble() -> ( (addr.to_string(), event_rx, handle) } +async fn start_mock_upstream_ws_resets_after_preamble_twice() -> ( + String, + tokio::sync::mpsc::UnboundedReceiver<(usize, String)>, + tokio::task::JoinHandle<()>, +) { + let listener = tokio::net::TcpListener::bind("127.0.0.1:0") + .await + .expect("bind repeated reset-after-preamble mock upstream"); + let addr = listener + .local_addr() + .expect("repeated reset-after-preamble mock upstream addr"); + let (event_tx, event_rx) = tokio::sync::mpsc::unbounded_channel(); + let handle = tokio::spawn(async move { + for round in 0..=2 { + let (stream, _) = listener + .accept() + .await + .expect("accept repeated reset-after-preamble upstream"); + let mut websocket = + accept_hdr_async(stream, |_: &Request, response: Response| Ok(response)) + .await + .expect("accept repeated reset-after-preamble websocket handshake"); + let text = match websocket.next().await { + Some(Ok(Message::Text(text))) => text.to_string(), + other => panic!( + "expected repeated reset-after-preamble response.create for round {round}, got {other:?}" + ), + }; + event_tx + .send((round, text)) + .expect("record repeated reset-after-preamble frame"); + websocket + .send(Message::Text( + serde_json::json!({ + "type": "response.created", + "response": { "id": format!("resp_ws_repeated_reset_{round}") } + }) + .to_string() + .into(), + )) + .await + .expect("send repeated reset-after-preamble response.created"); + + if round < 2 { + let raw_stream = websocket + .into_inner() + .into_std() + .expect("convert repeated reset-after-preamble stream"); + force_tcp_reset(&raw_stream); + drop(raw_stream); + continue; + } + + websocket + .send(Message::Text( + serde_json::json!({ + "type": "response.completed", + "response": { "id": "resp_ws_repeated_reset_completed" } + }) + .to_string() + .into(), + )) + .await + .expect("send repeated reset-after-preamble response.completed"); + let _ = websocket.next().await; + } + }); + (addr.to_string(), event_rx, handle) +} + async fn start_mock_upstream_ws_connection_limit_then_success() -> ( String, tokio::sync::mpsc::UnboundedReceiver<(usize, String)>, @@ -2627,6 +2784,125 @@ async fn official_responses_websocket_recovers_after_initial_upstream_send_failu .expect("join initial-send recovery mock upstream"); } +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +async fn official_responses_websocket_retries_when_reconnected_socket_breaks_before_send() { + let _guard = crate::test_env_guard(); + let _http_proxy = EnvGuard::clear("http_proxy"); + let _https_proxy = EnvGuard::clear("https_proxy"); + let _all_proxy = EnvGuard::clear("all_proxy"); + let _upper_http_proxy = EnvGuard::clear("HTTP_PROXY"); + let _upper_https_proxy = EnvGuard::clear("HTTPS_PROXY"); + let _upper_all_proxy = EnvGuard::clear("ALL_PROXY"); + let _no_proxy = EnvGuard::set("NO_PROXY", "127.0.0.1,localhost"); + let _lower_no_proxy = EnvGuard::clear("no_proxy"); + let db_path = new_test_db_path("codexmanager-proxy-runtime-ws-double-initial-send-recovery"); + let storage = init_test_storage(&db_path); + let _db_guard = EnvGuard::set("CODEXMANAGER_DB_PATH", db_path.to_string_lossy().as_ref()); + let (upstream_addr, mut upstream_events, upstream_handle) = + start_mock_upstream_ws_resets_twice_before_first_frame().await; + insert_api_key_record( + &storage, + "platform_key_ws_double_initial_send_recovery", + crate::apikey_profile::ROTATION_ACCOUNT, + Some(format!( + "http://{upstream_addr}/chatgpt.com/backend-api/codex" + )), + ); + insert_account_and_token(&storage); + tokio::task::spawn_blocking(|| { + crate::gateway::reload_runtime_config_from_env(); + let _ = crate::gateway::front_proxy_max_body_bytes(); + }) + .await + .expect("reload runtime config"); + + let state = ProxyState { + backend_base_url: "http://127.0.0.1:1".to_string(), + client: Client::new(), + }; + let (front_addr, shutdown_tx, server_handle) = start_front_proxy_test_server(state).await; + let request = build_ws_request( + &format!("ws://{front_addr}/v1/responses"), + "platform_key_ws_double_initial_send_recovery", + &[("OpenAI-Beta", "responses_websockets=2026-02-06")], + ); + let (mut client_ws, response) = connect_async(request).await.expect("websocket connects"); + assert_eq!(response.status(), StatusCode::SWITCHING_PROTOCOLS); + + let image_data = "A".repeat(TEST_LARGE_RESPONSES_WS_FRAME_BYTES); + let payload = serde_json::json!({ + "type": "response.create", + "model": "gpt-5.6-sol", + "store": true, + "input": [{ + "type": "message", + "role": "user", + "content": [ + { "type": "input_text", "text": "recover after two pre-send websocket resets" }, + { + "type": "input_image", + "image_url": format!("data:image/png;base64,{image_data}") + } + ] + }] + }) + .to_string(); + assert!(payload.len() > 16 * 1024 * 1024); + client_ws + .send(Message::Text(payload.into())) + .await + .expect("send double-reset image-heavy response.create"); + + let mut forwarded = None; + for _ in 0..3 { + let (round, text) = tokio::time::timeout(Duration::from_secs(10), upstream_events.recv()) + .await + .expect("double-reset recovery frame timeout") + .expect("double-reset recovery frame channel"); + match round { + 1 | 2 => assert!(text.is_empty()), + 3 => { + forwarded = Some(text); + break; + } + 0 => { + panic!("double-reset recovery did not attempt a final upstream connection: {text}") + } + other => panic!("unexpected double-reset upstream round {other}"), + } + } + let forwarded = forwarded.expect("final replacement must receive response.create"); + assert!(forwarded.contains("recover after two pre-send websocket resets")); + assert!(forwarded.contains("data:image/png;base64,")); + + loop { + let event = tokio::time::timeout(Duration::from_secs(10), client_ws.next()) + .await + .expect("double-reset recovery response timeout") + .expect("double-reset recovery client event") + .expect("double-reset recovery client event result"); + match event { + Message::Text(text) if text.contains("\"response.completed\"") => break, + Message::Text(text) if text.contains("\"type\":\"error\"") => { + panic!("double-reset recovery error escaped to client: {text}"); + } + Message::Text(_) => {} + other => panic!("unexpected double-reset recovery event: {other:?}"), + } + } + + let _ = client_ws.close(None).await; + let _ = shutdown_tx.send(()); + tokio::time::timeout(Duration::from_secs(10), server_handle) + .await + .expect("front proxy double-reset recovery shutdown timeout") + .expect("join double-reset recovery front proxy"); + tokio::time::timeout(Duration::from_secs(10), upstream_handle) + .await + .expect("mock upstream double-reset recovery shutdown timeout") + .expect("join double-reset recovery mock upstream"); +} + #[tokio::test(flavor = "multi_thread", worker_threads = 4)] async fn official_responses_websocket_keeps_idle_session_alive_with_heartbeat() { let _guard = crate::test_env_guard(); @@ -2903,6 +3179,131 @@ async fn official_responses_websocket_replays_after_upstream_reset_after_preambl .expect("join reset-after-preamble mock upstream"); } +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +async fn official_responses_websocket_recovers_after_repeated_preamble_disconnects() { + let _guard = crate::test_env_guard(); + let _http_proxy = EnvGuard::clear("http_proxy"); + let _https_proxy = EnvGuard::clear("https_proxy"); + let _all_proxy = EnvGuard::clear("all_proxy"); + let _upper_http_proxy = EnvGuard::clear("HTTP_PROXY"); + let _upper_https_proxy = EnvGuard::clear("HTTPS_PROXY"); + let _upper_all_proxy = EnvGuard::clear("ALL_PROXY"); + let _no_proxy = EnvGuard::set("NO_PROXY", "127.0.0.1,localhost"); + let _lower_no_proxy = EnvGuard::clear("no_proxy"); + let db_path = new_test_db_path("codexmanager-proxy-runtime-ws-repeated-preamble-recovery"); + let storage = init_test_storage(&db_path); + let _db_guard = EnvGuard::set("CODEXMANAGER_DB_PATH", db_path.to_string_lossy().as_ref()); + let (upstream_addr, mut upstream_events, upstream_handle) = + start_mock_upstream_ws_resets_after_preamble_twice().await; + insert_api_key_record( + &storage, + "platform_key_ws_repeated_preamble_recovery", + crate::apikey_profile::ROTATION_ACCOUNT, + Some(format!( + "http://{upstream_addr}/chatgpt.com/backend-api/codex" + )), + ); + insert_account_and_token(&storage); + tokio::task::spawn_blocking(|| { + crate::gateway::reload_runtime_config_from_env(); + let _ = crate::gateway::front_proxy_max_body_bytes(); + }) + .await + .expect("reload runtime config"); + + let state = ProxyState { + backend_base_url: "http://127.0.0.1:1".to_string(), + client: Client::new(), + }; + let (front_addr, shutdown_tx, server_handle) = start_front_proxy_test_server(state).await; + let request = build_ws_request( + &format!("ws://{front_addr}/v1/responses"), + "platform_key_ws_repeated_preamble_recovery", + &[ + ("OpenAI-Beta", "responses_websockets=2026-02-06"), + ("session_id", "session_ws_repeated_preamble_recovery"), + ], + ); + let (mut client_ws, response) = connect_async(request).await.expect("websocket connects"); + assert_eq!(response.status(), StatusCode::SWITCHING_PROTOCOLS); + + client_ws + .send(Message::Text( + serde_json::json!({ + "type": "response.create", + "model": "gpt-4.1", + "store": true, + "previous_response_id": "resp_prior_for_repeated_preamble", + "input": [{ + "type": "message", + "role": "user", + "content": [{ + "type": "input_text", + "text": "continue after repeated preamble disconnects" + }] + }] + }) + .to_string() + .into(), + )) + .await + .expect("send repeated preamble continuation"); + + for round in 0..=2 { + let (upstream_round, upstream_text) = + tokio::time::timeout(Duration::from_secs(10), upstream_events.recv()) + .await + .expect("repeated preamble upstream frame timeout") + .expect("repeated preamble upstream frame channel"); + assert_eq!(upstream_round, round); + assert!(upstream_text.contains("continue after repeated preamble disconnects")); + assert!(upstream_text.contains("resp_prior_for_repeated_preamble")); + } + + let mut created_events = 0; + loop { + let event = tokio::time::timeout(Duration::from_secs(10), client_ws.next()) + .await + .expect("repeated preamble client event timeout") + .expect("repeated preamble client event") + .expect("repeated preamble client event result"); + match event { + Message::Text(text) if text.contains("\"response.created\"") => { + created_events += 1; + } + Message::Text(text) if text.contains("\"response.completed\"") => break, + Message::Text(text) if text.contains("\"type\":\"error\"") => { + panic!("repeated preamble recovery error escaped to client: {text}"); + } + Message::Text(_) => {} + other => panic!("unexpected repeated preamble event: {other:?}"), + } + } + assert_eq!(created_events, 1, "replayed preambles must stay suppressed"); + + tokio::time::sleep(Duration::from_millis(100)).await; + let request_logs = storage + .list_request_logs(None, 10) + .expect("list repeated preamble request logs"); + let ws_logs = request_logs + .iter() + .filter(|item| item.request_type.as_deref() == Some("ws")) + .collect::>(); + assert_eq!(ws_logs.len(), 1); + assert_eq!(ws_logs[0].status_code, Some(200)); + + let _ = client_ws.close(None).await; + let _ = shutdown_tx.send(()); + tokio::time::timeout(Duration::from_secs(10), server_handle) + .await + .expect("front proxy repeated preamble shutdown timeout") + .expect("join repeated preamble front proxy"); + tokio::time::timeout(Duration::from_secs(10), upstream_handle) + .await + .expect("mock upstream repeated preamble shutdown timeout") + .expect("join repeated preamble mock upstream"); +} + #[tokio::test(flavor = "multi_thread", worker_threads = 4)] async fn official_responses_websocket_recovers_after_connection_limit_error() { let _guard = crate::test_env_guard(); diff --git a/crates/service/src/usage/usage_token_refresh.rs b/crates/service/src/usage/usage_token_refresh.rs index 411a05f7f..70320cc1d 100644 --- a/crates/service/src/usage/usage_token_refresh.rs +++ b/crates/service/src/usage/usage_token_refresh.rs @@ -97,7 +97,12 @@ pub(crate) fn refresh_and_persist_access_token( if let Some(id_token) = refreshed.id_token { token.id_token = id_token.clone(); - let exchange_client_id = token_refresh_client_id(token, refresh_client_id.as_str()); + // The refresh grant uses the access-token client id, while the API-key + // exchange uses the newly issued ID token as its subject. Keep the + // two client-id rules separate so an access-token audience cannot + // cause an ID-token exchange to be rejected. + let exchange_client_id = + crate::gateway::api_key_exchange_client_id(token, refresh_client_id.as_str()); if let Ok(api_key) = obtain_api_key(issuer, &exchange_client_id, &id_token) { token.api_key_access_token = Some(api_key); }