From 91bc31dca3a26cc4384ae1daee5e406a3f296389 Mon Sep 17 00:00:00 2001 From: Anton Panasenko Date: Mon, 29 Jun 2026 17:11:01 -0700 Subject: [PATCH] perf(thread-store): project append metadata asynchronously --- codex-rs/Cargo.lock | 1 + .../thread_goal_processor.rs | 2 +- codex-rs/core/src/codex_thread.rs | 20 + codex-rs/thread-store/Cargo.toml | 1 + codex-rs/thread-store/src/error.rs | 2 +- codex-rs/thread-store/src/live_thread.rs | 418 ++++++++++--- .../thread-store/src/live_thread_tests.rs | 567 ++++++++++++++++++ codex-rs/thread-store/src/local/mod.rs | 8 + .../thread-store/src/thread_metadata_sync.rs | 6 + 9 files changed, 954 insertions(+), 71 deletions(-) create mode 100644 codex-rs/thread-store/src/live_thread_tests.rs diff --git a/codex-rs/Cargo.lock b/codex-rs/Cargo.lock index 0d4d3c33bf..88639d4492 100644 --- a/codex-rs/Cargo.lock +++ b/codex-rs/Cargo.lock @@ -4062,6 +4062,7 @@ dependencies = [ "thiserror 2.0.18", "tokio", "tracing", + "tracing-subscriber", "uuid", ] diff --git a/codex-rs/app-server/src/request_processors/thread_goal_processor.rs b/codex-rs/app-server/src/request_processors/thread_goal_processor.rs index 5878047237..d30973a9cb 100644 --- a/codex-rs/app-server/src/request_processors/thread_goal_processor.rs +++ b/codex-rs/app-server/src/request_processors/thread_goal_processor.rs @@ -141,7 +141,7 @@ impl ThreadGoalRequestProcessor { // Live goal-first threads can be listed before any user turn is written. // Use the live path so JSONL and SQLite preview metadata stay in sync. thread - .append_rollout_items(&[outcome.thread_goal_updated_item()]) + .append_rollout_items_and_flush_metadata(&[outcome.thread_goal_updated_item()]) .await } Err(_) => Ok(()), diff --git a/codex-rs/core/src/codex_thread.rs b/codex-rs/core/src/codex_thread.rs index c80268e2f9..009279e573 100644 --- a/codex-rs/core/src/codex_thread.rs +++ b/codex-rs/core/src/codex_thread.rs @@ -572,6 +572,26 @@ impl CodexThread { live_thread.append_items(items).await } + /// Appends rollout items and waits until their derived thread metadata is queryable. + /// + /// Most turn-path appends intentionally project metadata in the background. Callers that + /// respond before a first turn exists, such as thread goals, need this stronger visibility + /// contract so an immediate thread read/list observes the new preview. + pub async fn append_rollout_items_and_flush_metadata( + &self, + items: &[RolloutItem], + ) -> ThreadStoreResult<()> { + let live_thread = self + .codex + .session + .live_thread_for_persistence("append rollout items and flush metadata") + .map_err(|err| ThreadStoreError::Internal { + message: err.to_string(), + })?; + live_thread.append_items(items).await?; + live_thread.flush().await + } + pub fn state_db(&self) -> Option { self.codex.state_db() } diff --git a/codex-rs/thread-store/Cargo.toml b/codex-rs/thread-store/Cargo.toml index 35b3919c60..91bb37c900 100644 --- a/codex-rs/thread-store/Cargo.toml +++ b/codex-rs/thread-store/Cargo.toml @@ -31,4 +31,5 @@ tracing = { workspace = true } pretty_assertions = { workspace = true } tempfile = { workspace = true } tokio = { workspace = true, features = ["macros", "rt-multi-thread"] } +tracing-subscriber = { workspace = true } uuid = { workspace = true } diff --git a/codex-rs/thread-store/src/error.rs b/codex-rs/thread-store/src/error.rs index bc5e8ed7c4..2445b07aa2 100644 --- a/codex-rs/thread-store/src/error.rs +++ b/codex-rs/thread-store/src/error.rs @@ -16,7 +16,7 @@ pub(crate) fn reject_paginated_history_mode( } /// Error type shared by thread-store implementations. -#[derive(Debug, thiserror::Error)] +#[derive(Clone, Debug, thiserror::Error)] pub enum ThreadStoreError { /// The requested thread does not exist in this store. #[error("thread {thread_id} not found")] diff --git a/codex-rs/thread-store/src/live_thread.rs b/codex-rs/thread-store/src/live_thread.rs index 793bbf9965..bf79f741d6 100644 --- a/codex-rs/thread-store/src/live_thread.rs +++ b/codex-rs/thread-store/src/live_thread.rs @@ -1,13 +1,17 @@ use std::path::PathBuf; use std::sync::Arc; +use std::time::Instant; use codex_protocol::ThreadId; use codex_protocol::protocol::RolloutItem; use codex_protocol::protocol::ThreadMemoryMode; +use codex_protocol::protocol::W3cTraceContext; use codex_rollout::RolloutPersistenceTelemetry; use codex_rollout::measure_and_filter_rollout_items; use codex_rollout::persisted_rollout_items; use tokio::sync::Mutex; +use tokio::sync::oneshot; +use tracing::Instrument; use tracing::warn; use crate::AppendThreadItemsParams; @@ -20,8 +24,10 @@ use crate::StoredThread; use crate::StoredThreadHistory; use crate::ThreadMetadataPatch; use crate::ThreadStore; +use crate::ThreadStoreError; use crate::ThreadStoreResult; use crate::UpdateThreadMetadataParams; +use crate::thread_metadata_sync::PendingThreadMetadataPatch; use crate::thread_metadata_sync::ThreadMetadataSync; /// Handle for an active thread's persistence lifecycle. @@ -33,10 +39,54 @@ use crate::thread_metadata_sync::ThreadMetadataSync; pub struct LiveThread { thread_id: ThreadId, thread_store: Arc, - metadata_sync: Arc>, + metadata_worker: Arc, persistence_telemetry: RolloutPersistenceTelemetry, } +struct MetadataUpdateWorker { + thread_id: ThreadId, + thread_store: Arc, + state: Mutex, +} + +struct MetadataUpdateWorkerState { + metadata_sync: ThreadMetadataSync, + applied_generation: u64, + running: bool, + in_flight_generation: Option, + queued_update: Option, + barriers: Vec, + discarding: bool, + idle_waiters: Vec>, +} + +struct QueuedMetadataUpdate { + update: PendingThreadMetadataPatch, + enqueued_at: Instant, + trace_context: Option, +} + +impl QueuedMetadataUpdate { + fn new(update: PendingThreadMetadataPatch) -> Self { + Self { + update, + enqueued_at: Instant::now(), + trace_context: codex_otel::current_span_w3c_trace_context(), + } + } +} + +struct MetadataBarrier { + target_generation: u64, + sender: oneshot::Sender>, +} + +#[derive(Clone, Copy)] +enum MetadataBarrierKind { + All, + ExistingHistory, +} + /// Owns a live thread while session initialization is still fallible. /// /// If initialization returns early after persistence has been opened, dropping this guard discards @@ -94,12 +144,7 @@ impl LiveThread { let thread_id = params.thread_id; let metadata_sync = ThreadMetadataSync::for_create(¶ms).await; thread_store.create_thread(params).await?; - Ok(Self { - thread_id, - thread_store, - metadata_sync: Arc::new(Mutex::new(metadata_sync)), - persistence_telemetry: RolloutPersistenceTelemetry::new(thread_id), - }) + Ok(Self::new(thread_id, thread_store, metadata_sync)) } pub async fn resume( @@ -130,12 +175,34 @@ impl LiveThread { } } } - Ok(Self { + Ok(Self::new(thread_id, thread_store, metadata_sync)) + } + + fn new( + thread_id: ThreadId, + thread_store: Arc, + metadata_sync: ThreadMetadataSync, + ) -> Self { + let metadata_worker = Arc::new(MetadataUpdateWorker { + thread_id, + thread_store: Arc::clone(&thread_store), + state: Mutex::new(MetadataUpdateWorkerState { + metadata_sync, + applied_generation: 0, + running: false, + in_flight_generation: None, + queued_update: None, + barriers: Vec::new(), + discarding: false, + idle_waiters: Vec::new(), + }), + }); + Self { thread_id, thread_store, - metadata_sync: Arc::new(Mutex::new(metadata_sync)), + metadata_worker, persistence_telemetry: RolloutPersistenceTelemetry::new(thread_id), - }) + } } #[tracing::instrument( @@ -148,17 +215,27 @@ impl LiveThread { if items.is_empty() { return Ok(()); } - let (canonical_items, measurement) = if self.persistence_telemetry.is_enabled() { - let (canonical_items, measurement) = measure_and_filter_rollout_items(items); - (canonical_items, Some(measurement)) - } else { - (persisted_rollout_items(items), None) - }; + let (canonical_items, measurement) = tracing::info_span!( + "thread_store.live_append.filter_items", + item_count = items.len(), + ) + .in_scope(|| { + if self.persistence_telemetry.is_enabled() { + let (canonical_items, measurement) = measure_and_filter_rollout_items(items); + (canonical_items, Some(measurement)) + } else { + (persisted_rollout_items(items), None) + } + }); self.thread_store .append_items(AppendThreadItemsParams { thread_id: self.thread_id, items: items.to_vec(), }) + .instrument(tracing::info_span!( + "thread_store.live_append.store_append", + item_count = items.len(), + )) .await?; if let Some(measurement) = measurement.as_ref() { self.persistence_telemetry.record_batch(items, measurement); @@ -166,45 +243,37 @@ impl LiveThread { if canonical_items.is_empty() { return Ok(()); } - let update = self - .metadata_sync - .lock() - .await - .observe_appended_items(canonical_items.as_slice()); - if let Some(update) = update { - self.thread_store - .update_thread_metadata(UpdateThreadMetadataParams { - thread_id: self.thread_id, - patch: update.patch.clone(), - include_archived: true, - }) - .await?; - self.metadata_sync - .lock() - .await - .mark_pending_update_applied(&update); - } + self.metadata_worker + .observe_appended_items(canonical_items.as_slice()) + .instrument(tracing::info_span!( + "thread_store.live_append.observe_metadata", + item_count = canonical_items.len(), + )) + .await; Ok(()) } pub async fn persist(&self) -> ThreadStoreResult<()> { self.thread_store.persist_thread(self.thread_id).await?; - self.flush_pending_metadata_update().await + self.metadata_worker.barrier(MetadataBarrierKind::All).await } pub async fn flush(&self) -> ThreadStoreResult<()> { self.thread_store.flush_thread(self.thread_id).await?; - self.flush_pending_metadata_update_for_existing_history() + self.metadata_worker + .barrier(MetadataBarrierKind::ExistingHistory) .await } pub async fn shutdown(&self) -> ThreadStoreResult<()> { - self.flush_pending_metadata_update_for_existing_history() + self.metadata_worker + .barrier(MetadataBarrierKind::ExistingHistory) .await?; self.thread_store.shutdown_thread(self.thread_id).await } pub async fn discard(&self) -> ThreadStoreResult<()> { + self.metadata_worker.discard_pending_and_wait().await?; self.thread_store.discard_thread(self.thread_id).await } @@ -239,7 +308,9 @@ impl LiveThread { mode: ThreadMemoryMode, include_archived: bool, ) -> ThreadStoreResult<()> { - self.flush_pending_metadata_update().await?; + self.metadata_worker + .barrier(MetadataBarrierKind::All) + .await?; self.thread_store .update_thread_metadata(UpdateThreadMetadataParams { thread_id: self.thread_id, @@ -258,7 +329,9 @@ impl LiveThread { patch: ThreadMetadataPatch, include_archived: bool, ) -> ThreadStoreResult { - self.flush_pending_metadata_update().await?; + self.metadata_worker + .barrier(MetadataBarrierKind::All) + .await?; self.thread_store .update_thread_metadata(UpdateThreadMetadataParams { thread_id: self.thread_id, @@ -284,39 +357,246 @@ impl LiveThread { .await .map(Some) } +} - async fn flush_pending_metadata_update(&self) -> ThreadStoreResult<()> { - let update = self.metadata_sync.lock().await.take_pending_update(); - self.apply_pending_metadata_update(update).await - } - - async fn flush_pending_metadata_update_for_existing_history(&self) -> ThreadStoreResult<()> { - let update = self - .metadata_sync - .lock() - .await - .take_pending_update_for_existing_history(); - self.apply_pending_metadata_update(update).await - } - - async fn apply_pending_metadata_update( - &self, - update: Option, - ) -> ThreadStoreResult<()> { - let Some(update) = update else { - return Ok(()); +impl MetadataUpdateWorker { + async fn observe_appended_items(self: &Arc, items: &[RolloutItem]) { + let should_spawn = { + let mut state = self.state.lock().await; + if state.discarding { + return; + } + let update = state.metadata_sync.observe_appended_items(items); + if let Some(update) = update { + state.queued_update = Some(QueuedMetadataUpdate::new(update)); + } + Self::start_if_needed(&mut state) }; - self.thread_store - .update_thread_metadata(UpdateThreadMetadataParams { - thread_id: self.thread_id, - patch: update.patch.clone(), - include_archived: true, + if should_spawn { + self.spawn(); + } + } + + async fn barrier(self: &Arc, kind: MetadataBarrierKind) -> ThreadStoreResult<()> { + let (receiver, should_spawn) = { + let mut state = self.state.lock().await; + if state.discarding { + return Err(ThreadStoreError::InvalidRequest { + message: "thread metadata update worker is being discarded".to_string(), + }); + } + let update = match kind { + MetadataBarrierKind::All => state.metadata_sync.take_pending_update(), + MetadataBarrierKind::ExistingHistory => state + .metadata_sync + .take_pending_update_for_existing_history(), + }; + let Some(update) = update else { + return Ok(()); + }; + let target_generation = update.generation(); + if target_generation <= state.applied_generation { + return Ok(()); + } + let generation_is_covered = state.in_flight_generation == Some(target_generation) + || state + .queued_update + .as_ref() + .is_some_and(|queued| queued.update.generation() >= target_generation); + if !generation_is_covered { + state.queued_update = Some(QueuedMetadataUpdate::new(update)); + } + let (sender, receiver) = oneshot::channel(); + state.barriers.push(MetadataBarrier { + target_generation, + sender, + }); + let should_spawn = Self::start_if_needed(&mut state); + (receiver, should_spawn) + }; + if should_spawn { + self.spawn(); + } + receiver.await.unwrap_or_else(|err| { + Err(ThreadStoreError::Internal { + message: format!("thread metadata update worker stopped before barrier: {err}"), }) - .await?; - self.metadata_sync - .lock() - .await - .mark_pending_update_applied(&update); + }) + } + + async fn discard_pending_and_wait(&self) -> ThreadStoreResult<()> { + let (receiver, barriers) = { + let mut state = self.state.lock().await; + state.discarding = true; + state.queued_update = None; + let barriers = std::mem::take(&mut state.barriers); + let receiver = if state.running { + let (sender, receiver) = oneshot::channel(); + state.idle_waiters.push(sender); + Some(receiver) + } else { + None + }; + (receiver, barriers) + }; + for barrier in barriers { + let _ = barrier.sender.send(Err(ThreadStoreError::InvalidRequest { + message: "thread metadata update worker was discarded".to_string(), + })); + } + if let Some(receiver) = receiver { + receiver.await.map_err(|err| ThreadStoreError::Internal { + message: format!("thread metadata update worker stopped before discard: {err}"), + })?; + } Ok(()) } + + fn start_if_needed(state: &mut MetadataUpdateWorkerState) -> bool { + if state.discarding || state.running || state.queued_update.is_none() { + return false; + } + state.running = true; + true + } + + fn spawn(self: &Arc) { + tokio::spawn( + Arc::clone(self) + .run() + .instrument(tracing::info_span!("thread_store.metadata_update_worker")), + ); + } + + async fn run(self: Arc) { + loop { + let queued = { + let mut state = self.state.lock().await; + let Some(queued) = state.queued_update.take() else { + state.running = false; + let barriers = std::mem::take(&mut state.barriers); + let idle_waiters = std::mem::take(&mut state.idle_waiters); + let applied_generation = state.applied_generation; + drop(state); + for barrier in barriers { + let result = if barrier.target_generation <= applied_generation { + Ok(()) + } else { + Err(ThreadStoreError::Internal { + message: + "metadata worker became idle before its barrier generation" + .to_string(), + }) + }; + let _ = barrier.sender.send(result); + } + for waiter in idle_waiters { + let _ = waiter.send(()); + } + return; + }; + state.in_flight_generation = Some(queued.update.generation()); + queued + }; + + let generation = queued.update.generation(); + let apply_span = tracing::info_span!( + "thread_store.metadata_update.apply", + generation, + queue_age_ms = queued.enqueued_at.elapsed().as_secs_f64() * 1_000.0, + ); + if let Some(trace_context) = queued.trace_context.as_ref() { + let _ = codex_otel::set_parent_from_w3c_trace_context(&apply_span, trace_context); + } + let result = async { + self.thread_store + .update_thread_metadata(UpdateThreadMetadataParams { + thread_id: self.thread_id, + patch: queued.update.patch.clone(), + include_archived: true, + }) + .instrument(tracing::info_span!( + "thread_store.metadata_update.update_thread_metadata" + )) + .await + } + .instrument(apply_span) + .await; + match result { + Ok(_) => { + let mut state = self.state.lock().await; + state + .metadata_sync + .mark_pending_update_applied(&queued.update); + state.applied_generation = + state.applied_generation.max(queued.update.generation()); + state.in_flight_generation = None; + let applied_generation = state.applied_generation; + let mut pending_barriers = Vec::new(); + let mut satisfied_barriers = Vec::new(); + for barrier in std::mem::take(&mut state.barriers) { + if barrier.target_generation <= applied_generation { + satisfied_barriers.push(barrier); + } else { + pending_barriers.push(barrier); + } + } + state.barriers = pending_barriers; + drop(state); + for barrier in satisfied_barriers { + let _ = barrier.sender.send(Ok(())); + } + } + Err(err) => { + warn!( + thread_id = %self.thread_id, + error = %err, + "failed to asynchronously update thread metadata" + ); + let (failed_barriers, idle_waiters, should_retry) = { + let mut state = self.state.lock().await; + state.in_flight_generation = None; + let retry_generation = if state.discarding { + None + } else { + state + .queued_update + .as_ref() + .map(|queued| queued.update.generation()) + }; + let should_retry = retry_generation.is_some(); + let (pending_barriers, failed_barriers) = + std::mem::take(&mut state.barriers) + .into_iter() + .partition(|barrier| { + retry_generation.is_some_and(|generation| { + barrier.target_generation <= generation + }) + }); + state.barriers = pending_barriers; + state.running = should_retry; + let idle_waiters = if should_retry { + Vec::new() + } else { + std::mem::take(&mut state.idle_waiters) + }; + (failed_barriers, idle_waiters, should_retry) + }; + for barrier in failed_barriers { + let _ = barrier.sender.send(Err(err.clone())); + } + for waiter in idle_waiters { + let _ = waiter.send(()); + } + if !should_retry { + return; + } + } + } + } + } } + +#[cfg(test)] +#[path = "live_thread_tests.rs"] +mod tests; diff --git a/codex-rs/thread-store/src/live_thread_tests.rs b/codex-rs/thread-store/src/live_thread_tests.rs new file mode 100644 index 0000000000..8554225edf --- /dev/null +++ b/codex-rs/thread-store/src/live_thread_tests.rs @@ -0,0 +1,567 @@ +use std::any::Any; +use std::sync::Arc; +use std::sync::Mutex as StdMutex; +use std::time::Duration; + +use codex_protocol::ThreadId; +use codex_protocol::models::BaseInstructions; +use codex_protocol::protocol::EventMsg; +use codex_protocol::protocol::RolloutItem; +use codex_protocol::protocol::SessionSource; +use codex_protocol::protocol::ThreadHistoryMode; +use codex_protocol::protocol::ThreadMemoryMode; +use codex_protocol::protocol::UserMessageEvent; +use pretty_assertions::assert_eq; +use tokio::sync::mpsc; +use tokio::sync::oneshot; +use tokio::time::timeout; +use tracing::Subscriber; +use tracing_subscriber::Layer; +use tracing_subscriber::layer::Context; +use tracing_subscriber::layer::SubscriberExt; +use tracing_subscriber::registry::LookupSpan; +use tracing_subscriber::util::SubscriberInitExt; +use uuid::Uuid; + +use super::LiveThread; +use crate::AppendThreadItemsParams; +use crate::ArchiveThreadParams; +use crate::CreateThreadParams; +use crate::DeleteThreadParams; +use crate::InMemoryThreadStore; +use crate::ItemPage; +use crate::ListItemsParams; +use crate::ListThreadsParams; +use crate::ListTurnsParams; +use crate::LoadThreadHistoryParams; +use crate::ReadThreadByRolloutPathParams; +use crate::ReadThreadParams; +use crate::ResumeThreadParams; +use crate::SearchThreadsParams; +use crate::StoredThread; +use crate::StoredThreadHistory; +use crate::ThreadPage; +use crate::ThreadPersistenceMetadata; +use crate::ThreadSearchPage; +use crate::ThreadStore; +use crate::ThreadStoreError; +use crate::ThreadStoreFuture; +use crate::ThreadStoreResult; +use crate::TurnPage; +use crate::UpdateThreadMetadataParams; + +const TEST_TIMEOUT: Duration = Duration::from_secs(1); + +struct MetadataUpdateAttempt { + params: UpdateThreadMetadataParams, + completion: oneshot::Sender>, +} + +struct ControlledThreadStore { + inner: InMemoryThreadStore, + update_sender: mpsc::UnboundedSender, +} + +impl ControlledThreadStore { + fn new() -> (Arc, mpsc::UnboundedReceiver) { + let (update_sender, update_receiver) = mpsc::unbounded_channel(); + ( + Arc::new(Self { + inner: InMemoryThreadStore::default(), + update_sender, + }), + update_receiver, + ) + } +} + +impl ThreadStore for ControlledThreadStore { + fn as_any(&self) -> &dyn Any { + self + } + + fn create_thread(&self, params: CreateThreadParams) -> ThreadStoreFuture<'_, ()> { + ThreadStore::create_thread(&self.inner, params) + } + + fn resume_thread(&self, params: ResumeThreadParams) -> ThreadStoreFuture<'_, ()> { + ThreadStore::resume_thread(&self.inner, params) + } + + fn append_items(&self, params: AppendThreadItemsParams) -> ThreadStoreFuture<'_, ()> { + ThreadStore::append_items(&self.inner, params) + } + + fn persist_thread(&self, thread_id: ThreadId) -> ThreadStoreFuture<'_, ()> { + ThreadStore::persist_thread(&self.inner, thread_id) + } + + fn flush_thread(&self, thread_id: ThreadId) -> ThreadStoreFuture<'_, ()> { + ThreadStore::flush_thread(&self.inner, thread_id) + } + + fn shutdown_thread(&self, thread_id: ThreadId) -> ThreadStoreFuture<'_, ()> { + ThreadStore::shutdown_thread(&self.inner, thread_id) + } + + fn discard_thread(&self, thread_id: ThreadId) -> ThreadStoreFuture<'_, ()> { + ThreadStore::discard_thread(&self.inner, thread_id) + } + + fn load_history( + &self, + params: LoadThreadHistoryParams, + ) -> ThreadStoreFuture<'_, StoredThreadHistory> { + ThreadStore::load_history(&self.inner, params) + } + + fn read_thread(&self, params: ReadThreadParams) -> ThreadStoreFuture<'_, StoredThread> { + ThreadStore::read_thread(&self.inner, params) + } + + fn read_thread_by_rollout_path( + &self, + params: ReadThreadByRolloutPathParams, + ) -> ThreadStoreFuture<'_, StoredThread> { + ThreadStore::read_thread_by_rollout_path(&self.inner, params) + } + + fn list_threads(&self, params: ListThreadsParams) -> ThreadStoreFuture<'_, ThreadPage> { + ThreadStore::list_threads(&self.inner, params) + } + + fn search_threads( + &self, + params: SearchThreadsParams, + ) -> ThreadStoreFuture<'_, ThreadSearchPage> { + ThreadStore::search_threads(&self.inner, params) + } + + fn list_turns(&self, params: ListTurnsParams) -> ThreadStoreFuture<'_, TurnPage> { + ThreadStore::list_turns(&self.inner, params) + } + + fn list_items(&self, params: ListItemsParams) -> ThreadStoreFuture<'_, ItemPage> { + ThreadStore::list_items(&self.inner, params) + } + + fn update_thread_metadata( + &self, + params: UpdateThreadMetadataParams, + ) -> ThreadStoreFuture<'_, StoredThread> { + Box::pin(async move { + let (completion, result) = oneshot::channel(); + self.update_sender + .send(MetadataUpdateAttempt { + params: params.clone(), + completion, + }) + .map_err(|_| ThreadStoreError::Internal { + message: "metadata update test receiver dropped".to_string(), + })?; + result.await.map_err(|err| ThreadStoreError::Internal { + message: format!("metadata update test completion dropped: {err}"), + })??; + ThreadStore::update_thread_metadata(&self.inner, params).await + }) + } + + fn archive_thread(&self, params: ArchiveThreadParams) -> ThreadStoreFuture<'_, ()> { + ThreadStore::archive_thread(&self.inner, params) + } + + fn unarchive_thread(&self, params: ArchiveThreadParams) -> ThreadStoreFuture<'_, StoredThread> { + ThreadStore::unarchive_thread(&self.inner, params) + } + + fn delete_thread(&self, params: DeleteThreadParams) -> ThreadStoreFuture<'_, ()> { + ThreadStore::delete_thread(&self.inner, params) + } +} + +#[tokio::test] +async fn append_returns_after_history_acceptance_while_metadata_is_blocked() { + let (store, mut updates) = ControlledThreadStore::new(); + let live_thread = create_live_thread(Arc::clone(&store)).await; + + timeout( + TEST_TIMEOUT, + live_thread.append_items(&[user_message("first message")]), + ) + .await + .expect("append should not wait for metadata persistence") + .expect("append should succeed"); + let attempt = next_update(&mut updates).await; + + assert_eq!(store.inner.calls().await.append_items, 1); + assert_eq!(store.inner.calls().await.update_thread_metadata, 0); + assert_eq!( + attempt.params.patch.preview.as_deref(), + Some("first message") + ); + + attempt + .completion + .send(Ok(())) + .expect("metadata worker should still be waiting"); + live_thread + .flush() + .await + .expect("flush should drain metadata"); + assert_eq!(store.inner.calls().await.update_thread_metadata, 1); +} + +#[tokio::test] +async fn appends_coalesce_while_metadata_update_is_in_flight() { + let (store, mut updates) = ControlledThreadStore::new(); + let live_thread = create_live_thread(Arc::clone(&store)).await; + + live_thread + .append_items(&[user_message("first message")]) + .await + .expect("first append should succeed"); + let first_attempt = next_update(&mut updates).await; + live_thread + .append_items(&[user_message("second message")]) + .await + .expect("second append should succeed"); + live_thread + .append_items(&[user_message("third message")]) + .await + .expect("third append should succeed"); + + first_attempt + .completion + .send(Ok(())) + .expect("metadata worker should still be waiting"); + let second_attempt = next_update(&mut updates).await; + let mut flush = tokio::spawn({ + let live_thread = live_thread.clone(); + async move { live_thread.flush().await } + }); + assert!( + timeout(Duration::from_millis(20), &mut flush) + .await + .is_err(), + "flush should wait for the in-flight coalesced update" + ); + live_thread + .append_items(&[user_message("accepted after flush barrier")]) + .await + .expect("later append should succeed"); + second_attempt + .completion + .send(Ok(())) + .expect("metadata worker should still be waiting"); + timeout(TEST_TIMEOUT, flush) + .await + .expect("flush should not wait for metadata accepted after its barrier") + .expect("flush task should not panic") + .expect("flush should succeed"); + let third_attempt = next_update(&mut updates).await; + third_attempt + .completion + .send(Ok(())) + .expect("metadata worker should still be waiting"); + live_thread + .flush() + .await + .expect("final flush should drain later metadata"); + + let calls = store.inner.calls().await; + assert_eq!(calls.append_items, 4); + assert_eq!(calls.update_thread_metadata, 3); +} + +#[tokio::test] +async fn queued_generation_keeps_its_barrier_when_in_flight_update_fails() { + let (store, mut updates) = ControlledThreadStore::new(); + let live_thread = create_live_thread(Arc::clone(&store)).await; + + live_thread + .append_items(&[user_message("first message")]) + .await + .expect("first append should succeed"); + let first_attempt = next_update(&mut updates).await; + live_thread + .append_items(&[user_message("second message")]) + .await + .expect("second append should succeed"); + let mut flush = tokio::spawn({ + let live_thread = live_thread.clone(); + async move { live_thread.flush().await } + }); + assert!( + timeout(Duration::from_millis(20), &mut flush) + .await + .is_err(), + "flush should wait for the queued generation" + ); + + first_attempt + .completion + .send(Err(test_error("superseded generation failed"))) + .expect("metadata worker should still be waiting"); + let retry_attempt = next_update(&mut updates).await; + assert!( + timeout(Duration::from_millis(20), &mut flush) + .await + .is_err(), + "a failure covered by the queued generation must not fail its barrier" + ); + retry_attempt + .completion + .send(Ok(())) + .expect("metadata worker should still be waiting"); + timeout(TEST_TIMEOUT, flush) + .await + .expect("flush should finish after the queued generation is applied") + .expect("flush task should not panic") + .expect("flush should succeed"); + assert_eq!(store.inner.calls().await.update_thread_metadata, 1); +} + +#[derive(Clone)] +struct WorkerSpanParentLayer { + parent_name: Arc>>, +} + +impl Layer for WorkerSpanParentLayer +where + S: Subscriber + for<'a> LookupSpan<'a>, +{ + fn on_new_span( + &self, + _attrs: &tracing::span::Attributes<'_>, + id: &tracing::span::Id, + ctx: Context<'_, S>, + ) { + let Some(span) = ctx.span(id) else { + return; + }; + if span.metadata().name() != "thread_store.metadata_update_worker" { + return; + } + *self.parent_name.lock().expect("parent name lock") = span + .parent() + .map(|parent| parent.metadata().name().to_string()); + } +} + +#[tokio::test(flavor = "current_thread")] +async fn metadata_worker_span_preserves_append_trace() { + let (store, mut updates) = ControlledThreadStore::new(); + let live_thread = create_live_thread(Arc::clone(&store)).await; + let parent_name = Arc::new(StdMutex::new(None)); + let _guard = tracing_subscriber::registry() + .with(WorkerSpanParentLayer { + parent_name: Arc::clone(&parent_name), + }) + .set_default(); + tracing::callsite::rebuild_interest_cache(); + + live_thread + .append_items(&[user_message("trace me")]) + .await + .expect("append should succeed"); + next_update(&mut updates) + .await + .completion + .send(Ok(())) + .expect("metadata worker should still be waiting"); + live_thread + .flush() + .await + .expect("flush should drain metadata"); + + assert_eq!( + parent_name.lock().expect("parent name lock").as_deref(), + Some("thread_store.live_append.observe_metadata") + ); +} + +#[tokio::test] +async fn failed_async_update_is_retried_and_reported_by_barriers() { + let (store, mut updates) = ControlledThreadStore::new(); + let live_thread = create_live_thread(Arc::clone(&store)).await; + + live_thread + .append_items(&[user_message("retry me")]) + .await + .expect("append should succeed independently of metadata"); + next_update(&mut updates) + .await + .completion + .send(Err(test_error("background failure"))) + .expect("metadata worker should still be waiting"); + wait_until_worker_stops(&live_thread).await; + + let retrying_flush = tokio::spawn({ + let live_thread = live_thread.clone(); + async move { live_thread.flush().await } + }); + next_update(&mut updates) + .await + .completion + .send(Err(ThreadStoreError::Conflict { + message: "barrier failure".to_string(), + })) + .expect("metadata worker should still be waiting"); + let error = retrying_flush + .await + .expect("flush task should not panic") + .expect_err("barrier should report the repeated failure"); + assert!(matches!( + error, + ThreadStoreError::Conflict { message } if message == "barrier failure" + )); + wait_until_worker_stops(&live_thread).await; + + let successful_flush = tokio::spawn({ + let live_thread = live_thread.clone(); + async move { live_thread.flush().await } + }); + next_update(&mut updates) + .await + .completion + .send(Ok(())) + .expect("metadata worker should still be waiting"); + successful_flush + .await + .expect("flush task should not panic") + .expect("later barrier should retry pending metadata successfully"); + assert_eq!(store.inner.calls().await.update_thread_metadata, 1); +} + +#[tokio::test] +async fn shutdown_waits_for_queued_metadata_before_closing_store() { + let (store, mut updates) = ControlledThreadStore::new(); + let live_thread = create_live_thread(Arc::clone(&store)).await; + + live_thread + .append_items(&[user_message("finish before shutdown")]) + .await + .expect("append should succeed"); + let attempt = next_update(&mut updates).await; + let mut shutdown = tokio::spawn({ + let live_thread = live_thread.clone(); + async move { live_thread.shutdown().await } + }); + assert!( + timeout(Duration::from_millis(20), &mut shutdown) + .await + .is_err(), + "shutdown should wait for the in-flight metadata update" + ); + assert_eq!(store.inner.calls().await.shutdown_thread, 0); + + attempt + .completion + .send(Ok(())) + .expect("metadata worker should still be waiting"); + shutdown + .await + .expect("shutdown task should not panic") + .expect("shutdown should succeed"); + assert_eq!(store.inner.calls().await.shutdown_thread, 1); +} + +#[tokio::test] +async fn discard_waits_for_in_flight_metadata_before_discarding_store() { + let (store, mut updates) = ControlledThreadStore::new(); + let live_thread = create_live_thread(Arc::clone(&store)).await; + + live_thread + .append_items(&[user_message("discard me")]) + .await + .expect("append should succeed"); + let attempt = next_update(&mut updates).await; + let mut discard = tokio::spawn({ + let live_thread = live_thread.clone(); + async move { live_thread.discard().await } + }); + assert!( + timeout(Duration::from_millis(20), &mut discard) + .await + .is_err(), + "discard should wait for the in-flight metadata update" + ); + assert_eq!(store.inner.calls().await.discard_thread, 0); + + attempt + .completion + .send(Ok(())) + .expect("metadata worker should still be waiting"); + discard + .await + .expect("discard task should not panic") + .expect("discard should succeed"); + assert_eq!(store.inner.calls().await.discard_thread, 1); +} + +async fn create_live_thread(store: Arc) -> LiveThread { + let thread_id = ThreadId::new(); + let thread_store: Arc = store; + LiveThread::create(thread_store, create_params(thread_id)) + .await + .expect("live thread should be created") +} + +fn create_params(thread_id: ThreadId) -> CreateThreadParams { + CreateThreadParams { + session_id: thread_id.into(), + thread_id, + extra_config: None, + forked_from_id: None, + parent_thread_id: None, + source: SessionSource::Exec, + thread_source: None, + originator: "test_originator".to_string(), + base_instructions: BaseInstructions::default(), + dynamic_tools: Vec::new(), + selected_capability_roots: Vec::new(), + multi_agent_version: None, + history_mode: ThreadHistoryMode::Legacy, + initial_window_id: Uuid::now_v7().to_string(), + metadata: ThreadPersistenceMetadata { + cwd: None, + model_provider: "test-provider".to_string(), + memory_mode: ThreadMemoryMode::Enabled, + }, + } +} + +fn user_message(message: &str) -> RolloutItem { + RolloutItem::EventMsg(EventMsg::UserMessage(UserMessageEvent { + message: message.to_string(), + ..Default::default() + })) +} + +async fn next_update( + updates: &mut mpsc::UnboundedReceiver, +) -> MetadataUpdateAttempt { + timeout(TEST_TIMEOUT, updates.recv()) + .await + .expect("metadata update should start") + .expect("metadata update channel should remain open") +} + +async fn wait_until_worker_stops(live_thread: &LiveThread) { + timeout(TEST_TIMEOUT, async { + loop { + if !live_thread.metadata_worker.state.lock().await.running { + return; + } + tokio::task::yield_now().await; + } + }) + .await + .expect("metadata worker should stop"); +} + +fn test_error(message: &str) -> ThreadStoreError { + ThreadStoreError::Internal { + message: message.to_string(), + } +} diff --git a/codex-rs/thread-store/src/local/mod.rs b/codex-rs/thread-store/src/local/mod.rs index 5939ea606e..3fb3b8064a 100644 --- a/codex-rs/thread-store/src/local/mod.rs +++ b/codex-rs/thread-store/src/local/mod.rs @@ -713,6 +713,10 @@ mod tests { .append_items(&[user_message_item("new live append")]) .await .expect("append after resume"); + live_thread + .flush() + .await + .expect("flush append-derived metadata"); let metadata = runtime .get_thread(thread_id) @@ -767,6 +771,10 @@ mod tests { .append_items(&[user_message_item("new external append")]) .await .expect("append after external resume"); + live_thread + .flush() + .await + .expect("flush append-derived metadata"); let metadata = runtime .get_thread(thread_id) diff --git a/codex-rs/thread-store/src/thread_metadata_sync.rs b/codex-rs/thread-store/src/thread_metadata_sync.rs index 1057715a4a..879f241cd6 100644 --- a/codex-rs/thread-store/src/thread_metadata_sync.rs +++ b/codex-rs/thread-store/src/thread_metadata_sync.rs @@ -50,6 +50,12 @@ pub(crate) struct PendingThreadMetadataPatch { generation: u64, } +impl PendingThreadMetadataPatch { + pub(crate) fn generation(&self) -> u64 { + self.generation + } +} + impl ThreadMetadataSync { pub(crate) async fn for_create(params: &CreateThreadParams) -> Self { let created_at = Utc::now();