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`
2026-07-01 00:31:56 +08:00 · 2026-06-09 11:16:27 -05:00
parent 18ce671fed
commit 6a9a49b334
3 changed files with 155 additions and 11 deletions
@@ -327,6 +327,17 @@ pub(crate) struct ThreadRequestProcessor {
    pub(super) skills_watcher: Arc<SkillsWatcher>,
 }

+/// 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<Box<StoredThread>>),
+}
+
 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,
                &params,
@@ -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<String>,
        app_server_client_version: Option<String>,
-    ) -> Result<bool, JSONRPCErrorError> {
+    ) -> Result<RunningThreadResumeResult, JSONRPCErrorError> {
        let running_thread = if params.history.is_some() {
            if let Ok(existing_thread_id) = ThreadId::from_string(&params.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(
@@ -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<Config>,
+    loader_overrides: LoaderOverrides,
+) -> std::io::Result<in_process::InProcessClientHandle> {
+    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
@@ -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<StoredThread> {
        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)
    }