harness: race in-flight unary model calls against cancellation

src/harness/agent_loop/model_call.rs

@@ -120,8 +120,23 @@ impl<State: Send + Sync, Ctx: Send + Sync> AgentHarness<State, Ctx> {
                         self.invoke_model_streaming_once(state, ctx, &model, request, call_id);
                     Self::with_call_budget(remaining, run_id.as_str(), "model call", fut).await
                 } else {
+                    // Race the wall-clock-bounded unary call against cooperative
+                    // cancellation, mirroring the streaming path: a cancel
+                    // requested while a buffered (non-streamed) provider call is
+                    // in flight drops the future — reqwest cancels the underlying
+                    // request — and unwinds with `Cancelled` instead of paying
+                    // for the call to run to completion. `cancelled()` is
+                    // cancel-safe, and the pre-call `is_cancelled()` check above
+                    // still short-circuits before the request is ever issued.
+                    let cancellation = ctx.cancellation.clone();
                     let fut = model.invoke(state, request.clone());
-                    Self::with_call_budget(remaining, run_id.as_str(), "model call", fut).await
+                    let budgeted =
+                        Self::with_call_budget(remaining, run_id.as_str(), "model call", fut);
+                    tokio::select! {
+                        biased;
+                        _ = cancellation.cancelled() => Err(TinyAgentsError::Cancelled),
+                        result = budgeted => result,
+                    }
                 };
                 match attempt_result {
                     Ok(response) => break Ok(response),

src/harness/agent_loop/test.rs

@@ -1267,6 +1267,58 @@ async fn cancelled_mid_run_stops_before_next_model_call() {
     assert_eq!(*invocations.lock().unwrap(), 1);
 }
 
+/// A model whose unary `invoke` never returns on its own, signalling once it has
+/// started so a test can cancel the run while the call is genuinely in flight.
+struct BlockForeverModel {
+    started: Arc<tokio::sync::Notify>,
+}
+
+#[async_trait]
+impl ChatModel<()> for BlockForeverModel {
+    async fn invoke(&self, _state: &(), _request: ModelRequest) -> Result<ModelResponse> {
+        self.started.notify_one();
+        // Simulate a long buffered (non-streamed) provider call that only ends
+        // when the caller drops this future. Without the loop racing
+        // cancellation against the in-flight call, the run would hang here.
+        std::future::pending::<()>().await;
+        unreachable!("pending future never resolves")
+    }
+}
+
+#[tokio::test]
+async fn cancelled_during_unary_model_call_drops_the_in_flight_call() {
+    let token = CancellationToken::new();
+    let started = Arc::new(tokio::sync::Notify::new());
+
+    let mut harness: AgentHarness<()> = AgentHarness::new();
+    harness.register_model(
+        "mock",
+        Arc::new(BlockForeverModel {
+            started: started.clone(),
+        }),
+    );
+
+    let ctx = RunContext::new(RunConfig::new("cancel-unary"), ()).with_cancellation(token.clone());
+
+    // Cancel only once the model call has actually begun, so the pre-call
+    // checkpoint cannot short-circuit and we exercise the in-flight race.
+    let canceller = tokio::spawn(async move {
+        started.notified().await;
+        token.cancel();
+    });
+
+    let err = tokio::time::timeout(
+        std::time::Duration::from_secs(5),
+        harness.invoke_in_context(&(), ctx, vec![Message::user("hi")]),
+    )
+    .await
+    .expect("run must not hang: a cancel mid unary call must drop the in-flight future")
+    .expect_err("a run cancelled mid model call must not complete");
+
+    assert!(matches!(err, TinyAgentsError::Cancelled), "got {err:?}");
+    canceller.await.unwrap();
+}
+
 // ── Per-model-call timeout ─────────────────────────────────────────────────────
 
 #[tokio::test]

src/harness/model/mod.rs

@@ -753,10 +753,10 @@ impl StreamAccumulator {
     fn push_tool_chunk(&mut self, call_id: &str, content: &str, tool_name: Option<&str>) {
         if let Some(entry) = self.tool_chunks.iter_mut().find(|(id, ..)| id == call_id) {
             entry.1.push_str(content);
-            if entry.2.is_none() {
-                if let Some(name) = tool_name.filter(|n| !n.is_empty()) {
-                    entry.2 = Some(name.to_string());
-                }
+            if entry.2.is_none()
+                && let Some(name) = tool_name.filter(|n| !n.is_empty())
+            {
+                entry.2 = Some(name.to_string());
             }
         } else {
             self.tool_chunks.push((

src/harness/stream/mod.rs

@@ -143,6 +143,5 @@ pub fn stream(chunks: &[StreamChunk], modes: &[StreamMode]) -> Vec<StreamChunk>
         .collect()
 }
 
-#[cfg(test)]
 #[cfg(test)]
 mod test;