From 77b18eec68282b17de9d3f77c10df64c4f57164f Mon Sep 17 00:00:00 2001 From: zyycn Date: Tue, 14 Jul 2026 05:13:48 +0800 Subject: [PATCH] fix(dispatch): handle deactivated and cyber policy failures Classify deactivated_workspace as banned while preserving ordinary 402 quota behavior. Track cyber_policy per client session so the next request avoids failed accounts, cap rotation at three accounts, guard success cleanup with ABA-safe token CAS, and document the Request/Attempt/Stream lifecycle controller target. Co-authored-by: Codex --- backend/src/api/client/errors.rs | 43 +- backend/src/api/client/responses/mod.rs | 8 +- backend/src/dispatch/affinity/service.rs | 33 +- backend/src/dispatch/affinity/store.rs | 115 +++++- backend/src/dispatch/affinity/types.rs | 20 + backend/src/dispatch/attempts.rs | 4 +- backend/src/dispatch/errors.rs | 17 +- .../src/dispatch/recovery/account_failure.rs | 32 +- backend/src/dispatch/recovery/cyber_policy.rs | 386 ++++++++++++++++++ backend/src/dispatch/recovery/mod.rs | 1 + backend/src/dispatch/service.rs | 53 ++- backend/src/dispatch/stream/lifecycle.rs | 51 ++- backend/src/dispatch/stream/live.rs | 102 ++++- .../src/upstream/openai/transport/client.rs | 33 +- backend/src/upstream/openai/transport/mod.rs | 3 +- backend/tests/dispatch/affinity/mod.rs | 82 +++- backend/tests/dispatch/service.rs | 36 ++ .../dispatch/service/responses_recovery.rs | 287 ++++++++++++- .../dispatch/service/responses_websocket.rs | 305 ++++++++++++++ .../upstream/openai/protocol/responses.rs | 13 + .../tests/upstream/openai/transport/client.rs | 52 ++- docs/architecture.md | 14 +- docs/request-lifecycle-controller-design.md | 228 +++++++++++ 23 files changed, 1870 insertions(+), 48 deletions(-) create mode 100644 backend/src/dispatch/recovery/cyber_policy.rs create mode 100644 docs/request-lifecycle-controller-design.md diff --git a/backend/src/api/client/errors.rs b/backend/src/api/client/errors.rs index ddb8ab70a..dee506a6e 100644 --- a/backend/src/api/client/errors.rs +++ b/backend/src/api/client/errors.rs @@ -8,7 +8,10 @@ use axum::{ use serde_json::{Value, json}; use crate::{ - dispatch::errors::{DispatchFailureClass, ResponseDispatchError}, + dispatch::{ + errors::{DispatchFailureClass, ResponseDispatchError}, + recovery::cyber_policy::is_cyber_policy_failure, + }, upstream::openai::protocol::sse::response_failed_sse_event as encode_response_failed_sse_event, upstream::openai::transport::CodexClientError, }; @@ -123,6 +126,20 @@ 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); + return openai_error_response( + status, + &failure.message, + kind.error_type, + failure.upstream_code.as_deref().unwrap_or(kind.code), + ) + .into_response(); + } + match error { ResponseDispatchError::NoActiveAccount => { responses_no_available_accounts_response().into_response() @@ -132,6 +149,17 @@ 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); + return encode_response_failed_sse_event( + kind.error_type, + failure.upstream_code.as_deref().unwrap_or(kind.code), + &failure.message, + ); + } let http_error = response_dispatch_http_error(error, ResponseDispatchMessageStyle::ResponsesStream); let kind = response_dispatch_error_kind(error); @@ -142,6 +170,19 @@ 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); + return responses_websocket_error_event( + status, + kind.error_type, + failure.upstream_code.as_deref().unwrap_or(kind.code), + &failure.message, + request_id, + ); + } if let Some(event) = upstream_client_websocket_error_event(error, request_id) { return event; } diff --git a/backend/src/api/client/responses/mod.rs b/backend/src/api/client/responses/mod.rs index 126f55952..56a413bed 100644 --- a/backend/src/api/client/responses/mod.rs +++ b/backend/src/api/client/responses/mod.rs @@ -109,9 +109,13 @@ pub fn build_codex_request( } fn apply_identity_fields(request: &mut CodexResponsesRequest, headers: &HeaderMap) { - request.client_conversation_id = request.prompt_cache_key().map(ToString::to_string); + request.client_conversation_id = identity_context_string(request.body(), "conversation_id") + .or_else(|| header_string(headers, "conversation-id")) + .or_else(|| header_string(headers, "conversation_id")) + .or_else(|| request.prompt_cache_key().map(ToString::to_string)); request.client_session_id = identity_context_string(request.body(), "session_id") - .or_else(|| header_string(headers, "session-id")); + .or_else(|| header_string(headers, "session-id")) + .or_else(|| header_string(headers, "session_id")); request.client_thread_id = identity_context_string(request.body(), "thread_id") .or_else(|| header_string(headers, "thread-id")); request.client_request_id = identity_context_string(request.body(), "x-client-request-id") diff --git a/backend/src/dispatch/affinity/service.rs b/backend/src/dispatch/affinity/service.rs index 236849b05..16a4e83c0 100644 --- a/backend/src/dispatch/affinity/service.rs +++ b/backend/src/dispatch/affinity/service.rs @@ -5,7 +5,7 @@ use thiserror::Error; use super::{ store::{RedisSessionAffinityStore, RedisSessionAffinityStoreError}, - types::SessionAffinityEntry, + types::{CyberPolicyFailureSnapshot, CyberPolicySessionState, SessionAffinityEntry}, }; const DEFAULT_SESSION_AFFINITY_TTL_SECS: i64 = 4 * 60 * 60; @@ -82,6 +82,37 @@ 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?) + } + async fn entry(&self, response_id: &str, now: DateTime) -> Option { match self.store.get(response_id, now, self.ttl).await { Ok(entry) => entry, diff --git a/backend/src/dispatch/affinity/store.rs b/backend/src/dispatch/affinity/store.rs index c852936c7..4d9d00dec 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::SessionAffinityEntry; +use super::types::{CyberPolicyFailureSnapshot, CyberPolicySessionState, SessionAffinityEntry}; /// 单会话保留的最近响应元数据数量。 pub const MAX_CONVERSATION_AFFINITIES: usize = 128; @@ -21,6 +21,65 @@ 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] @@ -461,6 +520,55 @@ 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()) } @@ -473,6 +581,11 @@ 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 bb4d8fec0..60955c054 100644 --- a/backend/src/dispatch/affinity/types.rs +++ b/backend/src/dispatch/affinity/types.rs @@ -5,6 +5,26 @@ 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/attempts.rs b/backend/src/dispatch/attempts.rs index aa901f76c..b9085e610 100644 --- a/backend/src/dispatch/attempts.rs +++ b/backend/src/dispatch/attempts.rs @@ -26,8 +26,10 @@ impl AccountAttemptLedger { pub(in crate::dispatch) async fn freeze( account_pool: &AccountPoolService, request: &AccountAcquireRequest, + excluded_account_ids: &BTreeSet, ) -> Self { - let candidate_ids = account_pool.candidate_snapshot(request).await; + let mut candidate_ids = account_pool.candidate_snapshot(request).await; + candidate_ids.retain(|account_id| !excluded_account_ids.contains(account_id)); Self { model: request.model.clone(), pending: candidate_ids.iter().cloned().collect(), diff --git a/backend/src/dispatch/errors.rs b/backend/src/dispatch/errors.rs index 6d966cbcb..efda88808 100644 --- a/backend/src/dispatch/errors.rs +++ b/backend/src/dispatch/errors.rs @@ -11,7 +11,10 @@ use crate::{ }, upstream::openai::{ protocol::{responses::ResponsesSseFailure, sse::SseError}, - transport::{CodexBackendTransport, CodexClientError, CodexUpstreamDiagnostics}, + transport::{ + CodexBackendTransport, CodexClientError, CodexUpstreamDiagnostics, + is_banned_upstream_error, is_cyber_policy_upstream_error, + }, }, }; @@ -189,6 +192,9 @@ 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()) @@ -196,6 +202,9 @@ pub(crate) fn is_rate_limit_upstream_error(error: &CodexClientError) -> bool { } 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, .. } @@ -223,7 +232,10 @@ pub(crate) fn is_retryable_account_transport_error(error: &CodexClientError) -> 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()) + CodexClientError::Upstream { status, .. } + if status_code_is_quota_exhausted(status.as_u16()) + && !is_banned_upstream_error(error) + && !is_cyber_policy_upstream_error(error) ) } @@ -233,6 +245,7 @@ pub(crate) fn is_model_unsupported_upstream_error(error: &CodexClientError) -> b 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) ) } diff --git a/backend/src/dispatch/recovery/account_failure.rs b/backend/src/dispatch/recovery/account_failure.rs index af9464561..da0aa22fb 100644 --- a/backend/src/dispatch/recovery/account_failure.rs +++ b/backend/src/dispatch/recovery/account_failure.rs @@ -26,7 +26,7 @@ use crate::{ fleet::{account::AccountStatus, pool::AccountPoolService}, upstream::openai::{ protocol::responses::ResponsesSseFailure, - transport::{CodexClientError, is_banned_upstream_error}, + transport::{CodexClientError, is_banned_upstream_error, is_cyber_policy_upstream_error}, }, }; @@ -41,6 +41,10 @@ pub(crate) async fn isolate_rotatable_account_failure( 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 @@ -49,6 +53,19 @@ pub(crate) async fn isolate_rotatable_account_failure( 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 @@ -90,19 +107,6 @@ pub(crate) async fn isolate_rotatable_account_failure( 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; - } - false } diff --git a/backend/src/dispatch/recovery/cyber_policy.rs b/backend/src/dispatch/recovery/cyber_policy.rs new file mode 100644 index 000000000..6673ab1bd --- /dev/null +++ b/backend/src/dispatch/recovery/cyber_policy.rs @@ -0,0 +1,386 @@ +//! `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/mod.rs b/backend/src/dispatch/recovery/mod.rs index 4c73a2913..8570c7539 100644 --- a/backend/src/dispatch/recovery/mod.rs +++ b/backend/src/dispatch/recovery/mod.rs @@ -3,5 +3,6 @@ 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/service.rs b/backend/src/dispatch/service.rs index cee820c70..4795a9990 100644 --- a/backend/src/dispatch/service.rs +++ b/backend/src/dispatch/service.rs @@ -27,6 +27,7 @@ use crate::{ recovery::{ account_failure::{isolate_rotatable_account_failure, isolate_sse_account_failure}, cloudflare::CloudflareRecovery, + cyber_policy::CyberPolicyRecovery, exhaustion::AccountExhaustionTracker, history::HistoryRecoveryPlan, }, @@ -167,15 +168,43 @@ impl ResponseDispatchService { prepare_variant_identity(&mut request); self.account_identity.prepare_local_identity(&mut request); let mut history = HistoryRecoveryPlan::load(&self.session_affinity, &request).await; - let preferred_account_id = self - .preferred_account_id_for_request(&request, &history, now) + 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), + ); + if let Some(failure) = cyber_policy_plan.exhausted_failure() { + let error = ResponseDispatchError::Failed(failure); + 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(), + 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 mut candidates = - AccountAttemptLedger::freeze(&self.account_pool, &acquire_request).await; + 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(); @@ -210,6 +239,11 @@ impl ResponseDispatchService { 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, @@ -218,7 +252,9 @@ impl ResponseDispatchService { started_at, ResponseDispatchErrorDetails { client_api_key_id: request.client_api_key_id.as_deref(), - account_id: exhausted_accounts.last_account_id(), + account_id: cyber_policy_plan + .last_account_id() + .or_else(|| exhausted_accounts.last_account_id()), stream: false, compact, transport: Some(backend_transport_name( @@ -339,6 +375,9 @@ impl ResponseDispatchService { 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) { @@ -362,6 +401,9 @@ impl ResponseDispatchService { 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, @@ -430,6 +472,7 @@ impl ResponseDispatchService { 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 { diff --git a/backend/src/dispatch/stream/lifecycle.rs b/backend/src/dispatch/stream/lifecycle.rs index 2970ecc78..c619e8616 100644 --- a/backend/src/dispatch/stream/lifecycle.rs +++ b/backend/src/dispatch/stream/lifecycle.rs @@ -18,6 +18,7 @@ use crate::{ }, recovery::{ account_failure::{isolate_rotatable_account_failure, isolate_sse_account_failure}, + cyber_policy::CyberPolicyRecovery, exhaustion::AccountExhaustionTracker, history::HistoryRecoveryPlan, }, @@ -63,15 +64,22 @@ impl ResponseDispatchService { prepare_variant_identity(&mut request); self.account_identity.prepare_local_identity(&mut request); let mut history = HistoryRecoveryPlan::load(&self.session_affinity, &request).await; - let preferred_account_id = self - .preferred_account_id_for_request(&request, &history, now) - .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 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 mut candidates = - AccountAttemptLedger::freeze(&self.account_pool, &acquire_request).await; + 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(); @@ -85,7 +93,9 @@ impl ResponseDispatchService { started_at, ResponseDispatchErrorDetails { client_api_key_id: request.client_api_key_id.as_deref(), - account_id: exhausted_accounts.last_account_id(), + account_id: cyber_policy_plan + .last_account_id() + .or_else(|| exhausted_accounts.last_account_id()), stream: true, compact, transport: Some(backend_transport_name( @@ -117,6 +127,15 @@ impl ResponseDispatchService { 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 @@ -143,6 +162,11 @@ impl ResponseDispatchService { 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); }; @@ -219,6 +243,13 @@ impl ResponseDispatchService { 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, @@ -314,6 +345,9 @@ impl ResponseDispatchService { } }; 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) { @@ -376,6 +410,8 @@ impl ResponseDispatchService { 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(), @@ -405,6 +441,9 @@ impl ResponseDispatchService { } 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, diff --git a/backend/src/dispatch/stream/live.rs b/backend/src/dispatch/stream/live.rs index 6310032ae..3db19f84a 100644 --- a/backend/src/dispatch/stream/live.rs +++ b/backend/src/dispatch/stream/live.rs @@ -18,7 +18,13 @@ use crate::{ insert_first_token_ms, live_response_rate_limit_headers, live_response_turn_state, record_live_response_stream_event, }, - recovery::{cloudflare::CloudflareRecovery, history::HistoryRecoveryPlan}, + recovery::{ + cloudflare::CloudflareRecovery, + cyber_policy::{ + CyberPolicyRecovery, CyberPolicyRoutingPlan, CyberPolicyStreamObserver, + }, + history::HistoryRecoveryPlan, + }, service::ResponseDispatchStream, }, fleet::pool::{AccountLease, AccountPoolService}, @@ -93,6 +99,15 @@ pub(in crate::dispatch) fn spawn_live_response_stream( .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( @@ -109,8 +124,14 @@ pub(in crate::dispatch) fn spawn_live_response_stream( return; } Err(LiveStreamWriteError::CaptureLimitExceeded) => { - terminate_oversized_live_stream(context, &sender, &mut body_bytes, first_token_ms) - .await; + terminate_oversized_live_stream( + context, + &sender, + &mut body_bytes, + first_token_ms, + cyber_policy_observer.state_transition_applied(), + ) + .await; return; } } @@ -128,6 +149,14 @@ pub(in crate::dispatch) fn spawn_live_response_stream( }; 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, @@ -148,6 +177,7 @@ pub(in crate::dispatch) fn spawn_live_response_stream( &sender, &mut body_bytes, first_token_ms, + cyber_policy_observer.state_transition_applied(), ) .await; return; @@ -155,6 +185,13 @@ pub(in crate::dispatch) fn spawn_live_response_stream( } } 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, @@ -173,6 +210,7 @@ pub(in crate::dispatch) fn spawn_live_response_stream( &sender, &mut body_bytes, first_token_ms, + cyber_policy_observer.state_transition_applied(), ) .await; return; @@ -186,12 +224,26 @@ pub(in crate::dispatch) fn spawn_live_response_stream( context.account_lease.complete().await; return; }; - finalize_live_response_stream(context, body_text, first_token_ms).await; + 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, @@ -205,8 +257,14 @@ pub(in crate::dispatch) fn spawn_live_response_stream( return; } Err(LiveStreamWriteError::CaptureLimitExceeded) => { - terminate_oversized_live_stream(context, &sender, &mut body_bytes, first_token_ms) - .await; + terminate_oversized_live_stream( + context, + &sender, + &mut body_bytes, + first_token_ms, + cyber_policy_observer.state_transition_applied(), + ) + .await; return; } } @@ -216,7 +274,13 @@ pub(in crate::dispatch) fn spawn_live_response_stream( return; }; - finalize_live_response_stream(context, body_text, first_token_ms).await; + finalize_live_response_stream( + context, + body_text, + first_token_ms, + cyber_policy_observer.state_transition_applied(), + ) + .await; }); ResponseDispatchStream { @@ -362,6 +426,7 @@ async fn terminate_oversized_live_stream( sender: &mpsc::Sender>, body_bytes: &mut Vec, first_token_ms: Option, + cyber_policy_state_transition_applied: bool, ) { tracing::warn!( request_id = %context.request_id, @@ -380,7 +445,13 @@ async fn terminate_oversized_live_stream( context.account_lease.complete().await; return; }; - finalize_live_response_stream(context, body_text, first_token_ms).await; + finalize_live_response_stream( + context, + body_text, + first_token_ms, + cyber_policy_state_transition_applied, + ) + .await; } async fn send_live_response_stream_tail( @@ -447,6 +518,8 @@ 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, @@ -508,6 +581,7 @@ 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 @@ -542,6 +616,12 @@ pub(in crate::dispatch) async fn finalize_live_response_stream( 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) @@ -582,6 +662,12 @@ pub(in crate::dispatch) async fn finalize_live_response_stream( .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); diff --git a/backend/src/upstream/openai/transport/client.rs b/backend/src/upstream/openai/transport/client.rs index 0dd067e6a..a031b0dc4 100644 --- a/backend/src/upstream/openai/transport/client.rs +++ b/backend/src/upstream/openai/transport/client.rs @@ -155,12 +155,43 @@ pub fn is_banned_auth_signal(value: &str) -> bool { || value.contains("banned") } +/// 判断 402 错误正文是否带有 OpenAI 工作区停用标记。 +pub fn is_deactivated_workspace_error_body(value: &str) -> bool { + serde_json::from_str::(value).is_ok_and(|value| { + value.pointer("/detail/code").and_then(Value::as_str) == Some("deactivated_workspace") + }) +} + +/// 判断错误正文是否带有 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) + if (status.as_u16() == 403 + && !is_html_error_body(body) + && !is_cyber_policy_upstream_error(error)) + || (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 20f971ce1..cffae19b3 100644 --- a/backend/src/upstream/openai/transport/mod.rs +++ b/backend/src/upstream/openai/transport/mod.rs @@ -18,7 +18,8 @@ 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_banned_auth_signal, is_banned_upstream_error, is_cyber_policy_error_body, + is_cyber_policy_upstream_error, is_deactivated_workspace_error_body, }, diagnostics::CodexUpstreamDiagnostics, endpoints::{ diff --git a/backend/tests/dispatch/affinity/mod.rs b/backend/tests/dispatch/affinity/mod.rs index 64f10b4b9..ecbc1fbdd 100644 --- a/backend/tests/dispatch/affinity/mod.rs +++ b/backend/tests/dispatch/affinity/mod.rs @@ -1,6 +1,7 @@ use chrono::{Duration, Utc}; use codex_proxy_rs::dispatch::affinity::{ - MAX_CONVERSATION_AFFINITIES, RedisSessionAffinityStore, SessionAffinityEntry, + CyberPolicyFailureSnapshot, MAX_CONVERSATION_AFFINITIES, RedisSessionAffinityStore, + SessionAffinityEntry, }; use redis::AsyncCommands; @@ -246,6 +247,85 @@ 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/service.rs b/backend/tests/dispatch/service.rs index b1254c5f8..cda149e19 100644 --- a/backend/tests/dispatch/service.rs +++ b/backend/tests/dispatch/service.rs @@ -991,6 +991,42 @@ async fn accept_websocket_response_with_authorization_and_message( payload } +async fn accept_cyber_policy_websocket_response_with_authorization( + listener: &TcpListener, + expected_authorization: &'static str, + response_id: &str, +) -> Value { + let (stream, _) = listener.accept().await.unwrap(); + let mut websocket = accept_websocket_with_authorization(stream, expected_authorization).await; + let message = websocket.next().await.unwrap().unwrap(); + let payload = serde_json::from_str::(&message.into_text().unwrap()) + .expect("websocket payload should be json"); + for event in [ + json!({ + "type": "response.created", + "response": {"id": response_id, "status": "in_progress"} + }), + json!({ + "type": "response.output_text.delta", + "delta": "partial output before policy failure" + }), + json!({ + "type": "error", + "error": { + "code": "cyber_policy", + "message": "This request has been flagged for possible cybersecurity risk." + } + }), + ] { + websocket + .send(Message::Text(event.to_string().into())) + .await + .unwrap(); + } + websocket.close(None).await.unwrap(); + payload +} + async fn accept_websocket_with_authorization( stream: TcpStream, expected_authorization: &'static str, diff --git a/backend/tests/dispatch/service/responses_recovery.rs b/backend/tests/dispatch/service/responses_recovery.rs index 2e2c7a77b..08dfc86ab 100644 --- a/backend/tests/dispatch/service/responses_recovery.rs +++ b/backend/tests/dispatch/service/responses_recovery.rs @@ -87,7 +87,187 @@ async fn responses_should_classify_sse_cyber_policy_failure_as_bad_request() { assert_eq!(status, StatusCode::BAD_REQUEST); assert_eq!(body["error"]["type"], "invalid_request_error"); - assert_eq!(body["error"]["code"], "codex_client_error"); + assert_eq!(body["error"]["code"], "cyber_policy"); + assert_eq!( + body["error"]["message"], + "This request has been flagged for possible cybersecurity risk." + ); +} + +#[tokio::test] +async fn responses_http_cyber_policy_should_keep_account_active_and_change_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(403).set_body_json(json!({ + "error": { + "code": "cyber_policy", + "message": "This request has been flagged for possible cybersecurity risk." + } + }))) + .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_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_cyber_next_request", + "input": [{"role": "user", "content": "Security assessment"}], + "stream": false, + "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 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"); + + let second_response = app + .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"); +} + +#[tokio::test] +async fn responses_http_rate_limited_cyber_policy_should_not_retry_in_same_request() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/codex/responses")) + .and(header("authorization", "Bearer access-primary")) + .respond_with(ResponseTemplate::new(429).set_body_json(json!({ + "error": { + "code": "cyber_policy", + "message": "This request has been flagged for possible cybersecurity risk." + } + }))) + .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_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_429_cyber_next_request", + "input": [], + "stream": false, + "use_websocket": false + }); + + let first_response = app + .clone() + .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" + ); + let account_status: (String, Option>) = + sqlx::query_as("select status, quota_cooldown_until from accounts where id = $1") + .bind("acct_primary") + .fetch_one(&pool) + .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] +async fn responses_cyber_rotation_should_return_current_quota_error_when_candidates_exhaust() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/codex/responses")) + .and(header("authorization", "Bearer access-primary")) + .respond_with(ResponseTemplate::new(403).set_body_json(json!({ + "error": {"code": "cyber_policy", "message": "cyber blocked"} + }))) + .expect(1) + .mount(&server) + .await; + Mock::given(method("POST")) + .and(path("/codex/responses")) + .and(header("authorization", "Bearer access-secondary")) + .respond_with(ResponseTemplate::new(402).set_body_json(json!({ + "error": {"code": "billing_limit", "message": "quota exhausted"} + }))) + .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_cyber_then_quota", + "input": [], + "stream": false, + "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 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") + }) + ); + let account_status: (String,) = sqlx::query_as("select status from accounts where id = $1") + .bind("acct_secondary") + .fetch_one(&pool) + .await + .unwrap(); + assert_eq!(account_status.0, "quota_exhausted"); } #[tokio::test] @@ -485,6 +665,111 @@ async fn responses_should_mark_banned_when_401_says_account_deactivated() { assert_eq!(account_status.0, "banned"); } +#[tokio::test] +async fn responses_should_mark_banned_after_402_deactivated_workspace_and_fallback() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/codex/responses")) + .and(header("authorization", "Bearer access-primary")) + .respond_with(ResponseTemplate::new(402).set_body_json(json!({ + "detail": {"code": "deactivated_workspace"} + }))) + .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_AFTER_402_SSE), + ) + .expect(1) + .mount(&server) + .await; + + let (app, api_key, pool, _dir) = test_app_with_two_accounts(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": [{"role": "user", "content": "Say hello"}], + "stream": false, + "use_websocket": false + }) + .to_string(), + )) + .unwrap(), + ) + .await + .unwrap(); + let body = response_json(response).await; + let account_status: (String,) = sqlx::query_as("select status from accounts where id = $1") + .bind("acct_primary") + .fetch_one(&pool) + .await + .unwrap(); + + assert_eq!(body["id"], "resp_after_402"); + assert_eq!(account_status.0, "banned"); +} + +#[tokio::test] +async fn responses_stream_should_mark_banned_after_402_deactivated_workspace_and_fallback() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/codex/responses")) + .and(header("authorization", "Bearer access-primary")) + .respond_with(ResponseTemplate::new(402).set_body_json(json!({ + "detail": {"code": "deactivated_workspace"} + }))) + .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 response = app + .oneshot(responses_json_request( + &api_key, + &json!({ + "model": "gpt-5.5", + "input": [{"role": "user", "content": "Say hello"}], + "stream": true, + "use_websocket": false + }), + )) + .await + .unwrap(); + assert_eq!(response.status(), StatusCode::OK); + let body = response_text(response).await; + let account_status: (String,) = sqlx::query_as("select status from accounts where id = $1") + .bind("acct_primary") + .fetch_one(&pool) + .await + .unwrap(); + + assert!(body.contains("resp_stream_after_402")); + assert_eq!(account_status.0, "banned"); +} + #[tokio::test] async fn responses_should_mark_banned_after_403_and_fallback() { let server = MockServer::start().await; diff --git a/backend/tests/dispatch/service/responses_websocket.rs b/backend/tests/dispatch/service/responses_websocket.rs index 92c0b6b67..f5f7738c5 100644 --- a/backend/tests/dispatch/service/responses_websocket.rs +++ b/backend/tests/dispatch/service/responses_websocket.rs @@ -136,6 +136,311 @@ async fn responses_websocket_stream_should_synthesize_response_failed_when_close assert!(body.contains("stream_disconnected")); } +#[tokio::test] +async fn responses_websocket_cyber_policy_should_change_account_on_next_request() { + 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_payload = accept_cyber_policy_websocket_response_with_authorization( + &listener, + "Bearer access-primary", + "resp_cyber_first", + ) + .await; + + let second_payload = accept_websocket_response_with_authorization_and_message( + &listener, + "Bearer access-secondary", + websocket_completed_response("resp_after_cyber_rotation", 3, 1), + ) + .await; + let third_payload = accept_websocket_response_with_authorization_and_message( + &listener, + "Bearer access-primary", + websocket_completed_response("resp_after_cyber_clear", 3, 1), + ) + .await; + (first_payload, second_payload, third_payload) + }); + let (app, state, api_key, pool, _dir) = + test_app_with_two_accounts_and_state_config(base_url, |config| { + config.auth.rotation_strategy = "round_robin".to_string(); + }) + .await; + let request_body = json!({ + "model": "gpt-5.5", + "prompt_cache_key": "conv_cyber_next_request", + "input": [{"role": "user", "content": "Security assessment"}], + "stream": true, + "use_websocket": true + }); + + let first_response = app + .clone() + .oneshot(responses_json_request(&api_key, &request_body)) + .await + .unwrap(); + assert_eq!(first_response.status(), StatusCode::OK); + let mut first_body_stream = first_response.into_body().into_data_stream(); + let mut first_body = String::new(); + while let Some(chunk) = first_body_stream.next().await { + first_body.push_str(&String::from_utf8_lossy(&chunk.unwrap())); + if first_body.contains("\"code\":\"cyber_policy\"") { + break; + } + } + drop(first_body_stream); + assert!(first_body.contains("partial output before policy failure")); + assert!(first_body.contains("\"code\":\"cyber_policy\"")); + let statuses: Vec<(String, String)> = + sqlx::query_as("select id, status from accounts order by id") + .fetch_all(&pool) + .await + .unwrap(); + assert_eq!( + statuses, + vec![ + ("acct_primary".to_string(), "active".to_string()), + ("acct_secondary".to_string(), "active".to_string()), + ] + ); + + let second_response = app + .clone() + .oneshot(responses_json_request(&api_key, &request_body)) + .await + .unwrap(); + let mut second_body_stream = second_response.into_body().into_data_stream(); + let mut second_body = String::new(); + while let Some(chunk) = second_body_stream.next().await { + second_body.push_str(&String::from_utf8_lossy(&chunk.unwrap())); + if second_body.contains("resp_after_cyber_rotation") { + break; + } + } + drop(second_body_stream); + assert!(second_body.contains("resp_after_cyber_rotation")); + + assert!( + state + .services + .account_pool + .set_status("acct_secondary", AccountStatus::Disabled) + .await + ); + let third_response = app + .oneshot(responses_json_request(&api_key, &request_body)) + .await + .unwrap(); + let third_body = response_text(third_response).await; + let (first_payload, second_payload, third_payload) = upstream.await.unwrap(); + + assert!(third_body.contains("resp_after_cyber_clear")); + assert_eq!(first_payload["type"], "response.create"); + assert_eq!(second_payload["type"], "response.create"); + assert_eq!(third_payload["type"], "response.create"); +} + +#[tokio::test] +async fn responses_websocket_cyber_policy_should_stop_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() + { + payloads.push( + accept_cyber_policy_websocket_response_with_authorization( + &listener, + authorization, + &format!("resp_cyber_cap_{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", + "input": [{"role": "user", "content": "Security assessment"}], + "stream": true, + "use_websocket": true + }); + + for _ in 0..3 { + let response = app + .clone() + .oneshot(responses_json_request(&api_key, &request_body)) + .await + .unwrap(); + assert_eq!(response.status(), StatusCode::OK); + assert!(response_text(response).await.contains("cyber_policy")); + } + three_attempts_rx.await.unwrap(); + + let fourth_response = app + .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." + ); + verify_no_fourth_tx.send(()).unwrap(); + assert_eq!(upstream.await.unwrap().len(), 3); +} + +#[tokio::test] +async fn responses_websocket_cyber_policy_state_should_isolate_session_and_api_key() { + 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 = accept_cyber_policy_websocket_response_with_authorization( + &listener, + "Bearer access-primary", + "resp_cyber_isolation", + ) + .await; + let second = accept_websocket_response_with_authorization_and_message( + &listener, + "Bearer access-primary", + websocket_completed_response("resp_other_session", 3, 1), + ) + .await; + let third = accept_websocket_response_with_authorization_and_message( + &listener, + "Bearer access-primary", + websocket_completed_response("resp_other_api_key", 3, 1), + ) + .await; + (first, second, third) + }); + let (app, state, first_api_key, pool, _dir) = + test_app_with_two_accounts_and_state(base_url).await; + assert!( + state + .services + .account_pool + .set_status("acct_secondary", AccountStatus::Disabled) + .await + ); + let first_session = json!({ + "model": "gpt-5.5", + "prompt_cache_key": "conv_cyber_isolation_a", + "input": [], + "stream": true, + "use_websocket": true + }); + let other_session = json!({ + "model": "gpt-5.5", + "prompt_cache_key": "conv_cyber_isolation_b", + "input": [], + "stream": true, + "use_websocket": true + }); + + let first_response = app + .clone() + .oneshot(responses_json_request(&first_api_key, &first_session)) + .await + .unwrap(); + assert!(response_text(first_response).await.contains("cyber_policy")); + + let other_session_response = app + .clone() + .oneshot(responses_json_request(&first_api_key, &other_session)) + .await + .unwrap(); + assert!( + response_text(other_session_response) + .await + .contains("resp_other_session") + ); + + let second_api_key = insert_client_api_key(&pool).await; + let other_key_response = app + .oneshot(responses_json_request(&second_api_key, &first_session)) + .await + .unwrap(); + assert!( + response_text(other_key_response) + .await + .contains("resp_other_api_key") + ); + let _ = upstream.await.unwrap(); +} + +#[tokio::test] +async fn responses_websocket_cyber_policy_should_not_move_external_previous_response() { + 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 = accept_cyber_policy_websocket_response_with_authorization( + &listener, + "Bearer access-primary", + "resp_cyber_external_previous", + ) + .await; + let second = accept_websocket_response_with_authorization_and_message( + &listener, + "Bearer access-primary", + websocket_completed_response("resp_same_external_previous_account", 3, 1), + ) + .await; + (first, second) + }); + let (app, state, api_key, _pool, _dir) = test_app_with_two_accounts_and_state(base_url).await; + assert!( + state + .services + .account_pool + .set_status("acct_secondary", AccountStatus::Disabled) + .await + ); + let request_body = json!({ + "model": "gpt-5.5", + "prompt_cache_key": "conv_cyber_external_previous", + "previous_response_id": "resp_external_unknown", + "input": [], + "stream": true, + "use_websocket": true + }); + + let first_response = app + .clone() + .oneshot(responses_json_request(&api_key, &request_body)) + .await + .unwrap(); + assert!(response_text(first_response).await.contains("cyber_policy")); + let second_response = app + .oneshot(responses_json_request(&api_key, &request_body)) + .await + .unwrap(); + assert!( + response_text(second_response) + .await + .contains("resp_same_external_previous_account") + ); + let _ = upstream.await.unwrap(); +} + #[tokio::test] async fn responses_websocket_stream_first_error_429_should_retry_fallback_account() { let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); diff --git a/backend/tests/upstream/openai/protocol/responses.rs b/backend/tests/upstream/openai/protocol/responses.rs index f7469f12c..9052792fd 100644 --- a/backend/tests/upstream/openai/protocol/responses.rs +++ b/backend/tests/upstream/openai/protocol/responses.rs @@ -399,6 +399,11 @@ fn build_codex_request_should_prefer_body_context_fields_then_fall_back_to_heade "x-codex-parent-thread-id", " parent-from-header ".parse().unwrap(), ); + headers.insert("session_id", " session-from-header ".parse().unwrap()); + headers.insert( + "conversation_id", + " conversation-from-header ".parse().unwrap(), + ); let codex = build_codex_request(body, &headers, None); @@ -407,6 +412,14 @@ fn build_codex_request_should_prefer_body_context_fields_then_fall_back_to_heade assert_eq!(codex.turn_metadata.as_deref(), Some("meta-from-header")); assert_eq!(codex.include_timing_metrics.as_deref(), Some("true")); assert_eq!(codex.codex_window_id.as_deref(), Some("window-from-header")); + assert_eq!( + codex.client_session_id.as_deref(), + Some("session-from-header") + ); + assert_eq!( + codex.client_conversation_id.as_deref(), + Some("conversation-from-header") + ); assert_eq!( codex.parent_thread_id.as_deref(), Some("parent-from-header") diff --git a/backend/tests/upstream/openai/transport/client.rs b/backend/tests/upstream/openai/transport/client.rs index 3b1c9588e..ecb125049 100644 --- a/backend/tests/upstream/openai/transport/client.rs +++ b/backend/tests/upstream/openai/transport/client.rs @@ -5,11 +5,61 @@ 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, CodexWebSocketPool}; +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, +}; use futures::{SinkExt, StreamExt}; +use reqwest::StatusCode; use tokio::net::TcpListener; use tokio_tungstenite::tungstenite::Message; +#[test] +fn deactivated_workspace_error_body_should_match_detail_code() { + let body = r#"{"detail":{"code":"deactivated_workspace"}}"#; + + assert!(is_deactivated_workspace_error_body(body)); +} + +#[test] +fn deactivated_workspace_error_body_should_not_match_other_json_paths() { + let body = r#"{"error":{"code":"deactivated_workspace"}}"#; + + 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/docs/architecture.md b/docs/architecture.md index 4dc11479d..061190b23 100644 --- a/docs/architecture.md +++ b/docs/architecture.md @@ -38,7 +38,9 @@ PostgreSQL 和 Redis 都是必需依赖。任一连接失败时进程不能完 │ └── dist/ 前端构建产物 ├── deploy/ Dockerfile、Compose 和部署模板 ├── docs/ -│ └── architecture.md 当前架构说明 +│ ├── architecture.md 当前架构说明 +│ └── request-lifecycle-controller-design.md +│ Responses 生命周期控制器审计与目标设计 ├── release/ 版本、目标平台和发布脚本 └── skills/ 仓库内开发约束 ``` @@ -232,6 +234,7 @@ dispatch/ │ ├── account_failure.rs │ ├── auth.rs │ ├── cloudflare.rs +│ ├── cyber_policy.rs │ ├── exhaustion.rs │ └── history.rs └── stream/ @@ -379,13 +382,15 @@ bootstrap/ 2. `client::auth` 从 `Authorization: Bearer sk_...` 取出完整 Key，并在 PostgreSQL 校验是否存在且启用。WebSocket 在 HTTP upgrade 前完成鉴权。 3. API 层解析 HTTP JSON 或官方 `response.create` 文本帧，提取客户端 IP、User-Agent 和 transport-only 参数。 4. `models` 解析模型别名；`dispatch` 建立会话变体标识和 previous response 恢复计划。 -5. `AccountAttemptLedger` 按当前调度策略冻结完整候选顺序。会话亲和账号可以排在前面，但不会绕过可用性检查。 +5. `AccountAttemptLedger` 按当前调度策略冻结完整候选顺序。会话亲和账号可以排在前面，但不会绕过可用性检查；会话级 `cyber_policy` 负亲和账号在冻结前排除。 6. 账号池租用并发槽位，必要时先刷新 quota；请求间隔在发送上游前执行。 7. `upstream_call` 注入账号 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，还不是统一的生命周期控制器模型。工程审计确认可以收敛为 Request、Attempt、Stream 三层嵌套洋葱；详细的现状证据、目标类型、顺序、迁移步骤和验收条件见 [Responses 请求生命周期控制器设计](request-lifecycle-controller-design.md)。该设计限定在 `dispatch` 内部，不新增顶层 `application` 层，也不替代鉴权、请求 ID、连接排空等 Tower/Axum 中间件。 + 会话压缩不使用独立端点。Codex 通过 `/v1/responses` 请求中的 `compaction_trigger` 和 `request_kind=compaction` 表达压缩语义；代理沿用同一条调度链路，并在遥测元数据中记录 `compact`、`requestKind` 和 `subagentKind`。 ## 账号调度与换号 @@ -408,6 +413,10 @@ bootstrap/ 这些失败只更新当前账号的状态、quota 或冷却时间，不启用独立熔断器，也不固定切到某一种策略。下一账号仍来自本次请求按当前设置生成的调度顺序。 +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` 的请求不使用该负亲和策略，避免把只有原账号持有的上游历史错误地路由到其他账号；客户端需要在新请求中重发完整历史后才能换号。 + 请求本身不可修复的 4xx、协议错误或已经向客户端提交内容后的流错误不会通过换号掩盖。候选耗尽后，返回最后一个有业务意义的账号错误；没有候选时返回无可用账号错误。 ## previous response 与流式恢复 @@ -486,6 +495,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: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 new file mode 100644 index 000000000..7344ab63f --- /dev/null +++ b/docs/request-lifecycle-controller-design.md @@ -0,0 +1,228 @@ +# Responses 请求生命周期控制器设计 + +状态：工程审计结论与目标设计，尚未实现。 + +本文回答一个具体问题：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、首输出和策略事实 | + +领域规则并非全部重复。`HistoryRecoveryPlan`、`AccountAttemptLedger`、quota 验证、账号失败隔离和 `CyberPolicyRecovery` 已经分别集中。真正分散的是“何时调用、如何排序、如何把结果传到下一阶段”的生命周期布线。新增一个跨阶段策略时，仍可能同时修改非流式主循环、流式建连、live finalizer 和 API encoder。 + +当前取消路径也说明普通 handler 后置钩子不够：HTTP body 或下游 WebSocket 被丢弃后，live task 会收到取消并释放账号槽位，但不会经过完整的流终态结算。此时取消事实、部分用量、Cookie、rate-limit、turn state 和遥测可能缺失。 + +## 设计边界 + +该控制器模型只属于 `dispatch`，不新增顶层 `application` 层。 + +- API 层继续负责鉴权提取、HTTP/WebSocket 协议解析和最终编码。 +- request ID、访问日志、连接排空继续使用 Tower/Axum middleware。 +- `fleet` 继续拥有账号池、quota 和账号状态。 +- `upstream` 继续拥有 Codex 协议、HTTP/SSE 和 WebSocket 传输。 +- `telemetry` 继续保存已经确定的事实，不参与调度决策。 +- controller 不成为运行时任意安装、任意排序的插件系统；顺序是静态架构契约。 + +## 三层生命周期 + +```text +RequestLifecycle + enter: normalize/model -> local identity -> history/session plan -> candidate snapshot + + AttemptLifecycle (每个候选账号一次) + enter: acquire -> quota verify -> history for account -> interval -> scoped identity -> call + leave: normalize observation -> classify -> account effects -> retry/fail decision -> trace + + StreamLifecycle (响应提交后) + enter: canonical decoder -> observers -> downstream encoder + leave: terminal/cancel -> policy effects -> affinity/usage/telemetry -> lease release + + start: produce DispatchStart for API + leave: finalize one FinalOutcome immediately or when StreamHandle ends +``` + +进入按声明顺序执行，退出按逆序执行。controller 在进入阶段短路时，只退出已经成功进入的 controller。 + +三层的责任不同： + +- Request 决定请求不变量和本次请求允许尝试的账号集合。 +- Attempt 决定单账号调用是否被接受、同账号恢复、换下一个账号或立即失败。 +- Stream 只处理已经越过提交边界的事件、下游交付和终态；此时不能再产生换号 decision。 + +流式请求有两个不同时间点，不能共用一个“最终结果”类型。API 在首帧提交边界只拿到 `DispatchStart::Streaming(StreamHandle)`，Request/Attempt 已进入 controller 的退出状态随 handle 转移给后台 lifecycle；真正的 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, + candidates: AccountAttemptLedger, + trace: ResponseDispatchTrace, +} + +struct AttemptContext { + lease: AccountLease, + attempt_request: CodexResponsesRequest, + transport: CodexBackendTransport, + attempt: ResponseDispatchAttempt, +} +``` + +过程控制使用显式枚举，不使用 `bool` 表达多种语义： + +```rust +enum AttemptOutcome { + Accepted(EstablishedResponse), + RetrySameAccount { reason: RetryReason }, + RetryNextCandidate { failure: AttemptFailure }, + ReturnFailure(DispatchFailure), +} + +enum StreamTerminal { + Completed, + Incomplete, + Failed(StreamFailure), + Disconnected(StreamFailure), + InvalidProtocol(StreamFailure), + DownstreamCancelled, + Shutdown, +} + +enum DispatchStart { + Immediate(ResponsePayload), + Streaming(StreamHandle), + Failed(ClientFailure), +} + +enum FinalOutcome { + Completed(ResponseFacts), + Incomplete(ResponseFacts), + Failed(ClientFailure), + Cancelled(CancelFacts), +} +``` + +账号状态、会话恢复和客户端错误映射只分类一次。API 立即消费 `DispatchStart`：非流式 payload 直接编码，stream handle 绑定到对应 transport body，失败由统一 `ClientFailure` 映射到 HTTP JSON、SSE 或 WebSocket。后台 lifecycle 只消费 `FinalOutcome` 做领域结算，不再重新识别 `cyber_policy`、quota、banned 或 history failure。 + +## controller 契约 + +controller 分为两类： + +- Decision controller：可以返回 retry、fail 或 commit 等显式 decision。 +- Observer controller：只消费已经确定的事实。其 Redis、数据库或遥测失败默认是 best-effort，不能改写已经确定或已经发给客户端的 outcome。 + +错误也分为两类： + +- Critical：无法建立正确请求语义，转为 `DispatchFailure` 并触发逆序退出。 +- BestEffort：记录警告并 fail-open，不阻塞正常请求。 + +同层 controller 不直接调用另一个 controller，也不自行选择下一账号。pipeline 是唯一解释 decision 的组件。这样模块只拥有规则，编排器只拥有顺序和控制流。 + +## 固定顺序 + +### Request enter + +1. 请求规范化、模型解析和本地身份。 +2. history plan。 +3. 普通 affinity 与会话 recovery 查询；无依赖项可在同一 stage 并行。 +4. 合并为 `RoutePlan`，明确 preferred、required、excluded 和 fail-fast。 +5. 冻结 `AccountAttemptLedger`。 +6. 建立请求级 trace/telemetry 外层作用域。 + +### Attempt enter/leave + +1. 获取候选账号 lease。 +2. 必要时验证 quota。 +3. 为该账号准备 history；不能跨账号时返回 typed failure。 +4. 等待账号请求间隔。 +5. 生成账号作用域 identity、Cookie 和上游上下文。 +6. 发起上游调用。 +7. 统一规范化 HTTP、SSE 首帧和 WebSocket 错误为 `AttemptObservation`。 +8. 按唯一优先级分类 history、session policy、账号隔离和 transport retry。 +9. 产生 `AttemptOutcome` 并记录 attempt trace。 + +classifier 优先级只能存在一处。例如 `cyber_policy` 必须先于通用 403、quota 或 5xx 分类；其他模块消费分类后的类型，不能再次调用 `is_cyber_policy_*`。 + +### Stream enter/leave + +1. 上游 SSE 或 WebSocket 事件转换为一次性的 `CanonicalStreamEvent`。 +2. 保留 raw bytes 供 HTTP SSE 透明转发；typed event 同时提供给 observer 和下游 WebSocket encoder。 +3. 在首个可转发事件处记录 commit；commit 后类型上不再允许 retry。 +4. 影响下一请求路由的关键状态在 terminal 向下游转发前完成有界更新。 +5. terminal、上游断开、协议错误、下游取消和 shutdown 都转换为 `StreamTerminal`。 +6. 统一 finalizer 逆序执行 controller，并且只执行一次。 +7. 最后释放账号 lease。 + +`Drop` 不能执行异步 Redis、数据库或遥测操作。response body 的 Drop 只发送取消信号；持有完整生命周期的受监督 task 必须显式调用 `finalize_once(StreamTerminal::DownstreamCancelled).await`。Drop guard 只作为 panic/abort 的资源释放兜底。 + +## Rust 落地选择 + +核心热路径不需要 `Vec>`： + +- 非流式与流式模式用泛型或 sealed enum 做静态 dispatch，共享一个 request prelude 和 `AttemptRunner`。 +- controller 集合使用具体 `ControllerSet` 或 sealed controller enum，顺序由代码和 contract test 固定。 +- 流 decoder 和状态机使用具体类型，不能对每个 token 走 async trait/vtable。 +- 只有未来确实需要异构、可配置 controller 集合时，才在注册边界使用 `Arc`；每请求私有状态仍应封装在 controller scope 内。 +- typestate 只用于 Prepared、Routed、Committed 等少数不可逆边界，不把每个小阶段都编码为泛型状态。 + +与 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。 + +## 性能约束 + +- 独立且都需要的入口查询并行执行，不能把多个 Redis RTT 串行叠加。 +- best-effort controller 必须有有界超时和明确 fail-open 行为。 +- 逐事件 observer 只做同步、低成本状态更新；Redis、数据库和遥测 I/O 只在必要的 terminal/finalize 阶段执行。 +- canonical event 同时保留 raw payload，避免为透明转发再次序列化。 +- 是否引入动态分发以基准和可维护性为依据，不凭感觉优化。 + +## 验收条件 + +- 新增一个跨阶段恢复策略时，只新增 controller、静态注册和本模块测试，不同时修改 `service.rs`、`stream/lifecycle.rs`、`stream/live.rs` 和 API handler。 +- 非流式与流式共享同一个 request prelude 和 `AttemptRunner`。 +- 同一上游失败在 complete、prefetch 和 live terminal 得到一致分类。 +- commit 后无法构造 retry decision，不能拼接两个账号的输出。 +- enter 顺序、leave 逆序、short-circuit 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/`。