From bdb868f9e24da74a2b344297fd2f51fd69db3218 Mon Sep 17 00:00:00 2001 From: jif-oai Date: Fri, 5 Jun 2026 16:19:57 +0200 Subject: [PATCH] fix: wake queued agent turns after slot release --- codex-rs/core/src/agent/control.rs | 31 +++++++- codex-rs/core/src/agent/control/spawn.rs | 4 + codex-rs/core/src/agent/registry.rs | 3 +- codex-rs/core/src/session/handlers.rs | 26 +++++-- codex-rs/core/src/tasks/mod.rs | 96 ++++++++++++++++-------- 5 files changed, 119 insertions(+), 41 deletions(-) diff --git a/codex-rs/core/src/agent/control.rs b/codex-rs/core/src/agent/control.rs index 60dbaa5456..32d01f518d 100644 --- a/codex-rs/core/src/agent/control.rs +++ b/codex-rs/core/src/agent/control.rs @@ -268,7 +268,7 @@ impl AgentControl { Ok(thread.subscribe_status()) } - pub(crate) async fn reserve_execution_slot_for_pending_turn( + pub(crate) async fn reserve_execution_slot_for_turn_start( &self, session: &Session, ) -> CodexResult> { @@ -305,8 +305,35 @@ impl AgentControl { } if session.active_turn.lock().await.is_none() && !session.input_queue.has_trigger_turn_mailbox_items().await + && self.state.release_execution_slot(session.thread_id) { - self.state.release_execution_slot(session.thread_id); + let Ok(state) = self.upgrade() else { + return; + }; + for agent in self.state.live_agents() { + let Some(thread_id) = agent.agent_id else { + continue; + }; + if thread_id == session.thread_id { + continue; + } + let Ok(thread) = state.get_thread(thread_id).await else { + continue; + }; + if thread + .codex + .session + .input_queue + .has_trigger_turn_mailbox_items() + .await + { + thread + .codex + .session + .maybe_start_turn_for_pending_work() + .await; + } + } } } diff --git a/codex-rs/core/src/agent/control/spawn.rs b/codex-rs/core/src/agent/control/spawn.rs index 7cf5d00b4d..a516908082 100644 --- a/codex-rs/core/src/agent/control/spawn.rs +++ b/codex-rs/core/src/agent/control/spawn.rs @@ -581,6 +581,10 @@ impl AgentControl { // Resumed threads are re-registered in-memory and need the same listener // attachment path as freshly spawned threads. state.notify_thread_created(resumed_thread.thread_id); + if multi_agent_version == MultiAgentVersion::V2 { + self.release_execution_slot_if_idle(resumed_thread.thread.codex.session.as_ref()) + .await; + } if multi_agent_version != MultiAgentVersion::V2 { let child_reference = agent_metadata .agent_path diff --git a/codex-rs/core/src/agent/registry.rs b/codex-rs/core/src/agent/registry.rs index 37e457e35c..75324f9ee6 100644 --- a/codex-rs/core/src/agent/registry.rs +++ b/codex-rs/core/src/agent/registry.rs @@ -110,7 +110,7 @@ impl AgentRegistry { } } - pub(crate) fn release_execution_slot(&self, thread_id: ThreadId) { + pub(crate) fn release_execution_slot(&self, thread_id: ThreadId) -> bool { let mut active_agents = self .active_agents .lock() @@ -119,6 +119,7 @@ impl AgentRegistry { if removed { self.total_count.fetch_sub(1, Ordering::AcqRel); } + removed } pub(crate) fn reserve_execution_slot( diff --git a/codex-rs/core/src/session/handlers.rs b/codex-rs/core/src/session/handlers.rs index 7730a30ad7..6a032f8166 100644 --- a/codex-rs/core/src/session/handlers.rs +++ b/codex-rs/core/src/session/handlers.rs @@ -22,6 +22,7 @@ use crate::realtime_conversation::prefix_realtime_v2_text; use crate::review_prompts::resolve_review_request; use crate::session::spawn_review_thread; use crate::tasks::CompactTask; +use crate::tasks::StartTaskOutcome; use crate::tasks::UserShellCommandMode; use crate::tasks::UserShellCommandTask; use crate::tasks::execute_user_shell_command; @@ -273,13 +274,24 @@ pub(super) async fn user_input_or_turn_inner( client_id: client_user_message_id, }); } - sess.spawn_task( - Arc::clone(¤t_context), - task_input, - crate::tasks::RegularTask::new(), - ) - .await; - Some(accepted_items) + match sess + .spawn_task( + Arc::clone(¤t_context), + task_input, + crate::tasks::RegularTask::new(), + ) + .await + { + StartTaskOutcome::Started => Some(accepted_items), + StartTaskOutcome::Rejected(err) => { + sess.send_event_raw(Event { + id: sub_id, + msg: EventMsg::Error(err.to_error_event(/*message_prefix*/ None)), + }) + .await; + None + } + } } Err(err) => { sess.send_event_raw(Event { diff --git a/codex-rs/core/src/tasks/mod.rs b/codex-rs/core/src/tasks/mod.rs index a3b9936ca9..ac0687cd1a 100644 --- a/codex-rs/core/src/tasks/mod.rs +++ b/codex-rs/core/src/tasks/mod.rs @@ -33,6 +33,7 @@ use crate::session::turn_context::TurnContext; use crate::state::ActiveTurn; use crate::state::RunningTask; use crate::state::TaskKind; +use crate::state::TurnState; use codex_analytics::TurnTokenUsageFact; use codex_login::AuthManager; use codex_models_manager::manager::SharedModelsManager; @@ -42,6 +43,7 @@ use codex_otel::TURN_MEMORY_METRIC; use codex_otel::TURN_NETWORK_PROXY_METRIC; use codex_otel::TURN_TOKEN_USAGE_METRIC; use codex_otel::TURN_TOOL_CALL_METRIC; +use codex_protocol::error::CodexErr; use codex_protocol::models::ResponseItem; use codex_protocol::protocol::EventMsg; use codex_protocol::protocol::MultiAgentVersion; @@ -70,6 +72,12 @@ pub(crate) enum InterruptedTurnHistoryMarker { Developer, } +#[derive(Debug)] +pub(crate) enum StartTaskOutcome { + Started, + Rejected(CodexErr), +} + impl InterruptedTurnHistoryMarker { pub(crate) fn from_config_and_version( config: &Config, @@ -307,7 +315,7 @@ impl Session { turn_context: Arc, input: Vec, task: T, - ) { + ) -> StartTaskOutcome { self.abort_all_tasks(TurnAbortReason::Replaced).await; self.clear_connector_selection().await; self.start_task(turn_context, input, task).await; @@ -318,7 +326,37 @@ impl Session { turn_context: Arc, input: Vec, task: T, - ) { + ) -> StartTaskOutcome { + let turn_state = { + let mut active = self.active_turn.lock().await; + let turn = active.get_or_insert_with(ActiveTurn::default); + if turn.task.is_some() { + return StartTaskOutcome::Rejected(CodexErr::InvalidRequest( + "thread already has an active turn".to_string(), + )); + } + Arc::clone(&turn.turn_state) + }; + let execution_reservation = match self + .services + .agent_control + .reserve_execution_slot_for_turn_start(self) + .await + { + Ok(reservation) => reservation, + Err(err) => { + self.clear_reserved_active_turn(&turn_state).await; + warn!( + thread_id = %self.thread_id, + "agent turn could not reserve an execution slot: {err}" + ); + return StartTaskOutcome::Rejected(err); + } + }; + if let Some(execution_reservation) = execution_reservation { + execution_reservation.commit(); + } + let task: Arc = Arc::new(task); let task_kind = task.kind(); let span_name = task.span_name(); @@ -351,12 +389,6 @@ impl Session { warn!("failed to apply goal runtime turn-start event: {err}"); } let pending_items = self.input_queue.get_pending_input(&self.active_turn).await; - let turn_state = { - let mut active = self.active_turn.lock().await; - let turn = active.get_or_insert_with(ActiveTurn::default); - debug_assert!(turn.task.is_none()); - Arc::clone(&turn.turn_state) - }; turn_state.lock().await.token_usage_at_turn_start = token_usage_at_turn_start.clone(); self.input_queue .extend_pending_input_for_turn_state(turn_state.as_ref(), pending_items) @@ -366,7 +398,12 @@ impl Session { let turn_extension_data = Arc::clone(&turn_context.extension_data); let mut active = self.active_turn.lock().await; - let turn = active.get_or_insert_with(ActiveTurn::default); + let Some(turn) = active.as_mut() else { + return StartTaskOutcome::Rejected(CodexErr::Fatal( + "active turn reservation was lost before task start".to_string(), + )); + }; + debug_assert!(Arc::ptr_eq(&turn.turn_state, &turn_state)); debug_assert!(turn.task.is_none()); let done_clone = Arc::clone(&done); let session_ctx = Arc::new(SessionTaskContext::new( @@ -442,6 +479,7 @@ impl Session { _timer: timer, }; turn.task = Some(running_task); + StartTaskOutcome::Started } /// Starts a regular turn when the session is idle and pending work is waiting. @@ -469,36 +507,22 @@ impl Session { return; } - let execution_reservation = match self - .services - .agent_control - .reserve_execution_slot_for_pending_turn(self) - .await - { - Ok(reservation) => reservation, - Err(err) => { - warn!( - thread_id = %self.thread_id, - "pending agent turn could not reserve an execution slot: {err}" - ); - return; - } - }; - - { + let turn_state = { let mut active_turn = self.active_turn.lock().await; if active_turn.is_some() { return; } - *active_turn = Some(ActiveTurn::default()); - } - if let Some(execution_reservation) = execution_reservation { - execution_reservation.commit(); - } + let active_turn = active_turn.get_or_insert_with(ActiveTurn::default); + Arc::clone(&active_turn.turn_state) + }; let turn_context = self.new_default_turn_with_sub_id(sub_id).await; self.maybe_emit_unknown_model_warning_for_turn(turn_context.as_ref()) .await; + if !self.input_queue.has_trigger_turn_mailbox_items().await { + self.clear_reserved_active_turn(&turn_state).await; + return; + } self.start_task(turn_context, Vec::new(), RegularTask::new()) .await; }) @@ -829,6 +853,16 @@ impl Session { active.take() } + async fn clear_reserved_active_turn(&self, turn_state: &Arc>) { + let mut active_turn_guard = self.active_turn.lock().await; + if let Some(active_turn) = active_turn_guard.as_ref() + && active_turn.task.is_none() + && Arc::ptr_eq(&active_turn.turn_state, turn_state) + { + *active_turn_guard = None; + } + } + pub(crate) async fn close_unified_exec_processes(&self) { self.services .unified_exec_manager