fix(runtime-sdk): implement WebSocket close semantics and app-task lifetime

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)
 

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

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