mirror of
https://github.com/openai/codex.git
synced 2026-09-04 15:08:45 +00:00
feat: add durable user message queue storage
This commit is contained in:
@@ -7,6 +7,7 @@ codex_rust_crate(
|
||||
"logs_migrations/**",
|
||||
"memory_migrations/**",
|
||||
"migrations/**",
|
||||
"queue_migrations/**",
|
||||
]),
|
||||
crate_name = "codex_state",
|
||||
)
|
||||
|
||||
18
codex-rs/state/queue_migrations/0001_queued_items.sql
Normal file
18
codex-rs/state/queue_migrations/0001_queued_items.sql
Normal file
@@ -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);
|
||||
@@ -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]
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
|
||||
@@ -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;
|
||||
|
||||
90
codex-rs/state/src/model/queued_item.rs
Normal file
90
codex-rs/state/src/model/queued_item.rs
Normal file
@@ -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<Self> {
|
||||
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<u8>,
|
||||
pub queue_order: i64,
|
||||
pub state: QueuedItemState,
|
||||
pub failure_jsonb: Option<Vec<u8>>,
|
||||
pub created_at: DateTime<Utc>,
|
||||
pub updated_at: DateTime<Utc>,
|
||||
}
|
||||
|
||||
#[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<u8>,
|
||||
pub queue_order: i64,
|
||||
pub state: String,
|
||||
pub failure_jsonb: Option<Vec<u8>>,
|
||||
pub created_at_ms: i64,
|
||||
pub updated_at_ms: i64,
|
||||
}
|
||||
|
||||
impl QueuedItemRow {
|
||||
pub(crate) fn try_from_row(row: &SqliteRow) -> Result<Self> {
|
||||
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<QueuedItemRow> for QueuedItemRecord {
|
||||
type Error = anyhow::Error;
|
||||
|
||||
fn try_from(row: QueuedItemRow) -> Result<Self> {
|
||||
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)?,
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -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<sqlx::SqlitePool>,
|
||||
thread_goals: GoalStore,
|
||||
memories: MemoryStore,
|
||||
thread_queue: QueueStore,
|
||||
thread_updated_at_millis: Arc<AtomicI64>,
|
||||
}
|
||||
|
||||
@@ -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<bool> {
|
||||
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<SqlitePool> {
|
||||
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",
|
||||
]
|
||||
|
||||
348
codex-rs/state/src/runtime/queued_items.rs
Normal file
348
codex-rs/state/src/runtime/queued_items.rs
Normal file
@@ -0,0 +1,348 @@
|
||||
use super::*;
|
||||
use uuid::Uuid;
|
||||
|
||||
#[derive(Clone)]
|
||||
pub struct QueueStore {
|
||||
pool: Arc<SqlitePool>,
|
||||
}
|
||||
|
||||
impl QueueStore {
|
||||
pub(crate) fn new(pool: Arc<SqlitePool>) -> 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<crate::QueuedItemRecord> {
|
||||
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<Vec<crate::QueuedItemRecord>> {
|
||||
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<bool> {
|
||||
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<bool> {
|
||||
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::<Vec<_>>();
|
||||
let visible_queue_orders = visible_rows
|
||||
.into_iter()
|
||||
.map(|(_, queue_order)| queue_order)
|
||||
.collect::<Vec<_>>();
|
||||
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<Option<crate::QueuedItemClaim>> {
|
||||
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<bool> {
|
||||
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<bool> {
|
||||
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<bool> {
|
||||
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<bool> {
|
||||
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<u64> {
|
||||
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<bool> {
|
||||
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::QueuedItemRecord> {
|
||||
crate::model::QueuedItemRow::try_from_row(row)?.try_into()
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
#[path = "queued_items_tests.rs"]
|
||||
mod tests;
|
||||
247
codex-rs/state/src/runtime/queued_items_tests.rs
Normal file
247
codex-rs/state/src/runtime/queued_items_tests.rs
Normal file
@@ -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<StateRuntime>, 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::<Vec<_>>()
|
||||
);
|
||||
}
|
||||
|
||||
#[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::<Vec<_>>()
|
||||
);
|
||||
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()
|
||||
);
|
||||
}
|
||||
@@ -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();
|
||||
|
||||
@@ -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",
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user