diff --git a/codex-rs/core/src/client.rs b/codex-rs/core/src/client.rs index 1f050df8de..a82e59f4d4 100644 --- a/codex-rs/core/src/client.rs +++ b/codex-rs/core/src/client.rs @@ -172,8 +172,6 @@ pub(crate) struct CompactConversationRequestSettings { struct ModelClientState { session_id: SessionId, thread_id: ThreadId, - // Monotonic context epoch used to avoid reusing remotely cached state after history rewrites. - // Forks may inherit a nonzero epoch when their initial history already contains rewrites. window_generation: AtomicU64, installation_id: String, provider: SharedModelProvider, @@ -187,7 +185,7 @@ struct ModelClientState { include_attestation: bool, attestation_provider: Option>, disable_websockets: AtomicBool, - cached_websocket_session: StdMutex, + cached_websocket_session: StdMutex, } /// Resolved API client setup for a single request attempt. @@ -244,7 +242,7 @@ pub struct ModelClient { pub struct ModelClientSession { client: ModelClient, websocket_session: WebsocketSession, - window_generation: u64, + websocket_cache_generation: u64, /// Turn state for sticky routing. /// /// This is an `OnceLock` that stores the turn state value received from the server @@ -273,6 +271,12 @@ struct WebsocketSession { connection_reused: StdMutex, } +#[derive(Debug, Default)] +struct CachedWebsocketSession { + generation: u64, + session: WebsocketSession, +} + impl WebsocketSession { fn set_connection_reused(&self, connection_reused: bool) { *self @@ -361,7 +365,7 @@ impl ModelClient { include_attestation, attestation_provider, disable_websockets: AtomicBool::new(false), - cached_websocket_session: StdMutex::new(WebsocketSession::default()), + cached_websocket_session: StdMutex::new(CachedWebsocketSession::default()), }), prompt_cache_key_override: None, } @@ -386,10 +390,11 @@ impl ModelClient { /// This constructor does not perform network I/O itself; the session opens a websocket lazily /// when the first stream request is issued. pub fn new_session(&self) -> ModelClientSession { + let (websocket_cache_generation, websocket_session) = self.take_cached_websocket_session(); ModelClientSession { client: self.clone(), - websocket_session: self.take_cached_websocket_session(), - window_generation: self.state.window_generation.load(Ordering::Relaxed), + websocket_session, + websocket_cache_generation, turn_state: Arc::new(OnceLock::new()), } } @@ -402,12 +407,12 @@ impl ModelClient { self.state .window_generation .store(window_generation, Ordering::Relaxed); - self.store_cached_websocket_session(WebsocketSession::default()); + self.invalidate_cached_websocket_session(); } pub(crate) fn advance_window_generation(&self) { self.state.window_generation.fetch_add(1, Ordering::Relaxed); - self.store_cached_websocket_session(WebsocketSession::default()); + self.invalidate_cached_websocket_session(); } pub(crate) fn current_window_id(&self) -> String { @@ -416,21 +421,36 @@ impl ModelClient { format!("{thread_id}:{window_generation}") } - fn take_cached_websocket_session(&self) -> WebsocketSession { + pub(crate) fn invalidate_cached_websocket_session(&self) { let mut cached_websocket_session = self .state .cached_websocket_session .lock() .unwrap_or_else(std::sync::PoisonError::into_inner); - std::mem::take(&mut *cached_websocket_session) + cached_websocket_session.generation = cached_websocket_session.generation.saturating_add(1); + cached_websocket_session.session = WebsocketSession::default(); } - fn store_cached_websocket_session(&self, websocket_session: WebsocketSession) { - *self + fn take_cached_websocket_session(&self) -> (u64, WebsocketSession) { + let mut cached_websocket_session = self .state .cached_websocket_session .lock() - .unwrap_or_else(std::sync::PoisonError::into_inner) = websocket_session; + .unwrap_or_else(std::sync::PoisonError::into_inner); + let generation = cached_websocket_session.generation; + let session = std::mem::take(&mut cached_websocket_session.session); + (generation, session) + } + + fn store_cached_websocket_session(&self, generation: u64, websocket_session: WebsocketSession) { + let mut cached_websocket_session = self + .state + .cached_websocket_session + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + if cached_websocket_session.generation == generation { + cached_websocket_session.session = websocket_session; + } } pub(crate) fn force_http_fallback( @@ -450,7 +470,7 @@ impl ModelClient { ); } - self.store_cached_websocket_session(WebsocketSession::default()); + self.invalidate_cached_websocket_session(); activated } @@ -979,13 +999,9 @@ impl ModelClient { impl Drop for ModelClientSession { fn drop(&mut self) { - // A history invalidation may advance the generation before an older turn session drops. - if self.window_generation != self.client.state.window_generation.load(Ordering::Relaxed) { - return; - } let websocket_session = std::mem::take(&mut self.websocket_session); self.client - .store_cached_websocket_session(websocket_session); + .store_cached_websocket_session(self.websocket_cache_generation, websocket_session); } } diff --git a/codex-rs/core/src/session/handlers.rs b/codex-rs/core/src/session/handlers.rs index 2dc7ac3916..f3daf58671 100644 --- a/codex-rs/core/src/session/handlers.rs +++ b/codex-rs/core/src/session/handlers.rs @@ -568,7 +568,9 @@ pub async fn thread_rollback(sess: &Arc, sub_id: String, num_turns: u32 sess.apply_rollout_reconstruction(turn_context.as_ref(), replay_items.as_slice()) .await; sess.recompute_token_usage(turn_context.as_ref()).await; - sess.services.model_client.advance_window_generation(); + sess.services + .model_client + .invalidate_cached_websocket_session(); sess.persist_rollout_items(&[RolloutItem::EventMsg(rollback_msg.clone())]) .await; diff --git a/codex-rs/core/src/session/session.rs b/codex-rs/core/src/session/session.rs index a920636201..37a501fb38 100644 --- a/codex-rs/core/src/session/session.rs +++ b/codex-rs/core/src/session/session.rs @@ -16,25 +16,6 @@ use codex_protocol::protocol::TurnEnvironmentSelections; use std::sync::OnceLock; use tokio::sync::Semaphore; -// This is a request-cache invalidation epoch, not a hash of logical context. A rollback must use a -// new generation even when it reconstructs a previously seen prefix, because the remote window -// with the old ID may already include requests from the rolled-back suffix. -fn window_generation_from_history_invalidations(items: &[RolloutItem]) -> u64 { - u64::try_from( - items - .iter() - .filter(|item| { - matches!( - item, - RolloutItem::Compacted(_) - | RolloutItem::EventMsg(EventMsg::ThreadRolledBack(_)) - ) - }) - .count(), - ) - .unwrap_or(u64::MAX) -} - /// Context for an initialized model agent /// /// A session has at most 1 running task at a time, and can be interrupted by user input. @@ -537,13 +518,15 @@ impl Session { InitialHistory::Resumed(resumed_history) => resumed_history.conversation_id, }; let window_generation = match &initial_history { - InitialHistory::Resumed(resumed_history) => { - window_generation_from_history_invalidations(&resumed_history.history) - } - InitialHistory::Forked(history) => { - window_generation_from_history_invalidations(history) - } - InitialHistory::New | InitialHistory::Cleared => 0, + InitialHistory::Resumed(resumed_history) => u64::try_from( + resumed_history + .history + .iter() + .filter(|item| matches!(item, RolloutItem::Compacted(_))) + .count(), + ) + .unwrap_or(u64::MAX), + InitialHistory::New | InitialHistory::Cleared | InitialHistory::Forked(_) => 0, }; // Kick off independent async setup tasks in parallel to reduce startup latency. // diff --git a/codex-rs/core/tests/suite/window_headers.rs b/codex-rs/core/tests/suite/window_headers.rs index bc89d7f717..51bb132799 100644 --- a/codex-rs/core/tests/suite/window_headers.rs +++ b/codex-rs/core/tests/suite/window_headers.rs @@ -21,8 +21,7 @@ use pretty_assertions::assert_eq; use std::sync::Arc; #[tokio::test(flavor = "multi_thread", worker_threads = 2)] -async fn window_id_advances_after_compact_persists_on_resume_and_resets_on_empty_fork() -> Result<()> -{ +async fn window_id_advances_after_compact_persists_on_resume_and_resets_on_fork() -> Result<()> { skip_if_no_network!(Ok(())); let server = start_mock_server().await; @@ -103,7 +102,7 @@ async fn window_id_advances_after_compact_persists_on_resume_and_resets_on_empty } #[tokio::test(flavor = "multi_thread", worker_threads = 2)] -async fn window_id_advances_after_rollback_and_persists_on_resume() -> Result<()> { +async fn window_id_stays_stable_after_rollback_and_resume() -> Result<()> { skip_if_no_network!(Ok(())); let server = start_mock_server().await; @@ -157,91 +156,8 @@ async fn window_id_advances_after_rollback_and_persists_on_resume() -> Result<() vec![ (initial_thread_id.clone(), 0), (initial_thread_id.clone(), 0), - (initial_thread_id.clone(), 1), - (initial_thread_id, 1), - ] - ); - - Ok(()) -} - -#[tokio::test(flavor = "multi_thread", worker_threads = 2)] -async fn forked_compacted_history_keeps_window_generation_on_resume() -> Result<()> { - skip_if_no_network!(Ok(())); - - let server = start_mock_server().await; - let request_log = mount_sse_sequence( - &server, - vec![ - sse(vec![ - ev_assistant_message("msg-1", "first reply"), - ev_completed("resp-1"), - ]), - sse(vec![ - ev_assistant_message("msg-2", "summary"), - ev_completed("resp-2"), - ]), - sse(vec![ev_completed("resp-3")]), - sse(vec![ev_completed("resp-4")]), - ], - ) - .await; - - let mut builder = test_codex().with_config(|config| { - config.model_provider.name = "Non-OpenAI Model provider".to_string(); - config.compact_prompt = Some(SUMMARIZATION_PROMPT.to_string()); - }); - let initial = builder.build(&server).await?; - let initial_thread = Arc::clone(&initial.codex); - let rollout_path = initial - .session_configured - .rollout_path - .clone() - .expect("rollout path"); - - submit_user_turn(&initial_thread, "before compact").await?; - submit_compact_turn(&initial_thread).await?; - shutdown_thread(&initial_thread).await?; - - let forked = initial - .thread_manager - .fork_thread( - /*snapshot*/ usize::MAX, - initial.config.clone(), - rollout_path, - /*thread_source*/ None, - /*persist_extended_history*/ false, - /*parent_trace*/ None, - ) - .await?; - let fork_rollout_path = forked - .session_configured - .rollout_path - .clone() - .expect("fork rollout path"); - submit_user_turn(&forked.thread, "after fork").await?; - shutdown_thread(&forked.thread).await?; - - let resumed = builder - .resume(&server, initial.home.clone(), fork_rollout_path) - .await?; - submit_user_turn(&resumed.codex, "after fork resume").await?; - shutdown_thread(&resumed.codex).await?; - - let requests = request_log.requests(); - assert_eq!(requests.len(), 4, "expected four model requests"); - - let window_ids = requests.iter().map(window_id_parts).collect::>(); - let initial_thread_id = window_ids[0].0.clone(); - let forked_thread_id = window_ids[2].0.clone(); - assert_ne!(forked_thread_id, initial_thread_id); - assert_eq!( - window_ids, - vec![ (initial_thread_id.clone(), 0), (initial_thread_id, 0), - (forked_thread_id.clone(), 1), - (forked_thread_id, 1), ] );