From 667afe1f3b50c1ca9e94a6430ddf8f311f1a59fd Mon Sep 17 00:00:00 2001 From: "Yuichiro Tachibana (Tsuchiya)" Date: Tue, 11 Aug 2026 15:35:23 +0900 Subject: [PATCH] fix(runtime-sdk): implement WebSocket close semantics and app-task lifetime An app-initiated websocket.close only logged a warning, the app task was never registered with the runtime through wait_until, and a task ending after accept without sending close left the transport open. Handle the close message, register the task, and close the transport (1011 on error) when the app ends without doing so. --- packages/runtime-sdk/src/asgi.py | 38 ++++++++++++++++++- .../tests/test_ws_disconnect.py | 16 ++++++++ .../workerd-test/asgi-ws-disconnect/worker.py | 35 ++++++++++++++++- 3 files changed, 87 insertions(+), 2 deletions(-) diff --git a/packages/runtime-sdk/src/asgi.py b/packages/runtime-sdk/src/asgi.py index cd4c6a1..c5ecaf5 100644 --- a/packages/runtime-sdk/src/asgi.py +++ b/packages/runtime-sdk/src/asgi.py @@ -348,7 +348,10 @@ def onmessage(evt): server.onclose = onclose server.onmessage = onmessage + app_closed = False + async def ws_send(got): + nonlocal app_closed if got["type"] == "websocket.send": b = got.get("bytes", None) s = got.get("text", None) @@ -360,6 +363,10 @@ async def ws_send(got): if s is not None: server.send(s) + elif got["type"] == "websocket.close": + app_closed = True + server.close(got.get("code", 1000), got.get("reason", "")) + else: logger.warning(" == Not implemented %s", got["type"]) @@ -368,7 +375,36 @@ async def ws_receive(): return received env = {} - run_in_background(app(request_to_scope(req, env, ws=True), ws_receive, ws_send)) + # The app task must be registered with the runtime via wait_until (as + # process_request does): waitUntil is the platform's mechanism for + # extending work past the response, so the task's lifetime after the 101 + # is otherwise unguaranteed. + from pyodide.ffi import create_proxy + + from workers import wait_until + + task = create_task(app(request_to_scope(req, env, ws=True), ws_receive, ws_send)) + background_tasks.add(task) + task_proxy = create_proxy(task) + + def _on_done(t): + background_tasks.discard(t) + exc = t.exception() if not t.cancelled() else None + if exc is not None: + logger.error("Exception in ASGI WebSocket application", exc_info=exc) + if not app_closed: + # Per the ASGI spec, the server closes the transport when the app + # task ends without sending websocket.close (1011 on error); + # otherwise the client keeps a half-open connection and hangs + # instead of reconnecting. + try: + server.close(1011 if exc is not None else 1000, "") + except Exception: + pass # the peer may already have closed the socket + task_proxy.destroy() + + task.add_done_callback(_on_done) + wait_until(task_proxy) return Response.new(None, status=101, webSocket=client) diff --git a/packages/runtime-sdk/tests/workerd-test/asgi-ws-disconnect/tests/test_ws_disconnect.py b/packages/runtime-sdk/tests/workerd-test/asgi-ws-disconnect/tests/test_ws_disconnect.py index 02112b6..6506832 100644 --- a/packages/runtime-sdk/tests/workerd-test/asgi-ws-disconnect/tests/test_ws_disconnect.py +++ b/packages/runtime-sdk/tests/workerd-test/asgi-ws-disconnect/tests/test_ws_disconnect.py @@ -140,3 +140,19 @@ async def test_empty_frames_reach_the_client(): assert messages[0] == "" assert messages[1].to_bytes() == b"" assert messages[2] == "done" + + +@pytest.mark.asyncio +async def test_app_initiated_close_reaches_the_client(): + async with _ws_session("/ws-app-close") as ws: + close = _listen(ws, "close") + evt = await asyncio.wait_for(close, TIMEOUT_S) + assert evt.code == 4001 + + +@pytest.mark.asyncio +async def test_app_crash_closes_the_transport(): + async with _ws_session("/ws-crash") as ws: + close = _listen(ws, "close") + evt = await asyncio.wait_for(close, TIMEOUT_S) + assert evt.code == 1011 diff --git a/packages/runtime-sdk/tests/workerd-test/asgi-ws-disconnect/worker.py b/packages/runtime-sdk/tests/workerd-test/asgi-ws-disconnect/worker.py index 4edf52b..3186000 100644 --- a/packages/runtime-sdk/tests/workerd-test/asgi-ws-disconnect/worker.py +++ b/packages/runtime-sdk/tests/workerd-test/asgi-ws-disconnect/worker.py @@ -101,16 +101,49 @@ async def __call__(self, scope, receive, send): await send({"type": "websocket.send", "text": "done"}) +class WSAppCloseApp: + """Accepts, then closes from the app side with a custom code.""" + + async def __call__(self, scope, receive, send): + if scope["type"] == "lifespan": + await _drain_lifespan(receive, send) + return + message = await receive() + assert message["type"] == "websocket.connect" + await send({"type": "websocket.accept"}) + await send({"type": "websocket.close", "code": 4001, "reason": "done"}) + + +class WSCrashApp: + """Accepts, then raises: the server must close the transport (1011).""" + + async def __call__(self, scope, receive, send): + if scope["type"] == "lifespan": + await _drain_lifespan(receive, send) + return + message = await receive() + assert message["type"] == "websocket.connect" + await send({"type": "websocket.accept"}) + raise RuntimeError("websocket app crashed") + + ws_app = WSWatchApp() echo_app = WSEchoApp() empty_frame_app = WSEmptyFrameApp() +app_close_app = WSAppCloseApp() +crash_app = WSCrashApp() class Default(WorkerEntrypoint): async def fetch(self, request): if (request.headers.get("upgrade") or "").lower() == "websocket": path = urlsplit(request.url).path - app = {"/ws-echo": echo_app, "/ws-empty": empty_frame_app}.get(path, ws_app) + app = { + "/ws-echo": echo_app, + "/ws-empty": empty_frame_app, + "/ws-app-close": app_close_app, + "/ws-crash": crash_app, + }.get(path, ws_app) return await asgi.websocket(app, request) import json