diff --git a/codex-rs/state/BUILD.bazel b/codex-rs/state/BUILD.bazel index ce4900a135..80537a0f7d 100644 --- a/codex-rs/state/BUILD.bazel +++ b/codex-rs/state/BUILD.bazel @@ -7,6 +7,7 @@ codex_rust_crate( "logs_migrations/**", "memory_migrations/**", "migrations/**", + "queue_migrations/**", "thread_history_migrations/**", ]), crate_name = "codex_state", diff --git a/codex-rs/state/queue_migrations/0001_queued_items.sql b/codex-rs/state/queue_migrations/0001_queued_items.sql new file mode 100644 index 0000000000..186cafefc0 --- /dev/null +++ b/codex-rs/state/queue_migrations/0001_queued_items.sql @@ -0,0 +1,11 @@ +CREATE TABLE queued_items ( + id TEXT PRIMARY KEY NOT NULL, + thread_id TEXT NOT NULL, + payload_json TEXT NOT NULL, + queue_order INTEGER NOT NULL, + created_at_ms INTEGER NOT NULL, + updated_at_ms INTEGER NOT NULL +); + +CREATE UNIQUE INDEX queued_items_thread_order_idx + ON queued_items(thread_id, queue_order); diff --git a/codex-rs/state/src/lib.rs b/codex-rs/state/src/lib.rs index 5e5233cb84..c852fa2355 100644 --- a/codex-rs/state/src/lib.rs +++ b/codex-rs/state/src/lib.rs @@ -23,6 +23,7 @@ pub use model::LogEntry; pub use model::LogQuery; pub use model::LogRow; pub use model::Phase2JobClaimOutcome; +pub use model::QueuedUserSubmissionRecord; /// Preferred entrypoint: owns configuration and metrics. pub use runtime::StateRuntime; pub use sqlite::SqliteConfig; @@ -65,6 +66,7 @@ pub use runtime::GoalUpdate; pub use runtime::MemoryStore; pub use runtime::RemoteControlEnrollmentRecord; pub use runtime::RuntimeDbBackup; +pub use runtime::SqliteQueueStore; pub use runtime::ThreadFilterOptions; pub use runtime::backup_runtime_db_for_fresh_start; pub use runtime::is_sqlite_corruption_error; @@ -80,6 +82,9 @@ pub use telemetry::install_process_db_telemetry; pub use telemetry::record_backfill_gate; pub use telemetry::record_fallback; +/// Maximum number of pending user submissions permitted for one thread. +pub const MAX_QUEUE_ITEMS: usize = 100; + /// Stable UUIDv7 identifying the built-in pinned thread section. pub const PINNED_THREAD_SECTION_ID: &str = "01984de2-8f74-7c91-a3b2-5c5e937cf318"; diff --git a/codex-rs/state/src/migrations.rs b/codex-rs/state/src/migrations.rs index 995ea4b642..07d36588f0 100644 --- a/codex-rs/state/src/migrations.rs +++ b/codex-rs/state/src/migrations.rs @@ -7,6 +7,7 @@ pub(crate) static STATE_MIGRATOR: Migrator = sqlx::migrate!("./migrations"); pub(crate) static LOGS_MIGRATOR: Migrator = sqlx::migrate!("./logs_migrations"); pub(crate) static GOALS_MIGRATOR: Migrator = sqlx::migrate!("./goals_migrations"); pub(crate) static MEMORIES_MIGRATOR: Migrator = sqlx::migrate!("./memory_migrations"); +pub(crate) static QUEUE_MIGRATOR: Migrator = sqlx::migrate!("./queue_migrations"); pub(crate) static THREAD_HISTORY_MIGRATOR: Migrator = sqlx::migrate!("./thread_history_migrations"); /// Allow an older Codex binary to open a database that has already been @@ -42,6 +43,10 @@ pub(crate) fn runtime_memories_migrator() -> Migrator { runtime_migrator(&MEMORIES_MIGRATOR) } +pub(crate) fn runtime_queue_migrator() -> Migrator { + runtime_migrator(&QUEUE_MIGRATOR) +} + // The paginated history projector will call this when it takes ownership of opening the database. #[allow(dead_code)] pub(crate) fn runtime_thread_history_migrator() -> Migrator { diff --git a/codex-rs/state/src/model/mod.rs b/codex-rs/state/src/model/mod.rs index d5b041214c..76ad9db254 100644 --- a/codex-rs/state/src/model/mod.rs +++ b/codex-rs/state/src/model/mod.rs @@ -2,6 +2,7 @@ mod backfill_state; mod graph; mod log; mod memories; +mod queued_item; mod thread_goal; mod thread_metadata; @@ -16,6 +17,7 @@ pub use memories::Stage1JobClaim; pub use memories::Stage1JobClaimOutcome; pub use memories::Stage1Output; pub use memories::Stage1StartupClaimParams; +pub use queued_item::QueuedUserSubmissionRecord; pub use thread_goal::ThreadGoal; pub use thread_goal::ThreadGoalStatus; pub use thread_metadata::Anchor; diff --git a/codex-rs/state/src/model/queued_item.rs b/codex-rs/state/src/model/queued_item.rs new file mode 100644 index 0000000000..ae7ec0c0bd --- /dev/null +++ b/codex-rs/state/src/model/queued_item.rs @@ -0,0 +1,22 @@ +use anyhow::Result; +use codex_protocol::ThreadId; +use sqlx::Row; +use sqlx::sqlite::SqliteRow; + +/// One durable, ordered user submission for a thread. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct QueuedUserSubmissionRecord { + pub id: String, + pub thread_id: ThreadId, + pub payload: String, +} + +impl QueuedUserSubmissionRecord { + pub(crate) fn try_from_row(row: &SqliteRow) -> Result { + Ok(Self { + id: row.try_get("id")?, + thread_id: ThreadId::try_from(row.try_get::("thread_id")?)?, + payload: row.try_get("payload_json")?, + }) + } +} diff --git a/codex-rs/state/src/runtime.rs b/codex-rs/state/src/runtime.rs index 433b16b84a..29528f3ce9 100644 --- a/codex-rs/state/src/runtime.rs +++ b/codex-rs/state/src/runtime.rs @@ -10,6 +10,7 @@ use crate::apply_rollout_item; use crate::migrations::runtime_goals_migrator; use crate::migrations::runtime_logs_migrator; use crate::migrations::runtime_memories_migrator; +use crate::migrations::runtime_queue_migrator; use crate::migrations::runtime_state_migrator; use crate::migrations::runtime_thread_history_migrator; use crate::model::ThreadRow; @@ -43,6 +44,7 @@ mod external_agent_config_imports; mod goals; mod logs; mod memories; +mod queued_items; mod recovery; mod remote_control; #[cfg(test)] @@ -60,6 +62,7 @@ pub use goals::GoalAccountingOutcome; pub use goals::GoalStore; pub use goals::GoalUpdate; pub use memories::MemoryStore; +pub use queued_items::SqliteQueueStore; pub use recovery::RuntimeDbBackup; pub(super) use recovery::RuntimeDbInitError; pub use recovery::backup_runtime_db_for_fresh_start; @@ -87,6 +90,7 @@ pub struct StateRuntime { logs_pool: Arc, thread_goals: GoalStore, memories: MemoryStore, + thread_queue: SqliteQueueStore, thread_updated_at_millis: Arc, thread_recency_at_millis: Arc, } @@ -121,10 +125,12 @@ impl StateRuntime { let logs_migrator = runtime_logs_migrator(); let goals_migrator = runtime_goals_migrator(); let memories_migrator = runtime_memories_migrator(); + let queue_migrator = runtime_queue_migrator(); let state_path = sqlite.state_db_path(); let logs_path = sqlite.logs_db_path(); let goals_path = sqlite.goals_db_path(); let memories_path = sqlite.memories_db_path(); + let queue_path = sqlite.queue_db_path(); let pool = match sqlite .open_state_db(&state_migrator, telemetry_override) .await @@ -171,6 +177,23 @@ impl StateRuntime { return Err(err); } }; + let queue_pool = match sqlite + .open_queue_db(&queue_migrator, telemetry_override) + .await + { + Ok(db) => Arc::new(db), + Err(err) => { + warn!("failed to open queue db at {}: {err}", queue_path.display()); + close_sqlite_pools(&[ + pool.as_ref(), + logs_pool.as_ref(), + goals_pool.as_ref(), + memories_pool.as_ref(), + ]) + .await; + return Err(err); + } + }; let started = Instant::now(); let backfill_state_result = ensure_backfill_state_row_in_pool(pool.as_ref()).await; crate::telemetry::record_init_result( @@ -186,6 +209,7 @@ impl StateRuntime { logs_pool.as_ref(), goals_pool.as_ref(), memories_pool.as_ref(), + queue_pool.as_ref(), ]) .await; return Err(err); @@ -214,6 +238,7 @@ impl StateRuntime { logs_pool.as_ref(), goals_pool.as_ref(), memories_pool.as_ref(), + queue_pool.as_ref(), ]) .await; return Err(err); @@ -224,6 +249,7 @@ impl StateRuntime { let runtime = Arc::new(Self { thread_goals: GoalStore::new(Arc::clone(&goals_pool)), memories: MemoryStore::new(Arc::clone(&memories_pool), Arc::clone(&pool)), + thread_queue: SqliteQueueStore::new(queue_pool), pool, logs_pool, sqlite, @@ -253,8 +279,14 @@ impl StateRuntime { &self.memories } + /// Return the durable, SQLite-backed user-message queue. + pub fn thread_queue(&self) -> &SqliteQueueStore { + &self.thread_queue + } + /// Close all SQLite pools and wait for outstanding pool workers to exit. pub async fn close(&self) { + self.thread_queue.close().await; self.memories.close().await; self.thread_goals.close().await; self.logs_pool.close().await; @@ -510,6 +542,8 @@ mod tests { "migrate_goals", "open_memories", "migrate_memories", + "open_queue", + "migrate_queue", "ensure_backfill_state", "post_init_query", ] diff --git a/codex-rs/state/src/runtime/queued_items.rs b/codex-rs/state/src/runtime/queued_items.rs new file mode 100644 index 0000000000..571d5eaa94 --- /dev/null +++ b/codex-rs/state/src/runtime/queued_items.rs @@ -0,0 +1,159 @@ +use super::*; +use crate::MAX_QUEUE_ITEMS; +use crate::QueuedUserSubmissionRecord; +use uuid::Uuid; + +/// SQLite-backed persistence for durable, thread-scoped user messages. +#[derive(Clone)] +pub struct SqliteQueueStore { + pool: Arc, +} + +impl SqliteQueueStore { + pub(crate) fn new(pool: Arc) -> Self { + Self { pool } + } + + pub(crate) async fn close(&self) { + self.pool.close().await; + } + + pub async fn enqueue( + &self, + thread_id: ThreadId, + payload_json: &str, + ) -> anyhow::Result { + let now_ms = datetime_to_epoch_millis(Utc::now()); + let row = sqlx::query( + "INSERT INTO queued_items ( + id, thread_id, payload_json, queue_order, + created_at_ms, updated_at_ms + ) + SELECT ?, ?, ?, + COALESCE((SELECT MAX(queue_order) FROM queued_items WHERE thread_id = ?), -1) + 1, + ?, ? + WHERE (SELECT COUNT(*) FROM queued_items WHERE thread_id = ?) < ? + RETURNING id, thread_id, payload_json", + ) + .bind(Uuid::now_v7().to_string()) + .bind(thread_id.to_string()) + .bind(payload_json) + .bind(thread_id.to_string()) + .bind(now_ms) + .bind(now_ms) + .bind(thread_id.to_string()) + .bind(i64::try_from(MAX_QUEUE_ITEMS)?) + .fetch_one(self.pool.as_ref()) + .await?; + QueuedUserSubmissionRecord::try_from_row(&row) + } + + pub async fn list_page( + &self, + thread_id: ThreadId, + offset: usize, + limit: usize, + ) -> anyhow::Result> { + let rows = sqlx::query( + "SELECT id, thread_id, payload_json + FROM queued_items + WHERE thread_id = ? + ORDER BY queue_order LIMIT ? OFFSET ?", + ) + .bind(thread_id.to_string()) + .bind(i64::try_from(limit)?) + .bind(i64::try_from(offset)?) + .fetch_all(self.pool.as_ref()) + .await?; + rows.iter() + .map(QueuedUserSubmissionRecord::try_from_row) + .collect() + } + + pub async fn update( + &self, + thread_id: ThreadId, + item_id: &str, + payload_json: &str, + ) -> anyhow::Result> { + let row = sqlx::query( + "UPDATE queued_items + SET payload_json = ?, updated_at_ms = ? + WHERE thread_id = ? AND id = ? + RETURNING id, thread_id, payload_json", + ) + .bind(payload_json) + .bind(datetime_to_epoch_millis(Utc::now())) + .bind(thread_id.to_string()) + .bind(item_id) + .fetch_optional(self.pool.as_ref()) + .await?; + row.as_ref() + .map(QueuedUserSubmissionRecord::try_from_row) + .transpose() + } + + pub async fn delete(&self, thread_id: ThreadId, item_id: &str) -> anyhow::Result { + Ok( + sqlx::query("DELETE FROM queued_items WHERE thread_id = ? AND id = ?") + .bind(thread_id.to_string()) + .bind(item_id) + .execute(self.pool.as_ref()) + .await? + .rows_affected() + > 0, + ) + } + + pub async fn reorder(&self, thread_id: ThreadId, ordered_ids: &[String]) -> anyhow::Result<()> { + let mut transaction = self.pool.begin().await?; + let rows: Vec<(String, i64)> = sqlx::query_as( + "SELECT id, queue_order FROM queued_items + WHERE thread_id = ? ORDER BY queue_order", + ) + .bind(thread_id.to_string()) + .fetch_all(transaction.as_mut()) + .await?; + let mut expected_ids = rows.iter().map(|(id, _)| id.clone()).collect::>(); + let mut requested_ids = ordered_ids.to_vec(); + expected_ids.sort(); + requested_ids.sort(); + if expected_ids != requested_ids { + transaction.rollback().await?; + return Err(std::io::Error::new( + std::io::ErrorKind::InvalidInput, + "queue reorder must include every queued item exactly once", + ) + .into()); + } + let now_ms = datetime_to_epoch_millis(Utc::now()); + let max_queue_order = rows.last().map_or(-1, |(_, queue_order)| *queue_order); + for (index, item_id) in ordered_ids.iter().enumerate() { + sqlx::query( + "UPDATE queued_items SET queue_order = ?, updated_at_ms = ? + WHERE thread_id = ? AND id = ?", + ) + .bind(max_queue_order + i64::try_from(index)? + 1) + .bind(now_ms) + .bind(thread_id.to_string()) + .bind(item_id) + .execute(transaction.as_mut()) + .await?; + } + transaction.commit().await?; + Ok(()) + } + + pub(crate) async fn delete_thread_queue(&self, thread_id: ThreadId) -> anyhow::Result { + Ok(sqlx::query("DELETE FROM queued_items WHERE thread_id = ?") + .bind(thread_id.to_string()) + .execute(self.pool.as_ref()) + .await? + .rows_affected() + > 0) + } +} + +#[cfg(test)] +#[path = "queued_items_tests.rs"] +mod tests; diff --git a/codex-rs/state/src/runtime/queued_items_tests.rs b/codex-rs/state/src/runtime/queued_items_tests.rs new file mode 100644 index 0000000000..1155832b1f --- /dev/null +++ b/codex-rs/state/src/runtime/queued_items_tests.rs @@ -0,0 +1,173 @@ +use super::*; +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; + +async fn runtime_with_thread() -> (Arc, ThreadId) { + let home = unique_temp_dir(); + let runtime = StateRuntime::init( + crate::SqliteConfig::new_for_testing(home.as_path().abs()), + "test-provider".to_string(), + ) + .await + .expect("state runtime"); + let thread_id = ThreadId::new(); + let metadata = test_thread_metadata(home.as_path(), thread_id, home.clone()); + runtime.upsert_thread(&metadata).await.unwrap(); + (runtime, thread_id) +} + +#[tokio::test] +async fn competing_runtimes_preserve_fifo_queue_order() { + let (runtime, thread_id) = runtime_with_thread().await; + let other = StateRuntime::init(runtime.sqlite().clone(), "test-provider".to_string()) + .await + .unwrap(); + let queue = runtime.thread_queue(); + let other_queue = other.thread_queue(); + let (first, second) = tokio::join!( + queue.enqueue(thread_id, r#"{"first":true}"#), + other_queue.enqueue(thread_id, r#"{"second":true}"#), + ); + let mut expected = vec![first.unwrap(), second.unwrap()]; + expected.sort_by(|first, second| first.id.cmp(&second.id)); + let mut actual = queue + .list_page(thread_id, /*offset*/ 0, /*limit*/ 2) + .await + .unwrap(); + actual.sort_by(|first, second| first.id.cmp(&second.id)); + assert_eq!(expected, actual); +} + +#[tokio::test] +async fn fifo_dispatch_preserves_edits_reordering_and_pagination() { + let (runtime, thread_id) = runtime_with_thread().await; + let queue = runtime.thread_queue(); + let first = queue.enqueue(thread_id, r#"{"n":1}"#).await.unwrap(); + let second = queue.enqueue(thread_id, r#"{"n":2}"#).await.unwrap(); + let third = queue.enqueue(thread_id, r#"{"n":3}"#).await.unwrap(); + + let updated = queue + .update(thread_id, &first.id, r#"{"n":"edited"}"#) + .await + .unwrap() + .unwrap(); + assert_eq!(first.id, updated.id); + let error = queue + .reorder(thread_id, std::slice::from_ref(&first.id)) + .await + .unwrap_err(); + assert_eq!( + std::io::ErrorKind::InvalidInput, + error.downcast_ref::().unwrap().kind() + ); + + let ordered_ids = vec![third.id, first.id, second.id]; + queue.reorder(thread_id, &ordered_ids).await.unwrap(); + + let items = queue + .list_page(thread_id, /*offset*/ 0, /*limit*/ 3) + .await + .unwrap(); + let page = queue + .list_page(thread_id, /*offset*/ 1, /*limit*/ 1) + .await + .unwrap(); + assert_eq!(vec![items[1].clone()], page); + assert_eq!(r#"{"n":"edited"}"#, items[1].payload); + + for item in items { + assert!(queue.delete(thread_id, &item.id).await.unwrap()); + } + assert!( + queue + .list_page(thread_id, /*offset*/ 0, /*limit*/ 1) + .await + .unwrap() + .is_empty() + ); +} + +#[tokio::test] +async fn queue_operations_cannot_mutate_another_threads_messages() { + let (runtime, thread_id) = runtime_with_thread().await; + let queue = runtime.thread_queue(); + let first = queue.enqueue(thread_id, r#"{"n":1}"#).await.unwrap(); + let other_thread_id = ThreadId::new(); + let other = queue.enqueue(other_thread_id, r#"{"n":2}"#).await.unwrap(); + let other_id = &other.id; + + assert_eq!( + None, + queue + .update(thread_id, other_id, r#"{"n":3}"#) + .await + .unwrap() + ); + assert!(!queue.delete(thread_id, other_id).await.unwrap()); + assert!( + queue + .reorder(thread_id, std::slice::from_ref(other_id)) + .await + .is_err() + ); + let (items, other_items) = tokio::join!( + queue.list_page(thread_id, /*offset*/ 0, /*limit*/ 1), + queue.list_page(other_thread_id, /*offset*/ 0, /*limit*/ 1), + ); + assert_eq!( + (vec![first], vec![other]), + (items.unwrap(), other_items.unwrap()) + ); +} + +#[tokio::test] +async fn deleting_a_thread_removes_its_queue() { + let (runtime, thread_id) = runtime_with_thread().await; + runtime + .thread_queue() + .enqueue(thread_id, r#"{"n":1}"#) + .await + .unwrap(); + + assert_eq!(1, runtime.delete_thread(thread_id).await.unwrap()); + assert!( + runtime + .thread_queue() + .list_page(thread_id, /*offset*/ 0, /*limit*/ 1) + .await + .unwrap() + .is_empty() + ); +} + +#[tokio::test] +async fn concurrent_inserts_enforce_the_queue_limit() { + let (runtime, thread_id) = runtime_with_thread().await; + let other = StateRuntime::init(runtime.sqlite().clone(), "test-provider".to_string()) + .await + .unwrap(); + + for _ in 0..MAX_QUEUE_ITEMS - 1 { + runtime + .thread_queue() + .enqueue(thread_id, r#"{"n":1}"#) + .await + .unwrap(); + } + let (first, second) = tokio::join!( + runtime.thread_queue().enqueue(thread_id, r#"{"n":2}"#), + other.thread_queue().enqueue(thread_id, r#"{"n":3}"#), + ); + assert_ne!(first.is_ok(), second.is_ok()); + assert_eq!( + MAX_QUEUE_ITEMS, + runtime + .thread_queue() + .list_page(thread_id, /*offset*/ 0, /*limit*/ MAX_QUEUE_ITEMS) + .await + .unwrap() + .len() + ); +} diff --git a/codex-rs/state/src/runtime/threads.rs b/codex-rs/state/src/runtime/threads.rs index d65c404070..100c2f8ff1 100644 --- a/codex-rs/state/src/runtime/threads.rs +++ b/codex-rs/state/src/runtime/threads.rs @@ -1064,6 +1064,7 @@ ON CONFLICT(id) DO UPDATE SET .bind(thread_id_string) .execute(self.logs_pool.as_ref()) .await?; + self.thread_queue.delete_thread_queue(*thread_id).await?; self.memories.delete_thread_memory(*thread_id).await?; self.thread_goals.delete_thread_goal(*thread_id).await?; } diff --git a/codex-rs/state/src/sqlite.rs b/codex-rs/state/src/sqlite.rs index 5fd1cf5357..4bb65c6e35 100644 --- a/codex-rs/state/src/sqlite.rs +++ b/codex-rs/state/src/sqlite.rs @@ -29,6 +29,7 @@ use std::time::Instant; const LOGS_DB_FILENAME: &str = "logs_2.sqlite"; const GOALS_DB_FILENAME: &str = "goals_1.sqlite"; const MEMORIES_DB_FILENAME: &str = "memories_1.sqlite"; +const QUEUE_DB_FILENAME: &str = "queue_1.sqlite"; const STATE_DB_FILENAME: &str = "state_5.sqlite"; const THREAD_HISTORY_DB_FILENAME: &str = "thread_history_1.sqlite"; @@ -79,6 +80,14 @@ const MEMORIES_DB: RuntimeDbSpec = RuntimeDbSpec { migrate_phase: "migrate_memories", }; +const QUEUE_DB: RuntimeDbSpec = RuntimeDbSpec { + label: "queue DB", + filename: QUEUE_DB_FILENAME, + kind: DbKind::Queue, + open_phase: "open_queue", + migrate_phase: "migrate_queue", +}; + const THREAD_HISTORY_DB: RuntimeDbSpec = RuntimeDbSpec { label: "thread history DB", filename: THREAD_HISTORY_DB_FILENAME, @@ -87,8 +96,14 @@ const THREAD_HISTORY_DB: RuntimeDbSpec = RuntimeDbSpec { migrate_phase: "migrate_thread_history", }; -const RUNTIME_DBS: [RuntimeDbSpec; 5] = - [STATE_DB, LOGS_DB, GOALS_DB, MEMORIES_DB, THREAD_HISTORY_DB]; +const RUNTIME_DBS: [RuntimeDbSpec; 6] = [ + STATE_DB, + LOGS_DB, + GOALS_DB, + MEMORIES_DB, + QUEUE_DB, + THREAD_HISTORY_DB, +]; #[derive(Clone, Debug, Eq, PartialEq)] pub struct RuntimeDbPath { @@ -135,6 +150,11 @@ impl SqliteConfig { MEMORIES_DB.path(self.home()) } + /// Return the path to the durable user-message queue database. + pub fn queue_db_path(&self) -> PathBuf { + QUEUE_DB.path(self.home()) + } + /// Return the path to the paginated thread-history database. pub fn thread_history_db_path(&self) -> PathBuf { THREAD_HISTORY_DB.path(self.home()) @@ -190,6 +210,15 @@ impl SqliteConfig { .await } + pub(super) async fn open_queue_db( + &self, + migrator: &Migrator, + telemetry_override: Option<&dyn DbTelemetry>, + ) -> anyhow::Result { + self.open_runtime_db(QUEUE_DB, migrator, telemetry_override) + .await + } + pub(super) async fn open_thread_history_db( &self, migrator: &Migrator, diff --git a/codex-rs/state/src/telemetry.rs b/codex-rs/state/src/telemetry.rs index 57a6c71314..1b14fc7d44 100644 --- a/codex-rs/state/src/telemetry.rs +++ b/codex-rs/state/src/telemetry.rs @@ -41,6 +41,7 @@ pub(crate) enum DbKind { Logs, Goals, Memories, + Queue, ThreadHistory, } @@ -51,6 +52,7 @@ impl DbKind { Self::Logs => "logs", Self::Goals => "goals", Self::Memories => "memories", + Self::Queue => "queue", Self::ThreadHistory => "thread_history", } } diff --git a/codex-rs/thread-store/src/lib.rs b/codex-rs/thread-store/src/lib.rs index 9cff38392c..817523bce8 100644 --- a/codex-rs/thread-store/src/lib.rs +++ b/codex-rs/thread-store/src/lib.rs @@ -8,11 +8,14 @@ mod error; mod in_memory; mod live_thread; mod local; +mod queue_store; mod store; mod thread_metadata_sync; mod thread_sections; mod types; +pub use codex_state::MAX_QUEUE_ITEMS; +pub use codex_state::QueuedUserSubmissionRecord; pub use error::ThreadStoreError; pub use error::ThreadStoreResult; pub use in_memory::InMemoryThreadStore; @@ -21,6 +24,8 @@ pub use live_thread::LiveThread; pub use live_thread::LiveThreadInitGuard; pub use local::LocalThreadStore; pub use local::LocalThreadStoreConfig; +pub use queue_store::LocalQueueStore; +pub use queue_store::QueueStore; pub use store::ThreadStore; pub use store::ThreadStoreFuture; pub use thread_sections::CreateThreadSectionParams; diff --git a/codex-rs/thread-store/src/queue_store.rs b/codex-rs/thread-store/src/queue_store.rs new file mode 100644 index 0000000000..6de7150b63 --- /dev/null +++ b/codex-rs/thread-store/src/queue_store.rs @@ -0,0 +1,134 @@ +use std::fmt::Display; +use std::future::Future; + +use codex_protocol::ThreadId; +use codex_rollout::StateDbHandle; +use codex_state::QueuedUserSubmissionRecord; +use codex_state::SqliteQueueStore; + +use crate::MAX_QUEUE_ITEMS; +use crate::ThreadStoreError; +use crate::ThreadStoreFuture; + +/// Storage-neutral persistence for ordered, thread-scoped user messages. +pub trait QueueStore: Send + Sync { + fn enqueue( + &self, + thread_id: ThreadId, + payload: String, + ) -> ThreadStoreFuture<'_, QueuedUserSubmissionRecord>; + + fn list_page( + &self, + thread_id: ThreadId, + offset: usize, + limit: usize, + ) -> ThreadStoreFuture<'_, Vec>; + + fn update( + &self, + thread_id: ThreadId, + item_id: String, + payload: String, + ) -> ThreadStoreFuture<'_, Option>; + + fn delete(&self, thread_id: ThreadId, item_id: String) -> ThreadStoreFuture<'_, bool>; + + /// Atomically replace queue order with every current item ID exactly once. + /// + /// Returns [`ThreadStoreError::InvalidRequest`] when `item_ids` is not a + /// permutation of the complete queue. + fn reorder(&self, thread_id: ThreadId, item_ids: Vec) -> ThreadStoreFuture<'_, ()>; +} + +/// Adapts the local state runtime to the shared queue-storage interface. +#[derive(Clone)] +pub struct LocalQueueStore { + state_db: StateDbHandle, +} + +impl LocalQueueStore { + pub fn new(state_db: StateDbHandle) -> Self { + Self { state_db } + } + + fn queue(&self) -> &SqliteQueueStore { + self.state_db.thread_queue() + } +} + +fn queue_future<'a, T, E>( + future: impl Future> + Send + 'a, +) -> ThreadStoreFuture<'a, T> +where + T: Send + 'a, + E: Display + Send + 'a, +{ + Box::pin(async move { + future.await.map_err(|error| ThreadStoreError::Internal { + message: format!("queue storage failed: {error}"), + }) + }) +} + +impl QueueStore for LocalQueueStore { + fn enqueue( + &self, + thread_id: ThreadId, + payload: String, + ) -> ThreadStoreFuture<'_, QueuedUserSubmissionRecord> { + Box::pin(async move { + self.queue() + .enqueue(thread_id, &payload) + .await + .map_err(|error| match error.downcast_ref::() { + Some(sqlx::Error::RowNotFound) => ThreadStoreError::InvalidRequest { + message: format!("queue cannot contain more than {MAX_QUEUE_ITEMS} items"), + }, + _ => ThreadStoreError::Internal { + message: format!("queue storage failed: {error}"), + }, + }) + }) + } + + fn list_page( + &self, + thread_id: ThreadId, + offset: usize, + limit: usize, + ) -> ThreadStoreFuture<'_, Vec> { + queue_future(self.queue().list_page(thread_id, offset, limit)) + } + + fn update( + &self, + thread_id: ThreadId, + item_id: String, + payload: String, + ) -> ThreadStoreFuture<'_, Option> { + queue_future(async move { self.queue().update(thread_id, &item_id, &payload).await }) + } + + fn delete(&self, thread_id: ThreadId, item_id: String) -> ThreadStoreFuture<'_, bool> { + queue_future(async move { self.queue().delete(thread_id, &item_id).await }) + } + + fn reorder(&self, thread_id: ThreadId, item_ids: Vec) -> ThreadStoreFuture<'_, ()> { + Box::pin(async move { + self.queue() + .reorder(thread_id, &item_ids) + .await + .map_err(|error| match error.downcast_ref::() { + Some(error) if error.kind() == std::io::ErrorKind::InvalidInput => { + ThreadStoreError::InvalidRequest { + message: error.to_string(), + } + } + _ => ThreadStoreError::Internal { + message: format!("queue storage failed: {error}"), + }, + }) + }) + } +}