From 83a78adda00a809adb7da75bc0e2f3d1432719bf Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 16 Oct 2026 18:27:56 -0700 Subject: [PATCH] test(integration): cover a crashed WebSocket turn check and failure-row tags A turn whose admission check raises gets a 500 server_error frame, the socket closes with 1011 and its one failure row carries the error, on both relays. The refused turn's failure row keeps the caller's tags, and HTTP failure rows on chat, responses and messages keep them too. --- ...test_responses_websocket_turn_admission.py | 124 +++++++++++++++--- .../spend/test_failed_dispatch_tokens.py | 42 ++++-- 2 files changed, 142 insertions(+), 24 deletions(-) diff --git a/tests/integration/routing/test_responses_websocket_turn_admission.py b/tests/integration/routing/test_responses_websocket_turn_admission.py index a5fc6566fa8..d6a8651c0cd 100644 --- a/tests/integration/routing/test_responses_websocket_turn_admission.py +++ b/tests/integration/routing/test_responses_websocket_turn_admission.py @@ -1,4 +1,5 @@ import asyncio +import contextlib import hashlib import os import uuid @@ -8,6 +9,7 @@ from dataclasses import dataclass from functools import partial from datetime import UTC, datetime from pathlib import Path +from types import MappingProxyType from typing import Final import httpx @@ -50,8 +52,23 @@ TURN_SPEND: Final = USAGE["input_tokens"] * INPUT_RATE + USAGE["output_tokens"] TURN_TOKENS: Final = USAGE["total_tokens"] MANAGED_PROVIDER: Final = "mistral" SPEND_ROWS_FOR_KEY: Final = ( - 'SELECT request_id, spend, status, model_group, metadata FROM "LiteLLM_SpendLogs" WHERE api_key = %s' + 'SELECT request_id, spend, status, model_group, metadata, request_tags FROM "LiteLLM_SpendLogs" WHERE api_key = %s' ) +TURN_CHECK_FAULT: Final = "turn_check_fault" +FAULTING_TURN_CHECK: Final = f""" +from litellm.integrations.custom_logger import CustomLogger + + +class FaultingTurnCheck(CustomLogger): + async def async_pre_call_check(self, deployment, parent_otel_span): + fault = (deployment.get("model_info") or {{}}).get("{TURN_CHECK_FAULT}") + if fault: + raise RuntimeError(fault) + return None + + +instance = FaultingTurnCheck() +""" FrameReader = Callable[[], int] @@ -70,13 +87,17 @@ def _redis() -> Redis: return Redis(host=os.environ["REDIS_HOST"], port=int(os.environ["REDIS_PORT"]), decode_responses=True) -def _config(directory: Path, router_settings: Mapping[str, JsonValue]) -> Path: +def _config(directory: Path, router_settings: Mapping[str, JsonValue], callbacks: tuple[str, ...]) -> Path: base: Final = JSON_OBJECT.validate_python(yaml.safe_load(PROXY_CONFIG.read_text())) path: Final = directory / "turn-admission.yaml" path.write_text( yaml.safe_dump( { **base, + "litellm_settings": { + **object_value(base["litellm_settings"]), + **({"callbacks": list(callbacks)} if callbacks else {}), + }, "router_settings": { **object_value(base["router_settings"]), "redis_host": os.environ["REDIS_HOST"], @@ -91,13 +112,17 @@ def _config(directory: Path, router_settings: Mapping[str, JsonValue]) -> Path: @contextmanager def _admitting_proxy( - tmp_path_factory: pytest.TempPathFactory, cert: tuple[Path, Path], router_settings: Mapping[str, JsonValue] + tmp_path_factory: pytest.TempPathFactory, + cert: tuple[Path, Path], + router_settings: Mapping[str, JsonValue], + callback_sources: Mapping[str, str] = MappingProxyType({}), ) -> Iterator[Gateway]: directory: Final = tmp_path_factory.mktemp("responses-ws-turn-admission") + for module, source in callback_sources.items(): + (directory / f"{module}.py").write_text(source) + config: Final = _config(directory, router_settings, tuple(f"{module}.instance" for module in callback_sources)) with gateway_from_environment() as base: - with owned_proxy( - base, directory, {"SSL_CERT_FILE": str(cert[0])}, config=_config(directory, router_settings) - ) as proxy: + with owned_proxy(base, directory, {"SSL_CERT_FILE": str(cert[0])}, config=config) as proxy: yield proxy @@ -112,6 +137,12 @@ def limited(tmp_path_factory: pytest.TempPathFactory, cert: tuple[Path, Path]) - yield proxy +@pytest.fixture(scope="module") +def faulting(tmp_path_factory: pytest.TempPathFactory, cert: tuple[Path, Path]) -> Iterator[Gateway]: + with _admitting_proxy(tmp_path_factory, cert, {}, {"faulting_turn_check": FAULTING_TURN_CHECK}) as proxy: + yield proxy + + @pytest.fixture(scope="module") def usage_routed(tmp_path_factory: pytest.TempPathFactory, cert: tuple[Path, Path]) -> Iterator[Gateway]: with _admitting_proxy(tmp_path_factory, cert, {"routing_strategy": "usage-based-routing-v2"}) as proxy: @@ -189,9 +220,13 @@ def _deployment_spend(deployment: Deployment) -> float: return float(redis.get(f"deployment_spend:{deployment.identity}:1d") or 0) -def _configured_tpm(gateway: Gateway, deployment: Deployment) -> JsonValue: +def _configured(gateway: Gateway, deployment: Deployment, section: str) -> dict[str, JsonValue]: entries: Final = list_value(gateway.get("/model/info", {"litellm_model_id": deployment.identity})["data"]) - return object_value(object_value(entries[0])["litellm_params"]).get("tpm") + return object_value(object_value(entries[0])[section]) + + +def _configured_tpm(gateway: Gateway, deployment: Deployment) -> JsonValue: + return _configured(gateway, deployment, "litellm_params").get("tpm") def _tokens_counted(identity: str, minute: str) -> int: @@ -222,6 +257,21 @@ async def _until(read: Callable[[], float], satisfied: Callable[[float], bool]) return await asyncio.to_thread(eventually, read, satisfied, 15) +async def _close_code(connection: websockets.ClientConnection) -> int | None: + with contextlib.suppress(TimeoutError): + await asyncio.wait_for(connection.wait_closed(), timeout=20) + return connection.close_code + + +def _failure_rows(key: str) -> list[dict[str, JsonValue]]: + rows: Final = eventually( + lambda: read_rows(SPEND_ROWS_FOR_KEY, (_key_hash(key),)), + lambda found: any(row["status"] == "failure" for row in found), + seconds=70, + ) + return [row for row in rows if row["status"] == "failure"] + + @pytest.mark.parametrize("relay", ("native", "managed")) def test_a_turn_past_the_deployment_budget_is_refused_on_the_open_socket_and_logged_once_as_its_failure( limited: Gateway, cert: tuple[Path, Path], relay: str @@ -230,9 +280,10 @@ def test_a_turn_past_the_deployment_budget_is_refused_on_the_open_socket_and_log budget: Final = {"max_budget": TURN_SPEND / 2, "budget_duration": "1d"} deployment: Final = _native(scenario, peer, **budget) if relay == "native" else _managed(scenario, **budget) key: Final = scenario.key(models=[deployment.name]) + tag: Final = f"ws-turn-{uuid.uuid4().hex}" async def session() -> tuple[tuple[dict[str, JsonValue], ...], ...]: - async with _socket(limited, key, deployment.name) as connection: + async with _socket(limited, key, deployment.name, {"x-litellm-tags": tag}) as connection: first: Final = await _say(connection, deployment.name) await _until(lambda: _deployment_spend(deployment), lambda spent: spent > TURN_SPEND / 2) return (first, await _say(connection, deployment.name), await _say(connection, deployment.name)) @@ -243,19 +294,60 @@ def test_a_turn_past_the_deployment_budget_is_refused_on_the_open_socket_and_log _assert_refused(second) _assert_refused(third) assert deployment.forwarded_turns() == 1 - rows: Final = eventually( - lambda: read_rows(SPEND_ROWS_FOR_KEY, (_key_hash(key),)), - lambda found: any(row["status"] == "failure" for row in found), - seconds=70, - ) - failures: Final = [row for row in rows if row["status"] == "failure"] - assert len(failures) == 1, rows + failures: Final = _failure_rows(key) + assert len(failures) == 1, failures assert failures[0]["model_group"] == deployment.name, failures + assert tag in list_value(failures[0]["request_tags"]), failures assert float(failures[0]["spend"]) == 0.0, failures error: Final = object_value(object_value(failures[0]["metadata"])["error_information"]) assert str(error["error_code"]) == "429", error +@pytest.mark.parametrize("relay", ("native", "managed")) +def test_a_turn_whose_admission_check_fails_gets_a_500_frame_and_ends_the_socket_as_its_failure( + faulting: Gateway, cert: tuple[Path, Path], relay: str +) -> None: + with responses_peer(cert) as peer, faulting.scenario() as scenario: + deployment: Final = _native(scenario, peer) if relay == "native" else _managed(scenario) + key: Final = scenario.key(models=[deployment.name]) + fault: Final = f"turn check unavailable {uuid.uuid4().hex}" + + def fail_the_turn_check() -> None: + faulting.request( + "PATCH", + f"/model/{deployment.identity}/update", + {"model_info": {"id": deployment.identity, TURN_CHECK_FAULT: fault}}, + ).raise_for_status() + eventually( + lambda: _configured(faulting, deployment, "model_info").get(TURN_CHECK_FAULT), + lambda configured: configured == fault, + seconds=15, + ) + + async def session() -> tuple[tuple[dict[str, JsonValue], ...], tuple[dict[str, JsonValue], ...], int | None]: + async with _socket(faulting, key, deployment.name) as connection: + first: Final = await _say(connection, deployment.name) + await asyncio.to_thread(fail_the_turn_check) + second: Final = await _say(connection, deployment.name) + return first, second, await _close_code(connection) + + first, second, close_code = asyncio.run(session()) + + completed(first) + assert [event.get("type") for event in second] == ["error"], second + assert second[0]["status"] == 500, second + error: Final = object_value(second[0]["error"]) + assert error["type"] == "server_error", second + assert fault in string_value(error["message"]), second + assert close_code == 1011, close_code + assert deployment.forwarded_turns() == 1 + failures: Final = _failure_rows(key) + assert len(failures) == 1, failures + assert failures[0]["model_group"] == deployment.name, failures + recorded: Final = object_value(object_value(failures[0]["metadata"])["error_information"]) + assert fault in string_value(recorded["error_message"]), recorded + + def test_a_socket_opened_after_websocket_turns_spent_the_budget_is_refused_and_so_is_http( limited: Gateway, cert: tuple[Path, Path] ) -> None: diff --git a/tests/integration/spend/test_failed_dispatch_tokens.py b/tests/integration/spend/test_failed_dispatch_tokens.py index 9ee1b0389c6..db5c49623eb 100644 --- a/tests/integration/spend/test_failed_dispatch_tokens.py +++ b/tests/integration/spend/test_failed_dispatch_tokens.py @@ -4,7 +4,7 @@ from collections.abc import Callable, Mapping from typing import Final import pytest -from integration._support.client import Gateway, eventually +from integration._support.client import Gateway, eventually, list_value from integration._support.database import read_rows from integration._support.wire import Reply, Request, wire_server from pydantic import JsonValue @@ -64,7 +64,7 @@ def _messages_body(model: str, text: str) -> Mapping[str, JsonValue]: return {"model": model, "max_tokens": 16, "messages": [{"role": "user", "content": text}]} -@pytest.mark.parametrize( +FAILURE_ROUTES: Final = pytest.mark.parametrize( ("route", "body"), ( ("/v1/chat/completions", _chat_body), @@ -73,9 +73,14 @@ def _messages_body(model: str, text: str) -> Mapping[str, JsonValue]: ), ids=("chat", "responses", "messages"), ) -def test_a_provider_500_failure_row_names_the_model_group_the_caller_asked_for( - gateway: Gateway, route: str, body: Callable[[str, str], Mapping[str, JsonValue]] -) -> None: + + +def _provider_500_failure_row( + gateway: Gateway, + route: str, + body: Callable[[str, str], Mapping[str, JsonValue]], + headers: Mapping[str, str] | None = None, +) -> tuple[str, dict[str, JsonValue]]: def provider(request: Request) -> Reply: assert request.method == "POST" return Reply( @@ -87,16 +92,37 @@ def test_a_provider_500_failure_row_names_the_model_group_the_caller_asked_for( model: Final = scenario.model(api_base=wire.url + "/v1", num_retries=0) key: Final = scenario.key(models=[model]) failed: Final = gateway.request( - "POST", route, body(model, f"failed dispatch group {uuid.uuid4().hex}"), key=key + "POST", route, body(model, f"failed dispatch group {uuid.uuid4().hex}"), key=key, headers=headers ) assert failed.status_code == 500, failed.text assert wire.drain() rows: Final = eventually( lambda: read_rows( - 'SELECT status, model_group FROM "LiteLLM_SpendLogs" WHERE request_id=%s', + 'SELECT status, model_group, request_tags FROM "LiteLLM_SpendLogs" WHERE request_id=%s', (failed.headers["x-litellm-call-id"],), ), lambda values: len(values) == 1, seconds=70, ) - assert rows[0]["status"] == "failure" and rows[0]["model_group"] == model, rows + assert rows[0]["status"] == "failure", rows + return model, rows[0] + + +@FAILURE_ROUTES +def test_a_provider_500_failure_row_names_the_model_group_the_caller_asked_for( + gateway: Gateway, route: str, body: Callable[[str, str], Mapping[str, JsonValue]] +) -> None: + model, row = _provider_500_failure_row(gateway, route, body) + + assert row["model_group"] == model, row + + +@FAILURE_ROUTES +def test_a_provider_500_failure_row_keeps_the_tags_the_caller_sent( + gateway: Gateway, route: str, body: Callable[[str, str], Mapping[str, JsonValue]] +) -> None: + tag: Final = f"failed-dispatch-{uuid.uuid4().hex}" + + _, row = _provider_500_failure_row(gateway, route, body, {"x-litellm-tags": tag}) + + assert tag in list_value(row["request_tags"]), row