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 b09e081798..117bd1a290 100644 --- a/codex-rs/app-server/src/request_processors/thread_lifecycle.rs +++ b/codex-rs/app-server/src/request_processors/thread_lifecycle.rs @@ -335,9 +335,12 @@ pub(super) async fn ensure_listener_task_running( fallback_model_provider.clone(), ) .await; - if matches!(&event.msg, EventMsg::TurnStarted(_)) { + if let EventMsg::TurnStarted(payload) = &event.msg { thread_queue_processor - .complete_dispatch_after_turn_started(conversation_id) + .complete_dispatch_after_turn_started( + conversation_id, + payload.turn_id.as_str(), + ) .await; } if matches!( 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 b2f5986e8f..36da796841 100644 --- a/codex-rs/app-server/src/request_processors/thread_processor.rs +++ b/codex-rs/app-server/src/request_processors/thread_processor.rs @@ -2597,7 +2597,7 @@ impl ThreadRequestProcessor { .emit_resume_goal_snapshot_and_continue(thread_id, codex_thread.as_ref()) .await; self.thread_queue_processor - .emit_resume_queue_snapshot_and_drain(thread_id) + .recover_resume_queue_snapshot_and_drain(thread_id) .await; } Err(err) => { diff --git a/codex-rs/app-server/src/request_processors/thread_queue_processor.rs b/codex-rs/app-server/src/request_processors/thread_queue_processor.rs index e00e6f73f4..7ee17e404b 100644 --- a/codex-rs/app-server/src/request_processors/thread_queue_processor.rs +++ b/codex-rs/app-server/src/request_processors/thread_queue_processor.rs @@ -122,7 +122,7 @@ impl ThreadQueueRequestProcessor { Ok(None) } - pub(crate) async fn emit_resume_queue_snapshot_and_drain(&self, thread_id: ThreadId) { + pub(crate) async fn recover_resume_queue_snapshot_and_drain(&self, thread_id: ThreadId) { let failure = turn_error("queued turn dispatch was interrupted while app-server restarted"); let failure_json = match serde_json::to_vec(&failure) { Ok(failure_json) => failure_json, @@ -146,14 +146,23 @@ impl ThreadQueueRequestProcessor { self.drain_thread_queue_if_idle(thread_id).await; } + pub(crate) async fn emit_resume_queue_snapshot_and_drain(&self, thread_id: ThreadId) { + self.emit_thread_queue_changed(thread_id).await; + self.drain_thread_queue_if_idle(thread_id).await; + } + pub(crate) async fn drain_thread_queue_after_terminal_turn(&self, thread_id: ThreadId) { self.drain_thread_queue_if_idle(thread_id).await; } - pub(crate) async fn complete_dispatch_after_turn_started(&self, thread_id: ThreadId) { + pub(crate) async fn complete_dispatch_after_turn_started( + &self, + thread_id: ThreadId, + turn_id: &str, + ) { match self .state_db - .remove_dispatching_thread_queued_turn(thread_id) + .remove_dispatching_thread_queued_turn(thread_id, turn_id) .await { Ok(true) | Ok(false) => {} @@ -199,16 +208,55 @@ impl ThreadQueueRequestProcessor { return; } }; - if let Err(err) = self.turn_processor.queued_turn_start(params).await { - self.fail_dispatch( - thread_id, - record.queued_turn_id.as_str(), - turn_error(format!( - "queued turn could not start: {message}", - message = err.message - )), - ) - .await; + match self.turn_processor.queued_turn_start(params).await { + Ok(response) => { + let turn_id = response.turn.id; + match self + .state_db + .set_dispatching_thread_queued_turn_turn_id( + record.queued_turn_id.as_str(), + turn_id.as_str(), + ) + .await + { + Ok(true) => {} + Ok(false) => { + tracing::warn!( + "queued turn {} lost its dispatch claim before turn {turn_id} was recorded", + record.queued_turn_id + ); + return; + } + Err(err) => { + tracing::warn!( + "failed to record dispatch turn {turn_id} for queued turn {}: {err}", + record.queued_turn_id + ); + return; + } + } + let thread_state = self.thread_state_manager.thread_state(thread_id).await; + if thread_state + .lock() + .await + .active_turn_snapshot() + .is_some_and(|turn| turn.id == turn_id) + { + self.complete_dispatch_after_turn_started(thread_id, turn_id.as_str()) + .await; + } + } + Err(err) => { + self.fail_dispatch( + thread_id, + record.queued_turn_id.as_str(), + turn_error(format!( + "queued turn could not start: {message}", + message = err.message + )), + ) + .await; + } } } diff --git a/codex-rs/state/migrations/0031_thread_queued_turns.sql b/codex-rs/state/migrations/0031_thread_queued_turns.sql index 0ba1f16ad8..6630553408 100644 --- a/codex-rs/state/migrations/0031_thread_queued_turns.sql +++ b/codex-rs/state/migrations/0031_thread_queued_turns.sql @@ -4,6 +4,7 @@ CREATE TABLE thread_queued_turns ( turn_start_params_jsonb BLOB NOT NULL, queue_order INTEGER NOT NULL, state TEXT NOT NULL CHECK(state IN ('pending', 'dispatching', 'failed')), + dispatch_turn_id TEXT, failure_jsonb BLOB, created_at_ms INTEGER NOT NULL, updated_at_ms INTEGER NOT NULL, diff --git a/codex-rs/state/src/model/thread_queued_turn.rs b/codex-rs/state/src/model/thread_queued_turn.rs index 81d84054db..665c2e1bb2 100644 --- a/codex-rs/state/src/model/thread_queued_turn.rs +++ b/codex-rs/state/src/model/thread_queued_turn.rs @@ -45,6 +45,7 @@ pub struct ThreadQueuedTurn { pub turn_start_params_jsonb: Vec, pub queue_order: i64, pub state: ThreadQueuedTurnState, + pub dispatch_turn_id: Option, pub failure_jsonb: Option>, pub created_at: DateTime, pub updated_at: DateTime, @@ -56,6 +57,7 @@ pub(crate) struct ThreadQueuedTurnRow { pub turn_start_params_jsonb: Vec, pub queue_order: i64, pub state: String, + pub dispatch_turn_id: Option, pub failure_jsonb: Option>, pub created_at_ms: i64, pub updated_at_ms: i64, @@ -69,6 +71,7 @@ impl ThreadQueuedTurnRow { turn_start_params_jsonb: row.try_get("turn_start_params_jsonb")?, queue_order: row.try_get("queue_order")?, state: row.try_get("state")?, + dispatch_turn_id: row.try_get("dispatch_turn_id")?, 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")?, @@ -86,6 +89,7 @@ impl TryFrom for ThreadQueuedTurn { turn_start_params_jsonb: row.turn_start_params_jsonb, queue_order: row.queue_order, state: ThreadQueuedTurnState::try_from(row.state.as_str())?, + dispatch_turn_id: row.dispatch_turn_id, 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/queued_turns.rs b/codex-rs/state/src/runtime/queued_turns.rs index 4be6ca20e3..7246ae1ad0 100644 --- a/codex-rs/state/src/runtime/queued_turns.rs +++ b/codex-rs/state/src/runtime/queued_turns.rs @@ -17,6 +17,7 @@ INSERT INTO thread_queued_turns ( turn_start_params_jsonb, queue_order, state, + dispatch_turn_id, failure_jsonb, created_at_ms, updated_at_ms @@ -28,6 +29,7 @@ SELECT COALESCE(MAX(queue_order), -1) + 1, 'pending', NULL, + NULL, ?, ? FROM thread_queued_turns @@ -38,6 +40,7 @@ RETURNING CAST(json(turn_start_params_jsonb) AS BLOB) AS turn_start_params_jsonb, queue_order, state, + dispatch_turn_id, CASE WHEN failure_jsonb IS NULL THEN NULL ELSE CAST(json(failure_jsonb) AS BLOB) @@ -70,6 +73,7 @@ SELECT CAST(json(turn_start_params_jsonb) AS BLOB) AS turn_start_params_jsonb, queue_order, state, + dispatch_turn_id, CASE WHEN failure_jsonb IS NULL THEN NULL ELSE CAST(json(failure_jsonb) AS BLOB) @@ -215,6 +219,7 @@ RETURNING CAST(json(turn_start_params_jsonb) AS BLOB) AS turn_start_params_jsonb, queue_order, state, + dispatch_turn_id, CASE WHEN failure_jsonb IS NULL THEN NULL ELSE CAST(json(failure_jsonb) AS BLOB) @@ -231,18 +236,44 @@ RETURNING row.map(|row| thread_queued_turn_from_row(&row)).transpose() } + pub async fn set_dispatching_thread_queued_turn_turn_id( + &self, + queued_turn_id: &str, + turn_id: &str, + ) -> anyhow::Result { + let now_ms = datetime_to_epoch_millis(Utc::now()); + let result = sqlx::query( + r#" +UPDATE thread_queued_turns +SET dispatch_turn_id = ?, updated_at_ms = ? +WHERE queued_turn_id = ? + AND state = 'dispatching' + "#, + ) + .bind(turn_id) + .bind(now_ms) + .bind(queued_turn_id) + .execute(self.pool.as_ref()) + .await?; + + Ok(result.rows_affected() > 0) + } + pub async fn remove_dispatching_thread_queued_turn( &self, thread_id: ThreadId, + turn_id: &str, ) -> anyhow::Result { let result = sqlx::query( r#" DELETE FROM thread_queued_turns WHERE thread_id = ? AND state = 'dispatching' + AND dispatch_turn_id = ? "#, ) .bind(thread_id.to_string()) + .bind(turn_id) .execute(self.pool.as_ref()) .await?; @@ -536,4 +567,46 @@ mod tests { .is_some() ); } + + #[tokio::test] + async fn dispatch_claim_clears_only_for_its_submitted_turn() { + let (runtime, thread_id) = runtime_with_thread().await; + let queued_turn = runtime + .append_thread_queued_turn(thread_id, br#"{"threadId":"t","input":[]}"#) + .await + .expect("append queued turn"); + runtime + .claim_head_thread_queued_turn(thread_id) + .await + .expect("claim queued turn") + .expect("claimed row"); + + assert!( + !runtime + .remove_dispatching_thread_queued_turn(thread_id, "regular-turn") + .await + .expect("unmatched turn must not clear claim") + ); + assert!( + runtime + .set_dispatching_thread_queued_turn_turn_id( + &queued_turn.queued_turn_id, + "queued-turn", + ) + .await + .expect("record submitted queued turn id") + ); + assert!( + !runtime + .remove_dispatching_thread_queued_turn(thread_id, "regular-turn") + .await + .expect("different started turn must not clear claim") + ); + assert!( + runtime + .remove_dispatching_thread_queued_turn(thread_id, "queued-turn") + .await + .expect("matching queued turn clears claim") + ); + } }