diff --git a/src/harness/agent_loop/model_call.rs b/src/harness/agent_loop/model_call.rs index b423c8f..59dadbe 100644 --- a/src/harness/agent_loop/model_call.rs +++ b/src/harness/agent_loop/model_call.rs @@ -120,8 +120,23 @@ impl AgentHarness { 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), diff --git a/src/harness/agent_loop/test.rs b/src/harness/agent_loop/test.rs index 9627f7a..cfc7376 100644 --- a/src/harness/agent_loop/test.rs +++ b/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, +} + +#[async_trait] +impl ChatModel<()> for BlockForeverModel { + async fn invoke(&self, _state: &(), _request: ModelRequest) -> Result { + 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] diff --git a/src/harness/model/mod.rs b/src/harness/model/mod.rs index aa24cb1..ea3e6b8 100644 --- a/src/harness/model/mod.rs +++ b/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(( diff --git a/src/harness/stream/mod.rs b/src/harness/stream/mod.rs index 14411f0..1712142 100644 --- a/src/harness/stream/mod.rs +++ b/src/harness/stream/mod.rs @@ -143,6 +143,5 @@ pub fn stream(chunks: &[StreamChunk], modes: &[StreamMode]) -> Vec .collect() } -#[cfg(test)] #[cfg(test)] mod test;