diff --git a/codex-rs/core/src/codex_thread.rs b/codex-rs/core/src/codex_thread.rs index d00c5fafe2..e4a6d9deeb 100644 --- a/codex-rs/core/src/codex_thread.rs +++ b/codex-rs/core/src/codex_thread.rs @@ -164,7 +164,7 @@ impl CodexThread { } pub async fn continue_idle_extension_work(&self) { - self.codex.session.maybe_start_idle_extension_turn().await; + crate::tasks::maybe_start_idle_extension_turn(Arc::clone(&self.codex.session)).await; } pub fn thread_extension_data(&self) -> &codex_extension_api::ExtensionData { diff --git a/codex-rs/core/src/context/extension_context.rs b/codex-rs/core/src/context/extension_context.rs index 4b9a6d0ad4..5f40992237 100644 --- a/codex-rs/core/src/context/extension_context.rs +++ b/codex-rs/core/src/context/extension_context.rs @@ -1,5 +1,6 @@ use super::ContextualUserFragment; +/// Hidden user-context fragment for extension-owned steering prompts. #[derive(Debug, Clone, PartialEq, Eq)] pub struct ExtensionContext { body: String, diff --git a/codex-rs/core/src/session/turn.rs b/codex-rs/core/src/session/turn.rs index 1ac59830be..8d8c0924fa 100644 --- a/codex-rs/core/src/session/turn.rs +++ b/codex-rs/core/src/session/turn.rs @@ -144,7 +144,8 @@ pub(crate) async fn run_turn( // diffs/full reinjection + user input) and trigger compaction preemptively // when they would push the thread over the compaction threshold. if let Err(err) = run_pre_sampling_compact(&sess, &turn_context, &mut client_session).await { - if err.to_codex_protocol_error() == CodexErrorInfo::UsageLimitExceeded { + let error_info = err.to_codex_protocol_error(); + if error_info == CodexErrorInfo::UsageLimitExceeded { sess.emit_turn_error_lifecycle( turn_context.as_ref(), CodexErrorInfo::UsageLimitExceeded, @@ -289,7 +290,8 @@ pub(crate) async fn run_turn( ) .await { - if err.to_codex_protocol_error() == CodexErrorInfo::UsageLimitExceeded { + let error_info = err.to_codex_protocol_error(); + if error_info == CodexErrorInfo::UsageLimitExceeded { sess.emit_turn_error_lifecycle( turn_context.as_ref(), CodexErrorInfo::UsageLimitExceeded, @@ -374,7 +376,8 @@ pub(crate) async fn run_turn( } Err(e) => { info!("Turn error: {e:#}"); - if e.to_codex_protocol_error() == CodexErrorInfo::UsageLimitExceeded { + let error_info = e.to_codex_protocol_error(); + if error_info == CodexErrorInfo::UsageLimitExceeded { sess.emit_turn_error_lifecycle( turn_context.as_ref(), CodexErrorInfo::UsageLimitExceeded, diff --git a/codex-rs/core/src/tasks/idle_extension.rs b/codex-rs/core/src/tasks/idle_extension.rs new file mode 100644 index 0000000000..3880379ba3 --- /dev/null +++ b/codex-rs/core/src/tasks/idle_extension.rs @@ -0,0 +1,102 @@ +//! 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 + .services + .extensions + .idle_turn_contributors() + .is_empty() + { + return; + } + + let session = Arc::clone(session); + let _handle = tokio::spawn(async move { + maybe_start_turn(session).await; + }); +} + +pub(crate) 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 b29ded1e19..72c6639da1 100644 --- a/codex-rs/core/src/tasks/mod.rs +++ b/codex-rs/core/src/tasks/mod.rs @@ -1,4 +1,5 @@ mod compact; +mod idle_extension; mod lifecycle; mod regular; mod review; @@ -32,7 +33,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; @@ -54,6 +54,7 @@ use codex_protocol::protocol::WarningEvent; use codex_features::Feature; use codex_protocol::models::ContentItem; pub(crate) use compact::CompactTask; +pub(crate) use idle_extension::maybe_start_turn as maybe_start_idle_extension_turn; pub(crate) use regular::RegularTask; pub(crate) use review::ReviewTask; pub(crate) use user_shell::UserShellCommandMode; @@ -486,107 +487,6 @@ impl Session { .await; } - pub(crate) async fn maybe_start_idle_extension_turn(self: &Arc) { - 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(request) = 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(request); - 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; @@ -869,14 +769,7 @@ impl Session { if !cleared_active_turn { return; } - self.schedule_idle_extension_turn(); - } - - fn schedule_idle_extension_turn(self: &Arc) { - let session = Arc::clone(self); - let _handle = tokio::spawn(async move { - session.maybe_start_idle_extension_turn().await; - }); + idle_extension::schedule_turn(self); } async fn take_active_turn(&self) -> Option { 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;