""" Red-green tests for the forwardedProps -> state.inputs splice. Background — the regression these tests pin down: The previous implementation (`ForwardedPropsMiddleware`, a `BaseHTTPMiddleware` subclass that called `await request.body()` and reinstalled `request._receive`) race-conditioned with Starlette's inner anyio TaskGroup wrapped_receive (see `starlette/middleware/base.py`). Stacking it under `HealthMiddleware` (also `BaseHTTPMiddleware`) caused a `RuntimeError: Unexpected message received: http.request` mid-stream — AG-UI clients saw RUN_STARTED and then the SSE stream aborted (`RUN_ERROR: INCOMPLETE_STREAM`), no TEXT_MESSAGE_* / RUN_FINISHED ever emitted. The fix replaces the racy `BaseHTTPMiddleware` subclass with a *raw* ASGI middleware (`ForwardedPropsASGIMiddleware`) that buffers the body once at the ASGI boundary and replays it via a fresh `receive` callable — no `request._receive` surgery, no inner-TaskGroup race. These tests pin the new contract: 1. POSTs to "/" stream cleanly with HealthMiddleware in the chain — no `Unexpected message received: http.request` RuntimeError, full body delivered to the route handler. 2. `forwardedProps.tone` / `expertise` / `responseLength` still flow through to `state.inputs` (the agent-config demo's contract). 3. Bodies WITHOUT agent-config props are left untouched (non-config demos keep their exact request bytes). 4. /health short-circuits without touching the body splice path. 5. The racy `BaseHTTPMiddleware` shape never reappears in agent_server.py (irreversible structural pin against the regression). """ from __future__ import annotations import json import os import sys import types from typing import Any import pytest from fastapi import FastAPI, Request from fastapi.testclient import TestClient from starlette.middleware.base import BaseHTTPMiddleware from starlette.responses import JSONResponse, StreamingResponse # Force-disable aimock import-time side-effects so we can import agent_server # in a unit-test process without dotenv loading a developer-local OPENAI key. os.environ.setdefault("AIMOCK_URL", "") # --------------------------------------------------------------------------- # # Helper: build a minimal app that replicates the production middleware stack # # (HealthMiddleware on the outside, the forwardedProps splice on the inside) # # plus a streaming route at "/" — the surface where the bug manifested. # # --------------------------------------------------------------------------- # class _HealthMiddleware(BaseHTTPMiddleware): """Mirror of agent_server.HealthMiddleware — same shape, same dispatch.""" async def dispatch(self, request, call_next): if request.url.path == "/health" and request.method == "GET": return JSONResponse({"status": "ok"}) return await call_next(request) # --------------------------------------------------------------------------- # # Inline copy of the splice helpers — matches what's in agent_server.py. # # Keeping a local copy lets the streaming-shape test run without importing # # crewai. The structural-pin test (`test_forwarded_props_middleware_class_.. # # _removed_from_agent_server`) reads the source file directly, and the # # real-module test (`test_real_agent_server_*`) imports agent_server with # # heavy deps stubbed — both keep the production helpers as the source of # # truth. If the helpers below drift, the real-module tests will catch it. # # --------------------------------------------------------------------------- # _AGENT_CONFIG_KEYS = ("tone", "expertise", "responseLength") _TONE_RULES = { "professional": "Use neutral, precise language. No emoji. Short sentences.", "casual": "Use friendly, conversational language. Contractions OK. Light humor welcome.", "enthusiastic": "Use upbeat, energetic language. Exclamation points OK. Emoji OK.", } _EXPERTISE_RULES = { "beginner": "Assume no prior knowledge. Define jargon. Use analogies.", "intermediate": "Assume common terms are understood; explain specialized terms.", "expert": "Assume technical fluency. Use precise terminology. Skip basics.", } _LENGTH_RULES = { "concise": "Respond in 1-3 sentences.", "detailed": "Respond in multiple paragraphs with examples where relevant.", } def _build_agent_config_guidance(tone, expertise, response_length): tone_rule = _TONE_RULES.get(str(tone), _TONE_RULES["professional"]) expertise_rule = _EXPERTISE_RULES.get( str(expertise), _EXPERTISE_RULES["intermediate"] ) length_rule = _LENGTH_RULES.get(str(response_length), _LENGTH_RULES["concise"]) return ( "Follow these style rules for your response to the user. " f"TONE: {tone_rule} " f"EXPERTISE: {expertise_rule} " f"LENGTH: {length_rule}" ) def _has_agent_config_props(props): if not isinstance(props, dict): return False return any(k in props for k in _AGENT_CONFIG_KEYS) def _splice_forwarded_props(body): """Reference splice — what the route handler should call post-fix.""" if not isinstance(body, dict): return body forwarded = body.get("forwardedProps") if not _has_agent_config_props(forwarded): return body # _has_agent_config_props guarantees forwarded is a dict, but narrow for the # type checker without using `assert` (which is stripped under `python -O`). # Mirrors the production-side pattern in agent_server.py. if not isinstance(forwarded, dict): return body existing_state = body.get("state") state: dict[str, Any] = existing_state if isinstance(existing_state, dict) else {} raw_inputs = state.get("inputs") inputs: dict[str, Any] = dict(raw_inputs) if isinstance(raw_inputs, dict) else {} tone = forwarded.get("tone") expertise = forwarded.get("expertise") response_length = forwarded.get("responseLength") if tone is not None: inputs["tone"] = tone if expertise is not None: inputs["expertise"] = expertise if response_length is not None: inputs["response_length"] = response_length inputs["agent_config_guidance"] = _build_agent_config_guidance( tone=tone, expertise=expertise, response_length=response_length ) state["inputs"] = inputs body["state"] = state return body # --------------------------------------------------------------------------- # # Fixture: stub heavy modules + (re)import agent_server for every test that # # needs the real module. Function-scoped + autouse so any of the real-module # # tests can run standalone or in any order (e.g. -k filters, pytest-xdist, # # pytest-randomly). Function scope is REQUIRED — the teardown purges # # `agent_server` from sys.modules so the next test re-binds against fresh # # stubs; module/session scope would only run teardown once per module and # # break ordering independence (sibling tests would share stale module state). # # The stub install is idempotent and the stale `agent_server` entry is # # purged so re-imports pick up our stubs. # # --------------------------------------------------------------------------- # @pytest.fixture(autouse=True) def _stub_agent_server_deps(): """Install stubs for ag_ui_crewai + agents.* so `import agent_server` succeeds without crewai installed, and reset the cached module so each test gets a fresh import bound to OUR stubs (not whatever sibling test happened to import first). """ # ag_ui_crewai.endpoint — provide a no-op `add_crewai_crew_fastapi_endpoint` # that mounts a streaming stub at the requested path. ag_ui_crewai = types.ModuleType("ag_ui_crewai") ag_ui_crewai_endpoint = types.ModuleType("ag_ui_crewai.endpoint") def _add_crewai_crew_fastapi_endpoint(app, crew, path): # Each call defines a fresh local function so closures don't collide # across the multiple `add_crewai_crew_fastapi_endpoint` calls in # agent_server.py. def _make_stub(): async def _stub(request: Request): body = await request.json() if request.method == "POST" else {} async def gen(): yield b"event: RUN_STARTED\ndata: {}\n\n" yield ( b"event: STATE\ndata: " + json.dumps(body.get("state", {})).encode() + b"\n\n" ) yield b"event: RUN_FINISHED\ndata: {}\n\n" return StreamingResponse(gen(), media_type="text/event-stream") return _stub app.post(path)(_make_stub()) setattr( ag_ui_crewai_endpoint, "add_crewai_crew_fastapi_endpoint", _add_crewai_crew_fastapi_endpoint, ) # Mirror crew-endpoint stub for raw `Flow` mounts. The shared-state-rw and # subagents demos use `add_crewai_flow_fastapi_endpoint` to register CrewAI # `Flow`s directly instead of `Crew`s; the test stub treats both the same # since neither needs a real LLM round-trip. setattr( ag_ui_crewai_endpoint, "add_crewai_flow_fastapi_endpoint", _add_crewai_crew_fastapi_endpoint, ) sys.modules["ag_ui_crewai"] = ag_ui_crewai sys.modules["ag_ui_crewai.endpoint"] = ag_ui_crewai_endpoint # Stub agents.* — agent_server imports several crew classes; replace them # with no-arg sentinels so import succeeds without crewai being installed. agents_pkg = types.ModuleType("agents") agents_pkg.__path__ = [] # mark as package sys.modules["agents"] = agents_pkg for name in ( "crew", "a2ui_fixed", "beautiful_chat", "byoc_hashbrown_agent", "byoc_json_render_agent", "declarative_gen_ui", "gen_ui_agent", "interrupt_crew", "mcp_apps_agent", "shared_state_read_write", "subagents", "tool_rendering", "reasoning_agent", "_header_forwarding", ): sys.modules[f"agents.{name}"] = types.ModuleType(f"agents.{name}") setattr(sys.modules["agents.crew"], "LatestAiDevelopment", lambda: object()) setattr(sys.modules["agents.a2ui_fixed"], "A2UIFixedSchema", lambda: object()) setattr(sys.modules["agents.beautiful_chat"], "BeautifulChat", lambda: object()) setattr( sys.modules["agents.byoc_hashbrown_agent"], "ByocHashbrown", lambda: object() ) setattr( sys.modules["agents.byoc_json_render_agent"], "ByocJsonRender", lambda: object() ) setattr( sys.modules["agents.declarative_gen_ui"], "DeclarativeGenUI", lambda: object() ) setattr( sys.modules["agents.interrupt_crew"], "InterruptScheduling", lambda: object() ) setattr(sys.modules["agents.mcp_apps_agent"], "MCPApps", lambda: object()) # Sentinel `Flow` instances for the raw-Flow endpoints (set_notes, # subagents). Real flows are CrewAI `Flow` subclasses; the stub only # needs to be a non-`None` value `add_crewai_flow_fastapi_endpoint` # accepts as its `flow` argument. `object()` works because the stub # endpoint mounter ignores the flow entirely. setattr( sys.modules["agents.shared_state_read_write"], "shared_state_read_write_flow", object(), ) setattr(sys.modules["agents.subagents"], "subagents_flow", object()) setattr(sys.modules["agents.tool_rendering"], "tool_rendering_flow", object()) setattr(sys.modules["agents.gen_ui_agent"], "gen_ui_agent_flow", object()) # agent_server mounts the reasoning sub-app at import time via # `app.mount("/reasoning", reasoning_app)`. `Starlette.mount` only STORES the # app on a Mount route (it isn't invoked unless a request hits /reasoning, # which these tests never do), so a bare sentinel is sufficient here — same # rationale as the Flow sentinels above. setattr(sys.modules["agents.reasoning_agent"], "reasoning_app", object()) # agent_server imports the header-forwarding helpers at top level and calls # install_global_httpx_hook() at import time, then mounts # HeaderForwardingHTTPMiddleware in the stack. Stub both: a no-arg no-op # callable for the hook, and a BaseHTTPMiddleware subclass for the # middleware. The real middleware's dispatch is a transparent pass-through # (copies x-* headers onto a ContextVar, then `return await # call_next(request)`), so the stub mirrors that pass-through — the bare # BaseHTTPMiddleware.dispatch raises NotImplementedError and would break # the real-agent_server streaming tests that exercise the mounted stack. async def _passthrough_dispatch(self, request, call_next): return await call_next(request) setattr( sys.modules["agents._header_forwarding"], "HeaderForwardingHTTPMiddleware", type( "HeaderForwardingHTTPMiddleware", (BaseHTTPMiddleware,), {"dispatch": _passthrough_dispatch}, ), ) setattr( sys.modules["agents._header_forwarding"], "install_global_httpx_hook", lambda: None, ) # Drop any stale agent_server import — we want the next `import agent_server` # to re-run module-init against OUR stubs. sys.modules.pop("agent_server", None) yield # Teardown: leave the stubs in place across tests (cheap), but drop # agent_server so a follow-up test re-binds against fresh stubs. sys.modules.pop("agent_server", None) # --------------------------------------------------------------------------- # # Tests # # --------------------------------------------------------------------------- # def test_streaming_post_with_forwarded_props_completes_without_runtimeerror(): """Stand-in regression shape: HealthMiddleware (BaseHTTPMiddleware) wraps a streaming route that consumes a JSON body and emits SSE frames. The pre-fix `ForwardedPropsMiddleware` (also BaseHTTPMiddleware) would have raced with Starlette's inner TaskGroup here and aborted the stream with `RuntimeError: Unexpected message received: http.request`. Post-fix the splice is a *raw* ASGI middleware (`ForwardedPropsASGIMiddleware` in agent_server.py) that buffers the body at the ASGI boundary and replays it via a fresh `receive` callable — no `BaseHTTPMiddleware` machinery in the body path, no race. The structural pin in `test_forwarded_props_middleware_class_removed_from_agent_server` enforces that the racy shape can never re-enter the source. This test exercises the *post-fix* contract: splice happens via the `_splice_forwarded_props` helper inside the route handler, the streaming response runs to completion under HealthMiddleware, and the spliced state.inputs reaches the handler intact. """ app = FastAPI() @app.post("/") async def root(request: Request): body = await request.json() spliced = _splice_forwarded_props(body) async def gen(): yield b"event: RUN_STARTED\ndata: {}\n\n" yield ( b"event: STATE\ndata: " + json.dumps(spliced.get("state", {})).encode() + b"\n\n" ) yield b"event: RUN_FINISHED\ndata: {}\n\n" return StreamingResponse(gen(), media_type="text/event-stream") app.add_middleware(_HealthMiddleware) payload = { "threadId": "t1", "runId": "r1", "messages": [{"role": "user", "content": "hi"}], "state": {"inputs": {}}, "forwardedProps": {"tone": "casual", "expertise": "beginner"}, } with TestClient(app) as client: response = client.post( "/", json=payload, headers={"content-type": "application/json"}, ) assert response.status_code == 200 text = response.text assert "RUN_STARTED" in text assert "RUN_FINISHED" in text # Splice landed: assert "casual" in text assert "agent_config_guidance" in text def test_forwarded_props_middleware_class_removed_from_agent_server(): """Structural regression pin: the racy `ForwardedPropsMiddleware` (BaseHTTPMiddleware subclass) MUST NOT exist in agent_server.py. The raw-ASGI replacement (`ForwardedPropsASGIMiddleware`) is fine. Pre-fix this fails. Post-fix the BaseHTTPMiddleware-subclass shape is gone. This test is the irreversible RED→GREEN pin — re-introducing the bad pattern (a BaseHTTPMiddleware that does body-replay surgery) trips it immediately. """ here = os.path.dirname(os.path.abspath(__file__)) src = os.path.normpath(os.path.join(here, "..", "..", "src", "agent_server.py")) with open(src) as f: source = f.read() assert "class ForwardedPropsMiddleware(BaseHTTPMiddleware" not in source, ( "ForwardedPropsMiddleware (BaseHTTPMiddleware subclass) must be " "removed — it caused a body-replay race against Starlette's inner " "TaskGroup, aborting AG-UI streams with " "`RuntimeError: Unexpected message received: http.request`. " "Use a raw ASGI middleware instead." ) # Also assert the helpers stay (so the route handler can call them): assert "_build_agent_config_guidance" in source assert "_has_agent_config_props" in source def test_real_agent_server_streams_post_root_without_runtimeerror(): """End-to-end: import the REAL `agent_server` module and POST to '/' to confirm the middleware stack no longer raises the body-replay RuntimeError. Stubs out the heavy crewai-backed route with a streaming one (via the autouse fixture) so the test stays unit-scoped (no LLM, no crewai install needed). """ import agent_server # noqa: F401 payload = { "threadId": "t1", "runId": "r1", "messages": [{"role": "user", "content": "hi"}], "state": {"inputs": {}}, "forwardedProps": {"tone": "casual", "expertise": "beginner"}, } with TestClient(agent_server.app) as client: response = client.post( "/", json=payload, headers={"content-type": "application/json"}, ) assert response.status_code == 200, response.text text = response.text assert "RUN_STARTED" in text assert "RUN_FINISHED" in text # Splice landed in state.inputs (this is the agent-config contract). # Parse the STATE event JSON rather than relying on bare substring # checks — substring checks pass even if the props land in the wrong # place (e.g. echoed in headers, or under a sibling key). state_payload = None for line in text.splitlines(): if line.startswith("data: ") and "inputs" in line: state_payload = json.loads(line[len("data: ") :]) break assert state_payload is not None, ( f"No STATE event with inputs found in: {text!r}" ) assert state_payload["inputs"]["tone"] == "casual" # `agent_config_guidance` expands enums to prose rules — assert the # well-known prefix shape rather than the raw enum value. assert "TONE:" in state_payload["inputs"]["agent_config_guidance"] def test_health_endpoint_short_circuits(): """/health must continue to short-circuit at the middleware layer.""" import agent_server # noqa: F401 with TestClient(agent_server.app) as client: r = client.get("/health") assert r.status_code == 200 assert r.json() == {"status": "ok"} def test_post_without_forwarded_props_passes_through_unchanged(): """Bodies without agent-config props must reach the handler with state unmodified — non-config demos keep their exact request bytes.""" import agent_server # noqa: F401 payload = { "threadId": "t1", "runId": "r1", "messages": [{"role": "user", "content": "hi"}], "state": {"inputs": {"some_demo_key": "preserved"}}, } with TestClient(agent_server.app) as client: response = client.post( "/", json=payload, headers={"content-type": "application/json"}, ) assert response.status_code == 200, response.text # Original state.inputs preserved, no agent_config_guidance injected: assert "preserved" in response.text assert "agent_config_guidance" not in response.text # --------------------------------------------------------------------------- # # Bucket (a) regression tests: ASGI middleware error-handling semantics. # # --------------------------------------------------------------------------- # # # These tests drive `ForwardedPropsASGIMiddleware` directly (bypassing # TestClient) because they need to inject specific receive() event sequences # and exception types that the test client's transport does not emit. def _get_middleware_class(): """Import the real `ForwardedPropsASGIMiddleware` from agent_server. Stubs for `ag_ui_crewai` and `agents.*` are installed by the module-scoped `_stub_agent_server_deps` autouse fixture before each test runs, so a fresh `import agent_server` here binds against those stubs. """ if "agent_server" not in sys.modules: import agent_server # noqa: F401 return sys.modules["agent_server"].ForwardedPropsASGIMiddleware def _build_http_scope(body_bytes: bytes) -> dict: return { "type": "http", "method": "POST", "path": "/", "headers": [ (b"content-type", b"application/json"), (b"content-length", str(len(body_bytes)).encode()), ], } @pytest.mark.anyio async def test_replay_receive_propagates_cancellederror(): """`replay_receive` MUST propagate `asyncio.CancelledError` so the outer task can be cancelled cleanly. The bare `except Exception` was a correctness hazard: even if Python 3.8+ correctly inherits CancelledError from BaseException, the broad except still swallows programming errors (KeyError, AttributeError) that should surface, not silently convert into a clean http.disconnect. Post-fix: CancelledError is caught explicitly and re-raised; any other unexpected exception is logged (observable) rather than silently swallowed. """ import asyncio MW = _get_middleware_class() body = json.dumps({"threadId": "t1"}).encode() async def inner_app(scope, receive, send): msg1 = await receive() assert msg1["type"] == "http.request" # Poll again — this is where CancelledError should propagate up. await receive() mw = MW(inner_app) state = {"calls": 0} async def receive(): state["calls"] += 1 if state["calls"] == 1: return {"type": "http.request", "body": body, "more_body": False} raise asyncio.CancelledError() async def send(_msg): pass scope = _build_http_scope(body) with pytest.raises(asyncio.CancelledError): await mw(scope, receive, send) def test_replay_receive_explicitly_handles_cancellederror(): """Source-level RED-GREEN: assert `replay_receive` catches `asyncio.CancelledError` explicitly (and re-raises) BEFORE the broader except. The bare `except Exception` left CancelledError propagation contingent on Python version and the BaseException hierarchy, which is fragile; the fix makes it explicit and version-independent.""" here = os.path.dirname(os.path.abspath(__file__)) src_path = os.path.normpath( os.path.join(here, "..", "..", "src", "agent_server.py") ) with open(src_path) as f: source = f.read() # Locate the replay_receive function body. marker = "async def replay_receive" idx = source.find(marker) assert idx != -1, "replay_receive function not found in agent_server.py" # Take a window large enough to include its body. window = source[idx : idx + 1500] assert "except asyncio.CancelledError" in window, ( "replay_receive must catch asyncio.CancelledError explicitly and " "re-raise it. A bare `except Exception` is too broad — it can " "swallow programming errors that should surface, and is fragile " "across Python versions where CancelledError's class hierarchy " "changes." ) @pytest.mark.anyio async def test_replay_disconnect_delivers_buffered_chunks_before_disconnect(): """RED pre-fix: when http.disconnect arrives mid-buffer, the middleware only forwards the disconnect — buffered body chunks are silently dropped. GREEN post-fix: the inner ASGI app must observe the buffered chunks (as a single http.request message) BEFORE the http.disconnect, so it sees the partial body that actually arrived. """ MW = _get_middleware_class() chunk1 = b'{"threadId":"t1",' chunk2 = b'"runId":"r1",' # No final chunk — disconnect arrives mid-stream. received: list[dict] = [] async def inner_app(scope, receive, send): # Drain receive until we see a disconnect. while True: msg = await receive() received.append(msg) if msg["type"] == "http.disconnect": return mw = MW(inner_app) state = {"step": 0} async def receive(): state["step"] += 1 if state["step"] == 1: return {"type": "http.request", "body": chunk1, "more_body": True} if state["step"] == 2: return {"type": "http.request", "body": chunk2, "more_body": True} if state["step"] == 3: return {"type": "http.disconnect"} # No further events. return {"type": "http.disconnect"} async def send(_msg): pass # Body bytes hint for content-length is approximate; middleware doesn't # validate it against actual chunks, so any value works. scope = _build_http_scope(chunk1 + chunk2) await mw(scope, receive, send) # The inner app must have seen at least one http.request with the # buffered chunks BEFORE the http.disconnect. request_msgs = [m for m in received if m["type"] == "http.request"] disconnect_msgs = [m for m in received if m["type"] == "http.disconnect"] assert request_msgs, ( "Inner ASGI app did not receive buffered body chunks before disconnect — " "they were silently dropped." ) assert disconnect_msgs, "Inner ASGI app never observed http.disconnect." # Buffered chunks must combine to chunk1+chunk2. combined = b"".join(m.get("body", b"") for m in request_msgs) assert combined == chunk1 + chunk2, ( f"Buffered body mismatch: expected {chunk1 + chunk2!r}, got {combined!r}" ) # Order: at least one request message must precede the first disconnect. first_disconnect_idx = next( i for i, m in enumerate(received) if m["type"] == "http.disconnect" ) first_request_idx = next( (i for i, m in enumerate(received) if m["type"] == "http.request"), None ) assert first_request_idx is not None assert first_request_idx < first_disconnect_idx @pytest.fixture def anyio_backend(): return "asyncio"