From ffda54f5b8d5595eb0963594de32e206f98cc77f Mon Sep 17 00:00:00 2001 From: Charlie Marsh Date: Thu, 22 Oct 2026 17:06:28 +0000 Subject: [PATCH] Reuse current model context across paginated resume config loading (#56224) ## Why Resuming a paginated thread rechecks stored history after loading configuration. Unchanged model context can be reused instead of loading it again. ## What changed Retain the prepared model context across configuration loading and reuse it when the thread ID, rollout path, and history revision still match. Add `ThreadStore::is_model_context_current` to validate revisions, with a fresh read when validation fails or is unavailable. ## Testing Add coverage for revision invalidation after history is appended or removed. Extend resume tests to cover paginated history, concurrent metadata and lifecycle changes, workspace-root changes, and preservation of saved user messages in model requests. GitOrigin-RevId: bad2ca0caddb63523264ac530ffb801db3961f2b --- .../request_processors/thread_processor.rs | 50 +++++++++-- .../app-server/tests/suite/v2/mcp_resource.rs | 90 +++++++++++++++---- codex-rs/thread-store/src/local/mod.rs | 19 ++++ .../src/local/model_context_tests.rs | 56 ++++++++++++ codex-rs/thread-store/src/store.rs | 13 +++ 5 files changed, 207 insertions(+), 21 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 4d104e2d5e..925feda55c 100644 --- a/codex-rs/app-server/src/request_processors/thread_processor.rs +++ b/codex-rs/app-server/src/request_processors/thread_processor.rs @@ -110,6 +110,7 @@ struct ResumeConfigState { struct PreparedResumeConfig { state: ResumeConfigState, config: Config, + model_context: Option, } struct ThreadRevertRuntimeSnapshot { @@ -2427,7 +2428,10 @@ impl ThreadRequestProcessor { ) .await?; let (thread_history, resume_source_thread) = self - .load_resume_initial_history_from_stored_thread(stored_thread) + .load_resume_initial_history_from_stored_thread( + stored_thread, + /*prepared_model_context*/ None, + ) .await?; let response_history = thread_history.clone(); let NewThread { @@ -3986,14 +3990,20 @@ impl ThreadRequestProcessor { } = params.clone(); let include_turns = !exclude_turns; + let prepared_model_context = prepared_config + .as_mut() + .and_then(|prepared| prepared.model_context.take()); 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.load_resume_initial_history_from_stored_thread(*stored_thread) - .await - .map(|(thread_history, stored_thread)| (thread_history, Some(stored_thread))) + self.load_resume_initial_history_from_stored_thread( + *stored_thread, + prepared_model_context, + ) + .await + .map(|(thread_history, stored_thread)| (thread_history, Some(stored_thread))) } else { match self .read_stored_thread_for_resume( @@ -4004,7 +4014,10 @@ impl ThreadRequestProcessor { .await { Ok(stored_thread) => self - .load_resume_initial_history_from_stored_thread(stored_thread) + .load_resume_initial_history_from_stored_thread( + stored_thread, + prepared_model_context, + ) .await .map(|(thread_history, stored_thread)| (thread_history, Some(stored_thread))), Err(error) => Err(error), @@ -4261,6 +4274,13 @@ impl ThreadRequestProcessor { *prepared_config = Some(PreparedResumeConfig { state: config_state, config, + model_context: if paginated_resume + && let InitialHistory::Resumed(history) = thread_history + { + Some(history) + } else { + None + }, }); // Recheck the thread and history after a possible archive, delete, or aliased resume. return Ok(ControlFlow::Continue(())); @@ -4893,8 +4913,27 @@ impl ThreadRequestProcessor { async fn load_resume_initial_history_from_stored_thread( &self, stored_thread: StoredThread, + prepared_model_context: Option, ) -> Result<(InitialHistory, StoredThread), JSONRPCErrorError> { if matches!(stored_thread.history_mode, ThreadHistoryMode::Paginated) { + if let Some(context) = prepared_model_context + && context.conversation_id == stored_thread.thread_id + && context.rollout_path == stored_thread.rollout_path + && let Some(revision) = context.history_revision.as_ref() + && self + .thread_store + .is_model_context_current( + StoreLoadThreadHistoryParams { + thread_id: stored_thread.thread_id, + include_archived: true, + }, + revision.clone(), + ) + .await + .unwrap_or(false) + { + return Ok((InitialHistory::Resumed(context), stored_thread)); + } let model_context = self .thread_store .load_latest_model_context(StoreLoadThreadHistoryParams { @@ -4912,6 +4951,7 @@ impl ThreadRequestProcessor { return Ok((history, stored_thread)); } + drop(prepared_model_context); let thread_id = stored_thread.thread_id.to_string(); let rollout_path = stored_thread.rollout_path.clone(); let mut stored_thread = self diff --git a/codex-rs/app-server/tests/suite/v2/mcp_resource.rs b/codex-rs/app-server/tests/suite/v2/mcp_resource.rs index cb5f8760dc..aa4eb5b3bb 100644 --- a/codex-rs/app-server/tests/suite/v2/mcp_resource.rs +++ b/codex-rs/app-server/tests/suite/v2/mcp_resource.rs @@ -13,6 +13,7 @@ use anyhow::Result; use app_test_support::ChatGptAuthFixture; use app_test_support::MockResponsesConfig; use app_test_support::TestAppServer; +use app_test_support::create_fake_paginated_rollout; use app_test_support::create_fake_rollout; use app_test_support::create_mock_responses_server_repeating_assistant; use app_test_support::rollout_path; @@ -78,6 +79,7 @@ use codex_protocol::models::PermissionProfile; use codex_protocol::models::ResponseItem; use codex_protocol::protocol::EventMsg; use codex_protocol::protocol::SessionSource; +use codex_protocol::protocol::ThreadHistoryMode; use codex_protocol::protocol::ThreadSettingsAppliedEvent; use codex_rollout::RolloutItem; use codex_rollout::RolloutRecorder; @@ -765,6 +767,7 @@ async fn metadata_and_mcp_requests_complete_while_unrelated_resume_loads_config( #[derive(Clone, Copy)] enum ResumeMutation { + Unchanged, GitMetadata, ProjectDeletion, Archive, @@ -772,20 +775,40 @@ enum ResumeMutation { DuplicateResume, } -#[test_case(ResumeMutation::GitMetadata; "git_metadata_and_project_assignment")] -#[test_case(ResumeMutation::ProjectDeletion; "project_deletion")] -#[test_case(ResumeMutation::Archive; "archive")] -#[test_case(ResumeMutation::Delete; "delete")] -#[test_case(ResumeMutation::DuplicateResume; "duplicate_resume")] +#[test_case(ResumeMutation::GitMetadata, ThreadHistoryMode::Legacy; "git_metadata_and_project_assignment")] +#[test_case(ResumeMutation::ProjectDeletion, ThreadHistoryMode::Legacy; "project_deletion")] +#[test_case(ResumeMutation::Archive, ThreadHistoryMode::Legacy; "archive")] +#[test_case(ResumeMutation::Delete, ThreadHistoryMode::Legacy; "delete")] +#[test_case(ResumeMutation::DuplicateResume, ThreadHistoryMode::Legacy; "duplicate_resume")] +#[test_case(ResumeMutation::Unchanged, ThreadHistoryMode::Paginated; "unchanged_paginated_history")] +#[test_case(ResumeMutation::GitMetadata, ThreadHistoryMode::Paginated; "paginated_git_metadata")] +#[test_case(ResumeMutation::ProjectDeletion, ThreadHistoryMode::Paginated; "paginated_project_deletion")] +#[test_case(ResumeMutation::Archive, ThreadHistoryMode::Paginated; "paginated_archive")] +#[test_case(ResumeMutation::Delete, ThreadHistoryMode::Paginated; "paginated_delete")] +#[test_case(ResumeMutation::DuplicateResume, ThreadHistoryMode::Paginated; "paginated_duplicate_resume")] #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn resume_revalidates_persisted_thread_after_config_load( mutation: ResumeMutation, + history_mode: ThreadHistoryMode, ) -> Result<()> { - let responses_server = create_mock_responses_server_repeating_assistant("Ready").await; + let responses_server = responses::start_mock_server().await; + let model_requests = responses::mount_sse_once( + &responses_server, + responses::sse(vec![ + responses::ev_response_created("resp-1"), + responses::ev_assistant_message("msg-1", "Ready"), + responses::ev_completed("resp-1"), + ]), + ) + .await; let codex_home = TempDir::new()?; MockResponsesConfig::new(&responses_server.uri()).write(codex_home.path())?; let filename_timestamp = "2025-01-05T12-00-00"; - let thread_id = create_fake_rollout( + let create_rollout = match history_mode { + ThreadHistoryMode::Legacy => create_fake_rollout, + ThreadHistoryMode::Paginated => create_fake_paginated_rollout, + }; + let thread_id = create_rollout( codex_home.path(), filename_timestamp, "2025-01-05T12:00:00Z", @@ -854,11 +877,19 @@ async fn resume_revalidates_persisted_thread_after_config_load( thread_id: "resume-by-path".to_string(), path: Some(path.canonicalize()?), cwd: Some(resume_cwd.path().to_string_lossy().into_owned()), + exclude_turns: matches!(history_mode, ThreadHistoryMode::Paginated), ..Default::default() }, })); let mut early_resume = poll!(&mut resume); let mutation_request = match mutation { + ResumeMutation::Unchanged => ClientRequest::ThreadRead { + request_id: RequestId::Integer(2), + params: ThreadReadParams { + thread_id: thread_id.clone(), + include_turns: false, + }, + }, ResumeMutation::GitMetadata => ClientRequest::ThreadMetadataUpdate { request_id: RequestId::Integer(2), params: ThreadMetadataUpdateParams { @@ -895,6 +926,7 @@ async fn resume_revalidates_persisted_thread_after_config_load( params: ThreadResumeParams { thread_id: thread_id.clone(), cwd: Some(resume_cwd.path().to_string_lossy().into_owned()), + exclude_turns: matches!(history_mode, ThreadHistoryMode::Paginated), ..Default::default() }, }, @@ -920,7 +952,9 @@ async fn resume_revalidates_persisted_thread_after_config_load( }; let resume_config_loads = blocked_resume.matching_loads.load(Ordering::Relaxed); let turn_result = match mutation { - ResumeMutation::GitMetadata | ResumeMutation::ProjectDeletion => { + ResumeMutation::Unchanged + | ResumeMutation::GitMetadata + | ResumeMutation::ProjectDeletion => { complete_in_process_turn(&mut client, RequestId::Integer(3), &thread_id).await } ResumeMutation::Archive | ResumeMutation::Delete | ResumeMutation::DuplicateResume => { @@ -947,7 +981,8 @@ async fn resume_revalidates_persisted_thread_after_config_load( ThreadLoadedListResponse { data: match mutation { ResumeMutation::Archive | ResumeMutation::Delete => Vec::new(), - ResumeMutation::GitMetadata + ResumeMutation::Unchanged + | ResumeMutation::GitMetadata | ResumeMutation::ProjectDeletion | ResumeMutation::DuplicateResume => vec![thread_id.clone()], }, @@ -955,7 +990,9 @@ async fn resume_revalidates_persisted_thread_after_config_load( } ); match mutation { - ResumeMutation::GitMetadata | ResumeMutation::ProjectDeletion => { + ResumeMutation::Unchanged + | ResumeMutation::GitMetadata + | ResumeMutation::ProjectDeletion => { let resumed: ThreadResumeResponse = serde_json::from_value(resumed.expect("resume updated thread"))?; let expected_git_info = @@ -973,11 +1010,23 @@ async fn resume_revalidates_persisted_thread_after_config_load( (resumed.thread.git_info, resumed.thread.project_id), (expected_git_info.clone(), expected_project_id.clone()) ); - assert!( - resume_config_loads >= 2, - "changed persisted metadata must reload config" - ); + if matches!(mutation, ResumeMutation::Unchanged) { + assert_eq!(resume_config_loads, 1); + } else { + assert!( + resume_config_loads >= 2, + "changed persisted metadata must reload config" + ); + } turn_result?; + assert!( + model_requests + .single_request() + .message_input_texts("user") + .iter() + .any(|text| text == "Saved user message"), + "resumed model context must retain the saved user message" + ); // Check persisted metadata after the resumed writer appends a turn // and shuts down; an in-memory response alone can hide stale writes. @@ -1033,13 +1082,21 @@ async fn resume_revalidates_persisted_thread_after_config_load( Ok(()) } +#[test_case(ThreadHistoryMode::Legacy; "legacy")] +#[test_case(ThreadHistoryMode::Paginated; "paginated")] #[tokio::test(flavor = "multi_thread", worker_threads = 2)] -async fn resume_reloads_config_when_saved_workspace_roots_change() -> Result<()> { +async fn resume_reloads_config_when_saved_workspace_roots_change( + history_mode: ThreadHistoryMode, +) -> Result<()> { let responses_server = create_mock_responses_server_repeating_assistant("Ready").await; let codex_home = TempDir::new()?; MockResponsesConfig::new(&responses_server.uri()).write(codex_home.path())?; let filename_timestamp = "2025-01-05T12-00-00"; - let thread_id = create_fake_rollout( + let create_rollout = match history_mode { + ThreadHistoryMode::Legacy => create_fake_rollout, + ThreadHistoryMode::Paginated => create_fake_paginated_rollout, + }; + let thread_id = create_rollout( codex_home.path(), filename_timestamp, "2025-01-05T12:00:00Z", @@ -1085,6 +1142,7 @@ async fn resume_reloads_config_when_saved_workspace_roots_change() -> Result<()> params: ThreadResumeParams { thread_id: thread_id.clone(), cwd: Some(cwd.to_string_lossy().into_owned()), + exclude_turns: matches!(history_mode, ThreadHistoryMode::Paginated), ..Default::default() }, })); diff --git a/codex-rs/thread-store/src/local/mod.rs b/codex-rs/thread-store/src/local/mod.rs index 17e58a6180..98c6ff70c2 100644 --- a/codex-rs/thread-store/src/local/mod.rs +++ b/codex-rs/thread-store/src/local/mod.rs @@ -616,6 +616,25 @@ impl ThreadStore for LocalThreadStore { Box::pin(async move { model_context::load_latest_model_context(self, params).await }) } + fn is_model_context_current( + &self, + params: LoadThreadHistoryParams, + revision: String, + ) -> ThreadStoreFuture<'_, bool> { + Box::pin(async move { + let resolved = if params.include_archived { + thread_rollout_resolver::resolve_current_including_archived(self, params.thread_id) + .await? + } else { + thread_rollout_resolver::resolve_current(self, params.thread_id).await? + }; + let Some(resolved) = resolved else { + return Ok(false); + }; + Ok(history_revision::read(&resolved.path).await.as_ref() == Some(&revision)) + }) + } + fn prepare_fork(&self, params: PrepareForkParams) -> ThreadStoreFuture<'_, PreparedFork> { Box::pin(async move { paginated_fork::prepare(self, params).await }) } diff --git a/codex-rs/thread-store/src/local/model_context_tests.rs b/codex-rs/thread-store/src/local/model_context_tests.rs index 683d66a9fb..25fca4e638 100644 --- a/codex-rs/thread-store/src/local/model_context_tests.rs +++ b/codex-rs/thread-store/src/local/model_context_tests.rs @@ -33,6 +33,62 @@ use crate::ThreadStore; use crate::local::test_support::test_config; use crate::local::test_support::write_session_file_with_history_mode; +#[tokio::test] +async fn model_context_revision_invalidates_after_append_or_removal() { + let home = TempDir::new().expect("temp dir"); + let uuid = Uuid::from_u128(/*v*/ 1009); + let thread_id = ThreadId::from_string(&uuid.to_string()).expect("thread id"); + let path = write_paginated_rollout( + home.path(), + "2025-01-03T13-00-08", + uuid, + [user_message("saved message")], + ); + let store = LocalThreadStore::new(test_config(home.path()), /*state_db*/ None); + let params = LoadThreadHistoryParams { + thread_id, + include_archived: false, + }; + let context = store + .load_latest_model_context(params.clone()) + .await + .expect("load model context"); + let revision = context.revision.expect("stable history revision"); + assert!( + store + .is_model_context_current(params.clone(), revision.clone()) + .await + .expect("validate unchanged context") + ); + + let new_item = user_message("message appended during configuration loading"); + append_items(&path, [new_item.clone()]); + assert!( + !store + .is_model_context_current(params.clone(), revision.clone()) + .await + .expect("invalidate stale context") + ); + let updated = store + .load_latest_model_context(params.clone()) + .await + .expect("load updated model context"); + let mut expected = context.items; + expected.push(new_item); + assert_eq!( + serde_json::to_value(updated.items).expect("serialize updated context"), + serde_json::to_value(expected).expect("serialize expected context") + ); + + std::fs::remove_file(path).expect("remove rollout"); + assert!( + !store + .is_model_context_current(params, revision) + .await + .expect("invalidate removed context") + ); +} + #[tokio::test] async fn stops_at_newest_usable_compaction_and_keeps_companions() { let home = TempDir::new().expect("temp dir"); diff --git a/codex-rs/thread-store/src/store.rs b/codex-rs/thread-store/src/store.rs index 762861e771..9a36132dff 100644 --- a/codex-rs/thread-store/src/store.rs +++ b/codex-rs/thread-store/src/store.rs @@ -268,6 +268,19 @@ pub trait ThreadStore: Any + Send + Sync { }) } + /// Checks whether a previously loaded model context still represents the stored history. + /// + /// Stores that cannot validate the opaque revision must return false. A false result or an + /// error requires a fresh read. This does not acquire writer ownership; + /// [`Self::resume_thread`] must still validate the snapshot after acquisition. + fn is_model_context_current( + &self, + _params: LoadThreadHistoryParams, + _revision: String, + ) -> ThreadStoreFuture<'_, bool> { + Box::pin(async { Ok(false) }) + } + /// Freezes source history and model context used to initialize a referenced fork. /// /// Stores without reference-backed fork support can retain this default implementation.