From 30b011282e1548797d6107aa6f2b91cf088e3354 Mon Sep 17 00:00:00 2001 From: Nick Franck <46548427+CyMule@users.noreply.github.com> Date: Wed, 5 Aug 2026 15:40:10 -0400 Subject: [PATCH] refactor(etl-uvicorn): extract reserved /invoke fields in a route dependency, drop the sealed-settings opt-in MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The body-replay middleware is replaced by bind_invocation_envelope, a FastAPI dependency that reads the framework's Starlette-cached body parse, so the /invoke body is buffered and decoded exactly once — no receive replay, no second copy of a large body. The request-size cap stays below the framework as InvokeBodyLimitMiddleware, a streaming byte counter that never buffers. /metadata now advertises invoke_with_sealed_dag_node_settings unconditionally: sealed per-invoke settings are the platform's required settings path, so the wrap_in_fastapi/generate_fast_api parameter and --sealed-dag-node-settings CLI flag are removed. --- CHANGELOG.md | 35 +- test/api/test_invocation_envelope.py | 446 +++++++++++++++++ test/api/test_invocation_middleware.py | 458 ------------------ test/api/test_invocation_settings.py | 22 +- .../etl_uvicorn/api_generator.py | 22 +- .../etl_uvicorn/main.py | 10 - .../invocation_settings.py | 335 ++++++++----- 7 files changed, 685 insertions(+), 643 deletions(-) create mode 100644 test/api/test_invocation_envelope.py delete mode 100644 test/api/test_invocation_middleware.py diff --git a/CHANGELOG.md b/CHANGELOG.md index b71f6b2..c8cb399 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,13 +1,14 @@ ## 0.1.0 * **This package now owns the `/invoke` transport for the reserved fields.** - `unstructured_platform_plugins.invocation_settings` holds the ASGI middleware, the `/metadata` - capability route, the request-scoped binding, and `http_status_for` — the HTTP spelling of the + `unstructured_platform_plugins.invocation_settings` holds the `/invoke` binding dependency and + body cap, the `/metadata` capability route, the request-scoped accessors, and `http_status_for` + — the HTTP spelling of the library's normative `blame` → status rule. It sits on `utic-invocation-settings >=0.4.0`, which owns the *settings contract* — which key carries settings, how a sealed envelope is told from plaintext, and what an absent field is allowed to mean. That split is deliberate: the - absence rule is a security decision and belongs next to the crypto it governs, while body - buffering and route registration belong here, where a web framework is already a dependency. + absence rule is a security decision and belongs next to the crypto it governs, while request + handling and route registration belong here, where a web framework is already a dependency. Nothing about the sealed-settings wire format is decided in this repository. * **This package now owns the `invocation_context` identity model.** `unstructured_platform_plugins.invocation_context` holds `InvocationContext`, @@ -23,17 +24,25 @@ values are exposed through `current_invocation_settings()` / `current_invocation_context()`. An absent field preserves the existing fallback behaviour; under `FF_REQUIRE_INVOKE_WITH_SEALED_DAG_NODE_SETTINGS` missing or plaintext settings fail closed. - Repeated installation is safe: the middleware installs once and the last `/metadata` + Repeated installation is safe: the dependency installs once and the last `/metadata` registration wins. -* **New opt-in `invoke_with_sealed_dag_node_settings` advertisement.** Pass - `invoke_with_sealed_dag_node_settings=True` to `wrap_in_fastapi` / `generate_fast_api` (or - `--sealed-dag-node-settings` on the CLI) only for a plugin that consumes per-invoke settings; - it advertises that the application accepts and consumes sealed per-invocation settings. - A plugin that serves a custom `/metadata` payload must register it via `add_metadata_route` - (which replaces the wrapper's route) — a plain `@app.get("/metadata")` added after construction - is shadowed by the wrapper's earlier registration. +* **Extraction is a route dependency.** `bind_invocation_envelope` reads the body the framework + buffered and parsed (`request.json()` is Starlette-cached), so the `/invoke` body is held and + decoded exactly once per request. `install_invocation_envelope` attaches the dependency to the + registered POST `/invoke` route(s), registers the failure response shape, and installs + `InvokeBodyLimitMiddleware`, a streaming byte counter that answers 413 over the cap without + buffering; it must be called after the `/invoke` route is registered and raises if none exists. +* **Sealed settings have no opt-out.** `/metadata` advertises `invocation_settings`, + `invocation_context` and `invoke_with_sealed_dag_node_settings` unconditionally: per-invoke + sealed settings are the platform's required settings path, and every plugin on this version is + expected to consume `current_invocation_settings()` in place of any other settings source. The + `invoke_with_sealed_dag_node_settings` parameter on `wrap_in_fastapi` / `generate_fast_api` and + the `--sealed-dag-node-settings` CLI flag are gone. A plugin that serves a custom `/metadata` + payload must register it via `add_metadata_route` (which replaces the wrapper's route) — a + plain `@app.get("/metadata")` added after construction is shadowed by the wrapper's earlier + registration. * **Resolution runs off the event loop.** A cold resolve is an RSA unwrap of a couple of - milliseconds and this middleware fronts every invoke on the pod, so it is dispatched with + milliseconds and this dependency fronts every invoke on the pod, so it is dispatched with `asyncio.to_thread` rather than blocking the loop. * **Failures map through the library's blame taxonomy**, not a flat 500: only a caller-fixable fault answers 422. Sealing drift, an envelope addressed to another recipient and a broken local diff --git a/test/api/test_invocation_envelope.py b/test/api/test_invocation_envelope.py new file mode 100644 index 0000000..aae30c8 --- /dev/null +++ b/test/api/test_invocation_envelope.py @@ -0,0 +1,446 @@ +"""The transport for the reserved /invoke fields: dependency binding, /metadata, the body cap. + +The *contract* these exercise — which shapes carry settings, what absence means — is owned and +tested in `utic_invocation_settings`. What is tested here is delivery: that the framework-parsed +body is resolved and bound for the handler, and that a payload which cannot be used fails the +request instead of reaching a handler as absence. +""" + +from __future__ import annotations + +import asyncio +import json +from base64 import b64decode, b64encode +from typing import Optional + +import pytest +from cryptography.hazmat.primitives import serialization +from cryptography.hazmat.primitives.asymmetric import rsa +from fastapi import FastAPI, Request +from fastapi.testclient import TestClient +from utic_invocation_settings import ( + DAG_NODE_SETTINGS_KEY, + REQUIRE_INVOKE_WITH_SEALED_DAG_NODE_SETTINGS_ENV_VAR, + default_resolver, + reset_workload_identity_cache, +) +from utic_invocation_settings.crypto import seal_settings + +from unstructured_platform_plugins.invocation_settings import ( + InvokeBodyLimitMiddleware, + add_metadata_route, + current_invocation_context, + current_invocation_settings, + install_invocation_envelope, +) + +SENTINEL_SECRET = "sealed-settings-sentinel-secret" + + +@pytest.fixture(scope="session") +def private_key() -> rsa.RSAPrivateKey: + return rsa.generate_private_key(public_exponent=65537, key_size=3072) + + +@pytest.fixture(autouse=True) +def isolated_identity(monkeypatch): + """The identity memo and the resolver caches both outlive a test; an inherited env var or a + stale entry would make these order-dependent.""" + for var in ( + "WORKLOAD_IDENTITY_DIR", + "INVOCATION_SETTINGS_KEY_DIR", + REQUIRE_INVOKE_WITH_SEALED_DAG_NODE_SETTINGS_ENV_VAR, + ): + monkeypatch.delenv(var, raising=False) + reset_workload_identity_cache() + default_resolver().clear_caches() + yield + reset_workload_identity_cache() + default_resolver().clear_caches() + + +@pytest.fixture +def key_dir(tmp_path, private_key, monkeypatch, isolated_identity): + (tmp_path / "tls.key").write_bytes( + private_key.private_bytes( + encoding=serialization.Encoding.PEM, + format=serialization.PrivateFormat.PKCS8, + encryption_algorithm=serialization.NoEncryption(), + ) + ) + monkeypatch.setenv("WORKLOAD_IDENTITY_DIR", str(tmp_path)) + reset_workload_identity_cache() + return tmp_path + + +def sealed_payload(private_key: rsa.RSAPrivateKey, settings: dict) -> dict: + return seal_settings(settings, private_key.public_key()).model_dump( + mode="json", exclude_none=True + ) + + +def tampered(sealed: dict) -> dict: + ciphertext = bytearray(b64decode(sealed["content_encryption"]["ciphertext"])) + ciphertext[0] ^= 0x01 + sealed["content_encryption"]["ciphertext"] = b64encode(bytes(ciphertext)).decode() + return sealed + + +ALL_CAPABILITIES = [ + "invocation_settings", + "invocation_context", + "invoke_with_sealed_dag_node_settings", +] + + +class TestMetadataRoute: + def test_advertises_every_capability_unconditionally(self): + # Sealed per-invoke settings are the required settings path, not an opt-in: a flag whose + # absence quietly kept a plugin on boot-time state would be the wrong-tenant hazard the + # sealed path exists to prevent. + app = FastAPI() + add_metadata_route(app, identifier="plugin.test") + + with TestClient(app) as client: + payload = client.get("/metadata").json() + + assert payload == { + "api_version": "3", + "identifier": "plugin.test", + "capabilities": ALL_CAPABILITIES, + } + + def test_last_call_wins(self): + # A host wrapper registers /metadata at construction; the plugin's later call with its own + # identifier must replace it, not be shadowed by route order. + app = FastAPI() + add_metadata_route(app, identifier="wrapper.default") + add_metadata_route(app, identifier="plugin.test") + + with TestClient(app) as client: + payload = client.get("/metadata").json() + + assert payload["identifier"] == "plugin.test" + + def test_replaces_a_directly_registered_metadata_route(self): + # A route the app registered itself would otherwise win by route order and pin its stale + # payload. + app = FastAPI() + + @app.get("/metadata") + async def stale_metadata() -> dict: + return {"api_version": "3", "identifier": "stale", "capabilities": []} + + add_metadata_route(app, identifier="plugin.test") + + with TestClient(app) as client: + payload = client.get("/metadata").json() + + assert payload["identifier"] == "plugin.test" + assert payload["capabilities"] == ALL_CAPABILITIES + + +class _Recorder: + """What the /invoke handler observed: the bound envelope, and whether it ran at all.""" + + def __init__(self): + self.called = False + self.seen_settings = "unset" + self.seen_context = "unset" + + +def _envelope_app(recorder: _Recorder, max_body_bytes: Optional[int] = None) -> FastAPI: + """A hand-rolled host app: the /invoke route reads the raw request, like a plugin that owns + its own route, so any body shape reaches the handler unless the dependency rejects it.""" + app = FastAPI() + + @app.post("/invoke") + async def invoke(request: Request) -> dict: + recorder.called = True + recorder.seen_settings = current_invocation_settings() + recorder.seen_context = current_invocation_context() + return {} + + @app.get("/schema") + async def schema() -> dict: + recorder.called = True + return {} + + if max_body_bytes is None: + install_invocation_envelope(app) + else: + install_invocation_envelope(app, max_body_bytes=max_body_bytes) + return app + + +def _post_invoke(payload, recorder: Optional[_Recorder] = None, **app_kwargs): + recorder = recorder if recorder is not None else _Recorder() + app = _envelope_app(recorder, **app_kwargs) + with TestClient(app, raise_server_exceptions=False) as client: + if isinstance(payload, bytes): + response = client.post("/invoke", content=payload) + else: + response = client.post("/invoke", json=payload) + return recorder, response + + +class TestInvocationEnvelopeBinding: + def test_binds_reserved_fields(self): + recorder, response = _post_invoke( + { + "element_dicts": "/in.json", + "invocation_settings": {"model": "m"}, + "invocation_context": {"schema_version": "1", "job_id": "job-1"}, + } + ) + + assert response.status_code == 200 + assert recorder.seen_settings == {"model": "m"} + assert recorder.seen_context.job_id == "job-1" + + def test_absent_fields_bind_none(self): + recorder, response = _post_invoke({"element_dicts": "/in.json"}) + + assert response.status_code == 200 + assert recorder.seen_settings is None + assert recorder.seen_context is None + + def test_non_dict_reserved_field_is_rejected(self): + recorder, response = _post_invoke({"invocation_settings": "not-a-dict"}) + + assert not recorder.called + assert response.status_code == 422 + body = response.json() + assert "invocation_settings" in body["detail"] + assert body["reason"] == "malformed_envelope" + + def test_context_with_an_unreadable_schema_version_fails_as_platform_error(self): + # Absence means "older caller, use the boot settings"; a context this plugin cannot read + # must not be downgraded to that. And it is deployment skew, not a caller fault — a 422 + # would let an upstream blame classifier pin version skew on the customer. + recorder, response = _post_invoke({"invocation_context": {"schema_version": "99"}}) + + assert not recorder.called + assert response.status_code == 500 + body = response.json() + assert "invocation_context" in body["detail"] + assert "UnsupportedContextVersionError" in body["detail"] + assert body["reason"] == "unsupported_context_version" + + def test_malformed_context_is_rejected(self): + recorder, response = _post_invoke({"invocation_context": "not-an-object"}) + + assert not recorder.called + assert response.status_code == 422 + + def test_non_invoke_routes_are_untouched(self): + recorder = _Recorder() + app = _envelope_app(recorder) + + with TestClient(app) as client: + response = client.get("/schema") + + assert response.status_code == 200 + assert recorder.called + + def test_envelope_does_not_leak_between_requests(self): + recorder = _Recorder() + app = _envelope_app(recorder) + + with TestClient(app) as client: + client.post("/invoke", json={"invocation_settings": {"model": "m"}}) + client.post("/invoke", json={"element_dicts": "/in.json"}) + + assert recorder.seen_settings is None + assert recorder.seen_context is None + + def test_install_without_an_invoke_route_fails_loudly(self): + # Silently installing nothing would leave the reserved fields unbound on every request — + # the plugin would run on boot-time settings while advertising otherwise. + with pytest.raises(RuntimeError, match="POST /invoke"): + install_invocation_envelope(FastAPI()) + + def test_oversized_body_is_rejected(self): + recorder, response = _post_invoke( + {"invocation_settings": {"pad": "x" * 64}}, max_body_bytes=16 + ) + + assert not recorder.called + assert response.status_code == 413 + + +class TestInvokeBodyLimitMiddleware: + """ASGI-level behavior of the streaming byte counter: no buffering, disconnect pass-through.""" + + @staticmethod + def _run(middleware, scope, chunks) -> tuple[list, list]: + received = [] + sent = [] + + async def receive(): + return chunks.pop(0) + + async def send(message): + sent.append(message) + + async def downstream(scope, receive, send): + while True: + message = await receive() + received.append(message) + if message["type"] != "http.request" or not message.get("more_body"): + break + await send({"type": "http.response.start", "status": 200, "headers": []}) + await send({"type": "http.response.body", "body": b"{}"}) + + middleware = middleware(downstream) + asyncio.run(middleware(scope, receive, send)) + return received, sent + + def test_chunked_body_over_the_cap_answers_413_and_cuts_downstream(self): + chunks = [ + {"type": "http.request", "body": b"x" * 10, "more_body": True}, + {"type": "http.request", "body": b"x" * 10, "more_body": False}, + ] + received, sent = self._run( + lambda app: InvokeBodyLimitMiddleware(app, max_body_bytes=16), + {"type": "http", "method": "POST", "path": "/invoke"}, + chunks, + ) + + assert received[-1] == {"type": "http.disconnect"} + assert sent[0]["status"] == 413 + + def test_client_disconnect_passes_through_uncounted(self): + chunks = [{"type": "http.disconnect"}] + received, sent = self._run( + lambda app: InvokeBodyLimitMiddleware(app, max_body_bytes=16), + {"type": "http", "method": "POST", "path": "/invoke"}, + chunks, + ) + + assert received == [{"type": "http.disconnect"}] + + def test_non_invoke_requests_are_not_capped(self): + chunks = [{"type": "http.request", "body": b"x" * 100, "more_body": False}] + received, sent = self._run( + lambda app: InvokeBodyLimitMiddleware(app, max_body_bytes=16), + {"type": "http", "method": "GET", "path": "/schema"}, + chunks, + ) + + assert sent[0]["status"] == 200 + + +class TestEnvelopeResolution: + """Sealed payloads through the binding dependency. The HTTP class comes from the library's + blame taxonomy, so only a caller-fixable fault is a 422.""" + + def test_sealed_envelope_binds_plaintext_settings(self, key_dir, private_key): + settings = {"api_key": SENTINEL_SECRET, "max_characters": 700} + sealed = sealed_payload(private_key, settings) + + recorder, response = _post_invoke({"invocation_settings": sealed}) + + assert response.status_code == 200 + assert recorder.seen_settings == settings + + def test_composite_dag_node_settings_member_is_opened(self, key_dir, private_key): + settings = {"api_key": SENTINEL_SECRET, "max_characters": 700} + composite = {DAG_NODE_SETTINGS_KEY: sealed_payload(private_key, settings)} + + recorder, response = _post_invoke({"invocation_settings": composite}) + + assert response.status_code == 200 + assert recorder.seen_settings == settings + + def test_plain_dict_settings_pass_through(self): + recorder, response = _post_invoke({"invocation_settings": {"model": "m"}}) + + assert response.status_code == 200 + assert recorder.seen_settings == {"model": "m"} + + def test_non_envelope_member_fails_as_platform_error(self, key_dir): + composite = {DAG_NODE_SETTINGS_KEY: {"model": "m"}} + + recorder, response = _post_invoke({"invocation_settings": composite}) + + assert response.status_code == 500 + assert "MalformedDagNodeSettingsError" in response.json()["detail"] + assert not recorder.called + + def test_undecryptable_envelope_fails_without_leaking_the_secret( + self, key_dir, private_key, caplog + ): + sealed = tampered(sealed_payload(private_key, {"api_key": SENTINEL_SECRET})) + + recorder, response = _post_invoke({"invocation_settings": sealed}) + + assert response.status_code == 500 + detail = response.json()["detail"] + assert "DecryptionError" in detail + assert SENTINEL_SECRET not in caplog.text + assert SENTINEL_SECRET not in detail + assert not recorder.called + + def test_unmounted_identity_fails_as_platform_error(self, tmp_path, monkeypatch, private_key): + monkeypatch.setenv("WORKLOAD_IDENTITY_DIR", str(tmp_path / "missing")) + sealed = sealed_payload(private_key, {"api_key": SENTINEL_SECRET}) + + recorder, response = _post_invoke({"invocation_settings": sealed}) + + assert response.status_code == 500 + assert "IdentityNotMountedError" in response.json()["detail"] + assert not recorder.called + + +class TestRequireSealedDagNodeSettings: + """A native pod's operator sets FF_REQUIRE_INVOKE_WITH_SEALED_DAG_NODE_SETTINGS in the same + decision that drops the init-secrets sidecar: without it, an invoke that arrived with no + envelope would fall back to a settings file that was never written.""" + + @pytest.fixture(autouse=True) + def _require_sealed(self, monkeypatch): + monkeypatch.setenv(REQUIRE_INVOKE_WITH_SEALED_DAG_NODE_SETTINGS_ENV_VAR, "true") + + def test_missing_settings_fail_as_platform_error(self): + recorder, response = _post_invoke({"element_dicts": "/tmp/x.json"}) + + assert response.status_code == 500 + assert "SealedDagNodeSettingsRequiredError" in response.json()["detail"] + assert not recorder.called + + def test_plaintext_settings_fail_as_platform_error(self): + recorder, response = _post_invoke({"invocation_settings": {"model": "m"}}) + + assert response.status_code == 500 + assert not recorder.called + + def test_sealed_settings_still_bind(self, key_dir, private_key): + settings = {"api_key": SENTINEL_SECRET} + composite = {DAG_NODE_SETTINGS_KEY: sealed_payload(private_key, settings)} + + recorder, response = _post_invoke({"invocation_settings": composite}) + + assert response.status_code == 200 + assert recorder.seen_settings == settings + + def test_bodyless_invoke_fails_as_platform_error(self): + # A bodyless invoke cannot carry the envelope; on a pod with no boot settings file it must + # not dispatch a handler with no settings source at all. + recorder, response = _post_invoke(b"") + + assert response.status_code == 500 + assert not recorder.called + + def test_non_object_json_body_fails_as_platform_error(self): + recorder, response = _post_invoke(json.dumps([{"element": 1}]).encode()) + + assert response.status_code == 500 + assert not recorder.called + + +def test_non_object_bodies_pass_through_when_sealed_settings_not_required(): + recorder, response = _post_invoke(b"") + + assert response.status_code == 200 + assert recorder.seen_settings is None diff --git a/test/api/test_invocation_middleware.py b/test/api/test_invocation_middleware.py deleted file mode 100644 index 9906d8c..0000000 --- a/test/api/test_invocation_middleware.py +++ /dev/null @@ -1,458 +0,0 @@ -"""The transport for the reserved /invoke fields: middleware, /metadata, request-scoped binding. - -The *contract* these exercise — which shapes carry settings, what absence means — is owned and -tested in `utic_invocation_settings`. What is tested here is delivery: that a raw body is read, -resolved, bound, and replayed intact, and that a payload which cannot be used fails the request -instead of reaching a handler as absence. -""" - -from __future__ import annotations - -import asyncio -import json -from base64 import b64decode, b64encode - -import pytest -from cryptography.hazmat.primitives import serialization -from cryptography.hazmat.primitives.asymmetric import rsa -from fastapi import FastAPI -from fastapi.testclient import TestClient -from utic_invocation_settings import ( - DAG_NODE_SETTINGS_KEY, - REQUIRE_INVOKE_WITH_SEALED_DAG_NODE_SETTINGS_ENV_VAR, - default_resolver, - reset_workload_identity_cache, -) -from utic_invocation_settings.crypto import seal_settings - -from unstructured_platform_plugins.invocation_settings import ( - InvocationEnvelopeMiddleware, - add_metadata_route, - current_invocation_context, - current_invocation_settings, -) - -SENTINEL_SECRET = "sealed-settings-sentinel-secret" - - -@pytest.fixture(scope="session") -def private_key() -> rsa.RSAPrivateKey: - return rsa.generate_private_key(public_exponent=65537, key_size=3072) - - -@pytest.fixture(autouse=True) -def isolated_identity(monkeypatch): - """The identity memo and the resolver caches both outlive a test; an inherited env var or a - stale entry would make these order-dependent.""" - for var in ("WORKLOAD_IDENTITY_DIR", "INVOCATION_SETTINGS_KEY_DIR", - REQUIRE_INVOKE_WITH_SEALED_DAG_NODE_SETTINGS_ENV_VAR): - monkeypatch.delenv(var, raising=False) - reset_workload_identity_cache() - default_resolver().clear_caches() - yield - reset_workload_identity_cache() - default_resolver().clear_caches() - - -@pytest.fixture -def key_dir(tmp_path, private_key, monkeypatch, isolated_identity): - (tmp_path / "tls.key").write_bytes( - private_key.private_bytes( - encoding=serialization.Encoding.PEM, - format=serialization.PrivateFormat.PKCS8, - encryption_algorithm=serialization.NoEncryption(), - ) - ) - monkeypatch.setenv("WORKLOAD_IDENTITY_DIR", str(tmp_path)) - reset_workload_identity_cache() - return tmp_path - - -def sealed_payload(private_key: rsa.RSAPrivateKey, settings: dict) -> dict: - return seal_settings(settings, private_key.public_key()).model_dump( - mode="json", exclude_none=True - ) - - -def tampered(sealed: dict) -> dict: - ciphertext = bytearray(b64decode(sealed["content_encryption"]["ciphertext"])) - ciphertext[0] ^= 0x01 - sealed["content_encryption"]["ciphertext"] = b64encode(bytes(ciphertext)).decode() - return sealed - - -class TestMetadataRoute: - def test_advertises_settings_and_context_capabilities(self): - app = FastAPI() - add_metadata_route(app, identifier="plugin.test") - - with TestClient(app) as client: - payload = client.get("/metadata").json() - - assert payload == { - "api_version": "3", - "identifier": "plugin.test", - "capabilities": ["invocation_settings", "invocation_context"], - } - - def test_sealed_dag_node_settings_flag_advertises_capability(self): - app = FastAPI() - add_metadata_route(app, invoke_with_sealed_dag_node_settings=True) - - with TestClient(app) as client: - payload = client.get("/metadata").json() - - assert payload["capabilities"] == [ - "invocation_settings", - "invocation_context", - "invoke_with_sealed_dag_node_settings", - ] - - def test_last_call_wins(self): - # A host wrapper registers /metadata with defaults at construction; the plugin's later call - # with the sealed capability must replace it, not be shadowed by route order. - app = FastAPI() - add_metadata_route(app, identifier="plugin.test") - add_metadata_route(app, identifier="plugin.test", invoke_with_sealed_dag_node_settings=True) - - with TestClient(app) as client: - payload = client.get("/metadata").json() - - assert "invoke_with_sealed_dag_node_settings" in payload["capabilities"] - - def test_replaces_a_directly_registered_metadata_route(self): - # A route the app registered itself would otherwise win by route order and pin its stale - # payload. - app = FastAPI() - - @app.get("/metadata") - async def stale_metadata() -> dict: - return {"api_version": "3", "identifier": "stale", "capabilities": []} - - add_metadata_route(app, identifier="plugin.test", invoke_with_sealed_dag_node_settings=True) - - with TestClient(app) as client: - payload = client.get("/metadata").json() - - assert payload["identifier"] == "plugin.test" - assert "invoke_with_sealed_dag_node_settings" in payload["capabilities"] - - -def _invoke_scope(path: str = "/invoke", method: str = "POST") -> dict: - return {"type": "http", "method": method, "path": path} - - -def _receive_for(body: bytes): - chunks = [ - {"type": "http.request", "body": body[: len(body) // 2], "more_body": True}, - {"type": "http.request", "body": body[len(body) // 2 :], "more_body": False}, - ] - - async def receive(): - return chunks.pop(0) - - return receive - - -async def _ok(send): - await send({"type": "http.response.start", "status": 200, "headers": []}) - await send({"type": "http.response.body", "body": b"{}"}) - - -class _DownstreamApp: - """Records the replayed body and the envelope bound while handling.""" - - def __init__(self): - self.body = None - self.called = False - self.seen_settings = "unset" - self.seen_context = "unset" - - async def __call__(self, scope, receive, send): - body = b"" - while True: - message = await receive() - body += message.get("body", b"") - if not message.get("more_body"): - break - self.body = body - self.called = True - self.seen_settings = current_invocation_settings() - self.seen_context = current_invocation_context() - await _ok(send) - - -def _run_middleware(body: bytes, scope: dict | None = None) -> tuple[_DownstreamApp, list]: - downstream = _DownstreamApp() - middleware = InvocationEnvelopeMiddleware(downstream) - sent = [] - - async def send(message): - sent.append(message) - - asyncio.run(middleware(scope or _invoke_scope(), _receive_for(body), send)) - return downstream, sent - - -class TestInvocationEnvelopeMiddleware: - def test_binds_reserved_fields_and_replays_body(self): - body = json.dumps( - { - "element_dicts": "/in.json", - "invocation_settings": {"model": "m"}, - "invocation_context": {"schema_version": "1", "job_id": "job-1"}, - } - ).encode() - - downstream, sent = _run_middleware(body) - - assert downstream.body == body - assert downstream.seen_settings == {"model": "m"} - assert downstream.seen_context.job_id == "job-1" - assert sent[0]["status"] == 200 - - def test_absent_fields_bind_none(self): - downstream, _ = _run_middleware(json.dumps({"element_dicts": "/in.json"}).encode()) - - assert downstream.seen_settings is None - assert downstream.seen_context is None - - def test_non_dict_reserved_field_is_rejected(self): - downstream, sent = _run_middleware( - json.dumps({"invocation_settings": "not-a-dict"}).encode() - ) - - assert downstream.body is None - assert sent[0]["status"] == 422 - body = json.loads(sent[1]["body"]) - assert "invocation_settings" in body["detail"] - assert body["reason"] == "malformed_envelope" - - def test_context_with_an_unreadable_schema_version_fails_as_platform_error(self): - # Absence means "older caller, use the boot settings"; a context this plugin cannot read - # must not be downgraded to that. And it is deployment skew, not a caller fault — a 422 - # would let an upstream blame classifier pin version skew on the customer. - downstream, sent = _run_middleware( - json.dumps({"invocation_context": {"schema_version": "99"}}).encode() - ) - - assert downstream.body is None - assert sent[0]["status"] == 500 - body = json.loads(sent[1]["body"]) - assert "invocation_context" in body["detail"] - assert "UnsupportedContextVersionError" in body["detail"] - assert body["reason"] == "unsupported_context_version" - - def test_malformed_context_is_rejected(self): - downstream, sent = _run_middleware( - json.dumps({"invocation_context": "not-an-object"}).encode() - ) - - assert downstream.body is None - assert sent[0]["status"] == 422 - - def test_non_invoke_requests_pass_through_untouched(self): - body = json.dumps({"invocation_settings": "not-a-dict"}).encode() - - downstream, sent = _run_middleware(body, scope=_invoke_scope(path="/schema", method="GET")) - - assert downstream.body == body - assert sent[0]["status"] == 200 - - def test_envelope_is_reset_after_request(self): - body = json.dumps({"invocation_settings": {"model": "m"}}).encode() - - async def scenario(): - middleware = InvocationEnvelopeMiddleware(_DownstreamApp()) - - async def send(_message): - pass - - await middleware(_invoke_scope(), _receive_for(body), send) - return current_invocation_settings(), current_invocation_context() - - assert asyncio.run(scenario()) == (None, None) - - def test_oversized_body_is_rejected(self): - body = json.dumps({"invocation_settings": {"pad": "x" * 64}}).encode() - - downstream = _DownstreamApp() - middleware = InvocationEnvelopeMiddleware(downstream, max_body_bytes=16) - sent = [] - - async def send(message): - sent.append(message) - - asyncio.run(middleware(_invoke_scope(), _receive_for(body), send)) - - assert downstream.body is None - assert sent[0]["status"] == 413 - - def test_drained_replay_proxies_disconnect(self): - body = json.dumps({"invocation_settings": {"model": "m"}}).encode() - - class _DisconnectWatcher: - def __init__(self): - self.saw_disconnect = False - - async def __call__(self, scope, receive, send): - while True: - message = await receive() - if message["type"] == "http.disconnect": - self.saw_disconnect = True - return - if not message.get("more_body"): - break - message = await receive() - self.saw_disconnect = message["type"] == "http.disconnect" - await _ok(send) - - chunks = [ - {"type": "http.request", "body": body, "more_body": False}, - {"type": "http.disconnect"}, - ] - - async def receive(): - return chunks.pop(0) - - watcher = _DisconnectWatcher() - middleware = InvocationEnvelopeMiddleware(watcher) - sent = [] - - async def send(message): - sent.append(message) - - asyncio.run(middleware(_invoke_scope(), receive, send)) - - assert watcher.saw_disconnect - - -class TestMiddlewareResolution: - """Sealed payloads through the middleware. The HTTP class comes from the library's blame - taxonomy, so only a caller-fixable fault is a 422.""" - - def test_sealed_envelope_binds_plaintext_settings(self, key_dir, private_key): - settings = {"api_key": SENTINEL_SECRET, "max_characters": 700} - sealed = sealed_payload(private_key, settings) - - downstream, sent = _run_middleware( - json.dumps({"invocation_settings": sealed}).encode() - ) - - assert sent[0]["status"] == 200 - assert downstream.seen_settings == settings - - def test_composite_dag_node_settings_member_is_opened(self, key_dir, private_key): - settings = {"api_key": SENTINEL_SECRET, "max_characters": 700} - composite = {DAG_NODE_SETTINGS_KEY: sealed_payload(private_key, settings)} - - downstream, sent = _run_middleware( - json.dumps({"invocation_settings": composite}).encode() - ) - - assert sent[0]["status"] == 200 - assert downstream.seen_settings == settings - - def test_plain_dict_settings_pass_through(self): - downstream, sent = _run_middleware( - json.dumps({"invocation_settings": {"model": "m"}}).encode() - ) - - assert sent[0]["status"] == 200 - assert downstream.seen_settings == {"model": "m"} - - def test_non_envelope_member_fails_as_platform_error(self, key_dir): - composite = {DAG_NODE_SETTINGS_KEY: {"model": "m"}} - - downstream, sent = _run_middleware( - json.dumps({"invocation_settings": composite}).encode() - ) - - assert sent[0]["status"] == 500 - assert "MalformedDagNodeSettingsError" in json.loads(sent[1]["body"])["detail"] - assert not downstream.called - - def test_undecryptable_envelope_fails_without_leaking_the_secret( - self, key_dir, private_key, caplog - ): - sealed = tampered(sealed_payload(private_key, {"api_key": SENTINEL_SECRET})) - - downstream, sent = _run_middleware( - json.dumps({"invocation_settings": sealed}).encode() - ) - - assert sent[0]["status"] == 500 - detail = json.loads(sent[1]["body"])["detail"] - assert "DecryptionError" in detail - assert SENTINEL_SECRET not in caplog.text - assert SENTINEL_SECRET not in detail - assert not downstream.called - - def test_unmounted_identity_fails_as_platform_error(self, tmp_path, monkeypatch, private_key): - monkeypatch.setenv("WORKLOAD_IDENTITY_DIR", str(tmp_path / "missing")) - sealed = sealed_payload(private_key, {"api_key": SENTINEL_SECRET}) - - downstream, sent = _run_middleware( - json.dumps({"invocation_settings": sealed}).encode() - ) - - assert sent[0]["status"] == 500 - assert "IdentityNotMountedError" in json.loads(sent[1]["body"])["detail"] - assert not downstream.called - - -class TestRequireSealedDagNodeSettings: - """A native pod's operator sets FF_REQUIRE_INVOKE_WITH_SEALED_DAG_NODE_SETTINGS in the same - decision that drops the init-secrets sidecar: without it, an invoke that arrived with no - envelope would fall back to a settings file that was never written.""" - - @pytest.fixture(autouse=True) - def _require_sealed(self, monkeypatch): - monkeypatch.setenv(REQUIRE_INVOKE_WITH_SEALED_DAG_NODE_SETTINGS_ENV_VAR, "true") - - def test_missing_settings_fail_as_platform_error(self): - downstream, sent = _run_middleware(json.dumps({"element_dicts": "/tmp/x.json"}).encode()) - - assert sent[0]["status"] == 500 - assert "SealedDagNodeSettingsRequiredError" in json.loads(sent[1]["body"])["detail"] - assert not downstream.called - - def test_plaintext_settings_fail_as_platform_error(self): - downstream, sent = _run_middleware( - json.dumps({"invocation_settings": {"model": "m"}}).encode() - ) - - assert sent[0]["status"] == 500 - assert not downstream.called - - def test_sealed_settings_still_bind(self, key_dir, private_key): - settings = {"api_key": SENTINEL_SECRET} - composite = {DAG_NODE_SETTINGS_KEY: sealed_payload(private_key, settings)} - - downstream, sent = _run_middleware( - json.dumps({"invocation_settings": composite}).encode() - ) - - assert sent[0]["status"] == 200 - assert downstream.seen_settings == settings - - def test_bodyless_invoke_fails_as_platform_error(self): - # A bodyless invoke cannot carry the envelope; on a pod with no boot settings file it must - # not dispatch a handler with no settings source at all. - downstream, sent = _run_middleware(b"") - - assert sent[0]["status"] == 500 - assert not downstream.called - - def test_non_object_json_body_fails_as_platform_error(self): - downstream, sent = _run_middleware(json.dumps([{"element": 1}]).encode()) - - assert sent[0]["status"] == 500 - assert not downstream.called - - -def test_non_object_bodies_pass_through_when_sealed_settings_not_required(): - downstream, sent = _run_middleware(b"") - - assert sent[0]["status"] == 200 - assert downstream.seen_settings is None diff --git a/test/api/test_invocation_settings.py b/test/api/test_invocation_settings.py index 5523b23..ecfd5d4 100644 --- a/test/api/test_invocation_settings.py +++ b/test/api/test_invocation_settings.py @@ -18,7 +18,7 @@ def _echo_settings(content: str) -> _Echo: return _Echo(content=content, settings=current_invocation_settings()) -def test_metadata_route_is_registered_with_default_capabilities(): +def test_metadata_route_advertises_every_capability(): client = TestClient(wrap_in_fastapi(func=_echo_settings, plugin_id="mock_plugin")) resp = client.get("/metadata") @@ -26,21 +26,11 @@ def test_metadata_route_is_registered_with_default_capabilities(): assert resp.status_code == 200 payload = resp.json() assert payload["identifier"] == "mock_plugin" - assert payload["capabilities"] == ["invocation_settings", "invocation_context"] - - -def test_sealed_capability_is_opt_in(): - client = TestClient( - wrap_in_fastapi( - func=_echo_settings, - plugin_id="mock_plugin", - invoke_with_sealed_dag_node_settings=True, - ) - ) - - payload = client.get("/metadata").json() - - assert "invoke_with_sealed_dag_node_settings" in payload["capabilities"] + assert payload["capabilities"] == [ + "invocation_settings", + "invocation_context", + "invoke_with_sealed_dag_node_settings", + ] def test_reserved_settings_field_binds_without_appearing_in_schema(): diff --git a/unstructured_platform_plugins/etl_uvicorn/api_generator.py b/unstructured_platform_plugins/etl_uvicorn/api_generator.py index 2426ecc..8fa0d25 100644 --- a/unstructured_platform_plugins/etl_uvicorn/api_generator.py +++ b/unstructured_platform_plugins/etl_uvicorn/api_generator.py @@ -128,14 +128,12 @@ def wrap_in_fastapi( func: Callable, plugin_id: str, precheck_func: Optional[Callable] = None, - invoke_with_sealed_dag_node_settings: bool = False, ) -> FastAPI: try: return _wrap_in_fastapi( func=func, plugin_id=plugin_id, precheck_func=precheck_func, - invoke_with_sealed_dag_node_settings=invoke_with_sealed_dag_node_settings, ) except Exception as e: logger.error(f"failed to wrap function in FastAPI: {e}", exc_info=True) @@ -146,7 +144,6 @@ def _wrap_in_fastapi( func: Callable, plugin_id: str, precheck_func: Optional[Callable] = None, - invoke_with_sealed_dag_node_settings: bool = False, ) -> FastAPI: if precheck_func is not None: check_precheck_func(precheck_func=precheck_func) @@ -371,17 +368,14 @@ async def get_id() -> str: except TypeError as e: raise TypeError(f"failed to validate function schema: {e}") from e - # The middleware handles the reserved /invoke fields (invocation_settings and + # The route dependency handles the reserved /invoke fields (invocation_settings and # invocation_context) outside the generated handler schema. It resolves sealed settings with - # the configured private key and exposes both values through request-scoped accessors. The - # sealed-settings capability remains opt-in because it asserts that the wrapped function - # consumes current_invocation_settings(), not merely that the host can resolve it. Repeated - # installation is safe: the middleware installs once and the last /metadata registration wins. - add_metadata_route( - fastapi_app, - identifier=plugin_id, - invoke_with_sealed_dag_node_settings=invoke_with_sealed_dag_node_settings, - ) + # the configured private key and exposes both values through request-scoped accessors. + # Sealed per-invoke settings are the platform's required settings path, so the capability is + # advertised unconditionally: every wrapped plugin is expected to consume + # current_invocation_settings(). Repeated installation is safe: the dependency installs once + # and the last /metadata registration wins. + add_metadata_route(fastapi_app, identifier=plugin_id) install_invocation_envelope(fastapi_app) FastAPIInstrumentor.instrument_app( @@ -398,7 +392,6 @@ def generate_fast_api( id_method: Optional[str] = None, precheck_str: Optional[str] = None, precheck_method: Optional[str] = None, - invoke_with_sealed_dag_node_settings: bool = False, ) -> FastAPI: instance = import_from_string(app) func = get_func(instance, method_name) @@ -421,5 +414,4 @@ def generate_fast_api( func=func, plugin_id=plugin_id, precheck_func=precheck_func, - invoke_with_sealed_dag_node_settings=invoke_with_sealed_dag_node_settings, ) diff --git a/unstructured_platform_plugins/etl_uvicorn/main.py b/unstructured_platform_plugins/etl_uvicorn/main.py index 62b067f..600e0c1 100644 --- a/unstructured_platform_plugins/etl_uvicorn/main.py +++ b/unstructured_platform_plugins/etl_uvicorn/main.py @@ -56,7 +56,6 @@ def api_wrapper( plugin_id_method: Optional[str] = None, precheck_app: Optional[str] = None, precheck_app_method: Optional[str] = None, - sealed_dag_node_settings: bool = False, **kwargs, ): # Make sure logging is configured before the call to run() so any setup has the same format @@ -74,7 +73,6 @@ def api_wrapper( id_method=plugin_id_method, precheck_str=precheck_app, precheck_method=precheck_app_method, - invoke_with_sealed_dag_node_settings=sealed_dag_node_settings, ) # Explicitly map values that are manipulated in the original # call to run(), preventing **kwargs reference @@ -132,14 +130,6 @@ def api_wrapper( "If precheck-app not provided, assumes method " "lives on main class passes in.", ), - click.Option( - ["--sealed-dag-node-settings"], - is_flag=True, - default=False, - help="Advertise the invoke_with_sealed_dag_node_settings capability on " - "/metadata. Set only for a plugin that consumes per-invoke settings " - "through current_invocation_settings().", - ), ] ) return cmd diff --git a/unstructured_platform_plugins/invocation_settings.py b/unstructured_platform_plugins/invocation_settings.py index 7b73d43..92e6a61 100644 --- a/unstructured_platform_plugins/invocation_settings.py +++ b/unstructured_platform_plugins/invocation_settings.py @@ -1,4 +1,4 @@ -"""Transport for the reserved `/invoke` fields: ASGI middleware, `/metadata`, request binding. +"""Transport for the reserved `/invoke` fields: request dependency, `/metadata`, body cap. The *settings contract* — which key carries settings, how a sealed envelope is told from plaintext, and what an absent field is allowed to mean — lives in @@ -9,10 +9,17 @@ off the wire and the results to the handler, and spelling the shared `blame` taxonomy as HTTP statuses. -The reserved fields are a first-class HTTP contract independent of the generated input schema. They -never appear in a plugin's declared signature, so `wrap_in_fastapi` keeps producing a handler model -built purely from the wrapped function, and a plugin reads the fields through +The reserved fields are a first-class HTTP contract independent of the generated input schema. +They never appear in a plugin's declared signature, so `wrap_in_fastapi` keeps producing a handler +model built purely from the wrapped function, and a plugin reads the fields through `current_invocation_settings()` / `current_invocation_context()` instead. + +Extraction runs in a route dependency (`bind_invocation_envelope`), so the `/invoke` body is +buffered and parsed exactly once: Starlette caches both the bytes and the parsed JSON on the +`Request`, and the dependency reads that cached parse. The request-size cap is the one concern +that must sit below the framework — neither Starlette nor uvicorn bounds request-body size — and +`InvokeBodyLimitMiddleware` enforces it by counting bytes as they stream through, without +buffering. """ from __future__ import annotations @@ -24,12 +31,16 @@ import threading import time from collections import OrderedDict -from collections.abc import Callable, Iterator, Mapping +from collections.abc import AsyncIterator, Callable, Iterator, Mapping from contextlib import contextmanager from contextvars import ContextVar from typing import Any, Optional, TypeVar -from fastapi import FastAPI +from fastapi import Depends, FastAPI, Request +from fastapi.dependencies.utils import get_parameterless_sub_dependant +from fastapi.routing import APIRoute +from starlette.requests import ClientDisconnect +from starlette.responses import JSONResponse from starlette.types import ASGIApp, Receive, Scope, Send from utic_invocation_settings import ( INVOKE_WITH_SEALED_DAG_NODE_SETTINGS_CAPABILITY, @@ -60,13 +71,14 @@ def http_status_for(error: BaseException) -> int: """ return 422 if getattr(error, "blame", None) is Blame.CALLER else 500 + T = TypeVar("T") _METADATA_PATH = "/metadata" _INVOKE_PATH = "/invoke" -# Bounds middleware body buffering; generous because batch invokes carry an array of file_data -# payloads. The framework buffers the same body afterward, so this cap is the only guard against +# Bounds the /invoke request body; generous because batch invokes carry an array of file_data +# payloads. Nothing below the framework buffers, so this cap is the only guard against # unbounded-memory requests. MAX_INVOKE_BODY_BYTES = 64 * 1024 * 1024 @@ -80,7 +92,7 @@ def current_invocation_settings() -> Optional[dict]: `None` means the field was genuinely absent — the only case in which a plugin may fall back to its boot-time settings. A field that arrived and could not be opened never reaches a handler: - the middleware fails the request first. + the binding dependency fails the request first. """ return _INVOCATION.get()[0] @@ -102,37 +114,33 @@ def invocation_envelope( _INVOCATION.reset(token) -def add_metadata_route( - app: FastAPI, - identifier: Optional[str] = None, - invoke_with_sealed_dag_node_settings: bool = False, -) -> None: +def add_metadata_route(app: FastAPI, identifier: Optional[str] = None) -> None: """Register GET /metadata advertising the reserved /invoke fields this plugin accepts. `/metadata` is the plugin API spec's own discovery surface (`PluginMetadataOutput`): capability flags are strings in its `capabilities` list, which is where the controller looks before forwarding the reserved fields — no controller-private probe route. - The two tiers make different claims. `invocation_settings` / `invocation_context` are - *transport-level* facts, advertised unconditionally because the installed middleware makes - them true for every wrapped app: the reserved fields will be received, resolved, and bound — - or the request failed. They say nothing about whether the handler reads the binding. - `invoke_with_sealed_dag_node_settings` is the *consumption* claim — this plugin opens sealed - `dag_node_settings` itself and its handler acts on the result — and stays a per-plugin opt-in - set in the same change that makes it true, because it is the flag that invites the controller - to seal settings to this pod in place of any other settings source. + All three capabilities — `invocation_settings`, `invocation_context`, and + `invoke_with_sealed_dag_node_settings` — are advertised unconditionally. Per-invoke sealed + settings are the platform's required settings path: every plugin on this package version + receives, resolves and binds the reserved fields, and is expected to consume + `current_invocation_settings()` in place of any other settings source. A plugin that ignores + the binding runs on stale boot-time state — the wrong-tenant hazard sealed settings exist to + prevent. Last call wins: the payload lives on `app.state` and every call overwrites it, while the route - is registered once. A host wrapper may register with default capabilities at app construction - and a plugin can still declare the sealed capability afterwards, with no route-order dependence. + is registered once. A host wrapper may register at app construction and a plugin can still + re-register with its own identifier afterwards, with no route-order dependence. """ - capabilities = [RESERVED_ENVELOPE_KEY, RESERVED_CONTEXT_KEY] - if invoke_with_sealed_dag_node_settings: - capabilities.append(INVOKE_WITH_SEALED_DAG_NODE_SETTINGS_CAPABILITY) app.state.plugin_metadata_payload = { "api_version": "3", "identifier": identifier, - "capabilities": capabilities, + "capabilities": [ + RESERVED_ENVELOPE_KEY, + RESERVED_CONTEXT_KEY, + INVOKE_WITH_SEALED_DAG_NODE_SETTINGS_CAPABILITY, + ], } if getattr(app.state, "plugin_metadata_route_installed", False): return @@ -147,6 +155,98 @@ async def plugin_metadata() -> dict: return app.state.plugin_metadata_payload +class UnusableInvocationEnvelope(Exception): + """A reserved /invoke field arrived but cannot be used. + + Raised by `bind_invocation_envelope` and answered by the handler + `install_invocation_envelope` registers, so the response shape — `detail` plus the library's + stable `reason` code as top-level siblings — stays what orchestrators parse, independent of + FastAPI's own error envelope. + """ + + def __init__(self, status_code: int, payload: dict): + super().__init__(payload.get("detail")) + self.status_code = status_code + self.payload = payload + + +async def _unusable_envelope_response( + _request: Request, exc: UnusableInvocationEnvelope +) -> JSONResponse: + return JSONResponse(status_code=exc.status_code, content=exc.payload) + + +async def bind_invocation_envelope(request: Request) -> AsyncIterator[None]: + """Resolve the reserved /invoke fields and bind them for the duration of the request. + + Runs as a route dependency, after the framework has read the body: `request.json()` is + Starlette-cached, so the parse is shared with the framework's own body handling. + A reserved field that is present but unusable fails the request rather than being + treated as absent, because absence is the signal to fall back to the boot-time settings file: + degrading a malformed field to absence would quietly answer a request configured for one + tenant with whatever the pod happened to boot with. Under + `FF_REQUIRE_INVOKE_WITH_SEALED_DAG_NODE_SETTINGS` there is no settings file to fall back to, + so absent or plaintext settings fail too — as does any /invoke whose body is not a JSON + object, since such a body cannot carry the envelope a native pod requires. + """ + try: + parsed = await request.json() + except ValueError: + # Empty or malformed body: the framework's own validation answers for the body itself; + # for the reserved fields it is absence, which resolve_invocation_settings still judges + # (absence is a failure on a native pod). + parsed = None + + raw_settings: Optional[Any] = None + if isinstance(parsed, dict): + raw_settings = parsed.get(RESERVED_ENVELOPE_KEY) + if RESERVED_ENVELOPE_KEY in parsed and not isinstance(raw_settings, dict): + raise UnusableInvocationEnvelope( + 422, + { + "detail": f"Invalid field: {RESERVED_ENVELOPE_KEY}", + "reason": MalformedEnvelopeError.reason, + }, + ) + try: + # Off the event loop: a cold resolve is an RSA unwrap of a couple of milliseconds, and + # this dependency fronts every invoke on the pod. + invocation_settings = await asyncio.to_thread(resolve_invocation_settings, raw_settings) + except Exception as exc: + # Class name only — never envelope contents, and never the exception's own message, + # which can embed request-controlled values. + logger.warning("unusable %s payload: %s", RESERVED_ENVELOPE_KEY, type(exc).__name__) + payload = {"detail": f"Unusable invocation settings: {type(exc).__name__}"} + reason = getattr(exc, "reason", None) + if isinstance(reason, str): + payload["reason"] = reason + raise UnusableInvocationEnvelope(http_status_for(exc), payload) from exc + + invocation_context: Optional[InvocationContext] = None + if isinstance(parsed, dict): + try: + invocation_context = extract_context(parsed) + except InvocationSettingsError as exc: + # A context this plugin cannot read fails loudly here rather than running with + # silently absent identity. Status comes from the blame taxonomy: a malformed + # field is the caller's 422, but an unreadable schema_version is deployment skew + # between platform components and must not read as a caller fault. The log line is + # truncated because the message can embed request-controlled values. + logger.warning("rejecting invalid %s: %.200s", RESERVED_CONTEXT_KEY, exc) + status = http_status_for(exc) + detail = ( + f"Invalid field: {RESERVED_CONTEXT_KEY}" + if status == 422 + else f"Unusable {RESERVED_CONTEXT_KEY}: {type(exc).__name__}" + ) + raise UnusableInvocationEnvelope( + status, {"detail": detail, "reason": exc.reason} + ) from exc + + with invocation_envelope(invocation_settings, invocation_context): + yield + + async def _send_json(send: Send, status_code: int, payload: dict) -> None: body = json.dumps(payload).encode() await send( @@ -162,21 +262,14 @@ async def _send_json(send: Send, status_code: int, payload: dict) -> None: await send({"type": "http.response.body", "body": body}) -class InvocationEnvelopeMiddleware: - """Extract the reserved envelope fields from the raw POST /invoke body. +class InvokeBodyLimitMiddleware: + """Reject a POST /invoke body over ``max_body_bytes`` with 413. - Pure ASGI rather than `BaseHTTPMiddleware`: the body has to be read before the framework parses - it and then replayed intact, which is exactly what the raw protocol allows and what a - request/response middleware would fight. - - A reserved field that is present but unusable fails the request rather than being treated as - absent, because absence is the signal to fall back to the boot-time settings file: degrading a - malformed field to absence would quietly answer a request configured for one tenant with - whatever the pod happened to boot with. Under `FF_REQUIRE_INVOKE_WITH_SEALED_DAG_NODE_SETTINGS` - there is no settings file to fall back to, so absent or plaintext settings fail too — as does - any /invoke whose body is not a JSON object, since such a body cannot carry the envelope a - native pod requires. A body over `max_body_bytes` is rejected with 413 before it can exhaust - memory. + Counts bytes as the framework consumes them; nothing is buffered here. When the count crosses + the cap the downstream read is answered with ``http.disconnect``, which aborts the framework's + body read before another byte is held, and the 413 is sent once the application has unwound. + This has to sit below the framework because neither Starlette nor uvicorn bounds request-body + size. """ def __init__(self, app: ASGIApp, max_body_bytes: int = MAX_INVOKE_BODY_BYTES): @@ -192,103 +285,83 @@ async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None: await self.app(scope, receive, send) return - messages = [] - buffered_bytes = 0 - while True: + seen = 0 + exceeded = False + response_started = False + + async def counting_receive() -> dict: + nonlocal seen, exceeded message = await receive() - messages.append(message) - buffered_bytes += len(message.get("body", b"")) - if buffered_bytes > self.max_body_bytes: - await _send_json(send, 413, {"detail": "Request body too large"}) + if message["type"] == "http.request": + seen += len(message.get("body", b"")) + if seen > self.max_body_bytes: + exceeded = True + return {"type": "http.disconnect"} + return message + + async def guarded_send(message: dict) -> None: + nonlocal response_started + if exceeded and not response_started: + # A response computed after the body was cut is answering a truncated request; + # drop it so the 413 below is what the caller sees. A response that started + # before the cap tripped keeps streaming — its start is already on the wire. return - if message["type"] != "http.request" or not message.get("more_body"): - break - body = b"".join(m.get("body", b"") for m in messages if m["type"] == "http.request") + if message["type"] == "http.response.start": + response_started = True + await send(message) try: - parsed = json.loads(body) if body else None - except ValueError: - # Malformed JSON: forward unchanged so the framework returns its own error. - parsed = None - - invocation_context: Optional[InvocationContext] = None - raw_settings: Optional[dict[str, Any]] = None - if isinstance(parsed, dict): - raw_settings = parsed.get(RESERVED_ENVELOPE_KEY) - if RESERVED_ENVELOPE_KEY in parsed and not isinstance(raw_settings, dict): - await _send_json( - send, - 422, - { - "detail": f"Invalid field: {RESERVED_ENVELOPE_KEY}", - "reason": MalformedEnvelopeError.reason, - }, - ) - return - # Resolved even when the field — or the whole JSON object — is absent: - # resolve_invocation_settings owns the FF_REQUIRE_INVOKE_WITH_SEALED_DAG_NODE_SETTINGS - # policy, under which a bodyless or non-object invoke cannot carry the envelope a native - # pod requires and is a failure, not a fallback signal. - try: - # Off the event loop: a cold resolve is an RSA unwrap of a couple of milliseconds, and - # this middleware sits in front of every invoke on the pod. - invocation_settings = await asyncio.to_thread(resolve_invocation_settings, raw_settings) - except Exception as exc: - # Class name only — never envelope contents, and never the exception's own message, - # which can embed request-controlled values. - logger.warning("unusable %s payload: %s", RESERVED_ENVELOPE_KEY, type(exc).__name__) - body = {"detail": f"Unusable invocation settings: {type(exc).__name__}"} - reason = getattr(exc, "reason", None) - if isinstance(reason, str): - body["reason"] = reason - await _send_json(send, http_status_for(exc), body) - return - if isinstance(parsed, dict): - try: - invocation_context = extract_context(parsed) - except InvocationSettingsError as exc: - # A context this plugin cannot read fails loudly here rather than running with - # silently absent identity. Status comes from the blame taxonomy: a malformed - # field is the caller's 422, but an unreadable schema_version is deployment skew - # between platform components and must not read as a caller fault. The log line is - # truncated because the message can embed request-controlled values. - logger.warning("rejecting invalid %s: %.200s", RESERVED_CONTEXT_KEY, exc) - status = http_status_for(exc) - detail = ( - f"Invalid field: {RESERVED_CONTEXT_KEY}" - if status == 422 - else f"Unusable {RESERVED_CONTEXT_KEY}: {type(exc).__name__}" - ) - await _send_json(send, status, {"detail": detail, "reason": exc.reason}) - return - - # The joined body and its parsed tree can be tens of MB and are not needed past this point; - # the framework re-buffers and re-parses the replayed messages downstream, so holding these - # through the handler would double peak memory. - del body, parsed, raw_settings - - async def replay() -> dict: - if messages: - return messages.pop(0) - # Buffer drained: proxy the original channel so downstream still observes - # http.disconnect. - return await receive() - - with invocation_envelope(invocation_settings, invocation_context): - await self.app(scope, replay, send) - - -def install_invocation_envelope(app: FastAPI) -> None: - """Install out-of-schema envelope extraction on a FastAPI app. - - Idempotent per app: a host wrapper may install at app construction while a plugin that predates - the wrapper's support still calls this itself, and a double install would buffer and replay the - request body twice. + await self.app(scope, counting_receive, guarded_send) + except ClientDisconnect: + if not exceeded: + raise + except Exception: + # The cut body stream can surface downstream as something other than + # ClientDisconnect; once the cap is the cause, the 413 below is the answer. + if not exceeded: + raise + if exceeded and not response_started: + await _send_json(send, 413, {"detail": "Request body too large"}) + + +def _invoke_routes(app: FastAPI) -> list[APIRoute]: + return [ + r + for r in app.router.routes + if isinstance(r, APIRoute) and r.path == _INVOKE_PATH and "POST" in (r.methods or set()) + ] + + +def install_invocation_envelope(app: FastAPI, max_body_bytes: int = MAX_INVOKE_BODY_BYTES) -> None: + """Install reserved-field binding on a FastAPI app's POST /invoke route. + + Attaches `bind_invocation_envelope` as a dependency of every registered POST /invoke route — + the same insertion `include_router` performs for router-level dependencies — installs the + body-size cap beneath the framework, and registers the failure response shape. Must be called + after the /invoke route is registered; a call that finds none raises rather than leaving the + app silently uncovered. + + Idempotent per app: a host wrapper may install at app construction while a plugin that + predates the wrapper's support still calls this itself, and a double install would resolve + settings twice per request. """ if getattr(app.state, "invocation_envelope_installed", False): return + routes = _invoke_routes(app) + if not routes: + raise RuntimeError( + "install_invocation_envelope must be called after the POST /invoke route is registered" + ) app.state.invocation_envelope_installed = True - app.add_middleware(InvocationEnvelopeMiddleware) + for route in routes: + route.dependant.dependencies.insert( + 0, + get_parameterless_sub_dependant( + depends=Depends(bind_invocation_envelope), path=route.path_format + ), + ) + app.add_middleware(InvokeBodyLimitMiddleware, max_body_bytes=max_body_bytes) + app.add_exception_handler(UnusableInvocationEnvelope, _unusable_envelope_response) def settings_cache_key(invocation_settings: Mapping[str, Any]) -> str: