fix: stop waiting on Queue.join() when the event consumer can die - #4199
fix: stop waiting on Queue.join() when the event consumer can die#4199abhay-codes07 wants to merge 1 commit into
Conversation
There was a problem hiding this comment.
Pull request overview
Fixes a deadlock in several streaming/background-consumer patterns where the producer waits on asyncio.Queue.join() even if the consumer task can terminate early (notably on BaseException such as asyncio.CancelledError), causing the run/tool call to hang indefinitely. The PR removes the redundant Queue.join() waits and relies on awaiting the consumer task (which already implies the sentinel has been consumed) to both avoid hangs and surface dead consumers.
Changes:
- Replace
await queue.join()with awaiting the consumer/worker task in three call sites to prevent permanent hangs when the consumer dies. - Add regression tests for
Agent.as_tool(on_stream=...)and Codex_consume_events()to ensureBaseException/CancelledErrorin handlers does not deadlock. - Add an ordering/property test to ensure dropping
join()does not let tool invocation return before streamed events are handled.
Reviewed changes
Copilot reviewed 5 out of 5 changed files in this pull request and generated 3 comments.
Show a summary per file
| File | Description |
|---|---|
src/agents/agent.py |
Stops waiting on Queue.join() for Agent.as_tool(on_stream=...) streaming dispatch; awaits the dispatcher task instead. |
src/agents/extensions/experimental/codex/codex_tool.py |
Removes Queue.join() from Codex event dispatch cleanup and awaits the dispatcher task to prevent deadlocks. |
src/agents/sandbox/memory/manager.py |
Removes Queue.join() from sandbox memory flush() and awaits the worker task after enqueueing _STOP. |
tests/test_agent_as_tool.py |
Adds regression coverage for non-hanging behavior on BaseException/CancelledError and a property test about handler completion before return. |
tests/extensions/experiemental/codex/test_codex_tool.py |
Adds a regression test ensuring Codex _consume_events() does not hang when the stream handler raises BaseException/CancelledError. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
| if event_queue is not None: | ||
| await event_queue.put(None) | ||
| await event_queue.join() | ||
| if dispatch_task is not None: | ||
| # Wait on the dispatcher itself rather than on the queue draining. The sentinel is | ||
| # the last item queued and the dispatcher only returns after consuming it, so | ||
| # awaiting the task already implies every event was handled. Waiting on `join()` | ||
| # instead would hang forever when a handler raised a BaseException such as | ||
| # CancelledError, because the dispatcher is then gone and the outstanding | ||
| # `task_done()` calls can never arrive -- stranding the open spans below too. | ||
| await dispatch_task |
| self._ensure_worker() | ||
| for rollout_file in rollout_files: | ||
| self._queue.put_nowait(rollout_file) | ||
| await self._queue.join() | ||
| if self._worker_task is not None: | ||
| # Wait on the worker rather than on the queue draining. The worker consumes | ||
| # FIFO and only returns on _STOP, which is queued behind every rollout, so | ||
| # awaiting it already implies each one was processed. Waiting on `join()` | ||
| # instead would hang forever if the worker died on a BaseException, because | ||
| # the outstanding `task_done()` calls can never arrive. | ||
| self._queue.put_nowait(_STOP) | ||
| await self._worker_task | ||
| self._worker_task = None |
| async def test_agent_as_tool_streaming_reports_every_event_before_returning() -> None: | ||
| """Dropping `join()` must not let the tool return before handlers have run.""" |
![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: 1f24d44dc4
ℹ️ 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".
| # instead would hang forever when a handler raised a BaseException such as | ||
| # CancelledError, because the dispatcher is then gone and the outstanding | ||
| # `task_done()` calls can never arrive -- stranding the open spans below too. | ||
| await dispatch_task |
There was a problem hiding this comment.
When on_stream raises a BaseException while active_spans still contains a command execution span, awaiting the failed dispatcher here re-raises before the cleanup block below runs. In a stream that then errors or ends before the matching item.completed, the “Ensure any open spans are closed even on failure” loop is skipped, so the command span remains unfinished and unexported; run that span cleanup in a finally around this await before re-raising the handler failure.
Useful? React with 👍 / 👎.
Three streaming dispatchers hand user events to a background task and then wait for asyncio.Queue.join() before returning. Each dispatcher catches only Exception, so a handler raising a BaseException -- CancelledError being the realistic one, from a timeout or cancellation inside the callback -- kills the consumer. The outstanding task_done() calls can then never arrive and join() blocks forever, hanging the run with no error and no timeout. In every case join() is also redundant. The sentinel is queued last and the consumer only returns after consuming it, so awaiting the consumer task already implies every event was handled, and it surfaces a dead consumer instead of waiting on it. - Agent.as_tool(on_stream=...) hung the tool invocation. - Codex tool _consume_events() hung and left its tracing spans unfinished, because the span cleanup runs after the join. - The sandbox memory manager flush hung before its phase-two consolidation. Same failure shape as the STT listener fix in openai#4170.
1f24d44 to
44ccfb2
Compare
|
Thanks for identifying this issue and for the solid core approach. I have taken the fix into a maintainer-authored follow-up #4201 that builds on your Your contribution is credited with a |
Summary
Three streaming dispatchers hand user events to a background consumer task and then wait on
asyncio.Queue.join()before returning. Each wraps the user callback inexcept Exception, so aBaseExceptionescapes and ends the consumer. The outstandingtask_done()calls can then never arrive, andjoin()blocks forever — the run hangs with no exception and no timeout.asyncio.CancelledErroris the realistic trigger: anyasyncio.timeout(),wait_for, or task cancellation inside a user's handler raises one.The
join()is also redundant in all three cases. The sentinel is the last item queued and the consumer only returns after consuming it, so awaiting the consumer task already implies every event was handled — and it surfaces a dead consumer instead of waiting on it. Each fix is therefore a removal, not a new mechanism.Agent.as_tool(on_stream=...)codex_tool._consume_events()span.finish()runs after the joinsandbox/memory/manager.flush()Ordinary
Exceptions from handlers are unaffected: they stay logged and swallowed exactly as before.This is the same failure shape as #4170, which added an error sentinel so the STT waiters could not hang on a dead websocket listener.
Test plan
New regression tests, each parametrized over a custom
BaseExceptionandasyncio.CancelledError:tests/test_agent_as_tool.py::test_agent_as_tool_streaming_does_not_hang_when_handler_raises_base_exceptiontests/extensions/experiemental/codex/test_codex_tool.py::test_codex_tool_consume_events_does_not_hang_on_base_exceptionPlus
test_agent_as_tool_streaming_reports_every_event_before_returning, which pins the property that made droppingjoin()safe: the handler still observes every event before the tool returns. It awaits inside the handler so a missing wait would let the tool return first.On
mainthe four hang tests fail on the 5s guard; with the fix the same selection drops from 10.9s to 0.4s:Verification from the repository root:
make formatmake lintmake mypymain, none in the touched filesmake pyrightmain(src/agents/sandbox/util/tar_utils.py:161)uv run pytest tests/test_agent_as_tool.py tests/extensions/experiemental/codex/make testsThe full-suite run was done on Windows, where some sandbox symlink and tracing/realtime timing tests fail independently of this change. I diffed the failing set against a clean
maincheckout in the same environment: the two sets are identical (56 vs 56, no differences either way).One verification gap, stated plainly:
tests/sandbox/test_memory.pycannot be collected on Windows (ImportError: UnixLocalSandbox is not supported on Windows), so I could not execute a regression test for thesandbox/memory/manager.pysite locally and did not add one I could not run. I did reproduce the deadlock for that site structurally — driving its exact worker/queue/sentinel shape, the worker dies on aBaseExceptionandjoin()blocks with two unfinished items — and the change there is the same removal as the other two. Happy to add a test for it if you would rather have one, or to drop that hunk from this PR and leave it for someone who can run the file.Issue number
Closes #4198
Checks
.agents/skills/code-change-verification/scripts/run.sh/reviewbefore submitting this PRThe verification script is a bash script that shells out to
make; I ran the underlying steps individually instead, with the results above.@seratch — this is the same shape you fixed for the STT listener in #4170, so I went looking for siblings of that pattern and these three were the ones where a consumer death can strand a
Queue.join().Worth your call on one point: dropping
join()means aBaseExceptionfrom a handler now propagates out of the tool instead of hanging, which is a behavior change at that boundary even though the previous behavior was an unrecoverable hang. If you would rather keep handler failures fully contained, the alternative is widening the wrapper toexcept BaseExceptionand re-raisingCancelledErroronly — I did not do that because swallowing cancellation is usually worse than surfacing it.