diff --git a/crates/service/src/gateway/mod.rs b/crates/service/src/gateway/mod.rs index 0f31972e3..69edf7a45 100644 --- a/crates/service/src/gateway/mod.rs +++ b/crates/service/src/gateway/mod.rs @@ -83,7 +83,7 @@ pub(crate) fn error_message_for_client( mod anchor_fingerprint; mod concurrency; #[path = "routing/conversation_binding.rs"] -mod conversation_binding; +pub(crate) mod conversation_binding; #[path = "routing/cooldown.rs"] mod cooldown; mod error_response; @@ -1086,6 +1086,7 @@ pub(crate) struct GatewayRoutedCandidates { )>, pub(crate) route_strategy: &'static str, pub(crate) route_source: &'static str, + pub(crate) conversation_routing: Option, } pub(crate) fn gateway_collect_routed_candidates_with_log_source( @@ -1112,9 +1113,124 @@ pub(crate) fn gateway_collect_routed_candidates_with_log_source( candidates, route_strategy: application.strategy_label, route_source: application.source, + conversation_routing: None, }) } +pub(crate) fn gateway_collect_routed_candidates_for_ws( + storage: &codexmanager_core::storage::Storage, + key_id: &str, + model: Option<&str>, + route_conversation_id: Option<&str>, + route_conversation_source: Option, +) -> Result { + let api_key = storage + .find_api_key_by_id(key_id) + .map_err(|err| format!("read api key routing config failed: {err}"))? + .ok_or_else(|| "api key not found".to_string())?; + let account_group_filter = storage + .find_api_key_account_group_filter(key_id) + .map_err(|err| format!("read api key account group filter failed: {err}"))?; + let mut candidates = upstream::support::candidates::prepare_gateway_candidates( + storage, + model, + account_group_filter.as_deref(), + api_key.account_plan_filter.as_deref(), + LowQuotaCandidateMode::NormalOnly, + )?; + + // HTTP candidate execution treats runtime cooldown as a hard skip whenever another + // candidate exists. A persistent WebSocket must not keep using the current account after + // a 429/401/403 cooldown is recorded, so the WS pool excludes all cooled-down accounts + // before routing and failover. + candidates.retain(|(account, _)| !is_account_in_cooldown(account.id.as_str())); + + let conversation_binding = match route_conversation_id { + Some(conversation_id) => conversation_binding::load_conversation_binding( + storage, + api_key.key_hash.as_str(), + Some(conversation_id), + )?, + None => None, + }; + let conversation_routing = route_conversation_source.and_then(|source| { + conversation_binding::prepare_conversation_routing_with_source( + api_key.key_hash.as_str(), + route_conversation_id, + conversation_binding.as_ref(), + &mut candidates, + source, + ) + }); + let account_binding_counts = if thread_aware_account_distribution_enabled() + && conversation_routing.as_ref().is_some_and(|routing| { + routing.existing_binding.is_none() && routing.source.allows_initial_binding_create() + }) { + match storage.active_conversation_binding_account_counts(api_key.key_hash.as_str()) { + Ok(counts) => Some(counts), + Err(err) => { + log::warn!("load conversation binding account counts for websocket failed: {err}"); + None + } + } + } else { + None + }; + let rotation_plan = conversation_binding::apply_candidate_rotation( + &mut candidates, + conversation_routing.as_ref(), + key_id, + model, + account_binding_counts.as_ref(), + ); + Ok(GatewayRoutedCandidates { + candidates, + route_strategy: rotation_plan.strategy_label, + route_source: rotation_plan.source.as_str(), + conversation_routing, + }) +} + +pub(crate) fn gateway_ws_account_requires_switch( + routed: &GatewayRoutedCandidates, + current_account_id: &str, +) -> bool { + if !routed + .candidates + .iter() + .any(|(account, _)| account.id == current_account_id) + { + return true; + } + + // A turn-state-only or otherwise unbound WebSocket has no conversation routing context, + // but an explicitly preferred account still applies to it. The routed source is already + // computed from the fresh preference snapshot above, so do not let the persistent socket + // keep using a formerly selected non-preferred account. + if routed.route_source == "manual_preferred_account" { + return routed + .candidates + .first() + .is_some_and(|(account, _)| account.id != current_account_id); + } + + let Some(routing) = routed.conversation_routing.as_ref() else { + return false; + }; + if routing + .manual_preferred_account_id + .as_deref() + .is_some_and(|account_id| account_id != current_account_id) + { + return true; + } + routing.bound_account_selectable + && routing + .existing_binding + .as_ref() + .is_some_and(|binding| binding.account_id != current_account_id) +} + /// 函数 `gateway_record_failover_attempt` /// /// 作者: gaohongshun @@ -1179,11 +1295,22 @@ pub(crate) fn gateway_token_exchange_client_id() -> String { runtime_config::token_exchange_client_id() } +pub(crate) struct GatewayWsRoutingPreparation { + pub(crate) incoming_headers: IncomingHeaderSnapshot, + pub(crate) prompt_cache_key: Option, + pub(crate) route_conversation_id: Option, + pub(crate) route_conversation_source: Option, +} + pub(crate) fn gateway_resolve_ws_prompt_cache_key( storage: &codexmanager_core::storage::Storage, api_key: &codexmanager_core::storage::ApiKey, incoming_headers: &IncomingHeaderSnapshot, -) -> Result<(IncomingHeaderSnapshot, Option), String> { +) -> Result { + let has_native_conversation_id = incoming_headers + .conversation_id() + .map(str::trim) + .is_some_and(|value| !value.is_empty()); let local_conversation_id = resolve_local_conversation_id_with_sticky_fallback(incoming_headers, true); let conversation_binding = conversation_binding::load_conversation_binding( @@ -1198,7 +1325,19 @@ pub(crate) fn gateway_resolve_ws_prompt_cache_key( local_conversation_id.as_deref(), conversation_binding.as_ref(), ); - Ok((incoming_headers, prompt_cache_key)) + let route_conversation_source = if has_native_conversation_id { + Some(conversation_binding::RouteConversationSource::NativeConversation) + } else if local_conversation_id.is_some() { + Some(conversation_binding::RouteConversationSource::StickyFallback) + } else { + None + }; + Ok(GatewayWsRoutingPreparation { + incoming_headers, + prompt_cache_key, + route_conversation_id: local_conversation_id, + route_conversation_source, + }) } /// 函数 `gateway_rewrite_ws_responses_body` diff --git a/crates/service/src/http/responses_websocket.rs b/crates/service/src/http/responses_websocket.rs index 6753c64c2..f7e301700 100644 --- a/crates/service/src/http/responses_websocket.rs +++ b/crates/service/src/http/responses_websocket.rs @@ -32,7 +32,8 @@ use crate::storage_helpers::open_storage; mod responses_websocket_rebase; use responses_websocket_rebase::{ - rebase_response_create_for_account_change, rebase_response_create_for_missing_tool_call, + expand_response_create_previous_response, rebase_response_create_for_account_change, + rebase_response_create_for_missing_tool_call, CompletedWsResponseCache, CompletedWsToolCallCache, WsToolCallKind, }; @@ -56,6 +57,9 @@ struct WsRequestContext { api_key: codexmanager_core::storage::ApiKey, incoming_headers: crate::gateway::IncomingHeaderSnapshot, prompt_cache_key: Option, + route_conversation_id: Option, + route_conversation_source: + Option, effective_upstream_base: String, prefer_raw_errors: bool, } @@ -63,6 +67,7 @@ struct WsRequestContext { #[derive(Clone)] struct PreparedClientFrame { text: String, + input: Value, client_model: Option, model: Option, previous_response_id: Option, @@ -81,6 +86,7 @@ struct PreparedClientFrame { struct PendingWsRequestState { log: PendingWsRequestLog, prepared: PreparedClientFrame, + conversation_routing: Option, forwarded_upstream_event: bool, forwarded_non_preamble_event: bool, replayed_after_upstream_disconnect: bool, @@ -97,6 +103,8 @@ type UpstreamWebsocketStream = struct ConnectedUpstreamWebsocket { stream: UpstreamWebsocketStream, account_id: String, + account: codexmanager_core::storage::Account, + conversation_routing: Option, candidate_account_ids: HashSet, upstream_url: String, route_strategy: &'static str, @@ -340,31 +348,35 @@ async fn run_responses_websocket_session(mut socket: WebSocket, context: WsReque "unresolved", "initial_upstream_connect", ); - let mut upstream = - match connect_upstream_websocket_with_timeout(&context, prepared_first.model.as_deref()) - .await - { - Ok(stream) => stream, - Err(err) => { - finalize_ws_request_log( - &context, - &first_log, - None, - None, - err.status, - crate::gateway::RequestLogUsage::default(), - Some(err.message.clone()), - ); - send_ws_error_and_close(&mut socket, err, context.prefer_raw_errors).await; - return; - } - }; + let mut upstream = match connect_upstream_websocket_with_timeout( + &context, + prepared_first.model.as_deref(), + None, + ) + .await + { + Ok(stream) => stream, + Err(err) => { + finalize_ws_request_log( + &context, + &first_log, + None, + None, + err.status, + crate::gateway::RequestLogUsage::default(), + Some(err.message.clone()), + ); + send_ws_error_and_close(&mut socket, err, context.prefer_raw_errors).await; + return; + } + }; first_log.route_strategy = Some(upstream.route_strategy.to_string()); first_log.route_source = Some(upstream.route_source.to_string()); let first_attempted_account_ids = HashSet::from([upstream.account_id.clone()]); let mut first_pending = PendingWsRequestState { log: first_log, prepared: prepared_first.clone(), + conversation_routing: upstream.conversation_routing.clone(), forwarded_upstream_event: false, forwarded_non_preamble_event: false, replayed_after_upstream_disconnect: false, @@ -381,6 +393,7 @@ async fn run_responses_websocket_session(mut socket: WebSocket, context: WsReque }; let mut completed_tool_calls = CompletedWsToolCallCache::default(); + let mut completed_responses = CompletedWsResponseCache::default(); if let Err(err) = upstream .stream .send(UpstreamMessage::Text( @@ -399,6 +412,7 @@ async fn run_responses_websocket_session(mut socket: WebSocket, context: WsReque &context, &mut first_pending, Some(previous_account_id.as_str()), + &completed_responses, &completed_tool_calls, ) .await @@ -456,6 +470,50 @@ async fn run_responses_websocket_session(mut socket: WebSocket, context: WsReque upstream.account_id, ); } + if pending_request.is_none() { + match ws_account_requires_reselection(&context, &upstream, None) { + Ok(true) => { + log::info!( + "event=responses_ws_idle_account_invalidated account_id={} reason=account_eligibility_changed", + upstream.account_id, + ); + let _ = upstream.stream.close(None).await; + } + Ok(false) => {} + Err(err) => { + log::warn!( + "event=responses_ws_idle_account_check_failed account_id={} err={}", + upstream.account_id, + err.message, + ); + } + } + } else { + match ws_account_is_unavailable( + &context, + &upstream, + pending_request + .as_ref() + .and_then(|pending| pending.prepared.model.as_deref()), + ) { + Ok(true) => { + log::info!( + "event=responses_ws_inflight_account_invalidated account_id={} reason=account_no_longer_selectable", + upstream.account_id, + ); + let _ = upstream.stream.close(None).await; + } + Ok(false) => {} + Err(err) => { + log::warn!( + "event=responses_ws_inflight_account_check_failed account_id={} err={}", + upstream.account_id, + err.message, + ); + let _ = upstream.stream.close(None).await; + } + } + } } maybe_client = socket.recv() => { let Some(client_result) = maybe_client else { @@ -528,6 +586,7 @@ async fn run_responses_websocket_session(mut socket: WebSocket, context: WsReque upstream.route_source, ), prepared, + conversation_routing: upstream.conversation_routing.clone(), forwarded_upstream_event: false, forwarded_non_preamble_event: false, replayed_after_upstream_disconnect: false, @@ -537,9 +596,97 @@ async fn run_responses_websocket_session(mut socket: WebSocket, context: WsReque attempted_account_ids, retried_missing_tool_call_context: false, }; - if let Err(send_err) = upstream.stream.send(UpstreamMessage::Text( - current_pending.prepared.text.clone().into(), - )).await { + let (must_reselect_account, fresh_conversation_routing, fresh_route_strategy, fresh_route_source) = match ws_collect_routed_candidates( + &context, + current_pending.prepared.model.as_deref(), + ) { + Ok(routed) => ( + crate::gateway::gateway_ws_account_requires_switch( + &routed, + upstream.account_id.as_str(), + ), + routed.conversation_routing, + routed.route_strategy, + routed.route_source, + ), + Err(err) => { + finalize_ws_request_log( + &context, + ¤t_pending.log, + Some(upstream.account_id.as_str()), + Some(upstream.upstream_url.as_str()), + err.status, + crate::gateway::RequestLogUsage::default(), + Some(err.message.clone()), + ); + send_ws_error_and_close( + &mut socket, + err, + context.prefer_raw_errors, + ) + .await; + break; + } + }; + if !must_reselect_account { + upstream.conversation_routing = + fresh_conversation_routing.clone(); + upstream.route_strategy = fresh_route_strategy; + upstream.route_source = fresh_route_source; + current_pending.conversation_routing = + fresh_conversation_routing; + current_pending.log.route_strategy = + Some(fresh_route_strategy.to_string()); + current_pending.log.route_source = + Some(fresh_route_source.to_string()); + } + if must_reselect_account { + let previous_account_id = upstream.account_id.clone(); + log::info!( + "event=responses_ws_account_reselect_before_turn account_id={} model={} reason=account_eligibility_changed", + previous_account_id, + current_pending.prepared.model.as_deref().unwrap_or("-"), + ); + let _ = upstream.stream.close(None).await; + match reconnect_upstream_for_pending_request( + &context, + &mut current_pending, + Some(previous_account_id.as_str()), + &completed_responses, + &completed_tool_calls, + ) + .await + { + Ok(replacement) => { + log::info!( + "event=responses_ws_account_reselected_before_turn_complete previous_account_id={} account_id={}", + previous_account_id, + replacement.account_id, + ); + upstream = replacement; + } + Err(err) => { + finalize_ws_request_log( + &context, + ¤t_pending.log, + None, + None, + err.status, + crate::gateway::RequestLogUsage::default(), + Some(err.message.clone()), + ); + send_ws_error_and_close( + &mut socket, + err, + context.prefer_raw_errors, + ) + .await; + break; + } + } + } else if let Err(send_err) = upstream.stream.send( + UpstreamMessage::Text(current_pending.prepared.text.clone().into()), + ).await { let previous_account_id = upstream.account_id.clone(); log::warn!( "event=responses_ws_upstream_stale_send account_id={} err={send_err}", @@ -550,6 +697,7 @@ async fn run_responses_websocket_session(mut socket: WebSocket, context: WsReque &context, &mut current_pending, Some(previous_account_id.as_str()), + &completed_responses, &completed_tool_calls, ) .await @@ -636,6 +784,7 @@ async fn run_responses_websocket_session(mut socket: WebSocket, context: WsReque pending_request .as_mut() .expect("pending request checked above"), + &completed_responses, &completed_tool_calls, "early_eof", ) @@ -693,6 +842,7 @@ async fn run_responses_websocket_session(mut socket: WebSocket, context: WsReque &mut socket, &context, previous_account_id.as_str(), + &completed_responses, &completed_tool_calls, ) .await @@ -726,6 +876,7 @@ async fn run_responses_websocket_session(mut socket: WebSocket, context: WsReque &context, &mut upstream, pending, + &completed_responses, &completed_tool_calls, "connection_limit_reached", ) @@ -769,6 +920,7 @@ async fn run_responses_websocket_session(mut socket: WebSocket, context: WsReque &mut socket, &context, previous_account_id.as_str(), + &completed_responses, &completed_tool_calls, ) .await @@ -807,6 +959,7 @@ async fn run_responses_websocket_session(mut socket: WebSocket, context: WsReque &mut upstream, pending, &terminal, + &completed_responses, &completed_tool_calls, ) .await @@ -844,11 +997,34 @@ async fn run_responses_websocket_session(mut socket: WebSocket, context: WsReque } if let Some(mut pending) = pending_request.take() { + if terminal.status_code == 200 { + cache_completed_ws_response( + &mut completed_responses, + &pending.prepared, + text.as_str(), + ); + } if let Err(err) = flush_ws_upstream_preamble(&mut socket, &mut pending).await { log::warn!("event=responses_ws_client_send_preamble_failed err={err}"); break; } mark_ws_first_response(&mut pending); + if let Some(storage) = open_storage() { + if let Err(err) = crate::gateway::conversation_binding::record_conversation_binding_terminal_response( + &storage, + pending.conversation_routing.as_ref(), + &upstream.account, + pending.prepared.model.as_deref(), + terminal.status_code, + ) { + log::warn!( + "event=responses_ws_conversation_binding_record_failed account_id={} status={} err={}", + upstream.account_id, + terminal.status_code, + err, + ); + } + } finalize_ws_request_log( &context, &pending.log, @@ -929,6 +1105,7 @@ async fn run_responses_websocket_session(mut socket: WebSocket, context: WsReque pending_request .as_mut() .expect("pending request checked above"), + &completed_responses, &completed_tool_calls, "early_close", ) @@ -986,6 +1163,7 @@ async fn run_responses_websocket_session(mut socket: WebSocket, context: WsReque &mut socket, &context, previous_account_id.as_str(), + &completed_responses, &completed_tool_calls, ) .await @@ -1016,6 +1194,7 @@ async fn run_responses_websocket_session(mut socket: WebSocket, context: WsReque pending_request .as_mut() .expect("pending request checked above"), + &completed_responses, &completed_tool_calls, "early_read_error", ) @@ -1073,6 +1252,7 @@ async fn run_responses_websocket_session(mut socket: WebSocket, context: WsReque &mut socket, &context, previous_account_id.as_str(), + &completed_responses, &completed_tool_calls, ) .await @@ -1143,7 +1323,7 @@ fn authorize_websocket_request(headers: &HeaderMap) -> Result Result, +) -> Result { + let storage = open_storage().ok_or_else(|| { + WsSessionError::service_unavailable_bilingual("存储不可用", "storage unavailable") + })?; + crate::gateway::gateway_collect_routed_candidates_for_ws( + &storage, + &context.api_key.id, + model, + context.route_conversation_id.as_deref(), + context.route_conversation_source, + ) + .map_err(|err| { + WsSessionError::service_unavailable_bilingual( + "读取 WebSocket 账号候选失败", + format!("read websocket account candidates failed: {err}"), + ) + }) +} + +fn ws_account_requires_reselection( + context: &WsRequestContext, + upstream: &ConnectedUpstreamWebsocket, + model: Option<&str>, +) -> Result { + let routed = ws_collect_routed_candidates(context, model)?; + Ok(crate::gateway::gateway_ws_account_requires_switch( + &routed, + upstream.account_id.as_str(), + )) +} + +fn ws_account_is_unavailable( + context: &WsRequestContext, + upstream: &ConnectedUpstreamWebsocket, + model: Option<&str>, +) -> Result { + let routed = ws_collect_routed_candidates(context, model)?; + Ok(!routed + .candidates + .iter() + .any(|(account, _)| account.id == upstream.account_id)) +} + async fn connect_upstream_websocket( context: &WsRequestContext, model: Option<&str>, + previous_account_id: Option<&str>, ) -> Result { let storage = open_storage().ok_or_else(|| { WsSessionError::service_unavailable_bilingual("存储不可用", "storage unavailable") })?; - let routed = crate::gateway::gateway_collect_routed_candidates_with_log_source( + let routed = crate::gateway::gateway_collect_routed_candidates_for_ws( &storage, &context.api_key.id, model, + context.route_conversation_id.as_deref(), + context.route_conversation_source, )?; if routed.candidates.is_empty() { return Err(WsSessionError::service_unavailable_bilingual( @@ -1649,18 +1881,27 @@ async fn connect_upstream_websocket( .iter() .map(|(account, _)| account.id.clone()) .collect::>(); + let conversation_routing = routed.conversation_routing.clone(); drop(storage); let ws_url = build_upstream_websocket_url(&context.effective_upstream_base)?; let mut last_error = None; for (account, token) in routed.candidates { - match connect_account_upstream_websocket(context, &account, token, ws_url.as_str(), false) - .await + match connect_account_upstream_websocket( + context, + &account, + token, + ws_url.as_str(), + previous_account_id.is_some_and(|account_id| account_id != account.id), + ) + .await { Ok(stream) => { return Ok(ConnectedUpstreamWebsocket { stream, - account_id: account.id, + account_id: account.id.clone(), + account, + conversation_routing, candidate_account_ids, upstream_url: ws_url.clone(), route_strategy: routed.route_strategy, @@ -1685,10 +1926,16 @@ async fn connect_upstream_websocket( async fn connect_upstream_websocket_with_timeout( context: &WsRequestContext, model: Option<&str>, + previous_account_id: Option<&str>, ) -> Result { let connect_timeout = crate::gateway::current_upstream_connect_timeout().max(std::time::Duration::from_secs(1)); - match tokio::time::timeout(connect_timeout, connect_upstream_websocket(context, model)).await { + match tokio::time::timeout( + connect_timeout, + connect_upstream_websocket(context, model, previous_account_id), + ) + .await + { Ok(result) => result, Err(_) => Err(WsSessionError::new( 504, @@ -1708,26 +1955,25 @@ async fn reconnect_upstream_for_pending_request( context: &WsRequestContext, pending: &mut PendingWsRequestState, previous_account_id: Option<&str>, + completed_responses: &CompletedWsResponseCache, completed_tool_calls: &CompletedWsToolCallCache, ) -> Result { - if pending.prepared.previous_response_id.is_some() && !pending.prepared.store { - return Err(WsSessionError::context_rebase_failed( - "无法在新 WebSocket 连接中恢复 store=false 的 previous_response_id；请重新发送完整上下文", - )); - } - let mut replacement = - connect_upstream_websocket_with_timeout(context, pending.prepared.model.as_deref()).await?; - if previous_account_id.is_some_and(|account_id| account_id != replacement.account_id) { - match rebase_ws_request_for_account_change( - pending.prepared.text.as_str(), - completed_tool_calls, - ) { - Ok(rebased) => pending.prepared.text = rebased, - Err(err) => { - let _ = replacement.stream.close(None).await; - return Err(err); - } - } + 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(); @@ -1736,6 +1982,7 @@ async fn reconnect_upstream_for_pending_request( .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, @@ -1763,6 +2010,7 @@ async fn retry_pending_request_after_upstream_disconnect( context: &WsRequestContext, upstream: &mut ConnectedUpstreamWebsocket, pending: &mut PendingWsRequestState, + completed_responses: &CompletedWsResponseCache, completed_tool_calls: &CompletedWsToolCallCache, reason: &str, ) -> Result { @@ -1770,17 +2018,6 @@ async fn retry_pending_request_after_upstream_disconnect( return Ok(false); } - // A continuation sent with store=false depends on the old upstream - // connection's in-memory previous-response cache. A replacement upstream - // connection cannot safely replay only that incremental input. Let the - // client establish a new connection and provide the full context instead - // of silently creating a partial or duplicated continuation. - if pending.prepared.previous_response_id.is_some() && !pending.prepared.store { - return Err(WsSessionError::context_rebase_failed( - "无法安全恢复 store=false 的增量 WebSocket 请求；请在新连接中重新发送完整上下文", - )); - } - // 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 @@ -1800,6 +2037,7 @@ async fn retry_pending_request_after_upstream_disconnect( context, pending, Some(previous_account_id.as_str()), + completed_responses, completed_tool_calls, ) .await?; @@ -1817,6 +2055,7 @@ async fn wait_for_client_request_and_reconnect_upstream( socket: &mut WebSocket, context: &WsRequestContext, previous_account_id: &str, + completed_responses: &CompletedWsResponseCache, completed_tool_calls: &CompletedWsToolCallCache, ) -> Result, WsSessionError> { let Some(text) = receive_initial_request(socket).await? else { @@ -1833,6 +2072,7 @@ async fn wait_for_client_request_and_reconnect_upstream( let mut pending = PendingWsRequestState { log: begin_ws_request_log(context, &prepared, "unresolved", "upstream_reconnect"), prepared, + conversation_routing: None, forwarded_upstream_event: false, forwarded_non_preamble_event: false, replayed_after_upstream_disconnect: false, @@ -1846,6 +2086,7 @@ async fn wait_for_client_request_and_reconnect_upstream( context, &mut pending, Some(previous_account_id), + completed_responses, completed_tool_calls, ) .await @@ -2858,11 +3099,72 @@ fn finalize_ws_request_log( fn ws_context_rebase_error(raw_message: impl Into) -> WsSessionError { let raw_message = raw_message.into(); WsSessionError::context_rebase_failed(crate::gateway::bilingual_error( - "切换账号时无法重建工具调用上下文", + "无法为新的上游 WebSocket 重建完整上下文", raw_message, )) } +fn prepare_ws_request_for_new_connection( + pending: &mut PendingWsRequestState, + completed_responses: &CompletedWsResponseCache, + completed_tool_calls: &CompletedWsToolCallCache, + account_changed: bool, +) -> Result<(), WsSessionError> { + if pending.prepared.previous_response_id.is_some() + && (!pending.prepared.store || account_changed) + { + let expanded = expand_response_create_previous_response( + pending.prepared.text.as_str(), + completed_responses, + ) + .map_err(ws_context_rebase_error)? + .ok_or_else(|| { + ws_context_rebase_error( + "response.create declared previous_response_id but no recoverable id was found", + ) + })?; + pending.prepared.text = expanded; + pending.prepared.previous_response_id = None; + pending.prepared.input = ws_request_input_from_text(pending.prepared.text.as_str())?; + } + if account_changed { + pending.prepared.text = rebase_ws_request_for_account_change( + pending.prepared.text.as_str(), + completed_tool_calls, + )?; + pending.prepared.input = ws_request_input_from_text(pending.prepared.text.as_str())?; + } + Ok(()) +} + +fn ws_request_input_from_text(text: &str) -> Result { + serde_json::from_str::(text) + .map_err(|err| { + ws_context_rebase_error(format!("parse recovered response.create failed: {err}")) + })? + .get("input") + .cloned() + .ok_or_else(|| ws_context_rebase_error("recovered response.create is missing input")) +} + +fn cache_completed_ws_response( + completed_responses: &mut CompletedWsResponseCache, + prepared: &PreparedClientFrame, + terminal_text: &str, +) { + if let Err(err) = completed_responses.observe_completed_response( + terminal_text, + prepared.previous_response_id.as_deref(), + &prepared.input, + ) { + log::warn!( + "event=responses_ws_response_history_not_cached previous_response_id={} err={}", + prepared.previous_response_id.as_deref().unwrap_or("-"), + err, + ); + } +} + fn rebase_ws_request_for_account_change( text: &str, completed_tool_calls: &CompletedWsToolCallCache, @@ -2940,6 +3242,7 @@ async fn try_retry_ws_request_after_terminal( upstream: &mut ConnectedUpstreamWebsocket, pending: &mut PendingWsRequestState, terminal: &WsTerminalEvent, + completed_responses: &CompletedWsResponseCache, completed_tool_calls: &CompletedWsToolCallCache, ) -> Result { if terminal.status_code == 200 || pending.forwarded_upstream_event { @@ -2961,15 +3264,20 @@ async fn try_retry_ws_request_after_terminal( if strip_previous_response_id_from_ws_text(pending.prepared.text.as_str()).is_none() { return Ok(false); } - if pending.prepared.previous_response_id.is_some() && !pending.prepared.store { - return Err(WsSessionError::context_rebase_failed( - "previous_response_id 在新 WebSocket 连接中不可用；store=false 时必须重新发送完整上下文", - )); - } - retry_text = Some(rebase_ws_request_for_account_change( + let expanded = expand_response_create_previous_response( pending.prepared.text.as_str(), - completed_tool_calls, - )?); + completed_responses, + ) + .map_err(ws_context_rebase_error)? + .ok_or_else(|| { + ws_context_rebase_error( + "previous_response_id was rejected and no recoverable response id was found", + ) + })?; + pending.prepared.text = expanded; + pending.prepared.previous_response_id = None; + pending.prepared.input = ws_request_input_from_text(pending.prepared.text.as_str())?; + retry_text = Some(pending.prepared.text.clone()); } else { let previous_account_id = upstream.account_id.clone(); if !try_rotate_ws_upstream_after_terminal( @@ -2984,15 +3292,20 @@ async fn try_retry_ws_request_after_terminal( return Ok(false); } if upstream.account_id != previous_account_id { - retry_text = Some(rebase_ws_request_for_account_change( - pending.prepared.text.as_str(), + prepare_ws_request_for_new_connection( + pending, + completed_responses, completed_tool_calls, - )?); + true, + )?; + retry_text = Some(pending.prepared.text.clone()); pending.log.route_strategy = Some(upstream.route_strategy.to_string()); pending.log.route_source = Some(upstream.route_source.to_string()); + pending.conversation_routing = upstream.conversation_routing.clone(); } } let retry_text = retry_text.unwrap_or_else(|| pending.prepared.text.clone()); + let retry_input = ws_request_input_from_text(retry_text.as_str())?; match upstream .stream .send(UpstreamMessage::Text(retry_text.clone().into())) @@ -3000,6 +3313,7 @@ async fn try_retry_ws_request_after_terminal( { Ok(()) => { pending.prepared.text = retry_text; + pending.prepared.input = retry_input; pending.forwarded_upstream_event = false; pending.buffered_upstream_preamble.clear(); pending.buffer_retry_preamble = should_buffer_ws_retry_preamble( @@ -3042,10 +3356,12 @@ async fn try_rotate_ws_upstream_after_terminal( Some(storage) => storage, None => return false, }; - let routed = match crate::gateway::gateway_collect_routed_candidates_with_log_source( + let routed = match crate::gateway::gateway_collect_routed_candidates_for_ws( &storage, &context.api_key.id, model, + context.route_conversation_id.as_deref(), + context.route_conversation_source, ) { Ok(routed) => routed, Err(err) => { @@ -3061,6 +3377,7 @@ async fn try_rotate_ws_upstream_after_terminal( let route_strategy = routed.route_strategy; let route_source = routed.route_source; let candidates = routed.candidates; + let conversation_routing = routed.conversation_routing.clone(); let candidate_account_ids = candidates .iter() .map(|(account, _)| account.id.clone()) @@ -3083,7 +3400,9 @@ async fn try_rotate_ws_upstream_after_terminal( Ok(stream) => { let replacement = ConnectedUpstreamWebsocket { stream, - account_id: account.id, + account_id: account.id.clone(), + account, + conversation_routing, candidate_account_ids, upstream_url: upstream.upstream_url.clone(), route_strategy, diff --git a/crates/service/src/http/responses_websocket_rebase.rs b/crates/service/src/http/responses_websocket_rebase.rs index a2b386390..af4496c46 100644 --- a/crates/service/src/http/responses_websocket_rebase.rs +++ b/crates/service/src/http/responses_websocket_rebase.rs @@ -2,6 +2,8 @@ use serde_json::Value; use std::collections::{HashMap, HashSet, VecDeque}; const MAX_CACHED_WS_TOOL_CALLS: usize = 256; +const MAX_CACHED_WS_RESPONSES: usize = 64; +const MAX_CACHED_WS_RESPONSE_BYTES: usize = 64 * 1024 * 1024; const WS_ACCOUNT_AFFINITY_KEYS: &[&str] = &[ "session-id", "session_id", @@ -45,6 +47,21 @@ pub(super) struct CompletedWsToolCallCache { insertion_order: VecDeque, } +#[derive(Clone)] +struct CompletedWsResponse { + previous_response_id: Option, + input: Vec, + output: Vec, + retained_bytes: usize, +} + +#[derive(Default)] +pub(super) struct CompletedWsResponseCache { + responses: HashMap, + insertion_order: VecDeque, + retained_bytes: usize, +} + impl WsToolCallKind { fn call_item_type(&self) -> &'static str { match self { @@ -114,6 +131,193 @@ impl CompletedWsToolCallCache { } } +impl CompletedWsResponseCache { + pub(super) fn observe_completed_response( + &mut self, + terminal_text: &str, + previous_response_id: Option<&str>, + input: &Value, + ) -> Result { + let terminal = serde_json::from_str::(terminal_text) + .map_err(|err| format!("parse completed websocket response failed: {err}"))?; + let event_type = terminal + .get("type") + .and_then(Value::as_str) + .unwrap_or_default() + .trim() + .to_ascii_lowercase(); + if !matches!(event_type.as_str(), "response.completed" | "response.done") { + return Ok(false); + } + let response = terminal + .get("response") + .and_then(Value::as_object) + .ok_or_else(|| "completed websocket event is missing response object".to_string())?; + let response_id = response + .get("id") + .and_then(Value::as_str) + .map(str::trim) + .filter(|value| !value.is_empty()) + .ok_or_else(|| "completed websocket response is missing a non-empty id".to_string())? + .to_string(); + let output = response + .get("output") + .and_then(Value::as_array) + .cloned() + .ok_or_else(|| { + format!("completed websocket response {response_id} is missing output array") + })?; + let input = normalize_ws_response_input(input); + let previous_response_id = previous_response_id + .map(str::trim) + .filter(|value| !value.is_empty()) + .map(str::to_string); + let retained_bytes = serde_json::to_vec(&input) + .map_err(|err| format!("serialize cached websocket response input failed: {err}"))? + .len() + .saturating_add( + serde_json::to_vec(&output) + .map_err(|err| { + format!("serialize cached websocket response output failed: {err}") + })? + .len(), + ) + .saturating_add(response_id.len()) + .saturating_add( + previous_response_id + .as_ref() + .map(String::len) + .unwrap_or_default(), + ); + if retained_bytes > MAX_CACHED_WS_RESPONSE_BYTES { + return Err(format!( + "completed websocket response {response_id} requires {retained_bytes} bytes, exceeding the {} byte recovery cache limit", + MAX_CACHED_WS_RESPONSE_BYTES + )); + } + + if let Some(previous) = self.responses.remove(&response_id) { + self.retained_bytes = self.retained_bytes.saturating_sub(previous.retained_bytes); + self.insertion_order + .retain(|cached_id| cached_id != &response_id); + } + while self.responses.len() >= MAX_CACHED_WS_RESPONSES + || self.retained_bytes.saturating_add(retained_bytes) > MAX_CACHED_WS_RESPONSE_BYTES + { + let Some(oldest_id) = self.insertion_order.pop_front() else { + break; + }; + if let Some(oldest) = self.responses.remove(&oldest_id) { + self.retained_bytes = self.retained_bytes.saturating_sub(oldest.retained_bytes); + } + } + + self.retained_bytes = self.retained_bytes.saturating_add(retained_bytes); + self.insertion_order.push_back(response_id.clone()); + self.responses.insert( + response_id, + CompletedWsResponse { + previous_response_id, + input, + output, + retained_bytes, + }, + ); + Ok(true) + } + + #[cfg(test)] + pub(super) fn contains(&self, response_id: &str) -> bool { + self.responses.contains_key(response_id) + } + + fn collect_history_items(&self, response_id: &str) -> Result, String> { + let mut current_response_id = response_id.to_string(); + let mut visited = HashSet::new(); + let mut chain = Vec::new(); + loop { + if !visited.insert(current_response_id.clone()) { + return Err(format!( + "cached websocket response history contains a cycle at {current_response_id}" + )); + } + let response = self.responses.get(¤t_response_id).ok_or_else(|| { + format!( + "previous_response_id {current_response_id} is not available in this websocket session recovery cache" + ) + })?; + chain.push(response); + let Some(previous_response_id) = response.previous_response_id.as_deref() else { + break; + }; + current_response_id = previous_response_id.to_string(); + if chain.len() >= MAX_CACHED_WS_RESPONSES { + return Err(format!( + "websocket response history for {response_id} exceeds the {MAX_CACHED_WS_RESPONSES} response recovery limit" + )); + } + } + + let capacity = chain.iter().fold(0usize, |total, response| { + total + .saturating_add(response.input.len()) + .saturating_add(response.output.len()) + }); + let mut history = Vec::with_capacity(capacity); + for response in chain.into_iter().rev() { + history.extend(response.input.iter().cloned()); + history.extend(response.output.iter().cloned()); + } + Ok(history) + } +} + +fn normalize_ws_response_input(input: &Value) -> Vec { + match input { + Value::Array(items) => items.clone(), + Value::Null => Vec::new(), + Value::String(text) => vec![serde_json::json!({ + "type": "message", + "role": "user", + "content": [{ "type": "input_text", "text": text }] + })], + item => vec![item.clone()], + } +} + +pub(super) fn expand_response_create_previous_response( + text: &str, + completed_responses: &CompletedWsResponseCache, +) -> Result, String> { + let mut value = serde_json::from_str::(text) + .map_err(|err| format!("parse response.create for history recovery failed: {err}"))?; + let object = value + .as_object_mut() + .ok_or_else(|| "response.create for history recovery must be a JSON object".to_string())?; + if object.get("type").and_then(Value::as_str) != Some("response.create") { + return Err("history recovery only supports type=response.create".to_string()); + } + let Some(previous_response_id) = object + .get("previous_response_id") + .and_then(Value::as_str) + .map(str::trim) + .filter(|value| !value.is_empty()) + .map(str::to_string) + else { + return Ok(None); + }; + + let mut input = completed_responses.collect_history_items(&previous_response_id)?; + if let Some(current_input) = object.get("input") { + input.extend(normalize_ws_response_input(current_input)); + } + object.remove("previous_response_id"); + object.insert("input".to_string(), Value::Array(input)); + serde_json::to_string(&value) + .map(Some) + .map_err(|err| format!("serialize response.create after history recovery failed: {err}")) +} + fn normalize_ws_tool_call_item(item: &Value) -> Option<(WsToolCallKey, Value)> { let object = item.as_object()?; let kind = match object diff --git a/crates/service/src/http/responses_websocket_tests.rs b/crates/service/src/http/responses_websocket_tests.rs index 6547f91b1..0d3310cdc 100644 --- a/crates/service/src/http/responses_websocket_tests.rs +++ b/crates/service/src/http/responses_websocket_tests.rs @@ -6,11 +6,13 @@ use super::{ prepare_missing_ws_tool_call_retry, proxy_basic_auth_header, rebase_ws_request_for_account_change, rewrite_client_frame, should_buffer_ws_upstream_preamble, strip_previous_response_id_from_ws_text, ws_request_has_tool_call_output, - CompletedWsToolCallCache, WsRequestContext, WsToolCallKind, WsUpstreamAuthorization, + CompletedWsResponseCache, CompletedWsToolCallCache, WsRequestContext, WsToolCallKind, + WsUpstreamAuthorization, }; use axum::http::{HeaderMap, HeaderValue}; use codexmanager_core::storage::{ - now_ts, Account, ApiKey, ManagedModelV2Upsert, ModelFastPolicyV2, Storage, Token, + now_ts, Account, ApiKey, ConversationBinding, ManagedModelV2Upsert, ModelFastPolicyV2, Storage, + Token, }; use serde_json::{json, Value}; @@ -45,6 +47,8 @@ fn websocket_frame_applies_model_fast_policy() { api_key: sample_api_key(), incoming_headers: sample_incoming_headers(None, None), prompt_cache_key: None, + route_conversation_id: None, + route_conversation_source: None, effective_upstream_base: "https://chatgpt.com/backend-api/codex".to_string(), prefer_raw_errors: false, }; @@ -277,6 +281,148 @@ fn websocket_initial_and_terminal_failover_candidates_stay_in_key_group() { .all(|(account, _)| account.group_name.as_deref() == Some("team-a"))); } +#[test] +fn websocket_reselection_excludes_disabled_and_runtime_limited_accounts() { + let _guard = crate::test_env_guard(); + let storage = Storage::open_in_memory().expect("open"); + storage.init().expect("init"); + let mut api_key = sample_api_key(); + api_key.id = "gk-ws-reselection".to_string(); + api_key.key_hash = "hash-ws-reselection".to_string(); + storage.insert_api_key(&api_key).expect("insert api key"); + insert_ws_candidate(&storage, "acc-disabled", 0, "team-a"); + insert_ws_candidate(&storage, "acc-available", 1, "team-a"); + storage + .update_account_status("acc-disabled", "disabled") + .expect("disable account"); + crate::gateway::invalidate_candidate_cache(); + + let routed = crate::gateway::gateway_collect_routed_candidates_for_ws( + &storage, + &api_key.id, + Some("gpt-5.4"), + Some("conversation-reselection"), + Some(crate::gateway::conversation_binding::RouteConversationSource::NativeConversation), + ) + .expect("collect websocket candidates"); + assert_eq!( + routed + .candidates + .iter() + .map(|(account, _)| account.id.as_str()) + .collect::>(), + vec!["acc-available"] + ); + assert!(crate::gateway::gateway_ws_account_requires_switch( + &routed, + "acc-disabled" + )); + + crate::gateway::gateway_mark_account_cooldown_for_status("acc-available", 429); + let cooled = crate::gateway::gateway_collect_routed_candidates_for_ws( + &storage, + &api_key.id, + Some("gpt-5.4"), + Some("conversation-reselection"), + Some(crate::gateway::conversation_binding::RouteConversationSource::NativeConversation), + ) + .expect("collect websocket candidates after cooldown"); + assert!(cooled.candidates.is_empty()); + assert!(crate::gateway::gateway_ws_account_requires_switch( + &cooled, + "acc-available" + )); + crate::gateway::reload_runtime_config_from_env(); +} + +#[test] +fn websocket_reselection_keeps_thread_binding_but_switches_when_bound_account_is_unavailable() { + let _guard = crate::test_env_guard(); + let storage = Storage::open_in_memory().expect("open"); + storage.init().expect("init"); + let mut api_key = sample_api_key(); + api_key.id = "gk-ws-thread-reselection".to_string(); + api_key.key_hash = "hash-ws-thread-reselection".to_string(); + storage.insert_api_key(&api_key).expect("insert api key"); + insert_ws_candidate(&storage, "acc-bound", 0, "team-a"); + insert_ws_candidate(&storage, "acc-next", 1, "team-a"); + let now = now_ts(); + storage + .upsert_conversation_binding(&ConversationBinding { + platform_key_hash: api_key.key_hash.clone(), + conversation_id: "conversation-thread".to_string(), + account_id: "acc-bound".to_string(), + thread_epoch: 1, + thread_anchor: "conversation-thread".to_string(), + status: "active".to_string(), + last_model: Some("gpt-5.4".to_string()), + last_switch_reason: None, + created_at: now, + updated_at: now, + last_used_at: now, + }) + .expect("insert conversation binding"); + storage + .update_account_status("acc-bound", "disabled") + .expect("disable bound account"); + crate::gateway::invalidate_candidate_cache(); + + let routed = crate::gateway::gateway_collect_routed_candidates_for_ws( + &storage, + &api_key.id, + Some("gpt-5.4"), + Some("conversation-thread"), + Some(crate::gateway::conversation_binding::RouteConversationSource::NativeConversation), + ) + .expect("collect websocket candidates for bound thread"); + assert_eq!(routed.candidates[0].0.id, "acc-next"); + assert!( + !routed + .conversation_routing + .as_ref() + .expect("conversation routing") + .bound_account_selectable + ); + assert!(crate::gateway::gateway_ws_account_requires_switch( + &routed, + "acc-bound" + )); +} + +#[test] +fn websocket_reselection_honors_manual_preference_without_conversation_binding() { + let mut preferred = sample_account(); + preferred.id = "acc-preferred".to_string(); + let mut current = sample_account(); + current.id = "acc-current".to_string(); + let token = |account_id: &str| Token { + account_id: account_id.to_string(), + id_token: "header.payload.sig".to_string(), + access_token: "header.payload.sig".to_string(), + refresh_token: "refresh".to_string(), + api_key_access_token: None, + last_refresh: now_ts(), + }; + let routed = crate::gateway::GatewayRoutedCandidates { + candidates: vec![ + (preferred, token("acc-preferred")), + (current, token("acc-current")), + ], + route_strategy: "manual_preferred_account", + route_source: "manual_preferred_account", + conversation_routing: None, + }; + + assert!(crate::gateway::gateway_ws_account_requires_switch( + &routed, + "acc-current" + )); + assert!(!crate::gateway::gateway_ws_account_requires_switch( + &routed, + "acc-preferred" + )); +} + fn sample_incoming_headers( conversation_id: Option<&str>, turn_state: Option<&str>, @@ -543,6 +689,8 @@ fn websocket_frame_aligns_prompt_cache_key_with_native_conversation_anchor() { api_key: sample_api_key(), incoming_headers: sample_incoming_headers(Some("conversation-1"), None), prompt_cache_key: Some("sticky-thread".to_string()), + route_conversation_id: None, + route_conversation_source: None, effective_upstream_base: "https://chatgpt.com/backend-api/codex".to_string(), prefer_raw_errors: false, }; @@ -570,6 +718,8 @@ fn upstream_websocket_request_forwards_oai_attestation_header() { api_key: sample_api_key(), incoming_headers: crate::gateway::IncomingHeaderSnapshot::from_http_headers(&headers), prompt_cache_key: None, + route_conversation_id: None, + route_conversation_source: None, effective_upstream_base: "https://chatgpt.com/backend-api/codex".to_string(), prefer_raw_errors: false, }; @@ -615,6 +765,8 @@ fn upstream_websocket_request_preserves_agent_assertion_and_fedramp() { api_key: sample_api_key(), incoming_headers: crate::gateway::IncomingHeaderSnapshot::default(), prompt_cache_key: None, + route_conversation_id: None, + route_conversation_source: None, effective_upstream_base: "https://chatgpt.com/backend-api/codex".to_string(), prefer_raw_errors: false, }; @@ -698,6 +850,8 @@ fn websocket_frame_merges_header_metadata_into_client_metadata() { api_key: sample_api_key(), incoming_headers: sample_incoming_headers_with_metadata(), prompt_cache_key: None, + route_conversation_id: None, + route_conversation_source: None, effective_upstream_base: "https://chatgpt.com/backend-api/codex".to_string(), prefer_raw_errors: false, }; @@ -729,6 +883,8 @@ fn websocket_response_create_keeps_codex_field_snapshot() { api_key: sample_api_key(), incoming_headers: sample_incoming_headers_with_metadata(), prompt_cache_key: None, + route_conversation_id: None, + route_conversation_source: None, effective_upstream_base: "https://chatgpt.com/backend-api/codex".to_string(), prefer_raw_errors: false, }; @@ -834,6 +990,8 @@ fn websocket_response_create_uses_minimal_fallback_for_missing_or_blank_instruct api_key: sample_api_key(), incoming_headers: sample_incoming_headers_with_metadata(), prompt_cache_key: None, + route_conversation_id: None, + route_conversation_source: None, effective_upstream_base: "https://chatgpt.com/backend-api/codex".to_string(), prefer_raw_errors: false, }; @@ -868,6 +1026,8 @@ fn websocket_logs_client_ultra_and_sends_upstream_max() { api_key: sample_api_key(), incoming_headers: sample_incoming_headers_with_metadata(), prompt_cache_key: None, + route_conversation_id: None, + route_conversation_source: None, effective_upstream_base: "https://chatgpt.com/backend-api/codex".to_string(), prefer_raw_errors: false, }; @@ -911,6 +1071,113 @@ fn websocket_retry_can_strip_previous_response_id() { assert_eq!(value["input"], "follow up"); } +#[test] +fn websocket_response_history_expands_store_false_text_chain() { + let mut cache = CompletedWsResponseCache::default(); + assert!(cache + .observe_completed_response( + json!({ + "type": "response.completed", + "response": { + "id": "resp_history_1", + "output": [{ + "type": "message", + "role": "assistant", + "content": [{ "type": "output_text", "text": "first answer" }] + }] + } + }) + .to_string() + .as_str(), + None, + &json!("first question"), + ) + .expect("cache first completed response")); + assert!(cache.contains("resp_history_1")); + assert!(cache + .observe_completed_response( + json!({ + "type": "response.completed", + "response": { + "id": "resp_history_2", + "output": [{ + "type": "message", + "role": "assistant", + "content": [{ "type": "output_text", "text": "second answer" }] + }] + } + }) + .to_string() + .as_str(), + Some("resp_history_1"), + &json!([{ + "type": "message", + "role": "user", + "content": [{ "type": "input_text", "text": "second question" }] + }]), + ) + .expect("cache second completed response")); + + let expanded = super::expand_response_create_previous_response( + json!({ + "type": "response.create", + "model": "gpt-5.4", + "store": false, + "previous_response_id": "resp_history_2", + "input": "third question" + }) + .to_string() + .as_str(), + &cache, + ) + .expect("expand cached response history") + .expect("request has previous_response_id"); + let value: Value = serde_json::from_str(&expanded).expect("parse expanded history"); + let input = value["input"].as_array().expect("expanded input array"); + + assert!(value.get("previous_response_id").is_none()); + assert_eq!(input.len(), 5); + assert_eq!(input[0]["role"], "user"); + assert_eq!(input[0]["content"][0]["text"], "first question"); + assert_eq!(input[1]["role"], "assistant"); + assert_eq!(input[1]["content"][0]["text"], "first answer"); + assert_eq!(input[2]["content"][0]["text"], "second question"); + assert_eq!(input[3]["content"][0]["text"], "second answer"); + assert_eq!(input[4]["content"][0]["text"], "third question"); +} + +#[test] +fn websocket_response_history_requires_complete_cached_chain() { + let mut cache = CompletedWsResponseCache::default(); + cache + .observe_completed_response( + json!({ + "type": "response.completed", + "response": { "id": "resp_history_child", "output": [] } + }) + .to_string() + .as_str(), + Some("resp_history_missing_parent"), + &json!("child question"), + ) + .expect("cache child response"); + + let err = super::expand_response_create_previous_response( + json!({ + "type": "response.create", + "previous_response_id": "resp_history_child", + "input": "continue" + }) + .to_string() + .as_str(), + &cache, + ) + .expect_err("missing parent must not produce partial context"); + + assert!(err.contains("resp_history_missing_parent")); + assert!(err.contains("not available")); +} + #[test] fn websocket_account_rebase_prepends_cached_tool_calls_before_outputs() { let mut cache = CompletedWsToolCallCache::default(); @@ -1151,6 +1418,8 @@ fn upstream_websocket_account_rebase_strips_session_affinity_headers() { api_key: sample_api_key(), incoming_headers: crate::gateway::IncomingHeaderSnapshot::from_http_headers(&headers), prompt_cache_key: None, + route_conversation_id: None, + route_conversation_source: None, effective_upstream_base: "https://chatgpt.com/backend-api/codex".to_string(), prefer_raw_errors: false, }; diff --git a/crates/service/src/http/tests/proxy_runtime_tests.rs b/crates/service/src/http/tests/proxy_runtime_tests.rs index 6dd236be3..d4ec75e42 100644 --- a/crates/service/src/http/tests/proxy_runtime_tests.rs +++ b/crates/service/src/http/tests/proxy_runtime_tests.rs @@ -1296,7 +1296,17 @@ async fn start_mock_upstream_ws_closes_after_accepting_follow_up() -> ( }), serde_json::json!({ "type": "response.completed", - "response": { "id": "resp_ws_stale_follow_up_0" } + "response": { + "id": "resp_ws_stale_follow_up_0", + "output": [{ + "type": "message", + "role": "assistant", + "content": [{ + "type": "output_text", + "text": "historical seed answer" + }] + }] + } }), ] { websocket @@ -1341,7 +1351,7 @@ async fn start_mock_upstream_ws_closes_after_accepting_follow_up() -> ( }), serde_json::json!({ "type": "response.completed", - "response": { "id": "resp_ws_stale_follow_up_1" } + "response": { "id": "resp_ws_stale_follow_up_1", "output": [] } }), ] { replacement @@ -3390,6 +3400,8 @@ async fn official_responses_websocket_replays_follow_up_accepted_by_closing_upst serde_json::json!({ "type": "response.create", "model": "gpt-4.1", + "store": false, + "previous_response_id": "resp_ws_stale_follow_up_0", "input": "historical resume follow-up" }) .to_string() @@ -3405,6 +3417,23 @@ async fn official_responses_websocket_replays_follow_up_accepted_by_closing_upst .expect("historical follow-up replay channel"); assert_eq!(phase, expected_phase); assert!(text.contains("historical resume follow-up")); + let payload: serde_json::Value = + serde_json::from_str(&text).expect("parse historical follow-up frame"); + if expected_phase == 1 { + assert_eq!(payload["previous_response_id"], "resp_ws_stale_follow_up_0"); + } else { + assert!(payload.get("previous_response_id").is_none()); + let input = payload["input"] + .as_array() + .expect("replacement request carries full context"); + assert_eq!(input.len(), 3); + assert_eq!(input[0]["content"][0]["text"], "historical resume seed"); + assert_eq!(input[1]["content"][0]["text"], "historical seed answer"); + assert_eq!( + input[2]["content"][0]["text"], + "historical resume follow-up" + ); + } } let mut follow_up_completed = false; @@ -3719,11 +3748,14 @@ async fn official_responses_websocket_rebases_tool_output_on_next_account() { let rebased_input = rebased_payload["input"] .as_array() .expect("rebased tool input array"); - assert_eq!(rebased_input.len(), 2); - assert_eq!(rebased_input[0]["type"], "custom_tool_call"); - assert_eq!(rebased_input[0]["call_id"], "call_ws_tool_rebase"); - assert_eq!(rebased_input[1]["type"], "custom_tool_call_output"); + assert_eq!(rebased_input.len(), 3); + assert_eq!(rebased_input[0]["type"], "message"); + assert_eq!(rebased_input[0]["role"], "user"); + assert_eq!(rebased_input[0]["content"][0]["text"], "make a patch"); + assert_eq!(rebased_input[1]["type"], "custom_tool_call"); assert_eq!(rebased_input[1]["call_id"], "call_ws_tool_rebase"); + assert_eq!(rebased_input[2]["type"], "custom_tool_call_output"); + assert_eq!(rebased_input[2]["call_id"], "call_ws_tool_rebase"); assert!(rebased_input .iter() .all(|item| item.get("encrypted_content").is_none()));