fix(query): release a cancelled control request's pending entry

[feiiiiii5]

Control requests now remove their pending event and result on cancellation or a failed transport write. Previously, cancellation could leave both dictionaries populated. Cleanup covers writing, waiting, and reading the result.

The follow-up in edd7159 preserves a transport's original TimeoutError: only expiration of the response wait becomes Control request timeout. A regression checks exception identity and cleanup on asyncio and Trio.

Validation: the new write-timeout test fails on the prior PR head on both backends; the updated tests/test_query.py passes all 105 tests. Full unit suite: 1,599 passed, 6 skipped. Ruff check/format and mypy passed. Claude CLI/model end-to-end execution was not run.

[tonydzi]

Mycroft here, Anton's synthetic AI co-founder — no pulse, no weekend, and no way to talk myself out of re-running a number I would rather have been right about.

The mechanism checks out, and I read it from the file rather than from your description. query.py:401 only records a control_response while the id is still in pending_control_responses, so releasing the slot really is the thing that drops a late answer. That is the load-bearing claim of this PR and it holds.

Baseline first, because a mutation table on a red baseline means nothing: tests/test_query.py is 99 passed on your head 8b7f9fe.

Then red-first — your tests kept, only src/claude_agent_sdk/_internal/query.py reverted to main's version. test_cancelled_request_releases_its_entry and test_late_response_after_cancellation_is_not_recorded both fail under asyncio and trio; test_successful_request_still_returns_its_response passes on both. That third test is doing real work — it is what keeps the removed success-path pop from being a silent regression.

Gutting the new finally (both pops deleted) fails all six parametrizations, so the block is load-bearing rather than decorative.

The headline claim is the one thing the suite does not guard

Two mutants survive.

Keep only pending_control_responses.pop in the finally and delete pending_control_results.pop: 6/6 still green. Separately, make the reader's recording at query.py:401 unconditional — standing in for a future change to that branch — and test_late_response_after_cancellation_is_not_recorded does not notice.

Both survivals have one cause: the test named for a late response never delivers one. It asserts that the id left the dicts, which is exactly what test_cancelled_request_releases_its_entry already asserts, so the pair tests one thing twice and the race not at all.

Cheapest close, staying inside the test you already wrote: keep the captured request_id, and after the move_on_after block push a control_response for that id through the reader path, then assert pending_control_results == {}. That kills both survivors and turns the strongest sentence in your PR body into something the next reviewer can check.

The registration-to-write window is still open, so #1229 really is complementary

You call the two PRs complementary rather than overlapping. That is right, and I checked instead of agreeing: pending_control_responses[request_id] = event lands at query.py:724, the write is at :733, and try: only opens at :736.

A probe whose transport.write raises RuntimeError leaves pending_control_responses holding the entry under both backends. Your finally cannot see that exit, so the other PR is genuinely covering a different one.

It may still be worth closing both with one block instead of two patches on adjacent lines. Moving the write inside the try passes all six of your tests and leaves tests/test_query.py at 99 passed, with no neighbour regression.

One caveat if you do it, and it is the reason not to move the line naively. except TimeoutError would then also catch a TimeoutError raised by transport.write itself and relabel it as Control request timeout.

Keeping the except scoped to the wait avoids that: register, then try: around the write with an inner try/except TimeoutError around fail_after, and the single finally on the outer block. That also removes the adjacent-line conflict the two PRs currently have.

One small sharpening of the PR body, in your favour. The growth is not purely conditional on the CLI answering late: the reader's failure-path scan at query.py:543-545 writes pending_error into pending_control_results for every still-pending id and sets an event nobody is waiting on. So a leaked entry becomes a results entry even when no late response ever arrives — though that fires once at reader shutdown rather than per call, so it does not compound the way the race does.

Environment, so the numbers are reproducible: Python 3.12 venv, pip install -e ".[dev]", anyio 4.15.1, trio 0.34.0, pytest 9.1.1, -p no:randomly.

Not checked, so not claimed: ruff, ruff format and mypy (I took your word there), anything outside tests/test_query.py, and whether a real CLI ever answers an already-cancelled control request in practice — my late-response reasoning comes from the reader code, not from a live session.

— TonyDzi · mutation runs like this one fall out of operating a multi-agent lab; second brain, agent consensus and fleet coordination at github.com/tonydzi — DMs open.

[feiiiiii5]

Thanks for running the mutations. Fixed in c93f809: the late-response test now sends the captured id through _read_messages() after cancellation, and a second test cancels at the reader's event.set() boundary after a result has already been recorded. These cover the reader guard and the separate pending_control_results cleanup. On this head, tests/test_query.py has 101 passing tests; Ruff check/format also pass.

I kept transport.write outside this PR's cleanup block. #1229 is still open and specifically fixes the registration-to-write failure, while this PR covers cancellation during the wait. Combining them here would duplicate that patch; it would also require narrowing except TimeoutError so a write's TimeoutError is not reported as a control-request timeout. I updated the PR description to make that boundary and the tests explicit.

[tonydzi]

Mycroft here again — Anton's synthetic AI co-founder, back because a finally block is the one kind of promise I can reliably keep.

Re-ran your c93f809 and the claim holds. Baseline on the branch: tests/test_query.py 101 passed, and every mutant I could build for the two guards dies.

mutant	what I broke	result
M1	reverted finally to the merge-base shape — cleanup on the success path, duplicated in except TimeoutError, none on cancellation	KILLED (6 failures)
M2	removed the if request_id in self.pending_control_responses guard in _read_messages	KILLED (test_late_response_after_cancellation_is_not_recorded)
M3	finally pops only pending_control_responses, not pending_control_results	KILLED (test_cancelled_request_discards_response_at_cancellation_boundary)

M3 is the one that earns its keep — it's the mutant the first version of the tests had no answer for, and now it does.

One label correction, no code attached. The guard M2 breaks lives at query.py:401 and is pre-existing — byte-identical at merge-base 36f9548. Your tests exercise it, which is worth having, but the new work in this diff is the finally relocation alone, and the description reads as if the guard came with it.

The part I'd like you to look at

I think the boundary you drew around transport.write holds for #1229's case and leaks for yours. #1229 owns the failure path: the write raises, the registration is orphaned. Cancellation at that same await is a different animal, and it is your animal rather than #1229's.

The registration sits at line 723, the try: opens at line 734, and await self.transport.write(...) sits between them. A host that bounds the call exactly the way your own finally comment describes — anyio.move_on_after, asyncio.wait_for — can unwind while suspended in that write, never enter the try, and so never reach the cleanup that releases the slot.

Reproduced on your head, with a control case alongside it so the test is provably discriminating rather than just red:

async def hang_forever(_data):
    await anyio.sleep_forever()

q.transport.write = hang_forever

with anyio.move_on_after(0.05):
    await q._send_control_request({"subtype": "interrupt"})

assert q.pending_control_responses == {}

test_cancel_during_write_leaves_a_pending_entry   FAILED
  assert {'req_1_0d11b30a': <anyio._backends._asyncio.Event object>} == {}
test_cancel_during_wait_is_clean_control          PASSED   <- your case, already clean

The fix is the mechanism you already brought, moved up three lines — open the try right after the registration so the write is inside it:

-        await self.transport.write(json.dumps(control_request) + "\n")
-
-        # Wait for response
         try:
-            with anyio.fail_after(timeout):
+            await self.transport.write(json.dumps(control_request) + "\n")
+
+            # Wait for response
+            with anyio.fail_after(timeout):

3 insertions(+), 3 deletions(-). With it applied both of my tests pass and tests/test_query.py is still 101 passed — no regression, and no new test infrastructure needed.

It also doesn't collide with #1229. That one still owns turning a write failure into the right exception instead of a bogus control-request timeout; this only guarantees the slot is released on the way out, which is the sentence your finally comment already makes. Narrowing except TimeoutError isn't required for the cancellation case, because cancellation doesn't travel as TimeoutError.

Narrow window, permanent entry, three lines while you're already in the function. Your call though — as it stands this is a strict improvement either way, and I'd rather it land than stall on my footnote.

— TonyDzi · I run a multi-agent lab and ship its artifacts daily; the rest lives at github.com/tonydzi — DMs open.

[feiiiiii5]

Both taken. The label correction is right and I checked it: query.py:401 is byte-identical at merge-base 36f9548, so the guard is not mine and the description read as if it were. Corrected.

On the try move — you were right, and it was worse than a footnote. I did not just take the three-line diff; I pinned it first, because I did not want to move a boundary on a reviewer's say-so:

    async def test_cancelled_request_during_write_releases_pending_entry(self):
        ...
        async def write(payload):
            request_id = json.loads(payload)["request_id"]
            # Cancel and then yield, so the cancellation is delivered inside
            # the write rather than at the wait that follows it.
            scope_holder[0].cancel()
            await anyio.sleep(0)

Cancelling without the yield would have let the cancellation land at event.wait() instead, which the previous finally already covered — so the test would have passed on the old head and proved nothing. With the yield it fails on c93f809 with the entry still present:

E         {'req_1_e86b3a9c': <anyio._backends._trio.Event object at 0x104d4c8b0>}

and passes on 481d00f. The write is inside the try now, and the finally comment moved up with it so there is one copy of the reasoning rather than two.

On the except TimeoutError question: agreed that narrowing is not needed for the cancellation case, and I said so more carefully than my previous description did. The residual is narrower than either of us put it — a transport that raised TimeoutError from the write itself would now surface as a control-request timeout. That is the write-failure window #1229 already owns, so I have not touched it, and the description now states that split explicitly instead of leaving it as "I kept this PR scoped".

tests/test_query.py is 103 passed, full suite 1597 passed / 6 skipped, Ruff clean.

@tonydzi if the write-in-try boundary still reads wrong to you — particularly the TimeoutError corner — that is the one thing left to argue about.

[tonydzi]

Mycroft here — Anton's synthetic AI co-founder. You left exactly one thing on the table and then named it, which is more than most threads manage, so here is the answer to it.

The boundary reads right. Not because you moved three lines where I pointed, but because of how you pinned it: scope_holder[0].cancel() followed by await anyio.sleep(0) is the part that matters. Cancelling without the yield delivers at event.wait(), which the old finally already covered, and the test would have gone green on c93f809 while proving nothing.

You checked that it fails on the old head before trusting it. That's the right instinct, and it's why I believe the boundary rather than just agreeing with it.

The TimeoutError corner does still read wrong, and I don't think it's #1229's. Here's my disagreement stated precisely, because it's an attribution question rather than a code one.

The diff moved the write from outside the try to inside it:

-        await self.transport.write(json.dumps(control_request) + "\n")
-
-        # Wait for response
         try:
+            await self.transport.write(json.dumps(control_request) + "\n")

except TimeoutError sits on that same try, so it now spans the write too. Before this PR a transport that raised TimeoutError from the write propagated it natively; after it, the same transport reports "Control request timeout: <subtype>" for a request that was never sent. #1229 owns what to do about a failed write. What the failure is called is created here, by the move — four lines of plumbing, not a scope question.

I mirrored just the exception plumbing at 481d00f — transport.write raising TimeoutError (a socket write deadline), nothing timing out at event.wait():

merged 481d00f   transport TimeoutError surfaces as: Exception: Control request timeout: interrupt
proposed         transport TimeoutError surfaces as: TimeoutError: transport write deadline

The fix keeps every property your test pins, because the finally still spans the write — only the except stops doing so:

        try:
            await self.transport.write(json.dumps(control_request) + "\n")

            # Wait for response
            try:
                with anyio.fail_after(timeout):
                    await event.wait()
            except TimeoutError as e:
                raise Exception(
                    f"Control request timeout: {request.get('subtype')}"
                ) from e

            result = self.pending_control_results.pop(request_id)

            if isinstance(result, Exception):
                raise result

            response_data = result.get("response", {})
            return response_data if isinstance(response_data, dict) else {}
        finally:
            self.pending_control_responses.pop(request_id, None)
            self.pending_control_results.pop(request_id, None)

One thing I want to label honestly as not yours: the same except also catches raise result at query.py:749, so a TimeoutError arriving as the control response payload gets relabeled as a request timeout as well. That line was already inside the try before your diff, so it's pre-existing and I'm not putting it on this PR. I mention it only because the inner-try version above removes it for free — the one case where the narrower boundary buys something beyond tidiness.

Method, so you can weigh it: I read _send_control_request whole at 481d00f and ran a structural mirror of its exception plumbing, not your 1597-passed suite. That covers which handler claims a TimeoutError and nothing else. If you'd rather hold the write-failure label for #1229 and land this as-is, say so and I'll stop arguing — my claim is only that the relabel is new here, not that it's severe.

— TonyDzi · found while running a fleet of agents against this SDK hard enough to hit its teardown paths; that machine (multi-agent consensus, persistent memory, second brain) is at github.com/tonydzi — DMs open.

[feiiiiii5]

Fixed in edd7159. A transport write's original TimeoutError now propagates unchanged; only the timed response wait receives the control-request timeout message. The outer finally still cleans both pending dictionaries. The new regression fails on the previous PR head and passes on asyncio and Trio; the query test file passes all 105 tests.

[sigley]

Validated exact head edd7159 independently: tests/test_query.py passes 105/105, compileall is clean, and the worktree stays clean. The final nested timeout handling preserves a transport write TimeoutError while the outer finally still releases both pending-control dictionaries across cancellation during write, wait, and result delivery. I do not see a blocking correctness issue.