diff --git a/crates/tinyagents-harness/src/turn_runner/mod_tests.rs b/crates/tinyagents-harness/src/turn_runner/mod_tests.rs index 16cb3e6cf..74a5077e0 100644 --- a/crates/tinyagents-harness/src/turn_runner/mod_tests.rs +++ b/crates/tinyagents-harness/src/turn_runner/mod_tests.rs @@ -2712,3 +2712,6 @@ async fn cancelled_run_unwinds_within_the_cancel_grace() { async fn cancelled_run_is_dropped_at_once_without_a_cancel_grace() { assert!(!cancel_unwinding_turn(std::time::Duration::ZERO).await); } + +#[path = "progress_delivery_tests.rs"] +mod progress_delivery; diff --git a/crates/tinyagents-harness/src/turn_runner/progress_delivery_tests.rs b/crates/tinyagents-harness/src/turn_runner/progress_delivery_tests.rs new file mode 100644 index 000000000..cc0875f42 --- /dev/null +++ b/crates/tinyagents-harness/src/turn_runner/progress_delivery_tests.rs @@ -0,0 +1,119 @@ +//! The runner keeps polling the turn while it hands a progress event to the +//! sink. +//! +//! A sink's delivery can depend on the turn moving: directly, or through a +//! resource the turn's future already holds or is queued for (a lock whose +//! permit was handed to a waiter inside the turn). If the runner awaited the +//! delivery without polling the turn, the two would wait on each other for +//! good, with every deadline inside the turn frozen with it. + +use super::*; +use crate::progress::AgentProgress; + +/// Sends one progress event, yields, then signals that it was polled again. +struct SignalsAfterProgress { + polled_again: Arc, +} + +#[async_trait::async_trait] +impl TurnHost for SignalsAfterProgress { + async fn checkout(&self, request: &TurnRequest) -> Result { + Ok(TurnSession { + session_id: request.thread_id.clone(), + fingerprint: Some("stable".into()), + }) + } + + async fn checkout_fork(&self, request: &TurnRequest) -> Result { + self.checkout(request).await + } + + async fn validate_fingerprint( + &self, + _request: &TurnRequest, + _session: &TurnSession, + ) -> Result<(), String> { + Ok(()) + } + + fn run_origin(&self, request: &TurnRequest) -> serde_json::Value { + serde_json::json!({"request": request.request_id}) + } + + async fn run( + &self, + _request: &TurnRequest, + _session: &TurnSession, + _origin: &serde_json::Value, + _cancellation: crate::CancellationToken, + _queue: Arc>, + progress: tokio::sync::mpsc::Sender, + ) -> Result { + progress + .send(TurnSignal::Progress(Box::new(AgentProgress::TurnStarted))) + .await + .map_err(|_| "progress channel closed".to_owned())?; + // Hand control back so the runner sees the event before the turn + // takes its next step. + tokio::task::yield_now().await; + self.polled_again.notify_one(); + Ok(TurnResult { + reply: "done".into(), + data: serde_json::Value::Null, + }) + } + + async fn checkin(&self, _request: &TurnRequest) {} + + async fn cancel_children(&self, _request: &TurnRequest) {} +} + +/// Delivers a progress event only once the turn has taken its next step. +struct WaitsForTheTurn { + polled_again: Arc, + terminal: Mutex>, +} + +#[async_trait::async_trait] +impl EventSink for WaitsForTheTurn { + async fn progress(&self, _request: &TurnRequest, _event: AgentProgress) { + self.polled_again.notified().await; + } + + async fn terminal(&self, _request: &TurnRequest, event: TerminalEvent) { + self.terminal.lock().unwrap().push(event); + } + + async fn collected(&self, _request: &TurnRequest, _items: Vec) {} +} + +#[tokio::test] +async fn a_progress_delivery_that_waits_on_the_turn_does_not_freeze_it() { + let polled_again = Arc::new(tokio::sync::Notify::new()); + let events = Arc::new(WaitsForTheTurn { + polled_again: polled_again.clone(), + terminal: Mutex::default(), + }); + let runner = TurnRunner::new( + Arc::new(SignalsAfterProgress { polled_again }), + Arc::new(Ingress::default()), + Arc::new(Replies::default()), + events.clone(), + Arc::new(Failures), + Arc::new(Summary), + TurnRunnerConfig::default(), + ); + + let outcome = tokio::time::timeout( + Duration::from_secs(10), + runner.run(request(QueueDisposition::Interrupt, false)), + ) + .await + .expect("the turn must keep running while its progress is delivered"); + + assert!(outcome.is_ok(), "{outcome:?}"); + assert!(matches!( + events.terminal.lock().unwrap().as_slice(), + [TerminalEvent::Completed { .. }] + )); +} diff --git a/crates/tinyagents-harness/src/turn_runner/types.rs b/crates/tinyagents-harness/src/turn_runner/types.rs index b737b930b..e9b61e9c3 100644 --- a/crates/tinyagents-harness/src/turn_runner/types.rs +++ b/crates/tinyagents-harness/src/turn_runner/types.rs @@ -1290,6 +1290,14 @@ where let mut work_in_flight = 0usize; let mut progress_open = true; tokio::pin!(run); + // The progress event being handed to the sink. It is a branch of the + // loop below, never awaited inline, so the turn keeps being polled + // while it is delivered: a delivery can need the turn to move (directly, + // or through a lock permit already handed to a waiter inside it), and + // awaiting it inline would freeze both, every deadline inside the turn + // included. No further progress is received until it lands, which keeps + // the bounded channel's backpressure on the host. + let mut delivering: Option> = None; loop { if deadline .at @@ -1328,47 +1336,42 @@ where result = &mut run => { // A host can emit progress and finish in the same poll. Close // the channel before draining so retained child senders - // cannot extend completion indefinitely, then deliver every - // event already accepted by the bounded queue. + // cannot extend completion indefinitely, then finish the + // event in flight and deliver every event already accepted + // by the bounded queue, in order. progress_rx.close(); + if let Some(delivery) = delivering.take() + && let Some(stopped) = self + .deliver_within_turn(request, delivery, &cancellation, &task_cancellation, &deadline) + .await + { + return stopped; + } while let Ok(signal) = progress_rx.try_recv() { - if let TurnSignal::Progress(event) = signal { - tokio::select! { - biased; - _ = cancellation.cancelled() => { - task_cancellation.cancel(); - self.host.cancel_children(request).await; - return Err("turn cancelled".into()); - } - _ = async { if let Some(deadline) = deadline.at { tokio::time::sleep_until(deadline).await } else { std::future::pending().await } }, if deadline.at.is_some() => { - task_cancellation.cancel(); - self.host.cancel_children(request).await; - return self.close_or_timeout(request, "wall-clock deadline", &cancellation).await; - } - _ = self.events.progress(request, *event) => {} - } + if let TurnSignal::Progress(event) = signal + && let Some(stopped) = self + .deliver_within_turn( + request, + self.events.progress(request, *event), + &cancellation, + &task_cancellation, + &deadline, + ) + .await + { + return stopped; } } return result; } - signal = progress_rx.recv(), if progress_open => match signal { + () = async { if let Some(delivery) = delivering.as_mut() { delivery.await } }, if delivering.is_some() => { + delivering = None; + } + signal = progress_rx.recv(), if progress_open && delivering.is_none() => match signal { Some(signal) => match signal { TurnSignal::Progress(event) => { last_progress = Instant::now(); - tokio::select! { - biased; - _ = cancellation.cancelled() => { - task_cancellation.cancel(); - self.host.cancel_children(request).await; - return Err("turn cancelled".into()); - } - _ = async { if let Some(deadline) = deadline.at { tokio::time::sleep_until(deadline).await } else { std::future::pending().await } }, if deadline.at.is_some() => { - task_cancellation.cancel(); - self.host.cancel_children(request).await; - return self.close_or_timeout(request, "wall-clock deadline", &cancellation).await; - } - _ = self.events.progress(request, *event) => {} - } + delivering = Some(self.events.progress(request, *event)); } TurnSignal::WorkStarted => work_in_flight = work_in_flight.saturating_add(1), TurnSignal::WorkFinished => { @@ -1380,7 +1383,7 @@ where }, None => progress_open = false, }, - _ = tokio::time::sleep(idle_left), if work_in_flight == 0 => { + _ = tokio::time::sleep(idle_left), if work_in_flight == 0 && delivering.is_none() => { task_cancellation.cancel(); self.host.cancel_children(request).await; return self.close_or_timeout(request, "idle timeout", &cancellation).await; @@ -1389,6 +1392,33 @@ where } } + /// Delivers one progress event after the turn finished, still honouring + /// cancellation and the wall-clock deadline. `Some` is the turn's outcome + /// when either stops it first. + async fn deliver_within_turn( + &self, + request: &TurnRequest

, + delivery: futures::future::BoxFuture<'_, ()>, + cancellation: &crate::CancellationToken, + task_cancellation: &crate::CancellationToken, + deadline: &TurnDeadline, + ) -> Option> { + tokio::select! { + biased; + _ = cancellation.cancelled() => { + task_cancellation.cancel(); + self.host.cancel_children(request).await; + Some(Err("turn cancelled".into())) + } + _ = async { if let Some(deadline) = deadline.at { tokio::time::sleep_until(deadline).await } else { std::future::pending().await } }, if deadline.at.is_some() => { + task_cancellation.cancel(); + self.host.cancel_children(request).await; + Some(self.close_or_timeout(request, "wall-clock deadline", cancellation).await) + } + () = delivery => None, + } + } + async fn close_or_timeout( &self, request: &TurnRequest

,