fix: handle ASGI WebSocket close events and environment propagation

packages/runtime-sdk/src/workers/asgi.py

@@ -311,7 +311,9 @@ async def run_app():
         app_task.destroy()
 
 
-async def process_websocket(app: Any, req: "Request | js.Request") -> js.Response:
+async def process_websocket(
+    app: Any, req: "Request | js.Request", env: Any = None
+) -> js.Response:
     from js import Response, WebSocketPair
 
     client, server = WebSocketPair.new().object_values()
@@ -352,6 +354,10 @@ def onmessage(evt):
     server.onmessage = onmessage
 
     async def ws_send(got):
+        if got["type"] == "websocket.accept":
+            # The Workers WebSocketPair is accepted before the upgrade response
+            # is returned, so there is no deferred accept operation here.
+            return
         if got["type"] == "websocket.send":
             b = got.get("bytes", None)
             s = got.get("text", None)
@@ -362,16 +368,23 @@ async def ws_send(got):
                     server.send(jsbytes)
             if s is not None:
                 server.send(s)
-
-        else:
-            logger.warning(" == Not implemented %s", got["type"])
+            return
+        if got["type"] == "websocket.close":
+            server.close(got.get("code", 1000), got.get("reason", ""))
+            return
+        logger.warning(" == Not implemented %s", got["type"])
 
     async def ws_receive():
         received = await queue.get()
         return received
 
-    env = {}
-    run_in_background(app(request_to_scope(req, env, ws=True), ws_receive, ws_send))
+    run_in_background(
+        app(
+            request_to_scope(req, env if env is not None else {}, ws=True),
+            ws_receive,
+            ws_send,
+        )
+    )
 
     return Response.new(None, status=101, webSocket=client)
 
@@ -390,8 +403,10 @@ async def fetch(
     return result
 
 
-async def websocket(app: Any, req: "Request | js.Request") -> js.Response:
-    return await process_websocket(app, req)
+async def websocket(
+    app: Any, req: "Request | js.Request", env: Any = None
+) -> js.Response:
+    return await process_websocket(app, req, env)
 
 
 def __getattr__(name):

packages/runtime-sdk/tests/workerd-test/asgi-ws-disconnect/tests/test_ws_disconnect.py

@@ -140,3 +140,22 @@ 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_application_close_reaches_client_with_code_and_reason():
+    async with _ws_session("/ws-close") as ws:
+        closed = _listen(ws, "close")
+        ws.send("close")
+        event = await asyncio.wait_for(closed, TIMEOUT_S)
+        assert event.code == 4001
+        assert event.reason == "application-close"
+
+
+@pytest.mark.asyncio
+async def test_environment_reaches_websocket_scope():
+    async with _ws_session("/ws-env") as ws:
+        message = _listen(ws, "message")
+        ws.send("env")
+        event = await asyncio.wait_for(message, TIMEOUT_S)
+        assert event.data == "worker-environment"

packages/runtime-sdk/tests/workerd-test/asgi-ws-disconnect/worker.py

@@ -101,16 +101,58 @@ async def __call__(self, scope, receive, send):
             await send({"type": "websocket.send", "text": "done"})
 
 
+class WSCloseApp:
+    """Closes the connection with an application-provided code and reason."""
+
+    async def __call__(self, scope, receive, send):
+        message = await receive()
+        assert message["type"] == "websocket.connect"
+        await send({"type": "websocket.accept"})
+        message = await receive()
+        if message["type"] == "websocket.disconnect":
+            return
+        await send(
+            {
+                "type": "websocket.close",
+                "code": 4001,
+                "reason": "application-close",
+            }
+        )
+
+
+class WSEnvApp:
+    """Sends a value supplied through the ASGI WebSocket scope environment."""
+
+    async def __call__(self, scope, receive, send):
+        message = await receive()
+        assert message["type"] == "websocket.connect"
+        await send({"type": "websocket.accept"})
+        message = await receive()
+        if message["type"] == "websocket.disconnect":
+            return
+        await send({"type": "websocket.send", "text": scope["env"]["marker"]})
+
+
 ws_app = WSWatchApp()
 echo_app = WSEchoApp()
 empty_frame_app = WSEmptyFrameApp()
+close_app = WSCloseApp()
+env_app = WSEnvApp()
 
 
 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)
+            if path == "/ws-env":
+                return await asgi.websocket(
+                    env_app, request, {"marker": "worker-environment"}
+                )
+            app = {
+                "/ws-close": close_app,
+                "/ws-echo": echo_app,
+                "/ws-empty": empty_frame_app,
+            }.get(path, ws_app)
             return await asgi.websocket(app, request)
         import json