From 0386181410d1f6fd595ecd3c55da1bd50666668e Mon Sep 17 00:00:00 2001 From: br-mbp Date: Fri, 16 Oct 2026 15:40:33 +0800 Subject: [PATCH 1/2] fix(ws): ignore non-object JSON frames in client transport --- docs/web-transport.md | 2 ++ src/acp/ws/client.py | 4 +++- tests/http/test_websocket.py | 41 ++++++++++++++++++++++++++++++++++++ 3 files changed, 46 insertions(+), 1 deletion(-) diff --git a/docs/web-transport.md b/docs/web-transport.md index ce77eba4..dcbf3416 100644 --- a/docs/web-transport.md +++ b/docs/web-transport.md @@ -12,6 +12,8 @@ The SDK can run ACP over two remote connectivity profiles in addition to stdio: - **WebSocket** — a `GET` upgrade on the same endpoint carrying full-duplex JSON-RPC text frames, one message per frame. Binary frames are ignored, as the transport RFD specifies; the server logs a warning on the first one. + Both endpoints ignore malformed JSON and JSON values that are not objects, + continuing to receive subsequent messages. Both reuse the existing JSON-RPC message format and ACP lifecycle (`initialize` → session methods → close). diff --git a/src/acp/ws/client.py b/src/acp/ws/client.py index 18acf5b1..f91cd67c 100644 --- a/src/acp/ws/client.py +++ b/src/acp/ws/client.py @@ -53,9 +53,11 @@ 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 + if isinstance(message, dict): + return message async def close(self) -> None: if self._closed: diff --git a/tests/http/test_websocket.py b/tests/http/test_websocket.py index eca6cbd5..f136c927 100644 --- a/tests/http/test_websocket.py +++ b/tests/http/test_websocket.py @@ -11,6 +11,7 @@ from starlette.websockets import WebSocket from websockets.asyncio.server import serve +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 +90,46 @@ 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_ignores_non_object_frames(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_non_object_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} + + # -- ASGI websocket server handler --------------------------------------------- From ca27aef8d90bff838c15a8566c9e3c16c5fe2fa7 Mon Sep 17 00:00:00 2001 From: Chojan Shang Date: Sun, 18 Oct 2026 16:46:39 +0800 Subject: [PATCH 2/2] fix(ws): report invalid frames without losing buffered messages --- docs/web-transport.md | 2 - src/acp/ws/client.py | 17 +++- src/acp/ws/server.py | 19 ++++- tests/http/test_websocket.py | 157 +++++++++++++++++++++++++++++++++-- 4 files changed, 180 insertions(+), 15 deletions(-) diff --git a/docs/web-transport.md b/docs/web-transport.md index dcbf3416..ce77eba4 100644 --- a/docs/web-transport.md +++ b/docs/web-transport.md @@ -12,8 +12,6 @@ The SDK can run ACP over two remote connectivity profiles in addition to stdio: - **WebSocket** — a `GET` upgrade on the same endpoint carrying full-duplex JSON-RPC text frames, one message per frame. Binary frames are ignored, as the transport RFD specifies; the server logs a warning on the first one. - Both endpoints ignore malformed JSON and JSON values that are not objects, - continuing to receive subsequent messages. Both reuse the existing JSON-RPC message format and ACP lifecycle (`initialize` → session methods → close). diff --git a/src/acp/ws/client.py b/src/acp/ws/client.py index f91cd67c..531739f7 100644 --- a/src/acp/ws/client.py +++ b/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 @@ -55,9 +56,19 @@ async def receive(self) -> dict[str, Any] | None: try: message = json.loads(frame) 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() + # 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: diff --git a/src/acp/ws/server.py b/src/acp/ws/server.py index a563732c..543a790d 100644 --- a/src/acp/ws/server.py +++ b/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: diff --git a/tests/http/test_websocket.py b/tests/http/test_websocket.py index f136c927..5bbddd27 100644 --- a/tests/http/test_websocket.py +++ b/tests/http/test_websocket.py @@ -3,6 +3,7 @@ from __future__ import annotations import asyncio +import contextlib import json import logging from typing import Any @@ -10,6 +11,7 @@ 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 @@ -91,8 +93,8 @@ async def send_binary_then_text(ws: Any) -> None: @pytest.mark.asyncio -@pytest.mark.parametrize("frame", ["null", "[]", "123", "1.5", '"x"', "true", "false", "invalid json", b"binary"]) -async def test_client_transport_ignores_non_object_frames(frame: str | bytes) -> None: +@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: @@ -111,8 +113,8 @@ async def send_frames(ws: Any) -> None: @pytest.mark.asyncio -@pytest.mark.parametrize("frame", ["null", "[]"]) -async def test_client_connection_recovers_after_non_object_frame(frame: str) -> None: +@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) @@ -130,6 +132,115 @@ async def handler(method: str, params: Any, is_notification: bool) -> None: 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 --------------------------------------------- @@ -227,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)) @@ -239,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)