diff --git a/codex-rs/state/BUILD.bazel b/codex-rs/state/BUILD.bazel index 8b7c454fdb..acbf5b7f28 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/**", ]), 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..e16022dbc6 --- /dev/null +++ b/codex-rs/state/queue_migrations/0001_queued_items.sql @@ -0,0 +1,18 @@ +CREATE TABLE queued_items ( + queued_item_id TEXT PRIMARY KEY NOT NULL, + thread_id TEXT NOT NULL, + payload_jsonb BLOB NOT NULL, + queue_order INTEGER NOT NULL, + state TEXT NOT NULL CHECK (state IN ('pending', 'claimed', 'failed')), + claim_token TEXT, + failure_jsonb BLOB, + created_at_ms INTEGER NOT NULL, + updated_at_ms INTEGER NOT NULL, + CHECK ((state = 'claimed') = (claim_token IS NOT NULL)) +); + +CREATE INDEX queued_items_thread_state_order_idx + ON queued_items(thread_id, state, queue_order); + +CREATE 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 4b8e3ec3f5..9bd230fcd4 100644 --- a/codex-rs/state/src/lib.rs +++ b/codex-rs/state/src/lib.rs @@ -45,6 +45,9 @@ pub use model::BackfillStats; pub use model::BackfillStatus; pub use model::DirectionalThreadSpawnEdgeStatus; pub use model::ExtractionOutcome; +pub use model::QueuedItemClaim; +pub use model::QueuedItemRecord; +pub use model::QueuedItemState; pub use model::SortDirection; pub use model::SortKey; pub use model::Stage1JobClaim; @@ -61,6 +64,7 @@ pub use runtime::GoalAccountingOutcome; pub use runtime::GoalStore; pub use runtime::GoalUpdate; pub use runtime::MemoryStore; +pub use runtime::QueueStore; pub use runtime::RemoteControlEnrollmentRecord; pub use runtime::RuntimeDbBackup; pub use runtime::RuntimeDbPath; @@ -92,6 +96,7 @@ pub const SQLITE_HOME_ENV: &str = "CODEX_SQLITE_HOME"; pub const LOGS_DB_FILENAME: &str = "logs_2.sqlite"; pub const GOALS_DB_FILENAME: &str = "goals_1.sqlite"; pub const MEMORIES_DB_FILENAME: &str = "memories_1.sqlite"; +pub const QUEUE_DB_FILENAME: &str = "queue_1.sqlite"; pub const STATE_DB_FILENAME: &str = "state_5.sqlite"; /// Errors encountered during DB operations. Tags: [stage] diff --git a/codex-rs/state/src/migrations.rs b/codex-rs/state/src/migrations.rs index 0e12f17c45..971d007722 100644 --- a/codex-rs/state/src/migrations.rs +++ b/codex-rs/state/src/migrations.rs @@ -6,6 +6,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"); /// Allow an older Codex binary to open a database that has already been /// migrated by a newer binary running in parallel. @@ -39,3 +40,7 @@ pub(crate) fn runtime_goals_migrator() -> Migrator { pub(crate) fn runtime_memories_migrator() -> Migrator { runtime_migrator(&MEMORIES_MIGRATOR) } + +pub(crate) fn runtime_queue_migrator() -> Migrator { + runtime_migrator(&QUEUE_MIGRATOR) +} diff --git a/codex-rs/state/src/model/mod.rs b/codex-rs/state/src/model/mod.rs index 2c725e50ab..d75507222a 100644 --- a/codex-rs/state/src/model/mod.rs +++ b/codex-rs/state/src/model/mod.rs @@ -3,6 +3,7 @@ mod backfill_state; mod graph; mod log; mod memories; +mod queued_item; mod thread_goal; mod thread_metadata; @@ -24,6 +25,9 @@ pub use memories::Stage1JobClaim; pub use memories::Stage1JobClaimOutcome; pub use memories::Stage1Output; pub use memories::Stage1StartupClaimParams; +pub use queued_item::QueuedItemClaim; +pub use queued_item::QueuedItemRecord; +pub use queued_item::QueuedItemState; pub use thread_goal::ThreadGoal; pub use thread_goal::ThreadGoalStatus; pub use thread_metadata::Anchor; @@ -37,6 +41,7 @@ pub use thread_metadata::ThreadsPage; pub(crate) use agent_job::AgentJobItemRow; pub(crate) use agent_job::AgentJobRow; +pub(crate) use queued_item::QueuedItemRow; pub(crate) use thread_goal::ThreadGoalRow; pub(crate) use thread_metadata::ThreadRow; pub(crate) use thread_metadata::anchor_from_item; 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..c08a4df01e --- /dev/null +++ b/codex-rs/state/src/model/queued_item.rs @@ -0,0 +1,90 @@ +use anyhow::Result; +use anyhow::anyhow; +use chrono::DateTime; +use chrono::Utc; +use codex_protocol::ThreadId; +use sqlx::Row; +use sqlx::sqlite::SqliteRow; + +use super::epoch_millis_to_datetime; + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum QueuedItemState { + Pending, + Claimed, + Failed, +} + +impl TryFrom<&str> for QueuedItemState { + type Error = anyhow::Error; + + fn try_from(value: &str) -> Result { + match value { + "pending" => Ok(Self::Pending), + "claimed" => Ok(Self::Claimed), + "failed" => Ok(Self::Failed), + other => Err(anyhow!("unknown queued item state `{other}`")), + } + } +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct QueuedItemRecord { + pub queued_item_id: String, + pub thread_id: ThreadId, + pub payload_jsonb: Vec, + pub queue_order: i64, + pub state: QueuedItemState, + pub failure_jsonb: Option>, + pub created_at: DateTime, + pub updated_at: DateTime, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct QueuedItemClaim { + pub item: QueuedItemRecord, + pub claim_token: String, +} + +pub(crate) struct QueuedItemRow { + pub queued_item_id: String, + pub thread_id: String, + pub payload_jsonb: Vec, + pub queue_order: i64, + pub state: String, + pub failure_jsonb: Option>, + pub created_at_ms: i64, + pub updated_at_ms: i64, +} + +impl QueuedItemRow { + pub(crate) fn try_from_row(row: &SqliteRow) -> Result { + Ok(Self { + queued_item_id: row.try_get("queued_item_id")?, + thread_id: row.try_get("thread_id")?, + payload_jsonb: row.try_get("payload_jsonb")?, + queue_order: row.try_get("queue_order")?, + state: row.try_get("state")?, + failure_jsonb: row.try_get("failure_jsonb")?, + created_at_ms: row.try_get("created_at_ms")?, + updated_at_ms: row.try_get("updated_at_ms")?, + }) + } +} + +impl TryFrom for QueuedItemRecord { + type Error = anyhow::Error; + + fn try_from(row: QueuedItemRow) -> Result { + Ok(Self { + queued_item_id: row.queued_item_id, + thread_id: ThreadId::try_from(row.thread_id)?, + payload_jsonb: row.payload_jsonb, + queue_order: row.queue_order, + state: QueuedItemState::try_from(row.state.as_str())?, + failure_jsonb: row.failure_jsonb, + created_at: epoch_millis_to_datetime(row.created_at_ms)?, + updated_at: epoch_millis_to_datetime(row.updated_at_ms)?, + }) + } +} diff --git a/codex-rs/state/src/runtime.rs b/codex-rs/state/src/runtime.rs index 2937bf0c2e..ea0fcfc6c5 100644 --- a/codex-rs/state/src/runtime.rs +++ b/codex-rs/state/src/runtime.rs @@ -11,6 +11,7 @@ use crate::LogEntry; use crate::LogQuery; use crate::LogRow; use crate::MEMORIES_DB_FILENAME; +use crate::QUEUE_DB_FILENAME; use crate::STATE_DB_FILENAME; use crate::SortKey; use crate::ThreadMetadata; @@ -20,6 +21,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::model::AgentJobRow; use crate::model::ThreadRow; @@ -62,6 +64,7 @@ mod backfill; mod goals; mod logs; mod memories; +mod queued_items; mod recovery; mod remote_control; #[cfg(test)] @@ -73,6 +76,7 @@ pub use goals::GoalAccountingOutcome; pub use goals::GoalStore; pub use goals::GoalUpdate; pub use memories::MemoryStore; +pub use queued_items::QueueStore; pub use recovery::RuntimeDbBackup; pub use recovery::backup_runtime_db_for_fresh_start; pub use recovery::is_sqlite_corruption_error; @@ -138,7 +142,15 @@ const MEMORIES_DB: RuntimeDbSpec = RuntimeDbSpec { migrate_phase: "migrate_memories", }; -const RUNTIME_DBS: [RuntimeDbSpec; 4] = [STATE_DB, LOGS_DB, GOALS_DB, MEMORIES_DB]; +const QUEUE_DB: RuntimeDbSpec = RuntimeDbSpec { + label: "queue DB", + filename: QUEUE_DB_FILENAME, + kind: DbKind::Queue, + open_phase: "open_queue", + migrate_phase: "migrate_queue", +}; + +const RUNTIME_DBS: [RuntimeDbSpec; 5] = [STATE_DB, LOGS_DB, GOALS_DB, MEMORIES_DB, QUEUE_DB]; #[derive(Clone, Debug, Eq, PartialEq)] pub struct RuntimeDbPath { @@ -154,6 +166,7 @@ pub struct StateRuntime { logs_pool: Arc, thread_goals: GoalStore, memories: MemoryStore, + thread_queue: QueueStore, thread_updated_at_millis: Arc, } @@ -191,10 +204,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 = STATE_DB.path(codex_home.as_path()); let logs_path = LOGS_DB.path(codex_home.as_path()); let goals_path = GOALS_DB.path(codex_home.as_path()); let memories_path = MEMORIES_DB.path(codex_home.as_path()); + let queue_path = QUEUE_DB.path(codex_home.as_path()); let pool = match open_state_sqlite(&state_path, &state_migrator, telemetry_override).await { Ok(db) => Arc::new(db), Err(err) => { @@ -237,6 +252,21 @@ impl StateRuntime { return Err(err); } }; + let queue_pool = + match open_queue_sqlite(&queue_path, &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( @@ -252,6 +282,7 @@ impl StateRuntime { logs_pool.as_ref(), goals_pool.as_ref(), memories_pool.as_ref(), + queue_pool.as_ref(), ]) .await; return Err(err); @@ -277,6 +308,7 @@ impl StateRuntime { logs_pool.as_ref(), goals_pool.as_ref(), memories_pool.as_ref(), + queue_pool.as_ref(), ]) .await; return Err(err); @@ -286,6 +318,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: QueueStore::new(Arc::clone(&queue_pool)), pool, logs_pool, codex_home, @@ -316,12 +349,17 @@ impl StateRuntime { /// 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; self.pool.close().await; } + pub fn thread_queue(&self) -> &QueueStore { + &self.thread_queue + } + pub async fn clear_memory_data_in_sqlite_home(sqlite_home: &Path) -> anyhow::Result { let memories_path = MEMORIES_DB.path(sqlite_home); if !tokio::fs::try_exists(&memories_path).await? { @@ -392,6 +430,14 @@ async fn open_memories_sqlite( open_sqlite(path, migrator, MEMORIES_DB, telemetry_override).await } +async fn open_queue_sqlite( + path: &Path, + migrator: &Migrator, + telemetry_override: Option<&dyn DbTelemetry>, +) -> anyhow::Result { + open_sqlite(path, migrator, QUEUE_DB, telemetry_override).await +} + async fn open_sqlite( path: &Path, migrator: &Migrator, @@ -708,6 +754,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..d54ea696a8 --- /dev/null +++ b/codex-rs/state/src/runtime/queued_items.rs @@ -0,0 +1,348 @@ +use super::*; +use uuid::Uuid; + +#[derive(Clone)] +pub struct QueueStore { + pool: Arc, +} + +impl QueueStore { + 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: &[u8], + ) -> anyhow::Result { + let payload_json = std::str::from_utf8(payload_json)?; + let queued_item_id = Uuid::new_v4().to_string(); + let now_ms = datetime_to_epoch_millis(Utc::now()); + let row = sqlx::query( + r#" +INSERT INTO queued_items ( + queued_item_id, + thread_id, + payload_jsonb, + queue_order, + state, + failure_jsonb, + created_at_ms, + updated_at_ms +) +SELECT + ?, + ?, + jsonb(?), + COALESCE(MAX(queue_order), -1) + 1, + 'pending', + NULL, + ?, + ? +FROM queued_items +WHERE thread_id = ? +RETURNING + queued_item_id, + thread_id, + CAST(json(payload_jsonb) AS BLOB) AS payload_jsonb, + queue_order, + state, + CASE + WHEN failure_jsonb IS NULL THEN NULL + ELSE CAST(json(failure_jsonb) AS BLOB) + END AS failure_jsonb, + created_at_ms, + updated_at_ms + "#, + ) + .bind(queued_item_id) + .bind(thread_id.to_string()) + .bind(payload_json) + .bind(now_ms) + .bind(now_ms) + .bind(thread_id.to_string()) + .fetch_one(self.pool.as_ref()) + .await?; + + queued_item_from_row(&row) + } + + pub async fn list_page( + &self, + thread_id: ThreadId, + offset: usize, + limit: usize, + ) -> anyhow::Result> { + let rows = sqlx::query( + r#" +SELECT + queued_item_id, + thread_id, + CAST(json(payload_jsonb) AS BLOB) AS payload_jsonb, + queue_order, + state, + CASE + WHEN failure_jsonb IS NULL THEN NULL + ELSE CAST(json(failure_jsonb) AS BLOB) + END AS failure_jsonb, + created_at_ms, + updated_at_ms +FROM queued_items +WHERE thread_id = ? + AND state IN ('pending', 'failed') +ORDER BY queue_order ASC +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(queued_item_from_row).collect() + } + + pub async fn delete(&self, thread_id: ThreadId, queued_item_id: &str) -> anyhow::Result { + let result = sqlx::query( + r#" +DELETE FROM queued_items +WHERE thread_id = ? + AND queued_item_id = ? + AND state IN ('pending', 'failed') + "#, + ) + .bind(thread_id.to_string()) + .bind(queued_item_id) + .execute(self.pool.as_ref()) + .await?; + + Ok(result.rows_affected() > 0) + } + + pub async fn reorder( + &self, + thread_id: ThreadId, + ordered_ids: &[String], + ) -> anyhow::Result { + let mut transaction = self.pool.begin().await?; + let visible_rows: Vec<(String, i64)> = sqlx::query_as( + r#" +SELECT queued_item_id, queue_order +FROM queued_items +WHERE thread_id = ? + AND state IN ('pending', 'failed') +ORDER BY queue_order ASC + "#, + ) + .bind(thread_id.to_string()) + .fetch_all(transaction.as_mut()) + .await?; + + let visible_ids = visible_rows + .iter() + .map(|(queued_item_id, _)| queued_item_id.clone()) + .collect::>(); + let visible_queue_orders = visible_rows + .into_iter() + .map(|(_, queue_order)| queue_order) + .collect::>(); + let mut expected_ids = visible_ids; + expected_ids.sort(); + let mut requested_ids = ordered_ids.to_vec(); + requested_ids.sort(); + if expected_ids != requested_ids { + transaction.rollback().await?; + return Ok(false); + } + + let now_ms = datetime_to_epoch_millis(Utc::now()); + for (queue_order, queued_item_id) in visible_queue_orders.into_iter().zip(ordered_ids) { + sqlx::query( + r#" +UPDATE queued_items +SET queue_order = ?, updated_at_ms = ? +WHERE thread_id = ? + AND queued_item_id = ? + AND state IN ('pending', 'failed') + "#, + ) + .bind(queue_order) + .bind(now_ms) + .bind(thread_id.to_string()) + .bind(queued_item_id) + .execute(transaction.as_mut()) + .await?; + } + transaction.commit().await?; + Ok(true) + } + + /// Atomically claims the pending FIFO head. A failed or already claimed + /// head blocks later items. + pub async fn claim_next( + &self, + thread_id: ThreadId, + ) -> anyhow::Result> { + let now_ms = datetime_to_epoch_millis(Utc::now()); + let claim_token = Uuid::new_v4().to_string(); + let row = sqlx::query( + r#" +UPDATE queued_items +SET state = 'claimed', claim_token = ?, updated_at_ms = ? +WHERE queued_item_id = ( + SELECT queued_item_id + FROM queued_items + WHERE thread_id = ? + ORDER BY queue_order ASC + LIMIT 1 +) + AND state = 'pending' +RETURNING + queued_item_id, + thread_id, + CAST(json(payload_jsonb) AS BLOB) AS payload_jsonb, + queue_order, + state, + CASE + WHEN failure_jsonb IS NULL THEN NULL + ELSE CAST(json(failure_jsonb) AS BLOB) + END AS failure_jsonb, + created_at_ms, + updated_at_ms + "#, + ) + .bind(&claim_token) + .bind(now_ms) + .bind(thread_id.to_string()) + .fetch_optional(self.pool.as_ref()) + .await?; + + row.map(|row| { + Ok(crate::QueuedItemClaim { + item: queued_item_from_row(&row)?, + claim_token, + }) + }) + .transpose() + } + + pub async fn has_claimed_item(&self, thread_id: ThreadId) -> anyhow::Result { + sqlx::query_scalar( + "SELECT EXISTS(SELECT 1 FROM queued_items WHERE thread_id = ? AND state = 'claimed')", + ) + .bind(thread_id.to_string()) + .fetch_one(self.pool.as_ref()) + .await + .map_err(Into::into) + } + + pub async fn release_claim( + &self, + queued_item_id: &str, + claim_token: &str, + ) -> anyhow::Result { + let now_ms = datetime_to_epoch_millis(Utc::now()); + let result = sqlx::query( + r#" +UPDATE queued_items +SET state = 'pending', claim_token = NULL, failure_jsonb = NULL, updated_at_ms = ? +WHERE queued_item_id = ? AND state = 'claimed' AND claim_token = ? + "#, + ) + .bind(now_ms) + .bind(queued_item_id) + .bind(claim_token) + .execute(self.pool.as_ref()) + .await?; + Ok(result.rows_affected() > 0) + } + + pub async fn complete_claim( + &self, + queued_item_id: &str, + claim_token: &str, + ) -> anyhow::Result { + let result = sqlx::query( + "DELETE FROM queued_items \ + WHERE queued_item_id = ? AND state = 'claimed' AND claim_token = ?", + ) + .bind(queued_item_id) + .bind(claim_token) + .execute(self.pool.as_ref()) + .await?; + Ok(result.rows_affected() > 0) + } + + pub async fn fail_claim( + &self, + queued_item_id: &str, + claim_token: &str, + failure_json: &[u8], + ) -> anyhow::Result { + let failure_json = std::str::from_utf8(failure_json)?; + let now_ms = datetime_to_epoch_millis(Utc::now()); + let result = sqlx::query( + r#" +UPDATE queued_items +SET state = 'failed', claim_token = NULL, failure_jsonb = jsonb(?), updated_at_ms = ? +WHERE queued_item_id = ? AND state = 'claimed' AND claim_token = ? + "#, + ) + .bind(failure_json) + .bind(now_ms) + .bind(queued_item_id) + .bind(claim_token) + .execute(self.pool.as_ref()) + .await?; + Ok(result.rows_affected() > 0) + } + + pub async fn recover_claims_as_failed_before( + &self, + thread_id: ThreadId, + stale_before_ms: i64, + failure_json: &[u8], + ) -> anyhow::Result { + let failure_json = std::str::from_utf8(failure_json)?; + let now_ms = datetime_to_epoch_millis(Utc::now()); + let result = sqlx::query( + r#" +UPDATE queued_items +SET state = 'failed', claim_token = NULL, failure_jsonb = jsonb(?), updated_at_ms = ? +WHERE thread_id = ? + AND state = 'claimed' + AND updated_at_ms <= ? + "#, + ) + .bind(failure_json) + .bind(now_ms) + .bind(thread_id.to_string()) + .bind(stale_before_ms) + .execute(self.pool.as_ref()) + .await?; + Ok(result.rows_affected()) + } + + pub(crate) async fn delete_thread_queue(&self, thread_id: ThreadId) -> anyhow::Result { + let result = sqlx::query("DELETE FROM queued_items WHERE thread_id = ?") + .bind(thread_id.to_string()) + .execute(self.pool.as_ref()) + .await?; + Ok(result.rows_affected() > 0) + } +} + +fn queued_item_from_row(row: &sqlx::sqlite::SqliteRow) -> anyhow::Result { + crate::model::QueuedItemRow::try_from_row(row)?.try_into() +} + +#[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..098be91f5e --- /dev/null +++ b/codex-rs/state/src/runtime/queued_items_tests.rs @@ -0,0 +1,247 @@ +use super::*; +use crate::runtime::test_support::test_thread_metadata; +use crate::runtime::test_support::unique_temp_dir; +use pretty_assertions::assert_eq; + +async fn runtime_with_thread() -> (Arc, ThreadId) { + let codex_home = unique_temp_dir(); + let runtime = StateRuntime::init(codex_home.clone(), "test-provider".to_string()) + .await + .expect("state runtime"); + let thread_id = ThreadId::new(); + runtime + .upsert_thread(&test_thread_metadata( + codex_home.as_path(), + thread_id, + codex_home.clone(), + )) + .await + .expect("insert thread"); + (runtime, thread_id) +} + +#[tokio::test] +async fn claim_next_has_one_winner_across_runtime_instances() { + let (runtime, thread_id) = runtime_with_thread().await; + let second_runtime = StateRuntime::init( + runtime.codex_home().to_path_buf(), + "test-provider".to_string(), + ) + .await + .expect("second state runtime"); + runtime + .thread_queue() + .enqueue(thread_id, br#"{"input":[]}"#) + .await + .expect("enqueue"); + + let (first, second) = tokio::join!( + runtime.thread_queue().claim_next(thread_id), + second_runtime.thread_queue().claim_next(thread_id), + ); + assert_eq!( + 1, + [first.expect("first claim"), second.expect("second claim")] + .into_iter() + .flatten() + .count() + ); +} + +#[tokio::test] +async fn stale_claim_owner_cannot_mutate_reclaimed_item() { + let (runtime, thread_id) = runtime_with_thread().await; + let queue = runtime.thread_queue(); + let item = queue + .enqueue(thread_id, br#"{"input":["hello"]}"#) + .await + .expect("enqueue"); + let first_claim = queue.claim_next(thread_id).await.unwrap().unwrap(); + assert!( + queue + .release_claim(&item.queued_item_id, &first_claim.claim_token) + .await + .unwrap() + ); + let second_claim = queue.claim_next(thread_id).await.unwrap().unwrap(); + + assert!( + !queue + .complete_claim(&item.queued_item_id, &first_claim.claim_token) + .await + .unwrap() + ); + assert!( + queue + .complete_claim(&item.queued_item_id, &second_claim.claim_token) + .await + .unwrap() + ); + assert!( + queue + .list_page(thread_id, /*offset*/ 0, /*limit*/ 1) + .await + .unwrap() + .is_empty() + ); +} + +#[tokio::test] +async fn failed_head_blocks_later_pending_item_until_removed() { + let (runtime, thread_id) = runtime_with_thread().await; + let queue = runtime.thread_queue(); + let first = queue.enqueue(thread_id, br#"{"n":1}"#).await.unwrap(); + let second = queue.enqueue(thread_id, br#"{"n":2}"#).await.unwrap(); + let claim = queue.claim_next(thread_id).await.unwrap().unwrap(); + queue + .fail_claim( + &first.queued_item_id, + &claim.claim_token, + br#"{"message":"nope"}"#, + ) + .await + .unwrap(); + + assert_eq!(None, queue.claim_next(thread_id).await.unwrap()); + assert!( + queue + .delete(thread_id, &first.queued_item_id) + .await + .unwrap() + ); + assert_eq!( + second.queued_item_id, + queue + .claim_next(thread_id) + .await + .unwrap() + .unwrap() + .item + .queued_item_id + ); +} + +#[tokio::test] +async fn recovery_only_marks_stale_claims_failed() { + let (runtime, thread_id) = runtime_with_thread().await; + let queue = runtime.thread_queue(); + let item = queue.enqueue(thread_id, br#"{"n":1}"#).await.unwrap(); + queue.claim_next(thread_id).await.unwrap().unwrap(); + + assert_eq!( + 0, + queue + .recover_claims_as_failed_before( + thread_id, + /*stale_before_ms*/ 0, + br#"{"message":"claim interrupted"}"#, + ) + .await + .unwrap() + ); + assert_eq!( + 1, + queue + .recover_claims_as_failed_before( + thread_id, + /*stale_before_ms*/ i64::MAX, + br#"{"message":"claim interrupted"}"#, + ) + .await + .unwrap() + ); + let visible = queue + .list_page(thread_id, /*offset*/ 0, /*limit*/ 1) + .await + .unwrap(); + assert_eq!(item.queued_item_id, visible[0].queued_item_id); + assert_eq!(crate::QueuedItemState::Failed, visible[0].state); + assert_eq!(None, queue.claim_next(thread_id).await.unwrap()); +} + +#[tokio::test] +async fn reorder_requires_every_visible_item() { + let (runtime, thread_id) = runtime_with_thread().await; + let queue = runtime.thread_queue(); + let first = queue.enqueue(thread_id, br#"{"n":1}"#).await.unwrap(); + let second = queue.enqueue(thread_id, br#"{"n":2}"#).await.unwrap(); + assert!( + !queue + .reorder(thread_id, std::slice::from_ref(&first.queued_item_id)) + .await + .unwrap() + ); + assert!( + queue + .reorder( + thread_id, + &[second.queued_item_id.clone(), first.queued_item_id.clone()], + ) + .await + .unwrap() + ); + let reordered = queue + .list_page(thread_id, /*offset*/ 0, /*limit*/ 2) + .await + .unwrap(); + assert_eq!( + vec![second.queued_item_id, first.queued_item_id], + reordered + .into_iter() + .map(|item| item.queued_item_id) + .collect::>() + ); +} + +#[tokio::test] +async fn list_page_preserves_fifo_order_and_payload() { + let (runtime, thread_id) = runtime_with_thread().await; + let queue = runtime.thread_queue(); + queue.enqueue(thread_id, br#"{"n":1}"#).await.unwrap(); + let second = queue.enqueue(thread_id, br#"{"n":2}"#).await.unwrap(); + let third = queue.enqueue(thread_id, br#"{"n":3}"#).await.unwrap(); + + let page = queue + .list_page(thread_id, /*offset*/ 1, /*limit*/ 2) + .await + .unwrap(); + assert_eq!( + vec![second.queued_item_id, third.queued_item_id], + page.iter() + .map(|item| item.queued_item_id.clone()) + .collect::>() + ); + assert_eq!(br#"{"n":2}"#, page[0].payload_jsonb.as_slice()); +} + +#[tokio::test] +async fn json_payloads_are_bound_as_text() { + let (runtime, thread_id) = runtime_with_thread().await; + let queue = runtime.thread_queue(); + queue.enqueue(thread_id, b"3456").await.unwrap(); + + let visible = queue + .list_page(thread_id, /*offset*/ 0, /*limit*/ 1) + .await + .unwrap(); + assert_eq!(b"3456", visible[0].payload_jsonb.as_slice()); +} + +#[tokio::test] +async fn deleting_thread_deletes_its_queue() { + let (runtime, thread_id) = runtime_with_thread().await; + runtime + .thread_queue() + .enqueue(thread_id, br#"{"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() + ); +} diff --git a/codex-rs/state/src/runtime/threads.rs b/codex-rs/state/src/runtime/threads.rs index 6cebf27e87..13913630c3 100644 --- a/codex-rs/state/src/runtime/threads.rs +++ b/codex-rs/state/src/runtime/threads.rs @@ -930,6 +930,7 @@ ON CONFLICT(id) DO UPDATE SET .await?; self.memories.delete_thread_memory(*thread_id).await?; self.thread_goals.delete_thread_goal(*thread_id).await?; + self.thread_queue.delete_thread_queue(*thread_id).await?; } let now = Utc::now().timestamp(); diff --git a/codex-rs/state/src/telemetry.rs b/codex-rs/state/src/telemetry.rs index 1a9e4f9951..f9045acb9d 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, } impl DbKind { @@ -50,6 +51,7 @@ impl DbKind { Self::Logs => "logs", Self::Goals => "goals", Self::Memories => "memories", + Self::Queue => "queue", } } }