diff --git a/packages/runtime-sdk/src/workers/wsgi.py b/packages/runtime-sdk/src/workers/wsgi.py index b9af4043..c43ed750 100644 --- a/packages/runtime-sdk/src/workers/wsgi.py +++ b/packages/runtime-sdk/src/workers/wsgi.py @@ -1,3 +1,4 @@ +import contextvars import io import logging import sys @@ -212,6 +213,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]]", @@ -321,15 +330,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 diff --git a/packages/runtime-sdk/tests/web-frameworks-test/flask-tests/pyproject.toml b/packages/runtime-sdk/tests/web-frameworks-test/flask-tests/pyproject.toml index 93514e48..979f2b4d 100644 --- a/packages/runtime-sdk/tests/web-frameworks-test/flask-tests/pyproject.toml +++ b/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" } diff --git a/packages/runtime-sdk/tests/web-frameworks-test/flask-tests/src/test_stream_with_context.py b/packages/runtime-sdk/tests/web-frameworks-test/flask-tests/src/test_stream_with_context.py new file mode 100644 index 00000000..4b42cf81 --- /dev/null +++ b/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"] diff --git a/packages/runtime-sdk/tests/workerd-test/wsgi/tests/test_wsgi.py b/packages/runtime-sdk/tests/workerd-test/wsgi/tests/test_wsgi.py index bfcec792..61f14751 100644 --- a/packages/runtime-sdk/tests/workerd-test/wsgi/tests/test_wsgi.py +++ b/packages/runtime-sdk/tests/workerd-test/wsgi/tests/test_wsgi.py @@ -1,11 +1,14 @@ +import asyncio import json import js import pytest from pyodide.ffi import to_js from worker import ( + CONTEXTVAR_CLOSE_RESULTS, STREAMING_CHUNK_SIZE, STREAMING_NUM_CHUNKS, + contextvar_streaming_app, crash_app, example_hdr, header_echo_app, @@ -103,6 +106,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") diff --git a/packages/runtime-sdk/tests/workerd-test/wsgi/worker.py b/packages/runtime-sdk/tests/workerd-test/wsgi/worker.py index 0b227f74..228dbfb8 100644 --- a/packages/runtime-sdk/tests/workerd-test/wsgi/worker.py +++ b/packages/runtime-sdk/tests/workerd-test/wsgi/worker.py @@ -1,5 +1,7 @@ import asyncio +import contextvars import os +from urllib.parse import parse_qs from pyodide.ffi import run_sync from testlib.entrypoint import run_pytest @@ -91,6 +93,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() + + def crash_app(environ, start_response): raise RuntimeError("app crash before response for testing") @@ -113,6 +149,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) return await wsgi.fetch(app, request, self.env)