From 2aa714e273a3c979fa2f7a3e672b193397d01b48 Mon Sep 17 00:00:00 2001 From: Edward Frazer Date: Mon, 15 Jun 2026 19:04:40 +0000 Subject: [PATCH] feat: dispatch queued user messages from thread idle --- codex-rs/Cargo.lock | 20 + codex-rs/Cargo.toml | 2 + codex-rs/app-server/Cargo.toml | 1 + codex-rs/app-server/src/extensions.rs | 8 +- .../thread_goal_processor.rs | 9 +- .../request_processors/thread_lifecycle.rs | 4 +- .../request_processors/thread_processor.rs | 5 +- codex-rs/core/src/session/turn.rs | 1 + codex-rs/ext/queue/BUILD.bazel | 6 + codex-rs/ext/queue/Cargo.toml | 32 ++ codex-rs/ext/queue/src/lib.rs | 22 ++ codex-rs/ext/queue/src/service.rs | 351 ++++++++++++++++++ codex-rs/ext/queue/src/types.rs | 51 +++ codex-rs/ext/queue/tests/queue_service.rs | 267 +++++++++++++ codex-rs/mcp-server/src/codex_tool_runner.rs | 1 + codex-rs/protocol/src/protocol.rs | 10 + codex-rs/rollout-trace/src/protocol_event.rs | 2 + codex-rs/rollout/src/policy.rs | 1 + 18 files changed, 780 insertions(+), 13 deletions(-) create mode 100644 codex-rs/ext/queue/BUILD.bazel create mode 100644 codex-rs/ext/queue/Cargo.toml create mode 100644 codex-rs/ext/queue/src/lib.rs create mode 100644 codex-rs/ext/queue/src/service.rs create mode 100644 codex-rs/ext/queue/src/types.rs create mode 100644 codex-rs/ext/queue/tests/queue_service.rs diff --git a/codex-rs/Cargo.lock b/codex-rs/Cargo.lock index d94daa0e8f..6e844cafaa 100644 --- a/codex-rs/Cargo.lock +++ b/codex-rs/Cargo.lock @@ -2003,6 +2003,7 @@ dependencies = [ "codex-otel", "codex-plugin", "codex-protocol", + "codex-queue-extension", "codex-rmcp-client", "codex-rollout", "codex-sandboxing", @@ -3660,6 +3661,25 @@ dependencies = [ "wildmatch", ] +[[package]] +name = "codex-queue-extension" +version = "0.0.0" +dependencies = [ + "anyhow", + "codex-core", + "codex-extension-api", + "codex-protocol", + "codex-state", + "core_test_support", + "pretty_assertions", + "serde", + "serde_json", + "tempfile", + "thiserror 2.0.18", + "tokio", + "tracing", +] + [[package]] name = "codex-realtime-webrtc" version = "0.0.0" diff --git a/codex-rs/Cargo.toml b/codex-rs/Cargo.toml index e5f5025e91..6f135a16d5 100644 --- a/codex-rs/Cargo.toml +++ b/codex-rs/Cargo.toml @@ -49,6 +49,7 @@ members = [ "execpolicy-legacy", "ext/extension-api", "ext/goal", + "ext/queue", "ext/guardian", "ext/image-generation", "ext/memories", @@ -175,6 +176,7 @@ codex-exec-server = { path = "exec-server" } codex-execpolicy = { path = "execpolicy" } codex-extension-api = { path = "ext/extension-api" } codex-goal-extension = { path = "ext/goal" } +codex-queue-extension = { path = "ext/queue" } codex-guardian = { path = "ext/guardian" } codex-image-generation-extension = { path = "ext/image-generation" } codex-external-agent-migration = { path = "external-agent-migration" } diff --git a/codex-rs/app-server/Cargo.toml b/codex-rs/app-server/Cargo.toml index 3d08059a63..a0e264c934 100644 --- a/codex-rs/app-server/Cargo.toml +++ b/codex-rs/app-server/Cargo.toml @@ -48,6 +48,7 @@ codex-file-watcher = { workspace = true } codex-hooks = { workspace = true } codex-otel = { workspace = true } codex-plugin = { workspace = true } +codex-queue-extension = { workspace = true } codex-shell-command = { workspace = true } codex-skills-extension = { workspace = true } codex-utils-cli = { workspace = true } diff --git a/codex-rs/app-server/src/extensions.rs b/codex-rs/app-server/src/extensions.rs index 8a6bc07f3d..c3042724ae 100644 --- a/codex-rs/app-server/src/extensions.rs +++ b/codex-rs/app-server/src/extensions.rs @@ -59,8 +59,14 @@ where executor_skill_provider, thread_store: _thread_store, } = dependencies; - let mut builder = ExtensionRegistryBuilder::::with_event_sink(event_sink); + let mut builder = ExtensionRegistryBuilder::::with_event_sink(Arc::clone(&event_sink)); if let Some(state_db) = state_db { + let queue_service = Arc::new(codex_queue_extension::QueuedItemService::new( + Arc::clone(&state_db), + thread_manager.clone(), + Arc::clone(&event_sink), + )); + codex_queue_extension::install(&mut builder, queue_service); codex_goal_extension::install_with_backend( &mut builder, state_db, diff --git a/codex-rs/app-server/src/request_processors/thread_goal_processor.rs b/codex-rs/app-server/src/request_processors/thread_goal_processor.rs index 2f40ce6602..85b75ea52a 100644 --- a/codex-rs/app-server/src/request_processors/thread_goal_processor.rs +++ b/codex-rs/app-server/src/request_processors/thread_goal_processor.rs @@ -63,18 +63,11 @@ impl ThreadGoalRequestProcessor { .map(|()| None) } - pub(crate) async fn emit_resume_goal_snapshot_and_continue( - &self, - thread_id: ThreadId, - thread: &CodexThread, - ) { + pub(crate) async fn emit_resume_goal_snapshot(&self, thread_id: ThreadId) { if !self.config.features.enabled(Feature::Goals) { return; } self.emit_thread_goal_snapshot(thread_id).await; - // App-server owns resume response and snapshot ordering, so wait until - // those are sent before letting extensions react to the idle thread. - thread.emit_thread_idle_lifecycle_if_idle().await; } pub(crate) async fn pending_resume_goal_state( diff --git a/codex-rs/app-server/src/request_processors/thread_lifecycle.rs b/codex-rs/app-server/src/request_processors/thread_lifecycle.rs index fd9e93e189..f3f28ff9a6 100644 --- a/codex-rs/app-server/src/request_processors/thread_lifecycle.rs +++ b/codex-rs/app-server/src/request_processors/thread_lifecycle.rs @@ -698,9 +698,7 @@ pub(super) async fn handle_pending_thread_resume_request( .await; // App-server owns resume response and snapshot ordering, so wait until // replay completes before letting extensions react to the idle thread. - if pending.emit_thread_goal_update { - conversation.emit_thread_idle_lifecycle_if_idle().await; - } + conversation.emit_thread_idle_lifecycle_if_idle().await; } pub(super) async fn send_thread_goal_snapshot_notification( diff --git a/codex-rs/app-server/src/request_processors/thread_processor.rs b/codex-rs/app-server/src/request_processors/thread_processor.rs index 09ff0dd698..1becc221dd 100644 --- a/codex-rs/app-server/src/request_processors/thread_processor.rs +++ b/codex-rs/app-server/src/request_processors/thread_processor.rs @@ -2814,8 +2814,11 @@ impl ThreadRequestProcessor { .await; } self.thread_goal_processor - .emit_resume_goal_snapshot_and_continue(thread_id, codex_thread.as_ref()) + .emit_resume_goal_snapshot(thread_id) .await; + // App-server owns resume response and snapshot ordering, so wait until + // those are sent before letting extensions react to the idle thread. + codex_thread.emit_thread_idle_lifecycle_if_idle().await; } Err(err) => { let error = internal_error(format!("error resuming thread: {err}")); diff --git a/codex-rs/core/src/session/turn.rs b/codex-rs/core/src/session/turn.rs index 3c783eb486..1d20846e59 100644 --- a/codex-rs/core/src/session/turn.rs +++ b/codex-rs/core/src/session/turn.rs @@ -1444,6 +1444,7 @@ pub(super) fn realtime_text_for_event(msg: &EventMsg) -> Option { | EventMsg::AgentReasoningSectionBreak(_) | EventMsg::SessionConfigured(_) | EventMsg::ThreadGoalUpdated(_) + | EventMsg::ThreadQueueChanged(_) | EventMsg::McpStartupUpdate(_) | EventMsg::McpStartupComplete(_) | EventMsg::McpToolCallBegin(_) diff --git a/codex-rs/ext/queue/BUILD.bazel b/codex-rs/ext/queue/BUILD.bazel new file mode 100644 index 0000000000..21f5afd2c9 --- /dev/null +++ b/codex-rs/ext/queue/BUILD.bazel @@ -0,0 +1,6 @@ +load("//:defs.bzl", "codex_rust_crate") + +codex_rust_crate( + name = "queue", + crate_name = "codex_queue_extension", +) diff --git a/codex-rs/ext/queue/Cargo.toml b/codex-rs/ext/queue/Cargo.toml new file mode 100644 index 0000000000..787f0fec51 --- /dev/null +++ b/codex-rs/ext/queue/Cargo.toml @@ -0,0 +1,32 @@ +[package] +edition.workspace = true +license.workspace = true +name = "codex-queue-extension" +version.workspace = true + +[lib] +name = "codex_queue_extension" +path = "src/lib.rs" +test = false +doctest = false + +[lints] +workspace = true + +[dependencies] +anyhow = { workspace = true } +codex-core = { workspace = true } +codex-extension-api = { workspace = true } +codex-protocol = { workspace = true } +codex-state = { workspace = true } +serde = { workspace = true, features = ["derive"] } +serde_json = { workspace = true } +thiserror = { workspace = true } +tokio = { workspace = true, features = ["rt", "sync", "time"] } +tracing = { workspace = true } + +[dev-dependencies] +core_test_support = { workspace = true } +pretty_assertions = { workspace = true } +tempfile = { workspace = true } +tokio = { workspace = true, features = ["macros", "rt-multi-thread"] } diff --git a/codex-rs/ext/queue/src/lib.rs b/codex-rs/ext/queue/src/lib.rs new file mode 100644 index 0000000000..a77574b580 --- /dev/null +++ b/codex-rs/ext/queue/src/lib.rs @@ -0,0 +1,22 @@ +//! Durable user-submission queue and idle-dispatch extension. + +use std::sync::Arc; + +use codex_extension_api::ExtensionRegistryBuilder; + +mod service; +mod types; + +pub use service::QueueServiceError; +pub use service::QueuedItemService; +pub use types::QueuedItem; +pub use types::QueuedItemProvenance; +pub use types::QueuedItemStatus; + +/// Installs durable queue dispatch into the host's thread and turn lifecycle. +pub fn install(registry: &mut ExtensionRegistryBuilder, service: Arc) +where + C: Send + Sync + 'static, +{ + registry.thread_lifecycle_contributor(service); +} diff --git a/codex-rs/ext/queue/src/service.rs b/codex-rs/ext/queue/src/service.rs new file mode 100644 index 0000000000..6641185197 --- /dev/null +++ b/codex-rs/ext/queue/src/service.rs @@ -0,0 +1,351 @@ +use std::collections::HashSet; +use std::sync::Arc; +use std::sync::Weak; +use std::time::Duration; +use std::time::SystemTime; +use std::time::UNIX_EPOCH; + +use codex_core::IdleTurnInput; +use codex_core::ThreadManager; +use codex_core::TryStartTurnIfIdleRejectionReason; +use codex_extension_api::ExtensionEventSink; +use codex_extension_api::ExtensionFuture; +use codex_extension_api::ThreadIdleInput; +use codex_extension_api::ThreadLifecycleContributor; +use codex_protocol::ThreadId; +use codex_protocol::protocol::Event; +use codex_protocol::protocol::EventMsg; +use codex_protocol::protocol::ThreadQueueChangedEvent; +use codex_protocol::protocol::UserSubmission; +use codex_state::QueuedItemRecord; +use codex_state::QueuedItemState; +use serde::Deserialize; +use serde::Serialize; +use thiserror::Error; +use tokio::sync::Mutex; + +use crate::types::QueuedItem; +use crate::types::QueuedItemProvenance; +use crate::types::QueuedItemStatus; +use crate::types::StoredQueuedItemPayload; + +const INTERRUPTED_CLAIM_MESSAGE: &str = "queue claim was interrupted before core accepted it"; +const CLAIM_STALE_AFTER: Duration = Duration::from_secs(5 * 60); +const IDLE_RETRY_AFTER: Duration = Duration::from_secs(1); + +#[derive(Debug, Error)] +pub enum QueueServiceError { + #[error("queue storage failed: {0}")] + Storage(#[from] anyhow::Error), + #[error("queued item payload is invalid: {0}")] + InvalidPayload(#[from] serde_json::Error), + #[error("queue reorder must include every visible queued item exactly once")] + InvalidReorder, +} + +#[derive(Clone)] +pub struct QueuedItemService { + state_dbs: Arc, + thread_manager: Weak, + event_sink: Arc, + scheduled_recoveries: Arc>>, +} + +impl QueuedItemService { + pub fn new( + state_dbs: Arc, + thread_manager: Weak, + event_sink: Arc, + ) -> Self { + Self { + state_dbs, + thread_manager, + event_sink, + scheduled_recoveries: Arc::new(Mutex::new(HashSet::new())), + } + } + + pub async fn enqueue( + &self, + thread_id: ThreadId, + submission: UserSubmission, + provenance: QueuedItemProvenance, + ) -> Result { + self.recover_stale_claims(thread_id).await?; + let payload = serde_json::to_vec(&StoredQueuedItemPayload::V1 { + submission, + provenance, + })?; + let record = self + .state_dbs + .thread_queue() + .enqueue(thread_id, &payload) + .await?; + let item = queued_item_from_record(record)?; + self.emit_changed(thread_id); + self.wake_if_loaded(thread_id).await; + Ok(item) + } + + pub async fn list( + &self, + thread_id: ThreadId, + offset: usize, + limit: usize, + ) -> Result, QueueServiceError> { + self.recover_stale_claims(thread_id).await?; + self.state_dbs + .thread_queue() + .list_page(thread_id, offset, limit) + .await? + .into_iter() + .map(queued_item_from_record) + .collect() + } + + pub async fn delete( + &self, + thread_id: ThreadId, + queued_item_id: &str, + ) -> Result { + self.recover_stale_claims(thread_id).await?; + let deleted = self + .state_dbs + .thread_queue() + .delete(thread_id, queued_item_id) + .await?; + if deleted { + self.emit_changed(thread_id); + self.wake_if_loaded(thread_id).await; + } + Ok(deleted) + } + + pub async fn reorder( + &self, + thread_id: ThreadId, + ordered_ids: &[String], + ) -> Result<(), QueueServiceError> { + self.recover_stale_claims(thread_id).await?; + if !self + .state_dbs + .thread_queue() + .reorder(thread_id, ordered_ids) + .await? + { + return Err(QueueServiceError::InvalidReorder); + } + self.emit_changed(thread_id); + self.wake_if_loaded(thread_id).await; + Ok(()) + } + + async fn dispatch_if_idle(&self, thread_id: ThreadId) -> Result<(), QueueServiceError> { + self.recover_stale_claims(thread_id).await?; + let (record, claim_token, payload) = loop { + let Some(claim) = self.state_dbs.thread_queue().claim_next(thread_id).await? else { + if self + .state_dbs + .thread_queue() + .has_claimed_item(thread_id) + .await? + { + self.schedule_recovery(thread_id).await; + } + return Ok(()); + }; + let claim_token = claim.claim_token; + let record = claim.item; + match serde_json::from_slice::(&record.payload_jsonb) { + Ok(payload) => break (record, claim_token, payload), + Err(err) => { + let queued_item_id = &record.queued_item_id; + if !self + .state_dbs + .thread_queue() + .complete_claim(queued_item_id, &claim_token) + .await? + { + tracing::warn!(%queued_item_id, "invalid queued item lost its storage claim"); + } + tracing::warn!(%queued_item_id, %err, "discarding invalid queued item payload"); + self.emit_changed(thread_id); + } + } + }; + let queued_item_id = record.queued_item_id.clone(); + let (submission, _) = payload.into_parts(); + let thread = match self.thread_manager.upgrade() { + Some(thread_manager) => thread_manager.get_thread(thread_id).await.ok(), + None => None, + }; + let Some(thread) = thread else { + self.state_dbs + .thread_queue() + .release_claim(&queued_item_id, &claim_token) + .await?; + return Ok(()); + }; + + match thread + .try_start_turn_if_idle(IdleTurnInput::UserSubmission(submission)) + .await + { + Ok(()) => { + if !self + .state_dbs + .thread_queue() + .complete_claim(&queued_item_id, &claim_token) + .await? + { + tracing::warn!(%queued_item_id, "accepted queued item lost its storage claim"); + } + self.emit_changed(thread_id); + Ok(()) + } + Err(err) => match err.reason() { + TryStartTurnIfIdleRejectionReason::Busy => { + self.state_dbs + .thread_queue() + .release_claim(&queued_item_id, &claim_token) + .await?; + let service = self.clone(); + tokio::spawn(async move { + tokio::time::sleep(IDLE_RETRY_AFTER).await; + service.wake_if_loaded(thread_id).await; + }); + Ok(()) + } + TryStartTurnIfIdleRejectionReason::PendingTriggerTurn => { + self.state_dbs + .thread_queue() + .release_claim(&queued_item_id, &claim_token) + .await?; + Ok(()) + } + TryStartTurnIfIdleRejectionReason::PlanMode => { + let failure = failure_json(&format!( + "core rejected queued user input: {:?}", + err.reason() + ))?; + self.state_dbs + .thread_queue() + .fail_claim(&queued_item_id, &claim_token, &failure) + .await?; + self.emit_changed(thread_id); + Ok(()) + } + }, + } + } + + async fn recover_stale_claims(&self, thread_id: ThreadId) -> Result<(), QueueServiceError> { + let stale_before_ms = SystemTime::now() + .duration_since(UNIX_EPOCH) + .unwrap_or_default() + .saturating_sub(CLAIM_STALE_AFTER) + .as_millis() + .try_into() + .unwrap_or(i64::MAX); + let failure = failure_json(INTERRUPTED_CLAIM_MESSAGE)?; + let recovered = self + .state_dbs + .thread_queue() + .recover_claims_as_failed_before(thread_id, stale_before_ms, &failure) + .await?; + if recovered > 0 { + self.emit_changed(thread_id); + } + Ok(()) + } + + async fn schedule_recovery(&self, thread_id: ThreadId) { + if !self.scheduled_recoveries.lock().await.insert(thread_id) { + return; + } + let service = self.clone(); + tokio::spawn(async move { + tokio::time::sleep(CLAIM_STALE_AFTER).await; + service.scheduled_recoveries.lock().await.remove(&thread_id); + if let Err(err) = service.recover_stale_claims(thread_id).await { + tracing::warn!(%thread_id, %err, "failed to recover stale queue claim"); + } + }); + } + + async fn wake_if_loaded(&self, thread_id: ThreadId) { + let Some(thread_manager) = self.thread_manager.upgrade() else { + return; + }; + let Ok(thread) = thread_manager.get_thread(thread_id).await else { + return; + }; + thread.emit_thread_idle_lifecycle_if_idle().await; + } + + fn emit_changed(&self, thread_id: ThreadId) { + self.event_sink.emit(Event { + id: thread_id.to_string(), + msg: EventMsg::ThreadQueueChanged(ThreadQueueChangedEvent { thread_id }), + }); + } +} + +impl ThreadLifecycleContributor for QueuedItemService +where + C: Send + Sync + 'static, +{ + fn on_thread_idle<'a>(&'a self, input: ThreadIdleInput<'a>) -> ExtensionFuture<'a, ()> { + Box::pin(async move { + let Ok(thread_id) = ThreadId::from_string(input.thread_store.level_id()) else { + tracing::warn!( + level_id = input.thread_store.level_id(), + "queue extension received an invalid thread id" + ); + return; + }; + if let Err(err) = self.dispatch_if_idle(thread_id).await { + tracing::warn!(%thread_id, "failed to dispatch queued item: {err}"); + } + }) + } +} + +#[derive(Deserialize, Serialize)] +struct QueuedItemFailure { + message: String, +} + +fn failure_json(message: &str) -> Result, QueueServiceError> { + Ok(serde_json::to_vec(&QueuedItemFailure { + message: message.to_string(), + })?) +} + +fn queued_item_from_record(record: QueuedItemRecord) -> Result { + let payload: StoredQueuedItemPayload = serde_json::from_slice(&record.payload_jsonb)?; + let (submission, provenance) = payload.into_parts(); + let status = match record.state { + QueuedItemState::Pending => QueuedItemStatus::Pending, + QueuedItemState::Failed => { + let error = record + .failure_jsonb + .as_deref() + .and_then(|json| serde_json::from_slice::(json).ok()) + .map(|failure| failure.message) + .unwrap_or_else(|| "queued item failed".to_string()); + QueuedItemStatus::Failed { error } + } + QueuedItemState::Claimed => { + return Err(QueueServiceError::Storage(anyhow::anyhow!( + "claimed queued items are not client-visible" + ))); + } + }; + Ok(QueuedItem { + id: record.queued_item_id, + submission, + provenance, + status, + }) +} diff --git a/codex-rs/ext/queue/src/types.rs b/codex-rs/ext/queue/src/types.rs new file mode 100644 index 0000000000..cb91e36f31 --- /dev/null +++ b/codex-rs/ext/queue/src/types.rs @@ -0,0 +1,51 @@ +use std::collections::HashMap; + +use codex_protocol::protocol::UserSubmission; +use serde::Deserialize; +use serde::Serialize; +use serde_json::Value; + +#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)] +#[serde(tag = "type", rename_all = "camelCase")] +pub enum QueuedItemProvenance { + User, + ExternalEvent { + source: String, + metadata: HashMap, + }, +} + +#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)] +#[serde(tag = "version")] +pub(crate) enum StoredQueuedItemPayload { + #[serde(rename = "1")] + V1 { + submission: UserSubmission, + provenance: QueuedItemProvenance, + }, +} + +#[derive(Clone, Debug, PartialEq)] +pub struct QueuedItem { + pub id: String, + pub submission: UserSubmission, + pub provenance: QueuedItemProvenance, + pub status: QueuedItemStatus, +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub enum QueuedItemStatus { + Pending, + Failed { error: String }, +} + +impl StoredQueuedItemPayload { + pub(crate) fn into_parts(self) -> (UserSubmission, QueuedItemProvenance) { + match self { + Self::V1 { + submission, + provenance, + } => (submission, provenance), + } + } +} diff --git a/codex-rs/ext/queue/tests/queue_service.rs b/codex-rs/ext/queue/tests/queue_service.rs new file mode 100644 index 0000000000..45f8a346c2 --- /dev/null +++ b/codex-rs/ext/queue/tests/queue_service.rs @@ -0,0 +1,267 @@ +use std::collections::HashMap; +use std::sync::Arc; +use std::sync::Weak; + +use codex_extension_api::ExtensionData; +use codex_extension_api::NoopExtensionEventSink; +use codex_extension_api::ThreadIdleInput; +use codex_extension_api::ThreadLifecycleContributor; +use codex_protocol::ThreadId; +use codex_protocol::protocol::EventMsg; +use codex_protocol::protocol::UserSubmission; +use codex_protocol::user_input::UserInput; +use codex_queue_extension::QueuedItemProvenance; +use codex_queue_extension::QueuedItemService; +use core_test_support::responses; +use core_test_support::responses::start_mock_server; +use core_test_support::test_codex::test_codex; +use core_test_support::wait_for_event_match; +use pretty_assertions::assert_eq; +use serde_json::json; + +async fn service() -> anyhow::Result { + let runtime = + codex_state::StateRuntime::init(tempfile::tempdir()?.keep(), "test-provider".to_string()) + .await?; + Ok(QueuedItemService::new( + runtime, + Weak::new(), + Arc::new(NoopExtensionEventSink), + )) +} + +fn submission(text: &str) -> UserSubmission { + UserSubmission { + items: vec![UserInput::Text { + text: text.to_string(), + text_elements: Vec::new(), + }], + ..Default::default() + } +} + +#[tokio::test] +async fn service_round_trips_external_provenance_metadata() -> anyhow::Result<()> { + let service = service().await?; + let thread_id = ThreadId::new(); + let provenance = QueuedItemProvenance::ExternalEvent { + source: "slack".to_string(), + metadata: HashMap::from([("channel".to_string(), json!("C123"))]), + }; + + let added = service + .enqueue(thread_id, submission("hello"), provenance.clone()) + .await?; + let listed = service.list(thread_id, /*offset*/ 0, /*limit*/ 10).await?; + + assert_eq!(added, listed[0]); + assert_eq!(provenance, listed[0].provenance); + Ok(()) +} + +#[tokio::test] +async fn delete_and_reorder_preserve_visible_order() -> anyhow::Result<()> { + let service = service().await?; + let thread_id = ThreadId::new(); + let first = service + .enqueue(thread_id, submission("first"), QueuedItemProvenance::User) + .await?; + let second = service + .enqueue(thread_id, submission("second"), QueuedItemProvenance::User) + .await?; + + service + .reorder(thread_id, &[second.id.clone(), first.id.clone()]) + .await?; + let reordered = service.list(thread_id, /*offset*/ 0, /*limit*/ 10).await?; + assert_eq!( + vec![second.id.clone(), first.id.clone()], + reordered + .iter() + .map(|item| item.id.clone()) + .collect::>() + ); + assert!(service.delete(thread_id, &second.id).await?); + assert_eq!( + vec![first], + service.list(thread_id, /*offset*/ 0, /*limit*/ 10).await? + ); + Ok(()) +} + +#[test] +fn dispatch_starts_a_live_core_turn_and_completes_the_claim() -> anyhow::Result<()> { + tokio::runtime::Builder::new_multi_thread() + .worker_threads(2) + .thread_stack_size(16 * 1024 * 1024) + .enable_all() + .build()? + .block_on(async { + let server = start_mock_server().await; + let response = + responses::mount_sse_once(&server, responses::sse_completed("queued-turn")).await; + let test = test_codex() + .with_config(|config| config.include_environment_context = false) + .build(&server) + .await?; + let thread_id = test.session_configured.thread_id; + let state_db = test.codex.state_db().expect("state db"); + let service = QueuedItemService::new( + state_db, + Arc::downgrade(&test.thread_manager), + Arc::new(NoopExtensionEventSink), + ); + + service + .enqueue( + thread_id, + submission("durable follow-up"), + QueuedItemProvenance::User, + ) + .await?; + let session_store = ExtensionData::new("session"); + let thread_store = ExtensionData::new(thread_id.to_string()); + >::on_thread_idle( + &service, + ThreadIdleInput { + session_store: &session_store, + thread_store: &thread_store, + }, + ) + .await; + assert!( + service + .list(thread_id, /*offset*/ 0, /*limit*/ 10) + .await? + .is_empty() + ); + + wait_for_event_match(test.codex.as_ref(), |event| match event { + EventMsg::TurnComplete(event) => Some(event.turn_id.clone()), + _ => None, + }) + .await; + assert_eq!( + response.single_request().message_input_texts("user"), + vec!["durable follow-up"] + ); + Ok::<(), anyhow::Error>(()) + }) +} + +#[test] +fn post_resume_idle_dispatches_an_item_queued_while_the_thread_was_unloaded() -> anyhow::Result<()> +{ + tokio::runtime::Builder::new_multi_thread() + .worker_threads(2) + .thread_stack_size(16 * 1024 * 1024) + .enable_all() + .build()? + .block_on(async { + let server = start_mock_server().await; + responses::mount_sse_once(&server, responses::sse_completed("resumed-turn")).await; + let test = test_codex().build(&server).await?; + let thread_id = test.session_configured.thread_id; + let state_db = test.codex.state_db().expect("state db"); + QueuedItemService::new( + state_db.clone(), + Weak::new(), + Arc::new(NoopExtensionEventSink), + ) + .enqueue( + thread_id, + submission("queued while unloaded"), + QueuedItemProvenance::User, + ) + .await?; + let service = QueuedItemService::new( + state_db, + Arc::downgrade(&test.thread_manager), + Arc::new(NoopExtensionEventSink), + ); + let session_store = ExtensionData::new("session"); + let thread_store = ExtensionData::new(thread_id.to_string()); + + >::on_thread_idle( + &service, + ThreadIdleInput { + session_store: &session_store, + thread_store: &thread_store, + }, + ) + .await; + + assert!( + service + .list(thread_id, /*offset*/ 0, /*limit*/ 10) + .await? + .is_empty() + ); + Ok::<(), anyhow::Error>(()) + }) +} + +#[test] +fn invalid_payload_is_discarded_instead_of_blocking_the_queue() -> anyhow::Result<()> { + tokio::runtime::Builder::new_multi_thread() + .worker_threads(2) + .thread_stack_size(16 * 1024 * 1024) + .enable_all() + .build()? + .block_on(async { + let server = start_mock_server().await; + let response = + responses::mount_sse_once(&server, responses::sse_completed("valid-turn")).await; + let test = test_codex() + .with_config(|config| config.include_environment_context = false) + .build(&server) + .await?; + let thread_id = test.session_configured.thread_id; + let state_db = test.codex.state_db().expect("state db"); + state_db + .thread_queue() + .enqueue(thread_id, br#"{"unsupported":true}"#) + .await?; + let service = QueuedItemService::new( + state_db.clone(), + Arc::downgrade(&test.thread_manager), + Arc::new(NoopExtensionEventSink), + ); + service + .enqueue( + thread_id, + submission("valid follow-up"), + QueuedItemProvenance::User, + ) + .await?; + let session_store = ExtensionData::new("session"); + let thread_store = ExtensionData::new(thread_id.to_string()); + + >::on_thread_idle( + &service, + ThreadIdleInput { + session_store: &session_store, + thread_store: &thread_store, + }, + ) + .await; + + assert!( + state_db + .thread_queue() + .list_page(thread_id, /*offset*/ 0, /*limit*/ 1) + .await? + .is_empty() + ); + wait_for_event_match(test.codex.as_ref(), |event| match event { + EventMsg::TurnComplete(event) => Some(event.turn_id.clone()), + _ => None, + }) + .await; + assert_eq!( + response.single_request().message_input_texts("user"), + vec!["valid follow-up"] + ); + Ok::<(), anyhow::Error>(()) + }) +} diff --git a/codex-rs/mcp-server/src/codex_tool_runner.rs b/codex-rs/mcp-server/src/codex_tool_runner.rs index 7ef7d1cbc8..71c2155278 100644 --- a/codex-rs/mcp-server/src/codex_tool_runner.rs +++ b/codex-rs/mcp-server/src/codex_tool_runner.rs @@ -388,6 +388,7 @@ async fn run_codex_tool_session_inner( | EventMsg::RealtimeConversationSdp(_) | EventMsg::RealtimeConversationRealtime(_) | EventMsg::RealtimeConversationClosed(_) + | EventMsg::ThreadQueueChanged(_) | EventMsg::DeprecationNotice(_) => { // For now, we do not do anything extra for these // events. Note that diff --git a/codex-rs/protocol/src/protocol.rs b/codex-rs/protocol/src/protocol.rs index b758fe2dc2..2ba4c35a1a 100644 --- a/codex-rs/protocol/src/protocol.rs +++ b/codex-rs/protocol/src/protocol.rs @@ -1275,6 +1275,9 @@ pub enum EventMsg { /// Updated long-running goal metadata for the thread. ThreadGoalUpdated(ThreadGoalUpdatedEvent), + /// The durable user-submission queue for a thread changed. + ThreadQueueChanged(ThreadQueueChangedEvent), + /// Incremental MCP startup progress updates. McpStartupUpdate(McpStartupUpdateEvent), @@ -3769,6 +3772,13 @@ pub struct ThreadGoalUpdatedEvent { pub goal: ThreadGoal, } +#[derive(Debug, Clone, Deserialize, Serialize, PartialEq, Eq, JsonSchema, TS)] +#[serde(rename_all = "camelCase")] +#[ts(export_to = "protocol/")] +pub struct ThreadQueueChangedEvent { + pub thread_id: ThreadId, +} + /// User's decision in response to an ExecApprovalRequest. #[derive(Debug, Default, Clone, Deserialize, Serialize, PartialEq, Eq, Display, JsonSchema, TS)] #[serde(rename_all = "snake_case")] diff --git a/codex-rs/rollout-trace/src/protocol_event.rs b/codex-rs/rollout-trace/src/protocol_event.rs index bbba245307..4d82871e45 100644 --- a/codex-rs/rollout-trace/src/protocol_event.rs +++ b/codex-rs/rollout-trace/src/protocol_event.rs @@ -236,6 +236,7 @@ pub(crate) fn tool_runtime_trace_event(event: &EventMsg) -> Option Option<&'static s | EventMsg::AgentReasoningRawContent(_) | EventMsg::AgentReasoningSectionBreak(_) | EventMsg::ThreadGoalUpdated(_) + | EventMsg::ThreadQueueChanged(_) | EventMsg::McpStartupUpdate(_) | EventMsg::McpStartupComplete(_) | EventMsg::McpToolCallBegin(_) diff --git a/codex-rs/rollout/src/policy.rs b/codex-rs/rollout/src/policy.rs index 169f4360d3..aa908568e7 100644 --- a/codex-rs/rollout/src/policy.rs +++ b/codex-rs/rollout/src/policy.rs @@ -128,6 +128,7 @@ pub fn should_persist_event_msg(ev: &EventMsg) -> bool { | EventMsg::RawResponseItem(_) | EventMsg::SessionConfigured(_) | EventMsg::ThreadSettingsApplied(_) + | EventMsg::ThreadQueueChanged(_) | EventMsg::McpToolCallBegin(_) | EventMsg::ExecCommandBegin(_) | EventMsg::TerminalInteraction(_)