From 0384e392890ff71a385b5276d1813a74bfe42b0a Mon Sep 17 00:00:00 2001 From: zyycn Date: Tue, 14 Jul 2026 17:55:20 +0800 Subject: [PATCH] refactor(dispatch): implement v1 lifecycle onion Co-authored-by: Codex --- backend/src/api/client/errors.rs | 82 +- backend/src/api/client/responses/websocket.rs | 217 +--- .../src/api/middleware/connection_drain.rs | 5 + backend/src/bootstrap/services.rs | 8 +- backend/src/dispatch/affinity/resolve.rs | 102 +- backend/src/dispatch/affinity/service.rs | 35 +- backend/src/dispatch/affinity/store.rs | 119 +-- backend/src/dispatch/affinity/types.rs | 20 - .../dispatch/controllers/account_failure.rs | 298 ++++++ backend/src/dispatch/controllers/affinity.rs | 160 +++ .../src/dispatch/controllers/cloudflare.rs | 237 +++++ .../dispatch/controllers/cyber_policy/mod.rs | 201 ++++ .../controllers/cyber_policy/store.rs | 119 +++ .../controllers/cyber_policy/types.rs | 9 + backend/src/dispatch/controllers/history.rs | 400 ++++++++ backend/src/dispatch/controllers/mod.rs | 927 ++++++++++++++++++ backend/src/dispatch/controllers/quota.rs | 346 +++++++ .../telemetry/events.rs} | 96 +- .../src/dispatch/controllers/telemetry/mod.rs | 349 +++++++ backend/src/dispatch/controllers/usage.rs | 106 ++ backend/src/dispatch/errors.rs | 197 +--- .../{recovery => failure}/exhaustion.rs | 137 +-- backend/src/dispatch/failure/mod.rs | 4 + backend/src/dispatch/failure/sse.rs | 78 ++ backend/src/dispatch/lifecycle/attempt.rs | 919 +++++++++++++++++ backend/src/dispatch/lifecycle/contract.rs | 282 ++++++ backend/src/dispatch/lifecycle/finalizer.rs | 202 ++++ backend/src/dispatch/lifecycle/mod.rs | 13 + backend/src/dispatch/lifecycle/pipeline.rs | 53 + backend/src/dispatch/lifecycle/request.rs | 101 ++ backend/src/dispatch/lifecycle/stream.rs | 84 ++ .../dispatch/{stream => lifecycle}/trace.rs | 2 +- backend/src/dispatch/mod.rs | 9 +- .../src/dispatch/recovery/account_failure.rs | 147 --- backend/src/dispatch/recovery/auth.rs | 22 - backend/src/dispatch/recovery/cloudflare.rs | 146 --- backend/src/dispatch/recovery/cyber_policy.rs | 386 -------- backend/src/dispatch/recovery/history.rs | 192 ---- backend/src/dispatch/recovery/mod.rs | 8 - .../{attempts.rs => routing/candidates.rs} | 0 backend/src/dispatch/routing/mod.rs | 3 + backend/src/dispatch/service.rs | 730 ++++---------- backend/src/dispatch/stream/lifecycle.rs | 672 ++++--------- backend/src/dispatch/stream/live.rs | 917 +++++------------ backend/src/dispatch/stream/mod.rs | 3 - backend/src/dispatch/stream/prefetch.rs | 36 - backend/src/dispatch/stream/sse_failure.rs | 221 ----- backend/src/dispatch/transport/account.rs | 96 ++ backend/src/dispatch/transport/canonical.rs | 424 ++++++++ backend/src/dispatch/transport/mod.rs | 6 + backend/src/dispatch/transport/observation.rs | 152 +++ backend/src/dispatch/transport/prefetch.rs | 72 ++ backend/src/dispatch/upstream_call.rs | 221 ----- .../src/upstream/openai/protocol/events.rs | 51 +- .../src/upstream/openai/protocol/responses.rs | 107 +- backend/src/upstream/openai/protocol/sse.rs | 26 +- .../src/upstream/openai/protocol/websocket.rs | 4 - .../openai/protocol/websocket_errors.rs | 154 --- .../src/upstream/openai/transport/client.rs | 24 +- backend/src/upstream/openai/transport/mod.rs | 3 +- .../upstream/openai/transport/websocket.rs | 101 +- .../openai/transport/websocket_frames.rs | 277 ++---- .../tests/api/client/responses/websocket.rs | 284 +++++- backend/tests/dispatch/affinity/mod.rs | 82 +- backend/tests/dispatch/controllers/mod.rs | 277 ++++++ backend/tests/dispatch/lifecycle/mod.rs | 295 ++++++ backend/tests/dispatch/lifecycle/stream.rs | 296 ++++++ backend/tests/dispatch/mod.rs | 3 + backend/tests/dispatch/service.rs | 64 +- .../tests/dispatch/service/responses_http.rs | 245 ++--- .../dispatch/service/responses_recovery.rs | 206 +++- .../dispatch/service/responses_websocket.rs | 424 +++++++- backend/tests/dispatch/transport/mod.rs | 70 ++ .../http_sse/failed_server_overloaded.sse | 2 + .../http_sse/live_stream_failed_auth.sse | 2 + .../http_sse/live_stream_failed_quota.sse | 2 + .../http_sse/output_then_cyber_policy.sse | 8 + ...utput_then_previous_response_not_found.sse | 8 + ...tural_then_previous_response_not_found.sse | 8 + .../tests/upstream/openai/protocol/events.rs | 26 + backend/tests/upstream/openai/protocol/mod.rs | 8 +- .../upstream/openai/protocol/responses.rs | 54 - .../upstream/openai/protocol/websocket.rs | 185 ---- .../tests/upstream/openai/transport/client.rs | 36 +- .../tests/upstream/openai/transport/mod.rs | 2 +- .../upstream/openai/transport/websocket.rs | 59 +- .../openai/transport/websocket_pool.rs | 105 +- docs/architecture.md | 76 +- docs/request-lifecycle-controller-design.md | 210 ++-- 89 files changed, 8925 insertions(+), 5230 deletions(-) create mode 100644 backend/src/dispatch/controllers/account_failure.rs create mode 100644 backend/src/dispatch/controllers/affinity.rs create mode 100644 backend/src/dispatch/controllers/cloudflare.rs create mode 100644 backend/src/dispatch/controllers/cyber_policy/mod.rs create mode 100644 backend/src/dispatch/controllers/cyber_policy/store.rs create mode 100644 backend/src/dispatch/controllers/cyber_policy/types.rs create mode 100644 backend/src/dispatch/controllers/history.rs create mode 100644 backend/src/dispatch/controllers/mod.rs create mode 100644 backend/src/dispatch/controllers/quota.rs rename backend/src/dispatch/{recording.rs => controllers/telemetry/events.rs} (89%) create mode 100644 backend/src/dispatch/controllers/telemetry/mod.rs create mode 100644 backend/src/dispatch/controllers/usage.rs rename backend/src/dispatch/{recovery => failure}/exhaustion.rs (65%) create mode 100644 backend/src/dispatch/failure/mod.rs create mode 100644 backend/src/dispatch/failure/sse.rs create mode 100644 backend/src/dispatch/lifecycle/attempt.rs create mode 100644 backend/src/dispatch/lifecycle/contract.rs create mode 100644 backend/src/dispatch/lifecycle/finalizer.rs create mode 100644 backend/src/dispatch/lifecycle/mod.rs create mode 100644 backend/src/dispatch/lifecycle/pipeline.rs create mode 100644 backend/src/dispatch/lifecycle/request.rs create mode 100644 backend/src/dispatch/lifecycle/stream.rs rename backend/src/dispatch/{stream => lifecycle}/trace.rs (94%) delete mode 100644 backend/src/dispatch/recovery/account_failure.rs delete mode 100644 backend/src/dispatch/recovery/auth.rs delete mode 100644 backend/src/dispatch/recovery/cloudflare.rs delete mode 100644 backend/src/dispatch/recovery/cyber_policy.rs delete mode 100644 backend/src/dispatch/recovery/history.rs delete mode 100644 backend/src/dispatch/recovery/mod.rs rename backend/src/dispatch/{attempts.rs => routing/candidates.rs} (100%) create mode 100644 backend/src/dispatch/routing/mod.rs delete mode 100644 backend/src/dispatch/stream/prefetch.rs delete mode 100644 backend/src/dispatch/stream/sse_failure.rs create mode 100644 backend/src/dispatch/transport/account.rs create mode 100644 backend/src/dispatch/transport/canonical.rs create mode 100644 backend/src/dispatch/transport/mod.rs create mode 100644 backend/src/dispatch/transport/observation.rs create mode 100644 backend/src/dispatch/transport/prefetch.rs delete mode 100644 backend/src/dispatch/upstream_call.rs delete mode 100644 backend/src/upstream/openai/protocol/websocket_errors.rs create mode 100644 backend/tests/dispatch/controllers/mod.rs create mode 100644 backend/tests/dispatch/lifecycle/mod.rs create mode 100644 backend/tests/dispatch/lifecycle/stream.rs create mode 100644 backend/tests/dispatch/transport/mod.rs create mode 100644 backend/tests/fixtures/responses/http_sse/failed_server_overloaded.sse create mode 100644 backend/tests/fixtures/responses/http_sse/live_stream_failed_auth.sse create mode 100644 backend/tests/fixtures/responses/http_sse/live_stream_failed_quota.sse create mode 100644 backend/tests/fixtures/responses/http_sse/output_then_cyber_policy.sse create mode 100644 backend/tests/fixtures/responses/http_sse/output_then_previous_response_not_found.sse create mode 100644 backend/tests/fixtures/responses/http_sse/structural_then_previous_response_not_found.sse diff --git a/backend/src/api/client/errors.rs b/backend/src/api/client/errors.rs index dee506a6e..f455a061a 100644 --- a/backend/src/api/client/errors.rs +++ b/backend/src/api/client/errors.rs @@ -8,11 +8,11 @@ use axum::{ use serde_json::{Value, json}; use crate::{ - dispatch::{ - errors::{DispatchFailureClass, ResponseDispatchError}, - recovery::cyber_policy::is_cyber_policy_failure, + dispatch::errors::{DispatchFailureClass, ResponseDispatchError}, + upstream::openai::protocol::{ + responses::ResponsesSseFailure, + sse::response_failed_sse_event as encode_response_failed_sse_event, }, - upstream::openai::protocol::sse::response_failed_sse_event as encode_response_failed_sse_event, upstream::openai::transport::CodexClientError, }; @@ -50,6 +50,13 @@ struct OpenAiErrorKind { code: &'static str, } +struct EncodedClientFailure<'a> { + status: StatusCode, + kind: OpenAiErrorKind, + code: &'a str, + message: &'a str, +} + const NO_AVAILABLE_ACCOUNTS_ERROR: OpenAiErrorKind = OpenAiErrorKind { error_type: "server_error", code: "no_available_accounts", @@ -126,16 +133,12 @@ pub fn responses_dispatch_error_response_ref(error: &ResponseDispatchError) -> R return (*status, Json(body)).into_response(); } - if let ResponseDispatchError::Failed(failure) = error - && is_cyber_policy_failure(failure) - { - let status = status_from_u16(error.client_http_status_code(), StatusCode::BAD_GATEWAY); - let kind = responses_error_kind_for_status(status); + if let Some(failure) = response_dispatch_client_failure(error) { return openai_error_response( - status, - &failure.message, - kind.error_type, - failure.upstream_code.as_deref().unwrap_or(kind.code), + failure.status, + failure.message, + failure.kind.error_type, + failure.code, ) .into_response(); } @@ -149,15 +152,11 @@ pub fn responses_dispatch_error_response_ref(error: &ResponseDispatchError) -> R } pub fn responses_stream_dispatch_failed_sse_event(error: &ResponseDispatchError) -> String { - if let ResponseDispatchError::Failed(failure) = error - && is_cyber_policy_failure(failure) - { - let status = status_from_u16(error.client_http_status_code(), StatusCode::BAD_GATEWAY); - let kind = responses_error_kind_for_status(status); + if let Some(failure) = response_dispatch_client_failure(error) { return encode_response_failed_sse_event( - kind.error_type, - failure.upstream_code.as_deref().unwrap_or(kind.code), - &failure.message, + failure.kind.error_type, + failure.code, + failure.message, ); } let http_error = @@ -170,16 +169,12 @@ pub(super) fn responses_websocket_dispatch_error_event( error: &ResponseDispatchError, request_id: &str, ) -> String { - if let ResponseDispatchError::Failed(failure) = error - && is_cyber_policy_failure(failure) - { - let status = status_from_u16(error.client_http_status_code(), StatusCode::BAD_GATEWAY); - let kind = responses_error_kind_for_status(status); + if let Some(failure) = response_dispatch_client_failure(error) { return responses_websocket_error_event( - status, - kind.error_type, - failure.upstream_code.as_deref().unwrap_or(kind.code), - &failure.message, + failure.status, + failure.kind.error_type, + failure.code, + failure.message, request_id, ); } @@ -354,6 +349,33 @@ fn responses_error_kind_for_status(status: StatusCode) -> OpenAiErrorKind { } } +fn response_dispatch_client_failure( + error: &ResponseDispatchError, +) -> Option> { + let client_failure = match error { + ResponseDispatchError::Failed(failure) => failure, + _ => return None, + }; + let failure = client_failure.exposed_upstream()?; + Some(encoded_client_failure( + failure, + status_from_u16(client_failure.status_code(), StatusCode::BAD_GATEWAY), + )) +} + +fn encoded_client_failure( + failure: &ResponsesSseFailure, + status: StatusCode, +) -> EncodedClientFailure<'_> { + let kind = responses_error_kind_for_status(status); + EncodedClientFailure { + status, + kind, + code: failure.upstream_code.as_deref().unwrap_or(kind.code), + message: &failure.message, + } +} + fn response_dispatch_message( error: &ResponseDispatchError, style: ResponseDispatchMessageStyle, diff --git a/backend/src/api/client/responses/websocket.rs b/backend/src/api/client/responses/websocket.rs index 0053df7c4..69bbfd2ff 100644 --- a/backend/src/api/client/responses/websocket.rs +++ b/backend/src/api/client/responses/websocket.rs @@ -1,7 +1,5 @@ //! 官方 Responses WebSocket 入站协议。 -use std::sync::Arc; - use axum::{ body::Body, extract::{Request, State}, @@ -24,11 +22,11 @@ use tokio_tungstenite::{ use crate::{ api::{AppState, middleware::request_id::RequestId}, - dispatch::{recovery::history::sanitize_replay_items, service::ResponseDispatchStream}, - upstream::openai::protocol::{ - sse::{SseEvent, SseEventDecoder}, - websocket::{is_terminal_websocket_event, websocket_event_type}, + dispatch::{ + service::{ConnectionTranscriptFacts, ResponseDispatchStream}, + transport::canonical::CanonicalResponseEvent, }, + upstream::openai::protocol::websocket::is_terminal_websocket_event, }; use super::{PreparedResponsesRequest, ResponsesRequestValidationError, prepare_responses_request}; @@ -133,129 +131,42 @@ struct ResponsesWebSocketSession { connection_id: String, } -struct ConnectionReplayState { - last_response_id: Option, - input: Option>>, -} - -enum PendingReplayUpdate { - Replace(Vec), - Append(Vec), - Unavailable, -} - -struct CompletedTurn { - response_id: String, - output: Vec, -} - #[derive(Default)] -struct CompletedTurnCollector { +struct ConnectionTranscriptCollector { completed_items: Vec, - completed: Option, + completed: Option, } enum ForwardDispatchOutcome { - Continue(Option), + Continue(Option), Disconnect, } -impl ConnectionReplayState { - fn new() -> Self { - Self { - last_response_id: None, - input: None, - } - } - - fn prepare( - &self, - request: &mut crate::upstream::openai::protocol::responses::CodexResponsesRequest, - ) -> PendingReplayUpdate { - let turn_input = request.input().to_vec(); - match request.previous_response_id() { - None => { - request.local_replay_input = None; - PendingReplayUpdate::Replace(turn_input) - } - Some(previous_response_id) - if self.last_response_id.as_deref() == Some(previous_response_id) - && self.input.is_some() => - { - request.local_replay_input.clone_from(&self.input); - PendingReplayUpdate::Append(turn_input) - } - Some(_) => { - request.local_replay_input = None; - PendingReplayUpdate::Unavailable - } - } - } - - fn commit(&mut self, pending: PendingReplayUpdate, completed: CompletedTurn) { - self.last_response_id = Some(completed.response_id); - match pending { - PendingReplayUpdate::Replace(input) => self.replace(input, completed.output), - PendingReplayUpdate::Append(input) => self.append(input, completed.output), - PendingReplayUpdate::Unavailable => self.clear_input(), - } - } - - fn replace(&mut self, mut input: Vec, mut output: Vec) { - sanitize_replay_items(&mut input); - sanitize_replay_items(&mut output); - input.append(&mut output); - self.input = Some(Arc::new(input)); - } - - fn append(&mut self, mut input: Vec, mut output: Vec) { - sanitize_replay_items(&mut input); - sanitize_replay_items(&mut output); - input.append(&mut output); - let Some(history) = self.input.as_mut() else { - return; - }; - Arc::make_mut(history).append(&mut input); - } - - fn clear_input(&mut self) { - self.input = None; - } -} - -impl CompletedTurnCollector { - fn observe(&mut self, event: &SseEvent) { - let Ok(value) = serde_json::from_str::(&event.data) else { - return; - }; - let event_type = event - .event - .as_deref() - .or_else(|| value.get("type").and_then(Value::as_str)); - match event_type { +impl ConnectionTranscriptCollector { + fn observe(&mut self, event: &CanonicalResponseEvent) { + match event.event_type() { Some("response.output_item.done") => { - if let Some(item) = value.get("item") { + if let Some(item) = event.data().get("item") { self.completed_items.push(item.clone()); } } Some("response.completed") => { - let Some(response_id) = value + let Some(response_id) = event + .data() .pointer("/response/id") .and_then(Value::as_str) .map(ToString::to_string) else { return; }; - let output = value + let output = event + .data() .pointer("/response/output") .and_then(Value::as_array) .filter(|output| !output.is_empty()) .cloned() .unwrap_or_else(|| self.completed_items.clone()); - self.completed = Some(CompletedTurn { - response_id, - output, - }); + self.completed = Some(ConnectionTranscriptFacts::new(response_id, output)); } _ => {} } @@ -271,7 +182,11 @@ async fn serve_responses_websocket( "Responses WebSocket connected" ); let mut request_count = 0u64; - let mut replay_state = ConnectionReplayState::new(); + let mut replay_snapshot = session + .state + .services + .responses + .connection_replay_snapshot(); while let Some(message) = socket.next().await { let payload = match message { @@ -361,11 +276,19 @@ async fn serve_responses_websocket( continue; } }; - let pending_replay = replay_state.prepare(&mut prepared.request); + let replay_plan = session + .state + .services + .responses + .prepare_connection_replay(&replay_snapshot, &mut prepared.request); match dispatch_response_create(&mut socket, &session, request_id, prepared).await { - ForwardDispatchOutcome::Continue(Some(completed)) => { - replay_state.commit(pending_replay, completed); + ForwardDispatchOutcome::Continue(Some(transcript)) => { + session.state.services.responses.commit_connection_replay( + &mut replay_snapshot, + replay_plan, + transcript, + ); } ForwardDispatchOutcome::Continue(None) => {} ForwardDispatchOutcome::Disconnect => break, @@ -424,6 +347,7 @@ async fn forward_dispatch_stream( ) -> ForwardDispatchOutcome { let ResponseDispatchStream { body, + mut canonical_events, response_headers, } = stream; if !send_text( @@ -436,9 +360,8 @@ async fn forward_dispatch_stream( } let mut body = body; - let mut decoder = SseEventDecoder::default(); let mut terminal_sent = false; - let mut completed = CompletedTurnCollector::default(); + let mut completed = ConnectionTranscriptCollector::default(); loop { let next = if terminal_sent { ActiveResponseInput::Upstream(body.next().await) @@ -450,28 +373,15 @@ async fn forward_dispatch_stream( }; match next { ActiveResponseInput::Upstream(Some(Ok(chunk))) => { - let events = match decoder.push(&chunk) { - Ok(events) => events, - Err(error) => { - tracing::warn!(request_id, error = %error, "Failed to decode Responses SSE for WebSocket forwarding"); - let sent = send_error( - socket, - StatusCode::BAD_GATEWAY, - "server_error", - "invalid_upstream_response", - INVALID_UPSTREAM_EVENT_MESSAGE, - request_id, - ) - .await; - drop(body); - return if sent { - ForwardDispatchOutcome::Continue(None) - } else { - ForwardDispatchOutcome::Disconnect - }; - } + let Some(events) = canonical_events.recv().await else { + tracing::warn!( + request_id, + "Canonical Responses event stream closed before body stream" + ); + return ForwardDispatchOutcome::Disconnect; }; - if !forward_sse_events( + drop(chunk); + if !forward_canonical_events( socket, events, request_id, @@ -504,37 +414,6 @@ async fn forward_dispatch_stream( }; } ActiveResponseInput::Upstream(None) => { - let events = match decoder.finish() { - Ok(events) => events, - Err(error) => { - tracing::warn!(request_id, error = %error, "Failed to finish Responses SSE decoding"); - return if send_error( - socket, - StatusCode::BAD_GATEWAY, - "server_error", - "invalid_upstream_response", - INVALID_UPSTREAM_EVENT_MESSAGE, - request_id, - ) - .await - { - ForwardDispatchOutcome::Continue(None) - } else { - ForwardDispatchOutcome::Disconnect - }; - } - }; - if !forward_sse_events( - socket, - events, - request_id, - &mut terminal_sent, - &mut completed, - ) - .await - { - return ForwardDispatchOutcome::Disconnect; - } if terminal_sent { return ForwardDispatchOutcome::Continue(completed.completed); } @@ -580,15 +459,15 @@ enum ActiveResponseInput { Downstream(Option>), } -async fn forward_sse_events( +async fn forward_canonical_events( socket: &mut ResponsesWebSocket, - events: Vec, + events: Vec, request_id: &str, terminal_sent: &mut bool, - completed: &mut CompletedTurnCollector, + completed: &mut ConnectionTranscriptCollector, ) -> bool { for event in events { - let Some(event_type) = websocket_event_type(&event.data) else { + let Some(event_type) = event.event_type() else { tracing::warn!( request_id, "Responses SSE event does not contain a valid JSON type" @@ -605,10 +484,10 @@ async fn forward_sse_events( return false; }; completed.observe(&event); - if !send_text(socket, event.data).await { + if !send_text(socket, event.data().to_string()).await { return false; } - *terminal_sent |= is_terminal_websocket_event(&event_type); + *terminal_sent |= is_terminal_websocket_event(event_type); } true } diff --git a/backend/src/api/middleware/connection_drain.rs b/backend/src/api/middleware/connection_drain.rs index c78f121b4..bd7941f45 100644 --- a/backend/src/api/middleware/connection_drain.rs +++ b/backend/src/api/middleware/connection_drain.rs @@ -28,6 +28,11 @@ impl Default for ConnectionDrain { } impl ConnectionDrain { + /// 为下层长生命周期执行器提供与连接排空一致的取消信号。 + pub(crate) fn cancellation_token(&self) -> CancellationToken { + self.cancellation.clone() + } + /// 注册需要随进程关闭而终止的长连接任务。 pub fn spawn(&self, future: impl Future + Send + 'static) { let cancellation = self.cancellation.clone(); diff --git a/backend/src/bootstrap/services.rs b/backend/src/bootstrap/services.rs index 2447f96c5..d8e357ce1 100644 --- a/backend/src/bootstrap/services.rs +++ b/backend/src/bootstrap/services.rs @@ -297,8 +297,10 @@ impl Services { let session_affinity = Arc::new(SessionAffinityService::new(stores.session_affinity.clone())); - let cloudflare_recovery = - crate::dispatch::recovery::cloudflare::CloudflareRecovery::new(stores.cookies.clone()); + let connection_drain = crate::api::middleware::connection_drain::ConnectionDrain::default(); + let cloudflare_recovery = crate::dispatch::controllers::cloudflare::CloudflareRecovery::new( + stores.cookies.clone(), + ); let responses = Arc::new(ResponseDispatchService::new(ResponseDispatchServiceParts { account_pool: account_pool.clone(), models: models.clone(), @@ -307,9 +309,9 @@ impl Services { recorder, account_identity, cloudflare: cloudflare_recovery, + shutdown: connection_drain.cancellation_token(), })); let system_update = Arc::new(SystemUpdateService::from_env()); - let connection_drain = crate::api::middleware::connection_drain::ConnectionDrain::default(); Ok(Self { database, diff --git a/backend/src/dispatch/affinity/resolve.rs b/backend/src/dispatch/affinity/resolve.rs index 3ce7d8fa1..491c8958a 100644 --- a/backend/src/dispatch/affinity/resolve.rs +++ b/backend/src/dispatch/affinity/resolve.rs @@ -1,107 +1,7 @@ -use chrono::{DateTime, Utc}; use serde_json::Value; use sha2::{Digest, Sha256}; -use crate::{ - dispatch::{ - affinity::{SessionAffinityEntry, SessionAffinityService}, - recovery::history::HistoryRecoveryPlan, - service::ResponseDispatchService, - }, - upstream::openai::protocol::{ - events::TokenUsage, - responses::{CodexResponsesRequest, completed_response_metadata}, - }, -}; - -impl ResponseDispatchService { - pub(in crate::dispatch) async fn preferred_account_id_for_request( - &self, - request: &CodexResponsesRequest, - history: &HistoryRecoveryPlan, - now: DateTime, - ) -> Option { - if let Some(account_id) = history.preferred_account_id() { - return Some(account_id.to_string()); - } - if request.previous_response_id().is_some() { - return None; - } - self.conversation_account_id(request, now).await - } - - async fn conversation_account_id( - &self, - request: &CodexResponsesRequest, - now: DateTime, - ) -> Option { - let conversation_id = request - .local_conversation_id - .as_deref() - .map(str::trim) - .filter(|value| !value.is_empty())?; - let variant_hash = compute_variant_hash(request); - self.session_affinity - .lookup_latest_account_by_conversation(conversation_id, None, Some(&variant_hash), now) - .await - } -} - -pub(in crate::dispatch) async fn record_response_affinity( - session_affinity: &SessionAffinityService, - history: &HistoryRecoveryPlan, - original_request: &CodexResponsesRequest, - account_id: &str, - body: &str, - turn_state: Option, - usage: Option, -) { - let metadata = match completed_response_metadata(body) { - Ok(Some(metadata)) => metadata, - Ok(None) => return, - Err(error) => { - tracing::warn!( - error = %error, - "Failed to parse completed response metadata for session affinity" - ); - return; - } - }; - - let conversation_id = history - .conversation_id(original_request) - .unwrap_or(&metadata.response_id) - .to_string(); - let variant_hash = compute_variant_hash(original_request); - let entry = SessionAffinityEntry { - account_id: account_id.to_string(), - conversation_id, - turn_state: turn_state - .filter(|value| !value.trim().is_empty()) - .or_else(|| original_request.turn_state.clone()), - instructions_hash: Some(hash_instructions(Some(original_request.instructions()))), - input_tokens: usage.map(|usage| usage.input_tokens), - function_call_ids: metadata.function_call_ids, - variant_hash: Some(variant_hash), - continuation_scope: if original_request.store() { - crate::upstream::openai::protocol::responses::PreviousResponseScope::Persisted - } else { - crate::upstream::openai::protocol::responses::PreviousResponseScope::ConnectionLocal - }, - created_at: Utc::now(), - }; - if let Err(error) = session_affinity - .record(metadata.response_id.clone(), entry) - .await - { - tracing::warn!( - error = %error, - response_id = %metadata.response_id, - account_id = %account_id, - "Failed to record session affinity" - ); - } -} +use crate::upstream::openai::protocol::responses::CodexResponsesRequest; // ==================================================================== // Variant identity helpers diff --git a/backend/src/dispatch/affinity/service.rs b/backend/src/dispatch/affinity/service.rs index 16a4e83c0..1d90a8352 100644 --- a/backend/src/dispatch/affinity/service.rs +++ b/backend/src/dispatch/affinity/service.rs @@ -3,9 +3,11 @@ use chrono::{DateTime, Duration, Utc}; use thiserror::Error; +use crate::infra::redis::RedisConnection; + use super::{ store::{RedisSessionAffinityStore, RedisSessionAffinityStoreError}, - types::{CyberPolicyFailureSnapshot, CyberPolicySessionState, SessionAffinityEntry}, + types::SessionAffinityEntry, }; const DEFAULT_SESSION_AFFINITY_TTL_SECS: i64 = 4 * 60 * 60; @@ -82,35 +84,8 @@ impl SessionAffinityService { } } - pub(crate) async fn load_cyber_policy_state( - &self, - session_key: &str, - ) -> Result, SessionAffinityError> { - Ok(self.store.cyber_policy_state(session_key).await?) - } - - pub(crate) async fn persist_cyber_policy_failure( - &self, - session_key: &str, - failure: &CyberPolicyFailureSnapshot, - max_accounts: usize, - ttl: Duration, - ) -> Result { - Ok(self - .store - .record_cyber_policy_failure(session_key, failure, max_accounts, ttl) - .await?) - } - - pub(crate) async fn delete_cyber_policy_state( - &self, - session_key: &str, - expected_revision: &str, - ) -> Result { - Ok(self - .store - .clear_cyber_policy_state(session_key, expected_revision) - .await?) + pub(crate) fn redis_connection(&self) -> RedisConnection { + self.store.redis_connection() } async fn entry(&self, response_id: &str, now: DateTime) -> Option { diff --git a/backend/src/dispatch/affinity/store.rs b/backend/src/dispatch/affinity/store.rs index 4d9d00dec..d11a7c901 100644 --- a/backend/src/dispatch/affinity/store.rs +++ b/backend/src/dispatch/affinity/store.rs @@ -6,7 +6,7 @@ use thiserror::Error; use crate::infra::redis::RedisConnection; -use super::types::{CyberPolicyFailureSnapshot, CyberPolicySessionState, SessionAffinityEntry}; +use super::types::SessionAffinityEntry; /// 单会话保留的最近响应元数据数量。 pub const MAX_CONVERSATION_AFFINITIES: usize = 128; @@ -21,65 +21,6 @@ const MAX_AFFINITY_METADATA_BYTES: usize = 64 * 1024; const GLOBAL_STALE_CLEANUP_BATCH: usize = 256; const AFFINITY_VERSION: &str = "affinity:v3"; -const RECORD_CYBER_POLICY_FAILURE_SCRIPT: &str = r#" -local raw = redis.call('GET', KEYS[1]) -local state = { failedAccountIds = {}, revision = '' } -if raw then - local ok, decoded = pcall(cjson.decode, raw) - if ok and type(decoded) == 'table' and type(decoded.failedAccountIds) == 'table' then - state = decoded - end -end - -local account_id = ARGV[1] -local max_accounts = tonumber(ARGV[5]) -local known = false -for _, failed_account_id in ipairs(state.failedAccountIds) do - if failed_account_id == account_id then - known = true - break - end -end - -local recorded = known -if not known and #state.failedAccountIds < max_accounts then - table.insert(state.failedAccountIds, account_id) - recorded = true -end - -if recorded then - local failure = { - accountId = account_id, - event = ARGV[2], - message = ARGV[4] - } - if ARGV[3] ~= '' then - failure.upstreamCode = ARGV[3] - end - state.lastFailure = failure - state.revision = ARGV[7] -end - -local encoded = cjson.encode(state) -redis.call('SET', KEYS[1], encoded, 'EX', tonumber(ARGV[6])) -return encoded -"#; - -const CLEAR_CYBER_POLICY_STATE_SCRIPT: &str = r#" -local raw = redis.call('GET', KEYS[1]) -if not raw then - return 0 -end -local ok, state = pcall(cjson.decode, raw) -if not ok or type(state) ~= 'table' then - return 0 -end -if tostring(state.revision or '') ~= ARGV[1] then - return 0 -end -return redis.call('DEL', KEYS[1]) -"#; - const UPSERT_SCRIPT: &str = r#" local response_id = ARGV[1] local metadata_json = ARGV[2] @@ -355,6 +296,10 @@ impl RedisSessionAffinityStore { Self { redis } } + pub(crate) fn redis_connection(&self) -> RedisConnection { + self.redis.clone() + } + pub async fn upsert( &self, response_id: &str, @@ -520,55 +465,6 @@ impl RedisSessionAffinityStore { Ok(forgotten) } - pub async fn cyber_policy_state( - &self, - session_key: &str, - ) -> RedisSessionAffinityStoreResult> { - let mut connection = self.redis.manager(); - let value: Option = connection.get(self.cyber_policy_key(session_key)).await?; - value - .map(|value| serde_json::from_str(&value)) - .transpose() - .map_err(Into::into) - } - - pub async fn record_cyber_policy_failure( - &self, - session_key: &str, - failure: &CyberPolicyFailureSnapshot, - max_accounts: usize, - ttl: Duration, - ) -> RedisSessionAffinityStoreResult { - let mut connection = self.redis.manager(); - let revision = uuid::Uuid::new_v4().to_string(); - let value: String = redis::Script::new(RECORD_CYBER_POLICY_FAILURE_SCRIPT) - .key(self.cyber_policy_key(session_key)) - .arg(&failure.account_id) - .arg(&failure.event) - .arg(failure.upstream_code.as_deref().unwrap_or_default()) - .arg(&failure.message) - .arg(max_accounts.max(1)) - .arg(ttl.num_seconds().max(1)) - .arg(revision) - .invoke_async(&mut connection) - .await?; - Ok(serde_json::from_str(&value)?) - } - - pub async fn clear_cyber_policy_state( - &self, - session_key: &str, - expected_revision: &str, - ) -> RedisSessionAffinityStoreResult { - let mut connection = self.redis.manager(); - let removed: usize = redis::Script::new(CLEAR_CYBER_POLICY_STATE_SCRIPT) - .key(self.cyber_policy_key(session_key)) - .arg(expected_revision) - .invoke_async(&mut connection) - .await?; - Ok(removed > 0) - } - fn response_key(&self, response_id: &str) -> String { format!("{}{response_id}", self.response_key_prefix()) } @@ -581,11 +477,6 @@ impl RedisSessionAffinityStore { format!("{}{account_id}", self.account_key_prefix()) } - fn cyber_policy_key(&self, session_key: &str) -> String { - self.redis - .key(&format!("{AFFINITY_VERSION}:cyber:{session_key}")) - } - fn global_index_key(&self) -> String { self.redis.key(&format!("{AFFINITY_VERSION}:global:index")) } diff --git a/backend/src/dispatch/affinity/types.rs b/backend/src/dispatch/affinity/types.rs index 60955c054..bb4d8fec0 100644 --- a/backend/src/dispatch/affinity/types.rs +++ b/backend/src/dispatch/affinity/types.rs @@ -5,26 +5,6 @@ use serde::{Deserialize, Serialize}; use crate::upstream::openai::protocol::responses::PreviousResponseScope; -/// 单个账号返回的 `cyber_policy` 证据。 -#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] -#[serde(rename_all = "camelCase")] -pub struct CyberPolicyFailureSnapshot { - pub account_id: String, - pub event: String, - pub message: String, - pub upstream_code: Option, -} - -/// 会话级 `cyber_policy` 换号状态。 -#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize)] -#[serde(rename_all = "camelCase")] -pub struct CyberPolicySessionState { - pub failed_account_ids: Vec, - pub last_failure: Option, - #[serde(default)] - pub revision: String, -} - /// 会话亲和性条目。 #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] #[serde(rename_all = "camelCase")] diff --git a/backend/src/dispatch/controllers/account_failure.rs b/backend/src/dispatch/controllers/account_failure.rs new file mode 100644 index 000000000..ff95fb968 --- /dev/null +++ b/backend/src/dispatch/controllers/account_failure.rs @@ -0,0 +1,298 @@ +//! 账号级失败分类、隔离副作用与换号决策 owner。 + +use crate::{ + dispatch::{ + controllers::ControllerFailureFact, + errors::ClientFailure, + failure::exhaustion::{AccountExhaustionRecord, ExhaustedAccountKind}, + lifecycle::contract::{ + AttemptDecision, AttemptObservation, AttemptObservationKind, CompleteResponseFacts, + }, + transport::observation::{UpstreamFailureFacts, UpstreamFailureKind}, + }, + fleet::{account::AccountStatus, pool::AccountPoolService}, + upstream::openai::protocol::responses::ResponsesSseFailure, +}; + +pub(super) struct AccountFailureController; + +pub(super) struct ClassifiedFailure { + exhaustion: AccountExhaustionRecord, + effect: AccountEffect, +} + +enum AccountEffect { + None, + SetStatus(AccountStatus), +} + +impl AccountFailureController { + pub(super) fn classify(observation: &AttemptObservation) -> Option { + let account_id = observation.account.as_ref()?.id.as_str(); + classify(observation, account_id) + } + + pub(super) fn classify_failure( + account_id: &str, + failure: ControllerFailureFact<'_>, + ) -> Option { + classify_failure(account_id, failure) + } + + pub(super) async fn apply_effect( + account_pool: &AccountPoolService, + classified: &ClassifiedFailure, + ) { + let Some(account_id) = classified.exhaustion.account_id.as_deref() else { + return; + }; + match &classified.effect { + AccountEffect::None => {} + AccountEffect::SetStatus(status) => { + account_pool.set_status(account_id, *status).await; + } + } + } + + pub(super) fn decision( + observation: &AttemptObservation, + classified: ClassifiedFailure, + ) -> AttemptDecision { + if observation.routing.can_retry_next_candidate { + AttemptDecision::RetryNextCandidate { + exhaustion: Some(classified.exhaustion), + on_exhaustion: None, + } + } else { + AttemptDecision::Return( + crate::dispatch::lifecycle::contract::AttemptReturnKind::Observed, + ) + } + } + + pub(super) fn is_retryable_transport(observation: &AttemptObservation) -> bool { + let AttemptObservationKind::UpstreamFailure(facts) = &observation.kind else { + return false; + }; + matches!( + facts.kind, + UpstreamFailureKind::HttpConnect + | UpstreamFailureKind::HttpTimeout + | UpstreamFailureKind::HttpTransport + | UpstreamFailureKind::StreamIdle + | UpstreamFailureKind::WebSocketTransport + | UpstreamFailureKind::WebSocketTimeout + ) || facts + .status_code + .is_some_and(|status| (500..=599).contains(&status)) + } + + pub(super) fn client_failure( + classified: &ClassifiedFailure, + failure: &ResponsesSseFailure, + ) -> ClientFailure { + let status_code = + classified + .exhaustion + .status_code + .unwrap_or(match classified.exhaustion.kind { + ExhaustedAccountKind::ModelUnsupported => 400, + ExhaustedAccountKind::Expired | ExhaustedAccountKind::Disabled => 401, + ExhaustedAccountKind::Banned => 403, + ExhaustedAccountKind::QuotaExhausted + | ExhaustedAccountKind::RateLimited + | ExhaustedAccountKind::CloudflareChallenge + | ExhaustedAccountKind::CloudflarePathBlocked + | ExhaustedAccountKind::UpstreamUnavailable => 502, + }); + ClientFailure::new(failure.clone(), status_code, false) + } + + pub(super) fn unowned_client_failure(failure: &ResponsesSseFailure) -> Option { + let code = failure + .upstream_code + .as_deref() + .unwrap_or_default() + .to_ascii_lowercase(); + let message = failure.message.to_ascii_lowercase(); + if is_model_unsupported(&code) + || is_model_unsupported(&message) + || code.contains("invalid_request") + || code.contains("not_found") + || code.contains("context_window") + || code.contains("invalid_prompt") + || code.contains("bad_request") + { + return Some(ClientFailure::new(failure.clone(), 400, false)); + } + if is_auth_failure(&code, &message) { + return Some(ClientFailure::new(failure.clone(), 401, false)); + } + if code.contains("forbidden") || code.contains("banned") { + return Some(ClientFailure::new(failure.clone(), 403, false)); + } + if code.contains("server_overloaded") { + return Some(ClientFailure::new(failure.clone(), 503, true)); + } + None + } +} + +fn classify(observation: &AttemptObservation, account_id: &str) -> Option { + match &observation.kind { + AttemptObservationKind::CompleteResponse(CompleteResponseFacts::Empty) => { + Some(ClassifiedFailure { + exhaustion: AccountExhaustionRecord::new( + account_id, + ExhaustedAccountKind::UpstreamUnavailable, + "upstream response did not include visible output", + ), + effect: AccountEffect::None, + }) + } + _ => ControllerFailureFact::from_attempt(observation) + .and_then(|failure| classify_failure(account_id, failure)), + } +} + +fn classify_failure( + account_id: &str, + failure: ControllerFailureFact<'_>, +) -> Option { + match failure { + ControllerFailureFact::Upstream(facts) => classify_upstream(account_id, facts), + ControllerFailureFact::Response(failure) => classify_sse(account_id, failure), + } +} + +fn classify_upstream(account_id: &str, facts: &UpstreamFailureFacts) -> Option { + let status = facts.status_code; + let body = facts.body.as_str(); + let lower = body.to_ascii_lowercase(); + let (kind, effect, message, status_code) = if (status == Some(403) && !is_html(body)) + || (status == Some(402) && is_deactivated_workspace(body)) + { + ( + ExhaustedAccountKind::Banned, + AccountEffect::SetStatus(AccountStatus::Banned), + body.to_string(), + status, + ) + } else if status == Some(401) { + let account_status = if lower.contains("banned") || lower.contains("deactivated") { + AccountStatus::Banned + } else { + AccountStatus::Expired + }; + ( + if account_status == AccountStatus::Banned { + ExhaustedAccountKind::Banned + } else { + ExhaustedAccountKind::Expired + }, + AccountEffect::SetStatus(account_status), + body.to_string(), + status, + ) + } else if status.is_some_and(|status| { + (400..=499).contains(&status) && !matches!(status, 401 | 402 | 403 | 404 | 429) + }) && is_model_unsupported(body) + { + ( + ExhaustedAccountKind::ModelUnsupported, + AccountEffect::None, + body.to_string(), + None, + ) + } else { + return None; + }; + let mut exhaustion = AccountExhaustionRecord::new(account_id, kind, message); + if let Some(status_code) = status_code { + exhaustion = exhaustion.with_status_code(status_code); + } + Some(ClassifiedFailure { exhaustion, effect }) +} + +fn classify_sse(account_id: &str, failure: &ResponsesSseFailure) -> Option { + let code = failure.upstream_code.as_deref().unwrap_or_default(); + let lower_message = failure.message.to_ascii_lowercase(); + let (kind, effect, status) = if is_model_unsupported(code) + || is_model_unsupported(&failure.message) + { + ( + ExhaustedAccountKind::ModelUnsupported, + AccountEffect::None, + None, + ) + } else if is_auth_failure(code, &lower_message) { + let account_status = if lower_message.contains("banned") || code == "account_deactivated" { + AccountStatus::Banned + } else { + AccountStatus::Expired + }; + ( + if account_status == AccountStatus::Banned { + ExhaustedAccountKind::Banned + } else { + ExhaustedAccountKind::Expired + }, + AccountEffect::SetStatus(account_status), + Some(if account_status == AccountStatus::Banned { + 403 + } else { + 401 + }), + ) + } else { + return None; + }; + let mut exhaustion = AccountExhaustionRecord::new( + account_id, + kind, + crate::dispatch::failure::sse::sse_failure_error_body(failure), + ); + if let Some(status) = status { + exhaustion = exhaustion.with_status_code(status); + } + Some(ClassifiedFailure { exhaustion, effect }) +} + +fn is_auth_failure(code: &str, message: &str) -> bool { + matches!( + code, + "token_invalid" + | "token_expired" + | "token_revoked" + | "account_deactivated" + | "unauthorized" + | "invalid_api_key" + ) || message.contains("token revoked") + || message.contains("token invalid") + || message.contains("token expired") +} + +fn is_model_unsupported(value: &str) -> bool { + let value = value.to_ascii_lowercase(); + value.contains("model_not_supported") + || value.contains("model_not_available") + || (value.contains("model") + && (value.contains("not supported") + || value.contains("not available") + || value.contains("not_supported") + || value.contains("not_available"))) +} + +fn is_html(value: &str) -> bool { + let value = value.trim_start().to_ascii_lowercase(); + value.starts_with(" bool { + serde_json::from_str::(value).is_ok_and(|value| { + value + .pointer("/detail/code") + .and_then(serde_json::Value::as_str) + == Some("deactivated_workspace") + }) +} diff --git a/backend/src/dispatch/controllers/affinity.rs b/backend/src/dispatch/controllers/affinity.rs new file mode 100644 index 000000000..50839d6bd --- /dev/null +++ b/backend/src/dispatch/controllers/affinity.rs @@ -0,0 +1,160 @@ +//! 会话亲和性路由与写入的唯一 feature owner。 + +use std::time::Duration; + +use chrono::{DateTime, Utc}; +use tokio::time::timeout; + +use crate::dispatch::affinity::{ + SessionAffinityEntry, SessionAffinityService, compute_variant_hash, hash_instructions, +}; +use crate::upstream::openai::protocol::{ + events::TokenUsage, + responses::{CodexResponsesRequest, completed_response_metadata}, +}; + +const REDIS_BEST_EFFORT_TIMEOUT: Duration = Duration::from_millis(100); + +pub(super) struct AffinityController; + +pub(super) struct StreamExit<'a> { + pub affinity: &'a SessionAffinityService, + pub conversation_id: Option<&'a str>, + pub request: &'a CodexResponsesRequest, + pub account_id: &'a str, + pub body: &'a str, + pub turn_state: Option, + pub usage: Option, + pub completed: bool, +} + +impl AffinityController { + pub(super) async fn preferred_account_id( + affinity: &SessionAffinityService, + request: &CodexResponsesRequest, + now: DateTime, + ) -> Option { + if request.previous_response_id().is_some() { + return None; + } + let conversation_id = request + .local_conversation_id + .as_deref() + .map(str::trim) + .filter(|value| !value.is_empty())?; + let variant_hash = compute_variant_hash(request); + match timeout( + REDIS_BEST_EFFORT_TIMEOUT, + affinity.lookup_latest_account_by_conversation( + conversation_id, + None, + Some(&variant_hash), + now, + ), + ) + .await + { + Ok(account_id) => account_id, + Err(_) => { + tracing::warn!( + conversation_id, + timeout_ms = REDIS_BEST_EFFORT_TIMEOUT.as_millis(), + "Timed out reading conversation affinity; continuing without preference" + ); + None + } + } + } + + pub(super) async fn leave_complete( + affinity: &SessionAffinityService, + conversation_id: Option<&str>, + request: &CodexResponsesRequest, + account_id: &str, + body: &str, + turn_state: Option, + usage: Option, + ) { + record_response_affinity( + affinity, + conversation_id, + request, + account_id, + body, + turn_state, + usage, + ) + .await; + } + + pub(super) async fn leave_stream(exit: StreamExit<'_>) { + if !exit.completed { + return; + } + record_response_affinity( + exit.affinity, + exit.conversation_id, + exit.request, + exit.account_id, + exit.body, + exit.turn_state, + exit.usage, + ) + .await; + } +} + +async fn record_response_affinity( + affinity: &SessionAffinityService, + conversation_id: Option<&str>, + request: &CodexResponsesRequest, + account_id: &str, + body: &str, + turn_state: Option, + usage: Option, +) { + let metadata = match completed_response_metadata(body) { + Ok(Some(metadata)) => metadata, + Ok(None) => return, + Err(error) => { + tracing::warn!(error = %error, "Failed to parse completed response metadata for session affinity"); + return; + } + }; + let entry = SessionAffinityEntry { + account_id: account_id.to_string(), + conversation_id: conversation_id.unwrap_or(&metadata.response_id).to_string(), + turn_state: turn_state + .filter(|value| !value.trim().is_empty()) + .or_else(|| request.turn_state.clone()), + instructions_hash: Some(hash_instructions(Some(request.instructions()))), + input_tokens: usage.map(|usage| usage.input_tokens), + function_call_ids: metadata.function_call_ids, + variant_hash: Some(compute_variant_hash(request)), + continuation_scope: if request.store() { + crate::upstream::openai::protocol::responses::PreviousResponseScope::Persisted + } else { + crate::upstream::openai::protocol::responses::PreviousResponseScope::ConnectionLocal + }, + created_at: Utc::now(), + }; + match timeout( + REDIS_BEST_EFFORT_TIMEOUT, + affinity.record(metadata.response_id.clone(), entry), + ) + .await + { + Ok(Ok(())) => {} + Ok(Err(error)) => { + tracing::warn!(error = %error, response_id = %metadata.response_id, account_id, "Failed to record session affinity"); + } + Err(_) => { + tracing::warn!( + response_id = %metadata.response_id, + account_id, + timeout_ms = REDIS_BEST_EFFORT_TIMEOUT.as_millis(), + "Timed out recording session affinity" + ); + } + } +} diff --git a/backend/src/dispatch/controllers/cloudflare.rs b/backend/src/dispatch/controllers/cloudflare.rs new file mode 100644 index 000000000..b7a4e82f8 --- /dev/null +++ b/backend/src/dispatch/controllers/cloudflare.rs @@ -0,0 +1,237 @@ +//! Cloudflare challenge、路径封禁与 Cookie 生命周期的唯一 feature owner。 + +use chrono::Utc; + +use crate::{ + dispatch::{ + failure::exhaustion::{AccountExhaustionRecord, ExhaustedAccountKind}, + lifecycle::contract::{ + AttemptDecision, AttemptObservation, AttemptObservationKind, AttemptReturnKind, + }, + lifecycle::stream::StreamTerminal, + transport::observation::UpstreamFailureFacts, + }, + fleet::{ + account::AccountStatus, + cookies::{CloudflareChallengeCooldownTracker, CloudflarePathBlockTracker, PgCookieStore}, + pool::AccountPoolService, + }, +}; + +pub(super) struct CloudflareController; + +pub(super) struct StreamExit<'a> { + pub recovery: &'a CloudflareRecovery, + pub account_id: &'a str, + pub set_cookie_headers: &'a [String], + pub terminal: &'a StreamTerminal, +} + +#[derive(Clone, Copy)] +enum CloudflareFailureKind { + Challenge, + PathBlocked, +} + +pub(super) struct CloudflareFailure { + kind: CloudflareFailureKind, + exhaustion: AccountExhaustionRecord, +} + +impl CloudflareController { + pub(super) fn classify(observation: &AttemptObservation) -> Option { + classify(observation) + } + + pub(super) async fn apply_effect( + recovery: &CloudflareRecovery, + account_pool: &AccountPoolService, + failure: &CloudflareFailure, + ) { + match failure.kind { + CloudflareFailureKind::Challenge => { + let Some(account_id) = failure.exhaustion.account_id.as_deref() else { + return; + }; + recovery.apply_challenge(account_pool, account_id).await; + } + CloudflareFailureKind::PathBlocked => { + let Some(account_id) = failure.exhaustion.account_id.as_deref() else { + return; + }; + recovery.apply_path_block(account_pool, account_id).await; + } + } + } + + pub(super) fn decision( + observation: &AttemptObservation, + failure: CloudflareFailure, + ) -> AttemptDecision { + if observation.routing.can_retry_next_candidate { + return AttemptDecision::RetryNextCandidate { + exhaustion: Some(failure.exhaustion), + on_exhaustion: None, + }; + } + AttemptDecision::Return(AttemptReturnKind::Observed) + } + + pub(super) async fn leave_complete(recovery: &CloudflareRecovery, account_id: &str) { + recovery.reset_account_recovery(account_id).await; + } + + pub(super) async fn leave_stream(exit: StreamExit<'_>) { + exit.recovery + .capture_set_cookie_headers(exit.account_id, exit.set_cookie_headers) + .await; + if matches!( + exit.terminal, + StreamTerminal::Completed { .. } | StreamTerminal::Incomplete { .. } + ) { + exit.recovery.reset_account_recovery(exit.account_id).await; + } + } +} + +fn classify(observation: &AttemptObservation) -> Option { + let account_id = observation.account.as_ref()?.id.as_str(); + let AttemptObservationKind::UpstreamFailure(facts) = &observation.kind else { + return None; + }; + classify_upstream(account_id, facts) +} + +fn classify_upstream(account_id: &str, facts: &UpstreamFailureFacts) -> Option { + if facts.status_code == Some(403) && is_challenge(&facts.body) { + return Some(CloudflareFailure { + kind: CloudflareFailureKind::Challenge, + exhaustion: AccountExhaustionRecord::new( + account_id, + ExhaustedAccountKind::CloudflareChallenge, + "Upstream blocked the request (Cloudflare challenge)", + ), + }); + } + if facts.status_code == Some(404) && facts.body.trim().is_empty() { + return Some(CloudflareFailure { + kind: CloudflareFailureKind::PathBlocked, + exhaustion: AccountExhaustionRecord::new( + account_id, + ExhaustedAccountKind::CloudflarePathBlocked, + "Upstream blocked the request (Cloudflare path-block)", + ), + }); + } + None +} + +fn is_challenge(value: &str) -> bool { + let value = value.to_ascii_lowercase(); + value.contains("cf-mitigated") + || value.contains("cf-chl-bypass") + || value.contains("_cf_chl") + || value.contains("cf_chl") + || value.contains("attention required") + || value.contains("just a moment") +} + +#[derive(Clone)] +pub struct CloudflareRecovery { + path_block_tracker: CloudflarePathBlockTracker, + challenge_tracker: CloudflareChallengeCooldownTracker, + cookie_store: PgCookieStore, +} + +impl CloudflareRecovery { + pub fn new(cookie_store: PgCookieStore) -> Self { + Self { + path_block_tracker: CloudflarePathBlockTracker::default(), + challenge_tracker: CloudflareChallengeCooldownTracker::default(), + cookie_store, + } + } + + pub async fn cookie_header_for_request(&self, account_id: &str, path: &str) -> Option { + self.cookie_store + .cookie_header_for_request(account_id, "chatgpt.com", path) + .await + .ok()? + } + + pub async fn capture_set_cookie_headers(&self, account_id: &str, headers: &[String]) { + for header in headers { + if let Err(error) = self + .cookie_store + .capture_set_cookie(account_id, header) + .await + { + tracing::warn!( + account_id, + error = %error, + "Failed to persist upstream set-cookie header" + ); + } + } + } + + pub async fn apply_challenge(&self, account_pool: &AccountPoolService, account_id: &str) { + let now = Utc::now(); + let cooldown = self + .challenge_tracker + .record_challenge(account_id, now) + .await; + account_pool + .set_cloudflare_cooldown_until(account_id, cooldown.cooldown_until) + .await; + if let Err(error) = self + .cookie_store + .expire_account_cookies_at(account_id, cooldown.cooldown_until) + .await + { + tracing::warn!( + account_id, + error = %error, + "Failed to persist Cloudflare challenge cookie cleanup deadline" + ); + } + } + + pub async fn apply_path_block(&self, account_pool: &AccountPoolService, account_id: &str) { + self.delete_account_cookies(account_id, "Cloudflare path-block") + .await; + let now = Utc::now(); + self.path_block_tracker + .record_path_block(account_id, now) + .await; + if self + .path_block_tracker + .should_disable(account_id, now) + .await + { + account_pool + .set_status(account_id, AccountStatus::Disabled) + .await; + } + } + + pub async fn reset_account_recovery(&self, account_id: &str) { + self.path_block_tracker.reset(account_id).await; + self.challenge_tracker.reset(account_id).await; + } + + async fn delete_account_cookies(&self, account_id: &str, reason: &str) { + delete_account_cookies(&self.cookie_store, account_id, reason).await; + } +} + +async fn delete_account_cookies(cookie_store: &PgCookieStore, account_id: &str, reason: &str) { + if let Err(error) = cookie_store.delete_account_cookies(account_id).await { + tracing::warn!( + account_id, + reason, + error = %error, + "Failed to delete account cookies after Cloudflare recovery signal" + ); + } +} diff --git a/backend/src/dispatch/controllers/cyber_policy/mod.rs b/backend/src/dispatch/controllers/cyber_policy/mod.rs new file mode 100644 index 000000000..ccb50fa7e --- /dev/null +++ b/backend/src/dispatch/controllers/cyber_policy/mod.rs @@ -0,0 +1,201 @@ +//! `cyber_policy` 会话级生命周期控制器。 + +use std::{collections::BTreeSet, sync::Arc, time::Duration as StdDuration}; + +use chrono::Duration; +use sha2::{Digest, Sha256}; +use tokio::time::timeout; + +mod store; +mod types; + +use crate::{ + dispatch::{ + affinity::SessionAffinityService, controllers::ControllerFailureFact, + errors::ClientFailure, transport::observation::UpstreamFailureFacts, + }, + upstream::openai::protocol::responses::{CodexResponsesRequest, ResponsesSseFailure}, +}; + +use self::{store::Store, types::SessionState}; + +const SESSION_STATE_TTL: Duration = Duration::hours(1); +const STATE_IO_TIMEOUT: StdDuration = StdDuration::from_millis(100); + +/// 单次请求使用的 `cyber_policy` 会话路由计划。 +pub(in crate::dispatch) struct CyberPolicyScope { + session_key: Option, + state: SessionState, +} + +impl CyberPolicyScope { + pub(in crate::dispatch) fn excluded_account_ids(&self) -> BTreeSet { + self.state.failed_account_ids.iter().cloned().collect() + } + + fn session_key(&self) -> Option<&str> { + self.session_key.as_deref() + } + + fn has_failures(&self) -> bool { + !self.state.failed_account_ids.is_empty() + } +} + +/// 识别失败、维护会话状态，并为下一次请求生成账号排除集合。 +#[derive(Clone)] +pub(in crate::dispatch) struct CyberPolicyController { + store: Store, +} + +impl CyberPolicyController { + pub(in crate::dispatch) fn new(session_affinity: Arc) -> Self { + Self { + store: Store::new(session_affinity.redis_connection()), + } + } + + pub(super) fn client_failure(failure: ResponsesSseFailure) -> ClientFailure { + ClientFailure::new(failure, 400, true) + } + + pub(super) fn attempt_client_failure(failure: ControllerFailureFact<'_>) -> ClientFailure { + match failure { + ControllerFailureFact::Upstream(facts) => ClientFailure::new( + failure_from_upstream_facts(facts), + facts.status_code.unwrap_or(400), + true, + ), + ControllerFailureFact::Response(failure) => Self::client_failure(failure.clone()), + } + } + + pub(in crate::dispatch) async fn prepare( + &self, + request: &CodexResponsesRequest, + ) -> CyberPolicyScope { + let Some(session_key) = session_key(request) else { + return CyberPolicyScope { + session_key: None, + state: SessionState::default(), + }; + }; + let state = match timeout(STATE_IO_TIMEOUT, self.store.load(&session_key)).await { + Ok(Ok(state)) => state.unwrap_or_default(), + Ok(Err(error)) => { + tracing::warn!(error = %error, "Failed to read cyber policy session state"); + SessionState::default() + } + Err(_) => { + tracing::warn!("Timed out reading cyber policy session state"); + SessionState::default() + } + }; + CyberPolicyScope { + session_key: Some(session_key), + state, + } + } + + pub(super) async fn exclude_account(&self, plan: &CyberPolicyScope, account_id: &str) { + let Some(session_key) = plan.session_key() else { + return; + }; + match timeout( + STATE_IO_TIMEOUT, + self.store + .record_failure(session_key, account_id, SESSION_STATE_TTL), + ) + .await + { + Ok(Ok(state)) => tracing::warn!( + account_id, + excluded_account_count = state.failed_account_ids.len(), + "Excluded cyber policy account from this session" + ), + Ok(Err(error)) => tracing::warn!( + account_id, + error = %error, + "Failed to record cyber policy session state" + ), + Err(_) => tracing::warn!(account_id, "Timed out recording cyber policy session state"), + } + } + + pub(super) fn owns_failure(&self, failure: ControllerFailureFact<'_>) -> bool { + match failure { + ControllerFailureFact::Upstream(facts) => { + facts + .status_code + .is_some_and(|status| (400..=499).contains(&status)) + && is_cyber_policy_code(facts.code.as_deref()) + } + ControllerFailureFact::Response(failure) => is_cyber_policy_failure(failure), + } + } + + pub(in crate::dispatch) async fn observe_success(&self, plan: &CyberPolicyScope) { + if !plan.has_failures() { + return; + } + let Some(session_key) = plan.session_key() else { + return; + }; + match timeout( + STATE_IO_TIMEOUT, + self.store.clear(session_key, &plan.state.revision), + ) + .await + { + Ok(Ok(_)) => {} + Ok(Err(error)) => { + tracing::warn!(error = %error, "Failed to clear cyber policy session state"); + } + Err(_) => tracing::warn!("Timed out clearing cyber policy session state"), + } + } +} + +fn session_key(request: &CodexResponsesRequest) -> Option { + if request.previous_response_id().is_some() { + return None; + } + let api_key_id = non_empty(request.client_api_key_id.as_deref())?; + let explicit_session_id = non_empty(request.client_session_id.as_deref()) + .or_else(|| non_empty(request.client_conversation_id.as_deref())) + .or_else(|| { + request + .explicit_prompt_cache_key + .then(|| non_empty(request.prompt_cache_key())) + .flatten() + })?; + let mut hasher = Sha256::new(); + hasher.update(b"cyber-policy-session\0"); + hasher.update(api_key_id.as_bytes()); + hasher.update(b"\0"); + hasher.update(explicit_session_id.as_bytes()); + Some(hex::encode(hasher.finalize())) +} + +fn is_cyber_policy_code(code: Option<&str>) -> bool { + code.is_some_and(|code| code.trim().eq_ignore_ascii_case("cyber_policy")) +} + +fn failure_from_upstream_facts(facts: &UpstreamFailureFacts) -> ResponsesSseFailure { + ResponsesSseFailure { + event: "error".to_string(), + message: facts.message.clone(), + upstream_code: facts.code.clone(), + upstream_type: None, + explicit_status_code: facts.status_code, + retry_after_seconds: facts.retry_after_seconds, + } +} + +pub(super) fn is_cyber_policy_failure(failure: &ResponsesSseFailure) -> bool { + is_cyber_policy_code(failure.upstream_code.as_deref()) +} + +fn non_empty(value: Option<&str>) -> Option<&str> { + value.map(str::trim).filter(|value| !value.is_empty()) +} diff --git a/backend/src/dispatch/controllers/cyber_policy/store.rs b/backend/src/dispatch/controllers/cyber_policy/store.rs new file mode 100644 index 000000000..68086a749 --- /dev/null +++ b/backend/src/dispatch/controllers/cyber_policy/store.rs @@ -0,0 +1,119 @@ +use chrono::Duration; +use redis::AsyncCommands; +use thiserror::Error; + +use crate::infra::redis::RedisConnection; + +use super::types::SessionState; + +const STATE_VERSION: &str = "cyber-policy:v2"; + +const RECORD_FAILURE_SCRIPT: &str = r#" +local raw = redis.call('GET', KEYS[1]) +local state = { failedAccountIds = {}, revision = '' } +if raw then + local ok, decoded = pcall(cjson.decode, raw) + if ok and type(decoded) == 'table' and type(decoded.failedAccountIds) == 'table' then + state = decoded + end +end + +local account_id = ARGV[1] +local known = false +for _, failed_account_id in ipairs(state.failedAccountIds) do + if failed_account_id == account_id then + known = true + break + end +end + +if not known then + table.insert(state.failedAccountIds, account_id) +end + +state.revision = ARGV[3] + +local encoded = cjson.encode(state) +redis.call('SET', KEYS[1], encoded, 'EX', tonumber(ARGV[2])) +return encoded +"#; + +const CLEAR_STATE_SCRIPT: &str = r#" +local raw = redis.call('GET', KEYS[1]) +if not raw then + return 0 +end +local ok, state = pcall(cjson.decode, raw) +if not ok or type(state) ~= 'table' then + return 0 +end +if tostring(state.revision or '') ~= ARGV[1] then + return 0 +end +return redis.call('DEL', KEYS[1]) +"#; + +#[derive(Clone)] +pub struct Store { + redis: RedisConnection, +} + +#[derive(Debug, Error)] +pub enum StoreError { + #[error("Redis cyber policy operation failed: {0}")] + Redis(#[from] redis::RedisError), + #[error("invalid Redis cyber policy value: {0}")] + Json(#[from] serde_json::Error), +} + +impl Store { + pub fn new(redis: RedisConnection) -> Self { + Self { redis } + } + + pub async fn load(&self, session_key: &str) -> Result, StoreError> { + let mut connection = self.redis.manager(); + let value: Option = connection.get(self.key(session_key)).await?; + value + .map(|value| serde_json::from_str(&value)) + .transpose() + .map_err(Into::into) + } + + pub async fn record_failure( + &self, + session_key: &str, + account_id: &str, + ttl: Duration, + ) -> Result { + let mut connection = self.redis.manager(); + let revision = uuid::Uuid::new_v4().to_string(); + let value: String = redis::Script::new(RECORD_FAILURE_SCRIPT) + .key(self.key(session_key)) + .arg(account_id) + .arg(ttl.num_seconds().max(1)) + .arg(revision) + .invoke_async(&mut connection) + .await?; + Ok(serde_json::from_str(&value)?) + } + + pub async fn clear( + &self, + session_key: &str, + expected_revision: &str, + ) -> Result { + let mut connection = self.redis.manager(); + let removed: usize = redis::Script::new(CLEAR_STATE_SCRIPT) + .key(self.key(session_key)) + .arg(expected_revision) + .invoke_async(&mut connection) + .await?; + Ok(removed > 0) + } + + fn key(&self, session_key: &str) -> String { + self.redis + .key(&format!("{STATE_VERSION}:session:{session_key}")) + } +} diff --git a/backend/src/dispatch/controllers/cyber_policy/types.rs b/backend/src/dispatch/controllers/cyber_policy/types.rs new file mode 100644 index 000000000..a6bf1cd79 --- /dev/null +++ b/backend/src/dispatch/controllers/cyber_policy/types.rs @@ -0,0 +1,9 @@ +use serde::{Deserialize, Serialize}; + +#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct SessionState { + pub failed_account_ids: Vec, + #[serde(default)] + pub revision: String, +} diff --git a/backend/src/dispatch/controllers/history.rs b/backend/src/dispatch/controllers/history.rs new file mode 100644 index 000000000..3a04f4df1 --- /dev/null +++ b/backend/src/dispatch/controllers/history.rs @@ -0,0 +1,400 @@ +//! previous response 所有权、重放与恢复规则的唯一 owner。 + +use std::{sync::Arc, time::Duration}; + +use chrono::Utc; +use serde_json::Value; +use tokio::time::timeout; + +use crate::{ + dispatch::{ + affinity::{SessionAffinityEntry, SessionAffinityService}, + lifecycle::contract::{ + AttemptDecision, AttemptObservation, AttemptObservationKind, AttemptReturnKind, + AttemptRoutingFacts, CompleteResponseFacts, PinnedCandidateAcquireFailureKind, + }, + transport::observation::{UpstreamFailureFacts, UpstreamFailureKind}, + }, + upstream::openai::{ + protocol::responses::{CodexResponsesRequest, StreamCommitPolicy}, + transport::websocket::PreviousResponseUnavailableReason, + }, +}; + +const REDIS_BEST_EFFORT_TIMEOUT: Duration = Duration::from_millis(100); +const CROSS_ACCOUNT_HISTORY_UNAVAILABLE_MESSAGE: &str = + "previous response history cannot be sent to another account"; + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum HistorySource { + None, + Managed, + ConnectionLocal, + ExternalUnknown, +} + +/// 单个请求内的 previous response 状态；策略只由 `HistoryController` 解释。 +pub(super) struct HistoryState { + source: HistorySource, + managed_entry: Option, + local_replay_input: Option>>, + external_attempted_account_id: Option, + replay_owner_account: bool, +} + +#[derive(Default)] +pub(in crate::dispatch) struct ConnectionReplaySnapshot { + last_response_id: Option, + input: Option>>, +} + +pub(in crate::dispatch) struct ConnectionReplayPlan { + update: ConnectionReplayUpdate, +} + +enum ConnectionReplayUpdate { + Replace(Vec), + Append { + history: Arc>, + input: Vec, + }, + Unavailable, +} + +pub(in crate::dispatch) struct HistoryController; + +impl HistoryController { + pub(in crate::dispatch) fn new_connection_replay_snapshot() -> ConnectionReplaySnapshot { + ConnectionReplaySnapshot::default() + } + + pub(in crate::dispatch) fn prepare_connection_replay( + snapshot: &ConnectionReplaySnapshot, + request: &mut CodexResponsesRequest, + ) -> ConnectionReplayPlan { + let turn_input = request.input().to_vec(); + let update = match request.previous_response_id() { + None => { + request.local_replay_input = None; + ConnectionReplayUpdate::Replace(turn_input) + } + Some(previous_response_id) + if snapshot.last_response_id.as_deref() == Some(previous_response_id) => + { + if let Some(history) = snapshot.input.as_ref() { + request.local_replay_input = Some(Arc::clone(history)); + ConnectionReplayUpdate::Append { + history: Arc::clone(history), + input: turn_input, + } + } else { + request.local_replay_input = None; + ConnectionReplayUpdate::Unavailable + } + } + Some(_) => { + request.local_replay_input = None; + ConnectionReplayUpdate::Unavailable + } + }; + ConnectionReplayPlan { update } + } + + pub(in crate::dispatch) fn commit_connection_replay( + snapshot: &mut ConnectionReplaySnapshot, + plan: ConnectionReplayPlan, + response_id: String, + output: Vec, + ) { + snapshot.last_response_id = Some(response_id); + match plan.update { + ConnectionReplayUpdate::Replace(input) => { + snapshot.input = Some(Arc::new(sanitized_transcript(input, output))); + } + ConnectionReplayUpdate::Append { mut history, input } => { + let mut turn = sanitized_transcript(input, output); + Arc::make_mut(&mut history).append(&mut turn); + snapshot.input = Some(history); + } + ConnectionReplayUpdate::Unavailable => snapshot.input = None, + } + } + + pub(super) async fn enter( + affinity: &SessionAffinityService, + request: &CodexResponsesRequest, + ) -> HistoryState { + let Some(previous_response_id) = request.previous_response_id() else { + return HistoryState::empty(); + }; + let managed_entry = match timeout( + REDIS_BEST_EFFORT_TIMEOUT, + affinity.lookup(previous_response_id, Utc::now()), + ) + .await + { + Ok(entry) => entry, + Err(_) => { + tracing::warn!( + previous_response_id, + timeout_ms = REDIS_BEST_EFFORT_TIMEOUT.as_millis(), + "Timed out reading previous response affinity; continuing without managed history" + ); + None + } + }; + let local_replay_input = request.local_replay_input.clone(); + let source = match (&managed_entry, &local_replay_input) { + (Some(_), _) => HistorySource::Managed, + (None, Some(_)) => HistorySource::ConnectionLocal, + (None, None) => HistorySource::ExternalUnknown, + }; + HistoryState { + source, + managed_entry, + local_replay_input, + external_attempted_account_id: None, + replay_owner_account: false, + } + } + + pub(super) fn preferred_account_id(state: &HistoryState) -> Option<&str> { + state + .managed_entry + .as_ref() + .map(|entry| entry.account_id.as_str()) + } + + pub(super) fn prepare_attempt( + state: &mut HistoryState, + request: &CodexResponsesRequest, + account_id: &str, + ) -> Result { + let prepared = match state.source { + HistorySource::None => Some(request.clone()), + HistorySource::ExternalUnknown => { + match state.external_attempted_account_id.as_deref() { + None => { + state.external_attempted_account_id = Some(account_id.to_string()); + Some(external_history_request(request)) + } + Some(first_account_id) if first_account_id == account_id => { + Some(external_history_request(request)) + } + Some(_) => None, + } + } + HistorySource::ConnectionLocal => state + .local_replay_input + .as_deref() + .map(|replay_input| full_replay_request(request, replay_input)), + HistorySource::Managed => { + let entry = state.managed_entry.as_ref(); + match entry { + Some(entry) + if entry.account_id == account_id && !state.replay_owner_account => + { + let mut prepared = request.clone(); + prepared.local_conversation_id = Some(entry.conversation_id.clone()); + prepared.stream_commit_policy = StreamCommitPolicy::UntilOutputOrTerminal; + prepared.previous_response_scope = Some(entry.continuation_scope); + if prepared.turn_state.as_deref().is_none_or(str::is_empty) { + prepared.turn_state.clone_from(&entry.turn_state); + } + Some(prepared) + } + Some(entry) => state.local_replay_input.as_deref().map(|replay_input| { + let mut prepared = full_replay_request(request, replay_input); + prepared.local_conversation_id = Some(entry.conversation_id.clone()); + prepared + }), + None => None, + } + } + }; + prepared.ok_or_else(|| CROSS_ACCOUNT_HISTORY_UNAVAILABLE_MESSAGE.to_string()) + } + + pub(super) fn prepare_same_account_retry(state: &mut HistoryState, account_id: &str) -> bool { + if !Self::can_retry_same_account(state, account_id) { + return false; + } + state.replay_owner_account = true; + true + } + + pub(super) fn routing_facts( + state: &HistoryState, + account_id: Option<&str>, + ) -> AttemptRoutingFacts { + AttemptRoutingFacts { + external_origin: state.source == HistorySource::ExternalUnknown, + can_retry_same_account: account_id + .is_some_and(|account_id| Self::can_retry_same_account(state, account_id)), + can_retry_next_candidate: match state.source { + HistorySource::None | HistorySource::ConnectionLocal => true, + HistorySource::Managed => state.local_replay_input.is_some(), + HistorySource::ExternalUnknown => false, + }, + } + } + + pub(super) fn conversation_id<'a>( + state: &'a HistoryState, + request: &'a CodexResponsesRequest, + ) -> Option<&'a str> { + state + .managed_entry + .as_ref() + .map(|entry| entry.conversation_id.as_str()) + .or_else(|| non_empty(request.local_conversation_id.as_deref())) + } + + pub(super) fn decide(observation: &AttemptObservation) -> Option { + if let AttemptObservationKind::PinnedCandidateUnavailable { kind, .. } = observation.kind { + return Some(match kind { + PinnedCandidateAcquireFailureKind::Busy => { + AttemptDecision::Return(AttemptReturnKind::ContinuationBusy) + } + PinnedCandidateAcquireFailureKind::Unavailable => { + AttemptDecision::Return(AttemptReturnKind::RouteUnavailable { + message: "the account that owns the previous response is unavailable for same-account recovery".to_string(), + }) + } + }); + } + let detail = match &observation.kind { + AttemptObservationKind::UpstreamFailure(facts) if is_continuation_busy(facts) => { + if observation.routing.can_retry_same_account { + return Some(AttemptDecision::RetrySameAccount); + } + return Some(AttemptDecision::Return(AttemptReturnKind::ContinuationBusy)); + } + AttemptObservationKind::UpstreamFailure(facts) if is_history_failure(facts) => { + facts.body.clone() + } + AttemptObservationKind::CompleteResponse(CompleteResponseFacts::Failed(failure)) + | AttemptObservationKind::StreamFailure(failure) + if is_history_code(failure.upstream_code.as_deref()) => + { + crate::dispatch::failure::sse::sse_failure_error_body(failure) + } + _ => return None, + }; + + if observation.routing.can_retry_same_account { + return Some(AttemptDecision::RetrySameAccount); + } + if observation.routing.external_origin { + return Some(AttemptDecision::Return(AttemptReturnKind::Observed)); + } + Some(AttemptDecision::Return( + AttemptReturnKind::RouteUnavailable { message: detail }, + )) + } + + fn can_retry_same_account(state: &HistoryState, account_id: &str) -> bool { + let Some(entry) = state.managed_entry.as_ref() else { + return false; + }; + entry.account_id == account_id + && state.local_replay_input.is_some() + && !state.replay_owner_account + } +} + +impl HistoryState { + fn empty() -> Self { + Self { + source: HistorySource::None, + managed_entry: None, + local_replay_input: None, + external_attempted_account_id: None, + replay_owner_account: false, + } + } +} + +fn external_history_request(original: &CodexResponsesRequest) -> CodexResponsesRequest { + let mut request = original.clone(); + request.stream_commit_policy = StreamCommitPolicy::UntilOutputOrTerminal; + request.previous_response_scope = + Some(crate::upstream::openai::protocol::responses::PreviousResponseScope::ExternalUnknown); + request +} + +fn full_replay_request( + original: &CodexResponsesRequest, + replay_input: &[Value], +) -> CodexResponsesRequest { + let mut request = original.clone(); + let mut input = Vec::with_capacity(replay_input.len() + original.input().len()); + input.extend_from_slice(replay_input); + input.extend_from_slice(original.input()); + sanitize_replay_items(&mut input); + request.set_input(input); + request.set_previous_response_id(None); + request.turn_state = None; + request.stream_commit_policy = StreamCommitPolicy::UntilOutputOrTerminal; + request +} + +fn sanitized_transcript(mut input: Vec, mut output: Vec) -> Vec { + sanitize_replay_items(&mut input); + sanitize_replay_items(&mut output); + input.append(&mut output); + input +} + +fn sanitize_replay_items(values: &mut [Value]) { + for value in values { + match value { + Value::Array(values) => sanitize_replay_items(values), + Value::Object(object) => { + object.remove("id"); + object.remove("encrypted_content"); + for value in object.values_mut() { + sanitize_replay_items(std::slice::from_mut(value)); + } + } + Value::Null | Value::Bool(_) | Value::Number(_) | Value::String(_) => {} + } + } +} + +fn is_continuation_busy(facts: &UpstreamFailureFacts) -> bool { + matches!( + facts.kind, + UpstreamFailureKind::ContinuationUnavailable( + PreviousResponseUnavailableReason::ConnectionBusy + ) + ) +} + +fn is_history_failure(facts: &UpstreamFailureFacts) -> bool { + facts + .code + .as_deref() + .is_some_and(|code| is_history_code(Some(code))) + || matches!( + facts.kind, + UpstreamFailureKind::ContinuationUnavailable(reason) + if reason != PreviousResponseUnavailableReason::ConnectionBusy + ) +} + +fn is_history_code(code: Option<&str>) -> bool { + matches!( + code, + Some( + "previous_response_not_found" + | "invalid_encrypted_content" + | "missing_tool_output" + | "no_tool_output" + ) + ) +} + +fn non_empty(value: Option<&str>) -> Option<&str> { + value.map(str::trim).filter(|value| !value.is_empty()) +} diff --git a/backend/src/dispatch/controllers/mod.rs b/backend/src/dispatch/controllers/mod.rs new file mode 100644 index 000000000..c9db57ad9 --- /dev/null +++ b/backend/src/dispatch/controllers/mod.rs @@ -0,0 +1,927 @@ +//! v1 请求生命周期控制器集合。 +//! +//! 控制器集合是静态类型的 feature owner 边界：生命周期只依赖这里暴露的 +//! typed hook，不直接知道具体功能名称或状态存储实现。 + +use std::{ + collections::BTreeSet, + future::Future, + sync::Arc, + time::{Duration, Instant}, +}; + +use chrono::{DateTime, Utc}; +use tokio::time::timeout; + +use crate::{ + dispatch::{ + affinity::{AccountIdentityService, SessionAffinityService}, + errors::{ClientFailure, ResponseDispatchError, upstream_error_set_cookie_headers}, + lifecycle::trace::{ResponseDispatchAttempt, ResponseDispatchTrace}, + }, + fleet::{account::Account, pool::AccountLease}, + models::service::ModelService, + telemetry::recorder::Recorder, + upstream::openai::protocol::events::TokenUsage, + upstream::openai::transport::{ + CodexBackendClient, CodexBackendResponse, CodexBackendTransport, CodexClientError, + CodexRateLimitHeaderUpdates, CodexResponseMetadata, CodexTurnStateUpdate, + CodexUpstreamDiagnostics, WebSocketPoolDecision, + }, +}; +use crate::{ + dispatch::{ + controllers::cloudflare::CloudflareRecovery, + controllers::cyber_policy::{CyberPolicyController, CyberPolicyScope}, + controllers::{ + account_failure::AccountFailureController, history::HistoryController, + history::HistoryState, quota::QuotaController, + }, + lifecycle::{ + contract::{ + AttemptDecision, AttemptObservation, AttemptObservationKind, AttemptReturnKind, + AttemptRoutingFacts, FinalOutcome, + }, + stream::StreamTerminal, + }, + transport::observation::UpstreamFailureFacts, + }, + fleet::pool::AccountPoolService, + upstream::openai::protocol::responses::{CodexResponsesRequest, ResponsesSseFailure}, +}; + +mod account_failure; +mod affinity; +pub(crate) mod cloudflare; +pub(crate) mod cyber_policy; +pub(in crate::dispatch) mod history; +mod quota; +pub(in crate::dispatch) mod telemetry; +mod usage; + +const BEST_EFFORT_IO_TIMEOUT: Duration = Duration::from_millis(100); + +pub(in crate::dispatch) struct CompleteExit<'a> { + pub request: &'a CodexResponsesRequest, + pub account_id: &'a str, + pub body: &'a str, + pub turn_state: Option, + pub usage: Option, + pub image_generation_requested: bool, + pub completed: bool, + pub recorder: &'a Recorder, + pub models: &'a ModelService, + pub request_id: &'a str, + pub route: &'a str, + pub requested_model: &'a str, + pub display_model: &'a str, + pub started_at: Instant, + pub response_id: Option<&'a str>, + pub response: &'a CodexBackendResponse, + pub trace: &'a ResponseDispatchTrace, + pub attempt: &'a ResponseDispatchAttempt, +} + +/// 已提交 stream 在 controller 洋葱退出阶段消费的稳定事实。 +/// +/// 该上下文不包含 live transport 执行器、账号 lease 或 shutdown 控制器；各 feature +/// owner 只能通过 `ControllerSet` 分发的窄 typed exit 读取所需字段。 +pub(in crate::dispatch) struct StreamControllerContext { + pub account_id: String, + pub account_plan_type: Option, + pub request_id: String, + pub route: String, + pub display_model: String, + pub requested_model: String, + pub request: CodexResponsesRequest, + pub transport: CodexBackendTransport, + pub set_cookie_headers: Vec, + pub rate_limit_headers: Vec<(String, String)>, + pub rate_limit_header_updates: Option, + pub turn_state_update: Option, + pub websocket_pool_decision: Option, + pub turn_state: Option, + pub diagnostics: CodexUpstreamDiagnostics, + pub response_metadata: CodexResponseMetadata, + pub attempt: ResponseDispatchAttempt, + pub attempts: Vec, + pub started_at: Instant, +} + +/// 单次 attempt 的原始上游失败上下文。 +/// +/// 生命周期只负责把事实交给控制器集合；Cloudflare 状态与遥测记录分别由各自 +/// controller owner 消费，避免失败信息在最终 exhaustion 折叠时丢失。 +pub(in crate::dispatch) struct AttemptUpstreamErrorContext<'a> { + pub request_id: &'a str, + pub account: &'a Account, + pub route: &'a str, + pub model: &'a str, + pub started_at: Instant, + pub stream: bool, + pub transport: crate::upstream::openai::transport::CodexBackendTransport, + pub request: &'a CodexResponsesRequest, + pub error: &'a CodexClientError, + pub trace: &'a ResponseDispatchTrace, + pub attempt: &'a ResponseDispatchAttempt, +} + +/// 请求在产生最终 dispatch error 后交给 telemetry owner 的稳定事实。 +pub(in crate::dispatch) struct DispatchErrorObservation<'a> { + pub request_id: &'a str, + pub client_api_key_id: Option<&'a str>, + pub account_id: Option<&'a str>, + pub route: &'a str, + pub model: &'a str, + pub started_at: Instant, + pub stream: bool, + pub compact: bool, + pub transport: Option<&'a str>, +} + +/// commit 前 stream failure 的 telemetry 事实；client status 已包含在 typed error 中。 +pub(in crate::dispatch) struct PrefetchedStreamFailureObservation<'a> { + pub request_id: &'a str, + pub account_id: &'a str, + pub route: &'a str, + pub model: &'a str, + pub requested_model: &'a str, + pub started_at: Instant, + pub transport: CodexBackendTransport, + pub request: &'a CodexResponsesRequest, + pub failure: &'a ResponsesSseFailure, + pub error: &'a ResponseDispatchError, + pub diagnostics: &'a CodexUpstreamDiagnostics, + pub rate_limit_headers: &'a [(String, String)], + pub prefetched: &'a [u8], + pub trace: &'a ResponseDispatchTrace, + pub attempt: &'a ResponseDispatchAttempt, +} + +/// 所有 controller owner 共同消费的失败事实。 +/// +/// attempt 与 commit 后的 stream finalizer 都必须先归一为该类型，避免同一种 +/// 上游失败因发生阶段不同而漏掉 feature 副作用。 +#[derive(Clone, Copy)] +enum ControllerFailureFact<'a> { + Upstream(&'a UpstreamFailureFacts), + Response(&'a ResponsesSseFailure), +} + +impl<'a> ControllerFailureFact<'a> { + fn from_attempt(observation: &'a AttemptObservation) -> Option { + match &observation.kind { + AttemptObservationKind::UpstreamFailure(facts) => Some(Self::Upstream(facts)), + AttemptObservationKind::CompleteResponse( + crate::dispatch::lifecycle::contract::CompleteResponseFacts::Failed(failure), + ) + | AttemptObservationKind::StreamFailure(failure) => Some(Self::Response(failure)), + _ => None, + } + } +} + +/// 单次 attempt 的唯一分类结果；副作用与 decision 必须共同消费这个值。 +enum AttemptClassification<'a> { + CyberPolicy(ControllerFailureFact<'a>), + Cloudflare(cloudflare::CloudflareFailure), + History(AttemptDecision), + AccountFailure(account_failure::ClassifiedFailure), + Quota(quota::ClassifiedQuotaFailure), + Default(AttemptDecision), +} + +enum StreamFailureClassification { + CyberPolicy(ClientFailure), + AccountFailure { + classified: account_failure::ClassifiedFailure, + client_failure: ClientFailure, + }, + Quota { + classified: quota::ClassifiedQuotaFailure, + client_failure: ClientFailure, + }, + Default(ClientFailure), +} + +impl StreamFailureClassification { + fn client_failure(&self) -> &ClientFailure { + match self { + Self::CyberPolicy(failure) + | Self::AccountFailure { + client_failure: failure, + .. + } + | Self::Quota { + client_failure: failure, + .. + } + | Self::Default(failure) => failure, + } + } +} + +/// 请求级控制器集合。使用具体类型而非动态 middleware 链，保证调用顺序和所有权可见。 +#[derive(Clone)] +pub(in crate::dispatch) struct ControllerSet { + cyber_policy: CyberPolicyController, + session_affinity: Arc, + account_pool: Arc, + cloudflare: CloudflareRecovery, + recorder: Arc, +} + +/// 控制器在单个请求中的不可变状态快照。 +pub(in crate::dispatch) struct ControllerRequestScope { + cyber_policy: CyberPolicyScope, + history: HistoryState, + usage: usage::UsageScope, +} + +/// 请求 enter 阶段的控制器结果。 +pub(in crate::dispatch) struct ControllerEnter { + pub scope: ControllerRequestScope, + pub preferred_account_id: Option, + pub excluded_account_ids: BTreeSet, +} + +pub(in crate::dispatch) struct AttemptAccountPreparationContext<'a> { + pub codex: &'a CodexBackendClient, + pub account_identity: &'a AccountIdentityService, + pub request_id: &'a str, +} + +pub(in crate::dispatch) struct PreparedAttemptAccount { + pub lease: AccountLease, + pub cookie_header: Option, +} + +pub(in crate::dispatch) enum AttemptAccountPreparation { + Ready(Box), + Rejected, +} + +pub(in crate::dispatch) enum AttemptRoutePreparation { + Ready(Box), + Unavailable { message: String }, +} + +impl ControllerSet { + pub(in crate::dispatch) async fn prepare_attempt_account( + &self, + context: AttemptAccountPreparationContext<'_>, + acquired: AccountLease, + ) -> AttemptAccountPreparation { + let usage_cookie = self + .cloudflare + .cookie_header_for_request(&acquired.account.id, "/codex/usage") + .await; + let acquired = match QuotaController::enter( + quota::QuotaEnterContext { + account_pool: &self.account_pool, + codex: context.codex, + cookie_header: usage_cookie.as_deref(), + account_identity: context.account_identity, + request_id: context.request_id, + }, + acquired, + ) + .await + { + quota::QuotaEnterOutcome::Ready(acquired) => *acquired, + quota::QuotaEnterOutcome::LimitReached => { + return AttemptAccountPreparation::Rejected; + } + }; + let cookie_header = self + .cloudflare + .cookie_header_for_request(&acquired.account.id, "/codex/responses") + .await; + AttemptAccountPreparation::Ready(Box::new(PreparedAttemptAccount { + lease: acquired, + cookie_header, + })) + } + + pub(in crate::dispatch) fn prepare_attempt_route( + &self, + scope: &mut ControllerRequestScope, + request: &CodexResponsesRequest, + account_id: &str, + ) -> AttemptRoutePreparation { + match HistoryController::prepare_attempt(&mut scope.history, request, account_id) { + Ok(request) => AttemptRoutePreparation::Ready(Box::new(request)), + Err(message) => AttemptRoutePreparation::Unavailable { message }, + } + } + + pub(in crate::dispatch) fn prepare_same_account_retry( + &self, + scope: &mut ControllerRequestScope, + account_id: &str, + ) -> bool { + HistoryController::prepare_same_account_retry(&mut scope.history, account_id) + } + + pub(in crate::dispatch) fn attempt_routing_facts( + &self, + scope: &ControllerRequestScope, + account_id: Option<&str>, + ) -> AttemptRoutingFacts { + HistoryController::routing_facts(&scope.history, account_id) + } + + pub(in crate::dispatch) async fn observe_complete_upstream( + &self, + account: &Account, + set_cookie_headers: &[String], + rate_limit_headers: &[(String, String)], + ) { + let _ = best_effort_controller_io( + "cookie.complete", + self.cloudflare + .capture_set_cookie_headers(&account.id, set_cookie_headers), + ) + .await; + let _ = best_effort_controller_io( + "quota.complete_headers", + self.account_pool + .sync_passive_rate_limit_headers(account, rate_limit_headers), + ) + .await; + } + + pub(in crate::dispatch) async fn observe_upstream_error( + &self, + context: AttemptUpstreamErrorContext<'_>, + ) { + let _ = best_effort_controller_io( + "cookie.attempt_error", + self.cloudflare.capture_set_cookie_headers( + &context.account.id, + upstream_error_set_cookie_headers(context.error), + ), + ) + .await; + let _ = best_effort_controller_io( + "telemetry.attempt_error", + telemetry::TelemetryController::observe_upstream_error(&self.recorder, &context), + ) + .await; + } + + pub(in crate::dispatch) async fn observe_dispatch_error( + &self, + observation: DispatchErrorObservation<'_>, + error: &ResponseDispatchError, + ) { + let _ = best_effort_controller_io( + "telemetry.dispatch_error", + telemetry::TelemetryController::observe_dispatch_error( + &self.recorder, + observation, + error, + ), + ) + .await; + } + + pub(in crate::dispatch) async fn observe_prefetched_stream_failure( + &self, + observation: PrefetchedStreamFailureObservation<'_>, + ) { + let _ = best_effort_controller_io( + "telemetry.prefetched_failure", + telemetry::TelemetryController::observe_prefetched_stream_failure( + &self.recorder, + observation, + ), + ) + .await; + } + + pub(in crate::dispatch) async fn finalize_complete( + &self, + scope: ControllerRequestScope, + exit: CompleteExit<'_>, + outcome: FinalOutcome, + ) { + let successful = matches!(outcome, FinalOutcome::Completed | FinalOutcome::Incomplete); + if successful { + let _ = best_effort_controller_io( + "cloudflare.complete", + cloudflare::CloudflareController::leave_complete(&self.cloudflare, exit.account_id), + ) + .await; + } + + if successful { + let _ = best_effort_controller_io( + "usage.complete", + usage::UsageController::leave_complete( + &self.account_pool, + exit.account_id, + exit.usage, + exit.image_generation_requested, + ), + ) + .await; + } + if exit.completed { + let conversation_id = HistoryController::conversation_id(&scope.history, exit.request); + let _ = best_effort_controller_io( + "affinity.complete", + affinity::AffinityController::leave_complete( + &self.session_affinity, + conversation_id, + exit.request, + exit.account_id, + exit.body, + exit.turn_state.clone(), + exit.usage, + ), + ) + .await; + } + if successful { + let _ = best_effort_controller_io( + "telemetry.complete", + telemetry::TelemetryController::leave_complete(&exit), + ) + .await; + } + exit.models + .observe_models_etag(exit.response.response_metadata.models_etag.as_deref()); + } + + pub(in crate::dispatch) fn new( + session_affinity: Arc, + account_pool: Arc, + cloudflare: CloudflareRecovery, + recorder: Arc, + ) -> Self { + Self { + cyber_policy: CyberPolicyController::new(Arc::clone(&session_affinity)), + session_affinity, + account_pool, + cloudflare, + recorder, + } + } + + pub(in crate::dispatch) async fn enter( + &self, + request: &CodexResponsesRequest, + now: DateTime, + ) -> ControllerEnter { + let usage = usage::UsageController::enter(request); + let (history, cyber_policy, affinity_preferred_account_id) = tokio::join!( + HistoryController::enter(&self.session_affinity, request), + self.cyber_policy.prepare(request), + affinity::AffinityController::preferred_account_id( + &self.session_affinity, + request, + now, + ), + ); + let preferred_account_id = HistoryController::preferred_account_id(&history) + .map(str::to_owned) + .or(affinity_preferred_account_id); + let excluded_account_ids = cyber_policy.excluded_account_ids(); + ControllerEnter { + scope: ControllerRequestScope { + cyber_policy, + history, + usage, + }, + preferred_account_id, + excluded_account_ids, + } + } + + async fn finalize_stream_success( + &self, + scope: &ControllerRequestScope, + terminal: &StreamTerminal, + ) { + match terminal { + StreamTerminal::Completed { response } | StreamTerminal::Incomplete { response } + if response.is_object() => + { + self.cyber_policy.observe_success(&scope.cyber_policy).await + } + StreamTerminal::Completed { .. } + | StreamTerminal::Incomplete { .. } + | StreamTerminal::Failed { .. } + | StreamTerminal::UpstreamClosed + | StreamTerminal::UpstreamError { .. } + | StreamTerminal::ProtocolError { .. } + | StreamTerminal::CaptureLimitExceeded + | StreamTerminal::Cancelled + | StreamTerminal::DownstreamClosed + | StreamTerminal::Shutdown => {} + } + } + + fn classify_stream_failure( + &self, + account_id: &str, + failure: &ResponsesSseFailure, + ) -> StreamFailureClassification { + let fact = ControllerFailureFact::Response(failure); + if self.cyber_policy.owns_failure(fact) { + return StreamFailureClassification::CyberPolicy( + CyberPolicyController::client_failure(failure.clone()), + ); + } + if let Some(classified) = AccountFailureController::classify_failure(account_id, fact) { + let client_failure = AccountFailureController::client_failure(&classified, failure); + return StreamFailureClassification::AccountFailure { + classified, + client_failure, + }; + } + if let Some(classified) = QuotaController::classify_failure(account_id, fact) { + return StreamFailureClassification::Quota { + classified, + client_failure: QuotaController::client_failure(failure), + }; + } + StreamFailureClassification::Default(default_response_client_failure(failure)) + } + + async fn apply_stream_failure( + &self, + scope: &ControllerRequestScope, + account_id: &str, + classification: &StreamFailureClassification, + ) { + match classification { + StreamFailureClassification::CyberPolicy(_) => { + self.cyber_policy + .exclude_account(&scope.cyber_policy, account_id) + .await; + } + StreamFailureClassification::AccountFailure { classified, .. } => { + let _ = best_effort_controller_io( + "account_failure.stream", + AccountFailureController::apply_effect(&self.account_pool, classified), + ) + .await; + } + StreamFailureClassification::Quota { classified, .. } => { + let _ = best_effort_controller_io( + "quota.stream", + QuotaController::apply_effect(&self.account_pool, classified), + ) + .await; + } + StreamFailureClassification::Default(_) => {} + } + } + + /// 按固定的反向 leave 顺序收敛 stream 终态；生命周期只调用这一 typed exit。 + pub(in crate::dispatch) async fn leave_stream( + &self, + scope: ControllerRequestScope, + context: StreamControllerContext, + summary: &crate::dispatch::lifecycle::stream::StreamSummary, + body: &str, + ) { + let stream_failure = match &summary.terminal { + StreamTerminal::Failed { failure } => { + Some(self.classify_stream_failure(&context.account_id, failure)) + } + _ => None, + }; + let client_failure_status = stream_failure + .as_ref() + .map(|failure| i64::from(failure.client_failure().status_code())); + let outcome = match &summary.terminal { + StreamTerminal::Completed { .. } => FinalOutcome::Completed, + StreamTerminal::Incomplete { .. } => FinalOutcome::Incomplete, + StreamTerminal::Failed { .. } + | StreamTerminal::UpstreamClosed + | StreamTerminal::UpstreamError { .. } + | StreamTerminal::ProtocolError { .. } + | StreamTerminal::CaptureLimitExceeded => FinalOutcome::Failed, + StreamTerminal::Cancelled | StreamTerminal::DownstreamClosed => FinalOutcome::Cancelled, + StreamTerminal::Shutdown => FinalOutcome::Shutdown, + }; + let _ = best_effort_controller_io( + "cloudflare", + cloudflare::CloudflareController::leave_stream(cloudflare::StreamExit { + recovery: &self.cloudflare, + account_id: &context.account_id, + set_cookie_headers: &context.set_cookie_headers, + terminal: &summary.terminal, + }), + ) + .await; + + let fallback = usage::StreamUsageExit { + rate_limit_headers: context.rate_limit_headers.clone(), + turn_state: context.turn_state.clone(), + }; + let mut usage_exit = best_effort_controller_io( + "usage", + usage::UsageController::leave_stream(usage::StreamExit { + account_pool: &self.account_pool, + account_id: &context.account_id, + account_plan_type: context.account_plan_type.as_deref(), + rate_limit_headers: &context.rate_limit_headers, + rate_limit_header_updates: context.rate_limit_header_updates.as_ref(), + turn_state_update: context.turn_state_update.as_ref(), + turn_state: context.turn_state.as_deref(), + usage: summary.usage, + }), + ) + .await + .unwrap_or(fallback); + + let conversation_id = HistoryController::conversation_id(&scope.history, &context.request); + affinity::AffinityController::leave_stream(affinity::StreamExit { + affinity: &self.session_affinity, + conversation_id, + request: &context.request, + account_id: &context.account_id, + body, + turn_state: usage_exit.turn_state.take(), + usage: summary.usage, + completed: matches!(summary.terminal, StreamTerminal::Completed { .. }), + }) + .await; + + let _ = best_effort_controller_io( + "telemetry", + telemetry::TelemetryController::leave_stream(telemetry::StreamExit { + context: telemetry::StreamContext { + recorder: &self.recorder, + account_id: &context.account_id, + request_id: &context.request_id, + route: &context.route, + display_model: &context.display_model, + requested_model: &context.requested_model, + request: &context.request, + transport: context.transport, + websocket_pool_decision: context.websocket_pool_decision, + diagnostics: &context.diagnostics, + response_metadata: &context.response_metadata, + attempt: &context.attempt, + attempts: &context.attempts, + started_at: context.started_at, + }, + summary, + rate_limit_headers: &usage_exit.rate_limit_headers, + failure_status: client_failure_status, + body, + }), + ) + .await; + + if !matches!(outcome, FinalOutcome::Cancelled | FinalOutcome::Shutdown) { + if let Some(classification) = stream_failure.as_ref() { + self.apply_stream_failure(&scope, &context.account_id, classification) + .await; + } else { + self.finalize_stream_success(&scope, &summary.terminal) + .await; + } + } + } + + /// 对一次 observation 只分类一次，并保证 feature effects 先于 decision 返回。 + pub(in crate::dispatch) async fn handle_attempt( + &self, + scope: &ControllerRequestScope, + observation: &AttemptObservation, + ) -> AttemptDecision { + let classification = self.classify_attempt(observation); + let _ = best_effort_controller_io( + "usage.attempt", + usage::UsageController::observe_attempt(&self.account_pool, &scope.usage, observation), + ) + .await; + + if matches!( + observation.kind, + AttemptObservationKind::CompleteResponse( + crate::dispatch::lifecycle::contract::CompleteResponseFacts::Completed + | crate::dispatch::lifecycle::contract::CompleteResponseFacts::Incomplete, + ) + ) { + self.cyber_policy.observe_success(&scope.cyber_policy).await; + } + + match &classification { + AttemptClassification::CyberPolicy(_) => { + if let Some(account_id) = observation + .account + .as_ref() + .map(|account| account.id.as_str()) + { + self.cyber_policy + .exclude_account(&scope.cyber_policy, account_id) + .await; + } + } + AttemptClassification::Cloudflare(failure) => { + let _ = best_effort_controller_io( + "cloudflare.attempt", + cloudflare::CloudflareController::apply_effect( + &self.cloudflare, + &self.account_pool, + failure, + ), + ) + .await; + } + AttemptClassification::AccountFailure(classified) => { + let _ = best_effort_controller_io( + "account_failure.attempt", + AccountFailureController::apply_effect(&self.account_pool, classified), + ) + .await; + } + AttemptClassification::Quota(classified) => { + let _ = best_effort_controller_io( + "quota.attempt", + QuotaController::apply_effect(&self.account_pool, classified), + ) + .await; + } + AttemptClassification::History(_) | AttemptClassification::Default(_) => {} + } + + self.attempt_decision(observation, classification) + } + + fn classify_attempt<'a>( + &self, + observation: &'a AttemptObservation, + ) -> AttemptClassification<'a> { + let failure = ControllerFailureFact::from_attempt(observation); + if let Some(failure) = failure + && self.cyber_policy.owns_failure(failure) + { + return AttemptClassification::CyberPolicy(failure); + } + if let Some(classified) = cloudflare::CloudflareController::classify(observation) { + return AttemptClassification::Cloudflare(classified); + } + if let Some(decision) = HistoryController::decide(observation) { + return AttemptClassification::History(decision); + } + if let Some(classified) = AccountFailureController::classify(observation) { + return AttemptClassification::AccountFailure(classified); + } + if let Some(classified) = QuotaController::classify(observation) { + return AttemptClassification::Quota(classified); + } + AttemptClassification::Default(Self::default_attempt_decision(observation)) + } + + fn attempt_decision( + &self, + observation: &AttemptObservation, + classification: AttemptClassification<'_>, + ) -> AttemptDecision { + match classification { + AttemptClassification::CyberPolicy(failure) => { + let client_failure = CyberPolicyController::attempt_client_failure(failure); + if observation.routing.can_retry_next_candidate { + AttemptDecision::RetryNextCandidate { + exhaustion: None, + on_exhaustion: Some(client_failure), + } + } else { + AttemptDecision::Return(AttemptReturnKind::Failed(client_failure)) + } + } + AttemptClassification::Cloudflare(classified) => { + cloudflare::CloudflareController::decision(observation, classified) + } + AttemptClassification::History(decision) => with_response_client_failure( + decision, + response_failure(observation) + .map(|failure| ClientFailure::new(failure.clone(), 400, false)), + ), + AttemptClassification::Default(decision) => decision, + AttemptClassification::AccountFailure(classified) => { + let client_failure = response_failure(observation) + .map(|failure| AccountFailureController::client_failure(&classified, failure)); + with_response_client_failure( + AccountFailureController::decision(observation, classified), + client_failure, + ) + } + AttemptClassification::Quota(classified) => { + let client_failure = + response_failure(observation).map(QuotaController::client_failure); + with_response_client_failure( + QuotaController::decision(observation, classified), + client_failure, + ) + } + } + } + + fn default_attempt_decision(observation: &AttemptObservation) -> AttemptDecision { + match &observation.kind { + AttemptObservationKind::CompleteResponse( + crate::dispatch::lifecycle::contract::CompleteResponseFacts::Completed + | crate::dispatch::lifecycle::contract::CompleteResponseFacts::Incomplete, + ) => AttemptDecision::Accept, + AttemptObservationKind::CompleteResponse( + crate::dispatch::lifecycle::contract::CompleteResponseFacts::Failed(failure), + ) => AttemptDecision::Return(AttemptReturnKind::Failed( + default_response_client_failure(failure), + )), + AttemptObservationKind::CompleteResponse( + crate::dispatch::lifecycle::contract::CompleteResponseFacts::MissingCompleted + | crate::dispatch::lifecycle::contract::CompleteResponseFacts::Empty, + ) => AttemptDecision::Return(AttemptReturnKind::Observed), + AttemptObservationKind::StreamFailure(failure) => AttemptDecision::Return( + AttemptReturnKind::Failed(default_response_client_failure(failure)), + ), + AttemptObservationKind::NoCandidate { .. } => { + AttemptDecision::Return(AttemptReturnKind::Observed) + } + AttemptObservationKind::CandidatePreparationRejected => { + unreachable!("candidate preparation controller owns it") + } + AttemptObservationKind::RoutePreparationRejected { message } => { + AttemptDecision::Return(AttemptReturnKind::RouteUnavailable { + message: message.clone(), + }) + } + AttemptObservationKind::PinnedCandidateUnavailable { .. } => { + unreachable!("route controller owns pinned-candidate acquisition failures") + } + AttemptObservationKind::ProtocolFailure { .. } => { + AttemptDecision::Return(AttemptReturnKind::Observed) + } + AttemptObservationKind::UpstreamFailure(facts) => { + if observation.routing.can_retry_next_candidate + && AccountFailureController::is_retryable_transport(observation) + { + return AttemptDecision::RetryNextCandidate { + exhaustion: observation.account.as_ref().map(|account| { + crate::dispatch::failure::exhaustion::AccountExhaustionRecord::new( + account.id.clone(), + crate::dispatch::failure::exhaustion::ExhaustedAccountKind::UpstreamUnavailable, + facts.message.clone(), + ) + }), + on_exhaustion: None, + }; + } + AttemptDecision::Return(AttemptReturnKind::Observed) + } + } + } +} + +fn response_failure(observation: &AttemptObservation) -> Option<&ResponsesSseFailure> { + match &observation.kind { + AttemptObservationKind::CompleteResponse( + crate::dispatch::lifecycle::contract::CompleteResponseFacts::Failed(failure), + ) + | AttemptObservationKind::StreamFailure(failure) => Some(failure), + _ => None, + } +} + +fn with_response_client_failure( + decision: AttemptDecision, + failure: Option, +) -> AttemptDecision { + match (decision, failure) { + (AttemptDecision::Return(AttemptReturnKind::Observed), Some(failure)) => { + AttemptDecision::Return(AttemptReturnKind::Failed(failure)) + } + (decision, _) => decision, + } +} + +fn default_response_client_failure(failure: &ResponsesSseFailure) -> ClientFailure { + AccountFailureController::unowned_client_failure(failure) + .unwrap_or_else(|| ClientFailure::new(failure.clone(), 502, false)) +} + +async fn best_effort_controller_io( + controller: &'static str, + operation: impl Future, +) -> Option { + match timeout(BEST_EFFORT_IO_TIMEOUT, operation).await { + Ok(output) => Some(output), + Err(_) => { + tracing::warn!( + controller, + timeout_ms = BEST_EFFORT_IO_TIMEOUT.as_millis(), + "Timed out running best-effort controller I/O" + ); + None + } + } +} diff --git a/backend/src/dispatch/controllers/quota.rs b/backend/src/dispatch/controllers/quota.rs new file mode 100644 index 000000000..a572c0220 --- /dev/null +++ b/backend/src/dispatch/controllers/quota.rs @@ -0,0 +1,346 @@ +//! 请求前配额验证规则 owner。 + +use chrono::{Duration, Utc}; +use serde_json::Value; + +use crate::{ + dispatch::{ + affinity::AccountIdentityService, + controllers::ControllerFailureFact, + errors::ClientFailure, + failure::exhaustion::{AccountExhaustionRecord, ExhaustedAccountKind}, + lifecycle::contract::{ + AttemptDecision, AttemptObservation, AttemptObservationKind, AttemptReturnKind, + }, + transport::observation::UpstreamFailureFacts, + }, + fleet::{ + pool::{AccountLease, AccountPoolService}, + quota::{ + quota_from_usage, quota_snapshot_limit_reached, quota_snapshot_limit_window_seconds, + quota_snapshot_reset_at, + }, + }, + upstream::openai::{ + protocol::responses::ResponsesSseFailure, + transport::{CodexBackendClient, CodexRequestContext}, + }, +}; + +const LIMIT_REACHED_MESSAGE: &str = "Upstream usage quota still reports limit_reached"; +const DEFAULT_RATE_LIMIT_RETRY_SECONDS: u64 = 60; + +pub(super) enum QuotaEnterOutcome { + Ready(Box), + LimitReached, +} + +pub(super) struct QuotaEnterContext<'a> { + pub account_pool: &'a AccountPoolService, + pub codex: &'a CodexBackendClient, + pub account_identity: &'a AccountIdentityService, + pub request_id: &'a str, + pub cookie_header: Option<&'a str>, +} + +pub(super) struct QuotaController; + +impl QuotaController { + pub(super) async fn enter( + context: QuotaEnterContext<'_>, + acquired: AccountLease, + ) -> QuotaEnterOutcome { + if !acquired.account.quota_verify_required { + return QuotaEnterOutcome::Ready(Box::new(acquired)); + } + + let account_id = acquired.account.id.clone(); + let identity = context + .account_identity + .scope_auxiliary(&account_id, context.request_id); + let usage = context + .codex + .fetch_usage(CodexRequestContext { + access_token: &acquired.account.access_token, + account_id: acquired.account.account_id.as_deref(), + request_id: context.request_id, + turn_state: None, + turn_metadata: None, + beta_features: None, + include_timing_metrics: None, + version: None, + codex_window_id: None, + parent_thread_id: None, + cookie_header: context.cookie_header, + installation_id: Some(&identity.installation_id), + session_id: None, + thread_id: None, + prompt_cache_key: None, + client_request_id: Some(&identity.client_request_id), + turn_id: None, + }) + .await; + + let raw = match usage { + Ok(raw) => raw, + Err(error) => { + tracing::warn!( + request_id = %context.request_id, + account_id = %account_id, + quota_verify_required = true, + quota_verify_result = "upstream_error", + retry_with_another_account = false, + error = %error, + "Failed to verify stale quota state before upstream request" + ); + return QuotaEnterOutcome::Ready(Box::new(acquired)); + } + }; + + let quota = quota_from_usage(&raw); + context + .account_pool + .apply_quota_snapshot(&account_id, "a) + .await; + if quota_snapshot_limit_reached("a) { + acquired.release_without_usage().await; + tracing::info!( + request_id = %context.request_id, + account_id = %account_id, + quota_verify_required = true, + quota_verify_result = "limit_reached", + retry_with_another_account = true, + "Quota verification reported exhausted account before upstream request" + ); + return QuotaEnterOutcome::LimitReached; + } + + QuotaEnterOutcome::Ready(Box::new(acquired_with_verified_quota(acquired, "a))) + } + + pub(super) fn classify(observation: &AttemptObservation) -> Option { + if matches!( + observation.kind, + AttemptObservationKind::CandidatePreparationRejected + ) { + return observation + .account + .as_ref() + .map(|account| ClassifiedQuotaFailure { + exhaustion: AccountExhaustionRecord::new( + account.id.clone(), + ExhaustedAccountKind::RateLimited, + LIMIT_REACHED_MESSAGE, + ), + effect: QuotaEffect::None, + }); + } + classify(observation) + } + + pub(super) fn classify_failure( + account_id: &str, + failure: ControllerFailureFact<'_>, + ) -> Option { + classify_failure(account_id, failure) + } + + pub(super) fn decision( + observation: &AttemptObservation, + classified: ClassifiedQuotaFailure, + ) -> AttemptDecision { + if observation.routing.can_retry_next_candidate { + return AttemptDecision::RetryNextCandidate { + exhaustion: Some(classified.exhaustion), + on_exhaustion: None, + }; + } + AttemptDecision::Return(AttemptReturnKind::Observed) + } + + pub(super) async fn apply_effect( + account_pool: &AccountPoolService, + classified: &ClassifiedQuotaFailure, + ) { + let Some(account_id) = classified.exhaustion.account_id.as_deref() else { + return; + }; + match classified.effect { + QuotaEffect::None => {} + QuotaEffect::SetExhausted => { + account_pool + .set_status( + account_id, + crate::fleet::account::AccountStatus::QuotaExhausted, + ) + .await; + } + QuotaEffect::MarkLimitedUntil(until) => { + account_pool + .mark_quota_limited_until(account_id, until) + .await; + } + } + } + + pub(super) fn client_failure(failure: &ResponsesSseFailure) -> ClientFailure { + ClientFailure::new(failure.clone(), 429, false) + } +} + +enum QuotaEffect { + None, + SetExhausted, + MarkLimitedUntil(chrono::DateTime), +} + +pub(super) struct ClassifiedQuotaFailure { + exhaustion: AccountExhaustionRecord, + effect: QuotaEffect, +} + +fn classify(observation: &AttemptObservation) -> Option { + let account_id = observation.account.as_ref()?.id.as_str(); + ControllerFailureFact::from_attempt(observation) + .and_then(|failure| classify_failure(account_id, failure)) +} + +fn classify_failure( + account_id: &str, + failure: ControllerFailureFact<'_>, +) -> Option { + match failure { + ControllerFailureFact::Upstream(facts) => classify_upstream(account_id, facts), + ControllerFailureFact::Response(failure) => classify_stream(account_id, failure), + } +} + +fn classify_upstream( + account_id: &str, + facts: &UpstreamFailureFacts, +) -> Option { + if facts.status_code == Some(429) { + return Some(rate_limited_failure( + account_id, + facts.body.clone(), + facts.retry_after_seconds, + )); + } + if facts.status_code == Some(402) { + return Some(quota_exhausted_failure(account_id, facts.body.clone())); + } + None +} + +fn classify_stream( + account_id: &str, + failure: &ResponsesSseFailure, +) -> Option { + let kind = match failure.explicit_status_code { + Some(429) => StreamQuotaKind::RateLimited, + Some(402) => StreamQuotaKind::QuotaExhausted, + Some(_) => return None, + None => classify_stream_signal(failure)?, + }; + let body = crate::dispatch::failure::sse::sse_failure_error_body(failure); + Some(match kind { + StreamQuotaKind::RateLimited => { + rate_limited_failure(account_id, body, failure.retry_after_seconds) + } + StreamQuotaKind::QuotaExhausted => quota_exhausted_failure(account_id, body), + }) +} + +#[derive(Clone, Copy)] +enum StreamQuotaKind { + RateLimited, + QuotaExhausted, +} + +fn classify_stream_signal(failure: &ResponsesSseFailure) -> Option { + failure + .upstream_code + .as_deref() + .and_then(classify_quota_signal) + .or_else(|| { + failure + .upstream_type + .as_deref() + .and_then(classify_quota_signal) + }) + .or_else(|| classify_quota_message(&failure.message)) +} + +fn classify_quota_signal(signal: &str) -> Option { + let signal = signal.trim().to_ascii_lowercase(); + match signal.as_str() { + "usage_limit_reached" + | "rate_limit_exceeded" + | "rate_limit_reached" + | "rate_limit_error" => Some(StreamQuotaKind::RateLimited), + "quota_exhausted" | "quota_exceeded" | "payment_required" | "insufficient_quota" => { + Some(StreamQuotaKind::QuotaExhausted) + } + signal if signal.starts_with("billing_limit") => Some(StreamQuotaKind::QuotaExhausted), + _ => None, + } +} + +fn classify_quota_message(message: &str) -> Option { + let message = message.to_ascii_lowercase(); + if message.contains("rate limit") || message.contains("usage limit") { + return Some(StreamQuotaKind::RateLimited); + } + if message.contains("quota") + || message.contains("payment required") + || message.contains("billing limit") + { + return Some(StreamQuotaKind::QuotaExhausted); + } + None +} + +fn rate_limited_failure( + account_id: &str, + message: String, + retry_after_seconds: Option, +) -> ClassifiedQuotaFailure { + let seconds = retry_after_seconds + .unwrap_or(DEFAULT_RATE_LIMIT_RETRY_SECONDS) + .min(i64::MAX as u64) as i64; + ClassifiedQuotaFailure { + exhaustion: AccountExhaustionRecord::new( + account_id, + ExhaustedAccountKind::RateLimited, + message, + ), + effect: QuotaEffect::MarkLimitedUntil(Utc::now() + Duration::seconds(seconds)), + } +} + +fn quota_exhausted_failure(account_id: &str, message: String) -> ClassifiedQuotaFailure { + ClassifiedQuotaFailure { + exhaustion: AccountExhaustionRecord::new( + account_id, + ExhaustedAccountKind::QuotaExhausted, + message, + ), + effect: QuotaEffect::SetExhausted, + } +} + +fn acquired_with_verified_quota(mut acquired: AccountLease, quota: &Value) -> AccountLease { + let limit_reached = quota_snapshot_limit_reached(quota); + acquired.account.quota_verify_required = false; + acquired.account.quota_limit_reached = limit_reached; + acquired.account.quota_cooldown_until = limit_reached + .then_some(quota_snapshot_reset_at(quota)) + .flatten(); + if let Some(reset_at) = quota_snapshot_reset_at(quota) { + acquired.account.window_reset_at = Some(reset_at); + if let Some(limit_window_seconds) = quota_snapshot_limit_window_seconds(quota) { + acquired.account.limit_window_seconds = Some(limit_window_seconds); + } + } + acquired +} diff --git a/backend/src/dispatch/recording.rs b/backend/src/dispatch/controllers/telemetry/events.rs similarity index 89% rename from backend/src/dispatch/recording.rs rename to backend/src/dispatch/controllers/telemetry/events.rs index 0abc2ea81..6c91dd545 100644 --- a/backend/src/dispatch/recording.rs +++ b/backend/src/dispatch/controllers/telemetry/events.rs @@ -27,17 +27,16 @@ use crate::{ }, }; -use super::{ +use crate::dispatch::{ errors::{ ResponseDispatchError, dispatch_error_metadata, enrich_response_dispatch_error_metadata, }, - stream::{ - live::{LiveResponseStreamContext, latest_response_id}, - sse_failure::{status_code_for_stream_failure, stream_failure_metadata}, - trace::{ResponseDispatchAttempt, ResponseDispatchTrace}, - }, + failure::sse::stream_failure_metadata, + lifecycle::trace::{ResponseDispatchAttempt, ResponseDispatchTrace}, }; +use super::StreamContext; + pub(super) struct ResponseUpstreamErrorEventRecord<'a> { pub(super) recorder: &'a Recorder, pub(super) request_id: &'a str, @@ -65,6 +64,7 @@ pub(super) struct ResponseStreamFailureEventRecord<'a> { pub(super) transport: CodexBackendTransport, pub(super) request: &'a CodexResponsesRequest, pub(super) failure: &'a ResponsesSseFailure, + pub(super) status_code: i64, pub(super) diagnostics: &'a CodexUpstreamDiagnostics, pub(super) rate_limit_headers: &'a [(String, String)], pub(super) prefetched: &'a [u8], @@ -86,14 +86,6 @@ pub(super) struct ResponseDispatchErrorEventRecord<'a> { pub(super) error: &'a ResponseDispatchError, } -pub(super) struct ResponseDispatchErrorDetails<'a> { - pub(super) client_api_key_id: Option<&'a str>, - pub(super) account_id: Option<&'a str>, - pub(super) stream: bool, - pub(super) compact: bool, - pub(super) transport: Option<&'a str>, -} - pub(super) async fn record_response_dispatch_error_event( record: ResponseDispatchErrorEventRecord<'_>, ) { @@ -188,7 +180,7 @@ pub(super) async fn record_response_upstream_error_event( pub(super) async fn record_prefetched_response_stream_failure_event( record: ResponseStreamFailureEventRecord<'_>, ) { - let event_status_code = status_code_for_stream_failure(record.failure); + let event_status_code = record.status_code; let mut metadata = stream_failure_metadata(record.failure, None); if let Some(object) = metadata.as_object_mut() { if record.transport == CodexBackendTransport::WebSocket { @@ -266,7 +258,7 @@ fn ensure_stream_metadata_flag(metadata: &mut Value) { } fn enrich_live_response_stream_metadata( - context: &LiveResponseStreamContext, + context: &StreamContext<'_>, rate_limit_headers: &[(String, String)], metadata: &mut Value, status_code: i64, @@ -292,11 +284,11 @@ fn enrich_live_response_stream_metadata( object, status_code, 200, - diagnostics_status_code(&context.diagnostics), + diagnostics_status_code(context.diagnostics), ); - insert_response_request_summary_object(object, &context.request, context.transport); - insert_upstream_diagnostics_metadata(object, &context.diagnostics); - insert_response_trace_metadata_object(object, &context.attempts, Some(&context.attempt)); + insert_response_request_summary_object(object, context.request, context.transport); + insert_upstream_diagnostics_metadata(object, context.diagnostics); + insert_response_trace_metadata_object(object, context.attempts, Some(context.attempt)); if context.recorder.captures_body() { object .entry("requestBody".to_string()) @@ -476,7 +468,7 @@ fn upstream_failure_client_status_code(upstream_status: i64) -> i64 { } pub(super) async fn record_live_response_stream_event( - context: &LiveResponseStreamContext, + context: &StreamContext<'_>, status_code: i64, level: UsageRecordLevel, message: &str, @@ -488,7 +480,7 @@ pub(super) async fn record_live_response_stream_event( .response_metadata .effective_model .as_deref() - .unwrap_or(&context.display_model); + .unwrap_or(context.display_model); ensure_stream_metadata_flag(&mut metadata); if let Some(object) = metadata.as_object_mut() { object.insert( @@ -508,7 +500,7 @@ pub(super) async fn record_live_response_stream_event( Value::Bool(context.response_metadata.reasoning_included), ); } - enrich_event_route_metadata(&mut metadata, &context.route); + enrich_event_route_metadata(&mut metadata, context.route); enrich_live_response_stream_metadata( context, rate_limit_headers, @@ -516,31 +508,31 @@ pub(super) async fn record_live_response_stream_event( status_code, body, ); - log_live_response_stream_finalized(context, status_code, level, message, &metadata, body); + log_live_response_stream_finalized(context, status_code, level, message, &metadata); enrich_usage_record_identity( &mut metadata, - Some(&context.requested_model), + Some(context.requested_model), effective_model, - context.client_ip.as_deref(), + context.request.client_ip.as_deref(), context.request.client_user_agent.as_deref(), - reasoning_effort_from_request(&context.request), + reasoning_effort_from_request(context.request), context.request.service_tier(), ); - enrich_response_request_semantics(&mut metadata, &context.request); + enrich_response_request_semantics(&mut metadata, context.request); if is_success_usage_event(level, status_code) { crate::telemetry::recorder::record_response_event( crate::telemetry::usage::types::ResponseUsageRecord { - recorder: &context.recorder, - request_id: &context.request_id, + recorder: context.recorder, + request_id: context.request_id, client_api_key_id: context.request.client_api_key_id.as_deref(), - account_id: &context.account_id, - route: &context.route, + account_id: context.account_id, + route: context.route, model: effective_model, - requested_model: Some(&context.requested_model), - client_ip: context.client_ip.as_deref(), + requested_model: Some(context.requested_model), + client_ip: context.request.client_ip.as_deref(), client_user_agent: context.request.client_user_agent.as_deref(), - reasoning_effort: reasoning_effort_from_request(&context.request), + reasoning_effort: reasoning_effort_from_request(context.request), service_tier: context.request.service_tier(), started_at: context.started_at, status_code, @@ -551,12 +543,12 @@ pub(super) async fn record_live_response_stream_event( ) .await; } else { - let mut event = OpsErrorLog::new(response_event_kind(&context.route), message); - event.request_id = Some(context.request_id.clone()); + let mut event = OpsErrorLog::new(response_event_kind(context.route), message); + event.request_id = Some(context.request_id.to_string()); event.client_api_key_id = context.request.client_api_key_id.clone(); event.provider = Some("openai".to_string()); - event.account_id = Some(context.account_id.clone()); - event.route = Some(context.route.clone()); + event.account_id = Some(context.account_id.to_string()); + event.route = Some(context.route.to_string()); event.model = Some(effective_model.to_string()); event.status_code = Some(status_code); event.latency_ms = Some(elapsed_millis_i64(context.started_at)); @@ -572,16 +564,13 @@ fn is_success_usage_event(level: UsageRecordLevel, status_code: i64) -> bool { } fn log_live_response_stream_finalized( - context: &LiveResponseStreamContext, + context: &StreamContext<'_>, status_code: i64, level: UsageRecordLevel, message: &str, metadata: &Value, - body: &str, ) { - let response_id = metadata_string_field(metadata, "responseId") - .map(ToString::to_string) - .or_else(|| latest_response_id(body)); + let response_id = metadata_string_field(metadata, "responseId").map(ToString::to_string); let first_token_ms = metadata.get("firstTokenMs").and_then(Value::as_i64); let websocket_pool_kind = context .websocket_pool_decision @@ -644,25 +633,6 @@ fn metadata_string_field<'a>(metadata: &'a Value, field: &str) -> Option<&'a str metadata.get(field).and_then(Value::as_str) } -pub(super) async fn live_response_rate_limit_headers( - context: &LiveResponseStreamContext, -) -> Vec<(String, String)> { - let mut headers = context.rate_limit_headers.clone(); - if let Some(updates) = &context.rate_limit_header_updates { - headers.extend(updates.lock().await.iter().cloned()); - } - headers -} - -pub(super) async fn live_response_turn_state( - context: &LiveResponseStreamContext, -) -> Option { - if let Some(update) = &context.turn_state_update { - return update.lock().await.clone(); - } - context.turn_state.clone() -} - pub(super) fn insert_first_token_ms(metadata: &mut Value, first_token_ms: Option) { let Some(first_token_ms) = first_token_ms else { return; diff --git a/backend/src/dispatch/controllers/telemetry/mod.rs b/backend/src/dispatch/controllers/telemetry/mod.rs new file mode 100644 index 000000000..efeb2d237 --- /dev/null +++ b/backend/src/dispatch/controllers/telemetry/mod.rs @@ -0,0 +1,349 @@ +//! 请求终态遥测的唯一 feature owner。 + +mod events; + +use serde_json::Value; + +use crate::{ + dispatch::{ + controllers::telemetry::events::{ + ResponseUpstreamErrorEventRecord, insert_first_token_ms, + insert_response_status_metadata, insert_response_trace_metadata, + insert_response_upstream_diagnostics, insert_websocket_pool_decision, + record_live_response_stream_event, record_response_upstream_error_event, + }, + controllers::{ + AttemptUpstreamErrorContext, CompleteExit, DispatchErrorObservation, + PrefetchedStreamFailureObservation, + }, + errors::{ResponseDispatchError, backend_transport_name}, + failure::sse::{ + STREAM_DISCONNECTED_CODE, STREAM_DISCONNECTED_MESSAGE, stream_failure_metadata, + stream_failure_source, synthetic_stream_disconnected_detail, + }, + lifecycle::{ + stream::{StreamSummary, StreamTerminal}, + trace::ResponseDispatchAttempt, + }, + }, + infra::time::elapsed_millis_i64, + telemetry::{ + recorder::{ + Recorder, enrich_response_request_semantics, reasoning_effort_from_request, + record_response_event, + }, + usage::types::{ResponseUsageRecord, UsageRecordLevel}, + }, + upstream::openai::{ + protocol::responses::{CodexResponsesRequest, ResponsesSseFailure}, + transport::{ + CodexBackendTransport, CodexResponseMetadata, CodexUpstreamDiagnostics, + WebSocketPoolDecision, + }, + }, +}; + +pub(super) struct TelemetryController; + +pub(super) struct StreamContext<'a> { + pub recorder: &'a Recorder, + pub account_id: &'a str, + pub request_id: &'a str, + pub route: &'a str, + pub display_model: &'a str, + pub requested_model: &'a str, + pub request: &'a CodexResponsesRequest, + pub transport: CodexBackendTransport, + pub websocket_pool_decision: Option, + pub diagnostics: &'a CodexUpstreamDiagnostics, + pub response_metadata: &'a CodexResponseMetadata, + pub attempt: &'a ResponseDispatchAttempt, + pub attempts: &'a [ResponseDispatchAttempt], + pub started_at: std::time::Instant, +} + +pub(super) struct StreamExit<'a> { + pub context: StreamContext<'a>, + pub summary: &'a StreamSummary, + pub rate_limit_headers: &'a [(String, String)], + pub failure_status: Option, + pub body: &'a str, +} + +impl TelemetryController { + pub(super) async fn observe_dispatch_error( + recorder: &Recorder, + observation: DispatchErrorObservation<'_>, + error: &ResponseDispatchError, + ) { + events::record_response_dispatch_error_event(events::ResponseDispatchErrorEventRecord { + recorder, + request_id: observation.request_id, + client_api_key_id: observation.client_api_key_id, + account_id: observation.account_id, + route: observation.route, + model: observation.model, + started_at: observation.started_at, + stream: observation.stream, + compact: observation.compact, + transport: observation.transport, + error, + }) + .await; + } + + pub(super) async fn observe_prefetched_stream_failure( + recorder: &Recorder, + observation: PrefetchedStreamFailureObservation<'_>, + ) { + events::record_prefetched_response_stream_failure_event( + events::ResponseStreamFailureEventRecord { + recorder, + request_id: observation.request_id, + account_id: observation.account_id, + route: observation.route, + model: observation.model, + requested_model: observation.requested_model, + started_at: observation.started_at, + transport: observation.transport, + request: observation.request, + failure: observation.failure, + status_code: i64::from(observation.error.client_http_status_code()), + diagnostics: observation.diagnostics, + rate_limit_headers: observation.rate_limit_headers, + prefetched: observation.prefetched, + trace: observation.trace, + attempt: observation.attempt, + }, + ) + .await; + } + + pub(super) async fn observe_upstream_error( + recorder: &Recorder, + context: &AttemptUpstreamErrorContext<'_>, + ) { + record_response_upstream_error_event(ResponseUpstreamErrorEventRecord { + recorder, + request_id: context.request_id, + account_id: &context.account.id, + account_email: context.account.email.as_deref(), + route: context.route, + model: context.model, + started_at: context.started_at, + stream: context.stream, + transport: context.transport, + request: context.request, + error: context.error, + trace: context.trace, + attempt: Some(context.attempt), + }) + .await; + } + + pub(super) async fn leave_complete(exit: &CompleteExit<'_>) { + let effective_model = exit + .response + .response_metadata + .effective_model + .as_deref() + .unwrap_or(exit.display_model); + let mut metadata = serde_json::json!({ + "responseId": exit.response_id, + "stream": false, + "completed": exit.completed, + "incomplete": !exit.completed, + "transport": backend_transport_name(exit.response.transport), + "firstTokenMs": exit.response.first_token_ms, + "usage": exit.response.usage, + "effectiveModel": effective_model, + "modelsEtag": exit.response.response_metadata.models_etag.as_deref(), + "reasoningIncluded": exit.response.response_metadata.reasoning_included, + }); + insert_response_status_metadata( + &mut metadata, + 200, + 200, + exit.response.diagnostics.status_code.map(i64::from), + ); + insert_response_upstream_diagnostics(&mut metadata, &exit.response.diagnostics); + insert_response_trace_metadata(&mut metadata, exit.trace, Some(exit.attempt)); + insert_websocket_pool_decision(&mut metadata, exit.response.websocket_pool_decision); + enrich_response_request_semantics(&mut metadata, exit.request); + record_response_event(ResponseUsageRecord { + recorder: exit.recorder, + request_id: exit.request_id, + client_api_key_id: exit.request.client_api_key_id.as_deref(), + account_id: exit.account_id, + route: exit.route, + model: effective_model, + requested_model: Some(exit.requested_model), + client_ip: exit.request.client_ip.as_deref(), + client_user_agent: exit.request.client_user_agent.as_deref(), + reasoning_effort: reasoning_effort_from_request(exit.request), + service_tier: exit.request.service_tier(), + started_at: exit.started_at, + status_code: 200, + message: if exit.completed { + "v1 responses completed" + } else { + "v1 responses incomplete" + }, + metadata, + rate_limit_headers: &exit.response.rate_limit_headers, + }) + .await; + } + + pub(super) async fn leave_stream(exit: StreamExit<'_>) { + let context = &exit.context; + let summary = exit.summary; + let rate_limit_headers = exit.rate_limit_headers; + let body = exit.body; + match &summary.terminal { + StreamTerminal::Completed { response } => { + record_success(context, summary, response, true, rate_limit_headers, body).await; + } + StreamTerminal::Incomplete { response } => { + record_success(context, summary, response, false, rate_limit_headers, body).await; + } + StreamTerminal::Failed { failure } => { + record_failure( + context, + summary, + failure, + exit.failure_status.unwrap_or(502), + rate_limit_headers, + body, + ) + .await; + } + terminal @ (StreamTerminal::UpstreamClosed + | StreamTerminal::UpstreamError { .. } + | StreamTerminal::ProtocolError { .. } + | StreamTerminal::CaptureLimitExceeded) => { + let failure = synthetic_failure(terminal); + record_failure(context, summary, &failure, 502, rate_limit_headers, body).await; + } + terminal @ (StreamTerminal::Cancelled + | StreamTerminal::DownstreamClosed + | StreamTerminal::Shutdown) => { + let status_code = if matches!(terminal, StreamTerminal::Shutdown) { + 503 + } else { + 499 + }; + let mut metadata = serde_json::json!({ + "stream": true, + "cancelled": true, + "terminal": terminal.name(), + "responseId": summary.last_response_id, + "usage": summary.usage, + }); + insert_first_token_ms(&mut metadata, summary.first_token_ms); + record_live_response_stream_event( + context, + status_code, + UsageRecordLevel::Warn, + "v1 responses stream cancelled", + metadata, + rate_limit_headers, + body, + ) + .await; + } + } + } +} + +async fn record_success( + context: &StreamContext<'_>, + summary: &StreamSummary, + response: &Value, + completed: bool, + rate_limit_headers: &[(String, String)], + body: &str, +) { + record_live_response_stream_event( + context, + 200, + UsageRecordLevel::Info, + if completed { + "v1 responses stream completed" + } else { + "v1 responses stream incomplete" + }, + serde_json::json!({ + "stream": true, + "completed": completed, + "incomplete": !completed, + "responseId": response.get("id").and_then(Value::as_str), + "firstTokenMs": summary.first_token_ms, + "usage": summary.usage, + }), + rate_limit_headers, + body, + ) + .await; +} + +fn synthetic_failure(terminal: &StreamTerminal) -> ResponsesSseFailure { + let detail = terminal.synthetic_failure_detail().unwrap_or_default(); + ResponsesSseFailure { + event: "response.failed".to_string(), + message: if detail.trim().is_empty() { + STREAM_DISCONNECTED_MESSAGE.to_string() + } else { + format!("{STREAM_DISCONNECTED_MESSAGE}: {detail}") + }, + upstream_code: Some(STREAM_DISCONNECTED_CODE.to_string()), + upstream_type: Some("server_error".to_string()), + explicit_status_code: None, + retry_after_seconds: None, + } +} + +async fn record_failure( + context: &StreamContext<'_>, + summary: &StreamSummary, + failure: &ResponsesSseFailure, + status_code: i64, + rate_limit_headers: &[(String, String)], + body: &str, +) { + let failure_source = stream_failure_source(failure); + let failure_detail = synthetic_stream_disconnected_detail(failure); + tracing::warn!( + account_id = %context.account_id, + request_id = %context.request_id, + response_id = summary.last_response_id.as_deref().unwrap_or(""), + transport = %backend_transport_name(context.transport), + websocket_pool_kind = ?context.websocket_pool_decision.map(WebSocketPoolDecision::kind), + websocket_pool_reason = ?context.websocket_pool_decision.and_then(WebSocketPoolDecision::reason), + first_token_ms = ?summary.first_token_ms, + latency_ms = elapsed_millis_i64(context.started_at), + event = %failure.event, + code = ?failure.upstream_code.as_deref(), + failure_source, + failure_detail = ?failure_detail.as_deref(), + "Live upstream stream ended with response.failed" + ); + let mut metadata = stream_failure_metadata(failure, summary.usage); + metadata["responseId"] = summary + .last_response_id + .as_ref() + .map_or(Value::Null, |response_id| { + Value::String(response_id.clone()) + }); + insert_first_token_ms(&mut metadata, summary.first_token_ms); + record_live_response_stream_event( + context, + status_code, + UsageRecordLevel::Error, + "v1 responses stream failed", + metadata, + rate_limit_headers, + body, + ) + .await; +} diff --git a/backend/src/dispatch/controllers/usage.rs b/backend/src/dispatch/controllers/usage.rs new file mode 100644 index 000000000..b132d7687 --- /dev/null +++ b/backend/src/dispatch/controllers/usage.rs @@ -0,0 +1,106 @@ +//! 配额响应头与 token 用量写入的唯一 feature owner。 + +use crate::{ + dispatch::lifecycle::contract::{ + AttemptObservation, AttemptObservationKind, CompleteResponseFacts, + }, + fleet::pool::AccountPoolService, + upstream::openai::{ + protocol::{events::TokenUsage, responses::CodexResponsesRequest}, + transport::{CodexRateLimitHeaderUpdates, CodexTurnStateUpdate}, + }, +}; + +pub(super) struct UsageScope { + image_generation_requested: bool, +} + +pub(super) struct StreamUsageExit { + pub(super) rate_limit_headers: Vec<(String, String)>, + pub(super) turn_state: Option, +} + +pub(super) struct StreamExit<'a> { + pub account_pool: &'a AccountPoolService, + pub account_id: &'a str, + pub account_plan_type: Option<&'a str>, + pub rate_limit_headers: &'a [(String, String)], + pub rate_limit_header_updates: Option<&'a CodexRateLimitHeaderUpdates>, + pub turn_state_update: Option<&'a CodexTurnStateUpdate>, + pub turn_state: Option<&'a str>, + pub usage: Option, +} + +pub(super) struct UsageController; + +impl UsageController { + pub(super) fn enter(request: &CodexResponsesRequest) -> UsageScope { + UsageScope { + image_generation_requested: request.expects_image_generation(), + } + } + + pub(super) async fn observe_attempt( + account_pool: &AccountPoolService, + scope: &UsageScope, + observation: &AttemptObservation, + ) { + if !matches!( + observation.kind, + AttemptObservationKind::CompleteResponse(CompleteResponseFacts::Empty) + ) { + return; + } + let Some(account_id) = observation + .account + .as_ref() + .map(|account| account.id.as_str()) + else { + return; + }; + account_pool + .record_empty_response_attempt(account_id, scope.image_generation_requested) + .await; + } + + pub(super) async fn leave_complete( + account_pool: &AccountPoolService, + account_id: &str, + usage: Option, + image_generation_requested: bool, + ) { + if let Some(usage) = usage { + account_pool + .record_response_usage(account_id, usage, image_generation_requested) + .await; + } + } + + pub(super) async fn leave_stream(exit: StreamExit<'_>) -> StreamUsageExit { + let mut rate_limit_headers = exit.rate_limit_headers.to_vec(); + if let Some(updates) = exit.rate_limit_header_updates { + rate_limit_headers.extend(updates.lock().await.iter().cloned()); + } + exit.account_pool + .sync_passive_rate_limit_headers_for_account( + exit.account_id, + exit.account_plan_type, + &rate_limit_headers, + ) + .await; + let turn_state = if let Some(update) = exit.turn_state_update { + update.lock().await.clone() + } else { + exit.turn_state.map(ToString::to_string) + }; + if let Some(usage) = exit.usage { + exit.account_pool + .record_token_usage(exit.account_id, &usage) + .await; + } + StreamUsageExit { + rate_limit_headers, + turn_state, + } + } +} diff --git a/backend/src/dispatch/errors.rs b/backend/src/dispatch/errors.rs index efda88808..f4b28c200 100644 --- a/backend/src/dispatch/errors.rs +++ b/backend/src/dispatch/errors.rs @@ -1,23 +1,54 @@ //! Shared upstream error classification for dispatch routes. -use chrono::{DateTime, Duration, Utc}; use serde_json::{Map, Value, json}; use thiserror::Error; use crate::{ - dispatch::{ - recovery::exhaustion::{ExhaustedAccount, ExhaustedAccountKind, ExhaustedAccountRef}, - stream::sse_failure::{sse_failure_error_body, stream_failure_http_status}, + dispatch::failure::{ + exhaustion::{ExhaustedAccount, ExhaustedAccountKind, ExhaustedAccountRef}, + sse::sse_failure_error_body, }, upstream::openai::{ protocol::{responses::ResponsesSseFailure, sse::SseError}, - transport::{ - CodexBackendTransport, CodexClientError, CodexUpstreamDiagnostics, - is_banned_upstream_error, is_cyber_policy_upstream_error, - }, + transport::{CodexBackendTransport, CodexClientError, CodexUpstreamDiagnostics}, }, }; +/// controller 已经完成分类、可直接交给 API 编码的客户端失败事实。 +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct ClientFailure { + failure: ResponsesSseFailure, + status_code: u16, + expose_upstream: bool, +} + +impl ClientFailure { + pub(in crate::dispatch) fn new( + failure: ResponsesSseFailure, + status_code: u16, + expose_upstream: bool, + ) -> Self { + Self { + failure, + status_code, + expose_upstream, + } + } + + pub(crate) fn failure(&self) -> &ResponsesSseFailure { + &self.failure + } + + pub(crate) fn status_code(&self) -> u16 { + self.status_code + } + + /// 只有 controller 明确允许时，API 才透传上游 code/message。 + pub(crate) fn exposed_upstream(&self) -> Option<&ResponsesSseFailure> { + self.expose_upstream.then_some(&self.failure) + } +} + #[derive(Clone, Debug)] pub(crate) struct DispatchErrorMetadata { pub failure_class: DispatchFailureClass, @@ -191,90 +222,7 @@ pub(crate) fn upstream_error_diagnostics( } } -pub(crate) fn is_rate_limit_upstream_error(error: &CodexClientError) -> bool { - if is_cyber_policy_upstream_error(error) { - return false; - } - matches!( - error, - CodexClientError::Upstream { status, .. } if status_code_is_rate_limited(status.as_u16()) - ) -} - -pub(crate) fn is_retryable_upstream_5xx_error(error: &CodexClientError) -> bool { - if is_cyber_policy_upstream_error(error) { - return false; - } - matches!( - error, - CodexClientError::Upstream { status, .. } - if status_code_is_transient_upstream(status.as_u16()) - ) -} - -pub(crate) fn is_retryable_account_transport_error(error: &CodexClientError) -> bool { - match error { - CodexClientError::Http(error) => error.is_connect() || error.is_timeout(), - CodexClientError::StreamIdleTimeout { .. } => true, - CodexClientError::WebSocket( - crate::upstream::openai::transport::websocket::CodexWebSocketExchangeError::Transport(_) - | crate::upstream::openai::transport::websocket::CodexWebSocketExchangeError::ConnectTimeout { .. } - | crate::upstream::openai::transport::websocket::CodexWebSocketExchangeError::SendTimeout { .. } - | crate::upstream::openai::transport::websocket::CodexWebSocketExchangeError::ClosedBeforeTerminal - | crate::upstream::openai::transport::websocket::CodexWebSocketExchangeError::ReceiveIdleTimeout { .. } - | crate::upstream::openai::transport::websocket::CodexWebSocketExchangeError::ReusedConnectionDiedBeforeFirstOutput { .. } - | crate::upstream::openai::transport::websocket::CodexWebSocketExchangeError::InitialEventTimeout { .. }, - ) => true, - _ => is_retryable_upstream_5xx_error(error), - } -} - -pub(crate) fn is_quota_exhausted_upstream_error(error: &CodexClientError) -> bool { - matches!( - error, - CodexClientError::Upstream { status, .. } - if status_code_is_quota_exhausted(status.as_u16()) - && !is_banned_upstream_error(error) - && !is_cyber_policy_upstream_error(error) - ) -} - -pub(crate) fn is_model_unsupported_upstream_error(error: &CodexClientError) -> bool { - matches!( - error, - CodexClientError::Upstream { status, body, .. } - if status.is_client_error() - && !matches!(status.as_u16(), 401 | 402 | 403 | 404 | 429) - && !is_cyber_policy_upstream_error(error) - && is_model_unsupported_signal(body) - ) -} - -pub(crate) fn is_history_recovery_upstream_error(error: &CodexClientError) -> bool { - upstream_error_code(error).is_some_and(|code| is_history_recovery_code(&code)) - || matches!( - error, - CodexClientError::WebSocket( - crate::upstream::openai::transport::websocket::CodexWebSocketExchangeError::ContinuationUnavailable { reason } - ) if !matches!( - reason, - crate::upstream::openai::transport::websocket::PreviousResponseUnavailableReason::ConnectionBusy - ) - ) -} - -pub(crate) fn is_continuation_busy_error(error: &CodexClientError) -> bool { - matches!( - error, - CodexClientError::WebSocket( - crate::upstream::openai::transport::websocket::CodexWebSocketExchangeError::ContinuationUnavailable { - reason: crate::upstream::openai::transport::websocket::PreviousResponseUnavailableReason::ConnectionBusy, - } - ) - ) -} - -pub(crate) fn upstream_error_body(error: &CodexClientError) -> String { +fn upstream_error_body(error: &CodexClientError) -> String { match error { CodexClientError::Upstream { body, .. } => body.clone(), error => error.to_string(), @@ -290,54 +238,6 @@ pub(crate) fn upstream_error_set_cookie_headers(error: &CodexClientError) -> &[S } } -pub(crate) fn is_model_unsupported_signal(value: &str) -> bool { - let value = value.to_ascii_lowercase(); - value.contains("model_not_supported") - || value.contains("model_not_available") - || (value.contains("model") - && (value.contains("not supported") - || value.contains("not available") - || value.contains("not_supported") - || value.contains("not_available"))) -} - -pub(crate) fn is_history_recovery_code(code: &str) -> bool { - matches!( - code, - "previous_response_not_found" - | "invalid_encrypted_content" - | "missing_tool_output" - | "no_tool_output" - ) -} - -pub(crate) fn upstream_error_code(error: &CodexClientError) -> Option { - let CodexClientError::Upstream { body, .. } = error else { - return None; - }; - let value = serde_json::from_str::(body).ok()?; - value - .pointer("/response/error/code") - .or_else(|| value.pointer("/error/code")) - .or_else(|| value.get("code")) - .and_then(Value::as_str) - .map(ToString::to_string) -} - -pub(crate) fn rate_limit_cooldown_until( - error: &CodexClientError, - now: DateTime, -) -> DateTime { - let retry_after_seconds = match error { - CodexClientError::Upstream { - retry_after_seconds, - .. - } => retry_after_seconds.unwrap_or(60), - _ => 60, - }; - now + Duration::seconds(retry_after_seconds.min(i64::MAX as u64) as i64) -} - pub(crate) fn upstream_error_http_status(error: &CodexClientError) -> u16 { match error { CodexClientError::Upstream { status, .. } => status.as_u16(), @@ -352,17 +252,6 @@ pub(crate) fn backend_transport_name(transport: CodexBackendTransport) -> &'stat } } -fn status_code_is_rate_limited(status_code: u16) -> bool { - status_code == 429 -} - -fn status_code_is_quota_exhausted(status_code: u16) -> bool { - status_code == 402 -} - -fn status_code_is_transient_upstream(status_code: u16) -> bool { - matches!(status_code, 500..=599) -} /// Responses 调度错误。 #[derive(Debug, Error)] pub enum ResponseDispatchError { @@ -427,7 +316,7 @@ pub enum ResponseDispatchError { #[error("upstream response did not include visible output")] EmptyUpstreamResponse, #[error("upstream response failed: {0:?}")] - Failed(ResponsesSseFailure), + Failed(ClientFailure), } impl ResponseDispatchError { @@ -487,7 +376,7 @@ impl ResponseDispatchError { | Self::EmptyUpstreamResponse => 502, Self::UpstreamUnavailable { .. } => 503, Self::ContinuationBusy => 503, - Self::Failed(failure) => stream_failure_http_status(failure), + Self::Failed(failure) => failure.status_code(), Self::ModelUnsupported { .. } | Self::HistoryUnavailable { .. } => 400, Self::Upstream(error) => upstream_error_http_status(error), } @@ -583,7 +472,7 @@ impl ResponseDispatchError { Self::Failed(failure) => DispatchErrorMetadata { failure_class: DispatchFailureClass::ResponseFailed, exhausted_count: None, - upstream_error: Some(sse_failure_error_body(failure)), + upstream_error: Some(sse_failure_error_body(failure.failure())), upstream_status: None, diagnostics: None, }, diff --git a/backend/src/dispatch/recovery/exhaustion.rs b/backend/src/dispatch/failure/exhaustion.rs similarity index 65% rename from backend/src/dispatch/recovery/exhaustion.rs rename to backend/src/dispatch/failure/exhaustion.rs index dea0491a3..ffd90ba05 100644 --- a/backend/src/dispatch/recovery/exhaustion.rs +++ b/backend/src/dispatch/failure/exhaustion.rs @@ -1,8 +1,6 @@ //! 账号耗尽状态聚合。 -use crate::fleet::account::AccountStatus; - -#[derive(Clone, Copy, Debug)] +#[derive(Clone, Copy, Debug, PartialEq, Eq)] pub(crate) enum ExhaustedAccountKind { QuotaExhausted, RateLimited, @@ -15,6 +13,35 @@ pub(crate) enum ExhaustedAccountKind { UpstreamUnavailable, } +/// Controller 交给 attempt runner 聚合的单账号耗尽事实。 +#[derive(Clone, Debug, PartialEq, Eq)] +pub(in crate::dispatch) struct AccountExhaustionRecord { + pub account_id: Option, + pub kind: ExhaustedAccountKind, + pub upstream_error: String, + pub status_code: Option, +} + +impl AccountExhaustionRecord { + pub(in crate::dispatch) fn new( + account_id: impl Into, + kind: ExhaustedAccountKind, + upstream_error: impl Into, + ) -> Self { + Self { + account_id: Some(account_id.into()), + kind, + upstream_error: upstream_error.into(), + status_code: None, + } + } + + pub(in crate::dispatch) fn with_status_code(mut self, status_code: u16) -> Self { + self.status_code = Some(status_code); + self + } +} + impl ExhaustedAccountKind { pub(crate) fn message_reason(self) -> &'static str { match self { @@ -31,7 +58,7 @@ impl ExhaustedAccountKind { } } -#[derive(Clone, Debug)] +#[derive(Clone, Debug, PartialEq, Eq)] pub(crate) struct ExhaustedAccount { pub kind: ExhaustedAccountKind, pub count: usize, @@ -69,100 +96,12 @@ struct ExhaustedCounter { } impl AccountExhaustionTracker { - pub(crate) fn last_account_id(&self) -> Option<&str> { - self.last_account_id.as_deref() - } - - pub(crate) fn record_rate_limited( - &mut self, - account_id: Option<&str>, - upstream_error: impl Into, - ) { + pub(in crate::dispatch) fn record_exhaustion(&mut self, record: AccountExhaustionRecord) { self.record( - account_id, - ExhaustedAccountKind::RateLimited, - upstream_error, - None, - ); - } - - pub(crate) fn record_quota_exhausted( - &mut self, - account_id: Option<&str>, - upstream_error: impl Into, - ) { - self.record( - account_id, - ExhaustedAccountKind::QuotaExhausted, - upstream_error, - None, - ); - } - - pub(crate) fn record_auth_failure( - &mut self, - account_id: Option<&str>, - account_status: AccountStatus, - upstream_error: impl Into, - status_code: Option, - ) { - let kind = match account_status { - AccountStatus::Disabled => ExhaustedAccountKind::Disabled, - AccountStatus::Banned => ExhaustedAccountKind::Banned, - _ => ExhaustedAccountKind::Expired, - }; - self.record(account_id, kind, upstream_error, status_code); - } - - pub(crate) fn record_cloudflare_challenge( - &mut self, - account_id: Option<&str>, - upstream_error: impl Into, - ) { - self.record( - account_id, - ExhaustedAccountKind::CloudflareChallenge, - upstream_error, - None, - ); - } - - pub(crate) fn record_cloudflare_path_blocked( - &mut self, - account_id: Option<&str>, - upstream_error: impl Into, - ) { - self.record( - account_id, - ExhaustedAccountKind::CloudflarePathBlocked, - upstream_error, - None, - ); - } - - pub(crate) fn record_model_unsupported( - &mut self, - account_id: Option<&str>, - upstream_error: impl Into, - ) { - self.record( - account_id, - ExhaustedAccountKind::ModelUnsupported, - upstream_error, - None, - ); - } - - pub(crate) fn record_upstream_unavailable( - &mut self, - account_id: Option<&str>, - upstream_error: impl Into, - ) { - self.record( - account_id, - ExhaustedAccountKind::UpstreamUnavailable, - upstream_error, - None, + record.account_id.as_deref(), + record.kind, + record.upstream_error, + record.status_code, ); } @@ -177,6 +116,10 @@ impl AccountExhaustionTracker { }) } + pub(crate) fn last_account_id(&self) -> Option<&str> { + self.last_account_id.as_deref() + } + fn record( &mut self, account_id: Option<&str>, diff --git a/backend/src/dispatch/failure/mod.rs b/backend/src/dispatch/failure/mod.rs new file mode 100644 index 000000000..b712e48a4 --- /dev/null +++ b/backend/src/dispatch/failure/mod.rs @@ -0,0 +1,4 @@ +//! 已分类账号失败的聚合、状态副作用和 SSE 终态事实。 + +pub(crate) mod exhaustion; +pub(in crate::dispatch) mod sse; diff --git a/backend/src/dispatch/failure/sse.rs b/backend/src/dispatch/failure/sse.rs new file mode 100644 index 000000000..692e145e1 --- /dev/null +++ b/backend/src/dispatch/failure/sse.rs @@ -0,0 +1,78 @@ +//! Responses SSE 失败事实和流终态元数据。 + +use serde_json::Value; + +use crate::upstream::openai::protocol::{events::TokenUsage, responses::ResponsesSseFailure}; + +pub(in crate::dispatch) const STREAM_DISCONNECTED_CODE: &str = "stream_disconnected"; +pub(in crate::dispatch) const STREAM_DISCONNECTED_MESSAGE: &str = + "Upstream stream closed before response.completed"; + +pub(in crate::dispatch) fn sse_failure_error_body(failure: &ResponsesSseFailure) -> String { + match failure.upstream_code.as_deref() { + Some(code) => serde_json::json!({ + "error": { + "code": code, + "message": failure.message.as_str(), + } + }) + .to_string(), + None => failure.message.clone(), + } +} + +pub(in crate::dispatch) fn stream_failure_metadata( + failure: &ResponsesSseFailure, + usage: Option, +) -> Value { + let mut metadata = serde_json::json!({ + "stream": true, + "failed": true, + "failureEvent": failure.event, + "failureMessage": failure.message, + "upstreamCode": failure.upstream_code, + "usage": usage, + }); + enrich_stream_failure_source_metadata(&mut metadata, failure); + metadata +} + +fn enrich_stream_failure_source_metadata(metadata: &mut Value, failure: &ResponsesSseFailure) { + let Some(object) = metadata.as_object_mut() else { + return; + }; + object.insert( + "failureSource".to_string(), + Value::String(stream_failure_source(failure).to_string()), + ); + if let Some(detail) = synthetic_stream_disconnected_detail(failure) { + object.insert("synthetic".to_string(), Value::Bool(true)); + if !detail.is_empty() { + object.insert("failureDetail".to_string(), Value::String(detail)); + } + } +} + +pub(in crate::dispatch) fn stream_failure_source(failure: &ResponsesSseFailure) -> &'static str { + if synthetic_stream_disconnected_detail(failure).is_some() { + "proxy" + } else { + "upstream" + } +} + +pub(in crate::dispatch) fn synthetic_stream_disconnected_detail( + failure: &ResponsesSseFailure, +) -> Option { + if failure.upstream_code.as_deref() != Some(STREAM_DISCONNECTED_CODE) { + return None; + } + let detail = failure + .message + .strip_prefix(STREAM_DISCONNECTED_MESSAGE)? + .strip_prefix(": ") + .unwrap_or_default() + .trim() + .to_string(); + Some(detail) +} diff --git a/backend/src/dispatch/lifecycle/attempt.rs b/backend/src/dispatch/lifecycle/attempt.rs new file mode 100644 index 000000000..921c056c9 --- /dev/null +++ b/backend/src/dispatch/lifecycle/attempt.rs @@ -0,0 +1,919 @@ +//! 单账号 attempt 的 stepwise 执行器。 + +use std::time::Instant; + +use serde_json::Value; + +use crate::{ + dispatch::{ + affinity::AccountIdentityService, + controllers::{ + AttemptAccountPreparation, AttemptAccountPreparationContext, AttemptRoutePreparation, + AttemptUpstreamErrorContext, ControllerRequestScope, ControllerSet, + }, + errors::{ClientFailure, ResponseDispatchError}, + failure::exhaustion::AccountExhaustionTracker, + routing::candidates::AccountAttemptLedger, + transport::{ + account::{ + AccountUpstreamContext, create_response_stream_with_account, + create_response_with_account, + }, + canonical::{CanonicalStreamDecoder, normalize_complete_response}, + observation::upstream_failure_facts, + prefetch::{PrefetchedStream, StreamPrefetchError, prefetch_until_commit}, + }, + }, + fleet::{ + account::Account, + pool::{AccountCandidateLease, AccountPoolService}, + }, + upstream::openai::{ + protocol::responses::{CodexResponsesRequest, CollectedResponse}, + transport::{ + CodexBackendClient, CodexBackendStreamingResponse, + backend_transport_for_response_request, + }, + }, +}; + +use super::{ + contract::{ + AttemptAccountFacts, AttemptApplyOutcome, AttemptContractError, AttemptDecision, + AttemptObservation, AttemptObservationKind, AttemptRejection, AttemptReturnKind, + CandidateLedgerFacts, CompleteResponseFacts, EstablishedAttemptContext, + EstablishedComplete, EstablishedCompleteBody, EstablishedResponse, EstablishedStream, + PendingAttempt, PendingCompleteResponse, PendingProtocolFailure, PendingStreamResponse, + PinnedCandidateAcquireFailureKind, ProtocolFailureKind, + }, + trace::ResponseDispatchTrace, +}; + +pub(in crate::dispatch) enum AttemptMode { + Complete { + started_at: Instant, + tuple_schema: Option, + }, + Stream { + tuple_schema: Option, + }, +} + +pub(in crate::dispatch) struct AttemptRunnerDependencies<'a> { + pub account_pool: &'a AccountPoolService, + pub codex: &'a CodexBackendClient, + pub controllers: &'a ControllerSet, + pub account_identity: &'a AccountIdentityService, + pub request_id: &'a str, + pub route: &'a str, + pub requested_model: &'a str, + pub started_at: Instant, +} + +/// 尚未越过提交边界的 attempt capability；只有该类型接受 retry decision。 +pub(in crate::dispatch) struct OpenAttempt<'runner, 'dependencies> { + runner: &'runner mut AttemptRunner<'dependencies>, + observation: Box, +} + +impl OpenAttempt<'_, '_> { + pub(in crate::dispatch) fn observation(&self) -> &AttemptObservation { + &self.observation + } + + pub(in crate::dispatch) fn controller_scope( + &self, + ) -> Result<&ControllerRequestScope, AttemptContractError> { + self.runner.controller_scope() + } + + pub(in crate::dispatch) async fn apply( + self, + decision: AttemptDecision, + ) -> Result { + self.runner.apply_open(decision).await + } +} + +/// 已越过提交边界的 attempt capability；该类型只允许建立 stream。 +pub(in crate::dispatch) struct CommittedAttempt { + established: EstablishedResponse, +} + +impl CommittedAttempt { + pub(in crate::dispatch) fn accept(self) -> EstablishedResponse { + self.established + } +} + +/// `next` 产出的提交 typestate；Committed 分支没有 decision/retry API。 +pub(in crate::dispatch) enum AttemptStep<'runner, 'dependencies> { + Open(OpenAttempt<'runner, 'dependencies>), + Committed(CommittedAttempt), +} + +pub(in crate::dispatch) struct AttemptRunner<'a> { + dependencies: AttemptRunnerDependencies<'a>, + mode: AttemptMode, + request: Option, + controller_scope: Option, + candidates: AccountAttemptLedger, + exhaustion: AccountExhaustionTracker, + on_exhaustion: Option, + trace: Option, + required_account_id: Option, + pending: Option, +} + +impl<'a> AttemptRunner<'a> { + pub(in crate::dispatch) fn new( + dependencies: AttemptRunnerDependencies<'a>, + mode: AttemptMode, + request: CodexResponsesRequest, + controller_scope: ControllerRequestScope, + candidates: AccountAttemptLedger, + ) -> Self { + Self { + dependencies, + mode, + request: Some(request), + controller_scope: Some(controller_scope), + candidates, + exhaustion: AccountExhaustionTracker::default(), + on_exhaustion: None, + trace: Some(ResponseDispatchTrace::default()), + required_account_id: None, + pending: None, + } + } + + pub(in crate::dispatch) async fn next( + &mut self, + ) -> Result, AttemptContractError> { + if self.pending.is_some() { + return Err(AttemptContractError::DecisionRequired); + } + + let request = self + .request + .as_ref() + .ok_or(AttemptContractError::Terminal)?; + let acquired = if let Some(account_id) = self.required_account_id.take() { + match self + .dependencies + .account_pool + .acquire_candidate(request.model(), &account_id, chrono::Utc::now()) + .await + { + AccountCandidateLease::Acquired(acquired) => Some(*acquired), + AccountCandidateLease::Busy => { + self.pending = Some(PendingAttempt::PinnedCandidateUnavailable { + account_id, + kind: PinnedCandidateAcquireFailureKind::Busy, + }); + return self.step(); + } + AccountCandidateLease::Unavailable => { + self.pending = Some(PendingAttempt::PinnedCandidateUnavailable { + account_id, + kind: PinnedCandidateAcquireFailureKind::Unavailable, + }); + return self.step(); + } + } + } else { + self.candidates + .acquire_next(self.dependencies.account_pool) + .await + }; + let Some(acquired) = acquired else { + self.pending = Some(PendingAttempt::NoCandidate); + return self.step(); + }; + + let account = acquired.account.clone(); + let controller_scope = self + .controller_scope + .as_mut() + .ok_or(AttemptContractError::Terminal)?; + let prepared_account = match self + .dependencies + .controllers + .prepare_attempt_account( + AttemptAccountPreparationContext { + codex: self.dependencies.codex, + account_identity: self.dependencies.account_identity, + request_id: self.dependencies.request_id, + }, + acquired, + ) + .await + { + AttemptAccountPreparation::Ready(prepared) => *prepared, + AttemptAccountPreparation::Rejected => { + self.pending = Some(PendingAttempt::CandidatePreparationRejected { account }); + return self.step(); + } + }; + let acquired = prepared_account.lease; + let cookie_header = prepared_account.cookie_header; + + let attempt_request = match self.dependencies.controllers.prepare_attempt_route( + controller_scope, + request, + &account.id, + ) { + AttemptRoutePreparation::Ready(request) => *request, + AttemptRoutePreparation::Unavailable { message } => { + acquired.release_without_usage().await; + self.pending = Some(PendingAttempt::RoutePreparationRejected { account, message }); + return self.step(); + } + }; + + self.dependencies + .account_pool + .wait_for_request_interval(&acquired) + .await; + let attempt = self + .trace + .as_mut() + .ok_or(AttemptContractError::Terminal)? + .start_attempt(&account.id); + let upstream_context = AccountUpstreamContext { + codex: self.dependencies.codex, + account_identity: self.dependencies.account_identity, + request_id: self.dependencies.request_id, + account: &account, + cookie_header: cookie_header.as_deref(), + }; + + match &self.mode { + AttemptMode::Complete { + started_at, + tuple_schema, + } => { + let result = + create_response_with_account(upstream_context, &attempt_request, *started_at) + .await; + acquired.complete().await; + self.pending = Some(match result { + Ok(response) => { + self.dependencies + .controllers + .observe_complete_upstream( + &account, + &response.set_cookie_headers, + &response.rate_limit_headers, + ) + .await; + match normalize_complete_response(&response.body, tuple_schema.as_ref()) { + Ok(collected) => PendingAttempt::CompleteResponse(Box::new( + PendingCompleteResponse { + account, + attempt_request, + attempt, + response, + collected, + }, + )), + Err(error) => PendingAttempt::ProtocolFailure { + account, + attempt_request, + attempt, + transport: response.transport, + error: PendingProtocolFailure::InvalidSse(error), + }, + } + } + Err(error) => { + let transport = backend_transport_for_response_request(request); + observe_upstream_error( + self.dependencies.controllers, + AttemptUpstreamErrorContext { + request_id: self.dependencies.request_id, + account: &account, + route: self.dependencies.route, + model: self.dependencies.requested_model, + started_at: self.dependencies.started_at, + stream: matches!(self.mode, AttemptMode::Stream { .. }), + transport, + request: &attempt_request, + error: &error, + trace: self.trace.as_ref().ok_or(AttemptContractError::Terminal)?, + attempt: &attempt, + }, + ) + .await; + PendingAttempt::UpstreamFailure { + account, + attempt_request, + attempt, + transport, + error, + } + } + }); + } + AttemptMode::Stream { tuple_schema } => { + let result = + create_response_stream_with_account(upstream_context, &attempt_request).await; + self.pending = Some(match result { + Err(error) => { + acquired.complete().await; + let transport = backend_transport_for_response_request(request); + observe_upstream_error( + self.dependencies.controllers, + AttemptUpstreamErrorContext { + request_id: self.dependencies.request_id, + account: &account, + route: self.dependencies.route, + model: self.dependencies.requested_model, + started_at: self.dependencies.started_at, + stream: matches!(self.mode, AttemptMode::Stream { .. }), + transport, + request: &attempt_request, + error: &error, + trace: self.trace.as_ref().ok_or(AttemptContractError::Terminal)?, + attempt: &attempt, + }, + ) + .await; + PendingAttempt::UpstreamFailure { + account, + attempt_request, + attempt, + transport, + error, + } + } + Ok(response) => { + let CodexBackendStreamingResponse { + body, + transport, + turn_state, + set_cookie_headers, + rate_limit_headers, + rate_limit_header_updates, + turn_state_update, + websocket_pool_decision, + diagnostics, + response_metadata, + } = response; + let mut decoder = CanonicalStreamDecoder::new(tuple_schema.clone()); + match prefetch_until_commit( + body, + &mut decoder, + attempt_request.stream_commit_policy, + ) + .await + { + Err(StreamPrefetchError::Upstream(error)) => { + acquired.complete().await; + observe_upstream_error( + self.dependencies.controllers, + AttemptUpstreamErrorContext { + request_id: self.dependencies.request_id, + account: &account, + route: self.dependencies.route, + model: self.dependencies.requested_model, + started_at: self.dependencies.started_at, + stream: matches!(self.mode, AttemptMode::Stream { .. }), + transport, + request: &attempt_request, + error: &error, + trace: self + .trace + .as_ref() + .ok_or(AttemptContractError::Terminal)?, + attempt: &attempt, + }, + ) + .await; + PendingAttempt::UpstreamFailure { + account, + attempt_request, + attempt, + transport, + error, + } + } + Err(StreamPrefetchError::Empty) => { + acquired.complete().await; + PendingAttempt::ProtocolFailure { + account, + attempt_request, + attempt, + transport, + error: PendingProtocolFailure::EmptyStream, + } + } + Err(StreamPrefetchError::NoCommitBoundary) => { + acquired.complete().await; + PendingAttempt::ProtocolFailure { + account, + attempt_request, + attempt, + transport, + error: PendingProtocolFailure::NoCommitBoundary, + } + } + Err(StreamPrefetchError::InvalidSse(error)) => { + acquired.complete().await; + PendingAttempt::ProtocolFailure { + account, + attempt_request, + attempt, + transport, + error: PendingProtocolFailure::InvalidSse(error), + } + } + Ok(PrefetchedStream { + bytes: prefetched, + initial_batch, + body, + }) => { + let first_failure = initial_batch.first_failure(); + PendingAttempt::StreamResponse(Box::new(PendingStreamResponse { + account, + attempt_request, + attempt, + lease: acquired, + response: CodexBackendStreamingResponse { + body, + transport, + turn_state, + set_cookie_headers, + rate_limit_headers, + rate_limit_header_updates, + turn_state_update, + websocket_pool_decision, + diagnostics, + response_metadata, + }, + prefetched, + decoder, + initial_batch, + first_failure, + })) + } + } + } + }); + } + } + self.step() + } + + pub(in crate::dispatch) fn controller_scope( + &self, + ) -> Result<&ControllerRequestScope, AttemptContractError> { + self.controller_scope + .as_ref() + .ok_or(AttemptContractError::Terminal) + } + + async fn apply_open( + &mut self, + decision: AttemptDecision, + ) -> Result { + let pending = self + .pending + .take() + .ok_or(AttemptContractError::DecisionRequired)?; + match decision { + AttemptDecision::Accept => self + .establish(pending) + .map(AttemptApplyOutcome::Established), + AttemptDecision::RetrySameAccount => { + let account_id = pending_account(&pending) + .ok_or_else(|| invalid("retry_same", &pending))? + .to_string(); + if !self.dependencies.controllers.prepare_same_account_retry( + self.controller_scope + .as_mut() + .ok_or(AttemptContractError::Terminal)?, + &account_id, + ) { + return Err(invalid("retry_same", &pending)); + } + release_pending_stream(pending).await; + self.required_account_id = Some(account_id); + Ok(AttemptApplyOutcome::Continue) + } + AttemptDecision::RetryNextCandidate { + exhaustion, + on_exhaustion, + } => { + release_pending_stream(pending).await; + self.on_exhaustion = on_exhaustion; + if let Some(exhaustion) = exhaustion { + self.exhaustion.record_exhaustion(exhaustion); + } + Ok(AttemptApplyOutcome::Continue) + } + AttemptDecision::Return(kind) => self.reject(pending, kind).await, + } + } + + fn step(&mut self) -> Result, AttemptContractError> { + let kind = { + let pending = self + .pending + .as_ref() + .ok_or(AttemptContractError::DecisionRequired)?; + observation_kind(pending, &self.candidates, &self.exhaustion) + }; + let Some(kind) = kind else { + let pending = self + .pending + .take() + .ok_or(AttemptContractError::DecisionRequired)?; + return Ok(AttemptStep::Committed(CommittedAttempt { + established: self.establish(pending)?, + })); + }; + let observation = self.observation(kind)?; + Ok(AttemptStep::Open(OpenAttempt { + runner: self, + observation: Box::new(observation), + })) + } + + fn observation( + &self, + kind: AttemptObservationKind, + ) -> Result { + let pending = self + .pending + .as_ref() + .ok_or(AttemptContractError::DecisionRequired)?; + let request = self + .request + .as_ref() + .ok_or(AttemptContractError::Terminal)?; + let controller_scope = self + .controller_scope + .as_ref() + .ok_or(AttemptContractError::Terminal)?; + let account = pending_account_value(pending); + let account_id = account.map(|account| account.id.as_str()); + Ok(AttemptObservation { + account: account.map(AttemptAccountFacts::from), + attempt: pending_attempt(pending).cloned(), + transport: pending_transport(pending) + .unwrap_or_else(|| backend_transport_for_response_request(request)), + routing: self + .dependencies + .controllers + .attempt_routing_facts(controller_scope, account_id), + kind, + }) + } + + fn establish( + &mut self, + pending: PendingAttempt, + ) -> Result { + let established = match pending { + PendingAttempt::CompleteResponse(complete) => { + let PendingCompleteResponse { + account, + attempt, + response, + collected, + .. + } = *complete; + let body = EstablishedCompleteBody::try_from(collected).map_err(|()| { + AttemptContractError::InvalidDecision { + decision: "establish", + observation: "non_successful_complete_response", + } + })?; + EstablishedResponse::Complete(Box::new(EstablishedComplete { + context: self.take_established_context(account, attempt)?, + response, + body, + })) + } + PendingAttempt::StreamResponse(stream) => { + let PendingStreamResponse { + account, + attempt, + lease, + response, + decoder, + initial_batch, + .. + } = *stream; + EstablishedResponse::Stream(Box::new(EstablishedStream { + context: self.take_established_context(account, attempt)?, + lease, + response, + decoder, + initial_batch, + })) + } + other => return Err(invalid("establish", &other)), + }; + Ok(established) + } + + async fn reject( + &mut self, + pending: PendingAttempt, + kind: AttemptReturnKind, + ) -> Result { + let account = pending_account_value(&pending).cloned(); + let account_id = account + .as_ref() + .map(|account| account.id.clone()) + .or_else(|| self.exhaustion.last_account_id().map(ToString::to_string)); + let attempt_request = pending_attempt_request(&pending).cloned(); + let attempt = pending_attempt(&pending).cloned(); + let request = self + .request + .as_ref() + .ok_or(AttemptContractError::Terminal)?; + let transport = pending_transport(&pending) + .unwrap_or_else(|| backend_transport_for_response_request(request)); + let stream_failure = match &pending { + PendingAttempt::StreamResponse(stream) => { + stream.first_failure.as_ref().map(|failure| { + super::contract::RejectedStreamFailure { + failure: failure.clone(), + prefetched: stream.prefetched.clone(), + diagnostics: stream.response.diagnostics.clone(), + rate_limit_headers: stream.response.rate_limit_headers.clone(), + } + }) + } + _ => None, + }; + let error = + rejection_error(pending, kind, &self.exhaustion, self.on_exhaustion.take()).await?; + let request = self.request.take().ok_or(AttemptContractError::Terminal)?; + self.controller_scope + .take() + .ok_or(AttemptContractError::Terminal)?; + let trace = self.trace.take().ok_or(AttemptContractError::Terminal)?; + Ok(AttemptApplyOutcome::Rejected(Box::new(AttemptRejection { + request, + trace, + account_id, + account, + attempt_request, + attempt, + transport, + stream_failure, + error, + }))) + } + + fn take_established_context( + &mut self, + account: Account, + attempt: super::trace::ResponseDispatchAttempt, + ) -> Result { + Ok(EstablishedAttemptContext { + request: self.request.take().ok_or(AttemptContractError::Terminal)?, + controller_scope: self + .controller_scope + .take() + .ok_or(AttemptContractError::Terminal)?, + trace: self.trace.take().ok_or(AttemptContractError::Terminal)?, + account, + attempt, + }) + } +} + +async fn observe_upstream_error( + controllers: &ControllerSet, + context: AttemptUpstreamErrorContext<'_>, +) { + controllers.observe_upstream_error(context).await; +} + +fn observation_kind( + pending: &PendingAttempt, + candidates: &AccountAttemptLedger, + exhaustion: &AccountExhaustionTracker, +) -> Option { + Some(match pending { + PendingAttempt::NoCandidate => AttemptObservationKind::NoCandidate { + ledger: CandidateLedgerFacts { + candidates: candidates.candidate_count(), + attempted: candidates.attempted_count(), + state_excluded: candidates.state_excluded_count(), + }, + last_exhausted: exhaustion.last_exhausted(), + }, + PendingAttempt::PinnedCandidateUnavailable { account_id, kind } => { + AttemptObservationKind::PinnedCandidateUnavailable { + account_id: account_id.clone(), + kind: *kind, + } + } + PendingAttempt::CandidatePreparationRejected { .. } => { + AttemptObservationKind::CandidatePreparationRejected + } + PendingAttempt::RoutePreparationRejected { message, .. } => { + AttemptObservationKind::RoutePreparationRejected { + message: message.clone(), + } + } + PendingAttempt::UpstreamFailure { error, .. } => { + AttemptObservationKind::UpstreamFailure(upstream_failure_facts(error)) + } + PendingAttempt::ProtocolFailure { error, .. } => AttemptObservationKind::ProtocolFailure { + kind: match error { + PendingProtocolFailure::InvalidSse(_) => ProtocolFailureKind::InvalidSse, + PendingProtocolFailure::EmptyStream => ProtocolFailureKind::EmptyStream, + PendingProtocolFailure::NoCommitBoundary => ProtocolFailureKind::NoCommitBoundary, + }, + message: match error { + PendingProtocolFailure::InvalidSse(error) => error.to_string(), + PendingProtocolFailure::EmptyStream => { + "upstream stream ended without data".to_string() + } + PendingProtocolFailure::NoCommitBoundary => { + "upstream stream ended before output or a terminal event".to_string() + } + }, + }, + PendingAttempt::CompleteResponse(complete) => { + AttemptObservationKind::CompleteResponse(match &complete.collected { + CollectedResponse::Completed(_) => CompleteResponseFacts::Completed, + CollectedResponse::Incomplete(_) => CompleteResponseFacts::Incomplete, + CollectedResponse::Failed(failure) => { + CompleteResponseFacts::Failed(failure.clone()) + } + CollectedResponse::MissingCompleted => CompleteResponseFacts::MissingCompleted, + CollectedResponse::Empty => CompleteResponseFacts::Empty, + }) + } + PendingAttempt::StreamResponse(stream) => { + return stream + .first_failure + .clone() + .map(AttemptObservationKind::StreamFailure); + } + }) +} + +async fn release_pending_stream(pending: PendingAttempt) { + if let PendingAttempt::StreamResponse(stream) = pending { + stream.lease.complete().await; + } +} + +async fn rejection_error( + pending: PendingAttempt, + kind: AttemptReturnKind, + exhaustion: &AccountExhaustionTracker, + on_exhaustion: Option, +) -> Result { + match kind { + AttemptReturnKind::Failed(failure) => { + release_pending_stream(pending).await; + Ok(ResponseDispatchError::Failed(failure)) + } + AttemptReturnKind::ContinuationBusy => Ok(ResponseDispatchError::ContinuationBusy), + AttemptReturnKind::RouteUnavailable { message } => { + Ok(ResponseDispatchError::HistoryUnavailable { + upstream_error: message, + }) + } + AttemptReturnKind::Observed => match pending { + PendingAttempt::NoCandidate | PendingAttempt::CandidatePreparationRejected { .. } => { + Ok(on_exhaustion.map_or_else( + || { + exhaustion + .last_exhausted() + .map(ResponseDispatchError::from_exhausted_account) + .unwrap_or(ResponseDispatchError::NoActiveAccount) + }, + ResponseDispatchError::Failed, + )) + } + PendingAttempt::RoutePreparationRejected { message, .. } => { + Ok(ResponseDispatchError::HistoryUnavailable { + upstream_error: message, + }) + } + PendingAttempt::PinnedCandidateUnavailable { .. } => { + Err(invalid("return", &PendingAttempt::NoCandidate)) + } + PendingAttempt::UpstreamFailure { error, .. } => { + Ok(ResponseDispatchError::Upstream(error)) + } + PendingAttempt::ProtocolFailure { error, .. } => Ok(match error { + PendingProtocolFailure::InvalidSse(error) => { + ResponseDispatchError::InvalidSse(error) + } + PendingProtocolFailure::EmptyStream => ResponseDispatchError::EmptyUpstreamResponse, + PendingProtocolFailure::NoCommitBoundary => ResponseDispatchError::MissingCompleted, + }), + PendingAttempt::CompleteResponse(complete) => match complete.collected { + CollectedResponse::Failed(_) => Err(invalid( + "return_untyped_failure", + &PendingAttempt::NoCandidate, + )), + CollectedResponse::MissingCompleted => Ok(ResponseDispatchError::MissingCompleted), + CollectedResponse::Empty => Ok(ResponseDispatchError::EmptyUpstreamResponse), + CollectedResponse::Completed(_) | CollectedResponse::Incomplete(_) => { + Err(invalid("return", &PendingAttempt::NoCandidate)) + } + }, + PendingAttempt::StreamResponse(stream) => match stream.first_failure { + Some(failure) => { + stream.lease.complete().await; + let _ = failure; + Err(invalid( + "return_untyped_failure", + &PendingAttempt::NoCandidate, + )) + } + None => Err(invalid("return", &PendingAttempt::NoCandidate)), + }, + }, + } +} + +fn pending_account(pending: &PendingAttempt) -> Option<&str> { + pending_account_value(pending).map(|account| account.id.as_str()) +} + +fn pending_account_value(pending: &PendingAttempt) -> Option<&Account> { + match pending { + PendingAttempt::NoCandidate | PendingAttempt::PinnedCandidateUnavailable { .. } => None, + PendingAttempt::CandidatePreparationRejected { account } + | PendingAttempt::RoutePreparationRejected { account, .. } + | PendingAttempt::UpstreamFailure { account, .. } + | PendingAttempt::ProtocolFailure { account, .. } => Some(account), + PendingAttempt::CompleteResponse(complete) => Some(&complete.account), + PendingAttempt::StreamResponse(stream) => Some(&stream.account), + } +} + +fn pending_attempt(pending: &PendingAttempt) -> Option<&super::trace::ResponseDispatchAttempt> { + match pending { + PendingAttempt::UpstreamFailure { attempt, .. } + | PendingAttempt::ProtocolFailure { attempt, .. } => Some(attempt), + PendingAttempt::CompleteResponse(complete) => Some(&complete.attempt), + PendingAttempt::StreamResponse(stream) => Some(&stream.attempt), + PendingAttempt::NoCandidate + | PendingAttempt::PinnedCandidateUnavailable { .. } + | PendingAttempt::CandidatePreparationRejected { .. } + | PendingAttempt::RoutePreparationRejected { .. } => None, + } +} + +fn pending_attempt_request(pending: &PendingAttempt) -> Option<&CodexResponsesRequest> { + match pending { + PendingAttempt::UpstreamFailure { + attempt_request, .. + } + | PendingAttempt::ProtocolFailure { + attempt_request, .. + } => Some(attempt_request), + PendingAttempt::CompleteResponse(complete) => Some(&complete.attempt_request), + PendingAttempt::StreamResponse(stream) => Some(&stream.attempt_request), + PendingAttempt::NoCandidate + | PendingAttempt::PinnedCandidateUnavailable { .. } + | PendingAttempt::CandidatePreparationRejected { .. } + | PendingAttempt::RoutePreparationRejected { .. } => None, + } +} + +fn pending_transport( + pending: &PendingAttempt, +) -> Option { + match pending { + PendingAttempt::UpstreamFailure { transport, .. } + | PendingAttempt::ProtocolFailure { transport, .. } => Some(*transport), + PendingAttempt::CompleteResponse(complete) => Some(complete.response.transport), + PendingAttempt::StreamResponse(stream) => Some(stream.response.transport), + PendingAttempt::NoCandidate + | PendingAttempt::PinnedCandidateUnavailable { .. } + | PendingAttempt::CandidatePreparationRejected { .. } + | PendingAttempt::RoutePreparationRejected { .. } => None, + } +} + +fn invalid(decision: &'static str, pending: &PendingAttempt) -> AttemptContractError { + AttemptContractError::InvalidDecision { + decision, + observation: match pending { + PendingAttempt::NoCandidate => "no_candidate", + PendingAttempt::PinnedCandidateUnavailable { .. } => "pinned_candidate_unavailable", + PendingAttempt::CandidatePreparationRejected { .. } => "candidate_preparation_rejected", + PendingAttempt::RoutePreparationRejected { .. } => "route_preparation_rejected", + PendingAttempt::UpstreamFailure { .. } => "upstream_failure", + PendingAttempt::ProtocolFailure { .. } => "protocol_failure", + PendingAttempt::CompleteResponse(_) => "complete_response", + PendingAttempt::StreamResponse(_) => "stream_response", + }, + } +} diff --git a/backend/src/dispatch/lifecycle/contract.rs b/backend/src/dispatch/lifecycle/contract.rs new file mode 100644 index 000000000..8a5af45be --- /dev/null +++ b/backend/src/dispatch/lifecycle/contract.rs @@ -0,0 +1,282 @@ +//! Attempt lifecycle 的类型化 observation、decision 与结果。 + +use bytes::Bytes; + +use crate::{ + dispatch::{ + errors::{ClientFailure, ResponseDispatchError}, + failure::exhaustion::{AccountExhaustionRecord, ExhaustedAccount}, + transport::{ + canonical::{CanonicalStreamBatch, CanonicalStreamDecoder}, + observation::UpstreamFailureFacts, + }, + }, + fleet::{account::Account, pool::AccountLease}, + upstream::openai::{ + protocol::{ + responses::{CodexResponsesRequest, CollectedResponse, ResponsesSseFailure}, + sse::SseError, + }, + transport::{ + CodexBackendResponse, CodexBackendStreamingResponse, CodexBackendTransport, + CodexClientError, CodexUpstreamDiagnostics, + }, + }, +}; + +use super::trace::{ResponseDispatchAttempt, ResponseDispatchTrace}; + +/// Request/Attempt/Stream 共同消费的最终结果类别。 +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub(in crate::dispatch) enum FinalOutcome { + Completed, + Incomplete, + Failed, + Cancelled, + Shutdown, +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub(in crate::dispatch) struct AttemptAccountFacts { + pub id: String, + pub email: Option, + pub plan_type: Option, +} + +impl From<&Account> for AttemptAccountFacts { + fn from(account: &Account) -> Self { + Self { + id: account.id.clone(), + email: account.email.clone(), + plan_type: account.plan_type.clone(), + } + } +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub(in crate::dispatch) struct AttemptRoutingFacts { + pub external_origin: bool, + pub can_retry_same_account: bool, + pub can_retry_next_candidate: bool, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub(in crate::dispatch) struct CandidateLedgerFacts { + pub candidates: usize, + pub attempted: usize, + pub state_excluded: usize, +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub(in crate::dispatch) enum CompleteResponseFacts { + Completed, + Incomplete, + Failed(ResponsesSseFailure), + MissingCompleted, + Empty, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub(in crate::dispatch) enum ProtocolFailureKind { + InvalidSse, + EmptyStream, + NoCommitBoundary, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub(in crate::dispatch) enum PinnedCandidateAcquireFailureKind { + Busy, + Unavailable, +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub(in crate::dispatch) enum AttemptObservationKind { + NoCandidate { + ledger: CandidateLedgerFacts, + last_exhausted: Option, + }, + CandidatePreparationRejected, + RoutePreparationRejected { + message: String, + }, + PinnedCandidateUnavailable { + account_id: String, + kind: PinnedCandidateAcquireFailureKind, + }, + UpstreamFailure(UpstreamFailureFacts), + ProtocolFailure { + kind: ProtocolFailureKind, + message: String, + }, + CompleteResponse(CompleteResponseFacts), + StreamFailure(ResponsesSseFailure), +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub(in crate::dispatch) struct AttemptObservation { + pub account: Option, + pub attempt: Option, + pub transport: CodexBackendTransport, + pub routing: AttemptRoutingFacts, + pub kind: AttemptObservationKind, +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub(in crate::dispatch) enum AttemptDecision { + Accept, + RetrySameAccount, + RetryNextCandidate { + exhaustion: Option, + on_exhaustion: Option, + }, + Return(AttemptReturnKind), +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub(in crate::dispatch) enum AttemptReturnKind { + Observed, + Failed(ClientFailure), + ContinuationBusy, + RouteUnavailable { message: String }, +} + +pub(in crate::dispatch) struct EstablishedAttemptContext { + pub request: CodexResponsesRequest, + pub controller_scope: crate::dispatch::controllers::ControllerRequestScope, + pub trace: ResponseDispatchTrace, + pub account: Account, + pub attempt: ResponseDispatchAttempt, +} + +pub(in crate::dispatch) struct EstablishedComplete { + pub context: EstablishedAttemptContext, + pub response: CodexBackendResponse, + pub body: EstablishedCompleteBody, +} + +pub(in crate::dispatch) enum EstablishedCompleteBody { + Completed(serde_json::Value), + Incomplete(serde_json::Value), +} + +impl TryFrom for EstablishedCompleteBody { + type Error = (); + + fn try_from(value: CollectedResponse) -> Result { + match value { + CollectedResponse::Completed(body) => Ok(Self::Completed(body)), + CollectedResponse::Incomplete(body) => Ok(Self::Incomplete(body)), + CollectedResponse::Failed(_) + | CollectedResponse::MissingCompleted + | CollectedResponse::Empty => Err(()), + } + } +} + +pub(in crate::dispatch) struct EstablishedStream { + pub context: EstablishedAttemptContext, + pub lease: AccountLease, + pub response: CodexBackendStreamingResponse, + pub decoder: CanonicalStreamDecoder, + pub initial_batch: CanonicalStreamBatch, +} + +pub(in crate::dispatch) enum EstablishedResponse { + Complete(Box), + Stream(Box), +} + +pub(in crate::dispatch) struct AttemptRejection { + pub request: CodexResponsesRequest, + pub trace: ResponseDispatchTrace, + pub account_id: Option, + pub account: Option, + pub attempt_request: Option, + pub attempt: Option, + pub transport: CodexBackendTransport, + pub stream_failure: Option, + pub error: ResponseDispatchError, +} + +pub(in crate::dispatch) struct RejectedStreamFailure { + pub failure: ResponsesSseFailure, + pub prefetched: Bytes, + pub diagnostics: CodexUpstreamDiagnostics, + pub rate_limit_headers: Vec<(String, String)>, +} + +pub(in crate::dispatch) enum AttemptApplyOutcome { + Continue, + Established(EstablishedResponse), + Rejected(Box), +} + +#[derive(Debug, thiserror::Error)] +pub(in crate::dispatch) enum AttemptContractError { + #[error("the previous attempt observation has not been decided")] + DecisionRequired, + #[error("attempt runner has already reached a terminal outcome")] + Terminal, + #[error("decision `{decision}` is invalid for observation `{observation}`")] + InvalidDecision { + decision: &'static str, + observation: &'static str, + }, +} + +pub(super) enum PendingProtocolFailure { + InvalidSse(SseError), + EmptyStream, + NoCommitBoundary, +} + +pub(super) struct PendingStreamResponse { + pub account: Account, + pub attempt_request: CodexResponsesRequest, + pub attempt: ResponseDispatchAttempt, + pub lease: AccountLease, + pub response: CodexBackendStreamingResponse, + pub prefetched: Bytes, + pub decoder: CanonicalStreamDecoder, + pub initial_batch: CanonicalStreamBatch, + pub first_failure: Option, +} + +pub(super) struct PendingCompleteResponse { + pub account: Account, + pub attempt_request: CodexResponsesRequest, + pub attempt: ResponseDispatchAttempt, + pub response: CodexBackendResponse, + pub collected: CollectedResponse, +} + +pub(super) enum PendingAttempt { + NoCandidate, + PinnedCandidateUnavailable { + account_id: String, + kind: PinnedCandidateAcquireFailureKind, + }, + CandidatePreparationRejected { + account: Account, + }, + RoutePreparationRejected { + account: Account, + message: String, + }, + UpstreamFailure { + account: Account, + attempt_request: CodexResponsesRequest, + attempt: ResponseDispatchAttempt, + transport: CodexBackendTransport, + error: CodexClientError, + }, + ProtocolFailure { + account: Account, + attempt_request: CodexResponsesRequest, + attempt: ResponseDispatchAttempt, + transport: CodexBackendTransport, + error: PendingProtocolFailure, + }, + CompleteResponse(Box), + StreamResponse(Box), +} diff --git a/backend/src/dispatch/lifecycle/finalizer.rs b/backend/src/dispatch/lifecycle/finalizer.rs new file mode 100644 index 000000000..da25bf4f4 --- /dev/null +++ b/backend/src/dispatch/lifecycle/finalizer.rs @@ -0,0 +1,202 @@ +//! Responses live stream 的唯一结算出口。 + +use std::time::{Duration, Instant}; + +use bytes::Bytes; +use tokio::sync::mpsc; +use tokio::time::timeout; + +use crate::{ + dispatch::{ + controllers::{ControllerRequestScope, ControllerSet, StreamControllerContext}, + errors::ResponseDispatchStreamError, + failure::sse::{STREAM_DISCONNECTED_CODE, STREAM_DISCONNECTED_MESSAGE}, + lifecycle::stream::StreamSummary, + transport::canonical::{CanonicalResponseChunk, CanonicalResponseEvent}, + }, + fleet::pool::AccountLease, + upstream::openai::protocol::sse::{ + DONE_SSE_FRAME, encode_sse_event, response_failed_sse_data_with_id, + }, +}; + +const LEASE_FINALIZE_TIMEOUT: Duration = Duration::from_millis(100); + +/// 消费 stream context，类型层面保证每个 live loop 只能结算一次。 +pub(in crate::dispatch) struct StreamFinalizer { + controllers: ControllerSet, + controller_scope: ControllerRequestScope, + controller_context: StreamControllerContext, + account_lease: AccountLease, +} + +impl StreamFinalizer { + pub(in crate::dispatch) fn new( + controllers: ControllerSet, + controller_scope: ControllerRequestScope, + controller_context: StreamControllerContext, + account_lease: AccountLease, + ) -> Self { + Self { + controllers, + controller_scope, + controller_context, + account_lease, + } + } + + pub(in crate::dispatch) fn response_headers(&self) -> &[(String, String)] { + &self.controller_context.response_metadata.client_headers + } + + pub(in crate::dispatch) fn started_at(&self) -> Instant { + self.controller_context.started_at + } + + pub(in crate::dispatch) async fn finalize( + self, + sender: &mpsc::Sender>, + event_sender: &mpsc::UnboundedSender>, + mut summary: StreamSummary, + ) { + let Self { + controllers, + controller_scope, + controller_context, + account_lease, + } = self; + let body = String::from_utf8_lossy(&summary.body).into_owned(); + controllers + .leave_stream(controller_scope, controller_context, &summary, &body) + .await; + if timeout(LEASE_FINALIZE_TIMEOUT, account_lease.complete()) + .await + .is_err() + { + tracing::warn!( + timeout_ms = LEASE_FINALIZE_TIMEOUT.as_millis(), + "Timed out finalizing response account lease" + ); + } + finish_client_stream(sender, event_sender, &mut summary).await; + } +} + +async fn finish_client_stream( + sender: &mpsc::Sender>, + event_sender: &mpsc::UnboundedSender>, + summary: &mut StreamSummary, +) { + if !summary.terminal.should_finish_client_stream() { + return; + } + if flush_terminal_chunks(sender, event_sender, &mut summary.terminal_chunks) + .await + .is_err() + { + return; + } + if let Some(detail) = summary.terminal.synthetic_failure_detail() { + if append_separator_if_needed(sender, event_sender, &mut summary.body) + .await + .is_err() + { + return; + } + let failure = response_failed_sse_data_with_id( + summary.last_response_id.as_deref(), + "server_error", + STREAM_DISCONNECTED_CODE, + &synthetic_failure_message(detail), + ); + let bytes = Bytes::from(encode_sse_event("response.failed", &failure.to_string())); + if append_client_chunk( + sender, + event_sender, + &mut summary.body, + CanonicalResponseChunk::new( + bytes, + vec![CanonicalResponseEvent::proxy_failure(failure)], + ), + ) + .await + .is_err() + { + return; + } + } + + if !body_has_done(&summary.body) { + let _ = append_client_chunk( + sender, + event_sender, + &mut summary.body, + CanonicalResponseChunk::new(Bytes::from_static(DONE_SSE_FRAME.as_bytes()), Vec::new()), + ) + .await; + } +} + +async fn flush_terminal_chunks( + sender: &mpsc::Sender>, + event_sender: &mpsc::UnboundedSender>, + terminal_chunks: &mut Vec, +) -> Result<(), ()> { + for chunk in terminal_chunks.drain(..) { + let (bytes, events) = chunk.into_parts(); + let _ = event_sender.send(events); + sender.send(Ok(bytes)).await.map_err(|_| ())?; + } + Ok(()) +} + +async fn append_separator_if_needed( + sender: &mpsc::Sender>, + event_sender: &mpsc::UnboundedSender>, + body: &mut Vec, +) -> Result<(), ()> { + let separator = if body.is_empty() + || body.ends_with(b"\n\n") + || body.ends_with(b"\r\n\r\n") + || body.ends_with(b"\r\r") + { + return Ok(()); + } else if body.ends_with(b"\n") || body.ends_with(b"\r") { + Bytes::from_static(b"\n") + } else { + Bytes::from_static(b"\n\n") + }; + append_client_chunk( + sender, + event_sender, + body, + CanonicalResponseChunk::new(separator, Vec::new()), + ) + .await +} + +async fn append_client_chunk( + sender: &mpsc::Sender>, + event_sender: &mpsc::UnboundedSender>, + body: &mut Vec, + chunk: CanonicalResponseChunk, +) -> Result<(), ()> { + body.extend_from_slice(chunk.bytes()); + let (bytes, events) = chunk.into_parts(); + let _ = event_sender.send(events); + sender.send(Ok(bytes)).await.map_err(|_| ()) +} + +fn body_has_done(body: &[u8]) -> bool { + String::from_utf8_lossy(body) + .trim_end_matches(['\r', '\n']) + .ends_with(DONE_SSE_FRAME.trim_end_matches(['\r', '\n'])) +} + +fn synthetic_failure_message(detail: &str) -> String { + if detail.trim().is_empty() { + STREAM_DISCONNECTED_MESSAGE.to_string() + } else { + format!("{STREAM_DISCONNECTED_MESSAGE}: {detail}") + } +} diff --git a/backend/src/dispatch/lifecycle/mod.rs b/backend/src/dispatch/lifecycle/mod.rs new file mode 100644 index 000000000..f624bca03 --- /dev/null +++ b/backend/src/dispatch/lifecycle/mod.rs @@ -0,0 +1,13 @@ +//! `v1/*` 请求的生命周期基础类型与入口准备。 + +pub(in crate::dispatch) mod attempt; +pub(in crate::dispatch) mod contract; +pub(in crate::dispatch) mod finalizer; +pub(in crate::dispatch) mod pipeline; +pub(in crate::dispatch) mod request; +pub(in crate::dispatch) mod stream; +pub(in crate::dispatch) mod trace; + +pub(in crate::dispatch) use request::{ + RequestContext, RequestEnterDependencies, RequestMode, enter_request, +}; diff --git a/backend/src/dispatch/lifecycle/pipeline.rs b/backend/src/dispatch/lifecycle/pipeline.rs new file mode 100644 index 000000000..237360995 --- /dev/null +++ b/backend/src/dispatch/lifecycle/pipeline.rs @@ -0,0 +1,53 @@ +//! Attempt 生命周期的唯一静态 pipeline。 + +use crate::{ + dispatch::{controllers::ControllerRequestScope, routing::candidates::AccountAttemptLedger}, + upstream::openai::protocol::responses::CodexResponsesRequest, +}; + +use super::{ + attempt::{AttemptMode, AttemptRunner, AttemptRunnerDependencies, AttemptStep}, + contract::{AttemptApplyOutcome, AttemptContractError, AttemptRejection, EstablishedResponse}, +}; + +pub(in crate::dispatch) enum AttemptPipelineOutcome { + Established(EstablishedResponse), + Rejected(Box), +} + +/// 执行完整的 attempt/retry/commit 状态机。 +/// +/// mode 适配器只能消费最终结果，不能自行复制 retry loop 或越过 commit 边界。 +pub(in crate::dispatch) async fn run_attempt_pipeline( + dependencies: AttemptRunnerDependencies<'_>, + mode: AttemptMode, + request: CodexResponsesRequest, + controller_scope: ControllerRequestScope, + candidates: AccountAttemptLedger, +) -> Result { + let controllers = dependencies.controllers; + let mut runner = AttemptRunner::new(dependencies, mode, request, controller_scope, candidates); + + loop { + let outcome = match runner.next().await? { + AttemptStep::Open(attempt) => { + let decision = controllers + .handle_attempt(attempt.controller_scope()?, attempt.observation()) + .await; + attempt.apply(decision).await? + } + AttemptStep::Committed(attempt) => { + return Ok(AttemptPipelineOutcome::Established(attempt.accept())); + } + }; + match outcome { + AttemptApplyOutcome::Continue => {} + AttemptApplyOutcome::Established(established) => { + return Ok(AttemptPipelineOutcome::Established(established)); + } + AttemptApplyOutcome::Rejected(rejection) => { + return Ok(AttemptPipelineOutcome::Rejected(rejection)); + } + } + } +} diff --git a/backend/src/dispatch/lifecycle/request.rs b/backend/src/dispatch/lifecycle/request.rs new file mode 100644 index 000000000..6e9c97aa3 --- /dev/null +++ b/backend/src/dispatch/lifecycle/request.rs @@ -0,0 +1,101 @@ +//! Responses 请求进入 dispatch 的共享生命周期阶段。 + +use chrono::Utc; +use serde_json::Value; + +use crate::{ + dispatch::{ + affinity::{AccountIdentityService, prepare_variant_identity}, + controllers::{ControllerEnter, ControllerRequestScope, ControllerSet}, + routing::candidates::AccountAttemptLedger, + }, + fleet::pool::{AccountAcquireRequest, AccountPoolService}, + models::service::ModelService, + upstream::openai::protocol::responses::CodexResponsesRequest, +}; + +/// Responses 请求的执行模式。 +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(in crate::dispatch) enum RequestMode { + Complete, + Stream, +} + +impl RequestMode { + fn stream(self) -> bool { + matches!(self, Self::Stream) + } +} + +/// 请求 enter 阶段所需的窄依赖集合。 +pub(in crate::dispatch) struct RequestEnterDependencies<'a> { + pub account_pool: &'a AccountPoolService, + pub models: &'a ModelService, + pub account_identity: &'a AccountIdentityService, + pub controllers: ControllerSet, +} + +/// 已完成请求级 enter 的 Responses 上下文。 +/// +/// 该类型只承载请求、会话恢复、会话策略和候选快照；账号 attempt +/// 与流生命周期由后续 pipeline 接管。 +pub(in crate::dispatch) struct RequestContext { + pub request: CodexResponsesRequest, + pub display_model: String, + pub compact: bool, + pub tuple_schema: Option, + pub image_generation_requested: bool, + pub controllers: ControllerSet, + pub controller_scope: ControllerRequestScope, + pub candidates: AccountAttemptLedger, +} + +/// 执行所有 Responses 模式共享的请求 enter 阶段。 +pub(in crate::dispatch) async fn enter_request( + dependencies: RequestEnterDependencies<'_>, + mut request: CodexResponsesRequest, + requested_model: &str, + mode: RequestMode, +) -> RequestContext { + let catalog = dependencies.models.catalog().await; + let display_model = catalog.resolve_model_id(requested_model); + request.set_model(display_model.clone()); + request.set_stream(mode.stream()); + + let compact = request.semantics().compact; + let tuple_schema = request.tuple_schema.clone(); + let image_generation_requested = request.expects_image_generation(); + let now = Utc::now(); + prepare_variant_identity(&mut request); + dependencies + .account_identity + .prepare_local_identity(&mut request); + let controllers = dependencies.controllers; + let ControllerEnter { + scope, + preferred_account_id, + excluded_account_ids, + } = controllers.enter(&request, now).await; + + let mut acquire_request = AccountAcquireRequest::new(request.model(), now); + if let Some(preferred_account_id) = preferred_account_id { + acquire_request = acquire_request.with_preferred_account_id(preferred_account_id); + } + let candidates = AccountAttemptLedger::freeze( + dependencies.account_pool, + &acquire_request, + &excluded_account_ids, + ) + .await; + + RequestContext { + request, + display_model, + compact, + tuple_schema, + image_generation_requested, + controllers, + controller_scope: scope, + candidates, + } +} diff --git a/backend/src/dispatch/lifecycle/stream.rs b/backend/src/dispatch/lifecycle/stream.rs new file mode 100644 index 000000000..9f969ad58 --- /dev/null +++ b/backend/src/dispatch/lifecycle/stream.rs @@ -0,0 +1,84 @@ +//! Responses live stream 的规范事件与终态。 + +use serde_json::Value; + +use crate::{ + dispatch::transport::canonical::{CanonicalResponseChunk, CanonicalStreamTerminal}, + upstream::openai::protocol::{events::TokenUsage, responses::ResponsesSseFailure}, +}; + +/// live stream 离开事件循环的唯一终态。 +#[derive(Debug, Clone, PartialEq, Eq)] +pub(in crate::dispatch) enum StreamTerminal { + Completed { response: Value }, + Incomplete { response: Value }, + Failed { failure: ResponsesSseFailure }, + UpstreamClosed, + UpstreamError { detail: String }, + ProtocolError { detail: String }, + CaptureLimitExceeded, + Cancelled, + DownstreamClosed, + Shutdown, +} + +/// live loop 交给唯一 finalizer 的完整事实快照。 +pub(in crate::dispatch) struct StreamSummary { + pub terminal: StreamTerminal, + pub body: Vec, + pub terminal_chunks: Vec, + pub first_token_ms: Option, + pub usage: Option, + pub last_response_id: Option, +} + +impl From for StreamTerminal { + fn from(terminal: CanonicalStreamTerminal) -> Self { + match terminal { + CanonicalStreamTerminal::Completed(response) => Self::Completed { response }, + CanonicalStreamTerminal::Incomplete(response) => Self::Incomplete { response }, + CanonicalStreamTerminal::Failed(failure) => Self::Failed { failure }, + } + } +} + +impl StreamTerminal { + /// 只有上游没有给出业务终态时，代理才需要合成 `response.failed`。 + pub(in crate::dispatch) fn synthetic_failure_detail(&self) -> Option<&str> { + match self { + Self::UpstreamClosed => Some(""), + Self::UpstreamError { detail } | Self::ProtocolError { detail } => Some(detail), + Self::CaptureLimitExceeded => { + Some("upstream response exceeded the 16 MiB proxy capture limit") + } + Self::Completed { .. } + | Self::Incomplete { .. } + | Self::Failed { .. } + | Self::Cancelled + | Self::DownstreamClosed + | Self::Shutdown => None, + } + } + + pub(in crate::dispatch) fn should_finish_client_stream(&self) -> bool { + !matches!( + self, + Self::Cancelled | Self::DownstreamClosed | Self::Shutdown + ) + } + + pub(in crate::dispatch) fn name(&self) -> &'static str { + match self { + Self::Completed { .. } => "completed", + Self::Incomplete { .. } => "incomplete", + Self::Failed { .. } => "failed", + Self::UpstreamClosed => "upstream_closed", + Self::UpstreamError { .. } => "upstream_error", + Self::ProtocolError { .. } => "protocol_error", + Self::CaptureLimitExceeded => "capture_limit_exceeded", + Self::Cancelled => "cancelled", + Self::DownstreamClosed => "downstream_closed", + Self::Shutdown => "shutdown", + } + } +} diff --git a/backend/src/dispatch/stream/trace.rs b/backend/src/dispatch/lifecycle/trace.rs similarity index 94% rename from backend/src/dispatch/stream/trace.rs rename to backend/src/dispatch/lifecycle/trace.rs index 936834b05..7e61f0ec6 100644 --- a/backend/src/dispatch/stream/trace.rs +++ b/backend/src/dispatch/lifecycle/trace.rs @@ -1,4 +1,4 @@ -//! Request-level trace state for Responses dispatch. +//! Responses 请求级账号尝试轨迹。 #[derive(Debug, Default)] pub(in crate::dispatch) struct ResponseDispatchTrace { diff --git a/backend/src/dispatch/mod.rs b/backend/src/dispatch/mod.rs index 73a5e94d0..4c2e8de83 100644 --- a/backend/src/dispatch/mod.rs +++ b/backend/src/dispatch/mod.rs @@ -1,10 +1,11 @@ //! 请求调度模块（账号选择、回退、恢复、会话亲和性）。 pub mod affinity; -mod attempts; +pub(crate) mod controllers; pub mod errors; -mod recording; -pub mod recovery; +pub(crate) mod failure; +pub(crate) mod lifecycle; +pub(crate) mod routing; pub(crate) mod service; pub(crate) mod stream; -pub mod upstream_call; +pub(crate) mod transport; diff --git a/backend/src/dispatch/recovery/account_failure.rs b/backend/src/dispatch/recovery/account_failure.rs deleted file mode 100644 index da0aa22fb..000000000 --- a/backend/src/dispatch/recovery/account_failure.rs +++ /dev/null @@ -1,147 +0,0 @@ -//! 可换号上游失败的账号风险剥离。 - -use chrono::Utc; - -use crate::{ - dispatch::{ - errors::{ - is_model_unsupported_upstream_error, is_quota_exhausted_upstream_error, - is_rate_limit_upstream_error, rate_limit_cooldown_until, upstream_error_body, - upstream_error_http_status, - }, - recovery::{ - auth::{auth_failure_account_status, is_auth_upstream_error}, - cloudflare::{ - CloudflareRecovery, cloudflare_challenge_error_message, - cloudflare_path_block_error_message, is_cloudflare_challenge_upstream_error, - is_cloudflare_path_block_upstream_error, - }, - exhaustion::AccountExhaustionTracker, - }, - stream::sse_failure::{ - auth_sse_failure_account_status, is_auth_sse_failure, is_model_unsupported_sse_failure, - is_quota_exhausted_sse_failure, sse_failure_error_body, stream_failure_http_status, - }, - }, - fleet::{account::AccountStatus, pool::AccountPoolService}, - upstream::openai::{ - protocol::responses::ResponsesSseFailure, - transport::{CodexClientError, is_banned_upstream_error, is_cyber_policy_upstream_error}, - }, -}; - -/// 处理需要隔离当前账号并交回调度器选择下一候选的上游失败。 -/// -/// 本函数不选择下一个账号，也不实现固定轮换策略；调用方继续使用请求级候选账本， -/// 因而始终服从当前运行时调度策略。 -pub(crate) async fn isolate_rotatable_account_failure( - account_pool: &AccountPoolService, - cloudflare: &CloudflareRecovery, - exhausted_accounts: &mut AccountExhaustionTracker, - account_id: &str, - error: &CodexClientError, -) -> bool { - if is_cyber_policy_upstream_error(error) { - return false; - } - - if is_rate_limit_upstream_error(error) { - exhausted_accounts.record_rate_limited(Some(account_id), upstream_error_body(error)); - account_pool - .mark_quota_limited_until(account_id, rate_limit_cooldown_until(error, Utc::now())) - .await; - return true; - } - - if is_banned_upstream_error(error) { - exhausted_accounts.record_auth_failure( - Some(account_id), - AccountStatus::Banned, - upstream_error_body(error), - Some(upstream_error_http_status(error)), - ); - account_pool - .set_status(account_id, AccountStatus::Banned) - .await; - return true; - } - - if is_quota_exhausted_upstream_error(error) { - exhausted_accounts.record_quota_exhausted(Some(account_id), upstream_error_body(error)); - account_pool - .set_status(account_id, AccountStatus::QuotaExhausted) - .await; - return true; - } - - if is_auth_upstream_error(error) { - let account_status = auth_failure_account_status(error); - exhausted_accounts.record_auth_failure( - Some(account_id), - account_status, - upstream_error_body(error), - Some(upstream_error_http_status(error)), - ); - account_pool.set_status(account_id, account_status).await; - return true; - } - - if is_cloudflare_challenge_upstream_error(error) { - exhausted_accounts - .record_cloudflare_challenge(Some(account_id), cloudflare_challenge_error_message()); - cloudflare.apply_challenge(account_pool, account_id).await; - return true; - } - - if is_cloudflare_path_block_upstream_error(error) { - exhausted_accounts.record_cloudflare_path_blocked( - Some(account_id), - cloudflare_path_block_error_message(), - ); - cloudflare.apply_path_block(account_pool, account_id).await; - return true; - } - - if is_model_unsupported_upstream_error(error) { - exhausted_accounts.record_model_unsupported(Some(account_id), upstream_error_body(error)); - return true; - } - - false -} - -/// 处理 SSE `response.failed` 中需要隔离当前账号的失败。 -pub(crate) async fn isolate_sse_account_failure( - account_pool: &AccountPoolService, - exhausted_accounts: &mut AccountExhaustionTracker, - account_id: &str, - failure: &ResponsesSseFailure, -) -> bool { - if is_model_unsupported_sse_failure(failure) { - exhausted_accounts - .record_model_unsupported(Some(account_id), sse_failure_error_body(failure)); - return true; - } - - if is_quota_exhausted_sse_failure(failure) { - exhausted_accounts.record_quota_exhausted(Some(account_id), failure.message.clone()); - account_pool - .set_status(account_id, AccountStatus::QuotaExhausted) - .await; - return true; - } - - if is_auth_sse_failure(failure) { - let account_status = auth_sse_failure_account_status(failure); - exhausted_accounts.record_auth_failure( - Some(account_id), - account_status, - sse_failure_error_body(failure), - Some(stream_failure_http_status(failure)), - ); - account_pool.set_status(account_id, account_status).await; - return true; - } - - false -} diff --git a/backend/src/dispatch/recovery/auth.rs b/backend/src/dispatch/recovery/auth.rs deleted file mode 100644 index 74709e815..000000000 --- a/backend/src/dispatch/recovery/auth.rs +++ /dev/null @@ -1,22 +0,0 @@ -//! 上游认证失败分类。 - -use crate::{ - fleet::account::AccountStatus, - upstream::openai::transport::{CodexClientError, is_banned_auth_signal}, -}; - -pub(crate) fn is_auth_upstream_error(error: &CodexClientError) -> bool { - matches!( - error, - CodexClientError::Upstream { status, .. } if status.as_u16() == 401 - ) -} - -pub(crate) fn auth_failure_account_status(error: &CodexClientError) -> AccountStatus { - match error { - CodexClientError::Upstream { body, .. } if is_banned_auth_signal(body) => { - AccountStatus::Banned - } - _ => AccountStatus::Expired, - } -} diff --git a/backend/src/dispatch/recovery/cloudflare.rs b/backend/src/dispatch/recovery/cloudflare.rs deleted file mode 100644 index d41b7df63..000000000 --- a/backend/src/dispatch/recovery/cloudflare.rs +++ /dev/null @@ -1,146 +0,0 @@ -//! Cloudflare challenge and path-block recovery shared by dispatch routes. - -use chrono::Utc; - -use crate::{ - fleet::{ - account::AccountStatus, - cookies::{CloudflareChallengeCooldownTracker, CloudflarePathBlockTracker, PgCookieStore}, - pool::AccountPoolService, - }, - upstream::openai::transport::CodexClientError, -}; - -#[derive(Clone)] -pub struct CloudflareRecovery { - path_block_tracker: CloudflarePathBlockTracker, - challenge_tracker: CloudflareChallengeCooldownTracker, - cookie_store: PgCookieStore, -} - -impl CloudflareRecovery { - pub fn new(cookie_store: PgCookieStore) -> Self { - Self { - path_block_tracker: CloudflarePathBlockTracker::default(), - challenge_tracker: CloudflareChallengeCooldownTracker::default(), - cookie_store, - } - } - - pub async fn cookie_header_for_request(&self, account_id: &str, path: &str) -> Option { - self.cookie_store - .cookie_header_for_request(account_id, "chatgpt.com", path) - .await - .ok()? - } - - pub async fn capture_set_cookie_headers(&self, account_id: &str, headers: &[String]) { - for header in headers { - if let Err(error) = self - .cookie_store - .capture_set_cookie(account_id, header) - .await - { - tracing::warn!( - account_id, - error = %error, - "Failed to persist upstream set-cookie header" - ); - } - } - } - - pub async fn apply_challenge(&self, account_pool: &AccountPoolService, account_id: &str) { - let now = Utc::now(); - let cooldown = self - .challenge_tracker - .record_challenge(account_id, now) - .await; - account_pool - .set_cloudflare_cooldown_until(account_id, cooldown.cooldown_until) - .await; - if let Err(error) = self - .cookie_store - .expire_account_cookies_at(account_id, cooldown.cooldown_until) - .await - { - tracing::warn!( - account_id, - error = %error, - "Failed to persist Cloudflare challenge cookie cleanup deadline" - ); - } - } - - pub async fn apply_path_block(&self, account_pool: &AccountPoolService, account_id: &str) { - self.delete_account_cookies(account_id, "Cloudflare path-block") - .await; - let now = Utc::now(); - self.path_block_tracker - .record_path_block(account_id, now) - .await; - if self - .path_block_tracker - .should_disable(account_id, now) - .await - { - account_pool - .set_status(account_id, AccountStatus::Disabled) - .await; - } - } - - pub async fn reset_account_recovery(&self, account_id: &str) { - self.path_block_tracker.reset(account_id).await; - self.challenge_tracker.reset(account_id).await; - } - - async fn delete_account_cookies(&self, account_id: &str, reason: &str) { - delete_account_cookies(&self.cookie_store, account_id, reason).await; - } -} - -async fn delete_account_cookies(cookie_store: &PgCookieStore, account_id: &str, reason: &str) { - if let Err(error) = cookie_store.delete_account_cookies(account_id).await { - tracing::warn!( - account_id, - reason, - error = %error, - "Failed to delete account cookies after Cloudflare recovery signal" - ); - } -} - -pub fn is_cloudflare_challenge_upstream_error(error: &CodexClientError) -> bool { - matches!( - error, - CodexClientError::Upstream { status, body, .. } - if status.as_u16() == 403 && is_cloudflare_challenge_signal(body) - ) -} - -pub fn is_cloudflare_path_block_upstream_error(error: &CodexClientError) -> bool { - matches!( - error, - CodexClientError::Upstream { status, body, .. } - if status.as_u16() == 404 && body.trim().is_empty() - ) -} - -pub fn cloudflare_challenge_error_message() -> &'static str { - "Upstream blocked the request (Cloudflare challenge)" -} - -pub fn cloudflare_path_block_error_message() -> &'static str { - "Upstream blocked the request (Cloudflare path-block)" -} - -fn is_cloudflare_challenge_signal(value: &str) -> bool { - let value = value.to_ascii_lowercase(); - value.contains("cf-mitigated") - || value.contains("cf-chl-bypass") - || value.contains("_cf_chl") - || value.contains("cf_chl") - || value.contains("attention required") - || value.contains("just a moment") -} diff --git a/backend/src/dispatch/recovery/cyber_policy.rs b/backend/src/dispatch/recovery/cyber_policy.rs deleted file mode 100644 index 6673ab1bd..000000000 --- a/backend/src/dispatch/recovery/cyber_policy.rs +++ /dev/null @@ -1,386 +0,0 @@ -//! `cyber_policy` 的会话级换号恢复。 - -use std::{collections::BTreeSet, sync::Arc, time::Duration as StdDuration}; - -use chrono::Duration; -use serde_json::Value; -use sha2::{Digest, Sha256}; -use tokio::time::timeout; - -use crate::{ - dispatch::affinity::{ - CyberPolicyFailureSnapshot, CyberPolicySessionState, SessionAffinityService, - }, - upstream::openai::{ - protocol::{ - responses::{CodexResponsesRequest, ResponsesSseFailure}, - sse::{SseEvent, SseEventDecoder}, - }, - transport::{CodexClientError, is_cyber_policy_upstream_error}, - }, -}; - -const MAX_ROTATED_ACCOUNTS: usize = 3; -const MAX_FAILURE_MESSAGE_BYTES: usize = 4 * 1024; -const SESSION_STATE_TTL: Duration = Duration::hours(1); -const STATE_IO_TIMEOUT: StdDuration = StdDuration::from_millis(100); - -/// 单次请求使用的 `cyber_policy` 会话路由计划。 -pub(in crate::dispatch) struct CyberPolicyRoutingPlan { - session_key: Option, - state: CyberPolicySessionState, -} - -impl CyberPolicyRoutingPlan { - pub(in crate::dispatch) fn excluded_account_ids(&self) -> BTreeSet { - self.state.failed_account_ids.iter().cloned().collect() - } - - pub(in crate::dispatch) fn exhausted_failure(&self) -> Option { - if self.state.failed_account_ids.len() < MAX_ROTATED_ACCOUNTS { - return None; - } - self.last_failure() - } - - pub(in crate::dispatch) fn last_failure(&self) -> Option { - self.state - .last_failure - .as_ref() - .map(snapshot_to_sse_failure) - } - - pub(in crate::dispatch) fn last_account_id(&self) -> Option<&str> { - self.state - .last_failure - .as_ref() - .map(|failure| failure.account_id.as_str()) - } - - fn session_key(&self) -> Option<&str> { - self.session_key.as_deref() - } - - fn enabled(&self) -> bool { - self.session_key.is_some() - } - - fn has_failures(&self) -> bool { - !self.state.failed_account_ids.is_empty() - } -} - -/// 识别失败、维护会话状态，并为下一次请求生成账号排除集合。 -#[derive(Clone)] -pub(in crate::dispatch) struct CyberPolicyRecovery { - session_affinity: Arc, -} - -/// 在 live SSE 转发前观察终止事件，保证下一请求看到最新会话状态。 -#[derive(Default)] -pub(in crate::dispatch) struct CyberPolicyStreamObserver { - decoder: SseEventDecoder, - decoder_stopped: bool, - state_transition_applied: bool, -} - -impl CyberPolicyStreamObserver { - pub(in crate::dispatch) async fn observe_chunk( - &mut self, - recovery: &CyberPolicyRecovery, - plan: &CyberPolicyRoutingPlan, - account_id: &str, - chunk: &[u8], - ) { - if self.decoder_stopped || self.state_transition_applied || !plan.enabled() { - return; - } - match self.decoder.push(chunk) { - Ok(events) => { - self.observe_events(recovery, plan, account_id, &events) - .await - } - Err(error) => { - self.decoder_stopped = true; - tracing::warn!(error = %error, "Failed to decode live SSE for cyber policy state"); - } - } - } - - pub(in crate::dispatch) async fn finish( - &mut self, - recovery: &CyberPolicyRecovery, - plan: &CyberPolicyRoutingPlan, - account_id: &str, - ) { - if self.decoder_stopped || self.state_transition_applied || !plan.enabled() { - return; - } - match self.decoder.finish() { - Ok(events) => { - self.observe_events(recovery, plan, account_id, &events) - .await - } - Err(error) => { - tracing::warn!(error = %error, "Failed to finish live SSE cyber policy state"); - } - } - } - - pub(in crate::dispatch) fn state_transition_applied(&self) -> bool { - self.state_transition_applied - } - - async fn observe_events( - &mut self, - recovery: &CyberPolicyRecovery, - plan: &CyberPolicyRoutingPlan, - account_id: &str, - events: &[SseEvent], - ) { - let mut completed = false; - for event in events { - match event.event.as_deref() { - Some("response.completed" | "response.incomplete") => { - completed |= serde_json::from_str::(&event.data) - .is_ok_and(|value| value.get("response").is_some_and(Value::is_object)); - continue; - } - Some(event_name @ ("error" | "response.failed")) => { - let Ok(value) = serde_json::from_str::(&event.data) else { - continue; - }; - let failure = ResponsesSseFailure::from_event(event_name, &value); - if is_cyber_policy_failure(&failure) { - recovery - .observe_sse_failure(plan, account_id, &failure) - .await; - self.state_transition_applied = true; - return; - } - continue; - } - Some(_) => continue, - None => {} - } - let Ok(value) = serde_json::from_str::(&event.data) else { - continue; - }; - match value.get("type").and_then(Value::as_str) { - Some(event_name @ ("error" | "response.failed")) => { - let failure = ResponsesSseFailure::from_event(event_name, &value); - if is_cyber_policy_failure(&failure) { - recovery - .observe_sse_failure(plan, account_id, &failure) - .await; - self.state_transition_applied = true; - return; - } - } - Some("response.completed" | "response.incomplete") => { - completed |= value.get("response").is_some_and(Value::is_object); - } - _ => {} - } - } - if completed { - recovery.observe_success(plan).await; - self.state_transition_applied = true; - } - } -} - -impl CyberPolicyRecovery { - pub(in crate::dispatch) fn new(session_affinity: Arc) -> Self { - Self { session_affinity } - } - - pub(in crate::dispatch) async fn prepare( - &self, - request: &CodexResponsesRequest, - ) -> CyberPolicyRoutingPlan { - let Some(session_key) = session_key(request) else { - return CyberPolicyRoutingPlan { - session_key: None, - state: CyberPolicySessionState::default(), - }; - }; - let state = match timeout( - STATE_IO_TIMEOUT, - self.session_affinity.load_cyber_policy_state(&session_key), - ) - .await - { - Ok(Ok(state)) => state.unwrap_or_default(), - Ok(Err(error)) => { - tracing::warn!(error = %error, "Failed to read cyber policy session state"); - CyberPolicySessionState::default() - } - Err(_) => { - tracing::warn!("Timed out reading cyber policy session state"); - CyberPolicySessionState::default() - } - }; - CyberPolicyRoutingPlan { - session_key: Some(session_key), - state, - } - } - - pub(in crate::dispatch) async fn observe_sse_failure( - &self, - plan: &CyberPolicyRoutingPlan, - account_id: &str, - failure: &ResponsesSseFailure, - ) { - if !is_cyber_policy_failure(failure) { - return; - } - let Some(session_key) = plan.session_key() else { - return; - }; - let snapshot = CyberPolicyFailureSnapshot { - account_id: account_id.to_string(), - event: failure.event.clone(), - message: truncate_utf8(&failure.message, MAX_FAILURE_MESSAGE_BYTES), - upstream_code: failure.upstream_code.clone(), - }; - match timeout( - STATE_IO_TIMEOUT, - self.session_affinity.persist_cyber_policy_failure( - session_key, - &snapshot, - MAX_ROTATED_ACCOUNTS, - SESSION_STATE_TTL, - ), - ) - .await - { - Ok(Ok(state)) => tracing::warn!( - account_id, - failed_account_count = state.failed_account_ids.len(), - "Recorded cyber policy failure for the next request in this session" - ), - Ok(Err(error)) => tracing::warn!( - account_id, - error = %error, - "Failed to record cyber policy session state" - ), - Err(_) => tracing::warn!(account_id, "Timed out recording cyber policy session state"), - } - } - - pub(in crate::dispatch) async fn observe_upstream_error( - &self, - plan: &CyberPolicyRoutingPlan, - account_id: &str, - error: &CodexClientError, - ) { - let Some(failure) = failure_from_upstream_error(error) else { - return; - }; - self.observe_sse_failure(plan, account_id, &failure).await; - } - - pub(in crate::dispatch) async fn observe_success(&self, plan: &CyberPolicyRoutingPlan) { - if !plan.has_failures() { - return; - } - let Some(session_key) = plan.session_key() else { - return; - }; - match timeout( - STATE_IO_TIMEOUT, - self.session_affinity - .delete_cyber_policy_state(session_key, &plan.state.revision), - ) - .await - { - Ok(Ok(_)) => {} - Ok(Err(error)) => { - tracing::warn!(error = %error, "Failed to clear cyber policy session state"); - } - Err(_) => tracing::warn!("Timed out clearing cyber policy session state"), - } - } -} - -fn session_key(request: &CodexResponsesRequest) -> Option { - if request.previous_response_id().is_some() { - return None; - } - let api_key_id = non_empty(request.client_api_key_id.as_deref())?; - let explicit_session_id = non_empty(request.client_session_id.as_deref()) - .or_else(|| non_empty(request.client_conversation_id.as_deref())) - .or_else(|| { - request - .explicit_prompt_cache_key - .then(|| non_empty(request.prompt_cache_key())) - .flatten() - })?; - let mut hasher = Sha256::new(); - hasher.update(b"cyber-policy-session\0"); - hasher.update(api_key_id.as_bytes()); - hasher.update(b"\0"); - hasher.update(explicit_session_id.as_bytes()); - Some(hex::encode(hasher.finalize())) -} - -fn failure_from_upstream_error(error: &CodexClientError) -> Option { - let CodexClientError::Upstream { body, .. } = error else { - return None; - }; - if !is_cyber_policy_upstream_error(error) { - return None; - } - let value = serde_json::from_str::(body).ok()?; - let code = value - .pointer("/response/error/code") - .or_else(|| value.pointer("/error/code")) - .or_else(|| value.get("code")) - .and_then(Value::as_str)?; - let message = value - .pointer("/response/error/message") - .or_else(|| value.pointer("/error/message")) - .or_else(|| value.get("message")) - .and_then(Value::as_str) - .filter(|message| !message.trim().is_empty()) - .unwrap_or("Codex upstream SSE failed") - .to_string(); - Some(ResponsesSseFailure { - event: "error".to_string(), - message, - upstream_code: Some(code.to_string()), - }) -} - -fn is_cyber_policy_code(code: Option<&str>) -> bool { - code.is_some_and(|code| code.trim().eq_ignore_ascii_case("cyber_policy")) -} - -pub(crate) fn is_cyber_policy_failure(failure: &ResponsesSseFailure) -> bool { - is_cyber_policy_code(failure.upstream_code.as_deref()) -} - -fn snapshot_to_sse_failure(snapshot: &CyberPolicyFailureSnapshot) -> ResponsesSseFailure { - ResponsesSseFailure { - event: snapshot.event.clone(), - message: snapshot.message.clone(), - upstream_code: snapshot.upstream_code.clone(), - } -} - -fn non_empty(value: Option<&str>) -> Option<&str> { - value.map(str::trim).filter(|value| !value.is_empty()) -} - -fn truncate_utf8(value: &str, max_bytes: usize) -> String { - if value.len() <= max_bytes { - return value.to_string(); - } - let mut end = max_bytes; - while !value.is_char_boundary(end) { - end -= 1; - } - value[..end].to_string() -} diff --git a/backend/src/dispatch/recovery/history.rs b/backend/src/dispatch/recovery/history.rs deleted file mode 100644 index 5b555be25..000000000 --- a/backend/src/dispatch/recovery/history.rs +++ /dev/null @@ -1,192 +0,0 @@ -//! previous response 所有权与连接内上下文重放。 - -use std::sync::Arc; - -use chrono::Utc; -use serde_json::Value; - -use crate::{ - dispatch::affinity::{SessionAffinityEntry, SessionAffinityService}, - upstream::openai::protocol::responses::{CodexResponsesRequest, StreamCommitPolicy}, -}; - -/// 当前请求的 previous response 来源。 -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -pub(in crate::dispatch) enum HistorySource { - None, - Managed, - ConnectionLocal, - ExternalUnknown, -} - -/// 请求级 previous response 恢复计划。 -pub(in crate::dispatch) struct HistoryRecoveryPlan { - source: HistorySource, - managed_entry: Option, - local_replay_input: Option>>, - external_attempted_account_id: Option, - replay_owner_account: bool, -} - -impl HistoryRecoveryPlan { - pub(in crate::dispatch) async fn load( - session_affinity: &SessionAffinityService, - request: &CodexResponsesRequest, - ) -> Self { - let Some(previous_response_id) = request.previous_response_id() else { - return Self { - source: HistorySource::None, - managed_entry: None, - local_replay_input: None, - external_attempted_account_id: None, - replay_owner_account: false, - }; - }; - let managed_entry = session_affinity - .lookup(previous_response_id, Utc::now()) - .await; - let local_replay_input = request.local_replay_input.clone(); - let source = match (&managed_entry, &local_replay_input) { - (Some(_), _) => HistorySource::Managed, - (None, Some(_)) => HistorySource::ConnectionLocal, - (None, None) => HistorySource::ExternalUnknown, - }; - Self { - source, - managed_entry, - local_replay_input, - external_attempted_account_id: None, - replay_owner_account: false, - } - } - - pub(in crate::dispatch) fn preferred_account_id(&self) -> Option<&str> { - self.managed_entry - .as_ref() - .map(|entry| entry.account_id.as_str()) - } - - pub(in crate::dispatch) fn prepare_attempt( - &mut self, - original: &CodexResponsesRequest, - account_id: &str, - ) -> Option { - match self.source { - HistorySource::None => Some(original.clone()), - HistorySource::ExternalUnknown => match self.external_attempted_account_id.as_deref() { - None => { - self.external_attempted_account_id = Some(account_id.to_string()); - Some(external_history_request(original)) - } - Some(first_account_id) if first_account_id == account_id => { - Some(external_history_request(original)) - } - Some(_) => None, - }, - HistorySource::ConnectionLocal => self - .local_replay_input - .as_deref() - .map(|replay_input| full_replay_request(original, replay_input)), - HistorySource::Managed => { - let entry = self.managed_entry.as_ref()?; - if entry.account_id == account_id && !self.replay_owner_account { - let mut request = original.clone(); - request.local_conversation_id = Some(entry.conversation_id.clone()); - request.stream_commit_policy = StreamCommitPolicy::UntilOutputOrTerminal; - request.previous_response_scope = Some(entry.continuation_scope); - if request.turn_state.as_deref().is_none_or(str::is_empty) { - request.turn_state.clone_from(&entry.turn_state); - } - Some(request) - } else { - self.local_replay_input.as_deref().map(|replay_input| { - let mut request = full_replay_request(original, replay_input); - request.local_conversation_id = Some(entry.conversation_id.clone()); - request - }) - } - } - } - } - - /// previous ID 在原连接不可用时，切换为同账号完整重放。 - pub(in crate::dispatch) fn recover_managed_history(&mut self, account_id: &str) -> bool { - let Some(entry) = self.managed_entry.as_ref() else { - return false; - }; - if entry.account_id != account_id - || self.local_replay_input.is_none() - || self.replay_owner_account - { - return false; - } - self.replay_owner_account = true; - true - } - - pub(in crate::dispatch) fn can_failover(&self) -> bool { - match self.source { - HistorySource::None | HistorySource::ConnectionLocal => true, - HistorySource::Managed => self.local_replay_input.is_some(), - HistorySource::ExternalUnknown => false, - } - } - - pub(in crate::dispatch) fn is_external_unknown(&self) -> bool { - self.source == HistorySource::ExternalUnknown - } - - pub(in crate::dispatch) fn conversation_id<'a>( - &'a self, - original: &'a CodexResponsesRequest, - ) -> Option<&'a str> { - self.managed_entry - .as_ref() - .map(|entry| entry.conversation_id.as_str()) - .or_else(|| non_empty(original.local_conversation_id.as_deref())) - } -} - -fn external_history_request(original: &CodexResponsesRequest) -> CodexResponsesRequest { - let mut request = original.clone(); - request.stream_commit_policy = StreamCommitPolicy::UntilOutputOrTerminal; - request.previous_response_scope = - Some(crate::upstream::openai::protocol::responses::PreviousResponseScope::ExternalUnknown); - request -} - -fn full_replay_request( - original: &CodexResponsesRequest, - replay_input: &[Value], -) -> CodexResponsesRequest { - let mut request = original.clone(); - let mut input = Vec::with_capacity(replay_input.len() + original.input().len()); - input.extend_from_slice(replay_input); - input.extend_from_slice(original.input()); - sanitize_replay_items(&mut input); - request.set_input(input); - request.set_previous_response_id(None); - request.turn_state = None; - request.stream_commit_policy = StreamCommitPolicy::UntilOutputOrTerminal; - request -} - -pub(crate) fn sanitize_replay_items(values: &mut [Value]) { - for value in values { - match value { - Value::Array(values) => sanitize_replay_items(values), - Value::Object(object) => { - object.remove("id"); - object.remove("encrypted_content"); - for value in object.values_mut() { - sanitize_replay_items(std::slice::from_mut(value)); - } - } - Value::Null | Value::Bool(_) | Value::Number(_) | Value::String(_) => {} - } - } -} - -fn non_empty(value: Option<&str>) -> Option<&str> { - value.map(str::trim).filter(|value| !value.is_empty()) -} diff --git a/backend/src/dispatch/recovery/mod.rs b/backend/src/dispatch/recovery/mod.rs deleted file mode 100644 index 8570c7539..000000000 --- a/backend/src/dispatch/recovery/mod.rs +++ /dev/null @@ -1,8 +0,0 @@ -//! 调度失败恢复策略。 - -pub(crate) mod account_failure; -pub(crate) mod auth; -pub mod cloudflare; -pub(crate) mod cyber_policy; -pub(crate) mod exhaustion; -pub(crate) mod history; diff --git a/backend/src/dispatch/attempts.rs b/backend/src/dispatch/routing/candidates.rs similarity index 100% rename from backend/src/dispatch/attempts.rs rename to backend/src/dispatch/routing/candidates.rs diff --git a/backend/src/dispatch/routing/mod.rs b/backend/src/dispatch/routing/mod.rs new file mode 100644 index 000000000..8359f7e09 --- /dev/null +++ b/backend/src/dispatch/routing/mod.rs @@ -0,0 +1,3 @@ +//! 请求级通用账号候选路由。 + +pub(in crate::dispatch) mod candidates; diff --git a/backend/src/dispatch/service.rs b/backend/src/dispatch/service.rs index 4795a9990..b742bfe70 100644 --- a/backend/src/dispatch/service.rs +++ b/backend/src/dispatch/service.rs @@ -10,70 +10,44 @@ use std::{pin::Pin, sync::Arc, time::Instant}; use bytes::Bytes; -use chrono::Utc; use futures::stream::Stream; -use serde_json::{Value, json}; +use serde_json::Value; +use tokio_util::sync::CancellationToken; use crate::{ dispatch::{ - affinity::resolve::{prepare_variant_identity, record_response_affinity}, affinity::{AccountIdentityService, SessionAffinityService}, - attempts::AccountAttemptLedger, - errors::{ - backend_transport_name, is_continuation_busy_error, is_history_recovery_upstream_error, - is_retryable_account_transport_error, upstream_error_body, - upstream_error_set_cookie_headers, - }, - recovery::{ - account_failure::{isolate_rotatable_account_failure, isolate_sse_account_failure}, + controllers::{ + CompleteExit, ControllerSet, DispatchErrorObservation, cloudflare::CloudflareRecovery, - cyber_policy::CyberPolicyRecovery, - exhaustion::AccountExhaustionTracker, - history::HistoryRecoveryPlan, + history::{ + ConnectionReplayPlan as HistoryConnectionReplayPlan, + ConnectionReplaySnapshot as HistoryConnectionReplaySnapshot, HistoryController, + }, }, - upstream_call::{ - AccountUpstreamContext, QUOTA_VERIFY_LIMIT_REACHED_MESSAGE, QuotaVerificationContext, - QuotaVerificationDecision, create_response_with_account, - verify_acquired_quota_if_required, + errors::backend_transport_name, + lifecycle::{ + RequestContext, RequestEnterDependencies, RequestMode, + attempt::{AttemptMode, AttemptRunnerDependencies}, + contract::{ + AttemptContractError, EstablishedAttemptContext, EstablishedComplete, + EstablishedCompleteBody, EstablishedResponse, FinalOutcome, + }, + enter_request, + pipeline::{AttemptPipelineOutcome, run_attempt_pipeline}, }, + transport::canonical::CanonicalResponseEvent, }, - fleet::{ - account::Account, - pool::{AccountAcquireRequest, AccountPoolService}, - }, + fleet::pool::AccountPoolService, models::service::ModelService, - telemetry::{ - recorder::{ - Recorder, enrich_response_request_semantics, reasoning_effort_from_request, - record_response_event, - }, - usage::types::ResponseUsageRecord, - }, + telemetry::recorder::Recorder, upstream::openai::{ - protocol::{ - events::TokenUsage, - responses::{CodexResponsesRequest, CollectedResponse, response_from_codex_sse}, - }, - transport::{ - CodexBackendClient, CodexBackendResponse, backend_transport_for_response_request, - }, + protocol::{responses::CodexResponsesRequest, sse::SseError}, + transport::CodexBackendClient, }, }; -use super::{ - errors::{ResponseDispatchError, ResponseDispatchStreamError}, - recording::{ - ResponseDispatchErrorDetails, ResponseDispatchErrorEventRecord, - ResponseUpstreamErrorEventRecord, insert_response_status_metadata, - insert_response_trace_metadata, insert_response_upstream_diagnostics, - insert_websocket_pool_decision, record_response_dispatch_error_event, - record_response_upstream_error_event, - }, - stream::{ - sse_failure::is_history_recovery_sse_failure, - trace::{ResponseDispatchAttempt, ResponseDispatchTrace}, - }, -}; +use super::errors::{ResponseDispatchError, ResponseDispatchStreamError}; /// OpenAI Responses 调度服务。 #[derive(Clone)] @@ -85,6 +59,7 @@ pub struct ResponseDispatchService { pub(in crate::dispatch) account_identity: Arc, pub(in crate::dispatch) recorder: Arc, pub(in crate::dispatch) cloudflare: CloudflareRecovery, + pub(in crate::dispatch) shutdown: CancellationToken, } pub(crate) struct ResponseDispatchServiceParts { @@ -95,6 +70,7 @@ pub(crate) struct ResponseDispatchServiceParts { pub account_identity: Arc, pub recorder: Arc, pub cloudflare: CloudflareRecovery, + pub shutdown: CancellationToken, } /// Responses live SSE 响应体流。 @@ -104,6 +80,7 @@ pub type ResponseBodyStream = /// Responses live SSE 调度结果。 pub struct ResponseDispatchStream { pub body: ResponseBodyStream, + pub(crate) canonical_events: tokio::sync::mpsc::UnboundedReceiver>, pub response_headers: Vec<(String, String)>, } @@ -113,6 +90,43 @@ pub struct ResponseDispatchResponse { pub response_headers: Vec<(String, String)>, } +/// 入站 Responses WebSocket 连接持有的窄重放快照。 +pub(crate) struct ConnectionReplaySnapshot { + inner: HistoryConnectionReplaySnapshot, +} + +/// 单次 `response.create` 对快照的待提交重放计划。 +pub(crate) struct ConnectionReplayPlan { + inner: HistoryConnectionReplayPlan, +} + +/// API 从 canonical 响应事件采集的原始 transcript 事实。 +pub(crate) struct ConnectionTranscriptFacts { + response_id: String, + output: Vec, +} + +impl ConnectionTranscriptFacts { + pub(crate) fn new(response_id: String, output: Vec) -> Self { + Self { + response_id, + output, + } + } +} + +fn attempt_contract_error(error: AttemptContractError) -> ResponseDispatchError { + ResponseDispatchError::InvalidSse(SseError::ParseError(format!( + "attempt lifecycle contract violation: {error}" + ))) +} + +fn attempt_mode_mismatch(expected: &str) -> ResponseDispatchError { + ResponseDispatchError::InvalidSse(SseError::ParseError(format!( + "attempt lifecycle returned a response for the wrong mode; expected {expected}" + ))) +} + impl ResponseDispatchService { pub(crate) fn new(parts: ResponseDispatchServiceParts) -> Self { Self { @@ -123,513 +137,175 @@ impl ResponseDispatchService { account_identity: parts.account_identity, recorder: parts.recorder, cloudflare: parts.cloudflare, + shutdown: parts.shutdown, } } - async fn record_response_affinity( + pub(crate) fn connection_replay_snapshot(&self) -> ConnectionReplaySnapshot { + ConnectionReplaySnapshot { + inner: HistoryController::new_connection_replay_snapshot(), + } + } + + pub(crate) fn prepare_connection_replay( &self, - history: &HistoryRecoveryPlan, - original_request: &CodexResponsesRequest, - account_id: &str, - body: &str, - turn_state: Option, - usage: Option, + snapshot: &ConnectionReplaySnapshot, + request: &mut CodexResponsesRequest, + ) -> ConnectionReplayPlan { + ConnectionReplayPlan { + inner: HistoryController::prepare_connection_replay(&snapshot.inner, request), + } + } + + pub(crate) fn commit_connection_replay( + &self, + snapshot: &mut ConnectionReplaySnapshot, + plan: ConnectionReplayPlan, + transcript: ConnectionTranscriptFacts, ) { - record_response_affinity( - &self.session_affinity, - history, - original_request, - account_id, - body, - turn_state, - usage, - ) - .await; + HistoryController::commit_connection_replay( + &mut snapshot.inner, + plan.inner, + transcript.response_id, + transcript.output, + ); } -} -impl ResponseDispatchService { /// 调度非流式 Responses 请求到 Codex Responses 上游。 pub async fn complete( &self, request_id: &str, route: &str, - mut request: CodexResponsesRequest, + request: CodexResponsesRequest, requested_model: &str, ) -> Result { let started_at = Instant::now(); - let catalog = self.models.catalog().await; - let display_model = catalog.resolve_model_id(requested_model); - request.set_model(display_model.clone()); - let compact = request.semantics().compact; - let tuple_schema = request.tuple_schema.clone(); - let image_generation_requested = request.expects_image_generation(); - let now = Utc::now(); - prepare_variant_identity(&mut request); - self.account_identity.prepare_local_identity(&mut request); - let mut history = HistoryRecoveryPlan::load(&self.session_affinity, &request).await; - let cyber_policy = CyberPolicyRecovery::new(Arc::clone(&self.session_affinity)); - let (cyber_policy_plan, preferred_account_id) = tokio::join!( - cyber_policy.prepare(&request), - self.preferred_account_id_for_request(&request, &history, now), + let controllers = ControllerSet::new( + Arc::clone(&self.session_affinity), + Arc::clone(&self.account_pool), + self.cloudflare.clone(), + Arc::clone(&self.recorder), ); - if let Some(failure) = cyber_policy_plan.exhausted_failure() { - let error = ResponseDispatchError::Failed(failure); - self.record_response_dispatch_error( + let context = enter_request( + RequestEnterDependencies { + account_pool: &self.account_pool, + models: &self.models, + account_identity: &self.account_identity, + controllers, + }, + request, + requested_model, + RequestMode::Complete, + ) + .await; + let RequestContext { + request, + display_model, + compact, + tuple_schema, + image_generation_requested, + controllers, + controller_scope, + candidates, + } = context; + let established = match run_attempt_pipeline( + AttemptRunnerDependencies { + account_pool: self.account_pool.as_ref(), + codex: self.codex.as_ref(), + controllers: &controllers, + account_identity: self.account_identity.as_ref(), request_id, route, requested_model, started_at, - ResponseDispatchErrorDetails { - client_api_key_id: request.client_api_key_id.as_deref(), - account_id: cyber_policy_plan.last_account_id(), - stream: false, - compact, - transport: Some(backend_transport_name( - backend_transport_for_response_request(&request), - )), - }, - &error, - ) - .await; - return Err(error); - } - let mut acquire_request = AccountAcquireRequest::new(request.model(), now); - if let Some(preferred_account_id) = preferred_account_id { - acquire_request = acquire_request.with_preferred_account_id(preferred_account_id); - } - let cyber_policy_excluded_accounts = cyber_policy_plan.excluded_account_ids(); - let mut candidates = AccountAttemptLedger::freeze( - &self.account_pool, - &acquire_request, - &cyber_policy_excluded_accounts, - ) - .await; - let mut exhausted_accounts = AccountExhaustionTracker::default(); - let mut next_required_account_id: Option = None; - let mut trace = ResponseDispatchTrace::default(); - let (account, response, collected_response, attempt): ( - Account, - CodexBackendResponse, - CollectedResponse, - ResponseDispatchAttempt, - ) = loop { - let acquired = if let Some(account_id) = next_required_account_id.take() { - match self - .account_pool - .acquire_with( - &AccountAcquireRequest::new(request.model(), Utc::now()) - .with_required_account_id(account_id), - ) - .await - { - Some(acquired) => Some(acquired), - None => candidates.acquire_next(&self.account_pool).await, - } - } else { - candidates.acquire_next(&self.account_pool).await - }; - let Some(acquired) = acquired else { - tracing::info!( - candidate_count = candidates.candidate_count(), - attempted = candidates.attempted_count(), - state_excluded = candidates.state_excluded_count(), - "Responses account candidate ledger exhausted" - ); - let error = exhausted_accounts - .last_exhausted() - .map(ResponseDispatchError::from_exhausted_account) - .or_else(|| { - cyber_policy_plan - .last_failure() - .map(ResponseDispatchError::Failed) - }) - .unwrap_or(ResponseDispatchError::NoActiveAccount); - self.record_response_dispatch_error( - request_id, - route, - requested_model, - started_at, - ResponseDispatchErrorDetails { - client_api_key_id: request.client_api_key_id.as_deref(), - account_id: cyber_policy_plan - .last_account_id() - .or_else(|| exhausted_accounts.last_account_id()), - stream: false, - compact, - transport: Some(backend_transport_name( - backend_transport_for_response_request(&request), - )), - }, - &error, - ) - .await; - return Err(error); - }; - let acquired_account_id = acquired.account.id.clone(); - // 配额验证 - let acquired = match verify_acquired_quota_if_required( - QuotaVerificationContext { - account_pool: self.account_pool.as_ref(), - codex: self.codex.as_ref(), - cloudflare: &self.cloudflare, - account_identity: &self.account_identity, - request_id, - }, - acquired, - ) - .await - { - QuotaVerificationDecision::Ready(acquired) => *acquired, - QuotaVerificationDecision::RetryWithAnotherAccount => { - exhausted_accounts.record_rate_limited( - Some(&acquired_account_id), - QUOTA_VERIFY_LIMIT_REACHED_MESSAGE, - ); - continue; - } - }; - - let Some(attempt_request) = history.prepare_attempt(&request, &acquired.account.id) - else { - acquired.release_without_usage().await; - return Err(ResponseDispatchError::HistoryUnavailable { - upstream_error: "previous response history cannot be sent to another account" - .to_string(), - }); - }; - self.account_pool.wait_for_request_interval(&acquired).await; - let account = acquired.account.clone(); - let release_account_id = account.id.clone(); - let attempt = trace.start_attempt(&release_account_id); - let response_result = create_response_with_account( - AccountUpstreamContext { - codex: &self.codex, - account_identity: &self.account_identity, - cloudflare: &self.cloudflare, - request_id, - account: &account, - }, - &attempt_request, + }, + AttemptMode::Complete { started_at, - ) - .await; - acquired.complete().await; - if let Err(error) = &response_result { - self.cloudflare - .capture_set_cookie_headers( - &release_account_id, - upstream_error_set_cookie_headers(error), + tuple_schema, + }, + request, + controller_scope, + candidates, + ) + .await + .map_err(attempt_contract_error)? + { + AttemptPipelineOutcome::Established(established) => match established { + EstablishedResponse::Complete(established) => *established, + EstablishedResponse::Stream(_) => { + return Err(attempt_mode_mismatch("complete")); + } + }, + AttemptPipelineOutcome::Rejected(rejection) => { + let rejection = *rejection; + controllers + .observe_dispatch_error( + DispatchErrorObservation { + request_id, + client_api_key_id: rejection.request.client_api_key_id.as_deref(), + account_id: rejection.account_id.as_deref(), + route, + model: requested_model, + started_at, + stream: false, + compact, + transport: Some(backend_transport_name(rejection.transport)), + }, + &rejection.error, ) .await; - } - - match response_result { - Ok(response) => { - self.cloudflare - .capture_set_cookie_headers( - &release_account_id, - &response.set_cookie_headers, - ) - .await; - self.account_pool - .sync_passive_rate_limit_headers(&account, &response.rate_limit_headers) - .await; - let collected_response = - match response_from_codex_sse(&response.body, tuple_schema.as_ref()) { - Ok(collected_response) => collected_response, - Err(error) => { - let error = ResponseDispatchError::InvalidSse(error); - self.record_response_dispatch_error( - request_id, - route, - requested_model, - started_at, - ResponseDispatchErrorDetails { - client_api_key_id: request.client_api_key_id.as_deref(), - account_id: Some(&release_account_id), - stream: false, - compact, - transport: Some(backend_transport_name(response.transport)), - }, - &error, - ) - .await; - return Err(error); - } - }; - if matches!(collected_response, CollectedResponse::Empty) { - self.account_pool - .record_empty_response_attempt( - &release_account_id, - image_generation_requested, - ) - .await; - exhausted_accounts.record_upstream_unavailable( - Some(&release_account_id), - "upstream response did not include visible output", - ); - if !history.can_failover() { - break (account, response, collected_response, attempt); - } - continue; - } - if let CollectedResponse::Failed(failure) = &collected_response { - cyber_policy - .observe_sse_failure(&cyber_policy_plan, &release_account_id, failure) - .await; - if is_history_recovery_sse_failure(failure) - && history.recover_managed_history(&release_account_id) - { - next_required_account_id = Some(release_account_id); - continue; - } - if isolate_sse_account_failure( - self.account_pool.as_ref(), - &mut exhausted_accounts, - &release_account_id, - failure, - ) - .await - { - if !history.can_failover() { - break (account, response, collected_response, attempt); - } - continue; - } - } - break (account, response, collected_response, attempt); - } - Err(error) => { - cyber_policy - .observe_upstream_error(&cyber_policy_plan, &release_account_id, &error) - .await; - if isolate_rotatable_account_failure( - self.account_pool.as_ref(), - &self.cloudflare, - &mut exhausted_accounts, - &release_account_id, - &error, - ) - .await - { - if !history.can_failover() { - return Err(ResponseDispatchError::Upstream(error)); - } - continue; - } - if is_continuation_busy_error(&error) { - if history.recover_managed_history(&release_account_id) { - next_required_account_id = Some(release_account_id); - continue; - } - return Err(ResponseDispatchError::ContinuationBusy); - } - let history_unavailable = is_history_recovery_upstream_error(&error); - if history_unavailable && history.recover_managed_history(&release_account_id) { - next_required_account_id = Some(release_account_id); - continue; - } - record_response_upstream_error_event(ResponseUpstreamErrorEventRecord { - recorder: &self.recorder, - request_id, - account_id: &release_account_id, - account_email: account.email.as_deref(), - route, - model: requested_model, - started_at, - stream: false, - transport: backend_transport_for_response_request(&attempt_request), - request: &attempt_request, - error: &error, - trace: &trace, - attempt: Some(&attempt), - }) - .await; - if is_retryable_account_transport_error(&error) { - if !history.can_failover() { - return Err(ResponseDispatchError::Upstream(error)); - } - exhausted_accounts.record_upstream_unavailable( - Some(&release_account_id), - upstream_error_body(&error), - ); - continue; - } - if history_unavailable { - if history.is_external_unknown() { - return Err(ResponseDispatchError::Upstream(error)); - } - return Err(ResponseDispatchError::HistoryUnavailable { - upstream_error: upstream_error_body(&error), - }); - } - return Err(ResponseDispatchError::Upstream(error)); - } + let error = rejection.error; + return Err(error); } }; - - let completed = matches!(&collected_response, CollectedResponse::Completed(_)); - match collected_response { - CollectedResponse::Completed(body) | CollectedResponse::Incomplete(body) => { - cyber_policy.observe_success(&cyber_policy_plan).await; - let response_id = body.get("id").and_then(Value::as_str); - self.cloudflare.reset_account_recovery(&account.id).await; - if let Some(usage) = response.usage { - self.account_pool - .record_response_usage(&account.id, usage, image_generation_requested) - .await; - } - if completed { - self.record_response_affinity( - &history, - &request, - &account.id, - &response.body, - response.turn_state.clone(), - response.usage, - ) - .await; - } - let effective_model = response - .response_metadata - .effective_model - .as_deref() - .unwrap_or(&display_model); - let mut metadata = json!({ - "responseId": response_id, - "stream": false, - "completed": completed, - "incomplete": !completed, - "transport": backend_transport_name(response.transport), - "firstTokenMs": response.first_token_ms, - "usage": response.usage, - "effectiveModel": effective_model, - "modelsEtag": response.response_metadata.models_etag.as_deref(), - "reasoningIncluded": response.response_metadata.reasoning_included, - }); - insert_response_status_metadata( - &mut metadata, - 200, - 200, - response.diagnostics.status_code.map(i64::from), - ); - insert_response_upstream_diagnostics(&mut metadata, &response.diagnostics); - insert_response_trace_metadata(&mut metadata, &trace, Some(&attempt)); - insert_websocket_pool_decision(&mut metadata, response.websocket_pool_decision); - enrich_response_request_semantics(&mut metadata, &request); - record_response_event(ResponseUsageRecord { - recorder: &self.recorder, - request_id, - client_api_key_id: request.client_api_key_id.as_deref(), + let EstablishedComplete { + context, + response, + body, + } = established; + let EstablishedAttemptContext { + request, + controller_scope, + trace, + account, + attempt, + .. + } = context; + let (outcome, body) = match body { + EstablishedCompleteBody::Completed(body) => (FinalOutcome::Completed, body), + EstablishedCompleteBody::Incomplete(body) => (FinalOutcome::Incomplete, body), + }; + let response_id = body.get("id").and_then(Value::as_str); + let completed = outcome == FinalOutcome::Completed; + controllers + .finalize_complete( + controller_scope, + CompleteExit { + request: &request, account_id: &account.id, - route, - model: effective_model, - requested_model: Some(requested_model), - client_ip: request.client_ip.as_deref(), - client_user_agent: request.client_user_agent.as_deref(), - reasoning_effort: reasoning_effort_from_request(&request), - service_tier: request.service_tier(), - started_at, - status_code: 200, - message: if completed { - "v1 responses completed" - } else { - "v1 responses incomplete" - }, - metadata, - rate_limit_headers: &response.rate_limit_headers, - }) - .await; - self.models - .observe_models_etag(response.response_metadata.models_etag.as_deref()); - Ok(ResponseDispatchResponse { - body, - response_headers: response.response_metadata.client_headers, - }) - } - CollectedResponse::Failed(failure) => { - let error = ResponseDispatchError::Failed(failure); - self.record_response_dispatch_error( + body: &response.body, + turn_state: response.turn_state.clone(), + usage: response.usage, + image_generation_requested, + completed, + recorder: &self.recorder, + models: &self.models, request_id, route, requested_model, + display_model: &display_model, started_at, - ResponseDispatchErrorDetails { - client_api_key_id: request.client_api_key_id.as_deref(), - account_id: Some(&account.id), - stream: false, - compact, - transport: Some(backend_transport_name(response.transport)), - }, - &error, - ) - .await; - Err(error) - } - CollectedResponse::MissingCompleted => { - let error = ResponseDispatchError::MissingCompleted; - self.record_response_dispatch_error( - request_id, - route, - requested_model, - started_at, - ResponseDispatchErrorDetails { - client_api_key_id: request.client_api_key_id.as_deref(), - account_id: Some(&account.id), - stream: false, - compact, - transport: Some(backend_transport_name(response.transport)), - }, - &error, - ) - .await; - Err(error) - } - CollectedResponse::Empty => { - let error = ResponseDispatchError::EmptyUpstreamResponse; - self.record_response_dispatch_error( - request_id, - route, - requested_model, - started_at, - ResponseDispatchErrorDetails { - client_api_key_id: request.client_api_key_id.as_deref(), - account_id: Some(&account.id), - stream: false, - compact, - transport: Some(backend_transport_name(response.transport)), - }, - &error, - ) - .await; - Err(error) - } - } - } - - pub(super) async fn record_response_dispatch_error( - &self, - request_id: &str, - route: &str, - requested_model: &str, - started_at: Instant, - details: ResponseDispatchErrorDetails<'_>, - error: &ResponseDispatchError, - ) { - record_response_dispatch_error_event(ResponseDispatchErrorEventRecord { - recorder: &self.recorder, - request_id, - client_api_key_id: details.client_api_key_id, - account_id: details.account_id, - route, - model: requested_model, - started_at, - stream: details.stream, - compact: details.compact, - transport: details.transport, - error, + response_id, + response: &response, + trace: &trace, + attempt: &attempt, + }, + outcome, + ) + .await; + Ok(ResponseDispatchResponse { + body, + response_headers: response.response_metadata.client_headers, }) - .await; } } diff --git a/backend/src/dispatch/stream/lifecycle.rs b/backend/src/dispatch/stream/lifecycle.rs index c619e8616..0da07ca81 100644 --- a/backend/src/dispatch/stream/lifecycle.rs +++ b/backend/src/dispatch/stream/lifecycle.rs @@ -1,47 +1,29 @@ -use std::{sync::Arc, time::Instant}; +//! 流式 Responses 在提交前共享的 request/attempt 生命周期。 -use chrono::Utc; +use std::{sync::Arc, time::Instant}; use crate::{ dispatch::{ - affinity::prepare_variant_identity, - attempts::AccountAttemptLedger, - errors::{ - ResponseDispatchError, backend_transport_name, is_continuation_busy_error, - is_history_recovery_upstream_error, is_retryable_account_transport_error, - upstream_error_body, upstream_error_set_cookie_headers, + controllers::{ + ControllerSet, DispatchErrorObservation, PrefetchedStreamFailureObservation, + StreamControllerContext, }, - recording::{ - ResponseDispatchErrorDetails, ResponseStreamFailureEventRecord, - ResponseUpstreamErrorEventRecord, record_prefetched_response_stream_failure_event, - record_response_upstream_error_event, - }, - recovery::{ - account_failure::{isolate_rotatable_account_failure, isolate_sse_account_failure}, - cyber_policy::CyberPolicyRecovery, - exhaustion::AccountExhaustionTracker, - history::HistoryRecoveryPlan, + errors::{ResponseDispatchError, backend_transport_name}, + lifecycle::{ + RequestContext, RequestEnterDependencies, RequestMode, + attempt::{AttemptMode, AttemptRunnerDependencies}, + contract::{ + AttemptContractError, EstablishedAttemptContext, EstablishedResponse, + EstablishedStream, + }, + enter_request, + finalizer::StreamFinalizer, + pipeline::{AttemptPipelineOutcome, run_attempt_pipeline}, }, service::{ResponseDispatchService, ResponseDispatchStream}, - stream::{ - live::{LiveResponseStreamContext, spawn_live_response_stream}, - prefetch::prefetch_first_sse_chunk, - sse_failure::{ - first_sse_failure, is_history_recovery_sse_failure, sse_failure_error_body, - }, - trace::ResponseDispatchTrace, - }, - upstream_call::{ - AccountUpstreamContext, QUOTA_VERIFY_LIMIT_REACHED_MESSAGE, QuotaVerificationContext, - QuotaVerificationDecision, create_response_stream_with_account, - verify_acquired_quota_if_required, - }, - }, - fleet::pool::AccountAcquireRequest, - upstream::openai::{ - protocol::responses::CodexResponsesRequest, - transport::backend_transport_for_response_request, + stream::live::spawn_live_response_stream, }, + upstream::openai::protocol::{responses::CodexResponsesRequest, sse::SseError}, }; impl ResponseDispatchService { @@ -50,475 +32,175 @@ impl ResponseDispatchService { &self, request_id: &str, route: &str, - mut request: CodexResponsesRequest, + request: CodexResponsesRequest, requested_model: &str, ) -> Result { let started_at = Instant::now(); - let catalog = self.models.catalog().await; - let display_model = catalog.resolve_model_id(requested_model); - request.set_model(display_model.clone()); - request.set_stream(true); - let compact = request.semantics().compact; - let tuple_schema = request.tuple_schema.clone(); - let now = Utc::now(); - prepare_variant_identity(&mut request); - self.account_identity.prepare_local_identity(&mut request); - let mut history = HistoryRecoveryPlan::load(&self.session_affinity, &request).await; - let cyber_policy = CyberPolicyRecovery::new(Arc::clone(&self.session_affinity)); - let (cyber_policy_plan, preferred_account_id) = tokio::join!( - cyber_policy.prepare(&request), - self.preferred_account_id_for_request(&request, &history, now), + let controllers = ControllerSet::new( + Arc::clone(&self.session_affinity), + Arc::clone(&self.account_pool), + self.cloudflare.clone(), + Arc::clone(&self.recorder), ); - let mut acquire_request = AccountAcquireRequest::new(request.model(), now); - if let Some(preferred_account_id) = preferred_account_id { - acquire_request = acquire_request.with_preferred_account_id(preferred_account_id); - } - let cyber_policy_excluded_accounts = cyber_policy_plan.excluded_account_ids(); - let mut candidates = AccountAttemptLedger::freeze( - &self.account_pool, - &acquire_request, - &cyber_policy_excluded_accounts, + let context = enter_request( + RequestEnterDependencies { + account_pool: &self.account_pool, + models: &self.models, + account_identity: &self.account_identity, + controllers, + }, + request, + requested_model, + RequestMode::Stream, ) .await; - let mut exhausted_accounts = AccountExhaustionTracker::default(); - let mut next_required_account_id: Option = None; - let mut trace = ResponseDispatchTrace::default(); - macro_rules! return_stream_dispatch_error { - ($error:expr) => {{ - let error = $error; - self.record_response_dispatch_error( - request_id, - route, - requested_model, - started_at, - ResponseDispatchErrorDetails { - client_api_key_id: request.client_api_key_id.as_deref(), - account_id: cyber_policy_plan - .last_account_id() - .or_else(|| exhausted_accounts.last_account_id()), - stream: true, - compact, - transport: Some(backend_transport_name( - backend_transport_for_response_request(&request), - )), - }, - &error, - ) - .await; - return Err(error); - }}; - ($error:expr, account_id: $account_id:expr, transport: $transport:expr) => {{ - let error = $error; - self.record_response_dispatch_error( - request_id, - route, - requested_model, - started_at, - ResponseDispatchErrorDetails { - client_api_key_id: request.client_api_key_id.as_deref(), - account_id: $account_id, - stream: true, - compact, - transport: $transport, - }, - &error, - ) - .await; - return Err(error); - }}; - } - if let Some(failure) = cyber_policy_plan.exhausted_failure() { - return_stream_dispatch_error!( - ResponseDispatchError::Failed(failure), - account_id: cyber_policy_plan.last_account_id(), - transport: Some(backend_transport_name( - backend_transport_for_response_request(&request) - )) - ); - } - loop { - let acquired = if let Some(account_id) = next_required_account_id.take() { - match self - .account_pool - .acquire_with( - &AccountAcquireRequest::new(request.model(), Utc::now()) - .with_required_account_id(account_id), - ) - .await - { - Some(acquired) => Some(acquired), - None => candidates.acquire_next(&self.account_pool).await, + let RequestContext { + request, + display_model, + compact, + tuple_schema, + controllers, + controller_scope, + candidates, + .. + } = context; + let established = match run_attempt_pipeline( + AttemptRunnerDependencies { + account_pool: self.account_pool.as_ref(), + codex: self.codex.as_ref(), + controllers: &controllers, + account_identity: self.account_identity.as_ref(), + request_id, + route, + requested_model, + started_at, + }, + AttemptMode::Stream { + tuple_schema: tuple_schema.clone(), + }, + request, + controller_scope, + candidates, + ) + .await + .map_err(attempt_contract_error)? + { + AttemptPipelineOutcome::Established(established) => match established { + EstablishedResponse::Stream(established) => *established, + EstablishedResponse::Complete(_) => return Err(attempt_mode_mismatch()), + }, + AttemptPipelineOutcome::Rejected(rejection) => { + let rejection = *rejection; + if let (Some(account), Some(attempt_request), Some(attempt), Some(stream_failure)) = ( + rejection.account.as_ref(), + rejection.attempt_request.as_ref(), + rejection.attempt.as_ref(), + rejection.stream_failure.as_ref(), + ) { + controllers + .observe_prefetched_stream_failure(PrefetchedStreamFailureObservation { + request_id, + account_id: &account.id, + route, + model: &display_model, + requested_model, + started_at, + transport: rejection.transport, + request: attempt_request, + failure: &stream_failure.failure, + error: &rejection.error, + diagnostics: &stream_failure.diagnostics, + rate_limit_headers: &stream_failure.rate_limit_headers, + prefetched: &stream_failure.prefetched, + trace: &rejection.trace, + attempt, + }) + .await; } - } else { - candidates.acquire_next(&self.account_pool).await - }; - let Some(acquired) = acquired else { - tracing::info!( - candidate_count = candidates.candidate_count(), - attempted = candidates.attempted_count(), - state_excluded = candidates.state_excluded_count(), - "Responses stream account candidate ledger exhausted" - ); - let error = exhausted_accounts - .last_exhausted() - .map(ResponseDispatchError::from_exhausted_account) - .or_else(|| { - cyber_policy_plan - .last_failure() - .map(ResponseDispatchError::Failed) - }) - .unwrap_or(ResponseDispatchError::NoActiveAccount); - return_stream_dispatch_error!(error); - }; - let acquired_account_id = acquired.account.id.clone(); - let acquired = match verify_acquired_quota_if_required( - QuotaVerificationContext { - account_pool: self.account_pool.as_ref(), - codex: self.codex.as_ref(), - cloudflare: &self.cloudflare, - account_identity: &self.account_identity, - request_id, - }, - acquired, - ) - .await - { - QuotaVerificationDecision::Ready(acquired) => *acquired, - QuotaVerificationDecision::RetryWithAnotherAccount => { - exhausted_accounts.record_rate_limited( - Some(&acquired_account_id), - QUOTA_VERIFY_LIMIT_REACHED_MESSAGE, - ); - continue; - } - }; - - let Some(attempt_request) = history.prepare_attempt(&request, &acquired.account.id) - else { - acquired.release_without_usage().await; - return_stream_dispatch_error!(ResponseDispatchError::HistoryUnavailable { - upstream_error: "previous response history cannot be sent to another account" - .to_string(), - }); - }; - - self.account_pool.wait_for_request_interval(&acquired).await; - let account = acquired.account.clone(); - let release_account_id = account.id.clone(); - let attempt = trace.start_attempt(&release_account_id); - let response_result = create_response_stream_with_account( - AccountUpstreamContext { - codex: &self.codex, - account_identity: &self.account_identity, - cloudflare: &self.cloudflare, - request_id, - account: &account, - }, - &attempt_request, - ) - .await; - if let Err(error) = &response_result { - self.cloudflare - .capture_set_cookie_headers( - &release_account_id, - upstream_error_set_cookie_headers(error), + controllers + .observe_dispatch_error( + DispatchErrorObservation { + request_id, + client_api_key_id: rejection.request.client_api_key_id.as_deref(), + account_id: rejection.account_id.as_deref(), + route, + model: requested_model, + started_at, + stream: true, + compact, + transport: Some(backend_transport_name(rejection.transport)), + }, + &rejection.error, ) .await; + let error = rejection.error; + return Err(error); } + }; - match response_result { - Ok(response) => { - let transport = response.transport; - let set_cookie_headers = response.set_cookie_headers; - let rate_limit_headers = response.rate_limit_headers; - let rate_limit_header_updates = response.rate_limit_header_updates; - let turn_state_update = response.turn_state_update; - let websocket_pool_decision = response.websocket_pool_decision; - let turn_state = response.turn_state; - let diagnostics = response.diagnostics; - let response_metadata = response.response_metadata; - self.models - .observe_models_etag(response_metadata.models_etag.as_deref()); - let (prefetched, body) = match prefetch_first_sse_chunk(response.body).await { - Ok(prefetched) => prefetched, - Err(ResponseDispatchError::Upstream(error)) => { - acquired.complete().await; - cyber_policy - .observe_upstream_error( - &cyber_policy_plan, - &release_account_id, - &error, - ) - .await; - if isolate_rotatable_account_failure( - self.account_pool.as_ref(), - &self.cloudflare, - &mut exhausted_accounts, - &release_account_id, - &error, - ) - .await - { - if history.is_external_unknown() { - return_stream_dispatch_error!( - ResponseDispatchError::Upstream(error), - account_id: Some(&release_account_id), - transport: Some(backend_transport_name(transport)) - ); - } - continue; - } - if is_continuation_busy_error(&error) { - if history.recover_managed_history(&release_account_id) { - next_required_account_id = Some(release_account_id); - continue; - } - return_stream_dispatch_error!( - ResponseDispatchError::ContinuationBusy, - account_id: Some(&release_account_id), - transport: Some(backend_transport_name(transport)) - ); - } - let history_unavailable = is_history_recovery_upstream_error(&error); - if history_unavailable - && history.recover_managed_history(&release_account_id) - { - next_required_account_id = Some(release_account_id); - continue; - } - record_response_upstream_error_event( - ResponseUpstreamErrorEventRecord { - recorder: &self.recorder, - request_id, - account_id: &release_account_id, - account_email: account.email.as_deref(), - route, - model: requested_model, - started_at, - stream: true, - transport, - request: &attempt_request, - error: &error, - trace: &trace, - attempt: Some(&attempt), - }, - ) - .await; - if is_retryable_account_transport_error(&error) { - exhausted_accounts.record_upstream_unavailable( - Some(&release_account_id), - upstream_error_body(&error), - ); - if history.is_external_unknown() { - return Err(ResponseDispatchError::Upstream(error)); - } - continue; - } - if history_unavailable { - if history.is_external_unknown() { - return Err(ResponseDispatchError::Upstream(error)); - } - return Err(ResponseDispatchError::HistoryUnavailable { - upstream_error: upstream_error_body(&error), - }); - } - return Err(ResponseDispatchError::Upstream(error)); - } - Err(error) => { - acquired.complete().await; - return_stream_dispatch_error!( - error, - account_id: Some(&release_account_id), - transport: Some(backend_transport_name(transport)) - ); - } - }; - let first_failure = match first_sse_failure(&prefetched) { - Ok(failure) => failure, - Err(error) => { - acquired.complete().await; - return_stream_dispatch_error!( - ResponseDispatchError::InvalidSse(error), - account_id: Some(&release_account_id), - transport: Some(backend_transport_name(transport)) - ); - } - }; - if let Some(failure) = first_failure { - cyber_policy - .observe_sse_failure(&cyber_policy_plan, &release_account_id, &failure) - .await; - if is_history_recovery_sse_failure(&failure) - && history.recover_managed_history(&release_account_id) - { - acquired.complete().await; - next_required_account_id = Some(release_account_id); - continue; - } - if isolate_sse_account_failure( - self.account_pool.as_ref(), - &mut exhausted_accounts, - &release_account_id, - &failure, - ) - .await - { - acquired.complete().await; - if history.is_external_unknown() { - return_stream_dispatch_error!( - ResponseDispatchError::Failed(failure), - account_id: Some(&release_account_id), - transport: Some(backend_transport_name(transport)) - ); - } - continue; - } - acquired.complete().await; - record_prefetched_response_stream_failure_event( - ResponseStreamFailureEventRecord { - recorder: &self.recorder, - request_id, - account_id: &release_account_id, - route, - model: &display_model, - requested_model, - started_at, - transport, - request: &attempt_request, - failure: &failure, - diagnostics: &diagnostics, - rate_limit_headers: &rate_limit_headers, - prefetched: &prefetched, - trace: &trace, - attempt: &attempt, - }, - ) - .await; - let error = if is_history_recovery_sse_failure(&failure) - && !history.is_external_unknown() - { - ResponseDispatchError::HistoryUnavailable { - upstream_error: sse_failure_error_body(&failure), - } - } else { - ResponseDispatchError::Failed(failure.clone()) - }; - return Err(error); - } - - let context = LiveResponseStreamContext { - account_pool: Arc::clone(&self.account_pool), - account_lease: acquired, - session_affinity: Arc::clone(&self.session_affinity), - cyber_policy, - cyber_policy_plan, - history, - recorder: Arc::clone(&self.recorder), - cloudflare: self.cloudflare.clone(), - account_id: account.id, - account_plan_type: account.plan_type, - request_id: request_id.to_string(), - route: route.to_string(), - display_model: display_model.clone(), - requested_model: requested_model.to_string(), - client_ip: request.client_ip.clone(), - request, - tuple_schema, - transport, - set_cookie_headers, - rate_limit_headers, - rate_limit_header_updates, - turn_state_update, - websocket_pool_decision, - turn_state, - diagnostics, - response_metadata, - attempt: attempt.clone(), - attempts: trace.attempts().to_vec(), - started_at, - }; - return Ok(spawn_live_response_stream(context, prefetched, body)); - } - Err(error) => { - acquired.complete().await; - cyber_policy - .observe_upstream_error(&cyber_policy_plan, &release_account_id, &error) - .await; - if isolate_rotatable_account_failure( - self.account_pool.as_ref(), - &self.cloudflare, - &mut exhausted_accounts, - &release_account_id, - &error, - ) - .await - { - if history.is_external_unknown() { - return_stream_dispatch_error!( - ResponseDispatchError::Upstream(error), - account_id: Some(&release_account_id), - transport: Some(backend_transport_name( - backend_transport_for_response_request(&attempt_request) - )) - ); - } - continue; - } - if is_continuation_busy_error(&error) { - if history.recover_managed_history(&release_account_id) { - next_required_account_id = Some(release_account_id); - continue; - } - return_stream_dispatch_error!( - ResponseDispatchError::ContinuationBusy, - account_id: Some(&release_account_id), - transport: Some(backend_transport_name( - backend_transport_for_response_request(&attempt_request) - )) - ); - } - let history_unavailable = is_history_recovery_upstream_error(&error); - if history_unavailable && history.recover_managed_history(&release_account_id) { - next_required_account_id = Some(release_account_id); - continue; - } - record_response_upstream_error_event(ResponseUpstreamErrorEventRecord { - recorder: &self.recorder, - request_id, - account_id: &release_account_id, - account_email: account.email.as_deref(), - route, - model: requested_model, - started_at, - stream: true, - transport: backend_transport_for_response_request(&attempt_request), - request: &attempt_request, - error: &error, - trace: &trace, - attempt: Some(&attempt), - }) - .await; - if is_retryable_account_transport_error(&error) { - exhausted_accounts.record_upstream_unavailable( - Some(&release_account_id), - upstream_error_body(&error), - ); - if history.is_external_unknown() { - return Err(ResponseDispatchError::Upstream(error)); - } - continue; - } - if history_unavailable { - if history.is_external_unknown() { - return Err(ResponseDispatchError::Upstream(error)); - } - return Err(ResponseDispatchError::HistoryUnavailable { - upstream_error: upstream_error_body(&error), - }); - } - return Err(ResponseDispatchError::Upstream(error)); - } - } - } + let EstablishedStream { + context, + lease, + response, + decoder, + initial_batch, + } = established; + let EstablishedAttemptContext { + request, + controller_scope, + trace, + account, + attempt, + .. + } = context; + self.models + .observe_models_etag(response.response_metadata.models_etag.as_deref()); + let attempts = trace.attempts().to_vec(); + let transport = response.transport; + let finalizer = StreamFinalizer::new( + controllers, + controller_scope, + StreamControllerContext { + account_id: account.id, + account_plan_type: account.plan_type, + request_id: request_id.to_string(), + route: route.to_string(), + display_model, + requested_model: requested_model.to_string(), + request, + transport, + set_cookie_headers: response.set_cookie_headers, + rate_limit_headers: response.rate_limit_headers, + rate_limit_header_updates: response.rate_limit_header_updates, + turn_state_update: response.turn_state_update, + websocket_pool_decision: response.websocket_pool_decision, + turn_state: response.turn_state, + diagnostics: response.diagnostics, + response_metadata: response.response_metadata, + attempt, + attempts, + started_at, + }, + lease, + ); + Ok(spawn_live_response_stream( + finalizer, + decoder, + initial_batch, + response.body, + self.shutdown.clone(), + )) } } + +fn attempt_contract_error(error: AttemptContractError) -> ResponseDispatchError { + ResponseDispatchError::InvalidSse(SseError::ParseError(format!( + "attempt lifecycle contract violation: {error}" + ))) +} + +fn attempt_mode_mismatch() -> ResponseDispatchError { + ResponseDispatchError::InvalidSse(SseError::ParseError( + "attempt lifecycle returned a non-stream response for stream mode".to_string(), + )) +} diff --git a/backend/src/dispatch/stream/live.rs b/backend/src/dispatch/stream/live.rs index 3db19f84a..cf7fad60a 100644 --- a/backend/src/dispatch/stream/live.rs +++ b/backend/src/dispatch/stream/live.rs @@ -1,67 +1,32 @@ use std::{ pin::Pin, - sync::Arc, task::{Context, Poll}, time::Instant, }; use bytes::Bytes; use futures::{StreamExt, stream::Stream}; -use serde_json::Value; use tokio::sync::{mpsc, oneshot}; +use tokio_util::sync::CancellationToken; use crate::{ dispatch::{ - affinity::{SessionAffinityService, resolve::record_response_affinity}, - errors::{ResponseDispatchStreamError, backend_transport_name}, - recording::{ - insert_first_token_ms, live_response_rate_limit_headers, live_response_turn_state, - record_live_response_stream_event, - }, - recovery::{ - cloudflare::CloudflareRecovery, - cyber_policy::{ - CyberPolicyRecovery, CyberPolicyRoutingPlan, CyberPolicyStreamObserver, - }, - history::HistoryRecoveryPlan, + errors::ResponseDispatchStreamError, + lifecycle::{ + finalizer::StreamFinalizer, + stream::{StreamSummary, StreamTerminal}, }, service::ResponseDispatchStream, + transport::canonical::{ + CanonicalResponseChunk, CanonicalResponseEvent, CanonicalStreamBatch, + CanonicalStreamDecoder, + }, }, - fleet::pool::{AccountLease, AccountPoolService}, infra::time::elapsed_millis_i64, - telemetry::{recorder::Recorder, usage::types::UsageRecordLevel}, - upstream::openai::{ - protocol::{ - events::extract_sse_usage, - responses::{ - CodexResponsesRequest, CollectedResponse, - reconvert_responses_sse_event_tuple_values, response_body_has_first_output, - response_from_codex_sse, response_sse_event_is_terminal, - }, - sse::{ - DONE_SSE_FRAME, encode_sse_event, parse_sse_events, - response_failed_sse_event_with_id, sse_body_has_done, sse_frame_end, - }, - }, - transport::{ - CodexBackendSseStream, CodexBackendTransport, CodexRateLimitHeaderUpdates, - CodexResponseMetadata, CodexTurnStateUpdate, CodexUpstreamDiagnostics, - WebSocketPoolDecision, - }, - }, + upstream::openai::transport::CodexBackendSseStream, }; const MAX_LIVE_RESPONSE_CAPTURE_BYTES: usize = 16 * 1024 * 1024; -const LIVE_RESPONSE_CAPTURE_LIMIT_DETAIL: &str = - "upstream response exceeded the 16 MiB proxy capture limit"; - -use super::{ - sse_failure::{ - STREAM_DISCONNECTED_CODE, STREAM_DISCONNECTED_MESSAGE, status_code_for_stream_failure, - stream_failure_metadata, stream_failure_source, synthetic_stream_disconnected_detail, - }, - trace::ResponseDispatchAttempt, -}; struct MpscResponseBodyStream { receiver: mpsc::Receiver>, @@ -85,202 +50,32 @@ impl Stream for MpscResponseBodyStream { } pub(in crate::dispatch) fn spawn_live_response_stream( - context: LiveResponseStreamContext, - prefetched: Bytes, - mut body: CodexBackendSseStream, + finalizer: StreamFinalizer, + decoder: CanonicalStreamDecoder, + initial_batch: CanonicalStreamBatch, + body: CodexBackendSseStream, + shutdown: CancellationToken, ) -> ResponseDispatchStream { - let response_headers = context.response_metadata.client_headers.clone(); + let response_headers = finalizer.response_headers().to_vec(); + let started_at = finalizer.started_at(); let (sender, receiver) = mpsc::channel(8); - let (cancel_sender, mut cancel_receiver) = oneshot::channel(); + let (event_sender, canonical_events) = mpsc::unbounded_channel(); + let (cancel_sender, cancel_receiver) = oneshot::channel(); tokio::spawn(async move { - let mut tuple_transformer = context - .tuple_schema - .clone() - .map(TupleSseEventTransformer::new); - let mut body_bytes = Vec::new(); - let mut first_token_ms = None; - let mut cyber_policy_observer = CyberPolicyStreamObserver::default(); - cyber_policy_observer - .observe_chunk( - &context.cyber_policy, - &context.cyber_policy_plan, - &context.account_id, - &prefetched, - ) - .await; - let mut first_output_detector = FirstOutputDetector::default(); - first_output_detector.observe(context.started_at, &prefetched, &mut first_token_ms); - match send_live_response_stream_chunk( + let summary = run_live_response_stream( &sender, - &mut body_bytes, - tuple_transformer.as_mut(), - prefetched, - ) - .await - { - Ok(()) => {} - Err(LiveStreamWriteError::DownstreamClosed) => { - context.account_lease.complete().await; - return; - } - Err(LiveStreamWriteError::CaptureLimitExceeded) => { - terminate_oversized_live_stream( - context, - &sender, - &mut body_bytes, - first_token_ms, - cyber_policy_observer.state_transition_applied(), - ) - .await; - return; - } - } - - loop { - let next = tokio::select! { - _ = &mut cancel_receiver => { - context.account_lease.complete().await; - return; - } - next = body.next() => next, - }; - let Some(next) = next else { - break; - }; - match next { - Ok(chunk) => { - cyber_policy_observer - .observe_chunk( - &context.cyber_policy, - &context.cyber_policy_plan, - &context.account_id, - &chunk, - ) - .await; - first_output_detector.observe(context.started_at, &chunk, &mut first_token_ms); - match send_live_response_stream_chunk( - &sender, - &mut body_bytes, - tuple_transformer.as_mut(), - chunk, - ) - .await - { - Ok(()) => {} - Err(LiveStreamWriteError::DownstreamClosed) => { - context.account_lease.complete().await; - return; - } - Err(LiveStreamWriteError::CaptureLimitExceeded) => { - terminate_oversized_live_stream( - context, - &sender, - &mut body_bytes, - first_token_ms, - cyber_policy_observer.state_transition_applied(), - ) - .await; - return; - } - } - } - Err(error) => { - cyber_policy_observer - .finish( - &context.cyber_policy, - &context.cyber_policy_plan, - &context.account_id, - ) - .await; - match flush_live_response_stream_transformer( - &sender, - &mut body_bytes, - tuple_transformer.as_mut(), - ) - .await - { - Ok(()) => {} - Err(LiveStreamWriteError::DownstreamClosed) => { - context.account_lease.complete().await; - return; - } - Err(LiveStreamWriteError::CaptureLimitExceeded) => { - terminate_oversized_live_stream( - context, - &sender, - &mut body_bytes, - first_token_ms, - cyber_policy_observer.state_transition_applied(), - ) - .await; - return; - } - } - let detail = error.to_string(); - let Some(body_text) = - send_live_response_stream_tail(&sender, &mut body_bytes, Some(&detail)) - .await - else { - context.account_lease.complete().await; - return; - }; - finalize_live_response_stream( - context, - body_text, - first_token_ms, - cyber_policy_observer.state_transition_applied(), - ) - .await; - return; - } - } - } - - cyber_policy_observer - .finish( - &context.cyber_policy, - &context.cyber_policy_plan, - &context.account_id, - ) - .await; - - match flush_live_response_stream_transformer( - &sender, - &mut body_bytes, - tuple_transformer.as_mut(), - ) - .await - { - Ok(()) => {} - Err(LiveStreamWriteError::DownstreamClosed) => { - context.account_lease.complete().await; - return; - } - Err(LiveStreamWriteError::CaptureLimitExceeded) => { - terminate_oversized_live_stream( - context, - &sender, - &mut body_bytes, - first_token_ms, - cyber_policy_observer.state_transition_applied(), - ) - .await; - return; - } - } - let Some(body_text) = send_live_response_stream_tail(&sender, &mut body_bytes, None).await - else { - context.account_lease.complete().await; - return; - }; - - finalize_live_response_stream( - context, - body_text, - first_token_ms, - cyber_policy_observer.state_transition_applied(), + &event_sender, + LiveRunInputs { + initial_batch, + body, + cancel: cancel_receiver, + decoder, + started_at, + shutdown, + }, ) .await; + finalizer.finalize(&sender, &event_sender, summary).await; }); ResponseDispatchStream { @@ -288,467 +83,239 @@ pub(in crate::dispatch) fn spawn_live_response_stream( receiver, cancel: Some(cancel_sender), }), + canonical_events, response_headers, } } -async fn send_live_response_stream_chunk( - sender: &mpsc::Sender>, - body_bytes: &mut Vec, - transformer: Option<&mut TupleSseEventTransformer>, - chunk: Bytes, -) -> Result<(), LiveStreamWriteError> { - let chunks = match transformer { - Some(transformer) => transformer.push(&chunk)?, - None => vec![chunk], - }; - send_live_response_stream_chunks(sender, body_bytes, chunks).await +struct LiveRunInputs { + initial_batch: CanonicalStreamBatch, + body: CodexBackendSseStream, + cancel: oneshot::Receiver<()>, + decoder: CanonicalStreamDecoder, + started_at: Instant, + shutdown: CancellationToken, } -async fn flush_live_response_stream_transformer( +async fn run_live_response_stream( sender: &mpsc::Sender>, - body_bytes: &mut Vec, - transformer: Option<&mut TupleSseEventTransformer>, -) -> Result<(), LiveStreamWriteError> { - let Some(transformer) = transformer else { - return Ok(()); - }; - send_live_response_stream_chunks(sender, body_bytes, transformer.finish()?).await -} + event_sender: &mpsc::UnboundedSender>, + inputs: LiveRunInputs, +) -> StreamSummary { + let LiveRunInputs { + initial_batch, + mut body, + mut cancel, + mut decoder, + started_at, + shutdown, + } = inputs; + let mut body_bytes = Vec::new(); + let mut terminal_chunks = Vec::new(); + let mut first_token_ms = None; -async fn send_live_response_stream_chunks( - sender: &mpsc::Sender>, - body_bytes: &mut Vec, - chunks: Vec, -) -> Result<(), LiveStreamWriteError> { - for chunk in chunks { - if body_bytes.len().saturating_add(chunk.len()) > MAX_LIVE_RESPONSE_CAPTURE_BYTES { - return Err(LiveStreamWriteError::CaptureLimitExceeded); + let decoded_terminal = take_stream_terminal(&mut decoder); + let terminal = if decoded_terminal.is_some() { + stage_terminal_batch(&mut body_bytes, &mut terminal_chunks, initial_batch) + .err() + .or(decoded_terminal) + } else { + forward_batch(sender, event_sender, &mut body_bytes, initial_batch) + .await + .err() + }; + update_first_token_ms(&decoder, started_at, &mut first_token_ms); + if let Some(terminal) = terminal { + return stream_summary( + terminal, + body_bytes, + terminal_chunks, + first_token_ms, + &decoder, + ); + } + + let mut shutdown = Box::pin(shutdown.cancelled_owned()); + let terminal = loop { + let next = tokio::select! { + _ = &mut cancel => break StreamTerminal::Cancelled, + () = &mut shutdown => break StreamTerminal::Shutdown, + next = body.next() => next, + }; + match next { + Some(Ok(chunk)) => match decoder.push(chunk) { + Ok(batch) => { + update_first_token_ms(&decoder, started_at, &mut first_token_ms); + if let Some(terminal) = take_stream_terminal(&mut decoder) { + match stage_terminal_batch(&mut body_bytes, &mut terminal_chunks, batch) { + Ok(()) => break terminal, + Err(terminal) => break terminal, + } + } + if let Err(terminal) = + forward_batch(sender, event_sender, &mut body_bytes, batch).await + { + break terminal; + } + } + Err(error) => { + break StreamTerminal::ProtocolError { + detail: error.to_string(), + }; + } + }, + Some(Err(error)) => { + break finish_decoder( + sender, + event_sender, + &mut body_bytes, + &mut terminal_chunks, + &mut decoder, + StreamTerminal::UpstreamError { + detail: error.to_string(), + }, + ) + .await; + } + None => { + break finish_decoder( + sender, + event_sender, + &mut body_bytes, + &mut terminal_chunks, + &mut decoder, + StreamTerminal::UpstreamClosed, + ) + .await; + } } - body_bytes.extend_from_slice(&chunk); - if sender.send(Ok(chunk)).await.is_err() { - return Err(LiveStreamWriteError::DownstreamClosed); + }; + + update_first_token_ms(&decoder, started_at, &mut first_token_ms); + stream_summary( + terminal, + body_bytes, + terminal_chunks, + first_token_ms, + &decoder, + ) +} + +async fn finish_decoder( + sender: &mpsc::Sender>, + event_sender: &mpsc::UnboundedSender>, + body_bytes: &mut Vec, + terminal_chunks: &mut Vec, + decoder: &mut CanonicalStreamDecoder, + fallback: StreamTerminal, +) -> StreamTerminal { + let batch = match decoder.finish() { + Ok(batch) => batch, + Err(error) => { + return StreamTerminal::ProtocolError { + detail: error.to_string(), + }; + } + }; + if let Some(terminal) = take_stream_terminal(decoder) { + return match stage_terminal_batch(body_bytes, terminal_chunks, batch) { + Ok(()) => terminal, + Err(terminal) => terminal, + }; + } + if let Err(terminal) = forward_batch(sender, event_sender, body_bytes, batch).await { + return terminal; + } + fallback +} + +fn take_stream_terminal(decoder: &mut CanonicalStreamDecoder) -> Option { + let terminal = decoder.take_terminal().map(StreamTerminal::from); + let transport_done = decoder.take_transport_done(); + terminal.or_else(|| transport_done.then_some(StreamTerminal::UpstreamClosed)) +} + +fn stage_terminal_batch( + body_bytes: &mut Vec, + terminal_chunks: &mut Vec, + batch: CanonicalStreamBatch, +) -> Result<(), StreamTerminal> { + let batch_bytes = batch.chunks.iter().fold(0usize, |total, chunk| { + total.saturating_add(chunk.bytes().len()) + }); + ensure_capture_capacity(body_bytes.len(), batch_bytes)?; + for chunk in batch.chunks { + body_bytes.extend_from_slice(chunk.bytes()); + terminal_chunks.push(chunk); + } + Ok(()) +} + +async fn forward_batch( + sender: &mpsc::Sender>, + event_sender: &mpsc::UnboundedSender>, + body_bytes: &mut Vec, + batch: CanonicalStreamBatch, +) -> Result<(), StreamTerminal> { + for chunk in batch.chunks { + capture_chunk(body_bytes, &chunk)?; + let (bytes, events) = chunk.into_parts(); + let _ = event_sender.send(events); + if sender.send(Ok(bytes)).await.is_err() { + return Err(StreamTerminal::DownstreamClosed); } } Ok(()) } -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -enum LiveStreamWriteError { - DownstreamClosed, - CaptureLimitExceeded, +fn capture_chunk( + body_bytes: &mut Vec, + chunk: &CanonicalResponseChunk, +) -> Result<(), StreamTerminal> { + ensure_capture_capacity(body_bytes.len(), chunk.bytes().len())?; + body_bytes.extend_from_slice(chunk.bytes()); + Ok(()) } -#[derive(Default)] -struct FirstOutputDetector { - pending: Vec, -} - -impl FirstOutputDetector { - fn observe(&mut self, started_at: Instant, chunk: &[u8], first_token_ms: &mut Option) { - if first_token_ms.is_some() { - return; - } - self.pending.extend_from_slice(chunk); - while let Some(frame_end) = sse_frame_end(&self.pending) { - let frame = self.pending.drain(..frame_end).collect::>(); - if response_body_has_first_output(&frame) { - *first_token_ms = Some(elapsed_millis_i64(started_at).max(1)); - self.pending.clear(); - return; - } - } - if self.pending.len() > MAX_LIVE_RESPONSE_CAPTURE_BYTES { - self.pending.clear(); - } - } -} - -struct TupleSseEventTransformer { - tuple_schema: Value, - pending: Vec, -} - -impl TupleSseEventTransformer { - fn new(tuple_schema: Value) -> Self { - Self { - tuple_schema, - pending: Vec::new(), - } - } - - fn push(&mut self, chunk: &[u8]) -> Result, LiveStreamWriteError> { - self.pending.extend_from_slice(chunk); - if self.pending.len() > MAX_LIVE_RESPONSE_CAPTURE_BYTES { - return Err(LiveStreamWriteError::CaptureLimitExceeded); - } - let mut chunks = Vec::new(); - while let Some(frame_end) = sse_frame_end(&self.pending) { - let frame = self.pending.drain(..frame_end).collect::>(); - chunks.push(self.transform_frame(&frame)); - } - Ok(chunks) - } - - fn finish(&mut self) -> Result, LiveStreamWriteError> { - if self.pending.is_empty() { - return Ok(Vec::new()); - } - let frame = std::mem::take(&mut self.pending); - Ok(vec![self.transform_frame(&frame)]) - } - - fn transform_frame(&self, frame: &[u8]) -> Bytes { - let frame_text = String::from_utf8_lossy(frame); - let Ok(events) = parse_sse_events(&frame_text) else { - return Bytes::copy_from_slice(frame); - }; - let [event] = events.as_slice() else { - return Bytes::copy_from_slice(frame); - }; - let Ok(data) = serde_json::from_str::(&event.data) else { - return Bytes::copy_from_slice(frame); - }; - let transformed = reconvert_responses_sse_event_tuple_values( - event.event.as_deref(), - data, - &self.tuple_schema, +fn ensure_capture_capacity( + captured_bytes: usize, + incoming_bytes: usize, +) -> Result<(), StreamTerminal> { + if captured_bytes.saturating_add(incoming_bytes) > MAX_LIVE_RESPONSE_CAPTURE_BYTES { + tracing::warn!( + captured_bytes, + incoming_bytes, + capture_limit_bytes = MAX_LIVE_RESPONSE_CAPTURE_BYTES, + "Live response capture limit exceeded" ); - Bytes::from(encode_sse_event( - event.event.as_deref().unwrap_or_default(), - &transformed.to_string(), - )) + return Err(StreamTerminal::CaptureLimitExceeded); + } + Ok(()) +} + +fn update_first_token_ms( + decoder: &CanonicalStreamDecoder, + started_at: Instant, + first_token_ms: &mut Option, +) { + if first_token_ms.is_none() && decoder.first_output_seen() { + *first_token_ms = Some(elapsed_millis_i64(started_at).max(1)); } } -async fn terminate_oversized_live_stream( - context: LiveResponseStreamContext, - sender: &mpsc::Sender>, - body_bytes: &mut Vec, +fn stream_summary( + terminal: StreamTerminal, + body: Vec, + terminal_chunks: Vec, first_token_ms: Option, - cyber_policy_state_transition_applied: bool, -) { - tracing::warn!( - request_id = %context.request_id, - account_id = %context.account_id, - captured_bytes = body_bytes.len(), - capture_limit_bytes = MAX_LIVE_RESPONSE_CAPTURE_BYTES, - "Live response capture limit exceeded" - ); - let Some(body_text) = send_live_response_stream_tail( - sender, - body_bytes, - Some(LIVE_RESPONSE_CAPTURE_LIMIT_DETAIL), - ) - .await - else { - context.account_lease.complete().await; - return; - }; - finalize_live_response_stream( - context, - body_text, + decoder: &CanonicalStreamDecoder, +) -> StreamSummary { + StreamSummary { + terminal, + body, + terminal_chunks, first_token_ms, - cyber_policy_state_transition_applied, - ) - .await; -} - -async fn send_live_response_stream_tail( - sender: &mpsc::Sender>, - body_bytes: &mut Vec, - failure_detail: Option<&str>, -) -> Option { - let mut body_text = String::from_utf8_lossy(body_bytes).to_string(); - if !sse_body_has_terminal_event(&body_text) { - if let Some(separator) = missing_sse_event_separator(&body_text) { - body_text.push_str(separator); - body_bytes.extend_from_slice(separator.as_bytes()); - if sender - .send(Ok(Bytes::copy_from_slice(separator.as_bytes()))) - .await - .is_err() - { - return None; - } - } - let failure = - premature_close_failed_event(latest_response_id(&body_text).as_deref(), failure_detail); - body_text.push_str(&failure); - body_bytes.extend_from_slice(failure.as_bytes()); - if sender.send(Ok(Bytes::from(failure))).await.is_err() { - return None; - } - } - - if !sse_body_has_done(&body_text) { - body_text.push_str(DONE_SSE_FRAME); - body_bytes.extend_from_slice(DONE_SSE_FRAME.as_bytes()); - if sender - .send(Ok(Bytes::from_static(DONE_SSE_FRAME.as_bytes()))) - .await - .is_err() - { - return None; - } - } - - Some(body_text) -} - -fn sse_body_has_terminal_event(body: &str) -> bool { - parse_sse_events(body).is_ok_and(|events| events.iter().any(response_sse_event_is_terminal)) -} - -fn missing_sse_event_separator(body: &str) -> Option<&'static str> { - if body.is_empty() - || body.ends_with("\n\n") - || body.ends_with("\r\n\r\n") - || body.ends_with("\r\r") - { - None - } else if body.ends_with('\n') || body.ends_with('\r') { - Some("\n") - } else { - Some("\n\n") + usage: decoder.usage(), + last_response_id: decoder.last_response_id().map(ToString::to_string), } } - -pub(in crate::dispatch) struct LiveResponseStreamContext { - pub(in crate::dispatch) account_pool: Arc, - pub(in crate::dispatch) account_lease: AccountLease, - pub(in crate::dispatch) session_affinity: Arc, - pub(in crate::dispatch) cyber_policy: CyberPolicyRecovery, - pub(in crate::dispatch) cyber_policy_plan: CyberPolicyRoutingPlan, - pub(in crate::dispatch) history: HistoryRecoveryPlan, - pub(in crate::dispatch) recorder: Arc, - pub(in crate::dispatch) cloudflare: CloudflareRecovery, - pub(in crate::dispatch) account_id: String, - pub(in crate::dispatch) account_plan_type: Option, - pub(in crate::dispatch) request_id: String, - pub(in crate::dispatch) route: String, - pub(in crate::dispatch) display_model: String, - pub(in crate::dispatch) requested_model: String, - pub(in crate::dispatch) client_ip: Option, - pub(in crate::dispatch) request: CodexResponsesRequest, - pub(in crate::dispatch) tuple_schema: Option, - pub(in crate::dispatch) transport: CodexBackendTransport, - pub(in crate::dispatch) set_cookie_headers: Vec, - pub(in crate::dispatch) rate_limit_headers: Vec<(String, String)>, - pub(in crate::dispatch) rate_limit_header_updates: Option, - pub(in crate::dispatch) turn_state_update: Option, - pub(in crate::dispatch) websocket_pool_decision: Option, - pub(in crate::dispatch) turn_state: Option, - pub(in crate::dispatch) diagnostics: CodexUpstreamDiagnostics, - pub(in crate::dispatch) response_metadata: CodexResponseMetadata, - pub(in crate::dispatch) attempt: ResponseDispatchAttempt, - pub(in crate::dispatch) attempts: Vec, - pub(in crate::dispatch) started_at: Instant, -} - -pub(in crate::dispatch) fn latest_response_id(body: &str) -> Option { - parse_sse_events(body).ok().and_then(|events| { - events.iter().rev().find_map(|event| { - serde_json::from_str::(&event.data) - .ok() - .and_then(|data| { - data.pointer("/response/id") - .and_then(Value::as_str) - .filter(|id| !id.trim().is_empty()) - .map(ToString::to_string) - }) - }) - }) -} - -pub(in crate::dispatch) fn premature_close_failed_event( - response_id: Option<&str>, - detail: Option<&str>, -) -> String { - let message = match detail.filter(|value| !value.trim().is_empty()) { - Some(detail) => format!("{STREAM_DISCONNECTED_MESSAGE}: {detail}"), - None => STREAM_DISCONNECTED_MESSAGE.to_string(), - }; - response_failed_sse_event_with_id( - response_id, - "server_error", - STREAM_DISCONNECTED_CODE, - &message, - ) -} - -pub(in crate::dispatch) async fn finalize_live_response_stream( - context: LiveResponseStreamContext, - body: String, - first_token_ms: Option, - cyber_policy_state_transition_applied: bool, -) { - context - .cloudflare - .capture_set_cookie_headers(&context.account_id, &context.set_cookie_headers) - .await; - let rate_limit_headers = live_response_rate_limit_headers(&context).await; - context - .account_pool - .sync_passive_rate_limit_headers_for_account( - &context.account_id, - context.account_plan_type.as_deref(), - &rate_limit_headers, - ) - .await; - let turn_state = live_response_turn_state(&context).await; - let usage = match extract_sse_usage(&body) { - Ok(Some(usage)) => { - context - .account_pool - .record_token_usage(&context.account_id, &usage) - .await; - Some(usage) - } - Ok(None) => None, - Err(error) => { - tracing::warn!(account_id = %context.account_id, error = %error, "Failed to extract streaming token usage"); - None - } - }; - - let collected_response = response_from_codex_sse(&body, context.tuple_schema.as_ref()); - let completed_terminal = matches!(&collected_response, Ok(CollectedResponse::Completed(_))); - match collected_response { - Ok(CollectedResponse::Completed(terminal) | CollectedResponse::Incomplete(terminal)) => { - if !cyber_policy_state_transition_applied { - context - .cyber_policy - .observe_success(&context.cyber_policy_plan) - .await; - } - context - .cloudflare - .reset_account_recovery(&context.account_id) - .await; - let response_id = terminal.get("id").and_then(Value::as_str); - if completed_terminal { - record_response_affinity( - &context.session_affinity, - &context.history, - &context.request, - &context.account_id, - &body, - turn_state, - usage, - ) - .await; - } - record_live_response_stream_event( - &context, - 200, - UsageRecordLevel::Info, - if completed_terminal { - "v1 responses stream completed" - } else { - "v1 responses stream incomplete" - }, - serde_json::json!({ - "stream": true, - "completed": completed_terminal, - "incomplete": !completed_terminal, - "responseId": response_id, - "firstTokenMs": first_token_ms, - "usage": usage, - }), - &rate_limit_headers, - &body, - ) - .await; - } - Ok(CollectedResponse::Failed(failure)) => { - if !cyber_policy_state_transition_applied { - context - .cyber_policy - .observe_sse_failure(&context.cyber_policy_plan, &context.account_id, &failure) - .await; - } - let response_id = latest_response_id(&body); - let latency_ms = elapsed_millis_i64(context.started_at); - let failure_source = stream_failure_source(&failure); - let failure_detail = synthetic_stream_disconnected_detail(&failure); - let websocket_pool_kind = context - .websocket_pool_decision - .map(|decision| decision.kind()); - let websocket_pool_reason = context - .websocket_pool_decision - .and_then(|decision| decision.reason()); - tracing::warn!( - account_id = %context.account_id, - request_id = %context.request_id, - response_id = response_id.as_deref().unwrap_or(""), - transport = %backend_transport_name(context.transport), - websocket_pool_kind = ?websocket_pool_kind, - websocket_pool_reason = ?websocket_pool_reason, - first_token_ms = ?first_token_ms, - latency_ms, - event = %failure.event, - code = ?failure.upstream_code.as_deref(), - failure_source = %failure_source, - failure_detail = ?failure_detail.as_deref(), - "Live upstream stream ended with response.failed" - ); - let mut metadata = stream_failure_metadata(&failure, usage); - insert_first_token_ms(&mut metadata, first_token_ms); - record_live_response_stream_event( - &context, - status_code_for_stream_failure(&failure), - UsageRecordLevel::Error, - "v1 responses stream failed", - metadata, - &rate_limit_headers, - &body, - ) - .await; - } - Ok(CollectedResponse::MissingCompleted | CollectedResponse::Empty) => { - tracing::warn!( - account_id = %context.account_id, - "Live upstream stream ended without response.completed" - ); - let mut metadata = serde_json::json!({ - "stream": true, - "failed": true, - "upstreamCode": "missing_completed", - "usage": usage, - }); - insert_first_token_ms(&mut metadata, first_token_ms); - record_live_response_stream_event( - &context, - 502, - UsageRecordLevel::Error, - "v1 responses stream ended without response.completed", - metadata, - &rate_limit_headers, - &body, - ) - .await; - } - Err(error) => { - tracing::warn!(account_id = %context.account_id, error = %error, "Failed to parse completed live stream"); - let mut metadata = serde_json::json!({ - "stream": true, - "sseParseError": error.to_string(), - "usage": usage, - }); - insert_first_token_ms(&mut metadata, first_token_ms); - record_live_response_stream_event( - &context, - 502, - UsageRecordLevel::Warn, - "v1 responses stream SSE response invalid", - metadata, - &rate_limit_headers, - &body, - ) - .await; - } - } - - context.account_lease.complete().await; -} diff --git a/backend/src/dispatch/stream/mod.rs b/backend/src/dispatch/stream/mod.rs index 589774e87..360811612 100644 --- a/backend/src/dispatch/stream/mod.rs +++ b/backend/src/dispatch/stream/mod.rs @@ -2,6 +2,3 @@ pub(crate) mod lifecycle; pub(crate) mod live; -pub(crate) mod prefetch; -pub(crate) mod sse_failure; -pub(crate) mod trace; diff --git a/backend/src/dispatch/stream/prefetch.rs b/backend/src/dispatch/stream/prefetch.rs deleted file mode 100644 index 9ceeb4e56..000000000 --- a/backend/src/dispatch/stream/prefetch.rs +++ /dev/null @@ -1,36 +0,0 @@ -use bytes::Bytes; -use futures::StreamExt; - -use crate::upstream::openai::{ - protocol::sse::{SseError, sse_frame_end}, - transport::CodexBackendSseStream, -}; - -use crate::dispatch::errors::ResponseDispatchError; - -const MAX_STREAM_PREFETCH_BYTES: usize = 64 * 1024; - -pub(in crate::dispatch) async fn prefetch_first_sse_chunk( - mut body: CodexBackendSseStream, -) -> Result<(Bytes, CodexBackendSseStream), ResponseDispatchError> { - let mut prefetched = Vec::new(); - while sse_frame_end(&prefetched).is_none() { - let Some(next) = body.next().await else { - if prefetched.is_empty() { - return Err(ResponseDispatchError::EmptyUpstreamResponse); - } - return Err(ResponseDispatchError::MissingCompleted); - }; - let chunk = next.map_err(ResponseDispatchError::Upstream)?; - prefetched.extend_from_slice(&chunk); - if prefetched.len() > MAX_STREAM_PREFETCH_BYTES { - return Err(ResponseDispatchError::InvalidSse( - SseError::BufferExceeded { - max_bytes: MAX_STREAM_PREFETCH_BYTES, - }, - )); - } - } - - Ok((Bytes::from(prefetched), body)) -} diff --git a/backend/src/dispatch/stream/sse_failure.rs b/backend/src/dispatch/stream/sse_failure.rs deleted file mode 100644 index 502d527a7..000000000 --- a/backend/src/dispatch/stream/sse_failure.rs +++ /dev/null @@ -1,221 +0,0 @@ -use serde_json::Value; - -use crate::{ - dispatch::errors::{is_history_recovery_code, is_model_unsupported_signal}, - fleet::account::AccountStatus, - upstream::openai::{ - protocol::{ - events::TokenUsage, - responses::{CollectedResponse, ResponsesSseFailure, response_from_codex_sse}, - sse::SseError, - }, - transport::is_banned_auth_signal, - }, -}; - -pub(in crate::dispatch) const STREAM_DISCONNECTED_CODE: &str = "stream_disconnected"; -pub(in crate::dispatch) const STREAM_DISCONNECTED_MESSAGE: &str = - "Upstream stream closed before response.completed"; - -pub(in crate::dispatch) fn sse_failure_error_body(failure: &ResponsesSseFailure) -> String { - match failure.upstream_code.as_deref() { - Some(code) => serde_json::json!({ - "error": { - "code": code, - "message": failure.message.as_str(), - } - }) - .to_string(), - None => failure.message.clone(), - } -} - -fn sse_failure_matches(failure: &ResponsesSseFailure, signal: F) -> bool -where - F: Fn(&str) -> bool, -{ - failure.upstream_code.as_deref().is_some_and(&signal) || signal(&failure.message) -} - -fn sse_failure_matches_parts( - failure: &ResponsesSseFailure, - code_signal: C, - message_signal: M, -) -> bool -where - C: Fn(&str) -> bool, - M: Fn(&str) -> bool, -{ - failure.upstream_code.as_deref().is_some_and(code_signal) || message_signal(&failure.message) -} - -pub(crate) fn is_quota_exhausted_sse_failure(failure: &ResponsesSseFailure) -> bool { - sse_failure_matches_parts( - failure, - |code| matches!(code, "quota_exceeded" | "insufficient_quota"), - |message| message.to_ascii_lowercase().contains("quota"), - ) -} - -fn is_auth_sse_failure_code(code: &str) -> bool { - let code = code.to_ascii_lowercase(); - matches!( - code.as_str(), - "token_invalid" - | "token_expired" - | "token_revoked" - | "account_deactivated" - | "unauthorized" - | "invalid_api_key" - ) -} - -fn is_auth_sse_failure_message(message: &str) -> bool { - let message = message.to_ascii_lowercase(); - message.contains("token revoked") - || message.contains("token invalid") - || message.contains("token expired") -} - -pub(crate) fn is_auth_sse_failure(failure: &ResponsesSseFailure) -> bool { - sse_failure_matches_parts( - failure, - is_auth_sse_failure_code, - is_auth_sse_failure_message, - ) -} - -pub(in crate::dispatch) fn is_model_unsupported_sse_failure(failure: &ResponsesSseFailure) -> bool { - sse_failure_matches(failure, is_model_unsupported_signal) -} - -pub(in crate::dispatch) fn is_history_recovery_sse_failure(failure: &ResponsesSseFailure) -> bool { - failure - .upstream_code - .as_deref() - .is_some_and(is_history_recovery_code) -} - -pub(crate) fn auth_sse_failure_account_status(failure: &ResponsesSseFailure) -> AccountStatus { - if sse_failure_matches(failure, is_banned_auth_signal) { - AccountStatus::Banned - } else { - AccountStatus::Expired - } -} - -pub(in crate::dispatch) fn first_sse_failure( - prefetched: &[u8], -) -> Result, SseError> { - let body = String::from_utf8_lossy(prefetched); - match response_from_codex_sse(&body, None)? { - CollectedResponse::Failed(failure) => Ok(Some(failure)), - CollectedResponse::Completed(_) - | CollectedResponse::Incomplete(_) - | CollectedResponse::MissingCompleted - | CollectedResponse::Empty => Ok(None), - } -} - -pub(in crate::dispatch) fn stream_failure_metadata( - failure: &ResponsesSseFailure, - usage: Option, -) -> Value { - let mut metadata = serde_json::json!({ - "stream": true, - "failed": true, - "failureEvent": failure.event, - "failureMessage": failure.message, - "upstreamCode": failure.upstream_code, - "usage": usage, - }); - enrich_stream_failure_source_metadata(&mut metadata, failure); - metadata -} - -fn enrich_stream_failure_source_metadata(metadata: &mut Value, failure: &ResponsesSseFailure) { - let Some(object) = metadata.as_object_mut() else { - return; - }; - object.insert( - "failureSource".to_string(), - Value::String(stream_failure_source(failure).to_string()), - ); - if let Some(detail) = synthetic_stream_disconnected_detail(failure) { - object.insert("synthetic".to_string(), Value::Bool(true)); - if !detail.is_empty() { - object.insert("failureDetail".to_string(), Value::String(detail)); - } - } -} - -pub(in crate::dispatch) fn stream_failure_source(failure: &ResponsesSseFailure) -> &'static str { - if synthetic_stream_disconnected_detail(failure).is_some() { - "proxy" - } else { - "upstream" - } -} - -pub(in crate::dispatch) fn synthetic_stream_disconnected_detail( - failure: &ResponsesSseFailure, -) -> Option { - if failure.upstream_code.as_deref() != Some(STREAM_DISCONNECTED_CODE) { - return None; - } - let detail = failure - .message - .strip_prefix(STREAM_DISCONNECTED_MESSAGE)? - .strip_prefix(": ") - .unwrap_or_default() - .trim() - .to_string(); - Some(detail) -} - -pub(in crate::dispatch) fn status_code_for_stream_failure(failure: &ResponsesSseFailure) -> i64 { - let code = failure - .upstream_code - .as_deref() - .unwrap_or("error") - .to_ascii_lowercase(); - if code.contains("model") && (code.contains("not_supported") || code.contains("not_available")) - { - return 400; - } - if code.contains("invalid_request") || code.contains("not_found") { - return 400; - } - if code.contains("context_window") - || code.contains("invalid_prompt") - || code.contains("cyber_policy") - || code.contains("bad_request") - { - return 400; - } - if code.contains("rate_limit") || code.contains("usage_limit") { - return 429; - } - if code.contains("unauthorized") - || code.contains("invalid_api_key") - || code == "token_invalid" - || code == "token_expired" - || code == "account_deactivated" - { - return 401; - } - if code.contains("forbidden") || code.contains("banned") { - return 403; - } - if code.contains("payment") || code.contains("quota") { - return 429; - } - if code.contains("server_overloaded") { - return 503; - } - 502 -} - -pub(in crate::dispatch) fn stream_failure_http_status(failure: &ResponsesSseFailure) -> u16 { - u16::try_from(status_code_for_stream_failure(failure)).unwrap_or(502) -} diff --git a/backend/src/dispatch/transport/account.rs b/backend/src/dispatch/transport/account.rs new file mode 100644 index 000000000..ee4c77cad --- /dev/null +++ b/backend/src/dispatch/transport/account.rs @@ -0,0 +1,96 @@ +//! 调度层共享的账号级 Codex 上游调用。 + +use std::time::Instant; + +use crate::{ + dispatch::affinity::AccountIdentityService, + fleet::account::Account, + upstream::openai::{ + protocol::responses::CodexResponsesRequest, + transport::{ + CodexBackendClient, CodexBackendResponse, CodexBackendStreamingResponse, + CodexClientError, CodexRequestContext, + }, + }, +}; + +#[derive(Clone, Copy)] +pub(in crate::dispatch) struct AccountUpstreamContext<'a> { + pub codex: &'a CodexBackendClient, + pub account_identity: &'a AccountIdentityService, + pub request_id: &'a str, + pub account: &'a Account, + pub cookie_header: Option<&'a str>, +} + +pub(in crate::dispatch) async fn create_response_with_account( + context: AccountUpstreamContext<'_>, + request: &CodexResponsesRequest, + started_at: Instant, +) -> Result { + let identity = context + .account_identity + .scope(request, &context.account.id, context.request_id); + context + .codex + .create_response_with_pool_account_started_at( + request, + CodexRequestContext { + access_token: &context.account.access_token, + account_id: context.account.account_id.as_deref(), + request_id: context.request_id, + turn_state: request.turn_state.as_deref(), + turn_metadata: request.turn_metadata.as_deref(), + beta_features: request.beta_features.as_deref(), + include_timing_metrics: request.include_timing_metrics.as_deref(), + version: request.version.as_deref(), + codex_window_id: identity.window_id.as_deref(), + parent_thread_id: identity.parent_thread_id.as_deref(), + cookie_header: context.cookie_header, + installation_id: Some(&identity.installation_id), + session_id: identity.session_id.as_deref(), + thread_id: identity.thread_id.as_deref(), + prompt_cache_key: identity.prompt_cache_key.as_deref(), + client_request_id: Some(&identity.client_request_id), + turn_id: identity.turn_id.as_deref(), + }, + Some(&context.account.id), + started_at, + ) + .await +} + +pub(in crate::dispatch) async fn create_response_stream_with_account( + context: AccountUpstreamContext<'_>, + request: &CodexResponsesRequest, +) -> Result { + let identity = context + .account_identity + .scope(request, &context.account.id, context.request_id); + context + .codex + .create_response_stream_with_pool_account( + request, + CodexRequestContext { + access_token: &context.account.access_token, + account_id: context.account.account_id.as_deref(), + request_id: context.request_id, + turn_state: request.turn_state.as_deref(), + turn_metadata: request.turn_metadata.as_deref(), + beta_features: request.beta_features.as_deref(), + include_timing_metrics: request.include_timing_metrics.as_deref(), + version: request.version.as_deref(), + codex_window_id: identity.window_id.as_deref(), + parent_thread_id: identity.parent_thread_id.as_deref(), + cookie_header: context.cookie_header, + installation_id: Some(&identity.installation_id), + session_id: identity.session_id.as_deref(), + thread_id: identity.thread_id.as_deref(), + prompt_cache_key: identity.prompt_cache_key.as_deref(), + client_request_id: Some(&identity.client_request_id), + turn_id: identity.turn_id.as_deref(), + }, + Some(&context.account.id), + ) + .await +} diff --git a/backend/src/dispatch/transport/canonical.rs b/backend/src/dispatch/transport/canonical.rs new file mode 100644 index 000000000..8ca11b4ec --- /dev/null +++ b/backend/src/dispatch/transport/canonical.rs @@ -0,0 +1,424 @@ +//! Responses SSE 原始帧到规范事件的唯一解码入口。 + +use bytes::Bytes; +use serde_json::Value; + +use crate::upstream::openai::protocol::{ + events::{TokenUsage, extract_usage}, + responses::{ + CollectedResponse, ResponsesSseFailure, StreamCommitPolicy, + reconvert_responses_sse_event_tuple_values, + }, + sse::{ + MAX_SSE_EVENT_BUFFER_BYTES, SseError, encode_sse_event, parse_sse_events, sse_frame_end, + sse_frame_is_done, + }, +}; + +/// 将 complete transport 的原始 SSE 归一化为类型化终态事实。 +pub(in crate::dispatch) fn normalize_complete_response( + body: &str, + tuple_schema: Option<&Value>, +) -> Result { + let events = parse_sse_events(body)?; + let mut output_text = String::new(); + let mut output_items = Vec::new(); + let mut terminal_response = None; + let mut terminal_incomplete = false; + let mut failed_response = None; + + for event in events { + let Ok(mut value) = serde_json::from_str::(&event.data) else { + continue; + }; + if let Some(tuple_schema) = tuple_schema { + value = reconvert_responses_sse_event_tuple_values( + event.event.as_deref(), + value, + tuple_schema, + ); + } + match event.event.as_deref() { + Some("response.output_text.delta") => { + if let Some(delta) = value.get("delta").and_then(Value::as_str) { + output_text.push_str(delta); + } + } + Some("response.output_item.done") => { + if let Some(item) = value.get("item") { + output_items.push(item.clone()); + } + } + Some("response.completed") => { + terminal_response = value.get("response").cloned(); + terminal_incomplete = false; + } + Some("response.incomplete") => { + terminal_response = value.get("response").cloned(); + terminal_incomplete = true; + } + Some(event_name @ ("error" | "response.failed")) if failed_response.is_none() => { + failed_response = Some(ResponsesSseFailure::from_event(event_name, &value)); + } + _ => {} + } + } + + if let Some(failure) = failed_response { + return Ok(CollectedResponse::Failed(failure)); + } + let Some(response) = terminal_response else { + return Ok(CollectedResponse::MissingCompleted); + }; + if terminal_incomplete { + return Ok(CollectedResponse::Incomplete(response)); + } + if complete_response_is_empty(&response, &output_text, &output_items) { + return Ok(CollectedResponse::Empty); + } + Ok(CollectedResponse::Completed(response)) +} + +fn complete_response_is_empty(response: &Value, output_text: &str, output_items: &[Value]) -> bool { + if !output_text.trim().is_empty() || !output_items.is_empty() { + return false; + } + if response.get("status").and_then(Value::as_str) == Some("incomplete") { + return false; + } + response + .pointer("/usage/output_tokens") + .and_then(Value::as_u64) + .unwrap_or_default() + == 0 +} + +/// 一次解码后可同时供 HTTP SSE、Responses WebSocket 与生命周期控制器消费的事件。 +#[derive(Debug, Clone, PartialEq, Eq)] +pub(crate) struct CanonicalResponseEvent { + event_type: Option, + data: Value, +} + +impl CanonicalResponseEvent { + pub(crate) fn event_type(&self) -> Option<&str> { + self.event_type.as_deref() + } + + pub(crate) fn data(&self) -> &Value { + &self.data + } + + pub(in crate::dispatch) fn response_id(&self) -> Option<&str> { + self.data + .pointer("/response/id") + .and_then(Value::as_str) + .filter(|value| !value.trim().is_empty()) + } + + pub(in crate::dispatch) fn usage(&self) -> Option { + self.data + .get("response") + .and_then(extract_usage) + .or_else(|| extract_usage(&self.data)) + } + + pub(in crate::dispatch) fn has_first_output(&self) -> bool { + match self.event_type() { + Some( + "response.output_text.delta" + | "response.reasoning_summary_text.delta" + | "response.reasoning_text.delta" + | "response.function_call_arguments.delta" + | "response.custom_tool_call_input.delta", + ) => self + .data + .get("delta") + .and_then(Value::as_str) + .is_some_and(|value| !value.is_empty()), + Some("response.output_text.done") => self + .data + .get("text") + .and_then(Value::as_str) + .is_some_and(|value| !value.is_empty()), + Some("response.function_call_arguments.done") => self + .data + .get("arguments") + .and_then(Value::as_str) + .is_some_and(|value| !value.is_empty()), + Some("response.output_item.done") => { + self.data.get("item").is_some_and(Value::is_object) + } + _ => false, + } + } + + pub(in crate::dispatch) fn terminal(&self) -> Option { + match self.event_type() { + Some("response.completed") => self + .data + .get("response") + .filter(|response| response.is_object()) + .cloned() + .map(CanonicalStreamTerminal::Completed), + Some("response.incomplete") => self + .data + .get("response") + .filter(|response| response.is_object()) + .cloned() + .map(CanonicalStreamTerminal::Incomplete), + Some(event @ ("response.failed" | "error")) => Some(CanonicalStreamTerminal::Failed( + ResponsesSseFailure::from_event(event, &self.data), + )), + _ => None, + } + } + + pub(in crate::dispatch) fn proxy_failure(data: Value) -> Self { + Self { + event_type: Some("response.failed".to_string()), + data, + } + } +} + +/// 上游明确发出的业务终态。 +#[derive(Debug, Clone, PartialEq, Eq)] +pub(in crate::dispatch) enum CanonicalStreamTerminal { + Completed(Value), + Incomplete(Value), + Failed(ResponsesSseFailure), +} + +/// 已解码事件与对应下游字节的单个 fanout 单元。 +#[derive(Debug, Clone)] +pub(crate) struct CanonicalResponseChunk { + bytes: Bytes, + events: Vec, +} + +impl CanonicalResponseChunk { + pub(in crate::dispatch) fn new(bytes: Bytes, events: Vec) -> Self { + Self { bytes, events } + } + + pub(in crate::dispatch) fn bytes(&self) -> &Bytes { + &self.bytes + } + + pub(in crate::dispatch) fn into_parts(self) -> (Bytes, Vec) { + (self.bytes, self.events) + } + + pub(crate) fn events(&self) -> &[CanonicalResponseEvent] { + &self.events + } +} + +/// 单次输入产生的所有下游 fanout 单元。 +#[derive(Debug)] +pub(in crate::dispatch) struct CanonicalStreamBatch { + pub chunks: Vec, +} + +impl CanonicalStreamBatch { + pub(in crate::dispatch) fn append(&mut self, mut other: Self) { + self.chunks.append(&mut other.chunks); + } + + /// 只把真实输出前出现的失败视为可在 attempt 阶段处理的失败。 + /// 同一上游 chunk 可能同时包含输出与终态，事件顺序不能被网络分块抹掉。 + pub(in crate::dispatch) fn first_failure(&self) -> Option { + for event in self.chunks.iter().flat_map(|chunk| chunk.events()) { + if event.has_first_output() { + return None; + } + if let Some(terminal) = event.terminal() { + return match terminal { + CanonicalStreamTerminal::Failed(failure) => Some(failure), + CanonicalStreamTerminal::Completed(_) + | CanonicalStreamTerminal::Incomplete(_) => None, + }; + } + } + None + } +} + +/// 原始 SSE 在整个 attempt/live 生命周期中唯一持有的增量 decoder。 +pub(in crate::dispatch) struct CanonicalStreamDecoder { + tuple_schema: Option, + pending: Vec, + forwardable_frame_seen: bool, + first_output_seen: bool, + terminal: Option, + transport_done_seen: bool, + completed_usage: Option, + fallback_usage: Option, + last_response_id: Option, +} + +impl CanonicalStreamDecoder { + pub(in crate::dispatch) fn new(tuple_schema: Option) -> Self { + Self { + tuple_schema, + pending: Vec::new(), + forwardable_frame_seen: false, + first_output_seen: false, + terminal: None, + transport_done_seen: false, + completed_usage: None, + fallback_usage: None, + last_response_id: None, + } + } + + pub(in crate::dispatch) fn push( + &mut self, + chunk: Bytes, + ) -> Result { + self.pending.extend_from_slice(&chunk); + let mut transformed_chunks = Vec::new(); + let mut consumed = 0usize; + + while let Some(frame_len) = sse_frame_end(&self.pending[consumed..]) { + let end = consumed.saturating_add(frame_len); + let frame = self.pending[consumed..end].to_vec(); + if let Some((transformed, events)) = self.normalize_frame(&frame)? { + transformed_chunks.push(CanonicalResponseChunk::new(transformed, events)); + } + consumed = end; + if self.terminal.is_some() || self.transport_done_seen { + consumed = self.pending.len(); + break; + } + } + + if consumed != 0 { + self.pending.drain(..consumed); + } + if self.pending.len() > MAX_SSE_EVENT_BUFFER_BYTES { + return Err(SseError::BufferExceeded { + max_bytes: MAX_SSE_EVENT_BUFFER_BYTES, + }); + } + + Ok(CanonicalStreamBatch { + chunks: transformed_chunks, + }) + } + + pub(in crate::dispatch) fn finish(&mut self) -> Result { + if self.pending.is_empty() { + return Ok(CanonicalStreamBatch { chunks: Vec::new() }); + } + let frame = std::mem::take(&mut self.pending); + let Some((transformed, events)) = self.normalize_frame(&frame)? else { + return Ok(CanonicalStreamBatch { chunks: Vec::new() }); + }; + Ok(CanonicalStreamBatch { + chunks: vec![CanonicalResponseChunk::new(transformed, events)], + }) + } + + pub(in crate::dispatch) fn take_terminal(&mut self) -> Option { + self.terminal.take() + } + + pub(in crate::dispatch) fn take_transport_done(&mut self) -> bool { + std::mem::take(&mut self.transport_done_seen) + } + + pub(in crate::dispatch) fn transport_done_seen(&self) -> bool { + self.transport_done_seen + } + + pub(in crate::dispatch) fn first_output_seen(&self) -> bool { + self.first_output_seen + } + + pub(in crate::dispatch) fn commit_boundary_reached(&self, policy: StreamCommitPolicy) -> bool { + match policy { + StreamCommitPolicy::FirstForwardableEvent => self.forwardable_frame_seen, + StreamCommitPolicy::UntilOutputOrTerminal => { + self.first_output_seen || self.terminal.is_some() + } + } + } + + pub(in crate::dispatch) fn usage(&self) -> Option { + self.completed_usage.or(self.fallback_usage) + } + + pub(in crate::dispatch) fn last_response_id(&self) -> Option<&str> { + self.last_response_id.as_deref() + } + + fn normalize_frame( + &mut self, + frame: &[u8], + ) -> Result)>, SseError> { + let frame_text = + std::str::from_utf8(frame).map_err(|error| SseError::ParseError(error.to_string()))?; + if sse_frame_is_done(frame_text) { + self.transport_done_seen = true; + return Ok(None); + } + let events = parse_sse_events(frame_text)?; + let mut canonical_events = Vec::new(); + let mut transformed = String::new(); + + for event in events { + let Ok(mut data) = serde_json::from_str::(&event.data) else { + continue; + }; + let event_type = event.event.or_else(|| { + data.get("type") + .and_then(Value::as_str) + .map(ToString::to_string) + }); + if let Some(tuple_schema) = self.tuple_schema.as_ref() { + data = reconvert_responses_sse_event_tuple_values( + event_type.as_deref(), + data, + tuple_schema, + ); + transformed.push_str(&encode_sse_event( + event_type.as_deref().unwrap_or_default(), + &data.to_string(), + )); + } + let canonical = CanonicalResponseEvent { event_type, data }; + self.observe(&canonical); + canonical_events.push(canonical); + } + + self.forwardable_frame_seen = true; + + Ok(Some(( + if self.tuple_schema.is_some() && !transformed.is_empty() { + Bytes::from(transformed) + } else { + Bytes::copy_from_slice(frame) + }, + canonical_events, + ))) + } + + fn observe(&mut self, event: &CanonicalResponseEvent) { + self.first_output_seen |= event.has_first_output(); + if let Some(response_id) = event.response_id() { + self.last_response_id = Some(response_id.to_string()); + } + if let Some(usage) = event.usage() { + if event.event_type() == Some("response.completed") { + self.completed_usage = Some(usage); + } else { + self.fallback_usage = Some(usage); + } + } + if self.terminal.is_none() { + self.terminal = event.terminal(); + } + } +} diff --git a/backend/src/dispatch/transport/mod.rs b/backend/src/dispatch/transport/mod.rs new file mode 100644 index 000000000..d178e8e92 --- /dev/null +++ b/backend/src/dispatch/transport/mod.rs @@ -0,0 +1,6 @@ +//! Attempt 阶段的账号作用域上游调用与事实规范化。 + +pub(in crate::dispatch) mod account; +pub(crate) mod canonical; +pub(in crate::dispatch) mod observation; +pub(in crate::dispatch) mod prefetch; diff --git a/backend/src/dispatch/transport/observation.rs b/backend/src/dispatch/transport/observation.rs new file mode 100644 index 000000000..4f66978f5 --- /dev/null +++ b/backend/src/dispatch/transport/observation.rs @@ -0,0 +1,152 @@ +//! 上游原始错误到 attempt 通用事实的单次规范化。 + +use serde_json::Value; + +use crate::upstream::openai::transport::{ + CodexClientError, + websocket::{CodexWebSocketExchangeError, PreviousResponseUnavailableReason}, +}; + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub(in crate::dispatch) enum UpstreamFailureKind { + HttpStatus, + HttpConnect, + HttpTimeout, + HttpTransport, + StreamIdle, + WebSocketUpstream, + WebSocketTransport, + WebSocketTimeout, + WebSocketProtocol, + ContinuationUnavailable(PreviousResponseUnavailableReason), + RequestEncoding, + Protocol, +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub(in crate::dispatch) struct UpstreamFailureFacts { + pub kind: UpstreamFailureKind, + pub status_code: Option, + pub code: Option, + pub message: String, + pub body: String, + pub retry_after_seconds: Option, +} + +pub(in crate::dispatch) fn upstream_failure_facts( + error: &CodexClientError, +) -> UpstreamFailureFacts { + match error { + CodexClientError::Http(error) => UpstreamFailureFacts { + kind: if error.is_connect() { + UpstreamFailureKind::HttpConnect + } else if error.is_timeout() { + UpstreamFailureKind::HttpTimeout + } else { + UpstreamFailureKind::HttpTransport + }, + status_code: error.status().map(|status| status.as_u16()), + code: None, + message: error.to_string(), + body: error.to_string(), + retry_after_seconds: None, + }, + CodexClientError::Upstream { + status, + body, + retry_after_seconds, + .. + } => { + let (code, message) = error_fields(body); + UpstreamFailureFacts { + kind: UpstreamFailureKind::HttpStatus, + status_code: Some(status.as_u16()), + code, + message: message.unwrap_or_else(|| body.clone()), + body: body.clone(), + retry_after_seconds: *retry_after_seconds, + } + } + CodexClientError::StreamIdleTimeout { .. } => { + simple_facts(UpstreamFailureKind::StreamIdle, error) + } + CodexClientError::WebSocket(CodexWebSocketExchangeError::Upstream(upstream)) => { + UpstreamFailureFacts { + kind: UpstreamFailureKind::WebSocketUpstream, + status_code: Some(upstream.status_code), + code: upstream + .code + .clone() + .or_else(|| error_fields(&upstream.body).0), + message: upstream + .message + .clone() + .or_else(|| error_fields(&upstream.body).1) + .unwrap_or_else(|| upstream.body.clone()), + body: upstream.body.clone(), + retry_after_seconds: upstream.retry_after_seconds, + } + } + CodexClientError::WebSocket(CodexWebSocketExchangeError::ContinuationUnavailable { + reason, + }) => simple_facts(UpstreamFailureKind::ContinuationUnavailable(*reason), error), + CodexClientError::WebSocket( + CodexWebSocketExchangeError::ConnectTimeout { .. } + | CodexWebSocketExchangeError::SendTimeout { .. } + | CodexWebSocketExchangeError::ReceiveIdleTimeout { .. } + | CodexWebSocketExchangeError::InitialEventTimeout { .. }, + ) => simple_facts(UpstreamFailureKind::WebSocketTimeout, error), + CodexClientError::WebSocket( + CodexWebSocketExchangeError::Transport(_) + | CodexWebSocketExchangeError::ClosedBeforeTerminal + | CodexWebSocketExchangeError::ReusedConnectionDiedBeforeFirstOutput { .. }, + ) => simple_facts(UpstreamFailureKind::WebSocketTransport, error), + CodexClientError::WebSocket( + CodexWebSocketExchangeError::InvalidSse(_) + | CodexWebSocketExchangeError::InvalidCompletedResponse { .. } + | CodexWebSocketExchangeError::UnexpectedBinaryEvent, + ) => simple_facts(UpstreamFailureKind::WebSocketProtocol, error), + CodexClientError::WebSocket(CodexWebSocketExchangeError::InvalidRequest(_)) + | CodexClientError::WebSocketEncode(_) + | CodexClientError::InvalidHeaderName(_) + | CodexClientError::InvalidHeaderValue(_) + | CodexClientError::CustomCa(_) => { + simple_facts(UpstreamFailureKind::RequestEncoding, error) + } + CodexClientError::InvalidSse(_) => simple_facts(UpstreamFailureKind::Protocol, error), + } +} + +fn simple_facts(kind: UpstreamFailureKind, error: &CodexClientError) -> UpstreamFailureFacts { + let message = error.to_string(); + UpstreamFailureFacts { + kind, + status_code: None, + code: None, + body: message.clone(), + message, + retry_after_seconds: None, + } +} + +fn error_fields(body: &str) -> (Option, Option) { + let Ok(value) = serde_json::from_str::(body) else { + return (None, None); + }; + let error = value + .pointer("/response/error") + .or_else(|| value.get("error")) + .unwrap_or(&value); + ( + error + .get("code") + .or_else(|| value.get("code")) + .and_then(Value::as_str) + .map(ToString::to_string), + error + .get("message") + .or_else(|| value.get("message")) + .and_then(Value::as_str) + .map(ToString::to_string), + ) +} diff --git a/backend/src/dispatch/transport/prefetch.rs b/backend/src/dispatch/transport/prefetch.rs new file mode 100644 index 000000000..694b945e3 --- /dev/null +++ b/backend/src/dispatch/transport/prefetch.rs @@ -0,0 +1,72 @@ +//! 上游 SSE 提交边界前的规范事件预取。 + +use bytes::Bytes; +use futures::StreamExt; + +use crate::{ + dispatch::transport::canonical::{CanonicalStreamBatch, CanonicalStreamDecoder}, + upstream::openai::{ + protocol::{responses::StreamCommitPolicy, sse::SseError}, + transport::{CodexBackendSseStream, CodexClientError}, + }, +}; + +const MAX_STREAM_PREFETCH_BYTES: usize = 64 * 1024; + +#[derive(Debug, thiserror::Error)] +pub(in crate::dispatch) enum StreamPrefetchError { + #[error("upstream stream failed before the commit boundary: {0}")] + Upstream(#[from] CodexClientError), + #[error("upstream stream ended without any data")] + Empty, + #[error("upstream stream ended before output or a terminal event")] + NoCommitBoundary, + #[error("invalid upstream SSE response: {0}")] + InvalidSse(#[from] SseError), +} + +pub(in crate::dispatch) struct PrefetchedStream { + pub bytes: Bytes, + pub initial_batch: CanonicalStreamBatch, + pub body: CodexBackendSseStream, +} + +pub(in crate::dispatch) async fn prefetch_until_commit( + mut body: CodexBackendSseStream, + decoder: &mut CanonicalStreamDecoder, + policy: StreamCommitPolicy, +) -> Result { + let mut prefetched = Vec::new(); + let mut initial_batch = CanonicalStreamBatch { chunks: Vec::new() }; + while !decoder.commit_boundary_reached(policy) { + let Some(next) = body.next().await else { + if prefetched.is_empty() { + return Err(StreamPrefetchError::Empty); + } + initial_batch.append(decoder.finish()?); + if decoder.commit_boundary_reached(policy) { + break; + } + return Err(StreamPrefetchError::NoCommitBoundary); + }; + let chunk = next?; + prefetched.extend_from_slice(&chunk); + let batch = decoder.push(chunk)?; + initial_batch.append(batch); + if decoder.transport_done_seen() && !decoder.commit_boundary_reached(policy) { + return Err(StreamPrefetchError::NoCommitBoundary); + } + if !decoder.commit_boundary_reached(policy) && prefetched.len() > MAX_STREAM_PREFETCH_BYTES + { + return Err(StreamPrefetchError::InvalidSse(SseError::BufferExceeded { + max_bytes: MAX_STREAM_PREFETCH_BYTES, + })); + } + } + + Ok(PrefetchedStream { + bytes: Bytes::from(prefetched), + initial_batch, + body, + }) +} diff --git a/backend/src/dispatch/upstream_call.rs b/backend/src/dispatch/upstream_call.rs deleted file mode 100644 index 46ebb0b15..000000000 --- a/backend/src/dispatch/upstream_call.rs +++ /dev/null @@ -1,221 +0,0 @@ -//! 调度层共享的账号级 Codex 上游调用。 - -use std::time::Instant; - -use serde_json::Value; - -use crate::{ - dispatch::{affinity::AccountIdentityService, recovery::cloudflare::CloudflareRecovery}, - fleet::{ - account::Account, - pool::{AccountLease, AccountPoolService}, - quota::{ - quota_from_usage, quota_snapshot_limit_reached, quota_snapshot_limit_window_seconds, - quota_snapshot_reset_at, - }, - }, - upstream::openai::{ - protocol::responses::CodexResponsesRequest, - transport::{ - CodexBackendClient, CodexBackendResponse, CodexBackendStreamingResponse, - CodexClientError, CodexRequestContext, - }, - }, -}; - -pub(crate) const QUOTA_VERIFY_LIMIT_REACHED_MESSAGE: &str = - "Upstream usage quota still reports limit_reached"; - -pub(crate) enum QuotaVerificationDecision { - Ready(Box), - RetryWithAnotherAccount, -} - -pub(crate) struct QuotaVerificationContext<'a> { - pub account_pool: &'a AccountPoolService, - pub codex: &'a CodexBackendClient, - pub cloudflare: &'a CloudflareRecovery, - pub account_identity: &'a AccountIdentityService, - pub request_id: &'a str, -} - -#[derive(Clone, Copy)] -pub(crate) struct AccountUpstreamContext<'a> { - pub codex: &'a CodexBackendClient, - pub account_identity: &'a AccountIdentityService, - pub cloudflare: &'a CloudflareRecovery, - pub request_id: &'a str, - pub account: &'a Account, -} - -pub(crate) async fn verify_acquired_quota_if_required( - context: QuotaVerificationContext<'_>, - acquired: AccountLease, -) -> QuotaVerificationDecision { - if !acquired.account.quota_verify_required { - return QuotaVerificationDecision::Ready(Box::new(acquired)); - } - - let account_id = acquired.account.id.clone(); - let identity = context - .account_identity - .scope_auxiliary(&account_id, context.request_id); - let cookie_header = context - .cloudflare - .cookie_header_for_request(&account_id, "/codex/usage") - .await; - let usage = context - .codex - .fetch_usage(CodexRequestContext { - access_token: &acquired.account.access_token, - account_id: acquired.account.account_id.as_deref(), - request_id: context.request_id, - turn_state: None, - turn_metadata: None, - beta_features: None, - include_timing_metrics: None, - version: None, - codex_window_id: None, - parent_thread_id: None, - cookie_header: cookie_header.as_deref(), - installation_id: Some(&identity.installation_id), - session_id: None, - thread_id: None, - prompt_cache_key: None, - client_request_id: Some(&identity.client_request_id), - turn_id: None, - }) - .await; - - let raw = match usage { - Ok(raw) => raw, - Err(error) => { - tracing::warn!( - request_id = %context.request_id, - account_id = %account_id, - quota_verify_required = true, - quota_verify_result = "upstream_error", - retry_with_another_account = false, - error = %error, - "Failed to verify stale quota state before upstream request" - ); - return QuotaVerificationDecision::Ready(Box::new(acquired)); - } - }; - - let quota = quota_from_usage(&raw); - context - .account_pool - .apply_quota_snapshot(&account_id, "a) - .await; - if quota_snapshot_limit_reached("a) { - acquired.release_without_usage().await; - tracing::info!( - request_id = %context.request_id, - account_id = %account_id, - quota_verify_required = true, - quota_verify_result = "limit_reached", - retry_with_another_account = true, - "Quota verification reported exhausted account before upstream request" - ); - return QuotaVerificationDecision::RetryWithAnotherAccount; - } - - QuotaVerificationDecision::Ready(Box::new(acquired_with_verified_quota(acquired, "a))) -} - -fn acquired_with_verified_quota(mut acquired: AccountLease, quota: &Value) -> AccountLease { - let limit_reached = quota_snapshot_limit_reached(quota); - acquired.account.quota_verify_required = false; - acquired.account.quota_limit_reached = limit_reached; - acquired.account.quota_cooldown_until = limit_reached - .then_some(quota_snapshot_reset_at(quota)) - .flatten(); - if let Some(reset_at) = quota_snapshot_reset_at(quota) { - acquired.account.window_reset_at = Some(reset_at); - if let Some(limit_window_seconds) = quota_snapshot_limit_window_seconds(quota) { - acquired.account.limit_window_seconds = Some(limit_window_seconds); - } - } - acquired -} - -pub(crate) async fn create_response_with_account( - context: AccountUpstreamContext<'_>, - request: &CodexResponsesRequest, - started_at: Instant, -) -> Result { - let cookie_header = context - .cloudflare - .cookie_header_for_request(&context.account.id, "/codex/responses") - .await; - let identity = context - .account_identity - .scope(request, &context.account.id, context.request_id); - context - .codex - .create_response_with_pool_account_started_at( - request, - CodexRequestContext { - access_token: &context.account.access_token, - account_id: context.account.account_id.as_deref(), - request_id: context.request_id, - turn_state: request.turn_state.as_deref(), - turn_metadata: request.turn_metadata.as_deref(), - beta_features: request.beta_features.as_deref(), - include_timing_metrics: request.include_timing_metrics.as_deref(), - version: request.version.as_deref(), - codex_window_id: identity.window_id.as_deref(), - parent_thread_id: identity.parent_thread_id.as_deref(), - cookie_header: cookie_header.as_deref(), - installation_id: Some(&identity.installation_id), - session_id: identity.session_id.as_deref(), - thread_id: identity.thread_id.as_deref(), - prompt_cache_key: identity.prompt_cache_key.as_deref(), - client_request_id: Some(&identity.client_request_id), - turn_id: identity.turn_id.as_deref(), - }, - Some(&context.account.id), - started_at, - ) - .await -} - -pub(crate) async fn create_response_stream_with_account( - context: AccountUpstreamContext<'_>, - request: &CodexResponsesRequest, -) -> Result { - let cookie_header = context - .cloudflare - .cookie_header_for_request(&context.account.id, "/codex/responses") - .await; - let identity = context - .account_identity - .scope(request, &context.account.id, context.request_id); - context - .codex - .create_response_stream_with_pool_account( - request, - CodexRequestContext { - access_token: &context.account.access_token, - account_id: context.account.account_id.as_deref(), - request_id: context.request_id, - turn_state: request.turn_state.as_deref(), - turn_metadata: request.turn_metadata.as_deref(), - beta_features: request.beta_features.as_deref(), - include_timing_metrics: request.include_timing_metrics.as_deref(), - version: request.version.as_deref(), - codex_window_id: identity.window_id.as_deref(), - parent_thread_id: identity.parent_thread_id.as_deref(), - cookie_header: cookie_header.as_deref(), - installation_id: Some(&identity.installation_id), - session_id: identity.session_id.as_deref(), - thread_id: identity.thread_id.as_deref(), - prompt_cache_key: identity.prompt_cache_key.as_deref(), - client_request_id: Some(&identity.client_request_id), - turn_id: identity.turn_id.as_deref(), - }, - Some(&context.account.id), - ) - .await -} diff --git a/backend/src/upstream/openai/protocol/events.rs b/backend/src/upstream/openai/protocol/events.rs index ad41083b4..4b1f27997 100644 --- a/backend/src/upstream/openai/protocol/events.rs +++ b/backend/src/upstream/openai/protocol/events.rs @@ -127,20 +127,26 @@ pub fn extract_sse_usage(body: &str) -> Result, SseError> { /// `try again in 11.054s` / `try again in 28ms` 文本。 pub fn retry_after_seconds_from_body(body: &str) -> Option { let value = serde_json::from_str::(body).ok()?; + retry_after_seconds_from_value(&value) +} + +/// 从已解析的上游错误事件中提取重试秒数，不做业务错误分类。 +pub(crate) fn retry_after_seconds_from_value(value: &Value) -> Option { let error = value .pointer("/response/error") .or_else(|| value.get("error")) - .unwrap_or(&value); - if let Some(seconds) = error - .get("resets_in_seconds") - .and_then(Value::as_u64) - .filter(|seconds| *seconds > 0) - { - return Some(seconds); - } - retry_after_seconds_from_resets_at(error) + .unwrap_or(value); + retry_after_seconds_field(error) + .or_else(|| { + error + .get("resets_in_seconds") + .and_then(Value::as_u64) + .filter(|seconds| *seconds > 0) + }) + .or_else(|| retry_after_seconds_from_resets_at(error)) + .or_else(|| retry_after_seconds_field(value)) + .or_else(|| retry_after_seconds_header(value)) .or_else(|| retry_after_seconds_from_rate_limit_message(error)) - .or_else(|| retry_after_seconds_field(&value)) } fn retry_after_seconds_field(value: &Value) -> Option { @@ -155,6 +161,31 @@ fn retry_after_seconds_field(value: &Value) -> Option { .filter(|seconds| *seconds > 0) } +fn retry_after_seconds_header(value: &Value) -> Option { + value + .get("headers") + .and_then(Value::as_object) + .and_then(|headers| { + headers.iter().find_map(|(name, value)| { + if name.eq_ignore_ascii_case("retry-after") { + json_value_as_positive_u64(value) + } else { + None + } + }) + }) +} + +fn json_value_as_positive_u64(value: &Value) -> Option { + let seconds = match value { + Value::Number(value) => value.as_u64()?, + Value::String(value) => value.trim().parse::().ok()?, + Value::Array(values) => values.first().and_then(json_value_as_positive_u64)?, + Value::Null | Value::Bool(_) | Value::Object(_) => return None, + }; + (seconds > 0).then_some(seconds) +} + /// 标准化的单个限流窗口。 #[derive(Debug, Clone, Copy, PartialEq)] pub struct RateLimitWindow { diff --git a/backend/src/upstream/openai/protocol/responses.rs b/backend/src/upstream/openai/protocol/responses.rs index 57ce00068..725390444 100644 --- a/backend/src/upstream/openai/protocol/responses.rs +++ b/backend/src/upstream/openai/protocol/responses.rs @@ -6,6 +6,7 @@ use serde_json::{Map, Value}; use crate::{ infra::time::elapsed_millis_i64, upstream::openai::protocol::{ + events, schema::reconvert_tuple_values, sse::{SseError, SseEvent, parse_sse_events}, }, @@ -250,6 +251,12 @@ pub struct ResponsesSseFailure { pub message: String, /// 上游错误码。 pub upstream_code: Option, + /// 上游显式错误类型；不从业务码推导。 + pub upstream_type: Option, + /// 上游显式状态码；不从业务码或错误类型推导。 + pub explicit_status_code: Option, + /// 上游显式重试间隔，或从官方限流消息中解析出的重试间隔。 + pub retry_after_seconds: Option, } impl ResponsesSseFailure { @@ -258,6 +265,9 @@ impl ResponsesSseFailure { event: event.to_string(), message: failure_message(value).unwrap_or_else(|| "Codex upstream SSE failed".into()), upstream_code: failure_code(value), + upstream_type: failure_type(value), + explicit_status_code: failure_explicit_status_code(value), + retry_after_seconds: events::retry_after_seconds_from_value(value), } } } @@ -273,66 +283,6 @@ pub struct CompletedResponseMetadata { pub output: Vec, } -/// 将 Codex Responses SSE 完成响应收集为非流式 Responses JSON。 -pub fn response_from_codex_sse( - body: &str, - tuple_schema: Option<&Value>, -) -> Result { - let events = parse_sse_events(body)?; - let mut output_text = String::new(); - let mut output_items = Vec::new(); - let mut terminal_response = None; - let mut terminal_incomplete = false; - let mut failed_response = None; - - for event in events { - let Ok(value) = serde_json::from_str::(&event.data) else { - continue; - }; - match event.event.as_deref() { - Some("response.output_text.delta") => { - if let Some(delta) = value.get("delta").and_then(Value::as_str) { - output_text.push_str(delta); - } - } - Some("response.output_item.done") => { - if let Some(item) = value.get("item") { - output_items.push(item.clone()); - } - } - Some("response.completed") => { - terminal_response = value.get("response").cloned(); - terminal_incomplete = false; - } - Some("response.incomplete") => { - terminal_response = value.get("response").cloned(); - terminal_incomplete = true; - } - Some(event_name @ ("error" | "response.failed")) if failed_response.is_none() => { - failed_response = Some(ResponsesSseFailure::from_event(event_name, &value)); - } - _ => {} - } - } - - if let Some(failure) = failed_response { - return Ok(CollectedResponse::Failed(failure)); - } - let Some(mut response) = terminal_response else { - return Ok(CollectedResponse::MissingCompleted); - }; - if terminal_incomplete { - reconvert_completed_response_tuple_values(&mut response, tuple_schema); - return Ok(CollectedResponse::Incomplete(response)); - } - if is_empty_response(&response, &output_text, &output_items) { - return Ok(CollectedResponse::Empty); - } - - reconvert_completed_response_tuple_values(&mut response, tuple_schema); - Ok(CollectedResponse::Completed(response)) -} - /// 从 Codex Responses SSE 中提取完成响应元数据。 pub fn completed_response_metadata( body: &str, @@ -437,21 +387,6 @@ fn reconvert_output_text_delta_tuple_values(data: &mut Value, tuple_schema: &Val data["delta"] = Value::String(reconverted.to_string()); } -fn is_empty_response(response: &Value, output_text: &str, output_items: &[Value]) -> bool { - if !output_text.trim().is_empty() || !output_items.is_empty() { - return false; - } - if response.get("status").and_then(Value::as_str) == Some("incomplete") { - return false; - } - - response - .pointer("/usage/output_tokens") - .and_then(Value::as_u64) - .unwrap_or_default() - == 0 -} - fn reconvert_completed_response_tuple_values(response: &mut Value, tuple_schema: Option<&Value>) { let Some(tuple_schema) = tuple_schema else { return; @@ -515,6 +450,28 @@ fn failure_code(value: &Value) -> Option { .map(ToString::to_string) } +fn failure_type(value: &Value) -> Option { + failure_error(value) + .and_then(|error| error.get("type")) + .and_then(Value::as_str) + .filter(|error_type| !error_type.trim().is_empty()) + .map(ToString::to_string) +} + +fn failure_explicit_status_code(value: &Value) -> Option { + value + .get("status") + .or_else(|| value.get("status_code")) + .and_then(Value::as_u64) + .and_then(|status| u16::try_from(status).ok()) +} + +fn failure_error(value: &Value) -> Option<&Value> { + value + .pointer("/response/error") + .or_else(|| value.get("error")) +} + impl CodexResponsesRequest { /// 从客户端原始 Responses JSON object 构造上游请求。 /// diff --git a/backend/src/upstream/openai/protocol/sse.rs b/backend/src/upstream/openai/protocol/sse.rs index 53bb17464..54238ec88 100644 --- a/backend/src/upstream/openai/protocol/sse.rs +++ b/backend/src/upstream/openai/protocol/sse.rs @@ -22,6 +22,16 @@ pub const MAX_SSE_EVENT_BUFFER_BYTES: usize = 64 * 1024 * 1024; /// SSE 流结束标记帧。 pub const DONE_SSE_FRAME: &str = "data: [DONE]\n\n"; +/// 判断一个完整 SSE frame 是否为传输层 `[DONE]` 控制帧。 +pub(crate) fn sse_frame_is_done(frame: &str) -> bool { + let mut data = frame.lines().filter_map(|raw_line| { + let line = raw_line.strip_suffix('\r').unwrap_or(raw_line); + let (field, value) = split_sse_field(line); + (field == "data").then_some(value) + }); + matches!((data.next(), data.next()), (Some("[DONE]"), None)) +} + /// 用于处理任意分块边界的增量 SSE 解码器。 #[derive(Debug, Default)] pub struct SseEventDecoder { @@ -78,6 +88,17 @@ pub fn response_failed_sse_event_with_id( code: &str, message: &str, ) -> String { + let data = response_failed_sse_data_with_id(response_id, error_type, code, message); + encode_sse_event("response.failed", &data.to_string()) +} + +/// 构造 OpenAI Responses `response.failed` 的 JSON 数据。 +pub fn response_failed_sse_data_with_id( + response_id: Option<&str>, + error_type: &str, + code: &str, + message: &str, +) -> serde_json::Value { let error = json!({ "type": error_type, "code": code, @@ -89,7 +110,7 @@ pub fn response_failed_sse_event_with_id( || format!("resp_proxy_{}", uuid::Uuid::new_v4().simple()), ToString::to_string, ); - let data = json!({ + json!({ "type": "response.failed", "response": { "id": response_id, @@ -97,8 +118,7 @@ pub fn response_failed_sse_event_with_id( "error": error, }, "error": error, - }); - encode_sse_event("response.failed", &data.to_string()) + }) } /// 判断 SSE 文本是否已经包含结束标记。 diff --git a/backend/src/upstream/openai/protocol/websocket.rs b/backend/src/upstream/openai/protocol/websocket.rs index 5412eda28..bedce7e82 100644 --- a/backend/src/upstream/openai/protocol/websocket.rs +++ b/backend/src/upstream/openai/protocol/websocket.rs @@ -181,10 +181,6 @@ pub fn websocket_response_completed_id(raw: &str) -> Result, Stri Ok(Some(completed.id)) } -#[path = "websocket_errors.rs"] -mod errors; -pub use errors::*; - /// 判断 WebSocket 事件是否会结束当前响应流。 pub fn is_terminal_websocket_event(event: &str) -> bool { matches!( diff --git a/backend/src/upstream/openai/protocol/websocket_errors.rs b/backend/src/upstream/openai/protocol/websocket_errors.rs deleted file mode 100644 index b6f3777e6..000000000 --- a/backend/src/upstream/openai/protocol/websocket_errors.rs +++ /dev/null @@ -1,154 +0,0 @@ -//! Responses WebSocket 错误帧解析。 - -use serde_json::{Map, Value}; - -use super::json_value_as_string; - -/// 从 WebSocket 错误帧推导出的上游错误分类。 -#[derive(Debug, Clone, PartialEq, Eq)] -pub struct ClassifiedWebSocketError { - /// 应映射给 HTTP 层的状态码。 - pub status_code: u16, - /// 上游错误类型。 - pub error_type: Option, - /// 上游错误码。 - pub code: Option, - /// 上游错误消息。 - pub message: Option, - /// 上游错误参数。 - pub param: Option, - /// 错误帧携带的安全响应头。 - pub headers: Vec<(String, String)>, -} - -/// 将 WebSocket 错误帧映射为结构化上游错误。 -pub fn classify_websocket_error_frame(raw: &str) -> Option { - let value = serde_json::from_str::(raw).ok()?; - let event_type = value.get("type").and_then(Value::as_str)?; - if event_type != "error" && event_type != "response.failed" { - return None; - } - - let error = websocket_error_value(&value); - let error_type = error - .and_then(|error| error.get("type")) - .and_then(Value::as_str) - .map(ToString::to_string); - let code = error - .and_then(|error| error.get("code")) - .and_then(Value::as_str) - .map(ToString::to_string); - let message = error - .and_then(|error| error.get("message")) - .and_then(Value::as_str) - .map(ToString::to_string); - let param = error - .and_then(|error| error.get("param")) - .and_then(Value::as_str) - .map(ToString::to_string); - let explicit_status = - explicit_error_status_code(&value).filter(|status_code| !(200..=299).contains(status_code)); - if explicit_error_status_code(&value).is_some() && explicit_status.is_none() { - return None; - } - let normalized_code = code.as_deref().map(str::to_ascii_lowercase); - let normalized_type = error_type.as_deref().map(str::to_ascii_lowercase); - let status_code = explicit_status - .or_else(|| normalized_code.as_deref().and_then(error_status_code)) - .or_else(|| normalized_type.as_deref().and_then(error_type_status_code))?; - - Some(ClassifiedWebSocketError { - status_code, - error_type, - code, - message, - param, - headers: websocket_error_headers(&value), - }) -} - -fn websocket_error_value(value: &Value) -> Option<&Map> { - value - .pointer("/response/error") - .or_else(|| value.get("error")) - .and_then(Value::as_object) -} - -fn explicit_error_status_code(value: &Value) -> Option { - value - .get("status") - .or_else(|| value.get("status_code")) - .and_then(Value::as_u64) - .and_then(|status| u16::try_from(status).ok()) -} - -fn error_status_code(code: &str) -> Option { - match code { - "usage_limit_reached" | "rate_limit_exceeded" | "rate_limit_reached" => Some(429), - "quota_exhausted" | "payment_required" => Some(402), - "unauthorized" | "token_invalid" | "token_expired" | "account_deactivated" => Some(401), - "forbidden" | "account_banned" | "banned" => Some(403), - "previous_response_not_found" - | "invalid_encrypted_content" - | "missing_tool_output" - | "no_tool_output" => Some(400), - "websocket_connection_limit_reached" => Some(503), - _ => None, - } -} - -fn error_type_status_code(error_type: &str) -> Option { - match error_type { - "invalid_request_error" => Some(400), - "authentication_error" => Some(401), - "permission_error" => Some(403), - "rate_limit_error" => Some(429), - "server_error" | "api_error" => Some(500), - _ => None, - } -} - -fn websocket_error_headers(value: &Value) -> Vec<(String, String)> { - value - .get("headers") - .and_then(Value::as_object) - .into_iter() - .flatten() - .filter_map(|(name, value)| json_value_as_string(value).map(|value| (name.clone(), value))) - .collect() -} - -/// 从包裹在 WebSocket `error.headers` 中的 retry-after 值提取秒数。 -pub fn retry_after_seconds_from_wrapped_error_headers(raw: &str) -> Option { - let value = serde_json::from_str::(raw).ok()?; - let header_retry_after = value - .get("headers") - .and_then(Value::as_object) - .and_then(|headers| { - headers.iter().find_map(|(name, value)| { - if name.eq_ignore_ascii_case("retry-after") { - retry_after_seconds_from_json_header_value(value) - } else { - None - } - }) - }); - header_retry_after.or_else(|| { - value - .pointer("/error/retry_after_seconds") - .or_else(|| value.get("retry_after_seconds")) - .and_then(Value::as_u64) - }) -} - -fn retry_after_seconds_from_json_header_value(value: &Value) -> Option { - let seconds = match value { - Value::String(value) => value.trim().parse::().ok()?, - Value::Number(value) => value.as_u64()?, - Value::Array(values) => values - .first() - .and_then(retry_after_seconds_from_json_header_value)?, - _ => return None, - }; - (seconds > 0).then_some(seconds) -} diff --git a/backend/src/upstream/openai/transport/client.rs b/backend/src/upstream/openai/transport/client.rs index a031b0dc4..f361196b9 100644 --- a/backend/src/upstream/openai/transport/client.rs +++ b/backend/src/upstream/openai/transport/client.rs @@ -162,35 +162,13 @@ pub fn is_deactivated_workspace_error_body(value: &str) -> bool { }) } -/// 判断错误正文是否带有 OpenAI `cyber_policy` 标记。 -pub fn is_cyber_policy_error_body(value: &str) -> bool { - serde_json::from_str::(value).is_ok_and(|value| { - value - .pointer("/response/error/code") - .or_else(|| value.pointer("/error/code")) - .or_else(|| value.get("code")) - .and_then(Value::as_str) - .is_some_and(|code| code.trim().eq_ignore_ascii_case("cyber_policy")) - }) -} - -/// 判断 Codex 上游错误是否为会话级 `cyber_policy`。 -pub fn is_cyber_policy_upstream_error(error: &CodexClientError) -> bool { - matches!( - error, - CodexClientError::Upstream { status, body, .. } - if status.is_client_error() && is_cyber_policy_error_body(body) - ) -} - /// 判断 Codex 上游错误是否表示账号已封禁或停用。 pub fn is_banned_upstream_error(error: &CodexClientError) -> bool { matches!( error, CodexClientError::Upstream { status, body, .. } if (status.as_u16() == 403 - && !is_html_error_body(body) - && !is_cyber_policy_upstream_error(error)) + && !is_html_error_body(body)) || (status.as_u16() == 402 && is_deactivated_workspace_error_body(body)) ) } diff --git a/backend/src/upstream/openai/transport/mod.rs b/backend/src/upstream/openai/transport/mod.rs index cffae19b3..0adce9f87 100644 --- a/backend/src/upstream/openai/transport/mod.rs +++ b/backend/src/upstream/openai/transport/mod.rs @@ -18,8 +18,7 @@ pub use self::{ CodexModelCatalogClient, CodexModelCatalogClientError, CodexModelCatalogRequest, CodexRateLimitHeaderUpdates, CodexRequestContext, CodexResponseMetadata, CodexTurnStateUpdate, backend_transport_for_response_request, build_reqwest_client, - is_banned_auth_signal, is_banned_upstream_error, is_cyber_policy_error_body, - is_cyber_policy_upstream_error, is_deactivated_workspace_error_body, + is_banned_auth_signal, is_banned_upstream_error, is_deactivated_workspace_error_body, }, diagnostics::CodexUpstreamDiagnostics, endpoints::{ diff --git a/backend/src/upstream/openai/transport/websocket.rs b/backend/src/upstream/openai/transport/websocket.rs index 30eaab2c6..aa67aa455 100644 --- a/backend/src/upstream/openai/transport/websocket.rs +++ b/backend/src/upstream/openai/transport/websocket.rs @@ -1,7 +1,6 @@ //! Codex WebSocket 连接建立（关键函数）。 use std::{ - collections::VecDeque, fmt, io, path::{Path, PathBuf}, pin::Pin, @@ -27,16 +26,14 @@ use tungstenite::{ use crate::infra::time::china_filename_timestamp_millis; use crate::upstream::openai::protocol::events::{self, TokenUsage}; use crate::upstream::openai::protocol::responses::{ - CodexResponsesRequest, PreviousResponseScope, StreamCommitPolicy, - response_body_has_first_output, + CodexResponsesRequest, PreviousResponseScope, response_body_has_first_output, }; use crate::upstream::openai::protocol::sse::SseError; use crate::upstream::openai::protocol::websocket::{ - ClassifiedWebSocketError, OpeningAuditHeader, OpeningAuditSnapshot, WebSocketAuditArtifact, - classify_websocket_error_frame, is_terminal_websocket_event, - retry_after_seconds_from_wrapped_error_headers, websocket_event_to_sse_frame, - websocket_event_type, websocket_metadata_headers, websocket_metadata_turn_state, - websocket_response_completed_id, websocket_response_create_payload_text, + OpeningAuditHeader, OpeningAuditSnapshot, WebSocketAuditArtifact, is_terminal_websocket_event, + websocket_event_to_sse_frame, websocket_event_type, websocket_metadata_headers, + websocket_metadata_turn_state, websocket_response_completed_id, + websocket_response_create_payload_text, }; use super::websocket_pool::{ @@ -105,7 +102,6 @@ pub struct CodexWebSocketRequest { pub(super) connection: CodexWebSocketConnection, pub(super) payload_text: String, pub(super) continuation: WebSocketContinuationRequirement, - pub(super) stream_commit_policy: StreamCommitPolicy, } /// 当前 WebSocket 请求对 previous response 状态的要求。 @@ -189,8 +185,7 @@ pub use frames::{ }; use frames::{ WebSocketStreamPoolReturn, WebSocketTerminalKind, audit_header_value, - collect_websocket_response, prefetch_stream_frames_until_output_or_terminal, - reusable_websocket_metadata, reused_stream_prefetch_error, stream_websocket_response, + collect_websocket_response, reusable_websocket_metadata, stream_websocket_response, websocket_audit_file_name, websocket_connection_metadata, }; @@ -209,10 +204,6 @@ impl CodexWebSocketRequest { pub fn continuation(&self) -> &WebSocketContinuationRequirement { &self.continuation } - - pub fn stream_commit_policy(&self) -> StreamCommitPolicy { - self.stream_commit_policy - } } impl CodexWebSocketConnection { @@ -260,7 +251,6 @@ impl CodexWebSocketConnection { connection: Self::responses(base_url, websocket_key, business_headers), payload_text: websocket_response_create_payload_text(request)?, continuation: WebSocketContinuationRequirement::from_request(request), - stream_commit_policy: request.stream_commit_policy, }) } @@ -463,33 +453,16 @@ async fn execute_fresh_response_create_request_stream( request: &CodexWebSocketRequest, initial_event_timeout: Option, ) -> Result { - let (mut websocket, response) = + let (websocket, response) = connect_pumped_websocket(request.connection(), PumpKeepalive::disabled()).await?; send_websocket_request(&websocket, request.payload_text()).await?; - let mut metadata = websocket_connection_metadata(&response); - let prefetched_frames = if should_prefetch_until_output_or_terminal(request) { - match prefetch_stream_frames_until_output_or_terminal( - &mut websocket, - &mut metadata, - initial_event_timeout, - ) - .await - { - Ok(prefetched_frames) => prefetched_frames, - Err(error) => { - websocket.close().await; - return Err(error); - } - } - } else { - Vec::new() - }; + let metadata = websocket_connection_metadata(&response); Ok(stream_websocket_response( websocket, metadata, None, - prefetched_frames, + false, initial_event_timeout, )) } @@ -761,7 +734,7 @@ async fn execute_fresh_pooled_response_create_request_stream( pool: &CodexWebSocketPool, lease: WebSocketPoolLease, ) -> Result { - let (mut websocket, response) = + let (websocket, response) = match connect_pumped_websocket(request.connection(), pool.keepalive()).await { Ok(connected) => connected, Err(error) => { @@ -775,25 +748,7 @@ async fn execute_fresh_pooled_response_create_request_stream( } let now = tokio::time::Instant::now(); - let mut metadata = websocket_connection_metadata(&response); - let prefetched_frames = if should_prefetch_until_output_or_terminal(request) { - match prefetch_stream_frames_until_output_or_terminal( - &mut websocket, - &mut metadata, - pool.initial_event_timeout(), - ) - .await - { - Ok(prefetched_frames) => prefetched_frames, - Err(error) => { - websocket.close().await; - lease.discard().await; - return Err(error); - } - } - } else { - Vec::new() - }; + let metadata = websocket_connection_metadata(&response); Ok(stream_websocket_response( websocket, metadata, @@ -802,7 +757,7 @@ async fn execute_fresh_pooled_response_create_request_stream( created_at: now, continuation: WebSocketContinuationState::default(), }), - prefetched_frames, + false, pool.initial_event_timeout(), )) } @@ -867,8 +822,8 @@ async fn execute_pooled_response_create_request_stream( lease: WebSocketPoolLease, connection: PooledWebSocketConnection, ) -> Result { - let mut websocket = connection.websocket; - let mut metadata = reusable_websocket_metadata(connection.metadata); + let websocket = connection.websocket; + let metadata = reusable_websocket_metadata(connection.metadata); let created_at = connection.created_at; let continuation = connection.continuation; if let Err(error) = send_websocket_request(&websocket, request.payload_text()).await { @@ -880,25 +835,6 @@ async fn execute_pooled_response_create_request_stream( ); } - let prefetched_frames = if should_prefetch_until_output_or_terminal(request) { - match prefetch_stream_frames_until_output_or_terminal( - &mut websocket, - &mut metadata, - pool.initial_event_timeout(), - ) - .await - { - Ok(prefetched_frames) => prefetched_frames, - Err(error) => { - lease.discard().await; - websocket.close().await; - return Err(reused_stream_prefetch_error(error)); - } - } - } else { - Vec::new() - }; - Ok(stream_websocket_response( websocket, metadata, @@ -907,14 +843,7 @@ async fn execute_pooled_response_create_request_stream( created_at, continuation, }), - prefetched_frames, + true, pool.initial_event_timeout(), )) } - -fn should_prefetch_until_output_or_terminal(request: &CodexWebSocketRequest) -> bool { - matches!( - request.stream_commit_policy(), - StreamCommitPolicy::UntilOutputOrTerminal - ) -} diff --git a/backend/src/upstream/openai/transport/websocket_frames.rs b/backend/src/upstream/openai/transport/websocket_frames.rs index 8825a404a..15f426103 100644 --- a/backend/src/upstream/openai/transport/websocket_frames.rs +++ b/backend/src/upstream/openai/transport/websocket_frames.rs @@ -184,27 +184,6 @@ impl CodexWebSocketExchangeError { diagnostics, })) } - - pub(super) fn upstream_frame( - classified: ClassifiedWebSocketError, - retry_after_seconds: Option, - body: String, - ) -> Self { - let diagnostics = - CodexUpstreamDiagnostics::from_pairs(Some(classified.status_code), &classified.headers); - Self::Upstream(Box::new(CodexWebSocketUpstreamError { - status_code: classified.status_code, - retry_after_seconds, - body, - error_type: classified.error_type, - code: classified.code, - message: classified.message, - param: classified.param, - headers: classified.headers, - set_cookie_headers: Vec::new(), - diagnostics, - })) - } } /// 执行一次 prepared Responses WebSocket 请求并聚合为 SSE。 @@ -275,15 +254,6 @@ pub(super) async fn collect_websocket_response( _ => continue, }; let raw = text.to_string(); - if let Some(classified) = classify_websocket_error_frame(&raw) { - let retry_after_seconds = retry_after_seconds_from_wrapped_error_headers(&raw) - .or_else(|| events::retry_after_seconds_from_body(&raw)); - return Err(CodexWebSocketExchangeError::upstream_frame( - classified, - retry_after_seconds, - raw, - )); - } if let Some(headers) = websocket_rate_limit_event_headers(&raw) { metadata.rate_limit_headers.extend(headers); continue; @@ -352,76 +322,7 @@ pub(super) async fn collect_websocket_response( Err(CodexWebSocketExchangeError::ClosedBeforeTerminal) } -pub(super) async fn prefetch_stream_frames_until_output_or_terminal( - websocket: &mut PumpedWebSocket, - metadata: &mut CodexWebSocketConnectionMetadata, - initial_event_timeout: Option, -) -> Result, CodexWebSocketExchangeError> { - let mut prefetched_frames = Vec::new(); - let mut saw_upstream_activity = false; - loop { - let receive_timeout = receive_idle_timeout(saw_upstream_activity, initial_event_timeout); - let message = match next_websocket_message(websocket, receive_timeout).await { - Ok(Some(message)) => message, - Ok(None) => return Err(CodexWebSocketExchangeError::ClosedBeforeTerminal), - Err(CodexWebSocketExchangeError::ReceiveIdleTimeout { timeout }) - if !saw_upstream_activity => - { - return Err(CodexWebSocketExchangeError::InitialEventTimeout { timeout }); - } - Err(error) => return Err(error), - }; - - let raw = match message { - Message::Text(text) => { - saw_upstream_activity = true; - text.to_string() - } - Message::Binary(_) => return Err(CodexWebSocketExchangeError::UnexpectedBinaryEvent), - Message::Close(_) => return Err(CodexWebSocketExchangeError::ClosedBeforeTerminal), - _ => continue, - }; - - if let Some(classified) = classify_websocket_error_frame(&raw) { - let retry_after_seconds = retry_after_seconds_from_wrapped_error_headers(&raw) - .or_else(|| events::retry_after_seconds_from_body(&raw)); - return Err(CodexWebSocketExchangeError::upstream_frame( - classified, - retry_after_seconds, - raw, - )); - } - if websocket_event_type(&raw).as_deref() == Some("response.completed") { - websocket_response_completed_id(&raw).map_err(|message| { - CodexWebSocketExchangeError::InvalidCompletedResponse { message } - })?; - } - if let Some(headers) = websocket_rate_limit_event_headers(&raw) { - metadata.rate_limit_headers.extend(headers); - continue; - } - update_websocket_response_metadata(metadata, &raw); - if let Some(metadata_turn_state) = websocket_metadata_turn_state(&raw) { - metadata.turn_state = Some(metadata_turn_state); - continue; - } - - let event = websocket_event_type(&raw); - let Some(frame) = websocket_event_to_sse_frame(&raw) else { - continue; - }; - let has_first_output = response_body_has_first_output(frame.as_bytes()); - let terminal = event.as_deref().is_some_and(is_terminal_websocket_event); - prefetched_frames.push(raw); - if has_first_output || terminal { - return Ok(prefetched_frames); - } - } -} - -pub(super) fn reused_stream_prefetch_error( - error: CodexWebSocketExchangeError, -) -> CodexWebSocketExchangeError { +fn reused_stream_receive_error(error: CodexWebSocketExchangeError) -> CodexWebSocketExchangeError { match error { CodexWebSocketExchangeError::ClosedBeforeTerminal | CodexWebSocketExchangeError::ReceiveIdleTimeout { .. } @@ -481,7 +382,6 @@ enum StreamWebSocketDiscardReason { InvalidCompletedResponse, UnexpectedBinaryEvent, UpstreamClosed, - UpstreamError, UpstreamReceiveFailed, } @@ -495,7 +395,6 @@ impl StreamWebSocketDiscardReason { Self::InvalidCompletedResponse => "invalid_completed_response", Self::UnexpectedBinaryEvent => "unexpected_binary_event", Self::UpstreamClosed => "upstream_closed", - Self::UpstreamError => "upstream_error", Self::UpstreamReceiveFailed => "upstream_receive_failed", } } @@ -505,7 +404,7 @@ pub(super) fn stream_websocket_response( websocket: PumpedWebSocket, metadata: CodexWebSocketConnectionMetadata, pool_return: Option, - prefetched_frames: Vec, + reused_connection: bool, initial_event_timeout: Option, ) -> CodexWebSocketStreamingExchange { let websocket_connection_id = websocket.connection_id(); @@ -520,7 +419,7 @@ pub(super) fn stream_websocket_response( websocket, metadata, pool_return, - prefetched_frames, + reused_connection, initial_event_timeout, rate_limit_header_updates: rate_limit_header_updates_for_task, turn_state_update: turn_state_update_for_task, @@ -551,7 +450,7 @@ struct WebSocketStreamForwardState { websocket: PumpedWebSocket, metadata: CodexWebSocketConnectionMetadata, pool_return: Option, - prefetched_frames: Vec, + reused_connection: bool, initial_event_timeout: Option, rate_limit_header_updates: CodexWebSocketRateLimitHeaderUpdates, turn_state_update: CodexWebSocketTurnStateUpdate, @@ -563,7 +462,7 @@ async fn forward_websocket_response_stream(state: WebSocketStreamForwardState) { mut websocket, mut metadata, pool_return, - prefetched_frames, + reused_connection, initial_event_timeout, rate_limit_header_updates, turn_state_update, @@ -574,95 +473,81 @@ async fn forward_websocket_response_stream(state: WebSocketStreamForwardState) { .as_mut() .map(|pool_return| std::mem::take(&mut pool_return.continuation)) .unwrap_or_default(); - let mut saw_upstream_activity = !prefetched_frames.is_empty(); - let mut prefetched_frames: VecDeque = prefetched_frames.into(); + let mut saw_upstream_activity = false; loop { - let raw = if let Some(text) = prefetched_frames.pop_front() { - text - } else { - let message = tokio::select! { - biased; - // 客户端断开下游 SSE 流：立即丢弃连接并释放池 slot， - // 不再傻等上游 idle 超时（否则同会话后续请求会一直 bypass/busy）。 - () = tx.closed() => { - discard_stream_websocket( - websocket, - pool_return, - StreamWebSocketDiscardReason::ClientCancelled, - ).await; - return; - } - message = next_websocket_message( - &mut websocket, - receive_idle_timeout(saw_upstream_activity, initial_event_timeout), - ) => message, - }; - let message = match message { - Ok(message) => message, - Err(error) => { - discard_stream_websocket( - websocket, - pool_return, - StreamWebSocketDiscardReason::UpstreamReceiveFailed, - ) - .await; - let _ = tx.send(Err(error)).await; - return; - } - }; - let Some(message) = message else { - break; - }; - match message { - Message::Text(text) => { - saw_upstream_activity = true; - text.to_string() - } - Message::Binary(_) => { - discard_stream_websocket( - websocket, - pool_return, - StreamWebSocketDiscardReason::UnexpectedBinaryEvent, - ) - .await; - let _ = tx - .send(Err(CodexWebSocketExchangeError::UnexpectedBinaryEvent)) - .await; - return; - } - Message::Close(_) => { - discard_stream_websocket( - websocket, - pool_return, - StreamWebSocketDiscardReason::UpstreamClosed, - ) - .await; - let _ = tx - .send(Err(CodexWebSocketExchangeError::ClosedBeforeTerminal)) - .await; - return; - } - _ => continue, + let message = tokio::select! { + biased; + // 客户端断开下游 SSE 流：立即丢弃连接并释放池 slot， + // 不再傻等上游 idle 超时（否则同会话后续请求会一直 bypass/busy）。 + () = tx.closed() => { + discard_stream_websocket( + websocket, + pool_return, + StreamWebSocketDiscardReason::ClientCancelled, + ).await; + return; + } + message = next_websocket_message( + &mut websocket, + receive_idle_timeout(saw_upstream_activity, initial_event_timeout), + ) => message, + }; + let message = match message { + Ok(message) => message, + Err(error) => { + discard_stream_websocket( + websocket, + pool_return, + StreamWebSocketDiscardReason::UpstreamReceiveFailed, + ) + .await; + let error = if reused_connection { + reused_stream_receive_error(error) + } else { + error + }; + let _ = tx.send(Err(error)).await; + return; } }; - if let Some(classified) = classify_websocket_error_frame(&raw) { - discard_stream_websocket( - websocket, - pool_return, - StreamWebSocketDiscardReason::UpstreamError, - ) - .await; - let retry_after_seconds = retry_after_seconds_from_wrapped_error_headers(&raw) - .or_else(|| events::retry_after_seconds_from_body(&raw)); - let _ = tx - .send(Err(CodexWebSocketExchangeError::upstream_frame( - classified, - retry_after_seconds, - raw, - ))) + let Some(message) = message else { + break; + }; + let raw = match message { + Message::Text(text) => { + saw_upstream_activity = true; + text.to_string() + } + Message::Binary(_) => { + discard_stream_websocket( + websocket, + pool_return, + StreamWebSocketDiscardReason::UnexpectedBinaryEvent, + ) .await; - return; - } + let _ = tx + .send(Err(CodexWebSocketExchangeError::UnexpectedBinaryEvent)) + .await; + return; + } + Message::Close(_) => { + discard_stream_websocket( + websocket, + pool_return, + StreamWebSocketDiscardReason::UpstreamClosed, + ) + .await; + let error = CodexWebSocketExchangeError::ClosedBeforeTerminal; + let error = if reused_connection { + reused_stream_receive_error(error) + } else { + error + }; + let _ = tx.send(Err(error)).await; + return; + } + _ => continue, + }; if let Some(headers) = websocket_rate_limit_event_headers(&raw) { metadata.rate_limit_headers.extend(headers.iter().cloned()); rate_limit_header_updates.lock().await.extend(headers); @@ -742,9 +627,13 @@ async fn forward_websocket_response_stream(state: WebSocketStreamForwardState) { StreamWebSocketDiscardReason::UpstreamClosed, ) .await; - let _ = tx - .send(Err(CodexWebSocketExchangeError::ClosedBeforeTerminal)) - .await; + let error = CodexWebSocketExchangeError::ClosedBeforeTerminal; + let error = if reused_connection { + reused_stream_receive_error(error) + } else { + error + }; + let _ = tx.send(Err(error)).await; } fn receive_idle_timeout( diff --git a/backend/tests/api/client/responses/websocket.rs b/backend/tests/api/client/responses/websocket.rs index 95dea7360..d60dfd485 100644 --- a/backend/tests/api/client/responses/websocket.rs +++ b/backend/tests/api/client/responses/websocket.rs @@ -1,6 +1,7 @@ -use std::{net::SocketAddr, time::Duration}; +use std::{fs, net::SocketAddr, path::Path, time::Duration}; use axum::Router; +use codex_proxy_rs::fleet::account::AccountStatus; use futures::{SinkExt, StreamExt}; use serde_json::{Value, json}; use tokio::{net::TcpListener, task::JoinHandle, time::timeout}; @@ -13,7 +14,9 @@ use tokio_tungstenite::{ }, }; -use crate::dispatch::service::accept_async as accept_upstream_websocket; +use crate::dispatch::service::{ + accept_async as accept_upstream_websocket, accept_websocket_with_authorization, +}; struct TestServer { address: SocketAddr, @@ -26,6 +29,57 @@ impl Drop for TestServer { } } +#[test] +fn responses_websocket_replay_policy_should_be_owned_by_history_controller() { + let manifest = Path::new(env!("CARGO_MANIFEST_DIR")); + let api = fs::read_to_string(manifest.join("src/api/client/responses/websocket.rs")).unwrap(); + let service = fs::read_to_string(manifest.join("src/dispatch/service.rs")).unwrap(); + let history = fs::read_to_string(manifest.join("src/dispatch/controllers/history.rs")).unwrap(); + let protocol = + fs::read_to_string(manifest.join("src/upstream/openai/protocol/responses.rs")).unwrap(); + + assert!(api.contains(".prepare_connection_replay(")); + assert!(api.contains(".commit_connection_replay(")); + for forbidden in [ + "ConnectionReplayState", + "PendingReplayUpdate", + "sanitize_replay_items", + ".previous_response_id()", + "dispatch::controllers", + ] { + assert!( + !api.contains(forbidden), + "Responses WebSocket API must not own replay policy {forbidden}" + ); + } + for owned in [ + "enum ConnectionReplayUpdate", + "ConnectionReplayUpdate::Replace", + "ConnectionReplayUpdate::Append", + "ConnectionReplayUpdate::Unavailable", + "sanitize_replay_items", + "snapshot.last_response_id.as_deref() == Some(previous_response_id)", + ] { + assert!( + history.contains(owned), + "HistoryController owner must contain {owned}" + ); + } + for boundary in [ + "pub(crate) struct ConnectionReplaySnapshot", + "pub(crate) struct ConnectionReplayPlan", + "pub(crate) struct ConnectionTranscriptFacts", + "fn prepare_connection_replay(", + "fn commit_connection_replay(", + ] { + assert!( + service.contains(boundary), + "ResponseDispatchService must expose narrow replay boundary {boundary}" + ); + } + assert!(!protocol.contains("sanitize_replay_items")); +} + #[tokio::test] async fn responses_websocket_should_reject_missing_client_api_key_during_handshake() { let (app, _api_key, _dir, _connections) = super::test_app_with_client_api_key().await; @@ -62,6 +116,56 @@ async fn responses_websocket_should_return_official_error_frame_when_accounts_ar ); } +#[tokio::test] +async fn responses_websocket_should_preserve_upstream_client_failure() { + let upstream_listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let upstream_base_url = format!("http://{}", upstream_listener.local_addr().unwrap()); + let upstream = tokio::spawn(async move { + let (stream, _) = upstream_listener.accept().await.unwrap(); + let mut websocket = accept_upstream_websocket(stream).await.unwrap(); + let _request = receive_upstream_request(&mut websocket).await; + websocket + .send(Message::Text( + json!({ + "type": "response.failed", + "response": { + "id": "resp_server_overloaded", + "status": "failed", + "error": { + "code": "server_overloaded", + "message": "upstream is temporarily overloaded", + }, + }, + }) + .to_string() + .into(), + )) + .await + .unwrap(); + websocket.close(None).await.unwrap(); + }); + let (app, api_key, _dir) = + crate::dispatch::service::test_app_with_account(upstream_base_url).await; + let server = spawn_app(app).await; + let mut websocket = connect_responses_websocket(server.address, &api_key).await; + + websocket + .send(Message::Text(response_create_payload("overload").into())) + .await + .unwrap(); + let events = receive_response(&mut websocket).await; + upstream.await.unwrap(); + let failure = events.last().unwrap(); + + assert_eq!(failure["type"], "error"); + assert_eq!(failure["status"], 503); + assert_eq!(failure["error"]["code"], "server_overloaded"); + assert_eq!( + failure["error"]["message"], + "upstream is temporarily overloaded" + ); +} + #[tokio::test] async fn responses_websocket_should_disconnect_when_connection_draining_starts() { let (app, api_key, _dir, connection_drain) = super::test_app_with_client_api_key().await; @@ -121,6 +225,91 @@ async fn responses_websocket_should_forward_multiple_response_create_requests_on assert_eq!(upstream_requests[1]["type"], "response.create"); } +#[tokio::test] +async fn responses_websocket_cyber_terminal_should_change_the_immediate_next_request() { + let upstream_listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let upstream_base_url = format!("http://{}", upstream_listener.local_addr().unwrap()); + let upstream = tokio::spawn(async move { + let (primary_stream, _) = upstream_listener.accept().await.unwrap(); + let mut primary = + accept_websocket_with_authorization(primary_stream, "Bearer access-primary").await; + let primary_request = receive_upstream_request(&mut primary).await; + for event in [ + json!({ + "type": "response.created", + "response": {"id": "resp_downstream_ws_cyber", "status": "in_progress"}, + }), + json!({ + "type": "response.output_text.delta", + "delta": "partial downstream WebSocket output", + }), + json!({ + "type": "response.failed", + "response": { + "id": "resp_downstream_ws_cyber", + "status": "failed", + "error": { + "code": "cyber_policy", + "message": "This request has been flagged for possible cybersecurity risk.", + }, + }, + }), + ] { + primary + .send(Message::Text(event.to_string().into())) + .await + .unwrap(); + } + primary.close(None).await.unwrap(); + + let (secondary_stream, _) = upstream_listener.accept().await.unwrap(); + let mut secondary = + accept_websocket_with_authorization(secondary_stream, "Bearer access-secondary").await; + let secondary_request = receive_upstream_request(&mut secondary).await; + send_upstream_response(&mut secondary, "resp_downstream_ws_rotated", "rotated").await; + (primary_request, secondary_request) + }); + let (app, state, api_key, _pool, _dir) = + crate::dispatch::service::test_app_with_two_accounts_and_state(upstream_base_url).await; + assert!( + state + .services + .account_pool + .set_status("acct_secondary", AccountStatus::Disabled) + .await + ); + let server = spawn_app(app).await; + let mut websocket = connect_responses_websocket(server.address, &api_key).await; + + websocket + .send(Message::Text(response_create_payload("first").into())) + .await + .unwrap(); + let first = receive_response(&mut websocket).await; + assert_eq!(first.last().unwrap()["type"], "response.failed"); + + assert!( + state + .services + .account_pool + .set_status("acct_secondary", AccountStatus::Active) + .await + ); + websocket + .send(Message::Text(response_create_payload("second").into())) + .await + .unwrap(); + let second = receive_response(&mut websocket).await; + let (primary_request, secondary_request) = upstream.await.unwrap(); + + assert_eq!( + second.last().unwrap()["response"]["id"], + "resp_downstream_ws_rotated" + ); + assert_eq!(primary_request["input"][0]["content"], "first"); + assert_eq!(secondary_request["input"][0]["content"], "second"); +} + #[tokio::test] async fn responses_websocket_should_replay_connection_history_after_previous_response_not_found() { let upstream_listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); @@ -197,6 +386,97 @@ async fn responses_websocket_should_replay_connection_history_after_previous_res assert!(!replayed_json.contains("assistant_output_resp_first")); } +#[tokio::test] +async fn responses_websocket_history_retry_should_not_fallback_when_owner_account_becomes_unavailable() + { + let upstream_listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let upstream_base_url = format!("http://{}", upstream_listener.local_addr().unwrap()); + let (continued_tx, continued_rx) = tokio::sync::oneshot::channel(); + let (release_failure_tx, release_failure_rx) = tokio::sync::oneshot::channel(); + let upstream = tokio::spawn(async move { + let (owner_stream, _) = upstream_listener.accept().await.unwrap(); + let mut owner = accept_upstream_websocket(owner_stream).await.unwrap(); + let initial = receive_upstream_request(&mut owner).await; + send_upstream_response(&mut owner, "resp_owner", "first answer").await; + + let continued = receive_upstream_request(&mut owner).await; + continued_tx.send(()).unwrap(); + release_failure_rx.await.unwrap(); + owner + .send(Message::Text( + json!({ + "type": "response.failed", + "response": { + "id": "resp_owner_missing", + "status": "failed", + "error": { + "code": "previous_response_not_found", + "message": "Previous response with id resp_owner was not found", + }, + }, + }) + .to_string() + .into(), + )) + .await + .unwrap(); + owner.close(None).await.unwrap(); + + let fallback = match timeout(Duration::from_millis(700), upstream_listener.accept()).await { + Ok(Ok((fallback_stream, _))) => { + let mut fallback = accept_upstream_websocket(fallback_stream).await.unwrap(); + let request = receive_upstream_request(&mut fallback).await; + send_upstream_response(&mut fallback, "resp_wrong_fallback", "wrong account").await; + Some(request) + } + Ok(Err(error)) => panic!("fallback listener failed: {error}"), + Err(_) => None, + }; + (initial, continued, fallback) + }); + let (app, state, api_key, _pool, _dir) = + crate::dispatch::service::test_app_with_two_accounts_and_state(upstream_base_url).await; + let server = spawn_app(app).await; + let mut websocket = connect_responses_websocket(server.address, &api_key).await; + + websocket + .send(Message::Text(response_create_payload("first").into())) + .await + .unwrap(); + let first = receive_response(&mut websocket).await; + let mut continued_payload: Value = + serde_json::from_str(&response_create_payload("second")).unwrap(); + continued_payload["previous_response_id"] = json!("resp_owner"); + websocket + .send(Message::Text(continued_payload.to_string().into())) + .await + .unwrap(); + continued_rx.await.unwrap(); + assert!( + state + .services + .account_pool + .set_status("acct_primary", AccountStatus::Expired) + .await + ); + release_failure_tx.send(()).unwrap(); + let second = receive_response(&mut websocket).await; + let (initial, continued, fallback) = upstream.await.unwrap(); + + assert_eq!(first.last().unwrap()["response"]["id"], "resp_owner"); + assert_eq!(initial["input"][0]["content"], "first"); + assert_eq!(continued["previous_response_id"], "resp_owner"); + assert_eq!(second.last().unwrap()["type"], "error"); + assert_eq!( + second.last().unwrap()["error"]["code"], + "previous_response_unavailable" + ); + assert!( + fallback.is_none(), + "same-account retry must not fall back to another candidate: {fallback:?}" + ); +} + #[tokio::test] async fn responses_websocket_should_allow_official_retry_after_midstream_upstream_disconnect() { let upstream_listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); diff --git a/backend/tests/dispatch/affinity/mod.rs b/backend/tests/dispatch/affinity/mod.rs index ecbc1fbdd..64f10b4b9 100644 --- a/backend/tests/dispatch/affinity/mod.rs +++ b/backend/tests/dispatch/affinity/mod.rs @@ -1,7 +1,6 @@ use chrono::{Duration, Utc}; use codex_proxy_rs::dispatch::affinity::{ - CyberPolicyFailureSnapshot, MAX_CONVERSATION_AFFINITIES, RedisSessionAffinityStore, - SessionAffinityEntry, + MAX_CONVERSATION_AFFINITIES, RedisSessionAffinityStore, SessionAffinityEntry, }; use redis::AsyncCommands; @@ -247,85 +246,6 @@ async fn redis_affinity_store_rejects_oversized_metadata() { assert!(error.to_string().contains("metadata is too large")); } -#[tokio::test] -async fn redis_affinity_store_caps_cyber_policy_failures_atomically() { - let store = RedisSessionAffinityStore::new(create_test_redis("affinity-cyber-cap").await); - for index in 1..=4 { - let failure = CyberPolicyFailureSnapshot { - account_id: format!("acct_{index}"), - event: "error".to_string(), - message: format!("blocked by account {index}"), - upstream_code: Some("cyber_policy".to_string()), - }; - store - .record_cyber_policy_failure("session-hash", &failure, 3, Duration::hours(1)) - .await - .unwrap(); - } - - let state = store - .cyber_policy_state("session-hash") - .await - .unwrap() - .unwrap(); - assert_eq!(state.failed_account_ids, vec!["acct_1", "acct_2", "acct_3"]); - let last_failure = state.last_failure.unwrap(); - assert_eq!(last_failure.account_id, "acct_3"); - assert_eq!(last_failure.message, "blocked by account 3"); -} - -#[tokio::test] -async fn redis_affinity_store_should_not_clear_recreated_cyber_policy_state_with_stale_revision() { - let store = RedisSessionAffinityStore::new(create_test_redis("affinity-cyber-cas").await); - let first = store - .record_cyber_policy_failure( - "session-hash", - &CyberPolicyFailureSnapshot { - account_id: "acct_1".to_string(), - event: "error".to_string(), - message: "first failure".to_string(), - upstream_code: Some("cyber_policy".to_string()), - }, - 3, - Duration::hours(1), - ) - .await - .unwrap(); - assert!( - store - .clear_cyber_policy_state("session-hash", &first.revision) - .await - .unwrap() - ); - let newer = store - .record_cyber_policy_failure( - "session-hash", - &CyberPolicyFailureSnapshot { - account_id: "acct_2".to_string(), - event: "error".to_string(), - message: "newer failure".to_string(), - upstream_code: Some("cyber_policy".to_string()), - }, - 3, - Duration::hours(1), - ) - .await - .unwrap(); - - assert!( - !store - .clear_cyber_policy_state("session-hash", &first.revision) - .await - .unwrap() - ); - assert!( - store - .clear_cyber_policy_state("session-hash", &newer.revision) - .await - .unwrap() - ); -} - fn affinity( account_id: &str, conversation_id: &str, diff --git a/backend/tests/dispatch/controllers/mod.rs b/backend/tests/dispatch/controllers/mod.rs new file mode 100644 index 000000000..f0d9e238c --- /dev/null +++ b/backend/tests/dispatch/controllers/mod.rs @@ -0,0 +1,277 @@ +use std::{fs, path::Path}; + +const CONTROLLER_OWNERS: &[(&str, &str)] = &[ + ("account_failure.rs", "AccountFailureController"), + ("affinity.rs", "AffinityController"), + ("cloudflare.rs", "CloudflareController"), + ("cyber_policy/mod.rs", "CyberPolicyController"), + ("history.rs", "HistoryController"), + ("quota.rs", "QuotaController"), + ("telemetry/mod.rs", "TelemetryController"), + ("usage.rs", "UsageController"), +]; + +#[test] +fn controller_owners_must_not_call_each_other() { + let root = Path::new(env!("CARGO_MANIFEST_DIR")).join("src/dispatch/controllers"); + + for (relative, own_controller) in CONTROLLER_OWNERS { + let source = read(&root.join(relative)); + for (_, sibling_controller) in CONTROLLER_OWNERS { + if sibling_controller == own_controller { + continue; + } + assert!( + !source.contains(sibling_controller), + "controller owner {relative} must not call sibling {sibling_controller}; compose them in ControllerSet" + ); + } + } +} + +#[test] +fn cyber_policy_semantics_must_stay_inside_its_owner() { + let dispatch = Path::new(env!("CARGO_MANIFEST_DIR")).join("src/dispatch"); + let mut violations = Vec::new(); + collect_cyber_policy_violations(&dispatch, &dispatch, &mut violations); + + assert!( + violations.is_empty(), + "cyber_policy feature semantics escaped its controller owner: {}", + violations.join(", ") + ); +} + +#[test] +fn request_enter_io_should_be_parallel_and_bounded() { + let root = Path::new(env!("CARGO_MANIFEST_DIR")).join("src/dispatch/controllers"); + let composition = read(&root.join("mod.rs")); + let history = read(&root.join("history.rs")); + let affinity = read(&root.join("affinity.rs")); + + assert!(composition.contains("tokio::join!(")); + for operation in [ + "HistoryController::enter(", + "self.cyber_policy.prepare(request)", + "AffinityController::preferred_account_id(", + ] { + assert!( + composition.contains(operation), + "request enter must run {operation} in its parallel I/O group" + ); + } + for (owner, source) in [("history", history), ("affinity", affinity)] { + assert!(source.contains("Duration::from_millis(100)")); + assert!( + source.contains("timeout("), + "{owner} Redis I/O must be best-effort bounded" + ); + } +} + +#[test] +fn controller_best_effort_io_should_share_one_bounded_owner() { + let manifest = Path::new(env!("CARGO_MANIFEST_DIR")); + let composition = read(&manifest.join("src/dispatch/controllers/mod.rs")); + let finalizer = read(&manifest.join("src/dispatch/lifecycle/finalizer.rs")); + + assert!( + composition + .contains("const BEST_EFFORT_IO_TIMEOUT: Duration = Duration::from_millis(100);") + ); + assert!(composition.contains("match timeout(BEST_EFFORT_IO_TIMEOUT, operation).await")); + for operation in [ + "\"cookie.complete\"", + "\"quota.complete_headers\"", + "\"cookie.attempt_error\"", + "\"telemetry.attempt_error\"", + "\"cloudflare.complete\"", + "\"usage.complete\"", + "\"telemetry.complete\"", + "\"affinity.complete\"", + "\"cloudflare\"", + "\"usage\"", + "\"telemetry\"", + ] { + assert!( + composition.contains(operation), + "controller I/O {operation} must use the shared best-effort timeout" + ); + } + assert!(finalizer.contains("timeout(LEASE_FINALIZE_TIMEOUT, account_lease.complete())")); +} + +#[test] +fn typed_client_failure_should_reach_api_without_reclassification() { + let manifest = Path::new(env!("CARGO_MANIFEST_DIR")); + let composition = read(&manifest.join("src/dispatch/controllers/mod.rs")); + let errors = read(&manifest.join("src/dispatch/errors.rs")); + let api = read(&manifest.join("src/api/client/errors.rs")); + + assert!(composition.contains("fn default_response_client_failure(")); + assert!(errors.contains("pub struct ClientFailure")); + assert!(api.contains("client_failure.status_code()")); + assert!(api.contains("client_failure.exposed_upstream()?")); + for forbidden in [ + "ControllerSet", + "classify_client_failure", + "client_failure_http_status", + ] { + assert!(!errors.contains(forbidden)); + assert!(!api.contains(forbidden)); + } +} + +#[test] +fn telemetry_exits_should_only_be_mounted_through_controller_set() { + let manifest = Path::new(env!("CARGO_MANIFEST_DIR")); + for relative in [ + "src/dispatch/service.rs", + "src/dispatch/stream/lifecycle.rs", + ] { + let adapter = read(&manifest.join(relative)); + for forbidden in [ + "controllers::telemetry", + "record_response_dispatch_error_event", + "record_prefetched_response_stream_failure_event", + ] { + assert!( + !adapter.contains(forbidden), + "{relative} must call a typed ControllerSet hook instead of {forbidden}" + ); + } + } +} + +#[test] +fn history_policy_should_not_escape_to_generic_routing() { + let manifest = Path::new(env!("CARGO_MANIFEST_DIR")); + let routing = manifest.join("src/dispatch/routing"); + assert!(!routing.join("history.rs").exists()); + let candidates = read(&routing.join("candidates.rs")); + for forbidden in ["previous_response", "History", "replay"] { + assert!( + !candidates.contains(forbidden), + "generic candidate routing must not own {forbidden} policy" + ); + } +} + +#[test] +fn canonical_failure_parser_should_extract_facts_without_business_code_mapping() { + let manifest = Path::new(env!("CARGO_MANIFEST_DIR")); + let parser = read(&manifest.join("src/upstream/openai/protocol/responses.rs")); + for raw_fact in [ + "failure_type", + "failure_explicit_status_code", + "retry_after_seconds_from_value", + ] { + assert!( + parser.contains(raw_fact), + "canonical failure parser must preserve {raw_fact}" + ); + } + for forbidden in [ + "usage_limit_reached", + "quota_exhausted", + "previous_response_not_found", + "account_banned", + ] { + assert!( + !parser.contains(forbidden), + "protocol parser must not classify business code {forbidden}" + ); + } +} + +#[test] +fn telemetry_business_logic_should_live_under_controller_owner() { + let manifest = Path::new(env!("CARGO_MANIFEST_DIR")); + assert!(!manifest.join("src/dispatch/recording.rs").exists()); + assert!( + manifest + .join("src/dispatch/controllers/telemetry/events.rs") + .exists() + ); +} + +#[test] +fn controllers_must_not_depend_on_live_stream_executor() { + let manifest = Path::new(env!("CARGO_MANIFEST_DIR")); + let controllers = manifest.join("src/dispatch/controllers"); + let mut violations = Vec::new(); + collect_live_executor_dependencies(&controllers, &controllers, &mut violations); + + assert!( + violations.is_empty(), + "controllers must consume typed exits instead of live executor context: {}", + violations.join(", ") + ); + + let live = read(&manifest.join("src/dispatch/stream/live.rs")); + for forbidden in [ + "LiveResponseStreamContext", + "ControllerSet", + "SessionAffinityService", + "AccountPoolService", + "CloudflareRecovery", + "Recorder", + ] { + assert!( + !live.contains(forbidden), + "live stream executor must only hold transport lifecycle state, not {forbidden}" + ); + } + assert!(live.contains("finalizer: StreamFinalizer")); +} + +fn collect_cyber_policy_violations(root: &Path, path: &Path, violations: &mut Vec) { + for entry in fs::read_dir(path).expect("dispatch source directory should be readable") { + let entry = entry.expect("dispatch source entry should be readable"); + let path = entry.path(); + if path.is_dir() { + collect_cyber_policy_violations(root, &path, violations); + continue; + } + if path.extension().is_none_or(|extension| extension != "rs") { + continue; + } + let relative = path.strip_prefix(root).expect("dispatch source path"); + if relative == Path::new("controllers/mod.rs") + || relative.starts_with("controllers/cyber_policy") + { + continue; + } + let source = read(&path); + if source.contains("cyber_policy") || source.contains("CyberPolicy") { + violations.push(relative.display().to_string()); + } + } +} + +fn collect_live_executor_dependencies(root: &Path, path: &Path, violations: &mut Vec) { + for entry in fs::read_dir(path).expect("controller source directory should be readable") { + let entry = entry.expect("controller source entry should be readable"); + let path = entry.path(); + if path.is_dir() { + collect_live_executor_dependencies(root, &path, violations); + continue; + } + if path.extension().is_none_or(|extension| extension != "rs") { + continue; + } + let source = read(&path); + if source.contains("stream::live") || source.contains("LiveResponseStreamContext") { + violations.push( + path.strip_prefix(root) + .expect("controller source path") + .display() + .to_string(), + ); + } + } +} + +fn read(path: &Path) -> String { + fs::read_to_string(path).unwrap_or_else(|error| panic!("read {}: {error}", path.display())) +} diff --git a/backend/tests/dispatch/lifecycle/mod.rs b/backend/tests/dispatch/lifecycle/mod.rs new file mode 100644 index 000000000..0312ef634 --- /dev/null +++ b/backend/tests/dispatch/lifecycle/mod.rs @@ -0,0 +1,295 @@ +mod stream; + +use std::{fs, path::Path}; + +#[test] +fn backend_src_should_not_contain_test_attributes() { + let source_root = Path::new(env!("CARGO_MANIFEST_DIR")).join("src"); + assert_no_test_attributes(&source_root); +} + +#[test] +fn live_responses_sse_should_have_one_canonical_decoder_owner() { + let manifest = Path::new(env!("CARGO_MANIFEST_DIR")); + let canonical = source(manifest, "src/dispatch/transport/canonical.rs"); + assert!( + canonical.contains("parse_sse_events"), + "canonical transport must own SSE decoding" + ); + + for relative in [ + "src/api/client/responses/websocket.rs", + "src/dispatch/stream/live.rs", + "src/dispatch/failure/sse.rs", + "src/dispatch/controllers/telemetry/events.rs", + ] { + let source = source(manifest, relative); + for forbidden in [ + "SseEventDecoder", + "parse_sse_events", + "response_from_codex_sse", + ] { + assert!( + !source.contains(forbidden), + "{relative} must consume canonical facts instead of {forbidden}" + ); + } + } +} + +#[test] +fn attempt_retry_contract_should_have_one_pipeline_owner() { + let manifest = Path::new(env!("CARGO_MANIFEST_DIR")); + let pipeline = source(manifest, "src/dispatch/lifecycle/pipeline.rs"); + for required in [ + "runner.next().await?", + ".handle_attempt(attempt.controller_scope()?, attempt.observation())", + "attempt.apply(decision).await?", + "AttemptStep::Committed(attempt)", + ] { + assert!( + pipeline.contains(required), + "attempt pipeline must own {required}" + ); + } + assert_eq!( + pipeline.matches(".handle_attempt(").count(), + 1, + "each attempt observation must enter the controller composition exactly once" + ); + for forbidden in [ + "observe_attempt(", + "controllers.decide(", + "controllers.last_failure(", + "ResponseDispatchError::NoActiveAccount", + ] { + assert!( + !pipeline.contains(forbidden), + "attempt pipeline must not split classification across {forbidden}" + ); + } + + for relative in [ + "src/dispatch/service.rs", + "src/dispatch/stream/lifecycle.rs", + ] { + let adapter = source(manifest, relative); + for forbidden in [ + "runner.next()", + ".handle_attempt(", + "observe_attempt(", + "controllers.decide(", + "runner.apply(", + "ResponseDispatchError::NoActiveAccount", + ] { + assert!( + !adapter.contains(forbidden), + "{relative} must map terminal results instead of owning {forbidden}" + ); + } + } +} + +#[test] +fn lifecycle_contract_should_not_leak_feature_owners() { + let manifest = Path::new(env!("CARGO_MANIFEST_DIR")); + for relative in [ + "src/dispatch/lifecycle/attempt.rs", + "src/dispatch/lifecycle/contract.rs", + "src/dispatch/lifecycle/pipeline.rs", + "src/dispatch/lifecycle/request.rs", + ] { + let lifecycle = source(manifest, relative); + for forbidden in [ + "HistoryRecoveryPlan", + "HistorySource", + "QuotaLimitReached", + "quota_limit_reached", + ] { + assert!( + !lifecycle.contains(forbidden), + "{relative} must expose generic lifecycle facts instead of {forbidden}" + ); + } + } +} + +#[test] +fn stream_commit_contract_should_decode_once_before_establishment() { + let manifest = Path::new(env!("CARGO_MANIFEST_DIR")); + let attempt = source(manifest, "src/dispatch/lifecycle/attempt.rs"); + let contract = source(manifest, "src/dispatch/lifecycle/contract.rs"); + let pipeline = source(manifest, "src/dispatch/lifecycle/pipeline.rs"); + let prefetch = source(manifest, "src/dispatch/transport/prefetch.rs"); + + assert!(attempt.contains("attempt_request.stream_commit_policy")); + assert!(!attempt.contains("decoder.push(prefetched")); + assert!(prefetch.contains("while !decoder.commit_boundary_reached(policy)")); + assert!(prefetch.contains("let batch = decoder.push(chunk)?")); + assert!(prefetch.contains("initial_batch.append(batch)")); + + let open = between(&attempt, "impl OpenAttempt<'_, '_> {", "/// 已越过提交边界"); + assert!(open.contains("decision: AttemptDecision")); + let committed = between( + &attempt, + "impl CommittedAttempt {", + "/// `next` 产出的提交 typestate", + ); + assert!(committed.contains("fn accept(self)")); + assert!(!committed.contains("AttemptDecision")); + assert!(!committed.contains("Retry")); + + for obsolete in [ + "AttemptCommitState", + "RetryAfterCommit", + "validate_attempt_decision", + "StreamResponseFacts::Ready", + ] { + assert!(!attempt.contains(obsolete)); + assert!(!contract.contains(obsolete)); + assert!(!pipeline.contains(obsolete)); + } +} + +#[test] +fn live_stream_exit_paths_should_share_one_consuming_finalizer() { + let manifest = Path::new(env!("CARGO_MANIFEST_DIR")); + let live = source(manifest, "src/dispatch/stream/live.rs"); + let finalizer = source(manifest, "src/dispatch/lifecycle/finalizer.rs"); + + assert_eq!( + live.matches(".finalize(").count(), + 1, + "the live stream may invoke exactly one finalizer" + ); + assert!( + finalizer.contains("async fn finalize(\n self,"), + "the finalizer must consume itself so it cannot run twice" + ); + for terminal_path in [ + "break StreamTerminal::Cancelled", + "break StreamTerminal::Shutdown", + "return Err(StreamTerminal::CaptureLimitExceeded)", + "return Err(StreamTerminal::DownstreamClosed)", + ] { + assert!( + live.contains(terminal_path), + "{terminal_path} must unwind through the shared finalizer" + ); + } + let leave = finalizer + .find(".leave_stream(") + .expect("controllers should leave from the finalizer"); + let lease = finalizer + .find("timeout(LEASE_FINALIZE_TIMEOUT, account_lease.complete())") + .expect("the account lease should finalize"); + let client = finalizer + .find("finish_client_stream(") + .expect("client termination should run after internal finalization"); + assert!(leave < lease && lease < client); +} + +#[test] +fn controller_exit_order_should_be_static_and_effect_only() { + let manifest = Path::new(env!("CARGO_MANIFEST_DIR")); + let controllers = source(manifest, "src/dispatch/controllers/mod.rs"); + let pipeline = source(manifest, "src/dispatch/lifecycle/pipeline.rs"); + let complete = source(manifest, "src/dispatch/service.rs"); + let stream = source(manifest, "src/dispatch/stream/lifecycle.rs"); + + for forbidden in [ + "EnteredControllerScopes", + "RequestControllerStage", + "AttemptControllerStage", + "ControllerScopeError", + "unwind_attempt", + "finalize_request", + "leave_next()", + "into_reverse()", + ] { + assert!( + !controllers.contains(forbidden), + "controller composition must not retain runtime marker path {forbidden}" + ); + } + assert!(!pipeline.contains("unwind_attempt")); + assert!(!complete.contains("finalize_request")); + assert!(!stream.contains("finalize_request")); + assert_eq!(complete.matches(".finalize_complete(").count(), 1); + + let complete_exit = between( + &controllers, + "async fn finalize_complete(", + "pub(in crate::dispatch) fn new(", + ); + assert_appears_in_order( + complete_exit, + &[ + "CloudflareController::leave_complete", + "UsageController::leave_complete", + "AffinityController::leave_complete", + "TelemetryController::leave_complete", + ], + ); + + let stream_exit = between( + &controllers, + "async fn leave_stream(", + "pub(in crate::dispatch) async fn handle_attempt(", + ); + assert_appears_in_order( + stream_exit, + &[ + "CloudflareController::leave_stream", + "UsageController::leave_stream", + "AffinityController::leave_stream", + "TelemetryController::leave_stream", + "self.apply_stream_failure", + ], + ); +} + +fn source(manifest: &Path, relative: &str) -> String { + fs::read_to_string(manifest.join(relative)).expect("Rust source should be readable") +} + +fn between<'a>(source: &'a str, start: &str, end: &str) -> &'a str { + let start = source + .find(start) + .expect("source start marker should exist"); + let tail = &source[start..]; + let end = tail.find(end).expect("source end marker should exist"); + &tail[..end] +} + +fn assert_appears_in_order(source: &str, required: &[&str]) { + let mut cursor = 0; + for marker in required { + let offset = source[cursor..] + .find(marker) + .unwrap_or_else(|| panic!("ordered source marker should exist: {marker}")); + cursor += offset + marker.len(); + } +} + +fn assert_no_test_attributes(path: &Path) { + let entries = fs::read_dir(path).expect("backend/src should be readable"); + for entry in entries { + let entry = entry.expect("backend/src entries should be readable"); + let path = entry.path(); + if path.is_dir() { + assert_no_test_attributes(&path); + continue; + } + if path.extension().is_some_and(|extension| extension == "rs") { + let source = fs::read_to_string(&path).expect("Rust source should be readable"); + for marker in ["#[cfg(test)]", "#[test]", "#[tokio::test]"] { + assert!( + !source.contains(marker), + "test marker {marker:?} must stay outside backend/src: {}", + path.display() + ); + } + } + } +} diff --git a/backend/tests/dispatch/lifecycle/stream.rs b/backend/tests/dispatch/lifecycle/stream.rs new file mode 100644 index 000000000..27bb6a14d --- /dev/null +++ b/backend/tests/dispatch/lifecycle/stream.rs @@ -0,0 +1,296 @@ +use std::time::Duration as StdDuration; + +use axum::{body::Body, http::Request}; +use futures::StreamExt; +use serde_json::json; +use tokio::time::timeout; +use tower::util::ServiceExt; + +use crate::dispatch::service::{ + collect_stream_body, spawn_chunked_sse_upstream, spawn_chunked_sse_upstream_bytes, + spawn_chunked_sse_upstream_then_abrupt_close, spawn_chunked_sse_upstream_then_clean_close, + test_app_with_account_and_pool, +}; + +const LIVE_OUTPUT: &str = + include_str!("../../fixtures/responses/http_sse/live_stream_hello_delta.sse"); + +#[tokio::test] +async fn stream_terminal_should_synthesize_failure_after_upstream_read_error() { + let (base_url, first_chunk_sent, close_upstream) = + spawn_chunked_sse_upstream_then_abrupt_close(include_str!( + "../../fixtures/responses/http_sse/partial_transport_failure.sse" + )) + .await; + + let (app, api_key, _pool, _dir) = test_app_with_account_and_pool(base_url).await; + let response_task = tokio::spawn(async move { + app.oneshot( + Request::builder() + .method("POST") + .uri("/v1/responses") + .header("authorization", format!("Bearer {api_key}")) + .header("content-type", "application/json") + .body(Body::from( + json!({ + "model": "gpt-5.5", + "input": [{"role": "user", "content": "Start then fail"}], + "stream": true, + "use_websocket": false + }) + .to_string(), + )) + .unwrap(), + ) + .await + .unwrap() + }); + + first_chunk_sent.await.unwrap(); + let response = timeout(StdDuration::from_millis(300), response_task) + .await + .expect("stream response should be returned before upstream closes") + .unwrap(); + let mut body_stream = response.into_body().into_data_stream(); + let first_chunk = timeout(StdDuration::from_millis(300), body_stream.next()) + .await + .expect("first proxied SSE chunk should arrive before upstream closes") + .unwrap() + .unwrap(); + assert!( + String::from_utf8(first_chunk.to_vec()) + .unwrap() + .contains("partial before transport failure") + ); + + close_upstream.send(()).unwrap(); + let rest = collect_stream_body(body_stream).await; + + assert!(rest.contains("event: response.failed")); + assert!(rest.contains("stream_disconnected")); + assert!(rest.ends_with("data: [DONE]\n\n")); +} + +#[tokio::test] +async fn stream_terminal_should_synthesize_failure_after_clean_upstream_close() { + let (base_url, first_chunk_sent, close_upstream) = spawn_chunked_sse_upstream_then_clean_close( + include_str!("../../fixtures/responses/http_sse/partial_clean_close.sse"), + ) + .await; + + let (app, api_key, _pool, _dir) = test_app_with_account_and_pool(base_url).await; + let response_task = tokio::spawn(async move { + app.oneshot( + Request::builder() + .method("POST") + .uri("/v1/responses") + .header("authorization", format!("Bearer {api_key}")) + .header("content-type", "application/json") + .body(Body::from( + json!({ + "model": "gpt-5.5", + "input": [{"role": "user", "content": "Start then close"}], + "stream": true, + "use_websocket": false + }) + .to_string(), + )) + .unwrap(), + ) + .await + .unwrap() + }); + + first_chunk_sent.await.unwrap(); + let response = timeout(StdDuration::from_millis(300), response_task) + .await + .expect("stream response should be returned before upstream closes") + .unwrap(); + let mut body_stream = response.into_body().into_data_stream(); + let before_close = timeout(StdDuration::from_millis(300), async { + let mut body = String::new(); + while let Some(chunk) = body_stream.next().await { + body.push_str(&String::from_utf8_lossy(&chunk.unwrap())); + if body.contains("partial before clean close") { + return body; + } + } + body + }) + .await + .expect("real output should arrive before upstream closes"); + assert!(before_close.contains("partial before clean close")); + + close_upstream.send(()).unwrap(); + let rest = collect_stream_body(body_stream).await; + + assert!(rest.contains("event: response.failed")); + assert!(rest.contains("stream_disconnected")); + assert!(rest.contains(r#""id":"resp_clean_close""#)); + assert!(rest.ends_with("data: [DONE]\n\n")); +} + +#[tokio::test] +async fn committed_quota_failure_should_apply_effect_and_finalize_once() { + let (base_url, first_chunk_sent, finish_upstream) = spawn_chunked_sse_upstream( + LIVE_OUTPUT, + include_str!("../../fixtures/responses/http_sse/live_stream_failed_quota.sse"), + ) + .await; + let (app, api_key, pool, _dir) = test_app_with_account_and_pool(base_url).await; + + let response = app + .oneshot( + Request::builder() + .method("POST") + .uri("/v1/responses") + .header("authorization", format!("Bearer {api_key}")) + .header("content-type", "application/json") + .body(Body::from( + json!({ + "model": "gpt-5.5", + "input": [{"role": "user", "content": "Fail after commit"}], + "stream": true, + "use_websocket": false + }) + .to_string(), + )) + .unwrap(), + ) + .await + .unwrap(); + first_chunk_sent.await.unwrap(); + let mut body = response.into_body().into_data_stream(); + let first = timeout(StdDuration::from_secs(1), body.next()) + .await + .expect("committed output should reach the client") + .expect("stream should yield committed output") + .expect("committed output should be readable"); + assert!(String::from_utf8_lossy(&first).contains("live stream hello")); + + finish_upstream.send(()).unwrap(); + let rest = collect_stream_body(body).await; + let account: (String,) = sqlx::query_as("select status from accounts where id = $1") + .bind("acct_chat") + .fetch_one(&pool) + .await + .unwrap(); + let usage: (i64,) = + sqlx::query_as("select request_count from account_usage where account_id = $1") + .bind("acct_chat") + .fetch_one(&pool) + .await + .unwrap(); + + assert!(rest.contains("live quota exhausted")); + assert_eq!((account.0.as_str(), usage.0), ("quota_exhausted", 1)); +} + +#[tokio::test] +async fn committed_auth_failure_should_expire_the_account() { + let (base_url, first_chunk_sent, finish_upstream) = spawn_chunked_sse_upstream( + LIVE_OUTPUT, + include_str!("../../fixtures/responses/http_sse/live_stream_failed_auth.sse"), + ) + .await; + let (app, api_key, pool, _dir) = test_app_with_account_and_pool(base_url).await; + + let response = app + .oneshot( + Request::builder() + .method("POST") + .uri("/v1/responses") + .header("authorization", format!("Bearer {api_key}")) + .header("content-type", "application/json") + .body(Body::from( + json!({ + "model": "gpt-5.5", + "input": [{"role": "user", "content": "Expire after commit"}], + "stream": true, + "use_websocket": false + }) + .to_string(), + )) + .unwrap(), + ) + .await + .unwrap(); + first_chunk_sent.await.unwrap(); + let mut body = response.into_body().into_data_stream(); + let first = timeout(StdDuration::from_secs(1), body.next()) + .await + .expect("committed output should reach the client") + .expect("stream should yield committed output") + .expect("committed output should be readable"); + assert!(String::from_utf8_lossy(&first).contains("live stream hello")); + + finish_upstream.send(()).unwrap(); + let rest = collect_stream_body(body).await; + let account: (String,) = sqlx::query_as("select status from accounts where id = $1") + .bind("acct_chat") + .fetch_one(&pool) + .await + .unwrap(); + + assert!(rest.contains("live token expired")); + assert_eq!(account.0, "expired"); +} + +#[tokio::test] +async fn capture_limit_terminal_should_finalize_once_and_fail_the_client_stream() { + let oversized_event = format!( + "event: response.output_text.delta\ndata: {}\n\n", + json!({ + "type": "response.output_text.delta", + "delta": "x".repeat(16 * 1024 * 1024), + }) + ) + .into_bytes(); + let (base_url, first_chunk_sent, finish_upstream) = + spawn_chunked_sse_upstream_bytes(LIVE_OUTPUT.as_bytes().to_vec(), oversized_event).await; + let (app, api_key, pool, _dir) = test_app_with_account_and_pool(base_url).await; + + let response = app + .oneshot( + Request::builder() + .method("POST") + .uri("/v1/responses") + .header("authorization", format!("Bearer {api_key}")) + .header("content-type", "application/json") + .body(Body::from( + json!({ + "model": "gpt-5.5", + "input": [{"role": "user", "content": "Exceed capture limit"}], + "stream": true, + "use_websocket": false + }) + .to_string(), + )) + .unwrap(), + ) + .await + .unwrap(); + first_chunk_sent.await.unwrap(); + let mut body = response.into_body().into_data_stream(); + let first = timeout(StdDuration::from_secs(1), body.next()) + .await + .expect("committed output should reach the client") + .expect("stream should yield committed output") + .expect("committed output should be readable"); + assert!(String::from_utf8_lossy(&first).contains("live stream hello")); + + finish_upstream.send(()).unwrap(); + let rest = timeout(StdDuration::from_secs(30), collect_stream_body(body)) + .await + .expect("capture-limit finalization should not stall"); + let usage: (i64,) = + sqlx::query_as("select request_count from account_usage where account_id = $1") + .bind("acct_chat") + .fetch_one(&pool) + .await + .unwrap(); + + assert!(rest.contains("upstream response exceeded the 16 MiB proxy capture limit")); + assert!(rest.ends_with("data: [DONE]\n\n")); + assert_eq!(usage.0, 1); +} diff --git a/backend/tests/dispatch/mod.rs b/backend/tests/dispatch/mod.rs index 7b4c0d241..78bc722b8 100644 --- a/backend/tests/dispatch/mod.rs +++ b/backend/tests/dispatch/mod.rs @@ -1,2 +1,5 @@ pub(crate) mod affinity; +mod controllers; +pub(crate) mod lifecycle; pub(crate) mod service; +mod transport; diff --git a/backend/tests/dispatch/service.rs b/backend/tests/dispatch/service.rs index cda149e19..2a31fc74a 100644 --- a/backend/tests/dispatch/service.rs +++ b/backend/tests/dispatch/service.rs @@ -274,7 +274,7 @@ pub(crate) async fn test_app_with_account( (router::router().with_state(state), api_key, dir) } -async fn test_app_with_account_and_pool( +pub(crate) async fn test_app_with_account_and_pool( base_url: String, ) -> ( axum::Router, @@ -495,7 +495,7 @@ async fn test_app_with_two_accounts_and_telemetry( (app, api_key, pool, dir) } -async fn test_app_with_two_accounts_and_state( +pub(crate) async fn test_app_with_two_accounts_and_state( base_url: String, ) -> ( axum::Router, @@ -552,6 +552,33 @@ async fn test_app_with_ranked_accounts( axum::Router, String, crate::support::storage::TestDatabaseGuard, +) { + let (app, _state, api_key, dir) = + test_app_with_ranked_accounts_and_state(base_url, account_count).await; + (app, api_key, dir) +} + +async fn test_app_with_ranked_accounts_and_state( + base_url: String, + account_count: usize, +) -> ( + axum::Router, + AppState, + String, + crate::support::storage::TestDatabaseGuard, +) { + test_app_with_ranked_accounts_and_state_config(base_url, account_count, |_| {}).await +} + +async fn test_app_with_ranked_accounts_and_state_config( + base_url: String, + account_count: usize, + configure: impl FnOnce(&mut AppConfig), +) -> ( + axum::Router, + AppState, + String, + crate::support::storage::TestDatabaseGuard, ) { let (pool, dir) = init_test_db("openai-responses-full-candidate-ledger").await; let api_key = insert_client_api_key(&pool).await; @@ -575,6 +602,7 @@ async fn test_app_with_ranked_accounts( } let mut config = test_config(test_database_url(), base_url); config.auth.rotation_strategy = "quota_reset_priority".to_string(); + configure(&mut config); let state = test_app_state_with_pool(&config, pool).await; state .services @@ -582,7 +610,12 @@ async fn test_app_with_ranked_accounts( .restore_from_store() .await .expect("account pool should restore"); - (router::router().with_state(state), api_key, dir) + ( + router::router().with_state(state.clone()), + state, + api_key, + dir, + ) } async fn seed_openai_admin_key(pool: &PgPool, key: &str) { @@ -1027,7 +1060,7 @@ async fn accept_cyber_policy_websocket_response_with_authorization( payload } -async fn accept_websocket_with_authorization( +pub(crate) async fn accept_websocket_with_authorization( stream: TcpStream, expected_authorization: &'static str, ) -> tokio_tungstenite::WebSocketStream { @@ -1469,9 +1502,20 @@ fn assert_rate_limit_header(metadata: &Value, name: &str, value: &str) { ); } -async fn spawn_chunked_sse_upstream( +pub(crate) async fn spawn_chunked_sse_upstream( first_frame: &'static str, final_frame: &'static str, +) -> (String, oneshot::Receiver<()>, oneshot::Sender<()>) { + spawn_chunked_sse_upstream_bytes( + first_frame.as_bytes().to_vec(), + final_frame.as_bytes().to_vec(), + ) + .await +} + +pub(crate) async fn spawn_chunked_sse_upstream_bytes( + first_frame: Vec, + final_frame: Vec, ) -> (String, oneshot::Receiver<()>, oneshot::Sender<()>) { let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); let base_url = format!("http://{}", listener.local_addr().unwrap()); @@ -1487,11 +1531,11 @@ async fn spawn_chunked_sse_upstream( ) .await .unwrap(); - write_http_chunk(&mut socket, first_frame.as_bytes()).await; + write_http_chunk(&mut socket, &first_frame).await; socket.flush().await.unwrap(); let _ = first_chunk_sent_tx.send(()); let _ = finish_rx.await; - write_http_chunk(&mut socket, final_frame.as_bytes()).await; + write_http_chunk(&mut socket, &final_frame).await; socket.write_all(b"0\r\n\r\n").await.unwrap(); socket.flush().await.unwrap(); }); @@ -1499,13 +1543,13 @@ async fn spawn_chunked_sse_upstream( (base_url, first_chunk_sent_rx, finish_tx) } -async fn spawn_chunked_sse_upstream_then_abrupt_close( +pub(crate) async fn spawn_chunked_sse_upstream_then_abrupt_close( first_frame: &'static str, ) -> (String, oneshot::Receiver<()>, oneshot::Sender<()>) { spawn_chunked_sse_upstream_then_close(first_frame, ChunkedSseCloseMode::Abrupt).await } -async fn spawn_chunked_sse_upstream_then_clean_close( +pub(crate) async fn spawn_chunked_sse_upstream_then_clean_close( first_frame: &'static str, ) -> (String, oneshot::Receiver<()>, oneshot::Sender<()>) { spawn_chunked_sse_upstream_then_close(first_frame, ChunkedSseCloseMode::Clean).await @@ -1572,7 +1616,7 @@ async fn spawn_chunked_sse_upstream_then_close_with_headers( (base_url, first_chunk_sent_rx, close_tx) } -async fn collect_stream_body(mut body_stream: S) -> String +pub(crate) async fn collect_stream_body(mut body_stream: S) -> String where S: futures::Stream> + Unpin, E: std::fmt::Debug, diff --git a/backend/tests/dispatch/service/responses_http.rs b/backend/tests/dispatch/service/responses_http.rs index 46da6047d..13d45722e 100644 --- a/backend/tests/dispatch/service/responses_http.rs +++ b/backend/tests/dispatch/service/responses_http.rs @@ -1,5 +1,107 @@ use super::*; +const RESPONSES_FAILED_SERVER_OVERLOADED_SSE: &str = concat!( + include_str!("../../fixtures/responses/http_sse/failed_server_overloaded.sse"), + "\n" +); + +#[tokio::test] +async fn responses_json_should_preserve_upstream_client_failure() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/codex/responses")) + .respond_with( + ResponseTemplate::new(200) + .insert_header("content-type", "text/event-stream") + .set_body_string(RESPONSES_FAILED_SERVER_OVERLOADED_SSE), + ) + .expect(1) + .mount(&server) + .await; + let (app, api_key, _dir) = test_app_with_account(server.uri()).await; + + let response = app + .oneshot( + Request::builder() + .method("POST") + .uri("/v1/responses") + .header("authorization", format!("Bearer {api_key}")) + .header("content-type", "application/json") + .body(Body::from( + json!({ + "model": "gpt-5.5", + "input": [], + "stream": false, + "use_websocket": false + }) + .to_string(), + )) + .unwrap(), + ) + .await + .unwrap(); + let status = response.status(); + let body = response_json(response).await; + + assert_eq!(status, StatusCode::SERVICE_UNAVAILABLE); + assert_eq!(body["error"]["code"], "server_overloaded"); + assert_eq!( + body["error"]["message"], + "upstream is temporarily overloaded" + ); +} + +#[tokio::test] +async fn responses_sse_should_preserve_upstream_client_failure() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/codex/responses")) + .respond_with( + ResponseTemplate::new(200) + .insert_header("content-type", "text/event-stream") + .set_body_string(RESPONSES_FAILED_SERVER_OVERLOADED_SSE), + ) + .expect(1) + .mount(&server) + .await; + let (app, api_key, _dir) = test_app_with_account(server.uri()).await; + + let response = app + .oneshot( + Request::builder() + .method("POST") + .uri("/v1/responses") + .header("authorization", format!("Bearer {api_key}")) + .header("content-type", "application/json") + .body(Body::from( + json!({ + "model": "gpt-5.5", + "input": [], + "stream": true, + "use_websocket": false + }) + .to_string(), + )) + .unwrap(), + ) + .await + .unwrap(); + let status = response.status(); + let content_type = response + .headers() + .get("content-type") + .and_then(|value| value.to_str().ok()) + .unwrap_or_default() + .to_string(); + let body = response_text(response).await; + + assert_eq!(status, StatusCode::OK); + assert!(content_type.starts_with("text/event-stream")); + assert!(body.contains(r#""code":"server_overloaded""#)); + assert!(body.contains("upstream is temporarily overloaded")); + assert!(body.ends_with("data: [DONE]\n\n")); +} + #[tokio::test] async fn responses_should_reject_invalid_json_without_upstream_request() { let server = MockServer::start().await; @@ -934,7 +1036,7 @@ async fn responses_should_scope_upstream_cookie_by_codex_response_path() { } #[tokio::test] -async fn responses_stream_should_close_http_sse_upstream_when_client_disconnects() { +async fn responses_stream_disconnect_should_cancel_upstream_and_finalize_once() { let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); let base_url = format!("http://{}", listener.local_addr().unwrap()); let upstream = tokio::spawn(async move { @@ -956,7 +1058,7 @@ async fn responses_stream_should_close_http_sse_upstream_when_client_disconnects .is_ok() }); - let (app, api_key, _dir) = test_app_with_account(base_url).await; + let (app, api_key, pool, _dir) = test_app_with_account_and_pool(base_url).await; let response = app .oneshot( Request::builder() @@ -996,6 +1098,32 @@ async fn responses_stream_should_close_http_sse_upstream_when_client_disconnects upstream.await.unwrap(), "dropping the downstream stream should close the HTTP SSE upstream socket" ); + let request_count = timeout(StdDuration::from_secs(2), async { + loop { + let request_count: Option<(i64,)> = + sqlx::query_as("select request_count from account_usage where account_id = $1") + .bind("acct_chat") + .fetch_optional(&pool) + .await + .unwrap(); + if let Some((request_count,)) = request_count { + break request_count; + } + tokio::task::yield_now().await; + } + }) + .await + .expect("cancelled stream should finalize its account lease"); + assert_eq!(request_count, 1); + + tokio::time::sleep(StdDuration::from_millis(50)).await; + let settled_request_count: (i64,) = + sqlx::query_as("select request_count from account_usage where account_id = $1") + .bind("acct_chat") + .fetch_one(&pool) + .await + .unwrap(); + assert_eq!(settled_request_count.0, 1); } #[tokio::test] @@ -1076,119 +1204,6 @@ async fn responses_stream_should_forward_first_chunk_before_upstream_completes() assert_eq!(usage, (1, 3, 4)); } -#[tokio::test] -async fn responses_stream_should_emit_failed_event_after_upstream_read_error_once_downstream_started() - { - let (base_url, first_chunk_sent, close_upstream) = - spawn_chunked_sse_upstream_then_abrupt_close(include_str!( - "../../fixtures/responses/http_sse/partial_transport_failure.sse" - )) - .await; - - let (app, api_key, _pool, _dir) = test_app_with_account_and_pool(base_url).await; - let response_task = tokio::spawn(async move { - app.oneshot( - Request::builder() - .method("POST") - .uri("/v1/responses") - .header("authorization", format!("Bearer {api_key}")) - .header("content-type", "application/json") - .body(Body::from( - json!({ - "model": "gpt-5.5", - "input": [{"role": "user", "content": "Start then fail"}], - "stream": true, - "use_websocket": false - }) - .to_string(), - )) - .unwrap(), - ) - .await - .unwrap() - }); - - first_chunk_sent.await.unwrap(); - let response = timeout(StdDuration::from_millis(300), response_task) - .await - .expect("stream response should be returned before upstream closes") - .unwrap(); - let mut body_stream = response.into_body().into_data_stream(); - let first_chunk = timeout(StdDuration::from_millis(300), body_stream.next()) - .await - .expect("first proxied SSE chunk should arrive before upstream closes") - .unwrap() - .unwrap(); - assert!( - String::from_utf8(first_chunk.to_vec()) - .unwrap() - .contains("partial before transport failure") - ); - - close_upstream.send(()).unwrap(); - let rest = collect_stream_body(body_stream).await; - - assert!(rest.contains("event: response.failed")); - assert!(rest.contains("stream_disconnected")); - assert!(rest.ends_with("data: [DONE]\n\n")); -} - -#[tokio::test] -async fn responses_stream_should_emit_failed_event_when_upstream_closes_without_completed() { - let (base_url, first_chunk_sent, close_upstream) = spawn_chunked_sse_upstream_then_clean_close( - include_str!("../../fixtures/responses/http_sse/partial_clean_close.sse"), - ) - .await; - - let (app, api_key, _pool, _dir) = test_app_with_account_and_pool(base_url).await; - let response_task = tokio::spawn(async move { - app.oneshot( - Request::builder() - .method("POST") - .uri("/v1/responses") - .header("authorization", format!("Bearer {api_key}")) - .header("content-type", "application/json") - .body(Body::from( - json!({ - "model": "gpt-5.5", - "input": [{"role": "user", "content": "Start then close"}], - "stream": true, - "use_websocket": false - }) - .to_string(), - )) - .unwrap(), - ) - .await - .unwrap() - }); - - first_chunk_sent.await.unwrap(); - let response = timeout(StdDuration::from_millis(300), response_task) - .await - .expect("stream response should be returned before upstream closes") - .unwrap(); - let mut body_stream = response.into_body().into_data_stream(); - let first_chunk = timeout(StdDuration::from_millis(300), body_stream.next()) - .await - .expect("first proxied SSE chunk should arrive before upstream closes") - .unwrap() - .unwrap(); - assert!( - String::from_utf8(first_chunk.to_vec()) - .unwrap() - .contains("partial before clean close") - ); - - close_upstream.send(()).unwrap(); - let rest = collect_stream_body(body_stream).await; - - assert!(rest.contains("event: response.failed")); - assert!(rest.contains("stream_disconnected")); - assert!(rest.contains(r#""id":"resp_clean_close""#)); - assert!(rest.ends_with("data: [DONE]\n\n")); -} - #[tokio::test] async fn responses_should_prefer_session_affinity_account_for_previous_response() { let (base_url, upstream) = spawn_single_websocket_completed_upstream("resp_affinity_ws").await; diff --git a/backend/tests/dispatch/service/responses_recovery.rs b/backend/tests/dispatch/service/responses_recovery.rs index 08dfc86ab..6379aed36 100644 --- a/backend/tests/dispatch/service/responses_recovery.rs +++ b/backend/tests/dispatch/service/responses_recovery.rs @@ -95,7 +95,7 @@ async fn responses_should_classify_sse_cyber_policy_failure_as_bad_request() { } #[tokio::test] -async fn responses_http_cyber_policy_should_keep_account_active_and_change_next_request() { +async fn responses_http_cyber_policy_should_keep_account_active_and_retry_next_account() { let server = MockServer::start().await; Mock::given(method("POST")) .and(path("/codex/responses")) @@ -123,20 +123,74 @@ async fn responses_http_cyber_policy_should_keep_account_active_and_change_next_ let (app, api_key, pool, _dir) = test_app_with_two_accounts(server.uri()).await; let request_body = json!({ "model": "gpt-5.5", - "prompt_cache_key": "conv_http_cyber_next_request", + "prompt_cache_key": "conv_http_cyber_retry_next_account", "input": [{"role": "user", "content": "Security assessment"}], "stream": false, "use_websocket": false }); + let response = app + .oneshot(responses_json_request(&api_key, &request_body)) + .await + .unwrap(); + assert_eq!(response.status(), StatusCode::OK); + assert_eq!(response_json(response).await["id"], "resp_after_402"); + let account_status: (String,) = sqlx::query_as("select status from accounts where id = $1") + .bind("acct_primary") + .fetch_one(&pool) + .await + .unwrap(); + assert_eq!(account_status.0, "active"); +} + +#[tokio::test] +async fn responses_http_sse_cyber_terminal_should_change_the_immediate_next_request() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/codex/responses")) + .and(header("authorization", "Bearer access-primary")) + .respond_with( + ResponseTemplate::new(200) + .insert_header("content-type", "text/event-stream") + .set_body_string(include_str!( + "../../fixtures/responses/http_sse/output_then_cyber_policy.sse" + )), + ) + .expect(1) + .mount(&server) + .await; + Mock::given(method("POST")) + .and(path("/codex/responses")) + .and(header("authorization", "Bearer access-secondary")) + .respond_with( + ResponseTemplate::new(200) + .insert_header("content-type", "text/event-stream") + .set_body_string(RESPONSES_STREAM_AFTER_402_SSE), + ) + .expect(1) + .mount(&server) + .await; + let (app, api_key, pool, _dir) = test_app_with_two_accounts(server.uri()).await; + let request_body = json!({ + "model": "gpt-5.5", + "prompt_cache_key": "conv_http_sse_cyber_immediate_next", + "input": [{"role": "user", "content": "Security assessment"}], + "stream": true, + "use_websocket": false + }); + let first_response = app .clone() .oneshot(responses_json_request(&api_key, &request_body)) .await .unwrap(); - assert_eq!(first_response.status(), StatusCode::FORBIDDEN); - let first_body = response_json(first_response).await; - assert_eq!(first_body["error"]["code"], "cyber_policy"); + let first_status = first_response.status(); + let first_body = response_text(first_response).await; + + assert_eq!(first_status, StatusCode::OK); + assert!(first_body.contains("partial output before HTTP SSE policy failure")); + assert!(first_body.contains("cyber_policy")); + let account_status: (String,) = sqlx::query_as("select status from accounts where id = $1") .bind("acct_primary") .fetch_one(&pool) @@ -148,12 +202,12 @@ async fn responses_http_cyber_policy_should_keep_account_active_and_change_next_ .oneshot(responses_json_request(&api_key, &request_body)) .await .unwrap(); - let second_body = response_json(second_response).await; - assert_eq!(second_body["id"], "resp_after_402"); + let second_body = response_text(second_response).await; + assert!(second_body.contains("resp_stream_after_402")); } #[tokio::test] -async fn responses_http_rate_limited_cyber_policy_should_not_retry_in_same_request() { +async fn responses_http_rate_limited_cyber_policy_should_retry_next_account() { let server = MockServer::start().await; Mock::given(method("POST")) .and(path("/codex/responses")) @@ -181,22 +235,18 @@ async fn responses_http_rate_limited_cyber_policy_should_not_retry_in_same_reque let (app, api_key, pool, _dir) = test_app_with_two_accounts(server.uri()).await; let request_body = json!({ "model": "gpt-5.5", - "prompt_cache_key": "conv_http_429_cyber_next_request", + "prompt_cache_key": "conv_http_429_cyber_retry_next_account", "input": [], "stream": false, "use_websocket": false }); - let first_response = app - .clone() + let response = app .oneshot(responses_json_request(&api_key, &request_body)) .await .unwrap(); - assert_eq!(first_response.status(), StatusCode::TOO_MANY_REQUESTS); - assert_eq!( - response_json(first_response).await["error"]["code"], - "cyber_policy" - ); + assert_eq!(response.status(), StatusCode::OK); + assert_eq!(response_json(response).await["id"], "resp_after_402"); let account_status: (String, Option>) = sqlx::query_as("select status, quota_cooldown_until from accounts where id = $1") .bind("acct_primary") @@ -204,12 +254,6 @@ async fn responses_http_rate_limited_cyber_policy_should_not_retry_in_same_reque .await .unwrap(); assert_eq!(account_status, ("active".to_string(), None)); - - let second_response = app - .oneshot(responses_json_request(&api_key, &request_body)) - .await - .unwrap(); - assert_eq!(response_json(second_response).await["id"], "resp_after_402"); } #[tokio::test] @@ -242,26 +286,16 @@ async fn responses_cyber_rotation_should_return_current_quota_error_when_candida "use_websocket": false }); - let first_response = app - .clone() + let response = app .oneshot(responses_json_request(&api_key, &request_body)) .await .unwrap(); - assert_eq!(first_response.status(), StatusCode::FORBIDDEN); - let second_response = app - .oneshot(responses_json_request(&api_key, &request_body)) - .await - .unwrap(); - assert_eq!(second_response.status(), StatusCode::TOO_MANY_REQUESTS); - let second_body = response_json(second_response).await; - assert_eq!(second_body["error"]["code"], "insufficient_quota"); - assert!( - second_body["error"]["message"] - .as_str() - .is_some_and(|message| { - message.contains("quota exhausted") && !message.contains("cyber blocked") - }) - ); + assert_eq!(response.status(), StatusCode::TOO_MANY_REQUESTS); + let body = response_json(response).await; + assert_eq!(body["error"]["code"], "insufficient_quota"); + assert!(body["error"]["message"].as_str().is_some_and(|message| { + message.contains("quota exhausted") && !message.contains("cyber blocked") + })); let account_status: (String,) = sqlx::query_as("select status from accounts where id = $1") .bind("acct_secondary") .fetch_one(&pool) @@ -1698,6 +1732,100 @@ async fn responses_stream_structural_events_should_not_commit_before_history_fai assert!(!captured.retry_attempted); } +#[tokio::test] +async fn responses_http_sse_structural_events_should_not_commit_before_history_failure() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/codex/responses")) + .respond_with( + ResponseTemplate::new(200) + .insert_header("content-type", "text/event-stream") + .set_body_string(include_str!( + "../../fixtures/responses/http_sse/structural_then_previous_response_not_found.sse" + )), + ) + .expect(1) + .mount(&server) + .await; + let (app, api_key, _pool, _dir) = + test_app_with_account_pool_and_affinity(server.uri(), "resp_stale").await; + let response = app + .oneshot( + Request::builder() + .method("POST") + .uri("/v1/responses") + .header("authorization", format!("Bearer {api_key}")) + .header("content-type", "application/json") + .body(Body::from( + json!({ + "model": "gpt-5.5", + "previous_response_id": "resp_stale", + "turnState": "turn-stale", + "input": [{"role": "user", "content": "Continue"}], + "stream": true, + "use_websocket": false + }) + .to_string(), + )) + .unwrap(), + ) + .await + .unwrap(); + let status = response.status(); + let body = response_text(response).await; + + assert_eq!(status, StatusCode::BAD_REQUEST); + assert!(body.contains("previous_response_unavailable")); + assert!(!body.contains("response.created")); +} + +#[tokio::test] +async fn responses_http_sse_should_commit_before_history_failure_after_real_output() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/codex/responses")) + .respond_with( + ResponseTemplate::new(200) + .insert_header("content-type", "text/event-stream") + .set_body_string(include_str!( + "../../fixtures/responses/http_sse/output_then_previous_response_not_found.sse" + )), + ) + .expect(1) + .mount(&server) + .await; + let (app, api_key, _pool, _dir) = + test_app_with_account_pool_and_affinity(server.uri(), "resp_stale").await; + let response = app + .oneshot( + Request::builder() + .method("POST") + .uri("/v1/responses") + .header("authorization", format!("Bearer {api_key}")) + .header("content-type", "application/json") + .body(Body::from( + json!({ + "model": "gpt-5.5", + "previous_response_id": "resp_stale", + "turnState": "turn-stale", + "input": [{"role": "user", "content": "Continue"}], + "stream": true, + "use_websocket": false + }) + .to_string(), + )) + .unwrap(), + ) + .await + .unwrap(); + let status = response.status(); + let body = response_text(response).await; + + assert_eq!(status, StatusCode::OK); + assert!(body.contains("visible HTTP SSE output")); + assert!(body.contains("response.failed")); +} + #[tokio::test] async fn responses_stream_should_not_retry_history_failure_after_real_output() { let (base_url, upstream) = spawn_single_websocket_sequence_upstream(vec![ diff --git a/backend/tests/dispatch/service/responses_websocket.rs b/backend/tests/dispatch/service/responses_websocket.rs index f5f7738c5..b9118c5e1 100644 --- a/backend/tests/dispatch/service/responses_websocket.rs +++ b/backend/tests/dispatch/service/responses_websocket.rs @@ -242,46 +242,41 @@ async fn responses_websocket_cyber_policy_should_change_account_on_next_request( } #[tokio::test] -async fn responses_websocket_cyber_policy_should_stop_after_three_accounts() { +async fn responses_websocket_cyber_policy_should_continue_after_three_accounts() { let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); let base_url = format!("http://{}", listener.local_addr().unwrap()); - let (three_attempts_tx, three_attempts_rx) = oneshot::channel(); - let (verify_no_fourth_tx, verify_no_fourth_rx) = oneshot::channel(); let upstream = tokio::spawn(async move { let mut payloads = Vec::new(); - for (index, authorization) in ["Bearer access-0", "Bearer access-1", "Bearer access-2"] - .into_iter() - .enumerate() + for (index, authorization) in [ + "Bearer access-0", + "Bearer access-1", + "Bearer access-2", + "Bearer access-3", + ] + .into_iter() + .enumerate() { payloads.push( accept_cyber_policy_websocket_response_with_authorization( &listener, authorization, - &format!("resp_cyber_cap_{index}"), + &format!("resp_cyber_rotation_{index}"), ) .await, ); } - three_attempts_tx.send(()).unwrap(); - verify_no_fourth_rx.await.unwrap(); - assert!( - timeout(StdDuration::from_millis(100), listener.accept()) - .await - .is_err(), - "a fourth upstream account must not be contacted" - ); payloads }); let (app, api_key, _dir) = test_app_with_ranked_accounts(base_url, 4).await; let request_body = json!({ "model": "gpt-5.5", - "prompt_cache_key": "conv_cyber_three_account_cap", + "prompt_cache_key": "conv_cyber_unbounded_rotation", "input": [{"role": "user", "content": "Security assessment"}], "stream": true, "use_websocket": true }); - for _ in 0..3 { + for _ in 0..4 { let response = app .clone() .oneshot(responses_json_request(&api_key, &request_body)) @@ -290,21 +285,265 @@ async fn responses_websocket_cyber_policy_should_stop_after_three_accounts() { assert_eq!(response.status(), StatusCode::OK); assert!(response_text(response).await.contains("cyber_policy")); } - three_attempts_rx.await.unwrap(); + assert_eq!(upstream.await.unwrap().len(), 4); +} - let fourth_response = app +#[tokio::test] +async fn responses_websocket_stale_success_must_not_clear_newer_cyber_failure() { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let base_url = format!("http://{}", listener.local_addr().unwrap()); + let (stale_accepted_tx, stale_accepted_rx) = oneshot::channel(); + let (release_stale_tx, release_stale_rx) = oneshot::channel(); + let (probe_started_tx, probe_started_rx) = oneshot::channel(); + let upstream = tokio::spawn(async move { + accept_cyber_policy_websocket_response_with_authorization( + &listener, + "Bearer access-0", + "resp_cyber_cas_initial", + ) + .await; + + let (stale_stream, _) = listener.accept().await.unwrap(); + let mut stale_websocket = + accept_websocket_with_authorization(stale_stream, "Bearer access-1").await; + let stale_payload = stale_websocket.next().await.unwrap().unwrap(); + let stale_payload = serde_json::from_str::(&stale_payload.into_text().unwrap()) + .expect("stale request payload should be json"); + stale_accepted_tx.send(()).unwrap(); + + accept_cyber_policy_websocket_response_with_authorization( + &listener, + "Bearer access-1", + "resp_cyber_cas_newer", + ) + .await; + + release_stale_rx.await.unwrap(); + stale_websocket + .send(Message::Text( + websocket_completed_response("resp_cyber_cas_stale_success", 3, 1).into(), + )) + .await + .unwrap(); + stale_websocket.close(None).await.unwrap(); + + probe_started_rx.await.unwrap(); + let unexpected_upstream = + match timeout(StdDuration::from_millis(500), listener.accept()).await { + Ok(Ok((stream, _))) => { + let mut websocket = + accept_websocket_with_authorization(stream, "Bearer access-0").await; + let _ = send_websocket_response_and_capture_payload( + &mut websocket, + websocket_completed_response("resp_cyber_cas_was_cleared", 3, 1), + ) + .await; + websocket.close(None).await.unwrap(); + true + } + Ok(Err(error)) => panic!("accept probe connection: {error}"), + Err(_) => false, + }; + (stale_payload, unexpected_upstream) + }); + let (app, api_key, _dir) = test_app_with_ranked_accounts(base_url, 2).await; + let request_body = json!({ + "model": "gpt-5.5", + "prompt_cache_key": "conv_cyber_revision_cas", + "input": [{"role": "user", "content": "Security assessment"}], + "stream": true, + "use_websocket": true + }); + + let initial = app + .clone() .oneshot(responses_json_request(&api_key, &request_body)) .await .unwrap(); - assert_eq!(fourth_response.status(), StatusCode::BAD_REQUEST); - let fourth_body = response_json(fourth_response).await; - assert_eq!(fourth_body["error"]["code"], "cyber_policy"); - assert_eq!( - fourth_body["error"]["message"], - "This request has been flagged for possible cybersecurity risk." + assert!(response_text(initial).await.contains("cyber_policy")); + + let stale_app = app.clone(); + let stale_api_key = api_key.clone(); + let stale_request_body = request_body.clone(); + let stale_request = tokio::spawn(async move { + stale_app + .oneshot(responses_json_request(&stale_api_key, &stale_request_body)) + .await + .unwrap() + }); + stale_accepted_rx.await.unwrap(); + + let newer_failure = app + .clone() + .oneshot(responses_json_request(&api_key, &request_body)) + .await + .unwrap(); + assert!(response_text(newer_failure).await.contains("cyber_policy")); + release_stale_tx.send(()).unwrap(); + + let stale_response = stale_request.await.unwrap(); + assert!( + response_text(stale_response) + .await + .contains("resp_cyber_cas_stale_success") ); - verify_no_fourth_tx.send(()).unwrap(); - assert_eq!(upstream.await.unwrap().len(), 3); + + probe_started_tx.send(()).unwrap(); + let probe = app + .oneshot(responses_json_request(&api_key, &request_body)) + .await + .unwrap(); + let probe_status = probe.status(); + let probe_body = response_text(probe).await; + let (stale_payload, unexpected_upstream) = upstream.await.unwrap(); + assert_eq!(stale_payload["type"], "response.create"); + assert!( + !unexpected_upstream, + "a stale success cleared the newer cyber failure revision: {probe_body}" + ); + assert_eq!(probe_status, StatusCode::OK); + assert!(probe_body.contains("no_available_accounts")); +} + +#[tokio::test] +async fn responses_websocket_recreated_cyber_state_must_get_a_new_revision() { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let base_url = format!("http://{}", listener.local_addr().unwrap()); + let (stale_accepted_tx, stale_accepted_rx) = oneshot::channel(); + let (release_stale_tx, release_stale_rx) = oneshot::channel(); + let upstream = tokio::spawn(async move { + accept_cyber_policy_websocket_response_with_authorization( + &listener, + "Bearer access-0", + "resp_cyber_aba_initial", + ) + .await; + + let (stale_stream, _) = listener.accept().await.unwrap(); + let mut stale_websocket = + accept_websocket_with_authorization(stale_stream, "Bearer access-1").await; + let stale_payload = stale_websocket.next().await.unwrap().unwrap(); + let stale_payload = serde_json::from_str::(&stale_payload.into_text().unwrap()) + .expect("stale request payload should be json"); + stale_accepted_tx.send(()).unwrap(); + + accept_websocket_response_with_authorization_and_message( + &listener, + "Bearer access-1", + websocket_completed_response("resp_cyber_aba_clear", 3, 1), + ) + .await; + accept_cyber_policy_websocket_response_with_authorization( + &listener, + "Bearer access-0", + "resp_cyber_aba_recreated", + ) + .await; + + release_stale_rx.await.unwrap(); + stale_websocket + .send(Message::Text( + websocket_completed_response("resp_cyber_aba_stale_success", 3, 1).into(), + )) + .await + .unwrap(); + stale_websocket.close(None).await.unwrap(); + + let probe_payload = accept_websocket_response_with_authorization_and_message( + &listener, + "Bearer access-1", + websocket_completed_response("resp_cyber_aba_probe", 3, 1), + ) + .await; + (stale_payload, probe_payload) + }); + let (app, state, api_key, _dir) = + test_app_with_ranked_accounts_and_state_config(base_url, 2, |config| { + config.ws_pool.enabled = false; + }) + .await; + let request_body = json!({ + "model": "gpt-5.5", + "prompt_cache_key": "conv_cyber_revision_aba", + "input": [{"role": "user", "content": "Security assessment"}], + "stream": true, + "use_websocket": true + }); + + let initial = app + .clone() + .oneshot(responses_json_request(&api_key, &request_body)) + .await + .unwrap(); + assert!(response_text(initial).await.contains("cyber_policy")); + + let stale_app = app.clone(); + let stale_api_key = api_key.clone(); + let stale_request_body = request_body.clone(); + let stale_request = tokio::spawn(async move { + stale_app + .oneshot(responses_json_request(&stale_api_key, &stale_request_body)) + .await + .unwrap() + }); + stale_accepted_rx.await.unwrap(); + + let clearing_success = app + .clone() + .oneshot(responses_json_request(&api_key, &request_body)) + .await + .unwrap(); + assert!( + response_text(clearing_success) + .await + .contains("resp_cyber_aba_clear") + ); + assert!( + state + .services + .session_affinity + .forget("resp_cyber_aba_clear") + .await + ); + + let recreated_failure = app + .clone() + .oneshot(responses_json_request(&api_key, &request_body)) + .await + .unwrap(); + assert!( + response_text(recreated_failure) + .await + .contains("cyber_policy") + ); + release_stale_tx.send(()).unwrap(); + + let stale_response = stale_request.await.unwrap(); + assert!( + response_text(stale_response) + .await + .contains("resp_cyber_aba_stale_success") + ); + assert!( + state + .services + .session_affinity + .forget("resp_cyber_aba_stale_success") + .await + ); + + let probe = app + .oneshot(responses_json_request(&api_key, &request_body)) + .await + .unwrap(); + let probe_body = response_text(probe).await; + assert!( + probe_body.contains("resp_cyber_aba_probe"), + "unexpected ABA probe response: {probe_body}" + ); + let (stale_payload, probe_payload) = upstream.await.unwrap(); + assert_eq!(stale_payload["type"], "response.create"); + assert_eq!(probe_payload["type"], "response.create"); } #[tokio::test] @@ -474,6 +713,7 @@ async fn responses_websocket_stream_first_error_429_should_retry_fallback_accoun (first_payload, second_payload) }); let (app, api_key, pool, _dir) = test_app_with_two_accounts(base_url).await; + let request_started_at = Utc::now(); let response = app .oneshot(responses_json_request( @@ -522,7 +762,11 @@ async fn responses_websocket_stream_first_error_429_should_retry_fallback_accoun assert_eq!(first_payload["type"], "response.create"); assert_eq!(second_payload["type"], "response.create"); assert!(primary_quota_state.0); - assert!(primary_quota_state.1.is_some()); + let cooldown_until = primary_quota_state + .1 + .expect("explicit retry delay should persist a quota cooldown"); + assert!(cooldown_until >= request_started_at + chrono::Duration::seconds(70)); + assert!(cooldown_until <= Utc::now() + chrono::Duration::seconds(80)); assert_eq!(primary_usage.0, 1); assert_eq!(secondary_usage.0, 1); } @@ -1332,6 +1576,130 @@ async fn responses_websocket_without_history_should_return_rate_limit_stream_err assert_eq!(secondary_usage.0, 1); } +#[tokio::test] +async fn responses_websocket_response_failed_server_overloaded_should_match_http_external_error() { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let base_url = format!("http://{}", listener.local_addr().unwrap()); + let upstream = tokio::spawn(async move { + let (stream, _) = listener.accept().await.unwrap(); + let mut websocket = + accept_websocket_with_authorization(stream, "Bearer access-secret").await; + let _payload = send_websocket_response_and_capture_payload( + &mut websocket, + json!({ + "type": "response.failed", + "response": { + "id": "resp_ws_server_overloaded", + "status": "failed", + "error": { + "type": "server_error", + "code": "server_overloaded", + "message": "upstream is temporarily overloaded" + } + } + }) + .to_string(), + ) + .await; + websocket.close(None).await.unwrap(); + }); + let (app, api_key, _pool, _dir) = test_app_with_account_and_pool(base_url).await; + + let response = app + .oneshot(responses_json_request( + &api_key, + &json!({ + "model": "gpt-5.5", + "input": [], + "stream": false, + "use_websocket": true + }), + )) + .await + .unwrap(); + let status = response.status(); + let body = response_json(response).await; + upstream.await.unwrap(); + + assert_eq!(status, StatusCode::SERVICE_UNAVAILABLE); + assert_eq!(body["error"]["code"], "server_overloaded"); + assert_eq!( + body["error"]["message"], + "upstream is temporarily overloaded" + ); +} + +#[tokio::test] +async fn responses_websocket_response_failed_rate_limit_should_preserve_cooldown_and_fallback() { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let base_url = format!("http://{}", listener.local_addr().unwrap()); + let upstream = tokio::spawn(async move { + let (first_stream, _) = listener.accept().await.unwrap(); + let mut first_websocket = + accept_websocket_with_authorization(first_stream, "Bearer access-primary").await; + let _first_payload = send_websocket_response_and_capture_payload( + &mut first_websocket, + json!({ + "type": "response.failed", + "response": { + "id": "resp_ws_rate_limited", + "status": "failed", + "error": { + "type": "rate_limit_error", + "code": "rate_limit_exceeded", + "message": "Rate limit reached. Please try again in 11.054s." + } + } + }) + .to_string(), + ) + .await; + first_websocket.close(None).await.unwrap(); + + let (second_stream, _) = listener.accept().await.unwrap(); + let mut second_websocket = + accept_websocket_with_authorization(second_stream, "Bearer access-secondary").await; + send_websocket_response_and_capture_payload( + &mut second_websocket, + websocket_completed_response("resp_after_ws_rate_limit", 3, 1), + ) + .await; + second_websocket.close(None).await.unwrap(); + }); + let (app, api_key, pool, _dir) = test_app_with_two_accounts(base_url).await; + let request_started_at = Utc::now(); + + let response = app + .oneshot(responses_json_request( + &api_key, + &json!({ + "model": "gpt-5.5", + "input": [], + "stream": false, + "use_websocket": true + }), + )) + .await + .unwrap(); + let body = response_json(response).await; + upstream.await.unwrap(); + let primary_quota_state: (bool, Option>) = sqlx::query_as( + "select quota_limit_reached, quota_cooldown_until from accounts where id = $1", + ) + .bind("acct_primary") + .fetch_one(&pool) + .await + .unwrap(); + + assert_eq!(body["id"], "resp_after_ws_rate_limit"); + assert!(primary_quota_state.0); + let cooldown_until = primary_quota_state + .1 + .expect("rate-limit failure should persist a quota cooldown"); + assert!(cooldown_until >= request_started_at + chrono::Duration::seconds(8)); + assert!(cooldown_until <= Utc::now() + chrono::Duration::seconds(20)); +} + #[tokio::test] async fn responses_websocket_response_failed_quota_should_retry_fallback_account() { let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); diff --git a/backend/tests/dispatch/transport/mod.rs b/backend/tests/dispatch/transport/mod.rs new file mode 100644 index 000000000..1217c80de --- /dev/null +++ b/backend/tests/dispatch/transport/mod.rs @@ -0,0 +1,70 @@ +use std::{fs, path::Path}; + +#[test] +fn transport_must_normalize_facts_without_owning_feature_decisions() { + let root = Path::new(env!("CARGO_MANIFEST_DIR")).join("src/dispatch/transport"); + + for path in rust_files(&root) { + let source = fs::read_to_string(&path) + .unwrap_or_else(|error| panic!("read {}: {error}", path.display())); + for forbidden in [ + "ControllerSet", + "controllers::", + "AttemptDecision", + "RetryNextCandidate", + "RetrySameCandidate", + "cyber_policy", + ] { + assert!( + !source.contains(forbidden), + "transport adapter {} must emit typed facts instead of owning {forbidden}", + path.display() + ); + } + } +} + +#[test] +fn canonical_stream_decoder_must_be_the_only_dispatch_sse_parser() { + let dispatch = Path::new(env!("CARGO_MANIFEST_DIR")).join("src/dispatch"); + let canonical = dispatch.join("transport/canonical.rs"); + let canonical_source = fs::read_to_string(&canonical).expect("canonical transport source"); + assert!(canonical_source.contains("parse_sse_events")); + assert!(canonical_source.contains("fn normalize_complete_response")); + let protocol = fs::read_to_string( + Path::new(env!("CARGO_MANIFEST_DIR")).join("src/upstream/openai/protocol/responses.rs"), + ) + .expect("responses protocol source"); + assert!(!protocol.contains("fn response_from_codex_sse")); + + for path in rust_files(&dispatch) { + if path == canonical { + continue; + } + let source = fs::read_to_string(&path) + .unwrap_or_else(|error| panic!("read {}: {error}", path.display())); + assert!( + !source.contains("parse_sse_events"), + "{} must consume canonical events instead of decoding SSE again", + path.display() + ); + } +} + +fn rust_files(root: &Path) -> Vec { + let mut files = Vec::new(); + collect_rust_files(root, &mut files); + files +} + +fn collect_rust_files(path: &Path, files: &mut Vec) { + for entry in fs::read_dir(path).expect("source directory should be readable") { + let entry = entry.expect("source entry should be readable"); + let path = entry.path(); + if path.is_dir() { + collect_rust_files(&path, files); + } else if path.extension().is_some_and(|extension| extension == "rs") { + files.push(path); + } + } +} diff --git a/backend/tests/fixtures/responses/http_sse/failed_server_overloaded.sse b/backend/tests/fixtures/responses/http_sse/failed_server_overloaded.sse new file mode 100644 index 000000000..28637c599 --- /dev/null +++ b/backend/tests/fixtures/responses/http_sse/failed_server_overloaded.sse @@ -0,0 +1,2 @@ +event: response.failed +data: {"response":{"id":"resp_server_overloaded","object":"response","status":"failed","error":{"code":"server_overloaded","message":"upstream is temporarily overloaded"}}} diff --git a/backend/tests/fixtures/responses/http_sse/live_stream_failed_auth.sse b/backend/tests/fixtures/responses/http_sse/live_stream_failed_auth.sse new file mode 100644 index 000000000..54e188961 --- /dev/null +++ b/backend/tests/fixtures/responses/http_sse/live_stream_failed_auth.sse @@ -0,0 +1,2 @@ +event: response.failed +data: {"response":{"id":"resp_live_auth_failed","object":"response","status":"failed","error":{"code":"token_invalid","message":"live token expired"}}} diff --git a/backend/tests/fixtures/responses/http_sse/live_stream_failed_quota.sse b/backend/tests/fixtures/responses/http_sse/live_stream_failed_quota.sse new file mode 100644 index 000000000..80fbd47d8 --- /dev/null +++ b/backend/tests/fixtures/responses/http_sse/live_stream_failed_quota.sse @@ -0,0 +1,2 @@ +event: response.failed +data: {"response":{"id":"resp_live_quota_failed","object":"response","status":"failed","error":{"code":"quota_exceeded","message":"live quota exhausted"}}} diff --git a/backend/tests/fixtures/responses/http_sse/output_then_cyber_policy.sse b/backend/tests/fixtures/responses/http_sse/output_then_cyber_policy.sse new file mode 100644 index 000000000..a36deca52 --- /dev/null +++ b/backend/tests/fixtures/responses/http_sse/output_then_cyber_policy.sse @@ -0,0 +1,8 @@ +event: response.created +data: {"type":"response.created","response":{"id":"resp_http_cyber_policy","status":"in_progress"}} + +event: response.output_text.delta +data: {"type":"response.output_text.delta","response_id":"resp_http_cyber_policy","output_index":0,"content_index":0,"delta":"partial output before HTTP SSE policy failure"} + +event: response.failed +data: {"type":"response.failed","response":{"id":"resp_http_cyber_policy","status":"failed","error":{"code":"cyber_policy","message":"This request has been flagged for possible cybersecurity risk."}}} diff --git a/backend/tests/fixtures/responses/http_sse/output_then_previous_response_not_found.sse b/backend/tests/fixtures/responses/http_sse/output_then_previous_response_not_found.sse new file mode 100644 index 000000000..8f86b0406 --- /dev/null +++ b/backend/tests/fixtures/responses/http_sse/output_then_previous_response_not_found.sse @@ -0,0 +1,8 @@ +event: response.created +data: {"type":"response.created","response":{"id":"resp_http_output_then_missing","status":"in_progress"}} + +event: response.output_text.delta +data: {"type":"response.output_text.delta","response_id":"resp_http_output_then_missing","output_index":0,"content_index":0,"delta":"visible HTTP SSE output"} + +event: response.failed +data: {"type":"response.failed","response":{"id":"resp_http_output_then_missing","status":"failed","error":{"code":"previous_response_not_found","message":"Previous response was not found"}}} diff --git a/backend/tests/fixtures/responses/http_sse/structural_then_previous_response_not_found.sse b/backend/tests/fixtures/responses/http_sse/structural_then_previous_response_not_found.sse new file mode 100644 index 000000000..a6096a063 --- /dev/null +++ b/backend/tests/fixtures/responses/http_sse/structural_then_previous_response_not_found.sse @@ -0,0 +1,8 @@ +event: response.created +data: {"type":"response.created","response":{"id":"resp_http_previous_missing","status":"in_progress"}} + +event: response.content_part.added +data: {"type":"response.content_part.added","part":{"type":"output_text","text":""}} + +event: response.failed +data: {"type":"response.failed","response":{"id":"resp_http_previous_missing","status":"failed","error":{"code":"previous_response_not_found","message":"Previous response was not found"}}} diff --git a/backend/tests/upstream/openai/protocol/events.rs b/backend/tests/upstream/openai/protocol/events.rs index 77e4179d5..93206cc1d 100644 --- a/backend/tests/upstream/openai/protocol/events.rs +++ b/backend/tests/upstream/openai/protocol/events.rs @@ -176,6 +176,32 @@ fn retry_after_seconds_from_body_should_read_structured_retry_delay() { assert_eq!(retry_after_seconds_from_body(&body), Some(45)); } +#[test] +fn retry_after_seconds_from_body_should_read_nested_error_retry_delay() { + let body = json!({ + "type": "error", + "error": { + "retry_after_seconds": 19 + } + }) + .to_string(); + + assert_eq!(retry_after_seconds_from_body(&body), Some(19)); +} + +#[test] +fn retry_after_seconds_from_body_should_read_case_insensitive_header_array() { + let body = json!({ + "type": "error", + "headers": { + "Retry-After": ["37"] + } + }) + .to_string(); + + assert_eq!(retry_after_seconds_from_body(&body), Some(37)); +} + #[test] fn retry_after_seconds_from_body_should_parse_rate_limit_message_seconds() { let body = json!({ diff --git a/backend/tests/upstream/openai/protocol/mod.rs b/backend/tests/upstream/openai/protocol/mod.rs index aae722fa0..09d8d55be 100644 --- a/backend/tests/upstream/openai/protocol/mod.rs +++ b/backend/tests/upstream/openai/protocol/mod.rs @@ -5,16 +5,14 @@ use codex_proxy_rs::upstream::openai::protocol::events::{ parse_rate_limits_event, rate_limit_quota, retry_after_seconds_from_body, }; use codex_proxy_rs::upstream::openai::protocol::responses::{ - CodexResponsesRequest, CodexTransport, CollectedResponse, http_sse_fallback_allowed, - response_body_has_first_output, response_from_codex_sse, response_sse_event_is_terminal, - transport_for_request, + CodexResponsesRequest, CodexTransport, http_sse_fallback_allowed, + response_body_has_first_output, response_sse_event_is_terminal, transport_for_request, }; use codex_proxy_rs::upstream::openai::protocol::sse::{ SseEventDecoder, parse_sse_events, response_failed_sse_event, }; use codex_proxy_rs::upstream::openai::protocol::websocket::{ - OpeningAuditHeader, OpeningAuditSnapshot, classify_websocket_error_frame, - is_terminal_websocket_event, retry_after_seconds_from_wrapped_error_headers, + OpeningAuditHeader, OpeningAuditSnapshot, is_terminal_websocket_event, websocket_audit_artifact_from_attempt, websocket_event_to_sse_frame, websocket_metadata_turn_state, websocket_payload_audit_snapshot, websocket_response_completed_id, diff --git a/backend/tests/upstream/openai/protocol/responses.rs b/backend/tests/upstream/openai/protocol/responses.rs index 9052792fd..a4556ddd5 100644 --- a/backend/tests/upstream/openai/protocol/responses.rs +++ b/backend/tests/upstream/openai/protocol/responses.rs @@ -70,60 +70,6 @@ fn codex_responses_transport_flags_should_not_serialize_to_upstream_json() { assert!(body.get("forceHttpSse").is_none()); } -#[test] -fn response_from_codex_sse_should_collect_completed_response_with_deltas() { - let body = include_str!("../../../fixtures/responses/http_sse/text_delta_completed_usage.sse"); - - let response = response_from_codex_sse(body, None).expect("conversion should succeed"); - - assert_eq!( - response, - CollectedResponse::Completed(json!({ - "id": "resp_1", - "object": "response", - "status": "completed", - "usage": { - "input_tokens": 2, - "output_tokens": 3 - } - })) - ); -} - -#[test] -fn response_from_codex_sse_should_collect_incomplete_response() { - let body = include_str!("../../../fixtures/responses/http_sse/chat_delta_incomplete_usage.sse"); - - let response = response_from_codex_sse(body, None).expect("conversion should succeed"); - - assert_eq!( - response, - CollectedResponse::Incomplete(json!({ - "id": "resp_incomplete", - "object": "response", - "status": "incomplete", - "incomplete_details": { - "reason": "max_output_tokens" - }, - "usage": { - "input_tokens": 2, - "output_tokens": 3 - }, - "output": [{ - "type": "message", - "status": "incomplete", - "role": "assistant", - "content": [{ - "type": "output_text", - "text": "hello", - "annotations": [] - }] - }], - "output_text": "hello" - })) - ); -} - #[test] fn response_sse_event_is_terminal_should_match_incomplete_response() { let body = include_str!("../../../fixtures/responses/http_sse/chat_delta_incomplete_usage.sse"); diff --git a/backend/tests/upstream/openai/protocol/websocket.rs b/backend/tests/upstream/openai/protocol/websocket.rs index 2f0ee0c81..226bf943f 100644 --- a/backend/tests/upstream/openai/protocol/websocket.rs +++ b/backend/tests/upstream/openai/protocol/websocket.rs @@ -195,191 +195,6 @@ fn codex_websocket_terminal_event_should_match_completed_failed_and_error() { assert!(!is_terminal_websocket_event("response.output_text.delta")); } -#[test] -fn codex_websocket_error_frame_should_classify_rotatable_codes() { - let frame = json!({ - "type": "response.failed", - "response": { - "error": { - "code": "usage_limit_reached", - "message": "Rate limit reached. Please try again in 11.054s." - } - } - }) - .to_string(); - - let error = classify_websocket_error_frame(&frame).expect("rate limit frame should classify"); - - assert_eq!(error.status_code, 429); -} - -#[test] -fn codex_websocket_error_frame_should_classify_ts_aligned_upstream_error_codes() { - let cases = [ - ("quota_exhausted", 402), - ("payment_required", 402), - ("unauthorized", 401), - ("token_invalid", 401), - ("token_expired", 401), - ("account_deactivated", 401), - ("forbidden", 403), - ("account_banned", 403), - ("banned", 403), - ("previous_response_not_found", 400), - ("invalid_encrypted_content", 400), - ("missing_tool_output", 400), - ("no_tool_output", 400), - ]; - - for (code, expected_status) in cases { - let frame = json!({ - "type": "response.failed", - "response": { - "error": { - "code": code, - "message": "classified by upstream error mapping" - } - } - }) - .to_string(); - - let error = classify_websocket_error_frame(&frame) - .expect("TS-aligned upstream code should classify"); - - assert_eq!(error.status_code, expected_status, "code: {code}"); - } -} - -#[test] -fn codex_websocket_error_frame_should_ignore_legacy_extension_codes() { - let cases = [ - "quota_exceeded", - "insufficient_quota", - "usage_not_included", - "invalid_plan", - "banned_unknown_charge", - "context_length_exceeded", - "invalid_prompt", - "cyber_policy", - "invalid_request", - "no_tool_output_found_for_function_call", - "server_is_overloaded", - "slow_down", - "temporarily_unavailable", - "over_capacity", - "server_error", - "upstream_error", - "rate_limited", - ]; - - for code in cases { - let frame = json!({ - "type": "response.failed", - "response": { - "error": { - "code": code, - "message": "not part of the TS WebSocket rotation allowlist" - } - } - }) - .to_string(); - - assert!( - classify_websocket_error_frame(&frame).is_none(), - "code should pass through: {code}" - ); - } -} - -#[test] -fn codex_websocket_error_frame_should_honor_explicit_error_status() { - let frame = json!({ - "type": "error", - "status_code": 403, - "error": { - "code": "forbidden" - } - }) - .to_string(); - - let error = classify_websocket_error_frame(&frame).expect("explicit status should classify"); - - assert_eq!(error.status_code, 403); -} - -#[test] -fn codex_websocket_error_frame_should_ignore_success_status_and_unmapped_error() { - let success_status = json!({ - "type": "error", - "status": 200, - "error": { - "code": "informational" - } - }) - .to_string(); - let unmapped_error = json!({ - "type": "error", - "error": { - "code": "not_a_rotatable_error" - } - }) - .to_string(); - - assert!(classify_websocket_error_frame(&success_status).is_none()); - assert!(classify_websocket_error_frame(&unmapped_error).is_none()); -} - -#[test] -fn codex_websocket_error_frame_should_ignore_unmapped_response_failed() { - let frame = json!({ - "type": "response.failed", - "response": { - "error": { - "code": "model_refusal", - "message": "The model refused the request" - } - } - }) - .to_string(); - - let classified = classify_websocket_error_frame(&frame); - - assert!(classified.is_none()); -} - -#[test] -fn codex_websocket_wrapped_error_headers_should_extract_retry_after() { - let frame = json!({ - "type": "error", - "headers": { - "retry-after": ["37"] - }, - "error": { - "code": "rate_limit_exceeded" - } - }) - .to_string(); - - assert_eq!( - retry_after_seconds_from_wrapped_error_headers(&frame), - Some(37) - ); - - let frame = json!({ - "type": "error", - "error": { - "code": "rate_limit_exceeded", - "retry_after_seconds": 19 - } - }) - .to_string(); - - assert_eq!( - retry_after_seconds_from_wrapped_error_headers(&frame), - Some(19) - ); -} - #[test] fn codex_websocket_response_completed_id_should_validate_shape() { let frame = json!({ diff --git a/backend/tests/upstream/openai/transport/client.rs b/backend/tests/upstream/openai/transport/client.rs index ecb125049..79d0e512b 100644 --- a/backend/tests/upstream/openai/transport/client.rs +++ b/backend/tests/upstream/openai/transport/client.rs @@ -6,12 +6,9 @@ use std::sync::{ use super::{accept_codex_test_websocket, completed_websocket_response, request_context}; use codex_proxy_rs::upstream::openai::protocol::responses::CodexResponsesRequest; use codex_proxy_rs::upstream::openai::transport::{ - CodexBackendClient, CodexClientError, CodexUpstreamDiagnostics, CodexWebSocketPool, - is_cyber_policy_error_body, is_cyber_policy_upstream_error, - is_deactivated_workspace_error_body, + CodexBackendClient, CodexWebSocketPool, is_deactivated_workspace_error_body, }; use futures::{SinkExt, StreamExt}; -use reqwest::StatusCode; use tokio::net::TcpListener; use tokio_tungstenite::tungstenite::Message; @@ -29,37 +26,6 @@ fn deactivated_workspace_error_body_should_not_match_other_json_paths() { assert!(!is_deactivated_workspace_error_body(body)); } -#[test] -fn cyber_policy_error_body_should_match_supported_error_paths_only() { - assert!(is_cyber_policy_error_body( - r#"{"error":{"code":"cyber_policy"}}"# - )); - assert!(is_cyber_policy_error_body( - r#"{"response":{"error":{"code":"cyber_policy"}}}"# - )); - assert!(!is_cyber_policy_error_body( - r#"{"detail":{"code":"cyber_policy"}}"# - )); -} - -#[test] -fn cyber_policy_upstream_error_should_require_client_error_status() { - let error = |status| CodexClientError::Upstream { - status, - body: r#"{"error":{"code":"cyber_policy"}}"#.to_string(), - retry_after_seconds: None, - diagnostics: CodexUpstreamDiagnostics::default(), - set_cookie_headers: Vec::new(), - }; - - assert!(is_cyber_policy_upstream_error(&error( - StatusCode::FORBIDDEN - ))); - assert!(!is_cyber_policy_upstream_error(&error( - StatusCode::INTERNAL_SERVER_ERROR - ))); -} - #[tokio::test] async fn codex_backend_client_should_apply_configured_websocket_pool() { let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); diff --git a/backend/tests/upstream/openai/transport/mod.rs b/backend/tests/upstream/openai/transport/mod.rs index 897c2297b..64f9b8fce 100644 --- a/backend/tests/upstream/openai/transport/mod.rs +++ b/backend/tests/upstream/openai/transport/mod.rs @@ -8,7 +8,7 @@ use std::{ }; use codex_proxy_rs::upstream::openai::{ - protocol::responses::{PreviousResponseScope, StreamCommitPolicy}, + protocol::responses::PreviousResponseScope, transport::websocket::{ CodexWebSocketConnection, CodexWebSocketExchangeError, CodexWebSocketRequest, execute_response_create_request, responses_websocket_endpoint, diff --git a/backend/tests/upstream/openai/transport/websocket.rs b/backend/tests/upstream/openai/transport/websocket.rs index 6b2f85cdc..93bcebb60 100644 --- a/backend/tests/upstream/openai/transport/websocket.rs +++ b/backend/tests/upstream/openai/transport/websocket.rs @@ -299,7 +299,7 @@ async fn websocket_execute_response_create_request_should_collect_completed_sse( } #[tokio::test] -async fn websocket_execute_response_create_request_should_surface_response_failed_as_upstream_error() +async fn websocket_execute_response_create_request_should_return_business_coded_response_failed_as_terminal_sse_fact() { let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); let addr = listener.local_addr().unwrap(); @@ -342,21 +342,23 @@ async fn websocket_execute_response_create_request_should_surface_response_faile ) .expect("payload should serialize"); - let error = execute_response_create_request(&prepared) + let response = execute_response_create_request(&prepared) .await - .expect_err("response.failed should be surfaced as upstream error"); + .expect("response.failed should remain a terminal SSE fact"); server.await.unwrap(); - let CodexWebSocketExchangeError::Upstream(error) = error else { - panic!("expected upstream websocket error"); - }; - assert_eq!(error.status_code, 429); - assert_eq!(error.retry_after_seconds, Some(12)); - assert!(error.body.contains("rate_limit_exceeded")); + assert!(response.body.contains("event: response.failed")); + assert!(response.body.contains("\"code\":\"rate_limit_exceeded\"")); + assert!( + response + .body + .contains("Rate limit reached. Please try again in 11.054s.") + ); } #[tokio::test] -async fn websocket_execute_response_create_request_should_pass_through_unmapped_response_failed() { +async fn websocket_execute_response_create_request_should_return_unknown_response_failed_as_terminal_sse_fact() + { let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); let addr = listener.local_addr().unwrap(); let server = tokio::spawn(async move { @@ -387,7 +389,7 @@ async fn websocket_execute_response_create_request_should_pass_through_unmapped_ let response = execute_response_create_request(&prepared) .await - .expect("unmapped response.failed should remain a terminal SSE frame"); + .expect("unknown response.failed should remain a terminal SSE fact"); server.await.unwrap(); assert!(response.body.contains("event: response.failed")); @@ -457,7 +459,7 @@ async fn websocket_execute_response_create_request_should_reject_binary_event() } #[tokio::test] -async fn websocket_execute_response_create_request_should_surface_wrapped_error_status_and_retry_after() +async fn websocket_execute_response_create_request_should_return_wrapped_error_as_terminal_sse_fact() { let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); let addr = listener.local_addr().unwrap(); @@ -487,21 +489,20 @@ async fn websocket_execute_response_create_request_should_surface_wrapped_error_ }); let prepared = prepared_websocket_request(&format!("http://{addr}")); - let error = execute_response_create_request(&prepared) + let response = execute_response_create_request(&prepared) .await - .expect_err("wrapped error should surface as upstream error"); + .expect("wrapped error should remain a terminal SSE fact"); server.await.unwrap(); - let CodexWebSocketExchangeError::Upstream(error) = error else { - panic!("expected wrapped upstream error"); - }; - assert_eq!(error.status_code, 409); - assert_eq!(error.retry_after_seconds, Some(17)); - assert!(error.body.contains("wrapped conflict")); + assert!(response.body.contains("event: error")); + assert!(response.body.contains("\"status\":409")); + assert!(response.body.contains("\"retry-after\":[\"17\"]")); + assert!(response.body.contains("wrapped conflict")); } #[tokio::test] -async fn websocket_execute_response_create_request_should_surface_connection_limit_as_503() { +async fn websocket_execute_response_create_request_should_return_connection_limit_as_terminal_sse_fact() + { let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); let addr = listener.local_addr().unwrap(); let server = tokio::spawn(async move { @@ -529,16 +530,13 @@ async fn websocket_execute_response_create_request_should_surface_connection_lim }); let prepared = prepared_websocket_request(&format!("http://{addr}")); - let error = execute_response_create_request(&prepared) + let response = execute_response_create_request(&prepared) .await - .expect_err("connection limit should surface as upstream error"); + .expect("connection limit should remain a terminal SSE fact"); server.await.unwrap(); - let CodexWebSocketExchangeError::Upstream(error) = error else { - panic!("expected connection limit upstream error"); - }; - assert_eq!(error.status_code, 503); - assert!(error.body.contains("websocket_connection_limit_reached")); + assert!(response.body.contains("event: response.failed")); + assert!(response.body.contains("websocket_connection_limit_reached")); } #[tokio::test] @@ -915,8 +913,7 @@ async fn websocket_execute_response_create_request_should_reject_completed_witho } #[tokio::test] -async fn websocket_execute_response_create_request_should_pass_through_unclassified_error_terminal() -{ +async fn websocket_execute_response_create_request_should_return_error_terminal_as_sse_fact() { let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); let addr = listener.local_addr().unwrap(); let server = tokio::spawn(async move { @@ -959,7 +956,7 @@ async fn websocket_execute_response_create_request_should_pass_through_unclassif let response = execute_response_create_request(&prepared) .await - .expect("unclassified error frame should pass through as terminal SSE"); + .expect("error frame should remain a terminal SSE fact"); server.await.unwrap(); assert!(response.body.contains("event: error")); diff --git a/backend/tests/upstream/openai/transport/websocket_pool.rs b/backend/tests/upstream/openai/transport/websocket_pool.rs index 86963c97f..a2afe3316 100644 --- a/backend/tests/upstream/openai/transport/websocket_pool.rs +++ b/backend/tests/upstream/openai/transport/websocket_pool.rs @@ -598,69 +598,6 @@ async fn websocket_pool_should_bypass_busy_key_with_one_shot_connections() { ); } -#[tokio::test(start_paused = true)] -async fn websocket_pool_should_release_fresh_reservation_when_prefetch_is_cancelled() { - let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); - let addr = listener.local_addr().unwrap(); - let (first_request_received_tx, first_request_received_rx) = tokio::sync::oneshot::channel(); - let server = tokio::spawn(async move { - let (first_stream, _) = listener.accept().await.unwrap(); - let mut first_websocket = accept_codex_test_websocket(first_stream).await; - let _first_message = first_websocket.next().await.unwrap().unwrap(); - first_request_received_tx.send(()).unwrap(); - - let (second_stream, _) = listener.accept().await.unwrap(); - let mut second_websocket = accept_codex_test_websocket(second_stream).await; - let _second_message = second_websocket.next().await.unwrap().unwrap(); - second_websocket - .send(Message::Text( - completed_websocket_response("resp_after_cancelled_prefetch", 2, 1).into(), - )) - .await - .unwrap(); - second_websocket.close(None).await.unwrap(); - drop(first_websocket); - }); - let pool = Arc::new(CodexWebSocketPool::default()); - let backend = CodexBackendClient::new( - reqwest::Client::builder().no_proxy().build().unwrap(), - format!("http://{addr}"), - crate::support::fingerprint::runtime_test_fingerprint(), - ) - .with_websocket_pool(pool); - let mut request = pooled_websocket_request("conversation-cancelled-prefetch"); - request.stream_commit_policy = StreamCommitPolicy::UntilOutputOrTerminal; - - { - let first_request = backend.create_response_stream( - &request, - request_context("req_cancelled_prefetch", Some("chatgpt-account")), - ); - tokio::pin!(first_request); - tokio::select! { - _ = &mut first_request => { - panic!("prefetch request unexpectedly completed"); - } - result = first_request_received_rx => { - result.expect("server should observe the first request"); - } - } - } - tokio::time::sleep(Duration::from_millis(100)).await; - - let second = backend - .create_response( - &request, - request_context("req_after_cancelled_prefetch", Some("chatgpt-account")), - ) - .await - .expect("request after cancelled prefetch should succeed"); - server.await.unwrap(); - - assert!(second.body.contains("resp_after_cancelled_prefetch")); - assert_eq!(second.websocket_pool_decision.unwrap().kind(), "new"); -} - #[tokio::test(start_paused = true)] async fn websocket_pool_should_release_slot_when_client_drops_stream() { let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); @@ -1415,7 +1352,7 @@ async fn codex_backend_client_should_close_idle_pooled_websocket_after_liveness_ } #[tokio::test] -async fn codex_backend_client_should_discard_pooled_websocket_after_upstream_error() { +async fn codex_backend_client_should_discard_pooled_websocket_after_error_terminal() { let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); let addr = listener.local_addr().unwrap(); let accepted_connections = Arc::new(AtomicUsize::new(0)); @@ -1428,13 +1365,10 @@ async fn codex_backend_client_should_discard_pooled_websocket_after_upstream_err first_websocket .send(Message::Text( json!({ - "type": "response.failed", - "response": { - "id": "resp_pool_rate_limit", - "error": { - "code": "rate_limit_exceeded", - "message": "Rate limit reached. Please try again in 1s." - } + "type": "error", + "error": { + "code": "rate_limit_exceeded", + "message": "Rate limit reached. Please try again in 1s." } }) .to_string() @@ -1487,13 +1421,13 @@ async fn codex_backend_client_should_discard_pooled_websocket_after_upstream_err request.previous_response_scope = Some(PreviousResponseScope::Persisted); request.set_prompt_cache_key(Some("conversation-pool".to_string())); - let first_error = backend + let first = backend .create_response( &request, request_context("req_pool_error", Some("chatgpt-account")), ) .await - .expect_err("first pooled websocket response should surface upstream error"); + .expect("error should be returned as a terminal SSE fact"); let second = backend .create_response( &request, @@ -1503,17 +1437,15 @@ async fn codex_backend_client_should_discard_pooled_websocket_after_upstream_err .expect("second pooled websocket response should use a fresh connection"); server.await.unwrap(); - let CodexClientError::Upstream { status, body, .. } = first_error else { - panic!("expected upstream error from first pooled websocket response"); - }; - assert_eq!(status, reqwest::StatusCode::TOO_MANY_REQUESTS); - assert!(body.contains("rate_limit_exceeded")); + assert!(first.body.contains("event: error")); + assert!(first.body.contains("\"code\":\"rate_limit_exceeded\"")); assert!(second.body.contains("resp_pool_after_error")); + assert_eq!(second.websocket_pool_decision.unwrap().kind(), "new"); assert_eq!(accepted_connections.load(Ordering::SeqCst), 2); } #[tokio::test] -async fn codex_backend_client_should_discard_pooled_websocket_after_unmapped_response_failed() { +async fn codex_backend_client_should_discard_pooled_websocket_after_unknown_response_failed() { let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); let addr = listener.local_addr().unwrap(); let accepted_connections = Arc::new(AtomicUsize::new(0)); @@ -1548,7 +1480,7 @@ async fn codex_backend_client_should_discard_pooled_websocket_after_unmapped_res let _second_message = second_websocket.next().await.unwrap().unwrap(); second_websocket .send(Message::Text( - completed_websocket_response("resp_pool_after_unmapped_failed", 5, 2).into(), + completed_websocket_response("resp_pool_after_unknown_failed", 5, 2).into(), )) .await .unwrap(); @@ -1561,26 +1493,27 @@ async fn codex_backend_client_should_discard_pooled_websocket_after_unmapped_res crate::support::fingerprint::runtime_test_fingerprint(), ) .with_websocket_pool(pool); - let request = pooled_websocket_request("conversation-pool-unmapped-failed"); + let request = pooled_websocket_request("conversation-pool-unknown-failed"); let first = backend .create_response( &request, - request_context("req_pool_unmapped_failed", Some("chatgpt-account")), + request_context("req_pool_unknown_failed", Some("chatgpt-account")), ) .await - .expect("unmapped response.failed should be returned as terminal SSE"); + .expect("unknown response.failed should be returned as a terminal SSE fact"); let second = backend .create_response( &request, - request_context("req_pool_after_unmapped_failed", Some("chatgpt-account")), + request_context("req_pool_after_unknown_failed", Some("chatgpt-account")), ) .await - .expect("terminal failed websocket should be replaced"); + .expect("failed websocket should be replaced"); server.await.unwrap(); + assert!(first.body.contains("event: response.failed")); assert!(first.body.contains("resp_pool_model_refusal")); - assert!(second.body.contains("resp_pool_after_unmapped_failed")); + assert!(second.body.contains("resp_pool_after_unknown_failed")); assert_eq!(second.websocket_pool_decision.unwrap().kind(), "new"); assert_eq!(accepted_connections.load(Ordering::SeqCst), 2); } diff --git a/docs/architecture.md b/docs/architecture.md index bad521f7a..9773db485 100644 --- a/docs/architecture.md +++ b/docs/architecture.md @@ -45,7 +45,7 @@ PostgreSQL 和 Redis 都是必需依赖。任一连接失败时进程不能完 └── skills/ 仓库内开发约束 ``` -后端生产代码只放在 `backend/src/`。后端测试统一放在 `backend/tests/`，不在生产模块中加入 `#[cfg(test)]` 测试代码。 +后端生产代码只放在 `backend/src/`。后端测试统一放在 `backend/tests/`，并按生产边界分层组织；`backend/src/` 禁止出现 `#[cfg(test)]`、`mod tests` 和 test-only helper。 ## 后端模块 @@ -120,8 +120,7 @@ upstream/openai/ │ ├── responses.rs │ ├── events.rs │ ├── sse.rs -│ ├── websocket.rs -│ └── websocket_errors.rs +│ └── websocket.rs ├── transport/ │ ├── client.rs │ ├── client_sse.rs @@ -217,43 +216,64 @@ fleet/ `dispatch` 是代理请求的编排边界。账号选择、历史恢复、上游调用、流生命周期和遥测归因在这里汇合。 +`v1/*` 使用 Request、Attempt、Stream 三层洋葱生命周期。核心代码位于 `lifecycle/`、`controllers/`、`routing/`、`transport/` 和 `failure/`；`service.rs` 与 `stream/lifecycle.rs` 只是 complete/stream 入口适配器，`stream/live.rs` 是响应提交后的流执行器，不承载功能策略。 + +架构硬规则是“单功能单归属”：功能阈值、状态机、优先级和 decision 只能由对应 controller 拥有；`lifecycle` 只管固定顺序、retry、commit、逆序退出和 exactly-once finalize。迁移不保留 compatibility re-export、feature flag 切换或新旧双路径，新路径接管后必须在同一次变更中删除旧入口和重复判断。 + ```text dispatch/ ├── service.rs -├── upstream_call.rs -├── attempts.rs ├── errors.rs -├── recording.rs +├── lifecycle/ +│ ├── request.rs +│ ├── attempt.rs +│ ├── contract.rs +│ ├── pipeline.rs +│ ├── stream.rs +│ ├── finalizer.rs +│ └── trace.rs +├── controllers/ +│ ├── account_failure.rs +│ ├── affinity.rs +│ ├── cloudflare.rs +│ ├── cyber_policy/ +│ ├── history.rs +│ ├── quota.rs +│ ├── telemetry/ +│ │ ├── mod.rs +│ │ └── events.rs +│ └── usage.rs +├── routing/ +│ └── candidates.rs +├── transport/ +│ ├── account.rs +│ ├── canonical.rs +│ ├── observation.rs +│ └── prefetch.rs +├── failure/ +│ ├── exhaustion.rs +│ └── sse.rs ├── affinity/ │ ├── identity.rs │ ├── resolve.rs │ ├── service.rs │ ├── store.rs │ └── types.rs -├── recovery/ -│ ├── account_failure.rs -│ ├── auth.rs -│ ├── cloudflare.rs -│ ├── cyber_policy.rs -│ ├── exhaustion.rs -│ └── history.rs └── stream/ ├── lifecycle.rs - ├── live.rs - ├── prefetch.rs - ├── sse_failure.rs - └── trace.rs + └── live.rs ``` -- `service.rs`：非流式 Responses 主循环。 -- `stream/lifecycle.rs`：流式 Responses 主循环。 -- `upstream_call.rs`：把账号凭据、账号作用域身份和 Cookie 交给上游客户端。 -- `attempts.rs`：在请求开始时冻结完整候选顺序，记录已尝试、忙碌和状态排除的账号。 +- `service.rs`、`stream/lifecycle.rs`：complete/stream 入口适配器，共用唯一静态 lifecycle pipeline。 +- `lifecycle`：固定阶段顺序、retry/commit、终态和 exactly-once finalize。 +- `controllers`：功能规则、状态机、副作用和 typed hook 的唯一 owner。 +- `routing`：只保存通用候选账本；previous response 所有权、重放和恢复全部归 `controllers/history.rs`。 +- `transport`：账号调用、complete 归一化和 SSE canonical facts；流事件只解码一次后扇出。 +- `failure`：通用失败事实与耗尽聚合，不持有功能阈值。 - `errors.rs`：错误分类与客户端错误映射。 -- `recording.rs`：把调度轨迹和上游诊断写入 `telemetry::Recorder`。 +- `controllers/telemetry`：唯一拥有调度轨迹、上游诊断和终态遥测构造，再写入 `telemetry::Recorder`。 - `affinity`：会话亲和、账号作用域身份和 Redis 响应归属索引。 -- `recovery`：账号风险隔离、Cloudflare 冷却、候选耗尽和 previous response 恢复。 -- `stream`：首段预取、SSE 失败识别、下游流转发和结束时结算。 +- `stream`：流式入口和已提交响应的执行循环；策略与结算归 controller/lifecycle。 ### telemetry @@ -384,12 +404,12 @@ bootstrap/ 4. `models` 解析模型别名；`dispatch` 建立会话变体标识和 previous response 恢复计划。 5. `AccountAttemptLedger` 按当前调度策略冻结完整候选顺序。会话亲和账号可以排在前面，但不会绕过可用性检查；会话级 `cyber_policy` 负亲和账号在冻结前排除。 6. 账号池租用并发槽位，必要时先刷新 quota；请求间隔在发送上游前执行。 -7. `upstream_call` 注入账号 token、`chatgpt-account-id`、账号作用域身份、指纹请求头和 Cookie。 +7. `transport/account` 注入账号 token、`chatgpt-account-id`、账号作用域身份、指纹请求头和 Cookie。 8. 上游通过 HTTP/SSE 或 WebSocket 返回结果。建连最长等待 15 秒；WebSocket 发送后 20 秒没有任何上游事件时放弃异常连接，收到首个事件后与 HTTP/SSE 一样使用 5 分钟流空闲窗口。账号级 5xx 和传输故障回到候选账本，不在同一账号内叠加重试。 9. 在结果尚未提交给客户端时，账号级失败可继续尝试候选账本中的下一个账号。 10. 完成后更新账号用量、会话亲和、调度反馈和遥测事实，并释放账号槽位。 -当前实现是由 `service.rs`、`stream/lifecycle.rs` 和 `stream/live.rs` 手工连接各阶段的 phase orchestration，还不是统一的生命周期控制器模型。工程审计确认可以收敛为只服务 `v1/*` 的 Request、Attempt、Stream 三层嵌套洋葱；详细的现状证据、单功能不越界原则、目标目录、顺序、迁移步骤和验收条件见 [Responses 请求生命周期控制器设计](request-lifecycle-controller-design.md)。该设计限定在 `dispatch` 内部，不新增顶层 `application` 层，也不替代鉴权、请求 ID、连接排空等 Tower/Axum 中间件。 +`dispatch/lifecycle` 只服务 `v1/*`：非流式与流式共享 request enter 和 stepwise `AttemptRunner`，live stream 使用单次 canonical decode、typed terminal 和唯一 finalizer。功能进入与退出由静态 `ControllerSet` 按固定顺序调用，不存在旧入口包装、兼容 re-export 或新旧双路径。详细契约见 [Responses 请求生命周期控制器设计](request-lifecycle-controller-design.md)。该设计限定在 `dispatch` 内部，不替代鉴权、请求 ID、连接排空等 Tower/Axum 中间件。 会话压缩不使用独立端点。Codex 通过 `/v1/responses` 请求中的 `compaction_trigger` 和 `request_kind=compaction` 表达压缩语义；代理沿用同一条调度链路，并在遥测元数据中记录 `compact`、`requestKind` 和 `subagentKind`。 @@ -415,7 +435,7 @@ bootstrap/ OpenAI 返回 HTTP 402 时，仅当 JSON 正文的 `detail.code` 精确等于 `deactivated_workspace` 才把账号标记为 `banned`；其他 402 仍按额度耗尽处理。 -`cyber_policy` 不属于账号全局故障。流已经提交时，当前请求保留上游错误并正常结束，不在服务端重放。对于带显式 `session_id`、`conversation_id` 或 `prompt_cache_key` 的请求，代理按客户端 API Key 隔离记录命中账号；下一次同会话请求排除该账号后重新进入当前调度策略。单会话最多记录 3 个不同账号，第三个仍命中后后续请求直接返回最后一次上游错误；任一替代账号成功后按读取到的不可复用 revision token 做 CAS 清除，避免旧成功请求删除并发写入或删除后重建的新失败。状态保留 1 小时，不修改账号状态，也不影响其他会话。状态读取与普通会话亲和查询并行，Redis 操作超过 100 毫秒时 fail-open。携带 `previous_response_id` 的请求不使用该负亲和策略，避免把只有原账号持有的上游历史错误地路由到其他账号；客户端需要在新请求中重发完整历史后才能换号。 +`cyber_policy` 不属于账号全局故障。提交边界前命中时，attempt 立即排除当前账号并尝试下一个冻结候选；全部候选都不可用时，由通用 attempt outcome 返回最后一个已分类失败，不存在功能私有的重试次数或“三账号”终态。流已经提交时不能透明重放，当前请求保留上游错误并正常结束；对于带显式 `session_id`、`conversation_id` 或 `prompt_cache_key` 的请求，代理按客户端 API Key 隔离、去重记录命中账号，让下一次同会话请求在冻结候选前排除它们。读取过该状态的替代请求成功后，按不可复用 revision token 做 CAS 清除，避免旧成功请求删除并发写入或删除后重建的新失败。状态保留 1 小时，不修改账号状态，也不影响其他会话。状态读取与普通会话亲和查询并行，Redis 操作超过 100 毫秒时 fail-open。携带 `previous_response_id` 的请求不写入该负亲和状态；只有 history controller 已持有可重放完整历史时才允许跨账号，否则保留当前失败。 请求本身不可修复的 4xx、协议错误或已经向客户端提交内容后的流错误不会通过换号掩盖。候选耗尽后，返回最后一个有业务意义的账号错误；没有候选时返回无可用账号错误。 @@ -495,7 +515,7 @@ Redis 使用统一的 `cpr:` 前缀。 | `cpr:affinity:v3:resp:` | 响应归属元数据 | 4 小时 | | `cpr:affinity:v3:conv:` | 会话响应 ZSET 索引，最多 128 项 | 4 小时 | | `cpr:affinity:v3:account:` | 账号响应 ZSET 索引，最多 16384 项 | 4 小时 | -| `cpr:affinity:v3:cyber:` | API Key 隔离的会话级 `cyber_policy` 账号排除记录，最多 3 个账号 | 1 小时 | +| `cpr:cyber-policy:v2:session:` | API Key 隔离、按账号去重且无功能级数量阈值的会话 `cyber_policy` 排除记录 | 1 小时 | | `cpr:affinity:v3:global:index` | 全局响应 ZSET 索引，最多 65536 项 | 4 小时 | | `cpr:affinity:v3:global:bytes` | 亲和元数据总字节数 | 4 小时 | | `cpr:affinity:v3:global:info` | 亲和索引清理所需的归属与字节信息 | 4 小时 | diff --git a/docs/request-lifecycle-controller-design.md b/docs/request-lifecycle-controller-design.md index e13c17bfe..a9b355b55 100644 --- a/docs/request-lifecycle-controller-design.md +++ b/docs/request-lifecycle-controller-design.md @@ -1,26 +1,25 @@ # Responses 请求生命周期控制器设计 -状态：工程审计结论与目标设计，尚未实现。 +状态：已实现。`v1/*` 已统一采用 Request → Attempt → Stream 三层洋葱生命周期；不存在新旧双路径或兼容层。 本文回答一个具体问题：Responses 代理能否像 Koa 洋葱模型一样，让身份、历史恢复、账号恢复、会话策略、用量和遥测等模块只处理自己的入口与出口，并按固定顺序组合，而不把同一规则散落在请求过程中。 结论是可以，而且值得做。但目标不能是单层 `before / await next / after`。当前代理存在响应提交前后的不可逆边界，并且一次请求可能尝试多个账号，因此需要 Request、Attempt、Stream 三层嵌套生命周期。 -## 当前审计结论 +## 重构结果 -当前实现是 phase orchestration，不是 controller onion： +重构前分散在非流式请求、流式建连、已提交长流和客户端错误编码中的生命周期编排，已经收敛到统一的类型化生命周期： -| 范围 | 当前归属 | 主要问题 | +| 范围 | 当前归属 | 责任 | | --- | --- | --- | -| 非流式请求 | `dispatch/service.rs::complete` | 自己完成请求准备、候选冻结、账号循环、失败优先级和最终结算 | -| 流式建连 | `dispatch/stream/lifecycle.rs::stream` | 与非流式重复请求准备、候选、quota、history 和上游失败处理 | -| 已提交长流 | `dispatch/stream/live.rs` | handler 返回后继续持有账号 lease、会话状态、诊断、用量和遥测，并在多个提前返回分支分别释放资源 | -| 客户端错误 | `api/client/errors.rs` | HTTP JSON、SSE 和下游 WebSocket 分别判断同一种领域失败 | -| 流事件 | `stream/live.rs`、`recovery/cyber_policy.rs`、下游 WebSocket adapter | 同一 SSE 被多次解码或扫描，以重新推断 terminal、首输出和策略事实 | +| Request | `dispatch/lifecycle/request.rs` | 统一模型、身份、history/session plan、controller enter 和候选冻结 | +| Attempt | `dispatch/lifecycle/attempt.rs`、`contract.rs` | 统一单账号尝试、typed observation、decision、重试与提交边界 | +| Stream | `dispatch/lifecycle/stream.rs`、`finalizer.rs` | 统一 canonical event、终态、取消、shutdown 和 exactly-once finalize | +| 功能规则 | `dispatch/controllers//` | 每个功能唯一拥有自己的规则、状态和 lifecycle hook | +| 协议事实 | `dispatch/transport/` | 将上游响应、首帧和流事件规范化为 typed facts | +| 路由与失败 | `dispatch/routing/`、`failure/` | 通用候选/history 路由，以及失败事实聚合 | -领域规则并非全部重复。`HistoryRecoveryPlan`、`AccountAttemptLedger`、quota 验证、账号失败隔离和 `CyberPolicyRecovery` 已经分别集中。真正分散的是“何时调用、如何排序、如何把结果传到下一阶段”的生命周期布线。新增一个跨阶段策略时，仍可能同时修改非流式主循环、流式建连、live finalizer 和 API encoder。 - -当前取消路径也说明普通 handler 后置钩子不够：HTTP body 或下游 WebSocket 被丢弃后，live task 会收到取消并释放账号槽位，但不会经过完整的流终态结算。此时取消事实、部分用量、Cookie、rate-limit、turn state 和遥测可能缺失。 +`dispatch/service.rs` 和 `dispatch/stream/lifecycle.rs` 是 complete 与 stream 两种入口适配器：它们选择请求模式、调用同一 `enter_request` 和 `AttemptRunner`，再把已建立的流交给统一 Stream lifecycle。两者不是两套业务生命周期，也不得拥有功能规则。`dispatch/stream/live.rs` 是已提交流的传输执行器，终态统一交给 lifecycle finalizer，不构成兼容路径。 ## 设计边界 @@ -41,67 +40,74 @@ > 单功能单归属：一个功能的业务语义只能由一个 controller 拥有。其他层只能传递事实、执行通用控制流、适配协议或持久化数据，不得再次识别、推断或决定该功能。 -以 `cyber_policy` 为例，阈值、适用请求、session key、失败状态机、短路条件、CAS 清理和状态转换全部属于 `controllers/cyber_policy/`。它可以暴露 Request、Attempt、Stream、Finalize 四类窄 hook，但不应把同一规则复制到 API、transport、pipeline 或 store。 +以 `cyber_policy` 为例，适用请求、session key、失败识别、账号排除、CAS 清理和状态转换全部属于 `controllers/cyber_policy/`。它可以暴露 Request、Attempt、Stream、Finalize 四类窄 hook，但不应把同一规则复制到 API、transport、lifecycle 或 store；是否还能 retry 只由 lifecycle 的通用 commit/candidate contract 决定。 边界必须满足以下规则： - `controller` 唯一拥有功能规则、状态机、优先级、阈值、typed facts、decision 和 scope；不能选择下一个账号、编码 HTTP/SSE/WebSocket 或直接调用其他 controller。 - `lifecycle` 只拥有固定顺序、enter/leave、短路、逆序 unwind、retry 循环、commit 边界和 exactly-once finalize；不得出现 `if cyber_policy` 或功能阈值。 -- `adapter` / `upstream` 只把原始 HTTP、SSE、WebSocket 输入规范化为通用事实，例如状态码、错误码、terminal、usage；不得解释事实的业务后果。 -- `api` 只解析请求并把 `DispatchStart`、`ClientFailure` 编码成对应 transport；不得重新识别 cyber、quota、history 等功能。 +- `transport` / `upstream` 只把原始 HTTP、SSE、WebSocket 输入规范化为通用事实，例如状态码、错误码、terminal、usage；不得解释事实的业务后果。 +- `api` 只解析请求，并编码 `ResponseDispatchResponse`、`ResponseDispatchStream` 或包含 `ClientFailure` 的 `ResponseDispatchError`；不得重新识别 cyber、quota、history 等功能。 - `store` 只提供原子读写、CAS、TTL 和数据映射；业务专用 store trait 可以由 controller 定义，基础设施实现不能拥有业务阈值和状态转换。 - controller 之间不得横向调用；共享内容只能是没有业务含义的协议事实和值对象，不能把共享业务规则藏进 `utils`、`common` 或 helper。 -允许的唯一挂载点是 composition root / static registry：它负责构造 controller、注入依赖并注册一次。之后 pipeline 自动把所有阶段 hook 分发给该 controller，业务代码不在各个主循环中手工挂载。 +允许的唯一挂载点是静态 composition root `ControllerSet`：它负责构造 controller、注入依赖，并通过类型化方法分发各阶段 hook。业务代码不在各个入口或尝试循环中手工挂载功能。 -## 目标目录 +## 已实现目录与边界 -洋葱重构完成后的 `dispatch` 目标目录如下。现有文件会按纵向切片逐步迁移，下面是目标边界，不代表当前工作树已经具备全部目录。 +当前源码采用以下真实目录；目录本身就是实现边界，不另设抽象目标树： ```text backend/src/dispatch/ ├── lifecycle/ │ ├── mod.rs 洋葱组合入口 -│ ├── pipeline.rs 固定顺序、短路和逆序退出 -│ ├── registry.rs controller 静态注册 -│ ├── context.rs Request/Attempt/Stream 窄 context -│ ├── outcome.rs DispatchStart、AttemptOutcome、FinalOutcome │ ├── request.rs Request enter/leave │ ├── attempt.rs AttemptRunner -│ └── stream.rs StreamRuntime 与 exactly-once finalizer +│ ├── pipeline.rs 唯一 retry/commit 编排器 +│ ├── contract.rs typed observation、decision 与提交契约 +│ ├── stream.rs canonical stream 状态与终态 +│ ├── finalizer.rs exactly-once finalizer +│ └── trace.rs 请求与尝试 trace ├── controllers/ │ ├── mod.rs │ ├── cyber_policy/ │ │ ├── mod.rs -│ │ ├── controller.rs 规则、状态机和四类 hook -│ │ ├── scope.rs 单请求状态 -│ │ ├── facts.rs typed observation -│ │ ├── decision.rs typed decision -│ │ └── store.rs CyberPolicyStore port -│ ├── history/ -│ ├── quota/ -│ ├── account_failure/ -│ ├── affinity/ -│ ├── usage/ +│ │ ├── types.rs scope 与状态值对象 +│ │ └── store.rs 原子状态存取 +│ ├── history.rs +│ ├── quota.rs +│ ├── account_failure.rs +│ ├── cloudflare.rs +│ ├── affinity.rs +│ ├── usage.rs │ └── telemetry/ -├── adapters/ -│ ├── attempt_observation.rs 上游错误/首帧规范化 -│ └── canonical_stream.rs SSE/WS 单次解码和事件封装 -├── ports/ dispatch 对 fleet/upstream/storage 的窄接口 +│ ├── mod.rs +│ └── events.rs +├── routing/ 通用候选 ledger；不拥有功能策略 +├── transport/ 账号调用、prefetch、observation、canonical event +├── failure/ 通用失败聚合与 SSE 终态事实 +├── affinity/ 通用会话亲和基础设施 +├── stream/ +│ ├── lifecycle.rs stream 入口适配器 +│ └── live.rs 已提交流传输执行器 ├── errors.rs 通用 DispatchFailure 与 ClientFailure -└── service.rs v1/* 对外的唯一 dispatch facade +└── service.rs v1/* complete facade 与共享依赖容器 ``` -目录规则是：功能代码向 `controllers//` 收敛；生命周期代码只能位于 `lifecycle/`；协议事实只能位于 `adapters/`；跨领域依赖只能通过 `ports/`；API 不得反向进入 controller 的内部实现。 +目录规则是：功能代码向 `controllers/` 收敛；生命周期契约和控制流只能位于 `lifecycle/`；协议事实规范化只能位于 `transport/`；通用候选与失败聚合分别位于 `routing/`、`failure/`；API 不得反向进入 controller 的内部实现。 + +生产源码和测试必须物理分离：`backend/src/` 不允许出现 `#[cfg(test)]`、`mod tests` 或 test-only helper；所有 contract、controller、transport 和集成测试只能放在 `backend/tests/`，并按生产边界组织。 + +本次重构没有保留兼容路径：不存在 compatibility re-export、feature flag 切换或新旧双路径。`service.rs`、`stream/lifecycle.rs` 与 `stream/live.rs` 仅按传输阶段承担入口或执行职责，不复制 lifecycle/controller 业务语义。 ## 不越界验收 - 修改功能阈值、状态转换或 session 规则时，只修改该 controller 目录及其测试。 -- 新增功能只允许新增一个 controller、静态 registry 注册项和本模块测试，不修改 request/attempt/stream 主循环的业务判断。 -- `pipeline` contract test 只验证顺序和控制流；controller test 只验证功能语义；adapter test 只验证事实规范化；API test 只验证统一 outcome 编码。 +- 新增功能只允许新增一个 controller、`ControllerSet` 静态组合项和本模块测试，不修改 request/attempt/stream 主循环的业务判断。 +- lifecycle contract test 只验证顺序和控制流；controller test 只验证功能语义；transport test 只验证事实规范化；API test 只验证统一 outcome 编码。 - 用 `rg` 检查功能关键词：除 controller、自有 store/schema、测试和文档外，生产代码不得出现功能专用判断。 -- 删除 controller 后，工程只能因 registry 或类型引用失败；pipeline、API、adapter 和 store 中不能残留该功能的决策逻辑。 -- 以下情况直接判定越界：功能名出现在 pipeline 条件分支、API 重判错误、store 根据业务阈值决策、多个 controller 互相调用、同一原始事件被多次解析。 +- 删除 controller 后，工程只能因静态 `ControllerSet` 或类型引用失败；lifecycle、API、transport 和 store 中不能残留该功能的决策逻辑。 +- 以下情况直接判定越界：功能名出现在 lifecycle 控制流条件分支、API 重判错误、store 根据业务阈值决策、多个 controller 互相调用、同一原始事件被多次解析。 ## 三层生命周期 @@ -115,13 +121,13 @@ RequestLifecycle StreamLifecycle (响应提交后) enter: canonical decoder -> observers -> downstream encoder - leave: terminal/cancel -> policy effects -> affinity/usage/telemetry -> lease release + leave: terminal/cancel -> Cloudflare -> usage -> affinity -> telemetry -> cyber -> lease release - start: produce DispatchStart for API - leave: finalize one FinalOutcome immediately or when StreamHandle ends + start: return ResponseDispatchResponse or ResponseDispatchStream to API + leave: finalize one FinalOutcome immediately or when the stream body ends ``` -进入按声明顺序执行，退出按逆序执行。controller 在进入阶段短路时，只退出已经成功进入的 controller。 +进入顺序和有副作用的退出顺序都由静态 `ControllerSet` 明确声明。`ControllerRequestScope` 只保存 cyber、history、usage 的真实请求状态，不维护运行时 marker 栈；retry 或 rejection 不执行空 unwind，complete 与已提交 stream 才按固定逆序调用实际存在 leave 副作用的 owner。 三层的责任不同： @@ -129,82 +135,57 @@ RequestLifecycle - Attempt 决定单账号调用是否被接受、同账号恢复、换下一个账号或立即失败。 - Stream 只处理已经越过提交边界的事件、下游交付和终态；此时不能再产生换号 decision。 -流式请求有两个不同时间点，不能共用一个“最终结果”类型。API 在首帧提交边界只拿到 `DispatchStart::Streaming(StreamHandle)`，Request/Attempt 已进入 controller 的退出状态随 handle 转移给后台 lifecycle；真正的 completed、failed 或 cancelled 要等 stream terminal 后才能形成 `FinalOutcome` 并逆序退出。 +流式请求有两个不同时间点，不能共用一个“最终结果”类型。API 在首帧提交边界拿到 `ResponseDispatchStream`；Request/Attempt 状态已经被消费并转移进后台任务持有的 `StreamFinalizer`，真正的 completed、failed 或 cancelled 要等 stream terminal 后才能形成 `FinalOutcome` 并逆序退出。 ## 类型化上下文与结果 controller 不接收一个可随意写入的 Koa 风格字典，也不使用 `HashMap` 或无边界的 `serde_json::Value` 扩展包。公共 context 只保存跨模块稳定事实，每个 controller 通过窄输入或 capability view 读取所需字段。 -建议的核心类型： +实际核心类型： ```rust -struct RequestEnvelope { - request_id: String, - route: String, - requested_model: String, - started_at: Instant, - client_api_key_id: String, -} - struct RequestContext { - envelope: RequestEnvelope, request: CodexResponsesRequest, display_model: String, - history: HistoryRecoveryPlan, - route_plan: RoutePlan, + controller_scope: ControllerRequestScope, candidates: AccountAttemptLedger, - trace: ResponseDispatchTrace, } -struct AttemptContext { - lease: AccountLease, - attempt_request: CodexResponsesRequest, - transport: CodexBackendTransport, - attempt: ResponseDispatchAttempt, +enum AttemptStep<'runner, 'dependencies> { + Open(OpenAttempt<'runner, 'dependencies>), + Committed(CommittedAttempt), } -``` -过程控制使用显式枚举，不使用 `bool` 表达多种语义： - -```rust -enum AttemptOutcome { - Accepted(EstablishedResponse), - RetrySameAccount { reason: RetryReason }, - RetryNextCandidate { failure: AttemptFailure }, - ReturnFailure(DispatchFailure), +enum AttemptDecision { + Accept, + RetrySameAccount, + RetryNextCandidate { /* typed exhaustion */ }, + Return(AttemptReturnKind), } enum StreamTerminal { - Completed, - Incomplete, - Failed(StreamFailure), - Disconnected(StreamFailure), - InvalidProtocol(StreamFailure), - DownstreamCancelled, + Completed { /* canonical response */ }, + Incomplete { /* canonical response */ }, + Failed { /* typed failure */ }, + UpstreamClosed, + ProtocolError { /* detail */ }, + DownstreamClosed, Shutdown, } -enum DispatchStart { - Immediate(ResponsePayload), - Streaming(StreamHandle), - Failed(ClientFailure), -} - -enum FinalOutcome { - Completed(ResponseFacts), - Incomplete(ResponseFacts), - Failed(ClientFailure), - Cancelled(CancelFacts), +enum AttemptPipelineOutcome { + Established(EstablishedResponse), + Rejected(AttemptRejection), } ``` -账号状态、会话恢复和客户端错误映射只分类一次。API 立即消费 `DispatchStart`：非流式 payload 直接编码，stream handle 绑定到对应 transport body，失败由统一 `ClientFailure` 映射到 HTTP JSON、SSE 或 WebSocket。后台 lifecycle 只消费 `FinalOutcome` 做领域结算，不再重新识别 `cyber_policy`、quota、banned 或 history failure。 +账号状态、会话恢复和客户端错误映射只分类一次。API 消费 `ResponseDispatchResponse`、`ResponseDispatchStream` 或 `ResponseDispatchError`：非流式 payload 直接编码，stream body 绑定到对应 transport，typed failure 映射到 HTTP JSON、SSE 或 WebSocket。后台 lifecycle 只消费 `StreamTerminal` / `FinalOutcome` 做领域结算，不再重新识别 `cyber_policy`、quota、banned 或 history failure。 ## controller 契约 controller 分为两类： -- Decision controller：可以返回 retry、fail 或 commit 等显式 decision。 +- Decision controller：只在 `OpenAttempt` 上返回 retry、accept 或 fail 等显式 decision；`CommittedAttempt` 不暴露 decision API。 - Observer controller：只消费已经确定的事实。其 Redis、数据库或遥测失败默认是 best-effort，不能改写已经确定或已经发给客户端的 outcome。 错误也分为两类： @@ -212,7 +193,7 @@ controller 分为两类： - Critical：无法建立正确请求语义，转为 `DispatchFailure` 并触发逆序退出。 - BestEffort：记录警告并 fail-open，不阻塞正常请求。 -同层 controller 不直接调用另一个 controller，也不自行选择下一账号。pipeline 是唯一解释 decision 的组件。这样模块只拥有规则，编排器只拥有顺序和控制流。 +同层 controller 不直接调用另一个 controller，也不自行选择下一账号。lifecycle 是唯一解释 decision 的组件。这样模块只拥有规则，编排器只拥有顺序和控制流。 ## 固定顺序 @@ -221,7 +202,7 @@ controller 分为两类： 1. 请求规范化、模型解析和本地身份。 2. history plan。 3. 普通 affinity 与会话 recovery 查询；无依赖项可在同一 stage 并行。 -4. 合并为 `RoutePlan`，明确 preferred、required、excluded 和 fail-fast。 +4. 合并 controller enter 结果，明确 preferred 和 excluded 账号集合。 5. 冻结 `AccountAttemptLedger`。 6. 建立请求级 trace/telemetry 外层作用域。 @@ -235,13 +216,13 @@ controller 分为两类： 6. 发起上游调用。 7. 统一规范化 HTTP、SSE 首帧和 WebSocket 错误为 `AttemptObservation`。 8. 按唯一优先级分类 history、session policy、账号隔离和 transport retry。 -9. 产生 `AttemptOutcome` 并记录 attempt trace。 +9. 产生 `AttemptDecision`，由唯一 pipeline 应用为 `AttemptApplyOutcome` 并记录 attempt trace。 classifier 优先级只能存在一处。例如 `cyber_policy` 必须先于通用 403、quota 或 5xx 分类；其他模块消费分类后的类型，不能再次调用 `is_cyber_policy_*`。 ### Stream enter/leave -1. 上游 SSE 或 WebSocket 事件转换为一次性的 `CanonicalStreamEvent`。 +1. 上游 SSE 或 WebSocket 事件转换为一次性的 `CanonicalResponseChunk` / `CanonicalResponseEvent`。 2. 保留 raw bytes 供 HTTP SSE 透明转发；typed event 同时提供给 observer 和下游 WebSocket encoder。 3. 在首个可转发事件处记录 commit；commit 后类型上不再允许 retry。 4. 影响下一请求路由的关键状态在 terminal 向下游转发前完成有界更新。 @@ -249,7 +230,7 @@ classifier 优先级只能存在一处。例如 `cyber_policy` 必须先于通 6. 统一 finalizer 逆序执行 controller，并且只执行一次。 7. 最后释放账号 lease。 -`Drop` 不能执行异步 Redis、数据库或遥测操作。response body 的 Drop 只发送取消信号；持有完整生命周期的受监督 task 必须显式调用 `finalize_once(StreamTerminal::DownstreamCancelled).await`。Drop guard 只作为 panic/abort 的资源释放兜底。 +`Drop` 不能执行异步 Redis、数据库或遥测操作。response body 的 Drop 只发送取消信号；持有完整生命周期的受监督 task 在 live loop 返回唯一 `StreamSummary` 后，消费式调用 `StreamFinalizer::finalize(...).await`。Drop guard 只作为 panic/abort 的资源释放兜底。 ## Rust 落地选择 @@ -259,23 +240,23 @@ classifier 优先级只能存在一处。例如 `cyber_policy` 必须先于通 - controller 集合使用具体 `ControllerSet` 或 sealed controller enum，顺序由代码和 contract test 固定。 - 流 decoder 和状态机使用具体类型，不能对每个 token 走 async trait/vtable。 - 只有未来确实需要异构、可配置 controller 集合时，才在注册边界使用 `Arc`；每请求私有状态仍应封装在 controller scope 内。 -- typestate 只用于 Prepared、Routed、Committed 等少数不可逆边界，不把每个小阶段都编码为泛型状态。 +- typestate 只用于 `OpenAttempt` → `CommittedAttempt` 这一不可逆边界，不把每个 controller 小阶段都编码为泛型状态。 与 Koa 相同的是嵌套进入、逆序退出和短路语义；不同的是 Rust 通过所有权、显式 outcome 和提交 typestate 保证边界，而不是依赖一个任意可变 context 和递归 `next()` future。 -## 迁移顺序 +## 已落地的实现约束 -迁移必须按纵向切片完成，每一步迁移后删除原调用点，不长期保留新旧双路径。 - -1. 增加 characterization tests，固定当前 HTTP JSON、HTTP SSE、Responses WebSocket、history 和错误优先级行为。 -2. 引入 `RequestEnvelope`、typed failure/outcome、统一错误映射与各 transport encoder，先消除三种客户端错误输出的领域特判。 -3. 抽取共享 request prelude，删除 `complete` 与 `stream` 开头的重复准备逻辑。 -4. 抽取 `AttemptRunner`，统一 quota、history、账号作用域调用、失败优先级和候选 decision。 -5. 把 live task 改成单出口状态机；所有路径返回 `StreamTerminal` 后调用一个 exactly-once finalizer。 -6. 引入单次 SSE decode 的 canonical event pipeline，逐个迁移 cyber、usage、affinity、Cloudflare、rate-limit/turn state 和 telemetry observer。 -7. 删除旧的 phase wiring 和专用 API 特判，更新本文件状态为已实现。 - -不建议一开始重写上游 WebSocket pool。它已有独立的 lease/Drop 边界和连接回收测试，应先让 dispatch 生命周期消费它的 typed transport outcome。 +- complete 与 stream 入口共享 `enter_request` 和 `AttemptRunner`，不存在各自维护的账号尝试主循环。 +- `ControllerSet` 是静态 composition root；controller 顺序和 typed hook 在编译期可见，不使用动态 middleware 链。 +- `ControllerRequestScope` 只保存 cyber、history、usage 的真实状态；不存在通用 entered-controller marker、容量错误或空 unwind hook。 +- complete 的 accepted response 归一为 `FinalOutcome` 后由 `finalize_complete` 消费 scope；retry 继续进入下一 attempt，rejection 在 controller effect 已完成后直接返回，不经过空 finalize。 +- complete 的真实退出顺序固定为 Cloudflare → Usage → Affinity → Telemetry；live stream 固定为 Cloudflare → Usage → Affinity → Telemetry → Cyber，无副作用的阶段不制造空分支。 +- 上游 HTTP、SSE 首帧与 live stream 统一转换为 observation 或 canonical event，再由 controller 消费 typed facts。 +- stream 的 completed、incomplete、failed、断开、协议错误、下游取消和 shutdown 均形成显式终态，并经过 exactly-once finalizer。 +- WebSocket upgrade 成功后的 `response.failed` / `error` 保留 code、type、显式 status 和 retry 事实进入 canonical lifecycle；transport 不得据此推导 HTTP 状态或业务后果。 +- 入站 Responses WebSocket API 只采集 canonical transcript facts 并持有 service 暴露的 opaque replay snapshot/plan；previous ID 匹配、Replace/Append/Unavailable、清洗与拼接全部由 `HistoryController` 唯一拥有。 +- 上游 WebSocket pool 保持独立的连接 lease 与回收职责，dispatch lifecycle 只消费其 typed transport outcome。 +- 测试全部位于 `backend/tests/`；`backend/src/` 中不存在 `#[cfg(test)]`、`#[test]`、`mod tests` 或 test-only helper。 ## 性能约束 @@ -291,8 +272,9 @@ classifier 优先级只能存在一处。例如 `cyber_policy` 必须先于通 - 非流式与流式共享同一个 request prelude 和 `AttemptRunner`。 - 同一上游失败在 complete、prefetch 和 live terminal 得到一致分类。 - commit 后无法构造 retry decision，不能拼接两个账号的输出。 -- enter 顺序、leave 逆序、short-circuit unwind 和 controller 顺序有 contract test。 +- enter 顺序、真实 leave 副作用的固定逆序、无运行时 marker/unwind 和 controller 顺序有 contract test。 - completed、incomplete、failed、upstream disconnect、decoder error、oversize、downstream cancel 和 shutdown 都 exactly-once finalize。 - HTTP JSON、HTTP SSE 与 Responses WebSocket 对相同 `ClientFailure` 保持语义一致。 - 并发同会话状态使用不可复用 revision token/CAS 测试，旧成功不能删除新失败，状态删除或过期后重建也不能出现 ABA。 -- 生产模块不加入 `#[cfg(test)]`；测试放在 `backend/tests/dispatch/lifecycle/`。 +- 生产模块不加入 `#[cfg(test)]`、`mod tests` 或 test-only helper；测试只能放在 `backend/tests/`。 +- 不存在 compatibility re-export、旧入口包装器、feature flag 双路径或重复的旧失败分类。