harness: race in-flight unary model calls against cancellation

[M3gA-Mind]

Summary

In the model-call retry loop, the streaming branch races the provider call against the run's cancellation token, but the unary (non-streamed) branch did not — it awaited model.invoke(...) bounded only by the wall-clock timeout. Cancelling a run during a long buffered call (structured-output extraction, summarizer sub-calls, any non-streaming provider) did not drop the in-flight request; it ran to completion and was billed until it returned or the deadline fired. Streaming turns cancelled promptly, unary turns didn't.

Change

Race the wall-clock-bounded unary call against ctx.cancellation.cancelled() with a biased select! (mirroring the streaming path), returning TinyAgentsError::Cancelled and dropping the future — reqwest cancels the underlying request. The pre-call is_cancelled() checkpoint still short-circuits before the request is issued.

Tests

cancelled_during_unary_model_call_drops_the_in_flight_call: a model whose unary invoke blocks on future::pending is cancelled mid-flight; without the race the run would hang (guarded by an outer 5s timeout), with it the run unwinds as Cancelled. cargo fmt --check, cargo clippy --lib --tests -- -D warnings, and cargo test --lib all pass.

Closes tinyhumansai/openhuman#4635

Summary by CodeRabbit

• Bug Fixes

  • Improved cancellation handling for in-flight model requests so runs stop immediately when cancelled instead of waiting for the request to finish.
  • Fixed unary model calls to return a cancelled result more reliably during long-running operations.

• Tests

  • Added regression coverage to verify cancellation works correctly while a unary model call is actively running.

[coderabbitai]

Warning

Review limit reached

You’ve reached a temporary PR review limit under our Fair Usage Limits Policy.

Your recent review volume is higher than typical usage, so adaptive limits are currently applied.

Next review available in: 33 minutes

Enable usage-based reviews in Billing to review now. Otherwise, wait until the next included review is available.
You're only billed for reviews past your plan's rate limits ($0.25/file).

How can I continue?

After more reviews become available, a review can be triggered using the @coderabbitai review command as a PR comment. Alternatively, push new commits to this PR.

To avoid repeated limits, reduce automatic review volume by pausing incremental auto-reviews earlier, using label-based review opt-in, excluding WIP or generated PR titles, or requesting reviews manually when the PR is ready. If your team needs uninterrupted high-volume reviews, an organization admin can enable usage-based reviews.

How do review limits work?

CodeRabbit enforces per-developer PR review limits for each organization. Most developers receive the normal plan review availability.

For paid Pro and Pro+ PR reviews, CodeRabbit uses adaptive limits for sustained high-volume activity. When a developer's recent PR review activity reaches the 95th percentile or higher among CodeRabbit users, additional reviews become available more gradually as earlier reviews age out of the rolling window.

Please refer docs for additional details.

Review details

⚙️ Run configuration

Configuration used: Organization UI

Review profile: CHILL

Plan: Pro

Run ID: 7d69ff0b-1fd1-44fd-90ea-9561cc3c567b

📥 Commits

Reviewing files that changed from the base of the PR and between 5508109 and 5b83664.

📒 Files selected for processing (2)

• src/harness/model/mod.rs
• src/harness/stream/mod.rs

📝 Walkthrough

Walkthrough

The unary (non-streaming) model invocation path now races the wall-clock budgeted call against cooperative cancellation via tokio::select!, returning Cancelled immediately instead of waiting for the in-flight call. A new regression test verifies this behavior using a model that blocks indefinitely until cancelled.

Changes

Unary cancellation fix and test

Layer / File(s)	Summary
Race unary call against cancellation src/harness/agent_loop/model_call.rs	Wraps the budgeted unary provider future in tokio::select! against ctx.cancellation.cancelled(), returning TinyAgentsError::Cancelled when cancellation fires during an in-flight non-streamed call.
Regression test for dropped in-flight call src/harness/agent_loop/test.rs	Adds BlockForeverModel that signals start then blocks forever, and a new test that cancels after invoke begins, asserting Cancelled is returned and the call does not hang.

Estimated code review effort: 2 (Simple) | ~15 minutes

Sequence Diagram(s)

sequenceDiagram
    participant RetryLoop as Model Call Retry Loop
    participant BudgetWrapper as with_call_budget Future
    participant Model as Model Provider
    participant Cancellation as ctx.cancellation

    RetryLoop->>BudgetWrapper: await budgeted invoke()
    BudgetWrapper->>Model: invoke(request)
    par racing
        Cancellation-->>RetryLoop: cancelled() fires
    and
        Model-->>BudgetWrapper: response (in-flight)
    end
    alt cancellation fired first
        RetryLoop->>RetryLoop: return TinyAgentsError::Cancelled
    else budget future completes first
        BudgetWrapper-->>RetryLoop: return response
    end

Loading

Possibly related issues

• [TinyAgents][T1] Unary model calls ignore mid-flight cancellation — billed to completion openhuman#4635: Directly implements the unary mid-flight cancellation fix described, racing the budgeted call against cancellation as requested.

Poem

A rabbit hopped mid-call one day,
"Cancelled!" it cried, and hopped away 🐇
No more billing while we wait,
Unary calls now share the fate
Of streams that stop when told "not yet" —
Fair cancellation, no regret!

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The title clearly matches the main change: unary model calls now race cancellation while in flight.
Linked Issues check	✅ Passed	The PR implements the linked issue by racing unary calls against cancellation and adding a regression test.
Out of Scope Changes check	✅ Passed	The changes stay focused on unary cancellation handling and its test, with no unrelated scope visible.
Docstring Coverage	✅ Passed	No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.

---

_{Comment @coderabbitai help to get the list of available commands.}

[M3gA-Mind]

CI guardian: the Rust SDK check was red on a pre-existing main clippy failure under stable rust 1.96 (a collapsible_if in StreamAccumulator::push_tool_chunk and a duplicated #[cfg(test)] on the stream test module), inherited by this branch. Pushed a behavior-neutral lint fix so CI can go green; no change to this PR's own logic.