From c9292d3af2fb9a8aab2111a26082e5852b61dc75 Mon Sep 17 00:00:00 2001 From: Yassin Kortam Date: Fri, 7 Aug 2026 09:57:30 -0700 Subject: [PATCH 1/4] fix(proxy): stop alerting on health probes that lose the planned engine-restart race (#36141) --- litellm/proxy/db/prisma_client.py | 32 ++ litellm/proxy/utils.py | 118 ++++- .../db/test_prisma_planned_engine_restart.py | 444 ++++++++++++++++++ tests/test_litellm/proxy/test_proxy_utils.py | 8 + 4 files changed, 593 insertions(+), 9 deletions(-) diff --git a/litellm/proxy/db/prisma_client.py b/litellm/proxy/db/prisma_client.py index 6863687081c..5f86490a474 100644 --- a/litellm/proxy/db/prisma_client.py +++ b/litellm/proxy/db/prisma_client.py @@ -486,6 +486,38 @@ class PrismaWrapper: os.environ[self._db_url_env_var] = _db_url return _db_url + @property + def engine_generation(self) -> int: + """How many query-engine replacements have completed on this wrapper. + + Bumped under `_reconnection_lock` only after a replacement engine has + connected, so a change across an await proves a *successful* planned + replacement happened in between — a replacement that failed (a real + outage) leaves it untouched. + """ + return self._engine_generation + + async def _reconnection_settled(self) -> None: + async with self._reconnection_lock: + pass + + async def wait_for_planned_engine_replacement(self, timeout_seconds: float) -> None: + """Wait, bounded, for an in-flight planned engine replacement to finish. + + Both replacement paths (`recreate_prisma_client` and + `_safe_refresh_token`) hold `_reconnection_lock` across their whole + kill/connect window, so re-acquiring it means the replacement has + settled one way or the other. Gives up silently on timeout: a caller + that stopped waiting must treat the replacement as not completed and + consult `engine_generation` rather than assume success. + """ + if timeout_seconds <= 0 or not self._reconnection_lock.locked(): + return + try: + await asyncio.wait_for(self._reconnection_settled(), timeout=timeout_seconds) + except asyncio.TimeoutError: + return + async def recreate_prisma_client( self, new_db_url: str, diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index cdd41d2ed42..605455a2f73 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -3000,6 +3000,13 @@ class PrismaClient: ] = [] # mutable-ok: drained queue, mirrors tool_usage_transactions _autorouter_turn_transactions_lock = asyncio.Lock() + # How long a health probe failure waits for an in-flight planned engine + # replacement to settle before deciding whether to report itself. Generous + # against a replacement that takes well under a second, and far short of the + # reconnect budget an outage-hung `connect()` runs under, so a real outage + # is never waited out. + PLANNED_ENGINE_REPLACEMENT_SETTLE_SECONDS: ClassVar[float] = 5.0 + def __init__( self, database_url: str, @@ -4970,6 +4977,101 @@ class PrismaClient: else: verbose_proxy_logger.debug("Prisma DB health watchdog observed non-DB error: %s", e) + def _probe_target_wrapper(self) -> PrismaWrapper: + """The Prisma wrapper a `SELECT 1` health probe actually reaches. + + `health_check()` issues `query_raw`, which `RoutingPrismaWrapper` sends + to the reader unless the reader is degraded. The writer's engine state + therefore says nothing about a probe that failed against the reader, so + the gate has to follow the same routing rule the probe did. + """ + if isinstance(self.db, RoutingPrismaWrapper): + return self.db.writer if self.db.reader_unavailable else self.db.reader + return self.db + + async def _run_health_probe(self, wrapper: PrismaWrapper) -> object: + """Issue the `SELECT 1` a health check is made of, against `wrapper`. + + Takes the wrapper rather than re-reading `self.db`, because routing is + re-resolved on every attribute access: a reader that recovers between + the caller picking its target and the query going out would send the + probe to a different engine than the one whose generation the caller is + about to check, and attribute the failure to the wrong replacement. + """ + sql_query: Final = "SELECT 1" + response: Final = await wrapper.query_raw(sql_query) + return response + + async def _probe_answers_now(self, wrapper: PrismaWrapper) -> bool: + try: + await self._run_health_probe(wrapper) + except Exception as probe_error: # noqa: BLE001 # any failure means the database is not answering + verbose_proxy_logger.debug("Prisma health_check() confirmation probe failed: %s", probe_error) + return False + return True + + async def _planned_engine_replacement_absorbed( + self, + e: Exception, + wrapper: PrismaWrapper, + generation_before: int, + ) -> bool: + """True iff `e` is a connection-class probe failure that a completed + planned query-engine replacement explains. + + Planned replacements (RDS IAM token refresh, guarded reconnect) kill the + running query engine and spawn a new one. A `SELECT 1` probe that races + that sub-second window fails with a transport error against the engine's + local HTTP port even though nothing is wrong with the database, and + reporting it drives a false-positive `db_exceptions` alert on every + replacement. + + Two things must both hold, because neither is sufficient alone. The + engine generation must have moved, which says a replacement completed + rather than merely being attempted: reconnect attempts during a real + outage hold the same lock for tens of seconds, so gating on an in-flight + replacement would swallow most of an outage's alerts. And a fresh probe + must succeed, because `Prisma.connect()` polls the query engine's own + `/status` endpoint rather than round-tripping to the database, so a + future engine that binds before it validates its connection pool would + let the generation advance with the database still unreachable. + + Waiting for an in-flight replacement to settle is what makes the + generation check meaningful, since the generation has not moved yet at + the instant the probe fails. The wait is generous against a replacement + that takes well under a second and short enough that an outage-hung + reconnect is not waited out; a replacement that has not settled by then + reports rather than stays silent. + """ + if not PrismaDBExceptionHandler.is_database_connection_error(e): + return False + await wrapper.wait_for_planned_engine_replacement(self.PLANNED_ENGINE_REPLACEMENT_SETTLE_SECONDS) + if wrapper.engine_generation == generation_before: + return False + return await self._probe_answers_now(wrapper) + + async def _report_health_check_failure( + self, + e: Exception, + duration: float, + traceback_str: str, + wrapper: PrismaWrapper, + generation_before: int, + ) -> None: + if await self._planned_engine_replacement_absorbed(e, wrapper, generation_before): + verbose_proxy_logger.info( + "Prisma health_check() connection error raced a planned query-engine replacement; " + "not reporting it as a DB exception: %s", + e, + ) + return + await self.proxy_logging_obj.failure_handler( + original_exception=e, + duration=duration, + call_type="health_check", + traceback_str=traceback_str, + ) + @backoff.on_exception( backoff.expo, Exception, @@ -4982,13 +5084,10 @@ class PrismaClient: Health check endpoint for the prisma client """ start_time: Final = time.time() + probe_wrapper: Final = self._probe_target_wrapper() + generation_before: Final = probe_wrapper.engine_generation try: - sql_query: Final = "SELECT 1" - - # Execute the raw query - # The asterisk before `user_id_list` unpacks the list into separate arguments - response: Final = await self.db.query_raw(sql_query) - return response + return await self._run_health_probe(probe_wrapper) except Exception as e: import traceback @@ -4998,11 +5097,12 @@ class PrismaClient: end_time: Final = time.time() _duration: Final = end_time - start_time asyncio.create_task( - self.proxy_logging_obj.failure_handler( - original_exception=e, + self._report_health_check_failure( + e=e, duration=_duration, - call_type="health_check", traceback_str=error_traceback, + wrapper=probe_wrapper, + generation_before=generation_before, ) ) raise e diff --git a/tests/test_litellm/proxy/db/test_prisma_planned_engine_restart.py b/tests/test_litellm/proxy/db/test_prisma_planned_engine_restart.py index 18a238b7545..c8e0338eeaa 100644 --- a/tests/test_litellm/proxy/db/test_prisma_planned_engine_restart.py +++ b/tests/test_litellm/proxy/db/test_prisma_planned_engine_restart.py @@ -6,13 +6,24 @@ subprocess), and the engine-death watcher / in-flight transport-error retries must not treat that planned restart as a crash and recreate the client a second time. +A planned restart is also invisible to the database itself, so a ``SELECT 1`` +health probe that races the kill/connect window fails with a transport error +against the engine's local HTTP port. That failure must not be reported as a +``db_exceptions`` DB failure, or every IAM refresh cycle raises a false alarm. + Symbols pinned here: - ``PrismaWrapper._expected_engine_deaths`` - ``PrismaWrapper._engine_generation`` + - ``PrismaWrapper.engine_generation`` + - ``PrismaWrapper.wait_for_planned_engine_replacement`` - ``PrismaWrapper.on_engine_replaced`` - ``PrismaWrapper.recreate_prisma_client`` (expected_generation guard) - ``PrismaWrapper._safe_refresh_token`` (refresh coalescing) - ``RoutingPrismaWrapper.recreate_prisma_client`` (guard forwarding) + - ``PrismaClient.health_check`` (planned-replacement alert suppression) + - ``PrismaClient._probe_target_wrapper`` + - ``PrismaClient._probe_answers_now`` + - ``PrismaClient._planned_engine_replacement_absorbed`` """ import asyncio @@ -21,22 +32,29 @@ import signal import sys import urllib.parse from datetime import datetime, timedelta +from typing import Any, List from unittest.mock import AsyncMock, MagicMock, call, patch +import httpx import pytest from prisma import Prisma as GeneratedPrisma +from prisma.engine.errors import EngineConnectionError sys.path.insert( 0, os.path.abspath("../../../..") ) # Adds the parent directory to the system path from litellm.proxy.db.prisma_client import PrismaWrapper +from litellm.proxy.utils import PrismaClient @pytest.fixture(autouse=True) def mock_prisma_binary(): """Mock prisma.Prisma to avoid requiring generated Prisma binaries for unit tests.""" mock_module = MagicMock() + # Production code isinstance-checks against this, which a bare MagicMock + # attribute cannot satisfy. + mock_module.engine.errors.EngineConnectionError = EngineConnectionError with patch.dict(sys.modules, {"prisma": mock_module}): yield mock_module @@ -49,6 +67,87 @@ def _make_wrapper(engine_pid: int = 111, iam: bool = False) -> PrismaWrapper: return PrismaWrapper(original_prisma=prisma, iam_token_db_auth=iam) +def _make_prisma_client(db: Any) -> PrismaClient: + """A ``PrismaClient`` whose ``db`` is a real wrapper and whose alerting + hook is observable.""" + proxy_logging_obj = MagicMock() + proxy_logging_obj.failure_handler = AsyncMock() + client = PrismaClient( + database_url="postgresql://user:pass@localhost:5432/db", + proxy_logging_obj=proxy_logging_obj, + ) + client.db = db + client._db_watchdog_reconnect_timeout_seconds = 5.0 + return client + + +_real_asyncio_sleep = asyncio.sleep + + +async def _yield_to_loop(times: int = 10) -> None: + """Let already-scheduled tasks make progress. + + Bound to the real ``asyncio.sleep`` at import time: the tests below patch + ``asyncio.sleep`` to skip the SIGTERM/SIGKILL grace, and an ``AsyncMock`` + stand-in never yields to the event loop, which would silently leave every + background task un-started and the assertions vacuous. + """ + for _ in range(times): + await _real_asyncio_sleep(0) + + +async def _await_health_check_reports() -> None: + """Await the fire-and-forget reporting tasks ``health_check()`` scheduled. + + Selected by coroutine qualname rather than by draining every pending task, + so an unrelated background task can never make these assertions pass by + accident. + """ + reports = [ + task + for task in asyncio.all_tasks() + if getattr(task.get_coro(), "__qualname__", "") + == "PrismaClient._report_health_check_failure" + ] + if reports: + await asyncio.gather(*reports, return_exceptions=True) + + +def _fails_then_answers(error: Exception, failures: int = 3) -> Any: + """Raise ``error`` for the first ``failures`` probes, then answer. + + ``health_check`` retries up to three times, so this exhausts the retries and + still lets the confirmation probe that decides suppression succeed. Without + that, a test would report for the wrong reason: the confirmation probe would + fail too, masking whether the error type was classified at all. + """ + seen: List[int] = [] + + async def _query_raw(_sql: str) -> Any: + seen.append(1) + if len(seen) <= failures: + raise error + return [{"?column?": 1}] + + return _query_raw + + +def _blocking_replacement(gate: asyncio.Event, fail: bool = False) -> MagicMock: + """A replacement Prisma whose ``connect()`` parks until ``gate`` is set. + + Holds ``_reconnection_lock`` open for as long as the test needs, which is + how a health probe is made to fail *while* a planned replacement is in + flight rather than after it. + """ + + async def _connect(*_: Any, **__: Any) -> None: + await gate.wait() + if fail: + raise ConnectionRefusedError("database is down") + + return MagicMock(connect=AsyncMock(side_effect=_connect)) + + def _token_db_url(created: datetime, expires_in: int = 900) -> str: """Build a DATABASE_URL whose password is a parseable RDS IAM token.""" token = ( @@ -603,3 +702,348 @@ async def test_recreate_caps_expected_engine_deaths_set(mock_prisma_binary): await wrapper.recreate_prisma_client("postgresql://new") assert wrapper._expected_engine_deaths == {111} + + +@pytest.mark.asyncio +async def test_wait_for_planned_engine_replacement_returns_once_recreate_settles( + mock_prisma_binary, +): + wrapper = _make_wrapper(engine_pid=111) + gate = asyncio.Event() + mock_prisma_binary.Prisma.return_value = _blocking_replacement(gate) + + with patch("os.kill"), patch("asyncio.sleep", new_callable=AsyncMock): + recreate = asyncio.create_task( + wrapper.recreate_prisma_client("postgresql://new") + ) + await _yield_to_loop() + assert wrapper._reconnection_lock.locked() is True + + waiter = asyncio.create_task(wrapper.wait_for_planned_engine_replacement(5.0)) + await _yield_to_loop() + blocked_while_in_flight = not waiter.done() + + gate.set() + await recreate + await waiter + + assert { + "blocked_while_in_flight": blocked_while_in_flight, + "generation": wrapper.engine_generation, + } == {"blocked_while_in_flight": True, "generation": 1} + + +@pytest.mark.asyncio +async def test_wait_for_planned_engine_replacement_gives_up_at_timeout( + mock_prisma_binary, +): + """A replacement that never settles must not stall the caller forever; the + caller then sees an unchanged generation and reports the failure.""" + wrapper = _make_wrapper(engine_pid=111) + gate = asyncio.Event() + mock_prisma_binary.Prisma.return_value = _blocking_replacement(gate) + + with patch("os.kill"), patch("asyncio.sleep", new_callable=AsyncMock): + recreate = asyncio.create_task( + wrapper.recreate_prisma_client("postgresql://new") + ) + await _yield_to_loop() + assert wrapper._reconnection_lock.locked() is True + + await asyncio.wait_for( + wrapper.wait_for_planned_engine_replacement(0.05), timeout=5.0 + ) + gave_up_with_replacement_still_in_flight = not recreate.done() + + gate.set() + await recreate + + assert gave_up_with_replacement_still_in_flight is True + + +@pytest.mark.asyncio +async def test_health_check_does_not_alert_when_probe_races_a_completed_replacement( + mock_prisma_binary, +): + """The reported bug: an IAM-refresh engine recreate makes a concurrent + readiness probe fail transiently, and that failure was alerting as a DB + exception on every refresh cycle. + + The reporting task is drained while the replacement is still in flight, + which is when it runs in production; a decision taken at that instant sees + an engine generation that has not moved yet. + """ + wrapper = _make_wrapper(engine_pid=111) + client = _make_prisma_client(wrapper) + wrapper.query_raw = AsyncMock( + side_effect=[ + httpx.ConnectError("All connection attempts failed"), + [{"?column?": 1}], + [{"?column?": 1}], + ] + ) + gate = asyncio.Event() + mock_prisma_binary.Prisma.return_value = _blocking_replacement(gate) + + with patch("os.kill"), patch("asyncio.sleep", new_callable=AsyncMock): + recreate = asyncio.create_task( + wrapper.recreate_prisma_client("postgresql://new") + ) + await _yield_to_loop() + assert wrapper._reconnection_lock.locked() is True + + probe_result = await client.health_check() + + drain = asyncio.create_task(_await_health_check_reports()) + await _yield_to_loop() + alerts_while_replacement_in_flight = ( + client.proxy_logging_obj.failure_handler.await_count + ) + + gate.set() + await recreate + await drain + + assert { + "probe_result": probe_result, + "probe_attempts": wrapper.query_raw.await_count, + "alerts_while_in_flight": alerts_while_replacement_in_flight, + "alerts": client.proxy_logging_obj.failure_handler.await_count, + } == { + "probe_result": [{"?column?": 1}], + "probe_attempts": 3, + "alerts_while_in_flight": 0, + "alerts": 0, + } + + +@pytest.mark.asyncio +async def test_health_check_alerts_when_a_completed_replacement_still_cannot_reach_the_database( + mock_prisma_binary, +): + """``Prisma.connect()`` polls the query engine's own ``/status`` endpoint + rather than round-tripping to the database, so a replacement can complete + against a database that is still unreachable. The engine generation alone + must not be enough to stay silent.""" + wrapper = _make_wrapper(engine_pid=111) + client = _make_prisma_client(wrapper) + wrapper.query_raw = AsyncMock( + side_effect=httpx.ConnectError("All connection attempts failed") + ) + gate = asyncio.Event() + mock_prisma_binary.Prisma.return_value = _blocking_replacement(gate) + + with patch("os.kill"), patch("asyncio.sleep", new_callable=AsyncMock): + recreate = asyncio.create_task( + wrapper.recreate_prisma_client("postgresql://new") + ) + await _yield_to_loop() + assert wrapper._reconnection_lock.locked() is True + + with pytest.raises(httpx.ConnectError): + await client.health_check() + + gate.set() + await recreate + await _await_health_check_reports() + + assert { + "replacement_completed": wrapper.engine_generation, + "alerted": client.proxy_logging_obj.failure_handler.await_count > 0, + } == {"replacement_completed": 1, "alerted": True} + + +@pytest.mark.asyncio +async def test_health_check_alerts_when_the_replacement_never_completes( + mock_prisma_binary, +): + """A real outage also has a replacement in flight, but it fails, so the + engine generation never advances and the probe failure must still alert.""" + wrapper = _make_wrapper(engine_pid=111) + client = _make_prisma_client(wrapper) + wrapper.query_raw = AsyncMock( + side_effect=httpx.ConnectError("All connection attempts failed") + ) + gate = asyncio.Event() + mock_prisma_binary.Prisma.return_value = _blocking_replacement(gate, fail=True) + + with patch("os.kill"), patch("asyncio.sleep", new_callable=AsyncMock): + recreate = asyncio.create_task( + wrapper.recreate_prisma_client("postgresql://new") + ) + await _yield_to_loop() + assert wrapper._reconnection_lock.locked() is True + + with pytest.raises(httpx.ConnectError): + await client.health_check() + + gate.set() + with pytest.raises(ConnectionRefusedError): + await recreate + await _await_health_check_reports() + + call_types: List[str] = [ + c.kwargs["call_type"] + for c in client.proxy_logging_obj.failure_handler.await_args_list + ] + assert { + "generation": wrapper.engine_generation, + "alerted": len(call_types) > 0, + "call_types": set(call_types), + } == {"generation": 0, "alerted": True, "call_types": {"health_check"}} + + +@pytest.mark.asyncio +async def test_health_check_alerts_for_non_connection_errors_during_a_replacement( + mock_prisma_binary, +): + """Suppression is scoped to transport failures. A query the database itself + rejected is a real defect and must alert even mid-replacement, and even + though the database is plainly reachable a moment later.""" + wrapper = _make_wrapper(engine_pid=111) + client = _make_prisma_client(wrapper) + wrapper.query_raw = AsyncMock(side_effect=_fails_then_answers(ValueError("malformed SELECT"))) + gate = asyncio.Event() + mock_prisma_binary.Prisma.return_value = _blocking_replacement(gate) + + with patch("os.kill"), patch("asyncio.sleep", new_callable=AsyncMock): + recreate = asyncio.create_task( + wrapper.recreate_prisma_client("postgresql://new") + ) + await _yield_to_loop() + assert wrapper._reconnection_lock.locked() is True + + with pytest.raises(ValueError): + await client.health_check() + + gate.set() + await recreate + await _await_health_check_reports() + + assert { + "generation": wrapper.engine_generation, + "alerted": client.proxy_logging_obj.failure_handler.await_count > 0, + } == {"generation": 1, "alerted": True} + + +@pytest.mark.asyncio +async def test_health_check_alerts_for_a_transient_failure_with_no_engine_replacement( + mock_prisma_binary, +): + """Suppression is scoped to failures a planned replacement explains. A + transport blip that self-heals with no engine replacement at all still + alerts, so the gate cannot be widened into silencing every failure whose + database happens to answer a moment later.""" + wrapper = _make_wrapper(engine_pid=111) + client = _make_prisma_client(wrapper) + wrapper.query_raw = AsyncMock( + side_effect=[ + httpx.ConnectError("All connection attempts failed"), + [{"?column?": 1}], + [{"?column?": 1}], + ] + ) + + probe_result = await client.health_check() + await _await_health_check_reports() + + assert { + "probe_result": probe_result, + "generation": wrapper.engine_generation, + "alerted": client.proxy_logging_obj.failure_handler.await_count > 0, + } == {"probe_result": [{"?column?": 1}], "generation": 0, "alerted": True} + + +@pytest.mark.asyncio +async def test_health_probe_stays_on_its_target_when_reader_availability_flips( + mock_prisma_binary, monkeypatch +): + """Routing is re-resolved on every attribute access, so a reader that + recovers mid-call would otherwise send the probe to a different engine than + the one whose generation is being checked, and blame the wrong replacement. + The probe follows the wrapper it was handed.""" + from litellm.proxy.db.routing_prisma_wrapper import RoutingPrismaWrapper + + monkeypatch.setenv("DATABASE_URL_READ_REPLICA", "postgresql://reader") + writer = _make_wrapper(engine_pid=111) + reader = _make_wrapper(engine_pid=222) + routing = RoutingPrismaWrapper(writer=writer, reader=reader) + client = _make_prisma_client(routing) + writer.query_raw = AsyncMock(return_value=[{"writer": 1}]) + reader.query_raw = AsyncMock(return_value=[{"reader": 1}]) + + routing._reader_unavailable = False + target = client._probe_target_wrapper() + routing._reader_unavailable = True + result = await client._run_health_probe(target) + + assert { + "target_is_reader": target is reader, + "result": result, + "reader_probes": reader.query_raw.await_count, + "writer_probes": writer.query_raw.await_count, + } == { + "target_is_reader": True, + "result": [{"reader": 1}], + "reader_probes": 1, + "writer_probes": 0, + } + + +@pytest.mark.asyncio +async def test_health_check_consults_the_reader_wrapper_under_read_replica_routing( + mock_prisma_binary, monkeypatch +): + """``query_raw`` is routed to the reader, so a reader-side planned + replacement is the one that explains a probe failure.""" + from litellm.proxy.db.routing_prisma_wrapper import RoutingPrismaWrapper + + monkeypatch.setenv("DATABASE_URL_READ_REPLICA", "postgresql://reader") + writer = _make_wrapper(engine_pid=111) + reader = _make_wrapper(engine_pid=222) + routing = RoutingPrismaWrapper(writer=writer, reader=reader) + client = _make_prisma_client(routing) + reader.query_raw = AsyncMock( + side_effect=[ + httpx.ConnectError("All connection attempts failed"), + [{"?column?": 1}], + [{"?column?": 1}], + ] + ) + gate = asyncio.Event() + mock_prisma_binary.Prisma.return_value = _blocking_replacement(gate) + + with patch("os.kill"), patch("asyncio.sleep", new_callable=AsyncMock): + recreate = asyncio.create_task( + reader.recreate_prisma_client("postgresql://new-reader") + ) + await _yield_to_loop() + assert reader._reconnection_lock.locked() is True + + probe_result = await client.health_check() + + drain = asyncio.create_task(_await_health_check_reports()) + await _yield_to_loop() + alerts_while_replacement_in_flight = ( + client.proxy_logging_obj.failure_handler.await_count + ) + + gate.set() + await recreate + await drain + + assert { + "probe_target_is_the_reader": client._probe_target_wrapper() is reader, + "writer_generation": writer.engine_generation, + "reader_generation": reader.engine_generation, + "probe_result": probe_result, + "alerts_while_in_flight": alerts_while_replacement_in_flight, + "alerts": client.proxy_logging_obj.failure_handler.await_count, + } == { + "probe_target_is_the_reader": True, + "writer_generation": 0, + "reader_generation": 1, + "probe_result": [{"?column?": 1}], + "alerts_while_in_flight": 0, + "alerts": 0, + } diff --git a/tests/test_litellm/proxy/test_proxy_utils.py b/tests/test_litellm/proxy/test_proxy_utils.py index abd6220144b..a8e81e92ebd 100644 --- a/tests/test_litellm/proxy/test_proxy_utils.py +++ b/tests/test_litellm/proxy/test_proxy_utils.py @@ -1096,6 +1096,7 @@ async def test_prisma_health_check_failure_names_itself_at_operator_visible_leve function and reads as "the check never ran", and reporting it only at debug level hides a database fault behind a flag nobody enables in production.""" import logging + from functools import partial from unittest.mock import AsyncMock from litellm.proxy.utils import PrismaClient @@ -1103,6 +1104,9 @@ async def test_prisma_health_check_failure_names_itself_at_operator_visible_leve client = MagicMock() client.db.query_raw = AsyncMock(side_effect=Exception("connection refused")) client.proxy_logging_obj.failure_handler = AsyncMock() + client._probe_target_wrapper = MagicMock(return_value=client.db) + client._run_health_probe = partial(PrismaClient._run_health_probe, client) + client._report_health_check_failure = AsyncMock() with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): with pytest.raises(Exception, match="connection refused"): @@ -1142,6 +1146,7 @@ async def test_prisma_health_check_failure_redacts_database_credentials(caplog): text can carry a full connection string, so the credential has to be gone from the emitted record.""" import logging + from functools import partial from unittest.mock import AsyncMock from litellm.proxy.utils import PrismaClient @@ -1151,6 +1156,9 @@ async def test_prisma_health_check_failure_redacts_database_credentials(caplog): side_effect=Exception("could not connect to postgresql://admin:hunter2@db.internal:5432/litellm") ) client.proxy_logging_obj.failure_handler = AsyncMock() + client._probe_target_wrapper = MagicMock(return_value=client.db) + client._run_health_probe = partial(PrismaClient._run_health_probe, client) + client._report_health_check_failure = AsyncMock() with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): with pytest.raises(Exception): From 6ba744b340fd0901f9bfab428a846bc1677ef3c4 Mon Sep 17 00:00:00 2001 From: Yassin Kortam Date: Fri, 7 Aug 2026 09:57:37 -0700 Subject: [PATCH 2/4] test(docker): gate the componentized gateway and backend images on an arbitrary-uid offline boot (#36136) --- .github/ci-coverage-allowlist.yml | 7 - .github/workflows/image-scan.yml | 65 ++++++ .../test_component_image_serves_offline.py | 187 ++++++++++++++++++ 3 files changed, 252 insertions(+), 7 deletions(-) create mode 100644 tests/proxy_migration_tests/test_component_image_serves_offline.py diff --git a/.github/ci-coverage-allowlist.yml b/.github/ci-coverage-allowlist.yml index 140f1155e3a..1423228e725 100644 --- a/.github/ci-coverage-allowlist.yml +++ b/.github/ci-coverage-allowlist.yml @@ -136,13 +136,6 @@ test_paths: - tests/litellm-proxy-extras/test_litellm_proxy_extras_utils.py dockerfiles: - - reason: >- - The componentized images the microservices chart deploys are built by no job; wiring both into - the scan workflow costs a full image build each and is deferred to a change that prices the - whole set - paths: - - backend/Dockerfile - - gateway/Dockerfile - reason: >- The dashboard container is a static Next.js export served by nginx, and the dashboard build and lint workflows already exercise that output, so building the image adds no signal about it diff --git a/.github/workflows/image-scan.yml b/.github/workflows/image-scan.yml index e0c0bfcedae..8faf3ef6229 100644 --- a/.github/workflows/image-scan.yml +++ b/.github/workflows/image-scan.yml @@ -12,6 +12,11 @@ on: - docker/Dockerfile.non_root - migrations/Dockerfile - migrations/run.py + - gateway/Dockerfile + - gateway/main.py + - backend/Dockerfile + - backend/main.py + - docker/component_entrypoint.sh - litellm-proxy-extras/** - tests/proxy_migration_tests/** - uv.lock @@ -147,3 +152,63 @@ jobs: run: | python -m pip install "pytest==9.0.3" python -m pytest tests/proxy_migration_tests/test_offline_image_migration.py -v + + gateway-image: + name: gateway-image + runs-on: ubuntu-latest + if: >- + github.event_name != 'pull_request' || + github.event.pull_request.head.repo.full_name == github.repository + timeout-minutes: 30 + permissions: + contents: read + steps: + - uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0 + with: + persist-credentials: false + + - name: Build gateway image + run: docker build -f gateway/Dockerfile -t litellm-gateway-scan:${{ github.sha }} . + + - name: Set up Python + uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0 + with: + python-version: "3.12" + + - name: Verify the gateway serves offline as a non-root uid + env: + LITELLM_IMAGE: litellm-gateway-scan:${{ github.sha }} + LITELLM_COMPONENT_PORT: "4000" + run: | + python -m pip install "pytest==9.0.3" + python -m pytest tests/proxy_migration_tests/test_component_image_serves_offline.py -v + + backend-image: + name: backend-image + runs-on: ubuntu-latest + if: >- + github.event_name != 'pull_request' || + github.event.pull_request.head.repo.full_name == github.repository + timeout-minutes: 30 + permissions: + contents: read + steps: + - uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0 + with: + persist-credentials: false + + - name: Build backend image + run: docker build -f backend/Dockerfile -t litellm-backend-scan:${{ github.sha }} . + + - name: Set up Python + uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0 + with: + python-version: "3.12" + + - name: Verify the backend serves offline as a non-root uid + env: + LITELLM_IMAGE: litellm-backend-scan:${{ github.sha }} + LITELLM_COMPONENT_PORT: "4001" + run: | + python -m pip install "pytest==9.0.3" + python -m pytest tests/proxy_migration_tests/test_component_image_serves_offline.py -v diff --git a/tests/proxy_migration_tests/test_component_image_serves_offline.py b/tests/proxy_migration_tests/test_component_image_serves_offline.py new file mode 100644 index 00000000000..528c2960072 --- /dev/null +++ b/tests/proxy_migration_tests/test_component_image_serves_offline.py @@ -0,0 +1,187 @@ +"""Image-level regression net for the prisma bake in the componentized images. + +The gateway and backend serve requests; they never shell out to the Prisma CLI +(``PrismaManager.setup_database`` is reachable only from ``proxy_cli.py``, which +uvicorn'ing ``gateway.main:app`` bypasses). What they do need is the generated +client's baked query engine, and prisma-python resolves those baked paths +eagerly, with an existence check that propagates EACCES rather than skipping the +candidate. An engine baked under a build-time ``HOME`` is therefore unreadable to +any other runtime uid, and the process dies during startup before +``PRISMA_QUERY_ENGINE_BINARY`` is ever consulted. + +That is what an OpenShift ``restricted-v2`` namespace produces: the image +``USER`` is ignored and an arbitrary uid in GID 0 is assigned instead. The +symptom is not a degraded proxy, it is a proxy that does not serve at all. + +Booting the image the way that deployment does, and requiring it to answer a +request with a live database connection, is what catches the whole class: +a boot as the default uid, or one that reaches the internet, passes even when +the bake is unusable everywhere it actually ships. + +Gated on LITELLM_IMAGE (the tag of the image to exercise) so it is skipped in +the normal unit-test run and exercised only where an image has been built (the +image-scan workflow). Requires a working docker CLI. +""" + +import json +import shutil +import subprocess +import time +import uuid + +import os +import pytest + +IMAGE = os.getenv("LITELLM_IMAGE") +POSTGRES_IMAGE = os.getenv("LITELLM_TEST_POSTGRES_IMAGE", "postgres:16-alpine") +CURL_IMAGE = os.getenv("LITELLM_TEST_CURL_IMAGE", "curlimages/curl:8.11.1") +COMPONENT_PORT = os.getenv("LITELLM_COMPONENT_PORT", "4000") +NON_ROOT_UID = "12345:0" +STARTUP_TIMEOUT_SECONDS = int(os.getenv("LITELLM_COMPONENT_STARTUP_TIMEOUT", "180")) + +pytestmark = [ + pytest.mark.skipif(IMAGE is None, reason="requires a built image (set LITELLM_IMAGE)"), + pytest.mark.skipif(shutil.which("docker") is None, reason="requires the docker CLI"), +] + + +def _docker(*args: str, check: bool = True) -> subprocess.CompletedProcess: + return subprocess.run( + ["docker", *args], capture_output=True, text=True, check=check + ) + + +@pytest.fixture() +def offline_stack(): + """A component container and a fresh Postgres on a network with no egress. + + NON_ROOT_UID is an arbitrary uid in GID 0, the shape OpenShift restricted-v2 + assigns. Postgres and curl are pulled while egress still exists, because the + ``--internal`` network below has none: that is what makes a prisma engine + download (binaries.prisma.sh / npm) fail rather than mask a bake that is not + self-contained. + + The container runs with DISABLE_SCHEMA_UPDATE, since applying the schema is + the migration job's responsibility in this topology and needs the Prisma CLI + these images deliberately omit, and with LITELLM_LOCAL_MODEL_COST_MAP, or the + proxy spends the whole startup budget timing out on a cost-map fetch over the + network it does not have. + + Yields (network_name, component_container). Both are torn down afterwards. + """ + run_id = f"componentserve-{uuid.uuid4().hex[:8]}" + network = f"{run_id}-net" + pg = f"{run_id}-pg" + component = f"{run_id}-app" + + _docker("pull", "--quiet", POSTGRES_IMAGE) + _docker("pull", "--quiet", CURL_IMAGE) + _docker("network", "create", "--internal", network) + try: + _docker( + "run", "-d", "--name", pg, "--network", network, + "-e", "POSTGRES_PASSWORD=pw", "-e", "POSTGRES_DB=litellm", + POSTGRES_IMAGE, + ) + _wait_until_postgres_ready(pg) + assert IMAGE is not None + _docker( + "run", "-d", "--name", component, "--network", network, + "--user", NON_ROOT_UID, + "-e", f"DATABASE_URL=postgresql://postgres:pw@{pg}:5432/litellm", + "-e", "LITELLM_MASTER_KEY=sk-component-serve-test", + "-e", "DISABLE_SCHEMA_UPDATE=true", + "-e", "LITELLM_LOCAL_MODEL_COST_MAP=True", + IMAGE, + ) + yield network, component + finally: + _docker("logs", component, check=False) + _docker("rm", "-f", component, check=False) + _docker("rm", "-f", pg, check=False) + _docker("network", "rm", network, check=False) + + +def _wait_until_postgres_ready(pg: str, attempts: int = 60) -> None: + for _ in range(attempts): + running = _docker( + "ps", "--filter", f"name={pg}", "--filter", "status=running", + "--format", "{{.Names}}", check=False, + ).stdout + if pg not in running: + logs = _docker("logs", pg, check=False) + pytest.fail(f"postgres container is not running:\n{logs.stdout}\n{logs.stderr}") + ready = _docker( + "exec", pg, "pg_isready", "-U", "postgres", "-d", "litellm", check=False + ) + if ready.returncode == 0: + return + time.sleep(1) + pytest.fail(f"postgres never became ready after {attempts}s") + + +def _container_logs(container: str) -> str: + logs = _docker("logs", container, check=False) + return f"stdout:\n{logs.stdout}\nstderr:\n{logs.stderr}" + + +def _is_running(container: str) -> bool: + return bool( + _docker( + "ps", "--filter", f"name={container}", "--filter", "status=running", + "--format", "{{.Names}}", check=False, + ).stdout.strip() + ) + + +def _readiness(network: str, component: str) -> subprocess.CompletedProcess: + """Ask the component for its readiness, from a peer on the same egress-less network.""" + return _docker( + "run", "--rm", "--network", network, CURL_IMAGE, + "--silent", "--max-time", "10", + f"http://{component}:{COMPONENT_PORT}/health/readiness", + check=False, + ) + + +def test_component_serves_offline_as_non_root_uid(offline_stack): + """The component answers a request with a live DB connection, offline, as an arbitrary uid. + + On the pre-fix image this never gets a response: the engine baked under + /home/nonroot (mode 0700, owned by uid 65532) raises + ``PermissionError: .../query-engine-linux-...`` out of pathlib and uvicorn + reports ``Application startup failed. Exiting.``. A bake at the fixed, + world-readable /opt/prisma is what lets any uid start the client. + + `db: connected` is the load-bearing part of the assertion: it means the + query engine binary was found, executed, and reached Postgres. A liveness + probe alone would pass on an image whose engine never resolved. + """ + network, component = offline_stack + + deadline = time.time() + STARTUP_TIMEOUT_SECONDS + probe = None + while time.time() < deadline: + if not _is_running(component): + pytest.fail( + f"the component exited during startup as uid {NON_ROOT_UID} with no egress. " + "The prisma bake is not readable to a uid other than the one that built it, " + "so the proxy does not serve at all.\n" + f"{_container_logs(component)}" + ) + probe = _readiness(network, component) + if probe.returncode == 0 and probe.stdout.strip(): + break + time.sleep(2) + + assert probe is not None and probe.returncode == 0 and probe.stdout.strip(), ( + f"/health/readiness never answered within {STARTUP_TIMEOUT_SECONDS}s as uid " + f"{NON_ROOT_UID} with no egress.\n{_container_logs(component)}" + ) + + payload = json.loads(probe.stdout) + assert payload.get("db") == "connected", ( + f"the component answered but its database is {payload.get('db')!r}, so the baked " + f"query engine did not resolve as uid {NON_ROOT_UID}.\nresponse: {probe.stdout}\n" + f"{_container_logs(component)}" + ) From ae1d1cb05ebc921591c9717acbb5f007f7571805 Mon Sep 17 00:00:00 2001 From: Yassin Kortam Date: Fri, 7 Aug 2026 11:05:59 -0700 Subject: [PATCH 3/4] fix(http): stop pooled clients persisting cookies on the aiohttp jar too (#36149) #35978 stopped the pooled A2A client replaying one upstream's Set-Cookie to another by installing a blocking policy on that client's httpx cookie jar. That covers only one of the two jars on the request path. AiohttpTransport is the default transport unless it is explicitly disabled, and the aiohttp ClientSession behind it keeps its own cookie jar which no httpx-level assertion can observe, so the leak is still live on the default path: a live proxy on that commit still delivers agent-alpha's session cookie to agent-beta's card fetch and JSON-RPC call. The reason it looked fixed is that aiohttp's default CookieJar is built with unsafe=False and refuses to store cookies for IP hosts, so a proof addressed to 127.0.0.1 comes back clean whether or not that jar is blocked. Cookie persistence is now blocked where the clients are built rather than at one call site: blocked_cookie_jar() gives every httpx client, async and sync, a jar whose DefaultCookiePolicy(allowed_domains=()) rejects every domain in both directions, and both ClientSession constructions litellm owns, the transport's session factory and the proxy's shared startup session, get a DummyCookieJar. LiteLLM reads a response cookie nowhere, and an explicitly supplied Cookie header still goes out, so passthrough forwarding and an agent's extra_headers are unaffected. The A2A-scoped policy #35978 added is removed, since it is now dead. The two suites that drive the aiohttp session factory synchronously mock ClientSession because a real one needs a running event loop; DummyCookieJar has the same requirement, so they mock it for the same reason. --- litellm/a2a_protocol/main.py | 4 - litellm/llms/custom_httpx/http_handler.py | 15 +++- litellm/proxy/proxy_server.py | 4 +- tests/test_litellm/a2a_protocol/test_main.py | 47 ++++++------ .../test_aiohttp_cleanup_closed.py | 4 +- .../custom_httpx/test_aiohttp_so_keepalive.py | 6 +- .../llms/custom_httpx/test_http_handler.py | 74 +++++++++++++++++++ 7 files changed, 116 insertions(+), 38 deletions(-) diff --git a/litellm/a2a_protocol/main.py b/litellm/a2a_protocol/main.py index fe0e6837586..322393cd9c4 100644 --- a/litellm/a2a_protocol/main.py +++ b/litellm/a2a_protocol/main.py @@ -13,7 +13,6 @@ import asyncio import datetime import uuid from collections.abc import AsyncIterator, Coroutine -from http.cookiejar import DefaultCookiePolicy from typing import TYPE_CHECKING, Any, Final, Optional, cast import litellm @@ -80,8 +79,6 @@ from litellm.a2a_protocol.exceptions import A2ALocalhostURLError # Use our custom resolver instead of the default A2A SDK resolver A2ACardResolver: Final = LiteLLMA2ACardResolver -_BLOCK_ALL_COOKIES: Final = DefaultCookiePolicy(allowed_domains=()) - def _set_usage_on_logging_obj( kwargs: dict[str, Any], @@ -770,7 +767,6 @@ async def create_a2a_client( params={"timeout": timeout}, ) httpx_client: Final = _async_handler.client - httpx_client.cookies.jar.set_policy(_BLOCK_ALL_COOKIES) if extra_headers: verbose_proxy_logger.debug("A2A client created with extra_headers=%s", list(extra_headers.keys())) diff --git a/litellm/llms/custom_httpx/http_handler.py b/litellm/llms/custom_httpx/http_handler.py index d586156b625..9ada3674d33 100644 --- a/litellm/llms/custom_httpx/http_handler.py +++ b/litellm/llms/custom_httpx/http_handler.py @@ -8,11 +8,12 @@ import sys import threading import time from collections.abc import Callable, Mapping +from http.cookiejar import CookieJar, DefaultCookiePolicy from typing import TYPE_CHECKING, Any, Final, Optional import certifi import httpx -from aiohttp import ClientSession, TCPConnector +from aiohttp import ClientSession, DummyCookieJar, TCPConnector from httpx import USE_CLIENT_DEFAULT, AsyncHTTPTransport, HTTPTransport from httpx._types import RequestFiles @@ -144,6 +145,15 @@ def _handler_may_close_client(client_refcount: int, owns_client: bool) -> bool: return owns_client and client_refcount <= _CLIENT_REFCOUNT_WHEN_HANDLER_IS_SOLE_REFERRER +def blocked_cookie_jar() -> CookieJar: + """A jar that stores no response cookie and sends none, for httpx clients. + + LiteLLM's outbound clients are pooled and shared by every caller, so a cookie one + upstream sets would be replayed to every other upstream on a matching domain. + """ + return CookieJar(policy=DefaultCookiePolicy(allowed_domains=())) + + _STREAMING_ERROR_BODY_READ_TIMEOUT_SECONDS: Final = 5.0 _STREAMING_ERROR_BODY_READ_EXECUTOR: Final = concurrent.futures.ThreadPoolExecutor( max_workers=50, @@ -587,6 +597,7 @@ class AsyncHTTPHandler: verify=ssl_config, cert=cert, headers=default_headers, + cookies=blocked_cookie_jar(), follow_redirects=True, ) @@ -1063,6 +1074,7 @@ class AsyncHTTPHandler: def session_factory() -> ClientSession: return ClientSession( connector=TCPConnector(**transport_connector_kwargs), + cookie_jar=DummyCookieJar(), trust_env=trust_env, ) @@ -1132,6 +1144,7 @@ class HTTPHandler: verify=ssl_config, cert=cert, headers=default_headers, + cookies=blocked_cookie_jar(), follow_redirects=True, ) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 84f8685f5dc..d016eb57dd4 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -871,7 +871,7 @@ async def proxy_shutdown_event(): async def _initialize_shared_aiohttp_session(): """Initialize shared aiohttp session for connection reuse with connection limits.""" try: - from aiohttp import ClientSession, TCPConnector + from aiohttp import ClientSession, DummyCookieJar, TCPConnector from litellm.llms.custom_httpx.http_handler import ( _build_aiohttp_keepalive_socket_factory, @@ -892,7 +892,7 @@ async def _initialize_shared_aiohttp_session(): connector_kwargs["socket_factory"] = socket_factory connector: Final = TCPConnector(**connector_kwargs) - session: Final = ClientSession(connector=connector) + session: Final = ClientSession(connector=connector, cookie_jar=DummyCookieJar()) verbose_proxy_logger.info( "SESSION REUSE: Created shared aiohttp session for connection pooling (ID: %s, limit=%s, limit_per_host=%s)", diff --git a/tests/test_litellm/a2a_protocol/test_main.py b/tests/test_litellm/a2a_protocol/test_main.py index 59cb8c2c438..08f6b9f25bb 100644 --- a/tests/test_litellm/a2a_protocol/test_main.py +++ b/tests/test_litellm/a2a_protocol/test_main.py @@ -174,21 +174,15 @@ _RPC_REPLY = { _AGENT_A_HEADERS = {"x-agent-token": "token-for-a", "x-tenant": "tenant-a"} _AGENT_B_HEADERS = {"x-agent-token": "token-for-b", "x-tenant": "tenant-b"} -_UPSTREAM_SESSION_COOKIE = "a2a_session=only-agent-a-may-hold-this; Path=/" class _RequestRecorder: - """Records the headers httpx put on the wire, per outbound request. + """Records the headers httpx put on the wire, per outbound request.""" - ``cookie_from_tenant`` makes that tenant's agent answer with a Set-Cookie, standing in - for an upstream that issues a session cookie. - """ - - def __init__(self, cookie_from_tenant: str | None = None): + def __init__(self): self.card_requests = [] self.rpc_requests = [] self.client = None - self.cookie_from_tenant = cookie_from_tenant def __call__(self, request: httpx.Request) -> httpx.Response: headers = {k.lower(): v for k, v in request.headers.items()} @@ -196,8 +190,6 @@ class _RequestRecorder: self.card_requests.append(headers) return httpx.Response(200, json=_AGENT_CARD) self.rpc_requests.append(headers) - if self.cookie_from_tenant is not None and headers.get("x-tenant") == self.cookie_from_tenant: - return httpx.Response(200, json=_RPC_REPLY, headers={"set-cookie": _UPSTREAM_SESSION_COOKIE}) return httpx.Response(200, json=_RPC_REPLY) @@ -205,15 +197,14 @@ def _a2a_client_cache_key(timeout: float) -> str: return "async_httpx_client" + f"timeout_{timeout}" + httpxSpecialProvider.A2AProvider -async def _seed_shared_a2a_client(cookie_from_tenant: str | None = None) -> _RequestRecorder: +async def _seed_shared_a2a_client() -> _RequestRecorder: """Put the one A2A client the cache will hand out behind a mock transport. Seeding has to happen on the test's own event loop, because the client cache keys on it. The injected client is a real httpx.AsyncClient, so the merge of per-request - headers over client defaults, and httpx's own cookie handling, which is what these - tests are about, stay real. + headers over client defaults, which is what these tests are about, stays real. """ - recorder = _RequestRecorder(cookie_from_tenant=cookie_from_tenant) + recorder = _RequestRecorder() handler = AsyncHTTPHandler(timeout=DEFAULT_A2A_AGENT_TIMEOUT) owned_client = handler.client handler.client = httpx.AsyncClient(transport=httpx.MockTransport(recorder)) @@ -333,17 +324,21 @@ async def test_agent_card_fetch_carries_the_callers_headers(isolated_client_cach @pytest.mark.asyncio -async def test_one_agents_session_cookie_never_reaches_another_agent(isolated_client_cache): - """One pooled client is also one httpx cookie jar. httpx stores every Set-Cookie on the - client and replays it on any later request to a matching domain, so an agent's session - cookie would ride along on a different agent's call to the same host.""" - recorder = await _seed_shared_a2a_client(cookie_from_tenant="tenant-a") +async def test_the_pooled_a2a_client_arrives_with_cookie_persistence_disabled(isolated_client_cache): + """create_a2a_client takes its client from the shared builder rather than building one, + and the builder is what refuses to persist cookies. This pins the join between those + two facts, so the A2A path cannot quietly start acquiring a client that keeps a jar. - client_a = await create_a2a_client(base_url="http://127.0.0.1:9", extra_headers=_AGENT_A_HEADERS) - await _send_message(client_a, _send_request("a")) - client_b = await create_a2a_client(base_url="http://127.0.0.1:9", extra_headers=_AGENT_B_HEADERS) - await _send_message(client_b, _send_request("b")) + test_callers_with_different_headers_reuse_one_pooled_client pins the other half, that + create_a2a_client hands back exactly this cached client.""" + handler = get_async_httpx_client( + llm_provider=httpxSpecialProvider.A2AProvider, + params={"timeout": DEFAULT_A2A_AGENT_TIMEOUT}, + ) + request = httpx.Request("GET", "https://agent-a.example.com/") + handler.client.cookies.extract_cookies( + httpx.Response(200, headers={"set-cookie": "SESSION=only-agent-a-may-hold-this"}, request=request) + ) - assert dict(recorder.client.cookies) == {}, "the shared client kept an agent's session cookie" - assert "cookie" not in recorder.card_requests[-1] - assert "cookie" not in recorder.rpc_requests[-1] + assert dict(handler.client.cookies) == {}, "the pooled A2A client kept an upstream's cookie" + await handler.close() diff --git a/tests/test_litellm/llms/custom_httpx/test_aiohttp_cleanup_closed.py b/tests/test_litellm/llms/custom_httpx/test_aiohttp_cleanup_closed.py index f279acfd60c..82010e82cea 100644 --- a/tests/test_litellm/llms/custom_httpx/test_aiohttp_cleanup_closed.py +++ b/tests/test_litellm/llms/custom_httpx/test_aiohttp_cleanup_closed.py @@ -13,7 +13,7 @@ def test_create_aiohttp_transport_sets_enable_cleanup_closed_when_needed(monkeyp ) as mock_tcp_connector: with patch.object( http_handler_module, "ClientSession", return_value=session_mock - ): + ), patch.object(http_handler_module, "DummyCookieJar", return_value=MagicMock(name="cookie_jar")): transport = http_handler_module.AsyncHTTPHandler._create_aiohttp_transport( shared_session=None ) @@ -36,7 +36,7 @@ def test_create_aiohttp_transport_omits_enable_cleanup_closed_when_not_needed( ) as mock_tcp_connector: with patch.object( http_handler_module, "ClientSession", return_value=session_mock - ): + ), patch.object(http_handler_module, "DummyCookieJar", return_value=MagicMock(name="cookie_jar")): transport = http_handler_module.AsyncHTTPHandler._create_aiohttp_transport( shared_session=None ) diff --git a/tests/test_litellm/llms/custom_httpx/test_aiohttp_so_keepalive.py b/tests/test_litellm/llms/custom_httpx/test_aiohttp_so_keepalive.py index 5a37e681c4a..0065bf8f4ef 100644 --- a/tests/test_litellm/llms/custom_httpx/test_aiohttp_so_keepalive.py +++ b/tests/test_litellm/llms/custom_httpx/test_aiohttp_so_keepalive.py @@ -34,7 +34,7 @@ def test_socket_factory_omitted_when_disabled(monkeypatch): ) as mock_tcp_connector: with patch.object( http_handler_module, "ClientSession", return_value=session_mock - ): + ), patch.object(http_handler_module, "DummyCookieJar", return_value=MagicMock(name="cookie_jar")): _invoke_connector_factory(http_handler_module) assert mock_tcp_connector.call_count >= 1 @@ -55,7 +55,7 @@ def test_socket_factory_attached_when_enabled(monkeypatch): ) as mock_tcp_connector: with patch.object( http_handler_module, "ClientSession", return_value=session_mock - ): + ), patch.object(http_handler_module, "DummyCookieJar", return_value=MagicMock(name="cookie_jar")): _invoke_connector_factory(http_handler_module) assert mock_tcp_connector.call_count >= 1 @@ -77,7 +77,7 @@ def test_socket_factory_skipped_on_old_aiohttp(monkeypatch): ) as mock_tcp_connector: with patch.object( http_handler_module, "ClientSession", return_value=session_mock - ): + ), patch.object(http_handler_module, "DummyCookieJar", return_value=MagicMock(name="cookie_jar")): _invoke_connector_factory(http_handler_module) assert mock_tcp_connector.call_count >= 1 diff --git a/tests/test_litellm/llms/custom_httpx/test_http_handler.py b/tests/test_litellm/llms/custom_httpx/test_http_handler.py index 86d097c6123..fa1c7308c6f 100644 --- a/tests/test_litellm/llms/custom_httpx/test_http_handler.py +++ b/tests/test_litellm/llms/custom_httpx/test_http_handler.py @@ -1171,3 +1171,77 @@ async def test_client_handed_out_by_async_cache_survives_eviction_and_collection assert not consumer_client.is_closed await consumer_client.aclose() + + +_SET_COOKIE = "SESSION=upstream-a-secret; Path=/" + + +def _cookie_recorder(): + """A transport that hands out a Set-Cookie once, and records what comes back.""" + seen = [] + + def handler(request: httpx.Request) -> httpx.Response: + seen.append(request.headers.get("cookie")) + if request.url.path == "/set": + return httpx.Response(200, headers={"set-cookie": _SET_COOKIE}) + return httpx.Response(200) + + return handler, seen + + +@pytest.mark.asyncio +async def test_async_client_never_replays_one_upstreams_cookie_to_another(): + """LiteLLM's async clients are pooled and shared by every caller, so a cookie one + upstream sets would be attached to every later request on a matching domain, reaching + a different tenant's upstream. The client must persist no response cookie.""" + handler, seen = _cookie_recorder() + http_handler = AsyncHTTPHandler() + client = http_handler.client + client._transport = httpx.MockTransport(handler) + + await client.get("https://upstream-a.example.com/set") + await client.get("https://upstream-b.example.com/rpc") + await client.aclose() + + assert dict(client.cookies) == {}, "the shared client stored an upstream's cookie" + assert seen == [None, None] + + +def test_sync_client_never_replays_one_upstreams_cookie_to_another(): + """Same invariant on the sync client, which is pooled the same way.""" + handler, seen = _cookie_recorder() + http_handler = HTTPHandler() + client = http_handler.client + client._transport = httpx.MockTransport(handler) + + client.get("https://upstream-a.example.com/set") + client.get("https://upstream-b.example.com/rpc") + client.close() + + assert dict(client.cookies) == {} + assert seen == [None, None] + + +@pytest.mark.asyncio +async def test_aiohttp_session_never_replays_one_upstreams_cookie_to_another(): + """The httpx jar is not the only one. AiohttpTransport is litellm's default transport + and the aiohttp ClientSession keeps its own cookie jar, which httpx-level assertions + cannot see, so blocking only the httpx jar leaves the leak intact on the real path. + + aiohttp's default jar refuses cookies for IP hosts, so this drives a hostname. An + IP-addressed check passes whether or not the session jar is blocked.""" + from aiohttp import DummyCookieJar + from yarl import URL + + http_handler = AsyncHTTPHandler(timeout=61.0) + transport = http_handler.client._transport + assert isinstance(transport, LiteLLMAiohttpTransport), "aiohttp is no longer the default transport" + + session = transport.client() if callable(transport.client) else transport.client + jar = session.cookie_jar + assert isinstance(jar, DummyCookieJar) + + jar.update_cookies({"SESSION": "upstream-a-secret"}, URL("https://upstream-a.example.com")) + assert len(jar) == 0 + assert dict(jar.filter_cookies(URL("https://upstream-a.example.com"))) == {} + await session.close() From 330a09235d1a8ca5cbd35acae2adc0e8319f4cf8 Mon Sep 17 00:00:00 2001 From: Yassin Kortam Date: Fri, 7 Aug 2026 11:07:09 -0700 Subject: [PATCH 4/4] fix(router): bound fallback-walk work and error-log volume (#36148) --- litellm/constants.py | 1 + litellm/router.py | 14 +- litellm/router_utils/common_utils.py | 17 ++ .../router_utils/fallback_event_handlers.py | 76 ++++++- litellm/types/utils.py | 1 + .../test_fallback_event_handlers.py | 204 ++++++++++++++++++ .../test_router_utils_common_utils.py | 26 +++ tests/test_litellm/test_router.py | 135 ++++++++++++ 8 files changed, 469 insertions(+), 5 deletions(-) diff --git a/litellm/constants.py b/litellm/constants.py index 75f12190b3e..6f0e9e7afe2 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -7,6 +7,7 @@ from litellm.litellm_core_utils.env_utils import get_env_int, get_env_int_or_non DEFAULT_HEALTH_CHECK_PROMPT: Final = str(os.getenv("DEFAULT_HEALTH_CHECK_PROMPT", "test from litellm")) AZURE_DEFAULT_RESPONSES_API_VERSION: Final = str(os.getenv("AZURE_DEFAULT_RESPONSES_API_VERSION", "preview")) ROUTER_MAX_FALLBACKS: Final = int(os.getenv("ROUTER_MAX_FALLBACKS", 5)) +ROUTER_FALLBACK_ERROR_DETAIL_MAX_CHARS: Final = 2000 DEFAULT_BATCH_SIZE: Final = int(os.getenv("DEFAULT_BATCH_SIZE", 512)) DEFAULT_FLUSH_INTERVAL_SECONDS: Final = int(os.getenv("DEFAULT_FLUSH_INTERVAL_SECONDS", 5)) DEFAULT_S3_FLUSH_INTERVAL_SECONDS: Final = int(os.getenv("DEFAULT_S3_FLUSH_INTERVAL_SECONDS", 10)) diff --git a/litellm/router.py b/litellm/router.py index 1edb80da7ce..1a76eb5d59b 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -109,6 +109,7 @@ from litellm.router_utils.common_utils import ( filter_team_based_models, filter_web_search_deployments, resolve_model_group_alias, + truncate_fallback_error_detail, ) from litellm.router_utils.cooldown_cache import CooldownCache from litellm.router_utils.cooldown_handlers import ( @@ -342,6 +343,12 @@ def _replay_live_router_model_cost() -> None: set_live_deployment_replay(_replay_live_router_model_cost) +# Kwargs that log_retry must not copy into a retry breadcrumb. The breadcrumbs reach spend +# logs and logging callbacks, and these carry either the request payload or router-internal +# walk state rather than anything that identifies the failed attempt. +RETRY_BREADCRUMB_EXCLUDED_KWARGS: Final = frozenset(("messages", "original_function", "attempted_targets")) + + class Router: model_names: set = set() cache_responses: bool | None = False @@ -6361,17 +6368,16 @@ class Router: return response except Exception as new_exception: parent_otel_span: Final = _get_parent_otel_span_from_kwargs(kwargs) - fallback_failure_exception_str = redact_string(str(new_exception)) + fallback_failure_exception_str = truncate_fallback_error_detail(redact_string(str(new_exception))) cooldown_info: Final = await _async_get_cooldown_deployments_with_debug_info( litellm_router_instance=self, parent_otel_span=parent_otel_span, ) verbose_router_logger.error( "litellm.router.py::async_function_with_fallbacks() - " - "Error occurred while trying to do fallbacks - %s\n%s\n" + "Error occurred while trying to do fallbacks - %s\n" "Debug Information:\nCooldown Deployments=%s", fallback_failure_exception_str, - redact_string(traceback.format_exc()), cooldown_info, ) @@ -7162,7 +7168,7 @@ class Router: k, v, ) in kwargs.items(): # log everything in kwargs except the old previous_models value - prevent nesting - if k not in [_metadata_var, "messages", "original_function"]: + if k != _metadata_var and k not in RETRY_BREADCRUMB_EXCLUDED_KWARGS: previous_model[k] = v elif k == _metadata_var and isinstance(v, dict): previous_model[_metadata_var] = {} diff --git a/litellm/router_utils/common_utils.py b/litellm/router_utils/common_utils.py index 1c628ef9ee3..6fad2dd31e9 100644 --- a/litellm/router_utils/common_utils.py +++ b/litellm/router_utils/common_utils.py @@ -7,6 +7,7 @@ if TYPE_CHECKING: from litellm.types.llms.openai import OpenAIFileObject from litellm._logging import verbose_logger +from litellm.constants import ROUTER_FALLBACK_ERROR_DETAIL_MAX_CHARS from litellm.exceptions import BadRequestError from litellm.types.router import CredentialLiteLLMParams @@ -43,6 +44,22 @@ def resolve_model_group_alias(model_group_alias: object, model: str) -> str | No return target +def truncate_fallback_error_detail(detail: str) -> str: + """ + Bound a fallback failure detail before it is logged or appended to an exception message. + + Each level of the fallback walk records the failure of the level below it, so an + untruncated detail carries every nested failure with it and grows superlinearly with + the number of attempted model groups. One deterministic pre-network failure walked + through a small fallback graph is enough to turn that into hundreds of megabytes of + output on the event-loop thread, which starves the process that produced it. + """ + if len(detail) <= ROUTER_FALLBACK_ERROR_DETAIL_MAX_CHARS: + return detail + dropped: Final = len(detail) - ROUTER_FALLBACK_ERROR_DETAIL_MAX_CHARS + return f"{detail[:ROUTER_FALLBACK_ERROR_DETAIL_MAX_CHARS]}... [truncated {dropped} characters]" + + def get_litellm_params_sensitive_credential_hash(litellm_params: dict) -> str: """ Hash of the credential params, used for mapping the file id to the right model diff --git a/litellm/router_utils/fallback_event_handlers.py b/litellm/router_utils/fallback_event_handlers.py index 1c6bb52ccb8..ef48ccc821e 100644 --- a/litellm/router_utils/fallback_event_handlers.py +++ b/litellm/router_utils/fallback_event_handlers.py @@ -1,3 +1,6 @@ +import hashlib +import json +from dataclasses import dataclass from enum import Enum from typing import TYPE_CHECKING, Any, Final @@ -19,6 +22,52 @@ else: LitellmRouter = Any +def fallback_attempt_key(fallback_target: object) -> str | None: + """ + Identity of one fallback attempt, so the same attempt is never made twice per request. + + A bare model group name and a `{"model": name}` entry describe the same attempt. An + entry carrying anything else describes a different one and keeps its own identity: a + client-side fallback list overrides request params such as `messages`, and the router + re-targets the group that just failed by attaching `_target_order` or + `_excluded_deployment_ids` to select a different set of deployments inside it. The + payload is hashed rather than kept, so a large `messages` override does not make the + request hold a second copy of itself. + + Returns None for a shape with no usable identity, which is never skipped. + """ + if isinstance(fallback_target, str): + return fallback_target + if not isinstance(fallback_target, dict): + return None + model: Final = fallback_target.get("model") + if tuple(fallback_target) == ("model",) and isinstance(model, str): + return model + serialized: Final = json.dumps(fallback_target, sort_keys=True, default=str) + return hashlib.sha256(serialized.encode()).hexdigest() + + +@dataclass(slots=True) +class AttemptedFallbackTargets: + """ + The fallback attempts a single request has already made. + + One instance is created on the first fallback hop and shared by reference for the rest + of the walk, so an attempt made in one branch is not repeated in a sibling branch. + Without it the walk enumerates paths rather than attempts: a fallback graph containing + a cycle retries one deterministic failure once per path through the cycle, and a + client-side fallback list is re-walked at every level of the recursion. + """ + + keys: frozenset[str] = frozenset() + + def __contains__(self, key: str) -> bool: + return key in self.keys + + def record(self, key: str) -> None: + self.keys = self.keys | frozenset((key,)) + + def _check_stripped_model_group(model_group: str, fallback_key: str) -> bool: """ Handles wildcard routing scenario @@ -106,7 +155,14 @@ async def run_async_fallback( fallback_model_group: List[str] of fallback model groups. example: ["gpt-4", "gpt-3.5-turbo"] original_model_group: The original model group. example: "gpt-3.5-turbo" original_exception: The original exception. - **kwargs: Keyword arguments. + **kwargs: Keyword arguments. `attempted_targets` carries the fallback attempts + already made for this request, created on the first hop and shared by reference + for the rest of the walk. A target already in it is skipped, so neither a + fallback graph that loops back on itself nor a client-side fallback list + re-walked at each level can repeat an attempt that has already failed. Identity + comes from `fallback_attempt_key`, so an entry that overrides request params or + re-targets the failed group with a different deployment selection stays distinct + from a bare name. Returns: The response from the successful fallback model group. @@ -120,10 +176,27 @@ async def run_async_fallback( error_from_fallbacks = original_exception fallback_errors = (get_fallback_error_info(original_exception),) + # Read out of kwargs and narrowed here rather than declared as a parameter: every caller + # reaches this function by spreading a loosely-typed kwargs dict, so a declared parameter + # would carry an annotation that no call site can actually be checked against. + carried_targets: Final = kwargs.get("attempted_targets") + attempted: Final = ( + carried_targets if isinstance(carried_targets, AttemptedFallbackTargets) else AttemptedFallbackTargets() + ) + attempted.record(original_model_group) for mg in fallback_model_group: if mg == original_model_group: continue + attempt_key = fallback_attempt_key(mg) + if attempt_key is not None: + if attempt_key in attempted: + verbose_router_logger.info( + "Skipping fallback to model_group = %s, already attempted for this request", + mask_sensitive_structure(mg), + ) + continue + attempted.record(attempt_key) try: # LOGGING kwargs = litellm_router.log_retry(kwargs=kwargs, e=original_exception) @@ -138,6 +211,7 @@ async def run_async_fallback( fallback_depth = fallback_depth + 1 kwargs["fallback_depth"] = fallback_depth kwargs["max_fallbacks"] = max_fallbacks + kwargs["attempted_targets"] = attempted if include_fallback_errors: kwargs["include_fallback_errors"] = include_fallback_errors response = await litellm_router.async_function_with_fallbacks(*args, **kwargs) diff --git a/litellm/types/utils.py b/litellm/types/utils.py index 329b10e72e7..35d4250782f 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -3479,6 +3479,7 @@ all_litellm_params = ( "user_continue_message", "fallback_depth", "max_fallbacks", + "attempted_targets", "max_budget", "budget_duration", "use_in_pass_through", diff --git a/tests/test_litellm/router_utils/test_fallback_event_handlers.py b/tests/test_litellm/router_utils/test_fallback_event_handlers.py index 98a34de295c..e2348d28701 100644 --- a/tests/test_litellm/router_utils/test_fallback_event_handlers.py +++ b/tests/test_litellm/router_utils/test_fallback_event_handlers.py @@ -3,6 +3,8 @@ import json import pytest from litellm.router_utils.fallback_event_handlers import ( + AttemptedFallbackTargets, + fallback_attempt_key, get_fallback_model_group, run_async_fallback, ) @@ -142,6 +144,208 @@ async def test_run_async_fallback_skips_original_model_group(): assert response._hidden_params["additional_headers"]["x-litellm-attempted-fallbacks"] == 1 +class RecordingFailRouter: + def __init__(self): + self.attempted_models = [] + + def log_retry(self, kwargs, e): + return kwargs + + async def async_function_with_fallbacks(self, *args, **kwargs): + self.attempted_models.append(kwargs.get("model")) + raise RuntimeError("fallback model also failed") + + +@pytest.mark.asyncio +async def test_run_async_fallback_skips_model_group_already_attempted(): + """A fallback graph that loops back on itself must not re-attempt a model group that + already failed for this request. Every group in a cycle fails identically, so + revisiting one multiplies the work and the error output without any chance of + succeeding.""" + router = RecordingFailRouter() + + with pytest.raises(RuntimeError, match="original failed"): + await run_async_fallback( + litellm_router=router, + fallback_model_group=["already-attempted"], + original_model_group="primary-model", + original_exception=RuntimeError("original failed"), + max_fallbacks=3, + fallback_depth=0, + attempted_targets=AttemptedFallbackTargets(frozenset({"already-attempted"})), + ) + + assert router.attempted_models == [] + + +@pytest.mark.asyncio +async def test_run_async_fallback_attempts_a_repeated_target_once(): + router = RecordingFailRouter() + + with pytest.raises(RuntimeError, match="fallback model also failed"): + await run_async_fallback( + litellm_router=router, + fallback_model_group=["fallback-model", "fallback-model", "other-model"], + original_model_group="primary-model", + original_exception=RuntimeError("original failed"), + max_fallbacks=5, + fallback_depth=0, + ) + + assert router.attempted_models == ["fallback-model", "other-model"] + + +@pytest.mark.asyncio +async def test_run_async_fallback_forwards_attempted_model_groups_to_nested_call(): + """The nested call is where the next hop of the walk decides what to skip, so the + accumulated set has to reach it, carrying both the group that just failed and the + target being attempted.""" + router = RecordingRouter() + + await run_async_fallback( + litellm_router=router, + fallback_model_group=["fallback-model"], + original_model_group="primary-model", + original_exception=RuntimeError("original failed"), + max_fallbacks=3, + fallback_depth=0, + attempted_targets=AttemptedFallbackTargets(frozenset({"earlier-model"})), + ) + + assert router.received_kwargs["attempted_targets"].keys == frozenset( + {"earlier-model", "primary-model", "fallback-model"} + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "entry", + [ + {"model": "primary-model", "_target_order": 2}, + {"model": "primary-model", "_excluded_deployment_ids": ["dep-1"]}, + ], +) +async def test_run_async_fallback_still_retargets_the_same_group_via_dict_entry(entry): + """Order-based fallback and weighted intra-group failover both re-target the group that + just failed, selecting a different set of deployments inside it. Those entries are dicts + rather than plain names and must survive a guard that skips already-attempted names.""" + router = RecordingRouter() + + await run_async_fallback( + litellm_router=router, + fallback_model_group=[entry], + original_model_group="primary-model", + original_exception=RuntimeError("original failed"), + max_fallbacks=3, + fallback_depth=0, + attempted_targets=AttemptedFallbackTargets(frozenset({"primary-model"})), + ) + + assert router.received_kwargs["model"] == "primary-model" + + +@pytest.mark.asyncio +async def test_run_async_fallback_skips_a_repeated_dict_target(): + """A client-side fallback list names its targets with dicts, and that list is re-walked + at every level of the recursion, so an entry that carries no request override has to be + recognised as the same attempt as the bare name.""" + router = RecordingFailRouter() + + with pytest.raises(RuntimeError, match="original failed"): + await run_async_fallback( + litellm_router=router, + fallback_model_group=[{"model": "already-attempted"}], + original_model_group="primary-model", + original_exception=RuntimeError("original failed"), + max_fallbacks=3, + fallback_depth=0, + attempted_targets=AttemptedFallbackTargets(frozenset({"already-attempted"})), + ) + + assert router.attempted_models == [] + + +@pytest.mark.asyncio +async def test_run_async_fallback_attempts_a_repeated_dict_target_once(): + router = RecordingFailRouter() + entry = {"model": "fallback-model", "messages": [{"role": "user", "content": "shorter"}]} + + with pytest.raises(RuntimeError, match="fallback model also failed"): + await run_async_fallback( + litellm_router=router, + fallback_model_group=[entry, entry, {"model": "other-model"}], + original_model_group="primary-model", + original_exception=RuntimeError("original failed"), + max_fallbacks=5, + fallback_depth=0, + ) + + assert router.attempted_models == ["fallback-model", "other-model"] + + +@pytest.mark.asyncio +async def test_run_async_fallback_keeps_a_request_override_distinct_from_the_bare_name(): + """The documented use of the client-side form is to retry a group with different request + params, so an entry carrying an override must survive even when the bare name of that + same group has already been attempted.""" + router = RecordingFailRouter() + + with pytest.raises(RuntimeError, match="fallback model also failed"): + await run_async_fallback( + litellm_router=router, + fallback_model_group=[ + {"model": "already-attempted", "messages": [{"role": "user", "content": "shorter"}]} + ], + original_model_group="primary-model", + original_exception=RuntimeError("original failed"), + max_fallbacks=3, + fallback_depth=0, + attempted_targets=AttemptedFallbackTargets(frozenset({"already-attempted"})), + ) + + assert router.attempted_models == ["already-attempted"] + + +@pytest.mark.parametrize( + "target, expected", + [ + ("group-a", "group-a"), + ({"model": "group-a"}, "group-a"), + (None, None), + (["group-a"], None), + ], +) +def test_fallback_attempt_key_identity(target, expected): + """A bare name and a `{"model": name}` entry are the same attempt. A shape with no + usable identity returns None and is never skipped, so an unrecognised entry keeps + today's behaviour rather than being silently dropped.""" + assert fallback_attempt_key(target) == expected + + +def test_fallback_attempt_key_gives_a_param_only_entry_its_own_identity(): + """An entry with no `model` re-targets the group currently being attempted with + different request params, so it is a distinct attempt and still needs an identity.""" + key = fallback_attempt_key({"messages": [{"role": "user", "content": "shorter"}]}) + + assert key is not None + assert key != fallback_attempt_key({"messages": [{"role": "user", "content": "other"}]}) + + +def test_fallback_attempt_key_separates_overrides_from_the_bare_name(): + bare = fallback_attempt_key("group-a") + override = fallback_attempt_key({"model": "group-a", "messages": [{"role": "user", "content": "x"}]}) + other_override = fallback_attempt_key({"model": "group-a", "messages": [{"role": "user", "content": "y"}]}) + order_retarget = fallback_attempt_key({"model": "group-a", "_target_order": 2}) + + assert len({bare, override, other_override, order_retarget}) == 4 + + +def test_fallback_attempt_key_is_stable_across_key_order(): + assert fallback_attempt_key({"model": "group-a", "_target_order": 2}) == fallback_attempt_key( + {"_target_order": 2, "model": "group-a"} + ) + + def test_get_fallback_model_group_does_not_mutate_fallbacks(): """A string fallback must be resolved without mutating the caller's fallbacks list, which is the live router config shared across requests.""" diff --git a/tests/test_litellm/router_utils/test_router_utils_common_utils.py b/tests/test_litellm/router_utils/test_router_utils_common_utils.py index 7d453c72652..0d063ad14f5 100644 --- a/tests/test_litellm/router_utils/test_router_utils_common_utils.py +++ b/tests/test_litellm/router_utils/test_router_utils_common_utils.py @@ -4,6 +4,7 @@ from unittest.mock import Mock import pytest from litellm import Router +from litellm.constants import ROUTER_FALLBACK_ERROR_DETAIL_MAX_CHARS from litellm.proxy._types import UserAPIKeyAuth from litellm.router_utils.common_utils import ( _deployment_supports_web_search, @@ -11,6 +12,7 @@ from litellm.router_utils.common_utils import ( filter_team_based_models, filter_web_search_deployments, resolve_model_group_alias, + truncate_fallback_error_detail, ) @@ -558,3 +560,27 @@ class TestResolveModelGroupAlias: assert router._get_model_from_alias("group-a") == "group-b" assert router._get_model_from_alias("group-item") == "group-b" assert router._get_model_from_alias("group-b") is None + + +class TestTruncateFallbackErrorDetail: + def test_short_detail_is_returned_unchanged(self): + assert truncate_fallback_error_detail("boom") == "boom" + + def test_detail_at_the_limit_is_returned_unchanged(self): + detail = "x" * ROUTER_FALLBACK_ERROR_DETAIL_MAX_CHARS + assert truncate_fallback_error_detail(detail) == detail + + def test_long_detail_is_bounded_and_reports_what_was_dropped(self): + detail = "x" * (ROUTER_FALLBACK_ERROR_DETAIL_MAX_CHARS + 500) + + truncated = truncate_fallback_error_detail(detail) + + assert truncated.startswith("x" * ROUTER_FALLBACK_ERROR_DETAIL_MAX_CHARS) + assert truncated.endswith("... [truncated 500 characters]") + assert len(truncated) < len(detail) + + def test_a_megabyte_of_detail_comes_back_small(self): + """The detail is what a fallback level records about the level below it, so it has + to stay small enough that a walk over many model groups cannot compound it into an + output volume that starves the process.""" + assert len(truncate_fallback_error_detail("x" * 1_000_000)) < 3_000 diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py index 4a3395a7d3f..da2d78edb73 100644 --- a/tests/test_litellm/test_router.py +++ b/tests/test_litellm/test_router.py @@ -1,6 +1,7 @@ import asyncio import copy import json +import logging import os import sys from unittest.mock import AsyncMock, MagicMock, patch @@ -14,6 +15,7 @@ sys.path.insert( import litellm from litellm.exceptions import MidStreamFallbackError +from litellm.integrations.custom_logger import CustomLogger def test_update_kwargs_does_not_mutate_defaults_and_merges_metadata(): @@ -7415,3 +7417,136 @@ class TestAutoRouterMaxInputCharsWiring: router = self._router() assert self._registered_auto_router(router).max_input_chars == DEFAULT_AUTO_ROUTER_MAX_INPUT_CHARS + + +class _LogCapture(logging.Handler): + def __init__(self, level): + super().__init__(level=level) + self._level = level + self.messages = [] + + def emit(self, record): + if record.levelno == self._level: + self.messages.append(record.getMessage()) + + +class _FallbackAttemptRecorder(CustomLogger): + def __init__(self): + super().__init__() + self.failed_targets = [] + + async def log_failure_fallback_event(self, original_model_group, kwargs, original_exception): + self.failed_targets.append(kwargs.get("model")) + + +def _cyclic_fallback_router(num_retries=0): + groups = ["group-a", "group-b", "group-c", "group-d"] + return litellm.Router( + model_list=[ + { + "model_name": group, + "litellm_params": { + "model": "openai/gpt-4o-mini", + "api_key": "sk-fake", + "mock_response": "litellm.InternalServerError", + }, + } + for group in groups + ], + fallbacks=[ + {"group-a": ["group-b", "group-c"]}, + {"group-b": ["group-a", "group-c"]}, + {"group-c": ["group-d"]}, + {"group-d": ["group-b", "group-a"]}, + ], + num_retries=num_retries, + ) + + +async def _drive_cyclic_fallback(router, capture, recorder=None, **request_kwargs): + router_logger = logging.getLogger("LiteLLM Router") + previous_level = router_logger.level + router_logger.setLevel(capture.level) + router_logger.addHandler(capture) + if recorder is not None: + litellm.callbacks.append(recorder) + try: + with pytest.raises(litellm.InternalServerError): + await router.acompletion( + model="group-a", messages=[{"role": "user", "content": "hi"}], **request_kwargs + ) + finally: + router_logger.removeHandler(capture) + router_logger.setLevel(previous_level) + if recorder is not None: + litellm.callbacks.remove(recorder) + + +@pytest.mark.asyncio +async def test_cyclic_fallback_graph_does_not_amplify_one_request(): + """A fallback graph whose entries loop back on each other is easy to build by accident, + and every group in the loop fails identically on a deterministic error, so the walk must + not revisit a group and must not re-emit a growing chained traceback at each level. Left + unbounded, one request blocks the event loop long enough for health probes to fail.""" + recorder = _FallbackAttemptRecorder() + capture = _LogCapture(logging.ERROR) + + await _drive_cyclic_fallback(_cyclic_fallback_router(), capture, recorder) + + assert sorted(set(recorder.failed_targets)) == ["group-b", "group-c", "group-d"] + assert len(recorder.failed_targets) == len(set(recorder.failed_targets)) + assert not any("Traceback (most recent call last)" in message for message in capture.messages) + assert sum(len(message) for message in capture.messages) < 5_000 + + +@pytest.mark.asyncio +async def test_retry_breadcrumbs_do_not_carry_the_walk_state(): + """log_retry copies every kwarg into previous_models, which reaches spend logs and + logging callbacks. The set of already-attempted groups is router-internal walk state + with no diagnostic value there, and it is the one entry that is not a plain scalar. + A retry has to be configured for the walk state to reach log_retry at all.""" + router = _cyclic_fallback_router(num_retries=1) + capture = _LogCapture(logging.ERROR) + + await _drive_cyclic_fallback(router, capture) + + assert router.previous_models, "no retry breadcrumbs were recorded" + assert any( + "fallback_depth" in breadcrumb for breadcrumb in router.previous_models + ), "no breadcrumb carried router walk state, so this test cannot see the leak" + for breadcrumb in router.previous_models: + assert "attempted_targets" not in breadcrumb + + +@pytest.mark.asyncio +async def test_fallback_traceback_stays_available_at_debug_level(): + """Dropping the stack from the ERROR line is only safe because the fallback path still + emits it once per level at DEBUG, which is what an operator needs to diagnose why every + fallback failed. This pins that remaining debug traceback.""" + capture = _LogCapture(logging.DEBUG) + + await _drive_cyclic_fallback(_cyclic_fallback_router(), capture) + + assert any("Traceback (most recent call last)" in message for message in capture.messages) + + +@pytest.mark.asyncio +async def test_fallback_failure_detail_from_upstream_is_bounded(): + """The detail each level records about the level below it is attacker-influenced, since + it carries whatever the upstream error said. It has to be bounded on its own, so a walk + over several groups cannot compound one large message into the log or into the message + handed back to the caller.""" + huge_message = "z" * 50_000 + capture = _LogCapture(logging.ERROR) + + await _drive_cyclic_fallback( + _cyclic_fallback_router(), + capture, + mock_response=litellm.InternalServerError( + message=huge_message, llm_provider="openai", model="group-a" + ), + ) + + assert capture.messages, "the fallback failure path did not log at ERROR" + assert huge_message not in "".join(capture.messages) + assert max(len(message) for message in capture.messages) < 5_000