mirror of
https://github.com/zyycn/codex-proxy-rs.git
synced 2026-11-06 08:06:19 +00:00
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 <noreply@openai.com>
This commit is contained in:
@@ -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;
|
||||
}
|
||||
|
||||
@@ -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")
|
||||
|
||||
@@ -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<Option<CyberPolicySessionState>, 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<CyberPolicySessionState, SessionAffinityError> {
|
||||
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<bool, SessionAffinityError> {
|
||||
Ok(self
|
||||
.store
|
||||
.clear_cyber_policy_state(session_key, expected_revision)
|
||||
.await?)
|
||||
}
|
||||
|
||||
async fn entry(&self, response_id: &str, now: DateTime<Utc>) -> Option<SessionAffinityEntry> {
|
||||
match self.store.get(response_id, now, self.ttl).await {
|
||||
Ok(entry) => entry,
|
||||
|
||||
@@ -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<Option<CyberPolicySessionState>> {
|
||||
let mut connection = self.redis.manager();
|
||||
let value: Option<String> = 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<CyberPolicySessionState> {
|
||||
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<bool> {
|
||||
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"))
|
||||
}
|
||||
|
||||
@@ -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<String>,
|
||||
}
|
||||
|
||||
/// 会话级 `cyber_policy` 换号状态。
|
||||
#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub struct CyberPolicySessionState {
|
||||
pub failed_account_ids: Vec<String>,
|
||||
pub last_failure: Option<CyberPolicyFailureSnapshot>,
|
||||
#[serde(default)]
|
||||
pub revision: String,
|
||||
}
|
||||
|
||||
/// 会话亲和性条目。
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
|
||||
@@ -26,8 +26,10 @@ impl AccountAttemptLedger {
|
||||
pub(in crate::dispatch) async fn freeze(
|
||||
account_pool: &AccountPoolService,
|
||||
request: &AccountAcquireRequest,
|
||||
excluded_account_ids: &BTreeSet<String>,
|
||||
) -> 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(),
|
||||
|
||||
@@ -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)
|
||||
)
|
||||
}
|
||||
|
||||
@@ -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
|
||||
}
|
||||
|
||||
|
||||
@@ -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<String>,
|
||||
state: CyberPolicySessionState,
|
||||
}
|
||||
|
||||
impl CyberPolicyRoutingPlan {
|
||||
pub(in crate::dispatch) fn excluded_account_ids(&self) -> BTreeSet<String> {
|
||||
self.state.failed_account_ids.iter().cloned().collect()
|
||||
}
|
||||
|
||||
pub(in crate::dispatch) fn exhausted_failure(&self) -> Option<ResponsesSseFailure> {
|
||||
if self.state.failed_account_ids.len() < MAX_ROTATED_ACCOUNTS {
|
||||
return None;
|
||||
}
|
||||
self.last_failure()
|
||||
}
|
||||
|
||||
pub(in crate::dispatch) fn last_failure(&self) -> Option<ResponsesSseFailure> {
|
||||
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<SessionAffinityService>,
|
||||
}
|
||||
|
||||
/// 在 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::<Value>(&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::<Value>(&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::<Value>(&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<SessionAffinityService>) -> 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<String> {
|
||||
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<ResponsesSseFailure> {
|
||||
let CodexClientError::Upstream { body, .. } = error else {
|
||||
return None;
|
||||
};
|
||||
if !is_cyber_policy_upstream_error(error) {
|
||||
return None;
|
||||
}
|
||||
let value = serde_json::from_str::<Value>(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()
|
||||
}
|
||||
@@ -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;
|
||||
|
||||
@@ -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<String> = 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 {
|
||||
|
||||
@@ -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<String> = 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,
|
||||
|
||||
@@ -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<Result<Bytes, ResponseDispatchStreamError>>,
|
||||
body_bytes: &mut Vec<u8>,
|
||||
first_token_ms: Option<i64>,
|
||||
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<AccountPoolService>,
|
||||
pub(in crate::dispatch) account_lease: AccountLease,
|
||||
pub(in crate::dispatch) session_affinity: Arc<SessionAffinityService>,
|
||||
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<Recorder>,
|
||||
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<i64>,
|
||||
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);
|
||||
|
||||
@@ -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>(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>(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))
|
||||
)
|
||||
}
|
||||
|
||||
|
||||
@@ -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::{
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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::<Value>(&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,
|
||||
|
||||
@@ -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<chrono::DateTime<Utc>>) =
|
||||
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;
|
||||
|
||||
@@ -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();
|
||||
|
||||
@@ -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")
|
||||
|
||||
@@ -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();
|
||||
|
||||
+12
-2
@@ -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:<response_id>` | 响应归属元数据 | 4 小时 |
|
||||
| `cpr:affinity:v3:conv:<conversation_id>` | 会话响应 ZSET 索引,最多 128 项 | 4 小时 |
|
||||
| `cpr:affinity:v3:account:<account_id>` | 账号响应 ZSET 索引,最多 16384 项 | 4 小时 |
|
||||
| `cpr:affinity:v3:cyber:<session_hash>` | 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 小时 |
|
||||
|
||||
@@ -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<String, Any>` 或无边界的 `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<Box<dyn Middleware>>`:
|
||||
|
||||
- 非流式与流式模式用泛型或 sealed enum 做静态 dispatch,共享一个 request prelude 和 `AttemptRunner`。
|
||||
- controller 集合使用具体 `ControllerSet` 或 sealed controller enum,顺序由代码和 contract test 固定。
|
||||
- 流 decoder 和状态机使用具体类型,不能对每个 token 走 async trait/vtable。
|
||||
- 只有未来确实需要异构、可配置 controller 集合时,才在注册边界使用 `Arc<dyn Controller>`;每请求私有状态仍应封装在 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/`。
|
||||
Reference in New Issue
Block a user