fix(runtime-sdk): add options argument to asgi and wsgi

packages/runtime-sdk/AGENTS.md

@@ -16,6 +16,7 @@ This package provides utilities for:
 - `_pyodide_entrypoint_helper` and `_cloudflare_compat_flags` are internal values that are imported from `workerd`.
   - `workerd` (https://github.com/cloudflare/workerd) is the runtime that Cloudflare Workers use to run Python code.
 - When using a feature from `workerd`, always check if we can implement it in this package instead to avoid tight coupling with the runtime.
+- Code shared by the ASGI (`workers/asgi.py`) and WSGI (`workers/wsgi.py`) adapters lives in the `workers/gateway/` directory.
 
 ## Testing
 

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

@@ -2,12 +2,14 @@
 from asyncio import Event, Future, Queue, create_task, ensure_future
 from collections.abc import Awaitable
 from contextlib import contextmanager
+from dataclasses import dataclass
 from typing import Any
 from urllib.parse import unquote
 
 import js
 
 from workers import Context, Request, WorkerEntrypoint
+from workers.gateway import BaseOptions
 from workers.utils import _to_js_headers
 
 ASGI = {"spec_version": "2.0", "version": "3.0"}
@@ -16,6 +18,14 @@
 background_tasks = set()
 
 
+@dataclass(frozen=True, kw_only=True)
+class Options(BaseOptions):
+    """Options for the ASGI adapter"""
+
+
+DEFAULT_OPTIONS = Options()
+
+
 def run_in_background(coro: Awaitable[Any]) -> Future:
     fut = ensure_future(coro)
     background_tasks.add(fut)
@@ -195,10 +205,12 @@ async def process_request(  # noqa: PLR0913
     # TODO(later): remove this parameter after unvendoring Python SDK from workerd
     ctx: Context | None,
     state: dict[str, Any] | None = None,
+    options: Options = DEFAULT_OPTIONS,
 ) -> tuple[js.Response, Future[None]]:
     from js import Response, TransformStream
     from pyodide.ffi import create_proxy
 
+    response_init = options.response_init()
     status = None
     headers = None
     result = Future()
@@ -262,7 +274,10 @@ async def send(got):
                 readable = transform_stream.readable
                 writer = transform_stream.writable.getWriter()
                 resp = Response.new(
-                    readable, headers=_to_js_headers(headers), status=status
+                    readable,
+                    headers=_to_js_headers(headers),
+                    status=status,
+                    **response_init,
                 )
                 result.set_result(resp)
                 with acquire_js_buffer(body) as jsbytes:
@@ -271,7 +286,10 @@ async def send(got):
                 # 101/103/204/205/304 must not carry a body per the Fetch spec.
                 # https://fetch.spec.whatwg.org/#null-body-status
                 resp = Response.new(
-                    None, headers=_to_js_headers(headers), status=status
+                    None,
+                    headers=_to_js_headers(headers),
+                    status=status,
+                    **response_init,
                 )
                 result.set_result(resp)
                 finished_response.set()
@@ -281,7 +299,10 @@ async def send(got):
                 buf = px.getBuffer()
                 px.destroy()
                 resp = Response.new(
-                    buf.data, headers=_to_js_headers(headers), status=status
+                    buf.data,
+                    headers=_to_js_headers(headers),
+                    status=status,
+                    **response_init,
                 )
                 result.set_result(resp)
                 finished_response.set()
@@ -436,12 +457,19 @@ 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,
+    *,
+    options: Options | None = None,
 ) -> js.Response:
     logger.debug("ASGI request: %s %s", req.method, req.url)
     shutdown, state = await start_application(app)
     try:
-        result, request_task = await process_request(app, req, env, ctx, state=state)
+        result, request_task = await process_request(
+            app, req, env, ctx, state=state, options=options or DEFAULT_OPTIONS
+        )
     except Exception:
         logger.exception("ASGI request failed")
         await shutdown()
@@ -469,14 +497,14 @@ async def websocket(
     return await process_websocket(app, req, env)
 
 
-def entrypoint(app: Any) -> type[WorkerEntrypoint]:
+def entrypoint(app: Any, *, options: Options | None = None) -> type[WorkerEntrypoint]:
     """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)
+            return await fetch(app, request, self.env, self.ctx, options=options)
 
     return Default
 

packages/runtime-sdk/src/workers/gateway/__init__.py

@@ -0,0 +1,5 @@
+"""Shared building blocks for the ASGI and WSGI adapters."""
+
+from .options import BaseOptions, EncodeBody
+
+__all__ = ["BaseOptions", "EncodeBody"]

packages/runtime-sdk/src/workers/gateway/options.py

@@ -0,0 +1,30 @@
+from dataclasses import dataclass
+from typing import Any, Literal, get_args
+
+EncodeBody = Literal["automatic", "manual"]
+
+
+@dataclass(frozen=True, kw_only=True)
+class BaseOptions:
+    """Options shared by the ASGI and WSGI adapters.
+
+    Downstream adapters should inherit from this class to define their own options.
+    """
+
+    # Whether to encode the body according to the ``Content-Encoding`` header
+    # See https://developers.cloudflare.com/workers/runtime-apis/response/#the-encodebody-option
+    encode_body: EncodeBody = "automatic"
+
+    def __post_init__(self) -> None:
+        if self.encode_body not in get_args(EncodeBody):
+            raise ValueError(
+                f"encode_body must be one of {get_args(EncodeBody)}, "
+                f"got {self.encode_body!r}"
+            )
+
+    def response_init(self) -> dict[str, Any]:
+        """Extra keyword arguments for ``Response.new``."""
+        init: dict[str, Any] = {}
+        if self.encode_body != "automatic":
+            init["encodeBody"] = self.encode_body
+        return init

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

@@ -2,17 +2,27 @@
 import logging
 import sys
 from collections.abc import Callable, Iterator
+from dataclasses import dataclass
 from typing import Any
 from urllib.parse import unquote, urlsplit
 
 import js
 
 from workers import Context, Request, WorkerEntrypoint
+from workers.gateway import BaseOptions
 
 logger = logging.getLogger("wsgi")
 NULL_BODY_STATUSES = frozenset({101, 103, 204, 205, 304})
 
 
+@dataclass(frozen=True, kw_only=True)
+class Options(BaseOptions):
+    """Options for the WSGI adapter, see ``BaseOptions`` for shared fields."""
+
+
+DEFAULT_OPTIONS = Options()
+
+
 def _wsgi_native_string(value: str) -> str:
     """Convert a (possibly non-ASCII) Python ``str`` into a WSGI "native string".
 
@@ -217,6 +227,7 @@ def _make_streaming_response(
     headers: "list[tuple[str, str]]",
     chunks: "Iterator[bytes]",
     on_close: "Callable[[], None]",
+    response_init: dict[str, Any],
 ) -> js.Response:
     """Build a ``js.Response`` whose body is pulled lazily from *chunks*.
 
@@ -232,7 +243,7 @@ def _make_streaming_response(
     # WSGI status is e.g. "200 OK"; split into code + reason phrase.
     code_str, _, reason = status.partition(" ")
     code = int(code_str)
-    options: dict[str, Any] = {"status": code}
+    options: dict[str, Any] = {"status": code, **response_init}
     if reason:
         options["statusText"] = reason
 
@@ -291,6 +302,7 @@ def process_request(
     req: "Request | js.Request",
     env: Any,
     body: "bytes | io.IOBase",
+    options: Options = DEFAULT_OPTIONS,
 ) -> js.Response:
     environ = build_environ(req, env, body)
 
@@ -357,7 +369,11 @@ def body_chunks() -> "Iterator[bytes]":
                 yield chunk
 
     return _make_streaming_response(
-        response_state["status"], response_state["headers"], body_chunks(), close_all
+        response_state["status"],
+        response_state["headers"],
+        body_chunks(),
+        close_all,
+        options.response_init(),
     )
 
 
@@ -367,6 +383,8 @@ async def fetch(
     env: Any,
     # Accepted for parity with asgi.fetch; WSGI has no use for it.
     ctx: Context | None = None,
+    *,
+    options: Options | None = None,
 ) -> js.Response:
     logger.debug("WSGI request: %s %s", req.method, req.url)
     # Prefer lazily streaming the body through `wsgi.input` (no full buffering);
@@ -375,17 +393,17 @@ async def fetch(
     if body is None:
         body = await _read_body(req)
     try:
-        return process_request(app, req, env, body)
+        return process_request(app, req, env, body, options or DEFAULT_OPTIONS)
     except Exception:
         logger.exception("WSGI request failed")
         raise
 
 
-def entrypoint(app: Any) -> type[WorkerEntrypoint]:
+def entrypoint(app: Any, *, options: Options | None = None) -> type[WorkerEntrypoint]:
     """Create the default Worker entrypoint for a WSGI application."""
 
     class Default(WorkerEntrypoint):
         async def fetch(self, request):
-            return await fetch(app, request, self.env)
+            return await fetch(app, request, self.env, options=options)
 
     return Default

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

@@ -10,10 +10,27 @@ const unitTests :Workerd.Config = (
         ],
         bindings = [
           ( name = "SELF", service = "python-asgi" ),
+          # SELF.fetch does not serialize the response over HTTP, so it cannot
+          # observe the runtime encoding the body. ASGI_HTTP goes through the
+          # loopback HTTP socket below instead.
+          ( name = "ASGI_HTTP", service = "asgi-http" ),
         ],
         compatibilityDate = "%COMPAT_DATE",
         compatibilityFlags = ["python_workers"],
       )
-    )
+    ),
+    ( name = "asgi-http",
+      external = (
+        address = "loopback:asgi-http",
+        http = (),
+      )
+    ),
+  ],
+  sockets = [
+    ( name = "asgi-http",
+      address = "loopback:asgi-http",
+      http = (),
+      service = "python-asgi",
+    ),
   ],
 );

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

@@ -1,11 +1,17 @@
 import asyncio
+import gzip
 import json
 import logging
 
 import js
 import pytest
 from pyodide.ffi import to_js
-from worker import STREAMING_CHUNK_SIZE, STREAMING_NUM_CHUNKS, example_hdr
+from worker import (
+    ENCODED_RESPONSE_PAYLOAD,
+    STREAMING_CHUNK_SIZE,
+    STREAMING_NUM_CHUNKS,
+    example_hdr,
+)
 
 import asgi
 from workers import Request, env
@@ -474,3 +480,35 @@ async def test_lifespan_preack_crash_is_logged_and_treated_as_unsupported():
         )
     finally:
         _remove_handler(handler)
+
+
+def test_options_validate_encode_body():
+    assert asgi.Options().response_init() == {}
+    assert asgi.Options(encode_body="manual").response_init() == {
+        "encodeBody": "manual"
+    }
+    with pytest.raises(ValueError, match="encode_body"):
+        asgi.Options(encode_body="gzip")
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("mode", ["buffered", "streaming"])
+@pytest.mark.parametrize("encoding", ["gzip", "identity"])
+@pytest.mark.parametrize("encode_body", ["automatic", "manual"])
+async def test_encode_body_option(encode_body, encoding, mode):
+    # SELF.fetch does not serialize the response over HTTP, so it cannot
+    # observe the runtime encoding the body; ASGI_HTTP goes over a socket.
+    response = await env.ASGI_HTTP.fetch(
+        f"http://example.com/encoded/{encoding}/{mode}/{encode_body}",
+        headers={"accept-encoding": "gzip"},
+    )
+    assert response.status == 200
+    assert response.headers["content-encoding"] == encoding
+
+    # fetch() decodes the declared Content-Encoding once, like a browser.
+    body = await response.bytes()
+    if encoding == "gzip" and encode_body == "automatic":
+        # The runtime gzips the app's already-gzipped bytes a second time.
+        assert gzip.decompress(body) == ENCODED_RESPONSE_PAYLOAD
+    else:
+        assert body == ENCODED_RESPONSE_PAYLOAD