From 6a9a49b334e8081756934b0ce7d909234b53aac7 Mon Sep 17 00:00:00 2001 From: Zanie Blue Date: Tue, 9 Jun 2026 11:16:27 -0500 Subject: [PATCH] Avoid rereading rollout history during cold resume (#27031) ## Summary - reuse the history-bearing `StoredThread` loaded while probing for a running thread - avoid rereading and reparsing the rollout when that probe finds no active process - reload after shutting down a loaded thread because shutdown may flush newer rollout items - add a regression test that verifies cold resume performs one history-bearing store read ## Problem `thread/resume` first reads the persisted thread with history while checking whether the thread is already running. When no running process exists, cold resume currently falls through to `resume_thread_from_rollout`, which reads and parses the same history again. That duplicate work grows with rollout size and remains on the synchronous resume path even when the caller requests `excludeTurns`. ## Background The duplicate read was introduced by #24528, which fixed resume overrides for idle cached threads. To support resumes specified by rollout path, `resume_running_thread` began loading the stored thread with history so it could resolve the canonical thread ID and determine whether a cached `CodexThread` was already loaded. That history is needed when the loaded-thread path handles the request. On a cold miss, however, the function's boolean result could only report that no loaded thread handled the request. It discarded the history-bearing `StoredThread`, and the normal cold-resume path immediately loaded and parsed the same rollout again. This change preserves the idle cached-thread behavior from #24528 while allowing the cold-resume path to reuse the probe result. ## Performance I benchmarked real retained rollouts using isolated `CODEX_HOME` directories, explicit rollout paths, debug builds of the commit and its exact parent, and alternating parent/patch order. The table below uses `thread/resume` with `excludeTurns: true`; response payload sizes were identical. | Rollout size | Records | Parent median | Patch median | Median paired saving | | ---: | ---: | ---: | ---: | ---: | | 6 MB | 3,574 | 541 ms | 441 ms | 132 ms | | 30 MB | 15,220 | 1.505 s | 1.041 s | 701 ms | | 60 MB | 31,453 | 2.644 s | 1.742 s | 970 ms | | 149 MB | 100,874 | 10.506 s | 7.156 s | 3.350 s | | 559 MB | 259,734 | 27.759 s | 16.725 s | 9.836 s | The absolute saving increases with thread size, as expected when removing one complete JSONL history read and parse. Total resume time is also content-dependent, so the relationship is not perfectly linear. I also tested full-history resume with `excludeTurns: false`. The response payload was byte-identical between variants, and the same size-dependent improvement remained visible: | Rollout size | Parent median | Patch median | Median paired saving | | ---: | ---: | ---: | ---: | | 6 MB | 1.052 s | 904 ms | 270 ms | | 30 MB | 2.667 s | 1.762 s | 924 ms | | 60 MB | 8.464 s | 6.272 s | 3.680 s | | 149 MB | 26.719 s | 12.118 s | 14.601 s | | 559 MB | 40.359 s | 25.475 s | 16.590 s | ## Validation - `just test -p codex-app-server cold_thread_resume_reuses_non_local_history_probe` - `just fix -p codex-app-server -p codex-thread-store` - `just fmt` --- .../request_processors/thread_processor.rs | 44 +++++-- .../tests/suite/v2/remote_thread_store.rs | 118 ++++++++++++++++++ codex-rs/thread-store/src/in_memory.rs | 4 + 3 files changed, 155 insertions(+), 11 deletions(-) diff --git a/codex-rs/app-server/src/request_processors/thread_processor.rs b/codex-rs/app-server/src/request_processors/thread_processor.rs index 4cc169f80..ec3bd57b0 100644 --- a/codex-rs/app-server/src/request_processors/thread_processor.rs +++ b/codex-rs/app-server/src/request_processors/thread_processor.rs @@ -327,6 +327,17 @@ pub(crate) struct ThreadRequestProcessor { pub(super) skills_watcher: Arc, } +/// Outcome of trying to satisfy a resume request from an already loaded thread. +enum RunningThreadResumeResult { + /// The request was delegated to the loaded thread. + Handled, + /// No loaded thread handled the request. + /// + /// The optional stored thread contains the history-bearing probe that cold + /// resume can reuse instead of reading the rollout again. + NotRunning(Option>), +} + impl ThreadRequestProcessor { #[allow(clippy::too_many_arguments)] pub(crate) fn new( @@ -2419,7 +2430,7 @@ impl ThreadRequestProcessor { return Ok(()); } }; - match self + let stored_thread_from_running_probe = match self .resume_running_thread( &request_id, ¶ms, @@ -2428,13 +2439,13 @@ impl ThreadRequestProcessor { ) .await { - Ok(true) => return Ok(()), - Ok(false) => {} + Ok(RunningThreadResumeResult::Handled) => return Ok(()), + Ok(RunningThreadResumeResult::NotRunning(stored_thread)) => stored_thread, Err(error) => { self.outgoing.send_error(request_id, error).await; return Ok(()); } - } + }; let ThreadResumeParams { thread_id, @@ -2458,15 +2469,20 @@ impl ThreadRequestProcessor { } = params; let include_turns = !exclude_turns; - let (thread_history, resume_source_thread) = match if let Some(history) = history { + let resume_result = if let Some(history) = history { self.resume_thread_from_history(history.as_slice()) .await .map(|thread_history| (thread_history, None)) + } else if let Some(stored_thread) = stored_thread_from_running_probe { + self.stored_thread_to_initial_history(&stored_thread) + .await + .map(|thread_history| (thread_history, Some(*stored_thread))) } else { self.resume_thread_from_rollout(&thread_id, path.as_ref()) .await .map(|(thread_history, stored_thread)| (thread_history, Some(stored_thread))) - } { + }; + let (thread_history, resume_source_thread) = match resume_result { Ok(value) => value, Err(error) => { self.outgoing.send_error(request_id, error).await; @@ -2707,7 +2723,7 @@ impl ThreadRequestProcessor { params: &ThreadResumeParams, app_server_client_name: Option, app_server_client_version: Option, - ) -> Result { + ) -> Result { let running_thread = if params.history.is_some() { if let Ok(existing_thread_id) = ThreadId::from_string(¶ms.thread_id) && self @@ -2743,7 +2759,11 @@ impl ThreadRequestProcessor { let existing_thread_id = source_thread.thread_id; match self.thread_manager.get_thread(existing_thread_id).await { Ok(existing_thread) => Some((existing_thread_id, existing_thread, source_thread)), - Err(_) => None, + Err(_) => { + return Ok(RunningThreadResumeResult::NotRunning(Some(Box::new( + source_thread, + )))); + } } }; @@ -2784,7 +2804,9 @@ impl ThreadRequestProcessor { ThreadShutdownResult::Complete => { self.thread_manager.remove_thread(&existing_thread_id).await; self.finalize_thread_teardown(existing_thread_id).await; - return Ok(false); + // Shutdown can flush newer rollout items, so reload the + // stored thread before starting the replacement session. + return Ok(RunningThreadResumeResult::NotRunning(None)); } ThreadShutdownResult::SubmitFailed => { warn!("failed to submit Shutdown to thread {existing_thread_id}"); @@ -2876,9 +2898,9 @@ impl ThreadRequestProcessor { "failed to enqueue running thread resume for thread {existing_thread_id}: thread listener command channel is closed" ))); } - return Ok(true); + return Ok(RunningThreadResumeResult::Handled); } - Ok(false) + Ok(RunningThreadResumeResult::NotRunning(None)) } async fn resume_thread_from_history( diff --git a/codex-rs/app-server/tests/suite/v2/remote_thread_store.rs b/codex-rs/app-server/tests/suite/v2/remote_thread_store.rs index 016e4c3df..ca2121dd4 100644 --- a/codex-rs/app-server/tests/suite/v2/remote_thread_store.rs +++ b/codex-rs/app-server/tests/suite/v2/remote_thread_store.rs @@ -29,6 +29,7 @@ use codex_app_server_protocol::RequestId; use codex_app_server_protocol::ServerNotification; use codex_app_server_protocol::ThreadListParams; use codex_app_server_protocol::ThreadListResponse; +use codex_app_server_protocol::ThreadResumeParams; use codex_app_server_protocol::ThreadStartParams; use codex_app_server_protocol::ThreadStartResponse; use codex_app_server_protocol::TurnStartParams; @@ -37,6 +38,7 @@ use codex_arg0::Arg0DispatchPaths; use codex_config::CloudConfigBundleLoader; use codex_config::LoaderOverrides; use codex_config::NoopThreadConfigLoader; +use codex_core::config::Config; use codex_core::config::ConfigBuilder; use codex_exec_server::EnvironmentManager; use codex_feedback::CodexFeedback; @@ -183,6 +185,122 @@ async fn thread_start_with_non_local_thread_store_does_not_create_local_persiste Ok(()) } +#[tokio::test] +async fn cold_thread_resume_reuses_non_local_history_probe() -> Result<()> { + let server = create_mock_responses_server_repeating_assistant("Done").await; + let codex_home = TempDir::new()?; + let store_id = Uuid::new_v4().to_string(); + create_config_toml_with_thread_store(codex_home.path(), &server.uri(), &store_id)?; + + let loader_overrides = LoaderOverrides::without_managed_config_for_tests(); + let config = Arc::new( + ConfigBuilder::default() + .codex_home(codex_home.path().to_path_buf()) + .fallback_cwd(Some(codex_home.path().to_path_buf())) + .loader_overrides(loader_overrides.clone()) + .build() + .await?, + ); + let thread_store = InMemoryThreadStore::for_id(store_id.clone()); + let _in_memory_store = InMemoryThreadStoreId { store_id }; + + let mut client = start_in_process_client(config.clone(), loader_overrides.clone()).await?; + let response = client + .request(ClientRequest::ThreadStart { + request_id: RequestId::Integer(1), + params: ThreadStartParams::default(), + }) + .await? + .expect("thread/start should succeed"); + let ThreadStartResponse { thread, .. } = serde_json::from_value(response)?; + + client + .request(ClientRequest::TurnStart { + request_id: RequestId::Integer(2), + params: TurnStartParams { + thread_id: thread.id.clone(), + client_user_message_id: None, + input: vec![V2UserInput::Text { + text: "Materialize the thread".to_string(), + text_elements: Vec::new(), + }], + ..Default::default() + }, + }) + .await? + .expect("turn/start should succeed"); + timeout(DEFAULT_READ_TIMEOUT, async { + loop { + let Some(event) = client.next_event().await else { + anyhow::bail!("in-process app-server stopped before turn/completed"); + }; + if let InProcessServerEvent::ServerNotification(ServerNotification::TurnCompleted( + completed, + )) = event + && completed.thread_id == thread.id + { + return Ok::<(), anyhow::Error>(()); + } + } + }) + .await??; + client.shutdown().await?; + + let client = start_in_process_client(config, loader_overrides).await?; + let reads_before_resume = thread_store.calls().await.read_thread_with_history; + // The in-memory store is pathless, so resume currently fails later while + // assembling the response. The history-bearing probe must still be reused. + let _resume_result = client + .request(ClientRequest::ThreadResume { + request_id: RequestId::Integer(3), + params: ThreadResumeParams { + thread_id: thread.id.clone(), + ..Default::default() + }, + }) + .await?; + + assert_eq!( + thread_store.calls().await.read_thread_with_history, + reads_before_resume + 1 + ); + + client.shutdown().await?; + Ok(()) +} + +async fn start_in_process_client( + config: Arc, + loader_overrides: LoaderOverrides, +) -> std::io::Result { + in_process::start(InProcessStartArgs { + arg0_paths: Arg0DispatchPaths::default(), + config, + cli_overrides: Vec::new(), + loader_overrides, + strict_config: false, + cloud_config_bundle: CloudConfigBundleLoader::default(), + thread_config_loader: Arc::new(NoopThreadConfigLoader), + feedback: CodexFeedback::new(), + log_db: None, + state_db: None, + environment_manager: Arc::new(EnvironmentManager::default_for_tests()), + config_warnings: Vec::new(), + session_source: SessionSource::Cli, + enable_codex_api_key_env: false, + initialize: InitializeParams { + client_info: ClientInfo { + name: "codex-app-server-tests".to_string(), + title: None, + version: "0.1.0".to_string(), + }, + capabilities: None, + }, + channel_capacity: in_process::DEFAULT_IN_PROCESS_CHANNEL_CAPACITY, + }) + .await +} + fn assert_no_local_persistence_artifacts(codex_home: &Path) -> Result<()> { // These are the observable tripwires for accidental local persistence. If a // future code path constructs a local rollout/session store or opens the diff --git a/codex-rs/thread-store/src/in_memory.rs b/codex-rs/thread-store/src/in_memory.rs index edc308a90..beed1d18f 100644 --- a/codex-rs/thread-store/src/in_memory.rs +++ b/codex-rs/thread-store/src/in_memory.rs @@ -109,6 +109,7 @@ pub struct InMemoryThreadStoreCalls { pub discard_thread: usize, pub load_history: usize, pub read_thread: usize, + pub read_thread_with_history: usize, pub read_thread_by_rollout_path: usize, pub list_threads: usize, pub update_thread_metadata: usize, @@ -262,6 +263,9 @@ impl ThreadStore for InMemoryThreadStore { async fn read_thread(&self, params: ReadThreadParams) -> ThreadStoreResult { let mut state = self.state.lock().await; state.calls.read_thread += 1; + if params.include_history { + state.calls.read_thread_with_history += 1; + } stored_thread_from_state(&state, params.thread_id, params.include_history) }