fix(run): cancel streamed models when input guardrails fail - #4004
Conversation
seratch
left a comment
There was a problem hiding this comment.
Thanks for the contribution. The underlying deadlock is valid, and the guardrail-first fix is directionally correct. One interleaving still needs to be addressed before merge.
If the model fails first, the run loop marks the result complete and enters its finally block, where it awaits the input-guardrail task. If that guardrail then raises, the new unconditional run_loop_task.cancel() cancels a run loop that is already finalizing, which can skip computer disposal and span/trace completion.
Please cancel the run loop only while streamed_result.is_complete is false, and add a controlled model-first/guardrail-second regression test that verifies finalization still completes and the run-loop task is not cancelled. Once that is covered, this should be ready for another review.
Summary
This pull request fixes a streamed-run deadlock when a parallel input guardrail raises an ordinary exception before the model emits an event.
The streamed guardrail task now cancels the in-flight run loop and wakes the event consumer after draining sibling guardrails. This lets
stream_events()propagate the original guardrail exception and await model cancellation instead of hanging on an empty queue and leaving an unretrieved task exception.Test plan
upstream/mainwith a one-second timeout instream_events()and an unretrievedValueError.uv run pytest -q tests/test_guardrails.py::test_parallel_guardrail_error_cancels_streaming_model(5 consecutive passes)uv run pytest -q tests/test_guardrails.py(44 passed)uv run pytest -q tests/test_agent_runner_streamed.py(55 passed)uv run pytest -q tests/test_cancel_streaming.py(11 passed)bash .agents/skills/code-change-verification/scripts/run.shmake format-checkmake tests-asyncio-stabilitymake coverage(5,806 passed, 7 skipped; 90% coverage)Issue number
N/A — no pre-filed issue is required for this focused bug fix.
Checks
.agents/skills/code-change-verification/scripts/run.sh/reviewbefore submitting this PR