test(integration): audit cells for passthrough error relay across routes, disconnect burst and unchanged llm endpoint normalization

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

tests/integration/observability/test_passthrough_upstream_error_chaos.py

@@ -2,12 +2,12 @@ import asyncio
 import json
 import re
 import signal
+import threading
 from pathlib import Path
 from typing import Final
 
 import httpx
 import psutil
-import pytest
 import yaml
 from integration._support.client import Gateway, eventually, object_value
 from integration._support.database import read_rows
@@ -148,3 +148,61 @@ async def test_passthrough_worker_sigkill_leaves_sibling_serving_and_logging(gat
                     _single_spend_row(response.headers["x-litellm-call-id"])
             error_information: Final = _error_information(follow_up.headers["x-litellm-call-id"])
             assert "not found for this scripted upstream" in str(error_information["error_message"]), follow_up.text
+
+
+_RATE_LIMITED_FRAMES: Final = (b'data: {"error":"rate limited"}\n\n', b"data: [DONE]\n\n")
+
+
+async def test_passthrough_disconnect_burst_logs_every_failure_once(gateway: Gateway, tmp_path: Path) -> None:
+    gate: Final = threading.Event()
+
+    def respond(request: Request) -> Reply:
+        if "streamGenerateContent" in request.target:
+            return Reply(
+                status=429, content_type="text/event-stream", chunks=_RATE_LIMITED_FRAMES, gate_after_first=gate
+            )
+        return Reply(status=200, body=json.dumps({"ok": True}).encode())
+
+    config: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text())
+    path: Final = tmp_path / "chaos-disconnect-burst.yaml"
+    with wire_server(respond) as wire:
+        config["environment_variables"] = {"GEMINI_API_BASE": wire.url, "GEMINI_API_KEY": "scripted"}
+        path.write_text(yaml.safe_dump(config))
+        with owned_proxy_process(gateway, tmp_path, {}, config=path, workers=2) as owned:
+            candidate: Final = owned.gateway
+            call_ids: Final = await asyncio.gather(
+                *(_first_frame_then_close(str(candidate.client.base_url), candidate.key) for _ in range(20))
+            )
+            assert len(set(call_ids)) == 20, call_ids
+            gate.set()
+            for call_id in call_ids:
+                _single_spend_row(call_id)
+                error_information: Final = _error_information(call_id)
+                assert error_information["error_code"] == "429", error_information
+            eventually(
+                lambda: tuple(line for line in owned.log.read_text().splitlines() if "returned 429" in line),
+                lambda lines: len(lines) == 20,
+                seconds=60,
+            )
+            follow_up: Final = candidate.request(
+                "POST",
+                "/gemini/v1beta/models/nope-9:generateContent",
+                _GENERATE_CONTENT,
+                headers={"x-goog-api-key": candidate.key},
+            )
+            assert follow_up.status_code == 200, follow_up.text
+
+
+async def _first_frame_then_close(base_url: str, key: str) -> str:
+    async with httpx.AsyncClient(base_url=base_url, timeout=httpx.Timeout(5, connect=5), trust_env=False) as client:
+        async with client.stream(
+            "POST",
+            "/gemini/v1beta/models/nope-9:streamGenerateContent",
+            params={"alt": "sse"},
+            json=_GENERATE_CONTENT,
+            headers={"Authorization": f"Bearer {key}", "x-goog-api-key": key},
+        ) as response:
+            assert response.status_code == 429, response.status_code
+            first: Final = await response.aiter_bytes().__anext__()
+            assert first.startswith(b'data: {"error":"rate limited"}'), first
+            return response.headers["x-litellm-call-id"]

tests/integration/observability/test_passthrough_upstream_error_visibility.py

@@ -1,6 +1,9 @@
 import gzip
 import json
 import threading
+import time
+import uuid
+from collections.abc import Callable
 from hashlib import sha256
 from pathlib import Path
 from typing import Final
@@ -11,7 +14,9 @@ import yaml
 from integration._support.client import Gateway, eventually, object_value
 from integration._support.database import read_rows
 from integration._support.process import owned_proxy_process
+from integration._support.upstream import delete_scenario, register_scenario
 from integration._support.wire import Reply, Request, wire_server
+from integration.cost_calculation.cost_tracking_case import JsonResponse
 from openai import AsyncOpenAI, NotFoundError, OpenAI
 from pydantic import JsonValue
 
@@ -720,3 +725,365 @@ def test_gemini_passthrough_streaming_429_client_disconnect_still_logs_failure(
             assert error_information["normalized_error"] == "500_UPSTREAM_PASSTHROUGH", error_information
             warnings: Final = _upstream_warnings(owned.log, "returned 429")
             assert len(warnings) == 1, warnings
+
+
+async def test_gemini_passthrough_async_streaming_429_first_frame_reaches_client_while_upstream_holds(
+    gateway: Gateway, tmp_path: Path
+) -> None:
+    gate: Final = threading.Event()
+    frames: Final = (b'data: {"error":"rate limited"}\n\n', b"data: [DONE]\n\n")
+
+    def respond(request: Request) -> Reply:
+        return Reply(status=429, content_type="text/event-stream", chunks=frames, gate_after_first=gate)
+
+    path: Final = tmp_path / "gemini-async-stream-429.yaml"
+    with wire_server(respond) as wire:
+        _gemini_config(path, wire.url)
+        with owned_proxy_process(gateway, tmp_path, {}, config=path, workers=2) as owned:
+            candidate: Final = owned.gateway
+            try:
+                async with httpx.AsyncClient(
+                    base_url=str(candidate.client.base_url), timeout=httpx.Timeout(3, connect=5), trust_env=False
+                ) as async_client:
+                    async with async_client.stream(
+                        "POST",
+                        _GEMINI_STREAM_PATH,
+                        params={"alt": "sse"},
+                        json=_GENERATE_CONTENT,
+                        headers=_gemini_headers(candidate),
+                    ) as response:
+                        assert response.status_code == 429, response.text
+                        iterator = response.aiter_bytes()
+                        first: Final = await iterator.__anext__()
+                        assert first.startswith(b'data: {"error":"rate limited"}'), first
+                        gate.set()
+                        rest: Final = b""
+                        async for chunk in iterator:
+                            rest += chunk
+                    assert first + rest == b"".join(frames), first + rest
+            finally:
+                gate.set()
+            warning: Final = _upstream_warning(owned.log)
+            assert "rate limited" in warning, warning
+            error_information: Final = _spend_error_information(response.headers["x-litellm-call-id"])
+            assert error_information["error_code"] == "429", error_information
+
+
+_PROVIDER_STREAM_SPECS: Final = (
+    (
+        "anthropic",
+        "/anthropic/v1/messages",
+        {"ANTHROPIC_API_BASE": None, "ANTHROPIC_API_KEY": "scripted"},
+        "/v1/messages",
+    ),
+    (
+        "cohere",
+        "/cohere/v2/chat",
+        {"COHERE_API_BASE": None, "COHERE_API_KEY": "scripted"},
+        "/v2/chat",
+    ),
+    (
+        "mistral",
+        "/mistral/v1/chat/completions",
+        {"MISTRAL_API_BASE": None, "MISTRAL_API_KEY": "scripted"},
+        "/v1/chat/completions",
+    ),
+    (
+        "openai",
+        "/openai/v1/chat/completions",
+        {"OPENAI_API_BASE": None, "OPENAI_API_KEY": "scripted"},
+        "/v1/chat/completions",
+    ),
+    (
+        "azure",
+        "/azure/openai/deployments/nope-9/chat/completions",
+        {"AZURE_API_BASE": None, "AZURE_API_KEY": "scripted"},
+        "/openai/deployments/nope-9/chat/completions",
+    ),
+    ("config-route", "/audit-pt", None, "/upstream"),
+)
+
+
+@pytest.mark.parametrize(
+    ("name", "proxy_path", "environment", "upstream_suffix"),
+    _PROVIDER_STREAM_SPECS,
+    ids=tuple(spec[0] for spec in _PROVIDER_STREAM_SPECS),
+)
+def test_provider_passthrough_streaming_429_first_frame_reaches_client_while_upstream_holds(
+    gateway: Gateway,
+    tmp_path: Path,
+    name: str,
+    proxy_path: str,
+    environment: dict[str, str | None] | None,
+    upstream_suffix: str,
+) -> None:
+    gate: Final = threading.Event()
+    frames: Final = (b'data: {"error":"rate limited"}\n\n', b"data: [DONE]\n\n")
+
+    def respond(request: Request) -> Reply:
+        return Reply(status=429, content_type="text/event-stream", chunks=frames, gate_after_first=gate)
+
+    config: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text())
+    path: Final = tmp_path / f"{name}-stream-429.yaml"
+    with wire_server(respond) as wire:
+        if environment is None:
+            config["general_settings"]["pass_through_endpoints"] = [
+                {"path": proxy_path, "target": f"{wire.url}{upstream_suffix}", "include_subpath": True}
+            ]
+        else:
+            config["environment_variables"] = {
+                key: (wire.url if value is None else value) for key, value in environment.items()
+            }
+        path.write_text(yaml.safe_dump(config))
+        with owned_proxy_process(gateway, tmp_path, {}, config=path, workers=2) as owned:
+            candidate: Final = owned.gateway
+            with candidate.client.stream(
+                "POST",
+                proxy_path,
+                json={"model": "nope-9", "messages": [{"role": "user", "content": "hi"}], "stream": True},
+                headers={"Authorization": f"Bearer {candidate.key}"},
+                timeout=httpx.Timeout(3, connect=5),
+            ) as response:
+                assert response.status_code == 429, response.text
+                iterator: Final = response.iter_bytes()
+                first: Final = next(iterator)
+                assert first.startswith(b'data: {"error":"rate limited"}'), first
+                gate.set()
+                rest: Final = b"".join(iterator)
+            assert first + rest == b"".join(frames), first + rest
+            received: Final = wire.drain()
+            assert any(request.target.endswith(upstream_suffix) for request in received), received
+            warning: Final = _upstream_warning(owned.log)
+            assert "rate limited" in warning, warning
+            error_information: Final = _spend_error_information(response.headers["x-litellm-call-id"])
+            assert error_information["error_code"] == "429", error_information
+
+
+async def test_gemini_passthrough_async_quota_wording_keeps_passthrough_normalized_error(
+    gateway: Gateway, tmp_path: Path
+) -> None:
+    body: Final[dict[str, JsonValue]] = {
+        "error": {
+            "message": "You exceeded your current quota, please check your plan and billing details. "
+            f"Budget for Key={_LEAKED_UPSTREAM_KEY} is spent",
+            "type": "insufficient_quota",
+        }
+    }
+
+    def respond(request: Request) -> Reply:
+        return Reply(status=429, body=json.dumps(body).encode())
+
+    path: Final = tmp_path / "gemini-async-quota-429.yaml"
+    with wire_server(respond) as wire:
+        _gemini_config(path, wire.url)
+        with owned_proxy_process(gateway, tmp_path, {}, config=path, workers=2) as owned:
+            candidate: Final = owned.gateway
+            async with httpx.AsyncClient(
+                base_url=str(candidate.client.base_url), timeout=15, trust_env=False
+            ) as async_client:
+                response: Final = await async_client.post(
+                    _GEMINI_MODEL_PATH, json=_GENERATE_CONTENT, headers=_gemini_headers(candidate)
+                )
+            assert response.status_code == 429, response.text
+            assert response.json() == body, response.text
+            error_information: Final = _spend_error_information(response.headers["x-litellm-call-id"])
+            assert error_information["normalized_error"] == "500_UPSTREAM_PASSTHROUGH", error_information
+            assert _LEAKED_UPSTREAM_KEY not in str(error_information["error_message"]), error_information
+            warning: Final = _upstream_warning(owned.log)
+            assert _LEAKED_UPSTREAM_KEY not in warning, warning
+
+
+def test_gemini_passthrough_near_miss_pattern_body_keeps_liveliness_and_bounded_preview(
+    gateway: Gateway, tmp_path: Path
+) -> None:
+    body: Final = (b"Upstream passthrough request faile " * 1873)[: 64 * 1024]
+    assert len(body) == 64 * 1024
+
+    def respond(request: Request) -> Reply:
+        return Reply(
+            status=429,
+            content_type="text/event-stream",
+            chunks=tuple(body[index : index + 512] for index in range(0, len(body), 512)),
+            pause_between_chunks=0.01,
+        )
+
+    liveliness_seconds: list[float] = []
+    liveliness_status: list[int] = []
+
+    def probe_liveliness(client: httpx.Client) -> None:
+        started: Final = time.monotonic()
+        probe: Final = client.get("/health/liveliness")
+        liveliness_seconds.append(time.monotonic() - started)
+        liveliness_status.append(probe.status_code)
+
+    path: Final = tmp_path / "gemini-near-miss-64k.yaml"
+    with wire_server(respond) as wire:
+        _gemini_config(path, wire.url)
+        with owned_proxy_process(gateway, tmp_path, {}, config=path, workers=2) as owned:
+            candidate: Final = owned.gateway
+            probe_thread: Final = threading.Thread(target=probe_liveliness, args=(candidate.client,))
+            with candidate.client.stream(
+                "POST",
+                _GEMINI_STREAM_PATH,
+                params={"alt": "sse"},
+                json=_GENERATE_CONTENT,
+                headers=_gemini_headers(candidate),
+                timeout=httpx.Timeout(30, connect=5),
+            ) as response:
+                assert response.status_code == 429, response.text
+                iterator: Final = response.iter_bytes()
+                first: Final = next(iterator)
+                probe_thread.start()
+                rest: Final = b"".join(iterator)
+            streamed: Final = first + rest
+            probe_thread.join(timeout=10)
+            assert sha256(streamed).hexdigest() == sha256(body).hexdigest(), "relayed body differs"
+            assert liveliness_status == [200], liveliness_status
+            assert liveliness_seconds[0] < 2, liveliness_seconds
+            warning: Final = _upstream_warning(owned.log)
+            assert warning.endswith("... (truncated at 4096 chars)"), warning
+            error_information: Final = _spend_error_information(response.headers["x-litellm-call-id"])
+            assert error_information["error_code"] == "429", error_information
+            assert error_information["normalized_error"] == "500_UPSTREAM_PASSTHROUGH", error_information
+
+
+def test_gemini_passthrough_streaming_429_upstream_abort_after_first_frame_still_logs_once(
+    gateway: Gateway, tmp_path: Path
+) -> None:
+    frames: Final = (b'data: {"error":"rate limited"}\n\n', b"data: second\n\n", b"data: [DONE]\n\n")
+
+    def respond(request: Request) -> Reply:
+        if "streamGenerateContent" in request.target:
+            return Reply(status=429, content_type="text/event-stream", chunks=frames, abort_after=1)
+        return Reply(status=200, body=json.dumps({"ok": True}).encode())
+
+    path: Final = tmp_path / "gemini-stream-429-abort.yaml"
+    with wire_server(respond) as wire:
+        _gemini_config(path, wire.url)
+        with owned_proxy_process(gateway, tmp_path, {}, config=path, workers=2) as owned:
+            candidate: Final = owned.gateway
+            received: Final = bytearray()
+            call_ids: Final[list[str]] = []
+            try:
+                with candidate.client.stream(
+                    "POST",
+                    _GEMINI_STREAM_PATH,
+                    params={"alt": "sse"},
+                    json=_GENERATE_CONTENT,
+                    headers=_gemini_headers(candidate),
+                    timeout=httpx.Timeout(15, connect=5),
+                ) as response:
+                    assert response.status_code == 429, response.text
+                    call_ids.append(response.headers["x-litellm-call-id"])
+                    for chunk in response.iter_bytes():
+                        received += chunk
+            except httpx.HTTPError:
+                pass
+            assert len(call_ids) == 1, call_ids
+            assert bytes(received).startswith(b'data: {"error":"rate limited"}'), bytes(received)
+            eventually(
+                lambda: _upstream_warnings(owned.log),
+                lambda lines: any("returned 429" in line or "read failed" in line for line in lines),
+                seconds=30,
+            )
+            matching: Final = tuple(
+                line for line in _upstream_warnings(owned.log) if "returned 429" in line or "read failed" in line
+            )
+            assert len(matching) == 1, matching
+            error_information: Final = _spend_error_information(call_ids[0])
+            assert error_information["error_code"] == "429", error_information
+            follow_up: Final = candidate.request(
+                "POST", _GEMINI_MODEL_PATH, _GENERATE_CONTENT, headers=_gemini_headers(candidate)
+            )
+            assert follow_up.status_code == 200, follow_up.text
+
+
+def test_gemini_passthrough_empty_streaming_429_still_logged(gateway: Gateway, tmp_path: Path) -> None:
+    def respond(request: Request) -> Reply:
+        return Reply(status=429, content_type="text/event-stream", chunks=())
+
+    path: Final = tmp_path / "gemini-stream-429-empty.yaml"
+    with wire_server(respond) as wire:
+        _gemini_config(path, wire.url)
+        with owned_proxy_process(gateway, tmp_path, {}, config=path, workers=2) as owned:
+            candidate: Final = owned.gateway
+            with candidate.client.stream(
+                "POST",
+                _GEMINI_STREAM_PATH,
+                params={"alt": "sse"},
+                json=_GENERATE_CONTENT,
+                headers=_gemini_headers(candidate),
+            ) as response:
+                assert response.status_code == 429, response.text
+                streamed: Final = response.read()
+            assert streamed == b"", streamed
+            warning: Final = _upstream_warning(owned.log)
+            assert "returned 429" in warning, warning
+            error_information: Final = _spend_error_information(response.headers["x-litellm-call-id"])
+            assert error_information["normalized_error"] == "500_UPSTREAM_PASSTHROUGH", error_information
+
+
+_LLM_429_ENDPOINTS: Final = (
+    (
+        "/v1/chat/completions",
+        "openai/gpt-4o-mini",
+        lambda model, stream: {"model": model, "messages": [{"role": "user", "content": "audit"}], "stream": stream},
+    ),
+    (
+        "/v1/messages",
+        "anthropic/claude-sonnet-4-5-20250929",
+        lambda model, stream: {
+            "model": model,
+            "max_tokens": 8,
+            "messages": [{"role": "user", "content": "audit"}],
+            "stream": stream,
+        },
+    ),
+    (
+        "/v1/responses",
+        "openai/gpt-4o-mini",
+        lambda model, stream: {"model": model, "input": "audit", "stream": stream},
+    ),
+)
+_LLM_429_BODY: Final[dict[str, JsonValue]] = {
+    "error": {
+        "message": "You exceeded your current quota, please check your plan and billing details.",
+        "type": "insufficient_quota",
+    }
+}
+_LLM_429_NORMALIZED: Final = "429_RATE_LIMIT_EXCEEDED"
+
+
+@pytest.mark.parametrize(
+    ("endpoint", "litellm_model", "build_body"),
+    _LLM_429_ENDPOINTS,
+    ids=tuple(spec[0] for spec in _LLM_429_ENDPOINTS),
+)
+@pytest.mark.parametrize("stream", [False, True])
+def test_llm_endpoint_upstream_quota_429_normalized_error_unchanged(
+    gateway: Gateway, endpoint: str, litellm_model: str, build_body: Callable[..., dict[str, JsonValue]], stream: bool
+) -> None:
+    scenario_id: Final = f"quota429-{uuid.uuid4().hex}"
+    handle: Final = register_scenario(
+        scenario_id, JsonResponse(content_type="application/json", body=_LLM_429_BODY, status=429)
+    )
+    try:
+        with gateway.scenario() as scenario:
+            api_base: Final = handle.api_base() if litellm_model.startswith("anthropic/") else f"{handle.api_base()}/v1"
+            model: Final = scenario.model(model=litellm_model, api_base=api_base)
+            if stream:
+                with gateway.client.stream(
+                    "POST",
+                    endpoint,
+                    json=build_body(model, True),
+                    headers={"Authorization": f"Bearer {gateway.key}"},
+                ) as response:
+                    assert response.status_code == 429, response.text
+                    response.read()
+            else:
+                response = gateway.request("POST", endpoint, build_body(model, False))
+                assert response.status_code == 429, response.text
+            error_information: Final = _spend_error_information(response.headers["x-litellm-call-id"])
+            assert error_information["normalized_error"] == _LLM_429_NORMALIZED, error_information
+    finally:
+        delete_scenario(handle)