diff --git a/packages/runtime-sdk/src/workers/asgi.py b/packages/runtime-sdk/src/workers/asgi.py index 463d5e62..eb0334a1 100644 --- a/packages/runtime-sdk/src/workers/asgi.py +++ b/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): 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 02112b6d..7cb035bf 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,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" 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 4edf52b0..5130d71d 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,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