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