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
This commit is contained in:
Charlie Marsh
2026-10-22 17:15:06 +00:00
committed by copyberry
parent 9cb547349e
commit ffda54f5b8
5 changed files with 207 additions and 21 deletions
@@ -110,6 +110,7 @@ struct ResumeConfigState {
struct PreparedResumeConfig {
state: ResumeConfigState,
config: Config,
model_context: Option<ResumedHistory>,
}
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<ResumedHistory>,
) -> 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
@@ -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()
},
}));
+19
View File
@@ -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 })
}
@@ -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");
+13
View File
@@ -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.