fix: support persistent lifespans in asgi

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

@@ -1,8 +1,8 @@
 import logging
-from asyncio import Event, Future, Queue, create_task, ensure_future
+from asyncio import Event, Future, Queue, create_task, ensure_future, shield
 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,47 +430,70 @@ 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
 
         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 shield(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
 

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."""

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
 # --------------------------------------------------------------------------- #

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":