fix: Make wsgi streaming response handlers async

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

@@ -1,3 +1,4 @@
+import contextvars
 import io
 import logging
 import sys
@@ -250,22 +251,24 @@ def _make_streaming_response(
 
     proxies: list[Any] = []
     done = False
+    ctx = contextvars.copy_context()
 
     def cleanup() -> None:
         nonlocal done
         if done:
             return
         done = True
         try:
-            on_close()
+            ctx.run(on_close)
         finally:
             for proxy in proxies:
                 proxy.destroy()
 
+    # Make proxies async so that it is possible to stack switch inside them.
     @create_proxy
-    def pull(controller: Any) -> None:
+    async def pull(controller: Any) -> None:
         try:
-            chunk = next(chunks, _END)
+            chunk = ctx.run(next, chunks, _END)
         except Exception as exc:  # noqa: BLE001 - forward app errors to the stream
             logger.exception("Exception while streaming WSGI response body")
             cleanup()
@@ -278,7 +281,7 @@ def pull(controller: Any) -> None:
         controller.enqueue(_to_js_uint8array(chunk))
 
     @create_proxy
-    def cancel(_reason: Any = None) -> None:
+    async def cancel(_reason: Any = None) -> None:
         cleanup()
 
     proxies = [pull, cancel]

packages/runtime-sdk/tests/workerd-test/wsgi/pyproject.toml

@@ -2,4 +2,4 @@
 name = "test"
 version = "0.0.0"
 requires-python = ">=3.12"
-dependencies = ["pytest"]
+dependencies = ["pytest", "pytest-asyncio<1.2.0"]

packages/runtime-sdk/tests/workerd-test/wsgi/tests/test_wsgi.py

@@ -0,0 +1,108 @@
+import json
+
+import js
+import pytest
+from pyodide.ffi import to_js
+from worker import (
+    STREAMING_CHUNK_SIZE,
+    STREAMING_NUM_CHUNKS,
+    crash_app,
+    example_hdr,
+)
+
+from workers import Request, env, wsgi
+
+
+@pytest.mark.asyncio
+async def test_headers():
+    response = await env.SELF.fetch("http://example.com/", headers=to_js(example_hdr))
+    assert response.status == 200
+    text = await response.text()
+    assert text == "Hello, World"
+    # Echoed-back headers should be present.
+    assert response.headers.get("header1") == "Value1"
+    assert response.headers.get("header2") == "Value2"
+
+
+@pytest.mark.asyncio
+async def test_echo_body():
+    response = await env.SELF.fetch(
+        "http://example.com/echo-body",
+        method="POST",
+        body="hello body",
+    )
+    assert response.status == 200
+    text = await response.text()
+    assert text == "hello body"
+
+
+@pytest.mark.asyncio
+async def test_meta():
+    response = await env.SELF.fetch("http://example.com/meta?foo=bar&baz=qux")
+    assert response.status == 200
+
+    payload = json.loads(await response.text())
+    assert payload["method"] == "GET"
+    assert payload["path"] == "/meta"
+    assert payload["query"] == "foo=bar&baz=qux"
+    assert payload["scheme"] == "http"
+    assert payload["has_env"] is True
+
+
+@pytest.mark.asyncio
+async def test_cookies():
+    response = await env.SELF.fetch("http://example.com/cookies")
+    assert response.status == 200
+    # `env.SELF.fetch` returns the SDK `FetchResponse`, whose `.headers` is an
+    # `http.client.HTTPMessage`. Repeated Set-Cookie headers are preserved as
+    # separate entries (see `python_request_headers_preserve_commas`), so use
+    # `get_all` to recover the individual values.
+    cookies = response.headers.get_all("Set-Cookie")
+    assert "a=1" in cookies
+    assert "b=2" in cookies
+
+
+@pytest.mark.parametrize(
+    "endpoint", ("stream", "stream-stack-switch", "stream-context")
+)
+@pytest.mark.asyncio
+async def test_streaming(endpoint):
+    response = await env.SELF.fetch("http://example.com/" + endpoint)
+    assert response.status == 200
+    assert response.headers.get("content-type") == "application/octet-stream"
+
+    reader = response.body.getReader()
+    body_bytes = b""
+    while True:
+        result = await reader.read()
+        if result.done:
+            break
+        body_bytes += result.value.to_bytes()
+
+    expected_size = STREAMING_CHUNK_SIZE * STREAMING_NUM_CHUNKS
+    assert len(body_bytes) == expected_size, (
+        f"Expected {expected_size} bytes, got {len(body_bytes)}"
+    )
+    for i in range(STREAMING_NUM_CHUNKS):
+        start = i * STREAMING_CHUNK_SIZE
+        end = start + STREAMING_CHUNK_SIZE
+        expected_byte = i % 256
+        assert all(b == expected_byte for b in body_bytes[start:end])
+
+
+@pytest.mark.asyncio
+async def test_app_exception_is_raised():
+    req = js.Request.new("http://example.com/crash-test")
+    with pytest.raises(RuntimeError, match="app crash before response for testing"):
+        await wsgi.fetch(crash_app, req, env)
+
+
+def test_build_environ_handles_js_and_python_requests():
+    # Verify `build_environ` handles JS-style and Python-style headers
+    # identically, mirroring the asgi `request_to_scope` check.
+    js_request = js.Request.new("http://example.com/", headers=to_js(example_hdr))
+    py_request = Request("http://example.com/", headers=example_hdr)
+    js_env = wsgi.build_environ(js_request, env, b"")
+    py_env = wsgi.build_environ(py_request, env, b"")
+    assert js_env["HTTP_HEADER1"] == py_env["HTTP_HEADER1"] == "Value1"
+    assert js_env["HTTP_HEADER2"] == py_env["HTTP_HEADER2"] == "Value2"

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

@@ -1,7 +1,29 @@
-import js
-from pyodide.ffi import to_js
+import asyncio
+import contextvars
+import os
+import sys
+
+import pytest
+from pyodide.webloop import WebLoop
+from pyodide.ffi import run_sync
+
+from workers import WorkerEntrypoint, wsgi
+
+
+async def noop(*args):
+    pass
+
+
+# pytest-asyncio relies on these but in Pyodide < 0.29 WebLoop does not implement them
+WebLoop.shutdown_asyncgens = noop
+WebLoop.shutdown_default_executor = noop
+
+# Pyodide 0.26.0a2's _cancel_all_tasks calls task.exception() on pending tasks,
+# which raises InvalidStateError under Pyodide's WebLoop.
+# Ignore this error to prevent pytest-asyncio from crashing.
+if sys.version_info < (3, 13):
+    asyncio.runners._cancel_all_tasks = lambda loop: None  # type: ignore[attr-defined]
 
-from workers import Request, WorkerEntrypoint, wsgi
 
 # ---------------------------------------------------------------------------
 # WSGI apps
@@ -76,130 +98,73 @@ def generate():
     return generate()
 
 
-def crash_app(environ, start_response):
-    raise RuntimeError("app crash before response for testing")
+def streaming_app_stack_switch(environ, start_response):
+    """WSGI app that returns multiple body chunks via a generator."""
+    start_response("200 OK", [("Content-Type", "application/octet-stream")])
 
+    def generate():
+        for i in range(STREAMING_NUM_CHUNKS):
+            run_sync(asyncio.sleep(0))
+            yield bytes([i % 256]) * STREAMING_CHUNK_SIZE
 
-example_hdr = {"Header1": "Value1", "Header2": "Value2"}
+    return generate()
 
+STREAMING_CONTEXT_VAR = contextvars.ContextVar("streaming_counter")
 
-class Default(WorkerEntrypoint):
-    async def fetch(self, request):
-        from js import URL
 
-        url = URL.new(request.url)
-        path = url.pathname
+def streaming_app_uses_context(environ, start_response):
+    """WSGI app whose body generator reads and writes a ContextVar.
 
-        if path == "/echo-body":
-            return await wsgi.fetch(echo_body_app, request, self.env)
-        elif path == "/meta":
-            return await wsgi.fetch(echo_meta_app, request, self.env)
-        elif path == "/cookies":
-            return await wsgi.fetch(cookies_app, request, self.env)
-        elif path == "/stream":
-            return await wsgi.fetch(streaming_app, request, self.env)
-
-        # Verify `build_environ` handles JS-style and Python-style headers
-        # identically, mirroring the asgi `request_to_scope` check.
-        js_request = js.Request.new("http://example.com/", headers=to_js(example_hdr))
-        py_request = Request("http://example.com/", headers=example_hdr)
-        js_env = wsgi.build_environ(js_request, self.env, b"")
-        py_env = wsgi.build_environ(py_request, self.env, b"")
-        assert js_env["HTTP_HEADER1"] == py_env["HTTP_HEADER1"] == "Value1"
-        assert js_env["HTTP_HEADER2"] == py_env["HTTP_HEADER2"] == "Value2"
-
-        return await wsgi.fetch(header_echo_app, request, self.env)
+    The generator is resumed from the `ReadableStream` pull callback, which runs
+    in a fresh context, so this only works if the server carries the request's
+    `contextvars.Context` into every pull. If the context is lost, `get()`
+    raises `LookupError` and the stream errors out; if a *fresh copy* is used
+    per pull, the mutations don't stick and every chunk repeats the same byte.
+    """
+    start_response("200 OK", [("Content-Type", "application/octet-stream")])
+    STREAMING_CONTEXT_VAR.set(0)
 
-    async def test(self, ctrl):
-        await test_headers(self.env)
-        await test_echo_body(self.env)
-        await test_meta(self.env)
-        await test_cookies(self.env)
-        await test_streaming(self.env)
-        await test_app_exception_is_raised(self.env)
+    def generate():
+        for _ in range(STREAMING_NUM_CHUNKS):
+            # Stack switch so each chunk is pulled from a separate callback.
+            run_sync(asyncio.sleep(0))
+            counter = STREAMING_CONTEXT_VAR.get()
+            STREAMING_CONTEXT_VAR.set(counter + 1)
+            yield bytes([counter % 256]) * STREAMING_CHUNK_SIZE
 
+    return generate()
 
-# ---------------------------------------------------------------------------
-# Tests
-# ---------------------------------------------------------------------------
 
+def crash_app(environ, start_response):
+    raise RuntimeError("app crash before response for testing")
 
-async def test_headers(env):
-    response = await env.SELF.fetch("http://example.com/", headers=to_js(example_hdr))
-    assert response.status == 200
-    text = await response.text()
-    assert text == "Hello, World"
-    # Echoed-back headers should be present.
-    assert response.headers.get("header1") == "Value1"
-    assert response.headers.get("header2") == "Value2"
 
+example_hdr = {"Header1": "Value1", "Header2": "Value2"}
 
-async def test_echo_body(env):
-    response = await env.SELF.fetch(
-        "http://example.com/echo-body",
-        method="POST",
-        body="hello body",
-    )
-    assert response.status == 200
-    text = await response.text()
-    assert text == "hello body"
 
+class Default(WorkerEntrypoint):
+    # Each path in this handler serves one of the WSGI apps above; the
+    # assertions live in tests/test_wsgi.py.
+    async def fetch(self, request):
+        from js import URL
 
-async def test_meta(env):
-    response = await env.SELF.fetch("http://example.com/meta?foo=bar&baz=qux")
-    assert response.status == 200
-    import json
+        url = URL.new(request.url)
+        path = url.pathname
 
-    payload = json.loads(await response.text())
-    assert payload["method"] == "GET"
-    assert payload["path"] == "/meta"
-    assert payload["query"] == "foo=bar&baz=qux"
-    assert payload["scheme"] == "http"
-    assert payload["has_env"] is True
-
-
-async def test_cookies(env):
-    response = await env.SELF.fetch("http://example.com/cookies")
-    assert response.status == 200
-    # `env.SELF.fetch` returns the SDK `FetchResponse`, whose `.headers` is an
-    # `http.client.HTTPMessage`. Repeated Set-Cookie headers are preserved as
-    # separate entries (see `python_request_headers_preserve_commas`), so use
-    # `get_all` to recover the individual values.
-    cookies = response.headers.get_all("Set-Cookie")
-    assert "a=1" in cookies
-    assert "b=2" in cookies
-
-
-async def test_streaming(env):
-    response = await env.SELF.fetch("http://example.com/stream")
-    assert response.status == 200
-    assert response.headers.get("content-type") == "application/octet-stream"
-
-    reader = response.body.getReader()
-    body_bytes = b""
-    while True:
-        result = await reader.read()
-        if result.done:
-            break
-        body_bytes += result.value.to_bytes()
-
-    expected_size = STREAMING_CHUNK_SIZE * STREAMING_NUM_CHUNKS
-    assert len(body_bytes) == expected_size, (
-        f"Expected {expected_size} bytes, got {len(body_bytes)}"
-    )
-    for i in range(STREAMING_NUM_CHUNKS):
-        start = i * STREAMING_CHUNK_SIZE
-        end = start + STREAMING_CHUNK_SIZE
-        expected_byte = i % 256
-        assert all(b == expected_byte for b in body_bytes[start:end])
-
-
-async def test_app_exception_is_raised(env):
-    req = js.Request.new("http://example.com/crash-test")
-    threw = False
-    try:
-        await wsgi.fetch(crash_app, req, env)
-    except RuntimeError as e:
-        threw = True
-        assert "app crash before response for testing" in str(e)
-    assert threw, "Expected RuntimeError to be raised from wsgi.fetch"
+        app = {
+            "/echo-body" : echo_body_app,
+            "/meta": echo_meta_app,
+            "/cookies": cookies_app,
+            "/stream": streaming_app,
+            "/stream-stack-switch": streaming_app_stack_switch,
+            "/stream-context": streaming_app_uses_context,
+        }.get(path, header_echo_app)
+
+        return await wsgi.fetch(app, request, self.env)
+
+    async def test(self, ctrl):
+        os.chdir("/session/metadata/tests")
+        args = [".", "-vv"]
+        if self.env.color:
+            args.append("--color=yes")
+        assert pytest.main(args) == 0

packages/runtime-sdk/tests/workerd-test/wsgi/wsgi.wd-test

@@ -10,6 +10,10 @@ const unitTests :Workerd.Config = (
         ],
         bindings = [
           ( name = "SELF", service = "python-wsgi" ),
+          (
+            name = "color",
+            json = "%COLOR"
+          ),
         ],
         compatibilityDate = "%COMPAT_DATE",
         compatibilityFlags = ["python_workers"],