From db1111720acb2909b0d94fde53075df6b803d411 Mon Sep 17 00:00:00 2001 From: Eric Traut Date: Wed, 27 May 2026 06:58:57 -0700 Subject: [PATCH] Simplify idle extension turn scheduling --- codex-rs/core/src/tasks/idle_extension.rs | 85 +++++++++++++- codex-rs/core/src/tasks/mod.rs | 105 ------------------ .../ext/extension-api/src/contributors.rs | 2 + 3 files changed, 86 insertions(+), 106 deletions(-) diff --git a/codex-rs/core/src/tasks/idle_extension.rs b/codex-rs/core/src/tasks/idle_extension.rs index 03a3953b34..71d15632b3 100644 --- a/codex-rs/core/src/tasks/idle_extension.rs +++ b/codex-rs/core/src/tasks/idle_extension.rs @@ -1,6 +1,12 @@ +//! Scheduling glue for extension-owned turns that should run when a session is idle. + use std::sync::Arc; +use crate::session::TurnInput; use crate::session::session::Session; +use crate::state::ActiveTurn; + +use super::RegularTask; pub(super) fn schedule_turn(session: &Arc) { if session @@ -14,6 +20,83 @@ pub(super) fn schedule_turn(session: &Arc) { let session = Arc::clone(session); let _handle = tokio::spawn(async move { - session.maybe_start_idle_extension_turn().await; + maybe_start_turn(session).await; }); } + +async fn maybe_start_turn(session: Arc) { + // Give queued user-visible work the first chance to wake the session. + session.maybe_start_turn_for_pending_work().await; + if has_active_or_pending_work(&session).await { + return; + } + + // Ask extensions for one idle turn only while the session still looks quiet. + let Some(input) = next_idle_turn_input(&session).await else { + return; + }; + + // Extension callbacks can race with user input or mailbox delivery. Re-check before + // claiming the idle slot, and let the normal pending-work path handle anything that appeared. + if has_active_or_pending_work(&session).await { + session.maybe_start_turn_for_pending_work().await; + return; + } + + // Reserve the active turn after the final quiet check so another task cannot start while + // the turn context is being built. + { + let mut active_turn = session.active_turn.lock().await; + if active_turn.is_some() { + return; + } + *active_turn = Some(ActiveTurn::default()); + } + + // Treat extension-provided items as the new turn's initial input rather than stashing them + // in turn-state pending input; this keeps rollback logic out of the scheduler. + let turn_context = session + .new_default_turn_with_sub_id(uuid::Uuid::new_v4().to_string()) + .await; + session + .maybe_emit_unknown_model_warning_for_turn(turn_context.as_ref()) + .await; + session + .start_task(turn_context, input, RegularTask::new()) + .await; +} + +async fn has_active_or_pending_work(session: &Session) -> bool { + session.active_turn.lock().await.is_some() + || session + .input_queue + .has_queued_response_items_for_next_turn() + .await + || session.input_queue.has_trigger_turn_mailbox_items().await +} + +async fn next_idle_turn_input(session: &Session) -> Option> { + let collaboration_mode = session.collaboration_mode().await; + for contributor in session.services.extensions.idle_turn_contributors() { + let Some(items) = contributor + .next_idle_turn(codex_extension_api::IdleTurnInput { + collaboration_mode: &collaboration_mode, + session_store: &session.services.session_extension_data, + thread_store: &session.services.thread_extension_data, + }) + .await + else { + continue; + }; + if !items.is_empty() { + return Some( + items + .into_iter() + .map(TurnInput::ResponseInputItem) + .collect(), + ); + } + } + + None +} diff --git a/codex-rs/core/src/tasks/mod.rs b/codex-rs/core/src/tasks/mod.rs index 3a5484fe75..bda60e0121 100644 --- a/codex-rs/core/src/tasks/mod.rs +++ b/codex-rs/core/src/tasks/mod.rs @@ -34,7 +34,6 @@ 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; @@ -497,110 +496,6 @@ impl Session { .await; } - pub(crate) async fn maybe_start_idle_extension_turn(self: &Arc) { - if self.services.extensions.idle_turn_contributors().is_empty() { - return; - } - self.maybe_start_turn_for_pending_work().await; - if self.active_turn.lock().await.is_some() { - return; - } - if self - .input_queue - .has_queued_response_items_for_next_turn() - .await - || self.input_queue.has_trigger_turn_mailbox_items().await - { - return; - } - - let items = { - let collaboration_mode = self.collaboration_mode().await; - let mut requested_items = None; - for contributor in self.services.extensions.idle_turn_contributors() { - if let Some(items) = contributor - .next_idle_turn(codex_extension_api::IdleTurnInput { - collaboration_mode: &collaboration_mode, - session_store: &self.services.session_extension_data, - thread_store: &self.services.thread_extension_data, - }) - .await - { - requested_items = Some(items); - break; - } - } - requested_items - }; - let Some(items) = items.filter(|items| !items.is_empty()) else { - return; - }; - - let turn_state = { - let mut active_turn = self.active_turn.lock().await; - if active_turn.is_some() { - return; - } - let active_turn = active_turn.get_or_insert_with(ActiveTurn::default); - Arc::clone(&active_turn.turn_state) - }; - if self - .input_queue - .has_queued_response_items_for_next_turn() - .await - || self.input_queue.has_trigger_turn_mailbox_items().await - { - self.clear_reserved_idle_extension_turn(&turn_state).await; - self.maybe_start_turn_for_pending_work().await; - return; - } - - self.input_queue - .extend_pending_input_for_turn_state( - turn_state.as_ref(), - items - .into_iter() - .map(TurnInput::ResponseInputItem) - .collect(), - ) - .await; - - let turn_context = self - .new_default_turn_with_sub_id(uuid::Uuid::new_v4().to_string()) - .await; - self.maybe_emit_unknown_model_warning_for_turn(turn_context.as_ref()) - .await; - let still_reserved = { - let active_turn = self.active_turn.lock().await; - active_turn.as_ref().is_some_and(|active_turn| { - active_turn.task.is_none() && Arc::ptr_eq(&active_turn.turn_state, &turn_state) - }) - }; - if !still_reserved { - self.input_queue - .take_pending_input_for_turn_state(turn_state.as_ref()) - .await; - self.clear_reserved_idle_extension_turn(&turn_state).await; - return; - } - drop(turn_state); - self.start_task(turn_context, Vec::new(), RegularTask::new()) - .await; - } - - async fn clear_reserved_idle_extension_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 async fn abort_all_tasks(self: &Arc, reason: TurnAbortReason) { let mut aborted_turn = false; let mut active_turn_to_clear = None; diff --git a/codex-rs/ext/extension-api/src/contributors.rs b/codex-rs/ext/extension-api/src/contributors.rs index ee15b82bc3..e01552c9c1 100644 --- a/codex-rs/ext/extension-api/src/contributors.rs +++ b/codex-rs/ext/extension-api/src/contributors.rs @@ -1,3 +1,5 @@ +//! Contributor traits and inputs that let extensions hook into host-owned workflows. + use std::future::Future; use std::pin::Pin; use std::sync::Arc;