fix(connection): keep request responses out of the handler's cancellation

Two races could still lose the response to a cancelled request:

- A `$/cancel_request` read in the same buffer as its request cancelled the
  request task before it first ran, so the coroutine that would have sent
  `-32800` never executed.
- A late or repeated `$/cancel_request` arriving while the response was being
  sent cancelled that send, so the peer got no response at all.

Each incoming request now runs its handler in its own task, which is the only
thing `$/cancel_request` targets. A separate supervised task waits for the
handler with `asyncio.wait` and sends its result, or `-32800` if the handler
was cancelled. `close()` still cancels both, so shutdown never waits on a
blocked send, and the `Task.uncancel()` / closed-send special cases are no
longer needed.

Also document request cancellation in the quickstart and state that a
cancelled handler answers with its result or `-32800`.

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>

Author: monody0007

docs/experimental-v2.md

@@ -152,6 +152,6 @@ when an agent selects v1.
 Client-side fallback is application controlled and may require opening a new
 transport. Protocol-level request cancellation (`$/cancel_request`) is handled
 by the shared connection layer, as in v1: cancelling the task awaiting a request
-notifies the peer, and an incoming cancellation cancels the handler task and
-answers with `-32800`. `session/cancel` remains available for cancelling active
-session work.
+sends a best-effort cancellation to the peer, and an incoming cancellation
+cancels the handler task, which answers with its result or `-32800`.
+`session/cancel` remains available for cancelling active session work.

docs/quickstart.md

@@ -228,6 +228,21 @@ MCP requests return the inner JSON result unchanged, including `null`. Use
 and optional `params`. These methods share the same connections and routers
 across stdio, HTTP, and WebSocket transports.
 
+## Request cancellation
+
+Connections handle the protocol's `$/cancel_request` notification. When the
+peer cancels one of its requests, the SDK cancels the task running your handler.
+The handler may catch `asyncio.CancelledError` and return a (partial) result;
+otherwise the peer receives a `-32800` "Request cancelled" error. Cancellations
+for unknown or already finished requests are ignored; a cancelled request still
+gets exactly one response unless the connection is closed first.
+
+Cancelling the task that awaits an outgoing request (for example through
+`asyncio.wait_for`) still raises `CancelledError` locally without waiting for the
+peer, and additionally sends a best-effort `$/cancel_request` for that request.
+Cancellation support is optional for peers. Use `session/cancel` to stop a
+prompt turn.
+
 ## Maintaining protocol routes
 
 The `Agent` and `Client` protocols in `src/acp/interfaces.py` are the source of

src/acp/connection.py

@@ -5,7 +5,6 @@
 import inspect
 import json
 import logging
-import sys
 from collections.abc import Awaitable, Callable
 from dataclasses import dataclass
 from enum import Enum
@@ -161,12 +160,14 @@ def _process_message(self, message: dict[str, Any]) -> None:
         if method is not None:  # this is a request or notification
             # {"jsonrpc": "2.0", "id": 1, "method": "foo", "params": {...}}  # request
             # {"jsonrpc": "2.0", "method": "foo", "params: {...}}  # notification
-            task = self._tasks.create(
-                self._run_request(message) if has_id else self._run_notification(message),
-                name="acp.Connection.request" if has_id else "acp.Connection.notification",
-            )
-            if has_id:
-                self._track_incoming(message["id"], task)
+            if not has_id:
+                self._tasks.create(self._run_notification(message), name="acp.Connection.notification")
+                return
+            # The handler gets its own task so ``$/cancel_request`` can cancel it, even before it
+            # starts, without also cancelling delivery of the response the peer still expects.
+            handler = self._tasks.create(self._execute_request(message), name="acp.Connection.request")
+            self._track_incoming(message["id"], handler)
+            self._tasks.create(self._run_request(message, handler), name="acp.Connection.response")
             return
         if has_id:  # this is a response, {"id", "result" | "error"}
             self._handle_response(message)
@@ -226,22 +227,19 @@ def _send_cancel_request(self, request_id: int) -> None:
     def _on_cancel_request_error(self, task: asyncio.Task[Any], exc: BaseException) -> None:
         logging.debug("Failed to send %s", _CANCEL_REQUEST_METHOD, exc_info=exc)
 
-    async def _run_request(self, message: dict[str, Any]) -> None:
-        try:
-            payload = await self._execute_request(message)
-        except asyncio.CancelledError:
-            if self._closed:
-                raise
-            # Cancelled by ``$/cancel_request`` or from inside the handler; either way the
-            # protocol still requires a response for the original request.
-            task = asyncio.current_task()
-            if sys.version_info >= (3, 11) and task is not None:
-                task.uncancel()
+    async def _run_request(
+        self, message: dict[str, Any], handler: asyncio.Future[dict[str, Any]] | None = None
+    ) -> None:
+        if handler is None:
+            handler = self._tasks.create(self._execute_request(message), name="acp.Connection.request")
+        # Unlike ``await handler``, ``asyncio.wait`` does not re-raise the handler's cancellation, and
+        # ``$/cancel_request`` only targets the handler, so only ``close()`` can abandon the response.
+        await asyncio.wait({handler})
+        if handler.cancelled():
+            # Cancelled by ``$/cancel_request`` or from inside the handler.
             payload = {"jsonrpc": "2.0", "id": message["id"], "error": RequestError.request_cancelled().to_error_obj()}
-        if self._closed:
-            # The transport is gone (e.g. the handler returned a result while being
-            # cancelled by ``close()``); sending would never complete.
-            return
+        else:
+            payload = handler.result()
         await self._transport.send(payload)
         self._notify_observers(StreamDirection.OUTGOING, payload)
 

tests/test_request_cancellation.py

@@ -10,6 +10,7 @@
 import pytest
 
 from acp import Agent
+from acp.connection import Connection
 from acp.core import AgentSideConnection, ClientSideConnection
 from acp.schema import PermissionOption, ToolCallUpdate
 from tests.conftest import TestAgent, TestClient
@@ -24,7 +25,7 @@ async def _read(reader: asyncio.StreamReader) -> dict[str, Any]:
     return json.loads(await asyncio.wait_for(reader.readline(), timeout=1))
 
 
-def _prompt(request_id: int | str) -> dict[str, Any]:
+def _prompt(request_id: int | str | None) -> dict[str, Any]:
     return {
         "jsonrpc": "2.0",
         "id": request_id,
@@ -33,7 +34,7 @@ def _prompt(request_id: int | str) -> dict[str, Any]:
     }
 
 
-def _cancel_request(request_id: int | str) -> dict[str, Any]:
+def _cancel_request(request_id: int | str | None) -> dict[str, Any]:
     return {"jsonrpc": "2.0", "method": "$/cancel_request", "params": {"requestId": request_id}}
 
 
@@ -77,6 +78,122 @@ async def test_cancel_request_cancels_handler_and_replies_request_cancelled(
         assert "$/cancel_request" not in caplog.text
 
 
+@pytest.mark.asyncio
+@pytest.mark.parametrize("request_id", [0, "req-0", None])
+async def test_cancel_request_before_the_handler_starts_still_replies(
+    server, caplog: pytest.LogCaptureFixture, request_id: int | str | None
+) -> None:
+    agent = _BlockingAgent()
+    async with AgentSideConnection(cast(Agent, agent), server.server_writer, server.server_reader, listening=True):
+        with caplog.at_level(logging.ERROR):
+            # One write, so the receive loop reads both frames before the handler task first runs.
+            server.client_writer.write(
+                (json.dumps(_prompt(request_id)) + "\n" + json.dumps(_cancel_request(request_id)) + "\n").encode()
+            )
+            await server.client_writer.drain()
+            response = await _read(server.client_reader)
+
+        assert response["id"] == request_id
+        assert response["error"]["code"] == -32800
+        assert not agent.started.is_set()
+        assert caplog.text == ""
+        with pytest.raises(asyncio.TimeoutError):
+            await asyncio.wait_for(server.client_reader.readline(), timeout=0.1)
+
+
+class _GatedTransport:
+    """Message transport whose sends block until ``release`` is set."""
+
+    def __init__(self) -> None:
+        self.incoming: asyncio.Queue[dict[str, Any]] = asyncio.Queue()
+        self.sent: list[dict[str, Any]] = []
+        self.sending = asyncio.Event()
+        self.release = asyncio.Event()
+        self.settled = asyncio.Event()
+        self.send_cancelled = False
+        self._receiving = False
+
+    async def send(self, message: dict[str, Any]) -> None:
+        self.sending.set()
+        try:
+            await self.release.wait()
+        except asyncio.CancelledError:
+            self.send_cancelled = True
+            raise
+        else:
+            self.sent.append(message)
+        finally:
+            self.settled.set()
+
+    async def receive(self) -> dict[str, Any] | None:
+        self._receiving = True
+        try:
+            return await self.incoming.get()
+        finally:
+            self._receiving = False
+
+    async def close(self) -> None:
+        pass
+
+    async def deliver(self, message: dict[str, Any]) -> None:
+        """Queue ``message`` and wait until the connection has processed it."""
+        await self.incoming.put(message)
+        for _ in range(100):
+            if self._receiving and self.incoming.empty():
+                return
+            await asyncio.sleep(0)
+        raise AssertionError("the connection did not process the message")
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("handler_cancelled", [False, True], ids=["result", "request_cancelled"])
+async def test_cancel_request_during_response_send_keeps_the_response(handler_cancelled: bool) -> None:
+    transport = _GatedTransport()
+    started = asyncio.Event()
+
+    async def handler(method: str, params: Any, is_notification: bool) -> Any:
+        started.set()
+        if handler_cancelled:
+            await asyncio.Event().wait()
+        return {"ok": True}
+
+    async with Connection(handler, transport):
+        await transport.deliver(_prompt(0))
+        await asyncio.wait_for(started.wait(), timeout=1)
+        if handler_cancelled:
+            await transport.deliver(_cancel_request(0))
+        await asyncio.wait_for(transport.sending.wait(), timeout=1)
+
+        # A late (or repeated) cancellation lands while the response is still being sent.
+        await transport.deliver(_cancel_request(0))
+        assert transport.sent == []
+        transport.release.set()
+        await asyncio.wait_for(transport.settled.wait(), timeout=1)
+        assert not transport.send_cancelled, "the cancellation aborted the response send"
+
+    if handler_cancelled:
+        assert [(m["id"], m["error"]["code"]) for m in transport.sent] == [(0, -32800)]
+    else:
+        assert transport.sent == [{"jsonrpc": "2.0", "id": 0, "result": {"ok": True}}]
+
+
+@pytest.mark.asyncio
+async def test_close_cancels_a_blocked_response_send() -> None:
+    transport = _GatedTransport()
+
+    async def handler(method: str, params: Any, is_notification: bool) -> Any:
+        return {"ok": True}
+
+    conn = Connection(handler, transport)
+    await transport.deliver(_prompt(0))
+    await asyncio.wait_for(transport.sending.wait(), timeout=1)
+
+    await asyncio.wait_for(conn.close(), timeout=1)
+
+    assert transport.send_cancelled
+    assert transport.sent == []
+
+
 @pytest.mark.asyncio
 async def test_cancel_request_only_affects_the_targeted_request(server, caplog: pytest.LogCaptureFixture) -> None:
     agent = _BlockingAgent()

tests/test_v2_runtime.py

@@ -550,3 +550,35 @@ async def handle_extension_request(self, method: str, params: Any) -> Any:
     finally:
         await client_connection.close()
         await agent_connection.close()
+
+
+@pytest.mark.asyncio
+async def test_cancel_request_before_the_handler_starts_still_replies() -> None:
+    handled: list[str] = []
+
+    class RecordingAgent(ExtensionAgent):
+        async def handle_extension_request(self, method: str, params: Any) -> Any:
+            handled.append(method)
+            return {}
+
+    peer, agent_transport = memory_transport_pair()
+    agent_connection = v2.AgentSideConnection(RecordingAgent(), agent_transport)
+
+    try:
+        await peer.send({
+            "jsonrpc": "2.0",
+            "id": 0,
+            "method": "initialize",
+            "params": initialize_request().model_dump(mode="json", by_alias=True, exclude_none=True),
+        })
+        assert "result" in await asyncio.wait_for(peer.receive(), timeout=1)
+
+        # Both frames are queued before the connection runs, so the cancel precedes the handler.
+        await peer.send({"jsonrpc": "2.0", "id": 1, "method": "_vendor/slow", "params": {}})
+        await peer.send({"jsonrpc": "2.0", "method": "$/cancel_request", "params": {"requestId": 1}})
+        response = await asyncio.wait_for(peer.receive(), timeout=1)
+
+        assert (response["id"], response["error"]["code"]) == (1, -32800)
+        assert handled == []
+    finally:
+        await agent_connection.close()