Avoid double polling idle turn contributors

This commit is contained in:
Eric Traut
2026-05-27 20:58:13 -07:00
parent dcf0f74199
commit d0727696da
2 changed files with 46 additions and 12 deletions

View File

@@ -52,8 +52,13 @@ pub(crate) async fn maybe_start_turn(session: Arc<Session>) {
return;
}
// Probe before reserving so a thread with no idle-turn work stays freely available.
if next_idle_turn_candidate(&session).await.is_none() {
// Produce at most one idle-turn request. The validation below is the stale-request check.
let Some(candidate) = next_idle_turn_candidate(&session).await else {
return;
};
if has_active_or_pending_work(&session).await {
session.maybe_start_turn_for_pending_work().await;
return;
}
@@ -63,13 +68,6 @@ pub(crate) async fn maybe_start_turn(session: Arc<Session>) {
return;
}
// Re-read the idle request after reservation. This drops any stale prompt that was produced
// before a goal/status/config change raced with the scheduler.
let Some(candidate) = next_idle_turn_candidate(&session).await else {
clear_reserved_idle_turn(&session).await;
return;
};
if has_pending_work(&session).await {
clear_reserved_idle_turn(&session).await;
session.maybe_start_turn_for_pending_work().await;
@@ -236,13 +234,30 @@ fn idle_extension_input(request: ThreadIdleRequest) -> TurnInput {
#[cfg(test)]
mod tests {
use std::sync::Arc;
use std::sync::atomic::AtomicUsize;
use std::sync::atomic::Ordering;
use codex_extension_api::HiddenContext;
use codex_extension_api::HiddenContextMarker;
use codex_protocol::models::ContentItem;
use codex_protocol::models::ResponseInputItem;
use pretty_assertions::assert_eq;
use super::*;
struct CountingIdleTurnContributor(AtomicUsize);
impl codex_extension_api::ThreadIdleTurnContributor for CountingIdleTurnContributor {
fn request_thread_idle_turn<'a>(
&'a self,
_input: codex_extension_api::ThreadIdleInput<'a>,
) -> codex_extension_api::ThreadIdleTurnRequestFuture<'a> {
self.0.fetch_add(1, Ordering::SeqCst);
Box::pin(std::future::ready(None))
}
}
#[test]
fn idle_extension_input_truncates_large_prompts() {
let prompt = format!(
@@ -269,4 +284,19 @@ mod tests {
assert!(text.contains("end"));
assert!(text.len() < original_len);
}
#[tokio::test]
async fn idle_scheduler_requests_idle_turn_once() {
let (mut session, _turn_context) = crate::session::tests::make_session_and_context().await;
let contributor = Arc::new(CountingIdleTurnContributor(AtomicUsize::new(0)));
let mut builder = codex_extension_api::ExtensionRegistryBuilder::new();
builder.thread_idle_turn_contributor(contributor.clone());
session.services.extensions = Arc::new(builder.build());
let session = Arc::new(session);
maybe_start_turn(Arc::clone(&session)).await;
assert_eq!(contributor.0.load(Ordering::SeqCst), 1);
assert!(session.active_turn.lock().await.is_none());
}
}

View File

@@ -81,9 +81,9 @@ pub type ThreadIdleTurnStartFuture<'a> = std::pin::Pin<Box<dyn Future<Output = b
///
/// Implementations should normally return hidden context. Raw items are
/// available for extensions that intentionally need unwrapped input. The host
/// owns hidden-context wrapping, applies the declared idle-turn policy before
/// calling into the contributor, and may ask the contributor to confirm the
/// request again immediately before the turn starts.
/// owns hidden-context wrapping, enforces the declared idle-turn policy before
/// calling into the contributor, and validates returned requests immediately
/// before starting a turn.
pub trait ThreadIdleTurnContributor: Send + Sync {
/// Returns the host scheduling policy for this contributor's idle turns.
fn idle_turn_policy(&self) -> IdleTurnPolicy {
@@ -92,6 +92,10 @@ pub trait ThreadIdleTurnContributor: Send + Sync {
/// Returns input to start an idle turn, if the extension still has work
/// that should run without user input.
///
/// The host calls this at most once per idle scheduling attempt. A returned
/// request can still be discarded if user work wins or validation fails, so
/// implementations should not treat this as a start notification.
fn request_thread_idle_turn<'a>(
&'a self,
_input: ThreadIdleInput<'a>,