fix: prevent request processing from blocking permission or cancel handling

Signed-off-by: Chojan Shang <[email protected]>

Author: Chojan Shang

src/acp/core.py

@@ -124,6 +124,7 @@ def __init__(
         self._reader = reader
         self._next_request_id = 0
         self._pending: dict[int, _Pending] = {}
+        self._inflight: set[asyncio.Task[Any]] = set()
         self._write_lock = asyncio.Lock()
         self._recv_task = asyncio.create_task(self._receive_loop())
 
@@ -132,6 +133,13 @@ async def close(self) -> None:
             self._recv_task.cancel()
             with contextlib.suppress(asyncio.CancelledError):
                 await self._recv_task
+        if self._inflight:
+            tasks = list(self._inflight)
+            for task in tasks:
+                task.cancel()
+            for task in tasks:
+                with contextlib.suppress(asyncio.CancelledError):
+                    await task
         # Do not close writer here; lifecycle owned by caller
 
     # --- IO loops ----------------------------------------------------------------
@@ -158,12 +166,28 @@ async def _process_message(self, message: dict) -> None:
         has_id = "id" in message
 
         if method is not None and has_id:
-            await self._handle_request(message)
-        elif method is not None and not has_id:
+            self._schedule(self._handle_request(message))
+            return
+        if method is not None and not has_id:
             await self._handle_notification(message)
-        elif has_id:
+            return
+        if has_id:
             await self._handle_response(message)
 
+    def _schedule(self, coro: Awaitable[Any]) -> None:
+        task = asyncio.create_task(coro)
+        self._inflight.add(task)
+        task.add_done_callback(self._task_done)
+
+    def _task_done(self, task: asyncio.Task[Any]) -> None:
+        self._inflight.discard(task)
+        if task.cancelled():
+            return
+        try:
+            task.result()
+        except Exception:
+            logging.exception("Unhandled error in JSON-RPC request handler")
+
     async def _handle_request(self, message: dict) -> None:
         """Handle JSON-RPC request."""
         payload = {"jsonrpc": "2.0", "id": message["id"]}

tests/test_cancel_prompt_flow.py

@@ -0,0 +1,62 @@
+import asyncio
+
+import pytest
+
+from acp import (
+    AgentSideConnection,
+    CancelNotification,
+    ClientSideConnection,
+    PromptRequest,
+    PromptResponse,
+)
+from acp.schema import TextContentBlock
+from tests.test_rpc import TestAgent, TestClient, _Server
+
+
+class LongRunningAgent(TestAgent):
+    """Agent variant whose prompt waits for a cancel notification."""
+
+    def __init__(self) -> None:
+        super().__init__()
+        self.prompt_started = asyncio.Event()
+        self.cancel_received = asyncio.Event()
+
+    async def prompt(self, params: PromptRequest) -> PromptResponse:
+        self.prompts.append(params)
+        self.prompt_started.set()
+        try:
+            await asyncio.wait_for(self.cancel_received.wait(), timeout=1.0)
+        except asyncio.TimeoutError as exc:
+            msg = "Cancel notification did not arrive while prompt pending"
+            raise AssertionError(msg) from exc
+        return PromptResponse(stopReason="cancelled")
+
+    async def cancel(self, params: CancelNotification) -> None:
+        await super().cancel(params)
+        self.cancel_received.set()
+
+
+@pytest.mark.asyncio
+async def test_cancel_reaches_agent_during_prompt() -> None:
+    async with _Server() as server:
+        agent = LongRunningAgent()
+        client = TestClient()
+        agent_conn = ClientSideConnection(lambda _conn: client, server.client_writer, server.client_reader)
+        _client_conn = AgentSideConnection(lambda _conn: agent, server.server_writer, server.server_reader)
+
+        prompt_request = PromptRequest(
+            sessionId="sess-xyz",
+            prompt=[TextContentBlock(type="text", text="hello")],
+        )
+        prompt_task = asyncio.create_task(agent_conn.prompt(prompt_request))
+
+        await agent.prompt_started.wait()
+        assert not prompt_task.done(), "Prompt finished before cancel was sent"
+
+        await agent_conn.cancel(CancelNotification(sessionId="sess-xyz"))
+
+        await asyncio.wait_for(agent.cancel_received.wait(), timeout=1.0)
+
+        response = await asyncio.wait_for(prompt_task, timeout=1.0)
+        assert response.stopReason == "cancelled"
+        assert agent.cancellations == ["sess-xyz"]

tests/test_permission_flow.py

@@ -0,0 +1,71 @@
+import asyncio
+
+import pytest
+
+from acp import (
+    AgentSideConnection,
+    ClientSideConnection,
+    PromptRequest,
+    PromptResponse,
+    RequestPermissionRequest,
+    RequestPermissionResponse,
+)
+from acp.schema import PermissionOption, TextContentBlock, ToolCallUpdate
+from tests.test_rpc import TestAgent, TestClient, _Server
+
+
+class PermissionRequestAgent(TestAgent):
+    """Agent that asks the client for permission during a prompt."""
+
+    def __init__(self, conn: AgentSideConnection) -> None:
+        super().__init__()
+        self._conn = conn
+        self.permission_responses: list[RequestPermissionResponse] = []
+
+    async def prompt(self, params: PromptRequest) -> PromptResponse:
+        permission = await self._conn.requestPermission(
+            RequestPermissionRequest(
+                sessionId=params.sessionId,
+                options=[
+                    PermissionOption(optionId="allow", name="Allow", kind="allow"),
+                    PermissionOption(optionId="deny", name="Deny", kind="deny"),
+                ],
+                toolCall=ToolCallUpdate(toolCallId="call-1", title="Write File"),
+            )
+        )
+        self.permission_responses.append(permission)
+        return await super().prompt(params)
+
+
+@pytest.mark.asyncio
+async def test_agent_request_permission_roundtrip() -> None:
+    async with _Server() as server:
+        client = TestClient()
+        client.queue_permission_selected("allow")
+
+        captured_agent: list[PermissionRequestAgent] = []
+
+        agent_conn = ClientSideConnection(lambda _conn: client, server.client_writer, server.client_reader)
+        _agent_conn = AgentSideConnection(
+            lambda conn: captured_agent.append(PermissionRequestAgent(conn)) or captured_agent[-1],
+            server.server_writer,
+            server.server_reader,
+        )
+
+        response = await asyncio.wait_for(
+            agent_conn.prompt(
+                PromptRequest(
+                    sessionId="sess-perm",
+                    prompt=[TextContentBlock(type="text", text="needs approval")],
+                )
+            ),
+            timeout=1.0,
+        )
+        assert response.stopReason == "end_turn"
+
+        assert captured_agent, "Agent was not constructed"
+        [agent] = captured_agent
+        assert agent.permission_responses, "Agent did not receive permission response"
+        permission_response = agent.permission_responses[0]
+        assert permission_response.outcome.outcome == "selected"
+        assert permission_response.outcome.optionId == "allow"