mirror of
https://github.com/zyycn/codex-proxy-rs.git
synced 2026-11-07 00:05:49 +00:00
Compare commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
71d87c5ab9 | ||
|
|
e23689ae22 |
@@ -29,7 +29,7 @@ Codex Proxy RS 为多个 Codex 账号提供统一的 Responses API。客户端
|
||||
- 单个账号发生额度、认证、模型或传输故障时,继续使用账号池中的其他可用账号。
|
||||
- 支持跨账号会话续接,并隔离不同账号的会话状态。
|
||||
- 提供账号、API Key、模型映射、运行参数、用量、错误记录和系统更新管理界面。
|
||||
- 使用 PostgreSQL 保存权威业务数据,使用 Redis 保存会话、租约、模型快照和会话亲和数据。
|
||||
- 使用 PostgreSQL 保存权威业务数据,使用 Redis 保存管理会话、租约、模型快照和响应归属。
|
||||
- 内置日志轮转、遥测保留期清理和 PostgreSQL/Redis 健康检查。
|
||||
|
||||
除账号调度和会话恢复外,请求与响应尽量保持原有语义。
|
||||
@@ -180,7 +180,7 @@ Compose 默认只监听宿主机 `127.0.0.1`。需要从其他设备访问时,
|
||||
| 数据 | 默认位置 |
|
||||
| --- | --- |
|
||||
| 账号、API Key、设置、用量和错误记录 | Compose 命名卷 `postgres-data` |
|
||||
| 管理会话、租约、模型快照和会话亲和 | Compose 命名卷 `redis-data` |
|
||||
| 管理会话、租约、模型快照和响应归属 | Compose 命名卷 `redis-data` |
|
||||
| 身份派生密钥和更新状态 | `.runtime/data` |
|
||||
| 文件日志 | `.runtime/logs` |
|
||||
| 启动配置 | `.runtime/config.yaml` |
|
||||
|
||||
@@ -1,5 +1,7 @@
|
||||
//! 官方 Responses WebSocket 入站协议。
|
||||
|
||||
use std::sync::Arc;
|
||||
|
||||
use axum::{
|
||||
body::Body,
|
||||
extract::{Request, State},
|
||||
@@ -22,7 +24,7 @@ use tokio_tungstenite::{
|
||||
|
||||
use crate::{
|
||||
api::{AppState, middleware::request_id::RequestId},
|
||||
dispatch::service::ResponseDispatchStream,
|
||||
dispatch::{recovery::history::sanitize_replay_items, service::ResponseDispatchStream},
|
||||
upstream::openai::protocol::{
|
||||
sse::{SseEvent, SseEventDecoder},
|
||||
websocket::{is_terminal_websocket_event, websocket_event_type},
|
||||
@@ -131,6 +133,135 @@ struct ResponsesWebSocketSession {
|
||||
connection_id: String,
|
||||
}
|
||||
|
||||
struct ConnectionReplayState {
|
||||
last_response_id: Option<String>,
|
||||
input: Option<Arc<Vec<Value>>>,
|
||||
}
|
||||
|
||||
enum PendingReplayUpdate {
|
||||
Replace(Vec<Value>),
|
||||
Append(Vec<Value>),
|
||||
Unavailable,
|
||||
}
|
||||
|
||||
struct CompletedTurn {
|
||||
response_id: String,
|
||||
output: Vec<Value>,
|
||||
}
|
||||
|
||||
#[derive(Default)]
|
||||
struct CompletedTurnCollector {
|
||||
completed_items: Vec<Value>,
|
||||
completed: Option<CompletedTurn>,
|
||||
}
|
||||
|
||||
enum ForwardDispatchOutcome {
|
||||
Continue(Option<CompletedTurn>),
|
||||
Disconnect,
|
||||
}
|
||||
|
||||
impl ConnectionReplayState {
|
||||
fn new() -> Self {
|
||||
Self {
|
||||
last_response_id: None,
|
||||
input: None,
|
||||
}
|
||||
}
|
||||
|
||||
fn prepare(
|
||||
&self,
|
||||
request: &mut crate::upstream::openai::protocol::responses::CodexResponsesRequest,
|
||||
) -> PendingReplayUpdate {
|
||||
let turn_input = request.input().to_vec();
|
||||
match request.previous_response_id() {
|
||||
None => {
|
||||
request.local_replay_input = None;
|
||||
PendingReplayUpdate::Replace(turn_input)
|
||||
}
|
||||
Some(previous_response_id)
|
||||
if self.last_response_id.as_deref() == Some(previous_response_id)
|
||||
&& self.input.is_some() =>
|
||||
{
|
||||
request.local_replay_input.clone_from(&self.input);
|
||||
PendingReplayUpdate::Append(turn_input)
|
||||
}
|
||||
Some(_) => {
|
||||
request.local_replay_input = None;
|
||||
PendingReplayUpdate::Unavailable
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn commit(&mut self, pending: PendingReplayUpdate, completed: CompletedTurn) {
|
||||
self.last_response_id = Some(completed.response_id);
|
||||
match pending {
|
||||
PendingReplayUpdate::Replace(input) => self.replace(input, completed.output),
|
||||
PendingReplayUpdate::Append(input) => self.append(input, completed.output),
|
||||
PendingReplayUpdate::Unavailable => self.clear_input(),
|
||||
}
|
||||
}
|
||||
|
||||
fn replace(&mut self, mut input: Vec<Value>, mut output: Vec<Value>) {
|
||||
sanitize_replay_items(&mut input);
|
||||
sanitize_replay_items(&mut output);
|
||||
input.append(&mut output);
|
||||
self.input = Some(Arc::new(input));
|
||||
}
|
||||
|
||||
fn append(&mut self, mut input: Vec<Value>, mut output: Vec<Value>) {
|
||||
sanitize_replay_items(&mut input);
|
||||
sanitize_replay_items(&mut output);
|
||||
input.append(&mut output);
|
||||
let Some(history) = self.input.as_mut() else {
|
||||
return;
|
||||
};
|
||||
Arc::make_mut(history).append(&mut input);
|
||||
}
|
||||
|
||||
fn clear_input(&mut self) {
|
||||
self.input = None;
|
||||
}
|
||||
}
|
||||
|
||||
impl CompletedTurnCollector {
|
||||
fn observe(&mut self, event: &SseEvent) {
|
||||
let Ok(value) = serde_json::from_str::<Value>(&event.data) else {
|
||||
return;
|
||||
};
|
||||
let event_type = event
|
||||
.event
|
||||
.as_deref()
|
||||
.or_else(|| value.get("type").and_then(Value::as_str));
|
||||
match event_type {
|
||||
Some("response.output_item.done") => {
|
||||
if let Some(item) = value.get("item") {
|
||||
self.completed_items.push(item.clone());
|
||||
}
|
||||
}
|
||||
Some("response.completed") => {
|
||||
let Some(response_id) = value
|
||||
.pointer("/response/id")
|
||||
.and_then(Value::as_str)
|
||||
.map(ToString::to_string)
|
||||
else {
|
||||
return;
|
||||
};
|
||||
let output = value
|
||||
.pointer("/response/output")
|
||||
.and_then(Value::as_array)
|
||||
.filter(|output| !output.is_empty())
|
||||
.cloned()
|
||||
.unwrap_or_else(|| self.completed_items.clone());
|
||||
self.completed = Some(CompletedTurn {
|
||||
response_id,
|
||||
output,
|
||||
});
|
||||
}
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
async fn serve_responses_websocket(
|
||||
mut socket: ResponsesWebSocket,
|
||||
session: ResponsesWebSocketSession,
|
||||
@@ -140,6 +271,7 @@ async fn serve_responses_websocket(
|
||||
"Responses WebSocket connected"
|
||||
);
|
||||
let mut request_count = 0u64;
|
||||
let mut replay_state = ConnectionReplayState::new();
|
||||
|
||||
while let Some(message) = socket.next().await {
|
||||
let payload = match message {
|
||||
@@ -190,7 +322,7 @@ async fn serve_responses_websocket(
|
||||
continue;
|
||||
}
|
||||
};
|
||||
let prepared = match prepare_responses_request(
|
||||
let mut prepared = match prepare_responses_request(
|
||||
&session.state,
|
||||
body,
|
||||
&session.headers,
|
||||
@@ -229,9 +361,14 @@ async fn serve_responses_websocket(
|
||||
continue;
|
||||
}
|
||||
};
|
||||
let pending_replay = replay_state.prepare(&mut prepared.request);
|
||||
|
||||
if !dispatch_response_create(&mut socket, &session, request_id, prepared).await {
|
||||
break;
|
||||
match dispatch_response_create(&mut socket, &session, request_id, prepared).await {
|
||||
ForwardDispatchOutcome::Continue(Some(completed)) => {
|
||||
replay_state.commit(pending_replay, completed);
|
||||
}
|
||||
ForwardDispatchOutcome::Continue(None) => {}
|
||||
ForwardDispatchOutcome::Disconnect => break,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -247,7 +384,7 @@ async fn dispatch_response_create(
|
||||
session: &ResponsesWebSocketSession,
|
||||
request_id: RequestId,
|
||||
prepared: PreparedResponsesRequest,
|
||||
) -> bool {
|
||||
) -> ForwardDispatchOutcome {
|
||||
let PreparedResponsesRequest {
|
||||
mut request,
|
||||
requested_model,
|
||||
@@ -270,7 +407,11 @@ async fn dispatch_response_create(
|
||||
Ok(stream) => stream,
|
||||
Err(error) => {
|
||||
let event = responses_websocket_dispatch_error_event(&error, request_id.as_str());
|
||||
return send_text(socket, event).await;
|
||||
return if send_text(socket, event).await {
|
||||
ForwardDispatchOutcome::Continue(None)
|
||||
} else {
|
||||
ForwardDispatchOutcome::Disconnect
|
||||
};
|
||||
}
|
||||
};
|
||||
forward_dispatch_stream(socket, stream, request_id.as_str()).await
|
||||
@@ -280,7 +421,7 @@ async fn forward_dispatch_stream(
|
||||
socket: &mut ResponsesWebSocket,
|
||||
stream: ResponseDispatchStream,
|
||||
request_id: &str,
|
||||
) -> bool {
|
||||
) -> ForwardDispatchOutcome {
|
||||
let ResponseDispatchStream {
|
||||
body,
|
||||
response_headers,
|
||||
@@ -291,12 +432,13 @@ async fn forward_dispatch_stream(
|
||||
)
|
||||
.await
|
||||
{
|
||||
return false;
|
||||
return ForwardDispatchOutcome::Disconnect;
|
||||
}
|
||||
|
||||
let mut body = body;
|
||||
let mut decoder = SseEventDecoder::default();
|
||||
let mut terminal_sent = false;
|
||||
let mut completed = CompletedTurnCollector::default();
|
||||
loop {
|
||||
let next = if terminal_sent {
|
||||
ActiveResponseInput::Upstream(body.next().await)
|
||||
@@ -322,19 +464,31 @@ async fn forward_dispatch_stream(
|
||||
)
|
||||
.await;
|
||||
drop(body);
|
||||
return sent;
|
||||
return if sent {
|
||||
ForwardDispatchOutcome::Continue(None)
|
||||
} else {
|
||||
ForwardDispatchOutcome::Disconnect
|
||||
};
|
||||
}
|
||||
};
|
||||
if !forward_sse_events(socket, events, request_id, &mut terminal_sent).await {
|
||||
return false;
|
||||
if !forward_sse_events(
|
||||
socket,
|
||||
events,
|
||||
request_id,
|
||||
&mut terminal_sent,
|
||||
&mut completed,
|
||||
)
|
||||
.await
|
||||
{
|
||||
return ForwardDispatchOutcome::Disconnect;
|
||||
}
|
||||
}
|
||||
ActiveResponseInput::Upstream(Some(Err(error))) => {
|
||||
tracing::warn!(request_id, error = %error, "Responses dispatch stream failed");
|
||||
if terminal_sent {
|
||||
return true;
|
||||
return ForwardDispatchOutcome::Continue(completed.completed);
|
||||
}
|
||||
return send_error(
|
||||
return if send_error(
|
||||
socket,
|
||||
StatusCode::BAD_GATEWAY,
|
||||
"server_error",
|
||||
@@ -342,14 +496,19 @@ async fn forward_dispatch_stream(
|
||||
STREAM_DISCONNECTED_MESSAGE,
|
||||
request_id,
|
||||
)
|
||||
.await;
|
||||
.await
|
||||
{
|
||||
ForwardDispatchOutcome::Continue(None)
|
||||
} else {
|
||||
ForwardDispatchOutcome::Disconnect
|
||||
};
|
||||
}
|
||||
ActiveResponseInput::Upstream(None) => {
|
||||
let events = match decoder.finish() {
|
||||
Ok(events) => events,
|
||||
Err(error) => {
|
||||
tracing::warn!(request_id, error = %error, "failed to finish Responses SSE decoding");
|
||||
return send_error(
|
||||
return if send_error(
|
||||
socket,
|
||||
StatusCode::BAD_GATEWAY,
|
||||
"server_error",
|
||||
@@ -357,16 +516,29 @@ async fn forward_dispatch_stream(
|
||||
INVALID_UPSTREAM_EVENT_MESSAGE,
|
||||
request_id,
|
||||
)
|
||||
.await;
|
||||
.await
|
||||
{
|
||||
ForwardDispatchOutcome::Continue(None)
|
||||
} else {
|
||||
ForwardDispatchOutcome::Disconnect
|
||||
};
|
||||
}
|
||||
};
|
||||
if !forward_sse_events(socket, events, request_id, &mut terminal_sent).await {
|
||||
return false;
|
||||
if !forward_sse_events(
|
||||
socket,
|
||||
events,
|
||||
request_id,
|
||||
&mut terminal_sent,
|
||||
&mut completed,
|
||||
)
|
||||
.await
|
||||
{
|
||||
return ForwardDispatchOutcome::Disconnect;
|
||||
}
|
||||
if terminal_sent {
|
||||
return true;
|
||||
return ForwardDispatchOutcome::Continue(completed.completed);
|
||||
}
|
||||
return send_error(
|
||||
return if send_error(
|
||||
socket,
|
||||
StatusCode::BAD_GATEWAY,
|
||||
"server_error",
|
||||
@@ -374,23 +546,30 @@ async fn forward_dispatch_stream(
|
||||
STREAM_DISCONNECTED_MESSAGE,
|
||||
request_id,
|
||||
)
|
||||
.await;
|
||||
.await
|
||||
{
|
||||
ForwardDispatchOutcome::Continue(None)
|
||||
} else {
|
||||
ForwardDispatchOutcome::Disconnect
|
||||
};
|
||||
}
|
||||
ActiveResponseInput::Downstream(Some(Ok(
|
||||
Message::Ping(_) | Message::Pong(_) | Message::Frame(_),
|
||||
))) => {}
|
||||
ActiveResponseInput::Downstream(Some(Ok(Message::Close(_))) | None) => return false,
|
||||
ActiveResponseInput::Downstream(Some(Ok(Message::Close(_))) | None) => {
|
||||
return ForwardDispatchOutcome::Disconnect;
|
||||
}
|
||||
ActiveResponseInput::Downstream(Some(Ok(Message::Text(_) | Message::Binary(_)))) => {
|
||||
close_for_protocol_violation(
|
||||
socket,
|
||||
"Only one response.create may be active per connection",
|
||||
)
|
||||
.await;
|
||||
return false;
|
||||
return ForwardDispatchOutcome::Disconnect;
|
||||
}
|
||||
ActiveResponseInput::Downstream(Some(Err(error))) => {
|
||||
tracing::info!(request_id, error = %error, "Responses WebSocket client disconnected during response");
|
||||
return false;
|
||||
return ForwardDispatchOutcome::Disconnect;
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -406,6 +585,7 @@ async fn forward_sse_events(
|
||||
events: Vec<SseEvent>,
|
||||
request_id: &str,
|
||||
terminal_sent: &mut bool,
|
||||
completed: &mut CompletedTurnCollector,
|
||||
) -> bool {
|
||||
for event in events {
|
||||
let Some(event_type) = websocket_event_type(&event.data) else {
|
||||
@@ -424,6 +604,7 @@ async fn forward_sse_events(
|
||||
.await;
|
||||
return false;
|
||||
};
|
||||
completed.observe(&event);
|
||||
if !send_text(socket, event.data).await {
|
||||
return false;
|
||||
}
|
||||
|
||||
@@ -17,6 +17,7 @@ use super::{
|
||||
const LOGIN_FAILURE_WINDOW: Duration = Duration::minutes(15);
|
||||
const LOGIN_LOCK_DURATION: Duration = Duration::minutes(15);
|
||||
const LOGIN_MAX_FAILURES: u32 = 5;
|
||||
const LOGIN_FAILURE_SOURCE_LIMIT: usize = 10_000;
|
||||
|
||||
/// 管理员会话服务。
|
||||
#[derive(Clone)]
|
||||
@@ -134,6 +135,9 @@ impl SessionService {
|
||||
let login_locked = {
|
||||
let mut failures = self.login_failures.lock().await;
|
||||
retain_active_login_failures(&mut failures, now);
|
||||
if !failures.contains_key(source) && failures.len() >= LOGIN_FAILURE_SOURCE_LIMIT {
|
||||
return Err(SessionError::LoginThrottled);
|
||||
}
|
||||
failures
|
||||
.get(source)
|
||||
.and_then(|state| state.locked_until)
|
||||
|
||||
@@ -88,7 +88,6 @@ pub(in crate::dispatch) async fn record_response_affinity(
|
||||
} else {
|
||||
crate::upstream::openai::protocol::responses::PreviousResponseScope::ConnectionLocal
|
||||
},
|
||||
replay: history.completed_replay(original_request, &metadata.output),
|
||||
created_at: Utc::now(),
|
||||
};
|
||||
if let Err(error) = session_affinity
|
||||
|
||||
@@ -9,9 +9,6 @@ use super::{
|
||||
};
|
||||
|
||||
const DEFAULT_SESSION_AFFINITY_TTL_SECS: i64 = 4 * 60 * 60;
|
||||
pub const MAX_REPLAY_SNAPSHOT_BYTES: u64 = 2 * 1024 * 1024;
|
||||
pub const MAX_REPLAY_SESSION_BYTES: u64 = 16 * 1024 * 1024;
|
||||
pub const MAX_REPLAY_DEPTH: u16 = 128;
|
||||
|
||||
/// 运行时会话亲和性服务。
|
||||
#[derive(Clone)]
|
||||
@@ -49,36 +46,6 @@ impl SessionAffinityService {
|
||||
self.entry(response_id, now).await
|
||||
}
|
||||
|
||||
pub async fn replay_input(
|
||||
&self,
|
||||
response_id: &str,
|
||||
entry: &SessionAffinityEntry,
|
||||
now: DateTime<Utc>,
|
||||
) -> Option<Vec<serde_json::Value>> {
|
||||
match self
|
||||
.store
|
||||
.replay_input(
|
||||
response_id,
|
||||
entry,
|
||||
now,
|
||||
self.ttl,
|
||||
MAX_REPLAY_DEPTH,
|
||||
MAX_REPLAY_SESSION_BYTES,
|
||||
)
|
||||
.await
|
||||
{
|
||||
Ok(input) => input,
|
||||
Err(error) => {
|
||||
tracing::warn!(
|
||||
response_id,
|
||||
error = %error,
|
||||
"failed to reconstruct Redis response replay history"
|
||||
);
|
||||
None
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn lookup_latest_account_by_conversation(
|
||||
&self,
|
||||
conversation_id: &str,
|
||||
|
||||
@@ -1,16 +1,277 @@
|
||||
//! Redis 会话亲和存储。
|
||||
|
||||
use std::collections::HashMap;
|
||||
|
||||
use chrono::{DateTime, Duration, Utc};
|
||||
use redis::AsyncCommands;
|
||||
use serde_json::Value;
|
||||
use thiserror::Error;
|
||||
|
||||
use crate::infra::redis::RedisConnection;
|
||||
|
||||
use super::types::SessionAffinityEntry;
|
||||
|
||||
/// 单会话保留的最近响应元数据数量。
|
||||
pub const MAX_CONVERSATION_AFFINITIES: usize = 128;
|
||||
/// 单账号保留的最近响应索引数量。
|
||||
pub const MAX_ACCOUNT_AFFINITIES: usize = 16_384;
|
||||
/// Redis 全局保留的最近响应元数据数量。
|
||||
pub const MAX_GLOBAL_AFFINITIES: usize = 65_536;
|
||||
/// Redis 全局响应元数据 JSON 字节预算。
|
||||
pub const MAX_GLOBAL_AFFINITY_BYTES: u64 = 128 * 1024 * 1024;
|
||||
|
||||
const MAX_AFFINITY_METADATA_BYTES: usize = 64 * 1024;
|
||||
const GLOBAL_STALE_CLEANUP_BATCH: usize = 256;
|
||||
const AFFINITY_VERSION: &str = "affinity:v3";
|
||||
|
||||
const UPSERT_SCRIPT: &str = r#"
|
||||
local response_id = ARGV[1]
|
||||
local metadata_json = ARGV[2]
|
||||
local score = tonumber(ARGV[3])
|
||||
local expires_at = tonumber(ARGV[4])
|
||||
local ttl_seconds = tonumber(ARGV[5])
|
||||
local cutoff_ms = tonumber(ARGV[6])
|
||||
local response_prefix = ARGV[7]
|
||||
local conversation_prefix = ARGV[8]
|
||||
local account_prefix = ARGV[9]
|
||||
local max_conversation = tonumber(ARGV[10])
|
||||
local max_account = tonumber(ARGV[11])
|
||||
local max_global = tonumber(ARGV[12])
|
||||
local max_global_bytes = tonumber(ARGV[13])
|
||||
local stale_cleanup_batch = tonumber(ARGV[14])
|
||||
|
||||
local current = cjson.decode(metadata_json)
|
||||
local conversation_id = current.conversationId
|
||||
local account_id = current.accountId
|
||||
|
||||
local function decrement_bytes(size)
|
||||
if size <= 0 then
|
||||
return
|
||||
end
|
||||
local remaining = redis.call('DECRBY', KEYS[5], size)
|
||||
if remaining <= 0 then
|
||||
redis.call('DEL', KEYS[5])
|
||||
end
|
||||
end
|
||||
|
||||
local function read_info(id)
|
||||
local raw = redis.call('HGET', KEYS[6], id)
|
||||
if not raw then
|
||||
return nil
|
||||
end
|
||||
local ok, info = pcall(cjson.decode, raw)
|
||||
if not ok then
|
||||
redis.call('HDEL', KEYS[6], id)
|
||||
return nil
|
||||
end
|
||||
return info
|
||||
end
|
||||
|
||||
local function remove_response(id)
|
||||
local response_key = response_prefix .. id
|
||||
local info = read_info(id)
|
||||
local metadata = nil
|
||||
local raw = redis.call('GET', response_key)
|
||||
if raw then
|
||||
local ok, decoded = pcall(cjson.decode, raw)
|
||||
if ok then
|
||||
metadata = decoded
|
||||
end
|
||||
end
|
||||
|
||||
local stored_conversation_id = info and info.conversationId or metadata and metadata.conversationId
|
||||
local stored_account_id = info and info.accountId or metadata and metadata.accountId
|
||||
local metadata_bytes = info and tonumber(info.metadataBytes or 0)
|
||||
or redis.call('STRLEN', response_key)
|
||||
|
||||
redis.call('UNLINK', response_key)
|
||||
redis.call('ZREM', KEYS[4], id)
|
||||
redis.call('HDEL', KEYS[6], id)
|
||||
redis.call('ZREM', KEYS[2], id)
|
||||
redis.call('ZREM', KEYS[3], id)
|
||||
decrement_bytes(metadata_bytes)
|
||||
if stored_conversation_id then
|
||||
redis.call('ZREM', conversation_prefix .. stored_conversation_id, id)
|
||||
end
|
||||
if stored_account_id then
|
||||
redis.call('ZREM', account_prefix .. stored_account_id, id)
|
||||
end
|
||||
end
|
||||
|
||||
local function trim_index(key, max_members)
|
||||
local excess = redis.call('ZCARD', key) - max_members
|
||||
if excess <= 0 then
|
||||
return
|
||||
end
|
||||
local ids = redis.call('ZRANGE', key, 0, excess - 1)
|
||||
for _, id in ipairs(ids) do
|
||||
remove_response(id)
|
||||
end
|
||||
end
|
||||
|
||||
local previous_info = read_info(response_id)
|
||||
if previous_info then
|
||||
if previous_info.conversationId and previous_info.conversationId ~= conversation_id then
|
||||
redis.call('ZREM', conversation_prefix .. previous_info.conversationId, response_id)
|
||||
end
|
||||
if previous_info.accountId and previous_info.accountId ~= account_id then
|
||||
redis.call('ZREM', account_prefix .. previous_info.accountId, response_id)
|
||||
end
|
||||
end
|
||||
local previous_bytes = previous_info and tonumber(previous_info.metadataBytes or 0)
|
||||
or redis.call('STRLEN', KEYS[1])
|
||||
local metadata_bytes = string.len(metadata_json)
|
||||
|
||||
redis.call('SET', KEYS[1], metadata_json, 'EXAT', expires_at)
|
||||
redis.call('ZADD', KEYS[2], score, response_id)
|
||||
redis.call('EXPIRE', KEYS[2], ttl_seconds)
|
||||
redis.call('ZADD', KEYS[3], score, response_id)
|
||||
redis.call('EXPIRE', KEYS[3], ttl_seconds)
|
||||
redis.call('ZADD', KEYS[4], score, response_id)
|
||||
redis.call('EXPIRE', KEYS[4], ttl_seconds)
|
||||
redis.call('INCRBY', KEYS[5], metadata_bytes - previous_bytes)
|
||||
redis.call('EXPIRE', KEYS[5], ttl_seconds)
|
||||
redis.call('HSET', KEYS[6], response_id, cjson.encode({
|
||||
conversationId = conversation_id,
|
||||
accountId = account_id,
|
||||
metadataBytes = metadata_bytes,
|
||||
}))
|
||||
redis.call('EXPIRE', KEYS[6], ttl_seconds)
|
||||
|
||||
local stale_conversation_ids = redis.call('ZRANGEBYSCORE', KEYS[2], '-inf', cutoff_ms)
|
||||
for _, id in ipairs(stale_conversation_ids) do
|
||||
remove_response(id)
|
||||
end
|
||||
local stale_account_ids = redis.call(
|
||||
'ZRANGEBYSCORE', KEYS[3], '-inf', cutoff_ms, 'LIMIT', 0, stale_cleanup_batch
|
||||
)
|
||||
for _, id in ipairs(stale_account_ids) do
|
||||
remove_response(id)
|
||||
end
|
||||
local stale_global_ids = redis.call(
|
||||
'ZRANGEBYSCORE', KEYS[4], '-inf', cutoff_ms, 'LIMIT', 0, stale_cleanup_batch
|
||||
)
|
||||
for _, id in ipairs(stale_global_ids) do
|
||||
remove_response(id)
|
||||
end
|
||||
|
||||
trim_index(KEYS[2], max_conversation)
|
||||
trim_index(KEYS[3], max_account)
|
||||
trim_index(KEYS[4], max_global)
|
||||
|
||||
while tonumber(redis.call('GET', KEYS[5]) or '0') > max_global_bytes do
|
||||
local oldest = redis.call('ZRANGE', KEYS[4], 0, 0)
|
||||
if #oldest == 0 then
|
||||
redis.call('DEL', KEYS[5])
|
||||
break
|
||||
end
|
||||
remove_response(oldest[1])
|
||||
end
|
||||
|
||||
return 1
|
||||
"#;
|
||||
|
||||
const FORGET_SCRIPT: &str = r#"
|
||||
local response_id = ARGV[1]
|
||||
local response_prefix = ARGV[2]
|
||||
local conversation_prefix = ARGV[3]
|
||||
local account_prefix = ARGV[4]
|
||||
local response_key = response_prefix .. response_id
|
||||
local info_json = redis.call('HGET', KEYS[3], response_id)
|
||||
local metadata_json = redis.call('GET', response_key)
|
||||
if not info_json and not metadata_json then
|
||||
return 0
|
||||
end
|
||||
|
||||
local info = nil
|
||||
if info_json then
|
||||
local ok, decoded = pcall(cjson.decode, info_json)
|
||||
if ok then
|
||||
info = decoded
|
||||
end
|
||||
end
|
||||
local metadata = nil
|
||||
if metadata_json then
|
||||
local ok, decoded = pcall(cjson.decode, metadata_json)
|
||||
if ok then
|
||||
metadata = decoded
|
||||
end
|
||||
end
|
||||
local conversation_id = info and info.conversationId or metadata and metadata.conversationId
|
||||
local account_id = info and info.accountId or metadata and metadata.accountId
|
||||
local metadata_bytes = info and tonumber(info.metadataBytes or 0)
|
||||
or redis.call('STRLEN', response_key)
|
||||
|
||||
redis.call('UNLINK', response_key)
|
||||
redis.call('ZREM', KEYS[1], response_id)
|
||||
redis.call('HDEL', KEYS[3], response_id)
|
||||
if metadata_bytes > 0 then
|
||||
local remaining = redis.call('DECRBY', KEYS[2], metadata_bytes)
|
||||
if remaining <= 0 then
|
||||
redis.call('DEL', KEYS[2])
|
||||
end
|
||||
end
|
||||
if conversation_id then
|
||||
redis.call('ZREM', conversation_prefix .. conversation_id, response_id)
|
||||
end
|
||||
if account_id then
|
||||
redis.call('ZREM', account_prefix .. account_id, response_id)
|
||||
end
|
||||
return 1
|
||||
"#;
|
||||
|
||||
const FORGET_ACCOUNT_SCRIPT: &str = r#"
|
||||
local account_id = ARGV[1]
|
||||
local response_prefix = ARGV[2]
|
||||
local conversation_prefix = ARGV[3]
|
||||
local account_prefix = ARGV[4]
|
||||
local account_key = account_prefix .. account_id
|
||||
local response_ids = redis.call('ZRANGE', account_key, 0, -1)
|
||||
|
||||
local function decrement_bytes(size)
|
||||
if size <= 0 then
|
||||
return
|
||||
end
|
||||
local remaining = redis.call('DECRBY', KEYS[3], size)
|
||||
if remaining <= 0 then
|
||||
redis.call('DEL', KEYS[3])
|
||||
end
|
||||
end
|
||||
|
||||
for _, response_id in ipairs(response_ids) do
|
||||
local response_key = response_prefix .. response_id
|
||||
local info_json = redis.call('HGET', KEYS[4], response_id)
|
||||
local metadata_json = redis.call('GET', response_key)
|
||||
local info = nil
|
||||
if info_json then
|
||||
local ok, decoded = pcall(cjson.decode, info_json)
|
||||
if ok then
|
||||
info = decoded
|
||||
end
|
||||
end
|
||||
local metadata = nil
|
||||
if metadata_json then
|
||||
local ok, decoded = pcall(cjson.decode, metadata_json)
|
||||
if ok then
|
||||
metadata = decoded
|
||||
end
|
||||
end
|
||||
local stored_account_id = info and info.accountId or metadata and metadata.accountId
|
||||
if stored_account_id == account_id then
|
||||
local conversation_id = info and info.conversationId or metadata and metadata.conversationId
|
||||
local metadata_bytes = info and tonumber(info.metadataBytes or 0)
|
||||
or redis.call('STRLEN', response_key)
|
||||
redis.call('UNLINK', response_key)
|
||||
redis.call('ZREM', KEYS[2], response_id)
|
||||
redis.call('HDEL', KEYS[4], response_id)
|
||||
decrement_bytes(metadata_bytes)
|
||||
if conversation_id then
|
||||
redis.call('ZREM', conversation_prefix .. conversation_id, response_id)
|
||||
end
|
||||
end
|
||||
end
|
||||
|
||||
redis.call('DEL', account_key)
|
||||
return #response_ids
|
||||
"#;
|
||||
|
||||
/// Redis 会话亲和性存储。
|
||||
#[derive(Clone)]
|
||||
pub struct RedisSessionAffinityStore {
|
||||
@@ -24,6 +285,8 @@ pub enum RedisSessionAffinityStoreError {
|
||||
Redis(#[from] redis::RedisError),
|
||||
#[error("invalid Redis session affinity value: {0}")]
|
||||
Json(#[from] serde_json::Error),
|
||||
#[error("session affinity metadata is too large: {bytes} bytes exceeds {max} bytes")]
|
||||
MetadataTooLarge { bytes: usize, max: usize },
|
||||
}
|
||||
|
||||
pub type RedisSessionAffinityStoreResult<T> = Result<T, RedisSessionAffinityStoreError>;
|
||||
@@ -39,7 +302,6 @@ impl RedisSessionAffinityStore {
|
||||
entry: &SessionAffinityEntry,
|
||||
ttl: Duration,
|
||||
) -> RedisSessionAffinityStoreResult<()> {
|
||||
let value = serde_json::to_string(entry)?;
|
||||
let expires_at = entry
|
||||
.created_at
|
||||
.checked_add_signed(ttl)
|
||||
@@ -48,50 +310,40 @@ impl RedisSessionAffinityStore {
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
let response_key = self.response_key(response_id);
|
||||
let conversation_key = self.conversation_key(&entry.conversation_id);
|
||||
let account_key = self.account_key(&entry.account_id);
|
||||
let metadata_json = serde_json::to_string(entry)?;
|
||||
if metadata_json.len() > MAX_AFFINITY_METADATA_BYTES {
|
||||
return Err(RedisSessionAffinityStoreError::MetadataTooLarge {
|
||||
bytes: metadata_json.len(),
|
||||
max: MAX_AFFINITY_METADATA_BYTES,
|
||||
});
|
||||
}
|
||||
|
||||
let ttl_seconds = ttl.num_seconds().max(1);
|
||||
let cutoff_ms = Utc::now().timestamp_millis() - ttl.num_milliseconds();
|
||||
let mut connection = self.redis.manager();
|
||||
let mut transaction = redis::pipe();
|
||||
transaction
|
||||
.atomic()
|
||||
.cmd("SET")
|
||||
.arg(response_key)
|
||||
.arg(value)
|
||||
.arg("EXAT")
|
||||
let _: i64 = redis::Script::new(UPSERT_SCRIPT)
|
||||
.key(self.response_key(response_id))
|
||||
.key(self.conversation_key(&entry.conversation_id))
|
||||
.key(self.account_key(&entry.account_id))
|
||||
.key(self.global_index_key())
|
||||
.key(self.global_bytes_key())
|
||||
.key(self.global_info_key())
|
||||
.arg(response_id)
|
||||
.arg(metadata_json)
|
||||
.arg(entry.created_at.timestamp_millis())
|
||||
.arg(expires_at.timestamp())
|
||||
.ignore()
|
||||
.cmd("ZADD")
|
||||
.arg(&conversation_key)
|
||||
.arg(entry.created_at.timestamp_millis())
|
||||
.arg(response_id)
|
||||
.ignore()
|
||||
.cmd("EXPIRE")
|
||||
.arg(&conversation_key)
|
||||
.arg(ttl_seconds)
|
||||
.ignore()
|
||||
.cmd("ZADD")
|
||||
.arg(&account_key)
|
||||
.arg(entry.created_at.timestamp_millis())
|
||||
.arg(response_id)
|
||||
.ignore()
|
||||
.cmd("EXPIRE")
|
||||
.arg(&account_key)
|
||||
.arg(ttl_seconds)
|
||||
.ignore()
|
||||
.cmd("ZREMRANGEBYSCORE")
|
||||
.arg(conversation_key)
|
||||
.arg("-inf")
|
||||
.arg(cutoff_ms)
|
||||
.ignore()
|
||||
.cmd("ZREMRANGEBYSCORE")
|
||||
.arg(&account_key)
|
||||
.arg("-inf")
|
||||
.arg(cutoff_ms)
|
||||
.ignore();
|
||||
let _: () = transaction.query_async(&mut connection).await?;
|
||||
.arg(self.response_key_prefix())
|
||||
.arg(self.conversation_key_prefix())
|
||||
.arg(self.account_key_prefix())
|
||||
.arg(MAX_CONVERSATION_AFFINITIES)
|
||||
.arg(MAX_ACCOUNT_AFFINITIES)
|
||||
.arg(MAX_GLOBAL_AFFINITIES)
|
||||
.arg(MAX_GLOBAL_AFFINITY_BYTES)
|
||||
.arg(GLOBAL_STALE_CLEANUP_BATCH)
|
||||
.invoke_async(&mut connection)
|
||||
.await?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
@@ -128,163 +380,14 @@ impl RedisSessionAffinityStore {
|
||||
) -> RedisSessionAffinityStoreResult<Option<(String, SessionAffinityEntry)>> {
|
||||
let conversation_key = self.conversation_key(conversation_id);
|
||||
let mut connection = self.redis.manager();
|
||||
let response_ids: Vec<String> = connection.zrevrange(&conversation_key, 0, -1).await?;
|
||||
for response_id in response_ids {
|
||||
let Some(entry) = self.get(&response_id, now, ttl).await? else {
|
||||
let _: usize = connection.zrem(&conversation_key, &response_id).await?;
|
||||
continue;
|
||||
};
|
||||
if entry.conversation_id != conversation_id {
|
||||
let _: usize = connection.zrem(&conversation_key, &response_id).await?;
|
||||
continue;
|
||||
}
|
||||
if variant_hash.is_some_and(|expected| entry.variant_hash.as_deref() != Some(expected))
|
||||
{
|
||||
continue;
|
||||
}
|
||||
if max_age.is_some_and(|max_age| now.signed_duration_since(entry.created_at) > max_age)
|
||||
{
|
||||
continue;
|
||||
}
|
||||
return Ok(Some((response_id, entry)));
|
||||
}
|
||||
Ok(None)
|
||||
}
|
||||
|
||||
pub async fn replay_input(
|
||||
&self,
|
||||
response_id: &str,
|
||||
head: &SessionAffinityEntry,
|
||||
now: DateTime<Utc>,
|
||||
ttl: Duration,
|
||||
max_depth: u16,
|
||||
max_total_bytes: u64,
|
||||
) -> RedisSessionAffinityStoreResult<Option<Vec<Value>>> {
|
||||
let Some(head_snapshot) = head.replay.as_ref() else {
|
||||
return Ok(None);
|
||||
};
|
||||
if head_snapshot.depth == 0
|
||||
|| head_snapshot.depth > max_depth
|
||||
|| head_snapshot.total_bytes > max_total_bytes
|
||||
{
|
||||
return Ok(None);
|
||||
}
|
||||
|
||||
let conversation_key = self.conversation_key(&head.conversation_id);
|
||||
let mut connection = self.redis.manager();
|
||||
let response_ids: Vec<String> = redis::cmd("ZREVRANGEBYSCORE")
|
||||
let response_ids: Vec<String> = redis::cmd("ZREVRANGE")
|
||||
.arg(&conversation_key)
|
||||
.arg(head.created_at.timestamp_millis())
|
||||
.arg("-inf")
|
||||
.arg("LIMIT")
|
||||
.arg(0)
|
||||
.arg(max_depth)
|
||||
.arg(MAX_CONVERSATION_AFFINITIES.saturating_sub(1))
|
||||
.query_async(&mut connection)
|
||||
.await?;
|
||||
if !response_ids
|
||||
.iter()
|
||||
.any(|candidate| candidate == response_id)
|
||||
{
|
||||
return Ok(None);
|
||||
}
|
||||
|
||||
let response_keys = response_ids
|
||||
.iter()
|
||||
.map(|candidate| self.response_key(candidate))
|
||||
.collect::<Vec<_>>();
|
||||
let values: Vec<Option<String>> = redis::cmd("MGET")
|
||||
.arg(&response_keys)
|
||||
.query_async(&mut connection)
|
||||
.await?;
|
||||
let mut entries = HashMap::with_capacity(values.len());
|
||||
for (candidate, value) in response_ids.into_iter().zip(values) {
|
||||
let Some(value) = value else {
|
||||
continue;
|
||||
};
|
||||
let entry: SessionAffinityEntry = serde_json::from_str(&value)?;
|
||||
if entry.conversation_id == head.conversation_id
|
||||
&& entry
|
||||
.created_at
|
||||
.checked_add_signed(ttl)
|
||||
.is_some_and(|expires_at| expires_at > now)
|
||||
{
|
||||
entries.insert(candidate, entry);
|
||||
}
|
||||
}
|
||||
|
||||
let mut current_id = response_id.to_string();
|
||||
let mut expected_depth = head_snapshot.depth;
|
||||
let mut expected_total_bytes = head_snapshot.total_bytes;
|
||||
let mut nodes = Vec::with_capacity(usize::from(expected_depth));
|
||||
loop {
|
||||
let Some(entry) = entries.get(¤t_id) else {
|
||||
return Ok(None);
|
||||
};
|
||||
let Some(snapshot) = entry.replay.as_ref() else {
|
||||
return Ok(None);
|
||||
};
|
||||
if snapshot.depth != expected_depth || snapshot.total_bytes != expected_total_bytes {
|
||||
return Ok(None);
|
||||
}
|
||||
let node_bytes = replay_node_bytes(snapshot)?;
|
||||
if node_bytes > expected_total_bytes {
|
||||
return Ok(None);
|
||||
}
|
||||
nodes.push((snapshot.turn_input.clone(), snapshot.turn_output.clone()));
|
||||
|
||||
match snapshot.parent_response_id.as_deref() {
|
||||
None if expected_depth == 1 && expected_total_bytes == node_bytes => break,
|
||||
Some(parent_response_id) if expected_depth > 1 => {
|
||||
expected_depth -= 1;
|
||||
expected_total_bytes -= node_bytes;
|
||||
current_id = parent_response_id.to_string();
|
||||
}
|
||||
None | Some(_) => return Ok(None),
|
||||
}
|
||||
}
|
||||
|
||||
nodes.reverse();
|
||||
let mut full_input = Vec::new();
|
||||
for (mut input, mut output) in nodes {
|
||||
full_input.append(&mut input);
|
||||
full_input.append(&mut output);
|
||||
}
|
||||
Ok(Some(full_input))
|
||||
}
|
||||
|
||||
pub async fn forget(&self, response_id: &str) -> RedisSessionAffinityStoreResult<bool> {
|
||||
let response_key = self.response_key(response_id);
|
||||
let mut connection = self.redis.manager();
|
||||
let value: Option<String> = connection.get(&response_key).await?;
|
||||
let Some(value) = value else {
|
||||
return Ok(false);
|
||||
};
|
||||
let entry: SessionAffinityEntry = serde_json::from_str(&value)?;
|
||||
let mut transaction = redis::pipe();
|
||||
transaction
|
||||
.atomic()
|
||||
.cmd("DEL")
|
||||
.arg(response_key)
|
||||
.ignore()
|
||||
.cmd("ZREM")
|
||||
.arg(self.conversation_key(&entry.conversation_id))
|
||||
.arg(response_id)
|
||||
.ignore()
|
||||
.cmd("ZREM")
|
||||
.arg(self.account_key(&entry.account_id))
|
||||
.arg(response_id)
|
||||
.ignore();
|
||||
let _: () = transaction.query_async(&mut connection).await?;
|
||||
Ok(true)
|
||||
}
|
||||
|
||||
pub async fn forget_account(&self, account_id: &str) -> RedisSessionAffinityStoreResult<u64> {
|
||||
let account_key = self.account_key(account_id);
|
||||
let mut connection = self.redis.manager();
|
||||
let response_ids: Vec<String> = connection.zrange(&account_key, 0, -1).await?;
|
||||
if response_ids.is_empty() {
|
||||
let _: usize = connection.del(account_key).await?;
|
||||
return Ok(0);
|
||||
return Ok(None);
|
||||
}
|
||||
|
||||
let response_keys = response_ids
|
||||
@@ -292,67 +395,105 @@ impl RedisSessionAffinityStore {
|
||||
.map(|response_id| self.response_key(response_id))
|
||||
.collect::<Vec<_>>();
|
||||
let values: Vec<Option<String>> = redis::cmd("MGET")
|
||||
.arg(&response_keys)
|
||||
.arg(response_keys)
|
||||
.query_async(&mut connection)
|
||||
.await?;
|
||||
let mut conversation_members = Vec::with_capacity(response_ids.len());
|
||||
for (response_id, value) in response_ids.iter().zip(values) {
|
||||
let mut stale_ids = Vec::new();
|
||||
let mut result = None;
|
||||
for (response_id, value) in response_ids.into_iter().zip(values) {
|
||||
let Some(value) = value else {
|
||||
stale_ids.push(response_id);
|
||||
continue;
|
||||
};
|
||||
match serde_json::from_str::<SessionAffinityEntry>(&value) {
|
||||
Ok(entry) if entry.account_id == account_id => {
|
||||
conversation_members.push((response_id, entry.conversation_id));
|
||||
}
|
||||
Ok(_) => {
|
||||
tracing::warn!(
|
||||
response_id,
|
||||
account_id,
|
||||
"account affinity index points to another account"
|
||||
);
|
||||
}
|
||||
Err(error) => {
|
||||
tracing::warn!(
|
||||
response_id,
|
||||
account_id,
|
||||
error = %error,
|
||||
"invalid account affinity entry removed during account cleanup"
|
||||
);
|
||||
}
|
||||
let entry: SessionAffinityEntry = serde_json::from_str(&value)?;
|
||||
if entry.conversation_id != conversation_id
|
||||
|| entry
|
||||
.created_at
|
||||
.checked_add_signed(ttl)
|
||||
.is_none_or(|expires_at| expires_at <= now)
|
||||
{
|
||||
stale_ids.push(response_id);
|
||||
continue;
|
||||
}
|
||||
if variant_hash.is_some_and(|expected| entry.variant_hash.as_deref() != Some(expected))
|
||||
|| max_age
|
||||
.is_some_and(|max_age| now.signed_duration_since(entry.created_at) > max_age)
|
||||
{
|
||||
continue;
|
||||
}
|
||||
result = Some((response_id, entry));
|
||||
break;
|
||||
}
|
||||
if !stale_ids.is_empty() {
|
||||
let _: usize = connection.zrem(conversation_key, stale_ids).await?;
|
||||
}
|
||||
Ok(result)
|
||||
}
|
||||
|
||||
let mut transaction = redis::pipe();
|
||||
transaction.atomic().cmd("DEL").arg(&response_keys).ignore();
|
||||
for (response_id, conversation_id) in &conversation_members {
|
||||
transaction
|
||||
.cmd("ZREM")
|
||||
.arg(self.conversation_key(conversation_id))
|
||||
.arg(response_id)
|
||||
.ignore();
|
||||
}
|
||||
transaction.cmd("DEL").arg(account_key).ignore();
|
||||
let _: () = transaction.query_async(&mut connection).await?;
|
||||
Ok(response_ids.len() as u64)
|
||||
pub async fn forget(&self, response_id: &str) -> RedisSessionAffinityStoreResult<bool> {
|
||||
let mut connection = self.redis.manager();
|
||||
let forgotten: i64 = redis::Script::new(FORGET_SCRIPT)
|
||||
.key(self.global_index_key())
|
||||
.key(self.global_bytes_key())
|
||||
.key(self.global_info_key())
|
||||
.arg(response_id)
|
||||
.arg(self.response_key_prefix())
|
||||
.arg(self.conversation_key_prefix())
|
||||
.arg(self.account_key_prefix())
|
||||
.invoke_async(&mut connection)
|
||||
.await?;
|
||||
Ok(forgotten > 0)
|
||||
}
|
||||
|
||||
pub async fn forget_account(&self, account_id: &str) -> RedisSessionAffinityStoreResult<u64> {
|
||||
let mut connection = self.redis.manager();
|
||||
let forgotten: u64 = redis::Script::new(FORGET_ACCOUNT_SCRIPT)
|
||||
.key(self.account_key(account_id))
|
||||
.key(self.global_index_key())
|
||||
.key(self.global_bytes_key())
|
||||
.key(self.global_info_key())
|
||||
.arg(account_id)
|
||||
.arg(self.response_key_prefix())
|
||||
.arg(self.conversation_key_prefix())
|
||||
.arg(self.account_key_prefix())
|
||||
.invoke_async(&mut connection)
|
||||
.await?;
|
||||
Ok(forgotten)
|
||||
}
|
||||
|
||||
fn response_key(&self, response_id: &str) -> String {
|
||||
self.redis.key(&format!("affinity:v2:resp:{response_id}"))
|
||||
format!("{}{response_id}", self.response_key_prefix())
|
||||
}
|
||||
|
||||
fn conversation_key(&self, conversation_id: &str) -> String {
|
||||
self.redis
|
||||
.key(&format!("affinity:v2:conv:{conversation_id}"))
|
||||
format!("{}{conversation_id}", self.conversation_key_prefix())
|
||||
}
|
||||
|
||||
fn account_key(&self, account_id: &str) -> String {
|
||||
self.redis.key(&format!("affinity:v2:account:{account_id}"))
|
||||
format!("{}{account_id}", self.account_key_prefix())
|
||||
}
|
||||
|
||||
fn global_index_key(&self) -> String {
|
||||
self.redis.key(&format!("{AFFINITY_VERSION}:global:index"))
|
||||
}
|
||||
|
||||
fn global_bytes_key(&self) -> String {
|
||||
self.redis.key(&format!("{AFFINITY_VERSION}:global:bytes"))
|
||||
}
|
||||
|
||||
fn global_info_key(&self) -> String {
|
||||
self.redis.key(&format!("{AFFINITY_VERSION}:global:info"))
|
||||
}
|
||||
|
||||
fn response_key_prefix(&self) -> String {
|
||||
self.redis.key(&format!("{AFFINITY_VERSION}:resp:"))
|
||||
}
|
||||
|
||||
fn conversation_key_prefix(&self) -> String {
|
||||
self.redis.key(&format!("{AFFINITY_VERSION}:conv:"))
|
||||
}
|
||||
|
||||
fn account_key_prefix(&self) -> String {
|
||||
self.redis.key(&format!("{AFFINITY_VERSION}:account:"))
|
||||
}
|
||||
}
|
||||
|
||||
fn replay_node_bytes(
|
||||
snapshot: &super::types::ResponseReplaySnapshot,
|
||||
) -> Result<u64, serde_json::Error> {
|
||||
serde_json::to_vec(&(&snapshot.turn_input, &snapshot.turn_output))
|
||||
.map(|value| value.len() as u64)
|
||||
}
|
||||
|
||||
@@ -2,7 +2,6 @@
|
||||
|
||||
use chrono::{DateTime, Utc};
|
||||
use serde::{Deserialize, Serialize};
|
||||
use serde_json::Value;
|
||||
|
||||
use crate::upstream::openai::protocol::responses::PreviousResponseScope;
|
||||
|
||||
@@ -18,22 +17,5 @@ pub struct SessionAffinityEntry {
|
||||
pub function_call_ids: Vec<String>,
|
||||
pub variant_hash: Option<String>,
|
||||
pub continuation_scope: PreviousResponseScope,
|
||||
pub replay: Option<ResponseReplaySnapshot>,
|
||||
pub created_at: DateTime<Utc>,
|
||||
}
|
||||
|
||||
/// 单轮 completed response 的有界、无凭据重放增量。
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub struct ResponseReplaySnapshot {
|
||||
/// 上一轮由本代理记录的 completed response。
|
||||
pub parent_response_id: Option<String>,
|
||||
/// 本轮客户端新增的输入,已剥离账号绑定字段。
|
||||
pub turn_input: Vec<Value>,
|
||||
/// 本轮上游输出,已剥离账号绑定字段和上游对象 ID。
|
||||
pub turn_output: Vec<Value>,
|
||||
/// 从会话起点到本节点的增量节点数。
|
||||
pub depth: u16,
|
||||
/// 从会话起点到本节点的 JSON 编码累计字节数。
|
||||
pub total_bytes: u64,
|
||||
}
|
||||
|
||||
@@ -3,11 +3,14 @@
|
||||
use std::collections::{BTreeSet, VecDeque};
|
||||
|
||||
use chrono::Utc;
|
||||
use tokio::time::{Duration, Instant, timeout_at};
|
||||
|
||||
use crate::fleet::pool::{
|
||||
AccountAcquireRequest, AccountCandidateLease, AccountLease, AccountPoolService,
|
||||
};
|
||||
|
||||
const ACCOUNT_CANDIDATE_WAIT_TIMEOUT: Duration = Duration::from_secs(10);
|
||||
|
||||
/// 请求开始时冻结的完整候选集,以及候选的最终处理状态。
|
||||
pub(in crate::dispatch) struct AccountAttemptLedger {
|
||||
model: String,
|
||||
@@ -16,6 +19,7 @@ pub(in crate::dispatch) struct AccountAttemptLedger {
|
||||
busy: VecDeque<String>,
|
||||
attempted: BTreeSet<String>,
|
||||
state_excluded: BTreeSet<String>,
|
||||
busy_wait_deadline: Instant,
|
||||
}
|
||||
|
||||
impl AccountAttemptLedger {
|
||||
@@ -31,6 +35,7 @@ impl AccountAttemptLedger {
|
||||
busy: VecDeque::new(),
|
||||
attempted: BTreeSet::new(),
|
||||
state_excluded: BTreeSet::new(),
|
||||
busy_wait_deadline: Instant::now() + ACCOUNT_CANDIDATE_WAIT_TIMEOUT,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -63,7 +68,22 @@ impl AccountAttemptLedger {
|
||||
return None;
|
||||
}
|
||||
|
||||
account_pool.wait_for_candidate_change().await;
|
||||
if timeout_at(
|
||||
self.busy_wait_deadline,
|
||||
account_pool.wait_for_candidate_change(),
|
||||
)
|
||||
.await
|
||||
.is_err()
|
||||
{
|
||||
tracing::warn!(
|
||||
candidate_count = self.candidate_ids.len(),
|
||||
attempted_count = self.attempted.len(),
|
||||
busy_count = self.busy.len(),
|
||||
timeout_ms = ACCOUNT_CANDIDATE_WAIT_TIMEOUT.as_millis(),
|
||||
"timed out waiting for a busy account candidate"
|
||||
);
|
||||
return None;
|
||||
}
|
||||
std::mem::swap(&mut self.pending, &mut self.busy);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -206,8 +206,11 @@ pub(crate) fn is_retryable_upstream_5xx_error(error: &CodexClientError) -> bool
|
||||
pub(crate) fn is_retryable_account_transport_error(error: &CodexClientError) -> bool {
|
||||
match error {
|
||||
CodexClientError::Http(error) => error.is_connect() || error.is_timeout(),
|
||||
CodexClientError::StreamIdleTimeout { .. } => true,
|
||||
CodexClientError::WebSocket(
|
||||
crate::upstream::openai::transport::websocket::CodexWebSocketExchangeError::Transport(_)
|
||||
| crate::upstream::openai::transport::websocket::CodexWebSocketExchangeError::ConnectTimeout { .. }
|
||||
| crate::upstream::openai::transport::websocket::CodexWebSocketExchangeError::SendTimeout { .. }
|
||||
| crate::upstream::openai::transport::websocket::CodexWebSocketExchangeError::ClosedBeforeTerminal
|
||||
| crate::upstream::openai::transport::websocket::CodexWebSocketExchangeError::ReceiveIdleTimeout { .. }
|
||||
| crate::upstream::openai::transport::websocket::CodexWebSocketExchangeError::ReusedConnectionDiedBeforeFirstOutput { .. }
|
||||
|
||||
@@ -204,11 +204,13 @@ pub(super) async fn record_prefetched_response_stream_failure_event(
|
||||
diagnostics_status_code(record.diagnostics),
|
||||
);
|
||||
insert_response_request_summary_object(object, record.request, record.transport);
|
||||
object.insert("requestBody".to_string(), json!(record.request));
|
||||
object.insert(
|
||||
"responseBody".to_string(),
|
||||
Value::String(String::from_utf8_lossy(record.prefetched).to_string()),
|
||||
);
|
||||
if record.recorder.captures_body() {
|
||||
object.insert("requestBody".to_string(), json!(record.request));
|
||||
object.insert(
|
||||
"responseBody".to_string(),
|
||||
Value::String(String::from_utf8_lossy(record.prefetched).to_string()),
|
||||
);
|
||||
}
|
||||
insert_upstream_diagnostics_metadata(object, record.diagnostics);
|
||||
insert_response_trace_metadata_object(
|
||||
object,
|
||||
@@ -295,12 +297,14 @@ fn enrich_live_response_stream_metadata(
|
||||
insert_response_request_summary_object(object, &context.request, context.transport);
|
||||
insert_upstream_diagnostics_metadata(object, &context.diagnostics);
|
||||
insert_response_trace_metadata_object(object, &context.attempts, Some(&context.attempt));
|
||||
object
|
||||
.entry("requestBody".to_string())
|
||||
.or_insert_with(|| serde_json::json!(context.request));
|
||||
object
|
||||
.entry("responseBody".to_string())
|
||||
.or_insert_with(|| Value::String(body.to_string()));
|
||||
if context.recorder.captures_body() {
|
||||
object
|
||||
.entry("requestBody".to_string())
|
||||
.or_insert_with(|| serde_json::json!(context.request));
|
||||
object
|
||||
.entry("responseBody".to_string())
|
||||
.or_insert_with(|| Value::String(body.to_string()));
|
||||
}
|
||||
}
|
||||
|
||||
pub(super) fn insert_response_upstream_diagnostics(
|
||||
|
||||
@@ -1,14 +1,12 @@
|
||||
//! previous response 所有权与完整上下文重放。
|
||||
//! previous response 所有权与连接内上下文重放。
|
||||
|
||||
use std::sync::Arc;
|
||||
|
||||
use chrono::Utc;
|
||||
use serde_json::Value;
|
||||
|
||||
use crate::{
|
||||
dispatch::affinity::{
|
||||
SessionAffinityEntry, SessionAffinityService,
|
||||
service::{MAX_REPLAY_DEPTH, MAX_REPLAY_SESSION_BYTES, MAX_REPLAY_SNAPSHOT_BYTES},
|
||||
types::ResponseReplaySnapshot,
|
||||
},
|
||||
dispatch::affinity::{SessionAffinityEntry, SessionAffinityService},
|
||||
upstream::openai::protocol::responses::{CodexResponsesRequest, StreamCommitPolicy},
|
||||
};
|
||||
|
||||
@@ -17,6 +15,7 @@ use crate::{
|
||||
pub(in crate::dispatch) enum HistorySource {
|
||||
None,
|
||||
Managed,
|
||||
ConnectionLocal,
|
||||
ExternalUnknown,
|
||||
}
|
||||
|
||||
@@ -24,7 +23,7 @@ pub(in crate::dispatch) enum HistorySource {
|
||||
pub(in crate::dispatch) struct HistoryRecoveryPlan {
|
||||
source: HistorySource,
|
||||
managed_entry: Option<SessionAffinityEntry>,
|
||||
replay_input: Option<Vec<Value>>,
|
||||
local_replay_input: Option<Arc<Vec<Value>>>,
|
||||
external_attempted_account_id: Option<String>,
|
||||
replay_owner_account: bool,
|
||||
}
|
||||
@@ -38,7 +37,7 @@ impl HistoryRecoveryPlan {
|
||||
return Self {
|
||||
source: HistorySource::None,
|
||||
managed_entry: None,
|
||||
replay_input: None,
|
||||
local_replay_input: None,
|
||||
external_attempted_account_id: None,
|
||||
replay_owner_account: false,
|
||||
};
|
||||
@@ -46,22 +45,16 @@ impl HistoryRecoveryPlan {
|
||||
let managed_entry = session_affinity
|
||||
.lookup(previous_response_id, Utc::now())
|
||||
.await;
|
||||
let replay_input = match managed_entry.as_ref() {
|
||||
Some(entry) => {
|
||||
session_affinity
|
||||
.replay_input(previous_response_id, entry, Utc::now())
|
||||
.await
|
||||
}
|
||||
None => None,
|
||||
let local_replay_input = request.local_replay_input.clone();
|
||||
let source = match (&managed_entry, &local_replay_input) {
|
||||
(Some(_), _) => HistorySource::Managed,
|
||||
(None, Some(_)) => HistorySource::ConnectionLocal,
|
||||
(None, None) => HistorySource::ExternalUnknown,
|
||||
};
|
||||
Self {
|
||||
source: if managed_entry.is_some() {
|
||||
HistorySource::Managed
|
||||
} else {
|
||||
HistorySource::ExternalUnknown
|
||||
},
|
||||
source,
|
||||
managed_entry,
|
||||
replay_input,
|
||||
local_replay_input,
|
||||
external_attempted_account_id: None,
|
||||
replay_owner_account: false,
|
||||
}
|
||||
@@ -83,23 +76,17 @@ impl HistoryRecoveryPlan {
|
||||
HistorySource::ExternalUnknown => match self.external_attempted_account_id.as_deref() {
|
||||
None => {
|
||||
self.external_attempted_account_id = Some(account_id.to_string());
|
||||
let mut request = original.clone();
|
||||
request.stream_commit_policy = StreamCommitPolicy::UntilOutputOrTerminal;
|
||||
request.previous_response_scope = Some(
|
||||
crate::upstream::openai::protocol::responses::PreviousResponseScope::ExternalUnknown,
|
||||
);
|
||||
Some(request)
|
||||
Some(external_history_request(original))
|
||||
}
|
||||
Some(first_account_id) if first_account_id == account_id => {
|
||||
let mut request = original.clone();
|
||||
request.stream_commit_policy = StreamCommitPolicy::UntilOutputOrTerminal;
|
||||
request.previous_response_scope = Some(
|
||||
crate::upstream::openai::protocol::responses::PreviousResponseScope::ExternalUnknown,
|
||||
);
|
||||
Some(request)
|
||||
Some(external_history_request(original))
|
||||
}
|
||||
Some(_) => None,
|
||||
},
|
||||
HistorySource::ConnectionLocal => self
|
||||
.local_replay_input
|
||||
.as_deref()
|
||||
.map(|replay_input| full_replay_request(original, replay_input)),
|
||||
HistorySource::Managed => {
|
||||
let entry = self.managed_entry.as_ref()?;
|
||||
if entry.account_id == account_id && !self.replay_owner_account {
|
||||
@@ -112,7 +99,7 @@ impl HistoryRecoveryPlan {
|
||||
}
|
||||
Some(request)
|
||||
} else {
|
||||
self.replay_input.as_ref().map(|replay_input| {
|
||||
self.local_replay_input.as_deref().map(|replay_input| {
|
||||
let mut request = full_replay_request(original, replay_input);
|
||||
request.local_conversation_id = Some(entry.conversation_id.clone());
|
||||
request
|
||||
@@ -128,7 +115,7 @@ impl HistoryRecoveryPlan {
|
||||
return false;
|
||||
};
|
||||
if entry.account_id != account_id
|
||||
|| self.replay_input.is_none()
|
||||
|| self.local_replay_input.is_none()
|
||||
|| self.replay_owner_account
|
||||
{
|
||||
return false;
|
||||
@@ -138,67 +125,17 @@ impl HistoryRecoveryPlan {
|
||||
}
|
||||
|
||||
pub(in crate::dispatch) fn can_failover(&self) -> bool {
|
||||
!matches!(self.source, HistorySource::ExternalUnknown)
|
||||
&& self
|
||||
.managed_entry
|
||||
.as_ref()
|
||||
.is_none_or(|_| self.replay_input.is_some())
|
||||
match self.source {
|
||||
HistorySource::None | HistorySource::ConnectionLocal => true,
|
||||
HistorySource::Managed => self.local_replay_input.is_some(),
|
||||
HistorySource::ExternalUnknown => false,
|
||||
}
|
||||
}
|
||||
|
||||
pub(in crate::dispatch) fn is_external_unknown(&self) -> bool {
|
||||
self.source == HistorySource::ExternalUnknown
|
||||
}
|
||||
|
||||
pub(in crate::dispatch) fn completed_replay(
|
||||
&self,
|
||||
original: &CodexResponsesRequest,
|
||||
output: &[Value],
|
||||
) -> Option<ResponseReplaySnapshot> {
|
||||
let (parent_response_id, parent_depth, parent_bytes) = match self.source {
|
||||
HistorySource::None => (None, 0, 0),
|
||||
HistorySource::Managed => {
|
||||
self.replay_input.as_ref()?;
|
||||
let replay = self.managed_entry.as_ref()?.replay.as_ref()?;
|
||||
(
|
||||
original.previous_response_id().map(ToString::to_string),
|
||||
replay.depth,
|
||||
replay.total_bytes,
|
||||
)
|
||||
}
|
||||
HistorySource::ExternalUnknown => return None,
|
||||
};
|
||||
let depth = parent_depth.checked_add(1)?;
|
||||
if depth > MAX_REPLAY_DEPTH {
|
||||
tracing::warn!(depth, "response replay depth limit reached");
|
||||
return None;
|
||||
}
|
||||
|
||||
let mut turn_input = original.input().to_vec();
|
||||
let mut turn_output = output.to_vec();
|
||||
sanitize_replay_items(&mut turn_input);
|
||||
sanitize_replay_items(&mut turn_output);
|
||||
let Some(node_bytes) = limited_replay_json_bytes(&turn_input, &turn_output) else {
|
||||
tracing::warn!("response replay snapshot byte limit reached");
|
||||
return None;
|
||||
};
|
||||
let total_bytes = parent_bytes.checked_add(node_bytes)?;
|
||||
if total_bytes > MAX_REPLAY_SESSION_BYTES {
|
||||
tracing::warn!(
|
||||
node_bytes,
|
||||
total_bytes,
|
||||
"response replay session byte limit reached"
|
||||
);
|
||||
return None;
|
||||
}
|
||||
Some(ResponseReplaySnapshot {
|
||||
parent_response_id,
|
||||
turn_input,
|
||||
turn_output,
|
||||
depth,
|
||||
total_bytes,
|
||||
})
|
||||
}
|
||||
|
||||
pub(in crate::dispatch) fn conversation_id<'a>(
|
||||
&'a self,
|
||||
original: &'a CodexResponsesRequest,
|
||||
@@ -210,12 +147,21 @@ impl HistoryRecoveryPlan {
|
||||
}
|
||||
}
|
||||
|
||||
fn external_history_request(original: &CodexResponsesRequest) -> CodexResponsesRequest {
|
||||
let mut request = original.clone();
|
||||
request.stream_commit_policy = StreamCommitPolicy::UntilOutputOrTerminal;
|
||||
request.previous_response_scope =
|
||||
Some(crate::upstream::openai::protocol::responses::PreviousResponseScope::ExternalUnknown);
|
||||
request
|
||||
}
|
||||
|
||||
fn full_replay_request(
|
||||
original: &CodexResponsesRequest,
|
||||
replay_input: &[Value],
|
||||
) -> CodexResponsesRequest {
|
||||
let mut request = original.clone();
|
||||
let mut input = replay_input.to_vec();
|
||||
let mut input = Vec::with_capacity(replay_input.len() + original.input().len());
|
||||
input.extend_from_slice(replay_input);
|
||||
input.extend_from_slice(original.input());
|
||||
sanitize_replay_items(&mut input);
|
||||
request.set_input(input);
|
||||
@@ -225,7 +171,7 @@ fn full_replay_request(
|
||||
request
|
||||
}
|
||||
|
||||
fn sanitize_replay_items(values: &mut [Value]) {
|
||||
pub(crate) fn sanitize_replay_items(values: &mut [Value]) {
|
||||
for value in values {
|
||||
match value {
|
||||
Value::Array(values) => sanitize_replay_items(values),
|
||||
@@ -233,7 +179,7 @@ fn sanitize_replay_items(values: &mut [Value]) {
|
||||
object.remove("id");
|
||||
object.remove("encrypted_content");
|
||||
for value in object.values_mut() {
|
||||
strip_encrypted_content(std::slice::from_mut(value));
|
||||
sanitize_replay_items(std::slice::from_mut(value));
|
||||
}
|
||||
}
|
||||
Value::Null | Value::Bool(_) | Value::Number(_) | Value::String(_) => {}
|
||||
@@ -241,53 +187,6 @@ fn sanitize_replay_items(values: &mut [Value]) {
|
||||
}
|
||||
}
|
||||
|
||||
fn strip_encrypted_content(values: &mut [Value]) {
|
||||
for value in values {
|
||||
match value {
|
||||
Value::Array(values) => strip_encrypted_content(values),
|
||||
Value::Object(object) => {
|
||||
object.remove("encrypted_content");
|
||||
for value in object.values_mut() {
|
||||
strip_encrypted_content(std::slice::from_mut(value));
|
||||
}
|
||||
}
|
||||
Value::Null | Value::Bool(_) | Value::Number(_) | Value::String(_) => {}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn limited_replay_json_bytes(input: &[Value], output: &[Value]) -> Option<u64> {
|
||||
let mut writer = ReplaySizeWriter::default();
|
||||
serde_json::to_writer(&mut writer, &(input, output)).ok()?;
|
||||
Some(writer.written)
|
||||
}
|
||||
|
||||
#[derive(Default)]
|
||||
struct ReplaySizeWriter {
|
||||
written: u64,
|
||||
}
|
||||
|
||||
impl std::io::Write for ReplaySizeWriter {
|
||||
fn write(&mut self, buffer: &[u8]) -> std::io::Result<usize> {
|
||||
let next = self
|
||||
.written
|
||||
.checked_add(buffer.len() as u64)
|
||||
.filter(|next| *next <= MAX_REPLAY_SNAPSHOT_BYTES)
|
||||
.ok_or_else(|| {
|
||||
std::io::Error::new(
|
||||
std::io::ErrorKind::FileTooLarge,
|
||||
"response replay snapshot byte limit reached",
|
||||
)
|
||||
})?;
|
||||
self.written = next;
|
||||
Ok(buffer.len())
|
||||
}
|
||||
|
||||
fn flush(&mut self) -> std::io::Result<()> {
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
fn non_empty(value: Option<&str>) -> Option<&str> {
|
||||
value.map(str::trim).filter(|value| !value.is_empty())
|
||||
}
|
||||
|
||||
@@ -32,7 +32,7 @@ use crate::{
|
||||
},
|
||||
upstream_call::{
|
||||
AccountUpstreamContext, QUOTA_VERIFY_LIMIT_REACHED_MESSAGE, QuotaVerificationContext,
|
||||
QuotaVerificationDecision, create_response_with_account_retrying_5xx,
|
||||
QuotaVerificationDecision, create_response_with_account,
|
||||
verify_acquired_quota_if_required,
|
||||
},
|
||||
},
|
||||
@@ -266,7 +266,7 @@ impl ResponseDispatchService {
|
||||
let account = acquired.account.clone();
|
||||
let release_account_id = account.id.clone();
|
||||
let attempt = trace.start_attempt(&release_account_id);
|
||||
let response_result = create_response_with_account_retrying_5xx(
|
||||
let response_result = create_response_with_account(
|
||||
AccountUpstreamContext {
|
||||
codex: &self.codex,
|
||||
account_identity: &self.account_identity,
|
||||
|
||||
@@ -32,7 +32,7 @@ use crate::{
|
||||
},
|
||||
upstream_call::{
|
||||
AccountUpstreamContext, QUOTA_VERIFY_LIMIT_REACHED_MESSAGE, QuotaVerificationContext,
|
||||
QuotaVerificationDecision, create_response_stream_with_account_retrying_5xx,
|
||||
QuotaVerificationDecision, create_response_stream_with_account,
|
||||
verify_acquired_quota_if_required,
|
||||
},
|
||||
},
|
||||
@@ -182,7 +182,7 @@ impl ResponseDispatchService {
|
||||
let account = acquired.account.clone();
|
||||
let release_account_id = account.id.clone();
|
||||
let attempt = trace.start_attempt(&release_account_id);
|
||||
let response_result = create_response_stream_with_account_retrying_5xx(
|
||||
let response_result = create_response_stream_with_account(
|
||||
AccountUpstreamContext {
|
||||
codex: &self.codex,
|
||||
account_identity: &self.account_identity,
|
||||
@@ -215,12 +215,6 @@ impl ResponseDispatchService {
|
||||
let response_metadata = response.response_metadata;
|
||||
self.models
|
||||
.observe_models_etag(response_metadata.models_etag.as_deref());
|
||||
self.cloudflare
|
||||
.capture_set_cookie_headers(&release_account_id, &set_cookie_headers)
|
||||
.await;
|
||||
self.account_pool
|
||||
.sync_passive_rate_limit_headers(&account, &rate_limit_headers)
|
||||
.await;
|
||||
let (prefetched, body) = match prefetch_first_sse_chunk(response.body).await {
|
||||
Ok(prefetched) => prefetched,
|
||||
Err(ResponseDispatchError::Upstream(error)) => {
|
||||
@@ -395,6 +389,7 @@ impl ResponseDispatchService {
|
||||
request,
|
||||
tuple_schema,
|
||||
transport,
|
||||
set_cookie_headers,
|
||||
rate_limit_headers,
|
||||
rate_limit_header_updates,
|
||||
turn_state_update,
|
||||
|
||||
@@ -29,8 +29,8 @@ use crate::{
|
||||
events::extract_sse_usage,
|
||||
responses::{
|
||||
CodexResponsesRequest, CollectedResponse,
|
||||
reconvert_responses_sse_event_tuple_values, response_from_codex_sse,
|
||||
response_sse_event_is_terminal, update_first_response_output_ms,
|
||||
reconvert_responses_sse_event_tuple_values, response_body_has_first_output,
|
||||
response_from_codex_sse, response_sse_event_is_terminal,
|
||||
},
|
||||
sse::{
|
||||
DONE_SSE_FRAME, encode_sse_event, parse_sse_events,
|
||||
@@ -45,6 +45,10 @@ use crate::{
|
||||
},
|
||||
};
|
||||
|
||||
const MAX_LIVE_RESPONSE_CAPTURE_BYTES: usize = 16 * 1024 * 1024;
|
||||
const LIVE_RESPONSE_CAPTURE_LIMIT_DETAIL: &str =
|
||||
"upstream response exceeded the 16 MiB proxy capture limit";
|
||||
|
||||
use super::{
|
||||
sse_failure::{
|
||||
STREAM_DISCONNECTED_CODE, STREAM_DISCONNECTED_MESSAGE, status_code_for_stream_failure,
|
||||
@@ -89,7 +93,9 @@ pub(in crate::dispatch) fn spawn_live_response_stream(
|
||||
.map(TupleSseEventTransformer::new);
|
||||
let mut body_bytes = Vec::new();
|
||||
let mut first_token_ms = None;
|
||||
if !send_live_response_stream_chunk(
|
||||
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(
|
||||
&sender,
|
||||
&mut body_bytes,
|
||||
tuple_transformer.as_mut(),
|
||||
@@ -97,10 +103,17 @@ pub(in crate::dispatch) fn spawn_live_response_stream(
|
||||
)
|
||||
.await
|
||||
{
|
||||
context.account_lease.complete().await;
|
||||
return;
|
||||
Ok(()) => {}
|
||||
Err(LiveStreamWriteError::DownstreamClosed) => {
|
||||
context.account_lease.complete().await;
|
||||
return;
|
||||
}
|
||||
Err(LiveStreamWriteError::CaptureLimitExceeded) => {
|
||||
terminate_oversized_live_stream(context, &sender, &mut body_bytes, first_token_ms)
|
||||
.await;
|
||||
return;
|
||||
}
|
||||
}
|
||||
update_first_response_output_ms(context.started_at, &body_bytes, &mut first_token_ms);
|
||||
|
||||
loop {
|
||||
let next = tokio::select! {
|
||||
@@ -115,7 +128,8 @@ pub(in crate::dispatch) fn spawn_live_response_stream(
|
||||
};
|
||||
match next {
|
||||
Ok(chunk) => {
|
||||
if !send_live_response_stream_chunk(
|
||||
first_output_detector.observe(context.started_at, &chunk, &mut first_token_ms);
|
||||
match send_live_response_stream_chunk(
|
||||
&sender,
|
||||
&mut body_bytes,
|
||||
tuple_transformer.as_mut(),
|
||||
@@ -123,25 +137,46 @@ pub(in crate::dispatch) fn spawn_live_response_stream(
|
||||
)
|
||||
.await
|
||||
{
|
||||
context.account_lease.complete().await;
|
||||
return;
|
||||
Ok(()) => {}
|
||||
Err(LiveStreamWriteError::DownstreamClosed) => {
|
||||
context.account_lease.complete().await;
|
||||
return;
|
||||
}
|
||||
Err(LiveStreamWriteError::CaptureLimitExceeded) => {
|
||||
terminate_oversized_live_stream(
|
||||
context,
|
||||
&sender,
|
||||
&mut body_bytes,
|
||||
first_token_ms,
|
||||
)
|
||||
.await;
|
||||
return;
|
||||
}
|
||||
}
|
||||
update_first_response_output_ms(
|
||||
context.started_at,
|
||||
&body_bytes,
|
||||
&mut first_token_ms,
|
||||
);
|
||||
}
|
||||
Err(error) => {
|
||||
if !flush_live_response_stream_transformer(
|
||||
match flush_live_response_stream_transformer(
|
||||
&sender,
|
||||
&mut body_bytes,
|
||||
tuple_transformer.as_mut(),
|
||||
)
|
||||
.await
|
||||
{
|
||||
context.account_lease.complete().await;
|
||||
return;
|
||||
Ok(()) => {}
|
||||
Err(LiveStreamWriteError::DownstreamClosed) => {
|
||||
context.account_lease.complete().await;
|
||||
return;
|
||||
}
|
||||
Err(LiveStreamWriteError::CaptureLimitExceeded) => {
|
||||
terminate_oversized_live_stream(
|
||||
context,
|
||||
&sender,
|
||||
&mut body_bytes,
|
||||
first_token_ms,
|
||||
)
|
||||
.await;
|
||||
return;
|
||||
}
|
||||
}
|
||||
let detail = error.to_string();
|
||||
let Some(body_text) =
|
||||
@@ -157,15 +192,23 @@ pub(in crate::dispatch) fn spawn_live_response_stream(
|
||||
}
|
||||
}
|
||||
|
||||
if !flush_live_response_stream_transformer(
|
||||
match flush_live_response_stream_transformer(
|
||||
&sender,
|
||||
&mut body_bytes,
|
||||
tuple_transformer.as_mut(),
|
||||
)
|
||||
.await
|
||||
{
|
||||
context.account_lease.complete().await;
|
||||
return;
|
||||
Ok(()) => {}
|
||||
Err(LiveStreamWriteError::DownstreamClosed) => {
|
||||
context.account_lease.complete().await;
|
||||
return;
|
||||
}
|
||||
Err(LiveStreamWriteError::CaptureLimitExceeded) => {
|
||||
terminate_oversized_live_stream(context, &sender, &mut body_bytes, first_token_ms)
|
||||
.await;
|
||||
return;
|
||||
}
|
||||
}
|
||||
let Some(body_text) = send_live_response_stream_tail(&sender, &mut body_bytes, None).await
|
||||
else {
|
||||
@@ -190,9 +233,9 @@ async fn send_live_response_stream_chunk(
|
||||
body_bytes: &mut Vec<u8>,
|
||||
transformer: Option<&mut TupleSseEventTransformer>,
|
||||
chunk: Bytes,
|
||||
) -> bool {
|
||||
) -> Result<(), LiveStreamWriteError> {
|
||||
let chunks = match transformer {
|
||||
Some(transformer) => transformer.push(&chunk),
|
||||
Some(transformer) => transformer.push(&chunk)?,
|
||||
None => vec![chunk],
|
||||
};
|
||||
send_live_response_stream_chunks(sender, body_bytes, chunks).await
|
||||
@@ -202,25 +245,59 @@ async fn flush_live_response_stream_transformer(
|
||||
sender: &mpsc::Sender<Result<Bytes, ResponseDispatchStreamError>>,
|
||||
body_bytes: &mut Vec<u8>,
|
||||
transformer: Option<&mut TupleSseEventTransformer>,
|
||||
) -> bool {
|
||||
) -> Result<(), LiveStreamWriteError> {
|
||||
let Some(transformer) = transformer else {
|
||||
return true;
|
||||
return Ok(());
|
||||
};
|
||||
send_live_response_stream_chunks(sender, body_bytes, transformer.finish()).await
|
||||
send_live_response_stream_chunks(sender, body_bytes, transformer.finish()?).await
|
||||
}
|
||||
|
||||
async fn send_live_response_stream_chunks(
|
||||
sender: &mpsc::Sender<Result<Bytes, ResponseDispatchStreamError>>,
|
||||
body_bytes: &mut Vec<u8>,
|
||||
chunks: Vec<Bytes>,
|
||||
) -> bool {
|
||||
) -> Result<(), LiveStreamWriteError> {
|
||||
for chunk in chunks {
|
||||
if body_bytes.len().saturating_add(chunk.len()) > MAX_LIVE_RESPONSE_CAPTURE_BYTES {
|
||||
return Err(LiveStreamWriteError::CaptureLimitExceeded);
|
||||
}
|
||||
body_bytes.extend_from_slice(&chunk);
|
||||
if sender.send(Ok(chunk)).await.is_err() {
|
||||
return false;
|
||||
return Err(LiveStreamWriteError::DownstreamClosed);
|
||||
}
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
enum LiveStreamWriteError {
|
||||
DownstreamClosed,
|
||||
CaptureLimitExceeded,
|
||||
}
|
||||
|
||||
#[derive(Default)]
|
||||
struct FirstOutputDetector {
|
||||
pending: Vec<u8>,
|
||||
}
|
||||
|
||||
impl FirstOutputDetector {
|
||||
fn observe(&mut self, started_at: Instant, chunk: &[u8], first_token_ms: &mut Option<i64>) {
|
||||
if first_token_ms.is_some() {
|
||||
return;
|
||||
}
|
||||
self.pending.extend_from_slice(chunk);
|
||||
while let Some(frame_end) = sse_frame_end(&self.pending) {
|
||||
let frame = self.pending.drain(..frame_end).collect::<Vec<_>>();
|
||||
if response_body_has_first_output(&frame) {
|
||||
*first_token_ms = Some(elapsed_millis_i64(started_at).max(1));
|
||||
self.pending.clear();
|
||||
return;
|
||||
}
|
||||
}
|
||||
if self.pending.len() > MAX_LIVE_RESPONSE_CAPTURE_BYTES {
|
||||
self.pending.clear();
|
||||
}
|
||||
}
|
||||
true
|
||||
}
|
||||
|
||||
struct TupleSseEventTransformer {
|
||||
@@ -236,22 +313,25 @@ impl TupleSseEventTransformer {
|
||||
}
|
||||
}
|
||||
|
||||
fn push(&mut self, chunk: &[u8]) -> Vec<Bytes> {
|
||||
fn push(&mut self, chunk: &[u8]) -> Result<Vec<Bytes>, LiveStreamWriteError> {
|
||||
self.pending.extend_from_slice(chunk);
|
||||
if self.pending.len() > MAX_LIVE_RESPONSE_CAPTURE_BYTES {
|
||||
return Err(LiveStreamWriteError::CaptureLimitExceeded);
|
||||
}
|
||||
let mut chunks = Vec::new();
|
||||
while let Some(frame_end) = sse_frame_end(&self.pending) {
|
||||
let frame = self.pending.drain(..frame_end).collect::<Vec<_>>();
|
||||
chunks.push(self.transform_frame(&frame));
|
||||
}
|
||||
chunks
|
||||
Ok(chunks)
|
||||
}
|
||||
|
||||
fn finish(&mut self) -> Vec<Bytes> {
|
||||
fn finish(&mut self) -> Result<Vec<Bytes>, LiveStreamWriteError> {
|
||||
if self.pending.is_empty() {
|
||||
return Vec::new();
|
||||
return Ok(Vec::new());
|
||||
}
|
||||
let frame = std::mem::take(&mut self.pending);
|
||||
vec![self.transform_frame(&frame)]
|
||||
Ok(vec![self.transform_frame(&frame)])
|
||||
}
|
||||
|
||||
fn transform_frame(&self, frame: &[u8]) -> Bytes {
|
||||
@@ -277,6 +357,32 @@ impl TupleSseEventTransformer {
|
||||
}
|
||||
}
|
||||
|
||||
async fn terminate_oversized_live_stream(
|
||||
context: LiveResponseStreamContext,
|
||||
sender: &mpsc::Sender<Result<Bytes, ResponseDispatchStreamError>>,
|
||||
body_bytes: &mut Vec<u8>,
|
||||
first_token_ms: Option<i64>,
|
||||
) {
|
||||
tracing::warn!(
|
||||
request_id = %context.request_id,
|
||||
account_id = %context.account_id,
|
||||
captured_bytes = body_bytes.len(),
|
||||
capture_limit_bytes = MAX_LIVE_RESPONSE_CAPTURE_BYTES,
|
||||
"live response capture limit exceeded"
|
||||
);
|
||||
let Some(body_text) = send_live_response_stream_tail(
|
||||
sender,
|
||||
body_bytes,
|
||||
Some(LIVE_RESPONSE_CAPTURE_LIMIT_DETAIL),
|
||||
)
|
||||
.await
|
||||
else {
|
||||
context.account_lease.complete().await;
|
||||
return;
|
||||
};
|
||||
finalize_live_response_stream(context, body_text, first_token_ms).await;
|
||||
}
|
||||
|
||||
async fn send_live_response_stream_tail(
|
||||
sender: &mpsc::Sender<Result<Bytes, ResponseDispatchStreamError>>,
|
||||
body_bytes: &mut Vec<u8>,
|
||||
@@ -354,6 +460,7 @@ pub(in crate::dispatch) struct LiveResponseStreamContext {
|
||||
pub(in crate::dispatch) request: CodexResponsesRequest,
|
||||
pub(in crate::dispatch) tuple_schema: Option<Value>,
|
||||
pub(in crate::dispatch) transport: CodexBackendTransport,
|
||||
pub(in crate::dispatch) set_cookie_headers: Vec<String>,
|
||||
pub(in crate::dispatch) rate_limit_headers: Vec<(String, String)>,
|
||||
pub(in crate::dispatch) rate_limit_header_updates: Option<CodexRateLimitHeaderUpdates>,
|
||||
pub(in crate::dispatch) turn_state_update: Option<CodexTurnStateUpdate>,
|
||||
@@ -402,6 +509,10 @@ pub(in crate::dispatch) async fn finalize_live_response_stream(
|
||||
body: String,
|
||||
first_token_ms: Option<i64>,
|
||||
) {
|
||||
context
|
||||
.cloudflare
|
||||
.capture_set_cookie_headers(&context.account_id, &context.set_cookie_headers)
|
||||
.await;
|
||||
let rate_limit_headers = live_response_rate_limit_headers(&context).await;
|
||||
context
|
||||
.account_pool
|
||||
|
||||
@@ -1,14 +1,11 @@
|
||||
//! 调度层共享的账号级 Codex 上游调用。
|
||||
|
||||
use std::{future::Future, time::Instant};
|
||||
use std::time::Instant;
|
||||
|
||||
use serde_json::Value;
|
||||
|
||||
use crate::{
|
||||
dispatch::{
|
||||
affinity::AccountIdentityService, errors::is_retryable_upstream_5xx_error,
|
||||
recovery::cloudflare::CloudflareRecovery,
|
||||
},
|
||||
dispatch::{affinity::AccountIdentityService, recovery::cloudflare::CloudflareRecovery},
|
||||
fleet::{
|
||||
account::Account,
|
||||
pool::{AccountLease, AccountPoolService},
|
||||
@@ -184,7 +181,7 @@ pub(crate) async fn create_response_with_account(
|
||||
.await
|
||||
}
|
||||
|
||||
async fn create_response_stream_with_account(
|
||||
pub(crate) async fn create_response_stream_with_account(
|
||||
context: AccountUpstreamContext<'_>,
|
||||
request: &CodexResponsesRequest,
|
||||
) -> Result<CodexBackendStreamingResponse, CodexClientError> {
|
||||
@@ -222,97 +219,3 @@ async fn create_response_stream_with_account(
|
||||
)
|
||||
.await
|
||||
}
|
||||
|
||||
const MAX_UPSTREAM_5XX_RETRIES_PER_ACCOUNT: usize = 2;
|
||||
|
||||
async fn retry_upstream_5xx<T, F, Fut>(
|
||||
request_id: &str,
|
||||
account_id: &str,
|
||||
endpoint: Option<&str>,
|
||||
retry_message: &'static str,
|
||||
mut operation: F,
|
||||
) -> Result<T, CodexClientError>
|
||||
where
|
||||
F: FnMut() -> Fut,
|
||||
Fut: Future<Output = Result<T, CodexClientError>>,
|
||||
{
|
||||
let mut retries = 0;
|
||||
loop {
|
||||
let result = operation().await;
|
||||
match result {
|
||||
Err(error)
|
||||
if is_retryable_upstream_5xx_error(&error)
|
||||
&& retries < MAX_UPSTREAM_5XX_RETRIES_PER_ACCOUNT =>
|
||||
{
|
||||
log_upstream_retry(
|
||||
request_id,
|
||||
account_id,
|
||||
endpoint,
|
||||
retries + 1,
|
||||
&error,
|
||||
retry_message,
|
||||
);
|
||||
retries += 1;
|
||||
}
|
||||
Ok(response) => return Ok(response),
|
||||
Err(error) => return Err(error),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn log_upstream_retry(
|
||||
request_id: &str,
|
||||
account_id: &str,
|
||||
endpoint: Option<&str>,
|
||||
retry: usize,
|
||||
error: &CodexClientError,
|
||||
message: &'static str,
|
||||
) {
|
||||
if let Some(endpoint) = endpoint {
|
||||
tracing::debug!(
|
||||
request_id,
|
||||
account_id = %account_id,
|
||||
endpoint,
|
||||
retry,
|
||||
error = %error,
|
||||
message
|
||||
);
|
||||
} else {
|
||||
tracing::debug!(
|
||||
request_id,
|
||||
account_id = %account_id,
|
||||
retry,
|
||||
error = %error,
|
||||
message
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn create_response_with_account_retrying_5xx(
|
||||
context: AccountUpstreamContext<'_>,
|
||||
request: &CodexResponsesRequest,
|
||||
started_at: Instant,
|
||||
) -> Result<CodexBackendResponse, CodexClientError> {
|
||||
retry_upstream_5xx(
|
||||
context.request_id,
|
||||
&context.account.id,
|
||||
None,
|
||||
"upstream response request failed with retryable 5xx",
|
||||
|| create_response_with_account(context, request, started_at),
|
||||
)
|
||||
.await
|
||||
}
|
||||
|
||||
pub(crate) async fn create_response_stream_with_account_retrying_5xx(
|
||||
context: AccountUpstreamContext<'_>,
|
||||
request: &CodexResponsesRequest,
|
||||
) -> Result<CodexBackendStreamingResponse, CodexClientError> {
|
||||
retry_upstream_5xx(
|
||||
context.request_id,
|
||||
&context.account.id,
|
||||
None,
|
||||
"upstream response stream request failed with retryable 5xx",
|
||||
|| create_response_stream_with_account(context, request),
|
||||
)
|
||||
.await
|
||||
}
|
||||
|
||||
@@ -321,6 +321,8 @@ impl CloudflarePathBlockTracker {
|
||||
/// 记录一次 path-block 失败,并返回当前未过期次数。
|
||||
pub async fn record_path_block(&self, account_id: &str, now: DateTime<Utc>) -> u32 {
|
||||
let mut counts = self.counts.write().await;
|
||||
counts
|
||||
.retain(|_, state| now.signed_duration_since(state.last_at) <= PATH_BLOCK_STALE_AFTER);
|
||||
let count = counts
|
||||
.get(account_id)
|
||||
.filter(|state| now.signed_duration_since(state.last_at) <= PATH_BLOCK_STALE_AFTER)
|
||||
@@ -377,6 +379,9 @@ impl CloudflareChallengeCooldownTracker {
|
||||
) -> CloudflareChallengeCooldown {
|
||||
let challenge_count = {
|
||||
let mut states = self.states.write().await;
|
||||
states.retain(|_, state| {
|
||||
now.signed_duration_since(state.updated_at) <= CHALLENGE_STALE_AFTER
|
||||
});
|
||||
let challenge_count = states
|
||||
.get(account_id)
|
||||
.filter(|state| {
|
||||
|
||||
@@ -17,6 +17,7 @@ const OPENAI_OAUTH_AUTHORIZE_ENDPOINT: &str = "https://auth.openai.com/oauth/aut
|
||||
const OPENAI_OAUTH_REDIRECT_URI: &str = "http://localhost:1455/auth/callback";
|
||||
const OPENAI_OAUTH_SCOPE: &str = "openid profile email offline_access";
|
||||
const OAUTH_SESSION_TTL_MINUTES: i64 = 30;
|
||||
const MAX_PENDING_OAUTH_SESSIONS: usize = 1_024;
|
||||
|
||||
#[derive(Clone)]
|
||||
pub struct AccountOAuthService {
|
||||
@@ -51,6 +52,14 @@ impl AccountOAuthService {
|
||||
|
||||
let mut sessions = self.sessions.lock().await;
|
||||
sessions.retain(|_, session| session.expires_at > now);
|
||||
if sessions.len() >= MAX_PENDING_OAUTH_SESSIONS
|
||||
&& let Some(oldest_id) = sessions
|
||||
.iter()
|
||||
.min_by_key(|(_, session)| session.expires_at)
|
||||
.map(|(session_id, _)| session_id.clone())
|
||||
{
|
||||
sessions.remove(&oldest_id);
|
||||
}
|
||||
sessions.insert(session_id.clone(), session);
|
||||
drop(sessions);
|
||||
|
||||
|
||||
@@ -1,6 +1,10 @@
|
||||
//! 运行时 quota refresh 业务服务。
|
||||
|
||||
use std::{collections::HashMap, sync::Arc, time::Duration};
|
||||
use std::{
|
||||
collections::{HashMap, HashSet},
|
||||
sync::Arc,
|
||||
time::Duration,
|
||||
};
|
||||
|
||||
use chrono::Utc;
|
||||
use thiserror::Error;
|
||||
@@ -8,7 +12,7 @@ use tokio::time::{Instant, sleep};
|
||||
use tracing::warn;
|
||||
|
||||
use crate::fleet::{
|
||||
account::AccountStatus,
|
||||
account::{Account, AccountStatus},
|
||||
pool::AccountPoolService,
|
||||
quota::{
|
||||
quota_from_usage, quota_snapshot_limit_reached, quota_snapshot_limit_window_seconds,
|
||||
@@ -120,13 +124,16 @@ impl QuotaRefreshService {
|
||||
let now = Instant::now();
|
||||
let now_utc = Utc::now();
|
||||
|
||||
let candidate_ids = accounts
|
||||
.iter()
|
||||
.filter(|account| quota_refresh_candidate(account))
|
||||
.map(|account| account.id.as_str())
|
||||
.collect::<HashSet<_>>();
|
||||
last_refreshed.retain(|account_id, _| candidate_ids.contains(account_id.as_str()));
|
||||
|
||||
let mut candidates = accounts
|
||||
.into_iter()
|
||||
.filter(|account| {
|
||||
account.status == AccountStatus::QuotaExhausted
|
||||
|| (account.status == AccountStatus::Active
|
||||
&& (account.quota_limit_reached || account.quota_verify_required))
|
||||
})
|
||||
.filter(quota_refresh_candidate)
|
||||
.peekable();
|
||||
|
||||
while let Some(account) = candidates.next() {
|
||||
@@ -245,6 +252,12 @@ impl QuotaRefreshService {
|
||||
}
|
||||
}
|
||||
|
||||
fn quota_refresh_candidate(account: &Account) -> bool {
|
||||
account.status == AccountStatus::QuotaExhausted
|
||||
|| (account.status == AccountStatus::Active
|
||||
&& (account.quota_limit_reached || account.quota_verify_required))
|
||||
}
|
||||
|
||||
fn quota_cooldown_ready(cooldown_until: chrono::DateTime<Utc>, now: chrono::DateTime<Utc>) -> bool {
|
||||
cooldown_until
|
||||
.checked_add_signed(chrono::Duration::seconds(COOLDOWN_REFRESH_GRACE_SECS))
|
||||
|
||||
@@ -20,6 +20,8 @@ use crate::{
|
||||
upstream::openai::protocol::responses::CodexResponsesRequest,
|
||||
};
|
||||
|
||||
const MAX_CAPTURED_BODY_BYTES: usize = 4 * 1024;
|
||||
|
||||
/// 成功与失败事实的唯一写入入口。
|
||||
#[derive(Clone)]
|
||||
pub struct Recorder {
|
||||
@@ -83,6 +85,10 @@ impl Recorder {
|
||||
self.ops_errors.append(&event).await?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub(crate) fn captures_body(&self) -> bool {
|
||||
self.capture_body
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) struct DispatchErrorLogRecord<'a> {
|
||||
@@ -175,27 +181,55 @@ fn is_usage_fact(event: &UsageRecord) -> bool {
|
||||
}
|
||||
|
||||
fn apply_success_capture_body_policy(event: &mut UsageRecord, capture_body: bool) {
|
||||
if capture_body {
|
||||
return;
|
||||
if !capture_body {
|
||||
remove_body_fields(&mut event.metadata);
|
||||
} else {
|
||||
limit_body_fields(&mut event.metadata);
|
||||
}
|
||||
remove_body_fields(&mut event.metadata);
|
||||
}
|
||||
|
||||
fn apply_error_capture_body_policy(event: &mut OpsErrorLog, capture_body: bool) {
|
||||
let Some(metadata) = event.metadata.as_object_mut() else {
|
||||
if !capture_body {
|
||||
remove_body_fields(&mut event.metadata);
|
||||
} else {
|
||||
limit_body_fields(&mut event.metadata);
|
||||
}
|
||||
}
|
||||
|
||||
fn limit_body_fields(metadata: &mut Value) {
|
||||
let Some(metadata) = metadata.as_object_mut() else {
|
||||
return;
|
||||
};
|
||||
for key in body_fields() {
|
||||
if capture_body {
|
||||
if let Some(Value::String(value)) = metadata.get_mut(key) {
|
||||
value.truncate(4096);
|
||||
let Some(value) = metadata.get_mut(key) else {
|
||||
continue;
|
||||
};
|
||||
match value {
|
||||
Value::String(value) => truncate_utf8(value, MAX_CAPTURED_BODY_BYTES),
|
||||
value => {
|
||||
let Ok(mut encoded) = serde_json::to_string(value) else {
|
||||
continue;
|
||||
};
|
||||
if encoded.len() > MAX_CAPTURED_BODY_BYTES {
|
||||
truncate_utf8(&mut encoded, MAX_CAPTURED_BODY_BYTES);
|
||||
*value = Value::String(encoded);
|
||||
}
|
||||
}
|
||||
} else {
|
||||
metadata.remove(key);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn truncate_utf8(value: &mut String, max_bytes: usize) {
|
||||
if value.len() <= max_bytes {
|
||||
return;
|
||||
}
|
||||
let mut end = max_bytes;
|
||||
while !value.is_char_boundary(end) {
|
||||
end -= 1;
|
||||
}
|
||||
value.truncate(end);
|
||||
}
|
||||
|
||||
fn remove_body_fields(metadata: &mut Value) {
|
||||
let Some(metadata) = metadata.as_object_mut() else {
|
||||
return;
|
||||
|
||||
@@ -1,4 +1,4 @@
|
||||
use std::{collections::BTreeSet, time::Instant};
|
||||
use std::{collections::BTreeSet, sync::Arc, time::Instant};
|
||||
|
||||
use serde::Serialize;
|
||||
use serde_json::{Map, Value};
|
||||
@@ -68,6 +68,8 @@ pub struct CodexResponsesRequest {
|
||||
pub parent_thread_id: Option<String>,
|
||||
/// 已知 previous response 的持久化范围,仅用于本地 transport 校验。
|
||||
pub previous_response_scope: Option<PreviousResponseScope>,
|
||||
/// 当前下游 WebSocket 持有的完整历史前缀,仅用于连接内换号恢复。
|
||||
pub local_replay_input: Option<Arc<Vec<Value>>>,
|
||||
/// 流式响应在交给下游前的本地提交策略。
|
||||
pub stream_commit_policy: StreamCommitPolicy,
|
||||
}
|
||||
@@ -537,6 +539,7 @@ impl CodexResponsesRequest {
|
||||
codex_window_id: None,
|
||||
parent_thread_id: None,
|
||||
previous_response_scope: None,
|
||||
local_replay_input: None,
|
||||
stream_commit_policy: StreamCommitPolicy::FirstForwardableEvent,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -9,7 +9,7 @@ use std::{
|
||||
|
||||
use async_trait::async_trait;
|
||||
use bytes::Bytes;
|
||||
use futures::{Stream, TryStreamExt};
|
||||
use futures::{Stream, StreamExt, TryStreamExt};
|
||||
use reqwest::{
|
||||
Client, Response as ReqwestResponse, StatusCode,
|
||||
header::{
|
||||
@@ -53,6 +53,8 @@ use super::websocket_pool::{
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
const MAX_UPSTREAM_ERROR_BODY_BYTES: usize = 1024 * 1024;
|
||||
const UPSTREAM_CONNECT_TIMEOUT: Duration = Duration::from_secs(15);
|
||||
const UPSTREAM_STREAM_IDLE_TIMEOUT: Duration = Duration::from_secs(5 * 60);
|
||||
const X_CODEX_WS_STREAM_REQUEST_START_MS_CLIENT_METADATA_KEY: &str =
|
||||
"x-codex-ws-stream-request-start-ms";
|
||||
type ReqwestClientCacheKey = (bool, Option<String>);
|
||||
@@ -75,6 +77,7 @@ pub fn build_reqwest_client(force_http11: bool) -> Result<Client, CustomCaError>
|
||||
.use_rustls_tls()
|
||||
.no_proxy()
|
||||
.pool_max_idle_per_host(4)
|
||||
.connect_timeout(UPSTREAM_CONNECT_TIMEOUT)
|
||||
.tcp_keepalive(Duration::from_secs(30))
|
||||
.gzip(true)
|
||||
.brotli(true)
|
||||
@@ -114,6 +117,12 @@ pub enum CodexClientError {
|
||||
/// SSE 响应解析失败。
|
||||
#[error("invalid upstream SSE response: {0}")]
|
||||
InvalidSse(#[from] SseError),
|
||||
/// HTTP/SSE 上游在空闲窗口内没有发送任何数据。
|
||||
#[error("upstream HTTP/SSE stream idle for {timeout:?}")]
|
||||
StreamIdleTimeout {
|
||||
/// 相邻数据块允许的最大空闲时间。
|
||||
timeout: Duration,
|
||||
},
|
||||
/// WebSocket 请求编码失败。
|
||||
#[error("failed to encode websocket request: {0}")]
|
||||
WebSocketEncode(#[source] serde_json::Error),
|
||||
|
||||
@@ -72,7 +72,7 @@ impl CodexBackendClient {
|
||||
|
||||
let mut body_bytes = Vec::new();
|
||||
let mut first_token_ms = None;
|
||||
let mut stream = response.bytes_stream();
|
||||
let mut stream = http_sse_stream(response);
|
||||
while let Some(chunk) = stream.try_next().await? {
|
||||
body_bytes.extend_from_slice(&chunk);
|
||||
response_meta::update_first_token_ms(started_at, &body_bytes, &mut first_token_ms);
|
||||
@@ -128,7 +128,7 @@ impl CodexBackendClient {
|
||||
}
|
||||
|
||||
Ok(CodexBackendStreamingResponse {
|
||||
body: Box::pin(response.bytes_stream().map_err(CodexClientError::Http)),
|
||||
body: http_sse_stream(response),
|
||||
transport: CodexBackendTransport::HttpSse,
|
||||
turn_state,
|
||||
set_cookie_headers,
|
||||
@@ -590,3 +590,21 @@ impl CodexBackendClient {
|
||||
Ok(headers)
|
||||
}
|
||||
}
|
||||
|
||||
fn http_sse_stream(response: ReqwestResponse) -> CodexBackendSseStream {
|
||||
let stream: CodexBackendSseStream =
|
||||
Box::pin(response.bytes_stream().map_err(CodexClientError::Http));
|
||||
Box::pin(futures::stream::unfold(Some(stream), |stream| async move {
|
||||
let mut stream = stream?;
|
||||
match tokio::time::timeout(UPSTREAM_STREAM_IDLE_TIMEOUT, stream.next()).await {
|
||||
Ok(Some(chunk)) => Some((chunk, Some(stream))),
|
||||
Ok(None) => None,
|
||||
Err(_) => Some((
|
||||
Err(CodexClientError::StreamIdleTimeout {
|
||||
timeout: UPSTREAM_STREAM_IDLE_TIMEOUT,
|
||||
}),
|
||||
None,
|
||||
)),
|
||||
}
|
||||
}))
|
||||
}
|
||||
|
||||
@@ -53,8 +53,9 @@ const CODEX_RESPONSES_PATH: &str = "/codex/responses";
|
||||
const WEBSOCKET_EXTENSIONS: &str = "permessage-deflate; client_max_window_bits";
|
||||
const WEBSOCKET_RECEIVE_IDLE_TIMEOUT: Duration = Duration::from_secs(20);
|
||||
const WEBSOCKET_ACTIVE_STREAM_IDLE_TIMEOUT: Duration = Duration::from_secs(5 * 60);
|
||||
const WEBSOCKET_CONNECT_TIMEOUT: Duration = Duration::from_secs(15);
|
||||
const WEBSOCKET_SEND_TIMEOUT: Duration = WEBSOCKET_ACTIVE_STREAM_IDLE_TIMEOUT;
|
||||
const WEBSOCKET_STREAM_BUFFER: usize = 16;
|
||||
const WEBSOCKET_FIRST_TOKEN_FRESH_RETRY_ATTEMPTS: usize = 2;
|
||||
/// WebSocket audit artifact 输出目录环境变量。
|
||||
pub const WS_AUDIT_DIR_ENV: &str = "CODEX_PROXY_WS_AUDIT_DIR";
|
||||
|
||||
@@ -188,10 +189,9 @@ pub use frames::{
|
||||
};
|
||||
use frames::{
|
||||
WebSocketStreamPoolReturn, WebSocketTerminalKind, audit_header_value,
|
||||
collect_websocket_response, is_initial_event_timeout,
|
||||
prefetch_stream_frames_until_output_or_terminal, reusable_websocket_metadata,
|
||||
reused_stream_prefetch_error, stream_websocket_response, websocket_audit_file_name,
|
||||
websocket_connection_metadata,
|
||||
collect_websocket_response, prefetch_stream_frames_until_output_or_terminal,
|
||||
reusable_websocket_metadata, reused_stream_prefetch_error, stream_websocket_response,
|
||||
websocket_audit_file_name, websocket_connection_metadata,
|
||||
};
|
||||
|
||||
impl CodexWebSocketRequest {
|
||||
@@ -327,13 +327,36 @@ async fn connect_websocket(
|
||||
let connector = super::tls::maybe_build_rustls_client_config_with_custom_ca()
|
||||
.map_err(|error| tungstenite::Error::Io(std::io::Error::other(error)))?
|
||||
.map(Connector::Rustls);
|
||||
match connect_async_tls_with_config(request, Some(websocket_config()), false, connector).await {
|
||||
let result = timeout(
|
||||
WEBSOCKET_CONNECT_TIMEOUT,
|
||||
connect_async_tls_with_config(request, Some(websocket_config()), false, connector),
|
||||
)
|
||||
.await
|
||||
.map_err(|_| CodexWebSocketExchangeError::ConnectTimeout {
|
||||
timeout: WEBSOCKET_CONNECT_TIMEOUT,
|
||||
})?;
|
||||
match result {
|
||||
Ok((websocket, response)) => Ok((websocket, response)),
|
||||
Err(tungstenite::Error::Http(response)) => Err(websocket_opening_error(response.as_ref())),
|
||||
Err(error) => Err(error.into()),
|
||||
}
|
||||
}
|
||||
|
||||
async fn send_websocket_request(
|
||||
websocket: &PumpedWebSocket,
|
||||
payload_text: &str,
|
||||
) -> Result<(), CodexWebSocketExchangeError> {
|
||||
timeout(
|
||||
WEBSOCKET_SEND_TIMEOUT,
|
||||
websocket.send(Message::Text(payload_text.to_string().into())),
|
||||
)
|
||||
.await
|
||||
.map_err(|_| CodexWebSocketExchangeError::SendTimeout {
|
||||
timeout: WEBSOCKET_SEND_TIMEOUT,
|
||||
})??;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// 建立连接并交给后台 pump 托管;`keepalive` 决定该连接是否做主动保活。
|
||||
///
|
||||
/// 非池化(即用即弃)连接用 [`PumpKeepalive::disabled`];池化连接用连接池派生的保活策略,
|
||||
@@ -420,9 +443,7 @@ async fn execute_fresh_response_create_request(
|
||||
) -> Result<CodexWebSocketExchange, CodexWebSocketExchangeError> {
|
||||
let (websocket, response) =
|
||||
connect_pumped_websocket(request.connection(), PumpKeepalive::disabled()).await?;
|
||||
websocket
|
||||
.send(Message::Text(request.payload_text().to_string().into()))
|
||||
.await?;
|
||||
send_websocket_request(&websocket, request.payload_text()).await?;
|
||||
|
||||
let metadata = websocket_connection_metadata(&response);
|
||||
let (exchange, websocket, _metadata, _continuation, _terminal) = collect_websocket_response(
|
||||
@@ -444,9 +465,7 @@ async fn execute_fresh_response_create_request_stream(
|
||||
) -> Result<CodexWebSocketStreamingExchange, CodexWebSocketExchangeError> {
|
||||
let (mut websocket, response) =
|
||||
connect_pumped_websocket(request.connection(), PumpKeepalive::disabled()).await?;
|
||||
websocket
|
||||
.send(Message::Text(request.payload_text().to_string().into()))
|
||||
.await?;
|
||||
send_websocket_request(&websocket, request.payload_text()).await?;
|
||||
|
||||
let mut metadata = websocket_connection_metadata(&response);
|
||||
let prefetched_frames = if should_prefetch_until_output_or_terminal(request) {
|
||||
@@ -475,41 +494,6 @@ async fn execute_fresh_response_create_request_stream(
|
||||
))
|
||||
}
|
||||
|
||||
async fn execute_fresh_response_create_request_with_retries(
|
||||
request: &CodexWebSocketRequest,
|
||||
started_at: Instant,
|
||||
initial_event_timeout: Option<Duration>,
|
||||
attempts: usize,
|
||||
) -> Result<CodexWebSocketExchange, CodexWebSocketExchangeError> {
|
||||
let attempts = attempts.max(1);
|
||||
for attempt in 1..=attempts {
|
||||
match execute_fresh_response_create_request(request, started_at, initial_event_timeout)
|
||||
.await
|
||||
{
|
||||
Ok(exchange) => return Ok(exchange),
|
||||
Err(error) if attempt < attempts && is_initial_event_timeout(&error) => continue,
|
||||
Err(error) => return Err(error),
|
||||
}
|
||||
}
|
||||
unreachable!("fresh websocket retry loop always returns");
|
||||
}
|
||||
|
||||
async fn execute_fresh_response_create_request_stream_with_retries(
|
||||
request: &CodexWebSocketRequest,
|
||||
initial_event_timeout: Option<Duration>,
|
||||
attempts: usize,
|
||||
) -> Result<CodexWebSocketStreamingExchange, CodexWebSocketExchangeError> {
|
||||
let attempts = attempts.max(1);
|
||||
for attempt in 1..=attempts {
|
||||
match execute_fresh_response_create_request_stream(request, initial_event_timeout).await {
|
||||
Ok(exchange) => return Ok(exchange),
|
||||
Err(error) if attempt < attempts && is_initial_event_timeout(&error) => continue,
|
||||
Err(error) => return Err(error),
|
||||
}
|
||||
}
|
||||
unreachable!("fresh websocket stream retry loop always returns");
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Pool-aware execution
|
||||
// ---------------------------------------------------------------------------
|
||||
@@ -526,11 +510,10 @@ pub(crate) async fn execute_response_create_request_with_pool(
|
||||
PreviousResponseUnavailableReason::PoolUnavailable,
|
||||
));
|
||||
}
|
||||
return execute_fresh_response_create_request_with_retries(
|
||||
return execute_fresh_response_create_request(
|
||||
request,
|
||||
started_at,
|
||||
fallback_initial_event_timeout,
|
||||
WEBSOCKET_FIRST_TOKEN_FRESH_RETRY_ATTEMPTS,
|
||||
)
|
||||
.await;
|
||||
};
|
||||
@@ -564,11 +547,10 @@ pub(crate) async fn execute_response_create_request_with_pool(
|
||||
PreviousResponseUnavailableReason::ReusedConnectionLost,
|
||||
));
|
||||
}
|
||||
let mut exchange = execute_fresh_response_create_request_with_retries(
|
||||
let mut exchange = execute_fresh_response_create_request(
|
||||
request,
|
||||
started_at,
|
||||
pool.initial_event_timeout(),
|
||||
WEBSOCKET_FIRST_TOKEN_FRESH_RETRY_ATTEMPTS,
|
||||
)
|
||||
.await?;
|
||||
exchange.pool_decision = Some(WebSocketPoolDecision::retry_after_stale_reuse());
|
||||
@@ -591,24 +573,10 @@ pub(crate) async fn execute_response_create_request_with_pool(
|
||||
let result =
|
||||
execute_fresh_pooled_response_create_request(request, pool, lease, started_at)
|
||||
.await;
|
||||
match result {
|
||||
Ok(mut exchange) => {
|
||||
exchange.pool_decision = Some(WebSocketPoolDecision::new());
|
||||
Ok(exchange)
|
||||
}
|
||||
Err(error) if is_initial_event_timeout(&error) => {
|
||||
let mut exchange = execute_fresh_response_create_request_with_retries(
|
||||
request,
|
||||
started_at,
|
||||
pool.initial_event_timeout(),
|
||||
1,
|
||||
)
|
||||
.await?;
|
||||
exchange.pool_decision = Some(WebSocketPoolDecision::new());
|
||||
Ok(exchange)
|
||||
}
|
||||
Err(error) => Err(error),
|
||||
}
|
||||
result.map(|mut exchange| {
|
||||
exchange.pool_decision = Some(WebSocketPoolDecision::new());
|
||||
exchange
|
||||
})
|
||||
}
|
||||
WebSocketPoolAcquire::Bypass(reason) => {
|
||||
if !request.continuation().permits_fresh_connection() {
|
||||
@@ -622,11 +590,10 @@ pub(crate) async fn execute_response_create_request_with_pool(
|
||||
}
|
||||
}));
|
||||
}
|
||||
let mut exchange = execute_fresh_response_create_request_with_retries(
|
||||
let mut exchange = execute_fresh_response_create_request(
|
||||
request,
|
||||
started_at,
|
||||
pool.initial_event_timeout(),
|
||||
WEBSOCKET_FIRST_TOKEN_FRESH_RETRY_ATTEMPTS,
|
||||
)
|
||||
.await?;
|
||||
exchange.pool_decision = Some(WebSocketPoolDecision::bypass(reason));
|
||||
@@ -646,10 +613,9 @@ pub(crate) async fn execute_response_create_request_stream_with_pool(
|
||||
PreviousResponseUnavailableReason::PoolUnavailable,
|
||||
));
|
||||
}
|
||||
return execute_fresh_response_create_request_stream_with_retries(
|
||||
return execute_fresh_response_create_request_stream(
|
||||
request,
|
||||
fallback_initial_event_timeout,
|
||||
WEBSOCKET_FIRST_TOKEN_FRESH_RETRY_ATTEMPTS,
|
||||
)
|
||||
.await;
|
||||
};
|
||||
@@ -678,10 +644,9 @@ pub(crate) async fn execute_response_create_request_stream_with_pool(
|
||||
PreviousResponseUnavailableReason::ReusedConnectionLost,
|
||||
));
|
||||
}
|
||||
let mut exchange = execute_fresh_response_create_request_stream_with_retries(
|
||||
let mut exchange = execute_fresh_response_create_request_stream(
|
||||
request,
|
||||
pool.initial_event_timeout(),
|
||||
WEBSOCKET_FIRST_TOKEN_FRESH_RETRY_ATTEMPTS,
|
||||
)
|
||||
.await?;
|
||||
exchange.pool_decision = Some(WebSocketPoolDecision::retry_after_stale_reuse());
|
||||
@@ -703,23 +668,10 @@ pub(crate) async fn execute_response_create_request_stream_with_pool(
|
||||
}
|
||||
let result =
|
||||
execute_fresh_pooled_response_create_request_stream(request, pool, lease).await;
|
||||
match result {
|
||||
Ok(mut exchange) => {
|
||||
exchange.pool_decision = Some(WebSocketPoolDecision::new());
|
||||
Ok(exchange)
|
||||
}
|
||||
Err(error) if is_initial_event_timeout(&error) => {
|
||||
let mut exchange = execute_fresh_response_create_request_stream_with_retries(
|
||||
request,
|
||||
pool.initial_event_timeout(),
|
||||
1,
|
||||
)
|
||||
.await?;
|
||||
exchange.pool_decision = Some(WebSocketPoolDecision::new());
|
||||
Ok(exchange)
|
||||
}
|
||||
Err(error) => Err(error),
|
||||
}
|
||||
result.map(|mut exchange| {
|
||||
exchange.pool_decision = Some(WebSocketPoolDecision::new());
|
||||
exchange
|
||||
})
|
||||
}
|
||||
WebSocketPoolAcquire::Bypass(reason) => {
|
||||
if !request.continuation().permits_fresh_connection() {
|
||||
@@ -733,12 +685,9 @@ pub(crate) async fn execute_response_create_request_stream_with_pool(
|
||||
}
|
||||
}));
|
||||
}
|
||||
let mut exchange = execute_fresh_response_create_request_stream_with_retries(
|
||||
request,
|
||||
pool.initial_event_timeout(),
|
||||
WEBSOCKET_FIRST_TOKEN_FRESH_RETRY_ATTEMPTS,
|
||||
)
|
||||
.await?;
|
||||
let mut exchange =
|
||||
execute_fresh_response_create_request_stream(request, pool.initial_event_timeout())
|
||||
.await?;
|
||||
exchange.pool_decision = Some(WebSocketPoolDecision::bypass(reason));
|
||||
Ok(exchange)
|
||||
}
|
||||
@@ -765,12 +714,9 @@ async fn execute_fresh_pooled_response_create_request(
|
||||
return Err(error);
|
||||
}
|
||||
};
|
||||
if let Err(error) = websocket
|
||||
.send(Message::Text(request.payload_text().to_string().into()))
|
||||
.await
|
||||
{
|
||||
if let Err(error) = send_websocket_request(&websocket, request.payload_text()).await {
|
||||
lease.discard().await;
|
||||
return Err(error.into());
|
||||
return Err(error);
|
||||
}
|
||||
|
||||
let now = tokio::time::Instant::now();
|
||||
@@ -823,12 +769,9 @@ async fn execute_fresh_pooled_response_create_request_stream(
|
||||
return Err(error);
|
||||
}
|
||||
};
|
||||
if let Err(error) = websocket
|
||||
.send(Message::Text(request.payload_text().to_string().into()))
|
||||
.await
|
||||
{
|
||||
if let Err(error) = send_websocket_request(&websocket, request.payload_text()).await {
|
||||
lease.discard().await;
|
||||
return Err(error.into());
|
||||
return Err(error);
|
||||
}
|
||||
|
||||
let now = tokio::time::Instant::now();
|
||||
@@ -871,10 +814,7 @@ async fn execute_pooled_response_create_request(
|
||||
connection: PooledWebSocketConnection,
|
||||
started_at: Instant,
|
||||
) -> Result<CodexWebSocketExchange, CodexWebSocketExchangeError> {
|
||||
if let Err(error) = connection
|
||||
.websocket
|
||||
.send(Message::Text(request.payload_text().to_string().into()))
|
||||
.await
|
||||
if let Err(error) = send_websocket_request(&connection.websocket, request.payload_text()).await
|
||||
{
|
||||
lease.discard().await;
|
||||
return Err(
|
||||
@@ -931,10 +871,7 @@ async fn execute_pooled_response_create_request_stream(
|
||||
let mut metadata = reusable_websocket_metadata(connection.metadata);
|
||||
let created_at = connection.created_at;
|
||||
let continuation = connection.continuation;
|
||||
if let Err(error) = websocket
|
||||
.send(Message::Text(request.payload_text().to_string().into()))
|
||||
.await
|
||||
{
|
||||
if let Err(error) = send_websocket_request(&websocket, request.payload_text()).await {
|
||||
lease.discard().await;
|
||||
return Err(
|
||||
CodexWebSocketExchangeError::ReusedConnectionDiedBeforeFirstOutput {
|
||||
|
||||
@@ -73,6 +73,18 @@ pub enum CodexWebSocketExchangeError {
|
||||
/// WebSocket 传输失败。
|
||||
#[error("websocket transport error: {0}")]
|
||||
Transport(#[from] tungstenite::Error),
|
||||
/// DNS、TCP、TLS 或 WebSocket upgrade 未在限定时间内完成。
|
||||
#[error("websocket connect timed out after {timeout:?}")]
|
||||
ConnectTimeout {
|
||||
/// 建连超时时长。
|
||||
timeout: Duration,
|
||||
},
|
||||
/// 请求帧未在限定时间内写入上游连接。
|
||||
#[error("websocket request send timed out after {timeout:?}")]
|
||||
SendTimeout {
|
||||
/// 发送超时时长。
|
||||
timeout: Duration,
|
||||
},
|
||||
/// SSE 聚合结果无法解析。
|
||||
#[error("invalid websocket SSE response: {0}")]
|
||||
InvalidSse(#[from] SseError),
|
||||
@@ -102,14 +114,14 @@ pub enum CodexWebSocketExchangeError {
|
||||
/// 上游返回非文本事件帧。
|
||||
#[error("unexpected binary websocket event")]
|
||||
UnexpectedBinaryEvent,
|
||||
/// 复用的池连接在首个真实输出前失效。
|
||||
#[error("reused websocket connection died before first output: {message}")]
|
||||
/// 复用的池连接在收到首个上游事件前失效。
|
||||
#[error("reused websocket connection died before first upstream event: {message}")]
|
||||
ReusedConnectionDiedBeforeFirstOutput {
|
||||
/// 底层失效原因。
|
||||
message: String,
|
||||
},
|
||||
/// 首个真实输出在配置的绝对超时内未到达(连接落到病态上游后端)。
|
||||
#[error("websocket first output not received within {timeout:?}")]
|
||||
/// 建连并发送后,上游在配置时间内没有产生任何事件。
|
||||
#[error("websocket first upstream event not received within {timeout:?}")]
|
||||
InitialEventTimeout {
|
||||
/// 首个上游事件超时时长。
|
||||
timeout: Duration,
|
||||
@@ -748,13 +760,6 @@ fn receive_idle_timeout(
|
||||
}
|
||||
}
|
||||
|
||||
pub(super) fn is_initial_event_timeout(error: &CodexWebSocketExchangeError) -> bool {
|
||||
matches!(
|
||||
error,
|
||||
CodexWebSocketExchangeError::InitialEventTimeout { .. }
|
||||
)
|
||||
}
|
||||
|
||||
async fn next_websocket_message(
|
||||
websocket: &mut PumpedWebSocket,
|
||||
receive_timeout: Duration,
|
||||
|
||||
@@ -85,7 +85,7 @@ pub struct CodexWebSocketPoolConfig {
|
||||
pub ping_timeout: Duration,
|
||||
/// idle socket 无活动多久后视为失活。
|
||||
pub liveness_timeout: Option<Duration>,
|
||||
/// 首个上游事件到达前的超时;`None` 表示禁用首事件超时。
|
||||
/// 建连并发送后首个上游事件到达前的超时;`None` 表示禁用。
|
||||
pub initial_event_timeout: Option<Duration>,
|
||||
}
|
||||
|
||||
@@ -144,7 +144,7 @@ impl CodexWebSocketPool {
|
||||
self.config.keepalive()
|
||||
}
|
||||
|
||||
/// 首个上游事件到达前的超时;`None` 表示禁用首事件超时。
|
||||
/// 建连并发送后首个上游事件到达前的超时;`None` 表示禁用。
|
||||
pub(crate) fn initial_event_timeout(&self) -> Option<Duration> {
|
||||
self.config.initial_event_timeout
|
||||
}
|
||||
|
||||
@@ -32,7 +32,6 @@ async fn admin_accounts_lifecycle_should_update_and_delete_accounts() {
|
||||
function_call_ids: Vec::new(),
|
||||
variant_hash: None,
|
||||
continuation_scope: codex_proxy_rs::upstream::openai::protocol::responses::PreviousResponseScope::ConnectionLocal,
|
||||
replay: None,
|
||||
created_at: chrono::Utc::now(),
|
||||
},
|
||||
chrono::Duration::hours(1),
|
||||
|
||||
@@ -121,6 +121,82 @@ async fn responses_websocket_should_forward_multiple_response_create_requests_on
|
||||
assert_eq!(upstream_requests[1]["type"], "response.create");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn responses_websocket_should_replay_connection_history_after_previous_response_not_found() {
|
||||
let upstream_listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
|
||||
let upstream_base_url = format!("http://{}", upstream_listener.local_addr().unwrap());
|
||||
let upstream = tokio::spawn(async move {
|
||||
let (first_stream, _) = upstream_listener.accept().await.unwrap();
|
||||
let mut first = accept_upstream_websocket(first_stream).await.unwrap();
|
||||
let initial = receive_upstream_request(&mut first).await;
|
||||
send_upstream_response(&mut first, "resp_first", "first answer").await;
|
||||
|
||||
let continued = receive_upstream_request(&mut first).await;
|
||||
first
|
||||
.send(Message::Text(
|
||||
json!({
|
||||
"type": "response.failed",
|
||||
"response": {
|
||||
"id": "resp_missing",
|
||||
"status": "failed",
|
||||
"error": {
|
||||
"code": "previous_response_not_found",
|
||||
"message": "Previous response with id resp_first was not found",
|
||||
},
|
||||
},
|
||||
})
|
||||
.to_string()
|
||||
.into(),
|
||||
))
|
||||
.await
|
||||
.unwrap();
|
||||
first.close(None).await.unwrap();
|
||||
|
||||
let (replay_stream, _) = upstream_listener.accept().await.unwrap();
|
||||
let mut replay = accept_upstream_websocket(replay_stream).await.unwrap();
|
||||
let replayed = receive_upstream_request(&mut replay).await;
|
||||
send_upstream_response(&mut replay, "resp_recovered", "recovered").await;
|
||||
(initial, continued, replayed)
|
||||
});
|
||||
let (app, api_key, _dir) =
|
||||
crate::dispatch::service::test_app_with_account(upstream_base_url).await;
|
||||
let server = spawn_app(app).await;
|
||||
let mut websocket = connect_responses_websocket(server.address, &api_key).await;
|
||||
|
||||
let mut first_payload: Value = serde_json::from_str(&response_create_payload("first")).unwrap();
|
||||
first_payload["input"][0]["id"] = json!("client_input_id");
|
||||
first_payload["input"][0]["encrypted_content"] = json!("client_secret");
|
||||
websocket
|
||||
.send(Message::Text(first_payload.to_string().into()))
|
||||
.await
|
||||
.unwrap();
|
||||
let first = receive_response(&mut websocket).await;
|
||||
assert_eq!(first.last().unwrap()["response"]["id"], "resp_first");
|
||||
|
||||
let mut second_payload: Value =
|
||||
serde_json::from_str(&response_create_payload("second")).unwrap();
|
||||
second_payload["previous_response_id"] = json!("resp_first");
|
||||
websocket
|
||||
.send(Message::Text(second_payload.to_string().into()))
|
||||
.await
|
||||
.unwrap();
|
||||
let second = receive_response(&mut websocket).await;
|
||||
let (initial, continued, replayed) = upstream.await.unwrap();
|
||||
|
||||
assert_eq!(second.last().unwrap()["response"]["id"], "resp_recovered");
|
||||
assert_eq!(initial["input"][0]["content"], "first");
|
||||
assert_eq!(continued["previous_response_id"], "resp_first");
|
||||
assert_eq!(continued["input"][0]["content"], "second");
|
||||
assert!(replayed.get("previous_response_id").is_none());
|
||||
assert_eq!(replayed["input"][0]["content"], "first");
|
||||
assert_eq!(replayed["input"][1]["role"], "assistant");
|
||||
assert_eq!(replayed["input"][2]["content"], "second");
|
||||
let replayed_json = replayed.to_string();
|
||||
assert!(!replayed_json.contains("client_input_id"));
|
||||
assert!(!replayed_json.contains("client_secret"));
|
||||
assert!(!replayed_json.contains("assistant_output_resp_first"));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn responses_websocket_should_allow_official_retry_after_midstream_upstream_disconnect() {
|
||||
let upstream_listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
|
||||
@@ -237,6 +313,12 @@ async fn send_upstream_response(
|
||||
"id": response_id,
|
||||
"object": "response",
|
||||
"status": "completed",
|
||||
"output": [{
|
||||
"type": "message",
|
||||
"id": format!("assistant_output_{response_id}"),
|
||||
"role": "assistant",
|
||||
"content": [{"type": "output_text", "text": text}],
|
||||
}],
|
||||
"usage": {
|
||||
"input_tokens": 3,
|
||||
"output_tokens": 1,
|
||||
@@ -252,6 +334,11 @@ async fn send_upstream_response(
|
||||
}
|
||||
}
|
||||
|
||||
async fn receive_upstream_request(websocket: &mut WebSocketStream<tokio::net::TcpStream>) -> Value {
|
||||
let request = websocket.next().await.unwrap().unwrap();
|
||||
serde_json::from_str(&request.into_text().unwrap()).unwrap()
|
||||
}
|
||||
|
||||
async fn receive_response<S>(websocket: &mut WebSocketStream<S>) -> Vec<Value>
|
||||
where
|
||||
S: tokio::io::AsyncRead + tokio::io::AsyncWrite + Unpin,
|
||||
|
||||
@@ -1,10 +1,8 @@
|
||||
use chrono::{Duration, Utc};
|
||||
use codex_proxy_rs::dispatch::affinity::{
|
||||
MAX_REPLAY_DEPTH, MAX_REPLAY_SESSION_BYTES, RedisSessionAffinityStore, ResponseReplaySnapshot,
|
||||
SessionAffinityEntry,
|
||||
MAX_CONVERSATION_AFFINITIES, RedisSessionAffinityStore, SessionAffinityEntry,
|
||||
};
|
||||
use redis::AsyncCommands;
|
||||
use serde_json::json;
|
||||
|
||||
use crate::support::storage::create_test_redis;
|
||||
|
||||
@@ -18,9 +16,10 @@ async fn redis_affinity_store_upserts_reads_and_forgets_entries() {
|
||||
.upsert("resp_1", &entry, Duration::hours(4))
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
let mut connection = redis.manager();
|
||||
let ttl_seconds: i64 = connection
|
||||
.ttl(redis.key("affinity:v2:resp:resp_1"))
|
||||
.ttl(redis.key("affinity:v3:resp:resp_1"))
|
||||
.await
|
||||
.unwrap();
|
||||
assert!((3 * 60 * 60..=4 * 60 * 60).contains(&ttl_seconds));
|
||||
@@ -96,8 +95,9 @@ async fn redis_affinity_store_returns_latest_matching_conversation_variant() {
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn redis_affinity_store_forgets_all_entries_for_account() {
|
||||
let store = RedisSessionAffinityStore::new(create_test_redis("affinity-account").await);
|
||||
async fn redis_affinity_store_forgets_account_entries_in_one_operation() {
|
||||
let redis = create_test_redis("affinity-account").await;
|
||||
let store = RedisSessionAffinityStore::new(redis.clone());
|
||||
let now = Utc::now();
|
||||
for response_id in ["resp_1", "resp_2"] {
|
||||
store
|
||||
@@ -110,7 +110,14 @@ async fn redis_affinity_store_forgets_all_entries_for_account() {
|
||||
.unwrap();
|
||||
}
|
||||
|
||||
let bytes_before: u64 = redis::cmd("GET")
|
||||
.arg(redis.key("affinity:v3:global:bytes"))
|
||||
.query_async(&mut redis.manager())
|
||||
.await
|
||||
.unwrap();
|
||||
assert!(bytes_before > 0);
|
||||
assert_eq!(store.forget_account("acct_delete").await.unwrap(), 2);
|
||||
|
||||
for response_id in ["resp_1", "resp_2"] {
|
||||
assert!(
|
||||
store
|
||||
@@ -120,99 +127,44 @@ async fn redis_affinity_store_forgets_all_entries_for_account() {
|
||||
.is_none()
|
||||
);
|
||||
}
|
||||
let bytes_after: Option<u64> = redis::cmd("GET")
|
||||
.arg(redis.key("affinity:v3:global:bytes"))
|
||||
.query_async(&mut redis.manager())
|
||||
.await
|
||||
.unwrap();
|
||||
assert!(bytes_after.is_none());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn redis_affinity_store_should_keep_replay_growth_linear_across_turns() {
|
||||
let redis = create_test_redis("affinity-replay-linear").await;
|
||||
async fn redis_affinity_store_caps_active_conversation_index() {
|
||||
let redis = create_test_redis("affinity-conversation-cap").await;
|
||||
let store = RedisSessionAffinityStore::new(redis.clone());
|
||||
let now = Utc::now();
|
||||
let mut parent_response_id = None;
|
||||
let mut total_bytes = 0;
|
||||
for turn in 0..16u16 {
|
||||
let response_id = format!("resp_{turn}");
|
||||
let turn_input = vec![json!({"role": "user", "content": format!("input-{turn}")})];
|
||||
let turn_output = vec![json!({"type": "message", "content": format!("output-{turn}")})];
|
||||
let node_bytes = replay_node_bytes(&turn_input, &turn_output);
|
||||
total_bytes += node_bytes;
|
||||
let mut entry = affinity(
|
||||
"acct_linear",
|
||||
"conv_linear",
|
||||
None,
|
||||
now + Duration::milliseconds(i64::from(turn)),
|
||||
);
|
||||
entry.replay = Some(ResponseReplaySnapshot {
|
||||
parent_response_id: parent_response_id.clone(),
|
||||
turn_input,
|
||||
turn_output,
|
||||
depth: turn + 1,
|
||||
total_bytes,
|
||||
});
|
||||
for index in 0..=MAX_CONVERSATION_AFFINITIES {
|
||||
store
|
||||
.upsert(&response_id, &entry, Duration::hours(4))
|
||||
.upsert(
|
||||
&format!("resp_cap_{index}"),
|
||||
&affinity(
|
||||
"acct_cap",
|
||||
"conv_cap",
|
||||
None,
|
||||
now + Duration::milliseconds(index as i64),
|
||||
),
|
||||
Duration::hours(4),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
parent_response_id = Some(response_id);
|
||||
}
|
||||
|
||||
let head = store
|
||||
.get("resp_15", now, Duration::hours(4))
|
||||
.await
|
||||
.unwrap()
|
||||
.unwrap();
|
||||
let replay = store
|
||||
.replay_input(
|
||||
"resp_15",
|
||||
&head,
|
||||
now,
|
||||
Duration::hours(4),
|
||||
MAX_REPLAY_DEPTH,
|
||||
MAX_REPLAY_SESSION_BYTES,
|
||||
)
|
||||
.await
|
||||
.unwrap()
|
||||
.unwrap();
|
||||
assert_eq!(replay.len(), 32);
|
||||
assert_eq!(replay[0]["content"], "input-0");
|
||||
assert_eq!(replay[31]["content"], "output-15");
|
||||
|
||||
let mut connection = redis.manager();
|
||||
let head_json: String = connection
|
||||
.get(redis.key("affinity:v2:resp:resp_15"))
|
||||
let count: usize = redis::cmd("ZCARD")
|
||||
.arg(redis.key("affinity:v3:conv:conv_cap"))
|
||||
.query_async(&mut redis.manager())
|
||||
.await
|
||||
.unwrap();
|
||||
assert!(head_json.contains("input-15"));
|
||||
assert!(!head_json.contains("input-0"));
|
||||
assert!(head_json.len() < 1_024);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn redis_affinity_store_should_reject_replay_metadata_over_capacity() {
|
||||
let store = RedisSessionAffinityStore::new(create_test_redis("affinity-replay-limit").await);
|
||||
let now = Utc::now();
|
||||
let mut entry = affinity("acct_limit", "conv_limit", None, now);
|
||||
entry.replay = Some(ResponseReplaySnapshot {
|
||||
parent_response_id: None,
|
||||
turn_input: vec![json!({"role": "user", "content": "small"})],
|
||||
turn_output: Vec::new(),
|
||||
depth: 1,
|
||||
total_bytes: MAX_REPLAY_SESSION_BYTES + 1,
|
||||
});
|
||||
store
|
||||
.upsert("resp_limit", &entry, Duration::hours(4))
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(count, MAX_CONVERSATION_AFFINITIES);
|
||||
assert!(
|
||||
store
|
||||
.replay_input(
|
||||
"resp_limit",
|
||||
&entry,
|
||||
now,
|
||||
Duration::hours(4),
|
||||
MAX_REPLAY_DEPTH,
|
||||
MAX_REPLAY_SESSION_BYTES,
|
||||
)
|
||||
.get("resp_cap_0", now, Duration::hours(4))
|
||||
.await
|
||||
.unwrap()
|
||||
.is_none()
|
||||
@@ -220,11 +172,11 @@ async fn redis_affinity_store_should_reject_replay_metadata_over_capacity() {
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn redis_affinity_store_should_use_pruned_account_zset_index() {
|
||||
async fn redis_affinity_store_prunes_expired_account_index_members_lazily() {
|
||||
let redis = create_test_redis("affinity-account-zset").await;
|
||||
let store = RedisSessionAffinityStore::new(redis.clone());
|
||||
let now = Utc::now();
|
||||
let account_key = redis.key("affinity:v2:account:acct_zset");
|
||||
let account_key = redis.key("affinity:v3:account:acct_zset");
|
||||
let mut connection = redis.manager();
|
||||
let _: usize = connection
|
||||
.zadd(
|
||||
@@ -255,7 +207,7 @@ async fn redis_affinity_store_should_use_pruned_account_zset_index() {
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn redis_affinity_store_should_not_persist_already_expired_entries() {
|
||||
async fn redis_affinity_store_does_not_persist_expired_entries() {
|
||||
let store = RedisSessionAffinityStore::new(create_test_redis("affinity-expired-write").await);
|
||||
let now = Utc::now();
|
||||
store
|
||||
@@ -280,6 +232,20 @@ async fn redis_affinity_store_should_not_persist_already_expired_entries() {
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn redis_affinity_store_rejects_oversized_metadata() {
|
||||
let store = RedisSessionAffinityStore::new(create_test_redis("affinity-metadata-limit").await);
|
||||
let now = Utc::now();
|
||||
let mut entry = affinity("acct_limit", "conv_limit", None, now);
|
||||
entry.function_call_ids = vec!["x".repeat(64 * 1024)];
|
||||
|
||||
let error = store
|
||||
.upsert("resp_limit", &entry, Duration::hours(4))
|
||||
.await
|
||||
.unwrap_err();
|
||||
assert!(error.to_string().contains("metadata is too large"));
|
||||
}
|
||||
|
||||
fn affinity(
|
||||
account_id: &str,
|
||||
conversation_id: &str,
|
||||
@@ -295,11 +261,6 @@ fn affinity(
|
||||
function_call_ids: vec!["call_1".to_string()],
|
||||
variant_hash: variant_hash.map(ToString::to_string),
|
||||
continuation_scope: codex_proxy_rs::upstream::openai::protocol::responses::PreviousResponseScope::ConnectionLocal,
|
||||
replay: None,
|
||||
created_at,
|
||||
}
|
||||
}
|
||||
|
||||
fn replay_node_bytes(input: &[serde_json::Value], output: &[serde_json::Value]) -> u64 {
|
||||
serde_json::to_vec(&(input, output)).unwrap().len() as u64
|
||||
}
|
||||
|
||||
@@ -16,7 +16,7 @@ use codex_proxy_rs::{
|
||||
api::router,
|
||||
bootstrap::config::AppConfig,
|
||||
bootstrap::services::{Services, UsageRecordOptions},
|
||||
dispatch::affinity::{ResponseReplaySnapshot, SessionAffinityEntry},
|
||||
dispatch::affinity::SessionAffinityEntry,
|
||||
fleet::{account::AccountStatus, cookies::PgCookieStore},
|
||||
infra::identity::hash_credential,
|
||||
telemetry::usage::store::{PgUsageRecordStore, UsageRecordFilter},
|
||||
@@ -63,19 +63,6 @@ fn websocket_accept_config() -> WebSocketConfig {
|
||||
config
|
||||
}
|
||||
|
||||
fn replay_snapshot(turn_input: Vec<Value>, turn_output: Vec<Value>) -> ResponseReplaySnapshot {
|
||||
let total_bytes = serde_json::to_vec(&(&turn_input, &turn_output))
|
||||
.unwrap()
|
||||
.len() as u64;
|
||||
ResponseReplaySnapshot {
|
||||
parent_response_id: None,
|
||||
turn_input,
|
||||
turn_output,
|
||||
depth: 1,
|
||||
total_bytes,
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn accept_async(stream: TcpStream) -> Result<WebSocketStream<TcpStream>, WsError> {
|
||||
accept_hdr_async(stream, |_request, _response| {}).await
|
||||
}
|
||||
@@ -287,32 +274,6 @@ pub(crate) async fn test_app_with_account(
|
||||
(router::router().with_state(state), api_key, dir)
|
||||
}
|
||||
|
||||
async fn test_app_with_account_and_state(
|
||||
base_url: String,
|
||||
) -> (
|
||||
axum::Router,
|
||||
AppState,
|
||||
String,
|
||||
crate::support::storage::TestDatabaseGuard,
|
||||
) {
|
||||
let (pool, dir) = init_test_db("openai-chat-upstream-state").await;
|
||||
let api_key = insert_client_api_key(&pool).await;
|
||||
insert_account(&pool).await;
|
||||
let state = test_app_state_with_pool(&test_config(test_database_url(), base_url), pool).await;
|
||||
state
|
||||
.services
|
||||
.account_pool
|
||||
.restore_from_store()
|
||||
.await
|
||||
.expect("account pool should restore");
|
||||
(
|
||||
router::router().with_state(state.clone()),
|
||||
state,
|
||||
api_key,
|
||||
dir,
|
||||
)
|
||||
}
|
||||
|
||||
async fn test_app_with_account_and_pool(
|
||||
base_url: String,
|
||||
) -> (
|
||||
@@ -492,13 +453,6 @@ async fn test_app_with_two_accounts_and_affinity(
|
||||
function_call_ids: Vec::new(),
|
||||
variant_hash: None,
|
||||
continuation_scope: codex_proxy_rs::upstream::openai::protocol::responses::PreviousResponseScope::Persisted,
|
||||
replay: Some(replay_snapshot(
|
||||
vec![json!({
|
||||
"role": "user",
|
||||
"content": "Original managed history"
|
||||
})],
|
||||
Vec::new(),
|
||||
)),
|
||||
created_at: now,
|
||||
},
|
||||
)
|
||||
@@ -513,79 +467,6 @@ async fn test_app_with_two_accounts_and_affinity(
|
||||
(router::router().with_state(state), api_key, dir)
|
||||
}
|
||||
|
||||
async fn test_app_with_two_accounts_and_affinity_status(
|
||||
base_url: String,
|
||||
affinity_account_status: &str,
|
||||
) -> (
|
||||
axum::Router,
|
||||
AppState,
|
||||
String,
|
||||
PgPool,
|
||||
crate::support::storage::TestDatabaseGuard,
|
||||
) {
|
||||
let (pool, dir) = init_test_db(&format!(
|
||||
"openai-responses-affinity-{affinity_account_status}"
|
||||
))
|
||||
.await;
|
||||
let api_key = insert_client_api_key(&pool).await;
|
||||
insert_named_account(&pool, "acct_primary", "access-primary", "chatgpt-primary").await;
|
||||
insert_named_account(
|
||||
&pool,
|
||||
"acct_affinity",
|
||||
"access-affinity",
|
||||
"chatgpt-affinity",
|
||||
)
|
||||
.await;
|
||||
sqlx::query("update accounts set status = $1 where id = $2")
|
||||
.bind(affinity_account_status)
|
||||
.bind("acct_affinity")
|
||||
.execute(&pool)
|
||||
.await
|
||||
.unwrap();
|
||||
let now = Utc::now();
|
||||
let state =
|
||||
test_app_state_with_pool(&test_config(test_database_url(), base_url), pool.clone()).await;
|
||||
state
|
||||
.services
|
||||
.session_affinity
|
||||
.record(
|
||||
"resp_affinity_risk".to_string(),
|
||||
SessionAffinityEntry {
|
||||
account_id: "acct_affinity".to_string(),
|
||||
conversation_id: "conv_affinity_risk".to_string(),
|
||||
turn_state: Some("turn_affinity_risk".to_string()),
|
||||
instructions_hash: None,
|
||||
input_tokens: None,
|
||||
function_call_ids: Vec::new(),
|
||||
variant_hash: None,
|
||||
continuation_scope: codex_proxy_rs::upstream::openai::protocol::responses::PreviousResponseScope::Persisted,
|
||||
replay: Some(replay_snapshot(
|
||||
vec![json!({
|
||||
"role": "user",
|
||||
"content": "Original affinity history"
|
||||
})],
|
||||
Vec::new(),
|
||||
)),
|
||||
created_at: now,
|
||||
},
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
state
|
||||
.services
|
||||
.account_pool
|
||||
.restore_from_store()
|
||||
.await
|
||||
.expect("account pool should restore");
|
||||
(
|
||||
router::router().with_state(state.clone()),
|
||||
state,
|
||||
api_key,
|
||||
pool,
|
||||
dir,
|
||||
)
|
||||
}
|
||||
|
||||
async fn test_app_with_two_accounts(
|
||||
base_url: String,
|
||||
) -> (
|
||||
@@ -807,7 +688,6 @@ async fn insert_session_affinity(
|
||||
function_call_ids: Vec::new(),
|
||||
variant_hash: None,
|
||||
continuation_scope: codex_proxy_rs::upstream::openai::protocol::responses::PreviousResponseScope::Persisted,
|
||||
replay: None,
|
||||
created_at: Utc::now(),
|
||||
},
|
||||
)
|
||||
|
||||
@@ -1226,109 +1226,3 @@ async fn responses_should_prefer_session_affinity_account_for_previous_response(
|
||||
);
|
||||
assert_eq!(captured.payload["previous_response_id"], "resp_previous");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn responses_should_replay_full_history_when_banned_affinity_switches_account() {
|
||||
let (base_url, upstream) =
|
||||
spawn_single_websocket_completed_upstream("resp_after_banned_affinity").await;
|
||||
let (app, state, api_key, _pool, _dir) =
|
||||
test_app_with_two_accounts_and_affinity_status(base_url, "banned").await;
|
||||
let response = app
|
||||
.oneshot(
|
||||
Request::builder()
|
||||
.method("POST")
|
||||
.uri("/v1/responses")
|
||||
.header("authorization", format!("Bearer {api_key}"))
|
||||
.header("content-type", "application/json")
|
||||
.body(Body::from(
|
||||
json!({
|
||||
"model": "gpt-5.5",
|
||||
"previous_response_id": "resp_affinity_risk",
|
||||
"input": [{"role": "user", "content": "Continue after ban"}],
|
||||
"stream": false
|
||||
})
|
||||
.to_string(),
|
||||
))
|
||||
.unwrap(),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
let captured = upstream.await.unwrap();
|
||||
let affinity = state
|
||||
.services
|
||||
.session_affinity
|
||||
.lookup("resp_affinity_risk", Utc::now())
|
||||
.await;
|
||||
|
||||
assert_eq!(response.status(), StatusCode::OK);
|
||||
assert_eq!(
|
||||
captured_header(&captured.headers, "authorization"),
|
||||
Some("Bearer access-primary")
|
||||
);
|
||||
assert_eq!(
|
||||
captured_header(&captured.headers, "chatgpt-account-id"),
|
||||
Some("chatgpt-primary")
|
||||
);
|
||||
assert!(captured_header(&captured.headers, "x-codex-turn-state").is_none());
|
||||
assert!(captured.payload.get("previous_response_id").is_none());
|
||||
assert_eq!(
|
||||
captured.payload["input"][0]["content"],
|
||||
"Original affinity history"
|
||||
);
|
||||
assert_eq!(
|
||||
captured.payload["input"][1]["content"],
|
||||
"Continue after ban"
|
||||
);
|
||||
assert!(affinity.is_some());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn responses_should_replay_full_history_when_quota_affinity_switches_account() {
|
||||
let (base_url, upstream) =
|
||||
spawn_single_websocket_completed_upstream("resp_after_quota_affinity").await;
|
||||
let (app, state, api_key, _pool, _dir) =
|
||||
test_app_with_two_accounts_and_affinity_status(base_url, "quota_exhausted").await;
|
||||
let response = app
|
||||
.oneshot(
|
||||
Request::builder()
|
||||
.method("POST")
|
||||
.uri("/v1/responses")
|
||||
.header("authorization", format!("Bearer {api_key}"))
|
||||
.header("content-type", "application/json")
|
||||
.body(Body::from(
|
||||
json!({
|
||||
"model": "gpt-5.5",
|
||||
"previous_response_id": "resp_affinity_risk",
|
||||
"input": [{"role": "user", "content": "Continue after quota"}],
|
||||
"stream": false
|
||||
})
|
||||
.to_string(),
|
||||
))
|
||||
.unwrap(),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
let captured = upstream.await.unwrap();
|
||||
let affinity = state
|
||||
.services
|
||||
.session_affinity
|
||||
.lookup("resp_affinity_risk", Utc::now())
|
||||
.await;
|
||||
|
||||
assert_eq!(response.status(), StatusCode::OK);
|
||||
assert_eq!(
|
||||
captured_header(&captured.headers, "authorization"),
|
||||
Some("Bearer access-primary")
|
||||
);
|
||||
assert!(captured_header(&captured.headers, "x-codex-turn-state").is_none());
|
||||
assert!(captured.payload.get("previous_response_id").is_none());
|
||||
assert_eq!(
|
||||
captured.payload["input"][0]["content"],
|
||||
"Original affinity history"
|
||||
);
|
||||
assert_eq!(
|
||||
captured.payload["input"][1]["content"],
|
||||
"Continue after quota"
|
||||
);
|
||||
assert!(affinity.is_some());
|
||||
}
|
||||
|
||||
@@ -1194,7 +1194,7 @@ async fn responses_external_previous_response_not_found_should_return_exact_upst
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn responses_managed_history_should_rebuild_full_replay_for_next_configured_candidate() {
|
||||
async fn responses_managed_history_without_local_transcript_should_not_fan_out() {
|
||||
let server = MockServer::start().await;
|
||||
Mock::given(method("POST"))
|
||||
.and(path("/codex/responses"))
|
||||
@@ -1205,17 +1205,6 @@ async fn responses_managed_history_should_rebuild_full_replay_for_next_configure
|
||||
.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_SUCCESS_SSE),
|
||||
)
|
||||
.expect(1)
|
||||
.mount(&server)
|
||||
.await;
|
||||
let (app, state, api_key, _pool, _dir) =
|
||||
test_app_with_two_accounts_and_state(server.uri()).await;
|
||||
state
|
||||
@@ -1232,18 +1221,6 @@ async fn responses_managed_history_should_rebuild_full_replay_for_next_configure
|
||||
function_call_ids: Vec::new(),
|
||||
variant_hash: None,
|
||||
continuation_scope: codex_proxy_rs::upstream::openai::protocol::responses::PreviousResponseScope::Persisted,
|
||||
replay: Some(replay_snapshot(
|
||||
vec![
|
||||
json!({"role": "user", "content": "Earlier question"}),
|
||||
json!({
|
||||
"type": "reasoning",
|
||||
"id": "reasoning-old",
|
||||
"encrypted_content": "account-bound-secret",
|
||||
"summary": []
|
||||
}),
|
||||
],
|
||||
Vec::new(),
|
||||
)),
|
||||
created_at: Utc::now(),
|
||||
},
|
||||
)
|
||||
@@ -1272,26 +1249,19 @@ async fn responses_managed_history_should_rebuild_full_replay_for_next_configure
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(response.status(), StatusCode::OK);
|
||||
assert_eq!(response.status(), StatusCode::UNAUTHORIZED);
|
||||
|
||||
let requests = server.received_requests().await.unwrap();
|
||||
let secondary = requests
|
||||
.iter()
|
||||
.find(|request| {
|
||||
request
|
||||
.headers
|
||||
.get("authorization")
|
||||
.and_then(|value| value.to_str().ok())
|
||||
== Some("Bearer access-secondary")
|
||||
})
|
||||
.expect("secondary request");
|
||||
let body: Value = serde_json::from_slice(&secondary.body).unwrap();
|
||||
assert!(body.get("previous_response_id").is_none());
|
||||
assert_eq!(body["input"][0]["content"], "Earlier question");
|
||||
assert!(body["input"][1].get("id").is_none());
|
||||
assert!(body["input"][1].get("encrypted_content").is_none());
|
||||
assert_eq!(body["input"][2]["content"], "Current question");
|
||||
assert!(secondary.headers.get("x-codex-turn-state").is_none());
|
||||
assert_eq!(requests.len(), 1);
|
||||
assert_eq!(
|
||||
requests[0]
|
||||
.headers
|
||||
.get("authorization")
|
||||
.and_then(|value| value.to_str().ok()),
|
||||
Some("Bearer access-primary")
|
||||
);
|
||||
let body: Value = serde_json::from_slice(&requests[0].body).unwrap();
|
||||
assert_eq!(body["previous_response_id"], "resp_managed");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
@@ -1956,7 +1926,7 @@ async fn responses_stream_should_return_cloudflare_path_block_error_when_404_fal
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn responses_should_retry_same_account_after_5xx_before_fallback() {
|
||||
async fn responses_should_rotate_account_after_5xx_without_same_account_retry() {
|
||||
let server = MockServer::start().await;
|
||||
Mock::given(method("POST"))
|
||||
.and(path("/codex/responses"))
|
||||
@@ -1964,12 +1934,12 @@ async fn responses_should_retry_same_account_after_5xx_before_fallback() {
|
||||
.respond_with(ResponseTemplate::new(500).set_body_json(json!({
|
||||
"error": {"message": "temporary upstream failure"}
|
||||
})))
|
||||
.up_to_n_times(2)
|
||||
.expect(1)
|
||||
.mount(&server)
|
||||
.await;
|
||||
Mock::given(method("POST"))
|
||||
.and(path("/codex/responses"))
|
||||
.and(header("authorization", "Bearer access-primary"))
|
||||
.and(header("authorization", "Bearer access-secondary"))
|
||||
.respond_with(
|
||||
ResponseTemplate::new(200)
|
||||
.insert_header("content-type", "text/event-stream")
|
||||
@@ -2014,8 +1984,8 @@ async fn responses_should_retry_same_account_after_5xx_before_fallback() {
|
||||
|
||||
assert_eq!(status, StatusCode::OK, "requests: {authorizations:?}");
|
||||
assert_eq!(body["id"], "resp_after_5xx_retry");
|
||||
assert_eq!(primary_requests, 3, "requests: {authorizations:?}");
|
||||
assert_eq!(secondary_requests, 0, "requests: {authorizations:?}");
|
||||
assert_eq!(primary_requests, 1, "requests: {authorizations:?}");
|
||||
assert_eq!(secondary_requests, 1, "requests: {authorizations:?}");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
@@ -2075,7 +2045,7 @@ async fn responses_should_traverse_all_configured_strategy_candidates_without_fi
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn responses_stream_should_retry_same_account_after_5xx_before_fallback() {
|
||||
async fn responses_stream_should_rotate_account_after_5xx_without_same_account_retry() {
|
||||
let server = MockServer::start().await;
|
||||
Mock::given(method("POST"))
|
||||
.and(path("/codex/responses"))
|
||||
@@ -2083,12 +2053,12 @@ async fn responses_stream_should_retry_same_account_after_5xx_before_fallback()
|
||||
.respond_with(ResponseTemplate::new(500).set_body_json(json!({
|
||||
"error": {"message": "temporary upstream failure"}
|
||||
})))
|
||||
.up_to_n_times(2)
|
||||
.expect(1)
|
||||
.mount(&server)
|
||||
.await;
|
||||
Mock::given(method("POST"))
|
||||
.and(path("/codex/responses"))
|
||||
.and(header("authorization", "Bearer access-primary"))
|
||||
.and(header("authorization", "Bearer access-secondary"))
|
||||
.respond_with(
|
||||
ResponseTemplate::new(200)
|
||||
.insert_header("content-type", "text/event-stream")
|
||||
@@ -2140,8 +2110,8 @@ async fn responses_stream_should_retry_same_account_after_5xx_before_fallback()
|
||||
assert_eq!(status, StatusCode::OK, "requests: {authorizations:?}");
|
||||
assert!(content_type.starts_with("text/event-stream"));
|
||||
assert!(body.contains("resp_stream_after_5xx_retry"));
|
||||
assert_eq!(primary_requests, 3, "requests: {authorizations:?}");
|
||||
assert_eq!(secondary_requests, 0, "requests: {authorizations:?}");
|
||||
assert_eq!(primary_requests, 1, "requests: {authorizations:?}");
|
||||
assert_eq!(secondary_requests, 1, "requests: {authorizations:?}");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
|
||||
@@ -354,305 +354,6 @@ async fn responses_websocket_should_reuse_connection_for_recorded_conversation()
|
||||
assert_eq!(second_payload["previous_response_id"], "resp_pool_first");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn responses_websocket_should_replay_history_when_reused_connection_dies_before_first_output()
|
||||
{
|
||||
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
|
||||
let base_url = format!("http://{}", listener.local_addr().unwrap());
|
||||
let upstream = tokio::spawn(async move {
|
||||
let (stream, _) = listener.accept().await.unwrap();
|
||||
let mut websocket = accept_async(stream).await.unwrap();
|
||||
let first_message = websocket.next().await.unwrap().unwrap();
|
||||
let first_payload = serde_json::from_str::<Value>(&first_message.into_text().unwrap())
|
||||
.expect("first websocket payload should be json");
|
||||
websocket
|
||||
.send(Message::Text(
|
||||
response_completed_websocket_message("resp_stale_reuse_first").into(),
|
||||
))
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
let stale_message = websocket.next().await.unwrap().unwrap();
|
||||
let stale_payload = serde_json::from_str::<Value>(&stale_message.into_text().unwrap())
|
||||
.expect("stale reused websocket payload should be json");
|
||||
websocket.close(None).await.unwrap();
|
||||
|
||||
let fresh_payload =
|
||||
accept_successful_websocket_response(&listener, "resp_after_stale_reuse").await;
|
||||
|
||||
(first_payload, stale_payload, fresh_payload)
|
||||
});
|
||||
let (app, api_key, _dir) = test_app_with_account(base_url).await;
|
||||
|
||||
let first_response = app
|
||||
.clone()
|
||||
.oneshot(responses_json_request(
|
||||
&api_key,
|
||||
&json!({
|
||||
"model": "gpt-5.5",
|
||||
"input": [{
|
||||
"role": "user",
|
||||
"content": "prime stale websocket reuse"
|
||||
}]
|
||||
}),
|
||||
))
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(first_response.status(), StatusCode::OK);
|
||||
assert!(
|
||||
response_text(first_response)
|
||||
.await
|
||||
.contains("\"id\":\"resp_stale_reuse_first\"")
|
||||
);
|
||||
|
||||
let second_response = app
|
||||
.oneshot(responses_json_request(
|
||||
&api_key,
|
||||
&json!({
|
||||
"model": "gpt-5.5",
|
||||
"input": [],
|
||||
"previous_response_id": "resp_stale_reuse_first"
|
||||
}),
|
||||
))
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(second_response.status(), StatusCode::OK);
|
||||
let second_body = response_text(second_response).await;
|
||||
assert!(
|
||||
second_body.contains("\"id\":\"resp_after_stale_reuse\""),
|
||||
"{second_body}"
|
||||
);
|
||||
let (first_payload, stale_payload, fresh_payload) = upstream.await.unwrap();
|
||||
|
||||
assert_eq!(
|
||||
stale_payload["previous_response_id"],
|
||||
"resp_stale_reuse_first"
|
||||
);
|
||||
assert!(fresh_payload.get("previous_response_id").is_none());
|
||||
assert_eq!(
|
||||
fresh_payload["input"][0]["content"],
|
||||
"prime stale websocket reuse"
|
||||
);
|
||||
assert!(first_payload.get("prompt_cache_key").is_none());
|
||||
assert!(fresh_payload.get("prompt_cache_key").is_none());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn responses_websocket_stream_should_replay_managed_history_after_previous_response_not_found()
|
||||
{
|
||||
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
|
||||
let base_url = format!("http://{}", listener.local_addr().unwrap());
|
||||
let upstream = tokio::spawn(async move {
|
||||
let (first_stream, _) = listener.accept().await.unwrap();
|
||||
let mut first_websocket = accept_async(first_stream).await.unwrap();
|
||||
let first_message = first_websocket.next().await.unwrap().unwrap();
|
||||
let first_payload = serde_json::from_str::<Value>(&first_message.into_text().unwrap())
|
||||
.expect("first websocket payload should be json");
|
||||
first_websocket
|
||||
.send(Message::Text(
|
||||
response_failed_websocket_message(
|
||||
"resp_previous_missing",
|
||||
"previous_response_not_found",
|
||||
"Previous response with id resp_managed_stale was not found",
|
||||
)
|
||||
.into(),
|
||||
))
|
||||
.await
|
||||
.unwrap();
|
||||
first_websocket.close(None).await.unwrap();
|
||||
|
||||
let replay_payload =
|
||||
accept_successful_websocket_response(&listener, "resp_replayed_success").await;
|
||||
(first_payload, replay_payload)
|
||||
});
|
||||
let (app, state, api_key, _dir) = test_app_with_account_and_state(base_url).await;
|
||||
state
|
||||
.services
|
||||
.session_affinity
|
||||
.record(
|
||||
"resp_managed_stale".to_string(),
|
||||
SessionAffinityEntry {
|
||||
account_id: "acct_chat".to_string(),
|
||||
conversation_id: "conversation-managed-stale".to_string(),
|
||||
turn_state: Some("account-bound-turn-state".to_string()),
|
||||
instructions_hash: None,
|
||||
input_tokens: Some(7),
|
||||
function_call_ids: Vec::new(),
|
||||
variant_hash: None,
|
||||
continuation_scope: codex_proxy_rs::upstream::openai::protocol::responses::PreviousResponseScope::Persisted,
|
||||
replay: Some(replay_snapshot(
|
||||
vec![json!({
|
||||
"role": "user",
|
||||
"content": "Earlier managed question"
|
||||
})],
|
||||
Vec::new(),
|
||||
)),
|
||||
created_at: Utc::now(),
|
||||
},
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
let response = app
|
||||
.oneshot(responses_json_request(
|
||||
&api_key,
|
||||
&json!({
|
||||
"model": "gpt-5.5",
|
||||
"previous_response_id": "resp_managed_stale",
|
||||
"input": [{"role": "user", "content": "Current managed question"}],
|
||||
"stream": true
|
||||
}),
|
||||
))
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(response.status(), StatusCode::OK);
|
||||
let body = response_text(response).await;
|
||||
let (first_payload, replay_payload) = upstream.await.unwrap();
|
||||
|
||||
assert!(body.contains("resp_replayed_success"), "{body}");
|
||||
assert!(!body.contains("resp_previous_missing"), "{body}");
|
||||
assert_eq!(first_payload["previous_response_id"], "resp_managed_stale");
|
||||
assert!(replay_payload.get("previous_response_id").is_none());
|
||||
assert_eq!(
|
||||
replay_payload["input"][0]["content"],
|
||||
"Earlier managed question"
|
||||
);
|
||||
assert_eq!(
|
||||
replay_payload["input"][1]["content"],
|
||||
"Current managed question"
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn responses_websocket_should_store_sanitized_incremental_replay_snapshot() {
|
||||
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
|
||||
let base_url = format!("http://{}", listener.local_addr().unwrap());
|
||||
let upstream = tokio::spawn(async move {
|
||||
let (stream, _) = listener.accept().await.unwrap();
|
||||
let mut websocket = accept_async(stream).await.unwrap();
|
||||
let _payload = websocket.next().await.unwrap().unwrap();
|
||||
let mut completed: Value = serde_json::from_str(include_str!(
|
||||
"../../fixtures/responses/websocket/completed_with_reasoning_replay.json"
|
||||
))
|
||||
.unwrap();
|
||||
completed["response"]["status"] = Value::String("completed".to_string());
|
||||
completed["response"]["output"]
|
||||
.as_array_mut()
|
||||
.unwrap()
|
||||
.push(json!({
|
||||
"type": "message",
|
||||
"id": "msg_sensitive_output",
|
||||
"role": "assistant",
|
||||
"content": [{"type": "output_text", "text": "done"}]
|
||||
}));
|
||||
websocket
|
||||
.send(Message::Text(completed.to_string().into()))
|
||||
.await
|
||||
.unwrap();
|
||||
websocket.close(None).await.unwrap();
|
||||
});
|
||||
let (app, state, api_key, _dir) = test_app_with_account_and_state(base_url).await;
|
||||
|
||||
let response = app
|
||||
.oneshot(responses_json_request(
|
||||
&api_key,
|
||||
&json!({
|
||||
"model": "gpt-5.5",
|
||||
"stream": false,
|
||||
"input": [
|
||||
{
|
||||
"type": "message",
|
||||
"id": "account-bound-input-id",
|
||||
"role": "user",
|
||||
"content": [{
|
||||
"type": "input_text",
|
||||
"text": "remember this",
|
||||
"encrypted_content": "nested-secret"
|
||||
}]
|
||||
},
|
||||
{
|
||||
"type": "function_call_output",
|
||||
"call_id": "call_must_survive",
|
||||
"output": "ok"
|
||||
}
|
||||
]
|
||||
}),
|
||||
))
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(response.status(), StatusCode::OK);
|
||||
upstream.await.unwrap();
|
||||
|
||||
let stored = state
|
||||
.services
|
||||
.session_affinity
|
||||
.lookup("resp_implicit_resume_first", Utc::now())
|
||||
.await
|
||||
.unwrap();
|
||||
let replay = stored.replay.unwrap();
|
||||
assert_eq!(replay.depth, 1);
|
||||
assert!(replay.parent_response_id.is_none());
|
||||
assert!(replay.turn_input[0].get("id").is_none());
|
||||
assert_eq!(replay.turn_input[1]["call_id"], "call_must_survive");
|
||||
assert!(replay.turn_output[0].get("id").is_none());
|
||||
let serialized = serde_json::to_string(&replay).unwrap();
|
||||
assert!(!serialized.contains("encrypted_content"));
|
||||
assert!(!serialized.contains("nested-secret"));
|
||||
assert!(!serialized.contains("enc_reasoning_replay"));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn responses_should_skip_oversized_replay_snapshot_without_failing_response() {
|
||||
let server = MockServer::start().await;
|
||||
let oversized = "x".repeat(2 * 1024 * 1024 + 1024);
|
||||
let completed = json!({
|
||||
"response": {
|
||||
"id": "resp_oversized_replay",
|
||||
"object": "response",
|
||||
"status": "completed",
|
||||
"output": [{
|
||||
"type": "message",
|
||||
"id": "msg_oversized",
|
||||
"role": "assistant",
|
||||
"content": [{"type": "output_text", "text": oversized}]
|
||||
}],
|
||||
"usage": {"input_tokens": 1, "output_tokens": 1, "total_tokens": 2}
|
||||
}
|
||||
});
|
||||
Mock::given(method("POST"))
|
||||
.and(path("/codex/responses"))
|
||||
.respond_with(
|
||||
ResponseTemplate::new(200)
|
||||
.insert_header("content-type", "text/event-stream")
|
||||
.set_body_string(format!("event: response.completed\ndata: {completed}\n\n")),
|
||||
)
|
||||
.expect(1)
|
||||
.mount(&server)
|
||||
.await;
|
||||
let (app, state, api_key, _dir) = test_app_with_account_and_state(server.uri()).await;
|
||||
|
||||
let response = app
|
||||
.oneshot(responses_json_request(
|
||||
&api_key,
|
||||
&json!({
|
||||
"model": "gpt-5.5",
|
||||
"input": [{"role": "user", "content": "small input"}],
|
||||
"stream": false,
|
||||
"use_websocket": false
|
||||
}),
|
||||
))
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(response.status(), StatusCode::OK);
|
||||
let stored = state
|
||||
.services
|
||||
.session_affinity
|
||||
.lookup("resp_oversized_replay", Utc::now())
|
||||
.await
|
||||
.unwrap();
|
||||
assert!(stored.replay.is_none());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn responses_websocket_should_not_reuse_connection_when_pool_is_disabled() {
|
||||
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
|
||||
@@ -767,8 +468,10 @@ async fn responses_websocket_should_not_reuse_connection_when_pool_is_disabled()
|
||||
.body(Body::from(
|
||||
json!({
|
||||
"model": "gpt-5.5",
|
||||
"input": [],
|
||||
"previous_response_id": "resp_disabled_pool_first"
|
||||
"input": [{
|
||||
"role": "user",
|
||||
"content": "open a second independent websocket"
|
||||
}]
|
||||
})
|
||||
.to_string(),
|
||||
))
|
||||
@@ -793,124 +496,10 @@ async fn responses_websocket_should_not_reuse_connection_when_pool_is_disabled()
|
||||
assert!(second_payload.get("previous_response_id").is_none());
|
||||
assert_eq!(
|
||||
second_payload["input"][0]["content"],
|
||||
"do not reuse this upstream websocket"
|
||||
"open a second independent websocket"
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn responses_websocket_stream_should_strip_turn_state_when_replaying_on_new_connection() {
|
||||
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
|
||||
let base_url = format!("http://{}", listener.local_addr().unwrap());
|
||||
let upstream = tokio::spawn(async move {
|
||||
let (first_stream, _) = listener.accept().await.unwrap();
|
||||
let mut first_websocket = accept_async(first_stream).await.unwrap();
|
||||
let first_message = first_websocket.next().await.unwrap().unwrap();
|
||||
let first_payload = serde_json::from_str::<Value>(&first_message.into_text().unwrap())
|
||||
.expect("first websocket payload should be json");
|
||||
first_websocket
|
||||
.send(Message::Text(
|
||||
json!({
|
||||
"type": "response.metadata",
|
||||
"headers": {
|
||||
"x-codex-turn-state": "turn-from-metadata"
|
||||
}
|
||||
})
|
||||
.to_string()
|
||||
.into(),
|
||||
))
|
||||
.await
|
||||
.unwrap();
|
||||
first_websocket
|
||||
.send(Message::Text(
|
||||
response_completed_websocket_message("resp_metadata_turn_state").into(),
|
||||
))
|
||||
.await
|
||||
.unwrap();
|
||||
first_websocket.close(None).await.unwrap();
|
||||
|
||||
let (second_stream, _) = listener.accept().await.unwrap();
|
||||
let second_headers = Arc::new(Mutex::new(Vec::new()));
|
||||
let second_headers_for_callback = Arc::clone(&second_headers);
|
||||
let mut second_websocket = accept_hdr_async(second_stream, move |request, _response| {
|
||||
*second_headers_for_callback.lock().unwrap() = request_headers(request);
|
||||
})
|
||||
.await
|
||||
.unwrap();
|
||||
let second_message = second_websocket.next().await.unwrap().unwrap();
|
||||
let second_payload = serde_json::from_str::<Value>(&second_message.into_text().unwrap())
|
||||
.expect("second websocket payload should be json");
|
||||
second_websocket
|
||||
.send(Message::Text(
|
||||
response_completed_websocket_message("resp_metadata_turn_state_next").into(),
|
||||
))
|
||||
.await
|
||||
.unwrap();
|
||||
second_websocket.close(None).await.unwrap();
|
||||
let second_headers = second_headers.lock().unwrap().clone();
|
||||
(first_payload, second_payload, second_headers)
|
||||
});
|
||||
let (app, api_key, _pool, _dir) = test_app_with_account_pool_config(base_url, |config| {
|
||||
config.ws_pool.enabled = false;
|
||||
})
|
||||
.await;
|
||||
|
||||
let first_response = app
|
||||
.clone()
|
||||
.oneshot(
|
||||
Request::builder()
|
||||
.method("POST")
|
||||
.uri("/v1/responses")
|
||||
.header("authorization", format!("Bearer {api_key}"))
|
||||
.header("content-type", "application/json")
|
||||
.body(Body::from(
|
||||
json!({
|
||||
"model": "gpt-5.5",
|
||||
"input": [],
|
||||
"stream": true,
|
||||
"use_websocket": true
|
||||
})
|
||||
.to_string(),
|
||||
))
|
||||
.unwrap(),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(first_response.status(), StatusCode::OK);
|
||||
let first_body = response_text(first_response).await;
|
||||
assert!(first_body.contains("\"id\":\"resp_metadata_turn_state\""));
|
||||
assert!(!first_body.contains("response.metadata"));
|
||||
let second_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": [],
|
||||
"previous_response_id": "resp_metadata_turn_state"
|
||||
})
|
||||
.to_string(),
|
||||
))
|
||||
.unwrap(),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(second_response.status(), StatusCode::OK);
|
||||
assert!(
|
||||
response_text(second_response)
|
||||
.await
|
||||
.contains("\"id\":\"resp_metadata_turn_state_next\"")
|
||||
);
|
||||
let (first_payload, second_payload, second_headers) = upstream.await.unwrap();
|
||||
|
||||
assert!(first_payload.get("previous_response_id").is_none());
|
||||
assert!(second_payload.get("previous_response_id").is_none());
|
||||
assert!(captured_header(&second_headers, "x-codex-turn-state").is_none());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn responses_websocket_pool_should_be_evicted_after_admin_account_status_cycle() {
|
||||
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
|
||||
@@ -991,8 +580,10 @@ async fn responses_websocket_pool_should_be_evicted_after_admin_account_status_c
|
||||
&api_key,
|
||||
&json!({
|
||||
"model": "gpt-5.5",
|
||||
"input": [],
|
||||
"previous_response_id": "resp_pool_status_first"
|
||||
"input": [{
|
||||
"role": "user",
|
||||
"content": "start again after the account status cycle"
|
||||
}]
|
||||
}),
|
||||
))
|
||||
.await
|
||||
|
||||
@@ -73,7 +73,7 @@ async fn responses_should_succeed_when_redis_affinity_write_fails() {
|
||||
.unwrap();
|
||||
let mut connection = redis.manager();
|
||||
let _: () = redis::cmd("SET")
|
||||
.arg(redis.key("affinity:v2:account:acct_chat"))
|
||||
.arg(redis.key("affinity:v3:account:acct_chat"))
|
||||
.arg("wrong-type")
|
||||
.query_async(&mut connection)
|
||||
.await
|
||||
@@ -908,7 +908,7 @@ async fn responses_stream_should_not_record_first_token_when_metadata_only_prefi
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn responses_should_record_request_count_when_5xx_retries_are_exhausted() {
|
||||
async fn responses_should_record_one_request_when_single_account_returns_5xx() {
|
||||
let server = MockServer::start().await;
|
||||
Mock::given(method("POST"))
|
||||
.and(path("/codex/responses"))
|
||||
@@ -950,7 +950,7 @@ async fn responses_should_record_request_count_when_5xx_retries_are_exhausted()
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(status, StatusCode::SERVICE_UNAVAILABLE);
|
||||
assert_eq!(authorizations.len(), 3, "requests: {authorizations:?}");
|
||||
assert_eq!(authorizations.len(), 1, "requests: {authorizations:?}");
|
||||
assert_eq!(failed_usage.map(|row| row.0).unwrap_or_default(), 1);
|
||||
let event = latest_response_upstream_ops_error_log(&pool).await;
|
||||
let metadata: Value = serde_json::from_str(&event.metadata_json).unwrap();
|
||||
@@ -976,7 +976,7 @@ async fn responses_should_record_attempt_trace_after_retrying_another_account()
|
||||
.respond_with(ResponseTemplate::new(500).set_body_json(json!({
|
||||
"error": {"message": "secondary failed"}
|
||||
})))
|
||||
.expect(3)
|
||||
.expect(1)
|
||||
.mount(&server)
|
||||
.await;
|
||||
|
||||
|
||||
@@ -212,6 +212,49 @@ async fn codex_backend_client_should_parse_retry_after_from_rate_limit_error_bod
|
||||
assert_eq!(retry_after_seconds, Some(12));
|
||||
}
|
||||
|
||||
#[tokio::test(start_paused = true)]
|
||||
async fn codex_backend_http_sse_should_fail_after_five_minutes_without_stream_data() {
|
||||
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
|
||||
let base_url = format!("http://{}", listener.local_addr().unwrap());
|
||||
let server = tokio::spawn(async move {
|
||||
let (mut stream, _) = listener.accept().await.unwrap();
|
||||
let _request = read_http_request(&mut stream).await;
|
||||
stream
|
||||
.write_all(
|
||||
b"HTTP/1.1 200 OK\r\ncontent-type: text/event-stream\r\ntransfer-encoding: chunked\r\n\r\n",
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
std::future::pending::<()>().await;
|
||||
});
|
||||
let client = CodexBackendClient::new(
|
||||
reqwest::Client::builder().no_proxy().build().unwrap(),
|
||||
base_url,
|
||||
crate::support::fingerprint::runtime_test_fingerprint(),
|
||||
);
|
||||
let mut request =
|
||||
codex_proxy_rs::upstream::openai::protocol::responses::CodexResponsesRequest::new_http_sse(
|
||||
"gpt-5.5",
|
||||
"",
|
||||
Vec::new(),
|
||||
);
|
||||
request.force_http_sse = true;
|
||||
|
||||
let mut response = client
|
||||
.create_response_stream(&request, request_context("req_http_idle", Some("acct")))
|
||||
.await
|
||||
.unwrap();
|
||||
let next = tokio::spawn(async move { response.body.next().await });
|
||||
tokio::task::yield_now().await;
|
||||
tokio::time::advance(Duration::from_secs(5 * 60)).await;
|
||||
|
||||
let Some(Err(CodexClientError::StreamIdleTimeout { timeout })) = next.await.unwrap() else {
|
||||
panic!("expected HTTP/SSE idle timeout");
|
||||
};
|
||||
assert_eq!(timeout, Duration::from_secs(5 * 60));
|
||||
server.abort();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn codex_backend_client_should_capture_only_safe_response_metadata() {
|
||||
let server = wiremock::MockServer::start().await;
|
||||
|
||||
+18
-19
@@ -244,11 +244,11 @@ dispatch/
|
||||
|
||||
- `service.rs`:非流式 Responses 主循环。
|
||||
- `stream/lifecycle.rs`:流式 Responses 主循环。
|
||||
- `upstream_call.rs`:把账号凭据、账号作用域身份和 Cookie 交给上游客户端,处理单账号有限重试。
|
||||
- `upstream_call.rs`:把账号凭据、账号作用域身份和 Cookie 交给上游客户端。
|
||||
- `attempts.rs`:在请求开始时冻结完整候选顺序,记录已尝试、忙碌和状态排除的账号。
|
||||
- `errors.rs`:错误分类与客户端错误映射。
|
||||
- `recording.rs`:把调度轨迹和上游诊断写入 `telemetry::Recorder`。
|
||||
- `affinity`:会话亲和、账号作用域身份、Redis 索引和响应历史快照。
|
||||
- `affinity`:会话亲和、账号作用域身份和 Redis 响应归属索引。
|
||||
- `recovery`:账号风险隔离、Cloudflare 冷却、候选耗尽和 previous response 恢复。
|
||||
- `stream`:首段预取、SSE 失败识别、下游流转发和结束时结算。
|
||||
|
||||
@@ -378,7 +378,7 @@ bootstrap/
|
||||
5. `AccountAttemptLedger` 按当前调度策略冻结完整候选顺序。会话亲和账号可以排在前面,但不会绕过可用性检查。
|
||||
6. 账号池租用并发槽位,必要时先刷新 quota;请求间隔在发送上游前执行。
|
||||
7. `upstream_call` 注入账号 token、`chatgpt-account-id`、账号作用域身份、指纹请求头和 Cookie。
|
||||
8. 上游通过 HTTP/SSE 或 WebSocket 返回结果。单账号可重试的 5xx 最多额外重试两次。
|
||||
8. 上游通过 HTTP/SSE 或 WebSocket 返回结果。建连最长等待 15 秒;WebSocket 发送后 20 秒没有任何上游事件时放弃异常连接,收到首个事件后与 HTTP/SSE 一样使用 5 分钟流空闲窗口。账号级 5xx 和传输故障回到候选账本,不在同一账号内叠加重试。
|
||||
9. 在结果尚未提交给客户端时,账号级失败可继续尝试候选账本中的下一个账号。
|
||||
10. 完成后更新账号用量、会话亲和、调度反馈和遥测事实,并释放账号槽位。
|
||||
|
||||
@@ -391,7 +391,7 @@ bootstrap/
|
||||
- `scheduler::candidates` 过滤非 active、配额不可用、模型不匹配、处于 Cloudflare 冷却或明确排除的账号。
|
||||
- 当前 `rotation_strategy` 对剩余账号排序。请求级快照保留所有符合条件的层级,层级只影响顺序,不缩小 failover 边界。
|
||||
|
||||
单个请求的候选顺序在开始时冻结。瞬时并发槽位占满的账号进入 busy 队列,等待槽位变化后再判断;状态已经变化的账号被排除。候选账本不会重复取出同一个账号。单账号内部的 5xx 重试,以及 previous response 失效后在原账号执行的完整重放,是账本之外的明确重试。
|
||||
单个请求的候选顺序在开始时冻结。瞬时并发槽位占满的账号进入 busy 队列,最多等待 10 秒;状态已经变化的账号被排除。候选账本不会重复取出同一个账号。复用的 WebSocket 已失效时允许新建一次连接;当前下游 WebSocket 持有完整历史时,previous response 恢复可以在同账号或后续候选账号上重放。
|
||||
|
||||
会触发账号隔离或换号的典型情况:
|
||||
|
||||
@@ -408,21 +408,17 @@ bootstrap/
|
||||
|
||||
## previous response 与流式恢复
|
||||
|
||||
成功响应写入 Redis 亲和记录,保存响应 ID、账号 ID、会话 ID、turn state、变体哈希和受限的 replay 节点。
|
||||
成功响应写入 Redis 亲和记录,保存响应 ID、账号 ID、会话 ID、turn state 和变体哈希。Redis 不保存请求正文、响应正文或完整对话历史。
|
||||
|
||||
previous response 分两类处理:
|
||||
完整历史只保存在当前下游 Responses WebSocket 的连接状态中。每次 `response.completed` 后,将本轮输入和上游输出追加到该连接的 transcript;写入前递归删除对象 `id` 和 `encrypted_content`,保留工具续接所需的 `call_id`。连接关闭时 transcript 随处理任务一并释放,不写入 Redis,也没有进程级共享 transcript 缓存。
|
||||
|
||||
- 受管历史:Redis 能找到响应归属。先使用原账号和原 previous response;若上游返回 `previous_response_not_found`、连接续接忙碌或原连接不可复用,先在原账号执行完整历史重放。需要换号且 replay 完整时,其他账号也使用完整历史重放。
|
||||
- 外部未知历史:Redis 没有该响应。只能在第一个账号原样尝试一次,因为服务端没有足够历史安全地转移到其他账号;失败时保留上游错误。
|
||||
previous response 按历史来源处理:
|
||||
|
||||
replay 按轮保存增量输入和输出,不保存每一轮累计的完整输入。写入前递归删除对象 ID 和 `encrypted_content`。容量边界:
|
||||
- 当前下游 WebSocket 有完整 transcript:先使用亲和账号和原 previous response。若上游返回 `previous_response_not_found`、连接忙碌或原连接失效,去掉 previous ID 后用完整 transcript 重放。账号级失败需要换号时,后续账号仍由当前调度策略选出,并接收同一份完整 transcript。
|
||||
- Redis 能找到归属但当前连接没有 transcript:只允许原账号按原请求续接;失败后不能把不完整上下文发给其他账号。
|
||||
- Redis 没有归属且当前连接也没有 transcript:在首个候选账号原样尝试一次,失败时保留上游错误。
|
||||
|
||||
- 单轮 snapshot:2 MiB
|
||||
- 单会话累计 replay:16 MiB
|
||||
- 最大深度:128 轮
|
||||
- Redis TTL:4 小时
|
||||
|
||||
超过任一边界时不再生成新的 replay 节点,正常代理响应不因此失败。
|
||||
新的下游 WebSocket 不继承上一条连接的 transcript。官方 Codex 客户端收到可重试断流后,从本地 Session 重新组装完整历史并建立新请求;该请求重新进入账号池调度。客户端单条 HTTP 请求或 WebSocket 文本帧上限为 16 MiB。
|
||||
|
||||
流式请求只有在响应尚未提交给客户端时才能透明恢复或换号。首段预取会识别终止事件和账号级失败;一旦真实输出已经发给客户端,代理会用 `response.failed` 结束当前 SSE 或 WebSocket response,不能在另一账号上重放后拼接到已有输出。
|
||||
|
||||
@@ -483,11 +479,14 @@ Redis 使用统一的 `cpr:` 前缀。
|
||||
| `cpr:admin:session:<hash>` | 管理员登录会话 | `admin.session_ttl_minutes` |
|
||||
| `cpr:lease:refresh:<account_id>` | token 刷新租约 | 租约 PX TTL |
|
||||
| `cpr:models:plan_snapshots` | 各订阅计划模型快照 HASH | 下一次完整刷新替换 |
|
||||
| `cpr:affinity:v2:resp:<response_id>` | 响应归属和 replay 节点 | 4 小时 |
|
||||
| `cpr:affinity:v2:conv:<conversation_id>` | 会话响应 ZSET 索引 | 4 小时 |
|
||||
| `cpr:affinity:v2:account:<account_id>` | 账号响应 ZSET 索引 | 4 小时 |
|
||||
| `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:global:index` | 全局响应 ZSET 索引,最多 65536 项 | 4 小时 |
|
||||
| `cpr:affinity:v3:global:bytes` | 亲和元数据总字节数 | 4 小时 |
|
||||
| `cpr:affinity:v3:global:info` | 亲和索引清理所需的归属与字节信息 | 4 小时 |
|
||||
|
||||
账号删除时会批量删除该账号的亲和响应并清理会话索引。Redis 写入失败不应改变已经得到的上游响应语义,但会失去对应的会话恢复、模型缓存或分布式协调能力。
|
||||
亲和元数据单条不超过 64 KiB,全局最多 128 MiB。写入时惰性清理过期成员并执行会话、账号和全局裁剪。账号删除使用一段 Redis Lua 操作批量删除该账号的响应归属并清理会话索引。Redis 写入失败不改变已经得到的上游响应语义,但会失去对应的账号亲和、模型缓存或分布式协调能力。
|
||||
|
||||
### 本地文件
|
||||
|
||||
|
||||
@@ -1 +1 @@
|
||||
version: "2.1.0"
|
||||
version: "2.2.0"
|
||||
|
||||
Reference in New Issue
Block a user