mirror of
https://github.com/openai/codex.git
synced 2026-08-23 13:09:46 +00:00
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:
@@ -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));
|
||||
}
|
||||
|
||||
@@ -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 {
|
||||
|
||||
@@ -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<()> {
|
||||
|
||||
@@ -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;
|
||||
@@ -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,
|
||||
|
||||
@@ -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;
|
||||
|
||||
@@ -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,
|
||||
|
||||
Reference in New Issue
Block a user