fix(runtime-sdk): fix contextvars not being shared in WSGI streaming response

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

@@ -1,3 +1,4 @@
+import contextvars
 import io
 import logging
 import sys
@@ -222,6 +223,14 @@ def _close_iterable(iterable: Any) -> None:
 _END = object()
 
 
+def _iter_in_context(
+    context: contextvars.Context, iterator: "Iterator[Any]"
+) -> "Iterator[Any]":
+    """Make sure every chunk is produced in the given context."""
+    while (chunk := context.run(next, iterator, _END)) is not _END:
+        yield chunk
+
+
 def _make_streaming_response(
     status: str,
     headers: "list[tuple[str, str]]",
@@ -333,15 +342,21 @@ def start_response(status, response_headers, exc_info=None):
         response_state["headers"] = response_headers
         return write
 
-    result = app(environ, start_response)
-    result_iter = iter(result)
+    # Make sure a single context is used per response.
+    # Otherwise, stack switching could cause context variables set by the app
+    # to be lost when the iterable is consumed.
+    context = contextvars.copy_context()
+    result = context.run(app, environ, start_response)
+    result_iter = _iter_in_context(context, context.run(iter, result))
 
     def close_all() -> None:
-        _close_iterable(result)
         try:
-            environ["wsgi.input"].close()
-        except Exception:  # noqa: BLE001 - best-effort cleanup
-            logger.exception("Failed to close wsgi.input")
+            context.run(_close_iterable, result)
+        finally:
+            try:
+                environ["wsgi.input"].close()
+            except Exception:  # noqa: BLE001 - best-effort cleanup
+                logger.exception("Failed to close wsgi.input")
 
     # WSGI apps must call start_response before yielding the first body chunk,
     # but some defer it until the first non-empty chunk is produced. Pull that

packages/runtime-sdk/tests/web-frameworks-test/flask-tests/pyproject.toml

@@ -2,7 +2,13 @@
 name = "flask-tests"
 version = "0.1.0"
 requires-python = ">=3.12"
-dependencies = ["flask", "pytest", "testlib", "workers-runtime-sdk"]
+dependencies = [
+    "flask",
+    "pytest",
+    "pytest-asyncio<1.2.0",
+    "testlib",
+    "workers-runtime-sdk",
+]
 
 [tool.uv.sources]
 testlib = { path = "../packages/testlib" }

packages/runtime-sdk/tests/web-frameworks-test/flask-tests/src/test_stream_with_context.py

@@ -0,0 +1,70 @@
+import asyncio
+
+import js
+import pytest
+from flask import (
+    Flask,
+    Response,
+    has_request_context,
+    request,
+    stream_with_context,
+)
+
+from workers import wsgi
+
+app = Flask(__name__)
+# Request values whose stream_with_context generator has run its cleanup.
+closed: list[str] = []
+
+
+@app.get("/stream")
+def stream():
+    @stream_with_context
+    def body():
+        value = request.args["value"]
+        try:
+            yield "flask:"
+            yield request.args["value"]
+        finally:
+            closed.append(value)
+
+    return Response(body(), content_type="text/plain")
+
+
+async def fetch(value):
+    req = js.Request.new(f"http://example.com/stream?value={value}")
+    return await wsgi.fetch(app, req, {})
+
+
+@pytest.fixture(autouse=True)
+def clear_closed():
+    closed.clear()
+
+
+@pytest.mark.asyncio
+async def test_stream_with_context():
+    response = await fetch("stream-body")
+    assert response.status == 200
+    # The request context pushed while priming the body must stay in the
+    # response's own context instead of leaking into the caller's.
+    assert not has_request_context()
+    assert await response.text() == "flask:stream-body"
+    assert closed == ["stream-body"]
+
+
+@pytest.mark.asyncio
+async def test_concurrent_stream_with_context():
+    responses = await asyncio.gather(fetch("first"), fetch("second"))
+    bodies = await asyncio.gather(*(response.text() for response in responses))
+    assert bodies == ["flask:first", "flask:second"]
+    assert sorted(closed) == ["first", "second"]
+
+
+@pytest.mark.asyncio
+async def test_cancelled_stream_with_context():
+    response = await fetch("cancel")
+    reader = response.body.getReader()
+    first = await reader.read()
+    assert first.value.to_bytes() == b"flask:"
+    await reader.cancel()
+    assert closed == ["cancel"]

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

@@ -1,13 +1,16 @@
+import asyncio
 import gzip
 import json
 
 import js
 import pytest
 from pyodide.ffi import to_js
 from worker import (
+    CONTEXTVAR_CLOSE_RESULTS,
     ENCODED_RESPONSE_PAYLOAD,
     STREAMING_CHUNK_SIZE,
     STREAMING_NUM_CHUNKS,
+    contextvar_streaming_app,
     crash_app,
     example_hdr,
     header_echo_app,
@@ -105,6 +108,55 @@ async def test_streaming(endpoint):
         assert all(b == expected_byte for b in body_bytes[start:end])
 
 
+async def read_all(response) -> bytes:
+    reader = response.body.getReader()
+    body = b""
+    while True:
+        result = await reader.read()
+        if result.done:
+            return body
+        body += result.value.to_bytes()
+
+
+@pytest.mark.asyncio
+async def test_streaming_preserves_contextvars():
+    response = await env.SELF.fetch(
+        "http://example.com/stream-contextvar?value=complete"
+    )
+    assert response.status == 200
+    assert await read_all(response) == b"value:" + b"complete" * STREAMING_NUM_CHUNKS
+    assert CONTEXTVAR_CLOSE_RESULTS.pop("complete") == ["ok"]
+
+
+@pytest.mark.asyncio
+async def test_concurrent_streams_keep_separate_contextvars():
+    responses = await asyncio.gather(
+        env.SELF.fetch("http://example.com/stream-contextvar?value=first"),
+        env.SELF.fetch("http://example.com/stream-contextvar?value=second"),
+    )
+    bodies = await asyncio.gather(*(read_all(response) for response in responses))
+
+    assert bodies == [
+        b"value:" + b"first" * STREAMING_NUM_CHUNKS,
+        b"value:" + b"second" * STREAMING_NUM_CHUNKS,
+    ]
+    assert CONTEXTVAR_CLOSE_RESULTS.pop("first") == ["ok"]
+    assert CONTEXTVAR_CLOSE_RESULTS.pop("second") == ["ok"]
+
+
+@pytest.mark.asyncio
+async def test_cancelled_stream_closes_in_app_context():
+    req = js.Request.new("http://example.com/stream-contextvar?value=cancel")
+    response = await wsgi.fetch(contextvar_streaming_app, req, env)
+
+    reader = response.body.getReader()
+    first = await reader.read()
+    assert first.value.to_bytes() == b"value:"
+    await reader.cancel()
+
+    assert CONTEXTVAR_CLOSE_RESULTS.pop("cancel") == ["ok"]
+
+
 @pytest.mark.asyncio
 async def test_app_exception_is_raised():
     req = js.Request.new("http://example.com/crash-test")

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

@@ -1,6 +1,8 @@
 import asyncio
+import contextvars
 import gzip
 import os
+from urllib.parse import parse_qs
 
 from pyodide.ffi import run_sync
 from testlib.entrypoint import run_pytest
@@ -92,6 +94,40 @@ def generate():
     return generate()
 
 
+REQUEST_VALUE: contextvars.ContextVar[str] = contextvars.ContextVar("request_value")
+# Outcome of resetting REQUEST_VALUE when each response generator closes, keyed
+# by request value: "ok", or the repr of the exception raised by reset().
+CONTEXTVAR_CLOSE_RESULTS: dict[str, list[str]] = {}
+
+
+def contextvar_streaming_app(environ, start_response):
+    """WSGI app whose body generator depends on a ContextVar set by the app.
+
+    Mirrors what Flask's ``stream_with_context`` does: state set in a ContextVar
+    while the app runs must stay visible while the body is iterated, and the
+    token must be reset in the same context when the iterable is closed.
+    """
+    value = parse_qs(environ["QUERY_STRING"])["value"][0]
+    token = REQUEST_VALUE.set(value)
+    start_response("200 OK", [("Content-Type", "text/plain")])
+
+    def generate():
+        try:
+            yield b"value:"
+            for _ in range(STREAMING_NUM_CHUNKS):
+                run_sync(asyncio.sleep(0))
+                yield REQUEST_VALUE.get().encode()
+        finally:
+            try:
+                REQUEST_VALUE.reset(token)
+                result = "ok"
+            except Exception as exc:  # noqa: BLE001 - recorded for the test
+                result = repr(exc)
+            CONTEXTVAR_CLOSE_RESULTS.setdefault(value, []).append(result)
+
+    return generate()
+
+
 ENCODED_RESPONSE_PAYLOAD = b'{"message":"already encoded by the app"}' * 64
 
 
@@ -139,6 +175,7 @@ async def fetch(self, request):
             "/cookies": cookies_app,
             "/stream": streaming_app,
             "/stream-stack-switch": streaming_app_stack_switch,
+            "/stream-contextvar": contextvar_streaming_app,
         }.get(path, header_echo_app)
 
         if path.startswith("/encoded/"):