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 2002b08a88..7f5f020263 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 @@ -32,6 +32,8 @@ use codex_app_server_protocol::ThreadDeleteParams; use codex_app_server_protocol::ThreadDeleteResponse; use codex_app_server_protocol::ThreadListParams; use codex_app_server_protocol::ThreadListResponse; +use codex_app_server_protocol::ThreadReadParams; +use codex_app_server_protocol::ThreadReadResponse; use codex_app_server_protocol::ThreadResumeParams; use codex_app_server_protocol::ThreadStartParams; use codex_app_server_protocol::ThreadStartResponse; @@ -143,8 +145,27 @@ async fn thread_delete_with_non_local_thread_store_does_not_create_local_persist assert_eq!(data.len(), 1); assert_eq!(data[0].id, thread.id); assert_eq!(data[0].path, None); + assert_eq!(data[0].preview, "Hello"); + assert_eq!(data[0].model_provider, "mock_provider"); - delete_thread(&client, /*request_id*/ 4, thread.id.clone()).await?; + let response = client + .request(ClientRequest::ThreadRead { + request_id: RequestId::Integer(4), + params: ThreadReadParams { + thread_id: thread.id.clone(), + include_turns: false, + }, + }) + .await? + .expect("thread/read should succeed"); + let ThreadReadResponse { + thread: read_thread, + } = serde_json::from_value(response).expect("thread/read response should parse"); + assert_eq!(read_thread.preview, "Hello"); + assert_eq!(read_thread.model_provider, "mock_provider"); + assert_eq!(read_thread.path, None); + + delete_thread(&client, /*request_id*/ 5, thread.id.clone()).await?; let unloaded_thread_id = ThreadId::from_string(&Uuid::new_v4().to_string())?; thread_store .create_thread(StoreCreateThreadParams { @@ -171,7 +192,7 @@ async fn thread_delete_with_non_local_thread_store_does_not_create_local_persist .await?; delete_thread( &client, - /*request_id*/ 5, + /*request_id*/ 6, unloaded_thread_id.to_string(), ) .await?; diff --git a/codex-rs/thread-store/src/in_memory.rs b/codex-rs/thread-store/src/in_memory.rs index cb3f198ff3..4ba1a346fc 100644 --- a/codex-rs/thread-store/src/in_memory.rs +++ b/codex-rs/thread-store/src/in_memory.rs @@ -38,6 +38,8 @@ use crate::ThreadStoreFuture; use crate::ThreadStoreResult; use crate::UpdateThreadMetadataParams; use crate::error::reject_paginated_history_mode; +use crate::thread_metadata_sync::PendingThreadMetadataPatch; +use crate::thread_metadata_sync::ThreadMetadataSync; use crate::types::canonical_history_mode_from_rollout_items; static IN_MEMORY_THREAD_STORES: OnceLock>>> = @@ -393,11 +395,47 @@ struct InMemoryThreadStoreState { calls: InMemoryThreadStoreCalls, created_threads: HashMap, histories: HashMap>, + metadata_syncs: HashMap, metadata_updates: HashMap, names: HashMap>, rollout_paths: HashMap, } +impl InMemoryThreadStoreState { + fn apply_pending_metadata_update( + &mut self, + thread_id: ThreadId, + update: Option, + ) { + let Some(update) = update else { + return; + }; + self.metadata_updates + .entry(thread_id) + .or_default() + .merge(update.patch.clone()); + if let Some(metadata_sync) = self.metadata_syncs.get_mut(&thread_id) { + metadata_sync.mark_pending_update_applied(&update); + } + } + + fn flush_pending_metadata_update(&mut self, thread_id: ThreadId) { + let update = self + .metadata_syncs + .get(&thread_id) + .and_then(ThreadMetadataSync::take_pending_update); + self.apply_pending_metadata_update(thread_id, update); + } + + fn flush_pending_metadata_update_for_existing_history(&mut self, thread_id: ThreadId) { + let update = self + .metadata_syncs + .get(&thread_id) + .and_then(ThreadMetadataSync::take_pending_update_for_existing_history); + self.apply_pending_metadata_update(thread_id, update); + } +} + impl InMemoryThreadStore { /// Returns the store associated with `id`, creating it if needed. pub fn for_id(id: impl Into) -> Arc { @@ -421,6 +459,7 @@ impl InMemoryThreadStore { async fn create_thread(&self, params: CreateThreadParams) -> ThreadStoreResult<()> { reject_paginated_history_mode(params.history_mode)?; + let metadata_sync = ThreadMetadataSync::for_create(¶ms).await; let mut state = self.state.lock().await; state.calls.create_thread += 1; let session_meta = SessionMeta { @@ -454,30 +493,36 @@ impl InMemoryThreadStore { meta: session_meta, git: None, })); - state.created_threads.insert(params.thread_id, params); + let thread_id = params.thread_id; + state.created_threads.insert(thread_id, params); + state.metadata_syncs.insert(thread_id, metadata_sync); Ok(()) } async fn resume_thread(&self, params: ResumeThreadParams) -> ThreadStoreResult<()> { + let mut metadata_sync = ThreadMetadataSync::for_resume(¶ms); let mut state = self.state.lock().await; state.calls.resume_thread += 1; + let thread_id = params.thread_id; let history_mode = params .history .as_deref() .map(Vec::as_slice) .map(canonical_history_mode_from_rollout_items) - .unwrap_or_else(|| history_mode_from_state(&state, params.thread_id)); + .unwrap_or_else(|| history_mode_from_state(&state, thread_id)); reject_paginated_history_mode(history_mode)?; if let Some(history) = params.history { state .histories - .insert(params.thread_id, Arc::unwrap_or_clone(history)); + .insert(thread_id, Arc::unwrap_or_clone(history)); } else { - state.histories.entry(params.thread_id).or_default(); + let history = state.histories.entry(thread_id).or_default(); + metadata_sync.record_resume_history(history); } if let Some(rollout_path) = params.rollout_path { - state.rollout_paths.insert(rollout_path, params.thread_id); + state.rollout_paths.insert(rollout_path, thread_id); } + state.metadata_syncs.insert(thread_id, metadata_sync); Ok(()) } @@ -488,11 +533,16 @@ impl InMemoryThreadStore { } let mut state = self.state.lock().await; state.calls.append_items += 1; + let update = state + .metadata_syncs + .get_mut(¶ms.thread_id) + .and_then(|metadata_sync| metadata_sync.observe_appended_items(&canonical_items)); state .histories .entry(params.thread_id) .or_default() .extend(canonical_items); + state.apply_pending_metadata_update(params.thread_id, update); Ok(()) } @@ -572,6 +622,7 @@ impl InMemoryThreadStore { ) -> ThreadStoreResult { let mut state = self.state.lock().await; state.calls.update_thread_metadata += 1; + state.flush_pending_metadata_update(params.thread_id); if let Some(name) = params.patch.name.clone() { state.names.insert(params.thread_id, name); } @@ -588,6 +639,7 @@ impl InMemoryThreadStore { state.calls.delete_thread += 1; let existed = state.histories.remove(¶ms.thread_id).is_some(); state.created_threads.remove(¶ms.thread_id); + state.metadata_syncs.remove(¶ms.thread_id); state.names.remove(¶ms.thread_id); state.metadata_updates.remove(¶ms.thread_id); state @@ -620,30 +672,39 @@ impl ThreadStore for InMemoryThreadStore { Box::pin(InMemoryThreadStore::append_items(self, params)) } - fn persist_thread(&self, _thread_id: ThreadId) -> ThreadStoreFuture<'_, ()> { + fn persist_thread(&self, thread_id: ThreadId) -> ThreadStoreFuture<'_, ()> { Box::pin(async move { - self.state.lock().await.calls.persist_thread += 1; + let mut state = self.state.lock().await; + state.calls.persist_thread += 1; + state.flush_pending_metadata_update(thread_id); Ok(()) }) } - fn flush_thread(&self, _thread_id: ThreadId) -> ThreadStoreFuture<'_, ()> { + fn flush_thread(&self, thread_id: ThreadId) -> ThreadStoreFuture<'_, ()> { Box::pin(async move { - self.state.lock().await.calls.flush_thread += 1; + let mut state = self.state.lock().await; + state.calls.flush_thread += 1; + state.flush_pending_metadata_update_for_existing_history(thread_id); Ok(()) }) } - fn shutdown_thread(&self, _thread_id: ThreadId) -> ThreadStoreFuture<'_, ()> { + fn shutdown_thread(&self, thread_id: ThreadId) -> ThreadStoreFuture<'_, ()> { Box::pin(async move { - self.state.lock().await.calls.shutdown_thread += 1; + let mut state = self.state.lock().await; + state.calls.shutdown_thread += 1; + state.flush_pending_metadata_update_for_existing_history(thread_id); + state.metadata_syncs.remove(&thread_id); Ok(()) }) } - fn discard_thread(&self, _thread_id: ThreadId) -> ThreadStoreFuture<'_, ()> { + fn discard_thread(&self, thread_id: ThreadId) -> ThreadStoreFuture<'_, ()> { Box::pin(async move { - self.state.lock().await.calls.discard_thread += 1; + let mut state = self.state.lock().await; + state.calls.discard_thread += 1; + state.metadata_syncs.remove(&thread_id); Ok(()) }) } diff --git a/codex-rs/thread-store/src/thread_metadata_sync.rs b/codex-rs/thread-store/src/thread_metadata_sync.rs index 9cd8bded17..256657a414 100644 --- a/codex-rs/thread-store/src/thread_metadata_sync.rs +++ b/codex-rs/thread-store/src/thread_metadata_sync.rs @@ -27,10 +27,10 @@ const THREAD_UPDATED_AT_TOUCH_INTERVAL: Duration = Duration::from_secs(5); #[cfg(test)] const THREAD_UPDATED_AT_TOUCH_INTERVAL: Duration = Duration::from_millis(50); -/// Local-store helper that derives SQLite metadata from canonical rollout items. +/// Store implementation helper that derives query metadata from canonical rollout items. /// -/// This is private implementation state for [`crate::LocalThreadStore`]. Other stores receive raw -/// history and may derive their own metadata, or none at all. +/// Stores may retain this as private per-thread state when they need the local metadata projection. +/// Other stores receive raw history and may derive different metadata, or none at all. pub(crate) struct ThreadMetadataSync { thread_id: ThreadId, cwd_seen: bool,