From f43a81b7e2ab9789b1ddde9b1e31d78471076934 Mon Sep 17 00:00:00 2001 From: hongv <> Date: Sat, 25 Jul 2026 00:27:10 +0800 Subject: [PATCH 1/2] fix: chunk oversized extension responses over WebSocket Large HTML captures were sent as a single response frame and tripped the default 1 MiB WebSocket limit, disconnecting the extension mid-read. Frame big responses as response-chunk messages and fall back to Playwright when an older extension still disconnects during read. Co-authored-by: Cursor --- AGENTS.md | 3 + browser-cli-extension/manifest.json | 2 +- browser-cli-extension/src/background.js | 3 +- .../src/background/response_framing.js | 31 ++++ browser-cli-extension/src/protocol.js | 2 + .../tests/response_framing.test.js | 55 ++++++++ src/browser_cli/daemon/browser_service.py | 3 + src/browser_cli/extension/__init__.py | 2 + src/browser_cli/extension/protocol.py | 2 + src/browser_cli/extension/session.py | 65 +++++++++ tests/unit/test_daemon_browser_service.py | 29 ++++ tests/unit/test_extension_transport.py | 133 ++++++++++++++++++ 12 files changed, 328 insertions(+), 2 deletions(-) create mode 100644 browser-cli-extension/src/background/response_framing.js create mode 100644 browser-cli-extension/tests/response_framing.test.js diff --git a/AGENTS.md b/AGENTS.md index 5131056..411808a 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -223,6 +223,9 @@ the implementation, and where should a change land first. `read` fails in extension mode with `Separator is found, but chunk is longer than limit` or `chunk exceed the limit` -> extension-side page extraction hit a large-result chunking limit on the current page -> inspect `src/browser_cli/daemon/browser_service.py::read_page`, `src/browser_cli/drivers/_extension/page_actions.py`, and `browser-cli-extension/src/background/page_actions.js`; safe fix is a one-shot read fallback to Playwright rather than changing CLI contracts + `read` fails with `Extension disconnected` and daemon log shows `message too big` / `exceeds limit of 1048576` + -> extension tried to send a single oversized `response` frame (common on large pages such as WeChat articles) + -> inspect `browser-cli-extension/src/background/response_framing.js`, `src/browser_cli/extension/session.py` response-chunk reassembly, and `browser_service.read_page` fallback patterns; large responses must be framed as `response-chunk` messages, and stale extensions without framing should fall back to Playwright on read - If the user reports popup/runtime observer drift: start at `src/browser_cli/daemon/runtime_presentation.py`, then `src/browser_cli/extension/session.py`, then `browser-cli-extension/src/background.js`, `browser-cli-extension/src/popup_view.js`, and `browser-cli-extension/src/popup.js`. - If the user reports long-run stability drift across repeated reloads, reconnects, or artifact runs: diff --git a/browser-cli-extension/manifest.json b/browser-cli-extension/manifest.json index db4eee8..f736c7f 100644 --- a/browser-cli-extension/manifest.json +++ b/browser-cli-extension/manifest.json @@ -1,7 +1,7 @@ { "manifest_version": 3, "name": "Browser CLI Bridge", - "version": "0.1.0", + "version": "0.1.1", "description": "Browser CLI real-Chrome backend extension.", "minimum_chrome_version": "116", "permissions": [ diff --git a/browser-cli-extension/src/background.js b/browser-cli-extension/src/background.js index 834a75a..b3adc68 100644 --- a/browser-cli-extension/src/background.js +++ b/browser-cli-extension/src/background.js @@ -14,6 +14,7 @@ import { createInputHandlers } from './background/input_actions.js'; import { createLocatorHandlers } from './background/locator_actions.js'; import { createObserveHandlers } from './background/observe_actions.js'; import { createPageHandlers } from './background/page_actions.js'; +import { sendFramedMessage } from './background/response_framing.js'; import { createTraceHandlers } from './background/trace_actions.js'; import { createVideoHandlers } from './background/video_actions.js'; import { createWorkspaceHandlers } from './background/workspace.js'; @@ -232,7 +233,7 @@ async function connect() { } const response = await handleRequest(payload); if (state.ws && state.ws.readyState === WebSocket.OPEN) { - state.ws.send(JSON.stringify(response)); + sendFramedMessage(state.ws, response); } }; socket.onclose = () => { diff --git a/browser-cli-extension/src/background/response_framing.js b/browser-cli-extension/src/background/response_framing.js new file mode 100644 index 0000000..2a9eeb7 --- /dev/null +++ b/browser-cli-extension/src/background/response_framing.js @@ -0,0 +1,31 @@ +import { RESPONSE_CHUNK_SIZE } from '../protocol.js'; + +/** + * Send a protocol message, chunking oversized JSON so the daemon WebSocket + * receiver never hits its default 1 MiB frame limit. + */ +export function sendFramedMessage(socket, message, chunkSize = RESPONSE_CHUNK_SIZE) { + if (!socket || socket.readyState !== WebSocket.OPEN) { + throw new Error('Extension socket is not connected.'); + } + const encoded = JSON.stringify(message); + const size = Number(chunkSize) > 0 ? Number(chunkSize) : RESPONSE_CHUNK_SIZE; + if (encoded.length <= size) { + socket.send(encoded); + return 1; + } + const id = String(message.id || ''); + let index = 0; + for (let offset = 0; offset < encoded.length; offset += size) { + const chunk = encoded.slice(offset, offset + size); + socket.send(JSON.stringify({ + type: 'response-chunk', + id, + index, + final: offset + size >= encoded.length, + chunk, + })); + index += 1; + } + return index; +} diff --git a/browser-cli-extension/src/protocol.js b/browser-cli-extension/src/protocol.js index 068e45d..fc3383f 100644 --- a/browser-cli-extension/src/protocol.js +++ b/browser-cli-extension/src/protocol.js @@ -2,6 +2,8 @@ export const PROTOCOL_VERSION = '1'; export const DEFAULT_DAEMON_HOST = '127.0.0.1'; export const DEFAULT_DAEMON_PORT = 19825; export const ARTIFACT_CHUNK_SIZE = 256 * 1024; +// Keep framed WebSocket messages comfortably under the default 1 MiB limit. +export const RESPONSE_CHUNK_SIZE = 256 * 1024; export const REQUIRED_CAPABILITIES = [ 'open', diff --git a/browser-cli-extension/tests/response_framing.test.js b/browser-cli-extension/tests/response_framing.test.js new file mode 100644 index 0000000..482b9a9 --- /dev/null +++ b/browser-cli-extension/tests/response_framing.test.js @@ -0,0 +1,55 @@ +import test from 'node:test' +import assert from 'node:assert/strict' + +import { RESPONSE_CHUNK_SIZE } from '../src/protocol.js' +import { sendFramedMessage } from '../src/background/response_framing.js' + +function installSocketStub() { + const sent = [] + return { + sent, + socket: { + readyState: 1, + send(payload) { + sent.push(payload) + }, + }, + } +} + +test('sendFramedMessage sends small responses as a single frame', () => { + const { sent, socket } = installSocketStub() + const message = { type: 'response', id: 'req-1', ok: true, data: { html: '

hi

' } } + assert.equal(sendFramedMessage(socket, message), 1) + assert.equal(sent.length, 1) + assert.deepEqual(JSON.parse(sent[0]), message) +}) + +test('sendFramedMessage chunks oversized responses under the frame limit', () => { + const { sent, socket } = installSocketStub() + const html = 'x'.repeat(RESPONSE_CHUNK_SIZE + 128) + const message = { type: 'response', id: 'req-2', ok: true, data: { html } } + const frameCount = sendFramedMessage(socket, message, 1024) + assert.ok(frameCount > 1) + assert.equal(sent.length, frameCount) + + const chunks = sent.map((raw) => JSON.parse(raw)) + for (const chunk of chunks) { + assert.equal(chunk.type, 'response-chunk') + assert.equal(chunk.id, 'req-2') + assert.ok(JSON.stringify(chunk).length < RESPONSE_CHUNK_SIZE) + } + assert.equal(chunks.at(-1)?.final, true) + const assembled = chunks + .sort((left, right) => left.index - right.index) + .map((chunk) => chunk.chunk) + .join('') + assert.deepEqual(JSON.parse(assembled), message) +}) + +test('sendFramedMessage rejects a closed socket', () => { + assert.throws( + () => sendFramedMessage({ readyState: 3, send() {} }, { type: 'response', id: 'x' }), + /not connected/, + ) +}) diff --git a/src/browser_cli/daemon/browser_service.py b/src/browser_cli/daemon/browser_service.py index 4b33f44..bb5b82a 100644 --- a/src/browser_cli/daemon/browser_service.py +++ b/src/browser_cli/daemon/browser_service.py @@ -65,6 +65,9 @@ class BrowserService: "Separator is not found, and chunk exceed the limit", "chunk is longer than limit", "chunk exceed the limit", + "Extension disconnected", + "message too big", + "exceeds limit of", ) def __init__( diff --git a/src/browser_cli/extension/__init__.py b/src/browser_cli/extension/__init__.py index b6a0641..6b50252 100644 --- a/src/browser_cli/extension/__init__.py +++ b/src/browser_cli/extension/__init__.py @@ -6,6 +6,7 @@ CORE_EXTENSION_CAPABILITIES, OPTIONAL_EXTENSION_CAPABILITIES, REQUIRED_EXTENSION_CAPABILITIES, + RESPONSE_CHUNK_SIZE, ExtensionArtifactBegin, ExtensionArtifactChunk, ExtensionArtifactEnd, @@ -21,6 +22,7 @@ "OPTIONAL_EXTENSION_CAPABILITIES", "ALL_EXTENSION_CAPABILITIES", "ARTIFACT_CHUNK_SIZE", + "RESPONSE_CHUNK_SIZE", "ExtensionHello", "ExtensionRequest", "ExtensionResponse", diff --git a/src/browser_cli/extension/protocol.py b/src/browser_cli/extension/protocol.py index 517a0b2..f217227 100644 --- a/src/browser_cli/extension/protocol.py +++ b/src/browser_cli/extension/protocol.py @@ -80,6 +80,8 @@ CORE_EXTENSION_CAPABILITIES = REQUIRED_EXTENSION_CAPABILITIES PROTOCOL_VERSION = "1" ARTIFACT_CHUNK_SIZE = 256 * 1024 +# Keep framed WebSocket messages comfortably under the default 1 MiB limit. +RESPONSE_CHUNK_SIZE = 256 * 1024 @dataclass(slots=True, frozen=True) diff --git a/src/browser_cli/extension/session.py b/src/browser_cli/extension/session.py index f9d8058..bf91cd0 100644 --- a/src/browser_cli/extension/session.py +++ b/src/browser_cli/extension/session.py @@ -58,6 +58,27 @@ def to_payload(self) -> dict[str, Any]: } +@dataclass(slots=True) +class _ResponseChunkBuffer: + chunks: dict[int, str] = field(default_factory=dict) + + def append(self, index: int, chunk: str) -> None: + self.chunks[index] = chunk + + def assemble(self) -> str: + if not self.chunks: + return "" + indexes = sorted(self.chunks) + expected = list(range(indexes[-1] + 1)) + if indexes != expected: + missing = [index for index in expected if index not in self.chunks] + raise OperationFailedError( + f"Extension response chunks are incomplete; missing indexes: {missing}", + error_code="EXTENSION_RESPONSE_CHUNK_INCOMPLETE", + ) + return "".join(self.chunks[index] for index in indexes) + + @dataclass(slots=True) class ExtensionSession: websocket: ServerConnection @@ -66,6 +87,7 @@ class ExtensionSession: _artifact_events: dict[str, asyncio.Event] = field(init=False, repr=False) _artifact_buffers: dict[tuple[str, str], _ArtifactBuffer] = field(init=False, repr=False) _completed_artifacts: dict[str, list[dict[str, Any]]] = field(init=False, repr=False) + _response_buffers: dict[str, _ResponseChunkBuffer] = field(init=False, repr=False) _lock: asyncio.Lock = field(init=False, repr=False) def __post_init__(self) -> None: @@ -73,6 +95,7 @@ def __post_init__(self) -> None: self._artifact_events = {} self._artifact_buffers = {} self._completed_artifacts = {} + self._response_buffers = {} self._lock = asyncio.Lock() @property @@ -123,6 +146,45 @@ def resolve_response(self, response: ExtensionResponse) -> None: return future.set_result(response) + def append_response_chunk(self, payload: dict[str, Any]) -> None: + request_id = str(payload.get("id") or "") + if not request_id: + return + future = self._pending.get(request_id) + if future is None or future.done(): + self._response_buffers.pop(request_id, None) + return + buffer = self._response_buffers.setdefault(request_id, _ResponseChunkBuffer()) + buffer.append(int(payload.get("index") or 0), str(payload.get("chunk") or "")) + if not bool(payload.get("final")): + return + try: + assembled = buffer.assemble() + message = json.loads(assembled) + if not isinstance(message, dict): + raise OperationFailedError( + "Extension chunked response is not a JSON object.", + error_code="EXTENSION_RESPONSE_CHUNK_INVALID", + ) + response = ExtensionResponse.from_message(message) + except OperationFailedError as exc: + self._response_buffers.pop(request_id, None) + if not future.done(): + future.set_exception(exc) + return + except (json.JSONDecodeError, TypeError, ValueError) as exc: + self._response_buffers.pop(request_id, None) + if not future.done(): + future.set_exception( + OperationFailedError( + f"Extension chunked response is invalid: {exc}", + error_code="EXTENSION_RESPONSE_CHUNK_INVALID", + ) + ) + return + self._response_buffers.pop(request_id, None) + self.resolve_response(response) + def fail_all(self, message: str) -> None: for future in list(self._pending.values()): if not future.done(): @@ -132,6 +194,7 @@ def fail_all(self, message: str) -> None: self._pending.clear() self._artifact_buffers.clear() self._completed_artifacts.clear() + self._response_buffers.clear() for event in list(self._artifact_events.values()): event.set() self._artifact_events.clear() @@ -302,6 +365,8 @@ async def _handle_websocket(self, websocket: ServerConnection) -> None: message_type = str(payload.get("type") or "") if message_type == "response": session.resolve_response(ExtensionResponse.from_message(dict(payload))) + elif message_type == "response-chunk": + session.append_response_chunk(dict(payload)) elif message_type == "artifact-begin": session.begin_artifact(ExtensionArtifactBegin.from_message(dict(payload))) elif message_type == "artifact-chunk": diff --git a/tests/unit/test_daemon_browser_service.py b/tests/unit/test_daemon_browser_service.py index 54ced48..d2b6b1f 100644 --- a/tests/unit/test_daemon_browser_service.py +++ b/tests/unit/test_daemon_browser_service.py @@ -639,6 +639,35 @@ async def _playwright_capture_html(page_id: str) -> dict[str, str]: asyncio.run(_scenario()) +def test_read_page_falls_back_to_playwright_when_extension_disconnects_mid_read( + _patched_browser_service: _FakeExtensionHub, +) -> None: + async def _scenario() -> None: + _patched_browser_service.connect() + service = browser_service_module.BrowserService(TabRegistry()) + await service.ensure_started() + assert service.active_driver_name == "extension" + + async def _extension_capture_html(page_id: str) -> dict[str, str]: + raise RuntimeError("Extension disconnected.") + + async def _playwright_capture_html(page_id: str) -> dict[str, str]: + return {"page_id": page_id, "html": "ok"} + + service._extension.capture_html = _extension_capture_html # type: ignore[method-assign] # noqa: SLF001 + service._playwright.capture_html = _playwright_capture_html # type: ignore[method-assign] # noqa: SLF001 + + payload = await service.read_page(url="https://example.com/large-doc", output_mode="html") + + assert payload["body"] == "ok" + assert payload["driver_fallback"]["reason"] == "extension-read-fallback" + assert "Extension disconnected" in payload["driver_fallback"]["error"] + assert service.active_driver_name == "playwright" + await service.stop() + + asyncio.run(_scenario()) + + def test_browser_service_downgrade_survives_previous_driver_stop_failure( _patched_browser_service: _FakeExtensionHub, monkeypatch: pytest.MonkeyPatch, diff --git a/tests/unit/test_extension_transport.py b/tests/unit/test_extension_transport.py index b8c8ed3..27db712 100644 --- a/tests/unit/test_extension_transport.py +++ b/tests/unit/test_extension_transport.py @@ -491,3 +491,136 @@ async def _scenario() -> None: await hub.stop() asyncio.run(_scenario()) + + +def test_extension_session_reassembles_chunked_responses(monkeypatch, tmp_path: Path) -> None: + async def _scenario() -> None: + monkeypatch.setenv(APP_HOME_ENV, str(tmp_path / ".browser-cli-runtime")) + monkeypatch.setenv(EXTENSION_PORT_ENV, str(_unused_port())) + + hub = ExtensionHub() + await hub.ensure_started() + app_paths = get_app_paths() + + async with websockets.connect(app_paths.extension_ws_url) as websocket: + await websocket.send( + json.dumps( + { + "type": "hello", + "protocol_version": PROTOCOL_VERSION, + "extension_version": "0.1.0-test", + "browser_name": "Chrome", + "browser_version": "146", + "capabilities": sorted(REQUIRED_EXTENSION_CAPABILITIES), + "workspace_window_state": {"connected": True}, + "extension_instance_id": "ext-test", + } + ) + ) + session = await hub.wait_for_session(timeout_seconds=1.0) + assert session is not None + + request_task = asyncio.create_task(session.send_request("capture-html", {})) + raw_request = json.loads(await websocket.recv()) + request_id = raw_request["id"] + html = "A" * 3000 + encoded = json.dumps( + { + "type": "response", + "id": request_id, + "ok": True, + "data": {"html": html}, + } + ) + await websocket.send( + json.dumps( + { + "type": "response-chunk", + "id": request_id, + "index": 0, + "final": False, + "chunk": encoded[:1200], + } + ) + ) + await websocket.send( + json.dumps( + { + "type": "response-chunk", + "id": request_id, + "index": 1, + "final": False, + "chunk": encoded[1200:2400], + } + ) + ) + await websocket.send( + json.dumps( + { + "type": "response-chunk", + "id": request_id, + "index": 2, + "final": True, + "chunk": encoded[2400:], + } + ) + ) + response = await request_task + assert response["html"] == html + assert session._response_buffers == {} + + await hub.stop() + + asyncio.run(_scenario()) + + +def test_extension_session_rejects_incomplete_chunked_responses( + monkeypatch, tmp_path: Path +) -> None: + async def _scenario() -> None: + monkeypatch.setenv(APP_HOME_ENV, str(tmp_path / ".browser-cli-runtime")) + monkeypatch.setenv(EXTENSION_PORT_ENV, str(_unused_port())) + + hub = ExtensionHub() + await hub.ensure_started() + app_paths = get_app_paths() + + async with websockets.connect(app_paths.extension_ws_url) as websocket: + await websocket.send( + json.dumps( + { + "type": "hello", + "protocol_version": PROTOCOL_VERSION, + "extension_version": "0.1.0-test", + "browser_name": "Chrome", + "browser_version": "146", + "capabilities": sorted(REQUIRED_EXTENSION_CAPABILITIES), + "workspace_window_state": {"connected": True}, + "extension_instance_id": "ext-test", + } + ) + ) + session = await hub.wait_for_session(timeout_seconds=1.0) + assert session is not None + + request_task = asyncio.create_task(session.send_request("capture-html", {})) + raw_request = json.loads(await websocket.recv()) + request_id = raw_request["id"] + await websocket.send( + json.dumps( + { + "type": "response-chunk", + "id": request_id, + "index": 1, + "final": True, + "chunk": '{"type":"response","id":"x","ok":true,"data":{}}', + } + ) + ) + with pytest.raises(OperationFailedError, match="incomplete"): + await request_task + assert session._response_buffers == {} + + await hub.stop() + + asyncio.run(_scenario()) From 6d4b41bce849f4d5e32ad61df06f3a5143cdbac1 Mon Sep 17 00:00:00 2001 From: hongv <> Date: Sat, 25 Jul 2026 00:55:54 +0800 Subject: [PATCH 2/2] fix: harden extension response-chunk validation Validate chunk indexes and reassembled response ids before resolving pending requests, clean response buffers on request completion, and cover the new failure paths in transport and read-fallback tests. Co-authored-by: Cursor --- src/browser_cli/extension/__init__.py | 2 + src/browser_cli/extension/protocol.py | 2 + src/browser_cli/extension/session.py | 47 +++++++- tests/unit/test_daemon_browser_service.py | 13 ++- tests/unit/test_extension_transport.py | 128 +++++++++++++++++++++- 5 files changed, 184 insertions(+), 8 deletions(-) diff --git a/src/browser_cli/extension/__init__.py b/src/browser_cli/extension/__init__.py index 6b50252..413cae9 100644 --- a/src/browser_cli/extension/__init__.py +++ b/src/browser_cli/extension/__init__.py @@ -4,6 +4,7 @@ ALL_EXTENSION_CAPABILITIES, ARTIFACT_CHUNK_SIZE, CORE_EXTENSION_CAPABILITIES, + MAX_RESPONSE_CHUNK_INDEX, OPTIONAL_EXTENSION_CAPABILITIES, REQUIRED_EXTENSION_CAPABILITIES, RESPONSE_CHUNK_SIZE, @@ -23,6 +24,7 @@ "ALL_EXTENSION_CAPABILITIES", "ARTIFACT_CHUNK_SIZE", "RESPONSE_CHUNK_SIZE", + "MAX_RESPONSE_CHUNK_INDEX", "ExtensionHello", "ExtensionRequest", "ExtensionResponse", diff --git a/src/browser_cli/extension/protocol.py b/src/browser_cli/extension/protocol.py index f217227..d187604 100644 --- a/src/browser_cli/extension/protocol.py +++ b/src/browser_cli/extension/protocol.py @@ -82,6 +82,8 @@ ARTIFACT_CHUNK_SIZE = 256 * 1024 # Keep framed WebSocket messages comfortably under the default 1 MiB limit. RESPONSE_CHUNK_SIZE = 256 * 1024 +# Cap chunk indexes so a hostile/malformed final index cannot force huge allocations. +MAX_RESPONSE_CHUNK_INDEX = 16_384 @dataclass(slots=True, frozen=True) diff --git a/src/browser_cli/extension/session.py b/src/browser_cli/extension/session.py index bf91cd0..7ea2bd2 100644 --- a/src/browser_cli/extension/session.py +++ b/src/browser_cli/extension/session.py @@ -25,6 +25,7 @@ from browser_cli.errors import ExtensionPortInUseError, OperationFailedError from .protocol import ( + MAX_RESPONSE_CHUNK_INDEX, ExtensionArtifactBegin, ExtensionArtifactChunk, ExtensionArtifactEnd, @@ -69,16 +70,38 @@ def assemble(self) -> str: if not self.chunks: return "" indexes = sorted(self.chunks) - expected = list(range(indexes[-1] + 1)) - if indexes != expected: - missing = [index for index in expected if index not in self.chunks] + min_index = indexes[0] + max_index = indexes[-1] + if min_index != 0 or len(indexes) != max_index + 1: raise OperationFailedError( - f"Extension response chunks are incomplete; missing indexes: {missing}", + "Extension response chunks are incomplete; " + f"received {len(indexes)} chunks covering {min_index}..{max_index}", error_code="EXTENSION_RESPONSE_CHUNK_INCOMPLETE", ) return "".join(self.chunks[index] for index in indexes) +def _parse_response_chunk_index(raw: Any) -> int: + if isinstance(raw, bool) or raw is None: + raise OperationFailedError( + f"Extension response chunk index is invalid: {raw!r}", + error_code="EXTENSION_RESPONSE_CHUNK_INVALID", + ) + try: + index = int(raw) + except (TypeError, ValueError) as exc: + raise OperationFailedError( + f"Extension response chunk index is invalid: {raw!r}", + error_code="EXTENSION_RESPONSE_CHUNK_INVALID", + ) from exc + if index < 0 or index > MAX_RESPONSE_CHUNK_INDEX: + raise OperationFailedError( + f"Extension response chunk index out of bounds: {index}", + error_code="EXTENSION_RESPONSE_CHUNK_INVALID", + ) + return index + + @dataclass(slots=True) class ExtensionSession: websocket: ServerConnection @@ -129,6 +152,7 @@ async def send_request( finally: self._pending.pop(request.id, None) self._artifact_events.pop(request.id, None) + self._response_buffers.pop(request.id, None) if not response.ok: raise OperationFailedError( response.error_message or f"Extension request failed: {action}", @@ -154,8 +178,15 @@ def append_response_chunk(self, payload: dict[str, Any]) -> None: if future is None or future.done(): self._response_buffers.pop(request_id, None) return + try: + index = _parse_response_chunk_index(payload.get("index")) + except OperationFailedError as exc: + self._response_buffers.pop(request_id, None) + if not future.done(): + future.set_exception(exc) + return buffer = self._response_buffers.setdefault(request_id, _ResponseChunkBuffer()) - buffer.append(int(payload.get("index") or 0), str(payload.get("chunk") or "")) + buffer.append(index, str(payload.get("chunk") or "")) if not bool(payload.get("final")): return try: @@ -167,6 +198,12 @@ def append_response_chunk(self, payload: dict[str, Any]) -> None: error_code="EXTENSION_RESPONSE_CHUNK_INVALID", ) response = ExtensionResponse.from_message(message) + if response.id != request_id: + raise OperationFailedError( + "Extension chunked response id mismatch: " + f"expected {request_id}, got {response.id}", + error_code="EXTENSION_RESPONSE_CHUNK_INVALID", + ) except OperationFailedError as exc: self._response_buffers.pop(request_id, None) if not future.done(): diff --git a/tests/unit/test_daemon_browser_service.py b/tests/unit/test_daemon_browser_service.py index d2b6b1f..becd192 100644 --- a/tests/unit/test_daemon_browser_service.py +++ b/tests/unit/test_daemon_browser_service.py @@ -639,8 +639,17 @@ async def _playwright_capture_html(page_id: str) -> dict[str, str]: asyncio.run(_scenario()) +@pytest.mark.parametrize( + "extension_error", + [ + "Extension disconnected.", + "message too big", + "exceeds limit of", + ], +) def test_read_page_falls_back_to_playwright_when_extension_disconnects_mid_read( _patched_browser_service: _FakeExtensionHub, + extension_error: str, ) -> None: async def _scenario() -> None: _patched_browser_service.connect() @@ -649,7 +658,7 @@ async def _scenario() -> None: assert service.active_driver_name == "extension" async def _extension_capture_html(page_id: str) -> dict[str, str]: - raise RuntimeError("Extension disconnected.") + raise RuntimeError(extension_error) async def _playwright_capture_html(page_id: str) -> dict[str, str]: return {"page_id": page_id, "html": "ok"} @@ -661,7 +670,7 @@ async def _playwright_capture_html(page_id: str) -> dict[str, str]: assert payload["body"] == "ok" assert payload["driver_fallback"]["reason"] == "extension-read-fallback" - assert "Extension disconnected" in payload["driver_fallback"]["error"] + assert payload["driver_fallback"]["error"] == extension_error assert service.active_driver_name == "playwright" await service.stop() diff --git a/tests/unit/test_extension_transport.py b/tests/unit/test_extension_transport.py index 27db712..5ffc1b2 100644 --- a/tests/unit/test_extension_transport.py +++ b/tests/unit/test_extension_transport.py @@ -11,7 +11,11 @@ from browser_cli.constants import APP_HOME_ENV, EXTENSION_PORT_ENV, get_app_paths from browser_cli.errors import OperationFailedError -from browser_cli.extension.protocol import PROTOCOL_VERSION, REQUIRED_EXTENSION_CAPABILITIES +from browser_cli.extension.protocol import ( + MAX_RESPONSE_CHUNK_INDEX, + PROTOCOL_VERSION, + REQUIRED_EXTENSION_CAPABILITIES, +) from browser_cli.extension.session import ExtensionHub @@ -624,3 +628,125 @@ async def _scenario() -> None: await hub.stop() asyncio.run(_scenario()) + + +def test_extension_session_rejects_mismatched_chunked_response_id( + monkeypatch, tmp_path: Path +) -> None: + async def _scenario() -> None: + monkeypatch.setenv(APP_HOME_ENV, str(tmp_path / ".browser-cli-runtime")) + monkeypatch.setenv(EXTENSION_PORT_ENV, str(_unused_port())) + + hub = ExtensionHub() + await hub.ensure_started() + app_paths = get_app_paths() + + async with websockets.connect(app_paths.extension_ws_url) as websocket: + await websocket.send( + json.dumps( + { + "type": "hello", + "protocol_version": PROTOCOL_VERSION, + "extension_version": "0.1.0-test", + "browser_name": "Chrome", + "browser_version": "146", + "capabilities": sorted(REQUIRED_EXTENSION_CAPABILITIES), + "workspace_window_state": {"connected": True}, + "extension_instance_id": "ext-test", + } + ) + ) + session = await hub.wait_for_session(timeout_seconds=1.0) + assert session is not None + + request_task = asyncio.create_task(session.send_request("capture-html", {})) + raw_request = json.loads(await websocket.recv()) + request_id = raw_request["id"] + encoded = json.dumps( + { + "type": "response", + "id": "other-request-id", + "ok": True, + "data": {"html": "

nope

"}, + } + ) + await websocket.send( + json.dumps( + { + "type": "response-chunk", + "id": request_id, + "index": 0, + "final": True, + "chunk": encoded, + } + ) + ) + with pytest.raises(OperationFailedError, match="id mismatch") as exc_info: + await request_task + assert exc_info.value.error_code == "EXTENSION_RESPONSE_CHUNK_INVALID" + assert session._response_buffers == {} + + await hub.stop() + + asyncio.run(_scenario()) + + +@pytest.mark.parametrize( + ("index", "match"), + [ + ("not-an-index", "invalid"), + (-1, "out of bounds"), + (MAX_RESPONSE_CHUNK_INDEX + 1, "out of bounds"), + ], +) +def test_extension_session_rejects_malformed_or_oversized_chunk_index( + monkeypatch, tmp_path: Path, index: object, match: str +) -> None: + async def _scenario() -> None: + monkeypatch.setenv(APP_HOME_ENV, str(tmp_path / ".browser-cli-runtime")) + monkeypatch.setenv(EXTENSION_PORT_ENV, str(_unused_port())) + + hub = ExtensionHub() + await hub.ensure_started() + app_paths = get_app_paths() + + async with websockets.connect(app_paths.extension_ws_url) as websocket: + await websocket.send( + json.dumps( + { + "type": "hello", + "protocol_version": PROTOCOL_VERSION, + "extension_version": "0.1.0-test", + "browser_name": "Chrome", + "browser_version": "146", + "capabilities": sorted(REQUIRED_EXTENSION_CAPABILITIES), + "workspace_window_state": {"connected": True}, + "extension_instance_id": "ext-test", + } + ) + ) + session = await hub.wait_for_session(timeout_seconds=1.0) + assert session is not None + + request_task = asyncio.create_task(session.send_request("capture-html", {})) + raw_request = json.loads(await websocket.recv()) + request_id = raw_request["id"] + await websocket.send( + json.dumps( + { + "type": "response-chunk", + "id": request_id, + "index": index, + "final": True, + "chunk": '{"type":"response","id":"x","ok":true,"data":{}}', + } + ) + ) + with pytest.raises(OperationFailedError, match=match) as exc_info: + await request_task + assert exc_info.value.error_code == "EXTENSION_RESPONSE_CHUNK_INVALID" + assert session._response_buffers == {} + + await hub.stop() + + asyncio.run(_scenario())