mirror of
https://github.com/zyycn/codex-proxy-rs.git
synced 2026-11-10 00:06:00 +00:00
Compare commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
bcccebc9fa | ||
|
|
d67b46a9e2 | ||
|
|
20c6d33cf4 | ||
|
|
cf0e84c548 |
@@ -16,6 +16,7 @@ const OPENAI_PROTOCOL: &str = "openai";
|
||||
const OPENAI_SUBAGENT_KEY: &str = "x-openai-subagent";
|
||||
const CODEX_TURN_METADATA_KEY: &str = "x-codex-turn-metadata";
|
||||
const PASSTHROUGH_HEADERS_CONTEXT_KEY: &str = "opaque_request_headers";
|
||||
const DOWNSTREAM_WEBSOCKET_CONNECTION_ID_CONTEXT_KEY: &str = "downstream_websocket_connection_id";
|
||||
|
||||
/// Responses 请求进入共享解码内核时的下游传输来源。
|
||||
#[derive(Clone, Copy)]
|
||||
@@ -35,6 +36,7 @@ pub struct OpenAiRequestHeaders {
|
||||
include_timing_metrics: Option<String>,
|
||||
codex_window_id: Option<String>,
|
||||
parent_thread_id: Option<String>,
|
||||
downstream_websocket_connection_id: Option<String>,
|
||||
|
||||
conversation_id: Option<String>,
|
||||
session_id: Option<String>,
|
||||
@@ -60,6 +62,7 @@ impl OpenAiRequestHeaders {
|
||||
include_timing_metrics: header_string(headers, "x-responsesapi-include-timing-metrics"),
|
||||
codex_window_id: header_string(headers, "x-codex-window-id"),
|
||||
parent_thread_id: header_string(headers, "x-codex-parent-thread-id"),
|
||||
downstream_websocket_connection_id: None,
|
||||
conversation_id: header_string(headers, "conversation-id")
|
||||
.or_else(|| header_string(headers, "conversation_id")),
|
||||
session_id: header_string(headers, "session-id")
|
||||
@@ -74,6 +77,16 @@ impl OpenAiRequestHeaders {
|
||||
}
|
||||
}
|
||||
|
||||
/// 绑定代理为当前下游 WebSocket 分配的连接身份。
|
||||
#[must_use]
|
||||
pub(super) fn with_downstream_websocket_connection_id(
|
||||
mut self,
|
||||
connection_id: impl Into<String>,
|
||||
) -> Self {
|
||||
self.downstream_websocket_connection_id = Some(connection_id.into());
|
||||
self
|
||||
}
|
||||
|
||||
fn apply_subagent(&self, body: &mut Map<String, Value>) {
|
||||
if let Some(subagent) = self.subagent.as_deref() {
|
||||
inject_subagent_metadata(body, subagent);
|
||||
@@ -109,6 +122,11 @@ impl OpenAiRequestHeaders {
|
||||
"parent_thread_id",
|
||||
self.parent_thread_id.as_ref(),
|
||||
);
|
||||
insert_protocol_context(
|
||||
&mut context,
|
||||
DOWNSTREAM_WEBSOCKET_CONNECTION_ID_CONTEXT_KEY,
|
||||
self.downstream_websocket_connection_id.as_ref(),
|
||||
);
|
||||
insert_protocol_context(
|
||||
&mut context,
|
||||
"conversation_id",
|
||||
|
||||
@@ -143,6 +143,8 @@ async fn serve_responses_websocket(socket: WebSocket, session: ResponsesWebSocke
|
||||
);
|
||||
let _connection_guard = connection_guard;
|
||||
let cancellation = lifecycle.cancellation();
|
||||
let request_headers =
|
||||
request_headers.with_downstream_websocket_connection_id(connection_id.clone());
|
||||
let mut connection = ResponsesWebSocketConnection::new(socket, connection_id, cancellation);
|
||||
let mut request_count = 0_u64;
|
||||
let mut replay = ConnectionReplaySnapshot::default();
|
||||
|
||||
@@ -2,7 +2,7 @@ mod connection;
|
||||
mod protocol;
|
||||
|
||||
use std::sync::{
|
||||
Arc,
|
||||
Arc, Mutex,
|
||||
atomic::{AtomicBool, AtomicUsize, Ordering},
|
||||
};
|
||||
use std::time::{Duration, SystemTime};
|
||||
@@ -280,6 +280,7 @@ struct AtomicFailureTrace {
|
||||
next_calls: AtomicUsize,
|
||||
committed: AtomicBool,
|
||||
finalized: AtomicBool,
|
||||
downstream_websocket_connection_ids: Mutex<Vec<String>>,
|
||||
}
|
||||
|
||||
struct AtomicFailureSession {
|
||||
@@ -365,6 +366,19 @@ impl ExecutionService for AtomicFailureExecution {
|
||||
request: StartExecution,
|
||||
) -> BoxFuture<'_, Result<StartedExecution, GatewayError>> {
|
||||
self.trace.starts.fetch_add(1, Ordering::AcqRel);
|
||||
if let Operation::Generate(generate) = &request.operation
|
||||
&& let Some(connection_id) = generate
|
||||
.protocol_payload()
|
||||
.context()
|
||||
.get("downstream_websocket_connection_id")
|
||||
.and_then(Value::as_str)
|
||||
{
|
||||
self.trace
|
||||
.downstream_websocket_connection_ids
|
||||
.lock()
|
||||
.expect("downstream connection trace lock")
|
||||
.push(connection_id.to_owned());
|
||||
}
|
||||
Box::pin(async move {
|
||||
Ok(StartedExecution {
|
||||
request_id: gateway_core::engine::ModelRequestId::new("req_ws_atomic")
|
||||
@@ -510,6 +524,14 @@ async fn websocket_atomic_upstream_failure_batch_should_be_forwarded_once() {
|
||||
assert!(trace.committed.load(Ordering::Acquire));
|
||||
assert_eq!(trace.next_calls.load(Ordering::Acquire), 2);
|
||||
assert!(trace.finalized.load(Ordering::Acquire));
|
||||
{
|
||||
let downstream_connection_ids = trace
|
||||
.downstream_websocket_connection_ids
|
||||
.lock()
|
||||
.expect("downstream connection trace lock");
|
||||
assert_eq!(downstream_connection_ids.len(), 1);
|
||||
assert!(downstream_connection_ids[0].starts_with("ws_"));
|
||||
}
|
||||
|
||||
socket.close(None).await.expect("close WebSocket");
|
||||
server.abort();
|
||||
|
||||
@@ -728,7 +728,8 @@ impl ModelRequestRepository for PgExecutionStore {
|
||||
upstream_connection_idle_ms = $47
|
||||
where id = $1 and outcome = 'running'
|
||||
returning id, client_api_key_ref, continuation_affinity_hash,
|
||||
continuation_requested, outcome, started_at, completed_at
|
||||
continuation_requested, provider_kind, upstream_transport,
|
||||
outcome, started_at, completed_at
|
||||
), recovery_target as (
|
||||
select prior.id
|
||||
from model_requests prior
|
||||
@@ -775,9 +776,47 @@ impl ModelRequestRepository for PgExecutionStore {
|
||||
where prior.id = target.id
|
||||
and prior.recovered_at is null
|
||||
returning prior.id
|
||||
), session_transport_recovery_target as (
|
||||
select prior.id
|
||||
from model_requests prior
|
||||
cross join finalized current
|
||||
where current.outcome = 'succeeded'
|
||||
and current.provider_kind = 'openai'
|
||||
and current.upstream_transport = 'http_sse'
|
||||
and current.continuation_affinity_hash is not null
|
||||
and prior.client_api_key_ref = current.client_api_key_ref
|
||||
and prior.continuation_affinity_hash = current.continuation_affinity_hash
|
||||
and prior.provider_kind = current.provider_kind
|
||||
and prior.outcome = 'failed'
|
||||
and prior.error_kind = 'upstream_unavailable'
|
||||
and prior.upstream_transport = 'websocket'
|
||||
and prior.upstream_send_state = 'ambiguous'
|
||||
and prior.downstream_committed_at is null
|
||||
and prior.recovered_at is null
|
||||
and current.started_at >= prior.completed_at
|
||||
and current.started_at <= prior.completed_at + interval '30 seconds'
|
||||
), session_transport_recovery_update as (
|
||||
update model_requests prior
|
||||
set recovery_attempt_count = prior.recovery_attempt_count + 1,
|
||||
recovery_request_id = current.id,
|
||||
recovered_at = current.completed_at,
|
||||
recovery_retry_delay_ms = greatest(
|
||||
0,
|
||||
floor(extract(epoch from (current.started_at - prior.completed_at)) * 1000)
|
||||
)::bigint,
|
||||
recovery_total_latency_ms = greatest(
|
||||
0,
|
||||
floor(extract(epoch from (current.completed_at - prior.completed_at)) * 1000)
|
||||
)::bigint
|
||||
from finalized current
|
||||
join session_transport_recovery_target target on true
|
||||
where prior.id = target.id
|
||||
and prior.recovered_at is null
|
||||
returning prior.id
|
||||
)
|
||||
select (select count(*) from finalized)::bigint
|
||||
+ (select count(*) * 0 from recovery_update)::bigint",
|
||||
+ (select count(*) * 0 from recovery_update)::bigint
|
||||
+ (select count(*) * 0 from session_transport_recovery_update)::bigint",
|
||||
)
|
||||
.bind(&finalization.model_request_id)
|
||||
.bind(finalization.outcome.as_str())
|
||||
|
||||
@@ -128,7 +128,6 @@ fn push_request_error_predicates(
|
||||
range: ObservabilityRange,
|
||||
filter: &OpsErrorFilter,
|
||||
) {
|
||||
statement.push(" and mr.recovered_at is null");
|
||||
push_range(statement, "mr.completed_at", range);
|
||||
for (column, value) in [
|
||||
("mr.client_api_key_ref", &filter.client_api_key_ref),
|
||||
@@ -171,7 +170,6 @@ fn push_ops_event_predicates(
|
||||
range: ObservabilityRange,
|
||||
filter: &OpsErrorFilter,
|
||||
) {
|
||||
statement.push(" and (mr.id is null or mr.recovered_at is null)");
|
||||
push_range(statement, "oe.created_at", range);
|
||||
for (column, value) in [
|
||||
("mr.client_api_key_ref", &filter.client_api_key_ref),
|
||||
|
||||
@@ -719,6 +719,165 @@ async fn successful_new_chain_should_mark_recent_continuation_failure_recovered(
|
||||
database.close().await;
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn successful_http_downgrade_should_mark_pending_websocket_failures_recovered() {
|
||||
let Some(database) = TestDatabase::create("execution_transport_downgrade_recovered").await
|
||||
else {
|
||||
return;
|
||||
};
|
||||
let affinity_hash = "d".repeat(64);
|
||||
let first_failure_at = Utc::now() - Duration::seconds(8);
|
||||
let second_failure_at = first_failure_at + Duration::seconds(1);
|
||||
seed_websocket_transport_failure(
|
||||
&database.pool,
|
||||
"req_websocket_failed_first",
|
||||
&affinity_hash,
|
||||
first_failure_at,
|
||||
)
|
||||
.await;
|
||||
seed_websocket_transport_failure(
|
||||
&database.pool,
|
||||
"req_websocket_failed_second",
|
||||
&affinity_hash,
|
||||
second_failure_at,
|
||||
)
|
||||
.await;
|
||||
seed_websocket_transport_failure(
|
||||
&database.pool,
|
||||
"req_websocket_failed_after_delivery",
|
||||
&affinity_hash,
|
||||
second_failure_at,
|
||||
)
|
||||
.await;
|
||||
sqlx::query(
|
||||
"update model_requests
|
||||
set downstream_committed_at = completed_at
|
||||
where id = 'req_websocket_failed_after_delivery'",
|
||||
)
|
||||
.execute(&database.pool)
|
||||
.await
|
||||
.expect("mark control failure delivered downstream");
|
||||
|
||||
let recovery_started_at = first_failure_at + Duration::seconds(4);
|
||||
seed_transport_recovery_request(
|
||||
&database.pool,
|
||||
"req_http_fallback_succeeded",
|
||||
&affinity_hash,
|
||||
recovery_started_at,
|
||||
)
|
||||
.await;
|
||||
let store = PgExecutionStore::new(database.pool.clone());
|
||||
let mut recovery = successful_core_finalization("req_http_fallback_succeeded");
|
||||
recovery.upstream_transport = Some("http_sse".to_owned());
|
||||
recovery.downstream_committed_at = Some(std::time::SystemTime::now());
|
||||
ExecutionStore::finalize_model_request(&store, recovery)
|
||||
.await
|
||||
.expect("finalize successful HTTP fallback");
|
||||
|
||||
type PersistedRecovery = (
|
||||
String,
|
||||
Option<String>,
|
||||
Option<chrono::DateTime<Utc>>,
|
||||
i32,
|
||||
Option<i64>,
|
||||
Option<i64>,
|
||||
);
|
||||
let recovered: Vec<PersistedRecovery> = sqlx::query_as(
|
||||
"select id, recovery_request_id, recovered_at, recovery_attempt_count,
|
||||
recovery_retry_delay_ms, recovery_total_latency_ms
|
||||
from model_requests
|
||||
where id like 'req_websocket_failed_%'
|
||||
order by id",
|
||||
)
|
||||
.fetch_all(&database.pool)
|
||||
.await
|
||||
.expect("load WebSocket recovery observations");
|
||||
|
||||
assert_eq!(recovered.len(), 3);
|
||||
let delivered = &recovered[0];
|
||||
assert_eq!(delivered.0, "req_websocket_failed_after_delivery");
|
||||
assert_eq!((&delivered.1, delivered.2, delivered.3), (&None, None, 0));
|
||||
for (row, expected_delay_ms) in recovered[1..].iter().zip([4_000, 3_000]) {
|
||||
assert_eq!(row.1.as_deref(), Some("req_http_fallback_succeeded"));
|
||||
assert!(row.2.is_some());
|
||||
assert_eq!(row.3, 1);
|
||||
assert_eq!(row.4, Some(expected_delay_ms));
|
||||
assert!(row.5.is_some_and(|latency| latency >= expected_delay_ms));
|
||||
}
|
||||
database.close().await;
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn websocket_success_or_late_http_success_should_not_hide_transport_failure() {
|
||||
let Some(database) = TestDatabase::create("execution_transport_downgrade_unresolved").await
|
||||
else {
|
||||
return;
|
||||
};
|
||||
let failure_at = Utc::now() - Duration::seconds(60);
|
||||
let store = PgExecutionStore::new(database.pool.clone());
|
||||
|
||||
let websocket_affinity = "e".repeat(64);
|
||||
seed_websocket_transport_failure(
|
||||
&database.pool,
|
||||
"req_websocket_control_failure",
|
||||
&websocket_affinity,
|
||||
failure_at,
|
||||
)
|
||||
.await;
|
||||
seed_transport_recovery_request(
|
||||
&database.pool,
|
||||
"req_websocket_control_success",
|
||||
&websocket_affinity,
|
||||
failure_at + Duration::seconds(4),
|
||||
)
|
||||
.await;
|
||||
ExecutionStore::finalize_model_request(
|
||||
&store,
|
||||
successful_core_finalization("req_websocket_control_success"),
|
||||
)
|
||||
.await
|
||||
.expect("finalize successful WebSocket control request");
|
||||
|
||||
let late_affinity = "f".repeat(64);
|
||||
seed_websocket_transport_failure(
|
||||
&database.pool,
|
||||
"req_websocket_late_failure",
|
||||
&late_affinity,
|
||||
failure_at,
|
||||
)
|
||||
.await;
|
||||
seed_transport_recovery_request(
|
||||
&database.pool,
|
||||
"req_http_late_success",
|
||||
&late_affinity,
|
||||
failure_at + Duration::seconds(31),
|
||||
)
|
||||
.await;
|
||||
let mut late_recovery = successful_core_finalization("req_http_late_success");
|
||||
late_recovery.upstream_transport = Some("http_sse".to_owned());
|
||||
ExecutionStore::finalize_model_request(&store, late_recovery)
|
||||
.await
|
||||
.expect("finalize late HTTP request");
|
||||
|
||||
type UnresolvedTransportFailure = (String, Option<String>, Option<chrono::DateTime<Utc>>, i32);
|
||||
let unresolved: Vec<UnresolvedTransportFailure> = sqlx::query_as(
|
||||
"select id, recovery_request_id, recovered_at, recovery_attempt_count
|
||||
from model_requests
|
||||
where id in ('req_websocket_control_failure', 'req_websocket_late_failure')
|
||||
order by id",
|
||||
)
|
||||
.fetch_all(&database.pool)
|
||||
.await
|
||||
.expect("load unresolved transport failures");
|
||||
assert_eq!(unresolved.len(), 2);
|
||||
assert!(
|
||||
unresolved
|
||||
.iter()
|
||||
.all(|row| row.1.is_none() && row.2.is_none() && row.3 == 0)
|
||||
);
|
||||
database.close().await;
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn failed_or_late_new_chain_should_leave_continuation_failure_unresolved() {
|
||||
let Some(database) = TestDatabase::create("execution_continuation_unresolved").await else {
|
||||
@@ -842,6 +1001,61 @@ async fn seed_recovery_request(
|
||||
.expect("prepare seeded recovery request");
|
||||
}
|
||||
|
||||
async fn seed_websocket_transport_failure(
|
||||
pool: &sqlx::PgPool,
|
||||
id: &str,
|
||||
affinity_hash: &str,
|
||||
completed_at: chrono::DateTime<Utc>,
|
||||
) {
|
||||
seed_running_request(pool, id)
|
||||
.await
|
||||
.expect("seed WebSocket transport failure request");
|
||||
sqlx::query(
|
||||
"update model_requests
|
||||
set client_api_key_ref = 'key_transport_recovery',
|
||||
client_transport = 'http_sse',
|
||||
continuation_affinity_hash = $2,
|
||||
attempt_count = 1, upstream_send_state = 'ambiguous',
|
||||
upstream_transport = 'websocket', outcome = 'failed',
|
||||
error_kind = 'upstream_unavailable',
|
||||
provider_error_code = 'websocket_close_1000',
|
||||
started_at = $3 - interval '1 second', deadline_at = $3,
|
||||
completed_at = $3
|
||||
where id = $1",
|
||||
)
|
||||
.bind(id)
|
||||
.bind(affinity_hash)
|
||||
.bind(completed_at)
|
||||
.execute(pool)
|
||||
.await
|
||||
.expect("finalize seeded WebSocket transport failure");
|
||||
}
|
||||
|
||||
async fn seed_transport_recovery_request(
|
||||
pool: &sqlx::PgPool,
|
||||
id: &str,
|
||||
affinity_hash: &str,
|
||||
started_at: chrono::DateTime<Utc>,
|
||||
) {
|
||||
seed_running_request(pool, id)
|
||||
.await
|
||||
.expect("seed transport recovery request");
|
||||
sqlx::query(
|
||||
"update model_requests
|
||||
set client_api_key_ref = 'key_transport_recovery',
|
||||
client_transport = 'http_sse',
|
||||
continuation_affinity_hash = $2,
|
||||
started_at = $3, deadline_at = $3 + interval '1 minute'
|
||||
where id = $1",
|
||||
)
|
||||
.bind(id)
|
||||
.bind(affinity_hash)
|
||||
.bind(started_at)
|
||||
.execute(pool)
|
||||
.await
|
||||
.expect("prepare seeded transport recovery request");
|
||||
}
|
||||
|
||||
async fn seed_running_request(pool: &sqlx::PgPool, id: &str) -> Result<(), sqlx::Error> {
|
||||
sqlx::query(
|
||||
"insert into model_requests (
|
||||
|
||||
@@ -258,7 +258,8 @@ async fn ops_search_should_treat_sql_wildcards_as_literals() {
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn recovered_continuation_failure_should_be_hidden_from_default_business_metrics() {
|
||||
async fn recovered_continuation_failure_should_be_visible_in_ops_but_hidden_from_business_metrics()
|
||||
{
|
||||
let Some(database) = TestDatabase::create("observability_recovered_continuation").await else {
|
||||
return;
|
||||
};
|
||||
@@ -318,8 +319,13 @@ async fn recovered_continuation_failure_should_be_hidden_from_default_business_m
|
||||
page_size: ObservabilityPageSize::new(10).expect("page size"),
|
||||
})
|
||||
.await
|
||||
.expect("ops errors without recovered intermediates");
|
||||
assert_eq!(errors.total, 0);
|
||||
.expect("ops errors retain recovered failures");
|
||||
assert_eq!(errors.total, 2);
|
||||
assert!(errors.items.iter().all(|error| {
|
||||
error.recovery_request_id.as_deref() == Some("req_observe_success")
|
||||
&& error.recovered_at.is_some()
|
||||
&& error.recovery_attempt_count == 1
|
||||
}));
|
||||
|
||||
let retained: (String, Option<String>) = sqlx::query_as(
|
||||
"select outcome, recovery_request_id
|
||||
|
||||
@@ -682,10 +682,12 @@ impl CodexBackendClient {
|
||||
.local_conversation_id
|
||||
.as_deref()
|
||||
.or(request.previous_response_id())?;
|
||||
Some(
|
||||
CodexWebSocketPoolKey::new(&self.base_url, account_id, conversation_id)
|
||||
.with_connection_profile(connection_profile),
|
||||
)
|
||||
let mut key = CodexWebSocketPoolKey::new(&self.base_url, account_id, conversation_id)
|
||||
.with_connection_profile(connection_profile);
|
||||
if let Some(connection_id) = request.downstream_websocket_connection_id.as_deref() {
|
||||
key = key.with_downstream_connection_id(connection_id);
|
||||
}
|
||||
Some(key)
|
||||
}
|
||||
|
||||
/// 获取后端模型目录条目。
|
||||
|
||||
@@ -82,6 +82,8 @@ pub struct CodexResponsesRequest {
|
||||
pub memgen_request: Option<String>,
|
||||
/// Codex window ID。
|
||||
pub codex_window_id: Option<String>,
|
||||
/// 代理分配的下游 WebSocket 连接 ID,仅用于隔离本地连接池通道。
|
||||
pub downstream_websocket_connection_id: Option<String>,
|
||||
/// 父线程 ID。
|
||||
pub parent_thread_id: Option<String>,
|
||||
/// 已知 previous response 的持久化范围,仅用于本地 transport 校验。
|
||||
@@ -519,6 +521,7 @@ impl CodexResponsesRequest {
|
||||
responses_lite: None,
|
||||
memgen_request: None,
|
||||
codex_window_id: None,
|
||||
downstream_websocket_connection_id: None,
|
||||
parent_thread_id: None,
|
||||
previous_response_scope: None,
|
||||
}
|
||||
|
||||
@@ -511,6 +511,9 @@ fn apply_protocol_context(request: &mut CodexResponsesRequest, context: &Map<Str
|
||||
.or_else(|| request.include_timing_metrics.take());
|
||||
request.codex_window_id =
|
||||
context_string(context, "codex_window_id").or_else(|| request.codex_window_id.take());
|
||||
request.downstream_websocket_connection_id =
|
||||
context_string(context, "downstream_websocket_connection_id")
|
||||
.or_else(|| request.downstream_websocket_connection_id.take());
|
||||
request.parent_thread_id =
|
||||
context_string(context, "parent_thread_id").or_else(|| request.parent_thread_id.take());
|
||||
let prompt_cache_key = request.prompt_cache_key().map(ToOwned::to_owned);
|
||||
|
||||
@@ -23,6 +23,7 @@ pub struct CodexWebSocketPoolKey {
|
||||
account_id: String,
|
||||
conversation_id: String,
|
||||
connection_profile: String,
|
||||
downstream_connection_id: String,
|
||||
}
|
||||
|
||||
impl CodexWebSocketPoolKey {
|
||||
@@ -37,6 +38,7 @@ impl CodexWebSocketPoolKey {
|
||||
account_id: account_id.into(),
|
||||
conversation_id: conversation_id.into(),
|
||||
connection_profile: String::new(),
|
||||
downstream_connection_id: String::new(),
|
||||
}
|
||||
}
|
||||
|
||||
@@ -46,6 +48,15 @@ impl CodexWebSocketPoolKey {
|
||||
self
|
||||
}
|
||||
|
||||
/// 隔离同一逻辑会话中由不同下游 WebSocket 驱动的并发响应链。
|
||||
pub(crate) fn with_downstream_connection_id(
|
||||
mut self,
|
||||
connection_id: impl Into<String>,
|
||||
) -> Self {
|
||||
self.downstream_connection_id = connection_id.into();
|
||||
self
|
||||
}
|
||||
|
||||
pub(crate) fn account_id(&self) -> &str {
|
||||
&self.account_id
|
||||
}
|
||||
@@ -60,6 +71,7 @@ impl CodexWebSocketPoolKey {
|
||||
self.account_id.as_str(),
|
||||
self.conversation_id.as_str(),
|
||||
self.connection_profile.as_str(),
|
||||
self.downstream_connection_id.as_str(),
|
||||
])
|
||||
}
|
||||
|
||||
|
||||
@@ -260,6 +260,35 @@ fn header_context_should_win_over_body_topline_aliases() {
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn encoder_should_preserve_downstream_websocket_connection_identity_outside_wire_body() {
|
||||
let payload = ProtocolPayload::json_object(
|
||||
"openai",
|
||||
Map::from_iter([
|
||||
("model".to_owned(), json!("client-model")),
|
||||
("input".to_owned(), json!("prompt")),
|
||||
]),
|
||||
)
|
||||
.expect("OpenAI payload")
|
||||
.with_context(Map::from_iter([(
|
||||
"downstream_websocket_connection_id".to_owned(),
|
||||
Value::String("ws_downstream_a".to_owned()),
|
||||
)]));
|
||||
let request = GenerateRequest::from_protocol_payload(payload);
|
||||
|
||||
let encoded = encode_generate_request(&request, "gpt-test").expect("encode");
|
||||
|
||||
assert_eq!(
|
||||
encoded.downstream_websocket_connection_id.as_deref(),
|
||||
Some("ws_downstream_a")
|
||||
);
|
||||
assert!(
|
||||
!encoded
|
||||
.body()
|
||||
.contains_key("downstream_websocket_connection_id")
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn body_topline_alias_should_only_fill_an_absent_header_context() {
|
||||
let payload = ProtocolPayload::json_object(
|
||||
|
||||
@@ -71,6 +71,138 @@ async fn codex_backend_client_should_reuse_pooled_websocket_for_same_account_and
|
||||
assert_eq!(accepted_connections.load(Ordering::SeqCst), 1);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn websocket_pool_should_isolate_concurrent_downstream_lanes_for_one_conversation() {
|
||||
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
|
||||
let addr = listener.local_addr().unwrap();
|
||||
let accepted_connections = Arc::new(AtomicUsize::new(0));
|
||||
let accepted_connections_for_server = Arc::clone(&accepted_connections);
|
||||
let (first_accepted_tx, first_accepted_rx) = tokio::sync::oneshot::channel();
|
||||
let server = tokio::spawn(async move {
|
||||
let (first_stream, _) = listener.accept().await.unwrap();
|
||||
accepted_connections_for_server.fetch_add(1, Ordering::SeqCst);
|
||||
let mut first = accept_codex_test_websocket(first_stream).await;
|
||||
first_accepted_tx.send(()).unwrap();
|
||||
|
||||
let (second_stream, _) = timeout(Duration::from_secs(2), listener.accept())
|
||||
.await
|
||||
.expect("the second downstream lane should open its own upstream websocket")
|
||||
.unwrap();
|
||||
accepted_connections_for_server.fetch_add(1, Ordering::SeqCst);
|
||||
let mut second = accept_codex_test_websocket(second_stream).await;
|
||||
|
||||
let (first_warmup, second_warmup) = tokio::join!(first.next(), second.next());
|
||||
first_warmup.unwrap().unwrap();
|
||||
second_warmup.unwrap().unwrap();
|
||||
first
|
||||
.send(Message::Text(
|
||||
completed_websocket_response("resp_lane_a", 1, 0).into(),
|
||||
))
|
||||
.await
|
||||
.unwrap();
|
||||
second
|
||||
.send(Message::Text(
|
||||
completed_websocket_response("resp_lane_b", 1, 0).into(),
|
||||
))
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
let _first_continuation = first.next().await.unwrap().unwrap();
|
||||
first
|
||||
.send(Message::Text(
|
||||
completed_websocket_response("resp_lane_a_next", 2, 1).into(),
|
||||
))
|
||||
.await
|
||||
.unwrap();
|
||||
let _second_continuation = second.next().await.unwrap().unwrap();
|
||||
second
|
||||
.send(Message::Text(
|
||||
completed_websocket_response("resp_lane_b_next", 2, 1).into(),
|
||||
))
|
||||
.await
|
||||
.unwrap();
|
||||
first.close(None).await.unwrap();
|
||||
second.close(None).await.unwrap();
|
||||
});
|
||||
let pool = Arc::new(CodexWebSocketPool::new(Duration::from_mins(1)));
|
||||
let backend = CodexBackendClient::new(
|
||||
reqwest::Client::builder().no_proxy().build().unwrap(),
|
||||
format!("http://{addr}"),
|
||||
test_wire_profile(),
|
||||
)
|
||||
.with_websocket_pool(Arc::clone(&pool));
|
||||
let warmup = |connection_id: &str| {
|
||||
let mut body = Map::new();
|
||||
body.insert("model".to_owned(), json!("gpt-5.5"));
|
||||
body.insert("input".to_owned(), json!([]));
|
||||
body.insert("generate".to_owned(), json!(false));
|
||||
body.insert("store".to_owned(), json!(false));
|
||||
let mut request = CodexResponsesRequest::from_body(body);
|
||||
request.use_websocket = true;
|
||||
request.local_conversation_id = Some("conversation-lanes".to_owned());
|
||||
request.downstream_websocket_connection_id = Some(connection_id.to_owned());
|
||||
request
|
||||
};
|
||||
|
||||
let first_backend = backend.clone();
|
||||
let first_request = warmup("ws_downstream_a");
|
||||
let first = tokio::spawn(async move {
|
||||
first_backend
|
||||
.create_response(
|
||||
&first_request,
|
||||
request_context("req_lane_a_warmup", Some("chatgpt-account")),
|
||||
)
|
||||
.await
|
||||
});
|
||||
first_accepted_rx.await.unwrap();
|
||||
let second = backend
|
||||
.create_response(
|
||||
&warmup("ws_downstream_b"),
|
||||
request_context("req_lane_b_warmup", Some("chatgpt-account")),
|
||||
)
|
||||
.await
|
||||
.expect("the second downstream lane should not contend with the first warmup");
|
||||
let first = first.await.unwrap().unwrap();
|
||||
|
||||
let mut first_continuation = pooled_websocket_request("conversation-lanes");
|
||||
first_continuation.downstream_websocket_connection_id = Some("ws_downstream_a".to_owned());
|
||||
first_continuation.set_previous_response_id(Some("resp_lane_a".to_owned()));
|
||||
first_continuation.previous_response_scope = Some(PreviousResponseScope::ConnectionLocal);
|
||||
let first_continuation = backend
|
||||
.create_response(
|
||||
&first_continuation,
|
||||
request_context("req_lane_a_next", Some("chatgpt-account")),
|
||||
)
|
||||
.await
|
||||
.expect("the first lane should retain its connection-local response owner");
|
||||
|
||||
let mut second_continuation = pooled_websocket_request("conversation-lanes");
|
||||
second_continuation.downstream_websocket_connection_id = Some("ws_downstream_b".to_owned());
|
||||
second_continuation.set_previous_response_id(Some("resp_lane_b".to_owned()));
|
||||
second_continuation.previous_response_scope = Some(PreviousResponseScope::ConnectionLocal);
|
||||
let second_continuation = backend
|
||||
.create_response(
|
||||
&second_continuation,
|
||||
request_context("req_lane_b_next", Some("chatgpt-account")),
|
||||
)
|
||||
.await
|
||||
.expect("the second lane should retain its connection-local response owner");
|
||||
server.await.unwrap();
|
||||
|
||||
assert_eq!(first.websocket_pool_decision.unwrap().kind(), "new");
|
||||
assert_eq!(second.websocket_pool_decision.unwrap().kind(), "new");
|
||||
assert_eq!(
|
||||
first_continuation.websocket_pool_decision.unwrap().kind(),
|
||||
"reuse"
|
||||
);
|
||||
assert_eq!(
|
||||
second_continuation.websocket_pool_decision.unwrap().kind(),
|
||||
"reuse"
|
||||
);
|
||||
assert_eq!(accepted_connections.load(Ordering::SeqCst), 2);
|
||||
pool.shutdown().await;
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn exact_continuation_reuses_owning_socket_after_connection_profile_updates() {
|
||||
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
|
||||
|
||||
@@ -3,3 +3,4 @@
|
||||
1b6ffab434c754814aebb0eee3b650bd9770010059472bccf5bc856b319880ca 0003_client_min_versions.sql
|
||||
22f912ab3e196de7983ef50a1dba5aee340bbea82bdd15d05863e87012000c83 0004_ops_error_send_state.sql
|
||||
194d7b4bedd83f0d76428a21699b0e946449bf2ba61a0fdf8710ebfd421af8b2 0005_openai_continuation_recovery_observability.sql
|
||||
3b4bfd5e3f5135fa2fad1088b0c2a287861e1250d66af31dbd40dc890751b567 0006_openai_session_transport_recovery.sql
|
||||
|
||||
@@ -0,0 +1,61 @@
|
||||
create index model_requests_pending_ws_transport_recovery_idx
|
||||
on model_requests (
|
||||
client_api_key_ref,
|
||||
continuation_affinity_hash,
|
||||
completed_at desc,
|
||||
id desc
|
||||
)
|
||||
where provider_kind = 'openai'
|
||||
and outcome = 'failed'
|
||||
and error_kind = 'upstream_unavailable'
|
||||
and upstream_transport = 'websocket'
|
||||
and upstream_send_state = 'ambiguous'
|
||||
and downstream_committed_at is null
|
||||
and recovered_at is null
|
||||
and continuation_affinity_hash is not null;
|
||||
|
||||
with session_transport_recoveries as (
|
||||
select
|
||||
prior.id as failed_request_id,
|
||||
recovered.id as recovery_request_id,
|
||||
recovered.started_at as recovery_started_at,
|
||||
recovered.completed_at as recovered_at
|
||||
from model_requests prior
|
||||
cross join lateral (
|
||||
select current.id, current.started_at, current.completed_at
|
||||
from model_requests current
|
||||
where current.client_api_key_ref = prior.client_api_key_ref
|
||||
and current.continuation_affinity_hash = prior.continuation_affinity_hash
|
||||
and current.provider_kind = prior.provider_kind
|
||||
and current.outcome = 'succeeded'
|
||||
and current.upstream_transport = 'http_sse'
|
||||
and current.completed_at is not null
|
||||
and current.started_at >= prior.completed_at
|
||||
and current.started_at <= prior.completed_at + interval '30 seconds'
|
||||
order by current.started_at, current.id
|
||||
limit 1
|
||||
) recovered
|
||||
where prior.provider_kind = 'openai'
|
||||
and prior.outcome = 'failed'
|
||||
and prior.error_kind = 'upstream_unavailable'
|
||||
and prior.upstream_transport = 'websocket'
|
||||
and prior.upstream_send_state = 'ambiguous'
|
||||
and prior.downstream_committed_at is null
|
||||
and prior.recovered_at is null
|
||||
and prior.continuation_affinity_hash is not null
|
||||
)
|
||||
update model_requests prior
|
||||
set recovery_attempt_count = prior.recovery_attempt_count + 1,
|
||||
recovery_request_id = recovery.recovery_request_id,
|
||||
recovered_at = recovery.recovered_at,
|
||||
recovery_retry_delay_ms = greatest(
|
||||
0,
|
||||
floor(extract(epoch from (recovery.recovery_started_at - prior.completed_at)) * 1000)
|
||||
)::bigint,
|
||||
recovery_total_latency_ms = greatest(
|
||||
0,
|
||||
floor(extract(epoch from (recovery.recovered_at - prior.completed_at)) * 1000)
|
||||
)::bigint
|
||||
from session_transport_recoveries recovery
|
||||
where prior.id = recovery.failed_request_id
|
||||
and prior.recovered_at is null;
|
||||
@@ -45,7 +45,8 @@ const thumbClasses = computed(() => [
|
||||
model.value ? 'left-[calc(100%-1.375rem)]' : 'left-0.5',
|
||||
])
|
||||
const inlineTextClasses = computed(() => [
|
||||
'pointer-events-none absolute inset-0 flex min-w-0 select-none items-center justify-center overflow-hidden text-[10px] leading-none font-heavy text-cp-white transition-[padding] duration-180 motion-reduce:transition-none',
|
||||
'pointer-events-none absolute inset-0 flex min-w-0 select-none items-center justify-center overflow-hidden text-[10px] leading-none font-heavy transition-[padding,color] duration-180 motion-reduce:transition-none',
|
||||
model.value ? 'text-cp-white' : 'text-cp-text-secondary',
|
||||
model.value ? 'pr-6 pl-1' : 'pr-1 pl-6',
|
||||
])
|
||||
const inactiveTextClasses = computed(() => [
|
||||
|
||||
@@ -1,5 +1,5 @@
|
||||
export const CODEX_DEFAULT_MODEL = 'gpt-5.6-terra'
|
||||
export const CODEX_WEBSOCKET_ENABLED_BY_DEFAULT = true
|
||||
export const CODEX_WEBSOCKET_ENABLED_BY_DEFAULT = false
|
||||
|
||||
export interface CodexConfigInput {
|
||||
apiKey: string
|
||||
|
||||
@@ -31,6 +31,10 @@ const accountLabel = computed(() => props.record?.accountEmail
|
||||
|| props.record?.accountId)
|
||||
|
||||
const errorFields = computed(() => visibleFields([
|
||||
{
|
||||
label: '处理状态',
|
||||
value: props.record?.metadata.recoveredAt ? '已自动恢复' : null,
|
||||
},
|
||||
{ label: 'Provider 错误码', value: props.record?.providerErrorCode, mono: true },
|
||||
{ label: '失败分类', value: failureClassText(props.record?.failureClass) },
|
||||
{ label: '分类代码', value: props.record?.failureClass, mono: true },
|
||||
|
||||
@@ -124,9 +124,17 @@ function upstreamSendStateText(value: string | null | undefined) {
|
||||
</template>
|
||||
<template #message="{ row }">
|
||||
<div class="min-w-0 py-0.5" :title="row.message || opsErrorSummary(row)">
|
||||
<code class="block max-w-full truncate font-mono text-cp-sm font-bold text-cp-error-text">
|
||||
{{ opsErrorSummary(row) }}
|
||||
</code>
|
||||
<div class="flex min-w-0 items-center gap-2">
|
||||
<code class="block min-w-0 flex-1 truncate font-mono text-cp-sm font-bold text-cp-error-text">
|
||||
{{ opsErrorSummary(row) }}
|
||||
</code>
|
||||
<span
|
||||
v-if="row.metadata.recoveredAt"
|
||||
class="inline-flex h-5 shrink-0 items-center rounded-full bg-cp-success-bg px-2 text-cp-xs leading-none font-heavy text-cp-success-text"
|
||||
>
|
||||
已自动恢复
|
||||
</span>
|
||||
</div>
|
||||
<p
|
||||
v-if="row.message"
|
||||
class="mt-1 mb-0 line-clamp-1 text-cp-xs leading-[1.45] font-emphasis text-cp-text-secondary"
|
||||
|
||||
@@ -1 +1 @@
|
||||
version: "3.3.20"
|
||||
version: "3.3.21"
|
||||
|
||||
Reference in New Issue
Block a user