Merge branch 'etraut/goal-extension-hooks-step1' into etraut/extract-goal-extension

# Conflicts:
#	codex-rs/core/src/context/contextual_user_message_tests.rs
#	codex-rs/core/src/context/extension_context.rs
#	codex-rs/core/src/session/turn.rs
#	codex-rs/core/src/tasks/mod.rs
This commit is contained in:
Eric Traut
2026-05-27 07:16:26 -07:00
6 changed files with 115 additions and 114 deletions

View File

@@ -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 {

View File

@@ -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,

View File

@@ -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,

View File

@@ -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<Session>) {
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<Session>) {
// 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<Vec<TurnInput>> {
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
}

View File

@@ -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>) {
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<tokio::sync::Mutex<TurnState>>,
) {
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<Self>, 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<Self>) {
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<ActiveTurn> {

View File

@@ -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;