fix: drain streaming cancel tasks before completion#3690
Conversation
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 9eb06ec020
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| elif self._cancel_mode == "immediate": | ||
| await self._drain_cancelled_tasks() | ||
| self._check_errors() | ||
| self.is_complete = True |
There was a problem hiding this comment.
When stream_events() itself is cancelled (for example an asyncio.wait_for(....__anext__()) timeout), the except asyncio.CancelledError path above calls self.cancel(), but cancelled stays true so this new immediate-cancel drain/completion branch is skipped. Since cancel() no longer sets is_complete synchronously and the final queue drain removes the sentinel it enqueued, a result whose run loop has not already posted its own sentinel can be left incomplete with an empty queue; following the documented advice to continue consuming stream_events() can then wait forever. Please also put the result into a terminal state in the cancelled path without blocking on the drain.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: fd017a0b70
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| if not cancelled and self._cancel_mode != "immediate": | ||
| await self._run_sandbox_cleanup() |
There was a problem hiding this comment.
When a sandboxed streamed run has already finished its run loop but is still in ensure_sandbox_cleanup_on_completion()'s wrapper awaiting _run_sandbox_cleanup(), calling cancel(mode="immediate") cancels run_loop_task; because that wrapper is awaiting the cleanup task, the cleanup task is cancelled too. This new guard then skips _run_sandbox_cleanup() in stream_events(), so runner-owned sandbox sessions can miss shutdown(), client delete(), and dependency close (for example, a fast model enqueues multiple events and the consumer cancels after the first event while cleanup is running). Immediate cancel should return promptly, but cleanup still needs to stay scheduled or be retried/shielded.
Useful? React with 👍 / 👎.
seratch
changed the title
2 participants
Summary
Tests
uv run pytest tests/test_cancel_streaming.py::test_cancel_immediate_drains_owned_tasks_before_marking_complete tests/sandbox/test_runtime.py::test_runner_streamed_immediate_cancel_skips_waiting_for_sandbox_cleanupuv run pytest tests/test_cancel_streaming.pyuv run ruff check src/agents/result.py tests/test_cancel_streaming.py