mirror of
https://github.com/openai/codex.git
synced 2026-10-19 17:28:46 +00:00
Add durable per-thread user submission queues (#36952)
## What changed - Add a storage-neutral `QueueStore` interface backed by a dedicated SQLite database for ordered, thread-scoped user submissions. - Support enqueueing, paginated listing, editing, deletion, and atomic reordering, with a limit of 100 pending items per thread. - Remove queued submissions when their thread is deleted. ## Testing - Cover concurrent inserts, ordering and pagination, edits, deletion, cross-thread isolation, queue limits, and thread cleanup. GitOrigin-RevId: 0a8e95ac71137afc6f3a039d4b5fc8603110c891
This commit is contained in:
@@ -7,6 +7,7 @@ codex_rust_crate(
|
||||
"logs_migrations/**",
|
||||
"memory_migrations/**",
|
||||
"migrations/**",
|
||||
"queue_migrations/**",
|
||||
"thread_history_migrations/**",
|
||||
]),
|
||||
crate_name = "codex_state",
|
||||
|
||||
@@ -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);
|
||||
@@ -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";
|
||||
|
||||
|
||||
@@ -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 {
|
||||
|
||||
@@ -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;
|
||||
|
||||
@@ -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<Self> {
|
||||
Ok(Self {
|
||||
id: row.try_get("id")?,
|
||||
thread_id: ThreadId::try_from(row.try_get::<String, _>("thread_id")?)?,
|
||||
payload: row.try_get("payload_json")?,
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -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<sqlx::SqlitePool>,
|
||||
thread_goals: GoalStore,
|
||||
memories: MemoryStore,
|
||||
thread_queue: SqliteQueueStore,
|
||||
thread_updated_at_millis: Arc<AtomicI64>,
|
||||
thread_recency_at_millis: Arc<AtomicI64>,
|
||||
}
|
||||
@@ -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",
|
||||
]
|
||||
|
||||
@@ -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<SqlitePool>,
|
||||
}
|
||||
|
||||
impl SqliteQueueStore {
|
||||
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: &str,
|
||||
) -> anyhow::Result<QueuedUserSubmissionRecord> {
|
||||
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<Vec<QueuedUserSubmissionRecord>> {
|
||||
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<Option<QueuedUserSubmissionRecord>> {
|
||||
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<bool> {
|
||||
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::<Vec<_>>();
|
||||
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<bool> {
|
||||
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;
|
||||
@@ -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<StateRuntime>, 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::<std::io::Error>().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()
|
||||
);
|
||||
}
|
||||
@@ -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?;
|
||||
}
|
||||
|
||||
@@ -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<SqlitePool> {
|
||||
self.open_runtime_db(QUEUE_DB, migrator, telemetry_override)
|
||||
.await
|
||||
}
|
||||
|
||||
pub(super) async fn open_thread_history_db(
|
||||
&self,
|
||||
migrator: &Migrator,
|
||||
|
||||
@@ -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",
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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;
|
||||
|
||||
@@ -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<QueuedUserSubmissionRecord>>;
|
||||
|
||||
fn update(
|
||||
&self,
|
||||
thread_id: ThreadId,
|
||||
item_id: String,
|
||||
payload: String,
|
||||
) -> ThreadStoreFuture<'_, Option<QueuedUserSubmissionRecord>>;
|
||||
|
||||
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<String>) -> 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<Output = Result<T, E>> + 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::<sqlx::Error>() {
|
||||
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<QueuedUserSubmissionRecord>> {
|
||||
queue_future(self.queue().list_page(thread_id, offset, limit))
|
||||
}
|
||||
|
||||
fn update(
|
||||
&self,
|
||||
thread_id: ThreadId,
|
||||
item_id: String,
|
||||
payload: String,
|
||||
) -> ThreadStoreFuture<'_, Option<QueuedUserSubmissionRecord>> {
|
||||
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<String>) -> ThreadStoreFuture<'_, ()> {
|
||||
Box::pin(async move {
|
||||
self.queue()
|
||||
.reorder(thread_id, &item_ids)
|
||||
.await
|
||||
.map_err(|error| match error.downcast_ref::<std::io::Error>() {
|
||||
Some(error) if error.kind() == std::io::ErrorKind::InvalidInput => {
|
||||
ThreadStoreError::InvalidRequest {
|
||||
message: error.to_string(),
|
||||
}
|
||||
}
|
||||
_ => ThreadStoreError::Internal {
|
||||
message: format!("queue storage failed: {error}"),
|
||||
},
|
||||
})
|
||||
})
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user