From f39bcb3632024bdebaea5a2397dec6a5802983b7 Mon Sep 17 00:00:00 2001 From: Venkat Ramachandran <268347452+venkat-uk@users.noreply.github.com> Date: Sat, 22 Aug 2026 13:59:06 +0100 Subject: [PATCH] Handle control-channel write failures without corrupting Query state Query writes to the CLI on both directions of the control channel, and neither write is guarded. _handle_control_request runs in a detached task and answers every hook, permission and SDK-MCP request. If the CLI has exited, or stdin has been closed, while one of those was in flight, the response write raises and the exception leaves the task: trio logs "Unhandled exception in detached trio task" with a traceback, asyncio logs "Task exception was never retrieved". Nothing is actually broken, but the logs read as an SDK crash during an ordinary shutdown race. Routing both writes through a helper also stops a failed success write from falling into the except clause, where it was answered with an error response for a request the handler had in fact completed. _send_control_request registers the waiter before writing, and the write sits outside the try that cleans up on timeout. A write that fails leaves the entry in pending_control_responses for the life of the Query, so a long-lived client whose interrupt() keeps failing grows both dicts. --- src/claude_agent_sdk/_internal/query.py | 41 ++++++++- tests/test_query.py | 117 ++++++++++++++++++++++++ 2 files changed, 155 insertions(+), 3 deletions(-) diff --git a/src/claude_agent_sdk/_internal/query.py b/src/claude_agent_sdk/_internal/query.py index 4d5f0070e..fa809dabe 100644 --- a/src/claude_agent_sdk/_internal/query.py +++ b/src/claude_agent_sdk/_internal/query.py @@ -577,7 +577,7 @@ async def _handle_control_request(self, request: SDKControlRequest) -> None: "response": response_data, }, } - await self.transport.write(json.dumps(success_response) + "\n") + await self._send_control_response(success_response) except anyio.get_cancelled_exc_class(): # Request was cancelled via control_cancel_request; the CLI has @@ -593,7 +593,32 @@ async def _handle_control_request(self, request: SDKControlRequest) -> None: "error": str(e), }, } - await self.transport.write(json.dumps(error_response) + "\n") + await self._send_control_response(error_response) + + async def _send_control_response(self, response: SDKControlResponse) -> None: + """Write a control response, tolerating a transport that has gone away. + + This runs in a detached task (see ``_spawn_control_request_handler``), + so an exception leaving it reaches no caller and surfaces as an + unhandled-task error instead: a warning from ``_task_compat`` on trio, + "Task exception was never retrieved" on asyncio. The write fails when + the CLI has exited, or stdin has been closed, while a hook, permission + or SDK-MCP request was still in flight. Nobody is left to answer at + that point, and the read loop reports the same transport failure to the + consumer, so log it and return. + + Keeping this out of ``_handle_control_request``'s ``except`` clause + also stops a failed *success* write from being reported to the CLI as + a handler error for a request the handler actually completed. + """ + try: + await self.transport.write(json.dumps(response) + "\n") + except Exception as e: + logger.debug( + "Could not send control response for %s: %s", + response["response"]["request_id"], + e, + ) async def _send_control_request( self, request: dict[str, Any], timeout: float = 60.0 @@ -622,7 +647,17 @@ async def _send_control_request( "request": request, } - await self.transport.write(json.dumps(control_request) + "\n") + try: + await self.transport.write(json.dumps(control_request) + "\n") + except BaseException: + # The request never reached the CLI, so no response can ever + # arrive for it. Drop the slot here: the write error already goes + # to the caller, and an entry left behind stays for the life of + # the Query (only the wait below and the read loop's failure path + # clean these up, and neither runs now). On a long-lived client + # a repeatedly failing interrupt() would grow both dicts. + self.pending_control_responses.pop(request_id, None) + raise # Wait for response try: diff --git a/tests/test_query.py b/tests/test_query.py index b5a254f33..c48f422a2 100644 --- a/tests/test_query.py +++ b/tests/test_query.py @@ -2026,3 +2026,120 @@ async def _test(): assert isinstance(q.pending_control_results["req_1"], ProcessError) anyio.run(_test) + + +class TestControlChannelWriteFailures: + """A failed transport write must not corrupt Query's control state. + + Both directions are covered. An outgoing request that is never written + must not leave a pending slot behind, and an incoming request whose + response cannot be delivered must not raise out of its detached task. + """ + + @staticmethod + def _dead_transport(): + """A transport whose write() fails the way a dead CLI subprocess does.""" + transport = AsyncMock() + transport.is_ready = Mock(return_value=True) + transport.write = AsyncMock( + side_effect=CLIConnectionError( + "Cannot write to terminated process (exit code: 1)" + ) + ) + return transport + + def test_failed_request_write_leaves_no_pending_entry(self): + """A control request that could not be written registers no waiter.""" + + async def _test(): + q = Query(transport=self._dead_transport(), is_streaming_mode=True) + + with pytest.raises(CLIConnectionError): + await q.interrupt() + + assert q.pending_control_responses == {} + assert q.pending_control_results == {} + + anyio.run(_test) + + def test_success_response_write_failure_does_not_raise(self): + """The callback ran; only its reply could not be delivered.""" + + async def _test(): + seen = [] + + async def hook(input_data, tool_use_id, context): + seen.append(input_data) + return {} + + q = Query(transport=self._dead_transport(), is_streaming_mode=True) + q.hook_callbacks["hook_0"] = hook + + await q._handle_control_request( + { + "type": "control_request", + "request_id": "req_1", + "request": { + "subtype": "hook_callback", + "callback_id": "hook_0", + "input": {"hook_event_name": "PreToolUse"}, + "tool_use_id": None, + }, + } + ) + + assert seen == [{"hook_event_name": "PreToolUse"}] + + anyio.run(_test) + + def test_error_response_write_failure_does_not_raise(self): + """The handler failed and the error reply could not be delivered.""" + + async def _test(): + q = Query(transport=self._dead_transport(), is_streaming_mode=True) + + await q._handle_control_request( + { + "type": "control_request", + "request_id": "req_1", + "request": { + "subtype": "hook_callback", + "callback_id": "no-such-callback", + "input": {}, + "tool_use_id": None, + }, + } + ) + + anyio.run(_test) + + def test_cancelled_request_still_skips_the_response(self): + """Cancellation keeps propagating; it is not swallowed as a write error.""" + + async def _test(): + transport = AsyncMock() + transport.is_ready = Mock(return_value=True) + + async def hook(input_data, tool_use_id, context): + raise anyio.get_cancelled_exc_class()() + + q = Query(transport=transport, is_streaming_mode=True) + q.hook_callbacks["hook_0"] = hook + + with pytest.raises(anyio.get_cancelled_exc_class()): + await q._handle_control_request( + { + "type": "control_request", + "request_id": "req_1", + "request": { + "subtype": "hook_callback", + "callback_id": "hook_0", + "input": {}, + "tool_use_id": None, + }, + } + ) + + transport.write.assert_not_awaited() + + anyio.run(_test)