From 32a4397066d4968544027776982e0665503d71f4 Mon Sep 17 00:00:00 2001 From: 0xallam Date: Thu, 19 Mar 2026 19:38:19 -0700 Subject: [PATCH] Fix CDP proxy and bump sandbox image --- strix/runtime/tool_server.py | 39 +++++++++++++++++++++--------------- 1 file changed, 23 insertions(+), 16 deletions(-) diff --git a/strix/runtime/tool_server.py b/strix/runtime/tool_server.py index f391558a..6b2f6350 100644 --- a/strix/runtime/tool_server.py +++ b/strix/runtime/tool_server.py @@ -10,7 +10,7 @@ from typing import Any import httpx import uvicorn import websockets -from fastapi import Depends, FastAPI, HTTPException, WebSocket, status +from fastapi import Depends, FastAPI, HTTPException, WebSocket, WebSocketDisconnect, status from fastapi.security import HTTPAuthorizationCredentials, HTTPBearer from pydantic import BaseModel, ValidationError @@ -174,6 +174,7 @@ async def health_check() -> dict[str, Any]: # WebSocket proxy to that exact path. No catch-all proxy — avoids SSRF. _cdp_ws_internal: str | None = None +_CDP_WS_MAX_SIZE = 16 * 1024 * 1024 async def _resolve_cdp_ws() -> None: @@ -214,28 +215,34 @@ async def cdp_proxy_ws(ws: WebSocket) -> None: await ws.accept() - async with websockets.connect(_cdp_ws_internal) as upstream: + async with websockets.connect(_cdp_ws_internal, max_size=_CDP_WS_MAX_SIZE) as upstream: async def relay_client_to_upstream() -> None: - async for msg in ws.iter_text(): - await upstream.send(msg) + try: + async for msg in ws.iter_text(): + await upstream.send(msg) + except (WebSocketDisconnect, websockets.ConnectionClosed): + return async def relay_upstream_to_client() -> None: - async for msg in upstream: - if isinstance(msg, str): - await ws.send_text(msg) - else: - await ws.send_bytes(msg) + try: + async for msg in upstream: + if isinstance(msg, str): + await ws.send_text(msg) + else: + await ws.send_bytes(msg) + except (WebSocketDisconnect, websockets.ConnectionClosed): + return - _, pending = await asyncio.wait( - [ - asyncio.create_task(relay_client_to_upstream()), - asyncio.create_task(relay_upstream_to_client()), - ], - return_when=asyncio.FIRST_COMPLETED, - ) + tasks = [ + asyncio.create_task(relay_client_to_upstream()), + asyncio.create_task(relay_upstream_to_client()), + ] + done, pending = await asyncio.wait(tasks, return_when=asyncio.FIRST_COMPLETED) for t in pending: t.cancel() + await asyncio.gather(*pending, return_exceptions=True) + await asyncio.gather(*done, return_exceptions=True) def signal_handler(_signum: int, _frame: Any) -> None: