From 9aa9a81224a4343a61e7f565dafd814e2b580cd6 Mon Sep 17 00:00:00 2001 From: Brent Traut Date: Wed, 24 Jun 2026 12:02:47 -0700 Subject: [PATCH] [codex] Add store-owned thread catalog changes --- codex-rs/thread-store/src/in_memory.rs | 139 +++++++++++++++++++++++-- codex-rs/thread-store/src/lib.rs | 1 + codex-rs/thread-store/src/local/mod.rs | 56 +++++++++- codex-rs/thread-store/src/store.rs | 8 ++ codex-rs/thread-store/src/types.rs | 9 ++ 5 files changed, 202 insertions(+), 11 deletions(-) diff --git a/codex-rs/thread-store/src/in_memory.rs b/codex-rs/thread-store/src/in_memory.rs index cb3f198ff3..3407880e61 100644 --- a/codex-rs/thread-store/src/in_memory.rs +++ b/codex-rs/thread-store/src/in_memory.rs @@ -17,6 +17,7 @@ use codex_protocol::protocol::SessionMetaLine; use codex_protocol::protocol::ThreadHistoryMode; use codex_protocol::protocol::ThreadMemoryMode; use codex_rollout::persisted_rollout_items; +use tokio::sync::broadcast; use crate::AppendThreadItemsParams; use crate::ArchiveThreadParams; @@ -29,6 +30,7 @@ use crate::ReadThreadParams; use crate::ResumeThreadParams; use crate::StoredThread; use crate::StoredThreadHistory; +use crate::ThreadCatalogChange; use crate::ThreadMetadataPatch; use crate::ThreadPage; use crate::ThreadRelationFilter; @@ -58,6 +60,7 @@ mod tests { use crate::ThreadSortKey; use codex_protocol::models::BaseInstructions; use codex_protocol::protocol::SessionSource; + use pretty_assertions::assert_eq; #[tokio::test] async fn default_turn_pagination_methods_return_unsupported() { @@ -348,6 +351,87 @@ mod tests { } )); } + + #[tokio::test] + async fn catalog_changes_follow_successful_mutations() { + let store = InMemoryThreadStore::default(); + let mut changes = store.subscribe_catalog_changes(); + let thread_id = ThreadId::default(); + + ThreadStore::create_thread( + &store, + create_thread_params(thread_id, ThreadHistoryMode::Legacy), + ) + .await + .expect("create thread"); + assert_eq!( + changes.recv().await.expect("create change"), + ThreadCatalogChange::Upsert { thread_id } + ); + + store + .read_thread(ReadThreadParams { + thread_id, + include_archived: false, + include_history: false, + }) + .await + .expect("read thread"); + ThreadStore::update_thread_metadata( + &store, + UpdateThreadMetadataParams { + thread_id, + patch: ThreadMetadataPatch::default(), + include_archived: false, + }, + ) + .await + .expect("apply empty metadata patch"); + assert!(matches!( + changes.try_recv(), + Err(broadcast::error::TryRecvError::Empty) + )); + + ThreadStore::update_thread_metadata( + &store, + UpdateThreadMetadataParams { + thread_id, + patch: ThreadMetadataPatch { + name: Some(Some("renamed".to_string())), + ..Default::default() + }, + include_archived: false, + }, + ) + .await + .expect("update metadata"); + store + .archive_thread(ArchiveThreadParams { thread_id }) + .await + .expect("archive thread"); + store + .unarchive_thread(ArchiveThreadParams { thread_id }) + .await + .expect("unarchive thread"); + ThreadStore::delete_thread(&store, DeleteThreadParams { thread_id }) + .await + .expect("delete thread"); + + assert_eq!( + [ + ThreadCatalogChange::Upsert { thread_id }, + ThreadCatalogChange::Upsert { thread_id }, + ThreadCatalogChange::Upsert { thread_id }, + ThreadCatalogChange::Delete { thread_id }, + ], + [ + changes.recv().await.expect("metadata change"), + changes.recv().await.expect("archive change"), + changes.recv().await.expect("unarchive change"), + changes.recv().await.expect("delete change"), + ] + ); + } } fn stores_guard() -> MutexGuard<'static, HashMap>> { @@ -383,9 +467,19 @@ pub struct InMemoryThreadStoreCalls { /// Test and debug configs can select this store by id, letting tests exercise /// config-driven non-local persistence without requiring the real remote gRPC /// service. -#[derive(Default)] pub struct InMemoryThreadStore { state: tokio::sync::Mutex, + catalog_changes_tx: broadcast::Sender, +} + +impl Default for InMemoryThreadStore { + fn default() -> Self { + let (catalog_changes_tx, _) = broadcast::channel(256); + Self { + state: tokio::sync::Mutex::new(InMemoryThreadStoreState::default()), + catalog_changes_tx, + } + } } #[derive(Default)] @@ -419,6 +513,10 @@ impl InMemoryThreadStore { self.state.lock().await.calls.clone() } + fn publish_catalog_change(&self, change: ThreadCatalogChange) { + let _ = self.catalog_changes_tx.send(change); + } + async fn create_thread(&self, params: CreateThreadParams) -> ThreadStoreResult<()> { reject_paginated_history_mode(params.history_mode)?; let mut state = self.state.lock().await; @@ -608,8 +706,17 @@ impl ThreadStore for InMemoryThreadStore { self } + fn subscribe_catalog_changes(&self) -> broadcast::Receiver { + self.catalog_changes_tx.subscribe() + } + fn create_thread(&self, params: CreateThreadParams) -> ThreadStoreFuture<'_, ()> { - Box::pin(InMemoryThreadStore::create_thread(self, params)) + let thread_id = params.thread_id; + Box::pin(async move { + InMemoryThreadStore::create_thread(self, params).await?; + self.publish_catalog_change(ThreadCatalogChange::Upsert { thread_id }); + Ok(()) + }) } fn resume_thread(&self, params: ResumeThreadParams) -> ThreadStoreFuture<'_, ()> { @@ -707,26 +814,46 @@ impl ThreadStore for InMemoryThreadStore { &self, params: UpdateThreadMetadataParams, ) -> ThreadStoreFuture<'_, StoredThread> { - Box::pin(InMemoryThreadStore::update_thread_metadata(self, params)) + let thread_id = params.thread_id; + let changed = !params.patch.is_empty(); + Box::pin(async move { + let thread = InMemoryThreadStore::update_thread_metadata(self, params).await?; + if changed { + self.publish_catalog_change(ThreadCatalogChange::Upsert { thread_id }); + } + Ok(thread) + }) } - fn archive_thread(&self, _params: ArchiveThreadParams) -> ThreadStoreFuture<'_, ()> { + fn archive_thread(&self, params: ArchiveThreadParams) -> ThreadStoreFuture<'_, ()> { + let thread_id = params.thread_id; Box::pin(async move { self.state.lock().await.calls.archive_thread += 1; + self.publish_catalog_change(ThreadCatalogChange::Upsert { thread_id }); Ok(()) }) } fn unarchive_thread(&self, params: ArchiveThreadParams) -> ThreadStoreFuture<'_, StoredThread> { + let thread_id = params.thread_id; Box::pin(async move { let mut state = self.state.lock().await; state.calls.unarchive_thread += 1; - stored_thread_from_state(&state, params.thread_id, /*include_history*/ false) + let thread = + stored_thread_from_state(&state, thread_id, /*include_history*/ false)?; + drop(state); + self.publish_catalog_change(ThreadCatalogChange::Upsert { thread_id }); + Ok(thread) }) } fn delete_thread(&self, params: DeleteThreadParams) -> ThreadStoreFuture<'_, ()> { - Box::pin(InMemoryThreadStore::delete_thread(self, params)) + let thread_id = params.thread_id; + Box::pin(async move { + InMemoryThreadStore::delete_thread(self, params).await?; + self.publish_catalog_change(ThreadCatalogChange::Delete { thread_id }); + Ok(()) + }) } } diff --git a/codex-rs/thread-store/src/lib.rs b/codex-rs/thread-store/src/lib.rs index 34a0ab2ba1..218611743c 100644 --- a/codex-rs/thread-store/src/lib.rs +++ b/codex-rs/thread-store/src/lib.rs @@ -47,6 +47,7 @@ pub use types::StoredTurn; pub use types::StoredTurnError; pub use types::StoredTurnItemsView; pub use types::StoredTurnStatus; +pub use types::ThreadCatalogChange; pub use types::ThreadMetadataPatch; pub use types::ThreadPage; pub use types::ThreadPersistenceMetadata; diff --git a/codex-rs/thread-store/src/local/mod.rs b/codex-rs/thread-store/src/local/mod.rs index 5939ea606e..3bbcd564aa 100644 --- a/codex-rs/thread-store/src/local/mod.rs +++ b/codex-rs/thread-store/src/local/mod.rs @@ -21,6 +21,7 @@ use std::collections::hash_map::Entry; use std::path::PathBuf; use std::sync::Arc; use tokio::sync::Mutex; +use tokio::sync::broadcast; use crate::AppendThreadItemsParams; use crate::ArchiveThreadParams; @@ -34,6 +35,7 @@ use crate::ResumeThreadParams; use crate::SearchThreadsParams; use crate::StoredThread; use crate::StoredThreadHistory; +use crate::ThreadCatalogChange; use crate::ThreadPage; use crate::ThreadSearchPage; use crate::ThreadStore; @@ -60,6 +62,7 @@ pub struct LocalThreadStore { pub(super) config: LocalThreadStoreConfig, live_recorders: Arc>>, state_db: Option, + catalog_changes_tx: broadcast::Sender, } struct LiveRecorderEntry { @@ -103,13 +106,19 @@ impl std::fmt::Debug for LocalThreadStore { impl LocalThreadStore { /// Create a local store using an already initialized state DB handle. pub fn new(config: LocalThreadStoreConfig, state_db: Option) -> Self { + let (catalog_changes_tx, _) = broadcast::channel(256); Self { config, live_recorders: Arc::new(Mutex::new(HashMap::new())), state_db, + catalog_changes_tx, } } + fn publish_catalog_change(&self, change: ThreadCatalogChange) { + let _ = self.catalog_changes_tx.send(change); + } + /// Return the state DB handle used by local rollout writers. pub async fn state_db(&self) -> Option { self.state_db.clone() @@ -242,6 +251,10 @@ impl ThreadStore for LocalThreadStore { self } + fn subscribe_catalog_changes(&self) -> broadcast::Receiver { + self.catalog_changes_tx.subscribe() + } + fn create_thread(&self, params: CreateThreadParams) -> ThreadStoreFuture<'_, ()> { Box::pin(async move { live_writer::create_thread(self, params).await }) } @@ -255,7 +268,11 @@ impl ThreadStore for LocalThreadStore { } fn persist_thread(&self, thread_id: ThreadId) -> ThreadStoreFuture<'_, ()> { - Box::pin(async move { live_writer::persist_thread(self, thread_id).await }) + Box::pin(async move { + live_writer::persist_thread(self, thread_id).await?; + self.publish_catalog_change(ThreadCatalogChange::Upsert { thread_id }); + Ok(()) + }) } fn flush_thread(&self, thread_id: ThreadId) -> ThreadStoreFuture<'_, ()> { @@ -305,19 +322,42 @@ impl ThreadStore for LocalThreadStore { &self, params: UpdateThreadMetadataParams, ) -> ThreadStoreFuture<'_, StoredThread> { - Box::pin(async move { update_thread_metadata::update_thread_metadata(self, params).await }) + let thread_id = params.thread_id; + let changed = !params.patch.is_empty(); + Box::pin(async move { + let thread = update_thread_metadata::update_thread_metadata(self, params).await?; + if changed { + self.publish_catalog_change(ThreadCatalogChange::Upsert { thread_id }); + } + Ok(thread) + }) } fn archive_thread(&self, params: ArchiveThreadParams) -> ThreadStoreFuture<'_, ()> { - Box::pin(async move { archive_thread::archive_thread(self, params).await }) + let thread_id = params.thread_id; + Box::pin(async move { + archive_thread::archive_thread(self, params).await?; + self.publish_catalog_change(ThreadCatalogChange::Upsert { thread_id }); + Ok(()) + }) } fn unarchive_thread(&self, params: ArchiveThreadParams) -> ThreadStoreFuture<'_, StoredThread> { - Box::pin(async move { unarchive_thread::unarchive_thread(self, params).await }) + let thread_id = params.thread_id; + Box::pin(async move { + let thread = unarchive_thread::unarchive_thread(self, params).await?; + self.publish_catalog_change(ThreadCatalogChange::Upsert { thread_id }); + Ok(thread) + }) } fn delete_thread(&self, params: DeleteThreadParams) -> ThreadStoreFuture<'_, ()> { - Box::pin(async move { delete_thread::delete_thread(self, params).await }) + let thread_id = params.thread_id; + Box::pin(async move { + delete_thread::delete_thread(self, params).await?; + self.publish_catalog_change(ThreadCatalogChange::Delete { thread_id }); + Ok(()) + }) } } @@ -339,6 +379,7 @@ mod tests { use codex_protocol::protocol::TurnCompleteEvent; use codex_protocol::protocol::TurnStartedEvent; use codex_protocol::protocol::UserMessageEvent; + use pretty_assertions::assert_eq; use tempfile::TempDir; use super::*; @@ -353,6 +394,7 @@ mod tests { async fn live_writer_lifecycle_writes_and_closes() { let home = TempDir::new().expect("temp dir"); let store = LocalThreadStore::new(test_config(home.path()), /*state_db*/ None); + let mut catalog_changes = store.subscribe_catalog_changes(); let thread_id = ThreadId::default(); store @@ -375,6 +417,10 @@ mod tests { .persist_thread(thread_id) .await .expect("persist live thread"); + assert_eq!( + catalog_changes.recv().await.expect("persist change"), + ThreadCatalogChange::Upsert { thread_id } + ); store .flush_thread(thread_id) .await diff --git a/codex-rs/thread-store/src/store.rs b/codex-rs/thread-store/src/store.rs index 67ee603f49..c972545718 100644 --- a/codex-rs/thread-store/src/store.rs +++ b/codex-rs/thread-store/src/store.rs @@ -19,6 +19,7 @@ use crate::ResumeThreadParams; use crate::SearchThreadsParams; use crate::StoredThread; use crate::StoredThreadHistory; +use crate::ThreadCatalogChange; use crate::ThreadPage; use crate::ThreadSearchPage; use crate::ThreadStoreError; @@ -42,6 +43,13 @@ pub trait ThreadStore: Any + Send + Sync { ThreadHistoryMode::Legacy } + /// Subscribes to future catalog-row invalidations. + /// + /// Events are not replayed and do not contain snapshots. Consumers should subscribe before + /// reading their initial catalog view, reread the affected thread after each event, and rebuild + /// that view if the receiver reports lag. + fn subscribe_catalog_changes(&self) -> tokio::sync::broadcast::Receiver; + /// Creates a new live thread. fn create_thread(&self, params: CreateThreadParams) -> ThreadStoreFuture<'_, ()>; diff --git a/codex-rs/thread-store/src/types.rs b/codex-rs/thread-store/src/types.rs index 4e688eca6f..6a2e225d2b 100644 --- a/codex-rs/thread-store/src/types.rs +++ b/codex-rs/thread-store/src/types.rs @@ -730,6 +730,15 @@ pub struct DeleteThreadParams { pub thread_id: ThreadId, } +/// Store-owned invalidation for a thread catalog row. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum ThreadCatalogChange { + /// The row was created or its metadata changed. + Upsert { thread_id: ThreadId }, + /// The row was deleted. + Delete { thread_id: ThreadId }, +} + #[cfg(test)] mod tests { use pretty_assertions::assert_eq;