fix(ws): ignore non-object JSON frames in client transport (#169)

* fix(ws): ignore non-object JSON frames in client transport

* fix(ws): report invalid frames without losing buffered messages

---------

Co-authored-by: br-mbp <brend.main@gmail.com>
Co-authored-by: Chojan Shang <psiace@apache.org>

Author: 3 people

src/acp/ws/client.py

@@ -13,6 +13,7 @@
 from typing import TYPE_CHECKING, Any
 
 from .._cookies import MemoryAcpCookieStore
+from ..exceptions import RequestError
 
 try:
     import websockets
@@ -53,9 +54,21 @@ async def receive(self) -> dict[str, Any] | None:
             if isinstance(frame, bytes):
                 continue
             try:
-                return json.loads(frame)
+                message = json.loads(frame)
             except json.JSONDecodeError:
-                continue
+                error = RequestError.parse_error()
+            else:
+                if isinstance(message, dict) and message:
+                    return message
+                if isinstance(message, list) and message:
+                    # Batch support is outside this message-level transport.
+                    continue
+                error = RequestError.invalid_request()
+            # Sending while closing waits for EOF, which may require draining
+            # buffered frames first. Keep receiving instead of waiting here.
+            if self._ws.state == websockets.State.OPEN:
+                with contextlib.suppress(websockets.ConnectionClosed):
+                    await self.send({"jsonrpc": "2.0", "id": None, "error": error.to_error_obj()})
 
     async def close(self) -> None:
         if self._closed:

src/acp/ws/server.py

@@ -2,14 +2,16 @@
 
 from __future__ import annotations
 
+import contextlib
 import json
 import logging
 import uuid
 from typing import TYPE_CHECKING, Any
 
-from starlette.websockets import WebSocket, WebSocketState
+from starlette.websockets import WebSocket, WebSocketDisconnect, WebSocketState
 
 from ..agent.connection import AgentSideConnection
+from ..exceptions import RequestError
 from ..http.protocol import CONNECTION_ID_HEADER
 
 if TYPE_CHECKING:
@@ -44,9 +46,18 @@ async def receive(self) -> dict[str, Any] | None:
             try:
                 message = json.loads(event["text"])
             except json.JSONDecodeError:
-                continue
-            if isinstance(message, dict):
-                return message
+                error = RequestError.parse_error()
+            else:
+                if isinstance(message, dict) and message:
+                    return message
+                if isinstance(message, list) and message:
+                    # Batch support is outside this message-level transport.
+                    continue
+                error = RequestError.invalid_request()
+            # A failed write must not discard valid messages buffered by ASGI.
+            if self._ws.application_state == WebSocketState.CONNECTED:
+                with contextlib.suppress(WebSocketDisconnect):
+                    await self.send({"jsonrpc": "2.0", "id": None, "error": error.to_error_obj()})
         return None
 
     async def close(self) -> None:

tests/http/test_websocket.py

@@ -3,14 +3,17 @@
 from __future__ import annotations
 
 import asyncio
+import contextlib
 import json
 import logging
 from typing import Any
 
 import pytest
 from starlette.websockets import WebSocket
 from websockets.asyncio.server import serve
+from websockets.protocol import State
 
+from acp.connection import Connection
 from acp.http.protocol import CONNECTION_ID_HEADER
 from acp.schema import NewSessionResponse, PromptResponse
 from acp.ws.client import create_websocket_stream
@@ -89,6 +92,155 @@ async def send_binary_then_text(ws: Any) -> None:
             await transport.close()
 
 
+@pytest.mark.asyncio
+@pytest.mark.parametrize("frame", ["null", "[]", "{}", "123", "1.5", '"x"', "true", "false", "invalid json", b"binary"])
+async def test_client_transport_receives_valid_message_after_invalid_frame(frame: str | bytes) -> None:
+    message = {"jsonrpc": "2.0", "id": 0, "result": {"items": [], "value": None}}
+
+    async def send_frames(ws: Any) -> None:
+        await ws.send(frame)
+        await ws.send(json.dumps(message))
+        await ws.close()
+
+    async with serve(send_frames, "127.0.0.1", 0) as server:
+        port = server.sockets[0].getsockname()[1]
+        transport = await create_websocket_stream(f"ws://127.0.0.1:{port}")
+        try:
+            assert await asyncio.wait_for(transport.receive(), timeout=1) == message
+            assert await asyncio.wait_for(transport.receive(), timeout=1) is None
+        finally:
+            await transport.close()
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("frame", ["null", "[]", "{}"])
+async def test_client_connection_recovers_after_invalid_frame(frame: str) -> None:
+    async def respond(ws: Any) -> None:
+        request = json.loads(await ws.recv())
+        await ws.send(frame)
+        await ws.send(json.dumps({"jsonrpc": "2.0", "id": request["id"], "result": {"protocolVersion": 1}}))
+        await ws.wait_closed()
+
+    async def handler(method: str, params: Any, is_notification: bool) -> None:
+        pass
+
+    async with serve(respond, "127.0.0.1", 0) as server:
+        port = server.sockets[0].getsockname()[1]
+        transport = await create_websocket_stream(f"ws://127.0.0.1:{port}")
+        async with Connection(handler, transport) as conn:
+            result = await asyncio.wait_for(conn.send_request("initialize", {"protocolVersion": 1}), timeout=1)
+            assert result == {"protocolVersion": 1}
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    ("frame", "code", "message"),
+    [
+        ("invalid json", -32700, "Parse error"),
+        ("null", -32600, "Invalid request"),
+        ("[]", -32600, "Invalid request"),
+        ("{}", -32600, "Invalid request"),
+        ("123", -32600, "Invalid request"),
+        ("1.5", -32600, "Invalid request"),
+        ('"x"', -32600, "Invalid request"),
+        ("true", -32600, "Invalid request"),
+        ("false", -32600, "Invalid request"),
+    ],
+)
+async def test_client_reports_protocol_error_then_receives_message(frame: str, code: int, message: str) -> None:
+    valid = {"jsonrpc": "2.0", "id": "response", "result": None}
+    replies = []
+
+    async def respond(ws: Any) -> None:
+        await ws.send(frame)
+        replies.append(json.loads(await ws.recv()))
+        await ws.send(json.dumps(valid))
+        await ws.wait_closed()
+
+    async with serve(respond, "127.0.0.1", 0) as server:
+        port = server.sockets[0].getsockname()[1]
+        transport = await create_websocket_stream(f"ws://127.0.0.1:{port}")
+        try:
+            assert await asyncio.wait_for(transport.receive(), timeout=1) == valid
+            assert replies == [
+                {"jsonrpc": "2.0", "id": None, "error": {"code": code, "message": message, "data": None}}
+            ]
+        finally:
+            await transport.close()
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("code", [1000, 1011])
+async def test_client_drains_buffered_response_when_error_reply_cannot_be_sent(code: int) -> None:
+    closed = asyncio.Event()
+    valid = {"jsonrpc": "2.0", "id": 0, "result": {"items": [], "value": None}}
+
+    async def respond(ws: Any) -> None:
+        await ws.send("invalid json")
+        await ws.send("null")
+        await ws.send(json.dumps(valid))
+        await ws.close(code=code)
+        closed.set()
+
+    async with serve(respond, "127.0.0.1", 0) as server:
+        port = server.sockets[0].getsockname()[1]
+        transport = await create_websocket_stream(f"ws://127.0.0.1:{port}")
+        try:
+            await asyncio.wait_for(closed.wait(), timeout=1)
+            assert await asyncio.wait_for(transport.receive(), timeout=1) == valid
+            assert await asyncio.wait_for(transport.receive(), timeout=1) is None
+        finally:
+            await transport.close()
+
+
+@pytest.mark.asyncio
+async def test_client_drains_buffered_response_before_close_handshake_finishes(monkeypatch: pytest.MonkeyPatch) -> None:
+    from websockets.asyncio.client import connect
+
+    closing = asyncio.Event()
+    release = asyncio.Event()
+    valid = {"jsonrpc": "2.0", "id": 0, "result": None}
+    connection: Any = None
+
+    async def open_socket(url: str, **kwargs: Any) -> Any:
+        nonlocal connection
+        connection = await connect(url, **kwargs)
+        return connection
+
+    monkeypatch.setattr("acp.ws.client.ws_connect", open_socket)
+
+    async def respond(ws: Any) -> None:
+        await ws.send("null")
+        await ws.send(json.dumps(valid))
+        ws.transport.pause_reading()
+        close = asyncio.create_task(ws.close())
+        try:
+            while ws.state == State.OPEN:
+                await asyncio.sleep(0)
+            closing.set()
+            await release.wait()
+        finally:
+            ws.transport.resume_reading()
+            await close
+
+    async with serve(respond, "127.0.0.1", 0) as server:
+        port = server.sockets[0].getsockname()[1]
+        transport = await create_websocket_stream(f"ws://127.0.0.1:{port}")
+        try:
+
+            async def wait_for_closing() -> None:
+                await closing.wait()
+                while connection.state == State.OPEN:
+                    await asyncio.sleep(0)
+
+            await asyncio.wait_for(wait_for_closing(), timeout=1)
+            assert await asyncio.wait_for(transport.receive(), timeout=1) == valid
+        finally:
+            release.set()
+            with contextlib.suppress(Exception):
+                await transport.close()
+
+
 # -- ASGI websocket server handler ---------------------------------------------
 
 
@@ -186,10 +338,12 @@ async def prompt(self, session_id: str, prompt: Any = None, **kwargs: Any) -> Pr
 
 
 @pytest.mark.asyncio
-async def test_asgi_websocket_ignores_non_message_frames() -> None:
+async def test_asgi_websocket_reports_invalid_frames_and_ignores_binary() -> None:
     socket = _FakeAsgiSocket()
     socket.client_connect()
-    for frame in ({"bytes": b"binary"}, {"text": "invalid json"}, {"text": "[]"}, {"text": "null"}):
+    frames = [b"binary", "invalid json", "[]", "{}", "null", "123", "1.5", '"x"', "true", "false"]
+    for frame in frames:
+        frame = {"bytes": frame} if isinstance(frame, bytes) else {"text": frame}
         socket._incoming.put_nowait({"type": "websocket.receive", **frame})
     handler = asyncio.create_task(
         handle_websocket(lambda conn: _Agent(), WebSocket({"type": "websocket"}, socket.receive, socket.send))
@@ -198,6 +352,38 @@ async def test_asgi_websocket_ignores_non_message_frames() -> None:
         socket.client_send_text({"jsonrpc": "2.0", "id": 0, "method": "initialize", "params": {"protocolVersion": 1}})
         response = await socket.wait_for_send(lambda m: m.get("id") == 0)
         assert "result" in response
+        assert [m["error"]["code"] for m in socket.sent if "error" in m] == [-32700, *([-32600] * 8)]
+        assert all(m["id"] is None and m["jsonrpc"] == "2.0" for m in socket.sent if "error" in m)
+    finally:
+        socket.client_disconnect()
+        await asyncio.wait_for(handler, timeout=1)
+
+
+@pytest.mark.asyncio
+async def test_asgi_websocket_delivers_notification_after_error_reply_fails() -> None:
+    received = asyncio.Event()
+
+    class Agent(_Agent):
+        async def cancel(self, session_id: str, **kwargs: Any) -> None:
+            assert session_id == "sess-ws"
+            received.set()
+
+    socket = _FakeAsgiSocket()
+    socket.client_connect()
+
+    async def failed_send(message: dict[str, Any]) -> None:
+        if message["type"] == "websocket.send":
+            raise OSError("Peer closed")
+        await socket.send(message)
+
+    for text in ("invalid json", "null", "{}"):
+        socket._incoming.put_nowait({"type": "websocket.receive", "text": text})
+    socket.client_send_text({"jsonrpc": "2.0", "method": "session/cancel", "params": {"sessionId": "sess-ws"}})
+    handler = asyncio.create_task(
+        handle_websocket(lambda conn: Agent(), WebSocket({"type": "websocket"}, socket.receive, failed_send))
+    )
+    try:
+        await asyncio.wait_for(received.wait(), timeout=1)
     finally:
         socket.client_disconnect()
         await asyncio.wait_for(handler, timeout=1)