From 58f8c45278cf35b7c5f8e97d52a39b954ad5fb93 Mon Sep 17 00:00:00 2001 From: Albert Zhang Date: Sun, 2 Aug 2026 11:01:22 +0800 Subject: [PATCH 1/4] feat(codex): isolate agent continuation owners Co-Authored-By: Claude --- src/lib.rs | 1 + src/provider.rs | 12 + src/providers/codex/client.rs | 118 +++---- src/providers/codex/continuation.rs | 468 ++++++++++++++++++---------- src/providers/codex/mod.rs | 174 ++++++++--- src/providers/codex/websocket.rs | 321 +++++++++++-------- src/request_identity.rs | 338 ++++++++++++++++++++ src/server.rs | 61 +++- tests/server.rs | 189 ++++++++++- 9 files changed, 1271 insertions(+), 411 deletions(-) create mode 100644 src/request_identity.rs diff --git a/src/lib.rs b/src/lib.rs index 0bfa1ce8..4d8f86b7 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -9,6 +9,7 @@ pub mod project; pub mod provider; pub mod providers; pub mod registry; +pub mod request_identity; pub mod retry; pub mod server; pub mod session; diff --git a/src/provider.rs b/src/provider.rs index 2e57fbc5..2631ce93 100644 --- a/src/provider.rs +++ b/src/provider.rs @@ -1,5 +1,6 @@ use crate::anthropic::schema::MessagesRequest; use crate::monitor::MonitorHandle; +use crate::request_identity::ConversationIdentity; use crate::traffic::TrafficCapture; use anyhow::Result; use async_trait::async_trait; @@ -26,6 +27,17 @@ pub trait Provider: Send + Sync { fn supported_models(&self) -> Vec; fn cli(&self) -> &'static dyn CliHandlers; async fn handle_messages(&self, body: MessagesRequest, ctx: RequestContext) -> Response; + + async fn handle_messages_with_conversation_identity( + &self, + body: MessagesRequest, + ctx: RequestContext, + conversation_identity: Option, + ) -> Response { + let _ = conversation_identity; + self.handle_messages(body, ctx).await + } + async fn handle_count_tokens(&self, body: MessagesRequest, ctx: RequestContext) -> Response; async fn generate_anthropic_stream( diff --git a/src/providers/codex/client.rs b/src/providers/codex/client.rs index fcbfe33f..ed12344e 100644 --- a/src/providers/codex/client.rs +++ b/src/providers/codex/client.rs @@ -5,6 +5,7 @@ use crate::anthropic::sse::parse_sse_events; use crate::config; use crate::logging::create_logger; use crate::provider::RequestContext; +use crate::request_identity::ConversationIdentity; use crate::retry::{compute_backoff_delay, should_retry_status, sleep}; use crate::traffic::TrafficCapture; @@ -964,12 +965,12 @@ impl CodexHttpClient { origin: CodexErrorOrigin::Auth, })?; - let initial_pool_key = websocket_pool_key(ctx, continuation); + let initial_pool_owner = websocket_pool_owner(continuation).cloned(); if should_reset_websocket_pool(continuation) - && let Some(key) = initial_pool_key + && let Some(owner) = initial_pool_owner.as_ref() { super::websocket::invalidate_codex_websocket_pool_turn( - key, + owner, continuation.and_then(|candidate| candidate.turn_id), ); } @@ -979,7 +980,7 @@ impl CodexHttpClient { let mut auth_refresh_attempted = false; let mut transport_failures = 0u32; loop { - let pool_key = websocket_pool_key(ctx, active_continuation.as_ref()); + let pool_owner = websocket_pool_owner(active_continuation.as_ref()).cloned(); let result = match transport { CodexTransport::Http => { let body_json = serde_json::to_string(body).map_err(|e| CodexError { @@ -1006,7 +1007,7 @@ impl CodexHttpClient { &ws_body, ctx, ctx.traffic.as_deref(), - pool_key, + pool_owner.as_ref(), super::websocket::WEBSOCKET_CONNECT_TIMEOUT_MS, super::websocket::WEBSOCKET_IDLE_TIMEOUT_MS, active_continuation.as_ref(), @@ -1028,7 +1029,7 @@ impl CodexHttpClient { &ws_body, ctx, ctx.traffic.as_deref(), - pool_key, + pool_owner.as_ref(), super::websocket::WEBSOCKET_CONNECT_TIMEOUT_MS, super::websocket::WEBSOCKET_IDLE_TIMEOUT_MS, active_continuation.as_ref(), @@ -1067,8 +1068,8 @@ impl CodexHttpClient { match self.auth_manager.force_refresh(&auth.access).await { Ok(new_auth) => { auth = new_auth; - if let Some(key) = pool_key { - super::websocket::invalidate_codex_websocket_pool_turn(key, turn_id); + if let Some(owner) = pool_owner.as_ref() { + super::websocket::invalidate_codex_websocket_pool_turn(owner, turn_id); } active_continuation = full_context_continuation(active_continuation.as_ref()); @@ -1245,8 +1246,8 @@ impl CodexHttpClient { Err(err) if should_retry_without_continuation(&err, active_continuation.as_ref()) => { - if let Some(key) = pool_key { - super::websocket::invalidate_codex_websocket_pool_turn(key, turn_id); + if let Some(owner) = pool_owner.as_ref() { + super::websocket::invalidate_codex_websocket_pool_turn(owner, turn_id); } active_continuation = full_context_continuation(active_continuation.as_ref()); continue; @@ -1303,11 +1304,11 @@ impl CodexHttpClient { })?; let turn_id = continuation.and_then(|candidate| candidate.turn_id); - let pool_key = websocket_pool_key(ctx, continuation).map(str::to_string); + let pool_owner = websocket_pool_owner(continuation).cloned(); if should_reset_websocket_pool(continuation) - && let Some(key) = pool_key.as_deref() + && let Some(owner) = pool_owner.as_ref() { - super::websocket::invalidate_codex_websocket_pool_turn(key, turn_id); + super::websocket::invalidate_codex_websocket_pool_turn(owner, turn_id); } let client = self.clone(); @@ -1317,7 +1318,7 @@ impl CodexHttpClient { let (tx, rx) = tokio::sync::mpsc::channel(64); tokio::spawn(async move { client - .coordinate_live_websocket_events(body, ctx, continuation, auth, pool_key, tx) + .coordinate_live_websocket_events(body, ctx, continuation, auth, pool_owner, tx) .await; }); @@ -1330,12 +1331,15 @@ impl CodexHttpClient { ctx: RequestContext, mut continuation: Option, mut auth: StoredAuth, - pool_key: Option, + pool_owner: Option, tx: tokio::sync::mpsc::Sender>, ) { let turn_id = continuation .as_ref() .and_then(|candidate| candidate.turn_id); + let continuation_owner = continuation + .as_ref() + .and_then(|candidate| candidate.owner.clone()); let mut auth_refresh_attempted = false; let mut continuation_retry_available = continuation .as_ref() @@ -1361,16 +1365,16 @@ impl CodexHttpClient { &ws_body, &ctx, ctx.traffic.clone(), - pool_key.as_deref(), + pool_owner.as_ref(), super::websocket::WEBSOCKET_CONNECT_TIMEOUT_MS, super::websocket::WEBSOCKET_IDLE_TIMEOUT_MS, continuation.as_ref(), ); let mut stream = tokio::select! { _ = tx.closed() => { - super::continuation::abort_continuation(ctx.session_id.as_deref(), turn_id); - if let Some(key) = pool_key.as_deref() { - super::websocket::invalidate_codex_websocket_pool_turn(key, turn_id); + super::continuation::abort_continuation(continuation_owner.as_ref(), turn_id); + if let Some(owner) = pool_owner.as_ref() { + super::websocket::invalidate_codex_websocket_pool_turn(owner, turn_id); } return; } @@ -1378,8 +1382,8 @@ impl CodexHttpClient { Ok(stream) => stream, Err(err) if err.status == 401 && !auth_refresh_attempted && !forwarded_any => { auth_refresh_attempted = true; - if let Some(key) = pool_key.as_deref() { - super::websocket::invalidate_codex_websocket_pool_turn(key, turn_id); + if let Some(owner) = pool_owner.as_ref() { + super::websocket::invalidate_codex_websocket_pool_turn(owner, turn_id); } let refresh = self.auth_manager.force_refresh(&auth.access); auth = match refresh.await { @@ -1399,8 +1403,8 @@ impl CodexHttpClient { } Err(err) if continuation_retry_available && is_continuation_retry_error(&err) => { continuation_retry_available = false; - if let Some(key) = pool_key.as_deref() { - super::websocket::invalidate_codex_websocket_pool_turn(key, turn_id); + if let Some(owner) = pool_owner.as_ref() { + super::websocket::invalidate_codex_websocket_pool_turn(owner, turn_id); } continuation = full_context_continuation(continuation.as_ref()); continue 'attempt; @@ -1415,8 +1419,8 @@ impl CodexHttpClient { loop { let item = tokio::select! { _ = tx.closed() => { - if let Some(key) = pool_key.as_deref() { - super::websocket::invalidate_codex_websocket_pool_turn(key, turn_id); + if let Some(owner) = pool_owner.as_ref() { + super::websocket::invalidate_codex_websocket_pool_turn(owner, turn_id); } return; } @@ -1432,8 +1436,8 @@ impl CodexHttpClient { }; if unauthorized && !auth_refresh_attempted && !forwarded_any { auth_refresh_attempted = true; - if let Some(key) = pool_key.as_deref() { - super::websocket::invalidate_codex_websocket_pool_turn(key, turn_id); + if let Some(owner) = pool_owner.as_ref() { + super::websocket::invalidate_codex_websocket_pool_turn(owner, turn_id); } let refresh = self.auth_manager.force_refresh(&auth.access); auth = match refresh.await { @@ -1457,8 +1461,8 @@ impl CodexHttpClient { && !forwarded_any { continuation_retry_available = false; - if let Some(key) = pool_key.as_deref() { - super::websocket::invalidate_codex_websocket_pool_turn(key, turn_id); + if let Some(owner) = pool_owner.as_ref() { + super::websocket::invalidate_codex_websocket_pool_turn(owner, turn_id); } continuation = full_context_continuation(continuation.as_ref()); continue 'attempt; @@ -1468,9 +1472,9 @@ impl CodexHttpClient { forwarded_any = true; } if tx.send(item).await.is_err() { - super::continuation::abort_continuation(ctx.session_id.as_deref(), turn_id); - if let Some(key) = pool_key.as_deref() { - super::websocket::invalidate_codex_websocket_pool_turn(key, turn_id); + super::continuation::abort_continuation(continuation_owner.as_ref(), turn_id); + if let Some(owner) = pool_owner.as_ref() { + super::websocket::invalidate_codex_websocket_pool_turn(owner, turn_id); } return; } @@ -2024,6 +2028,7 @@ fn full_context_continuation( continuation: Option<&super::continuation::ContinuationCandidate>, ) -> Option { continuation.map(|candidate| super::continuation::ContinuationCandidate { + owner: candidate.owner.clone(), turn_id: candidate.turn_id, previous_response_id: None, input_delta: None, @@ -2048,16 +2053,14 @@ fn is_continuation_retry_error(err: &CodexError) -> bool { ) } -fn websocket_pool_key<'a>( - ctx: &'a RequestContext, +fn websocket_pool_owner( continuation: Option<&super::continuation::ContinuationCandidate>, -) -> Option<&'a str> { - let session_id = ctx.session_id.as_deref()?; +) -> Option<&ConversationIdentity> { let continuation = continuation?; if continuation.disabled_reason.as_deref() == Some("disabled") { return None; } - Some(session_id) + continuation.owner.as_ref() } fn should_reset_websocket_pool( @@ -3048,16 +3051,10 @@ mod tests { } #[test] - fn websocket_pool_key_tracks_continuation_opt_in() { - let ctx = RequestContext { - req_id: "r".into(), - session_id: Some("session".into()), - session_seq: None, - provider: "codex".into(), - traffic: None, - monitor: None, - }; + fn websocket_pool_owner_tracks_typed_continuation_opt_in() { + let owner = ConversationIdentity::Agent("session".into(), "agent".into()); let disabled = super::super::continuation::ContinuationCandidate { + owner: Some(owner.clone()), turn_id: None, previous_response_id: None, input_delta: None, @@ -3065,31 +3062,40 @@ mod tests { disabled_reason: Some("disabled".into()), }; let first_enabled = super::super::continuation::ContinuationCandidate { - turn_id: None, + owner: Some(owner.clone()), + turn_id: Some(1), previous_response_id: None, input_delta: None, input_delta_count: 1, disabled_reason: Some("missing_state".into()), }; let append = super::super::continuation::ContinuationCandidate { - turn_id: None, + owner: Some(owner.clone()), + turn_id: Some(2), previous_response_id: Some("resp_1".into()), input_delta: None, input_delta_count: 1, disabled_reason: None, }; + let missing_identity = super::super::continuation::ContinuationCandidate { + owner: None, + turn_id: None, + previous_response_id: None, + input_delta: None, + input_delta_count: 1, + disabled_reason: Some("missing_identity".into()), + }; - assert_eq!(websocket_pool_key(&ctx, Some(&disabled)), None); - assert_eq!( - websocket_pool_key(&ctx, Some(&first_enabled)), - Some("session") - ); - assert_eq!(websocket_pool_key(&ctx, Some(&append)), Some("session")); + assert_eq!(websocket_pool_owner(Some(&disabled)), None); + assert_eq!(websocket_pool_owner(Some(&first_enabled)), Some(&owner)); + assert_eq!(websocket_pool_owner(Some(&append)), Some(&owner)); + assert_eq!(websocket_pool_owner(Some(&missing_identity)), None); } #[test] fn websocket_pool_reset_clears_initial_stale_state() { let missing_state = super::super::continuation::ContinuationCandidate { + owner: Some(ConversationIdentity::Main("session".into())), turn_id: None, previous_response_id: None, input_delta: None, @@ -3097,6 +3103,7 @@ mod tests { disabled_reason: Some("missing_state".into()), }; let disabled = super::super::continuation::ContinuationCandidate { + owner: Some(ConversationIdentity::Main("session".into())), turn_id: None, previous_response_id: None, input_delta: None, @@ -3104,6 +3111,7 @@ mod tests { disabled_reason: Some("disabled".into()), }; let prompt_changed = super::super::continuation::ContinuationCandidate { + owner: Some(ConversationIdentity::Main("session".into())), turn_id: None, previous_response_id: None, input_delta: None, @@ -3239,6 +3247,7 @@ mod tests { #[test] fn continuation_retry_requires_previous_response_id() { let append = super::super::continuation::ContinuationCandidate { + owner: Some(ConversationIdentity::Main("session".into())), turn_id: None, previous_response_id: Some("resp_1".into()), input_delta: None, @@ -3246,6 +3255,7 @@ mod tests { disabled_reason: None, }; let initial = super::super::continuation::ContinuationCandidate { + owner: Some(ConversationIdentity::Main("session".into())), turn_id: None, previous_response_id: None, input_delta: None, diff --git a/src/providers/codex/continuation.rs b/src/providers/codex/continuation.rs index 99faa761..9c45e64b 100644 --- a/src/providers/codex/continuation.rs +++ b/src/providers/codex/continuation.rs @@ -2,11 +2,13 @@ use std::collections::HashMap; use std::sync::Mutex; use std::sync::atomic::{AtomicU64, Ordering}; +use crate::request_identity::ConversationIdentity; + use super::translate::request::{ResponsesInputItem, ResponsesRequest}; const TTL_MS: u64 = 30 * 60 * 1000; const MAX_STATES: usize = 10_000; -const MAX_SESSION_TRANSCRIPT_BYTES: u64 = 2_000_000; +const MAX_OWNER_TRANSCRIPT_BYTES: u64 = 2_000_000; const MAX_TOTAL_TRANSCRIPT_BYTES: u64 = 20_000_000; #[derive(Clone)] @@ -18,7 +20,7 @@ struct ContinuationState { updated_at: u64, } -struct SessionState { +struct OwnerState { current_turn: u64, continuation: Option, updated_at: u64, @@ -26,7 +28,7 @@ struct SessionState { #[derive(Default)] struct ContinuationRegistry { - sessions: HashMap, + owners: HashMap, total_transcript_bytes: u64, } @@ -35,6 +37,7 @@ static NEXT_TURN_ID: AtomicU64 = AtomicU64::new(1); #[derive(Clone)] pub struct ContinuationCandidate { + pub owner: Option, pub turn_id: Option, pub previous_response_id: Option, pub input_delta: Option>, @@ -50,12 +53,13 @@ fn now_ms() -> u64 { } pub fn continuation_candidate( - session_id: Option<&str>, + owner: Option<&ConversationIdentity>, body: &ResponsesRequest, enabled: bool, ) -> ContinuationCandidate { if !enabled { return ContinuationCandidate { + owner: owner.cloned(), turn_id: None, previous_response_id: None, input_delta: None, @@ -64,13 +68,14 @@ pub fn continuation_candidate( }; } - let Some(session_id) = session_id else { + let Some(owner) = owner else { return ContinuationCandidate { + owner: None, turn_id: None, previous_response_id: None, input_delta: None, input_delta_count: body.input.len(), - disabled_reason: Some("missing_session".to_string()), + disabled_reason: Some("missing_identity".to_string()), }; }; @@ -79,17 +84,17 @@ pub fn continuation_candidate( let (state, superseded_turn) = { let mut guard = REGISTRY.lock().unwrap(); let registry = guard.get_or_insert_with(ContinuationRegistry::default); - let existing = registry.sessions.remove(session_id); + let existing = registry.owners.remove(owner); let superseded_turn = existing.is_some(); - let state = existing.and_then(|session| session.continuation); + let state = existing.and_then(|owner| owner.continuation); if let Some(state) = &state { registry.total_transcript_bytes = registry .total_transcript_bytes .saturating_sub(state.transcript_bytes); } - registry.sessions.insert( - session_id.to_string(), - SessionState { + registry.owners.insert( + owner.clone(), + OwnerState { current_turn: turn_id, continuation: None, updated_at: now, @@ -99,10 +104,11 @@ pub fn continuation_candidate( (state, superseded_turn) }; - continuation_candidate_from_state(turn_id, body, state, superseded_turn, now) + continuation_candidate_from_state(owner, turn_id, body, state, superseded_turn, now) } fn continuation_candidate_from_state( + owner: &ConversationIdentity, turn_id: u64, body: &ResponsesRequest, state: Option, @@ -113,6 +119,7 @@ fn continuation_candidate_from_state( Some(state) if now.saturating_sub(state.updated_at) <= TTL_MS => state, Some(_) | None => { return ContinuationCandidate { + owner: Some(owner.clone()), turn_id: Some(turn_id), previous_response_id: None, input_delta: None, @@ -129,6 +136,7 @@ fn continuation_candidate_from_state( let signature = prompt_signature(body); if signature != state.prompt_signature { return ContinuationCandidate { + owner: Some(owner.clone()), turn_id: Some(turn_id), previous_response_id: None, input_delta: None, @@ -139,6 +147,7 @@ fn continuation_candidate_from_state( let Some(suffix) = input_suffix_after_prefix(&body.input, &state.transcript) else { return ContinuationCandidate { + owner: Some(owner.clone()), turn_id: Some(turn_id), previous_response_id: None, input_delta: None, @@ -149,6 +158,7 @@ fn continuation_candidate_from_state( if suffix.is_empty() { return ContinuationCandidate { + owner: Some(owner.clone()), turn_id: Some(turn_id), previous_response_id: None, input_delta: None, @@ -158,6 +168,7 @@ fn continuation_candidate_from_state( } ContinuationCandidate { + owner: Some(owner.clone()), turn_id: Some(turn_id), previous_response_id: Some(state.response_id), input_delta_count: suffix.len(), @@ -167,21 +178,21 @@ fn continuation_candidate_from_state( } pub fn record_continuation( - session_id: Option<&str>, + owner: Option<&ConversationIdentity>, turn_id: Option, request_body: &ResponsesRequest, response_id: Option<&str>, output_items: &[ResponsesInputItem], ) { - let (session_id, turn_id) = match (session_id, turn_id) { - (Some(session_id), Some(turn_id)) => (session_id, turn_id), + let (owner, turn_id) = match (owner, turn_id) { + (Some(owner), Some(turn_id)) => (owner, turn_id), _ => return, }; let response_id = match response_id { Some(id) => id.to_string(), None => { - abort_continuation(Some(session_id), Some(turn_id)); + abort_continuation(Some(owner), Some(turn_id)); return; } }; @@ -192,8 +203,8 @@ pub fn record_continuation( let transcript_json = serde_json::to_string(&transcript).unwrap_or_default(); let transcript_bytes = transcript_json.len() as u64; - if transcript_bytes > MAX_SESSION_TRANSCRIPT_BYTES { - abort_continuation(Some(session_id), Some(turn_id)); + if transcript_bytes > MAX_OWNER_TRANSCRIPT_BYTES { + abort_continuation(Some(owner), Some(turn_id)); return; } @@ -209,13 +220,13 @@ pub fn record_continuation( let Some(registry) = guard.as_mut() else { return; }; - let Some(session) = registry.sessions.get_mut(session_id) else { + let Some(owner_state) = registry.owners.get_mut(owner) else { return; }; - if session.current_turn != turn_id { + if owner_state.current_turn != turn_id { return; } - if let Some(existing) = session.continuation.replace(state) { + if let Some(existing) = owner_state.continuation.replace(state) { registry.total_transcript_bytes = registry .total_transcript_bytes .saturating_sub(existing.transcript_bytes); @@ -224,8 +235,8 @@ pub fn record_continuation( evict_oldest(registry); } -pub fn abort_continuation(session_id: Option<&str>, turn_id: Option) { - let (Some(session_id), Some(turn_id)) = (session_id, turn_id) else { +pub fn abort_continuation(owner: Option<&ConversationIdentity>, turn_id: Option) { + let (Some(owner), Some(turn_id)) = (owner, turn_id) else { return; }; let mut guard = REGISTRY.lock().unwrap(); @@ -233,76 +244,76 @@ pub fn abort_continuation(session_id: Option<&str>, turn_id: Option) { return; }; if registry - .sessions - .get(session_id) - .is_some_and(|session| session.current_turn == turn_id) - && let Some(session) = registry.sessions.remove(session_id) - && let Some(state) = session.continuation + .owners + .get(owner) + .is_some_and(|state| state.current_turn == turn_id) + && let Some(state) = registry.owners.remove(owner) + && let Some(continuation) = state.continuation { registry.total_transcript_bytes = registry .total_transcript_bytes - .saturating_sub(state.transcript_bytes); + .saturating_sub(continuation.transcript_bytes); } } pub fn if_current_turn( - session_id: Option<&str>, + owner: Option<&ConversationIdentity>, turn_id: Option, action: impl FnOnce() -> T, ) -> Option { - let (Some(session_id), Some(turn_id)) = (session_id, turn_id) else { - return Some(action()); + let (Some(owner), Some(turn_id)) = (owner, turn_id) else { + return None; }; let guard = REGISTRY.lock().unwrap(); let current = guard .as_ref() - .and_then(|registry| registry.sessions.get(session_id)) - .is_some_and(|session| session.current_turn == turn_id); + .and_then(|registry| registry.owners.get(owner)) + .is_some_and(|state| state.current_turn == turn_id); current.then(action) } pub fn with_current_turn( - session_id: Option<&str>, + owner: Option<&ConversationIdentity>, turn_id: Option, action: impl FnOnce(), ) -> bool { - if_current_turn(session_id, turn_id, action).is_some() + if_current_turn(owner, turn_id, action).is_some() } -pub fn is_current_turn(session_id: Option<&str>, turn_id: Option) -> bool { - let (Some(session_id), Some(turn_id)) = (session_id, turn_id) else { +pub fn is_current_turn(owner: Option<&ConversationIdentity>, turn_id: Option) -> bool { + let (Some(owner), Some(turn_id)) = (owner, turn_id) else { return false; }; let guard = REGISTRY.lock().unwrap(); guard .as_ref() - .and_then(|registry| registry.sessions.get(session_id)) - .is_some_and(|session| session.current_turn == turn_id) + .and_then(|registry| registry.owners.get(owner)) + .is_some_and(|state| state.current_turn == turn_id) } -pub fn clear_continuation(session_id: Option<&str>) { - let Some(session_id) = session_id else { +pub fn clear_continuation(owner: Option<&ConversationIdentity>) { + let Some(owner) = owner else { return; }; let mut guard = REGISTRY.lock().unwrap(); let Some(registry) = guard.as_mut() else { return; }; - if let Some(session) = registry.sessions.remove(session_id) - && let Some(state) = session.continuation + if let Some(state) = registry.owners.remove(owner) + && let Some(continuation) = state.continuation { registry.total_transcript_bytes = registry .total_transcript_bytes - .saturating_sub(state.transcript_bytes); + .saturating_sub(continuation.transcript_bytes); } } -pub fn has_continuation_for_tests(session_id: &str) -> bool { +pub fn has_continuation_for_tests(owner: &ConversationIdentity) -> bool { let guard = REGISTRY.lock().unwrap(); guard .as_ref() - .and_then(|registry| registry.sessions.get(session_id)) - .is_some_and(|session| session.continuation.is_some()) + .and_then(|registry| registry.owners.get(owner)) + .is_some_and(|state| state.continuation.is_some()) } pub fn clear_all_continuations_for_tests() { @@ -376,23 +387,23 @@ fn stable_json(value: &serde_json::Value) -> String { } fn evict_oldest(registry: &mut ContinuationRegistry) { - while registry.sessions.len() > MAX_STATES + while registry.owners.len() > MAX_STATES || registry.total_transcript_bytes > MAX_TOTAL_TRANSCRIPT_BYTES { - let key = registry - .sessions + let owner = registry + .owners .iter() - .min_by_key(|(_, session)| session.updated_at) - .map(|(key, _)| key.clone()); - let Some(key) = key else { + .min_by_key(|(_, state)| state.updated_at) + .map(|(owner, _)| owner.clone()); + let Some(owner) = owner else { break; }; - if let Some(session) = registry.sessions.remove(&key) - && let Some(state) = session.continuation + if let Some(state) = registry.owners.remove(&owner) + && let Some(continuation) = state.continuation { registry.total_transcript_bytes = registry .total_transcript_bytes - .saturating_sub(state.transcript_bytes); + .saturating_sub(continuation.transcript_bytes); } } } @@ -402,6 +413,33 @@ mod tests { use super::*; use serde_json::json; + static TEST_REGISTRY_LOCK: Mutex<()> = Mutex::new(()); + + fn lock_registry() -> std::sync::MutexGuard<'static, ()> { + let guard = TEST_REGISTRY_LOCK.lock().unwrap(); + clear_all_continuations_for_tests(); + guard + } + + fn main_owner(session_id: &str) -> ConversationIdentity { + ConversationIdentity::Main(session_id.to_string()) + } + + fn agent_owner(session_id: &str, agent_id: &str) -> ConversationIdentity { + ConversationIdentity::Agent(session_id.to_string(), agent_id.to_string()) + } + + fn input(text: &str) -> ResponsesInputItem { + ResponsesInputItem::Message { + role: "user".to_string(), + content: vec![ + super::super::translate::request::ResponsesContentPart::InputText { + text: text.to_string(), + }, + ], + } + } + fn request_with_input( input: Vec, extra: Option, @@ -416,142 +454,226 @@ mod tests { if let Some(extras) = extra && let Some(obj) = extras.as_object() { - for (k, v) in obj { - fields.insert(k.clone(), v.clone()); + for (key, value) in obj { + fields.insert(key.clone(), value.clone()); } } serde_json::from_value(serde_json::Value::Object(fields)).unwrap() } - fn start_and_record(session_id: &str, request: &ResponsesRequest, response_id: Option<&str>) { - let candidate = continuation_candidate(Some(session_id), request, true); + fn start_and_record( + owner: &ConversationIdentity, + request: &ResponsesRequest, + response_id: &str, + ) { + let candidate = continuation_candidate(Some(owner), request, true); record_continuation( - Some(session_id), + candidate.owner.as_ref(), candidate.turn_id, request, - response_id, + Some(response_id), &[], ); } #[test] - fn continuation_behaviors() { - // All tests run in sequence to avoid global state interference + fn disabled_and_missing_identity_requests_are_stateless() { + let _registry_guard = lock_registry(); + let request = request_with_input(vec![input("one")], None); + let owner = main_owner("session-a"); + + let disabled = continuation_candidate(Some(&owner), &request, false); + assert_eq!(disabled.owner.as_ref(), Some(&owner)); + assert_eq!(disabled.turn_id, None); + assert_eq!(disabled.input_delta_count, request.input.len()); + assert_eq!(disabled.disabled_reason.as_deref(), Some("disabled")); + + let missing = continuation_candidate(None, &request, true); + assert_eq!(missing.owner, None); + assert_eq!(missing.turn_id, None); + assert_eq!(missing.input_delta_count, request.input.len()); + assert_eq!(missing.disabled_reason.as_deref(), Some("missing_identity")); + } - // disabled_when_not_enabled - clear_all_continuations_for_tests(); - let input = vec![ResponsesInputItem::Message { - role: "user".to_string(), - content: vec![ - super::super::translate::request::ResponsesContentPart::InputText { - text: "one".to_string(), - }, - ], - }]; - let req = request_with_input(input, None); - let result = continuation_candidate(Some("s1"), &req, false); - assert_eq!(result.disabled_reason, Some("disabled".to_string())); - assert_eq!(result.input_delta_count, 1); + #[test] + fn sibling_agents_reserve_and_publish_independently() { + let _registry_guard = lock_registry(); + let sibling_one = agent_owner("session-a", "agent-one"); + let sibling_two = agent_owner("session-a", "agent-two"); + let first_request = request_with_input(vec![input("one")], None); + + let first = continuation_candidate(Some(&sibling_one), &first_request, true); + let second = continuation_candidate(Some(&sibling_two), &first_request, true); + assert_ne!(first.turn_id, second.turn_id); + record_continuation( + first.owner.as_ref(), + first.turn_id, + &first_request, + Some("resp_one"), + &[], + ); + record_continuation( + second.owner.as_ref(), + second.turn_id, + &first_request, + Some("resp_two"), + &[], + ); + assert!(has_continuation_for_tests(&sibling_one)); + assert!(has_continuation_for_tests(&sibling_two)); + + let next_request = request_with_input(vec![input("one"), input("two")], None); + let first_next = continuation_candidate(Some(&sibling_one), &next_request, true); + let second_next = continuation_candidate(Some(&sibling_two), &next_request, true); + assert_eq!(first_next.previous_response_id.as_deref(), Some("resp_one")); + assert_eq!( + second_next.previous_response_id.as_deref(), + Some("resp_two") + ); + } - // missing_session - clear_all_continuations_for_tests(); - let input = vec![ResponsesInputItem::Message { - role: "user".to_string(), - content: vec![ - super::super::translate::request::ResponsesContentPart::InputText { - text: "one".to_string(), - }, - ], - }]; - let req = request_with_input(input, None); - let result = continuation_candidate(None, &req, true); - assert_eq!(result.disabled_reason, Some("missing_session".to_string())); + #[test] + fn different_owner_completion_order_cannot_interfere() { + let _registry_guard = lock_registry(); + let main = main_owner("session-a"); + let agent = agent_owner("session-a", "agent-a"); + let request = request_with_input(vec![input("one")], None); + let main_candidate = continuation_candidate(Some(&main), &request, true); + let agent_candidate = continuation_candidate(Some(&agent), &request, true); - // uses_previous_response_id_for_append_only - clear_all_continuations_for_tests(); - let input = vec![ResponsesInputItem::Message { - role: "user".to_string(), - content: vec![ - super::super::translate::request::ResponsesContentPart::InputText { - text: "one".to_string(), - }, - ], - }]; - let req = request_with_input(input, None); - start_and_record("s1", &req, Some("resp_1")); - - let input2 = vec![ - ResponsesInputItem::Message { - role: "user".to_string(), - content: vec![ - super::super::translate::request::ResponsesContentPart::InputText { - text: "one".to_string(), - }, - ], - }, - ResponsesInputItem::Message { - role: "user".to_string(), - content: vec![ - super::super::translate::request::ResponsesContentPart::InputText { - text: "two".to_string(), - }, - ], - }, - ]; - let req2 = request_with_input(input2, None); - let result = continuation_candidate(Some("s1"), &req2, true); - assert_eq!(result.previous_response_id, Some("resp_1".to_string())); - assert_eq!(result.input_delta_count, 1); + record_continuation( + agent_candidate.owner.as_ref(), + agent_candidate.turn_id, + &request, + Some("resp_agent"), + &[], + ); + record_continuation( + main_candidate.owner.as_ref(), + main_candidate.turn_id, + &request, + Some("resp_main"), + &[], + ); + abort_continuation(Some(&main), agent_candidate.turn_id); + + assert!(has_continuation_for_tests(&main)); + assert!(has_continuation_for_tests(&agent)); + let next = request_with_input(vec![input("one"), input("two")], None); + assert_eq!( + continuation_candidate(Some(&agent), &next, true) + .previous_response_id + .as_deref(), + Some("resp_agent") + ); + } - // clears_state_when_prompt_signature_changes - clear_all_continuations_for_tests(); - let input = vec![ResponsesInputItem::Message { - role: "user".to_string(), - content: vec![ - super::super::translate::request::ResponsesContentPart::InputText { - text: "one".to_string(), - }, - ], - }]; - let req = request_with_input(input.clone(), None); - start_and_record("s1", &req, Some("resp_1")); + #[test] + fn missing_response_id_aborts_only_the_current_owner() { + let _registry_guard = lock_registry(); + let owner = main_owner("session-a"); + let sibling = agent_owner("session-a", "agent-a"); + let request = request_with_input(vec![input("one")], None); + start_and_record(&owner, &request, "resp_main"); + start_and_record(&sibling, &request, "resp_agent"); + + let candidate = continuation_candidate(Some(&owner), &request, true); + record_continuation( + candidate.owner.as_ref(), + candidate.turn_id, + &request, + None, + &[], + ); - let req2 = request_with_input(input, Some(json!({"service_tier": "flex"}))); - let result = continuation_candidate(Some("s1"), &req2, true); - assert_eq!(result.disabled_reason, Some("prompt_changed".to_string())); - assert!(!has_continuation_for_tests("s1")); + assert!(!has_continuation_for_tests(&owner)); + assert!(has_continuation_for_tests(&sibling)); + } - // clears_state_when_missing_response_id - clear_all_continuations_for_tests(); - let input = vec![ResponsesInputItem::Message { - role: "user".to_string(), - content: vec![ - super::super::translate::request::ResponsesContentPart::InputText { - text: "one".to_string(), - }, - ], - }]; - let req = request_with_input(input.clone(), None); - start_and_record("s1", &req, Some("resp_1")); - assert!(has_continuation_for_tests("s1")); + #[test] + fn same_owner_stale_turn_cannot_publish_clear_or_run_actions() { + let _registry_guard = lock_registry(); + let owner = main_owner("session-a"); + let request = request_with_input(vec![input("one")], None); + start_and_record(&owner, &request, "resp_1"); + + let stale = continuation_candidate(Some(&owner), &request, true); + let current = continuation_candidate(Some(&owner), &request, true); + assert_eq!(current.disabled_reason.as_deref(), Some("superseded_turn")); + record_continuation( + stale.owner.as_ref(), + stale.turn_id, + &request, + Some("resp_stale"), + &[], + ); + assert!(!has_continuation_for_tests(&owner)); + record_continuation( + current.owner.as_ref(), + current.turn_id, + &request, + Some("resp_current"), + &[], + ); + assert!(has_continuation_for_tests(&owner)); + abort_continuation(stale.owner.as_ref(), stale.turn_id); + assert!(has_continuation_for_tests(&owner)); + + let mut ran = false; + assert_eq!( + if_current_turn(stale.owner.as_ref(), stale.turn_id, || ran = true), + None + ); + assert!(!ran); + } + + #[test] + fn missing_owner_or_turn_mutations_are_hard_noops() { + let _registry_guard = lock_registry(); + let owner = main_owner("session-a"); + let request = request_with_input(vec![input("one")], None); + start_and_record(&owner, &request, "resp_1"); + + record_continuation(None, Some(1), &request, Some("ignored"), &[]); + record_continuation(Some(&owner), None, &request, Some("ignored"), &[]); + abort_continuation(None, Some(1)); + abort_continuation(Some(&owner), None); + clear_continuation(None); + assert!(has_continuation_for_tests(&owner)); + + let mut runs = 0; + assert_eq!(if_current_turn(None, Some(1), || runs += 1), None); + assert_eq!(if_current_turn(Some(&owner), None, || runs += 1), None); + assert!(!with_current_turn(None, Some(1), || runs += 1)); + assert!(!with_current_turn(Some(&owner), None, || runs += 1)); + assert_eq!(runs, 0); + } - let candidate = continuation_candidate(Some("s1"), &req, true); - record_continuation(Some("s1"), candidate.turn_id, &req, None, &[]); - assert!(!has_continuation_for_tests("s1")); + #[test] + fn append_only_and_prompt_guards_stay_owner_scoped() { + let _registry_guard = lock_registry(); + let owner = main_owner("session-a"); + let request = request_with_input(vec![input("one")], None); + start_and_record(&owner, &request, "resp_1"); - // stale turns cannot publish or clear a newer turn - clear_all_continuations_for_tests(); - let first = continuation_candidate(Some("s1"), &req, true); - record_continuation(Some("s1"), first.turn_id, &req, Some("resp_1"), &[]); - let second = continuation_candidate(Some("s1"), &req, true); - let third = continuation_candidate(Some("s1"), &req, true); - assert_eq!(third.disabled_reason.as_deref(), Some("superseded_turn")); - - record_continuation(Some("s1"), second.turn_id, &req, Some("resp_2"), &[]); - assert!(!has_continuation_for_tests("s1")); - record_continuation(Some("s1"), third.turn_id, &req, Some("resp_3"), &[]); - assert!(has_continuation_for_tests("s1")); - abort_continuation(Some("s1"), second.turn_id); - assert!(has_continuation_for_tests("s1")); + let appended = request_with_input(vec![input("one"), input("two")], None); + let candidate = continuation_candidate(Some(&owner), &appended, true); + assert_eq!(candidate.previous_response_id.as_deref(), Some("resp_1")); + assert_eq!(candidate.input_delta_count, 1); + + record_continuation( + candidate.owner.as_ref(), + candidate.turn_id, + &appended, + Some("resp_2"), + &[], + ); + let changed = request_with_input( + vec![input("one"), input("two"), input("three")], + Some(json!({"service_tier": "flex"})), + ); + let candidate = continuation_candidate(Some(&owner), &changed, true); + assert_eq!(candidate.disabled_reason.as_deref(), Some("prompt_changed")); + assert!(!has_continuation_for_tests(&owner)); } } diff --git a/src/providers/codex/mod.rs b/src/providers/codex/mod.rs index d2534250..ec40268a 100644 --- a/src/providers/codex/mod.rs +++ b/src/providers/codex/mod.rs @@ -30,6 +30,7 @@ use crate::logging::create_logger; use crate::monitor::usage_from_anthropic_sse; use crate::provider::{CliHandlers, Provider, RequestContext}; use crate::registry; +use crate::request_identity::ConversationIdentity; use crate::retry::{compute_backoff_delay, sleep}; use self::auth::browser_login::run_browser_login; @@ -87,30 +88,13 @@ impl CodexProvider { } } -#[async_trait] -impl Provider for CodexProvider { - fn name(&self) -> &'static str { - "codex" - } - - fn supported_models(&self) -> Vec { - let mut models: Vec = registry::CODEX_MODELS - .iter() - .map(|m| m.to_string()) - .collect(); - for m in registry::CODEX_MODELS { - models.push(format!("{m}-fast")); - } - models.sort_unstable(); - models.dedup(); - models - } - - fn cli(&self) -> &'static dyn CliHandlers { - &CODEX_CLI - } - - async fn handle_messages(&self, body: MessagesRequest, ctx: RequestContext) -> Response { +impl CodexProvider { + async fn handle_messages_inner( + &self, + body: MessagesRequest, + ctx: RequestContext, + conversation_identity: Option, + ) -> Response { let message_id = format!("msg_{}", uuid::Uuid::new_v4().to_string().replace('-', "")); let want_stream = body.stream; let model = body.model.as_deref().unwrap_or("gpt-5.6-sol"); @@ -292,10 +276,11 @@ impl Provider for CodexProvider { // Check continuation let previous_response_id_enabled = config::codex_previous_response_id(); let continuation = continuation_candidate( - ctx.session_id.as_deref(), + conversation_identity.as_ref(), &translated, previous_response_id_enabled, ); + let continuation_owner = continuation.owner.clone(); let turn_id = continuation.turn_id; // Post to upstream with continuation @@ -330,7 +315,7 @@ impl Provider for CodexProvider { Ok(r) => r, Err(e) => { abort_compaction_attempt(ctx.session_id.as_deref(), compaction_attempt); - abort_continuation(ctx.session_id.as_deref(), turn_id); + abort_continuation(continuation_owner.as_ref(), turn_id); return map_codex_error_to_response(&e); } }; @@ -343,13 +328,13 @@ impl Provider for CodexProvider { drop_live_continuation_for_retry(&mut continuation); if attempt >= MAX_EMPTY_COMPLETION_RETRIES { abort_compaction_attempt(ctx.session_id.as_deref(), compaction_attempt); - abort_continuation(ctx.session_id.as_deref(), turn_id); + abort_continuation(continuation_owner.as_ref(), turn_id); return map_codex_error_to_response(&error); } let delay = compute_backoff_delay(attempt, None); if delay.exceeds_budget { abort_compaction_attempt(ctx.session_id.as_deref(), compaction_attempt); - abort_continuation(ctx.session_id.as_deref(), turn_id); + abort_continuation(continuation_owner.as_ref(), turn_id); return map_codex_error_to_response(&error); } attempt += 1; @@ -368,7 +353,7 @@ impl Provider for CodexProvider { Ok(b) => b, Err(e) => { abort_compaction_attempt(ctx.session_id.as_deref(), compaction_attempt); - abort_continuation(ctx.session_id.as_deref(), turn_id); + abort_continuation(continuation_owner.as_ref(), turn_id); return map_codex_failure_to_response(&format!( "Stream translation error: {e}" )); @@ -386,6 +371,7 @@ impl Provider for CodexProvider { } update_continuation_from_upstream( ctx.session_id.as_deref(), + continuation_owner.as_ref(), turn_id, compaction_attempt, &translated, @@ -417,6 +403,7 @@ impl Provider for CodexProvider { } update_continuation_from_upstream( ctx.session_id.as_deref(), + continuation_owner.as_ref(), turn_id, compaction_attempt, &translated, @@ -427,12 +414,50 @@ impl Provider for CodexProvider { } Err(e) => { abort_compaction_attempt(ctx.session_id.as_deref(), compaction_attempt); - abort_continuation(ctx.session_id.as_deref(), turn_id); + abort_continuation(continuation_owner.as_ref(), turn_id); map_codex_failure_to_response(&format!("Accumulation error: {e}")) } } } } +} + +#[async_trait] +impl Provider for CodexProvider { + fn name(&self) -> &'static str { + "codex" + } + + fn supported_models(&self) -> Vec { + let mut models: Vec = registry::CODEX_MODELS + .iter() + .map(|m| m.to_string()) + .collect(); + for m in registry::CODEX_MODELS { + models.push(format!("{m}-fast")); + } + models.sort_unstable(); + models.dedup(); + models + } + + fn cli(&self) -> &'static dyn CliHandlers { + &CODEX_CLI + } + + async fn handle_messages(&self, body: MessagesRequest, ctx: RequestContext) -> Response { + self.handle_messages_inner(body, ctx, None).await + } + + async fn handle_messages_with_conversation_identity( + &self, + body: MessagesRequest, + ctx: RequestContext, + conversation_identity: Option, + ) -> Response { + self.handle_messages_inner(body, ctx, conversation_identity) + .await + } async fn handle_count_tokens(&self, body: MessagesRequest, ctx: RequestContext) -> Response { let model = body.model.as_deref().unwrap_or("gpt-5.6-sol"); @@ -521,11 +546,12 @@ fn log_compaction_event( fn abort_request_state( session_id: Option<&str>, + continuation_owner: Option<&ConversationIdentity>, turn_id: Option, compaction_attempt: Option, ) { abort_compaction_attempt(session_id, compaction_attempt); - abort_continuation(session_id, turn_id); + abort_continuation(continuation_owner, turn_id); } enum LiveStreamStart { @@ -549,6 +575,7 @@ async fn live_stream_response( compaction: LiveStreamCompaction, ) -> Response { let model = model.to_string(); + let continuation_owner = continuation.owner.clone(); let turn_id = continuation.turn_id; let mut attempt = 0_u32; let mut continuation = Some(continuation); @@ -566,12 +593,22 @@ async fn live_stream_response( continue; } if attempt >= MAX_RETRYABLE_LIVE_STREAM_RETRIES { - abort_request_state(ctx.session_id.as_deref(), turn_id, compaction.attempt); + abort_request_state( + ctx.session_id.as_deref(), + continuation_owner.as_ref(), + turn_id, + compaction.attempt, + ); return map_codex_error_to_response(&err); } let delay = compute_backoff_delay(attempt, err.retry_after.as_deref()); if delay.exceeds_budget { - abort_request_state(ctx.session_id.as_deref(), turn_id, compaction.attempt); + abort_request_state( + ctx.session_id.as_deref(), + continuation_owner.as_ref(), + turn_id, + compaction.attempt, + ); return map_codex_error_to_response(&err); } attempt += 1; @@ -579,7 +616,12 @@ async fn live_stream_response( continue; } Err(err) => { - abort_request_state(ctx.session_id.as_deref(), turn_id, compaction.attempt); + abort_request_state( + ctx.session_id.as_deref(), + continuation_owner.as_ref(), + turn_id, + compaction.attempt, + ); return map_codex_error_to_response(&err); } }; @@ -589,6 +631,7 @@ async fn live_stream_response( message_id.clone(), &model, ctx.clone(), + continuation_owner.clone(), turn_id, request_body.clone(), compaction, @@ -603,12 +646,22 @@ async fn live_stream_response( continue; } if attempt >= MAX_RETRYABLE_LIVE_STREAM_RETRIES { - abort_request_state(ctx.session_id.as_deref(), turn_id, compaction.attempt); + abort_request_state( + ctx.session_id.as_deref(), + continuation_owner.as_ref(), + turn_id, + compaction.attempt, + ); return map_codex_error_to_response(&error); } let delay = compute_backoff_delay(attempt, error.retry_after.as_deref()); if delay.exceeds_budget { - abort_request_state(ctx.session_id.as_deref(), turn_id, compaction.attempt); + abort_request_state( + ctx.session_id.as_deref(), + continuation_owner.as_ref(), + turn_id, + compaction.attempt, + ); return map_codex_error_to_response(&error); } attempt += 1; @@ -618,11 +671,13 @@ async fn live_stream_response( } } +#[allow(clippy::too_many_arguments)] async fn live_stream_response_once( mut upstream_events: websocket::CodexWebSocketEventReceiver, message_id: String, model: &str, ctx: RequestContext, + continuation_owner: Option, turn_id: Option, request_body: translate::request::ResponsesRequest, compaction: LiveStreamCompaction, @@ -646,7 +701,12 @@ async fn live_stream_response_once( if retryable_live_start_codex_error(&err) { return LiveStreamStart::Retry { error: err }; } - abort_request_state(ctx.session_id.as_deref(), turn_id, compaction.attempt); + abort_request_state( + ctx.session_id.as_deref(), + continuation_owner.as_ref(), + turn_id, + compaction.attempt, + ); return LiveStreamStart::Response(map_codex_error_to_response(&err)); } }; @@ -694,7 +754,12 @@ async fn live_stream_response_once( }, }; } - abort_request_state(ctx.session_id.as_deref(), turn_id, compaction.attempt); + abort_request_state( + ctx.session_id.as_deref(), + continuation_owner.as_ref(), + turn_id, + compaction.attempt, + ); return LiveStreamStart::Response(map_codex_failure_to_response(&message)); } }; @@ -713,6 +778,7 @@ async fn live_stream_response_once( if terminal { update_continuation_from_upstream( ctx.session_id.as_deref(), + continuation_owner.as_ref(), turn_id, compaction.attempt, &request_body, @@ -726,6 +792,7 @@ async fn live_stream_response_once( translator, pending_chunk, ctx, + continuation_owner, turn_id, request_body, upstream_sse_body, @@ -735,6 +802,7 @@ async fn live_stream_response_once( if terminal { update_continuation_from_upstream( ctx.session_id.as_deref(), + continuation_owner.as_ref(), turn_id, compaction.attempt, &request_body, @@ -835,6 +903,7 @@ fn remaining_live_stream_response( mut translator: LiveStreamTranslator, first_chunk: Vec, ctx: RequestContext, + continuation_owner: Option, turn_id: Option, request_body: translate::request::ResponsesRequest, mut upstream_sse_body: Vec, @@ -843,7 +912,12 @@ fn remaining_live_stream_response( let (tx, rx) = tokio::sync::mpsc::channel::>(64); tokio::spawn(async move { if tx.send(Ok(Bytes::from(first_chunk))).await.is_err() { - abort_request_state(ctx.session_id.as_deref(), turn_id, compaction.attempt); + abort_request_state( + ctx.session_id.as_deref(), + continuation_owner.as_ref(), + turn_id, + compaction.attempt, + ); return; } while let Some(item) = upstream_events.recv().await { @@ -859,6 +933,7 @@ fn remaining_live_stream_response( Err(message) => { abort_request_state( ctx.session_id.as_deref(), + continuation_owner.as_ref(), turn_id, compaction.attempt, ); @@ -879,6 +954,7 @@ fn remaining_live_stream_response( if tx.send(Ok(Bytes::from(chunk))).await.is_err() { abort_request_state( ctx.session_id.as_deref(), + continuation_owner.as_ref(), turn_id, compaction.attempt, ); @@ -888,6 +964,7 @@ fn remaining_live_stream_response( if terminal { update_continuation_from_upstream( ctx.session_id.as_deref(), + continuation_owner.as_ref(), turn_id, compaction.attempt, &request_body, @@ -898,7 +975,12 @@ fn remaining_live_stream_response( } } Err(err) => { - abort_request_state(ctx.session_id.as_deref(), turn_id, compaction.attempt); + abort_request_state( + ctx.session_id.as_deref(), + continuation_owner.as_ref(), + turn_id, + compaction.attempt, + ); let chunk = translator.finish_after_closed_completed_tool_call(ctx.traffic.as_deref()); if !chunk.is_empty() { @@ -921,7 +1003,12 @@ fn remaining_live_stream_response( } } - abort_request_state(ctx.session_id.as_deref(), turn_id, compaction.attempt); + abort_request_state( + ctx.session_id.as_deref(), + continuation_owner.as_ref(), + turn_id, + compaction.attempt, + ); let chunk = translator.finish_after_closed_completed_tool_call(ctx.traffic.as_deref()); if !chunk.is_empty() { record_live_stream_progress(&ctx, &chunk); @@ -1096,6 +1183,7 @@ fn codex_stream_error_type(err: &client::CodexError) -> &'static str { fn update_continuation_from_upstream( session_id: Option<&str>, + continuation_owner: Option<&ConversationIdentity>, turn_id: Option, compaction_attempt: Option, request_body: &translate::request::ResponsesRequest, @@ -1113,7 +1201,7 @@ fn update_continuation_from_upstream( ); } record_continuation( - session_id, + continuation_owner, turn_id, request_body, finish.response_id.as_deref(), @@ -1122,7 +1210,7 @@ fn update_continuation_from_upstream( } _ => { abort_compaction_attempt(session_id, compaction_attempt); - abort_continuation(session_id, turn_id); + abort_continuation(continuation_owner, turn_id); } } } diff --git a/src/providers/codex/websocket.rs b/src/providers/codex/websocket.rs index cfa919e7..20f6d820 100644 --- a/src/providers/codex/websocket.rs +++ b/src/providers/codex/websocket.rs @@ -24,6 +24,7 @@ use tokio_tungstenite::{ use crate::logging::create_logger; use crate::provider::RequestContext; +use crate::request_identity::ConversationIdentity; use crate::retry::sleep as retry_sleep; use crate::traffic::TrafficCapture; @@ -242,7 +243,7 @@ impl PoolEntry { } static POOL_ACTIVITY_SEQUENCE: AtomicU64 = AtomicU64::new(1); -static WS_POOL: once_cell::sync::Lazy>>> = +static WS_POOL: once_cell::sync::Lazy>>> = once_cell::sync::Lazy::new(|| Mutex::new(HashMap::new())); static WS_CONNECT_GATE: once_cell::sync::Lazy = once_cell::sync::Lazy::new(|| WebSocketConnectGate::new(WEBSOCKET_CONNECT_START_SPACING)); @@ -263,73 +264,80 @@ pub fn clear_codex_websocket_pool_for_tests() { guard.clear(); } -pub fn invalidate_codex_websocket_pool_key(session_id: &str) { +pub fn invalidate_codex_websocket_pool_owner(owner: &ConversationIdentity) { let mut guard = WS_POOL.lock().unwrap(); - guard.remove(session_id); + guard.remove(owner); } -pub fn invalidate_codex_websocket_pool_turn(session_id: &str, turn_id: Option) { - super::continuation::with_current_turn(Some(session_id), turn_id, || { - invalidate_codex_websocket_pool_key(session_id) +pub fn invalidate_codex_websocket_pool_turn(owner: &ConversationIdentity, turn_id: Option) { + super::continuation::with_current_turn(Some(owner), turn_id, || { + invalidate_codex_websocket_pool_owner(owner) }); } -fn invalidate_pool_entry(session_id: &str, entry: &Arc) { +fn invalidate_pool_entry(owner: &ConversationIdentity, entry: &Arc) { let mut guard = WS_POOL.lock().unwrap(); if guard - .get(session_id) + .get(owner) .is_some_and(|pooled| Arc::ptr_eq(pooled, entry)) { - guard.remove(session_id); + guard.remove(owner); } } -fn invalidate_pool_owner(pool_key: Option<&str>, entry: Option<&Arc>) { - let Some(session_id) = pool_key else { +fn invalidate_pool_owner(owner: Option<&ConversationIdentity>, entry: Option<&Arc>) { + let Some(owner) = owner else { return; }; match entry { - Some(entry) => invalidate_pool_entry(session_id, entry), - None => invalidate_codex_websocket_pool_key(session_id), + Some(entry) => invalidate_pool_entry(owner, entry), + None => invalidate_codex_websocket_pool_owner(owner), } } -fn pool_get_for_turn(key: &str, turn_id: Option) -> Option> { - super::continuation::if_current_turn(Some(key), turn_id, || { - let entry = WS_POOL.lock().ok()?.get(key).cloned()?; +fn pool_get_for_turn(owner: &ConversationIdentity, turn_id: Option) -> Option> { + super::continuation::if_current_turn(Some(owner), turn_id, || { + let entry = WS_POOL.lock().ok()?.get(owner).cloned()?; entry.touch(); Some(entry) }) .flatten() } -fn pool_take_for_turn(key: &str, turn_id: Option) -> Option> { - super::continuation::if_current_turn(Some(key), turn_id, || WS_POOL.lock().ok()?.remove(key)) - .flatten() +fn pool_take_for_turn( + owner: &ConversationIdentity, + turn_id: Option, +) -> Option> { + super::continuation::if_current_turn(Some(owner), turn_id, || { + WS_POOL.lock().ok()?.remove(owner) + }) + .flatten() } -fn pool_insert_for_turn(key: String, entry: Arc, turn_id: Option) { - let session_id = key.clone(); - super::continuation::with_current_turn(Some(&session_id), turn_id, || pool_insert(key, entry)); +fn pool_insert_for_turn(owner: ConversationIdentity, entry: Arc, turn_id: Option) { + let current_owner = owner.clone(); + super::continuation::with_current_turn(Some(¤t_owner), turn_id, || { + pool_insert(owner, entry) + }); } -fn pool_remove_entry(key: &str, entry: &Arc) { - invalidate_pool_entry(key, entry); +fn pool_remove_entry(owner: &ConversationIdentity, entry: &Arc) { + invalidate_pool_entry(owner, entry); } -fn pool_insert(key: String, entry: Arc) { +fn pool_insert(owner: ConversationIdentity, entry: Arc) { entry.touch(); let mut guard = WS_POOL.lock().unwrap(); // Evict oldest if at capacity if guard.len() >= MAX_POOL_ENTRIES - && let Some(oldest_key) = guard.keys().next().cloned() + && let Some(oldest_owner) = guard.keys().next().cloned() { - guard.remove(&oldest_key); + guard.remove(&oldest_owner); } // Evict expired entries let now = now_ms(); - guard.retain(|_, e| now.saturating_sub(e.created_at) < POOL_IDLE_TTL_MS); - guard.insert(key, entry); + guard.retain(|_, entry| now.saturating_sub(entry.created_at) < POOL_IDLE_TTL_MS); + guard.insert(owner, entry); } fn cleanup_pool_before_connect() { @@ -343,15 +351,13 @@ fn cleanup_pool_before_connect() { let mut candidates: Vec<_> = guard .iter() .filter(|(_, entry)| Arc::strong_count(entry) == 1) - .map(|(key, entry)| (key.clone(), entry.last_activity.load(Ordering::Relaxed))) + .map(|(owner, entry)| (owner.clone(), entry.last_activity.load(Ordering::Relaxed))) .collect(); - candidates.sort_unstable_by(|left, right| { - left.1.cmp(&right.1).then_with(|| left.0.cmp(&right.0)) - }); + candidates.sort_unstable_by_key(|(_, activity)| *activity); candidates .into_iter() .take(remove_count) - .filter_map(|(key, _)| guard.remove(&key)) + .filter_map(|(owner, _)| guard.remove(&owner)) .collect::>() }; drop(removed); @@ -537,7 +543,7 @@ pub(super) async fn codex_websocket_request( body_value: &serde_json::Value, _ctx: &RequestContext, traffic: Option<&TrafficCapture>, - pool_key: Option<&str>, + pool_owner: Option<&ConversationIdentity>, connect_timeout_ms: u64, idle_timeout_ms: u64, continuation: Option<&ContinuationCandidate>, @@ -576,7 +582,7 @@ pub(super) async fn codex_websocket_request( let started_at = Instant::now(); // Check pool for existing connection - let pooled = pool_key.and_then(|key| { + let pooled = pool_owner.and_then(|key| { pool_get_for_turn(key, continuation.and_then(|candidate| candidate.turn_id)) }); @@ -585,7 +591,7 @@ pub(super) async fn codex_websocket_request( let mut ws_guard = entry.ws.lock().await; // Check if connection is still alive by sending a ping if ws_guard.send(Message::Ping(vec![])).await.is_err() { - invalidate_pool_entry(pool_key.unwrap(), &entry); + invalidate_pool_entry(pool_owner.unwrap(), &entry); // Fall through to new connection connect_with_timeout( websocket_client, @@ -599,7 +605,7 @@ pub(super) async fn codex_websocket_request( // Connection is alive, send the request through it let ws_msg = Message::Text(body_json.clone()); ws_guard.send(ws_msg).await.map_err(|e| { - if let Some(key) = pool_key { + if let Some(key) = pool_owner { invalidate_pool_entry(key, &entry); } CodexError { @@ -615,7 +621,7 @@ pub(super) async fn codex_websocket_request( let (sse_body, terminal_event) = collect_ws_events( &mut ws_guard, idle_timeout_ms, - pool_key, + pool_owner, Some(&entry), traffic, ) @@ -645,7 +651,7 @@ pub(super) async fn codex_websocket_request( // Write traffic metadata if let Some(tc) = traffic { - write_websocket_metadata_capture(tc, &ws_url, pool_key, continuation, true); + write_websocket_metadata_capture(tc, &ws_url, pool_owner, continuation, true); write_websocket_response_capture(tc, status, started_at.elapsed(), &sse_body); } @@ -685,7 +691,7 @@ pub(super) async fn codex_websocket_request( let (sse_body, terminal_event) = collect_ws_events( &mut ws_guard, idle_timeout_ms, - pool_key, + pool_owner, Some(&entry), traffic, ) @@ -696,7 +702,7 @@ pub(super) async fn codex_websocket_request( }; if is_previous_response_missing(&terminal_event.payload) { - if let Some(key) = pool_key { + if let Some(key) = pool_owner { invalidate_pool_entry(key, &entry); } return Err(CodexError { @@ -709,11 +715,11 @@ pub(super) async fn codex_websocket_request( } // Pool the connection if we have a key and it was successful - if let Some(key) = pool_key { + if let Some(key) = pool_owner { let should_pool = terminal_event.event_type == "response.completed"; if should_pool { pool_insert_for_turn( - key.to_string(), + key.clone(), entry.clone(), continuation.and_then(|candidate| candidate.turn_id), ); @@ -728,7 +734,7 @@ pub(super) async fn codex_websocket_request( // Write traffic metadata if let Some(tc) = traffic { - write_websocket_metadata_capture(tc, &ws_url, pool_key, continuation, false); + write_websocket_metadata_capture(tc, &ws_url, pool_owner, continuation, false); write_websocket_response_capture(tc, status, started_at.elapsed(), &sse_body); } @@ -746,7 +752,7 @@ pub(super) struct ReadyWebSocket { guard: OwnedMutexGuard, entry: Arc, used_pooled: bool, - pool_key: Option, + pool_owner: Option, turn_id: Option, traffic: Option>, idle_timeout_ms: u64, @@ -759,7 +765,7 @@ pub(super) async fn prepare_codex_websocket( url: &str, headers: &HeaderMap, traffic: Option>, - pool_key: Option<&str>, + pool_owner: Option<&ConversationIdentity>, turn_id: Option, connect_timeout_ms: u64, idle_timeout_ms: u64, @@ -771,7 +777,7 @@ pub(super) async fn prepare_codex_websocket( retry_after: None, origin: CodexErrorOrigin::WebSocketHandshake, })?; - let pooled = pool_key.and_then(|key| pool_take_for_turn(key, turn_id)); + let pooled = pool_owner.and_then(|key| pool_take_for_turn(key, turn_id)); let used_pooled = pooled.is_some(); let entry = if let Some(entry) = pooled { entry @@ -818,7 +824,7 @@ pub(super) async fn prepare_codex_websocket( guard, entry, used_pooled, - pool_key: pool_key.map(str::to_string), + pool_owner: pool_owner.cloned(), turn_id, traffic, idle_timeout_ms, @@ -846,7 +852,7 @@ pub(super) fn start_codex_websocket_events( write_websocket_metadata_capture( tc, &ready.ws_url, - ready.pool_key.as_deref(), + ready.pool_owner.as_ref(), continuation, ready.used_pooled, ); @@ -868,13 +874,13 @@ pub(super) fn start_codex_websocket_events( mut guard, entry, used_pooled: _, - pool_key, + pool_owner, turn_id, traffic, idle_timeout_ms, } = ready; if let Err(error) = guard.send(Message::Text(body_json)).await { - if let Some(key) = pool_key.as_deref() { + if let Some(key) = pool_owner.as_ref() { pool_remove_entry(key, &entry); } let _ = tx @@ -891,16 +897,16 @@ pub(super) fn start_codex_websocket_events( let reusable = stream_ws_events( &mut guard, idle_timeout_ms, - pool_key.as_deref(), + pool_owner.as_ref(), Some(&entry), traffic, tx, ) .await; drop(guard); - if let Some(key) = pool_key.as_deref() { + if let Some(key) = pool_owner.as_ref() { if reusable { - pool_insert_for_turn(key.to_string(), entry, turn_id); + pool_insert_for_turn(key.clone(), entry, turn_id); } else { pool_remove_entry(key, &entry); } @@ -918,7 +924,7 @@ pub(super) async fn codex_websocket_event_stream( body_value: &serde_json::Value, _ctx: &RequestContext, traffic: Option>, - pool_key: Option<&str>, + pool_owner: Option<&ConversationIdentity>, connect_timeout_ms: u64, idle_timeout_ms: u64, continuation: Option<&ContinuationCandidate>, @@ -936,7 +942,7 @@ pub(super) async fn codex_websocket_event_stream( url, headers, traffic, - pool_key, + pool_owner, continuation.and_then(|candidate| candidate.turn_id), connect_timeout_ms, idle_timeout_ms, @@ -974,7 +980,7 @@ fn response_start_timeout_error(timeout_ms: u64) -> CodexError { fn write_websocket_metadata_capture( traffic: &TrafficCapture, ws_url: &str, - pool_key: Option<&str>, + pool_owner: Option<&ConversationIdentity>, continuation: Option<&ContinuationCandidate>, pooled: bool, ) { @@ -984,7 +990,7 @@ fn write_websocket_metadata_capture( "provider": "codex", "transport": "websocket", "url": ws_url, - "poolKey": pool_key, + "poolingEnabled": pool_owner.is_some(), "pooled": pooled, "continuation": { "previousResponseId": continuation @@ -1740,7 +1746,7 @@ struct WsEvent { async fn collect_ws_events( ws: &mut WebSocketStream, idle_timeout_ms: u64, - pool_key: Option<&str>, + pool_owner: Option<&ConversationIdentity>, pool_entry: Option<&Arc>, traffic: Option<&TrafficCapture>, ) -> Result<(Vec, Option), CodexError> @@ -1764,7 +1770,7 @@ where match response_event_budget.checked_sub(response_deadline_started.elapsed()) { Some(remaining) if !remaining.is_zero() => remaining, _ => { - invalidate_pool_owner(pool_key, pool_entry); + invalidate_pool_owner(pool_owner, pool_entry); return Err(CodexError { status: 0, message: format!("WebSocket idle timeout after {idle_timeout_ms}ms"), @@ -1778,7 +1784,7 @@ where match response_event_budget.checked_sub(response_deadline_started.elapsed()) { Some(remaining) if !remaining.is_zero() => remaining, _ => { - invalidate_pool_owner(pool_key, pool_entry); + invalidate_pool_owner(pool_owner, pool_entry); return Err(response_start_timeout_error(idle_timeout_ms)); } } @@ -1787,7 +1793,7 @@ where let timeout = tokio::time::timeout(read_timeout, ws.next()); let frame = timeout.await.map_err(|_| { - invalidate_pool_owner(pool_key, pool_entry); + invalidate_pool_owner(pool_owner, pool_entry); if response_started { CodexError { status: 0, @@ -1848,7 +1854,7 @@ where } Some(Ok(Message::Binary(_))) => { // Reject binary frames - invalidate_pool_owner(pool_key, pool_entry); + invalidate_pool_owner(pool_owner, pool_entry); return Err(CodexError { status: 0, message: "WebSocket binary frames not supported".to_string(), @@ -1871,12 +1877,12 @@ where } Some(Ok(Message::Close(_))) => { // Connection closed - invalidate pool - invalidate_pool_owner(pool_key, pool_entry); + invalidate_pool_owner(pool_owner, pool_entry); break; } Some(Err(e)) => { // Stream error - invalidate pool - invalidate_pool_owner(pool_key, pool_entry); + invalidate_pool_owner(pool_owner, pool_entry); return Err(CodexError { status: 0, message: format!("WebSocket stream error: {e}"), @@ -1887,7 +1893,7 @@ where } None => { // Stream ended - invalidate pool - invalidate_pool_owner(pool_key, pool_entry); + invalidate_pool_owner(pool_owner, pool_entry); break; } } @@ -1899,7 +1905,7 @@ where async fn stream_ws_events( ws: &mut WebSocketStream, idle_timeout_ms: u64, - pool_key: Option<&str>, + pool_owner: Option<&ConversationIdentity>, pool_entry: Option<&Arc>, traffic: Option>, tx: mpsc::Sender>, @@ -1926,7 +1932,7 @@ where match response_event_budget.checked_sub(response_deadline_started.elapsed()) { Some(remaining) if !remaining.is_zero() => remaining, _ => { - invalidate_pool_owner(pool_key, pool_entry); + invalidate_pool_owner(pool_owner, pool_entry); let err = if response_started { CodexError { status: 0, @@ -1946,7 +1952,7 @@ where let frame = match tokio::time::timeout(read_timeout, ws.next()).await { Ok(frame) => frame, Err(_) => { - invalidate_pool_owner(pool_key, pool_entry); + invalidate_pool_owner(pool_owner, pool_entry); let err = if response_started { CodexError { status: 0, @@ -1997,7 +2003,7 @@ where } let terminal = is_terminal_event(&parsed); if terminal && is_previous_response_missing(&parsed) { - invalidate_pool_owner(pool_key, pool_entry); + invalidate_pool_owner(pool_owner, pool_entry); let _ = tx .send(Err(CodexError { status: 0, @@ -2015,7 +2021,7 @@ where .unwrap_or("unknown") .to_string(); if tx.send(Ok(parsed)).await.is_err() { - invalidate_pool_owner(pool_key, pool_entry); + invalidate_pool_owner(pool_owner, pool_entry); break; } if terminal { @@ -2024,7 +2030,7 @@ where } } Some(Ok(Message::Binary(_))) => { - invalidate_pool_owner(pool_key, pool_entry); + invalidate_pool_owner(pool_owner, pool_entry); let _ = tx .send(Err(CodexError { status: 0, @@ -2041,12 +2047,12 @@ where } Some(Ok(Message::Pong(_))) | Some(Ok(Message::Frame(_))) => {} Some(Ok(Message::Close(_))) | None => { - invalidate_pool_owner(pool_key, pool_entry); + invalidate_pool_owner(pool_owner, pool_entry); let _ = tx.send(Err(missing_terminal_error())).await; break; } Some(Err(e)) => { - invalidate_pool_owner(pool_key, pool_entry); + invalidate_pool_owner(pool_owner, pool_entry); let _ = tx .send(Err(CodexError { status: 0, @@ -2108,6 +2114,14 @@ mod tests { WS_POOL_TEST_LOCK.lock().await } + fn main_owner(session_id: &str) -> ConversationIdentity { + ConversationIdentity::Main(session_id.to_string()) + } + + fn agent_owner(session_id: &str, agent_id: &str) -> ConversationIdentity { + ConversationIdentity::Agent(session_id.to_string(), agent_id.to_string()) + } + fn test_websocket_client() -> reqwest::Client { reqwest::Client::builder() .http1_only() @@ -2195,7 +2209,10 @@ mod tests { { let mut guard = WS_POOL.lock().unwrap(); for index in 0..POOL_CONNECT_CLEANUP_THRESHOLD { - guard.insert(format!("entry-{index:02}"), shared_pool_entry(&ws)); + guard.insert( + main_owner(&format!("entry-{index:02}")), + shared_pool_entry(&ws), + ); } } @@ -2208,21 +2225,31 @@ mod tests { WS_POOL .lock() .unwrap() - .insert("entry-50".to_string(), shared_pool_entry(&ws)); - drop(pool_get_for_turn("entry-00", None).unwrap()); - let leased = WS_POOL.lock().unwrap().get("entry-01").unwrap().clone(); + .insert(main_owner("entry-50"), shared_pool_entry(&ws)); + WS_POOL + .lock() + .unwrap() + .get(&main_owner("entry-00")) + .unwrap() + .touch(); + let leased = WS_POOL + .lock() + .unwrap() + .get(&main_owner("entry-01")) + .unwrap() + .clone(); cleanup_pool_before_connect(); let guard = WS_POOL.lock().unwrap(); assert_eq!(guard.len(), POOL_CONNECT_CLEANUP_TARGET); - assert!(guard.contains_key("entry-00")); - assert!(guard.contains_key("entry-01")); + assert!(guard.contains_key(&main_owner("entry-00"))); + assert!(guard.contains_key(&main_owner("entry-01"))); for index in 2..=12 { - assert!(!guard.contains_key(&format!("entry-{index:02}"))); + assert!(!guard.contains_key(&main_owner(&format!("entry-{index:02}")))); } - assert!(guard.contains_key("entry-13")); - assert!(guard.contains_key("entry-50")); + assert!(guard.contains_key(&main_owner("entry-13"))); + assert!(guard.contains_key(&main_owner("entry-50"))); drop(guard); drop(leased); clear_codex_websocket_pool_for_tests(); @@ -2240,11 +2267,14 @@ mod tests { { let mut guard = WS_POOL.lock().unwrap(); guard.insert( - "entry-00".to_string(), + main_owner("entry-00"), Arc::new(PoolEntry::new(probe_stream)), ); for index in 1..=POOL_CONNECT_CLEANUP_THRESHOLD { - guard.insert(format!("entry-{index:02}"), shared_pool_entry(&shared_ws)); + guard.insert( + main_owner(&format!("entry-{index:02}")), + shared_pool_entry(&shared_ws), + ); } } @@ -2777,23 +2807,51 @@ mod tests { assert!(!is_previous_response_missing(&unrelated)); } + #[test] + fn websocket_metadata_does_not_serialize_typed_owner() { + let temp = tempfile::tempdir().unwrap(); + let traffic = crate::traffic::test_capture(temp.path().join("traffic")); + let owner = agent_owner("session-secret", "agent-secret"); + + write_websocket_metadata_capture( + &traffic, + "wss://example.invalid/responses", + Some(&owner), + None, + false, + ); + + let artifact = std::fs::read_dir(traffic.root()) + .unwrap() + .next() + .unwrap() + .unwrap() + .path(); + let captured = std::fs::read_to_string(artifact).unwrap(); + assert!(captured.contains("poolingEnabled")); + assert!(!captured.contains("poolKey")); + assert!(!captured.contains("session-secret")); + assert!(!captured.contains("agent-secret")); + } + #[tokio::test] async fn pool_checkout_is_exclusive_and_removal_is_identity_safe() { let _pool_test_guard = lock_ws_pool_tests().await; clear_codex_websocket_pool_for_tests(); let first = Arc::new(PoolEntry::new(create_dummy_stream_async().await)); - pool_insert("exclusive".to_string(), first.clone()); + let owner = main_owner("exclusive"); + pool_insert(owner.clone(), first.clone()); assert!(Arc::ptr_eq( - &WS_POOL.lock().unwrap().remove("exclusive").unwrap(), + &WS_POOL.lock().unwrap().remove(&owner).unwrap(), &first )); - assert!(WS_POOL.lock().unwrap().remove("exclusive").is_none()); + assert!(WS_POOL.lock().unwrap().remove(&owner).is_none()); let replacement = Arc::new(PoolEntry::new(create_dummy_stream_async().await)); - pool_insert("exclusive".to_string(), replacement.clone()); - pool_remove_entry("exclusive", &first); + pool_insert(owner.clone(), replacement.clone()); + pool_remove_entry(&owner, &first); assert!(Arc::ptr_eq( - WS_POOL.lock().unwrap().get("exclusive").unwrap(), + WS_POOL.lock().unwrap().get(&owner).unwrap(), &replacement )); clear_codex_websocket_pool_for_tests(); @@ -2803,17 +2861,25 @@ mod tests { async fn pool_invalidation() { let _pool_test_guard = lock_ws_pool_tests().await; clear_codex_websocket_pool_for_tests(); - // Verify pool operations work through the public API - // We insert an entry directly into the pool, then invalidate it - let stream = create_dummy_stream_async().await; + let first_stream = create_dummy_stream_async().await; + let second_stream = create_dummy_stream_async().await; + let first_owner = agent_owner("test-session", "first-agent"); + let sibling_owner = agent_owner("test-session", "sibling-agent"); { let mut guard = WS_POOL.lock().unwrap(); - guard.insert("test-session".to_string(), Arc::new(PoolEntry::new(stream))); + guard.insert(first_owner.clone(), Arc::new(PoolEntry::new(first_stream))); + guard.insert( + sibling_owner.clone(), + Arc::new(PoolEntry::new(second_stream)), + ); } - assert!(WS_POOL.lock().unwrap().contains_key("test-session")); + assert!(WS_POOL.lock().unwrap().contains_key(&first_owner)); + assert!(WS_POOL.lock().unwrap().contains_key(&sibling_owner)); - invalidate_codex_websocket_pool_key("test-session"); - assert!(!WS_POOL.lock().unwrap().contains_key("test-session")); + invalidate_codex_websocket_pool_owner(&first_owner); + assert!(!WS_POOL.lock().unwrap().contains_key(&first_owner)); + assert!(WS_POOL.lock().unwrap().contains_key(&sibling_owner)); + clear_codex_websocket_pool_for_tests(); } #[tokio::test] @@ -3132,16 +3198,14 @@ mod tests { } #[tokio::test] - async fn binary_frame_invalidates_pool_key() { + async fn binary_frame_invalidates_pool_owner() { let _pool_test_guard = lock_ws_pool_tests().await; clear_codex_websocket_pool_for_tests(); let pooled_stream = create_dummy_stream_async().await; + let owner = agent_owner("binary-session", "binary-agent"); { let mut guard = WS_POOL.lock().unwrap(); - guard.insert( - "binary-session".to_string(), - Arc::new(PoolEntry::new(pooled_stream)), - ); + guard.insert(owner.clone(), Arc::new(PoolEntry::new(pooled_stream))); } let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); @@ -3155,14 +3219,13 @@ mod tests { let (mut ws, _) = tokio_tungstenite::connect_async(format!("ws://{addr}/")) .await .unwrap(); - let err = match collect_ws_events(&mut ws, 1_000, Some("binary-session"), None, None).await - { + let err = match collect_ws_events(&mut ws, 1_000, Some(&owner), None, None).await { Ok(_) => panic!("expected binary frame to fail"), Err(err) => err, }; assert!(err.message.contains("binary frames")); - assert!(!WS_POOL.lock().unwrap().contains_key("binary-session")); + assert!(!WS_POOL.lock().unwrap().contains_key(&owner)); } #[tokio::test] @@ -3170,12 +3233,10 @@ mod tests { let _pool_test_guard = lock_ws_pool_tests().await; clear_codex_websocket_pool_for_tests(); let pooled_stream = create_dummy_stream_async().await; + let owner = main_owner("start-timeout-session"); { let mut guard = WS_POOL.lock().unwrap(); - guard.insert( - "start-timeout-session".to_string(), - Arc::new(PoolEntry::new(pooled_stream)), - ); + guard.insert(owner.clone(), Arc::new(PoolEntry::new(pooled_stream))); } let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); @@ -3199,22 +3260,16 @@ mod tests { let (mut ws, _) = tokio_tungstenite::connect_async(format!("ws://{addr}/")) .await .unwrap(); - let err = - match collect_ws_events(&mut ws, 50, Some("start-timeout-session"), None, None).await { - Ok(_) => panic!("expected response start timeout"), - Err(err) => err, - }; + let err = match collect_ws_events(&mut ws, 50, Some(&owner), None, None).await { + Ok(_) => panic!("expected response start timeout"), + Err(err) => err, + }; assert_eq!( err.detail.as_deref(), Some(WEBSOCKET_RESPONSE_START_TIMEOUT_DETAIL) ); - assert!( - !WS_POOL - .lock() - .unwrap() - .contains_key("start-timeout-session") - ); + assert!(!WS_POOL.lock().unwrap().contains_key(&owner)); } #[tokio::test] @@ -3222,12 +3277,10 @@ mod tests { let _pool_test_guard = lock_ws_pool_tests().await; clear_codex_websocket_pool_for_tests(); let pooled_stream = create_dummy_stream_async().await; + let owner = main_owner("response-idle-session"); { let mut guard = WS_POOL.lock().unwrap(); - guard.insert( - "response-idle-session".to_string(), - Arc::new(PoolEntry::new(pooled_stream)), - ); + guard.insert(owner.clone(), Arc::new(PoolEntry::new(pooled_stream))); } let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); @@ -3252,20 +3305,14 @@ mod tests { let (mut ws, _) = tokio_tungstenite::connect_async(format!("ws://{addr}/")) .await .unwrap(); - let err = - match collect_ws_events(&mut ws, 50, Some("response-idle-session"), None, None).await { - Ok(_) => panic!("expected response idle timeout"), - Err(err) => err, - }; + let err = match collect_ws_events(&mut ws, 50, Some(&owner), None, None).await { + Ok(_) => panic!("expected response idle timeout"), + Err(err) => err, + }; assert!(err.message.contains("idle timeout")); assert_eq!(err.detail, None); - assert!( - !WS_POOL - .lock() - .unwrap() - .contains_key("response-idle-session") - ); + assert!(!WS_POOL.lock().unwrap().contains_key(&owner)); } async fn create_dummy_stream_async() -> CodexWebSocketStream { diff --git a/src/request_identity.rs b/src/request_identity.rs new file mode 100644 index 00000000..0df0a13f --- /dev/null +++ b/src/request_identity.rs @@ -0,0 +1,338 @@ +use http::HeaderMap; + +pub const CLAUDE_SESSION_HEADER: &str = "x-claude-code-session-id"; +pub const CLAUDE_AGENT_HEADER: &str = "x-claude-code-agent-id"; +pub const CLAUDE_PARENT_AGENT_HEADER: &str = "x-claude-code-parent-agent-id"; + +const MAX_IDENTITY_LEN: usize = 512; + +#[derive(Debug, Clone, PartialEq, Eq, Hash)] +pub enum ConversationIdentity { + Main(String), + Agent(String, String), +} + +impl ConversationIdentity { + pub fn from_headers(headers: &HeaderMap) -> Option { + let session = read_identity_header(headers, CLAUDE_SESSION_HEADER); + let agent = read_identity_header(headers, CLAUDE_AGENT_HEADER); + let parent = read_identity_header(headers, CLAUDE_PARENT_AGENT_HEADER); + + if session.is_invalid() || agent.is_invalid() || parent.is_invalid() { + return None; + } + + match (session.value(), agent.value(), parent.value()) { + (Some(session_id), Some(agent_id), _) => { + Some(Self::Agent(session_id.to_string(), agent_id.to_string())) + } + (Some(session_id), None, None) => Some(Self::Main(session_id.to_string())), + _ => None, + } + } +} + +#[derive(Debug)] +enum ParsedHeader<'a> { + Missing, + Valid(&'a str), + Invalid, +} + +impl ParsedHeader<'_> { + fn value(&self) -> Option<&str> { + match self { + Self::Valid(value) => Some(value), + Self::Missing | Self::Invalid => None, + } + } + + fn is_invalid(&self) -> bool { + matches!(self, Self::Invalid) + } +} + +fn read_identity_header<'a>(headers: &'a HeaderMap, name: &str) -> ParsedHeader<'a> { + let mut values = headers.get_all(name).iter(); + let Some(value) = values.next() else { + return ParsedHeader::Missing; + }; + if values.next().is_some() { + return ParsedHeader::Invalid; + } + + let Ok(value) = value.to_str() else { + return ParsedHeader::Invalid; + }; + let value = value.trim_matches(|character| matches!(character, ' ' | '\t')); + if value.is_empty() + || value.len() > MAX_IDENTITY_LEN + || value.contains(',') + || !value.bytes().all(|byte| byte.is_ascii_graphic()) + { + return ParsedHeader::Invalid; + } + + ParsedHeader::Valid(value) +} + +#[cfg(test)] +mod tests { + use super::*; + use http::{HeaderName, HeaderValue}; + + fn headers(values: &[(&str, &str)]) -> HeaderMap { + let mut headers = HeaderMap::new(); + for (name, value) in values { + headers.append( + HeaderName::from_bytes(name.as_bytes()).unwrap(), + HeaderValue::from_str(value).unwrap(), + ); + } + headers + } + + #[test] + fn parses_main_agent_and_lineage_shapes() { + let cases = [ + ( + "main", + vec![(CLAUDE_SESSION_HEADER, "session-a")], + Some(ConversationIdentity::Main("session-a".to_string())), + ), + ( + "direct agent without parent", + vec![ + (CLAUDE_SESSION_HEADER, "session-a"), + (CLAUDE_AGENT_HEADER, "agent-a"), + ], + Some(ConversationIdentity::Agent( + "session-a".to_string(), + "agent-a".to_string(), + )), + ), + ( + "nested direct child", + vec![ + (CLAUDE_SESSION_HEADER, "session-a"), + (CLAUDE_AGENT_HEADER, "agent-child"), + (CLAUDE_PARENT_AGENT_HEADER, "agent-parent"), + ], + Some(ConversationIdentity::Agent( + "session-a".to_string(), + "agent-child".to_string(), + )), + ), + ( + "sibling one", + vec![ + (CLAUDE_SESSION_HEADER, "session-a"), + (CLAUDE_AGENT_HEADER, "agent-sibling-one"), + (CLAUDE_PARENT_AGENT_HEADER, "agent-parent"), + ], + Some(ConversationIdentity::Agent( + "session-a".to_string(), + "agent-sibling-one".to_string(), + )), + ), + ( + "sibling two", + vec![ + (CLAUDE_SESSION_HEADER, "session-a"), + (CLAUDE_AGENT_HEADER, "agent-sibling-two"), + (CLAUDE_PARENT_AGENT_HEADER, "agent-parent"), + ], + Some(ConversationIdentity::Agent( + "session-a".to_string(), + "agent-sibling-two".to_string(), + )), + ), + ( + "same agent in another session", + vec![ + (CLAUDE_SESSION_HEADER, "session-b"), + (CLAUDE_AGENT_HEADER, "agent-a"), + ], + Some(ConversationIdentity::Agent( + "session-b".to_string(), + "agent-a".to_string(), + )), + ), + ( + "outer space and tab", + vec![ + (CLAUDE_SESSION_HEADER, " \tsession-a\t "), + (CLAUDE_AGENT_HEADER, "\tagent-a "), + (CLAUDE_PARENT_AGENT_HEADER, " agent-parent\t"), + ], + Some(ConversationIdentity::Agent( + "session-a".to_string(), + "agent-a".to_string(), + )), + ), + ]; + + for (name, values, expected) in cases { + assert_eq!( + ConversationIdentity::from_headers(&headers(&values)), + expected, + "{name}" + ); + } + } + + #[test] + fn parent_is_validation_only_and_never_changes_the_owner() { + let direct = ConversationIdentity::from_headers(&headers(&[ + (CLAUDE_SESSION_HEADER, "session-a"), + (CLAUDE_AGENT_HEADER, "agent-child"), + ])); + let nested = ConversationIdentity::from_headers(&headers(&[ + (CLAUDE_SESSION_HEADER, "session-a"), + (CLAUDE_AGENT_HEADER, "agent-child"), + (CLAUDE_PARENT_AGENT_HEADER, "agent-parent"), + ])); + let reparented = ConversationIdentity::from_headers(&headers(&[ + (CLAUDE_SESSION_HEADER, "session-a"), + (CLAUDE_AGENT_HEADER, "agent-child"), + (CLAUDE_PARENT_AGENT_HEADER, "another-parent"), + ])); + + assert_eq!(direct, nested); + assert_eq!(nested, reparented); + } + + #[test] + fn ambiguous_or_absent_tuples_are_stateless() { + let cases = [ + ("all missing", vec![]), + ( + "agent without session", + vec![(CLAUDE_AGENT_HEADER, "agent-a")], + ), + ( + "parent without direct agent", + vec![ + (CLAUDE_SESSION_HEADER, "session-a"), + (CLAUDE_PARENT_AGENT_HEADER, "agent-parent"), + ], + ), + ( + "parent alone", + vec![(CLAUDE_PARENT_AGENT_HEADER, "agent-parent")], + ), + ( + "agent and parent without session", + vec![ + (CLAUDE_AGENT_HEADER, "agent-a"), + (CLAUDE_PARENT_AGENT_HEADER, "agent-parent"), + ], + ), + ]; + + for (name, values) in cases { + assert_eq!( + ConversationIdentity::from_headers(&headers(&values)), + None, + "{name}" + ); + } + } + + #[test] + fn rejects_malformed_text_in_every_identity_field() { + let malformed = [ + ("empty", ""), + ("spaces only", " "), + ("tabs only", "\t\t"), + ("internal space", "two values"), + ("internal tab", "two\tvalues"), + ("leading comma", ",value"), + ("trailing comma", "value,"), + ("coalesced", "first, second"), + ("oversize", "oversize-placeholder"), + ]; + + for field in [ + CLAUDE_SESSION_HEADER, + CLAUDE_AGENT_HEADER, + CLAUDE_PARENT_AGENT_HEADER, + ] { + for (shape, placeholder) in malformed { + let value = if shape == "oversize" { + "x".repeat(MAX_IDENTITY_LEN + 1) + } else { + placeholder.to_string() + }; + let mut values = vec![ + (CLAUDE_SESSION_HEADER, "session-a"), + (CLAUDE_AGENT_HEADER, "agent-a"), + (CLAUDE_PARENT_AGENT_HEADER, "agent-parent"), + ]; + values + .iter_mut() + .find(|(name, _)| *name == field) + .unwrap() + .1 = &value; + assert_eq!( + ConversationIdentity::from_headers(&headers(&values)), + None, + "field={field} shape={shape}" + ); + } + } + } + + #[test] + fn rejects_duplicate_headers_in_every_identity_field() { + for field in [ + CLAUDE_SESSION_HEADER, + CLAUDE_AGENT_HEADER, + CLAUDE_PARENT_AGENT_HEADER, + ] { + let mut values = vec![ + (CLAUDE_SESSION_HEADER, "session-a"), + (CLAUDE_AGENT_HEADER, "agent-a"), + (CLAUDE_PARENT_AGENT_HEADER, "agent-parent"), + ]; + values.push((field, "duplicate")); + assert_eq!( + ConversationIdentity::from_headers(&headers(&values)), + None, + "field={field}" + ); + } + } + + #[test] + fn rejects_nontext_headers_in_every_identity_field() { + for field in [ + CLAUDE_SESSION_HEADER, + CLAUDE_AGENT_HEADER, + CLAUDE_PARENT_AGENT_HEADER, + ] { + let mut values = headers(&[ + (CLAUDE_SESSION_HEADER, "session-a"), + (CLAUDE_AGENT_HEADER, "agent-a"), + (CLAUDE_PARENT_AGENT_HEADER, "agent-parent"), + ]); + values.insert(field, HeaderValue::from_bytes(&[0x80]).unwrap()); + assert_eq!( + ConversationIdentity::from_headers(&values), + None, + "field={field}" + ); + } + } + + #[test] + fn malformed_agent_cannot_downgrade_a_valid_session_to_main() { + for malformed_agent in ["", "agent one", "agent-a,agent-b"] { + let identity = ConversationIdentity::from_headers(&headers(&[ + (CLAUDE_SESSION_HEADER, "session-a"), + (CLAUDE_AGENT_HEADER, malformed_agent), + ])); + assert_eq!(identity, None, "agent={malformed_agent:?}"); + } + } +} diff --git a/src/server.rs b/src/server.rs index c31894b1..bb563d9e 100644 --- a/src/server.rs +++ b/src/server.rs @@ -25,6 +25,7 @@ use crate::{ }, }, registry::{Registry, normalize_incoming_model}, + request_identity::{CLAUDE_AGENT_HEADER, CLAUDE_PARENT_AGENT_HEADER, ConversationIdentity}, session::{self, SessionState}, traffic::{TrafficCaptureOptions, create_traffic_capture}, }; @@ -1382,6 +1383,9 @@ async fn dispatch_request( let method = req.method().clone(); let uri = req.uri().clone(); let headers = req.headers().clone(); + let conversation_identity = (!count_tokens) + .then(|| ConversationIdentity::from_headers(&headers)) + .flatten(); let path = uri.path().to_string(); let query = redacted_query(&uri); let endpoint = if count_tokens { @@ -1721,7 +1725,17 @@ async fn dispatch_request( let response = if count_tokens { provider.handle_count_tokens(body, context).await } else { - provider.handle_messages(body, context).await + provider + .handle_messages_with_conversation_identity( + body, + context, + if auto_review_route.is_some() { + None + } else { + conversation_identity + }, + ) + .await }; log_request_completed( &log, @@ -2040,7 +2054,15 @@ fn headers_to_record(headers: &http::HeaderMap) -> Value { let mut out = Map::new(); for (key, value) in headers { if let Ok(raw) = value.to_str() { - out.insert(key.as_str().to_string(), Value::String(raw.to_string())); + let value = if matches!( + key.as_str(), + CLAUDE_AGENT_HEADER | CLAUDE_PARENT_AGENT_HEADER + ) { + format!("[redacted len={}]", raw.len()) + } else { + raw.to_string() + }; + out.insert(key.as_str().to_string(), Value::String(value)); } } Value::Object(out) @@ -2124,8 +2146,10 @@ fn _unused(session_state: Option<&SessionState>) { #[cfg(test)] mod auto_review_tests { - use super::{apply_auto_review_model, is_claude_auto_review_request}; + use super::{apply_auto_review_model, headers_to_record, is_claude_auto_review_request}; use crate::anthropic::schema::MessagesRequest; + use crate::request_identity::{CLAUDE_AGENT_HEADER, CLAUDE_PARENT_AGENT_HEADER}; + use http::{HeaderMap, HeaderValue}; use serde_json::json; fn request(system: &str, stream: bool, tools: serde_json::Value) -> MessagesRequest { @@ -2208,6 +2232,37 @@ mod auto_review_tests { assert_eq!(classifier.model.as_deref(), Some("grok-4.5")); } + #[test] + fn traffic_headers_redact_agent_lineage_ids() { + let mut headers = HeaderMap::new(); + headers.insert( + "x-claude-code-session-id", + HeaderValue::from_static("session-visible"), + ); + headers.insert( + CLAUDE_AGENT_HEADER, + HeaderValue::from_static("agent-secret"), + ); + headers.insert( + CLAUDE_PARENT_AGENT_HEADER, + HeaderValue::from_static("parent-secret"), + ); + + let captured = headers_to_record(&headers); + assert_eq!( + captured["x-claude-code-session-id"], + json!("session-visible") + ); + assert_eq!(captured[CLAUDE_AGENT_HEADER], json!("[redacted len=12]")); + assert_eq!( + captured[CLAUDE_PARENT_AGENT_HEADER], + json!("[redacted len=13]") + ); + let serialized = captured.to_string(); + assert!(!serialized.contains("agent-secret")); + assert!(!serialized.contains("parent-secret")); + } + #[test] fn count_tokens_keeps_requested_model() { let mut classifier = request( diff --git a/tests/server.rs b/tests/server.rs index 9087afff..339078ce 100644 --- a/tests/server.rs +++ b/tests/server.rs @@ -8,13 +8,14 @@ use claude_code_proxy::{ monitor::{MonitorHandle, RequestStatus}, provider::{CliHandlers, Generation, GenerationBody, Provider, ProviderError, RequestContext}, registry::Registry, + request_identity::ConversationIdentity, server::{ AppFeatures, app, app_with_features, app_with_monitor, app_with_options, bind_proxy_listener, }, }; use serde_json::{Value, json}; -use std::sync::Arc; +use std::sync::{Arc, Mutex}; use tower::util::ServiceExt; fn body_string(json: &str) -> Body { @@ -95,6 +96,56 @@ impl Provider for FakeProvider { } } +type CapturedIdentity = (Option, Option); + +struct IdentityCaptureProvider { + captured: Arc>>, +} + +#[async_trait] +impl Provider for IdentityCaptureProvider { + fn name(&self) -> &'static str { + "codex" + } + + fn supported_models(&self) -> Vec { + vec!["gpt-5.5".to_string(), "gpt-5.6-luna".to_string()] + } + + fn cli(&self) -> &'static dyn CliHandlers { + &FAKE_CLI + } + + async fn handle_messages( + &self, + _body: MessagesRequest, + _ctx: RequestContext, + ) -> axum::response::Response { + (StatusCode::INTERNAL_SERVER_ERROR, "legacy path").into_response() + } + + async fn handle_messages_with_conversation_identity( + &self, + _body: MessagesRequest, + ctx: RequestContext, + conversation_identity: Option, + ) -> axum::response::Response { + self.captured + .lock() + .unwrap() + .push((conversation_identity, ctx.session_id)); + (StatusCode::OK, "captured").into_response() + } + + async fn handle_count_tokens( + &self, + _body: MessagesRequest, + _ctx: RequestContext, + ) -> axum::response::Response { + (StatusCode::OK, "counted").into_response() + } +} + fn routed_registry() -> Arc { Arc::new(Registry::from_providers( AliasProvider::Kimi, @@ -115,6 +166,142 @@ fn routed_registry() -> Arc { )) } +async fn call_identity_ingress( + app: &axum::Router, + path: &str, + headers: &[(&str, &str)], + body: Value, +) -> StatusCode { + let mut request = Request::builder() + .method(Method::POST) + .uri(path) + .header("content-type", "application/json"); + for (name, value) in headers { + request = request.header(*name, *value); + } + app.clone() + .oneshot(request.body(Body::from(body.to_string())).unwrap()) + .await + .unwrap() + .status() +} + +#[tokio::test] +async fn messages_ingress_forwards_only_strict_conversation_identity() { + let captured = Arc::new(Mutex::new(Vec::new())); + let provider = Arc::new(IdentityCaptureProvider { + captured: captured.clone(), + }) as Arc; + let app = app(Arc::new(Registry::from_providers( + AliasProvider::Codex, + [provider], + ))); + let normal_body = || { + json!({ + "model": "gpt-5.5", + "max_tokens": 32, + "messages": [{"role": "user", "content": "hello"}] + }) + }; + + let cases = [ + ( + vec![("x-claude-code-session-id", "session-main")], + normal_body(), + ), + ( + vec![ + ("x-claude-code-session-id", "session-agent"), + ("x-claude-code-agent-id", "agent-direct"), + ], + normal_body(), + ), + ( + vec![ + ("x-claude-code-session-id", "session-nested"), + ("x-claude-code-agent-id", "agent-child"), + ("x-claude-code-parent-agent-id", "agent-parent"), + ], + normal_body(), + ), + ( + vec![ + ("x-claude-code-session-id", "session-malformed-agent"), + ("x-claude-code-agent-id", "malformed agent"), + ], + normal_body(), + ), + (vec![], normal_body()), + ( + vec![("x-claude-code-session-id", " \tsession-trimmed\t ")], + normal_body(), + ), + ( + vec![ + ("x-claude-code-session-id", "session-auto-review"), + ("x-claude-code-agent-id", "agent-auto-review"), + ], + json!({ + "model": "gpt-5.5", + "max_tokens": 32, + "messages": [{"role": "user", "content": "review"}], + "system": [{ + "type": "text", + "text": "You are a security monitor for autonomous AI coding agents. Review this turn." + }] + }), + ), + ]; + + for (headers, body) in cases { + assert_eq!( + call_identity_ingress(&app, "/v1/messages", &headers, body).await, + StatusCode::OK + ); + } + assert_eq!( + call_identity_ingress( + &app, + "/v1/messages/count_tokens", + &[("x-claude-code-session-id", "session-count")], + normal_body(), + ) + .await, + StatusCode::OK + ); + + assert_eq!( + *captured.lock().unwrap(), + vec![ + ( + Some(ConversationIdentity::Main("session-main".to_string())), + Some("session-main".to_string()), + ), + ( + Some(ConversationIdentity::Agent( + "session-agent".to_string(), + "agent-direct".to_string(), + )), + Some("session-agent".to_string()), + ), + ( + Some(ConversationIdentity::Agent( + "session-nested".to_string(), + "agent-child".to_string(), + )), + Some("session-nested".to_string()), + ), + (None, Some("session-malformed-agent".to_string())), + (None, None), + ( + Some(ConversationIdentity::Main("session-trimmed".to_string())), + Some(" \tsession-trimmed\t ".to_string()), + ), + (None, Some("session-auto-review".to_string())), + ] + ); +} + #[tokio::test] async fn bind_error_names_address_and_port() { let occupied = std::net::TcpListener::bind("127.0.0.1:0").unwrap(); From e6374cf01cebe07505c6d99806d010515763f651 Mon Sep 17 00:00:00 2001 From: Albert Zhang Date: Sun, 2 Aug 2026 12:47:49 +0800 Subject: [PATCH 2/4] fix(codex): require the continuation origin socket Co-Authored-By: Claude --- src/providers/codex/client.rs | 1389 ++++++++++++++++++++++++--- src/providers/codex/compaction.rs | 2 +- src/providers/codex/continuation.rs | 646 +++++++++---- src/providers/codex/events.rs | 10 +- src/providers/codex/mod.rs | 738 +++++++++++--- src/providers/codex/websocket.rs | 1361 +++++++++++++++++++------- tests/public_codex_api_compat.rs | 94 ++ 7 files changed, 3453 insertions(+), 787 deletions(-) create mode 100644 tests/public_codex_api_compat.rs diff --git a/src/providers/codex/client.rs b/src/providers/codex/client.rs index ed12344e..16c2f9fd 100644 --- a/src/providers/codex/client.rs +++ b/src/providers/codex/client.rs @@ -351,6 +351,32 @@ pub struct CodexResponse { pub transport: ActualTransport, } +pub(crate) struct OwnerAwareCodexResponse { + response: CodexResponse, + pub(crate) socket_id: Option, +} + +impl OwnerAwareCodexResponse { + pub(crate) fn new(response: CodexResponse, socket_id: Option) -> Self { + Self { + response, + socket_id, + } + } + + fn into_response(self) -> CodexResponse { + self.response + } +} + +impl std::ops::Deref for OwnerAwareCodexResponse { + type Target = CodexResponse; + + fn deref(&self) -> &Self::Target { + &self.response + } +} + // --------------------------------------------------------------------------- // Client // --------------------------------------------------------------------------- @@ -882,6 +908,24 @@ impl CodexHttpClient { ctx: &RequestContext, continuation: Option<&super::continuation::ContinuationCandidate>, ) -> Result { + let reservation = + continuation.map(super::continuation::ContinuationReservation::from_public_candidate); + self.post_codex_with_transport( + body, + ctx, + reservation.as_ref(), + crate::config::codex_transport(), + ) + .await + .map(OwnerAwareCodexResponse::into_response) + } + + pub(crate) async fn post_codex_for_owner( + &self, + body: &ResponsesRequest, + ctx: &RequestContext, + continuation: Option<&super::continuation::ContinuationReservation>, + ) -> Result { self.post_codex_with_transport(body, ctx, continuation, crate::config::codex_transport()) .await } @@ -952,9 +996,9 @@ impl CodexHttpClient { &self, body: &ResponsesRequest, ctx: &RequestContext, - continuation: Option<&super::continuation::ContinuationCandidate>, + continuation: Option<&super::continuation::ContinuationReservation>, transport: crate::config::CodexTransport, - ) -> Result { + ) -> Result { use crate::config::CodexTransport; let mut auth = self.auth_manager.get_auth().await.map_err(|e| CodexError { @@ -969,18 +1013,16 @@ impl CodexHttpClient { if should_reset_websocket_pool(continuation) && let Some(owner) = initial_pool_owner.as_ref() { - super::websocket::invalidate_codex_websocket_pool_turn( + super::websocket::invalidate_codex_websocket_pool_turn_for_owner( owner, - continuation.and_then(|candidate| candidate.turn_id), + continuation.and_then(super::continuation::ContinuationReservation::turn_id), ); } - let turn_id = continuation.and_then(|candidate| candidate.turn_id); let mut active_continuation = continuation.cloned(); let mut auth_refresh_attempted = false; let mut transport_failures = 0u32; loop { - let pool_owner = websocket_pool_owner(active_continuation.as_ref()).cloned(); let result = match transport { CodexTransport::Http => { let body_json = serde_json::to_string(body).map_err(|e| CodexError { @@ -992,12 +1034,18 @@ impl CodexHttpClient { })?; self.attempt_post_http(&auth, &body_json, ctx, body.client_metadata.is_some()) .await + .map(|response| OwnerAwareCodexResponse::new(response, None)) } CodexTransport::WebSocket => { let ws_headers = build_codex_headers(&auth, ctx, body.client_metadata.is_some())?; let ws_headers = super::websocket::codex_websocket_headers(&ws_headers); - let ws_body = build_websocket_request(body, active_continuation.as_ref()); + let ws_body = build_websocket_request( + body, + active_continuation + .as_ref() + .map(super::continuation::ContinuationReservation::candidate), + ); super::websocket::codex_websocket_request( &self.websocket_client, @@ -1007,7 +1055,6 @@ impl CodexHttpClient { &ws_body, ctx, ctx.traffic.as_deref(), - pool_owner.as_ref(), super::websocket::WEBSOCKET_CONNECT_TIMEOUT_MS, super::websocket::WEBSOCKET_IDLE_TIMEOUT_MS, active_continuation.as_ref(), @@ -1018,7 +1065,12 @@ impl CodexHttpClient { let ws_headers = build_codex_headers(&auth, ctx, body.client_metadata.is_some())?; let ws_headers = super::websocket::codex_websocket_headers(&ws_headers); - let ws_body = build_websocket_request(body, active_continuation.as_ref()); + let ws_body = build_websocket_request( + body, + active_continuation + .as_ref() + .map(super::continuation::ContinuationReservation::candidate), + ); // Try WebSocket first let ws_result = super::websocket::codex_websocket_request( @@ -1029,7 +1081,6 @@ impl CodexHttpClient { &ws_body, ctx, ctx.traffic.as_deref(), - pool_owner.as_ref(), super::websocket::WEBSOCKET_CONNECT_TIMEOUT_MS, super::websocket::WEBSOCKET_IDLE_TIMEOUT_MS, active_continuation.as_ref(), @@ -1038,6 +1089,16 @@ impl CodexHttpClient { match ws_result { Ok(response) => Ok(response), + Err(err) + if should_retry_without_continuation( + &err, + active_continuation.as_ref(), + ) => + { + // Drop stale continuation state before considering a + // replacement transport or connection. + Err(err) + } Err(err) if self.auto_http_fallback_enabled && should_fallback_to_http(&err) => { @@ -1057,6 +1118,7 @@ impl CodexHttpClient { body.client_metadata.is_some(), ) .await + .map(|response| OwnerAwareCodexResponse::new(response, None)) } Err(err) => Err(err), } @@ -1068,9 +1130,7 @@ impl CodexHttpClient { match self.auth_manager.force_refresh(&auth.access).await { Ok(new_auth) => { auth = new_auth; - if let Some(owner) = pool_owner.as_ref() { - super::websocket::invalidate_codex_websocket_pool_turn(owner, turn_id); - } + invalidate_live_continuation_pool(active_continuation.as_ref()); active_continuation = full_context_continuation(active_continuation.as_ref()); continue; @@ -1215,7 +1275,7 @@ impl CodexHttpClient { .map(|(_, value)| value.as_str()); let delay = compute_backoff_delay(transport_failures, retry_after); if delay.exceeds_budget { - return Err(codex_status_error(response)); + return Err(codex_status_error(response.into_response())); } log_buffered_retry( ctx, @@ -1237,18 +1297,15 @@ impl CodexHttpClient { "upstream", "retryable upstream status", ); - return Err(codex_status_error(response)); + return Err(codex_status_error(response.into_response())); } Ok(response) if !(200..300).contains(&response.status) => { - return Err(codex_status_error(response)); + return Err(codex_status_error(response.into_response())); } Ok(response) => return Ok(response), Err(err) if should_retry_without_continuation(&err, active_continuation.as_ref()) => { - if let Some(owner) = pool_owner.as_ref() { - super::websocket::invalidate_codex_websocket_pool_turn(owner, turn_id); - } active_continuation = full_context_continuation(active_continuation.as_ref()); continue; } @@ -1295,6 +1352,19 @@ impl CodexHttpClient { ctx: &RequestContext, continuation: Option<&super::continuation::ContinuationCandidate>, ) -> Result { + let reservation = + continuation.map(super::continuation::ContinuationReservation::from_public_candidate); + self.stream_codex_websocket_events_for_owner(body, ctx, reservation.as_ref()) + .await + .map(super::websocket::CodexWebSocketEventStream::into_receiver) + } + + pub(crate) async fn stream_codex_websocket_events_for_owner( + self: &Arc, + body: &ResponsesRequest, + ctx: &RequestContext, + continuation: Option<&super::continuation::ContinuationReservation>, + ) -> Result { let auth = self.auth_manager.get_auth().await.map_err(|e| CodexError { status: 401, message: "Auth error".to_string(), @@ -1303,12 +1373,11 @@ impl CodexHttpClient { origin: CodexErrorOrigin::Auth, })?; - let turn_id = continuation.and_then(|candidate| candidate.turn_id); - let pool_owner = websocket_pool_owner(continuation).cloned(); + let turn_id = continuation.and_then(super::continuation::ContinuationReservation::turn_id); if should_reset_websocket_pool(continuation) - && let Some(owner) = pool_owner.as_ref() + && let Some(owner) = websocket_pool_owner(continuation) { - super::websocket::invalidate_codex_websocket_pool_turn(owner, turn_id); + super::websocket::invalidate_codex_websocket_pool_turn_for_owner(owner, turn_id); } let client = self.clone(); @@ -1316,47 +1385,61 @@ impl CodexHttpClient { let ctx = ctx.clone(); let continuation = continuation.cloned(); let (tx, rx) = tokio::sync::mpsc::channel(64); + let (rx, socket_id_publisher) = super::websocket::CodexWebSocketEventStream::pending(rx); tokio::spawn(async move { client - .coordinate_live_websocket_events(body, ctx, continuation, auth, pool_owner, tx) + .coordinate_live_websocket_events( + body, + ctx, + continuation, + auth, + tx, + socket_id_publisher, + ) .await; }); Ok(rx) } + #[allow(clippy::too_many_arguments)] async fn coordinate_live_websocket_events( &self, body: ResponsesRequest, ctx: RequestContext, - mut continuation: Option, + mut continuation: Option, mut auth: StoredAuth, - pool_owner: Option, tx: tokio::sync::mpsc::Sender>, + socket_id_publisher: super::websocket::CodexWebSocketSocketIdPublisher, ) { - let turn_id = continuation - .as_ref() - .and_then(|candidate| candidate.turn_id); - let continuation_owner = continuation - .as_ref() - .and_then(|candidate| candidate.owner.clone()); let mut auth_refresh_attempted = false; let mut continuation_retry_available = continuation .as_ref() - .and_then(|candidate| candidate.previous_response_id.as_deref()) + .and_then(|reservation| reservation.candidate().previous_response_id.as_deref()) .is_some(); let mut forwarded_any = false; 'attempt: loop { + socket_id_publisher.publish(None); let ws_headers = match build_codex_headers(&auth, &ctx, body.client_metadata.is_some()) { Ok(headers) => super::websocket::codex_websocket_headers(&headers), Err(err) => { - let _ = tx.send(Err(err)).await; + if tx.send(Err(err)).await.is_err() { + abort_abandoned_live_continuation( + continuation.as_ref(), + &socket_id_publisher, + ); + } return; } }; - let ws_body = build_websocket_request(&body, continuation.as_ref()); + let ws_body = build_websocket_request( + &body, + continuation + .as_ref() + .map(super::continuation::ContinuationReservation::candidate), + ); let start = super::websocket::codex_websocket_event_stream( &self.websocket_client, &self.websocket_proxy_config, @@ -1365,52 +1448,66 @@ impl CodexHttpClient { &ws_body, &ctx, ctx.traffic.clone(), - pool_owner.as_ref(), super::websocket::WEBSOCKET_CONNECT_TIMEOUT_MS, super::websocket::WEBSOCKET_IDLE_TIMEOUT_MS, continuation.as_ref(), ); let mut stream = tokio::select! { + biased; _ = tx.closed() => { - super::continuation::abort_continuation(continuation_owner.as_ref(), turn_id); - if let Some(owner) = pool_owner.as_ref() { - super::websocket::invalidate_codex_websocket_pool_turn(owner, turn_id); - } + abort_abandoned_live_continuation( + continuation.as_ref(), + &socket_id_publisher, + ); return; } result = start => match result { Ok(stream) => stream, Err(err) if err.status == 401 && !auth_refresh_attempted && !forwarded_any => { auth_refresh_attempted = true; - if let Some(owner) = pool_owner.as_ref() { - super::websocket::invalidate_codex_websocket_pool_turn(owner, turn_id); - } + invalidate_live_continuation_pool(continuation.as_ref()); let refresh = self.auth_manager.force_refresh(&auth.access); auth = match refresh.await { Ok(auth) => { if tx.is_closed() { + abort_abandoned_live_continuation( + continuation.as_ref(), + &socket_id_publisher, + ); return; } auth }, Err(refresh_err) => { - let _ = tx.send(Err(auth_refresh_error(refresh_err))).await; + if tx.send(Err(auth_refresh_error(refresh_err))).await.is_err() { + abort_abandoned_live_continuation( + continuation.as_ref(), + &socket_id_publisher, + ); + } return; } }; + if continuation_retry_available { + socket_id_publisher.mark_full_context_retry(); + } continuation = full_context_continuation(continuation.as_ref()); + continuation_retry_available = false; continue 'attempt; } Err(err) if continuation_retry_available && is_continuation_retry_error(&err) => { - continuation_retry_available = false; - if let Some(owner) = pool_owner.as_ref() { - super::websocket::invalidate_codex_websocket_pool_turn(owner, turn_id); - } + socket_id_publisher.mark_full_context_retry(); continuation = full_context_continuation(continuation.as_ref()); + continuation_retry_available = false; continue 'attempt; } Err(err) => { - let _ = tx.send(Err(err)).await; + if tx.send(Err(err)).await.is_err() { + abort_abandoned_live_continuation( + continuation.as_ref(), + &socket_id_publisher, + ); + } return; } } @@ -1418,17 +1515,36 @@ impl CodexHttpClient { loop { let item = tokio::select! { + biased; _ = tx.closed() => { - if let Some(owner) = pool_owner.as_ref() { - super::websocket::invalidate_codex_websocket_pool_turn(owner, turn_id); + if let Some(reservation) = continuation.as_ref() { + super::websocket::invalidate_codex_websocket_pool_socket( + reservation, + stream.socket_id(), + ); } + abort_abandoned_live_continuation( + continuation.as_ref(), + &socket_id_publisher, + ); return; } item = stream.recv() => item, }; + if tx.is_closed() { + if let Some(reservation) = continuation.as_ref() { + super::websocket::invalidate_codex_websocket_pool_socket( + reservation, + stream.socket_id(), + ); + } + abort_abandoned_live_continuation(continuation.as_ref(), &socket_id_publisher); + return; + } let Some(item) = item else { return; }; + socket_id_publisher.publish(stream.socket_id()); let unauthorized = match &item { Err(err) => err.status == 401, @@ -1436,22 +1552,46 @@ impl CodexHttpClient { }; if unauthorized && !auth_refresh_attempted && !forwarded_any { auth_refresh_attempted = true; - if let Some(owner) = pool_owner.as_ref() { - super::websocket::invalidate_codex_websocket_pool_turn(owner, turn_id); - } + invalidate_live_continuation_pool(continuation.as_ref()); let refresh = self.auth_manager.force_refresh(&auth.access); auth = match refresh.await { Ok(auth) => { if tx.is_closed() { + if let Some(reservation) = continuation.as_ref() { + super::websocket::invalidate_codex_websocket_pool_socket( + reservation, + stream.socket_id(), + ); + } + abort_abandoned_live_continuation( + continuation.as_ref(), + &socket_id_publisher, + ); return; } auth } Err(refresh_err) => { - let _ = tx.send(Err(auth_refresh_error(refresh_err))).await; + if tx.send(Err(auth_refresh_error(refresh_err))).await.is_err() { + if let Some(reservation) = continuation.as_ref() { + super::websocket::invalidate_codex_websocket_pool_socket( + reservation, + stream.socket_id(), + ); + } + abort_abandoned_live_continuation( + continuation.as_ref(), + &socket_id_publisher, + ); + } return; } }; + if continuation_retry_available { + socket_id_publisher.mark_full_context_retry(); + } + continuation = full_context_continuation(continuation.as_ref()); + continuation_retry_available = false; continue 'attempt; } @@ -1460,22 +1600,28 @@ impl CodexHttpClient { && is_continuation_retry_error(err) && !forwarded_any { - continuation_retry_available = false; - if let Some(owner) = pool_owner.as_ref() { - super::websocket::invalidate_codex_websocket_pool_turn(owner, turn_id); - } + socket_id_publisher.mark_full_context_retry(); continuation = full_context_continuation(continuation.as_ref()); + continuation_retry_available = false; continue 'attempt; } if item.as_ref().is_ok_and(event_closes_live_retry_window) { forwarded_any = true; } + let terminal = item.as_ref().is_err() + || item.as_ref().is_ok_and(super::websocket::is_terminal_event); if tx.send(item).await.is_err() { - super::continuation::abort_continuation(continuation_owner.as_ref(), turn_id); - if let Some(owner) = pool_owner.as_ref() { - super::websocket::invalidate_codex_websocket_pool_turn(owner, turn_id); + if let Some(reservation) = continuation.as_ref() { + super::websocket::invalidate_codex_websocket_pool_socket( + reservation, + stream.socket_id(), + ); } + abort_abandoned_live_continuation(continuation.as_ref(), &socket_id_publisher); + return; + } + if terminal { return; } } @@ -1987,7 +2133,7 @@ fn is_retryable_reqwest_error(err: &reqwest::Error) -> bool { } fn should_refresh_after_unauthorized( - result: &Result, + result: &Result, auth_refresh_attempted: bool, transport: crate::config::CodexTransport, ) -> bool { @@ -2012,10 +2158,10 @@ fn should_fallback_to_http(err: &CodexError) -> bool { fn should_retry_without_continuation( err: &CodexError, - continuation: Option<&super::continuation::ContinuationCandidate>, + continuation: Option<&super::continuation::ContinuationReservation>, ) -> bool { if continuation - .and_then(|c| c.previous_response_id.as_deref()) + .and_then(|reservation| reservation.candidate().previous_response_id.as_deref()) .is_none() { return false; @@ -2025,16 +2171,36 @@ fn should_retry_without_continuation( } fn full_context_continuation( - continuation: Option<&super::continuation::ContinuationCandidate>, -) -> Option { - continuation.map(|candidate| super::continuation::ContinuationCandidate { - owner: candidate.owner.clone(), - turn_id: candidate.turn_id, - previous_response_id: None, - input_delta: None, - input_delta_count: candidate.input_delta_count, - disabled_reason: Some("full_context_retry".to_string()), - }) + continuation: Option<&super::continuation::ContinuationReservation>, +) -> Option { + continuation.map(super::continuation::ContinuationReservation::full_context_retry) +} + +fn abort_live_continuation(continuation: Option<&super::continuation::ContinuationReservation>) { + if let Some(continuation) = continuation { + super::continuation::abort_continuation_for_owner(continuation); + } +} + +fn abort_abandoned_live_continuation( + continuation: Option<&super::continuation::ContinuationReservation>, + socket_id_publisher: &super::websocket::CodexWebSocketSocketIdPublisher, +) { + if !socket_id_publisher.is_provider_retry_handoff() { + abort_live_continuation(continuation); + } +} + +fn invalidate_live_continuation_pool( + continuation: Option<&super::continuation::ContinuationReservation>, +) { + let Some(continuation) = continuation else { + return; + }; + let Some(owner) = websocket_pool_owner(Some(continuation)) else { + return; + }; + super::websocket::invalidate_codex_websocket_pool_turn_for_owner(owner, continuation.turn_id()); } fn event_closes_live_retry_window(payload: &serde_json::Value) -> bool { @@ -2044,29 +2210,32 @@ fn event_closes_live_retry_window(payload: &serde_json::Value) -> bool { ) } -fn is_continuation_retry_error(err: &CodexError) -> bool { +pub(super) fn is_continuation_retry_error(err: &CodexError) -> bool { matches!( err.detail.as_deref(), Some("previous_response_not_found") + | Some(super::websocket::WEBSOCKET_CONTINUATION_SOCKET_MISSING_DETAIL) | Some(super::websocket::WEBSOCKET_RESPONSE_START_TIMEOUT_DETAIL) | Some(super::websocket::WEBSOCKET_MISSING_TERMINAL_DETAIL) ) } fn websocket_pool_owner( - continuation: Option<&super::continuation::ContinuationCandidate>, + continuation: Option<&super::continuation::ContinuationReservation>, ) -> Option<&ConversationIdentity> { let continuation = continuation?; - if continuation.disabled_reason.as_deref() == Some("disabled") { + if continuation.candidate().disabled_reason.as_deref() == Some("disabled") { return None; } - continuation.owner.as_ref() + continuation.owner() } fn should_reset_websocket_pool( - continuation: Option<&super::continuation::ContinuationCandidate>, + continuation: Option<&super::continuation::ContinuationReservation>, ) -> bool { - let Some(reason) = continuation.and_then(|c| c.disabled_reason.as_deref()) else { + let Some(reason) = + continuation.and_then(|continuation| continuation.candidate().disabled_reason.as_deref()) + else { return false; }; reason != "disabled" @@ -2075,9 +2244,30 @@ fn should_reset_websocket_pool( #[cfg(test)] mod tests { use super::*; + use futures_util::{SinkExt, StreamExt}; use tokio::io::{AsyncReadExt, AsyncWriteExt}; use tokio::net::TcpListener; + fn test_continuation( + owner: Option, + turn_id: Option, + previous_response_id: Option<&str>, + origin_socket_id: Option, + disabled_reason: Option<&str>, + ) -> super::super::continuation::ContinuationReservation { + super::super::continuation::ContinuationReservation::new( + super::super::continuation::ContinuationCandidate { + turn_id, + previous_response_id: previous_response_id.map(str::to_string), + input_delta: None, + input_delta_count: 1, + disabled_reason: disabled_reason.map(str::to_string), + }, + owner, + origin_socket_id, + ) + } + #[test] fn normalizes_supported_proxy_urls() { assert_eq!( @@ -2180,6 +2370,84 @@ mod tests { } } + fn buffered_request_with_texts(texts: &[&str]) -> ResponsesRequest { + let mut request = buffered_test_request(); + request.input = texts + .iter() + .map( + |text| super::super::translate::request::ResponsesInputItem::Message { + role: "user".to_string(), + content: vec![ + super::super::translate::request::ResponsesContentPart::InputText { + text: (*text).to_string(), + }, + ], + }, + ) + .collect(); + request + } + + async fn next_websocket_json( + websocket: &mut tokio_tungstenite::WebSocketStream, + ) -> serde_json::Value { + loop { + match websocket.next().await { + Some(Ok(tokio_tungstenite::tungstenite::Message::Ping(payload))) => { + websocket + .send(tokio_tungstenite::tungstenite::Message::Pong(payload)) + .await + .unwrap(); + } + Some(Ok(tokio_tungstenite::tungstenite::Message::Text(text))) => { + return serde_json::from_str(&text).unwrap(); + } + other => panic!("unexpected WebSocket request frame: {other:?}"), + } + } + } + + async fn send_completed_websocket_response( + websocket: &mut tokio_tungstenite::WebSocketStream, + response_id: &str, + ) { + websocket + .send(tokio_tungstenite::tungstenite::Message::Text( + serde_json::json!({ + "type": "response.completed", + "response": { + "id": response_id, + "status": "completed", + "output": [] + } + }) + .to_string(), + )) + .await + .unwrap(); + } + + async fn send_nested_previous_response_missing( + websocket: &mut tokio_tungstenite::WebSocketStream, + ) { + websocket + .send(tokio_tungstenite::tungstenite::Message::Text( + serde_json::json!({ + "type": "response.failed", + "response": { + "status": "failed", + "error": { + "code": "previous_response_not_found", + "message": "Previous response not found" + } + } + }) + .to_string(), + )) + .await + .unwrap(); + } + fn authenticated_http_test_client(base_url: String) -> CodexHttpClient { let client = http_test_client(base_url, 100); client.auth_manager().set_test_auth(http_test_auth()); @@ -2476,6 +2744,849 @@ mod tests { ); } + #[tokio::test] + async fn buffered_missing_origin_retries_full_context_and_rebinds_exact_socket() { + let _registry_guard = + super::super::continuation::lock_continuation_registry_for_async_tests().await; + let _pool_guard = super::super::websocket::lock_codex_websocket_pool_for_tests().await; + let owner = ConversationIdentity::Agent( + "buffered-recovery-session".to_string(), + "buffered-recovery-agent".to_string(), + ); + super::super::continuation::clear_continuation_for_owner(Some(&owner)); + super::super::websocket::invalidate_codex_websocket_pool_owner(&owner); + + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = listener.local_addr().unwrap(); + let (request_tx, mut request_rx) = tokio::sync::mpsc::unbounded_channel(); + let server = tokio::spawn(async move { + let (first_socket, _) = listener.accept().await.unwrap(); + let mut first_websocket = tokio_tungstenite::accept_async(first_socket).await.unwrap(); + request_tx + .send(next_websocket_json(&mut first_websocket).await) + .unwrap(); + send_completed_websocket_response(&mut first_websocket, "resp_a").await; + drop(first_websocket); + + let (second_socket, _) = listener.accept().await.unwrap(); + let mut second_websocket = tokio_tungstenite::accept_async(second_socket) + .await + .unwrap(); + request_tx + .send(next_websocket_json(&mut second_websocket).await) + .unwrap(); + send_completed_websocket_response(&mut second_websocket, "resp_b").await; + request_tx + .send(next_websocket_json(&mut second_websocket).await) + .unwrap(); + send_completed_websocket_response(&mut second_websocket, "resp_c").await; + }); + + let client = authenticated_http_test_client(format!("http://{addr}/responses")); + let context = http_test_context(); + let first_request = buffered_request_with_texts(&["one"]); + let first_candidate = super::super::continuation::continuation_candidate_for_owner( + Some(&owner), + &first_request, + true, + ); + let first_response = client + .post_codex_with_transport( + &first_request, + &context, + Some(&first_candidate), + crate::config::CodexTransport::WebSocket, + ) + .await + .unwrap(); + let first_socket_id = first_response + .socket_id + .expect("first socket must be reusable"); + super::super::update_continuation_from_upstream( + None, + &first_candidate, + None, + &first_request, + &first_response.body, + first_response.socket_id, + false, + ); + + let second_request = buffered_request_with_texts(&["one", "two"]); + let second_candidate = super::super::continuation::continuation_candidate_for_owner( + Some(&owner), + &second_request, + true, + ); + assert_eq!(second_candidate.origin_socket_id(), Some(first_socket_id)); + super::super::websocket::invalidate_codex_websocket_pool_owner(&owner); + let second_response = client + .post_codex_with_transport( + &second_request, + &context, + Some(&second_candidate), + crate::config::CodexTransport::WebSocket, + ) + .await + .unwrap(); + let second_socket_id = second_response + .socket_id + .expect("full-context retry socket must be reusable"); + assert_ne!(second_socket_id, first_socket_id); + super::super::update_continuation_from_upstream( + None, + &second_candidate, + None, + &second_request, + &second_response.body, + second_response.socket_id, + false, + ); + + let third_request = buffered_request_with_texts(&["one", "two", "three"]); + let third_candidate = super::super::continuation::continuation_candidate_for_owner( + Some(&owner), + &third_request, + true, + ); + assert_eq!( + third_candidate.candidate().previous_response_id.as_deref(), + Some("resp_b") + ); + assert_eq!(third_candidate.origin_socket_id(), Some(second_socket_id)); + let third_response = client + .post_codex_with_transport( + &third_request, + &context, + Some(&third_candidate), + crate::config::CodexTransport::WebSocket, + ) + .await + .unwrap(); + assert_eq!(third_response.socket_id, Some(second_socket_id)); + + let first_payload = request_rx.recv().await.unwrap(); + let retry_payload = request_rx.recv().await.unwrap(); + let continued_payload = request_rx.recv().await.unwrap(); + assert!(first_payload.get("previous_response_id").is_none()); + assert_eq!(first_payload["input"].as_array().unwrap().len(), 1); + assert!(retry_payload.get("previous_response_id").is_none()); + assert_eq!(retry_payload["input"].as_array().unwrap().len(), 2); + assert_eq!(continued_payload["previous_response_id"], "resp_b"); + assert_eq!(continued_payload["input"].as_array().unwrap().len(), 1); + server.await.unwrap(); + + super::super::websocket::invalidate_codex_websocket_pool_owner(&owner); + super::super::continuation::abort_continuation_for_owner(&third_candidate); + } + + #[tokio::test] + async fn buffered_nested_missing_response_retries_once_without_stale_previous_id() { + let _registry_guard = + super::super::continuation::lock_continuation_registry_for_async_tests().await; + let _pool_guard = super::super::websocket::lock_codex_websocket_pool_for_tests().await; + let owner = ConversationIdentity::Main("buffered-nested-recovery-session".to_string()); + super::super::continuation::clear_continuation_for_owner(Some(&owner)); + super::super::websocket::invalidate_codex_websocket_pool_owner(&owner); + + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = listener.local_addr().unwrap(); + let (request_tx, mut request_rx) = tokio::sync::mpsc::unbounded_channel(); + let server = tokio::spawn(async move { + let (first_socket, _) = listener.accept().await.unwrap(); + let mut first_websocket = tokio_tungstenite::accept_async(first_socket).await.unwrap(); + request_tx + .send(next_websocket_json(&mut first_websocket).await) + .unwrap(); + send_completed_websocket_response(&mut first_websocket, "resp_nested_origin").await; + request_tx + .send(next_websocket_json(&mut first_websocket).await) + .unwrap(); + send_nested_previous_response_missing(&mut first_websocket).await; + + let (retry_socket, _) = listener.accept().await.unwrap(); + let mut retry_websocket = tokio_tungstenite::accept_async(retry_socket).await.unwrap(); + request_tx + .send(next_websocket_json(&mut retry_websocket).await) + .unwrap(); + send_nested_previous_response_missing(&mut retry_websocket).await; + assert!( + tokio::time::timeout(Duration::from_millis(1_200), listener.accept()) + .await + .is_err(), + "nested missing-response recovery must not open a third socket" + ); + }); + + let client = authenticated_http_test_client(format!("http://{addr}/responses")); + let context = http_test_context(); + let first_request = buffered_request_with_texts(&["one"]); + let first_candidate = super::super::continuation::continuation_candidate_for_owner( + Some(&owner), + &first_request, + true, + ); + let first_response = client + .post_codex_with_transport( + &first_request, + &context, + Some(&first_candidate), + crate::config::CodexTransport::WebSocket, + ) + .await + .unwrap(); + super::super::update_continuation_from_upstream( + None, + &first_candidate, + None, + &first_request, + &first_response.body, + first_response.socket_id, + false, + ); + + let second_request = buffered_request_with_texts(&["one", "two"]); + let second_candidate = super::super::continuation::continuation_candidate_for_owner( + Some(&owner), + &second_request, + true, + ); + assert_eq!( + second_candidate.candidate().previous_response_id.as_deref(), + Some("resp_nested_origin") + ); + let error = match client + .post_codex_with_transport( + &second_request, + &context, + Some(&second_candidate), + crate::config::CodexTransport::WebSocket, + ) + .await + { + Ok(_) => { + panic!("the bounded full-context retry must surface a repeated nested failure") + } + Err(error) => error, + }; + assert_eq!(error.detail.as_deref(), Some("previous_response_not_found")); + + let first_payload = request_rx.recv().await.unwrap(); + let continued_payload = request_rx.recv().await.unwrap(); + let retry_payload = request_rx.recv().await.unwrap(); + assert!(first_payload.get("previous_response_id").is_none()); + assert_eq!( + continued_payload["previous_response_id"], + "resp_nested_origin" + ); + assert_eq!(continued_payload["input"].as_array().unwrap().len(), 1); + assert!(retry_payload.get("previous_response_id").is_none()); + assert_eq!(retry_payload["input"].as_array().unwrap().len(), 2); + assert!(request_rx.try_recv().is_err()); + server.await.unwrap(); + + super::super::websocket::invalidate_codex_websocket_pool_owner(&owner); + super::super::continuation::abort_continuation_for_owner(&second_candidate); + } + + #[tokio::test] + async fn live_nested_missing_response_retries_once_without_stale_previous_id() { + let _registry_guard = + super::super::continuation::lock_continuation_registry_for_async_tests().await; + let _pool_guard = super::super::websocket::lock_codex_websocket_pool_for_tests().await; + let owner = ConversationIdentity::Agent( + "live-nested-recovery-session".to_string(), + "live-nested-recovery-agent".to_string(), + ); + super::super::continuation::clear_continuation_for_owner(Some(&owner)); + super::super::websocket::invalidate_codex_websocket_pool_owner(&owner); + + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = listener.local_addr().unwrap(); + let (request_tx, mut request_rx) = tokio::sync::mpsc::unbounded_channel(); + let server = tokio::spawn(async move { + let (first_socket, _) = listener.accept().await.unwrap(); + let mut first_websocket = tokio_tungstenite::accept_async(first_socket).await.unwrap(); + request_tx + .send(next_websocket_json(&mut first_websocket).await) + .unwrap(); + send_completed_websocket_response(&mut first_websocket, "resp_live_nested_origin") + .await; + request_tx + .send(next_websocket_json(&mut first_websocket).await) + .unwrap(); + send_nested_previous_response_missing(&mut first_websocket).await; + + let (retry_socket, _) = listener.accept().await.unwrap(); + let mut retry_websocket = tokio_tungstenite::accept_async(retry_socket).await.unwrap(); + request_tx + .send(next_websocket_json(&mut retry_websocket).await) + .unwrap(); + send_nested_previous_response_missing(&mut retry_websocket).await; + assert!( + tokio::time::timeout(Duration::from_millis(1_200), listener.accept()) + .await + .is_err(), + "live nested missing-response recovery must not open a third socket" + ); + }); + + let client = Arc::new(authenticated_http_test_client(format!( + "http://{addr}/responses" + ))); + let context = http_test_context(); + let first_request = buffered_request_with_texts(&["one"]); + let first_candidate = super::super::continuation::continuation_candidate_for_owner( + Some(&owner), + &first_request, + true, + ); + let first_response = client + .post_codex_with_transport( + &first_request, + &context, + Some(&first_candidate), + crate::config::CodexTransport::WebSocket, + ) + .await + .unwrap(); + super::super::update_continuation_from_upstream( + None, + &first_candidate, + None, + &first_request, + &first_response.body, + first_response.socket_id, + false, + ); + + let second_request = buffered_request_with_texts(&["one", "two"]); + let second_candidate = super::super::continuation::continuation_candidate_for_owner( + Some(&owner), + &second_request, + true, + ); + assert_eq!( + second_candidate.candidate().previous_response_id.as_deref(), + Some("resp_live_nested_origin") + ); + let mut events = client + .stream_codex_websocket_events_for_owner( + &second_request, + &context, + Some(&second_candidate), + ) + .await + .unwrap(); + let error = events.recv().await.unwrap().unwrap_err(); + assert_eq!(error.detail.as_deref(), Some("previous_response_not_found")); + assert!(events.used_full_context_retry()); + assert_eq!(events.socket_id(), None); + + let first_payload = request_rx.recv().await.unwrap(); + let continued_payload = request_rx.recv().await.unwrap(); + let retry_payload = request_rx.recv().await.unwrap(); + assert!(first_payload.get("previous_response_id").is_none()); + assert_eq!( + continued_payload["previous_response_id"], + "resp_live_nested_origin" + ); + assert_eq!(continued_payload["input"].as_array().unwrap().len(), 1); + assert!(retry_payload.get("previous_response_id").is_none()); + assert_eq!(retry_payload["input"].as_array().unwrap().len(), 2); + assert!(request_rx.try_recv().is_err()); + server.await.unwrap(); + + super::super::websocket::invalidate_codex_websocket_pool_owner(&owner); + super::super::continuation::abort_continuation_for_owner(&second_candidate); + } + + #[tokio::test] + async fn live_missing_origin_retries_once_with_full_context_and_actual_socket() { + let _registry_guard = + super::super::continuation::lock_continuation_registry_for_async_tests().await; + let _pool_guard = super::super::websocket::lock_codex_websocket_pool_for_tests().await; + let owner = ConversationIdentity::Main("live-recovery-session".to_string()); + super::super::continuation::clear_continuation_for_owner(Some(&owner)); + super::super::websocket::invalidate_codex_websocket_pool_owner(&owner); + let request = buffered_request_with_texts(&["one", "two"]); + let reserved = super::super::continuation::continuation_candidate_for_owner( + Some(&owner), + &request, + true, + ); + let continuation = super::super::continuation::ContinuationReservation::new( + super::super::continuation::ContinuationCandidate { + turn_id: reserved.turn_id(), + previous_response_id: Some("resp_missing".to_string()), + input_delta: Some(vec![request.input.last().unwrap().clone()]), + input_delta_count: 1, + disabled_reason: None, + }, + Some(owner.clone()), + Some(u64::MAX), + ); + + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = listener.local_addr().unwrap(); + let (payload_tx, payload_rx) = tokio::sync::oneshot::channel(); + let server = tokio::spawn(async move { + let (socket, _) = listener.accept().await.unwrap(); + let mut websocket = tokio_tungstenite::accept_async(socket).await.unwrap(); + payload_tx + .send(next_websocket_json(&mut websocket).await) + .unwrap(); + send_completed_websocket_response(&mut websocket, "resp_live_retry").await; + assert!( + tokio::time::timeout(Duration::from_millis(100), listener.accept()) + .await + .is_err() + ); + }); + + let client = Arc::new(authenticated_http_test_client(format!( + "http://{addr}/responses" + ))); + let mut events = client + .stream_codex_websocket_events_for_owner( + &request, + &http_test_context(), + Some(&continuation), + ) + .await + .unwrap(); + let terminal = events.recv().await.unwrap().unwrap(); + assert_eq!(terminal["type"], "response.completed"); + assert!(events.used_full_context_retry()); + let socket_id = events + .socket_id() + .expect("successful internal retry must publish its socket"); + let payload = payload_rx.await.unwrap(); + assert!(payload.get("previous_response_id").is_none()); + assert_eq!(payload["input"].as_array().unwrap().len(), 2); + server.await.unwrap(); + assert_eq!( + super::super::websocket::pooled_socket_id_for_tests(&owner), + Some(socket_id) + ); + + super::super::websocket::invalidate_codex_websocket_pool_owner(&owner); + super::super::continuation::abort_continuation_for_owner(&reserved); + } + + #[tokio::test] + async fn public_ownerless_candidate_never_writes_previous_id_to_websocket() { + let _pool_guard = super::super::websocket::lock_codex_websocket_pool_for_tests().await; + super::super::websocket::clear_codex_websocket_pool_for_tests(); + let request = buffered_request_with_texts(&["one", "two"]); + let continuation = super::super::continuation::ContinuationCandidate { + turn_id: Some(17), + previous_response_id: Some("resp_unproven".to_string()), + input_delta: Some(vec![request.input.last().unwrap().clone()]), + input_delta_count: 1, + disabled_reason: None, + }; + + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = listener.local_addr().unwrap(); + let (payload_tx, payload_rx) = tokio::sync::oneshot::channel(); + let server = tokio::spawn(async move { + let (socket, _) = listener.accept().await.unwrap(); + let mut websocket = tokio_tungstenite::accept_async(socket).await.unwrap(); + payload_tx + .send(next_websocket_json(&mut websocket).await) + .unwrap(); + send_completed_websocket_response(&mut websocket, "resp_public_full_context").await; + assert!( + tokio::time::timeout(Duration::from_millis(100), listener.accept()) + .await + .is_err(), + "an ownerless stale candidate must use one full-context socket" + ); + }); + + let client = Arc::new(authenticated_http_test_client(format!( + "http://{addr}/responses" + ))); + let mut context = http_test_context(); + context.session_id = Some("must-not-be-derived-as-owner".to_string()); + let mut events = client + .stream_codex_websocket_events(&request, &context, Some(&continuation)) + .await + .unwrap(); + let terminal = events.recv().await.unwrap().unwrap(); + assert_eq!(terminal["type"], "response.completed"); + + let payload = payload_rx.await.unwrap(); + assert!(payload.get("previous_response_id").is_none()); + assert_eq!(payload["input"].as_array().unwrap().len(), 2); + server.await.unwrap(); + super::super::websocket::clear_codex_websocket_pool_for_tests(); + } + + #[tokio::test] + async fn live_missing_origin_does_not_enter_second_full_context_loop() { + let _registry_guard = + super::super::continuation::lock_continuation_registry_for_async_tests().await; + let _pool_guard = super::super::websocket::lock_codex_websocket_pool_for_tests().await; + let owner = ConversationIdentity::Main("live-bounded-recovery-session".to_string()); + super::super::continuation::clear_continuation_for_owner(Some(&owner)); + super::super::websocket::invalidate_codex_websocket_pool_owner(&owner); + let request = buffered_request_with_texts(&["one", "two"]); + let reserved = super::super::continuation::continuation_candidate_for_owner( + Some(&owner), + &request, + true, + ); + let continuation = super::super::continuation::ContinuationReservation::new( + super::super::continuation::ContinuationCandidate { + turn_id: reserved.turn_id(), + previous_response_id: Some("resp_missing".to_string()), + input_delta: Some(vec![request.input.last().unwrap().clone()]), + input_delta_count: 1, + disabled_reason: None, + }, + Some(owner.clone()), + Some(u64::MAX), + ); + + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = listener.local_addr().unwrap(); + let (payload_tx, payload_rx) = tokio::sync::oneshot::channel(); + let server = tokio::spawn(async move { + let (socket, _) = listener.accept().await.unwrap(); + let mut websocket = tokio_tungstenite::accept_async(socket).await.unwrap(); + payload_tx + .send(next_websocket_json(&mut websocket).await) + .unwrap(); + websocket.close(None).await.unwrap(); + assert!( + tokio::time::timeout(Duration::from_millis(150), listener.accept()) + .await + .is_err() + ); + }); + + let client = Arc::new(authenticated_http_test_client(format!( + "http://{addr}/responses" + ))); + let mut events = client + .stream_codex_websocket_events_for_owner( + &request, + &http_test_context(), + Some(&continuation), + ) + .await + .unwrap(); + let error = events.recv().await.unwrap().unwrap_err(); + assert_eq!( + error.detail.as_deref(), + Some(super::super::websocket::WEBSOCKET_MISSING_TERMINAL_DETAIL) + ); + assert!(events.used_full_context_retry()); + assert_eq!(events.socket_id(), None); + let payload = payload_rx.await.unwrap(); + assert!(payload.get("previous_response_id").is_none()); + assert_eq!(payload["input"].as_array().unwrap().len(), 2); + server.await.unwrap(); + + super::super::websocket::invalidate_codex_websocket_pool_owner(&owner); + super::super::continuation::abort_continuation_for_owner(&reserved); + } + + #[tokio::test] + async fn dropping_live_receiver_clears_reserved_turn() { + let _registry_guard = + super::super::continuation::lock_continuation_registry_for_async_tests().await; + let _pool_guard = super::super::websocket::lock_codex_websocket_pool_for_tests().await; + let owner = ConversationIdentity::Main("live-drop-session".to_string()); + super::super::continuation::clear_continuation_for_owner(Some(&owner)); + super::super::websocket::invalidate_codex_websocket_pool_owner(&owner); + let request = buffered_request_with_texts(&["one"]); + let continuation = super::super::continuation::continuation_candidate_for_owner( + Some(&owner), + &request, + true, + ); + + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = listener.local_addr().unwrap(); + let (request_seen_tx, request_seen_rx) = tokio::sync::oneshot::channel(); + let server = tokio::spawn(async move { + let (socket, _) = listener.accept().await.unwrap(); + let mut websocket = tokio_tungstenite::accept_async(socket).await.unwrap(); + let _ = next_websocket_json(&mut websocket).await; + request_seen_tx.send(()).unwrap(); + tokio::time::sleep(Duration::from_millis(100)).await; + let _ = websocket.close(None).await; + }); + + let events = Arc::new(authenticated_http_test_client(format!( + "http://{addr}/responses" + ))) + .stream_codex_websocket_events_for_owner( + &request, + &http_test_context(), + Some(&continuation), + ) + .await + .unwrap(); + request_seen_rx.await.unwrap(); + drop(events); + + tokio::time::timeout(Duration::from_secs(1), async { + while super::super::continuation::is_current_turn_for_owner(&continuation) { + tokio::task::yield_now().await; + } + }) + .await + .expect("dropping the live receiver must clear its reserved turn"); + server.await.unwrap(); + super::super::websocket::invalidate_codex_websocket_pool_owner(&owner); + } + + #[tokio::test] + async fn dropping_retry_handoff_receiver_preserves_reserved_turn() { + let _registry_guard = + super::super::continuation::lock_continuation_registry_for_async_tests().await; + let _pool_guard = super::super::websocket::lock_codex_websocket_pool_for_tests().await; + let owner = ConversationIdentity::Main("live-retry-handoff-session".to_string()); + super::super::continuation::clear_continuation_for_owner(Some(&owner)); + super::super::websocket::invalidate_codex_websocket_pool_owner(&owner); + let request = buffered_request_with_texts(&["one"]); + let continuation = super::super::continuation::continuation_candidate_for_owner( + Some(&owner), + &request, + true, + ); + + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = listener.local_addr().unwrap(); + let (request_seen_tx, request_seen_rx) = tokio::sync::oneshot::channel(); + let (socket_closed_tx, socket_closed_rx) = tokio::sync::oneshot::channel(); + let server = tokio::spawn(async move { + let (socket, _) = listener.accept().await.unwrap(); + let mut websocket = tokio_tungstenite::accept_async(socket).await.unwrap(); + let _ = next_websocket_json(&mut websocket).await; + request_seen_tx.send(()).unwrap(); + while websocket.next().await.is_some() {} + socket_closed_tx.send(()).unwrap(); + }); + + let events = Arc::new(authenticated_http_test_client(format!( + "http://{addr}/responses" + ))) + .stream_codex_websocket_events_for_owner( + &request, + &http_test_context(), + Some(&continuation), + ) + .await + .unwrap(); + request_seen_rx.await.unwrap(); + events.mark_provider_retry_handoff(); + drop(events); + + tokio::time::timeout(Duration::from_secs(1), socket_closed_rx) + .await + .expect("marked receiver drop must close only the abandoned attempt socket") + .expect("socket-close acknowledgement sender dropped"); + assert!(super::super::continuation::is_current_turn_for_owner( + &continuation + )); + server.await.unwrap(); + super::super::continuation::abort_continuation_for_owner(&continuation); + super::super::websocket::invalidate_codex_websocket_pool_owner(&owner); + } + + #[tokio::test] + async fn delayed_retry_handoff_cleanup_preserves_replacement_state_and_socket() { + let _registry_guard = + super::super::continuation::lock_continuation_registry_for_async_tests().await; + let _pool_guard = super::super::websocket::lock_codex_websocket_pool_for_tests().await; + let owner = ConversationIdentity::Main("live-retry-cleanup-race-session".to_string()); + super::super::continuation::clear_continuation_for_owner(Some(&owner)); + super::super::websocket::invalidate_codex_websocket_pool_owner(&owner); + let request = buffered_request_with_texts(&["one"]); + let continuation = super::super::continuation::continuation_candidate_for_owner( + Some(&owner), + &request, + true, + ); + + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = listener.local_addr().unwrap(); + let (release_replacement_tx, release_replacement_rx) = tokio::sync::oneshot::channel(); + let server = tokio::spawn(async move { + let (first_socket, _) = listener.accept().await.unwrap(); + let mut first_websocket = tokio_tungstenite::accept_async(first_socket).await.unwrap(); + let _ = next_websocket_json(&mut first_websocket).await; + send_completed_websocket_response(&mut first_websocket, "resp_attempt_a").await; + drop(first_websocket); + + let (replacement_socket, _) = listener.accept().await.unwrap(); + let mut replacement_websocket = tokio_tungstenite::accept_async(replacement_socket) + .await + .unwrap(); + let _ = next_websocket_json(&mut replacement_websocket).await; + send_completed_websocket_response(&mut replacement_websocket, "resp_attempt_b").await; + let _ = release_replacement_rx.await; + }); + let client = Arc::new(authenticated_http_test_client(format!( + "http://{addr}/responses" + ))); + + let mut attempt_a = client + .stream_codex_websocket_events_for_owner( + &request, + &http_test_context(), + Some(&continuation), + ) + .await + .unwrap(); + let terminal_a = attempt_a.recv().await.unwrap().unwrap(); + assert_eq!(terminal_a["response"]["id"], "resp_attempt_a"); + let socket_a = attempt_a.socket_id().expect("attempt A socket ID"); + super::super::websocket::invalidate_codex_websocket_pool_socket( + &continuation, + Some(socket_a), + ); + + let mut attempt_b = client + .stream_codex_websocket_events_for_owner( + &request, + &http_test_context(), + Some(&continuation), + ) + .await + .unwrap(); + let terminal_b = attempt_b.recv().await.unwrap().unwrap(); + assert_eq!(terminal_b["response"]["id"], "resp_attempt_b"); + let socket_b = attempt_b.socket_id().expect("attempt B socket ID"); + assert_ne!(socket_a, socket_b); + super::super::continuation::record_continuation_for_owner( + &continuation, + &request, + Some("resp_attempt_b"), + Some(socket_b), + &[], + ); + + let (_handoff_tx, handoff_rx) = + tokio::sync::mpsc::channel::>(1); + let (handoff_stream, handoff_publisher) = + super::super::websocket::CodexWebSocketEventStream::pending(handoff_rx); + handoff_stream.mark_provider_retry_handoff(); + let cleanup_barrier = Arc::new(tokio::sync::Barrier::new(2)); + let cleanup_task_barrier = cleanup_barrier.clone(); + let cleanup_continuation = continuation.clone(); + let cleanup = tokio::spawn(async move { + cleanup_task_barrier.wait().await; + super::super::websocket::invalidate_codex_websocket_pool_socket( + &cleanup_continuation, + Some(socket_a), + ); + abort_abandoned_live_continuation(Some(&cleanup_continuation), &handoff_publisher); + }); + + cleanup_barrier.wait().await; + cleanup.await.unwrap(); + assert!(super::super::continuation::is_current_turn_for_owner( + &continuation + )); + assert!(super::super::continuation::has_continuation_for_owner_for_tests(&owner)); + assert_eq!( + super::super::websocket::pooled_socket_id_for_tests(&owner), + Some(socket_b) + ); + + let _ = release_replacement_tx.send(()); + server.await.unwrap(); + super::super::continuation::abort_continuation_for_owner(&continuation); + super::super::websocket::invalidate_codex_websocket_pool_owner(&owner); + } + + #[tokio::test] + async fn auto_clears_missing_origin_before_ordinary_http_fallback() { + let _registry_guard = + super::super::continuation::lock_continuation_registry_for_async_tests().await; + let _pool_guard = super::super::websocket::lock_codex_websocket_pool_for_tests().await; + let owner = ConversationIdentity::Agent( + "auto-recovery-session".to_string(), + "auto-recovery-agent".to_string(), + ); + super::super::continuation::clear_continuation_for_owner(Some(&owner)); + super::super::websocket::invalidate_codex_websocket_pool_owner(&owner); + let request = buffered_request_with_texts(&["one", "two"]); + let reserved = super::super::continuation::continuation_candidate_for_owner( + Some(&owner), + &request, + true, + ); + let continuation = super::super::continuation::ContinuationReservation::new( + super::super::continuation::ContinuationCandidate { + turn_id: reserved.turn_id(), + previous_response_id: Some("resp_missing".to_string()), + input_delta: Some(vec![request.input.last().unwrap().clone()]), + input_delta_count: 1, + disabled_reason: None, + }, + Some(owner.clone()), + Some(u64::MAX), + ); + + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = listener.local_addr().unwrap(); + let server = tokio::spawn(async move { + let (mut websocket, _) = listener.accept().await.unwrap(); + let websocket_request = read_http_request(&mut websocket).await; + assert!(String::from_utf8_lossy(&websocket_request).starts_with("GET ")); + websocket + .write_all( + b"HTTP/1.1 400 Bad Request\r\ncontent-length: 0\r\nconnection: close\r\n\r\n", + ) + .await + .unwrap(); + drop(websocket); + + let (mut http, _) = listener.accept().await.unwrap(); + let http_request = read_http_request(&mut http).await; + let body_start = http_request + .windows(4) + .position(|part| part == b"\r\n\r\n") + .unwrap() + + 4; + let body: serde_json::Value = + serde_json::from_slice(&http_request[body_start..]).unwrap(); + assert!(body.get("previous_response_id").is_none()); + assert_eq!(body["input"].as_array().unwrap().len(), 2); + let response_body = + b"data: {\"type\":\"response.completed\",\"response\":{\"id\":\"resp_http\"}}\n\n"; + let response = format!( + "HTTP/1.1 200 OK\r\ncontent-length: {}\r\nconnection: close\r\n\r\n", + response_body.len() + ); + http.write_all(response.as_bytes()).await.unwrap(); + http.write_all(response_body).await.unwrap(); + }); + + let response = authenticated_http_test_client(format!("http://{addr}/responses")) + .post_codex_with_transport( + &request, + &http_test_context(), + Some(&continuation), + crate::config::CodexTransport::Auto, + ) + .await + .unwrap(); + assert_eq!(response.transport, ActualTransport::Http); + assert_eq!(response.socket_id, None); + server.await.unwrap(); + + super::super::websocket::invalidate_codex_websocket_pool_owner(&owner); + super::super::continuation::abort_continuation_for_owner(&reserved); + } + #[tokio::test] async fn buffered_http_retries_retryable_status() { let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); @@ -3053,38 +4164,16 @@ mod tests { #[test] fn websocket_pool_owner_tracks_typed_continuation_opt_in() { let owner = ConversationIdentity::Agent("session".into(), "agent".into()); - let disabled = super::super::continuation::ContinuationCandidate { - owner: Some(owner.clone()), - turn_id: None, - previous_response_id: None, - input_delta: None, - input_delta_count: 1, - disabled_reason: Some("disabled".into()), - }; - let first_enabled = super::super::continuation::ContinuationCandidate { - owner: Some(owner.clone()), - turn_id: Some(1), - previous_response_id: None, - input_delta: None, - input_delta_count: 1, - disabled_reason: Some("missing_state".into()), - }; - let append = super::super::continuation::ContinuationCandidate { - owner: Some(owner.clone()), - turn_id: Some(2), - previous_response_id: Some("resp_1".into()), - input_delta: None, - input_delta_count: 1, - disabled_reason: None, - }; - let missing_identity = super::super::continuation::ContinuationCandidate { - owner: None, - turn_id: None, - previous_response_id: None, - input_delta: None, - input_delta_count: 1, - disabled_reason: Some("missing_identity".into()), - }; + let disabled = test_continuation(Some(owner.clone()), None, None, None, Some("disabled")); + let first_enabled = test_continuation( + Some(owner.clone()), + Some(1), + None, + None, + Some("missing_state"), + ); + let append = test_continuation(Some(owner.clone()), Some(2), Some("resp_1"), Some(1), None); + let missing_identity = test_continuation(None, None, None, None, Some("missing_identity")); assert_eq!(websocket_pool_owner(Some(&disabled)), None); assert_eq!(websocket_pool_owner(Some(&first_enabled)), Some(&owner)); @@ -3094,30 +4183,11 @@ mod tests { #[test] fn websocket_pool_reset_clears_initial_stale_state() { - let missing_state = super::super::continuation::ContinuationCandidate { - owner: Some(ConversationIdentity::Main("session".into())), - turn_id: None, - previous_response_id: None, - input_delta: None, - input_delta_count: 1, - disabled_reason: Some("missing_state".into()), - }; - let disabled = super::super::continuation::ContinuationCandidate { - owner: Some(ConversationIdentity::Main("session".into())), - turn_id: None, - previous_response_id: None, - input_delta: None, - input_delta_count: 1, - disabled_reason: Some("disabled".into()), - }; - let prompt_changed = super::super::continuation::ContinuationCandidate { - owner: Some(ConversationIdentity::Main("session".into())), - turn_id: None, - previous_response_id: None, - input_delta: None, - input_delta_count: 1, - disabled_reason: Some("prompt_changed".into()), - }; + let owner = Some(ConversationIdentity::Main("session".into())); + let missing_state = + test_continuation(owner.clone(), None, None, None, Some("missing_state")); + let disabled = test_continuation(owner.clone(), None, None, None, Some("disabled")); + let prompt_changed = test_continuation(owner, None, None, None, Some("prompt_changed")); assert!(should_reset_websocket_pool(Some(&missing_state))); assert!(!should_reset_websocket_pool(Some(&disabled))); @@ -3165,12 +4235,15 @@ mod tests { #[test] fn unauthorized_retry_distinguishes_auto_and_strict_websocket_handshakes() { - let http_unauthorized = Ok(CodexResponse { - body: Vec::new(), - status: 401, - headers: Vec::new(), - transport: ActualTransport::Http, - }); + let http_unauthorized = Ok(OwnerAwareCodexResponse::new( + CodexResponse { + body: Vec::new(), + status: 401, + headers: Vec::new(), + transport: ActualTransport::Http, + }, + None, + )); let websocket_unauthorized = Err(CodexError { status: 401, message: "WebSocket connect error".to_string(), @@ -3246,22 +4319,11 @@ mod tests { #[test] fn continuation_retry_requires_previous_response_id() { - let append = super::super::continuation::ContinuationCandidate { - owner: Some(ConversationIdentity::Main("session".into())), - turn_id: None, - previous_response_id: Some("resp_1".into()), - input_delta: None, - input_delta_count: 1, - disabled_reason: None, - }; - let initial = super::super::continuation::ContinuationCandidate { - owner: Some(ConversationIdentity::Main("session".into())), - turn_id: None, - previous_response_id: None, - input_delta: None, - input_delta_count: 1, - disabled_reason: Some("missing_state".into()), - }; + let owner = ConversationIdentity::Main("session".into()); + let append = + test_continuation(Some(owner.clone()), Some(17), Some("resp_1"), Some(1), None); + let initial = + test_continuation(Some(owner.clone()), None, None, None, Some("missing_state")); let timeout = CodexError { status: 0, message: "WebSocket response start timeout after 60000ms".to_string(), @@ -3286,6 +4348,15 @@ mod tests { origin: CodexErrorOrigin::WebSocket, }; + let full_context = full_context_continuation(Some(&append)).unwrap(); + assert_eq!(full_context.owner(), Some(&owner)); + assert_eq!(full_context.turn_id(), Some(17)); + assert_eq!(full_context.candidate().previous_response_id, None); + assert_eq!(full_context.origin_socket_id(), None); + assert_eq!( + full_context.candidate().disabled_reason.as_deref(), + Some("full_context_retry") + ); assert!(should_retry_without_continuation(&timeout, Some(&append))); assert!(should_retry_without_continuation(&missing, Some(&append))); assert!(!should_retry_without_continuation(&idle, Some(&append))); diff --git a/src/providers/codex/compaction.rs b/src/providers/codex/compaction.rs index 2db3f363..12033d02 100644 --- a/src/providers/codex/compaction.rs +++ b/src/providers/codex/compaction.rs @@ -83,7 +83,7 @@ pub async fn request_compaction( compaction_request.include = Some(vec!["reasoning.encrypted_content".to_string()]); let response = client - .post_codex(&compaction_request, ctx, None) + .post_codex_for_owner(&compaction_request, ctx, None) .await .map_err(CompactionError::Upstream)?; let compaction = parse_compaction_response(&response.body)?; diff --git a/src/providers/codex/continuation.rs b/src/providers/codex/continuation.rs index 9c45e64b..0ea93181 100644 --- a/src/providers/codex/continuation.rs +++ b/src/providers/codex/continuation.rs @@ -14,6 +14,7 @@ const MAX_TOTAL_TRANSCRIPT_BYTES: u64 = 20_000_000; #[derive(Clone)] struct ContinuationState { response_id: String, + socket_id: u64, prompt_signature: String, transcript: Vec, transcript_bytes: u64, @@ -35,9 +36,22 @@ struct ContinuationRegistry { static REGISTRY: Mutex> = Mutex::new(None); static NEXT_TURN_ID: AtomicU64 = AtomicU64::new(1); +#[cfg(test)] +static TEST_REGISTRY_LOCK: tokio::sync::Mutex<()> = tokio::sync::Mutex::const_new(()); + +#[cfg(test)] +pub(crate) fn lock_continuation_registry_for_tests() -> tokio::sync::MutexGuard<'static, ()> { + TEST_REGISTRY_LOCK.blocking_lock() +} + +#[cfg(test)] +pub(crate) async fn lock_continuation_registry_for_async_tests() +-> tokio::sync::MutexGuard<'static, ()> { + TEST_REGISTRY_LOCK.lock().await +} + #[derive(Clone)] pub struct ContinuationCandidate { - pub owner: Option, pub turn_id: Option, pub previous_response_id: Option, pub input_delta: Option>, @@ -45,6 +59,76 @@ pub struct ContinuationCandidate { pub disabled_reason: Option, } +#[derive(Clone)] +pub(crate) struct ContinuationReservation { + candidate: ContinuationCandidate, + owner: Option, + origin_socket_id: Option, +} + +impl ContinuationReservation { + pub(crate) fn new( + candidate: ContinuationCandidate, + owner: Option, + origin_socket_id: Option, + ) -> Self { + Self { + candidate, + owner, + origin_socket_id, + } + } + + pub(crate) fn from_public_candidate(candidate: &ContinuationCandidate) -> Self { + Self::new(candidate.clone(), None, None) + } + + pub(crate) fn for_owner_turn( + owner: Option<&ConversationIdentity>, + turn_id: Option, + ) -> Self { + Self::new( + ContinuationCandidate { + turn_id, + previous_response_id: None, + input_delta: None, + input_delta_count: 0, + disabled_reason: None, + }, + owner.cloned(), + None, + ) + } + + pub(crate) fn candidate(&self) -> &ContinuationCandidate { + &self.candidate + } + + pub(crate) fn owner(&self) -> Option<&ConversationIdentity> { + self.owner.as_ref() + } + + pub(crate) fn turn_id(&self) -> Option { + self.candidate.turn_id + } + + pub(crate) fn origin_socket_id(&self) -> Option { + self.origin_socket_id + } + + pub(crate) fn into_candidate(self) -> ContinuationCandidate { + self.candidate + } + + pub(crate) fn full_context_retry(&self) -> Self { + let mut candidate = self.candidate.clone(); + candidate.previous_response_id = None; + candidate.input_delta = None; + candidate.disabled_reason = Some("full_context_retry".to_string()); + Self::new(candidate, self.owner.clone(), None) + } +} + fn now_ms() -> u64 { std::time::SystemTime::now() .duration_since(std::time::UNIX_EPOCH) @@ -52,31 +136,56 @@ fn now_ms() -> u64 { .as_millis() as u64 } +#[deprecated(note = "use the owner-aware provider flow for typed conversation ownership")] pub fn continuation_candidate( - owner: Option<&ConversationIdentity>, + session_id: Option<&str>, body: &ResponsesRequest, enabled: bool, ) -> ContinuationCandidate { + let owner = session_id.map(|session_id| ConversationIdentity::Main(session_id.to_owned())); + continuation_candidate_inner(owner.as_ref(), body, enabled, "missing_session").into_candidate() +} + +pub(crate) fn continuation_candidate_for_owner( + owner: Option<&ConversationIdentity>, + body: &ResponsesRequest, + enabled: bool, +) -> ContinuationReservation { + continuation_candidate_inner(owner, body, enabled, "missing_identity") +} + +fn continuation_candidate_inner( + owner: Option<&ConversationIdentity>, + body: &ResponsesRequest, + enabled: bool, + missing_owner_reason: &str, +) -> ContinuationReservation { if !enabled { - return ContinuationCandidate { - owner: owner.cloned(), - turn_id: None, - previous_response_id: None, - input_delta: None, - input_delta_count: body.input.len(), - disabled_reason: Some("disabled".to_string()), - }; + return ContinuationReservation::new( + ContinuationCandidate { + turn_id: None, + previous_response_id: None, + input_delta: None, + input_delta_count: body.input.len(), + disabled_reason: Some("disabled".to_string()), + }, + owner.cloned(), + None, + ); } let Some(owner) = owner else { - return ContinuationCandidate { - owner: None, - turn_id: None, - previous_response_id: None, - input_delta: None, - input_delta_count: body.input.len(), - disabled_reason: Some("missing_identity".to_string()), - }; + return ContinuationReservation::new( + ContinuationCandidate { + turn_id: None, + previous_response_id: None, + input_delta: None, + input_delta_count: body.input.len(), + disabled_reason: Some(missing_owner_reason.to_string()), + }, + None, + None, + ); }; let turn_id = NEXT_TURN_ID.fetch_add(1, Ordering::Relaxed); @@ -114,89 +223,128 @@ fn continuation_candidate_from_state( state: Option, superseded_turn: bool, now: u64, -) -> ContinuationCandidate { +) -> ContinuationReservation { let state = match state { Some(state) if now.saturating_sub(state.updated_at) <= TTL_MS => state, Some(_) | None => { - return ContinuationCandidate { - owner: Some(owner.clone()), - turn_id: Some(turn_id), - previous_response_id: None, - input_delta: None, - input_delta_count: body.input.len(), - disabled_reason: Some(if superseded_turn { - "superseded_turn".to_string() - } else { - "missing_state".to_string() - }), - }; + return ContinuationReservation::new( + ContinuationCandidate { + turn_id: Some(turn_id), + previous_response_id: None, + input_delta: None, + input_delta_count: body.input.len(), + disabled_reason: Some(if superseded_turn { + "superseded_turn".to_string() + } else { + "missing_state".to_string() + }), + }, + Some(owner.clone()), + None, + ); } }; let signature = prompt_signature(body); if signature != state.prompt_signature { - return ContinuationCandidate { - owner: Some(owner.clone()), - turn_id: Some(turn_id), - previous_response_id: None, - input_delta: None, - input_delta_count: body.input.len(), - disabled_reason: Some("prompt_changed".to_string()), - }; + return ContinuationReservation::new( + ContinuationCandidate { + turn_id: Some(turn_id), + previous_response_id: None, + input_delta: None, + input_delta_count: body.input.len(), + disabled_reason: Some("prompt_changed".to_string()), + }, + Some(owner.clone()), + None, + ); } let Some(suffix) = input_suffix_after_prefix(&body.input, &state.transcript) else { - return ContinuationCandidate { - owner: Some(owner.clone()), - turn_id: Some(turn_id), - previous_response_id: None, - input_delta: None, - input_delta_count: body.input.len(), - disabled_reason: Some("not_append_only".to_string()), - }; + return ContinuationReservation::new( + ContinuationCandidate { + turn_id: Some(turn_id), + previous_response_id: None, + input_delta: None, + input_delta_count: body.input.len(), + disabled_reason: Some("not_append_only".to_string()), + }, + Some(owner.clone()), + None, + ); }; if suffix.is_empty() { - return ContinuationCandidate { - owner: Some(owner.clone()), - turn_id: Some(turn_id), - previous_response_id: None, - input_delta: None, - input_delta_count: 0, - disabled_reason: Some("empty_delta".to_string()), - }; + return ContinuationReservation::new( + ContinuationCandidate { + turn_id: Some(turn_id), + previous_response_id: None, + input_delta: None, + input_delta_count: 0, + disabled_reason: Some("empty_delta".to_string()), + }, + Some(owner.clone()), + None, + ); } - ContinuationCandidate { - owner: Some(owner.clone()), - turn_id: Some(turn_id), - previous_response_id: Some(state.response_id), - input_delta_count: suffix.len(), - input_delta: Some(suffix), - disabled_reason: None, - } + ContinuationReservation::new( + ContinuationCandidate { + turn_id: Some(turn_id), + previous_response_id: Some(state.response_id), + input_delta_count: suffix.len(), + input_delta: Some(suffix), + disabled_reason: None, + }, + Some(owner.clone()), + Some(state.socket_id), + ) } +#[deprecated(note = "recording without typed socket provenance is not reusable")] pub fn record_continuation( - owner: Option<&ConversationIdentity>, + session_id: Option<&str>, turn_id: Option, request_body: &ResponsesRequest, response_id: Option<&str>, output_items: &[ResponsesInputItem], ) { - let (owner, turn_id) = match (owner, turn_id) { + let owner = session_id.map(|session_id| ConversationIdentity::Main(session_id.to_owned())); + let reservation = ContinuationReservation::new( + ContinuationCandidate { + turn_id, + previous_response_id: None, + input_delta: None, + input_delta_count: request_body.input.len(), + disabled_reason: Some("legacy_recording_without_socket".to_string()), + }, + owner, + None, + ); + record_continuation_for_owner(&reservation, request_body, response_id, None, output_items); +} + +pub(crate) fn record_continuation_for_owner( + reservation: &ContinuationReservation, + request_body: &ResponsesRequest, + response_id: Option<&str>, + socket_id: Option, + output_items: &[ResponsesInputItem], +) { + let (owner, turn_id) = match (reservation.owner(), reservation.turn_id()) { (Some(owner), Some(turn_id)) => (owner, turn_id), _ => return, }; - let response_id = match response_id { - Some(id) => id.to_string(), - None => { - abort_continuation(Some(owner), Some(turn_id)); + let (response_id, socket_id) = match (response_id, socket_id) { + (Some(response_id), Some(socket_id)) if socket_id != 0 => { + (response_id.to_string(), socket_id) + } + _ => { + abort_continuation_inner(Some(owner), Some(turn_id)); return; } }; - let mut transcript: Vec = request_body.input.clone(); transcript.extend_from_slice(output_items); @@ -204,12 +352,13 @@ pub fn record_continuation( let transcript_bytes = transcript_json.len() as u64; if transcript_bytes > MAX_OWNER_TRANSCRIPT_BYTES { - abort_continuation(Some(owner), Some(turn_id)); + abort_continuation_inner(Some(owner), Some(turn_id)); return; } let state = ContinuationState { response_id, + socket_id, prompt_signature: prompt_signature(request_body), transcript, transcript_bytes, @@ -235,7 +384,17 @@ pub fn record_continuation( evict_oldest(registry); } -pub fn abort_continuation(owner: Option<&ConversationIdentity>, turn_id: Option) { +#[deprecated(note = "use the owner-aware provider flow for typed conversation ownership")] +pub fn abort_continuation(session_id: Option<&str>, turn_id: Option) { + let owner = session_id.map(|session_id| ConversationIdentity::Main(session_id.to_owned())); + abort_continuation_inner(owner.as_ref(), turn_id); +} + +pub(crate) fn abort_continuation_for_owner(reservation: &ContinuationReservation) { + abort_continuation_inner(reservation.owner(), reservation.turn_id()); +} + +fn abort_continuation_inner(owner: Option<&ConversationIdentity>, turn_id: Option) { let (Some(owner), Some(turn_id)) = (owner, turn_id) else { return; }; @@ -256,7 +415,24 @@ pub fn abort_continuation(owner: Option<&ConversationIdentity>, turn_id: Option< } } +#[deprecated(note = "use the owner-aware provider flow for typed conversation ownership")] pub fn if_current_turn( + session_id: Option<&str>, + turn_id: Option, + action: impl FnOnce() -> T, +) -> Option { + let owner = session_id.map(|session_id| ConversationIdentity::Main(session_id.to_owned())); + if_current_turn_inner(owner.as_ref(), turn_id, action) +} + +pub(crate) fn if_current_turn_for_owner( + reservation: &ContinuationReservation, + action: impl FnOnce() -> T, +) -> Option { + if_current_turn_inner(reservation.owner(), reservation.turn_id(), action) +} + +fn if_current_turn_inner( owner: Option<&ConversationIdentity>, turn_id: Option, action: impl FnOnce() -> T, @@ -272,15 +448,35 @@ pub fn if_current_turn( current.then(action) } +#[deprecated(note = "use the owner-aware provider flow for typed conversation ownership")] pub fn with_current_turn( - owner: Option<&ConversationIdentity>, + session_id: Option<&str>, turn_id: Option, action: impl FnOnce(), ) -> bool { - if_current_turn(owner, turn_id, action).is_some() + let owner = session_id.map(|session_id| ConversationIdentity::Main(session_id.to_owned())); + if_current_turn_inner(owner.as_ref(), turn_id, action).is_some() +} + +pub(crate) fn with_current_turn_for_owner( + reservation: &ContinuationReservation, + action: impl FnOnce(), +) -> bool { + if_current_turn_for_owner(reservation, action).is_some() +} + +#[deprecated(note = "use the owner-aware provider flow for typed conversation ownership")] +pub fn is_current_turn(session_id: Option<&str>, turn_id: Option) -> bool { + let owner = session_id.map(|session_id| ConversationIdentity::Main(session_id.to_owned())); + is_current_turn_inner(owner.as_ref(), turn_id) } -pub fn is_current_turn(owner: Option<&ConversationIdentity>, turn_id: Option) -> bool { +#[allow(dead_code)] +pub(crate) fn is_current_turn_for_owner(reservation: &ContinuationReservation) -> bool { + is_current_turn_inner(reservation.owner(), reservation.turn_id()) +} + +fn is_current_turn_inner(owner: Option<&ConversationIdentity>, turn_id: Option) -> bool { let (Some(owner), Some(turn_id)) = (owner, turn_id) else { return false; }; @@ -291,7 +487,13 @@ pub fn is_current_turn(owner: Option<&ConversationIdentity>, turn_id: Option) { +#[deprecated(note = "use the owner-aware provider flow for typed conversation ownership")] +pub fn clear_continuation(session_id: Option<&str>) { + let owner = session_id.map(|session_id| ConversationIdentity::Main(session_id.to_owned())); + clear_continuation_for_owner(owner.as_ref()); +} + +pub(crate) fn clear_continuation_for_owner(owner: Option<&ConversationIdentity>) { let Some(owner) = owner else { return; }; @@ -308,7 +510,13 @@ pub fn clear_continuation(owner: Option<&ConversationIdentity>) { } } -pub fn has_continuation_for_tests(owner: &ConversationIdentity) -> bool { +#[deprecated(note = "use the owner-aware test helper for typed conversation ownership")] +pub fn has_continuation_for_tests(session_id: &str) -> bool { + let owner = ConversationIdentity::Main(session_id.to_owned()); + has_continuation_for_owner_for_tests(&owner) +} + +pub(crate) fn has_continuation_for_owner_for_tests(owner: &ConversationIdentity) -> bool { let guard = REGISTRY.lock().unwrap(); guard .as_ref() @@ -413,10 +621,8 @@ mod tests { use super::*; use serde_json::json; - static TEST_REGISTRY_LOCK: Mutex<()> = Mutex::new(()); - - fn lock_registry() -> std::sync::MutexGuard<'static, ()> { - let guard = TEST_REGISTRY_LOCK.lock().unwrap(); + fn lock_registry() -> tokio::sync::MutexGuard<'static, ()> { + let guard = lock_continuation_registry_for_tests(); clear_all_continuations_for_tests(); guard } @@ -466,33 +672,40 @@ mod tests { request: &ResponsesRequest, response_id: &str, ) { - let candidate = continuation_candidate(Some(owner), request, true); - record_continuation( - candidate.owner.as_ref(), - candidate.turn_id, - request, - Some(response_id), - &[], - ); + let reservation = continuation_candidate_for_owner(Some(owner), request, true); + record_continuation_for_owner(&reservation, request, Some(response_id), Some(1), &[]); } #[test] + #[allow(deprecated)] fn disabled_and_missing_identity_requests_are_stateless() { let _registry_guard = lock_registry(); let request = request_with_input(vec![input("one")], None); let owner = main_owner("session-a"); - let disabled = continuation_candidate(Some(&owner), &request, false); - assert_eq!(disabled.owner.as_ref(), Some(&owner)); - assert_eq!(disabled.turn_id, None); - assert_eq!(disabled.input_delta_count, request.input.len()); - assert_eq!(disabled.disabled_reason.as_deref(), Some("disabled")); + let disabled = continuation_candidate_for_owner(Some(&owner), &request, false); + assert_eq!(disabled.owner(), Some(&owner)); + assert_eq!(disabled.turn_id(), None); + assert_eq!(disabled.candidate().input_delta_count, request.input.len()); + assert_eq!( + disabled.candidate().disabled_reason.as_deref(), + Some("disabled") + ); - let missing = continuation_candidate(None, &request, true); - assert_eq!(missing.owner, None); - assert_eq!(missing.turn_id, None); - assert_eq!(missing.input_delta_count, request.input.len()); - assert_eq!(missing.disabled_reason.as_deref(), Some("missing_identity")); + let missing = continuation_candidate_for_owner(None, &request, true); + assert_eq!(missing.owner(), None); + assert_eq!(missing.turn_id(), None); + assert_eq!(missing.candidate().input_delta_count, request.input.len()); + assert_eq!( + missing.candidate().disabled_reason.as_deref(), + Some("missing_identity") + ); + + let legacy_missing = continuation_candidate(None, &request, true); + assert_eq!( + legacy_missing.disabled_reason.as_deref(), + Some("missing_session") + ); } #[test] @@ -502,34 +715,27 @@ mod tests { let sibling_two = agent_owner("session-a", "agent-two"); let first_request = request_with_input(vec![input("one")], None); - let first = continuation_candidate(Some(&sibling_one), &first_request, true); - let second = continuation_candidate(Some(&sibling_two), &first_request, true); - assert_ne!(first.turn_id, second.turn_id); - record_continuation( - first.owner.as_ref(), - first.turn_id, - &first_request, - Some("resp_one"), - &[], - ); - record_continuation( - second.owner.as_ref(), - second.turn_id, - &first_request, - Some("resp_two"), - &[], - ); - assert!(has_continuation_for_tests(&sibling_one)); - assert!(has_continuation_for_tests(&sibling_two)); + let first = continuation_candidate_for_owner(Some(&sibling_one), &first_request, true); + let second = continuation_candidate_for_owner(Some(&sibling_two), &first_request, true); + assert_ne!(first.turn_id(), second.turn_id()); + record_continuation_for_owner(&first, &first_request, Some("resp_one"), Some(11), &[]); + record_continuation_for_owner(&second, &first_request, Some("resp_two"), Some(22), &[]); + assert!(has_continuation_for_owner_for_tests(&sibling_one)); + assert!(has_continuation_for_owner_for_tests(&sibling_two)); let next_request = request_with_input(vec![input("one"), input("two")], None); - let first_next = continuation_candidate(Some(&sibling_one), &next_request, true); - let second_next = continuation_candidate(Some(&sibling_two), &next_request, true); - assert_eq!(first_next.previous_response_id.as_deref(), Some("resp_one")); + let first_next = continuation_candidate_for_owner(Some(&sibling_one), &next_request, true); + let second_next = continuation_candidate_for_owner(Some(&sibling_two), &next_request, true); + assert_eq!( + first_next.candidate().previous_response_id.as_deref(), + Some("resp_one") + ); + assert_eq!(first_next.origin_socket_id(), Some(11)); assert_eq!( - second_next.previous_response_id.as_deref(), + second_next.candidate().previous_response_id.as_deref(), Some("resp_two") ); + assert_eq!(second_next.origin_socket_id(), Some(22)); } #[test] @@ -538,30 +744,35 @@ mod tests { let main = main_owner("session-a"); let agent = agent_owner("session-a", "agent-a"); let request = request_with_input(vec![input("one")], None); - let main_candidate = continuation_candidate(Some(&main), &request, true); - let agent_candidate = continuation_candidate(Some(&agent), &request, true); + let main_reservation = continuation_candidate_for_owner(Some(&main), &request, true); + let agent_reservation = continuation_candidate_for_owner(Some(&agent), &request, true); - record_continuation( - agent_candidate.owner.as_ref(), - agent_candidate.turn_id, + record_continuation_for_owner( + &agent_reservation, &request, Some("resp_agent"), + Some(1), &[], ); - record_continuation( - main_candidate.owner.as_ref(), - main_candidate.turn_id, - &request, - Some("resp_main"), - &[], - ); - abort_continuation(Some(&main), agent_candidate.turn_id); + record_continuation_for_owner(&main_reservation, &request, Some("resp_main"), Some(1), &[]); + abort_continuation_for_owner(&ContinuationReservation::new( + ContinuationCandidate { + turn_id: agent_reservation.turn_id(), + previous_response_id: None, + input_delta: None, + input_delta_count: 0, + disabled_reason: None, + }, + Some(main.clone()), + None, + )); - assert!(has_continuation_for_tests(&main)); - assert!(has_continuation_for_tests(&agent)); + assert!(has_continuation_for_owner_for_tests(&main)); + assert!(has_continuation_for_owner_for_tests(&agent)); let next = request_with_input(vec![input("one"), input("two")], None); assert_eq!( - continuation_candidate(Some(&agent), &next, true) + continuation_candidate_for_owner(Some(&agent), &next, true) + .candidate() .previous_response_id .as_deref(), Some("resp_agent") @@ -577,75 +788,127 @@ mod tests { start_and_record(&owner, &request, "resp_main"); start_and_record(&sibling, &request, "resp_agent"); - let candidate = continuation_candidate(Some(&owner), &request, true); - record_continuation( - candidate.owner.as_ref(), - candidate.turn_id, + let reservation = continuation_candidate_for_owner(Some(&owner), &request, true); + record_continuation_for_owner(&reservation, &request, None, Some(1), &[]); + + assert!(!has_continuation_for_owner_for_tests(&owner)); + assert!(has_continuation_for_owner_for_tests(&sibling)); + } + + #[test] + fn missing_socket_id_does_not_publish_reusable_state() { + let _registry_guard = lock_registry(); + let owner = main_owner("session-no-socket"); + let request = request_with_input(vec![input("one")], None); + let reservation = continuation_candidate_for_owner(Some(&owner), &request, true); + + record_continuation_for_owner( + &reservation, &request, + Some("resp_without_socket"), None, &[], ); - assert!(!has_continuation_for_tests(&owner)); - assert!(has_continuation_for_tests(&sibling)); + assert!(!has_continuation_for_owner_for_tests(&owner)); + let next = continuation_candidate_for_owner(Some(&owner), &request, true); + assert_eq!(next.candidate().previous_response_id, None); + assert_eq!(next.origin_socket_id(), None); } #[test] - fn same_owner_stale_turn_cannot_publish_clear_or_run_actions() { + #[allow(deprecated)] + fn legacy_recording_without_provenance_publishes_no_reusable_state() { let _registry_guard = lock_registry(); - let owner = main_owner("session-a"); + let session_id = "legacy-no-provenance"; let request = request_with_input(vec![input("one")], None); - start_and_record(&owner, &request, "resp_1"); + let candidate = continuation_candidate(Some(session_id), &request, true); - let stale = continuation_candidate(Some(&owner), &request, true); - let current = continuation_candidate(Some(&owner), &request, true); - assert_eq!(current.disabled_reason.as_deref(), Some("superseded_turn")); record_continuation( - stale.owner.as_ref(), - stale.turn_id, - &request, - Some("resp_stale"), - &[], - ); - assert!(!has_continuation_for_tests(&owner)); - record_continuation( - current.owner.as_ref(), - current.turn_id, + Some(session_id), + candidate.turn_id, &request, - Some("resp_current"), + Some("resp_legacy"), &[], ); - assert!(has_continuation_for_tests(&owner)); - abort_continuation(stale.owner.as_ref(), stale.turn_id); - assert!(has_continuation_for_tests(&owner)); - let mut ran = false; + assert!(!has_continuation_for_tests(session_id)); + let next = continuation_candidate(Some(session_id), &request, true); + assert_eq!(next.previous_response_id, None); + } + + #[test] + fn same_owner_stale_turn_cannot_publish_clear_or_run_actions() { + let _registry_guard = lock_registry(); + let owner = main_owner("session-a"); + let request = request_with_input(vec![input("one")], None); + start_and_record(&owner, &request, "resp_1"); + + let stale = continuation_candidate_for_owner(Some(&owner), &request, true); + let current = continuation_candidate_for_owner(Some(&owner), &request, true); assert_eq!( - if_current_turn(stale.owner.as_ref(), stale.turn_id, || ran = true), - None + current.candidate().disabled_reason.as_deref(), + Some("superseded_turn") ); + record_continuation_for_owner(&stale, &request, Some("resp_stale"), Some(1), &[]); + assert!(!has_continuation_for_owner_for_tests(&owner)); + record_continuation_for_owner(¤t, &request, Some("resp_current"), Some(1), &[]); + assert!(has_continuation_for_owner_for_tests(&owner)); + abort_continuation_for_owner(&stale); + assert!(has_continuation_for_owner_for_tests(&owner)); + + let mut ran = false; + assert_eq!(if_current_turn_for_owner(&stale, || ran = true), None); assert!(!ran); } #[test] + #[allow(deprecated)] fn missing_owner_or_turn_mutations_are_hard_noops() { let _registry_guard = lock_registry(); let owner = main_owner("session-a"); let request = request_with_input(vec![input("one")], None); start_and_record(&owner, &request, "resp_1"); - record_continuation(None, Some(1), &request, Some("ignored"), &[]); - record_continuation(Some(&owner), None, &request, Some("ignored"), &[]); - abort_continuation(None, Some(1)); - abort_continuation(Some(&owner), None); - clear_continuation(None); - assert!(has_continuation_for_tests(&owner)); + let missing_owner = ContinuationReservation::new( + ContinuationCandidate { + turn_id: Some(1), + previous_response_id: None, + input_delta: None, + input_delta_count: 1, + disabled_reason: None, + }, + None, + None, + ); + let missing_turn = ContinuationReservation::new( + ContinuationCandidate { + turn_id: None, + previous_response_id: None, + input_delta: None, + input_delta_count: 1, + disabled_reason: None, + }, + Some(owner.clone()), + None, + ); + record_continuation_for_owner(&missing_owner, &request, Some("ignored"), Some(1), &[]); + record_continuation_for_owner(&missing_turn, &request, Some("ignored"), Some(1), &[]); + abort_continuation_for_owner(&missing_owner); + abort_continuation_for_owner(&missing_turn); + clear_continuation_for_owner(None); + assert!(has_continuation_for_owner_for_tests(&owner)); let mut runs = 0; + assert_eq!( + if_current_turn_for_owner(&missing_owner, || runs += 1), + None + ); + assert_eq!(if_current_turn_for_owner(&missing_turn, || runs += 1), None); + assert!(!with_current_turn_for_owner(&missing_owner, || runs += 1)); + assert!(!with_current_turn_for_owner(&missing_turn, || runs += 1)); assert_eq!(if_current_turn(None, Some(1), || runs += 1), None); - assert_eq!(if_current_turn(Some(&owner), None, || runs += 1), None); assert!(!with_current_turn(None, Some(1), || runs += 1)); - assert!(!with_current_turn(Some(&owner), None, || runs += 1)); assert_eq!(runs, 0); } @@ -657,23 +920,30 @@ mod tests { start_and_record(&owner, &request, "resp_1"); let appended = request_with_input(vec![input("one"), input("two")], None); - let candidate = continuation_candidate(Some(&owner), &appended, true); - assert_eq!(candidate.previous_response_id.as_deref(), Some("resp_1")); - assert_eq!(candidate.input_delta_count, 1); - - record_continuation( - candidate.owner.as_ref(), - candidate.turn_id, - &appended, - Some("resp_2"), - &[], + let reservation = continuation_candidate_for_owner(Some(&owner), &appended, true); + assert_eq!( + reservation.candidate().previous_response_id.as_deref(), + Some("resp_1") ); + assert_eq!(reservation.candidate().input_delta_count, 1); + + let full_context = reservation.full_context_retry(); + assert_eq!(full_context.owner(), Some(&owner)); + assert_eq!(full_context.turn_id(), reservation.turn_id()); + assert_eq!(full_context.candidate().previous_response_id, None); + assert!(full_context.candidate().input_delta.is_none()); + assert_eq!(full_context.origin_socket_id(), None); + + record_continuation_for_owner(&reservation, &appended, Some("resp_2"), Some(1), &[]); let changed = request_with_input( vec![input("one"), input("two"), input("three")], Some(json!({"service_tier": "flex"})), ); - let candidate = continuation_candidate(Some(&owner), &changed, true); - assert_eq!(candidate.disabled_reason.as_deref(), Some("prompt_changed")); - assert!(!has_continuation_for_tests(&owner)); + let reservation = continuation_candidate_for_owner(Some(&owner), &changed, true); + assert_eq!( + reservation.candidate().disabled_reason.as_deref(), + Some("prompt_changed") + ); + assert!(!has_continuation_for_owner_for_tests(&owner)); } } diff --git a/src/providers/codex/events.rs b/src/providers/codex/events.rs index 21f7e6da..51befdf1 100644 --- a/src/providers/codex/events.rs +++ b/src/providers/codex/events.rs @@ -39,6 +39,12 @@ pub(crate) fn is_terminal_rate_limit_event(payload: &Value) -> bool { != Some(true) } +pub(crate) fn event_error(payload: &Value) -> Option<&Value> { + payload + .get("error") + .or_else(|| payload.pointer("/response/error")) +} + pub(crate) fn classify_event_failure(payload: &Value) -> Option { let event_type = payload.get("type").and_then(Value::as_str)?; if event_type == "codex.rate_limits" { @@ -57,9 +63,7 @@ pub(crate) fn classify_event_failure(payload: &Value) -> Option r, Err(e) => { abort_compaction_attempt(ctx.session_id.as_deref(), compaction_attempt); - abort_continuation(continuation_owner.as_ref(), turn_id); + abort_continuation_for_owner(&request_continuation); return map_codex_error_to_response(&e); } }; @@ -328,13 +328,13 @@ impl CodexProvider { drop_live_continuation_for_retry(&mut continuation); if attempt >= MAX_EMPTY_COMPLETION_RETRIES { abort_compaction_attempt(ctx.session_id.as_deref(), compaction_attempt); - abort_continuation(continuation_owner.as_ref(), turn_id); + abort_continuation_for_owner(&request_continuation); return map_codex_error_to_response(&error); } let delay = compute_backoff_delay(attempt, None); if delay.exceeds_budget { abort_compaction_attempt(ctx.session_id.as_deref(), compaction_attempt); - abort_continuation(continuation_owner.as_ref(), turn_id); + abort_continuation_for_owner(&request_continuation); return map_codex_error_to_response(&error); } attempt += 1; @@ -353,7 +353,7 @@ impl CodexProvider { Ok(b) => b, Err(e) => { abort_compaction_attempt(ctx.session_id.as_deref(), compaction_attempt); - abort_continuation(continuation_owner.as_ref(), turn_id); + abort_continuation_for_owner(&request_continuation); return map_codex_failure_to_response(&format!( "Stream translation error: {e}" )); @@ -371,11 +371,11 @@ impl CodexProvider { } update_continuation_from_upstream( ctx.session_id.as_deref(), - continuation_owner.as_ref(), - turn_id, + &request_continuation, compaction_attempt, &translated, &upstream.body, + upstream.socket_id, compact_boundary, ); @@ -403,18 +403,18 @@ impl CodexProvider { } update_continuation_from_upstream( ctx.session_id.as_deref(), - continuation_owner.as_ref(), - turn_id, + &request_continuation, compaction_attempt, &translated, &upstream.body, + upstream.socket_id, compact_boundary, ); (StatusCode::OK, Json(json)).into_response() } Err(e) => { abort_compaction_attempt(ctx.session_id.as_deref(), compaction_attempt); - abort_continuation(continuation_owner.as_ref(), turn_id); + abort_continuation_for_owner(&request_continuation); map_codex_failure_to_response(&format!("Accumulation error: {e}")) } } @@ -546,17 +546,68 @@ fn log_compaction_event( fn abort_request_state( session_id: Option<&str>, - continuation_owner: Option<&ConversationIdentity>, - turn_id: Option, + continuation: &ContinuationReservation, compaction_attempt: Option, ) { abort_compaction_attempt(session_id, compaction_attempt); - abort_continuation(continuation_owner, turn_id); + abort_continuation_for_owner(continuation); +} + +struct LiveRequestStateCleanup { + continuation: ContinuationReservation, + session_id: Option, + compaction_attempt: Option, + armed: bool, +} + +impl LiveRequestStateCleanup { + fn new( + continuation: ContinuationReservation, + session_id: Option, + compaction_attempt: Option, + ) -> Self { + Self { + continuation, + session_id, + compaction_attempt, + armed: true, + } + } + + fn abort(&mut self) { + if self.armed { + abort_request_state( + self.session_id.as_deref(), + &self.continuation, + self.compaction_attempt, + ); + self.armed = false; + } + } + + fn disarm(&mut self) { + self.armed = false; + } +} + +impl Drop for LiveRequestStateCleanup { + fn drop(&mut self) { + if self.armed { + abort_request_state( + self.session_id.as_deref(), + &self.continuation, + self.compaction_attempt, + ); + } + } } enum LiveStreamStart { Response(Response), - Retry { error: client::CodexError }, + Retry { + error: client::CodexError, + full_context_retry_attempted: bool, + }, } #[derive(Clone, Copy)] @@ -571,18 +622,22 @@ async fn live_stream_response( model: &str, ctx: RequestContext, request_body: translate::request::ResponsesRequest, - continuation: ContinuationCandidate, + continuation: ContinuationReservation, compaction: LiveStreamCompaction, ) -> Response { let model = model.to_string(); - let continuation_owner = continuation.owner.clone(); - let turn_id = continuation.turn_id; + let request_continuation = continuation.clone(); + let mut cleanup = LiveRequestStateCleanup::new( + request_continuation.clone(), + ctx.session_id.clone(), + compaction.attempt, + ); let mut attempt = 0_u32; let mut continuation = Some(continuation); loop { let upstream_events = match client - .stream_codex_websocket_events(&request_body, &ctx, continuation.as_ref()) + .stream_codex_websocket_events_for_owner(&request_body, &ctx, continuation.as_ref()) .await { Ok(events) => events, @@ -593,22 +648,12 @@ async fn live_stream_response( continue; } if attempt >= MAX_RETRYABLE_LIVE_STREAM_RETRIES { - abort_request_state( - ctx.session_id.as_deref(), - continuation_owner.as_ref(), - turn_id, - compaction.attempt, - ); + cleanup.abort(); return map_codex_error_to_response(&err); } let delay = compute_backoff_delay(attempt, err.retry_after.as_deref()); if delay.exceeds_budget { - abort_request_state( - ctx.session_id.as_deref(), - continuation_owner.as_ref(), - turn_id, - compaction.attempt, - ); + cleanup.abort(); return map_codex_error_to_response(&err); } attempt += 1; @@ -616,12 +661,7 @@ async fn live_stream_response( continue; } Err(err) => { - abort_request_state( - ctx.session_id.as_deref(), - continuation_owner.as_ref(), - turn_id, - compaction.attempt, - ); + cleanup.abort(); return map_codex_error_to_response(&err); } }; @@ -631,37 +671,36 @@ async fn live_stream_response( message_id.clone(), &model, ctx.clone(), - continuation_owner.clone(), - turn_id, + request_continuation.clone(), request_body.clone(), compaction, ) .await { - LiveStreamStart::Response(response) => return response, - LiveStreamStart::Retry { error } => { + LiveStreamStart::Response(response) => { + cleanup.disarm(); + return response; + } + LiveStreamStart::Retry { + error, + full_context_retry_attempted, + } => { let dropped = drop_live_continuation_for_retry(&mut continuation); + if full_context_retry_attempted && client::is_continuation_retry_error(&error) { + cleanup.abort(); + return map_codex_error_to_response(&error); + } if dropped && is_missing_previous_response_error(&error) { attempt += 1; continue; } if attempt >= MAX_RETRYABLE_LIVE_STREAM_RETRIES { - abort_request_state( - ctx.session_id.as_deref(), - continuation_owner.as_ref(), - turn_id, - compaction.attempt, - ); + cleanup.abort(); return map_codex_error_to_response(&error); } let delay = compute_backoff_delay(attempt, error.retry_after.as_deref()); if delay.exceeds_budget { - abort_request_state( - ctx.session_id.as_deref(), - continuation_owner.as_ref(), - turn_id, - compaction.attempt, - ); + cleanup.abort(); return map_codex_error_to_response(&error); } attempt += 1; @@ -671,14 +710,25 @@ async fn live_stream_response( } } +fn provider_retry( + upstream_events: &websocket::CodexWebSocketEventStream, + error: client::CodexError, +) -> LiveStreamStart { + let full_context_retry_attempted = upstream_events.used_full_context_retry(); + upstream_events.mark_provider_retry_handoff(); + LiveStreamStart::Retry { + error, + full_context_retry_attempted, + } +} + #[allow(clippy::too_many_arguments)] async fn live_stream_response_once( - mut upstream_events: websocket::CodexWebSocketEventReceiver, + mut upstream_events: websocket::CodexWebSocketEventStream, message_id: String, model: &str, ctx: RequestContext, - continuation_owner: Option, - turn_id: Option, + request_continuation: ContinuationReservation, request_body: translate::request::ResponsesRequest, compaction: LiveStreamCompaction, ) -> LiveStreamStart { @@ -699,12 +749,11 @@ async fn live_stream_response_once( Ok(payload) => payload, Err(err) => { if retryable_live_start_codex_error(&err) { - return LiveStreamStart::Retry { error: err }; + return provider_retry(&upstream_events, err); } abort_request_state( ctx.session_id.as_deref(), - continuation_owner.as_ref(), - turn_id, + &request_continuation, compaction.attempt, ); return LiveStreamStart::Response(map_codex_error_to_response(&err)); @@ -744,20 +793,20 @@ async fn live_stream_response_once( 503 } }); - return LiveStreamStart::Retry { - error: client::CodexError { + return provider_retry( + &upstream_events, + client::CodexError { status, message: message.clone(), detail: Some(message), retry_after: retry_after_from_live_payload(&payload), origin: client::CodexErrorOrigin::WebSocket, }, - }; + ); } abort_request_state( ctx.session_id.as_deref(), - continuation_owner.as_ref(), - turn_id, + &request_continuation, compaction.attempt, ); return LiveStreamStart::Response(map_codex_failure_to_response(&message)); @@ -768,9 +817,7 @@ async fn live_stream_response_once( && is_codex_success_terminal_event(&payload) && !translator.has_semantic_output() { - return LiveStreamStart::Retry { - error: empty_live_completion_error(), - }; + return provider_retry(&upstream_events, empty_live_completion_error()); } if translator.has_semantic_output() && !pending_chunk.is_empty() { record_live_stream_downstream_capture(&ctx, &pending_chunk); @@ -778,11 +825,11 @@ async fn live_stream_response_once( if terminal { update_continuation_from_upstream( ctx.session_id.as_deref(), - continuation_owner.as_ref(), - turn_id, + &request_continuation, compaction.attempt, &request_body, &upstream_sse_body, + upstream_events.socket_id(), compaction.compact_boundary, ); return LiveStreamStart::Response(single_live_stream_response(pending_chunk)); @@ -792,8 +839,7 @@ async fn live_stream_response_once( translator, pending_chunk, ctx, - continuation_owner, - turn_id, + request_continuation, request_body, upstream_sse_body, compaction, @@ -802,11 +848,11 @@ async fn live_stream_response_once( if terminal { update_continuation_from_upstream( ctx.session_id.as_deref(), - continuation_owner.as_ref(), - turn_id, + &request_continuation, compaction.attempt, &request_body, &upstream_sse_body, + upstream_events.socket_id(), compaction.compact_boundary, ); if pending_chunk.is_empty() { @@ -818,15 +864,16 @@ async fn live_stream_response_once( } } - LiveStreamStart::Retry { - error: client::CodexError { + provider_retry( + &upstream_events, + client::CodexError { status: 0, message: "WebSocket connection closed before terminal Codex response event".to_string(), detail: Some(websocket::WEBSOCKET_MISSING_TERMINAL_DETAIL.to_string()), retry_after: None, origin: client::CodexErrorOrigin::WebSocket, }, - } + ) } fn empty_live_completion_error() -> client::CodexError { @@ -899,12 +946,11 @@ fn empty_live_stream_response() -> Response { #[allow(clippy::too_many_arguments)] fn remaining_live_stream_response( - mut upstream_events: websocket::CodexWebSocketEventReceiver, + mut upstream_events: websocket::CodexWebSocketEventStream, mut translator: LiveStreamTranslator, first_chunk: Vec, ctx: RequestContext, - continuation_owner: Option, - turn_id: Option, + request_continuation: ContinuationReservation, request_body: translate::request::ResponsesRequest, mut upstream_sse_body: Vec, compaction: LiveStreamCompaction, @@ -914,13 +960,27 @@ fn remaining_live_stream_response( if tx.send(Ok(Bytes::from(first_chunk))).await.is_err() { abort_request_state( ctx.session_id.as_deref(), - continuation_owner.as_ref(), - turn_id, + &request_continuation, compaction.attempt, ); return; } - while let Some(item) = upstream_events.recv().await { + loop { + let item = tokio::select! { + biased; + _ = tx.closed() => { + abort_request_state( + ctx.session_id.as_deref(), + &request_continuation, + compaction.attempt, + ); + return; + } + item = upstream_events.recv() => item, + }; + let Some(item) = item else { + break; + }; match item { Ok(payload) => { append_upstream_sse_payload(&mut upstream_sse_body, &payload); @@ -933,8 +993,7 @@ fn remaining_live_stream_response( Err(message) => { abort_request_state( ctx.session_id.as_deref(), - continuation_owner.as_ref(), - turn_id, + &request_continuation, compaction.attempt, ); let chunk = translator.error_chunk( @@ -954,8 +1013,7 @@ fn remaining_live_stream_response( if tx.send(Ok(Bytes::from(chunk))).await.is_err() { abort_request_state( ctx.session_id.as_deref(), - continuation_owner.as_ref(), - turn_id, + &request_continuation, compaction.attempt, ); return; @@ -964,11 +1022,11 @@ fn remaining_live_stream_response( if terminal { update_continuation_from_upstream( ctx.session_id.as_deref(), - continuation_owner.as_ref(), - turn_id, + &request_continuation, compaction.attempt, &request_body, &upstream_sse_body, + upstream_events.socket_id(), compaction.compact_boundary, ); return; @@ -977,8 +1035,7 @@ fn remaining_live_stream_response( Err(err) => { abort_request_state( ctx.session_id.as_deref(), - continuation_owner.as_ref(), - turn_id, + &request_continuation, compaction.attempt, ); let chunk = @@ -1005,8 +1062,7 @@ fn remaining_live_stream_response( abort_request_state( ctx.session_id.as_deref(), - continuation_owner.as_ref(), - turn_id, + &request_continuation, compaction.attempt, ); let chunk = translator.finish_after_closed_completed_tool_call(ctx.traffic.as_deref()); @@ -1124,22 +1180,24 @@ fn retryable_live_start_codex_error(err: &client::CodexError) -> bool { } fn is_missing_previous_response_error(err: &client::CodexError) -> bool { - err.detail.as_deref() == Some("previous_response_not_found") + matches!( + err.detail.as_deref(), + Some("previous_response_not_found") + | Some(websocket::WEBSOCKET_CONTINUATION_SOCKET_MISSING_DETAIL) + ) } -fn drop_live_continuation_for_retry(continuation: &mut Option) -> bool { +fn drop_live_continuation_for_retry(continuation: &mut Option) -> bool { if continuation .as_ref() - .and_then(|candidate| candidate.previous_response_id.as_deref()) + .and_then(|reservation| reservation.candidate().previous_response_id.as_deref()) .is_none() { return false; } - if let Some(candidate) = continuation.as_mut() { - candidate.previous_response_id = None; - candidate.input_delta = None; - candidate.disabled_reason = Some("full_context_retry".to_string()); + if let Some(reservation) = continuation.as_ref() { + *continuation = Some(reservation.full_context_retry()); } true } @@ -1181,13 +1239,14 @@ fn codex_stream_error_type(err: &client::CodexError) -> &'static str { } } +#[allow(clippy::too_many_arguments)] fn update_continuation_from_upstream( session_id: Option<&str>, - continuation_owner: Option<&ConversationIdentity>, - turn_id: Option, + continuation: &ContinuationReservation, compaction_attempt: Option, request_body: &translate::request::ResponsesRequest, upstream_body: &[u8], + socket_id: Option, compact_boundary: bool, ) { match finish_metadata_from_upstream(upstream_body) { @@ -1200,17 +1259,17 @@ fn update_continuation_from_upstream( &finish.output_items, ); } - record_continuation( - continuation_owner, - turn_id, + record_continuation_for_owner( + continuation, request_body, finish.response_id.as_deref(), + socket_id, &finish.output_items, ); } _ => { abort_compaction_attempt(session_id, compaction_attempt); - abort_continuation(continuation_owner, turn_id); + abort_continuation_for_owner(continuation); } } } @@ -1397,8 +1456,94 @@ fn format_auth_saved_output(auth_path: &str, account_id: Option<&str>) -> String #[cfg(test)] mod tests { + use futures_util::{SinkExt, StreamExt}; + use http_body_util::BodyExt; + use tokio::net::{TcpListener, TcpStream}; + use tokio_tungstenite::tungstenite::Message; + use super::*; + fn live_test_request(text: &str) -> translate::request::ResponsesRequest { + translate::request::ResponsesRequest { + model: "gpt-5.6-sol".to_string(), + instructions: None, + input: vec![translate::request::ResponsesInputItem::Message { + role: "user".to_string(), + content: vec![translate::request::ResponsesContentPart::InputText { + text: text.to_string(), + }], + }], + tools: None, + tool_choice: None, + store: false, + stream: true, + parallel_tool_calls: true, + include: None, + client_metadata: None, + service_tier: None, + prompt_cache_key: None, + text: translate::request::ResponsesText { + verbosity: None, + format: None, + }, + reasoning: None, + } + } + + fn live_test_context(session_id: &str) -> RequestContext { + RequestContext { + req_id: format!("request-{session_id}"), + session_id: Some(session_id.to_string()), + session_seq: None, + provider: "codex".to_string(), + traffic: None, + monitor: None, + } + } + + fn authenticated_live_test_client(base_url: String) -> Arc { + let client = CodexHttpClient::new_for_test( + reqwest::Client::builder().no_proxy().build().unwrap(), + base_url, + 1_000, + 1_000, + 0, + ); + client + .auth_manager() + .set_test_auth(auth::token_store::StoredAuth { + access: "test".to_string(), + refresh: String::new(), + expires: u64::MAX, + account_id: Some("acct".to_string()), + }); + Arc::new(client) + } + + async fn next_live_websocket_request( + websocket: &mut tokio_tungstenite::WebSocketStream, + ) -> serde_json::Value { + loop { + match websocket.next().await { + Some(Ok(Message::Ping(payload))) => { + websocket.send(Message::Pong(payload)).await.unwrap(); + } + Some(Ok(Message::Text(text))) => return serde_json::from_str(&text).unwrap(), + other => panic!("unexpected WebSocket request frame: {other:?}"), + } + } + } + + async fn emit_live_event( + websocket: &mut tokio_tungstenite::WebSocketStream, + event: &serde_json::Value, + ) { + websocket + .send(Message::Text(event.to_string())) + .await + .unwrap(); + } + fn upstream_sse(events: &[serde_json::Value]) -> Vec { let mut bytes = Vec::new(); for event in events { @@ -1739,4 +1884,381 @@ mod tests { "bad request", )); } + + async fn run_live_failure_case( + session_id: &str, + event: serde_json::Value, + expected_attempts: usize, + ) -> StatusCode { + let owner = ConversationIdentity::Main(session_id.to_string()); + continuation::clear_continuation_for_owner(Some(&owner)); + websocket::invalidate_codex_websocket_pool_owner(&owner); + let request = live_test_request("one"); + let continuation = continuation_candidate_for_owner(Some(&owner), &request, true); + let compaction_attempt = begin_compaction(session_id, &request.model); + + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = listener.local_addr().unwrap(); + let server = tokio::spawn(async move { + for _ in 0..expected_attempts { + let (socket, _) = listener.accept().await.unwrap(); + let mut websocket = tokio_tungstenite::accept_async(socket).await.unwrap(); + let _ = next_live_websocket_request(&mut websocket).await; + emit_live_event(&mut websocket, &event).await; + drop(websocket); + } + }); + let client = authenticated_live_test_client(format!("http://{addr}/responses")); + let response = tokio::time::timeout( + std::time::Duration::from_secs(15), + live_stream_response( + client, + "message".to_string(), + &request.model, + live_test_context(session_id), + request.clone(), + continuation.clone(), + LiveStreamCompaction { + compact_boundary: false, + attempt: Some(compaction_attempt), + }, + ), + ) + .await + .expect("live failure case timed out"); + tokio::time::timeout(std::time::Duration::from_secs(5), server) + .await + .expect("live failure server timed out") + .expect("live failure server failed"); + + assert!(!continuation::is_current_turn_for_owner(&continuation)); + assert!(!store_compaction( + session_id, + compaction_attempt, + Vec::new() + )); + websocket::invalidate_codex_websocket_pool_owner(&owner); + response.status() + } + + #[tokio::test] + async fn dropping_live_stream_during_retry_backoff_aborts_request_state() { + let _registry_guard = continuation::lock_continuation_registry_for_async_tests().await; + let _pool_guard = websocket::lock_codex_websocket_pool_for_tests().await; + let session_id = "live-retry-backoff-cleanup"; + let owner = ConversationIdentity::Main(session_id.to_string()); + continuation::clear_continuation_for_owner(Some(&owner)); + websocket::invalidate_codex_websocket_pool_owner(&owner); + let request = live_test_request("one"); + let continuation = continuation_candidate_for_owner(Some(&owner), &request, true); + let compaction_attempt = begin_compaction(session_id, &request.model); + + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = listener.local_addr().unwrap(); + let (event_sent_tx, event_sent_rx) = tokio::sync::oneshot::channel(); + let (socket_closed_tx, socket_closed_rx) = tokio::sync::oneshot::channel(); + let server = tokio::spawn(async move { + let (socket, _) = listener.accept().await.unwrap(); + let mut websocket = tokio_tungstenite::accept_async(socket).await.unwrap(); + let _ = next_live_websocket_request(&mut websocket).await; + emit_live_event( + &mut websocket, + &serde_json::json!({ + "type": "codex.rate_limits", + "rate_limits": {"allowed": false, "limit_reached": true} + }), + ) + .await; + event_sent_tx.send(()).unwrap(); + drop(websocket); + socket_closed_tx.send(()).unwrap(); + }); + let client = authenticated_live_test_client(format!("http://{addr}/responses")); + let task_request = request.clone(); + let task_continuation = continuation.clone(); + let response_task = tokio::spawn(async move { + let model = task_request.model.clone(); + live_stream_response( + client, + "message".to_string(), + &model, + live_test_context(session_id), + task_request, + task_continuation, + LiveStreamCompaction { + compact_boundary: false, + attempt: Some(compaction_attempt), + }, + ) + .await + }); + + event_sent_rx.await.unwrap(); + tokio::time::timeout(std::time::Duration::from_secs(1), socket_closed_rx) + .await + .expect("retry handoff did not close the abandoned attempt socket") + .expect("retry handoff socket-close sender dropped"); + for _ in 0..16 { + tokio::task::yield_now().await; + } + assert!( + !response_task.is_finished(), + "logical request must still be waiting in retry backoff" + ); + response_task.abort(); + assert!(response_task.await.unwrap_err().is_cancelled()); + + assert!(!continuation::is_current_turn_for_owner(&continuation)); + assert!(!store_compaction( + session_id, + compaction_attempt, + Vec::new() + )); + server.await.unwrap(); + websocket::invalidate_codex_websocket_pool_owner(&owner); + } + + #[tokio::test] + async fn dropping_live_response_body_after_first_chunk_aborts_request_state() { + let _registry_guard = continuation::lock_continuation_registry_for_async_tests().await; + let _pool_guard = websocket::lock_codex_websocket_pool_for_tests().await; + let session_id = "live-response-body-drop-cleanup"; + let owner = ConversationIdentity::Main(session_id.to_string()); + continuation::clear_continuation_for_owner(Some(&owner)); + websocket::invalidate_codex_websocket_pool_owner(&owner); + let request = live_test_request("one"); + let continuation = continuation_candidate_for_owner(Some(&owner), &request, true); + let compaction_attempt = begin_compaction(session_id, &request.model); + + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = listener.local_addr().unwrap(); + let (socket_closed_tx, socket_closed_rx) = tokio::sync::oneshot::channel(); + let server = tokio::spawn(async move { + let (socket, _) = listener.accept().await.unwrap(); + let mut websocket = tokio_tungstenite::accept_async(socket).await.unwrap(); + let _ = next_live_websocket_request(&mut websocket).await; + emit_live_event( + &mut websocket, + &serde_json::json!({ + "type": "response.output_item.added", + "output_index": 0, + "item": {"type": "message", "id": "msg-partial"} + }), + ) + .await; + emit_live_event( + &mut websocket, + &serde_json::json!({ + "type": "response.output_text.delta", + "output_index": 0, + "delta": "partial" + }), + ) + .await; + while websocket.next().await.is_some() {} + socket_closed_tx.send(()).unwrap(); + }); + let client = authenticated_live_test_client(format!("http://{addr}/responses")); + + let response = tokio::time::timeout( + std::time::Duration::from_secs(2), + live_stream_response( + client, + "message".to_string(), + &request.model, + live_test_context(session_id), + request.clone(), + continuation.clone(), + LiveStreamCompaction { + compact_boundary: false, + attempt: Some(compaction_attempt), + }, + ), + ) + .await + .expect("live response did not publish the first chunk"); + let mut body = response.into_body(); + tokio::time::timeout(std::time::Duration::from_secs(1), body.frame()) + .await + .expect("first downstream chunk timed out") + .expect("live response body ended before the first chunk") + .expect("first downstream chunk failed"); + drop(body); + + tokio::time::timeout(std::time::Duration::from_secs(1), socket_closed_rx) + .await + .expect("dropping the downstream body did not close the upstream socket") + .expect("socket-close acknowledgement sender dropped"); + assert!(!continuation::is_current_turn_for_owner(&continuation)); + assert!(!store_compaction( + session_id, + compaction_attempt, + Vec::new() + )); + server.await.unwrap(); + websocket::invalidate_codex_websocket_pool_owner(&owner); + } + + #[tokio::test] + async fn stale_request_cleanup_preserves_newer_turn_and_compaction_attempt() { + let _registry_guard = continuation::lock_continuation_registry_for_async_tests().await; + let session_id = "stale-live-request-cleanup"; + let owner = ConversationIdentity::Main(session_id.to_string()); + continuation::clear_continuation_for_owner(Some(&owner)); + let request = live_test_request("one"); + let stale_continuation = continuation_candidate_for_owner(Some(&owner), &request, true); + let stale_compaction = begin_compaction(session_id, &request.model); + let stale_cleanup = LiveRequestStateCleanup::new( + stale_continuation, + Some(session_id.to_string()), + Some(stale_compaction), + ); + + let newer_continuation = continuation_candidate_for_owner(Some(&owner), &request, true); + let newer_compaction = begin_compaction(session_id, &request.model); + drop(stale_cleanup); + + assert!(continuation::is_current_turn_for_owner(&newer_continuation)); + assert!(store_compaction(session_id, newer_compaction, Vec::new())); + abort_request_state( + Some(session_id), + &newer_continuation, + Some(newer_compaction), + ); + } + + #[tokio::test] + async fn retry_exhaustion_aborts_live_request_state_after_eleven_attempts() { + let _registry_guard = continuation::lock_continuation_registry_for_async_tests().await; + let _pool_guard = websocket::lock_codex_websocket_pool_for_tests().await; + let status = run_live_failure_case( + "live-retry-exhaustion-cleanup", + serde_json::json!({ + "type": "codex.rate_limits", + "rate_limits": { + "allowed": false, + "limit_reached": true, + "primary": {"reset_after_seconds": 0} + } + }), + 11, + ) + .await; + assert_eq!(status, StatusCode::TOO_MANY_REQUESTS); + } + + #[tokio::test] + async fn excessive_retry_after_aborts_live_request_state() { + let _registry_guard = continuation::lock_continuation_registry_for_async_tests().await; + let _pool_guard = websocket::lock_codex_websocket_pool_for_tests().await; + let status = run_live_failure_case( + "live-excessive-retry-after-cleanup", + serde_json::json!({ + "type": "codex.rate_limits", + "rate_limits": { + "allowed": false, + "limit_reached": true, + "primary": {"reset_after_seconds": 31} + } + }), + 1, + ) + .await; + assert_eq!(status, StatusCode::TOO_MANY_REQUESTS); + } + + #[tokio::test] + async fn nonretryable_live_error_aborts_request_state() { + let _registry_guard = continuation::lock_continuation_registry_for_async_tests().await; + let _pool_guard = websocket::lock_codex_websocket_pool_for_tests().await; + let status = run_live_failure_case( + "live-nonretryable-cleanup", + serde_json::json!({ + "type": "response.failed", + "response": { + "status": "failed", + "error": {"message": "invalid request"} + } + }), + 1, + ) + .await; + assert_eq!(status, StatusCode::BAD_GATEWAY); + } + + #[tokio::test] + async fn cancellation_while_replacement_startup_is_blocked_aborts_request_state() { + let _registry_guard = continuation::lock_continuation_registry_for_async_tests().await; + let _pool_guard = websocket::lock_codex_websocket_pool_for_tests().await; + let session_id = "live-blocked-replacement-cleanup"; + let owner = ConversationIdentity::Main(session_id.to_string()); + continuation::clear_continuation_for_owner(Some(&owner)); + websocket::invalidate_codex_websocket_pool_owner(&owner); + let request = live_test_request("one"); + let continuation = continuation_candidate_for_owner(Some(&owner), &request, true); + let compaction_attempt = begin_compaction(session_id, &request.model); + + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = listener.local_addr().unwrap(); + let (replacement_accepted_tx, replacement_accepted_rx) = tokio::sync::oneshot::channel(); + let (release_replacement_tx, release_replacement_rx) = tokio::sync::oneshot::channel(); + let server = tokio::spawn(async move { + let (first_socket, _) = listener.accept().await.unwrap(); + let mut first_websocket = tokio_tungstenite::accept_async(first_socket).await.unwrap(); + let _ = next_live_websocket_request(&mut first_websocket).await; + emit_live_event( + &mut first_websocket, + &serde_json::json!({ + "type": "codex.rate_limits", + "rate_limits": { + "allowed": false, + "limit_reached": true, + "primary": {"reset_after_seconds": 0} + } + }), + ) + .await; + drop(first_websocket); + + let (_replacement_socket, _) = listener.accept().await.unwrap(); + replacement_accepted_tx.send(()).unwrap(); + let _ = release_replacement_rx.await; + }); + let client = authenticated_live_test_client(format!("http://{addr}/responses")); + let task_request = request.clone(); + let task_continuation = continuation.clone(); + let response_task = tokio::spawn(async move { + let model = task_request.model.clone(); + live_stream_response( + client, + "message".to_string(), + &model, + live_test_context(session_id), + task_request, + task_continuation, + LiveStreamCompaction { + compact_boundary: false, + attempt: Some(compaction_attempt), + }, + ) + .await + }); + + tokio::time::timeout(std::time::Duration::from_secs(2), replacement_accepted_rx) + .await + .expect("replacement startup did not reach the blocked handshake") + .expect("replacement startup acknowledgement sender dropped"); + response_task.abort(); + assert!(response_task.await.unwrap_err().is_cancelled()); + let _ = release_replacement_tx.send(()); + server.await.unwrap(); + + assert!(!continuation::is_current_turn_for_owner(&continuation)); + assert!(!store_compaction( + session_id, + compaction_attempt, + Vec::new() + )); + websocket::invalidate_codex_websocket_pool_owner(&owner); + } } diff --git a/src/providers/codex/websocket.rs b/src/providers/codex/websocket.rs index 20f6d820..f1cce1de 100644 --- a/src/providers/codex/websocket.rs +++ b/src/providers/codex/websocket.rs @@ -1,6 +1,6 @@ use std::collections::HashMap; use std::pin::Pin; -use std::sync::atomic::{AtomicU64, Ordering}; +use std::sync::atomic::{AtomicBool, AtomicU64, Ordering}; use std::sync::{Arc, Mutex}; use std::task::{Context, Poll}; use std::time::{Duration, Instant}; @@ -28,8 +28,10 @@ use crate::request_identity::ConversationIdentity; use crate::retry::sleep as retry_sleep; use crate::traffic::TrafficCapture; -use super::client::{ActualTransport, CodexError, CodexErrorOrigin, CodexResponse}; -use super::continuation::ContinuationCandidate; +use super::client::{ + ActualTransport, CodexError, CodexErrorOrigin, CodexResponse, OwnerAwareCodexResponse, +}; +use super::continuation::ContinuationReservation; // --------------------------------------------------------------------------- // Constants @@ -40,6 +42,8 @@ pub const WEBSOCKET_CONNECT_TIMEOUT_MS: u64 = 15_000; pub const WEBSOCKET_IDLE_TIMEOUT_MS: u64 = 300_000; pub const WEBSOCKET_RESPONSE_START_TIMEOUT_DETAIL: &str = "websocket_response_start_timeout"; pub const WEBSOCKET_MISSING_TERMINAL_DETAIL: &str = "websocket_missing_terminal"; +pub const WEBSOCKET_CONTINUATION_SOCKET_MISSING_DETAIL: &str = + "websocket_continuation_socket_missing"; pub(super) const WEBSOCKET_PROXY_TUNNEL_REJECTED_DETAIL: &str = "websocket_proxy_tunnel_rejected"; const POOL_IDLE_TTL_MS: u64 = 30 * 60 * 1000; @@ -58,7 +62,83 @@ const TERMINAL_EVENTS: &[&str] = &[ "error", ]; -pub type CodexWebSocketEventReceiver = mpsc::Receiver>; +pub type CodexWebSocketEventReceiver = + tokio::sync::mpsc::Receiver>; + +pub(crate) struct CodexWebSocketEventStream { + receiver: CodexWebSocketEventReceiver, + socket_id: Arc, + full_context_retry: Arc, + provider_retry_handoff: Arc, +} + +#[derive(Clone)] +pub(crate) struct CodexWebSocketSocketIdPublisher { + socket_id: Arc, + full_context_retry: Arc, + provider_retry_handoff: Arc, +} + +impl CodexWebSocketEventStream { + pub(crate) fn pending( + receiver: CodexWebSocketEventReceiver, + ) -> (Self, CodexWebSocketSocketIdPublisher) { + let socket_id = Arc::new(AtomicU64::new(0)); + let full_context_retry = Arc::new(AtomicBool::new(false)); + let provider_retry_handoff = Arc::new(AtomicBool::new(false)); + ( + Self { + receiver, + socket_id: socket_id.clone(), + full_context_retry: full_context_retry.clone(), + provider_retry_handoff: provider_retry_handoff.clone(), + }, + CodexWebSocketSocketIdPublisher { + socket_id, + full_context_retry, + provider_retry_handoff, + }, + ) + } + + pub(crate) async fn recv(&mut self) -> Option> { + self.receiver.recv().await + } + + pub(crate) fn socket_id(&self) -> Option { + match self.socket_id.load(Ordering::Acquire) { + 0 => None, + socket_id => Some(socket_id), + } + } + + pub(crate) fn used_full_context_retry(&self) -> bool { + self.full_context_retry.load(Ordering::Acquire) + } + + pub(crate) fn mark_provider_retry_handoff(&self) { + self.provider_retry_handoff.store(true, Ordering::Release); + } + + pub(crate) fn into_receiver(self) -> CodexWebSocketEventReceiver { + self.receiver + } +} + +impl CodexWebSocketSocketIdPublisher { + pub(super) fn publish(&self, socket_id: Option) { + self.socket_id + .store(socket_id.unwrap_or(0), Ordering::Release); + } + + pub(super) fn mark_full_context_retry(&self) { + self.full_context_retry.store(true, Ordering::Release); + } + + pub(super) fn is_provider_retry_handoff(&self) -> bool { + self.provider_retry_handoff.load(Ordering::Acquire) + } +} trait WebSocketIo: AsyncRead + AsyncWrite + Unpin + Send {} @@ -223,6 +303,7 @@ impl std::fmt::Display for CodexWebSocketError { struct PoolEntry { ws: Arc>, + socket_id: u64, created_at: u64, last_activity: AtomicU64, } @@ -231,6 +312,7 @@ impl PoolEntry { fn new(ws: CodexWebSocketStream) -> Self { Self { ws: Arc::new(AsyncMutex::new(ws)), + socket_id: next_monotonic_nonzero(&NEXT_SOCKET_ID, "WebSocket ID"), created_at: now_ms(), last_activity: AtomicU64::new(next_pool_activity()), } @@ -242,12 +324,25 @@ impl PoolEntry { } } +static NEXT_SOCKET_ID: AtomicU64 = AtomicU64::new(0); +static POOLED_VALIDATION_SEQUENCE: AtomicU64 = AtomicU64::new(0); static POOL_ACTIVITY_SEQUENCE: AtomicU64 = AtomicU64::new(1); static WS_POOL: once_cell::sync::Lazy>>> = once_cell::sync::Lazy::new(|| Mutex::new(HashMap::new())); +#[cfg(test)] +static WS_POOL_TEST_LOCK: AsyncMutex<()> = AsyncMutex::const_new(()); static WS_CONNECT_GATE: once_cell::sync::Lazy = once_cell::sync::Lazy::new(|| WebSocketConnectGate::new(WEBSOCKET_CONNECT_START_SPACING)); +fn next_monotonic_nonzero(sequence: &AtomicU64, label: &str) -> u64 { + let previous = sequence + .fetch_update(Ordering::Relaxed, Ordering::Relaxed, |value| { + value.checked_add(1) + }) + .unwrap_or_else(|_| panic!("{label} sequence exhausted")); + previous + 1 +} + fn next_pool_activity() -> u64 { POOL_ACTIVITY_SEQUENCE.fetch_add(1, Ordering::Relaxed) } @@ -264,17 +359,50 @@ pub fn clear_codex_websocket_pool_for_tests() { guard.clear(); } +#[cfg(test)] +pub(crate) async fn lock_codex_websocket_pool_for_tests() -> tokio::sync::MutexGuard<'static, ()> { + WS_POOL_TEST_LOCK.lock().await +} + +#[cfg(test)] +pub(crate) fn pooled_socket_id_for_tests(owner: &ConversationIdentity) -> Option { + WS_POOL + .lock() + .unwrap() + .get(owner) + .map(|entry| entry.socket_id) +} + pub fn invalidate_codex_websocket_pool_owner(owner: &ConversationIdentity) { let mut guard = WS_POOL.lock().unwrap(); guard.remove(owner); } -pub fn invalidate_codex_websocket_pool_turn(owner: &ConversationIdentity, turn_id: Option) { - super::continuation::with_current_turn(Some(owner), turn_id, || { +#[deprecated(note = "use typed conversation ownership internally")] +pub fn invalidate_codex_websocket_pool_key(session_id: &str) { + let mut guard = WS_POOL.lock().unwrap(); + guard.retain(|owner, _| match owner { + ConversationIdentity::Main(owner_session_id) + | ConversationIdentity::Agent(owner_session_id, _) => owner_session_id != session_id, + }); +} + +pub(crate) fn invalidate_codex_websocket_pool_turn_for_owner( + owner: &ConversationIdentity, + turn_id: Option, +) { + let reservation = ContinuationReservation::for_owner_turn(Some(owner), turn_id); + super::continuation::with_current_turn_for_owner(&reservation, || { invalidate_codex_websocket_pool_owner(owner) }); } +#[deprecated(note = "use typed conversation ownership internally")] +pub fn invalidate_codex_websocket_pool_turn(session_id: &str, turn_id: Option) { + let owner = ConversationIdentity::Main(session_id.to_owned()); + invalidate_codex_websocket_pool_turn_for_owner(&owner, turn_id); +} + fn invalidate_pool_entry(owner: &ConversationIdentity, entry: &Arc) { let mut guard = WS_POOL.lock().unwrap(); if guard @@ -295,36 +423,105 @@ fn invalidate_pool_owner(owner: Option<&ConversationIdentity>, entry: Option<&Ar } } -fn pool_get_for_turn(owner: &ConversationIdentity, turn_id: Option) -> Option> { - super::continuation::if_current_turn(Some(owner), turn_id, || { - let entry = WS_POOL.lock().ok()?.get(owner).cloned()?; - entry.touch(); - Some(entry) - }) - .flatten() +fn reservation_pool_owner( + reservation: Option<&ContinuationReservation>, +) -> Option<&ConversationIdentity> { + let reservation = reservation?; + if reservation.candidate().disabled_reason.as_deref() == Some("disabled") { + return None; + } + reservation.owner() } -fn pool_take_for_turn( - owner: &ConversationIdentity, - turn_id: Option, -) -> Option> { - super::continuation::if_current_turn(Some(owner), turn_id, || { +fn pool_take_for_turn(reservation: &ContinuationReservation) -> Option> { + let owner = reservation.owner()?; + super::continuation::if_current_turn_for_owner(reservation, || { WS_POOL.lock().ok()?.remove(owner) }) .flatten() } -fn pool_insert_for_turn(owner: ConversationIdentity, entry: Arc, turn_id: Option) { - let current_owner = owner.clone(); - super::continuation::with_current_turn(Some(¤t_owner), turn_id, || { - pool_insert(owner, entry) - }); +fn take_pool_entry_for_request( + reservation: Option<&ContinuationReservation>, +) -> Result>, CodexError> { + let candidate = reservation.map(ContinuationReservation::candidate); + let requires_origin = candidate + .and_then(|candidate| candidate.previous_response_id.as_deref()) + .is_some(); + let expected_socket_id = reservation.and_then(ContinuationReservation::origin_socket_id); + let pool_owner = reservation_pool_owner(reservation); + let pooled = reservation + .filter(|_| pool_owner.is_some()) + .and_then(pool_take_for_turn); + + if requires_origin + && (pool_owner.is_none() + || expected_socket_id.is_none() + || pooled + .as_ref() + .is_none_or(|entry| Some(entry.socket_id) != expected_socket_id)) + { + if let (Some(owner), Some(entry)) = (pool_owner, pooled.as_ref()) { + pool_insert_if_vacant_or_same(owner.clone(), entry.clone()); + } + return Err(continuation_socket_missing_error()); + } + + Ok(pooled) +} + +fn pool_insert_for_turn(reservation: &ContinuationReservation, entry: Arc) -> bool { + let Some(owner) = reservation_pool_owner(Some(reservation)).cloned() else { + return false; + }; + super::continuation::if_current_turn_for_owner(reservation, || { + pool_insert_if_vacant_or_same(owner, entry) + }) + .unwrap_or(false) } fn pool_remove_entry(owner: &ConversationIdentity, entry: &Arc) { invalidate_pool_entry(owner, entry); } +pub(super) fn invalidate_codex_websocket_pool_socket( + reservation: &ContinuationReservation, + socket_id: Option, +) { + let Some(owner) = reservation_pool_owner(Some(reservation)) else { + return; + }; + let Some(socket_id) = socket_id else { + return; + }; + let entry = WS_POOL.lock().ok().and_then(|pool| { + pool.get(owner) + .filter(|entry| entry.socket_id == socket_id) + .cloned() + }); + if let Some(entry) = entry { + pool_remove_entry(owner, &entry); + } +} + +fn pool_insert_if_vacant_or_same(owner: ConversationIdentity, entry: Arc) -> bool { + entry.touch(); + let mut guard = WS_POOL.lock().unwrap(); + if let Some(existing) = guard.get(&owner) { + return Arc::ptr_eq(existing, &entry); + } + if guard.len() >= MAX_POOL_ENTRIES + && let Some(oldest_owner) = guard.keys().next().cloned() + { + guard.remove(&oldest_owner); + } + let now = now_ms(); + guard.retain(|_, pooled| now.saturating_sub(pooled.created_at) < POOL_IDLE_TTL_MS); + guard.insert(owner, entry); + true +} + +#[cfg(test)] fn pool_insert(owner: ConversationIdentity, entry: Arc) { entry.touch(); let mut guard = WS_POOL.lock().unwrap(); @@ -479,7 +676,7 @@ fn encode_sse(text: &str) -> Vec { // Terminal event detection // --------------------------------------------------------------------------- -fn is_terminal_event(payload: &serde_json::Value) -> bool { +pub(super) fn is_terminal_event(payload: &serde_json::Value) -> bool { match payload.get("type").and_then(|v| v.as_str()) { Some(t) => TERMINAL_EVENTS.contains(&t), None => false, @@ -495,19 +692,18 @@ fn is_response_event(payload: &serde_json::Value) -> bool { } fn is_previous_response_missing(payload: &serde_json::Value) -> bool { - if let Some(code) = payload - .get("error") - .and_then(|e| e.get("code")) - .and_then(|v| v.as_str()) - && code == "previous_response_not_found" + let error = super::events::event_error(payload); + if error + .and_then(|error| error.get("code")) + .and_then(|value| value.as_str()) + == Some("previous_response_not_found") { return true; } // Case-insensitive message check - if let Some(msg) = payload - .get("error") - .and_then(|e| e.get("message")) - .and_then(|v| v.as_str()) + if let Some(msg) = error + .and_then(|error| error.get("message")) + .and_then(|value| value.as_str()) { let lower = msg.to_lowercase(); if lower.contains("previous response") && lower.contains("not found") { @@ -543,11 +739,12 @@ pub(super) async fn codex_websocket_request( body_value: &serde_json::Value, _ctx: &RequestContext, traffic: Option<&TrafficCapture>, - pool_owner: Option<&ConversationIdentity>, connect_timeout_ms: u64, idle_timeout_ms: u64, - continuation: Option<&ContinuationCandidate>, -) -> Result { + reservation: Option<&ContinuationReservation>, +) -> Result { + let continuation = reservation.map(ContinuationReservation::candidate); + let pool_owner = reservation_pool_owner(reservation); let ws_url = to_websocket_url(url).map_err(|e| CodexError { status: 0, message: e.message, @@ -581,18 +778,15 @@ pub(super) async fn codex_websocket_request( } let started_at = Instant::now(); - // Check pool for existing connection - let pooled = pool_owner.and_then(|key| { - pool_get_for_turn(key, continuation.and_then(|candidate| candidate.turn_id)) - }); - - let ws_stream = if let Some(entry) = pooled { - // Use pooled connection - let mut ws_guard = entry.ws.lock().await; - // Check if connection is still alive by sending a ping - if ws_guard.send(Message::Ping(vec![])).await.is_err() { - invalidate_pool_entry(pool_owner.unwrap(), &entry); - // Fall through to new connection + let requires_origin = continuation + .and_then(|candidate| candidate.previous_response_id.as_deref()) + .is_some(); + let pooled = take_pool_entry_for_request(reservation)?; + let mut used_pooled = pooled.is_some(); + let mut entry = if let Some(entry) = pooled { + entry + } else { + Arc::new(PoolEntry::new( connect_with_timeout( websocket_client, proxy_config, @@ -600,151 +794,157 @@ pub(super) async fn codex_websocket_request( headers, connect_timeout_ms, ) - .await? - } else { - // Connection is alive, send the request through it - let ws_msg = Message::Text(body_json.clone()); - ws_guard.send(ws_msg).await.map_err(|e| { - if let Some(key) = pool_owner { - invalidate_pool_entry(key, &entry); - } - CodexError { - status: 0, - message: format!("WebSocket send error: {e}"), - detail: None, - retry_after: None, - origin: CodexErrorOrigin::WebSocket, - } - })?; - - // Collect events - let (sse_body, terminal_event) = collect_ws_events( - &mut ws_guard, - idle_timeout_ms, - pool_owner, - Some(&entry), - traffic, - ) - .await?; - entry.touch(); - let Some(terminal_event) = terminal_event else { - return Err(missing_terminal_error()); - }; - - // Handle previous response missing - if is_previous_response_missing(&terminal_event.payload) { - return Err(CodexError { - status: 0, - message: "Previous response not found".to_string(), - detail: Some("previous_response_not_found".to_string()), - retry_after: None, - origin: CodexErrorOrigin::WebSocket, - }); - } - - // Extract status from error events - let status = if terminal_event.event_type == "error" { - event_error_status(&terminal_event.payload).unwrap_or(500) - } else { - 200 - }; - - // Write traffic metadata - if let Some(tc) = traffic { - write_websocket_metadata_capture(tc, &ws_url, pool_owner, continuation, true); - write_websocket_response_capture(tc, status, started_at.elapsed(), &sse_body); - } - - return Ok(CodexResponse { - body: sse_body, - status, - headers: vec![], - transport: ActualTransport::WebSocket, - }); - } - } else { - connect_with_timeout( - websocket_client, - proxy_config, - &ws_url, - headers, - connect_timeout_ms, - ) - .await? + .await?, + )) }; + let mut guard = entry.ws.clone().lock_owned().await; - // New connection path (not pooled or pool miss) - let entry = Arc::new(PoolEntry::new(ws_stream)); - - // Send the request - let msg = Message::Text(body_json); + if used_pooled + && validate_pooled_websocket(&mut guard, connect_timeout_ms) + .await + .is_err() { - let mut ws_guard = entry.ws.lock().await; - ws_guard.send(msg).await.map_err(|e| CodexError { - status: 0, - message: format!("WebSocket send error: {e}"), - detail: None, - retry_after: None, - origin: CodexErrorOrigin::WebSocket, - })?; - - let (sse_body, terminal_event) = collect_ws_events( - &mut ws_guard, - idle_timeout_ms, - pool_owner, - Some(&entry), - traffic, - ) - .await?; - drop(ws_guard); - let Some(terminal_event) = terminal_event else { - return Err(missing_terminal_error()); - }; + drop(guard); + if let Some(owner) = pool_owner { + pool_remove_entry(owner, &entry); + } + if requires_origin { + return Err(continuation_socket_missing_error()); + } + entry = Arc::new(PoolEntry::new( + connect_with_timeout( + websocket_client, + proxy_config, + &ws_url, + headers, + connect_timeout_ms, + ) + .await?, + )); + guard = entry.ws.clone().lock_owned().await; + used_pooled = false; + } - if is_previous_response_missing(&terminal_event.payload) { - if let Some(key) = pool_owner { - invalidate_pool_entry(key, &entry); + guard + .send(Message::Text(body_json)) + .await + .map_err(|error| { + if let Some(owner) = pool_owner { + pool_remove_entry(owner, &entry); } - return Err(CodexError { + CodexError { status: 0, - message: "Previous response not found".to_string(), - detail: Some("previous_response_not_found".to_string()), + message: format!("WebSocket send error: {error}"), + detail: None, retry_after: None, origin: CodexErrorOrigin::WebSocket, - }); - } + } + })?; - // Pool the connection if we have a key and it was successful - if let Some(key) = pool_owner { - let should_pool = terminal_event.event_type == "response.completed"; - if should_pool { - pool_insert_for_turn( - key.clone(), - entry.clone(), - continuation.and_then(|candidate| candidate.turn_id), - ); + let collected = collect_ws_events( + &mut guard, + idle_timeout_ms, + pool_owner, + Some(&entry), + traffic, + ) + .await; + drop(guard); + let (sse_body, terminal_event) = match collected { + Ok(result) => result, + Err(error) => { + if let Some(owner) = pool_owner { + pool_remove_entry(owner, &entry); } + return Err(error); } + }; + let Some(terminal_event) = terminal_event else { + if let Some(owner) = pool_owner { + pool_remove_entry(owner, &entry); + } + return Err(missing_terminal_error()); + }; - let status = if terminal_event.event_type == "error" { - event_error_status(&terminal_event.payload).unwrap_or(500) - } else { - 200 - }; + if is_previous_response_missing(&terminal_event.payload) { + if let Some(owner) = pool_owner { + pool_remove_entry(owner, &entry); + } + return Err(CodexError { + status: 0, + message: "Previous response not found".to_string(), + detail: Some("previous_response_not_found".to_string()), + retry_after: None, + origin: CodexErrorOrigin::WebSocket, + }); + } - // Write traffic metadata - if let Some(tc) = traffic { - write_websocket_metadata_capture(tc, &ws_url, pool_owner, continuation, false); - write_websocket_response_capture(tc, status, started_at.elapsed(), &sse_body); + let completed = terminal_event.event_type == "response.completed"; + let origin_reinserted = if completed { + reservation.is_some_and(|reservation| pool_insert_for_turn(reservation, entry.clone())) + } else { + if let Some(owner) = pool_owner { + pool_remove_entry(owner, &entry); } + false + }; + let status = if terminal_event.event_type == "error" { + event_error_status(&terminal_event.payload).unwrap_or(500) + } else { + 200 + }; - Ok(CodexResponse { + if let Some(tc) = traffic { + write_websocket_metadata_capture(tc, &ws_url, reservation, used_pooled); + write_websocket_response_capture(tc, status, started_at.elapsed(), &sse_body); + } + + Ok(OwnerAwareCodexResponse::new( + CodexResponse { body: sse_body, status, headers: vec![], transport: ActualTransport::WebSocket, - }) - } + }, + origin_reinserted.then_some(entry.socket_id), + )) +} + +async fn validate_pooled_websocket( + websocket: &mut WebSocketStream, + timeout_ms: u64, +) -> Result<(), String> +where + S: AsyncRead + AsyncWrite + Unpin, +{ + let nonce = next_monotonic_nonzero(&POOLED_VALIDATION_SEQUENCE, "pooled validation") + .to_be_bytes() + .to_vec(); + websocket + .send(Message::Ping(nonce.clone())) + .await + .map_err(|error| error.to_string())?; + tokio::time::timeout(Duration::from_millis(timeout_ms), async { + loop { + match websocket.next().await { + Some(Ok(Message::Pong(payload))) if payload == nonce => return Ok(()), + Some(Ok(Message::Ping(payload))) => websocket + .send(Message::Pong(payload)) + .await + .map_err(|error| error.to_string())?, + Some(Ok(Message::Pong(_))) => { + return Err("unexpected Pong during pooled validation".to_string()); + } + Some(Ok(_)) => { + return Err("unexpected frame during pooled validation".to_string()); + } + Some(Err(error)) => return Err(error.to_string()), + None => return Err("connection closed during pooled validation".to_string()), + } + } + }) + .await + .map_err(|_| "validation timeout".to_string())? } pub(super) struct ReadyWebSocket { @@ -752,8 +952,7 @@ pub(super) struct ReadyWebSocket { guard: OwnedMutexGuard, entry: Arc, used_pooled: bool, - pool_owner: Option, - turn_id: Option, + reservation: Option, traffic: Option>, idle_timeout_ms: u64, } @@ -765,11 +964,12 @@ pub(super) async fn prepare_codex_websocket( url: &str, headers: &HeaderMap, traffic: Option>, - pool_owner: Option<&ConversationIdentity>, - turn_id: Option, + reservation: Option<&ContinuationReservation>, connect_timeout_ms: u64, idle_timeout_ms: u64, ) -> Result { + let pool_owner = reservation_pool_owner(reservation); + let continuation = reservation.map(ContinuationReservation::candidate); let ws_url = to_websocket_url(url).map_err(|error| CodexError { status: 0, message: error.message, @@ -777,7 +977,10 @@ pub(super) async fn prepare_codex_websocket( retry_after: None, origin: CodexErrorOrigin::WebSocketHandshake, })?; - let pooled = pool_owner.and_then(|key| pool_take_for_turn(key, turn_id)); + let requires_origin = continuation + .and_then(|candidate| candidate.previous_response_id.as_deref()) + .is_some(); + let pooled = take_pool_entry_for_request(reservation)?; let used_pooled = pooled.is_some(); let entry = if let Some(entry) = pooled { entry @@ -793,39 +996,25 @@ pub(super) async fn prepare_codex_websocket( Arc::new(PoolEntry::new(stream)) }; let mut guard = entry.ws.clone().lock_owned().await; - if used_pooled { - let nonce = b"codex-ready".to_vec(); - guard - .send(Message::Ping(nonce.clone())) - .await - .map_err(|error| pooled_validation_error(error.to_string()))?; - let validation = tokio::time::timeout(Duration::from_millis(connect_timeout_ms), async { - loop { - match guard.next().await { - Some(Ok(Message::Pong(payload))) if payload == nonce => return Ok(()), - Some(Ok(Message::Ping(payload))) => guard - .send(Message::Pong(payload)) - .await - .map_err(|error| error.to_string())?, - Some(Ok(_)) => { - return Err("unexpected frame during pooled validation".to_string()); - } - Some(Err(error)) => return Err(error.to_string()), - None => return Err("connection closed during pooled validation".to_string()), - } - } - }) - .await - .map_err(|_| pooled_validation_error("validation timeout".to_string()))?; - validation.map_err(pooled_validation_error)?; + if used_pooled + && let Err(detail) = validate_pooled_websocket(&mut guard, connect_timeout_ms).await + { + drop(guard); + if let Some(owner) = pool_owner { + pool_remove_entry(owner, &entry); + } + return Err(if requires_origin { + continuation_socket_missing_error() + } else { + pooled_validation_error(detail) + }); } Ok(ReadyWebSocket { ws_url, guard, entry, used_pooled, - pool_owner: pool_owner.cloned(), - turn_id, + reservation: reservation.cloned(), traffic, idle_timeout_ms, }) @@ -846,16 +1035,10 @@ pub(super) fn start_codex_websocket_events( body_value: &serde_json::Value, body_json: String, headers: &HeaderMap, - continuation: Option<&ContinuationCandidate>, -) -> CodexWebSocketEventReceiver { + reservation: Option<&ContinuationReservation>, +) -> CodexWebSocketEventStream { if let Some(tc) = ready.traffic.as_deref() { - write_websocket_metadata_capture( - tc, - &ready.ws_url, - ready.pool_owner.as_ref(), - continuation, - ready.used_pooled, - ); + write_websocket_metadata_capture(tc, &ready.ws_url, reservation, ready.used_pooled); tc.write_json("020-upstream-request", body_value); tc.write_json( "021-upstream-request-metadata", @@ -868,21 +1051,24 @@ pub(super) fn start_codex_websocket_events( ); } let (tx, rx) = mpsc::channel(64); + let (receiver, socket_id_publisher) = CodexWebSocketEventStream::pending(rx); tokio::spawn(async move { let ReadyWebSocket { ws_url: _, mut guard, entry, used_pooled: _, - pool_owner, - turn_id, + reservation, traffic, idle_timeout_ms, } = ready; + let pool_owner = reservation_pool_owner(reservation.as_ref()); if let Err(error) = guard.send(Message::Text(body_json)).await { - if let Some(key) = pool_owner.as_ref() { - pool_remove_entry(key, &entry); + drop(guard); + if let Some(owner) = pool_owner { + pool_remove_entry(owner, &entry); } + socket_id_publisher.publish(None); let _ = tx .send(Err(CodexError { status: 0, @@ -894,25 +1080,30 @@ pub(super) fn start_codex_websocket_events( .await; return; } - let reusable = stream_ws_events( - &mut guard, - idle_timeout_ms, - pool_owner.as_ref(), - Some(&entry), - traffic, - tx, - ) - .await; + let (reusable, terminal_item) = + stream_ws_events(&mut guard, idle_timeout_ms, traffic, &tx).await; drop(guard); - if let Some(key) = pool_owner.as_ref() { - if reusable { - pool_insert_for_turn(key.clone(), entry, turn_id); - } else { - pool_remove_entry(key, &entry); + + let origin_reinserted = if reusable { + reservation + .as_ref() + .is_some_and(|reservation| pool_insert_for_turn(reservation, entry.clone())) + } else { + if let Some(owner) = pool_owner { + pool_remove_entry(owner, &entry); } + false + }; + socket_id_publisher.publish(origin_reinserted.then_some(entry.socket_id)); + if let Some(item) = terminal_item + && tx.send(item).await.is_err() + && origin_reinserted + && let Some(owner) = pool_owner + { + pool_remove_entry(owner, &entry); } }); - rx + receiver } #[allow(clippy::too_many_arguments)] @@ -924,11 +1115,10 @@ pub(super) async fn codex_websocket_event_stream( body_value: &serde_json::Value, _ctx: &RequestContext, traffic: Option>, - pool_owner: Option<&ConversationIdentity>, connect_timeout_ms: u64, idle_timeout_ms: u64, - continuation: Option<&ContinuationCandidate>, -) -> Result { + reservation: Option<&ContinuationReservation>, +) -> Result { let body_json = serde_json::to_string(body_value).map_err(|error| CodexError { status: 500, message: "Failed to serialize WebSocket request".to_string(), @@ -942,8 +1132,7 @@ pub(super) async fn codex_websocket_event_stream( url, headers, traffic, - pool_owner, - continuation.and_then(|candidate| candidate.turn_id), + reservation, connect_timeout_ms, idle_timeout_ms, ) @@ -953,10 +1142,20 @@ pub(super) async fn codex_websocket_event_stream( body_value, body_json, headers, - continuation, + reservation, )) } +fn continuation_socket_missing_error() -> CodexError { + CodexError { + status: 0, + message: "Previous response socket is no longer available".to_string(), + detail: Some(WEBSOCKET_CONTINUATION_SOCKET_MISSING_DETAIL.to_string()), + retry_after: None, + origin: CodexErrorOrigin::WebSocketHandshake, + } +} + fn missing_terminal_error() -> CodexError { CodexError { status: 0, @@ -980,10 +1179,11 @@ fn response_start_timeout_error(timeout_ms: u64) -> CodexError { fn write_websocket_metadata_capture( traffic: &TrafficCapture, ws_url: &str, - pool_owner: Option<&ConversationIdentity>, - continuation: Option<&ContinuationCandidate>, + reservation: Option<&ContinuationReservation>, pooled: bool, ) { + let pool_owner = reservation_pool_owner(reservation); + let continuation = reservation.map(ContinuationReservation::candidate); traffic.write_json( "022-upstream-websocket-metadata", &serde_json::json!({ @@ -1905,11 +2105,9 @@ where async fn stream_ws_events( ws: &mut WebSocketStream, idle_timeout_ms: u64, - pool_owner: Option<&ConversationIdentity>, - pool_entry: Option<&Arc>, traffic: Option>, - tx: mpsc::Sender>, -) -> bool + tx: &mpsc::Sender>, +) -> (bool, Option>) where S: AsyncRead + AsyncWrite + Unpin, { @@ -1921,6 +2119,7 @@ where let mut response_started = false; let mut status = 200u16; let mut reusable = false; + let mut terminal_item = None; loop { let response_deadline_started = if response_started { @@ -1932,8 +2131,7 @@ where match response_event_budget.checked_sub(response_deadline_started.elapsed()) { Some(remaining) if !remaining.is_zero() => remaining, _ => { - invalidate_pool_owner(pool_owner, pool_entry); - let err = if response_started { + terminal_item = Some(Err(if response_started { CodexError { status: 0, message: format!("WebSocket idle timeout after {idle_timeout_ms}ms"), @@ -1943,36 +2141,37 @@ where } } else { response_start_timeout_error(idle_timeout_ms) - }; - let _ = tx.send(Err(err)).await; + })); break; } }; - let frame = match tokio::time::timeout(read_timeout, ws.next()).await { - Ok(frame) => frame, - Err(_) => { - invalidate_pool_owner(pool_owner, pool_entry); - let err = if response_started { - CodexError { - status: 0, - message: format!("WebSocket idle timeout after {idle_timeout_ms}ms"), - detail: None, - retry_after: None, - origin: CodexErrorOrigin::WebSocket, - } - } else { - response_start_timeout_error(idle_timeout_ms) - }; - let _ = tx.send(Err(err)).await; - break; - } + let frame = tokio::select! { + biased; + _ = tx.closed() => break, + frame = tokio::time::timeout(read_timeout, ws.next()) => match frame { + Ok(frame) => frame, + Err(_) => { + terminal_item = Some(Err(if response_started { + CodexError { + status: 0, + message: format!("WebSocket idle timeout after {idle_timeout_ms}ms"), + detail: None, + retry_after: None, + origin: CodexErrorOrigin::WebSocket, + } + } else { + response_start_timeout_error(idle_timeout_ms) + })); + break; + } + }, }; match frame { Some(Ok(Message::Text(text))) => { let parsed: serde_json::Value = match serde_json::from_str(&text) { - Ok(v) => v, + Ok(value) => value, Err(_) => { if let Some(tc) = traffic.as_deref() { tc.write_json_event( @@ -1997,49 +2196,39 @@ where response_started = true; last_response_event_at = Instant::now(); } - - if parsed.get("type").and_then(|v| v.as_str()) == Some("error") { + if parsed.get("type").and_then(|value| value.as_str()) == Some("error") { status = event_error_status(&parsed).unwrap_or(500); } - let terminal = is_terminal_event(&parsed); - if terminal && is_previous_response_missing(&parsed) { - invalidate_pool_owner(pool_owner, pool_entry); - let _ = tx - .send(Err(CodexError { + + if is_terminal_event(&parsed) { + if is_previous_response_missing(&parsed) { + terminal_item = Some(Err(CodexError { status: 0, message: "Previous response not found".to_string(), detail: Some("previous_response_not_found".to_string()), retry_after: None, origin: CodexErrorOrigin::WebSocket, - })) - .await; + })); + } else { + reusable = parsed.get("type").and_then(|value| value.as_str()) + == Some("response.completed"); + terminal_item = Some(Ok(parsed)); + } break; } - let event_type = parsed - .get("type") - .and_then(|v| v.as_str()) - .unwrap_or("unknown") - .to_string(); + if tx.send(Ok(parsed)).await.is_err() { - invalidate_pool_owner(pool_owner, pool_entry); - break; - } - if terminal { - reusable = event_type == "response.completed"; break; } } Some(Ok(Message::Binary(_))) => { - invalidate_pool_owner(pool_owner, pool_entry); - let _ = tx - .send(Err(CodexError { - status: 0, - message: "WebSocket binary frames not supported".to_string(), - detail: None, - retry_after: None, - origin: CodexErrorOrigin::WebSocket, - })) - .await; + terminal_item = Some(Err(CodexError { + status: 0, + message: "WebSocket binary frames not supported".to_string(), + detail: None, + retry_after: None, + origin: CodexErrorOrigin::WebSocket, + })); break; } Some(Ok(Message::Ping(data))) => { @@ -2047,21 +2236,17 @@ where } Some(Ok(Message::Pong(_))) | Some(Ok(Message::Frame(_))) => {} Some(Ok(Message::Close(_))) | None => { - invalidate_pool_owner(pool_owner, pool_entry); - let _ = tx.send(Err(missing_terminal_error())).await; + terminal_item = Some(Err(missing_terminal_error())); break; } - Some(Err(e)) => { - invalidate_pool_owner(pool_owner, pool_entry); - let _ = tx - .send(Err(CodexError { - status: 0, - message: format!("WebSocket stream error: {e}"), - detail: None, - retry_after: None, - origin: CodexErrorOrigin::WebSocket, - })) - .await; + Some(Err(error)) => { + terminal_item = Some(Err(CodexError { + status: 0, + message: format!("WebSocket stream error: {error}"), + detail: None, + retry_after: None, + origin: CodexErrorOrigin::WebSocket, + })); break; } } @@ -2070,7 +2255,7 @@ where if let Some(tc) = traffic.as_deref() { write_websocket_response_capture(tc, status, started_at.elapsed(), &sse_body); } - reusable + (reusable, terminal_item) } fn headers_to_json(headers: &HeaderMap) -> serde_json::Value { @@ -2108,10 +2293,18 @@ mod tests { use super::*; - static WS_POOL_TEST_LOCK: AsyncMutex<()> = AsyncMutex::const_new(()); + #[test] + fn provider_retry_handoff_is_attempt_local() { + let (_tx_a, rx_a) = mpsc::channel(1); + let (stream_a, publisher_a) = CodexWebSocketEventStream::pending(rx_a); + let (_tx_b, rx_b) = mpsc::channel(1); + let (_stream_b, publisher_b) = CodexWebSocketEventStream::pending(rx_b); - async fn lock_ws_pool_tests() -> tokio::sync::MutexGuard<'static, ()> { - WS_POOL_TEST_LOCK.lock().await + assert!(!publisher_a.is_provider_retry_handoff()); + assert!(!publisher_b.is_provider_retry_handoff()); + stream_a.mark_provider_retry_handoff(); + assert!(publisher_a.is_provider_retry_handoff()); + assert!(!publisher_b.is_provider_retry_handoff()); } fn main_owner(session_id: &str) -> ConversationIdentity { @@ -2122,6 +2315,47 @@ mod tests { ConversationIdentity::Agent(session_id.to_string(), agent_id.to_string()) } + fn test_continuation( + owner: Option, + turn_id: Option, + previous_response_id: Option<&str>, + origin_socket_id: Option, + ) -> ContinuationReservation { + ContinuationReservation::new( + super::super::continuation::ContinuationCandidate { + turn_id, + previous_response_id: previous_response_id.map(str::to_string), + input_delta: Some(vec![]), + input_delta_count: 0, + disabled_reason: None, + }, + owner, + origin_socket_id, + ) + } + + fn continuation_request() -> super::super::translate::request::ResponsesRequest { + super::super::translate::request::ResponsesRequest { + model: "gpt-5.6-sol".to_string(), + instructions: None, + input: vec![], + tools: None, + tool_choice: None, + store: false, + stream: true, + parallel_tool_calls: true, + include: None, + client_metadata: None, + service_tier: None, + prompt_cache_key: None, + text: super::super::translate::request::ResponsesText { + verbosity: None, + format: None, + }, + reasoning: None, + } + } + fn test_websocket_client() -> reqwest::Client { reqwest::Client::builder() .http1_only() @@ -2196,6 +2430,7 @@ mod tests { fn shared_pool_entry(ws: &Arc>) -> Arc { Arc::new(PoolEntry { ws: ws.clone(), + socket_id: next_monotonic_nonzero(&NEXT_SOCKET_ID, "WebSocket ID"), created_at: now_ms(), last_activity: AtomicU64::new(next_pool_activity()), }) @@ -2203,7 +2438,7 @@ mod tests { #[tokio::test] async fn connect_cleanup_uses_threshold_target_lru_and_skips_leases() { - let _pool_test_guard = lock_ws_pool_tests().await; + let _pool_test_guard = lock_codex_websocket_pool_for_tests().await; clear_codex_websocket_pool_for_tests(); let ws = Arc::new(AsyncMutex::new(raw_test_stream(None).await)); { @@ -2257,7 +2492,7 @@ mod tests { #[tokio::test] async fn connect_cleanup_drops_final_socket_owners_after_unlocking_pool() { - let _pool_test_guard = lock_ws_pool_tests().await; + let _pool_test_guard = lock_codex_websocket_pool_for_tests().await; clear_codex_websocket_pool_for_tests(); let dropped = Arc::new(AtomicBool::new(false)); let pool_was_unlocked = Arc::new(AtomicBool::new(false)); @@ -2803,6 +3038,17 @@ mod tests { }); assert!(is_previous_response_missing(&by_msg)); + let nested = serde_json::json!({ + "type": "response.failed", + "response": { + "error": { + "code": "previous_response_not_found", + "message": "Previous response not found" + } + } + }); + assert!(is_previous_response_missing(&nested)); + let unrelated = serde_json::json!({"type": "error", "error": {"message": "rate limited"}}); assert!(!is_previous_response_missing(&unrelated)); } @@ -2812,12 +3058,12 @@ mod tests { let temp = tempfile::tempdir().unwrap(); let traffic = crate::traffic::test_capture(temp.path().join("traffic")); let owner = agent_owner("session-secret", "agent-secret"); + let reservation = test_continuation(Some(owner), None, None, None); write_websocket_metadata_capture( &traffic, "wss://example.invalid/responses", - Some(&owner), - None, + Some(&reservation), false, ); @@ -2836,7 +3082,7 @@ mod tests { #[tokio::test] async fn pool_checkout_is_exclusive_and_removal_is_identity_safe() { - let _pool_test_guard = lock_ws_pool_tests().await; + let _pool_test_guard = lock_codex_websocket_pool_for_tests().await; clear_codex_websocket_pool_for_tests(); let first = Arc::new(PoolEntry::new(create_dummy_stream_async().await)); let owner = main_owner("exclusive"); @@ -2850,6 +3096,8 @@ mod tests { let replacement = Arc::new(PoolEntry::new(create_dummy_stream_async().await)); pool_insert(owner.clone(), replacement.clone()); pool_remove_entry(&owner, &first); + let reservation = test_continuation(Some(owner.clone()), None, None, None); + invalidate_codex_websocket_pool_socket(&reservation, Some(first.socket_id)); assert!(Arc::ptr_eq( WS_POOL.lock().unwrap().get(&owner).unwrap(), &replacement @@ -2857,9 +3105,393 @@ mod tests { clear_codex_websocket_pool_for_tests(); } + #[tokio::test] + async fn pooled_validation_uses_unique_nonce_and_rejects_stale_pong() { + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = listener.local_addr().unwrap(); + let (queue_stale_tx, queue_stale_rx) = tokio::sync::oneshot::channel(); + let (stale_queued_tx, stale_queued_rx) = tokio::sync::oneshot::channel(); + let (release_peer_tx, release_peer_rx) = tokio::sync::oneshot::channel(); + let peer = tokio::spawn(async move { + let (socket, _) = listener.accept().await.unwrap(); + let mut websocket = tokio_tungstenite::accept_async(socket).await.unwrap(); + let prior_nonce = match websocket.next().await { + Some(Ok(Message::Ping(payload))) => payload, + other => panic!("unexpected first validation frame: {other:?}"), + }; + assert_eq!(prior_nonce.len(), std::mem::size_of::()); + assert_ne!(prior_nonce, 0_u64.to_be_bytes()); + websocket + .send(Message::Pong(prior_nonce.clone())) + .await + .unwrap(); + + queue_stale_rx.await.unwrap(); + websocket + .send(Message::Pong(prior_nonce.clone())) + .await + .unwrap(); + stale_queued_tx.send(()).unwrap(); + + let current_nonce = match websocket.next().await { + Some(Ok(Message::Ping(payload))) => payload, + other => panic!("unexpected second validation frame: {other:?}"), + }; + assert_ne!(current_nonce, prior_nonce); + assert_ne!(current_nonce, 0_u64.to_be_bytes()); + release_peer_rx.await.unwrap(); + }); + let (mut websocket, _) = tokio_tungstenite::connect_async(format!("ws://{addr}/")) + .await + .unwrap(); + + validate_pooled_websocket(&mut websocket, 1_000) + .await + .unwrap(); + queue_stale_tx.send(()).unwrap(); + stale_queued_rx.await.unwrap(); + let error = validate_pooled_websocket(&mut websocket, 100) + .await + .unwrap_err(); + assert_eq!(error, "unexpected Pong during pooled validation"); + + release_peer_tx.send(()).unwrap(); + peer.await.unwrap(); + } + + #[tokio::test] + async fn pool_entries_receive_monotonic_nonzero_socket_ids() { + let first = PoolEntry::new(raw_test_stream(None).await); + let second = PoolEntry::new(raw_test_stream(None).await); + + assert_ne!(first.socket_id, 0); + assert!(second.socket_id > first.socket_id); + } + + #[tokio::test] + async fn continuation_rejects_and_preserves_same_owner_replacement() { + let _registry_guard = + super::super::continuation::lock_continuation_registry_for_async_tests().await; + let _pool_test_guard = lock_codex_websocket_pool_for_tests().await; + let owner = agent_owner("replacement-session", "replacement-agent"); + super::super::continuation::clear_continuation_for_owner(Some(&owner)); + invalidate_codex_websocket_pool_owner(&owner); + let request = continuation_request(); + let reserved = super::super::continuation::continuation_candidate_for_owner( + Some(&owner), + &request, + true, + ); + let replacement = Arc::new(PoolEntry::new(raw_test_stream(None).await)); + pool_insert(owner.clone(), replacement.clone()); + let continuation = test_continuation( + Some(owner.clone()), + reserved.turn_id(), + Some("resp_origin"), + Some(replacement.socket_id.checked_add(1).unwrap()), + ); + + let error = match prepare_codex_websocket( + &test_websocket_client(), + &WebSocketProxyConfig::direct(), + "ws://127.0.0.1:9/responses", + &HeaderMap::new(), + None, + Some(&continuation), + 50, + 50, + ) + .await + { + Ok(_) => panic!("replacement socket must not satisfy continuation provenance"), + Err(error) => error, + }; + + assert_eq!( + error.detail.as_deref(), + Some(WEBSOCKET_CONTINUATION_SOCKET_MISSING_DETAIL) + ); + assert!(!error.message.contains("replacement-session")); + assert!(!error.message.contains("replacement-agent")); + assert!(Arc::ptr_eq( + WS_POOL.lock().unwrap().get(&owner).unwrap(), + &replacement + )); + + invalidate_codex_websocket_pool_owner(&owner); + super::super::continuation::abort_continuation_for_owner(&reserved); + } + + #[tokio::test] + async fn dead_exact_origin_removes_only_that_arc_and_preserves_replacement() { + let _registry_guard = + super::super::continuation::lock_continuation_registry_for_async_tests().await; + let _pool_test_guard = lock_codex_websocket_pool_for_tests().await; + let owner = main_owner("dead-origin-session"); + super::super::continuation::clear_continuation_for_owner(Some(&owner)); + invalidate_codex_websocket_pool_owner(&owner); + let request = continuation_request(); + let reserved = super::super::continuation::continuation_candidate_for_owner( + Some(&owner), + &request, + true, + ); + let exact = Arc::new(PoolEntry::new(raw_test_stream(None).await)); + let replacement = Arc::new(PoolEntry::new(raw_test_stream(None).await)); + pool_insert(owner.clone(), exact.clone()); + let continuation = test_continuation( + Some(owner.clone()), + reserved.turn_id(), + Some("resp_exact"), + Some(exact.socket_id), + ); + + let replacement_owner = owner.clone(); + let replacement_for_task = replacement.clone(); + let mut insert_replacement = tokio::spawn(async move { + tokio::time::timeout(Duration::from_secs(1), async move { + loop { + if !WS_POOL.lock().unwrap().contains_key(&replacement_owner) { + pool_insert(replacement_owner, replacement_for_task); + return; + } + tokio::task::yield_now().await; + } + }) + .await + }); + let error = match prepare_codex_websocket( + &test_websocket_client(), + &WebSocketProxyConfig::direct(), + "ws://127.0.0.1:9/responses", + &HeaderMap::new(), + None, + Some(&continuation), + 50, + 50, + ) + .await + { + Ok(_) => panic!("dead continuation origin must be rejected before request send"), + Err(error) => error, + }; + match tokio::time::timeout(Duration::from_secs(2), &mut insert_replacement).await { + Ok(Ok(Ok(()))) => {} + Ok(Ok(Err(_))) => panic!( + "replacement polling timed out; currently pooled socket ID: {:?}", + pooled_socket_id_for_tests(&owner) + ), + Ok(Err(error)) => panic!( + "replacement polling task failed ({error}); currently pooled socket ID: {:?}", + pooled_socket_id_for_tests(&owner) + ), + Err(_) => { + insert_replacement.abort(); + let abort_result = insert_replacement.await; + panic!( + "replacement polling join timed out ({abort_result:?}); currently pooled socket ID: {:?}", + pooled_socket_id_for_tests(&owner) + ); + } + } + + assert_eq!( + error.detail.as_deref(), + Some(WEBSOCKET_CONTINUATION_SOCKET_MISSING_DETAIL) + ); + assert!(Arc::ptr_eq( + WS_POOL.lock().unwrap().get(&owner).unwrap(), + &replacement + )); + assert!(!Arc::ptr_eq( + WS_POOL.lock().unwrap().get(&owner).unwrap(), + &exact + )); + + invalidate_codex_websocket_pool_owner(&owner); + super::super::continuation::abort_continuation_for_owner(&reserved); + } + + #[tokio::test] + async fn completed_terminal_is_published_after_origin_returns_to_pool() { + let _registry_guard = + super::super::continuation::lock_continuation_registry_for_async_tests().await; + let _pool_test_guard = lock_codex_websocket_pool_for_tests().await; + let owner = agent_owner("terminal-order-session", "terminal-order-agent"); + super::super::continuation::clear_continuation_for_owner(Some(&owner)); + invalidate_codex_websocket_pool_owner(&owner); + let request = continuation_request(); + let continuation = super::super::continuation::continuation_candidate_for_owner( + Some(&owner), + &request, + true, + ); + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = listener.local_addr().unwrap(); + let (release_response_tx, release_response_rx) = tokio::sync::oneshot::channel(); + let server = tokio::spawn(async move { + let (socket, _) = listener.accept().await.unwrap(); + let mut websocket = tokio_tungstenite::accept_async(socket).await.unwrap(); + while let Some(Ok(message)) = websocket.next().await { + match message { + Message::Ping(payload) => { + websocket.send(Message::Pong(payload)).await.unwrap(); + } + Message::Text(_) => { + release_response_rx.await.unwrap(); + websocket + .send(Message::Text( + serde_json::json!({ + "type": "response.completed", + "response": {"id": "resp_terminal_order"} + }) + .to_string(), + )) + .await + .unwrap(); + return; + } + _ => {} + } + } + }); + let context = RequestContext { + req_id: "terminal-order-request".to_string(), + session_id: Some("header-session-is-not-pool-owner".to_string()), + session_seq: None, + provider: "codex".to_string(), + traffic: None, + monitor: None, + }; + let mut events = codex_websocket_event_stream( + &test_websocket_client(), + &WebSocketProxyConfig::direct(), + &format!("http://{addr}/responses"), + &HeaderMap::new(), + &serde_json::json!({"type":"response.create","input":[]}), + &context, + None, + 1_000, + 1_000, + Some(&continuation), + ) + .await + .unwrap(); + assert_eq!(events.socket_id(), None); + release_response_tx.send(()).unwrap(); + let terminal = events.recv().await.unwrap().unwrap(); + + assert_eq!( + terminal.get("type").and_then(serde_json::Value::as_str), + Some("response.completed") + ); + let pooled = WS_POOL + .lock() + .unwrap() + .get(&owner) + .cloned() + .expect("origin must be reusable before terminal publication"); + assert_eq!(events.socket_id(), Some(pooled.socket_id)); + server.await.unwrap(); + + invalidate_codex_websocket_pool_owner(&owner); + super::super::continuation::abort_continuation_for_owner(&continuation); + } + + #[tokio::test] + async fn dropped_receiver_removes_completed_origin_after_terminal_publication_fails() { + let _registry_guard = + super::super::continuation::lock_continuation_registry_for_async_tests().await; + let _pool_test_guard = lock_codex_websocket_pool_for_tests().await; + let owner = main_owner("dropped-terminal-receiver-session"); + super::super::continuation::clear_continuation_for_owner(Some(&owner)); + invalidate_codex_websocket_pool_owner(&owner); + let request = continuation_request(); + let continuation = super::super::continuation::continuation_candidate_for_owner( + Some(&owner), + &request, + true, + ); + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = listener.local_addr().unwrap(); + let (request_seen_tx, request_seen_rx) = tokio::sync::oneshot::channel(); + let (release_response_tx, release_response_rx) = tokio::sync::oneshot::channel(); + let server = tokio::spawn(async move { + let (socket, _) = listener.accept().await.unwrap(); + let mut websocket = tokio_tungstenite::accept_async(socket).await.unwrap(); + while let Some(Ok(message)) = websocket.next().await { + match message { + Message::Ping(payload) => { + websocket.send(Message::Pong(payload)).await.unwrap(); + } + Message::Text(_) => { + request_seen_tx.send(()).unwrap(); + release_response_rx.await.unwrap(); + websocket + .send(Message::Text( + serde_json::json!({ + "type": "response.completed", + "response": {"id": "resp_dropped_receiver"} + }) + .to_string(), + )) + .await + .unwrap(); + return; + } + _ => {} + } + } + }); + let headers = HeaderMap::new(); + let body = serde_json::json!({"type":"response.create","input":[]}); + let ready = prepare_codex_websocket( + &test_websocket_client(), + &WebSocketProxyConfig::direct(), + &format!("http://{addr}/responses"), + &headers, + None, + Some(&continuation), + 1_000, + 1_000, + ) + .await + .unwrap(); + let exact = ready.entry.clone(); + let events = start_codex_websocket_events( + ready, + &body, + serde_json::to_string(&body).unwrap(), + &headers, + Some(&continuation), + ); + + request_seen_rx.await.unwrap(); + drop(events); + release_response_tx.send(()).unwrap(); + server.await.unwrap(); + tokio::time::timeout(Duration::from_secs(1), async { + while Arc::strong_count(&exact) != 1 { + tokio::task::yield_now().await; + } + }) + .await + .expect("failed terminal publication must release the exact completed pool entry"); + + assert!( + WS_POOL + .lock() + .unwrap() + .get(&owner) + .is_none_or(|pooled| !Arc::ptr_eq(pooled, &exact)), + "the exact completed socket must not remain pooled" + ); + super::super::continuation::abort_continuation_for_owner(&continuation); + } + #[tokio::test] async fn pool_invalidation() { - let _pool_test_guard = lock_ws_pool_tests().await; + let _pool_test_guard = lock_codex_websocket_pool_for_tests().await; clear_codex_websocket_pool_for_tests(); let first_stream = create_dummy_stream_async().await; let second_stream = create_dummy_stream_async().await; @@ -2882,9 +3514,82 @@ mod tests { clear_codex_websocket_pool_for_tests(); } + #[tokio::test] + #[allow(deprecated)] + async fn session_key_invalidation_removes_main_and_agents_only_for_exact_session() { + let _pool_test_guard = lock_codex_websocket_pool_for_tests().await; + clear_codex_websocket_pool_for_tests(); + let session_id = "session-key"; + let main = main_owner(session_id); + let first_agent = agent_owner(session_id, "first-agent"); + let second_agent = agent_owner(session_id, "second-agent"); + let other_session = main_owner("session-key-other"); + for owner in [ + main.clone(), + first_agent.clone(), + second_agent.clone(), + other_session.clone(), + ] { + pool_insert( + owner, + Arc::new(PoolEntry::new(create_dummy_stream_async().await)), + ); + } + + invalidate_codex_websocket_pool_key(session_id); + + let pool = WS_POOL.lock().unwrap(); + assert!(!pool.contains_key(&main)); + assert!(!pool.contains_key(&first_agent)); + assert!(!pool.contains_key(&second_agent)); + assert!(pool.contains_key(&other_session)); + drop(pool); + clear_codex_websocket_pool_for_tests(); + } + + #[tokio::test] + async fn missing_owner_or_turn_cannot_mutate_pool_state() { + let _registry_guard = + super::super::continuation::lock_continuation_registry_for_async_tests().await; + let _pool_test_guard = lock_codex_websocket_pool_for_tests().await; + let owner = main_owner("missing-owner-turn-pool"); + super::super::continuation::clear_continuation_for_owner(Some(&owner)); + clear_codex_websocket_pool_for_tests(); + let request = continuation_request(); + let current = super::super::continuation::continuation_candidate_for_owner( + Some(&owner), + &request, + true, + ); + let pooled = Arc::new(PoolEntry::new(create_dummy_stream_async().await)); + pool_insert(owner.clone(), pooled.clone()); + let missing_owner = test_continuation(None, current.turn_id(), None, None); + let missing_turn = test_continuation(Some(owner.clone()), None, None, None); + + assert!(pool_take_for_turn(&missing_owner).is_none()); + assert!(pool_take_for_turn(&missing_turn).is_none()); + assert!(!pool_insert_for_turn( + &missing_owner, + Arc::new(PoolEntry::new(create_dummy_stream_async().await)), + )); + assert!(!pool_insert_for_turn( + &missing_turn, + Arc::new(PoolEntry::new(create_dummy_stream_async().await)), + )); + invalidate_codex_websocket_pool_turn_for_owner(&owner, None); + invalidate_codex_websocket_pool_socket(&missing_owner, Some(pooled.socket_id)); + + assert!(Arc::ptr_eq( + WS_POOL.lock().unwrap().get(&owner).unwrap(), + &pooled + )); + clear_codex_websocket_pool_for_tests(); + super::super::continuation::abort_continuation_for_owner(¤t); + } + #[tokio::test] async fn websocket_connect_401_is_pre_request_handshake_error() { - let _pool_test_guard = lock_ws_pool_tests().await; + let _pool_test_guard = lock_codex_websocket_pool_for_tests().await; use tokio::io::{AsyncReadExt, AsyncWriteExt}; use tokio::net::TcpListener; @@ -2921,7 +3626,7 @@ mod tests { #[tokio::test] async fn websocket_connect_502_preserves_retry_metadata() { - let _pool_test_guard = lock_ws_pool_tests().await; + let _pool_test_guard = lock_codex_websocket_pool_for_tests().await; use tokio::io::{AsyncReadExt, AsyncWriteExt}; use tokio::net::TcpListener; @@ -2961,7 +3666,7 @@ mod tests { #[tokio::test] async fn websocket_connects_through_explicit_http_proxy() { - let _pool_test_guard = lock_ws_pool_tests().await; + let _pool_test_guard = lock_codex_websocket_pool_for_tests().await; use tokio::io::{AsyncReadExt, AsyncWriteExt}; use tokio::net::TcpListener; @@ -3132,7 +3837,7 @@ mod tests { #[tokio::test] async fn websocket_wss_uses_http_connect_without_leaking_proxy_credentials() { - let _pool_test_guard = lock_ws_pool_tests().await; + let _pool_test_guard = lock_codex_websocket_pool_for_tests().await; use tokio::io::{AsyncReadExt, AsyncWriteExt}; use tokio::net::TcpListener; @@ -3199,7 +3904,7 @@ mod tests { #[tokio::test] async fn binary_frame_invalidates_pool_owner() { - let _pool_test_guard = lock_ws_pool_tests().await; + let _pool_test_guard = lock_codex_websocket_pool_for_tests().await; clear_codex_websocket_pool_for_tests(); let pooled_stream = create_dummy_stream_async().await; let owner = agent_owner("binary-session", "binary-agent"); @@ -3230,7 +3935,7 @@ mod tests { #[tokio::test] async fn response_start_timeout_ignores_rate_limits_and_pings() { - let _pool_test_guard = lock_ws_pool_tests().await; + let _pool_test_guard = lock_codex_websocket_pool_for_tests().await; clear_codex_websocket_pool_for_tests(); let pooled_stream = create_dummy_stream_async().await; let owner = main_owner("start-timeout-session"); @@ -3274,7 +3979,7 @@ mod tests { #[tokio::test] async fn response_idle_timeout_ignores_pings_after_response_event() { - let _pool_test_guard = lock_ws_pool_tests().await; + let _pool_test_guard = lock_codex_websocket_pool_for_tests().await; clear_codex_websocket_pool_for_tests(); let pooled_stream = create_dummy_stream_async().await; let owner = main_owner("response-idle-session"); diff --git a/tests/public_codex_api_compat.rs b/tests/public_codex_api_compat.rs new file mode 100644 index 00000000..e3cff8a2 --- /dev/null +++ b/tests/public_codex_api_compat.rs @@ -0,0 +1,94 @@ +#![allow(deprecated)] + +use std::pin::Pin; +use std::sync::Arc; +use std::task::{Context, Poll}; + +use claude_code_proxy::provider::RequestContext; +use claude_code_proxy::providers::codex::client::{ + ActualTransport, CodexError, CodexHttpClient, CodexResponse, +}; +use claude_code_proxy::providers::codex::continuation::{ + ContinuationCandidate, abort_continuation, clear_continuation, continuation_candidate, + has_continuation_for_tests, if_current_turn, is_current_turn, record_continuation, + with_current_turn, +}; +use claude_code_proxy::providers::codex::translate::request::ResponsesRequest; +use claude_code_proxy::providers::codex::websocket::{ + CodexWebSocketEventReceiver, invalidate_codex_websocket_pool_key, + invalidate_codex_websocket_pool_turn, +}; + +fn baseline_continuation_forms(session_id: Option<&str>, body: &ResponsesRequest) { + let candidate = continuation_candidate(session_id, body, true); + record_continuation( + session_id, + candidate.turn_id, + body, + candidate.previous_response_id.as_deref(), + &[], + ); + abort_continuation(session_id, candidate.turn_id); + let _: Option = if_current_turn(session_id, candidate.turn_id, || 1); + let _: bool = with_current_turn(session_id, candidate.turn_id, || {}); + let _: bool = is_current_turn(session_id, candidate.turn_id); + clear_continuation(session_id); + let _: bool = has_continuation_for_tests("compat-session"); +} + +fn baseline_pool_invalidation_forms() { + invalidate_codex_websocket_pool_key("compat-session"); + invalidate_codex_websocket_pool_turn("compat-session", Some(1)); +} + +async fn baseline_client_forms( + client: &Arc, + body: &ResponsesRequest, + ctx: &RequestContext, + candidate: &ContinuationCandidate, +) { + let _: Result = client.post_codex(body, ctx, Some(candidate)).await; + let _ = client + .stream_codex_websocket_events(body, ctx, Some(candidate)) + .await; +} + +#[test] +fn baseline_codex_public_api_forms_compile() { + let (_tx, receiver) = tokio::sync::mpsc::channel::>(1); + let mut receiver: CodexWebSocketEventReceiver = receiver; + + let _ = receiver.try_recv(); + receiver.close(); + let _: usize = receiver.len(); + let waker = futures_util::task::noop_waker(); + let mut context = Context::from_waker(&waker); + let _: Poll>> = + Pin::new(&mut receiver).poll_recv(&mut context); + + let response = CodexResponse { + body: vec![1, 2, 3], + status: 200, + headers: vec![("content-type".to_string(), "text/event-stream".to_string())], + transport: ActualTransport::WebSocket, + }; + assert_eq!(response.body, vec![1, 2, 3]); + assert_eq!(response.status, 200); + assert_eq!( + response.headers, + vec![("content-type".to_string(), "text/event-stream".to_string())] + ); + assert_eq!(response.transport, ActualTransport::WebSocket); + + let _candidate = ContinuationCandidate { + turn_id: Some(1), + previous_response_id: Some("resp_compat".to_string()), + input_delta: None, + input_delta_count: 0, + disabled_reason: None, + }; + + let _: fn(Option<&str>, &ResponsesRequest) = baseline_continuation_forms; + let _: fn() = baseline_pool_invalidation_forms; + let _ = baseline_client_forms; +} From a0ac9f69d6065089c079ce09d29726f24d54d98a Mon Sep 17 00:00:00 2001 From: Albert Zhang Date: Sun, 2 Aug 2026 13:42:11 +0800 Subject: [PATCH 3/4] test(codex): cover agent continuation over HTTP ingress Co-Authored-By: Claude --- tests/codex_agent_continuation.rs | 1893 +++++++++++++++++++++++++++++ 1 file changed, 1893 insertions(+) create mode 100644 tests/codex_agent_continuation.rs diff --git a/tests/codex_agent_continuation.rs b/tests/codex_agent_continuation.rs new file mode 100644 index 00000000..5a99c1a8 --- /dev/null +++ b/tests/codex_agent_continuation.rs @@ -0,0 +1,1893 @@ +use std::collections::{HashMap, HashSet}; +use std::ffi::{OsStr, OsString}; +use std::path::Path; +use std::sync::{Arc, Mutex, OnceLock}; +use std::time::Duration; + +use claude_code_proxy::providers::codex::websocket::{ + WEBSOCKET_PROTOCOL_HEADER, invalidate_codex_websocket_pool_owner, +}; +use claude_code_proxy::request_identity::ConversationIdentity; +use claude_code_proxy::server; +use futures_util::{SinkExt, StreamExt}; +use http::{HeaderMap, StatusCode}; +use serde_json::{Value, json}; +use tempfile::TempDir; +use tokio::net::{TcpListener, TcpStream}; +use tokio::sync::{mpsc, oneshot}; +use tokio::task::{JoinHandle, JoinSet}; +use tokio_tungstenite::tungstenite::Message; +use tokio_tungstenite::{WebSocketStream, accept_hdr_async}; +use uuid::Uuid; + +const REQUEST_TIMEOUT: Duration = Duration::from_secs(15); +const SESSION_HEADER: &str = "x-claude-code-session-id"; +const AGENT_HEADER: &str = "x-claude-code-agent-id"; +const PARENT_AGENT_HEADER: &str = "x-claude-code-parent-agent-id"; + +type ProbeLog = Arc)>>>; + +static ENV_LOCK: OnceLock> = OnceLock::new(); + +fn env_lock() -> std::sync::MutexGuard<'static, ()> { + ENV_LOCK + .get_or_init(|| Mutex::new(())) + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()) +} + +struct EnvGuard { + key: &'static str, + previous: Option, +} + +impl EnvGuard { + fn set(key: &'static str, value: impl AsRef) -> Self { + let previous = std::env::var_os(key); + unsafe { + std::env::set_var(key, value); + } + Self { key, previous } + } + + fn unset(key: &'static str) -> Self { + let previous = std::env::var_os(key); + unsafe { + std::env::remove_var(key); + } + Self { key, previous } + } +} + +impl Drop for EnvGuard { + fn drop(&mut self) { + unsafe { + match self.previous.take() { + Some(value) => std::env::set_var(self.key, value), + None => std::env::remove_var(self.key), + } + } + } +} + +fn configure_environment(config_dir: &Path, upstream_url: &str) -> Vec { + let mut guards = [ + "HTTP_PROXY", + "http_proxy", + "HTTPS_PROXY", + "https_proxy", + "ALL_PROXY", + "all_proxy", + "REQUEST_METHOD", + "CCP_AUTO_REVIEW_MODEL", + "CCP_CODEX_MODEL", + "CCP_CODEX_SERVICE_TIER", + "CCP_CODEX_EFFORT", + "CCP_CODEX_REASONING_SUMMARY", + "CCP_CODEX_ORIGINATOR", + "CCP_CODEX_USER_AGENT", + "CCP_USER_AGENT", + ] + .into_iter() + .map(EnvGuard::unset) + .collect::>(); + guards.extend([ + EnvGuard::set("NO_PROXY", "127.0.0.1,localhost"), + EnvGuard::set("no_proxy", "127.0.0.1,localhost"), + EnvGuard::set("CCP_CONFIG_DIR", config_dir), + EnvGuard::set("CCP_ALIAS_PROVIDER", "codex"), + EnvGuard::set("CCP_CODEX_TRANSPORT", "websocket"), + EnvGuard::set("CCP_CODEX_BASE_URL", upstream_url), + EnvGuard::set("CCP_CODEX_PREVIOUS_RESPONSE_ID", "1"), + EnvGuard::set("CCP_CODEX_SERVER_COMPACTION", "0"), + ]); + guards +} + +fn write_codex_auth(config_dir: &Path) { + let auth_dir = config_dir.join("codex"); + std::fs::create_dir_all(&auth_dir).unwrap(); + std::fs::write( + auth_dir.join("auth.json"), + serde_json::to_vec(&json!({ + "access": "test-access", + "refresh": "test-refresh", + "expires": 4_102_444_800_000_i64, + "account_id": "acct_test" + })) + .unwrap(), + ) + .unwrap(); +} + +#[derive(Debug, Clone)] +struct CapturedRequest { + socket_ordinal: usize, + headers: HeaderMap, + body: Value, +} + +impl CapturedRequest { + fn marker(&self) -> &str { + self.body["input"] + .as_array() + .and_then(|input| input.last()) + .and_then(|item| item.get("content")) + .and_then(Value::as_array) + .and_then(|content| content.last()) + .and_then(|part| part.get("text")) + .and_then(Value::as_str) + .unwrap_or_else(|| panic!("request has no final text marker: {}", self.body)) + } + + fn previous_response_id(&self) -> Option<&str> { + self.body + .get("previous_response_id") + .and_then(Value::as_str) + } + + fn assert_protocol_headers(&self, expected_session: Option<&str>) { + assert_eq!( + self.headers + .get(http::header::AUTHORIZATION) + .and_then(|value| value.to_str().ok()), + Some("Bearer test-access"), + "socket {} authorization header", + self.socket_ordinal + ); + assert_eq!( + self.headers + .get("chatgpt-account-id") + .and_then(|value| value.to_str().ok()), + Some("acct_test"), + "socket {} account header", + self.socket_ordinal + ); + assert_eq!( + self.headers + .get("openai-beta") + .and_then(|value| value.to_str().ok()), + Some(WEBSOCKET_PROTOCOL_HEADER), + "socket {} websocket protocol header", + self.socket_ordinal + ); + assert_eq!( + self.headers + .get("session_id") + .and_then(|value| value.to_str().ok()), + expected_session, + "socket {} session_id header", + self.socket_ordinal + ); + assert_eq!( + self.headers + .get("x-client-request-id") + .and_then(|value| value.to_str().ok()), + expected_session, + "socket {} x-client-request-id header", + self.socket_ordinal + ); + } +} + +enum MockOutcome { + Completion { + response_id: String, + text: String, + close_after: bool, + acknowledged: oneshot::Sender<()>, + }, + RawEvent { + event: Value, + acknowledged: oneshot::Sender<()>, + }, +} + +struct PendingRequest { + captured: CapturedRequest, + outcome: oneshot::Sender, +} + +impl PendingRequest { + async fn respond(self, response_id: &str, text: &str, close_after: bool) -> CapturedRequest { + let (acknowledged, acknowledgement) = oneshot::channel(); + self.outcome + .send(MockOutcome::Completion { + response_id: response_id.to_string(), + text: text.to_string(), + close_after, + acknowledged, + }) + .unwrap_or_else(|_| { + panic!( + "upstream socket closed before responding to {}", + self.captured.marker() + ) + }); + tokio::time::timeout(REQUEST_TIMEOUT, acknowledgement) + .await + .expect("mock response acknowledgement timed out") + .expect("mock response acknowledgement sender dropped"); + self.captured + } + + async fn respond_rate_limited(self) -> CapturedRequest { + let (acknowledged, acknowledgement) = oneshot::channel(); + self.outcome + .send(MockOutcome::RawEvent { + event: json!({ + "type": "codex.rate_limits", + "rate_limits": {"allowed": false, "limit_reached": true} + }), + acknowledged, + }) + .unwrap_or_else(|_| { + panic!( + "upstream socket closed before rate-limiting {}", + self.captured.marker() + ) + }); + tokio::time::timeout(REQUEST_TIMEOUT, acknowledgement) + .await + .expect("mock rate-limit acknowledgement timed out") + .expect("mock rate-limit acknowledgement sender dropped"); + self.captured + } +} + +struct InstrumentedUpstream { + base_url: String, + requests: mpsc::UnboundedReceiver, + captures: Arc>>, + probes: ProbeLog, + shutdown: Option>, + task: JoinHandle<()>, +} + +impl InstrumentedUpstream { + async fn spawn() -> Self { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let address = listener.local_addr().unwrap(); + let (requests_tx, requests) = mpsc::unbounded_channel(); + let captures = Arc::new(Mutex::new(Vec::new())); + let probes = Arc::new(Mutex::new(Vec::new())); + let (shutdown, shutdown_rx) = oneshot::channel(); + let task_captures = captures.clone(); + let task_probes = probes.clone(); + let task = tokio::spawn(async move { + run_upstream( + listener, + requests_tx, + task_captures, + task_probes, + shutdown_rx, + ) + .await; + }); + Self { + base_url: format!("http://{address}/backend-api/codex/responses"), + requests, + captures, + probes, + shutdown: Some(shutdown), + task, + } + } + + async fn next_request( + &mut self, + expected_marker: &str, + expected_session: Option<&str>, + ) -> PendingRequest { + let pending = self.next_any_request(expected_session).await; + assert_eq!(pending.captured.marker(), expected_marker); + pending + } + + async fn next_any_request(&mut self, expected_session: Option<&str>) -> PendingRequest { + let pending = tokio::time::timeout(REQUEST_TIMEOUT, self.requests.recv()) + .await + .expect("timed out waiting for response.create") + .expect("mock upstream stopped before response.create"); + pending.captured.assert_protocol_headers(expected_session); + pending + } + + fn snapshot(&self) -> Vec { + self.captures + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()) + .clone() + } + + async fn shutdown(mut self) { + if let Some(shutdown) = self.shutdown.take() { + let _ = shutdown.send(()); + } + tokio::time::timeout(REQUEST_TIMEOUT, self.task) + .await + .expect("mock upstream shutdown timed out") + .expect("mock upstream task failed"); + + let captures = self + .captures + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()); + let mut requests_per_socket = HashMap::new(); + let mut expected_probe_sites = captures + .iter() + .filter_map(|capture| { + let prior_requests = requests_per_socket + .entry(capture.socket_ordinal) + .or_insert(0usize); + let site = + (*prior_requests > 0).then_some((capture.socket_ordinal, *prior_requests)); + *prior_requests += 1; + site + }) + .collect::>(); + drop(captures); + + let probes = self + .probes + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()); + let mut actual_probe_sites = probes + .iter() + .map(|(socket_ordinal, prior_requests, _)| (*socket_ordinal, *prior_requests)) + .collect::>(); + expected_probe_sites.sort_unstable(); + actual_probe_sites.sort_unstable(); + assert_eq!( + actual_probe_sites, expected_probe_sites, + "each successful pooled reuse must have one probe at its exact socket-local request boundary: {probes:?}" + ); + let unique = probes + .iter() + .map(|(_, _, payload)| payload.clone()) + .collect::>(); + assert_eq!( + unique.len(), + probes.len(), + "pooled validation probes must use unique payloads: {probes:?}" + ); + assert!( + probes.iter().all(|(_, _, payload)| payload.len() == 8), + "pooled validation probes must preserve their eight-byte nonce: {probes:?}" + ); + } +} + +async fn run_upstream( + listener: TcpListener, + requests: mpsc::UnboundedSender, + captures: Arc>>, + probes: ProbeLog, + mut shutdown: oneshot::Receiver<()>, +) { + let mut sockets = JoinSet::new(); + let mut next_socket_ordinal = 1usize; + loop { + tokio::select! { + _ = &mut shutdown => break, + accepted = listener.accept() => { + let Ok((stream, _)) = accepted else { + break; + }; + let socket_ordinal = next_socket_ordinal; + next_socket_ordinal += 1; + sockets.spawn(handle_socket( + stream, + socket_ordinal, + requests.clone(), + captures.clone(), + probes.clone(), + )); + } + completed = sockets.join_next(), if !sockets.is_empty() => { + if let Some(Err(error)) = completed { + panic!("mock upstream socket task failed: {error}"); + } + } + } + } + + sockets.abort_all(); + while sockets.join_next().await.is_some() {} +} + +#[allow(clippy::result_large_err)] +async fn handle_socket( + stream: TcpStream, + socket_ordinal: usize, + requests: mpsc::UnboundedSender, + captures: Arc>>, + probes: ProbeLog, +) { + let handshake_headers = Arc::new(Mutex::new(None)); + let callback_headers = handshake_headers.clone(); + let Ok(mut websocket) = + accept_hdr_async(stream, move |request: &http::Request<()>, response| { + *callback_headers + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()) = Some(request.headers().clone()); + Ok(response) + }) + .await + else { + return; + }; + let headers = handshake_headers + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()) + .take() + .expect("websocket handshake headers"); + let mut requests_seen = 0usize; + + loop { + let Some(frame) = websocket.next().await else { + return; + }; + match frame { + Ok(Message::Ping(payload)) => { + probes + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()) + .push((socket_ordinal, requests_seen, payload.clone())); + if websocket.send(Message::Pong(payload)).await.is_err() { + return; + } + } + Ok(Message::Pong(_)) | Ok(Message::Binary(_)) | Ok(Message::Frame(_)) => {} + Ok(Message::Close(_)) | Err(_) => return, + Ok(Message::Text(text)) => { + let body: Value = serde_json::from_str(&text).unwrap_or_else(|error| { + panic!("invalid response.create JSON: {error}: {text}") + }); + assert_eq!(body["type"], "response.create"); + assert!(body.get("stream").is_none(), "websocket payload: {body}"); + let captured = CapturedRequest { + socket_ordinal, + headers: headers.clone(), + body, + }; + captures + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()) + .push(captured.clone()); + requests_seen += 1; + let (outcome, response) = oneshot::channel(); + if requests.send(PendingRequest { captured, outcome }).is_err() { + return; + } + let outcome = match tokio::time::timeout(REQUEST_TIMEOUT, response).await { + Ok(Ok(outcome)) => outcome, + Ok(Err(_)) | Err(_) => return, + }; + match outcome { + MockOutcome::Completion { + response_id, + text, + close_after, + acknowledged, + } => { + if emit_completion(&mut websocket, &response_id, &text) + .await + .is_err() + { + let _ = acknowledged.send(()); + return; + } + if close_after { + let _ = websocket.send(Message::Close(None)).await; + drop(websocket); + let _ = acknowledged.send(()); + return; + } + let _ = acknowledged.send(()); + } + MockOutcome::RawEvent { + event, + acknowledged, + } => { + if websocket + .send(Message::Text(event.to_string())) + .await + .is_err() + { + let _ = acknowledged.send(()); + return; + } + let _ = acknowledged.send(()); + } + } + } + } + } +} + +async fn emit_completion( + websocket: &mut WebSocketStream, + response_id: &str, + text: &str, +) -> Result<(), tokio_tungstenite::tungstenite::Error> { + let events = [ + json!({ + "type": "response.output_item.added", + "output_index": 0, + "item": {"type": "message", "id": format!("msg-{response_id}")} + }), + json!({ + "type": "response.output_text.delta", + "output_index": 0, + "delta": text + }), + json!({ + "type": "response.output_item.done", + "output_index": 0, + "item": {"type": "message", "id": format!("msg-{response_id}")} + }), + json!({ + "type": "response.completed", + "response": { + "id": response_id, + "usage": {"input_tokens": 5, "output_tokens": 2} + } + }), + ]; + for event in events { + websocket.send(Message::Text(event.to_string())).await?; + } + Ok(()) +} + +#[derive(Debug, Clone, Default)] +struct IdentityHeaders { + values: Vec<(&'static str, String)>, + upstream_session: Option, +} + +impl IdentityHeaders { + fn main(session: &str) -> Self { + Self { + values: vec![(SESSION_HEADER, session.to_string())], + upstream_session: Some(session.to_string()), + } + } + + fn agent(session: &str, agent: &str, parent: Option<&str>) -> Self { + let mut values = vec![ + (SESSION_HEADER, session.to_string()), + (AGENT_HEADER, agent.to_string()), + ]; + if let Some(parent) = parent { + values.push((PARENT_AGENT_HEADER, parent.to_string())); + } + Self { + values, + upstream_session: Some(session.to_string()), + } + } + + fn malformed_agent(session: &str, raw_agent: &str) -> Self { + Self { + values: vec![ + (SESSION_HEADER, session.to_string()), + (AGENT_HEADER, raw_agent.to_string()), + ], + upstream_session: Some(session.to_string()), + } + } +} + +struct DrainedResponse { + status: StatusCode, + body: String, +} + +impl DrainedResponse { + fn assert_success(&self, expected_text: &str) { + assert_eq!( + self.status, + StatusCode::OK, + "downstream body: {}", + self.body + ); + assert!( + self.body.contains(expected_text), + "downstream response did not contain {expected_text:?}: {}", + self.body + ); + } +} + +struct TestHarness { + client: reqwest::Client, + proxy_url: String, + upstream: InstrumentedUpstream, + server_shutdown: Option>, + server_task: JoinHandle>, + _config_dir: TempDir, + _environment: Vec, +} + +impl TestHarness { + async fn start() -> Self { + let upstream = InstrumentedUpstream::spawn().await; + let config_dir = TempDir::new().unwrap(); + write_codex_auth(config_dir.path()); + let environment = configure_environment(config_dir.path(), &upstream.base_url); + + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let proxy_address = listener.local_addr().unwrap(); + let (server_shutdown, shutdown) = oneshot::channel(); + let server_task = tokio::spawn(server::serve_listener(listener, None, async move { + let _ = shutdown.await; + })); + let client = reqwest::Client::builder() + .http1_only() + .no_proxy() + .build() + .unwrap(); + + Self { + client, + proxy_url: format!("http://{proxy_address}/v1/messages"), + upstream, + server_shutdown: Some(server_shutdown), + server_task, + _config_dir: config_dir, + _environment: environment, + } + } + + fn start_request(&self, body: Value, identity: IdentityHeaders) -> JoinHandle { + let client = self.client.clone(); + let url = self.proxy_url.clone(); + tokio::spawn(async move { + let mut request = client.post(url).json(&body); + for (name, value) in identity.values { + request = request.header(name, value); + } + let response = tokio::time::timeout(REQUEST_TIMEOUT, request.send()) + .await + .expect("proxy response headers timed out") + .expect("proxy request failed"); + let status = response.status(); + let body = tokio::time::timeout(REQUEST_TIMEOUT, response.bytes()) + .await + .expect("proxy response body timed out") + .expect("proxy response body failed"); + DrainedResponse { + status, + body: String::from_utf8_lossy(&body).into_owned(), + } + }) + } + + async fn pending(&mut self, marker: &str, identity: &IdentityHeaders) -> PendingRequest { + self.upstream + .next_request(marker, identity.upstream_session.as_deref()) + .await + } + + async fn round_trip( + &mut self, + body: Value, + identity: IdentityHeaders, + marker: &str, + response_id: &str, + reply: &str, + ) -> CapturedRequest { + let request = self.start_request(body, identity.clone()); + let pending = self.pending(marker, &identity).await; + resolve_request(pending, request, response_id, reply, false).await + } + + async fn shutdown(mut self) { + drop(self.client); + if let Some(shutdown) = self.server_shutdown.take() { + let _ = shutdown.send(()); + } + tokio::time::timeout(REQUEST_TIMEOUT, self.server_task) + .await + .expect("proxy shutdown timed out") + .expect("proxy server task failed") + .expect("proxy server returned an error"); + self.upstream.shutdown().await; + } +} + +async fn resolve_request( + pending: PendingRequest, + request: JoinHandle, + response_id: &str, + reply: &str, + close_after: bool, +) -> CapturedRequest { + let captured = pending.respond(response_id, reply, close_after).await; + let response = tokio::time::timeout(REQUEST_TIMEOUT, request) + .await + .expect("downstream request task timed out") + .expect("downstream request task failed"); + response.assert_success(reply); + captured +} + +fn unique(label: &str) -> String { + format!("{label}-{}", Uuid::new_v4()) +} + +fn tagged(case: &str, label: &str) -> String { + format!("{case}-{label}") +} + +fn message(role: &str, text: &str) -> Value { + json!({"role": role, "content": text}) +} + +fn messages_body(stream: bool, messages: Vec) -> Value { + json!({ + "model": "gpt-5.6-sol", + "max_tokens": 64, + "stream": stream, + "messages": messages + }) +} + +fn upstream_item(role: &str, text: &str) -> Value { + let content_type = if role == "assistant" { + "output_text" + } else { + "input_text" + }; + json!({ + "type": "message", + "role": role, + "content": [{"type": content_type, "text": text}] + }) +} + +fn assert_full_input(request: &CapturedRequest, expected: &[(&str, &str)]) { + assert!( + request.body.get("previous_response_id").is_none(), + "{} must omit previous_response_id on socket {}: {}", + request.marker(), + request.socket_ordinal, + request.body + ); + let expected = expected + .iter() + .map(|(role, text)| upstream_item(role, text)) + .collect::>(); + assert_eq!( + request.body["input"].as_array(), + Some(&expected), + "{} must send complete input on socket {}", + request.marker(), + request.socket_ordinal + ); +} + +fn assert_delta_input( + request: &CapturedRequest, + previous_response_id: &str, + socket_ordinal: usize, + delta_text: &str, +) { + assert_eq!( + request.previous_response_id(), + Some(previous_response_id), + "{} previous_response_id", + request.marker() + ); + assert_eq!( + request.socket_ordinal, + socket_ordinal, + "{} originating socket", + request.marker() + ); + assert_eq!( + request.body["input"].as_array(), + Some(&vec![upstream_item("user", delta_text)]), + "{} must send exactly one appended input item", + request.marker() + ); +} + +fn pending_pair( + first: PendingRequest, + second: PendingRequest, + first_marker: &str, + second_marker: &str, +) -> (PendingRequest, PendingRequest) { + match ( + first.captured.marker() == first_marker, + second.captured.marker() == second_marker, + ) { + (true, true) => (first, second), + (false, false) + if first.captured.marker() == second_marker + && second.captured.marker() == first_marker => + { + (second, first) + } + _ => panic!( + "expected pending markers {first_marker:?}/{second_marker:?}, got {:?}/{:?}", + first.captured.marker(), + second.captured.marker() + ), + } +} + +#[allow(clippy::await_holding_lock)] +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +async fn main_child_and_nested_child_interleave_independent_continuations() { + let _environment_lock = env_lock(); + let mut harness = TestHarness::start().await; + let case = unique("lineage"); + let session = tagged(&case, "session"); + let agent_a = tagged(&case, "agent-a"); + let agent_n = tagged(&case, "agent-n"); + let main = IdentityHeaders::main(&session); + let child = IdentityHeaders::agent(&session, &agent_a, None); + let nested = IdentityHeaders::agent(&session, &agent_n, Some(&agent_a)); + + let m1 = tagged(&case, "main-1"); + let a1 = tagged(&case, "a-1"); + let n1 = tagged(&case, "n-1"); + let m2 = tagged(&case, "main-2"); + let n2 = tagged(&case, "n-2"); + let a2 = tagged(&case, "a-2"); + let rm1 = tagged(&case, "reply-main-1"); + let ra1 = tagged(&case, "reply-a-1"); + let rn1 = tagged(&case, "reply-n-1"); + let resp_m1 = tagged(&case, "resp-main-1"); + let resp_a1 = tagged(&case, "resp-a-1"); + let resp_n1 = tagged(&case, "resp-n-1"); + + let main_first = harness + .round_trip( + messages_body(false, vec![message("user", &m1)]), + main.clone(), + &m1, + &resp_m1, + &rm1, + ) + .await; + let child_first = harness + .round_trip( + messages_body(false, vec![message("user", &a1)]), + child.clone(), + &a1, + &resp_a1, + &ra1, + ) + .await; + let nested_first = harness + .round_trip( + messages_body(false, vec![message("user", &n1)]), + nested.clone(), + &n1, + &resp_n1, + &rn1, + ) + .await; + + assert_full_input(&main_first, &[("user", &m1)]); + assert_full_input(&child_first, &[("user", &a1)]); + assert_full_input(&nested_first, &[("user", &n1)]); + assert_eq!(main_first.socket_ordinal, 1); + assert_eq!(child_first.socket_ordinal, 2); + assert_eq!(nested_first.socket_ordinal, 3); + + let main_second = harness + .round_trip( + messages_body( + false, + vec![ + message("user", &m1), + message("assistant", &rm1), + message("user", &m2), + ], + ), + main, + &m2, + &tagged(&case, "resp-main-2"), + &tagged(&case, "reply-main-2"), + ) + .await; + let nested_second = harness + .round_trip( + messages_body( + false, + vec![ + message("user", &n1), + message("assistant", &rn1), + message("user", &n2), + ], + ), + nested, + &n2, + &tagged(&case, "resp-n-2"), + &tagged(&case, "reply-n-2"), + ) + .await; + let child_second = harness + .round_trip( + messages_body( + false, + vec![ + message("user", &a1), + message("assistant", &ra1), + message("user", &a2), + ], + ), + child, + &a2, + &tagged(&case, "resp-a-2"), + &tagged(&case, "reply-a-2"), + ) + .await; + + assert_delta_input(&main_second, &resp_m1, main_first.socket_ordinal, &m2); + assert_delta_input(&nested_second, &resp_n1, nested_first.socket_ordinal, &n2); + assert_delta_input(&child_second, &resp_a1, child_first.socket_ordinal, &a2); + assert_ne!(nested_second.socket_ordinal, child_first.socket_ordinal); + assert_eq!(harness.upstream.snapshot().len(), 6); + harness.shutdown().await; +} + +#[allow(clippy::await_holding_lock)] +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +async fn compatible_prefix_sibling_cannot_steal_owner_continuation() { + let _environment_lock = env_lock(); + let mut harness = TestHarness::start().await; + let case = unique("prefix-attack"); + let session = tagged(&case, "session"); + let a_headers = IdentityHeaders::agent(&session, &tagged(&case, "agent-a"), None); + let b_headers = IdentityHeaders::agent(&session, &tagged(&case, "agent-b"), None); + let a1 = tagged(&case, "a-1"); + let a_reply = tagged(&case, "a-reply"); + let attack = tagged(&case, "b-compatible-suffix"); + let a2 = tagged(&case, "a-2"); + let resp_a1 = tagged(&case, "resp-a-1"); + + let first = harness + .round_trip( + messages_body(false, vec![message("user", &a1)]), + a_headers.clone(), + &a1, + &resp_a1, + &a_reply, + ) + .await; + let sibling_attack = harness + .round_trip( + messages_body( + false, + vec![ + message("user", &a1), + message("assistant", &a_reply), + message("user", &attack), + ], + ), + b_headers, + &attack, + &tagged(&case, "resp-b-1"), + &tagged(&case, "b-reply"), + ) + .await; + let second = harness + .round_trip( + messages_body( + false, + vec![ + message("user", &a1), + message("assistant", &a_reply), + message("user", &a2), + ], + ), + a_headers, + &a2, + &tagged(&case, "resp-a-2"), + &tagged(&case, "a-reply-2"), + ) + .await; + + assert_full_input(&first, &[("user", &a1)]); + assert_full_input( + &sibling_attack, + &[("user", &a1), ("assistant", &a_reply), ("user", &attack)], + ); + assert_eq!(first.socket_ordinal, 1); + assert_eq!(sibling_attack.socket_ordinal, 2); + assert_delta_input(&second, &resp_a1, first.socket_ordinal, &a2); + assert_ne!(sibling_attack.socket_ordinal, second.socket_ordinal); + assert_eq!(harness.upstream.snapshot().len(), 3); + harness.shutdown().await; +} + +#[allow(clippy::await_holding_lock)] +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +async fn divergent_siblings_retain_their_own_response_and_socket() { + let _environment_lock = env_lock(); + let mut harness = TestHarness::start().await; + let case = unique("divergent-siblings"); + let session = tagged(&case, "session"); + let a_headers = IdentityHeaders::agent(&session, &tagged(&case, "agent-a"), None); + let b_headers = IdentityHeaders::agent(&session, &tagged(&case, "agent-b"), None); + let a1 = tagged(&case, "a-1"); + let b1 = tagged(&case, "b-1"); + let a2 = tagged(&case, "a-2"); + let b2 = tagged(&case, "b-2"); + let ar1 = tagged(&case, "a-reply-1"); + let br1 = tagged(&case, "b-reply-1"); + let a_resp = tagged(&case, "resp-a-1"); + let b_resp = tagged(&case, "resp-b-1"); + + let a_first = harness + .round_trip( + messages_body(false, vec![message("user", &a1)]), + a_headers.clone(), + &a1, + &a_resp, + &ar1, + ) + .await; + let b_first = harness + .round_trip( + messages_body(false, vec![message("user", &b1)]), + b_headers.clone(), + &b1, + &b_resp, + &br1, + ) + .await; + let a_second = harness + .round_trip( + messages_body( + false, + vec![ + message("user", &a1), + message("assistant", &ar1), + message("user", &a2), + ], + ), + a_headers, + &a2, + &tagged(&case, "resp-a-2"), + &tagged(&case, "a-reply-2"), + ) + .await; + let b_second = harness + .round_trip( + messages_body( + false, + vec![ + message("user", &b1), + message("assistant", &br1), + message("user", &b2), + ], + ), + b_headers, + &b2, + &tagged(&case, "resp-b-2"), + &tagged(&case, "b-reply-2"), + ) + .await; + + assert_full_input(&a_first, &[("user", &a1)]); + assert_full_input(&b_first, &[("user", &b1)]); + assert_eq!(a_first.socket_ordinal, 1); + assert_eq!(b_first.socket_ordinal, 2); + assert_delta_input(&a_second, &a_resp, a_first.socket_ordinal, &a2); + assert_delta_input(&b_second, &b_resp, b_first.socket_ordinal, &b2); + assert_ne!(a_second.socket_ordinal, b_second.socket_ordinal); + assert_eq!(harness.upstream.snapshot().len(), 4); + harness.shutdown().await; +} + +#[allow(clippy::await_holding_lock)] +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +async fn same_agent_id_in_different_sessions_is_independent() { + let _environment_lock = env_lock(); + let mut harness = TestHarness::start().await; + let case = unique("same-agent-different-sessions"); + let agent = tagged(&case, "shared-agent"); + let session_one = tagged(&case, "session-one"); + let session_two = tagged(&case, "session-two"); + let one_headers = IdentityHeaders::agent(&session_one, &agent, None); + let two_headers = IdentityHeaders::agent(&session_two, &agent, None); + let one1 = tagged(&case, "one-1"); + let two1 = tagged(&case, "two-1"); + let one2 = tagged(&case, "one-2"); + let two2 = tagged(&case, "two-2"); + let one_reply = tagged(&case, "one-reply"); + let two_reply = tagged(&case, "two-reply"); + let one_response = tagged(&case, "resp-one-1"); + let two_response = tagged(&case, "resp-two-1"); + + let one_first = harness + .round_trip( + messages_body(false, vec![message("user", &one1)]), + one_headers.clone(), + &one1, + &one_response, + &one_reply, + ) + .await; + let two_first = harness + .round_trip( + messages_body(false, vec![message("user", &two1)]), + two_headers.clone(), + &two1, + &two_response, + &two_reply, + ) + .await; + + assert_full_input(&one_first, &[("user", &one1)]); + assert_full_input(&two_first, &[("user", &two1)]); + let one_second = harness + .round_trip( + messages_body( + false, + vec![ + message("user", &one1), + message("assistant", &one_reply), + message("user", &one2), + ], + ), + one_headers, + &one2, + &tagged(&case, "resp-one-2"), + &tagged(&case, "one-reply-2"), + ) + .await; + let two_second = harness + .round_trip( + messages_body( + false, + vec![ + message("user", &two1), + message("assistant", &two_reply), + message("user", &two2), + ], + ), + two_headers, + &two2, + &tagged(&case, "resp-two-2"), + &tagged(&case, "two-reply-2"), + ) + .await; + + assert_eq!(one_first.socket_ordinal, 1); + assert_eq!(two_first.socket_ordinal, 2); + assert_ne!(one_first.socket_ordinal, two_first.socket_ordinal); + assert_delta_input(&one_second, &one_response, one_first.socket_ordinal, &one2); + assert_delta_input(&two_second, &two_response, two_first.socket_ordinal, &two2); + assert_eq!(harness.upstream.snapshot().len(), 4); + harness.shutdown().await; +} + +#[allow(clippy::await_holding_lock)] +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +async fn malformed_identity_sandwich_is_stateless_and_preserves_main() { + let _environment_lock = env_lock(); + let mut harness = TestHarness::start().await; + let case = unique("malformed-sandwich"); + let session = tagged(&case, "session"); + let main_headers = IdentityHeaders::main(&session); + let malformed_headers = IdentityHeaders::malformed_agent( + &session, + &format!( + "{}, {}", + tagged(&case, "raw-agent"), + tagged(&case, "forged-agent") + ), + ); + let no_identity = IdentityHeaders::default(); + let main1 = tagged(&case, "main-1"); + let main_reply = tagged(&case, "main-reply"); + let main2 = tagged(&case, "main-2"); + let main_response = tagged(&case, "resp-main-1"); + let malformed1 = tagged(&case, "malformed-1"); + let malformed_reply = tagged(&case, "malformed-reply"); + let malformed2 = tagged(&case, "malformed-2"); + let control = tagged(&case, "no-identity-control"); + + let main_first = harness + .round_trip( + messages_body(false, vec![message("user", &main1)]), + main_headers.clone(), + &main1, + &main_response, + &main_reply, + ) + .await; + let malformed_first = harness + .round_trip( + messages_body(false, vec![message("user", &malformed1)]), + malformed_headers.clone(), + &malformed1, + &tagged(&case, "resp-malformed-1"), + &malformed_reply, + ) + .await; + let malformed_second = harness + .round_trip( + messages_body( + false, + vec![ + message("user", &malformed1), + message("assistant", &malformed_reply), + message("user", &malformed2), + ], + ), + malformed_headers, + &malformed2, + &tagged(&case, "resp-malformed-2"), + &tagged(&case, "malformed-reply-2"), + ) + .await; + let control_request = harness + .round_trip( + messages_body(false, vec![message("user", &control)]), + no_identity, + &control, + &tagged(&case, "resp-control"), + &tagged(&case, "control-reply"), + ) + .await; + let main_second = harness + .round_trip( + messages_body( + false, + vec![ + message("user", &main1), + message("assistant", &main_reply), + message("user", &main2), + ], + ), + main_headers, + &main2, + &tagged(&case, "resp-main-2"), + &tagged(&case, "main-reply-2"), + ) + .await; + + assert_full_input(&main_first, &[("user", &main1)]); + assert_full_input(&malformed_first, &[("user", &malformed1)]); + assert_full_input( + &malformed_second, + &[ + ("user", &malformed1), + ("assistant", &malformed_reply), + ("user", &malformed2), + ], + ); + assert_full_input(&control_request, &[("user", &control)]); + assert_eq!(main_first.socket_ordinal, 1); + assert_eq!(malformed_first.socket_ordinal, 2); + assert_eq!(malformed_second.socket_ordinal, 3); + assert_eq!(control_request.socket_ordinal, 4); + assert_ne!( + malformed_first.socket_ordinal, malformed_second.socket_ordinal, + "the raw malformed agent value must never become a pool owner" + ); + assert_delta_input( + &main_second, + &main_response, + main_first.socket_ordinal, + &main2, + ); + assert_eq!(harness.upstream.snapshot().len(), 5); + harness.shutdown().await; +} + +#[allow(clippy::await_holding_lock)] +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +async fn auto_review_with_agent_headers_is_stateless() { + let _environment_lock = env_lock(); + let mut harness = TestHarness::start().await; + let case = unique("auto-review"); + let session = tagged(&case, "session"); + let headers = IdentityHeaders::agent(&session, &tagged(&case, "agent-a"), None); + let a1 = tagged(&case, "a-1"); + let a_reply = tagged(&case, "a-reply"); + let review = tagged(&case, "review-command"); + let review_system = "You are a security monitor for autonomous AI coding agents.\n\n## Context"; + let a2 = tagged(&case, "a-2"); + let a_response = tagged(&case, "resp-a-1"); + + let first = harness + .round_trip( + messages_body(false, vec![message("user", &a1)]), + headers.clone(), + &a1, + &a_response, + &a_reply, + ) + .await; + let classifier = harness + .round_trip( + json!({ + "model": "gpt-5.6-sol", + "max_tokens": 64, + "stream": false, + "system": [{ + "type": "text", + "text": review_system + }], + "messages": [{"role": "user", "content": review}], + "tools": [] + }), + headers.clone(), + &review, + &tagged(&case, "resp-review"), + &tagged(&case, "review-reply"), + ) + .await; + let second = harness + .round_trip( + messages_body( + false, + vec![ + message("user", &a1), + message("assistant", &a_reply), + message("user", &a2), + ], + ), + headers, + &a2, + &tagged(&case, "resp-a-2"), + &tagged(&case, "a-reply-2"), + ) + .await; + + assert_full_input(&first, &[("user", &a1)]); + assert_full_input( + &classifier, + &[("developer", review_system), ("user", &review)], + ); + assert_eq!(classifier.body["model"], "gpt-5.6-luna"); + assert_eq!(first.socket_ordinal, 1); + assert_eq!(classifier.socket_ordinal, 2); + assert_delta_input(&second, &a_response, first.socket_ordinal, &a2); + assert_ne!(classifier.socket_ordinal, second.socket_ordinal); + assert_eq!(harness.upstream.snapshot().len(), 3); + harness.shutdown().await; +} + +#[allow(clippy::await_holding_lock)] +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +async fn missing_and_dead_origin_retry_full_context_once_and_republish_for_both_modes() { + let _environment_lock = env_lock(); + let mut harness = TestHarness::start().await; + let case = unique("origin-recovery"); + + for (index, (origin, delivery, stream)) in [ + ("missing", "buffered", false), + ("missing", "streaming", true), + ("dead", "buffered", false), + ("dead", "streaming", true), + ] + .into_iter() + .enumerate() + { + let label = format!("{origin}-{delivery}"); + let session = tagged(&case, &format!("{label}-session")); + let agent = tagged(&case, &format!("{label}-agent")); + let owner = ConversationIdentity::Agent(session.clone(), agent.clone()); + let headers = IdentityHeaders::agent(&session, &agent, None); + let one = tagged(&case, &format!("{label}-1")); + let one_reply = tagged(&case, &format!("{label}-reply-1")); + let two = tagged(&case, &format!("{label}-2")); + let two_reply = tagged(&case, &format!("{label}-reply-2")); + let three = tagged(&case, &format!("{label}-3")); + let response_one = tagged(&case, &format!("resp-{label}-1")); + let response_two = tagged(&case, &format!("resp-{label}-2")); + + let first_request = harness.start_request( + messages_body(stream, vec![message("user", &one)]), + headers.clone(), + ); + let first_pending = harness.pending(&one, &headers).await; + let first = resolve_request( + first_pending, + first_request, + &response_one, + &one_reply, + origin == "dead", + ) + .await; + assert_full_input(&first, &[("user", &one)]); + assert_eq!(first.socket_ordinal, index * 2 + 1); + + if origin == "missing" { + invalidate_codex_websocket_pool_owner(&owner); + } + + let second_request = harness.start_request( + messages_body( + stream, + vec![ + message("user", &one), + message("assistant", &one_reply), + message("user", &two), + ], + ), + headers.clone(), + ); + let second_pending = harness.pending(&two, &headers).await; + let second = resolve_request( + second_pending, + second_request, + &response_two, + &two_reply, + false, + ) + .await; + assert_eq!(second.socket_ordinal, index * 2 + 2); + assert_full_input( + &second, + &[("user", &one), ("assistant", &one_reply), ("user", &two)], + ); + + let third = harness + .round_trip( + messages_body( + stream, + vec![ + message("user", &one), + message("assistant", &one_reply), + message("user", &two), + message("assistant", &two_reply), + message("user", &three), + ], + ), + headers, + &three, + &tagged(&case, &format!("resp-{label}-3")), + &tagged(&case, &format!("{label}-reply-3")), + ) + .await; + assert_delta_input(&third, &response_two, second.socket_ordinal, &three); + + let captures = harness.upstream.snapshot(); + assert_eq!( + captures + .iter() + .filter(|capture| capture.marker() == two) + .count(), + 1, + "{label}: compatible turn 2 must send full context exactly once after origin recovery" + ); + assert_eq!( + captures.len(), + (index + 1) * 3, + "{label}: bounded fallback emitted an unexpected send" + ); + } + + harness.shutdown().await; +} + +#[allow(clippy::await_holding_lock)] +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +async fn rate_limited_live_retry_republishes_successful_continuation() { + let _environment_lock = env_lock(); + let mut harness = TestHarness::start().await; + let case = unique("rate-limit-retry"); + let session = tagged(&case, "session"); + let agent = tagged(&case, "agent"); + let headers = IdentityHeaders::agent(&session, &agent, None); + let first_message = tagged(&case, "first-message"); + let first_reply = tagged(&case, "first-reply"); + let appended_message = tagged(&case, "appended-message"); + + let request = harness.start_request( + messages_body(true, vec![message("user", &first_message)]), + headers.clone(), + ); + let first_attempt = harness.pending(&first_message, &headers).await; + assert_eq!(first_attempt.captured.socket_ordinal, 1); + assert_full_input(&first_attempt.captured, &[("user", &first_message)]); + first_attempt.respond_rate_limited().await; + + let replacement = harness.pending(&first_message, &headers).await; + assert_eq!(replacement.captured.socket_ordinal, 2); + assert_full_input(&replacement.captured, &[("user", &first_message)]); + let replacement = resolve_request( + replacement, + request, + "resp-retry-success", + &first_reply, + false, + ) + .await; + + let appended = harness + .round_trip( + messages_body( + true, + vec![ + message("user", &first_message), + message("assistant", &first_reply), + message("user", &appended_message), + ], + ), + headers, + &appended_message, + &tagged(&case, "resp-appended"), + &tagged(&case, "appended-reply"), + ) + .await; + assert_delta_input( + &appended, + "resp-retry-success", + replacement.socket_ordinal, + &appended_message, + ); + assert_eq!(replacement.socket_ordinal, 2); + assert_eq!(harness.upstream.snapshot().len(), 3); + + harness.shutdown().await; +} + +#[allow(clippy::await_holding_lock)] +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +async fn consecutive_rate_limit_handoffs_are_attempt_local() { + let _environment_lock = env_lock(); + let mut harness = TestHarness::start().await; + let case = unique("consecutive-rate-limit-retry"); + let session = tagged(&case, "session"); + let agent = tagged(&case, "agent"); + let headers = IdentityHeaders::agent(&session, &agent, None); + let first_message = tagged(&case, "first-message"); + let first_reply = tagged(&case, "first-reply"); + let appended_message = tagged(&case, "appended-message"); + + let request = harness.start_request( + messages_body(true, vec![message("user", &first_message)]), + headers.clone(), + ); + for expected_socket in 1..=2 { + let attempt = harness.pending(&first_message, &headers).await; + assert_eq!(attempt.captured.socket_ordinal, expected_socket); + assert_full_input(&attempt.captured, &[("user", &first_message)]); + attempt.respond_rate_limited().await; + } + + let replacement = harness.pending(&first_message, &headers).await; + assert_eq!(replacement.captured.socket_ordinal, 3); + assert_full_input(&replacement.captured, &[("user", &first_message)]); + let replacement = resolve_request( + replacement, + request, + "resp-consecutive-retry-success", + &first_reply, + false, + ) + .await; + + let appended = harness + .round_trip( + messages_body( + true, + vec![ + message("user", &first_message), + message("assistant", &first_reply), + message("user", &appended_message), + ], + ), + headers, + &appended_message, + &tagged(&case, "resp-appended"), + &tagged(&case, "appended-reply"), + ) + .await; + assert_delta_input( + &appended, + "resp-consecutive-retry-success", + replacement.socket_ordinal, + &appended_message, + ); + assert_eq!(replacement.socket_ordinal, 3); + assert_eq!(harness.upstream.snapshot().len(), 4); + + harness.shutdown().await; +} + +#[allow(clippy::await_holding_lock)] +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +async fn same_owner_stale_completion_cannot_overwrite_newer_turn() { + let _environment_lock = env_lock(); + let mut harness = TestHarness::start().await; + let case = unique("stale-completion"); + let session = tagged(&case, "session"); + let headers = IdentityHeaders::agent(&session, &tagged(&case, "agent"), None); + let base = tagged(&case, "base"); + let base_reply = tagged(&case, "base-reply"); + let old = tagged(&case, "old-turn"); + let old_reply = tagged(&case, "old-reply"); + let newer = tagged(&case, "newer-turn"); + let newer_reply = tagged(&case, "newer-reply"); + let follow = tagged(&case, "follow"); + let base_response = tagged(&case, "resp-base"); + let newer_response = tagged(&case, "resp-newer"); + + let baseline = harness + .round_trip( + messages_body(false, vec![message("user", &base)]), + headers.clone(), + &base, + &base_response, + &base_reply, + ) + .await; + assert_full_input(&baseline, &[("user", &base)]); + assert_eq!(baseline.socket_ordinal, 1); + + let old_request = harness.start_request( + messages_body( + false, + vec![ + message("user", &base), + message("assistant", &base_reply), + message("user", &old), + ], + ), + headers.clone(), + ); + let old_pending = harness.pending(&old, &headers).await; + assert_delta_input( + &old_pending.captured, + &base_response, + baseline.socket_ordinal, + &old, + ); + + let newer_request = harness.start_request( + messages_body( + false, + vec![ + message("user", &base), + message("assistant", &base_reply), + message("user", &newer), + ], + ), + headers.clone(), + ); + let newer_pending = harness.pending(&newer, &headers).await; + assert_eq!(newer_pending.captured.socket_ordinal, 2); + assert_full_input( + &newer_pending.captured, + &[ + ("user", &base), + ("assistant", &base_reply), + ("user", &newer), + ], + ); + let newer_capture = resolve_request( + newer_pending, + newer_request, + &newer_response, + &newer_reply, + false, + ) + .await; + + let old_capture = resolve_request( + old_pending, + old_request, + &tagged(&case, "resp-old"), + &old_reply, + false, + ) + .await; + assert_eq!(old_capture.socket_ordinal, baseline.socket_ordinal); + + let following = harness + .round_trip( + messages_body( + false, + vec![ + message("user", &base), + message("assistant", &base_reply), + message("user", &newer), + message("assistant", &newer_reply), + message("user", &follow), + ], + ), + headers, + &follow, + &tagged(&case, "resp-follow"), + &tagged(&case, "follow-reply"), + ) + .await; + assert_delta_input( + &following, + &newer_response, + newer_capture.socket_ordinal, + &follow, + ); + assert_ne!(following.socket_ordinal, old_capture.socket_ordinal); + assert_eq!(harness.upstream.snapshot().len(), 4); + harness.shutdown().await; +} + +#[allow(clippy::await_holding_lock)] +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +async fn cross_owner_completion_order_is_independent() { + let _environment_lock = env_lock(); + let mut harness = TestHarness::start().await; + let case = unique("cross-owner-order"); + let session = tagged(&case, "session"); + let a_headers = IdentityHeaders::agent(&session, &tagged(&case, "agent-a"), None); + let b_headers = IdentityHeaders::agent(&session, &tagged(&case, "agent-b"), None); + let a1 = tagged(&case, "a-1"); + let b1 = tagged(&case, "b-1"); + let a1_reply = tagged(&case, "a-reply-1"); + let b1_reply = tagged(&case, "b-reply-1"); + let a1_response = tagged(&case, "resp-a-1"); + let b1_response = tagged(&case, "resp-b-1"); + + let a1_request = harness.start_request( + messages_body(false, vec![message("user", &a1)]), + a_headers.clone(), + ); + let b1_request = harness.start_request( + messages_body(false, vec![message("user", &b1)]), + b_headers.clone(), + ); + let pending_one = harness.upstream.next_any_request(Some(&session)).await; + let pending_two = harness.upstream.next_any_request(Some(&session)).await; + let (a1_pending, b1_pending) = pending_pair(pending_one, pending_two, &a1, &b1); + assert_full_input(&a1_pending.captured, &[("user", &a1)]); + assert_full_input(&b1_pending.captured, &[("user", &b1)]); + assert_eq!( + HashSet::from([ + a1_pending.captured.socket_ordinal, + b1_pending.captured.socket_ordinal, + ]), + HashSet::from([1, 2]) + ); + + let b1_capture = resolve_request(b1_pending, b1_request, &b1_response, &b1_reply, false).await; + let a1_capture = resolve_request(a1_pending, a1_request, &a1_response, &a1_reply, false).await; + + let a2 = tagged(&case, "a-2"); + let b2 = tagged(&case, "b-2"); + let a2_reply = tagged(&case, "a-reply-2"); + let b2_reply = tagged(&case, "b-reply-2"); + let a2_response = tagged(&case, "resp-a-2"); + let b2_response = tagged(&case, "resp-b-2"); + let a2_request = harness.start_request( + messages_body( + false, + vec![ + message("user", &a1), + message("assistant", &a1_reply), + message("user", &a2), + ], + ), + a_headers.clone(), + ); + let b2_request = harness.start_request( + messages_body( + false, + vec![ + message("user", &b1), + message("assistant", &b1_reply), + message("user", &b2), + ], + ), + b_headers.clone(), + ); + let pending_one = harness.upstream.next_any_request(Some(&session)).await; + let pending_two = harness.upstream.next_any_request(Some(&session)).await; + let (a2_pending, b2_pending) = pending_pair(pending_one, pending_two, &a2, &b2); + assert_delta_input( + &a2_pending.captured, + &a1_response, + a1_capture.socket_ordinal, + &a2, + ); + assert_delta_input( + &b2_pending.captured, + &b1_response, + b1_capture.socket_ordinal, + &b2, + ); + + let a2_capture = resolve_request(a2_pending, a2_request, &a2_response, &a2_reply, false).await; + let b2_capture = resolve_request(b2_pending, b2_request, &b2_response, &b2_reply, false).await; + + let a3 = tagged(&case, "a-3"); + let b3 = tagged(&case, "b-3"); + let a3_request = harness.start_request( + messages_body( + false, + vec![ + message("user", &a1), + message("assistant", &a1_reply), + message("user", &a2), + message("assistant", &a2_reply), + message("user", &a3), + ], + ), + a_headers, + ); + let b3_request = harness.start_request( + messages_body( + false, + vec![ + message("user", &b1), + message("assistant", &b1_reply), + message("user", &b2), + message("assistant", &b2_reply), + message("user", &b3), + ], + ), + b_headers, + ); + let pending_one = harness.upstream.next_any_request(Some(&session)).await; + let pending_two = harness.upstream.next_any_request(Some(&session)).await; + let (a3_pending, b3_pending) = pending_pair(pending_one, pending_two, &a3, &b3); + assert_delta_input( + &a3_pending.captured, + &a2_response, + a2_capture.socket_ordinal, + &a3, + ); + assert_delta_input( + &b3_pending.captured, + &b2_response, + b2_capture.socket_ordinal, + &b3, + ); + + let a3_capture = resolve_request( + a3_pending, + a3_request, + &tagged(&case, "resp-a-3"), + &tagged(&case, "a-reply-3"), + false, + ) + .await; + let b3_capture = resolve_request( + b3_pending, + b3_request, + &tagged(&case, "resp-b-3"), + &tagged(&case, "b-reply-3"), + false, + ) + .await; + assert_eq!(a3_capture.socket_ordinal, a1_capture.socket_ordinal); + assert_eq!(b3_capture.socket_ordinal, b1_capture.socket_ordinal); + assert_ne!(a3_capture.socket_ordinal, b3_capture.socket_ordinal); + assert_eq!(harness.upstream.snapshot().len(), 6); + harness.shutdown().await; +} From a5d4e4c355de0b4a5dcc56feae011a4f17e7ff79 Mon Sep 17 00:00:00 2001 From: Albert Zhang Date: Sun, 2 Aug 2026 13:48:04 +0800 Subject: [PATCH 4/4] docs(codex): document agent continuation isolation Co-Authored-By: Claude --- docs/src/content/docs/providers/codex.md | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/docs/src/content/docs/providers/codex.md b/docs/src/content/docs/providers/codex.md index 76082580..2069977f 100644 --- a/docs/src/content/docs/providers/codex.md +++ b/docs/src/content/docs/providers/codex.md @@ -54,7 +54,11 @@ WebSocket is the default transport. Set `CCP_CODEX_TRANSPORT=http` for HTTP SSE, WebSocket setup honors `HTTP_PROXY` for `ws://`, `HTTPS_PROXY` for the default `wss://` endpoint, `ALL_PROXY` as a fallback, and `NO_PROXY` exclusions. A normal HTTP proxy can therefore carry the default WebSocket connection with CONNECT; TUN mode is not required. Set proxy variables before starting the process and restart after changing them. For example, setting `HTTPS_PROXY` to `http://127.0.0.1:7890` sends HTTPS/WSS destinations through the HTTP proxy at port 7890; it does not require an `https://` proxy URL. -`CCP_CODEX_PREVIOUS_RESPONSE_ID=1` enables append-only WebSocket continuation. It reuses a session connection and sends `previous_response_id` only when the translated request shape and transcript extension are safe. State is in memory, keyed by Claude Code session ID. +`CCP_CODEX_PREVIOUS_RESPONSE_ID=1` enables append-only WebSocket continuation. A valid identity containing only a Claude Code session ID owns the Main continuation for that session. Each valid direct Agent ID owns an independent continuation and reusable WebSocket within the same session. Nested Agents are keyed by their direct child ID; the parent ID is validated but does not become part of the owner key. The proxy sends `previous_response_id` only when the translated request shape and transcript extension are safe, and only on the exact live WebSocket that produced that response. + +An absent, malformed, or ambiguous identity does not reject the HTTP request; that request proceeds without continuation or WebSocket reuse. If the originating socket is missing, dead, or has been replaced, the proxy retries once with the full translated input and without the stale response ID. Continuation and connection state is held only in memory and is lost when the proxy restarts. + +Detected auto-review classifier subrequests are intentionally stateless even when valid session and Agent headers are present. They neither consume nor publish continuation or WebSocket ownership. ## Server compaction