Handle control-channel write failures without corrupting Query state

src/claude_agent_sdk/_internal/query.py

@@ -685,7 +685,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
@@ -701,7 +701,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
@@ -730,7 +755,22 @@ 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 write failed or was cancelled, so nobody will wait on this
+            # request. Drop its state here: the 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.
+            #
+            # Clear the result as well as the slot. Some bytes may have reached
+            # the CLI before the failure, and if its response raced in before
+            # this handler ran, _read_messages has already stored a result.
+            self.pending_control_responses.pop(request_id, None)
+            self.pending_control_results.pop(request_id, None)
+            raise
 
         # Wait for response
         try:

tests/test_query.py

@@ -2660,3 +2660,170 @@ 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_response_racing_a_failed_write_leaves_no_pending_result(self):
+        """The CLI answers before the write error reaches the cleanup.
+
+        The request bytes reached the CLI, its response came back and the
+        real read loop stored the result and set the event, and only then did
+        the write raise. The cleanup must drop the stored result as well as
+        the waiter slot, or it stays in pending_control_results for the life
+        of the Query.
+        """
+
+        async def _test():
+            responses: list[dict[str, Any]] = []
+
+            async def read_messages():
+                for message in responses:
+                    yield message
+
+            transport = AsyncMock()
+            transport.is_ready = Mock(return_value=True)
+            transport.read_messages = read_messages
+            q = Query(transport=transport, is_streaming_mode=True)
+
+            async def write_then_fail(data: str) -> None:
+                request_id = json.loads(data)["request_id"]
+                responses.append(
+                    {
+                        "type": "control_response",
+                        "response": {
+                            "subtype": "success",
+                            "request_id": request_id,
+                            "response": {},
+                        },
+                    }
+                )
+                # Let the read loop route the response before the write fails.
+                await q._read_messages()
+                assert request_id in q.pending_control_results
+                assert q.pending_control_responses[request_id].is_set()
+                raise CLIConnectionError("stdin closed after the write")
+
+            transport.write = AsyncMock(side_effect=write_then_fail)
+
+            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)