From eeb82a156d1b3944dca4234c3043296529ec5837 Mon Sep 17 00:00:00 2001 From: Eric Traut Date: Mon, 17 Aug 2026 16:40:04 +0000 Subject: [PATCH] Dispatch queued messages written by other processes (#39034) ## Why Durable thread queues can be updated through another SQLite connection, but loaded idle threads need to notice those updates before they can dispatch the new messages. ## What changed - Track a durable revision for each thread queue, including backfilling existing queues and updating revisions on inserts, updates, and deletes. - Poll SQLite's data version and query revisions to find changed queues for loaded threads. - Wake idle threads with pending external messages, discover queued work when threads are loaded or resumed, and retry each thread independently so one blocked queue does not stall others. ## Testing - Cover cross-runtime queue writes, edits, independent dispatch, wake retries, and resumed threads. - Cover migration backfills and revision tracking for queue updates and deletes. GitOrigin-RevId: 906d902bbf6a760be206a1b2c08fd71427b63c46 --- codex-rs/ext/queue/src/lib.rs | 2 + codex-rs/ext/queue/src/service.rs | 175 ++++++++++++++++++ codex-rs/ext/queue/tests/queue_service.rs | 160 +++++++++++++++- .../0002_queued_thread_revisions.sql | 34 ++++ codex-rs/state/src/runtime/queued_items.rs | 58 +++++- .../state/src/runtime/queued_items_tests.rs | 105 +++++++++++ codex-rs/thread-store/src/queue_store.rs | 22 +++ 7 files changed, 552 insertions(+), 4 deletions(-) create mode 100644 codex-rs/state/queue_migrations/0002_queued_thread_revisions.sql diff --git a/codex-rs/ext/queue/src/lib.rs b/codex-rs/ext/queue/src/lib.rs index 1f7be82988..47fa504595 100644 --- a/codex-rs/ext/queue/src/lib.rs +++ b/codex-rs/ext/queue/src/lib.rs @@ -15,5 +15,7 @@ pub fn install(registry: &mut ExtensionRegistryBuilder, service: Arc, event_sink: Arc, dispatch_locks: Arc>>>>, + resumed_threads: Arc>>, } impl QueuedItemService { @@ -76,6 +81,165 @@ impl QueuedItemService { thread_manager, event_sink, dispatch_locks: Arc::new(StdMutex::new(HashMap::new())), + resumed_threads: Arc::new(StdMutex::new(HashSet::new())), + } + } + + // Check SQLite's inexpensive data version every 10 seconds, then use the + // durable revision index to discover only changed threads. Independent + // dispatch tasks keep a blocked or failed thread from starving other queues. + pub(crate) async fn watch_external_messages(service: Weak) { + let mut last_version = None; + let mut last_revision = 0; + let mut dispatches: HashMap> = HashMap::new(); + let mut interval = tokio::time::interval(Duration::from_secs(/*secs*/ 10)); + let mut manager_initialized = false; + let mut thread_created = None; + let mut newly_loaded_threads = HashSet::new(); + loop { + interval.tick().await; + let Some(service) = service.upgrade() else { + return; + }; + let Some(manager) = service.thread_manager.upgrade() else { + if manager_initialized { + return; + } + drop(service); + tokio::time::sleep(Duration::from_millis(/*millis*/ 1)).await; + interval.reset_immediately(); + continue; + }; + manager_initialized = true; + let thread_created = + thread_created.get_or_insert_with(|| manager.subscribe_thread_created()); + loop { + match thread_created.try_recv() { + Ok(thread_id) => { + newly_loaded_threads.insert(thread_id); + } + Err(TryRecvError::Lagged(_)) => { + newly_loaded_threads.extend(manager.list_thread_ids().await); + } + Err(TryRecvError::Empty) => break, + Err(TryRecvError::Closed) => return, + } + } + newly_loaded_threads.extend( + service + .resumed_threads + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .drain(), + ); + + let version = match service.queue.change_version().await { + Ok(version) => version, + Err(error) => { + tracing::warn!(%error, "failed to check queue change version"); + continue; + } + }; + let version_changed = last_version != Some(version); + if !version_changed && newly_loaded_threads.is_empty() { + continue; + } + + let thread_ids = manager.list_thread_ids().await; + let mut changes = Vec::new(); + let mut observed_revision = last_revision; + if version_changed { + match service + .queue + .changes_since(last_revision, &thread_ids) + .await + { + Ok(changed_threads) => { + if let Some((_, revision)) = changed_threads.last() { + observed_revision = *revision; + } + changes.extend(changed_threads); + } + Err(error) => { + tracing::warn!(%error, "failed to discover changed thread queues"); + continue; + } + } + } + if !newly_loaded_threads.is_empty() { + let created_threads = thread_ids + .iter() + .copied() + .filter(|thread_id| newly_loaded_threads.contains(thread_id)) + .collect::>(); + match service + .queue + .changes_since(/*revision*/ 0, &created_threads) + .await + { + Ok(changed_threads) => changes.extend(changed_threads), + Err(error) => { + tracing::warn!(%error, "failed to discover newly loaded thread queues"); + continue; + } + } + } + last_version = Some(version); + last_revision = observed_revision; + newly_loaded_threads.clear(); + dispatches.retain(|_, dispatch| !dispatch.is_finished()); + + let mut changed_threads = HashSet::new(); + for (thread_id, _) in changes { + if !changed_threads.insert(thread_id) { + continue; + } + service.emit_changed(thread_id); + if dispatches + .get(&thread_id) + .is_some_and(|dispatch| !dispatch.is_finished()) + { + continue; + } + let service = Arc::downgrade(&service); + let dispatch = tokio::spawn(async move { + loop { + { + let Some(service) = service.upgrade() else { + return; + }; + let Some(manager) = service.thread_manager.upgrade() else { + return; + }; + let Ok(thread) = manager.get_thread(thread_id).await else { + return; + }; + if matches!( + thread.agent_status().await, + AgentStatus::Running + | AgentStatus::Interrupted + | AgentStatus::Shutdown + | AgentStatus::NotFound + ) { + return; + } + match service + .queue + .list_page(thread_id, /*offset*/ 0, /*limit*/ 1) + .await + { + Ok(items) if items.is_empty() => return, + Ok(_) => service.wake_if_loaded(thread_id).await, + Err(error) => { + tracing::warn!(%thread_id, %error, "failed to check queued user input"); + } + } + } + tokio::time::sleep(Duration::from_secs(/*secs*/ 10)).await; + } + }); + dispatches.insert(thread_id, dispatch); + } } } @@ -362,6 +526,17 @@ impl ThreadLifecycleContributor for QueuedItemService where C: Send + Sync + 'static, { + fn on_thread_resume<'a>(&'a self, input: ThreadResumeInput<'a>) -> ExtensionFuture<'a, ()> { + Box::pin(async move { + if let Ok(thread_id) = ThreadId::from_string(input.thread_store.level_id()) { + self.resumed_threads + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .insert(thread_id); + } + }) + } + fn on_thread_idle<'a>(&'a self, input: ThreadIdleInput<'a>) -> ExtensionFuture<'a, ()> { Box::pin(async move { if input.cause == ThreadIdleCause::Interrupted { diff --git a/codex-rs/ext/queue/tests/queue_service.rs b/codex-rs/ext/queue/tests/queue_service.rs index db63540b4c..875dae3292 100644 --- a/codex-rs/ext/queue/tests/queue_service.rs +++ b/codex-rs/ext/queue/tests/queue_service.rs @@ -8,6 +8,7 @@ use std::time::Duration; use anyhow::Context; use codex_core::NotSubmittedReason; use codex_core::StartIfIdleSubmission; +use codex_core::StartThreadOptions; use codex_core::TurnInput; use codex_core::TurnInputRequest; use codex_core::TurnInputSubmission; @@ -21,6 +22,7 @@ use codex_extension_api::NoopExtensionEventSink; use codex_extension_api::ThreadIdleCause; use codex_extension_api::ThreadIdleInput; use codex_extension_api::ThreadLifecycleContributor; +use codex_extension_api::ThreadResumeInput; use codex_protocol::ThreadId; use codex_protocol::items::TurnItem; use codex_protocol::models::ImageDetail; @@ -45,6 +47,7 @@ use core_test_support::streaming_sse::start_streaming_sse_server; use core_test_support::test_codex::TestCodex; use core_test_support::test_codex::test_codex; use core_test_support::wait_for_event_match; +use core_test_support::wait_for_event_with_timeout; use pretty_assertions::assert_eq; use tempfile::TempDir; use tokio::sync::oneshot; @@ -73,11 +76,32 @@ impl ExtensionEventSink for RecordingEventSink { } #[derive(Default)] -struct InstalledQueue(OnceLock>); +struct InstalledQueue { + service: OnceLock>, + skip_next_idle: Mutex>, +} impl ThreadLifecycleContributor for InstalledQueue { + fn on_thread_resume<'a>(&'a self, input: ThreadResumeInput<'a>) -> ExtensionFuture<'a, ()> { + match self.service.get() { + Some(service) => >::on_thread_resume(service.as_ref(), input), + None => Box::pin(async {}), + } + } + fn on_thread_idle<'a>(&'a self, input: ThreadIdleInput<'a>) -> ExtensionFuture<'a, ()> { - match self.0.get() { + if self + .skip_next_idle + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .take_if(|thread_id| thread_id.to_string() == input.thread_store.level_id()) + .is_some() + { + return Box::pin(async {}); + } + match self.service.get() { Some(service) => >::on_thread_idle(service.as_ref(), input), @@ -105,7 +129,7 @@ fn install_registered_queue( Arc::downgrade(&test.thread_manager), Arc::new(NoopExtensionEventSink), )); - assert!(installed.0.set(Arc::clone(&service)).is_ok()); + assert!(installed.service.set(Arc::clone(&service)).is_ok()); Ok(service) } @@ -485,6 +509,136 @@ async fn registered_queue_lifecycle_starts_messages_in_fifo_order() -> anyhow::R Ok(()) } +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn externally_changed_queues_dispatch_independently_and_retry_failed_wakes() +-> anyhow::Result<()> { + let server = start_mock_server().await; + let model_responses = responses::mount_sse_sequence( + &server, + [ + "independent-queued-turn", + "external-queued-turn", + "resumed-queued-turn", + ] + .into_iter() + .map(responses::sse_completed) + .collect(), + ) + .await; + let (installed, extensions) = registered_queue_extensions(); + let test = test_codex() + .with_extensions(extensions) + .with_config(|config| config.include_environment_context = false) + .build_with_auto_env(&server) + .await?; + let thread_id = test.session_configured.thread_id; + let queue = install_registered_queue(&test, installed.as_ref())?; + let independent_thread = test + .thread_manager + .start_thread(StartThreadOptions::new(test.config.clone())) + .await?; + let external_runtime = StateRuntime::init( + test.codex + .state_db() + .context("state runtime unavailable")? + .sqlite() + .clone(), + "test-provider".to_string(), + ) + .await?; + let external_queue = QueuedItemService::new( + Arc::new(LocalQueueStore::new(external_runtime)), + Weak::new(), + Arc::new(NoopExtensionEventSink), + ); + let mut watcher_extensions = ExtensionRegistryBuilder::::new(); + codex_queue_extension::install(&mut watcher_extensions, Arc::clone(&queue)); + + tokio::time::sleep(Duration::from_secs(/*secs*/ 1)).await; + assert!(model_responses.requests().is_empty()); + *installed + .skip_next_idle + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) = Some(thread_id); + let first = external_queue + .enqueue(thread_id, user_input("written by another process")) + .await?; + let updated = queue + .update( + thread_id, + first.id, + user_input("locally edited external message"), + ) + .await? + .context("external queue item disappeared")?; + external_queue + .enqueue( + independent_thread.thread_id, + user_input("independent thread"), + ) + .await?; + + wait_for_event_with_timeout( + independent_thread.thread.as_ref(), + |event| matches!(event, EventMsg::TurnComplete(_)), + Duration::from_secs(/*secs*/ 25), + ) + .await; + assert_eq!(1, model_responses.requests().len()); + assert_eq!(vec![updated], queue.list(thread_id).await?); + + wait_for_event_with_timeout( + test.codex.as_ref(), + |event| matches!(event, EventMsg::TurnComplete(_)), + Duration::from_secs(/*secs*/ 25), + ) + .await; + tokio::time::sleep(Duration::from_secs(/*secs*/ 11)).await; + + assert!(queue.list(thread_id).await?.is_empty()); + assert!(queue.list(independent_thread.thread_id).await?.is_empty()); + + let rollout_path = test.codex.rollout_path().context("rollout path missing")?; + test.codex.shutdown_and_wait().await?; + test.thread_manager.remove_thread(&thread_id).await; + external_queue + .enqueue(thread_id, user_input("queued before ordinary resume")) + .await?; + tokio::time::sleep(Duration::from_secs(/*secs*/ 11)).await; + let resumed = test + .thread_manager + .resume_thread_from_rollout( + test.config.clone(), + rollout_path, + test.thread_manager.auth_manager(), + /*parent_trace*/ None, + Default::default(), + ) + .await?; + wait_for_event_with_timeout( + resumed.thread.as_ref(), + |event| matches!(event, EventMsg::TurnComplete(_)), + Duration::from_secs(/*secs*/ 25), + ) + .await; + assert!(queue.list(thread_id).await?.is_empty()); + + let prompts = model_responses + .requests() + .into_iter() + .filter_map(|request| request.message_input_texts("user").pop()) + .collect::>(); + assert_eq!( + vec![ + "independent thread", + "locally edited external message", + "queued before ordinary resume", + ], + prompts + ); + Ok(()) +} + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn rejected_queue_messages_are_consumed_without_retrying_or_blocking_followups() -> anyhow::Result<()> { diff --git a/codex-rs/state/queue_migrations/0002_queued_thread_revisions.sql b/codex-rs/state/queue_migrations/0002_queued_thread_revisions.sql new file mode 100644 index 0000000000..355352c9cc --- /dev/null +++ b/codex-rs/state/queue_migrations/0002_queued_thread_revisions.sql @@ -0,0 +1,34 @@ +CREATE TABLE queued_thread_revisions ( + revision INTEGER PRIMARY KEY AUTOINCREMENT, + thread_id TEXT NOT NULL UNIQUE +); + +INSERT INTO queued_thread_revisions (thread_id) +SELECT DISTINCT thread_id FROM queued_items ORDER BY thread_id; + +CREATE TRIGGER queued_items_revision_after_insert +AFTER INSERT ON queued_items +BEGIN + INSERT INTO queued_thread_revisions (thread_id) + VALUES (NEW.thread_id) + ON CONFLICT(thread_id) DO UPDATE + SET revision = (SELECT COALESCE(MAX(revision), 0) + 1 FROM queued_thread_revisions); +END; + +CREATE TRIGGER queued_items_revision_after_update +AFTER UPDATE ON queued_items +BEGIN + INSERT INTO queued_thread_revisions (thread_id) + VALUES (NEW.thread_id) + ON CONFLICT(thread_id) DO UPDATE + SET revision = (SELECT COALESCE(MAX(revision), 0) + 1 FROM queued_thread_revisions); +END; + +CREATE TRIGGER queued_items_revision_after_delete +AFTER DELETE ON queued_items +BEGIN + INSERT INTO queued_thread_revisions (thread_id) + VALUES (OLD.thread_id) + ON CONFLICT(thread_id) DO UPDATE + SET revision = (SELECT COALESCE(MAX(revision), 0) + 1 FROM queued_thread_revisions); +END; diff --git a/codex-rs/state/src/runtime/queued_items.rs b/codex-rs/state/src/runtime/queued_items.rs index 85f0c1f016..3cbf32b6ee 100644 --- a/codex-rs/state/src/runtime/queued_items.rs +++ b/codex-rs/state/src/runtime/queued_items.rs @@ -1,23 +1,79 @@ use super::*; use crate::MAX_QUEUE_ITEMS; use crate::QueuedUserSubmissionRecord; +use sqlx::Connection; +use tokio::sync::Mutex; use uuid::Uuid; /// SQLite-backed persistence for durable, thread-scoped user messages. #[derive(Clone)] pub struct SqliteQueueStore { pool: Arc, + change_version_connection: Arc>>, } impl SqliteQueueStore { pub(crate) fn new(pool: Arc) -> Self { - Self { pool } + Self { + pool, + change_version_connection: Arc::new(Mutex::new(None)), + } } pub(crate) async fn close(&self) { + let connection = self.change_version_connection.lock().await.take(); + if let Some(connection) = connection + && let Err(error) = connection.close().await + { + tracing::warn!(%error, "failed to close queue change-version connection"); + } self.pool.close().await; } + /// Observe queue-database commits through one stable SQLite connection. + pub async fn change_version(&self) -> anyhow::Result { + let mut connection = Arc::clone(&self.change_version_connection) + .lock_owned() + .await; + if connection.is_none() { + *connection = Some(self.pool.acquire().await?.detach()); + } + let Some(connection) = connection.as_mut() else { + unreachable!("queue change-version connection was initialized"); + }; + Ok(sqlx::query_scalar("PRAGMA data_version") + .fetch_one(connection) + .await?) + } + + /// Return changed revisions only for the supplied loaded thread IDs. + pub async fn changes_since( + &self, + revision: i64, + thread_ids: &[ThreadId], + ) -> anyhow::Result> { + if thread_ids.is_empty() { + return Ok(Vec::new()); + } + + let mut query = QueryBuilder::::new( + "SELECT thread_id, revision FROM queued_thread_revisions WHERE revision > ", + ); + query.push_bind(revision).push(" AND thread_id IN ("); + let mut separated = query.separated(", "); + for thread_id in thread_ids { + separated.push_bind(thread_id.to_string()); + } + separated.push_unseparated(") ORDER BY revision"); + let rows = query + .build_query_as::<(String, i64)>() + .fetch_all(self.pool.as_ref()) + .await?; + rows.into_iter() + .map(|(thread_id, revision)| Ok((ThreadId::try_from(thread_id)?, revision))) + .collect() + } + pub async fn enqueue( &self, thread_id: ThreadId, diff --git a/codex-rs/state/src/runtime/queued_items_tests.rs b/codex-rs/state/src/runtime/queued_items_tests.rs index 1155832b1f..3f97eb3bde 100644 --- a/codex-rs/state/src/runtime/queued_items_tests.rs +++ b/codex-rs/state/src/runtime/queued_items_tests.rs @@ -1,8 +1,11 @@ use super::*; +use crate::migrations::QUEUE_MIGRATOR; use crate::runtime::test_support::test_thread_metadata; use crate::runtime::test_support::unique_temp_dir; use codex_utils_absolute_path::test_support::PathExt; use pretty_assertions::assert_eq; +use sqlx::migrate::Migrator; +use std::borrow::Cow; async fn runtime_with_thread() -> (Arc, ThreadId) { let home = unique_temp_dir(); @@ -40,6 +43,108 @@ async fn competing_runtimes_preserve_fifo_queue_order() { assert_eq!(expected, actual); } +#[tokio::test] +async fn migrating_existing_queue_backfills_thread_revisions() { + let home = unique_temp_dir(); + tokio::fs::create_dir_all(&home).await.unwrap(); + let sqlite = crate::SqliteConfig::new_for_testing(home.as_path().abs()); + let queue_path = sqlite.queue_db_path(); + let old_queue_migrator = Migrator { + migrations: Cow::Owned(vec![QUEUE_MIGRATOR.migrations[0].clone()]), + ignore_missing: false, + locking: true, + no_tx: false, + table_name: QUEUE_MIGRATOR.table_name.clone(), + create_schemas: QUEUE_MIGRATOR.create_schemas.clone(), + }; + let pool = sqlite.open_read_write_pool(&queue_path).await.unwrap(); + old_queue_migrator.run(&pool).await.unwrap(); + + let thread_id = ThreadId::new(); + let queued = QueuedUserSubmissionRecord { + id: Uuid::now_v7().to_string(), + thread_id, + payload: r#"{"existing":true}"#.to_string(), + }; + sqlx::query( + "INSERT INTO queued_items + (id, thread_id, payload_json, queue_order, created_at_ms, updated_at_ms) + VALUES (?, ?, ?, 0, 0, 0)", + ) + .bind(&queued.id) + .bind(thread_id.to_string()) + .bind(&queued.payload) + .execute(&pool) + .await + .unwrap(); + pool.close().await; + + let runtime = StateRuntime::init(sqlite, "test-provider".to_string()) + .await + .unwrap(); + let queue = runtime.thread_queue(); + assert_eq!( + vec![(thread_id, 1)], + queue + .changes_since(/*revision*/ 0, &[thread_id]) + .await + .unwrap() + ); + assert_eq!( + vec![queued], + queue + .list_page(thread_id, /*offset*/ 0, /*limit*/ 1) + .await + .unwrap() + ); +} + +#[tokio::test] +async fn queue_revisions_identify_changed_threads_after_updates_and_deletions() { + let (runtime, thread_id) = runtime_with_thread().await; + let queue = runtime.thread_queue(); + let first = queue.enqueue(thread_id, r#"{"first":true}"#).await.unwrap(); + let first_revision = queue + .changes_since(/*revision*/ 0, &[thread_id]) + .await + .unwrap()[0] + .1; + queue + .update(thread_id, &first.id, r#"{"updated":true}"#) + .await + .unwrap(); + let updated_revision = queue + .changes_since(first_revision, &[thread_id]) + .await + .unwrap()[0] + .1; + let other_thread_id = ThreadId::new(); + queue + .enqueue(other_thread_id, r#"{"other":true}"#) + .await + .unwrap(); + let newly_loaded_changes = queue + .changes_since(/*revision*/ 0, &[other_thread_id]) + .await + .unwrap(); + assert_eq!( + vec![(thread_id, updated_revision), newly_loaded_changes[0]], + queue + .changes_since(first_revision, &[thread_id, other_thread_id]) + .await + .unwrap() + ); + assert!(queue.delete(thread_id, &first.id).await.unwrap()); + assert!( + queue + .changes_since(updated_revision, &[thread_id]) + .await + .unwrap() + .iter() + .any(|(changed_thread, _)| *changed_thread == thread_id) + ); +} + #[tokio::test] async fn fifo_dispatch_preserves_edits_reordering_and_pagination() { let (runtime, thread_id) = runtime_with_thread().await; diff --git a/codex-rs/thread-store/src/queue_store.rs b/codex-rs/thread-store/src/queue_store.rs index 606afb1014..74fc2fc970 100644 --- a/codex-rs/thread-store/src/queue_store.rs +++ b/codex-rs/thread-store/src/queue_store.rs @@ -12,6 +12,16 @@ use crate::ThreadStoreFuture; /// Storage-neutral persistence for ordered, thread-scoped user messages. pub trait QueueStore: Send + Sync { + /// Return a stable revision that changes when another connection updates the queue. + fn change_version(&self) -> ThreadStoreFuture<'_, i64>; + + /// Return changed, loaded thread IDs and their durable revisions after `revision`. + fn changes_since<'a>( + &'a self, + revision: i64, + thread_ids: &'a [ThreadId], + ) -> ThreadStoreFuture<'a, Vec<(ThreadId, i64)>>; + fn enqueue( &self, thread_id: ThreadId, @@ -72,6 +82,18 @@ where } impl QueueStore for LocalQueueStore { + fn change_version(&self) -> ThreadStoreFuture<'_, i64> { + queue_future(self.queue().change_version()) + } + + fn changes_since<'a>( + &'a self, + revision: i64, + thread_ids: &'a [ThreadId], + ) -> ThreadStoreFuture<'a, Vec<(ThreadId, i64)>> { + queue_future(self.queue().changes_since(revision, thread_ids)) + } + fn enqueue( &self, thread_id: ThreadId,