From afea9e1844ae293cf9703844645540fe97a8acd4 Mon Sep 17 00:00:00 2001 From: Dominik Picheta Date: Thu, 27 Aug 2026 16:51:59 +0100 Subject: [PATCH] fix: support persistent lifespans in asgi --- packages/runtime-sdk/src/workers/asgi.py | 64 ++++++++++++----- .../src/test_platform_lifecycle.py | 25 ++++++- .../fastapi-tests/src/worker.py | 8 +++ .../workerd-test/asgi/tests/test_asgi.py | 69 +++++++++++++++++++ 4 files changed, 146 insertions(+), 20 deletions(-) diff --git a/packages/runtime-sdk/src/workers/asgi.py b/packages/runtime-sdk/src/workers/asgi.py index 76d2a5fb..1d6c9eec 100644 --- a/packages/runtime-sdk/src/workers/asgi.py +++ b/packages/runtime-sdk/src/workers/asgi.py @@ -2,7 +2,7 @@ from asyncio import Event, Future, Queue, create_task, ensure_future from collections.abc import Awaitable from contextlib import contextmanager -from typing import Any +from typing import Any, ClassVar from urllib.parse import unquote import js @@ -323,8 +323,11 @@ async def run_app(): return response, request_task -async def process_websocket( - app: Any, req: "Request | js.Request", env: Any = None +async def websocket( + app: Any, + req: "Request | js.Request", + env: Any = None, + state: dict[str, Any] | None = None, ) -> js.Response: from js import Response, WebSocketPair @@ -400,7 +403,7 @@ async def ws_receive(): task = run_in_background( app( - request_to_scope(req, env if env is not None else {}, ws=True), + request_to_scope(req, env if env is not None else {}, ws=True, state=state), ws_receive, ws_send, ) @@ -427,19 +430,30 @@ def _close_transport(finished): async def fetch( - app: Any, req: "Request | js.Request", env: Any, ctx: Context | None = None + app: Any, + req: "Request | js.Request", + env: Any, + ctx: Context | None = None, + state: dict[str, Any] | None = None, ) -> js.Response: + if (req.headers.get("upgrade") or "").lower() == "websocket": + return await websocket(app, req, env, state=state) + logger.debug("ASGI request: %s %s", req.method, req.url) - shutdown, state = await start_application(app) + shutdown = None + if state is None: + shutdown, state = await start_application(app) try: result, request_task = await process_request(app, req, env, ctx, state=state) except Exception: logger.exception("ASGI request failed") - await shutdown() + if shutdown is not None: + await shutdown() raise if request_task.done(): - await shutdown() + if shutdown is not None: + await shutdown() else: from workers import wait_until # noqa: PLC0415 @@ -447,27 +461,39 @@ async def finalize_request(): try: await request_task finally: - await shutdown() + if shutdown is not None: + await shutdown() wait_until(run_in_background(finalize_request())) return result -async def websocket( - app: Any, req: "Request | js.Request", env: Any = None -) -> js.Response: - return await process_websocket(app, req, env) +class AsgiWorkerEntrypoint(WorkerEntrypoint): + """Worker entrypoint for an ASGI application.""" + + app: Any + # WorkerEntrypoint instances are invocation-scoped, so this must live on the + # concrete class to preserve lifespan state across requests in the isolate. + _start_future: ClassVar[Future[Any] | None] = None + + async def lifespan_state(self): + start_future = self._start_future + if start_future is None: + start_future = create_task(start_application(self.app)) + type(self)._start_future = start_future + return (await start_future)[1] + + async def fetch(self, request): + state = await self.lifespan_state() + return await fetch(self.app, request, self.env, state=state) -def entrypoint(app: Any) -> type[WorkerEntrypoint]: +def entrypoint(app_: Any) -> type[AsgiWorkerEntrypoint]: """Create the default Worker entrypoint for an ASGI application.""" - class Default(WorkerEntrypoint): - async def fetch(self, request): - if (request.headers.get("upgrade") or "").lower() == "websocket": - return await websocket(app, request, self.env) - return await fetch(app, request, self.env, self.ctx) + class Default(AsgiWorkerEntrypoint): + app = app_ return Default diff --git a/packages/runtime-sdk/tests/web-frameworks-test/fastapi-tests/src/test_platform_lifecycle.py b/packages/runtime-sdk/tests/web-frameworks-test/fastapi-tests/src/test_platform_lifecycle.py index 1a2e1215..5bc51e3f 100644 --- a/packages/runtime-sdk/tests/web-frameworks-test/fastapi-tests/src/test_platform_lifecycle.py +++ b/packages/runtime-sdk/tests/web-frameworks-test/fastapi-tests/src/test_platform_lifecycle.py @@ -3,13 +3,16 @@ import asyncio import pytest -from _client import fetch, get_json +from _client import fetch, get_json, read_json from worker import ( + Default, _platform_events, _platform_shutdown_complete, reset_platform_events, ) +from workers import Request + @pytest.mark.asyncio async def test_lifespan_state_reaches_request(fastapi_app): @@ -19,6 +22,26 @@ async def test_lifespan_state_reaches_request(fastapi_app): assert data == {"available": True, "closed": False} +@pytest.mark.asyncio +async def test_generated_entrypoint_persists_mutable_lifespan_state(): + """Mutable lifespan resources are shared by generated-entrypoint requests.""" + entrypoint = object.__new__(Default) + entrypoint.env = {} + entrypoint.ctx = None + + first = await entrypoint.fetch( + Request("http://testserver/platform/lifespan-mutation") + ) + second = await entrypoint.fetch( + Request("http://testserver/platform/lifespan-mutation") + ) + first_data = await read_json(first) + second_data = await read_json(second) + + assert first_data == {"requests": 1, "closed": False} + assert second_data == {"requests": 2, "closed": False} + + @pytest.mark.asyncio async def test_stream_finishes_before_cleanup_and_shutdown(fastapi_app): """Streaming, background work, and dependency cleanup precede shutdown.""" diff --git a/packages/runtime-sdk/tests/web-frameworks-test/fastapi-tests/src/worker.py b/packages/runtime-sdk/tests/web-frameworks-test/fastapi-tests/src/worker.py index 86fecc5a..802529d2 100644 --- a/packages/runtime-sdk/tests/web-frameworks-test/fastapi-tests/src/worker.py +++ b/packages/runtime-sdk/tests/web-frameworks-test/fastapi-tests/src/worker.py @@ -51,6 +51,7 @@ async def _noop(*args): class PlatformResource: def __init__(self): self.closed = False + self.requests = 0 _platform_events: list[str] = [] @@ -471,6 +472,13 @@ async def platform_lifespan_state(request: Request): return {"available": True, "closed": resource.closed} +@app.get("/platform/lifespan-mutation") +async def platform_lifespan_mutation(request: Request): + resource = request.state.platform_resource + resource.requests += 1 + return {"requests": resource.requests, "closed": resource.closed} + + # --------------------------------------------------------------------------- # # Pyodide and Workers platform boundary routes # --------------------------------------------------------------------------- # diff --git a/packages/runtime-sdk/tests/workerd-test/asgi/tests/test_asgi.py b/packages/runtime-sdk/tests/workerd-test/asgi/tests/test_asgi.py index 28c4479d..624a7c75 100644 --- a/packages/runtime-sdk/tests/workerd-test/asgi/tests/test_asgi.py +++ b/packages/runtime-sdk/tests/workerd-test/asgi/tests/test_asgi.py @@ -347,6 +347,75 @@ async def test_lifespan_full_cycle(): assert lifespan_app.events == ["startup", "shutdown"] +class _PersistentLifespanApp: + def __init__(self): + self.startups = 0 + self.shutdowns = 0 + + async def __call__(self, scope, receive, send): + if scope["type"] == "lifespan": + while True: + message = await receive() + if message["type"] == "lifespan.startup": + self.startups += 1 + scope["state"]["resource"] = {"requests": 0} + # Let a concurrent first request reach entrypoint startup. + await asyncio.sleep(0) + await send({"type": "lifespan.startup.complete"}) + elif message["type"] == "lifespan.shutdown": + self.shutdowns += 1 + await send({"type": "lifespan.shutdown.complete"}) + return + + resource = scope["state"]["resource"] + resource["requests"] += 1 + await receive() + await send( + { + "type": "http.response.start", + "status": 200, + "headers": [(b"content-type", b"application/json")], + } + ) + await send( + { + "type": "http.response.body", + "body": json.dumps(resource).encode(), + } + ) + + +@pytest.mark.asyncio +async def test_generated_entrypoint_reuses_lifespan_state(): + lifespan_app = _PersistentLifespanApp() + Default = asgi.entrypoint(lifespan_app) + + class Custom(asgi.AsgiWorkerEntrypoint): + app = lifespan_app + + async def other_handler(self): + return "other response" + + first_entrypoint = object.__new__(Custom) + first_entrypoint.env = {} + first_entrypoint.ctx = None + second_entrypoint = object.__new__(Custom) + second_entrypoint.env = {} + second_entrypoint.ctx = None + + responses = await asyncio.gather( + first_entrypoint.fetch(js.Request.new("http://example.com/persistent/one")), + second_entrypoint.fetch(js.Request.new("http://example.com/persistent/two")), + ) + data = [json.loads(await response.text()) for response in responses] + + assert issubclass(Default, asgi.AsgiWorkerEntrypoint) + assert await first_entrypoint.other_handler() == "other response" + assert sorted(item["requests"] for item in data) == [1, 2] + assert lifespan_app.startups == 1 + assert lifespan_app.shutdowns == 0 + + class _StartupFailApp: async def __call__(self, scope, receive, send): if scope["type"] == "lifespan":