[Tests][Refactor] P0.1: Add API server surface guardrails (#6202)

Signed-off-by: herotai214 <herotai214@gmail.com>

tests/entrypoints/openai_api/test_api_server_guards.py

@@ -0,0 +1,810 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""API-server surface guards for the OpenAI Omni entrypoint.
+
+Purpose
+-------
+Freeze the *observable* server contract so helper moves (P0.2), route
+extraction (P0.3), and later endpoint-family PRs cannot silently change:
+
+* which routes exist,
+* how upstream routes are overridden at assembly time,
+* unavailable / missing-handler envelopes,
+* shared engine-error JSON fields,
+* app-state keys wired by ``omni_init_app_state``.
+
+CPU-only (fakes / TestClient). No checked-in OpenAPI golden JSON.
+
+Common wrong changes these tests are meant to catch
+----------------------------------------------------
+* Route accidentally removed, renamed, method-changed, or double-registered
+  after a move (census / assembled-app checks).
+* OpenAPI drops an Omni HTTP path while the router still has it (or the reverse).
+* Assembly forgets to strip upstream ``/v1/chat/completions`` (batch) /
+  ``/v1/models``, or also deletes unrelated upstream routes.
+* Profiler / Omni router not mounted; storage not started; Omni engine
+  exception handlers not registered.
+* Timestamp middleware stops stamping HTTP, or starts mutating WebSocket scopes.
+* Missing-handler WS/HTTP responses change status, JSON shape, or wording
+  (clients treat these as capability probes).
+* Engine error payloads lose ``request_id`` / ``error_stage_id``.
+* Pure-diffusion vs multi-stage init drops an ``app.state`` key, or leaves a
+  handler as ``None`` while the route is still registered.
+
+Not covered here (add elsewhere if needed)
+------------------------------------------
+* Deep per-endpoint success behavior (family UTs).
+* Byte-identical helper bodies after a move (P0.2 AST integrity).
+* Live GPU / speech byte-identity checks.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import json
+from contextlib import asynccontextmanager
+from types import SimpleNamespace
+
+import pytest
+from fastapi import FastAPI, HTTPException
+from fastapi.routing import APIRoute, APIWebSocketRoute
+from fastapi.testclient import TestClient
+from starlette.datastructures import State
+from starlette.requests import Request
+from starlette.websockets import WebSocketDisconnect
+from vllm.v1.engine.exceptions import EngineDeadError, EngineGenerateError
+
+from vllm_omni.entrypoints.openai import api_server
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+# FastAPI/Starlette auto-adds HEAD on every GET route. Including HEAD here
+# would invent ("HEAD", path) entries for every GET and fail the current census.
+# TODO: extend this set and the expected census if @router.head / @router.options appear.
+_HTTP_METHODS = {"GET", "POST", "DELETE", "PUT", "PATCH"}
+
+# Exact Omni router census on main. Update intentionally when routes change.
+_EXPECTED_ROUTER_ROUTES = {
+    ("POST", "/v1/chat/completions"),
+    ("POST", "/v1/chat/completions/batch"),
+    ("POST", "/v1/audio/speech"),
+    ("POST", "/v1/audio/speech/batch"),
+    ("POST", "/v1/audio/generate"),
+    ("GET", "/v1/audio/voices"),
+    ("POST", "/v1/audio/voices"),
+    ("DELETE", "/v1/audio/voices/{name}"),
+    ("WEBSOCKET", "/v1/audio/speech/stream"),
+    ("WEBSOCKET", "/v1/video/chat/stream"),
+    ("WEBSOCKET", "/v1/realtime/video"),
+    ("WEBSOCKET", "/v1/realtime"),
+    ("WEBSOCKET", "/v1/realtime/robot/openpi"),
+    ("WEBSOCKET", "/v1/duplex"),
+    ("GET", "/health"),
+    ("GET", "/v1/models"),
+    ("POST", "/v1/images/generations"),
+    ("POST", "/v1/images/edits"),
+    ("POST", "/v1/videos"),
+    ("POST", "/v1/videos/sync"),
+    ("GET", "/v1/videos"),
+    ("GET", "/v1/videos/{video_id}"),
+    ("DELETE", "/v1/videos/{video_id}"),
+    ("GET", "/v1/videos/{video_id}/content"),
+    ("POST", "/v1/omni/sleep"),
+    ("POST", "/v1/omni/wakeup"),
+}
+_EXPECTED_PROFILER_ROUTES = {
+    ("POST", "/start_profile"),
+    ("POST", "/stop_profile"),
+}
+
+# App-state keys that must remain available across refactor phases.
+# Presence is not enough: "not wired" is written as ``state.X = None``.
+_DIFFUSION_APP_STATE_KEYS = {
+    "engine_client",
+    "log_stats",
+    "args",
+    "sleeping_stages",
+    "stage_configs",
+    "vllm_config",
+    "diffusion_engine",
+    "openai_serving_models",
+    "serving_tokenization",
+    "openai_serving_chat",
+    "openai_serving_chat_batch",
+    "openai_serving_speech",
+    "openai_serving_audio_generate",
+    "openai_serving_video",
+    "openai_streaming_video_output",
+    "openai_serving_duplex",
+    "openai_streaming_speech",
+    "openai_streaming_video",
+    "openai_serving_realtime_robot",
+    "enable_server_load_tracking",
+    "server_load_metrics",
+}
+_DIFFUSION_MUST_BE_NONE = {
+    "vllm_config",
+    "serving_tokenization",
+    "openai_serving_duplex",
+    "openai_streaming_speech",
+    "openai_streaming_video",
+}
+_DIFFUSION_MUST_BE_WIRED = _DIFFUSION_APP_STATE_KEYS - _DIFFUSION_MUST_BE_NONE
+_MULTISTAGE_APP_STATE_KEYS = {
+    "engine_client",
+    "log_stats",
+    "args",
+    "sleeping_stages",
+    "stage_configs",
+    "vllm_config",
+    "openai_serving_models",
+    "serving_tokenization",
+    "online_renderer",
+    "openai_serving_chat",
+    "openai_serving_chat_batch",
+    "openai_serving_speech",
+    "openai_serving_audio_generate",
+    "openai_streaming_speech",
+    "openai_streaming_video",
+    "openai_serving_duplex",
+    "openai_serving_realtime",
+    "openai_serving_video",
+    "openai_serving_realtime_robot",
+    "enable_server_load_tracking",
+    "server_load_metrics",
+}
+_MULTISTAGE_MUST_BE_NONE = {
+    "openai_serving_duplex",
+    "openai_serving_realtime_robot",
+}
+_MULTISTAGE_MUST_BE_WIRED = _MULTISTAGE_APP_STATE_KEYS - _MULTISTAGE_MUST_BE_NONE
+
+
+def _route_entries(routes) -> list[tuple[str, str]]:
+    entries: list[tuple[str, str]] = []
+    for route in routes:
+        if isinstance(route, APIRoute):
+            for method in sorted((route.methods or set()) & _HTTP_METHODS):
+                entries.append((method, route.path))
+        elif isinstance(route, APIWebSocketRoute):
+            entries.append(("WEBSOCKET", route.path))
+    return entries
+
+
+def _route_manifest(routes) -> set[tuple[str, str]]:
+    return set(_route_entries(routes))
+
+
+def _assert_unique_route_entries(routes) -> None:
+    entries = _route_entries(routes)
+    duplicates = sorted({item for item in entries if entries.count(item) > 1})
+    assert not duplicates, f"duplicate route registrations: {duplicates}"
+
+
+def _matching_http_routes(routes, method: str, path: str) -> list[APIRoute]:
+    return [
+        route
+        for route in routes
+        if isinstance(route, APIRoute) and route.path == path and method in (route.methods or set())
+    ]
+
+
+def _single_http_route(routes, method: str, path: str) -> APIRoute:
+    matches = _matching_http_routes(routes, method, path)
+    assert len(matches) == 1, f"expected exactly one {method} {path}, got {len(matches)}"
+    return matches[0]
+
+
+def _assert_app_state_snapshot(
+    state,
+    *,
+    expected_keys: set[str],
+    must_be_wired: set[str],
+    must_be_none: set[str],
+) -> None:
+    assert must_be_wired.isdisjoint(must_be_none)
+    assert must_be_wired | must_be_none <= expected_keys
+    present = {key for key in expected_keys if hasattr(state, key)}
+    assert present == expected_keys
+    not_wired = sorted(key for key in must_be_wired if getattr(state, key) is None)
+    assert not not_wired, f"app.state keys registered but not wired: {not_wired}"
+    unexpectedly_set = sorted(key for key in must_be_none if getattr(state, key) is not None)
+    assert not unexpectedly_set, f"app.state keys expected to stay None: {unexpectedly_set}"
+
+
+def _minimal_args(**overrides) -> SimpleNamespace:
+    args = SimpleNamespace(
+        tool_parser_plugin=None,
+        reasoning_parser_plugin=None,
+        reasoning_parser=None,
+        structured_outputs_config=SimpleNamespace(reasoning_parser=None),
+        enable_ssl_refresh=False,
+        host="127.0.0.1",
+        port=8000,
+        uvicorn_log_level="info",
+        disable_uvicorn_access_log=True,
+        ssl_keyfile=None,
+        ssl_certfile=None,
+        ssl_ca_certs=None,
+        ssl_cert_reqs=0,
+        ssl_ciphers=None,
+        h11_max_incomplete_event_size=None,
+        h11_max_header_count=None,
+        log_config_file=None,
+        disable_access_log_for_endpoints=None,
+        disable_log_stats=True,
+        enable_log_requests=False,
+        max_log_len=None,
+        served_model_name=None,
+        model="demo-model",
+        enable_server_load_tracking=False,
+        log_error_stack=False,
+        chat_template=None,
+        chat_template_content_format="auto",
+        trust_request_chat_template=False,
+        tool_server=None,
+        lora_modules=None,
+        response_role="assistant",
+        default_chat_template_kwargs=None,
+        return_tokens_as_token_ids=False,
+        enable_auto_tool_choice=False,
+        exclude_tools_when_tool_choice_none=False,
+        tool_call_parser=None,
+        enable_prompt_tokens_details=False,
+        enable_force_include_usage=False,
+        enable_log_outputs=False,
+        enable_log_deltas=False,
+        tokens_only=False,
+        shutdown_timeout=0,
+        stage_configs_path=None,
+        deploy_config=None,
+    )
+    for key, value in overrides.items():
+        setattr(args, key, value)
+    return args
+
+
+def _request_for(app: FastAPI, *, method: str = "GET", path: str = "/") -> Request:
+    return Request(
+        {
+            "type": "http",
+            "method": method,
+            "path": path,
+            "headers": [],
+            "app": app,
+            "state": {},
+        }
+    )
+
+
+class _FakeSocket:
+    def __init__(self) -> None:
+        self.closed = False
+
+    def close(self) -> None:
+        self.closed = True
+
+
+class _FakeEngineClient:
+    def __init__(self, *, stage_configs=None, endpoint_restrictions=None, vllm_config=None) -> None:
+        self.stage_configs = stage_configs if stage_configs is not None else []
+        self.endpoint_restrictions = endpoint_restrictions if endpoint_restrictions is not None else {}
+        self.model_config = SimpleNamespace()
+        self.input_processor = object()
+        self.renderer = object()
+        self.errored = False
+        self.vllm_config = vllm_config
+
+    async def get_vllm_config(self):
+        return self.vllm_config
+
+    async def get_supported_tasks(self) -> tuple[str, ...]:
+        return ("generate",)
+
+    async def get_tokenizer(self):
+        return SimpleNamespace(chat_template="dummy")
+
+    async def check_health(self) -> None:
+        return None
+
+
+def _marker(name: str):
+    return SimpleNamespace(name=name)
+
+
+def test_router_route_manifest_is_explicit_and_complete() -> None:
+    """Lock the Omni + profiler router census.
+
+    Fails if a route is added/removed/renamed, a method changes, or a
+    WebSocket path disappears during helper/router moves. Also fails if the
+    same public route is registered more than once, which a set manifest alone
+    would collapse.
+    """
+    # TODO(P0.3): retarget this object-level census when routes move onto family
+    # routers. The assembled-app census is the client-facing lock that must survive.
+    _assert_unique_route_entries(api_server.router.routes)
+    _assert_unique_route_entries(api_server.profiler_router.routes)
+    assert _route_manifest(api_server.router.routes) == _EXPECTED_ROUTER_ROUTES
+    assert _route_manifest(api_server.profiler_router.routes) == _EXPECTED_PROFILER_ROUTES
+
+
+def test_router_openapi_paths_cover_http_manifest() -> None:
+    """Lock that OpenAPI advertises every HTTP census route/method.
+
+    Fails if a decorator still registers a route but OpenAPI loses the path
+    (or method), e.g. after moving route bodies / response models.
+    """
+    app = FastAPI()
+    app.include_router(api_server.router)
+    app.include_router(api_server.profiler_router)
+    paths = app.openapi().get("paths", {})
+
+    http_routes = {
+        (method, path) for method, path in _EXPECTED_ROUTER_ROUTES | _EXPECTED_PROFILER_ROUTES if method != "WEBSOCKET"
+    }
+    for method, path in sorted(http_routes):
+        assert path in paths, f"missing OpenAPI path {path}"
+        assert method.lower() in paths[path], f"missing OpenAPI method {method} {path}"
+
+
+@pytest.mark.asyncio
+async def test_api_server_assembly_replaces_upstream_routes_and_mounts_omni_router(monkeypatch) -> None:
+    """Lock worker assembly: override, mount, storage, handlers, full census.
+
+    Fails if assembly stops replacing upstream chat/batch/models, deletes
+    unrelated upstream routes, forgets Omni/profiler mounts, skips storage
+    start, or drops Omni ``EngineDeadError`` / ``EngineGenerateError`` handlers.
+    """
+    captured: dict[str, object] = {}
+
+    def fake_build_openai_app(args, supported_tasks):
+        app = FastAPI()
+
+        @app.post("/v1/chat/completions")
+        async def upstream_chat():
+            return {"owner": "upstream"}
+
+        @app.post("/v1/chat/completions/batch")
+        async def upstream_batch():
+            return {"owner": "upstream"}
+
+        @app.get("/v1/models")
+        async def upstream_models():
+            return {"owner": "upstream"}
+
+        @app.get("/v1/upstream-only")
+        async def upstream_only():
+            return {"owner": "upstream"}
+
+        captured["supported_tasks"] = supported_tasks
+        return app
+
+    @asynccontextmanager
+    async def fake_build_async_omni(*_args, **_kwargs):
+        yield _FakeEngineClient(
+            stage_configs=[{"engine_args": {"profiler_config": {"profiler": "torch"}}}],
+            vllm_config=SimpleNamespace(parallel_config=SimpleNamespace(_api_process_rank=0)),
+        )
+
+    async def fake_omni_init_app_state(engine_client, state, args):
+        state.engine_client = engine_client
+        state.initialized_by_omni = True
+
+    async def fake_storage_start():
+        captured["storage_started"] = True
+
+    async def fake_serve_http(app, **_kwargs):
+        captured["served_app"] = app
+        task = asyncio.create_task(asyncio.sleep(0))
+        await asyncio.sleep(0)
+        return task
+
+    monkeypatch.setattr(api_server, "build_openai_app", fake_build_openai_app)
+    monkeypatch.setattr(api_server, "build_async_omni", fake_build_async_omni)
+    monkeypatch.setattr(api_server, "omni_init_app_state", fake_omni_init_app_state)
+    monkeypatch.setattr(
+        api_server,
+        "shutdown_unsupported_routes",
+        lambda app, restrictions: captured.setdefault("restrictions", restrictions),
+    )
+    monkeypatch.setattr(api_server, "STORAGE_MANAGER", SimpleNamespace(start=fake_storage_start))
+    monkeypatch.setattr(api_server, "serve_http", fake_serve_http)
+
+    sock = _FakeSocket()
+    await api_server.omni_run_server_worker("127.0.0.1:8000", sock, _minimal_args())
+
+    served_wrapper = captured["served_app"]
+    # Timestamp middleware wraps the FastAPI app.
+    assert hasattr(served_wrapper, "_inner")
+    served_app = served_wrapper._inner
+    routes = served_app.routes
+    _assert_unique_route_entries(routes)
+    manifest = _route_manifest(routes)
+
+    assert captured["supported_tasks"] == ("generate",)
+    assert captured["storage_started"] is True
+    assert captured["restrictions"] == {}
+    assert sock.closed is True
+    assert served_app.state.initialized_by_omni is True
+
+    # Upstream override: unrelated route kept; Omni owns overridden pairs once.
+    # Endpoint identity ensures the surviving route is the Omni handler, not
+    # the fake upstream handler with the same method/path.
+    assert ("GET", "/v1/upstream-only") in manifest
+    assert _single_http_route(routes, "POST", "/v1/chat/completions").endpoint is api_server.create_chat_completion
+    assert (
+        _single_http_route(routes, "POST", "/v1/chat/completions/batch").endpoint
+        is api_server.create_batch_chat_completion
+    )
+    assert _single_http_route(routes, "GET", "/v1/models").endpoint is api_server.show_available_models
+
+    # Assembled app includes Omni router HTTP/WS surface + profiler.
+    for item in _EXPECTED_ROUTER_ROUTES:
+        assert item in manifest, f"assembled app missing {item}"
+    assert ("POST", "/start_profile") in manifest
+    assert ("POST", "/stop_profile") in manifest
+
+    assert served_app.exception_handlers[EngineGenerateError]
+    assert served_app.exception_handlers[EngineDeadError]
+
+    # Light OpenAPI check on the assembled FastAPI app (not the ASGI wrapper).
+    openapi_paths = served_app.openapi().get("paths", {})
+    assert "/v1/chat/completions" in openapi_paths
+    assert "/v1/models" in openapi_paths
+    assert "/health" in openapi_paths
+
+
+@pytest.mark.asyncio
+async def test_timestamp_middleware_stamps_http_and_passes_websocket(monkeypatch) -> None:
+    """Lock outermost timestamp middleware behavior.
+
+    Fails if HTTP requests lose ``request_timestamp``, or WebSocket scopes
+    start getting stamped / mutated incorrectly.
+    """
+    captured: dict[str, object] = {}
+
+    def fake_build_openai_app(args, supported_tasks):
+        return FastAPI()
+
+    @asynccontextmanager
+    async def fake_build_async_omni(*_args, **_kwargs):
+        yield _FakeEngineClient()
+
+    async def fake_omni_init_app_state(engine_client, state, args):
+        state.engine_client = engine_client
+
+    async def fake_serve_http(app, **_kwargs):
+        captured["served_app"] = app
+        task = asyncio.create_task(asyncio.sleep(0))
+        await asyncio.sleep(0)
+        return task
+
+    monkeypatch.setattr(api_server, "build_openai_app", fake_build_openai_app)
+    monkeypatch.setattr(api_server, "build_async_omni", fake_build_async_omni)
+    monkeypatch.setattr(api_server, "omni_init_app_state", fake_omni_init_app_state)
+    monkeypatch.setattr(api_server, "shutdown_unsupported_routes", lambda *_a, **_k: None)
+    monkeypatch.setattr(api_server, "STORAGE_MANAGER", SimpleNamespace(start=lambda: asyncio.sleep(0)))
+    monkeypatch.setattr(api_server, "serve_http", fake_serve_http)
+
+    await api_server.omni_run_server_worker("127.0.0.1:8000", _FakeSocket(), _minimal_args())
+    middleware = captured["served_app"]
+
+    http_scope = {"type": "http", "state": {}}
+    ws_scope = {"type": "websocket"}
+
+    async def _receive():
+        return {"type": "http.disconnect"}
+
+    async def _send(_message):
+        return None
+
+    called = {"http": False, "websocket": False}
+
+    async def inner(scope, receive, send):
+        called[scope["type"]] = True
+
+    middleware._inner = inner
+    await middleware(http_scope, _receive, _send)
+    await middleware(ws_scope, _receive, _send)
+
+    assert called["http"] is True
+    assert called["websocket"] is True
+    assert "request_timestamp" in http_scope["state"]
+    assert isinstance(http_scope["state"]["request_timestamp"], float)
+    assert "request_timestamp" not in ws_scope.get("state", {})
+
+
+@pytest.mark.parametrize(
+    ("path", "payload"),
+    [
+        ("/v1/audio/speech/stream", {"type": "error", "message": "Streaming speech is not available"}),
+        ("/v1/video/chat/stream", {"type": "error", "message": "Streaming video chat is not available"}),
+        ("/v1/realtime/video", {"type": "error", "message": "Streaming video generation is not available"}),
+        ("/v1/realtime", {"type": "error", "error": "Realtime API is not available", "code": "unsupported"}),
+        (
+            "/v1/realtime/robot/openpi",
+            {"type": "error", "error": "Robot policy not available", "code": "unsupported"},
+        ),
+        ("/v1/duplex", {"type": "error", "error": "Duplex API is not available", "code": "unsupported"}),
+    ],
+)
+def test_websocket_routes_emit_stable_unavailable_frames_and_close(path: str, payload: dict[str, str]) -> None:
+    """Lock missing-handler WebSocket capability probes.
+
+    Fails if unavailable JSON changes (keys/wording/code) or the socket stays
+    open after the error frame — common when WS wrappers move between modules.
+    """
+    app = FastAPI()
+    app.include_router(api_server.router)
+
+    with TestClient(app) as client:
+        with client.websocket_connect(path) as websocket:
+            assert websocket.receive_json() == payload
+            with pytest.raises(WebSocketDisconnect):
+                websocket.receive_text()
+
+
+def test_health_without_engine_returns_stable_unhealthy_response() -> None:
+    """Lock ``/health`` with no engine initialized.
+
+    Fails if status drifts from 503 or the unhealthy payload shape/reason
+    string changes during health-route extraction.
+    """
+    app = FastAPI()
+    response = asyncio.run(api_server.health(_request_for(app, path="/health")))
+
+    assert response.status_code == 503
+    assert json.loads(response.body) == {
+        "status": "unhealthy",
+        "reason": "No engine initialized",
+    }
+
+
+def test_models_without_handler_returns_empty_openai_list() -> None:
+    """Lock ``/v1/models`` when ``openai_serving_models`` is unset.
+
+    Fails if the empty listing becomes an error, or the OpenAI list envelope
+    changes (``object`` / ``data``).
+    """
+    app = FastAPI()
+    response = asyncio.run(api_server.show_available_models(_request_for(app, path="/v1/models")))
+
+    assert response.status_code == 200
+    assert json.loads(response.body) == {"object": "list", "data": []}
+
+
+def test_images_generation_without_engine_preserves_service_unavailable_error() -> None:
+    """Lock images generation when the multi-stage engine is missing.
+
+    Fails if the 503 detail string changes — clients/tests use it as a
+    capability signal when bootstrap failed.
+    """
+    app = FastAPI()
+    raw_request = _request_for(app, method="POST", path="/v1/images/generations")
+    request = api_server.ImageGenerationRequest(prompt="a cat", model="demo-model")
+
+    with pytest.raises(HTTPException) as exc_info:
+        asyncio.run(api_server.generate_images(request, raw_request))
+
+    assert exc_info.value.status_code == 503
+    assert exc_info.value.detail == "Multi-stage engine not initialized. Start server with a multi-stage omni model."
+
+
+def test_images_generation_without_multistage_chat_handler_preserves_unavailable_error(monkeypatch) -> None:
+    """Lock images generation when chat serving was not wired.
+
+    Fails if multi-stage images stop requiring ``openai_serving_chat``, or the
+    503 detail changes after images/chat ownership splits.
+    """
+    app = FastAPI()
+    app.state.engine_client = SimpleNamespace(stage_configs=[object(), object()])
+    app.state.stage_configs = app.state.engine_client.stage_configs
+    app.state.openai_serving_models = SimpleNamespace(base_model_paths=[SimpleNamespace(name="demo-model")])
+    app.state.openai_serving_chat = None
+    app.state.args = SimpleNamespace(max_generated_image_size=None)
+
+    raw_request = _request_for(app, method="POST", path="/v1/images/generations")
+    request = api_server.ImageGenerationRequest(prompt="a cat", model="demo-model")
+    monkeypatch.setattr(api_server, "get_stage_type", lambda _stage_cfg: "diffusion")
+
+    with pytest.raises(HTTPException) as exc_info:
+        asyncio.run(api_server.generate_images(request, raw_request))
+
+    assert exc_info.value.status_code == 503
+    assert exc_info.value.detail == "openai_serving_chat is not initialized for multi-stage image generation."
+
+
+def test_speech_without_handler_preserves_not_found_http_error() -> None:
+    """Lock speech HTTP when ``openai_serving_speech`` is unset.
+
+    Fails if missing speech becomes a different status/detail after TTS route
+    wrappers move out of ``api_server``.
+    """
+    app = FastAPI()
+    app.state.openai_serving_speech = None
+    raw_request = _request_for(app, method="POST", path="/v1/audio/speech")
+    request = api_server.OpenAICreateSpeechRequest(input="hello", model="demo-model", voice="alloy")
+
+    with pytest.raises(HTTPException) as exc_info:
+        asyncio.run(api_server.create_speech(request, raw_request))
+
+    assert exc_info.value.status_code == 404
+    assert exc_info.value.detail == "The model does not support Speech API"
+
+
+def test_engine_error_json_response_includes_request_and_stage_fields(monkeypatch) -> None:
+    """Lock shared ``EngineGenerateError`` JSON fields used by Omni handlers.
+
+    Fails if ``request_id`` / ``error_stage_id`` are dropped when error helpers
+    move (P0.2) or exception-handler wiring changes.
+
+    Exercises the envelope through the registered handler, not the private
+    ``_create_engine_error_json_response`` helper, so relocating that helper
+    in P0.2 is not a false red.
+    """
+    app = FastAPI()
+    # TODO(P0.2): retarget this call if ``_register_omni_exception_handlers``
+    # leaves api_server (red here is a rename, not a surface regression).
+    api_server._register_omni_exception_handlers(app)
+    app.state.engine_client = SimpleNamespace(errored=True, engine=SimpleNamespace(is_alive=lambda: False))
+    app.state.server = object()
+    app.state.args = SimpleNamespace(log_error_stack=False)
+
+    req = _request_for(app, method="POST", path="/v1/chat/completions")
+    req.state.request_metadata = SimpleNamespace(request_id="req-123")
+
+    monkeypatch.setattr(api_server, "terminate_if_errored", lambda **_kwargs: None)
+
+    exc = EngineGenerateError("boom")
+    exc.error_stage_id = 2  # type: ignore[attr-defined]
+    handler = app.exception_handlers[EngineGenerateError]
+    response = asyncio.run(handler(req, exc))
+    payload = json.loads(response.body)
+
+    assert "error" in payload
+    assert payload["error"]["request_id"] == "req-123"
+    assert payload["error"]["error_stage_id"] == 2
+    assert response.status_code >= 400
+
+
+def test_engine_dead_error_handler_registered_returns_json(monkeypatch) -> None:
+    """Lock Omni ``EngineDeadError`` handler registration + JSON envelope.
+
+    Fails if registration is skipped or the handler stops returning OpenAI-style
+    JSON with ``request_id``.
+    """
+    app = FastAPI()
+    # TODO(P0.2): retarget this call if ``_register_omni_exception_handlers``
+    # leaves api_server (red here is a rename, not a surface regression).
+    api_server._register_omni_exception_handlers(app)
+    app.state.engine_client = SimpleNamespace(errored=True, engine=SimpleNamespace(is_alive=lambda: False))
+    app.state.server = object()
+    app.state.args = SimpleNamespace(log_error_stack=False)
+
+    monkeypatch.setattr(api_server, "terminate_if_errored", lambda **_kwargs: None)
+
+    handler = app.exception_handlers[EngineDeadError]
+    req = _request_for(app)
+    req.state.request_metadata = SimpleNamespace(request_id="dead-1")
+    response = asyncio.run(handler(req, EngineDeadError()))
+    payload = json.loads(response.body)
+    assert payload["error"]["request_id"] == "dead-1"
+    assert response.status_code >= 400
+
+
+@pytest.mark.asyncio
+async def test_pure_diffusion_app_state_key_snapshot(monkeypatch) -> None:
+    """Lock pure-diffusion ``app.state`` keys after init, including live vs None.
+
+    Fails if bootstrap drops keys route owners still read (e.g. video/speech/
+    models), wires a handler that must stay None, or leaves a required
+    handler as None.
+    """
+    stage = SimpleNamespace(engine_args={})
+    engine = _FakeEngineClient(stage_configs=[stage])
+
+    def _for_diffusion_factory(label: str):
+        @classmethod
+        def _factory(cls, *args, **kwargs):
+            return _marker(label)
+
+        return _factory
+
+    monkeypatch.setattr(api_server, "get_stage_type", lambda _cfg: "diffusion")
+    monkeypatch.setattr(api_server.OmniOpenAIServingChat, "for_diffusion", _for_diffusion_factory("chat"))
+    monkeypatch.setattr(api_server.OmniOpenAIServingChatBatch, "for_diffusion", _for_diffusion_factory("chat_batch"))
+    monkeypatch.setattr(
+        api_server.OmniOpenAIServingAudioGenerate,
+        "for_diffusion",
+        _for_diffusion_factory("audio_generate"),
+    )
+    monkeypatch.setattr(api_server.OmniOpenAIServingVideo, "for_diffusion", _for_diffusion_factory("video"))
+    monkeypatch.setattr(api_server.OmniStreamingVideoOutputHandler, "__init__", lambda self, *a, **k: None)
+    monkeypatch.setattr(api_server.OmniOpenAIServingSpeech, "for_diffusion", _for_diffusion_factory("speech"))
+    monkeypatch.setattr(
+        api_server.ServingRealtimeRobotOpenPI,
+        "create_policy_server",
+        classmethod(lambda cls, *a, **k: _marker("openpi")),
+    )
+
+    state = State()
+    await api_server.omni_init_app_state(engine, state, _minimal_args())
+
+    _assert_app_state_snapshot(
+        state,
+        expected_keys=_DIFFUSION_APP_STATE_KEYS,
+        must_be_wired=_DIFFUSION_MUST_BE_WIRED,
+        must_be_none=_DIFFUSION_MUST_BE_NONE,
+    )
+    assert state.diffusion_engine is engine
+
+
+@pytest.mark.asyncio
+async def test_multistage_app_state_key_snapshot(monkeypatch) -> None:
+    """Lock multi-stage ``app.state`` keys after init, including live vs None.
+
+    Fails if chat/speech/video/realtime/tokenization keys disappear or are
+    left as None after init refactors — classic “route still registered,
+    handler never wired” failure mode.
+    """
+    engine = _FakeEngineClient(
+        stage_configs=[object(), object()],
+        vllm_config=SimpleNamespace(
+            lora_config=None,
+            model_config=SimpleNamespace(),
+            parallel_config=SimpleNamespace(_api_process_rank=0),
+        ),
+    )
+
+    class _FakeModels:
+        def __init__(self, *args, **kwargs):
+            self.base_model_paths = kwargs.get("base_model_paths") or []
+
+        async def init_static_loras(self):
+            return None
+
+    class _FakeCtor:
+        def __init__(self, *args, **kwargs):
+            pass
+
+        def warmup(self):
+            return None
+
+    class _FakeSpeech(_FakeCtor):
+        async def warmup(self):
+            return None
+
+    monkeypatch.setattr(api_server, "load_chat_template", lambda *_a, **_k: None)
+    monkeypatch.setattr(api_server, "process_lora_modules", lambda modules, _defaults: modules or [])
+    monkeypatch.setattr(api_server, "OpenAIServingModels", _FakeModels)
+    monkeypatch.setattr(api_server, "OnlineRenderer", _FakeCtor)
+    monkeypatch.setattr(api_server, "OpenAIServingResponses", _FakeCtor)
+    monkeypatch.setattr(api_server, "OmniOpenAIServingChat", _FakeCtor)
+    monkeypatch.setattr(api_server, "OmniOpenAIServingChatBatch", _FakeCtor)
+    monkeypatch.setattr(api_server, "OpenAIServingCompletion", _FakeCtor)
+    monkeypatch.setattr(api_server, "ServingPooling", _FakeCtor)
+    monkeypatch.setattr(api_server, "OpenAIServingEmbedding", _FakeCtor)
+    monkeypatch.setattr(api_server, "ServingClassification", _FakeCtor)
+    monkeypatch.setattr(api_server, "ServingScores", _FakeCtor)
+    monkeypatch.setattr(api_server, "ServingTokenization", _FakeCtor)
+    monkeypatch.setattr(api_server, "OpenAIServingTranscription", _FakeCtor)
+    monkeypatch.setattr(api_server, "OpenAIServingTranslation", _FakeCtor)
+    monkeypatch.setattr(api_server, "AnthropicServingMessages", _FakeCtor)
+    monkeypatch.setattr(api_server, "ServingTokens", _FakeCtor)
+    monkeypatch.setattr(api_server, "OmniOpenAIServingSpeech", _FakeSpeech)
+    monkeypatch.setattr(api_server, "OmniOpenAIServingAudioGenerate", _FakeCtor)
+    monkeypatch.setattr(api_server, "OmniStreamingSpeechHandler", _FakeCtor)
+    monkeypatch.setattr(api_server, "create_streaming_video_handler", lambda **_k: _marker("streaming_video"))
+    monkeypatch.setattr(api_server, "OpenAIServingRealtime", _FakeCtor)
+    monkeypatch.setattr(api_server, "OmniOpenAIServingVideo", _FakeCtor)
+    monkeypatch.setattr(api_server, "should_enable_duplex_endpoint", lambda *_a, **_k: False)
+
+    state = State()
+    await api_server.omni_init_app_state(engine, state, _minimal_args())
+
+    _assert_app_state_snapshot(
+        state,
+        expected_keys=_MULTISTAGE_APP_STATE_KEYS,
+        must_be_wired=_MULTISTAGE_MUST_BE_WIRED,
+        must_be_none=_MULTISTAGE_MUST_BE_NONE,
+    )