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
This commit is contained in:
Eric Traut
2026-08-17 16:40:04 +00:00
committed by copyberry
parent a8525b139c
commit eeb82a156d
7 changed files with 552 additions and 4 deletions

View File

@@ -15,5 +15,7 @@ pub fn install<C>(registry: &mut ExtensionRegistryBuilder<C>, service: Arc<Queue
where
C: Send + Sync + 'static,
{
let watcher = Arc::downgrade(&service);
registry.thread_lifecycle_contributor(service);
tokio::spawn(QueuedItemService::watch_external_messages(watcher));
}

View File

@@ -1,7 +1,9 @@
use std::collections::HashMap;
use std::collections::HashSet;
use std::sync::Arc;
use std::sync::Mutex as StdMutex;
use std::sync::Weak;
use std::time::Duration;
use codex_core::CodexThread;
use codex_core::StartIfIdleSubmission;
@@ -13,6 +15,7 @@ use codex_extension_api::ExtensionFuture;
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::error::CodexErr;
use codex_protocol::models::snapshot_local_user_input;
@@ -30,6 +33,7 @@ use codex_thread_store::ThreadStoreError;
use thiserror::Error;
use tokio::sync::Mutex;
use tokio::sync::OwnedMutexGuard;
use tokio::sync::broadcast::error::TryRecvError;
use uuid::Uuid;
/// One user message waiting to start on its thread.
@@ -63,6 +67,7 @@ pub struct QueuedItemService {
thread_manager: Weak<ThreadManager>,
event_sink: Arc<dyn ExtensionEventSink>,
dispatch_locks: Arc<StdMutex<HashMap<ThreadId, Weak<Mutex<()>>>>>,
resumed_threads: Arc<StdMutex<HashSet<ThreadId>>>,
}
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<Self>) {
let mut last_version = None;
let mut last_revision = 0;
let mut dispatches: HashMap<ThreadId, tokio::task::JoinHandle<()>> = 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::<Vec<_>>();
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<C> ThreadLifecycleContributor<C> 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 {

View File

@@ -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<Arc<QueuedItemService>>);
struct InstalledQueue {
service: OnceLock<Arc<QueuedItemService>>,
skip_next_idle: Mutex<Option<ThreadId>>,
}
impl ThreadLifecycleContributor<codex_core::config::Config> for InstalledQueue {
fn on_thread_resume<'a>(&'a self, input: ThreadResumeInput<'a>) -> ExtensionFuture<'a, ()> {
match self.service.get() {
Some(service) => <QueuedItemService as ThreadLifecycleContributor<
codex_core::config::Config,
>>::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) => <QueuedItemService as ThreadLifecycleContributor<
codex_core::config::Config,
>>::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::<codex_core::config::Config>::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::<Vec<_>>();
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<()> {

View File

@@ -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;

View File

@@ -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<SqlitePool>,
change_version_connection: Arc<Mutex<Option<SqliteConnection>>>,
}
impl SqliteQueueStore {
pub(crate) fn new(pool: Arc<SqlitePool>) -> 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<i64> {
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<Vec<(ThreadId, i64)>> {
if thread_ids.is_empty() {
return Ok(Vec::new());
}
let mut query = QueryBuilder::<Sqlite>::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,

View File

@@ -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<StateRuntime>, 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;

View File

@@ -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,