fix(turn-runner): keep polling the turn while delivering progress

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;

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<tokio::sync::Notify>,
+}
+
+#[async_trait::async_trait]
+impl TurnHost for SignalsAfterProgress {
+    async fn checkout(&self, request: &TurnRequest) -> Result<TurnSession, String> {
+        Ok(TurnSession {
+            session_id: request.thread_id.clone(),
+            fingerprint: Some("stable".into()),
+        })
+    }
+
+    async fn checkout_fork(&self, request: &TurnRequest) -> Result<TurnSession, String> {
+        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<RunQueue<TurnRequest>>,
+        progress: tokio::sync::mpsc::Sender<TurnSignal>,
+    ) -> Result<TurnResult, String> {
+        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<tokio::sync::Notify>,
+    terminal: Mutex<Vec<TerminalEvent>>,
+}
+
+#[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<TurnRequest>) {}
+}
+
+#[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 { .. }]
+    ));
+}

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<futures::future::BoxFuture<'_, ()>> = 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<P>,
+        delivery: futures::future::BoxFuture<'_, ()>,
+        cancellation: &crate::CancellationToken,
+        task_cancellation: &crate::CancellationToken,
+        deadline: &TurnDeadline,
+    ) -> Option<Result<TurnResult, String>> {
+        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<P>,