fix(runtime-sdk): ensure finalize_request is alive inside wait_until call

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

@@ -6,8 +6,9 @@
 from urllib.parse import unquote
 
 import js
+from pyodide.ffi import create_proxy
 
-from workers import Context, Request, WorkerEntrypoint
+from workers import Context, Request, WorkerEntrypoint, wait_until
 from workers.utils import _to_js_headers
 
 ASGI = {"spec_version": "2.0", "version": "3.0"}
@@ -450,15 +451,15 @@ async def fetch(
     if request_task.done():
         await shutdown()
     else:
-        from workers import wait_until  # noqa: PLC0415
-
         async def finalize_request():
             try:
                 await request_task
             finally:
                 await shutdown()
 
-        wait_until(run_in_background(finalize_request()))
+        finalizer_task = run_in_background(finalize_request())
+        task_proxy = create_proxy(finalizer_task)
+        wait_until(task_proxy)
 
     return result
 

packages/runtime-sdk/tests/test_in_workerd.py

@@ -15,6 +15,7 @@
 WORKERS_PY = TEST_DIR.parent.parent / "cli"
 WORKERS_RUNTIME_SDK = TEST_DIR.parent / "src"
 DISK_SERVICE_NAME = "TEST_TMPDIR"
+BORROWED_PROXY_ERROR = "This borrowed proxy was automatically destroyed"
 
 
 def discover_workerd_tests():
@@ -62,6 +63,7 @@ def bundle_cache_dir(tmp_path_factory):
 )
 @pytest.mark.parametrize("test_dir, wd_test_file", discover_workerd_tests())
 def test_in_workerd(  # noqa: PLR0913, PLR0917  (too-many-arguments)
+    capfd,
     tmp_path,
     test_dir,
     wd_test_file,
@@ -153,3 +155,10 @@ def test_in_workerd(  # noqa: PLR0913, PLR0917  (too-many-arguments)
         cwd=target,
         check=True,
     )
+
+    # This happens in the background so it is not captured
+    # inside the worker. We need to look at the worker's logs
+    # to see if there are any errors.
+    captured = capfd.readouterr()
+    output = captured.out + captured.err
+    assert BORROWED_PROXY_ERROR not in output

packages/runtime-sdk/tests/workerd-test/asgi/tests/test_asgi.py

@@ -5,7 +5,12 @@
 import js
 import pytest
 from pyodide.ffi import to_js
-from worker import STREAMING_CHUNK_SIZE, STREAMING_NUM_CHUNKS, example_hdr
+from worker import (
+    STREAMING_CHUNK_SIZE,
+    STREAMING_NUM_CHUNKS,
+    delayed_streaming_app,
+    example_hdr,
+)
 
 import asgi
 from workers import Request, env
@@ -89,6 +94,26 @@ async def test_streaming():
         assert all(b == i % 256 for b in chunk)
 
 
+@pytest.mark.asyncio
+async def test_streaming_finalizer_survives_fetch_return():
+    delayed_streaming_app.reset()
+    response = await asyncio.wait_for(
+        env.SELF.fetch("http://example.com/delayed-stream"), timeout=5
+    )
+
+    # The first chunk makes fetch resolve while the ASGI task is still blocked.
+    assert not delayed_streaming_app.shutdown_complete.is_set()
+    reader = response.body.getReader()
+    first = await asyncio.wait_for(reader.read(), timeout=5)
+    assert first.value.to_bytes() == b"first"
+
+    delayed_streaming_app.release_stream.set()
+    second = await asyncio.wait_for(reader.read(), timeout=5)
+    assert second.value.to_bytes() == b"second"
+    assert (await asyncio.wait_for(reader.read(), timeout=5)).done
+    await asyncio.wait_for(delayed_streaming_app.shutdown_complete.wait(), timeout=5)
+
+
 class _ListHandler(logging.Handler):
     """A logging handler that captures records into a list for assertions."""
 

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

@@ -177,6 +177,37 @@ async def __call__(self, scope, receive, send):
             )
 
 
+class DelayedStreamingApp:
+    """Streams one chunk, then waits so fetch returns before finalization."""
+
+    def __init__(self):
+        self.release_stream = asyncio.Event()
+        self.shutdown_complete = asyncio.Event()
+
+    def reset(self):
+        self.release_stream = asyncio.Event()
+        self.shutdown_complete = asyncio.Event()
+
+    async def __call__(self, scope, receive, send):
+        if scope["type"] == "lifespan":
+            while True:
+                message = await receive()
+                if message["type"] == "lifespan.startup":
+                    await send({"type": "lifespan.startup.complete"})
+                elif message["type"] == "lifespan.shutdown":
+                    await send({"type": "lifespan.shutdown.complete"})
+                    self.shutdown_complete.set()
+                    return
+
+        await receive()
+        await send({"type": "http.response.start", "status": 200, "headers": []})
+        await send(
+            {"type": "http.response.body", "body": b"first", "more_body": True}
+        )
+        await self.release_stream.wait()
+        await send({"type": "http.response.body", "body": b"second"})
+
+
 # ---------------------------------------------------------------------------
 # App instances and constants
 # ---------------------------------------------------------------------------
@@ -259,6 +290,7 @@ async def __call__(self, scope, receive, send):
 app = HeaderEchoApp()
 sse_app = SSEApp()
 streaming_app = StreamingApp()
+delayed_streaming_app = DelayedStreamingApp()
 scope_echo_app = ScopeEchoApp()
 late_failure_stream_app = LateFailureStreamApp()
 multi_cookie_app = MultiCookieApp()
@@ -277,6 +309,8 @@ async def fetch(self, request):
             return await asgi.fetch(sse_app, request, self.env, self.ctx)
         elif path == "/stream":
             return await asgi.fetch(streaming_app, request, self.env, self.ctx)
+        elif path == "/delayed-stream":
+            return await asgi.fetch(delayed_streaming_app, request, self.env, self.ctx)
         elif path.startswith("/scope"):
             return await asgi.fetch(scope_echo_app, request, self.env, self.ctx)
         elif path == "/stream-late-failure":