From fbcdacc44659813b05dd2f9d145ec63575e2843d Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Tue, 21 Apr 2026 23:29:17 -0700 Subject: [PATCH 01/57] [Fix] Proxy: reconnect Prisma DB without blocking the event loop When the DB becomes unreachable the reconnect path calls `prisma.disconnect()`, which ultimately invokes prisma-client-py's synchronous `subprocess.Popen.wait()` on the query engine subprocess. That call does not yield to asyncio, so the event loop freezes for however long the Rust engine takes to shut down (30-120+ seconds in production when the engine is stuck on TCP close). During the freeze `/health/liveliness` becomes unresponsive, and in Kubernetes the liveness probe fails and the pod is SIGKILL'd. Replace `disconnect()` in the reconnect paths with a direct, non-blocking kill of the engine subprocess (SIGTERM -> 0.5s asyncio-yielding sleep -> SIGKILL) followed by a fresh Prisma client and a new `connect()`. Both `recreate_prisma_client` and the formerly-separate "direct reconnect" path go through the same kill-then-recreate flow. Also validate `_get_engine_pid` returns an int (defensive; prevents a MagicMock leak under unit-test mocking). Tests that encoded the old blocking behavior are updated or removed; the deleted `test_lightweight_reconnect_skips_kill_on_successful_disconnect` invariant ("don't kill on successful disconnect") was part of the bug. --- litellm/proxy/db/prisma_client.py | 19 +-- litellm/proxy/utils.py | 23 ++-- .../proxy/test_prisma_engine_watchdog.py | 77 +++++++----- .../proxy/db/test_prisma_self_heal.py | 119 +++++++++--------- 4 files changed, 135 insertions(+), 103 deletions(-) diff --git a/litellm/proxy/db/prisma_client.py b/litellm/proxy/db/prisma_client.py index 73735796eb..a8942f92a1 100644 --- a/litellm/proxy/db/prisma_client.py +++ b/litellm/proxy/db/prisma_client.py @@ -52,7 +52,9 @@ class PrismaWrapper: engine = self._original_prisma._engine process = getattr(engine, "process", None) if engine is not None else None if process is not None: - return process.pid + pid = process.pid + if isinstance(pid, int): + return pid except (AttributeError, TypeError): pass return 0 @@ -217,15 +219,18 @@ class PrismaWrapper: async def recreate_prisma_client( self, new_db_url: str, http_client: Optional[Any] = None ): - """Disconnect and reconnect the Prisma client with a new database URL.""" + """Disconnect and reconnect the Prisma client with a new database URL. + + Kills the old engine subprocess directly (SIGTERM → SIGKILL) rather than + calling `disconnect()`. prisma-client-py's `disconnect()` calls a + synchronous `subprocess.Popen.wait()` that can freeze the asyncio event + loop for 30-120+ seconds when the engine is stuck on TCP close, + breaking `/health/liveliness` and causing Kubernetes pod restarts. + """ from prisma import Prisma # type: ignore old_engine_pid = self._get_engine_pid() - - try: - await self._original_prisma.disconnect() - except Exception as e: - verbose_proxy_logger.warning(f"Failed to disconnect Prisma client: {e}") + if old_engine_pid > 0: await self._kill_engine_process(old_engine_pid) if http_client is not None: diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index f21a729f55..48931a0c28 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -4129,17 +4129,20 @@ class PrismaClient: ) async def _do_direct_reconnect() -> None: - old_pid = self._get_engine_pid() - try: - await self.db.disconnect() - except Exception as disconnect_err: - verbose_proxy_logger.warning( - "Prisma DB disconnect before reconnect failed: %s", - disconnect_err, + db_url = os.getenv("DATABASE_URL", "") + if not db_url: + verbose_proxy_logger.error( + "DATABASE_URL not set; cannot reconnect Prisma client." ) - await PrismaWrapper._kill_engine_process(old_pid) - - await self.db.connect() + raise RuntimeError("DATABASE_URL not set") + # Fresh Prisma client + new engine subprocess. The previous + # "lightweight" path called `disconnect()` which blocks the + # event loop on `subprocess.Popen.wait()`; since that call + # ends up killing the engine anyway, we do it non-blockingly + # via `_kill_engine_process` inside `recreate_prisma_client`. + self._cleanup_engine_watcher() + await self.db.recreate_prisma_client(db_url) + await self._start_engine_watcher() await self.db.query_raw("SELECT 1") await asyncio.wait_for(_do_direct_reconnect(), timeout=effective_timeout) diff --git a/tests/litellm/proxy/test_prisma_engine_watchdog.py b/tests/litellm/proxy/test_prisma_engine_watchdog.py index 786167b948..0d241f7574 100644 --- a/tests/litellm/proxy/test_prisma_engine_watchdog.py +++ b/tests/litellm/proxy/test_prisma_engine_watchdog.py @@ -254,32 +254,48 @@ async def test_run_reconnect_cycle_uses_heavy_path_when_confirmed_dead( @pytest.mark.asyncio -async def test_run_reconnect_cycle_uses_lightweight_path_when_engine_alive( +async def test_run_reconnect_cycle_uses_direct_path_when_engine_alive( engine_client, ) -> None: - """_run_reconnect_cycle uses disconnect/connect when engine is alive.""" - engine_client._engine_pid = 1234 + """Direct reconnect (engine alive) calls recreate_prisma_client + SELECT 1. - with patch.object(engine_client, "_is_engine_alive", return_value=True): + The old "lightweight" path called `disconnect()` + `connect()`, which + blocks the event loop on the sync `process.wait()` inside aclose(). + The fix routes both engine-alive and engine-dead paths through + `recreate_prisma_client`, which non-blockingly kills the old engine. + """ + engine_client._engine_pid = 1234 + engine_client._start_engine_watcher = AsyncMock() + + with ( + patch.object(engine_client, "_is_engine_alive", return_value=True), + patch.dict(os.environ, {"DATABASE_URL": "postgresql://test"}), + ): await engine_client._run_reconnect_cycle(timeout_seconds=5.0) - engine_client.db.connect.assert_awaited_once() + engine_client.db.recreate_prisma_client.assert_awaited_once_with( + "postgresql://test" + ) engine_client.db.query_raw.assert_awaited_once_with("SELECT 1") - engine_client.db.recreate_prisma_client.assert_not_awaited() + engine_client.db.disconnect.assert_not_awaited() @pytest.mark.asyncio -async def test_run_reconnect_cycle_uses_lightweight_path_when_pid_unknown( +async def test_run_reconnect_cycle_uses_direct_path_when_pid_unknown( engine_client, ) -> None: - """_run_reconnect_cycle uses lightweight path when engine PID is not tracked.""" + """When the engine PID is not tracked, direct reconnect still runs.""" engine_client._engine_pid = 0 + engine_client._start_engine_watcher = AsyncMock() - await engine_client._run_reconnect_cycle(timeout_seconds=5.0) + with patch.dict(os.environ, {"DATABASE_URL": "postgresql://test"}): + await engine_client._run_reconnect_cycle(timeout_seconds=5.0) - engine_client.db.connect.assert_awaited_once() + engine_client.db.recreate_prisma_client.assert_awaited_once_with( + "postgresql://test" + ) engine_client.db.query_raw.assert_awaited_once_with("SELECT 1") - engine_client.db.recreate_prisma_client.assert_not_awaited() + engine_client.db.disconnect.assert_not_awaited() @pytest.mark.asyncio @@ -473,36 +489,38 @@ def test_on_engine_death_from_thread_ignores_stale_pid(engine_client): @pytest.mark.asyncio -async def test_escalation_after_consecutive_lightweight_failures(engine_client): - """After N consecutive lightweight reconnect failures, _engine_confirmed_dead +async def test_escalation_after_consecutive_direct_reconnect_failures(engine_client): + """After N consecutive direct reconnect failures, _engine_confirmed_dead is set to True so _run_reconnect_cycle takes the heavy reconnect path.""" engine_client._reconnect_escalation_threshold = 3 engine_client._consecutive_reconnect_failures = 0 engine_client._db_reconnect_cooldown_seconds = 0 # disable cooldown for test + engine_client._start_engine_watcher = AsyncMock(return_value=None) - # Make lightweight reconnect fail every time - engine_client.db.disconnect = AsyncMock(return_value=None) - engine_client.db.connect = AsyncMock(side_effect=Exception("connect failed")) + # Make direct reconnect fail every time + engine_client.db.recreate_prisma_client = AsyncMock( + side_effect=Exception("recreate failed") + ) # Run 3 failed reconnect attempts - for i in range(3): - result = await engine_client._attempt_reconnect_inside_lock( - force=True, reason="test", timeout_seconds=5.0 - ) - assert result is False + with patch.dict(os.environ, {"DATABASE_URL": "postgresql://test"}): + for _ in range(3): + result = await engine_client._attempt_reconnect_inside_lock( + force=True, reason="test", timeout_seconds=5.0 + ) + assert result is False assert engine_client._consecutive_reconnect_failures == 3 - # Next attempt should escalate: _engine_confirmed_dead set to True before _run_reconnect_cycle + # Next attempt should escalate to the heavy path (recreate_prisma_client still + # the call, but via the _engine_confirmed_dead branch that also re-arms the watcher). engine_client.db.recreate_prisma_client = AsyncMock(return_value=None) - engine_client._start_engine_watcher = AsyncMock(return_value=None) with patch.dict(os.environ, {"DATABASE_URL": "postgresql://test"}): result = await engine_client._attempt_reconnect_inside_lock( force=True, reason="test_escalation", timeout_seconds=5.0 ) - # Heavy reconnect should have been attempted (recreate_prisma_client called) engine_client.db.recreate_prisma_client.assert_awaited_once() @@ -511,15 +529,16 @@ async def test_successful_reconnect_resets_failure_counter(engine_client): """A successful reconnect resets _consecutive_reconnect_failures to 0.""" engine_client._consecutive_reconnect_failures = 2 engine_client._db_reconnect_cooldown_seconds = 0 + engine_client._start_engine_watcher = AsyncMock() # Make reconnect succeed - engine_client.db.disconnect = AsyncMock(return_value=None) - engine_client.db.connect = AsyncMock(return_value=None) + engine_client.db.recreate_prisma_client = AsyncMock(return_value=None) engine_client.db.query_raw = AsyncMock(return_value=[{"result": 1}]) - result = await engine_client._attempt_reconnect_inside_lock( - force=True, reason="test", timeout_seconds=5.0 - ) + with patch.dict(os.environ, {"DATABASE_URL": "postgresql://test"}): + result = await engine_client._attempt_reconnect_inside_lock( + force=True, reason="test", timeout_seconds=5.0 + ) assert result is True assert engine_client._consecutive_reconnect_failures == 0 diff --git a/tests/test_litellm/proxy/db/test_prisma_self_heal.py b/tests/test_litellm/proxy/db/test_prisma_self_heal.py index fb215e5477..57693a1a17 100644 --- a/tests/test_litellm/proxy/db/test_prisma_self_heal.py +++ b/tests/test_litellm/proxy/db/test_prisma_self_heal.py @@ -34,18 +34,18 @@ async def test_attempt_db_reconnect_should_succeed(mock_proxy_logging): client = PrismaClient( database_url="mock://test", proxy_logging_obj=mock_proxy_logging ) - client.db.disconnect = AsyncMock(return_value=None) - client.db.connect = AsyncMock(return_value=None) + client.db.recreate_prisma_client = AsyncMock(return_value=None) client.db.query_raw = AsyncMock(return_value=[{"result": 1}]) + client._start_engine_watcher = AsyncMock() - result = await client.attempt_db_reconnect( - reason="unit_test_reconnect_success", - force=True, - ) + with patch.dict(os.environ, {"DATABASE_URL": "postgresql://test"}): + result = await client.attempt_db_reconnect( + reason="unit_test_reconnect_success", + force=True, + ) assert result is True - client.db.disconnect.assert_awaited_once() - client.db.connect.assert_awaited_once() + client.db.recreate_prisma_client.assert_awaited_once_with("postgresql://test") client.db.query_raw.assert_awaited_once_with("SELECT 1") @@ -140,15 +140,19 @@ async def test_attempt_db_reconnect_should_set_cooldown_after_attempt( ) client._db_last_reconnect_attempt_ts = 0.0 client._db_reconnect_cooldown_seconds = 10 - client.db.disconnect = AsyncMock(return_value=None) - client.db.connect = AsyncMock(return_value=None) + client.db.recreate_prisma_client = AsyncMock(return_value=None) client.db.query_raw = AsyncMock(return_value=[{"result": 1}]) + client._start_engine_watcher = AsyncMock() # Use a counter-based mock to avoid StopIteration when time.time() is called # more times than expected (varies by Python version / internal code paths). fake_clock = iter(range(100, 10000)) - with patch( - "litellm.proxy.utils.time.time", side_effect=lambda: float(next(fake_clock)) + with ( + patch( + "litellm.proxy.utils.time.time", + side_effect=lambda: float(next(fake_clock)), + ), + patch.dict(os.environ, {"DATABASE_URL": "postgresql://test"}), ): result = await client.attempt_db_reconnect( reason="unit_test_cooldown_timestamp_after_attempt", @@ -162,23 +166,28 @@ async def test_attempt_db_reconnect_should_set_cooldown_after_attempt( @pytest.mark.asyncio -async def test_run_reconnect_cycle_watchdog_should_use_direct_db_ops( +async def test_run_reconnect_cycle_watchdog_should_use_recreate_prisma_client( mock_proxy_logging, ): + """Direct reconnect goes through recreate_prisma_client (which non-blockingly + kills the old engine) instead of calling disconnect() — see issue #26191. + """ client = PrismaClient( database_url="mock://test", proxy_logging_obj=mock_proxy_logging ) - client.disconnect = AsyncMock(side_effect=AssertionError("wrapper disconnect used")) - client.connect = AsyncMock(side_effect=AssertionError("wrapper connect used")) - client.db.disconnect = AsyncMock(return_value=None) - client.db.connect = AsyncMock(return_value=None) + client.db.disconnect = AsyncMock( + side_effect=AssertionError("disconnect must not be called") + ) + client.db.recreate_prisma_client = AsyncMock(return_value=None) client.db.query_raw = AsyncMock(return_value=[{"result": 1}]) + client._start_engine_watcher = AsyncMock() - await client._run_reconnect_cycle(timeout_seconds=None) + with patch.dict(os.environ, {"DATABASE_URL": "postgresql://test"}): + await client._run_reconnect_cycle(timeout_seconds=None) - client.db.disconnect.assert_awaited_once() - client.db.connect.assert_awaited_once() + client.db.recreate_prisma_client.assert_awaited_once_with("postgresql://test") client.db.query_raw.assert_awaited_once_with("SELECT 1") + client.db.disconnect.assert_not_awaited() @pytest.mark.asyncio @@ -189,19 +198,22 @@ async def test_run_reconnect_cycle_watchdog_should_use_default_timeout_budget( database_url="mock://test", proxy_logging_obj=mock_proxy_logging ) client._db_watchdog_reconnect_timeout_seconds = 0.1 - client.db.disconnect = AsyncMock(return_value=None) + client._start_engine_watcher = AsyncMock() - async def _slow_connect(): + async def _slow_recreate(_db_url): await asyncio.sleep(0.08) async def _slow_query(_query: str): await asyncio.sleep(0.08) return [{"result": 1}] - client.db.connect = AsyncMock(side_effect=_slow_connect) + client.db.recreate_prisma_client = AsyncMock(side_effect=_slow_recreate) client.db.query_raw = AsyncMock(side_effect=_slow_query) - with pytest.raises(asyncio.TimeoutError): + with ( + pytest.raises(asyncio.TimeoutError), + patch.dict(os.environ, {"DATABASE_URL": "postgresql://test"}), + ): await client._run_reconnect_cycle(timeout_seconds=None) @@ -212,19 +224,22 @@ async def test_run_reconnect_cycle_timeout_should_use_single_overall_budget( client = PrismaClient( database_url="mock://test", proxy_logging_obj=mock_proxy_logging ) - client.db.disconnect = AsyncMock(return_value=None) + client._start_engine_watcher = AsyncMock() - async def _slow_connect(): + async def _slow_recreate(_db_url): await asyncio.sleep(0.08) async def _slow_query(_query: str): await asyncio.sleep(0.08) return [{"result": 1}] - client.db.connect = AsyncMock(side_effect=_slow_connect) + client.db.recreate_prisma_client = AsyncMock(side_effect=_slow_recreate) client.db.query_raw = AsyncMock(side_effect=_slow_query) - with pytest.raises(asyncio.TimeoutError): + with ( + pytest.raises(asyncio.TimeoutError), + patch.dict(os.environ, {"DATABASE_URL": "postgresql://test"}), + ): await client._run_reconnect_cycle(timeout_seconds=0.1) @@ -319,42 +334,32 @@ async def test_db_health_watchdog_start_stop_lifecycle(mock_proxy_logging): @pytest.mark.asyncio -async def test_lightweight_reconnect_kills_engine_on_disconnect_failure( +async def test_recreate_prisma_client_kills_old_engine_without_disconnect( mock_proxy_logging, ): - """Lightweight reconnect must kill the old engine PID when disconnect() fails.""" + """recreate_prisma_client SIGTERMs the old engine PID directly rather than + calling `disconnect()`, which blocks the asyncio event loop on the sync + `subprocess.Popen.wait()` inside prisma-client-py — see issue #26191. + """ client = PrismaClient( database_url="mock://test", proxy_logging_obj=mock_proxy_logging ) - client.db.disconnect = AsyncMock(side_effect=Exception("disconnect failed")) - client.db.connect = AsyncMock(return_value=None) - client.db.query_raw = AsyncMock(return_value=[{"result": 1}]) + disconnect_mock = AsyncMock( + side_effect=AssertionError("disconnect must not be called on reconnect path") + ) + client.db._original_prisma.disconnect = disconnect_mock with ( - patch.object(client, "_get_engine_pid", return_value=9999), - patch("os.kill") as mock_kill, - patch("asyncio.sleep", new_callable=AsyncMock), + patch.object(client.db, "_get_engine_pid", return_value=9999), + patch("litellm.proxy.db.prisma_client.os.kill") as mock_kill, + patch("litellm.proxy.db.prisma_client.asyncio.sleep", new_callable=AsyncMock), ): - await client._run_reconnect_cycle(timeout_seconds=5.0) + # Return a Prisma instance whose connect() is awaitable. + fake_new_prisma = MagicMock() + fake_new_prisma.connect = AsyncMock(return_value=None) + with patch("prisma.Prisma", return_value=fake_new_prisma): + await client.db.recreate_prisma_client("postgresql://test") mock_kill.assert_any_call(9999, signal.SIGTERM) - client.db.connect.assert_awaited_once() - client.db.query_raw.assert_awaited_once_with("SELECT 1") - - -@pytest.mark.asyncio -async def test_lightweight_reconnect_skips_kill_on_successful_disconnect( - mock_proxy_logging, -): - """Lightweight reconnect must NOT kill when disconnect() succeeds.""" - client = PrismaClient( - database_url="mock://test", proxy_logging_obj=mock_proxy_logging - ) - client.db.disconnect = AsyncMock(return_value=None) - client.db.connect = AsyncMock(return_value=None) - client.db.query_raw = AsyncMock(return_value=[{"result": 1}]) - - with patch("os.kill") as mock_kill: - await client._run_reconnect_cycle(timeout_seconds=5.0) - - mock_kill.assert_not_called() + disconnect_mock.assert_not_awaited() + fake_new_prisma.connect.assert_awaited_once() From 2c3c8aa4ea2c0bdbee1e7d80e60ea1923b8c7f99 Mon Sep 17 00:00:00 2001 From: Ryan Crabbe Date: Thu, 23 Apr 2026 21:04:13 -0700 Subject: [PATCH 02/57] Move "Store Prompts in Spend Logs" toggle to Admin Settings MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Previously, the "Store Prompts in Spend Logs" and "Maximum Spend Logs Retention Period" settings were surfaced via a gear-icon modal on the Logs page. The gear was visible to every authenticated user even though the backend endpoints (/config/update, /config/list) require PROXY_ADMIN — so non-admins could open the modal but the request would 403 on load and save, giving a confusing UX. Move the controls into a new "Logging Settings" tab under Admin Settings, which is already gated to admins at the sidebar. Remove the gear button and the onOpenSettings prop chain (ConfigInfoMessage → LogDetailContent → LogDetailsDrawer). ConfigInfoMessage now points users to "Admin Settings → Logging Settings" inline. --- .../src/components/AdminPanel.tsx | 6 + .../LoggingSettings/LoggingSettings.test.tsx} | 255 +++++------------- .../LoggingSettings/LoggingSettings.tsx | 150 +++++++++++ .../view_logs/ConfigInfoMessage.test.tsx | 22 +- .../view_logs/ConfigInfoMessage.tsx | 17 +- .../LogDetailContent.test.tsx | 21 -- .../LogDetailsDrawer/LogDetailContent.tsx | 5 +- .../LogDetailsDrawer/LogDetailsDrawer.tsx | 3 - .../SpendLogsSettingsModal.tsx | 156 ----------- .../src/components/view_logs/index.tsx | 19 +- 10 files changed, 228 insertions(+), 426 deletions(-) rename ui/litellm-dashboard/src/components/{view_logs/SpendLogsSettingsModal/SpendLogsSettingsModal.test.tsx => Settings/AdminSettings/LoggingSettings/LoggingSettings.test.tsx} (53%) create mode 100644 ui/litellm-dashboard/src/components/Settings/AdminSettings/LoggingSettings/LoggingSettings.tsx delete mode 100644 ui/litellm-dashboard/src/components/view_logs/SpendLogsSettingsModal/SpendLogsSettingsModal.tsx diff --git a/ui/litellm-dashboard/src/components/AdminPanel.tsx b/ui/litellm-dashboard/src/components/AdminPanel.tsx index 7528b081ff..e4aec7706a 100644 --- a/ui/litellm-dashboard/src/components/AdminPanel.tsx +++ b/ui/litellm-dashboard/src/components/AdminPanel.tsx @@ -21,6 +21,7 @@ import { useBaseUrl } from "./constants"; import NotificationsManager from "./molecules/notifications_manager"; import { addAllowedIP, deleteAllowedIP, getAllowedIPs, getSSOSettings } from "./networking"; import SCIMConfig from "./SCIM"; +import LoggingSettings from "./Settings/AdminSettings/LoggingSettings/LoggingSettings"; import SSOSettings from "./Settings/AdminSettings/SSOSettings/SSOSettings"; import UISettings from "./Settings/AdminSettings/UISettings/UISettings"; import HashicorpVault from "./Settings/AdminSettings/HashicorpVault/HashicorpVault"; @@ -362,6 +363,11 @@ const AdminPanel: React.FC = ({ proxySettings }) => { ), children: , }, + { + key: "logging-settings", + label: "Logging Settings", + children: , + }, { key: "hashicorp-vault", label: "Hashicorp Vault", diff --git a/ui/litellm-dashboard/src/components/view_logs/SpendLogsSettingsModal/SpendLogsSettingsModal.test.tsx b/ui/litellm-dashboard/src/components/Settings/AdminSettings/LoggingSettings/LoggingSettings.test.tsx similarity index 53% rename from ui/litellm-dashboard/src/components/view_logs/SpendLogsSettingsModal/SpendLogsSettingsModal.test.tsx rename to ui/litellm-dashboard/src/components/Settings/AdminSettings/LoggingSettings/LoggingSettings.test.tsx index 40d06d9046..228e899a3a 100644 --- a/ui/litellm-dashboard/src/components/view_logs/SpendLogsSettingsModal/SpendLogsSettingsModal.test.tsx +++ b/ui/litellm-dashboard/src/components/Settings/AdminSettings/LoggingSettings/LoggingSettings.test.tsx @@ -5,11 +5,20 @@ import { parseErrorMessage } from "@/components/shared/errorUtils"; import { screen, waitFor } from "@testing-library/react"; import userEvent from "@testing-library/user-event"; import { beforeEach, describe, expect, it, vi } from "vitest"; -import { renderWithProviders } from "../../../../tests/test-utils"; -import SpendLogsSettingsModal from "./SpendLogsSettingsModal"; +import { renderWithProviders } from "../../../../../tests/test-utils"; +import LoggingSettings from "./LoggingSettings"; vi.mock("@/app/(dashboard)/hooks/storeRequestInSpendLogs/useStoreRequestInSpendLogs"); -vi.mock("@/app/(dashboard)/hooks/proxyConfig/useProxyConfig"); +vi.mock("@/app/(dashboard)/hooks/proxyConfig/useProxyConfig", async () => { + const actual = await vi.importActual( + "@/app/(dashboard)/hooks/proxyConfig/useProxyConfig", + ); + return { + ...actual, + useProxyConfig: vi.fn(), + useDeleteProxyConfigField: vi.fn(), + }; +}); vi.mock("@/components/molecules/notifications_manager", () => ({ default: { success: vi.fn(), @@ -26,19 +35,11 @@ const mockUseDeleteProxyConfigField = vi.mocked(useDeleteProxyConfigField); const mockNotificationsManager = vi.mocked(NotificationsManager); const mockParseErrorMessage = vi.mocked(parseErrorMessage); -describe("SpendLogsSettingsModal", () => { - const mockOnCancel = vi.fn(); - const mockOnSuccess = vi.fn(); +describe("LoggingSettings", () => { const mockMutateAsync = vi.fn(); const mockDeleteField = vi.fn(); const mockRefetch = vi.fn(); - const defaultProps = { - isVisible: true, - onCancel: mockOnCancel, - onSuccess: mockOnSuccess, - }; - beforeEach(() => { vi.clearAllMocks(); mockUseStoreRequestInSpendLogs.mockReturnValue({ @@ -57,50 +58,19 @@ describe("SpendLogsSettingsModal", () => { mockParseErrorMessage.mockImplementation((error: any) => error?.message || String(error)); }); - it("should render the modal", () => { - renderWithProviders(); - expect(screen.getByRole("dialog")).toBeInTheDocument(); - expect(screen.getByText("Spend Logs Settings")).toBeInTheDocument(); - }); - - it("should render form fields with initial values", () => { - renderWithProviders(); + it("should render the card with title and form fields", () => { + renderWithProviders(); + expect(screen.getByText("Logging Settings")).toBeInTheDocument(); expect(screen.getByText("Store Prompts in Spend Logs")).toBeInTheDocument(); expect(screen.getByLabelText("Maximum Spend Logs Retention Period (Optional)")).toBeInTheDocument(); expect(screen.getByPlaceholderText("e.g., 7d, 30d")).toBeInTheDocument(); - }); - - it("should render cancel and save buttons", () => { - renderWithProviders(); - - expect(screen.getByRole("button", { name: "Cancel" })).toBeInTheDocument(); expect(screen.getByRole("button", { name: "Save Settings" })).toBeInTheDocument(); }); - it("should call onCancel when cancel button is clicked", async () => { - const user = userEvent.setup(); - renderWithProviders(); - - const cancelButton = screen.getByRole("button", { name: "Cancel" }); - await user.click(cancelButton); - - expect(mockOnCancel).toHaveBeenCalledTimes(1); - }); - - it("should call onCancel when modal close button is clicked", async () => { - const user = userEvent.setup(); - renderWithProviders(); - - const closeButton = screen.getByRole("button", { name: /close/i }); - await user.click(closeButton); - - expect(mockOnCancel).toHaveBeenCalledTimes(1); - }); - it("should toggle store prompts switch", async () => { const user = userEvent.setup(); - renderWithProviders(); + renderWithProviders(); const switchElement = screen.getByRole("switch"); expect(switchElement).not.toBeChecked(); @@ -114,7 +84,7 @@ describe("SpendLogsSettingsModal", () => { it("should update retention period input", async () => { const user = userEvent.setup(); - renderWithProviders(); + renderWithProviders(); const retentionInput = screen.getByPlaceholderText("e.g., 7d, 30d"); await user.type(retentionInput, "30d"); @@ -124,13 +94,13 @@ describe("SpendLogsSettingsModal", () => { it("should submit form with store prompts enabled and retention period", async () => { const user = userEvent.setup(); - mockMutateAsync.mockImplementation(async (params, options) => { + mockMutateAsync.mockImplementation(async (_params, options) => { await Promise.resolve(); options?.onSuccess?.(); return { message: "Success" }; }); - renderWithProviders(); + renderWithProviders(); const switchElement = screen.getByRole("switch"); await user.click(switchElement); @@ -148,21 +118,21 @@ describe("SpendLogsSettingsModal", () => { store_prompts_in_spend_logs: true, maximum_spend_logs_retention_period: "30d", }, - expect.any(Object) + expect.any(Object), ); }); }); - it("should submit form with store prompts disabled and no retention period", async () => { + it("should delete retention period field when left empty on submit", async () => { const user = userEvent.setup(); mockDeleteField.mockResolvedValue({ message: "Field deleted successfully" }); - mockMutateAsync.mockImplementation(async (params, options) => { + mockMutateAsync.mockImplementation(async (_params, options) => { await Promise.resolve(); options?.onSuccess?.(); return { message: "Success" }; }); - renderWithProviders(); + renderWithProviders(); const saveButton = screen.getByRole("button", { name: "Save Settings" }); await user.click(saveButton); @@ -173,207 +143,106 @@ describe("SpendLogsSettingsModal", () => { { store_prompts_in_spend_logs: false, }, - expect.any(Object) + expect.any(Object), ); }); }); - it("should show success notification and call onSuccess on successful submission", async () => { + it("should show success notification on successful submission", async () => { const user = userEvent.setup(); mockDeleteField.mockResolvedValue({ message: "Field deleted successfully" }); - mockMutateAsync.mockImplementation(async (params, options) => { + mockMutateAsync.mockImplementation(async (_params, options) => { await Promise.resolve(); options?.onSuccess?.(); return { message: "Success" }; }); - renderWithProviders(); + renderWithProviders(); const saveButton = screen.getByRole("button", { name: "Save Settings" }); await user.click(saveButton); await waitFor(() => { expect(mockNotificationsManager.success).toHaveBeenCalledWith("Spend logs settings updated successfully"); - expect(mockRefetch).toHaveBeenCalled(); - expect(mockOnSuccess).toHaveBeenCalledTimes(1); }); }); - it("should show error notification when submission fails", async () => { + it("should show error notification when submission throws", async () => { const user = userEvent.setup(); const error = new Error("Network error"); mockMutateAsync.mockRejectedValue(error); mockParseErrorMessage.mockReturnValue("Network error"); - renderWithProviders(); + renderWithProviders(); const saveButton = screen.getByRole("button", { name: "Save Settings" }); await user.click(saveButton); await waitFor(() => { - expect(mockNotificationsManager.fromBackend).toHaveBeenCalledWith("Failed to save spend logs settings: Network error"); + expect(mockNotificationsManager.fromBackend).toHaveBeenCalledWith( + "Failed to save spend logs settings: Network error", + ); }); }); - it("should show error notification from onError callback", async () => { + it("should show error notification via onError callback", async () => { const user = userEvent.setup(); const error = new Error("Backend error"); - mockMutateAsync.mockImplementation((params, options) => { + mockMutateAsync.mockImplementation((_params, options) => { options?.onError?.(error); return Promise.reject(error); }); mockParseErrorMessage.mockReturnValue("Backend error"); - renderWithProviders(); + renderWithProviders(); const saveButton = screen.getByRole("button", { name: "Save Settings" }); await user.click(saveButton); await waitFor(() => { - expect(mockNotificationsManager.fromBackend).toHaveBeenCalledWith("Failed to save spend logs settings: Backend error"); + expect(mockNotificationsManager.fromBackend).toHaveBeenCalledWith( + "Failed to save spend logs settings: Backend error", + ); }); }); - it("should disable cancel button when pending", () => { + it("should show loading state on save button when update pending", () => { mockUseStoreRequestInSpendLogs.mockReturnValue({ mutateAsync: mockMutateAsync, isPending: true, } as any); - renderWithProviders(); + renderWithProviders(); - const cancelButton = screen.getByRole("button", { name: "Cancel" }); - expect(cancelButton).toBeDisabled(); + const saveButton = screen.getByRole("button", { name: /Saving/i }); + expect(saveButton).toBeInTheDocument(); + expect(saveButton.className).toContain("ant-btn-loading"); }); - it("should disable cancel button when deleting field", () => { + it("should show loading state on save button when delete pending", () => { mockUseDeleteProxyConfigField.mockReturnValue({ mutateAsync: mockDeleteField, isPending: true, } as any); - renderWithProviders(); + renderWithProviders(); - const cancelButton = screen.getByRole("button", { name: "Cancel" }); - expect(cancelButton).toBeDisabled(); + const saveButton = screen.getByRole("button", { name: /Saving/i }); + expect(saveButton).toBeInTheDocument(); + expect(saveButton.className).toContain("ant-btn-loading"); }); - it("should disable cancel button when loading config", () => { + it("should disable save button while config is loading", () => { mockUseProxyConfig.mockReturnValue({ data: undefined, isLoading: true, refetch: mockRefetch, } as any); - renderWithProviders(); - - const cancelButton = screen.getByRole("button", { name: "Cancel" }); - expect(cancelButton).toBeDisabled(); - }); - - it("should show loading state on save button when pending", () => { - mockUseStoreRequestInSpendLogs.mockReturnValue({ - mutateAsync: mockMutateAsync, - isPending: true, - } as any); - - renderWithProviders(); - - const saveButton = screen.getByRole("button", { name: /Saving/i }); - expect(saveButton).toBeInTheDocument(); - expect(saveButton.className).toContain("ant-btn-loading"); - }); - - it("should show loading state on save button when deleting field", () => { - mockUseDeleteProxyConfigField.mockReturnValue({ - mutateAsync: mockDeleteField, - isPending: true, - } as any); - - renderWithProviders(); - - const saveButton = screen.getByRole("button", { name: /Saving/i }); - expect(saveButton).toBeInTheDocument(); - expect(saveButton.className).toContain("ant-btn-loading"); - }); - - it("should call onCancel when cancel button is clicked after modifying form", async () => { - const user = userEvent.setup(); - renderWithProviders(); - - const switchElement = screen.getByRole("switch"); - await user.click(switchElement); - - const retentionInput = screen.getByPlaceholderText("e.g., 7d, 30d"); - await user.type(retentionInput, "30d"); - - expect(switchElement).toBeChecked(); - expect(retentionInput).toHaveValue("30d"); - - const cancelButton = screen.getByRole("button", { name: "Cancel" }); - await user.click(cancelButton); - - expect(mockOnCancel).toHaveBeenCalledTimes(1); - }); - - it("should call refetch after successful submission", async () => { - const user = userEvent.setup(); - mockDeleteField.mockResolvedValue({ message: "Field deleted successfully" }); - mockMutateAsync.mockImplementation(async (params, options) => { - await Promise.resolve(); - options?.onSuccess?.(); - return { message: "Success" }; - }); - - renderWithProviders(); - - const switchElement = screen.getByRole("switch"); - await user.click(switchElement); - - const retentionInput = screen.getByPlaceholderText("e.g., 7d, 30d"); - await user.type(retentionInput, "30d"); - - expect(switchElement).toBeChecked(); - expect(retentionInput).toHaveValue("30d"); + renderWithProviders(); const saveButton = screen.getByRole("button", { name: "Save Settings" }); - await user.click(saveButton); - - await waitFor(() => { - expect(mockNotificationsManager.success).toHaveBeenCalled(); - expect(mockRefetch).toHaveBeenCalled(); - }); - }); - - it("should not call onSuccess when it is not provided", async () => { - const user = userEvent.setup(); - mockDeleteField.mockResolvedValue({ message: "Field deleted successfully" }); - mockMutateAsync.mockImplementation(async (params, options) => { - await Promise.resolve(); - options?.onSuccess?.(); - return { message: "Success" }; - }); - - renderWithProviders(); - - const saveButton = screen.getByRole("button", { name: "Save Settings" }); - await user.click(saveButton); - - await waitFor(() => { - expect(mockNotificationsManager.success).toHaveBeenCalled(); - }); - }); - - it("should not render modal when isVisible is false", () => { - renderWithProviders(); - - expect(screen.queryByRole("dialog")).not.toBeInTheDocument(); - }); - - it("should call refetch when modal opens", () => { - renderWithProviders(); - - expect(mockRefetch).toHaveBeenCalledTimes(1); + expect(saveButton).toBeDisabled(); }); it("should render form with initial values from config data", () => { @@ -400,7 +269,7 @@ describe("SpendLogsSettingsModal", () => { refetch: mockRefetch, } as any); - renderWithProviders(); + renderWithProviders(); const switchElement = screen.getByRole("switch"); const retentionInput = screen.getByPlaceholderText("e.g., 7d, 30d"); @@ -416,13 +285,11 @@ describe("SpendLogsSettingsModal", () => { refetch: mockRefetch, } as any); - renderWithProviders(); + renderWithProviders(); - // Check that switch and input are not present when loading (skeletons are shown instead) expect(screen.queryByRole("switch")).not.toBeInTheDocument(); expect(screen.queryByPlaceholderText("e.g., 7d, 30d")).not.toBeInTheDocument(); - // Check for skeleton elements (Ant Design Skeleton.Input renders with ant-skeleton class) const skeletons = document.querySelectorAll(".ant-skeleton"); expect(skeletons.length).toBeGreaterThan(0); }); @@ -431,13 +298,13 @@ describe("SpendLogsSettingsModal", () => { const user = userEvent.setup(); const deleteError = new Error("Field does not exist"); mockDeleteField.mockRejectedValue(deleteError); - mockMutateAsync.mockImplementation(async (params, options) => { + mockMutateAsync.mockImplementation(async (_params, options) => { await Promise.resolve(); options?.onSuccess?.(); return { message: "Success" }; }); - renderWithProviders(); + renderWithProviders(); const saveButton = screen.getByRole("button", { name: "Save Settings" }); await user.click(saveButton); @@ -448,22 +315,22 @@ describe("SpendLogsSettingsModal", () => { { store_prompts_in_spend_logs: false, }, - expect.any(Object) + expect.any(Object), ); expect(mockNotificationsManager.success).toHaveBeenCalled(); }); }); - it("should submit form with only store prompts enabled and no retention period", async () => { + it("should submit with only store prompts enabled when retention is empty", async () => { const user = userEvent.setup(); mockDeleteField.mockResolvedValue({ message: "Field deleted successfully" }); - mockMutateAsync.mockImplementation(async (params, options) => { + mockMutateAsync.mockImplementation(async (_params, options) => { await Promise.resolve(); options?.onSuccess?.(); return { message: "Success" }; }); - renderWithProviders(); + renderWithProviders(); const switchElement = screen.getByRole("switch"); await user.click(switchElement); @@ -477,7 +344,7 @@ describe("SpendLogsSettingsModal", () => { { store_prompts_in_spend_logs: true, }, - expect.any(Object) + expect.any(Object), ); }); }); diff --git a/ui/litellm-dashboard/src/components/Settings/AdminSettings/LoggingSettings/LoggingSettings.tsx b/ui/litellm-dashboard/src/components/Settings/AdminSettings/LoggingSettings/LoggingSettings.tsx new file mode 100644 index 0000000000..c3aaebd3bf --- /dev/null +++ b/ui/litellm-dashboard/src/components/Settings/AdminSettings/LoggingSettings/LoggingSettings.tsx @@ -0,0 +1,150 @@ +"use client"; + +import { + ConfigType, + GeneralSettingsFieldName, + useDeleteProxyConfigField, + useProxyConfig, +} from "@/app/(dashboard)/hooks/proxyConfig/useProxyConfig"; +import { + StoreRequestInSpendLogsParams, + useStoreRequestInSpendLogs, +} from "@/app/(dashboard)/hooks/storeRequestInSpendLogs/useStoreRequestInSpendLogs"; +import NotificationsManager from "@/components/molecules/notifications_manager"; +import { parseErrorMessage } from "@/components/shared/errorUtils"; +import { ClockCircleOutlined } from "@ant-design/icons"; +import { Button, Card, Form, Input, Skeleton, Space, Switch, Typography } from "antd"; +import React, { useMemo } from "react"; + +const LoggingSettings: React.FC = () => { + const [form] = Form.useForm(); + const { mutateAsync, isPending } = useStoreRequestInSpendLogs(); + const { mutateAsync: deleteField, isPending: isDeletingField } = useDeleteProxyConfigField(); + const { data: proxyConfigData, isLoading: isLoadingConfig } = useProxyConfig(ConfigType.GENERAL_SETTINGS); + const storePromptsValue = Form.useWatch("store_prompts_in_spend_logs", form); + + const initialValues = useMemo(() => { + if (!proxyConfigData) { + return { + store_prompts_in_spend_logs: false, + maximum_spend_logs_retention_period: undefined, + }; + } + + const storePromptsField = proxyConfigData.find((field) => field.field_name === "store_prompts_in_spend_logs"); + const retentionPeriodField = proxyConfigData.find( + (field) => field.field_name === "maximum_spend_logs_retention_period", + ); + + return { + store_prompts_in_spend_logs: storePromptsField?.field_value ?? false, + maximum_spend_logs_retention_period: retentionPeriodField?.field_value ?? undefined, + }; + }, [proxyConfigData]); + + const handleFormSubmit = async (formValues: StoreRequestInSpendLogsParams) => { + try { + const retentionPeriodValue = formValues.maximum_spend_logs_retention_period; + const shouldDeleteRetentionPeriod = + !retentionPeriodValue || + (typeof retentionPeriodValue === "string" && retentionPeriodValue.trim() === ""); + + if (shouldDeleteRetentionPeriod) { + try { + await deleteField({ + config_type: ConfigType.GENERAL_SETTINGS, + field_name: GeneralSettingsFieldName.MAXIMUM_SPEND_LOGS_RETENTION_PERIOD, + }); + } catch (deleteError) { + console.warn("Failed to delete retention period field (may not exist):", deleteError); + } + } + + const updateParams: StoreRequestInSpendLogsParams = { + store_prompts_in_spend_logs: formValues.store_prompts_in_spend_logs, + ...(retentionPeriodValue && + typeof retentionPeriodValue === "string" && + retentionPeriodValue.trim() !== "" && { + maximum_spend_logs_retention_period: retentionPeriodValue, + }), + }; + + await mutateAsync(updateParams, { + onSuccess: () => { + NotificationsManager.success("Spend logs settings updated successfully"); + }, + onError: (error) => { + NotificationsManager.fromBackend("Failed to save spend logs settings: " + parseErrorMessage(error)); + }, + }); + } catch (error) { + NotificationsManager.fromBackend("Failed to save spend logs settings: " + parseErrorMessage(error)); + } + }; + + return ( + + + + Proxy-wide settings that control how request and response data are written to spend logs. + + +
+ f.field_name === "store_prompts_in_spend_logs")?.field_description || + "When enabled, prompts will be stored in spend logs for tracking and analysis purposes." + } + valuePropName="checked" + > + {isLoadingConfig ? ( + + ) : ( + form.setFieldValue("store_prompts_in_spend_logs", checked)} + /> + )} + + + f.field_name === "maximum_spend_logs_retention_period") + ?.field_description || + "Set the maximum retention period for spend logs (e.g., '7d' for 7 days, '30d' for 30 days). Leave empty for no limit." + } + > + {isLoadingConfig ? ( + + ) : ( + } /> + )} + + + + + +
+
+
+ ); +}; + +export default LoggingSettings; diff --git a/ui/litellm-dashboard/src/components/view_logs/ConfigInfoMessage.test.tsx b/ui/litellm-dashboard/src/components/view_logs/ConfigInfoMessage.test.tsx index 9e28b27cec..ad9b724ba8 100644 --- a/ui/litellm-dashboard/src/components/view_logs/ConfigInfoMessage.test.tsx +++ b/ui/litellm-dashboard/src/components/view_logs/ConfigInfoMessage.test.tsx @@ -1,6 +1,5 @@ import { render, screen } from "@testing-library/react"; -import { describe, expect, it, vi } from "vitest"; -import userEvent from "@testing-library/user-event"; +import { describe, expect, it } from "vitest"; import { ConfigInfoMessage } from "./ConfigInfoMessage"; describe("ConfigInfoMessage", () => { @@ -19,23 +18,8 @@ describe("ConfigInfoMessage", () => { expect(screen.getByText(/store_prompts_in_spend_logs: true/)).toBeInTheDocument(); }); - it("should render the settings button when onOpenSettings is provided", () => { - render( {}} />); - expect(screen.getByText("open the settings")).toBeInTheDocument(); - }); - - it("should not render the settings button when onOpenSettings is omitted", () => { + it("should reference Admin Settings \u2192 Logging Settings", () => { render(); - expect(screen.queryByText("open the settings")).not.toBeInTheDocument(); - }); - - it("should call onOpenSettings when the settings button is clicked", async () => { - const user = userEvent.setup(); - const onOpenSettings = vi.fn(); - - render(); - await user.click(screen.getByText("open the settings")); - - expect(onOpenSettings).toHaveBeenCalledOnce(); + expect(screen.getByText(/Admin Settings → Logging Settings/)).toBeInTheDocument(); }); }); diff --git a/ui/litellm-dashboard/src/components/view_logs/ConfigInfoMessage.tsx b/ui/litellm-dashboard/src/components/view_logs/ConfigInfoMessage.tsx index 509b1c73de..a231a09971 100644 --- a/ui/litellm-dashboard/src/components/view_logs/ConfigInfoMessage.tsx +++ b/ui/litellm-dashboard/src/components/view_logs/ConfigInfoMessage.tsx @@ -2,10 +2,9 @@ import React from "react"; interface ConfigInfoMessageProps { show: boolean; - onOpenSettings?: () => void; } -export const ConfigInfoMessage: React.FC = ({ show, onOpenSettings }) => { +export const ConfigInfoMessage: React.FC = ({ show }) => { if (!show) return null; return ( @@ -31,18 +30,8 @@ export const ConfigInfoMessage: React.FC = ({ show, onOp

Request/Response Data Not Available

To view request and response details, enable prompt storage in your LiteLLM configuration by adding the - following to your proxy_config.yaml file - {onOpenSettings && ( - <> or{" "} - - {" "}to configure this directly. - - )} + following to your proxy_config.yaml file, or toggle + the setting in Admin Settings → Logging Settings.

           {`general_settings:
diff --git a/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/LogDetailContent.test.tsx b/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/LogDetailContent.test.tsx
index 992063fb69..a2da513675 100644
--- a/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/LogDetailContent.test.tsx
+++ b/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/LogDetailContent.test.tsx
@@ -171,27 +171,6 @@ describe("LogDetailContent", () => {
     expect(screen.queryByText("Request/Response Data Not Available")).not.toBeInTheDocument();
   });
 
-  it("should call onOpenSettings when user clicks open settings in ConfigInfoMessage", async () => {
-    const onOpenSettings = vi.fn();
-    const user = userEvent.setup();
-
-    render(
-      ,
-    );
-
-    const settingsButton = screen.getByRole("button", { name: /open the settings/i });
-    await user.click(settingsButton);
-
-    expect(onOpenSettings).toHaveBeenCalledTimes(1);
-  });
-
   it("should display loading state when isLoadingDetails is true", () => {
     render(
        void;
   /** When true, log details (messages/response) are still being lazy-loaded. */
   isLoadingDetails?: boolean;
   accessToken?: string | null;
@@ -51,7 +50,7 @@ export interface LogDetailContentProps {
  * Designed to be placed inside LogDetailsDrawer's right panel so it can
  * be reused for both single-log and session-mode views.
  */
-export function LogDetailContent({ logEntry, onOpenSettings, isLoadingDetails = false, accessToken }: LogDetailContentProps) {
+export function LogDetailContent({ logEntry, isLoadingDetails = false, accessToken }: LogDetailContentProps) {
   const metadata = logEntry.metadata || {};
   const hasError = metadata.status === "failure";
   const errorInfo = hasError ? metadata.error_information : null;
@@ -153,7 +152,7 @@ export function LogDetailContent({ logEntry, onOpenSettings, isLoadingDetails =
       {/* Configuration Info Message */}
       {missingData && (
         
- +
)} diff --git a/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/LogDetailsDrawer.tsx b/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/LogDetailsDrawer.tsx index 036a24c045..39315b0b58 100644 --- a/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/LogDetailsDrawer.tsx +++ b/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/LogDetailsDrawer.tsx @@ -26,7 +26,6 @@ export interface LogDetailsDrawerProps { logEntry: LogEntry | null; sessionId?: string | null; accessToken?: string | null; - onOpenSettings?: () => void; allLogs?: LogEntry[]; onSelectLog?: (log: LogEntry) => void; startTime?: string; @@ -109,7 +108,6 @@ export function LogDetailsDrawer({ logEntry, sessionId, accessToken, - onOpenSettings, allLogs = [], onSelectLog, startTime, @@ -399,7 +397,6 @@ export function LogDetailsDrawer({
diff --git a/ui/litellm-dashboard/src/components/view_logs/SpendLogsSettingsModal/SpendLogsSettingsModal.tsx b/ui/litellm-dashboard/src/components/view_logs/SpendLogsSettingsModal/SpendLogsSettingsModal.tsx deleted file mode 100644 index b3f73af060..0000000000 --- a/ui/litellm-dashboard/src/components/view_logs/SpendLogsSettingsModal/SpendLogsSettingsModal.tsx +++ /dev/null @@ -1,156 +0,0 @@ -"use client"; - -import { ConfigType, GeneralSettingsFieldName, useDeleteProxyConfigField, useProxyConfig } from "@/app/(dashboard)/hooks/proxyConfig/useProxyConfig"; -import { StoreRequestInSpendLogsParams, useStoreRequestInSpendLogs } from "@/app/(dashboard)/hooks/storeRequestInSpendLogs/useStoreRequestInSpendLogs"; -import NotificationsManager from "@/components/molecules/notifications_manager"; -import { parseErrorMessage } from "@/components/shared/errorUtils"; -import { ClockCircleOutlined } from "@ant-design/icons"; -import { Button, Form, Input, Modal, Skeleton, Space, Switch, Typography } from "antd"; -import React, { useEffect, useMemo } from "react"; - -interface SpendLogsSettingsModalProps { - isVisible: boolean; - onCancel: () => void; - onSuccess?: () => void; -} - -const SpendLogsSettingsModal: React.FC = ({ isVisible, onCancel, onSuccess }) => { - const [form] = Form.useForm(); - const { mutateAsync, isPending } = useStoreRequestInSpendLogs(); - const { mutateAsync: deleteField, isPending: isDeletingField } = useDeleteProxyConfigField(); - const { data: proxyConfigData, isLoading: isLoadingConfig, refetch } = useProxyConfig(ConfigType.GENERAL_SETTINGS); - const storePromptsValue = Form.useWatch('store_prompts_in_spend_logs', form); - - // Refetch config when modal opens to ensure we have the latest values - useEffect(() => { - if (isVisible) { - refetch(); - } - }, [isVisible, refetch]); - - // Compute initial values from fetched config data - const initialValues = useMemo(() => { - if (!proxyConfigData) { - return { - store_prompts_in_spend_logs: false, - maximum_spend_logs_retention_period: undefined, - }; - } - - const storePromptsField = proxyConfigData.find(field => field.field_name === 'store_prompts_in_spend_logs'); - const retentionPeriodField = proxyConfigData.find(field => field.field_name === 'maximum_spend_logs_retention_period'); - - return { - store_prompts_in_spend_logs: storePromptsField?.field_value ?? false, - maximum_spend_logs_retention_period: retentionPeriodField?.field_value ?? undefined, - }; - }, [proxyConfigData]); - - const handleFormSubmit = async (formValues: StoreRequestInSpendLogsParams) => { - try { - // If maximum_spend_logs_retention_period is empty/null, delete the field first - const retentionPeriodValue = formValues.maximum_spend_logs_retention_period; - const shouldDeleteRetentionPeriod = - !retentionPeriodValue || - (typeof retentionPeriodValue === "string" && retentionPeriodValue.trim() === ""); - - if (shouldDeleteRetentionPeriod) { - try { - await deleteField({ - config_type: ConfigType.GENERAL_SETTINGS, - field_name: GeneralSettingsFieldName.MAXIMUM_SPEND_LOGS_RETENTION_PERIOD, - }); - } catch (deleteError) { - // If field doesn't exist, that's okay - continue with update - console.warn("Failed to delete retention period field (may not exist):", deleteError); - } - } - - // Update the settings (excluding maximum_spend_logs_retention_period if it's empty) - const updateParams: StoreRequestInSpendLogsParams = { - store_prompts_in_spend_logs: formValues.store_prompts_in_spend_logs, - ...(retentionPeriodValue && - typeof retentionPeriodValue === "string" && - retentionPeriodValue.trim() !== "" && { - maximum_spend_logs_retention_period: retentionPeriodValue, - }), - }; - - await mutateAsync(updateParams, { - onSuccess: () => { - NotificationsManager.success("Spend logs settings updated successfully"); - refetch(); // Refetch config to get updated values - onSuccess?.(); - }, - onError: (error) => { - NotificationsManager.fromBackend("Failed to save spend logs settings: " + parseErrorMessage(error)); - }, - }); - } catch (error) { - NotificationsManager.fromBackend("Failed to save spend logs settings: " + parseErrorMessage(error)); - } - }; - - const handleCancel = () => { - form.resetFields(); - onCancel(); - }; - - return ( - Spend Logs Settings} - open={isVisible} - footer={ - - - - - } - onCancel={handleCancel} - > - -
- f.field_name === 'store_prompts_in_spend_logs')?.field_description || - "When enabled, prompts will be stored in spend logs for tracking and analysis purposes." - } - valuePropName="checked" - > -
- - {isLoadingConfig ? : form.setFieldValue('store_prompts_in_spend_logs', checked)} />} -
-
- - f.field_name === 'maximum_spend_logs_retention_period')?.field_description || - "Set the maximum retention period for spend logs (e.g., '7d' for 7 days, '30d' for 30 days). Leave empty for no limit." - } - > - {isLoadingConfig ? : } - />} - -
-
- ); -}; - -export default SpendLogsSettingsModal; diff --git a/ui/litellm-dashboard/src/components/view_logs/index.tsx b/ui/litellm-dashboard/src/components/view_logs/index.tsx index 8205c0b4b8..8a015d8305 100644 --- a/ui/litellm-dashboard/src/components/view_logs/index.tsx +++ b/ui/litellm-dashboard/src/components/view_logs/index.tsx @@ -4,7 +4,7 @@ import { useCallback, useDeferredValue, useEffect, useRef, useState } from "reac import GuardrailViewer from "@/components/view_logs/GuardrailViewer/GuardrailViewer"; import { formatNumberWithCommas } from "@/utils/dataUtils"; import { truncateString } from "@/utils/textUtils"; -import { SettingOutlined, SyncOutlined } from "@ant-design/icons"; +import { SyncOutlined } from "@ant-design/icons"; import { Row } from "@tanstack/react-table"; import { Switch, Tab, TabGroup, TabList, TabPanel, TabPanels } from "@tremor/react"; import { Button, Tag, Tooltip } from "antd"; @@ -28,7 +28,6 @@ import { useLogFilterLogic } from "./log_filter_logic"; import { LogDetailsDrawer } from "./LogDetailsDrawer"; import { getTimeRangeDisplay } from "./logs_utils"; import { RequestResponsePanel } from "./RequestResponsePanel"; -import SpendLogsSettingsModal from "./SpendLogsSettingsModal/SpendLogsSettingsModal"; import { DataTable } from "./table"; import { VectorStoreViewer } from "./VectorStoreViewer"; @@ -85,7 +84,6 @@ export default function SpendLogsTable({ const [selectedLog, setSelectedLog] = useState(null); const [isDrawerOpen, setIsDrawerOpen] = useState(false); const [selectedSessionId, setSelectedSessionId] = useState(null); - const [isSpendLogsSettingsModalVisible, setIsSpendLogsSettingsModalVisible] = useState(false); const [sortBy, setSortBy] = useState("startTime"); const [sortOrder, setSortOrder] = useState<"asc" | "desc">("desc"); @@ -490,11 +488,6 @@ export default function SpendLogsTable({

Request Logs

-
{selectedKeyInfo && selectedKeyIdInfoView && selectedKeyInfo.api_key === selectedKeyIdInfoView ? ( - setIsSpendLogsSettingsModalVisible(false)} - onSuccess={() => setIsSpendLogsSettingsModalVisible(false)} - />
@@ -725,7 +713,6 @@ export default function SpendLogsTable({ logEntry={selectedLog} sessionId={selectedSessionId} accessToken={accessToken} - onOpenSettings={() => setIsSpendLogsSettingsModalVisible(true)} allLogs={filteredData} onSelectLog={handleSelectLog} startTime={moment(startTime).utc().format("YYYY-MM-DD HH:mm:ss")} @@ -734,7 +721,7 @@ export default function SpendLogsTable({ ); } -export function RequestViewer({ row, onOpenSettings }: { row: Row; onOpenSettings?: () => void }) { +export function RequestViewer({ row }: { row: Row }) { // Helper function to clean metadata by removing specific fields const formatData = (input: any) => { if (typeof input === "string") { @@ -961,7 +948,7 @@ export function RequestViewer({ row, onOpenSettings }: { row: Row; onO /> {/* Configuration Info Message - Show when data is missing */} - + {/* Request/Response Panel */}
From d21e90f6831eebde5eb8f8d42604f5b57116d05e Mon Sep 17 00:00:00 2001 From: Mateo Wang <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 24 Apr 2026 14:10:42 -0700 Subject: [PATCH 03/57] [Feat] Day-0 support for GPT-5.5 and GPT-5.5 Pro (#26449) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * feat(openai): day-0 support for GPT-5.5 and GPT-5.5 Pro Add pricing + capability entries for the new GPT-5.5 family launched by OpenAI on 2026-04-24: - gpt-5.5 / gpt-5.5-2026-04-23 (chat): $5/$30/$0.50 per 1M input/output/cached input - gpt-5.5-pro / gpt-5.5-pro-2026-04-23 (responses-only): $60/$360/$6 per 1M input/output/cached input Other fees (long-context >272k, flex, batches, priority, cache discounts) follow the same ratios as GPT-5.4, with context window retained at 1.05M input / 128K output. No transformation / classifier code changes are required: OpenAIGPT5Config.is_model_gpt_5_4_plus_model() already matches 5.5+ via numeric version parsing, and model registration is driven from the JSON. The existing responses-API bridge for tools + reasoning_effort (litellm/main.py:970) already covers gpt-5.5-pro. Tests: - GPT5_MODELS regression list now covers gpt-5.5-pro and dated variants - New test_generic_cost_per_token_gpt55_pro cost-calc test - Updated test_generic_cost_per_token_gpt55 for long-context fields * fix(openai): mirror reasoning_effort flags onto gpt-5.5 dated variants gpt-5.5-2026-04-23 and gpt-5.5-pro-2026-04-23 were missing the supports_none_reasoning_effort, supports_xhigh_reasoning_effort, and supports_minimal_reasoning_effort flags that their non-dated counterparts define. Reasoning-effort routing in OpenAIGPT5Config is fully capability-driven from these JSON flags — since an absent flag is treated as False for opt-in levels (xhigh), users pinning to a dated snapshot would silently lose xhigh support and diverge from the base alias on logprobs + flexible temperature handling. Copy the flags onto both dated variants so every dated snapshot inherits the base model's reasoning-effort capability profile. Adds a parametrized regression test that asserts supports_{none,minimal,xhigh}_reasoning_effort parity between each dated variant and its non-dated counterpart, preventing future drift when new snapshots are added. --- ...odel_prices_and_context_window_backup.json | 148 +++++++++++++++++- model_prices_and_context_window.json | 148 +++++++++++++++++- .../llm_cost_calc/test_llm_cost_calc_utils.py | 86 +++++++++- .../llms/openai/test_is_model_gpt_5_model.py | 3 + 4 files changed, 382 insertions(+), 3 deletions(-) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 1cf7c1f6c7..49ce5022c5 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -19275,13 +19275,24 @@ }, "gpt-5.5": { "cache_read_input_token_cost": 5e-07, + "cache_read_input_token_cost_above_272k_tokens": 1e-06, + "cache_read_input_token_cost_flex": 2.5e-07, + "cache_read_input_token_cost_priority": 1e-06, "input_cost_per_token": 5e-06, + "input_cost_per_token_above_272k_tokens": 1e-05, + "input_cost_per_token_flex": 2.5e-06, + "input_cost_per_token_batches": 2.5e-06, + "input_cost_per_token_priority": 1e-05, "litellm_provider": "openai", - "max_input_tokens": 272000, + "max_input_tokens": 1050000, "max_output_tokens": 128000, "max_tokens": 128000, "mode": "chat", "output_cost_per_token": 3e-05, + "output_cost_per_token_above_272k_tokens": 4.5e-05, + "output_cost_per_token_flex": 1.5e-05, + "output_cost_per_token_batches": 1.5e-05, + "output_cost_per_token_priority": 6e-05, "supported_endpoints": [ "/v1/chat/completions", "/v1/batch", @@ -19305,10 +19316,145 @@ "supports_tool_choice": true, "supports_service_tier": true, "supports_vision": true, + "supports_web_search": true, "supports_none_reasoning_effort": true, "supports_xhigh_reasoning_effort": true, "supports_minimal_reasoning_effort": true }, + "gpt-5.5-2026-04-23": { + "cache_read_input_token_cost": 5e-07, + "cache_read_input_token_cost_above_272k_tokens": 1e-06, + "cache_read_input_token_cost_flex": 2.5e-07, + "cache_read_input_token_cost_priority": 1e-06, + "input_cost_per_token": 5e-06, + "input_cost_per_token_above_272k_tokens": 1e-05, + "input_cost_per_token_flex": 2.5e-06, + "input_cost_per_token_batches": 2.5e-06, + "input_cost_per_token_priority": 1e-05, + "litellm_provider": "openai", + "max_input_tokens": 1050000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 3e-05, + "output_cost_per_token_above_272k_tokens": 4.5e-05, + "output_cost_per_token_flex": 1.5e-05, + "output_cost_per_token_batches": 1.5e-05, + "output_cost_per_token_priority": 6e-05, + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/batch", + "/v1/responses" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_native_streaming": true, + "supports_parallel_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_service_tier": true, + "supports_vision": true, + "supports_web_search": true, + "supports_none_reasoning_effort": true, + "supports_xhigh_reasoning_effort": true, + "supports_minimal_reasoning_effort": true + }, + "gpt-5.5-pro": { + "cache_read_input_token_cost": 6e-06, + "cache_read_input_token_cost_above_272k_tokens": 1.2e-05, + "input_cost_per_token": 6e-05, + "input_cost_per_token_above_272k_tokens": 0.00012, + "input_cost_per_token_flex": 3e-05, + "input_cost_per_token_batches": 3e-05, + "litellm_provider": "openai", + "max_input_tokens": 1050000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "responses", + "output_cost_per_token": 0.00036, + "output_cost_per_token_above_272k_tokens": 0.00054, + "output_cost_per_token_flex": 0.00018, + "output_cost_per_token_batches": 0.00018, + "supported_endpoints": [ + "/v1/responses", + "/v1/batch" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_native_streaming": true, + "supports_parallel_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": false, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_service_tier": true, + "supports_vision": true, + "supports_web_search": true, + "supports_none_reasoning_effort": false, + "supports_xhigh_reasoning_effort": true, + "supports_minimal_reasoning_effort": true + }, + "gpt-5.5-pro-2026-04-23": { + "cache_read_input_token_cost": 6e-06, + "cache_read_input_token_cost_above_272k_tokens": 1.2e-05, + "input_cost_per_token": 6e-05, + "input_cost_per_token_above_272k_tokens": 0.00012, + "input_cost_per_token_flex": 3e-05, + "input_cost_per_token_batches": 3e-05, + "litellm_provider": "openai", + "max_input_tokens": 1050000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "responses", + "output_cost_per_token": 0.00036, + "output_cost_per_token_above_272k_tokens": 0.00054, + "output_cost_per_token_flex": 0.00018, + "output_cost_per_token_batches": 0.00018, + "supported_endpoints": [ + "/v1/responses", + "/v1/batch" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_native_streaming": true, + "supports_parallel_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": false, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_service_tier": true, + "supports_vision": true, + "supports_web_search": true, + "supports_none_reasoning_effort": false, + "supports_xhigh_reasoning_effort": true, + "supports_minimal_reasoning_effort": true + }, "gpt-5.4": { "cache_read_input_token_cost": 2.5e-07, "cache_read_input_token_cost_above_272k_tokens": 5e-07, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 8dcd52cae2..3733f07a30 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -19289,13 +19289,24 @@ }, "gpt-5.5": { "cache_read_input_token_cost": 5e-07, + "cache_read_input_token_cost_above_272k_tokens": 1e-06, + "cache_read_input_token_cost_flex": 2.5e-07, + "cache_read_input_token_cost_priority": 1e-06, "input_cost_per_token": 5e-06, + "input_cost_per_token_above_272k_tokens": 1e-05, + "input_cost_per_token_flex": 2.5e-06, + "input_cost_per_token_batches": 2.5e-06, + "input_cost_per_token_priority": 1e-05, "litellm_provider": "openai", - "max_input_tokens": 272000, + "max_input_tokens": 1050000, "max_output_tokens": 128000, "max_tokens": 128000, "mode": "chat", "output_cost_per_token": 3e-05, + "output_cost_per_token_above_272k_tokens": 4.5e-05, + "output_cost_per_token_flex": 1.5e-05, + "output_cost_per_token_batches": 1.5e-05, + "output_cost_per_token_priority": 6e-05, "supported_endpoints": [ "/v1/chat/completions", "/v1/batch", @@ -19319,10 +19330,145 @@ "supports_tool_choice": true, "supports_service_tier": true, "supports_vision": true, + "supports_web_search": true, "supports_none_reasoning_effort": true, "supports_xhigh_reasoning_effort": true, "supports_minimal_reasoning_effort": true }, + "gpt-5.5-2026-04-23": { + "cache_read_input_token_cost": 5e-07, + "cache_read_input_token_cost_above_272k_tokens": 1e-06, + "cache_read_input_token_cost_flex": 2.5e-07, + "cache_read_input_token_cost_priority": 1e-06, + "input_cost_per_token": 5e-06, + "input_cost_per_token_above_272k_tokens": 1e-05, + "input_cost_per_token_flex": 2.5e-06, + "input_cost_per_token_batches": 2.5e-06, + "input_cost_per_token_priority": 1e-05, + "litellm_provider": "openai", + "max_input_tokens": 1050000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 3e-05, + "output_cost_per_token_above_272k_tokens": 4.5e-05, + "output_cost_per_token_flex": 1.5e-05, + "output_cost_per_token_batches": 1.5e-05, + "output_cost_per_token_priority": 6e-05, + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/batch", + "/v1/responses" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_native_streaming": true, + "supports_parallel_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_service_tier": true, + "supports_vision": true, + "supports_web_search": true, + "supports_none_reasoning_effort": true, + "supports_xhigh_reasoning_effort": true, + "supports_minimal_reasoning_effort": true + }, + "gpt-5.5-pro": { + "cache_read_input_token_cost": 6e-06, + "cache_read_input_token_cost_above_272k_tokens": 1.2e-05, + "input_cost_per_token": 6e-05, + "input_cost_per_token_above_272k_tokens": 0.00012, + "input_cost_per_token_flex": 3e-05, + "input_cost_per_token_batches": 3e-05, + "litellm_provider": "openai", + "max_input_tokens": 1050000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "responses", + "output_cost_per_token": 0.00036, + "output_cost_per_token_above_272k_tokens": 0.00054, + "output_cost_per_token_flex": 0.00018, + "output_cost_per_token_batches": 0.00018, + "supported_endpoints": [ + "/v1/responses", + "/v1/batch" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_native_streaming": true, + "supports_parallel_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": false, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_service_tier": true, + "supports_vision": true, + "supports_web_search": true, + "supports_none_reasoning_effort": false, + "supports_xhigh_reasoning_effort": true, + "supports_minimal_reasoning_effort": true + }, + "gpt-5.5-pro-2026-04-23": { + "cache_read_input_token_cost": 6e-06, + "cache_read_input_token_cost_above_272k_tokens": 1.2e-05, + "input_cost_per_token": 6e-05, + "input_cost_per_token_above_272k_tokens": 0.00012, + "input_cost_per_token_flex": 3e-05, + "input_cost_per_token_batches": 3e-05, + "litellm_provider": "openai", + "max_input_tokens": 1050000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "responses", + "output_cost_per_token": 0.00036, + "output_cost_per_token_above_272k_tokens": 0.00054, + "output_cost_per_token_flex": 0.00018, + "output_cost_per_token_batches": 0.00018, + "supported_endpoints": [ + "/v1/responses", + "/v1/batch" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_native_streaming": true, + "supports_parallel_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": false, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_service_tier": true, + "supports_vision": true, + "supports_web_search": true, + "supports_none_reasoning_effort": false, + "supports_xhigh_reasoning_effort": true, + "supports_minimal_reasoning_effort": true + }, "gpt-5.4": { "cache_read_input_token_cost": 2.5e-07, "cache_read_input_token_cost_above_272k_tokens": 5e-07, diff --git a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py index 7144279ad0..5e37e2a342 100644 --- a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py +++ b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py @@ -343,7 +343,10 @@ def test_generic_cost_per_token_gpt55(): assert model_cost_map["cache_read_input_token_cost"] == 5e-7 assert model_cost_map["litellm_provider"] == "openai" assert model_cost_map["mode"] == "chat" - assert model_cost_map["max_input_tokens"] == 272000 + # gpt-5.5 inherits GPT-5.4's long-context window + tiered pricing. + assert model_cost_map["max_input_tokens"] == 1050000 + assert model_cost_map["input_cost_per_token_above_272k_tokens"] == 1e-5 + assert model_cost_map["output_cost_per_token_above_272k_tokens"] == 4.5e-5 prompt_tokens = 1000 completion_tokens = 500 @@ -365,6 +368,87 @@ def test_generic_cost_per_token_gpt55(): ) +def test_generic_cost_per_token_gpt55_pro(): + """gpt-5.5-pro: responses-only model — $60/1M input, $360/1M output, $6/1M cached input.""" + model = "gpt-5.5-pro" + custom_llm_provider = "openai" + os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" + litellm.model_cost = litellm.get_model_cost_map(url="") + + model_cost_map = litellm.model_cost[model] + + # Sanity-check the map values match OpenAI's published pricing. + assert model_cost_map["input_cost_per_token"] == 6e-5 + assert model_cost_map["output_cost_per_token"] == 3.6e-4 + assert model_cost_map["cache_read_input_token_cost"] == 6e-6 + assert model_cost_map["litellm_provider"] == "openai" + # gpt-5.5-pro is a responses-only model (no /v1/chat/completions endpoint). + assert model_cost_map["mode"] == "responses" + assert "/v1/chat/completions" not in model_cost_map["supported_endpoints"] + assert "/v1/responses" in model_cost_map["supported_endpoints"] + # Inherits GPT-5.4-pro's long-context window + tiered pricing (scaled 2x). + assert model_cost_map["max_input_tokens"] == 1050000 + assert model_cost_map["input_cost_per_token_above_272k_tokens"] == 1.2e-4 + assert model_cost_map["output_cost_per_token_above_272k_tokens"] == 5.4e-4 + + prompt_tokens = 1000 + completion_tokens = 500 + usage = Usage( + prompt_tokens=prompt_tokens, + completion_tokens=completion_tokens, + total_tokens=prompt_tokens + completion_tokens, + ) + prompt_cost, completion_cost = generic_cost_per_token( + model=model, + usage=usage, + custom_llm_provider=custom_llm_provider, + ) + assert round(prompt_cost, 10) == round( + model_cost_map["input_cost_per_token"] * prompt_tokens, 10 + ) + assert round(completion_cost, 10) == round( + model_cost_map["output_cost_per_token"] * completion_tokens, 10 + ) + + +@pytest.mark.parametrize( + "base_model,dated_model", + [ + ("gpt-5.5", "gpt-5.5-2026-04-23"), + ("gpt-5.5-pro", "gpt-5.5-pro-2026-04-23"), + ], +) +def test_gpt55_dated_variants_match_base_reasoning_effort_capabilities( + base_model, dated_model +): + """Dated snapshots must carry the same reasoning_effort capability flags as + their non-dated counterparts. + + Regression guard: ``supports_{none,minimal,xhigh}_reasoning_effort`` gate + downstream routing in ``OpenAIGPT5Config`` — a missing flag is treated as + ``False`` for opt-in levels (e.g. ``xhigh``), which silently diverges + behavior between ``gpt-5.5`` and ``gpt-5.5-2026-04-23``. Pinning to a + dated variant must never lose capabilities relative to the base alias. + """ + os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" + litellm.model_cost = litellm.get_model_cost_map(url="") + + base = litellm.model_cost[base_model] + dated = litellm.model_cost[dated_model] + + for flag in ( + "supports_none_reasoning_effort", + "supports_minimal_reasoning_effort", + "supports_xhigh_reasoning_effort", + ): + assert dated.get(flag) == base.get(flag), ( + f"{dated_model} has {flag}={dated.get(flag)!r}, " + f"but {base_model} has {flag}={base.get(flag)!r}. " + f"Dated snapshots must inherit the base model's reasoning_effort " + f"capability profile." + ) + + def test_generic_cost_per_token_anthropic_prompt_caching(): model = "claude-sonnet-4@20250514" usage = Usage( diff --git a/tests/test_litellm/llms/openai/test_is_model_gpt_5_model.py b/tests/test_litellm/llms/openai/test_is_model_gpt_5_model.py index e611d5e6b7..02dd9dade0 100644 --- a/tests/test_litellm/llms/openai/test_is_model_gpt_5_model.py +++ b/tests/test_litellm/llms/openai/test_is_model_gpt_5_model.py @@ -47,6 +47,9 @@ GPT5_MODELS = [ "gpt-5.3", "gpt-5.4", "gpt-5.5", + "gpt-5.5-pro", + "gpt-5.5-2026-04-23", # dated variant + "gpt-5.5-pro-2026-04-23", # dated variant "gpt-5.1-chat", # versioned chat — THE KEY REGRESSION CASE "gpt-5.2-chat", # versioned chat — also a regression case "gpt-5.3-chat", # versioned chat — THE KEY REGRESSION CASE From 21cf42f5681776d7ecb9cef6b84f58d65e1b8338 Mon Sep 17 00:00:00 2001 From: Milan Date: Sat, 25 Apr 2026 01:57:22 +0300 Subject: [PATCH 04/57] Add expired UI session key cleanup job Made-with: Cursor --- litellm/constants.py | 10 + .../expired_ui_session_key_cleanup_manager.py | 116 +++++++++++ litellm/proxy/proxy_server.py | 83 +++++++- ..._expired_ui_session_key_cleanup_manager.py | 186 ++++++++++++++++++ 4 files changed, 393 insertions(+), 2 deletions(-) create mode 100644 litellm/proxy/common_utils/expired_ui_session_key_cleanup_manager.py create mode 100644 tests/test_litellm/proxy/common_utils/test_expired_ui_session_key_cleanup_manager.py diff --git a/litellm/constants.py b/litellm/constants.py index 012599ab6a..ceda523637 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -1396,6 +1396,15 @@ LITELLM_KEY_ROTATION_LOCK_TTL_SECONDS = int( os.getenv("LITELLM_KEY_ROTATION_LOCK_TTL_SECONDS", 600) ) # 10 minutes default — caps the deadlock window if a pod crashes mid-rotation UI_SESSION_TOKEN_TEAM_ID = "litellm-dashboard" +LITELLM_EXPIRED_UI_SESSION_KEY_CLEANUP_ENABLED = os.getenv( + "LITELLM_EXPIRED_UI_SESSION_KEY_CLEANUP_ENABLED", "false" +) +LITELLM_EXPIRED_UI_SESSION_KEY_CLEANUP_INTERVAL_SECONDS = int( + os.getenv("LITELLM_EXPIRED_UI_SESSION_KEY_CLEANUP_INTERVAL_SECONDS", 86400) +) # 24 hours default +LITELLM_EXPIRED_UI_SESSION_KEY_CLEANUP_BATCH_SIZE = int( + os.getenv("LITELLM_EXPIRED_UI_SESSION_KEY_CLEANUP_BATCH_SIZE", 1000) +) LITELLM_PROXY_ADMIN_NAME = "default_user_id" ########################### CLI SSO AUTHENTICATION CONSTANTS ########################### @@ -1425,6 +1434,7 @@ CLOUDZERO_MAX_FETCHED_DATA_RECORDS = int( ) SPEND_LOG_CLEANUP_JOB_NAME = "spend_log_cleanup" KEY_ROTATION_JOB_NAME = "litellm_key_rotation_job" +EXPIRED_UI_SESSION_KEY_CLEANUP_JOB_NAME = "litellm_expired_ui_session_key_cleanup_job" SPEND_LOG_RUN_LOOPS = int(os.getenv("SPEND_LOG_RUN_LOOPS", 500)) SPEND_LOG_CLEANUP_BATCH_SIZE = int(os.getenv("SPEND_LOG_CLEANUP_BATCH_SIZE", 1000)) SPEND_LOG_QUEUE_SIZE_THRESHOLD = int(os.getenv("SPEND_LOG_QUEUE_SIZE_THRESHOLD", 100)) diff --git a/litellm/proxy/common_utils/expired_ui_session_key_cleanup_manager.py b/litellm/proxy/common_utils/expired_ui_session_key_cleanup_manager.py new file mode 100644 index 0000000000..f8e9cd3d07 --- /dev/null +++ b/litellm/proxy/common_utils/expired_ui_session_key_cleanup_manager.py @@ -0,0 +1,116 @@ +""" +Expired UI session key cleanup manager. + +Deletes expired virtual keys created for LiteLLM dashboard sessions. +""" + +from datetime import datetime, timezone +from typing import List + +from litellm._logging import verbose_proxy_logger +from litellm.caching import DualCache +from litellm.constants import ( + EXPIRED_UI_SESSION_KEY_CLEANUP_JOB_NAME, + LITELLM_EXPIRED_UI_SESSION_KEY_CLEANUP_BATCH_SIZE, + LITELLM_INTERNAL_JOBS_SERVICE_ACCOUNT_NAME, + UI_SESSION_TOKEN_TEAM_ID, +) +from litellm.proxy._types import KeyRequest, LiteLLM_VerificationToken, UserAPIKeyAuth +from litellm.proxy.hooks.key_management_event_hooks import KeyManagementEventHooks +from litellm.proxy.management_endpoints.key_management_endpoints import ( + delete_verification_tokens, +) +from litellm.proxy.utils import PrismaClient + + +class ExpiredUISessionKeyCleanupManager: + """ + Cleans up expired UI session keys. + """ + + def __init__( + self, + prisma_client: PrismaClient, + user_api_key_cache: DualCache, + pod_lock_manager=None, + ): + self.prisma_client = prisma_client + self.user_api_key_cache = user_api_key_cache + self.pod_lock_manager = pod_lock_manager + + async def cleanup_expired_keys(self) -> int: + """ + Main entry point for deleting expired UI session keys. + Uses PodLockManager to ensure only one pod runs cleanup in multi-pod deployments. + """ + lock_acquired = False + try: + if self.pod_lock_manager and self.pod_lock_manager.redis_cache: + lock_acquired = ( + await self.pod_lock_manager.acquire_lock( + cronjob_id=EXPIRED_UI_SESSION_KEY_CLEANUP_JOB_NAME, + ) + or False + ) + if not lock_acquired: + verbose_proxy_logger.warning( + "Expired UI session key cleanup: another pod is already " + "running cleanup or Redis lock acquisition failed - " + "skipping this cycle." + ) + return 0 + + verbose_proxy_logger.info("Starting expired UI session key cleanup...") + + expired_keys = await self._find_expired_ui_session_keys() + if not expired_keys: + verbose_proxy_logger.debug("No expired UI session keys found") + return 0 + + tokens = [key.token for key in expired_keys if key.token is not None] + if not tokens: + return 0 + + system_user = UserAPIKeyAuth.get_litellm_internal_jobs_user_api_key_auth() + response, keys_being_deleted = await delete_verification_tokens( + tokens=tokens, + user_api_key_cache=self.user_api_key_cache, + user_api_key_dict=system_user, + litellm_changed_by=LITELLM_INTERNAL_JOBS_SERVICE_ACCOUNT_NAME, + ) + await KeyManagementEventHooks.async_key_deleted_hook( + data=KeyRequest(keys=tokens), + keys_being_deleted=keys_being_deleted, + response=response or {}, + user_api_key_dict=system_user, + litellm_changed_by=LITELLM_INTERNAL_JOBS_SERVICE_ACCOUNT_NAME, + ) + verbose_proxy_logger.info( + "Deleted %s expired UI session key(s)", len(tokens) + ) + return len(tokens) + except Exception as e: + verbose_proxy_logger.error(f"Expired UI session key cleanup failed: {e}") + return 0 + finally: + if ( + lock_acquired + and self.pod_lock_manager + and self.pod_lock_manager.redis_cache + ): + await self.pod_lock_manager.release_lock( + cronjob_id=EXPIRED_UI_SESSION_KEY_CLEANUP_JOB_NAME, + ) + + async def _find_expired_ui_session_keys(self) -> List[LiteLLM_VerificationToken]: + """ + Find expired LiteLLM dashboard session keys. + """ + now = datetime.now(timezone.utc) + return await self.prisma_client.db.litellm_verificationtoken.find_many( + where={ + "team_id": UI_SESSION_TOKEN_TEAM_ID, + "expires": {"lt": now}, + }, + take=LITELLM_EXPIRED_UI_SESSION_KEY_CLEANUP_BATCH_SIZE, + ) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 00c2cf9e3d..bb7f31f1c2 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -6737,6 +6737,10 @@ class ProxyStartupEvent: Args: scheduler: The scheduler to add the background jobs to """ + global prisma_client + global proxy_logging_obj + global user_api_key_cache + ######################################################## # CloudZero Background Job ######################################################## @@ -6810,8 +6814,6 @@ class ProxyStartupEvent: ) # Get prisma_client and proxy_logging_obj from global scope - global prisma_client - global proxy_logging_obj if prisma_client is not None: # Reuse the PodLockManager from db_spend_update_writer pod_lock_manager = ( @@ -6841,6 +6843,83 @@ class ProxyStartupEvent: "Key rotation disabled (set LITELLM_KEY_ROTATION_ENABLED=true to enable)" ) + await cls._initialize_expired_ui_session_key_cleanup_background_job( + scheduler=scheduler + ) + + @classmethod + async def _initialize_expired_ui_session_key_cleanup_background_job( + cls, scheduler: AsyncIOScheduler + ): + """ + Initialize the expired UI session key cleanup background job. + """ + global prisma_client + global proxy_logging_obj + global user_api_key_cache + + ######################################################## + # Expired UI Session Key Cleanup Background Job + ######################################################## + from litellm.constants import ( + EXPIRED_UI_SESSION_KEY_CLEANUP_JOB_NAME, + LITELLM_EXPIRED_UI_SESSION_KEY_CLEANUP_ENABLED, + LITELLM_EXPIRED_UI_SESSION_KEY_CLEANUP_INTERVAL_SECONDS, + ) + + expired_ui_session_key_cleanup_enabled: Optional[bool] = str_to_bool( + LITELLM_EXPIRED_UI_SESSION_KEY_CLEANUP_ENABLED + ) + verbose_proxy_logger.debug( + "expired_ui_session_key_cleanup_enabled: " + f"{expired_ui_session_key_cleanup_enabled}" + ) + + if expired_ui_session_key_cleanup_enabled is True: + try: + from litellm.proxy.common_utils.expired_ui_session_key_cleanup_manager import ( + ExpiredUISessionKeyCleanupManager, + ) + + if prisma_client is not None: + pod_lock_manager = ( + proxy_logging_obj.db_spend_update_writer.pod_lock_manager + ) + expired_ui_session_key_cleanup_manager = ( + ExpiredUISessionKeyCleanupManager( + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + pod_lock_manager=pod_lock_manager, + ) + ) + verbose_proxy_logger.debug( + "Expired UI session key cleanup background job scheduled " + "every " + f"{LITELLM_EXPIRED_UI_SESSION_KEY_CLEANUP_INTERVAL_SECONDS} " + "seconds " + "(LITELLM_EXPIRED_UI_SESSION_KEY_CLEANUP_ENABLED=true)" + ) + scheduler.add_job( + expired_ui_session_key_cleanup_manager.cleanup_expired_keys, + "interval", + seconds=LITELLM_EXPIRED_UI_SESSION_KEY_CLEANUP_INTERVAL_SECONDS, + id=EXPIRED_UI_SESSION_KEY_CLEANUP_JOB_NAME, + ) + else: + verbose_proxy_logger.warning( + "Expired UI session key cleanup enabled but prisma_client " + "not available" + ) + except Exception as e: + verbose_proxy_logger.warning( + f"Failed to setup expired UI session key cleanup job: {e}" + ) + else: + verbose_proxy_logger.debug( + "Expired UI session key cleanup disabled (set " + "LITELLM_EXPIRED_UI_SESSION_KEY_CLEANUP_ENABLED=true to enable)" + ) + @classmethod async def _initialize_slack_alerting_jobs( cls, diff --git a/tests/test_litellm/proxy/common_utils/test_expired_ui_session_key_cleanup_manager.py b/tests/test_litellm/proxy/common_utils/test_expired_ui_session_key_cleanup_manager.py new file mode 100644 index 0000000000..8663a2136f --- /dev/null +++ b/tests/test_litellm/proxy/common_utils/test_expired_ui_session_key_cleanup_manager.py @@ -0,0 +1,186 @@ +""" +Test expired UI session key cleanup manager functionality. +""" + +import os +import sys +from datetime import datetime, timedelta, timezone +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest + +sys.path.insert(0, os.path.abspath("../../../..")) + +from litellm.constants import ( + EXPIRED_UI_SESSION_KEY_CLEANUP_JOB_NAME, + LITELLM_EXPIRED_UI_SESSION_KEY_CLEANUP_BATCH_SIZE, + LITELLM_INTERNAL_JOBS_SERVICE_ACCOUNT_NAME, + UI_SESSION_TOKEN_TEAM_ID, +) +from litellm.proxy._types import LiteLLM_VerificationToken +from litellm.proxy.common_utils.expired_ui_session_key_cleanup_manager import ( + ExpiredUISessionKeyCleanupManager, +) + + +class TestExpiredUISessionKeyCleanupManager: + """Test the ExpiredUISessionKeyCleanupManager class functionality.""" + + @pytest.mark.asyncio + async def test_find_expired_ui_session_keys_filters_dashboard_team_and_expiry(self): + mock_prisma_client = AsyncMock() + mock_cache = MagicMock() + manager = ExpiredUISessionKeyCleanupManager( + prisma_client=mock_prisma_client, + user_api_key_cache=mock_cache, + ) + + now = datetime(2026, 4, 25, 12, 0, 0, tzinfo=timezone.utc) + mock_keys = [ + LiteLLM_VerificationToken( + token="expired-dashboard-token", + team_id=UI_SESSION_TOKEN_TEAM_ID, + expires=now - timedelta(seconds=1), + ) + ] + mock_prisma_client.db.litellm_verificationtoken.find_many.return_value = ( + mock_keys + ) + + with patch( + "litellm.proxy.common_utils.expired_ui_session_key_cleanup_manager.datetime" + ) as mock_datetime: + mock_datetime.now.return_value = now + mock_datetime.side_effect = lambda *args, **kwargs: datetime( + *args, **kwargs + ) + + keys = await manager._find_expired_ui_session_keys() + + mock_prisma_client.db.litellm_verificationtoken.find_many.assert_called_once_with( + where={ + "team_id": UI_SESSION_TOKEN_TEAM_ID, + "expires": {"lt": now}, + }, + take=LITELLM_EXPIRED_UI_SESSION_KEY_CLEANUP_BATCH_SIZE, + ) + assert keys == mock_keys + + @pytest.mark.asyncio + async def test_cleanup_expired_keys_uses_existing_delete_path(self): + mock_prisma_client = AsyncMock() + mock_cache = MagicMock() + manager = ExpiredUISessionKeyCleanupManager( + prisma_client=mock_prisma_client, + user_api_key_cache=mock_cache, + ) + expired_key = LiteLLM_VerificationToken( + token="expired-dashboard-token", + team_id=UI_SESSION_TOKEN_TEAM_ID, + expires=datetime.now(timezone.utc) - timedelta(seconds=1), + ) + manager._find_expired_ui_session_keys = AsyncMock(return_value=[expired_key]) + + with patch( + "litellm.proxy.common_utils.expired_ui_session_key_cleanup_manager.delete_verification_tokens", + new_callable=AsyncMock, + ) as mock_delete_verification_tokens: + mock_delete_verification_tokens.return_value = ( + {"deleted_keys": ["expired-dashboard-token"], "failed_tokens": []}, + [expired_key], + ) + with patch( + "litellm.proxy.common_utils.expired_ui_session_key_cleanup_manager.KeyManagementEventHooks.async_key_deleted_hook", + new_callable=AsyncMock, + ) as mock_key_deleted_hook: + deleted_count = await manager.cleanup_expired_keys() + + assert deleted_count == 1 + mock_delete_verification_tokens.assert_called_once() + call_kwargs = mock_delete_verification_tokens.call_args.kwargs + assert call_kwargs["tokens"] == ["expired-dashboard-token"] + assert call_kwargs["user_api_key_cache"] == mock_cache + assert ( + call_kwargs["litellm_changed_by"] + == LITELLM_INTERNAL_JOBS_SERVICE_ACCOUNT_NAME + ) + assert call_kwargs["user_api_key_dict"].user_id == "system" + mock_key_deleted_hook.assert_called_once() + hook_kwargs = mock_key_deleted_hook.call_args.kwargs + assert hook_kwargs["data"].keys == ["expired-dashboard-token"] + assert hook_kwargs["keys_being_deleted"] == [expired_key] + assert hook_kwargs["response"] == { + "deleted_keys": ["expired-dashboard-token"], + "failed_tokens": [], + } + assert ( + hook_kwargs["litellm_changed_by"] + == LITELLM_INTERNAL_JOBS_SERVICE_ACCOUNT_NAME + ) + + @pytest.mark.asyncio + async def test_cleanup_expired_keys_noops_when_no_keys_found(self): + mock_prisma_client = AsyncMock() + mock_cache = MagicMock() + manager = ExpiredUISessionKeyCleanupManager( + prisma_client=mock_prisma_client, + user_api_key_cache=mock_cache, + ) + manager._find_expired_ui_session_keys = AsyncMock(return_value=[]) + + with patch( + "litellm.proxy.common_utils.expired_ui_session_key_cleanup_manager.delete_verification_tokens", + new_callable=AsyncMock, + ) as mock_delete_verification_tokens: + deleted_count = await manager.cleanup_expired_keys() + + assert deleted_count == 0 + mock_delete_verification_tokens.assert_not_called() + + @pytest.mark.asyncio + async def test_cleanup_expired_keys_skips_when_lock_held(self): + mock_prisma_client = AsyncMock() + mock_cache = MagicMock() + mock_pod_lock_manager = MagicMock() + mock_pod_lock_manager.redis_cache = MagicMock() + mock_pod_lock_manager.acquire_lock = AsyncMock(return_value=False) + mock_pod_lock_manager.release_lock = AsyncMock() + + manager = ExpiredUISessionKeyCleanupManager( + prisma_client=mock_prisma_client, + user_api_key_cache=mock_cache, + pod_lock_manager=mock_pod_lock_manager, + ) + manager._find_expired_ui_session_keys = AsyncMock() + + deleted_count = await manager.cleanup_expired_keys() + + assert deleted_count == 0 + mock_pod_lock_manager.acquire_lock.assert_called_once_with( + cronjob_id=EXPIRED_UI_SESSION_KEY_CLEANUP_JOB_NAME, + ) + manager._find_expired_ui_session_keys.assert_not_called() + mock_pod_lock_manager.release_lock.assert_not_called() + + @pytest.mark.asyncio + async def test_cleanup_expired_keys_releases_acquired_lock(self): + mock_prisma_client = AsyncMock() + mock_cache = MagicMock() + mock_pod_lock_manager = MagicMock() + mock_pod_lock_manager.redis_cache = MagicMock() + mock_pod_lock_manager.acquire_lock = AsyncMock(return_value=True) + mock_pod_lock_manager.release_lock = AsyncMock() + + manager = ExpiredUISessionKeyCleanupManager( + prisma_client=mock_prisma_client, + user_api_key_cache=mock_cache, + pod_lock_manager=mock_pod_lock_manager, + ) + manager._find_expired_ui_session_keys = AsyncMock(return_value=[]) + + deleted_count = await manager.cleanup_expired_keys() + + assert deleted_count == 0 + mock_pod_lock_manager.release_lock.assert_called_once_with( + cronjob_id=EXPIRED_UI_SESSION_KEY_CLEANUP_JOB_NAME, + ) From 60f6a5dcfa9c2643bae78d7c2bb76bb2622201eb Mon Sep 17 00:00:00 2001 From: Milan Date: Sat, 25 Apr 2026 02:12:31 +0300 Subject: [PATCH 05/57] Tune expired UI session cleanup lock logging Made-with: Cursor --- .../common_utils/expired_ui_session_key_cleanup_manager.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/litellm/proxy/common_utils/expired_ui_session_key_cleanup_manager.py b/litellm/proxy/common_utils/expired_ui_session_key_cleanup_manager.py index f8e9cd3d07..61872603fb 100644 --- a/litellm/proxy/common_utils/expired_ui_session_key_cleanup_manager.py +++ b/litellm/proxy/common_utils/expired_ui_session_key_cleanup_manager.py @@ -53,7 +53,7 @@ class ExpiredUISessionKeyCleanupManager: or False ) if not lock_acquired: - verbose_proxy_logger.warning( + verbose_proxy_logger.debug( "Expired UI session key cleanup: another pod is already " "running cleanup or Redis lock acquisition failed - " "skipping this cycle." From 69c5840e5fe821e202e5efbee86fb274f4fc4b4d Mon Sep 17 00:00:00 2001 From: Milan Date: Sat, 25 Apr 2026 02:21:28 +0300 Subject: [PATCH 06/57] Test cleanup of multiple expired UI session keys Made-with: Cursor --- ..._expired_ui_session_key_cleanup_manager.py | 44 +++++++++++++++++++ 1 file changed, 44 insertions(+) diff --git a/tests/test_litellm/proxy/common_utils/test_expired_ui_session_key_cleanup_manager.py b/tests/test_litellm/proxy/common_utils/test_expired_ui_session_key_cleanup_manager.py index 8663a2136f..85f4e9dd6f 100644 --- a/tests/test_litellm/proxy/common_utils/test_expired_ui_session_key_cleanup_manager.py +++ b/tests/test_litellm/proxy/common_utils/test_expired_ui_session_key_cleanup_manager.py @@ -118,6 +118,50 @@ class TestExpiredUISessionKeyCleanupManager: == LITELLM_INTERNAL_JOBS_SERVICE_ACCOUNT_NAME ) + @pytest.mark.asyncio + async def test_cleanup_expired_keys_deletes_multiple_keys(self): + mock_prisma_client = AsyncMock() + mock_cache = MagicMock() + manager = ExpiredUISessionKeyCleanupManager( + prisma_client=mock_prisma_client, + user_api_key_cache=mock_cache, + ) + expired_keys = [ + LiteLLM_VerificationToken( + token="expired-dashboard-token-1", + team_id=UI_SESSION_TOKEN_TEAM_ID, + expires=datetime.now(timezone.utc) - timedelta(seconds=1), + ), + LiteLLM_VerificationToken( + token="expired-dashboard-token-2", + team_id=UI_SESSION_TOKEN_TEAM_ID, + expires=datetime.now(timezone.utc) - timedelta(seconds=1), + ), + ] + tokens = [key.token for key in expired_keys] + manager._find_expired_ui_session_keys = AsyncMock(return_value=expired_keys) + + with patch( + "litellm.proxy.common_utils.expired_ui_session_key_cleanup_manager.delete_verification_tokens", + new_callable=AsyncMock, + ) as mock_delete_verification_tokens: + mock_delete_verification_tokens.return_value = ( + {"deleted_keys": tokens, "failed_tokens": []}, + expired_keys, + ) + with patch( + "litellm.proxy.common_utils.expired_ui_session_key_cleanup_manager.KeyManagementEventHooks.async_key_deleted_hook", + new_callable=AsyncMock, + ) as mock_key_deleted_hook: + deleted_count = await manager.cleanup_expired_keys() + + assert deleted_count == 2 + assert mock_delete_verification_tokens.call_args.kwargs["tokens"] == tokens + hook_kwargs = mock_key_deleted_hook.call_args.kwargs + assert hook_kwargs["data"].keys == tokens + assert hook_kwargs["keys_being_deleted"] == expired_keys + assert hook_kwargs["response"] == {"deleted_keys": tokens, "failed_tokens": []} + @pytest.mark.asyncio async def test_cleanup_expired_keys_noops_when_no_keys_found(self): mock_prisma_client = AsyncMock() From 22439a119e601c25cbe54dc78970015c6a839fba Mon Sep 17 00:00:00 2001 From: Milan Date: Sat, 25 Apr 2026 03:01:34 +0300 Subject: [PATCH 07/57] Handle cleanup delete races and accurate counts Made-with: Cursor --- .../expired_ui_session_key_cleanup_manager.py | 48 ++++++- ..._expired_ui_session_key_cleanup_manager.py | 118 ++++++++++++++++++ 2 files changed, 162 insertions(+), 4 deletions(-) diff --git a/litellm/proxy/common_utils/expired_ui_session_key_cleanup_manager.py b/litellm/proxy/common_utils/expired_ui_session_key_cleanup_manager.py index 61872603fb..c25d853312 100644 --- a/litellm/proxy/common_utils/expired_ui_session_key_cleanup_manager.py +++ b/litellm/proxy/common_utils/expired_ui_session_key_cleanup_manager.py @@ -5,7 +5,7 @@ Deletes expired virtual keys created for LiteLLM dashboard sessions. """ from datetime import datetime, timezone -from typing import List +from typing import Any, Dict, List, Optional from litellm._logging import verbose_proxy_logger from litellm.caching import DualCache @@ -85,11 +85,22 @@ class ExpiredUISessionKeyCleanupManager: user_api_key_dict=system_user, litellm_changed_by=LITELLM_INTERNAL_JOBS_SERVICE_ACCOUNT_NAME, ) - verbose_proxy_logger.info( - "Deleted %s expired UI session key(s)", len(tokens) + deleted_count = self._get_deleted_token_count( + tokens=tokens, + response=response, ) - return len(tokens) + verbose_proxy_logger.info( + "Deleted %s expired UI session key(s)", deleted_count + ) + return deleted_count except Exception as e: + if getattr(e, "status_code", None) == 404: + verbose_proxy_logger.debug( + "Expired UI session key cleanup skipped because selected keys " + "were already deleted: %s", + e, + ) + return 0 verbose_proxy_logger.error(f"Expired UI session key cleanup failed: {e}") return 0 finally: @@ -102,6 +113,35 @@ class ExpiredUISessionKeyCleanupManager: cronjob_id=EXPIRED_UI_SESSION_KEY_CLEANUP_JOB_NAME, ) + @staticmethod + def _get_deleted_token_count( + tokens: List[str], + response: Optional[Dict[str, Any]], + ) -> int: + """ + Return the number of tokens actually deleted from the delete helper response. + """ + if response is None: + return len(tokens) + + deleted_keys = response.get("deleted_keys") + if isinstance(deleted_keys, list): + return len(deleted_keys) + if isinstance(deleted_keys, int): + return deleted_keys + if isinstance(deleted_keys, dict): + nested_deleted_keys = deleted_keys.get("deleted_keys") + if isinstance(nested_deleted_keys, list): + return len(nested_deleted_keys) + if isinstance(nested_deleted_keys, int): + return nested_deleted_keys + + failed_tokens = response.get("failed_tokens") or [] + if failed_tokens: + return max(len(tokens) - len(set(failed_tokens)), 0) + + return len(tokens) + async def _find_expired_ui_session_keys(self) -> List[LiteLLM_VerificationToken]: """ Find expired LiteLLM dashboard session keys. diff --git a/tests/test_litellm/proxy/common_utils/test_expired_ui_session_key_cleanup_manager.py b/tests/test_litellm/proxy/common_utils/test_expired_ui_session_key_cleanup_manager.py index 85f4e9dd6f..3efeeee9a2 100644 --- a/tests/test_litellm/proxy/common_utils/test_expired_ui_session_key_cleanup_manager.py +++ b/tests/test_litellm/proxy/common_utils/test_expired_ui_session_key_cleanup_manager.py @@ -8,6 +8,7 @@ from datetime import datetime, timedelta, timezone from unittest.mock import AsyncMock, MagicMock, patch import pytest +from fastapi import HTTPException, status sys.path.insert(0, os.path.abspath("../../../..")) @@ -162,6 +163,123 @@ class TestExpiredUISessionKeyCleanupManager: assert hook_kwargs["keys_being_deleted"] == expired_keys assert hook_kwargs["response"] == {"deleted_keys": tokens, "failed_tokens": []} + @pytest.mark.asyncio + async def test_cleanup_expired_keys_returns_successful_delete_count(self): + mock_prisma_client = AsyncMock() + mock_cache = MagicMock() + manager = ExpiredUISessionKeyCleanupManager( + prisma_client=mock_prisma_client, + user_api_key_cache=mock_cache, + ) + expired_keys = [ + LiteLLM_VerificationToken( + token="expired-dashboard-token-1", + team_id=UI_SESSION_TOKEN_TEAM_ID, + expires=datetime.now(timezone.utc) - timedelta(seconds=1), + ), + LiteLLM_VerificationToken( + token="expired-dashboard-token-2", + team_id=UI_SESSION_TOKEN_TEAM_ID, + expires=datetime.now(timezone.utc) - timedelta(seconds=1), + ), + ] + tokens = [key.token for key in expired_keys] + manager._find_expired_ui_session_keys = AsyncMock(return_value=expired_keys) + + with patch( + "litellm.proxy.common_utils.expired_ui_session_key_cleanup_manager.delete_verification_tokens", + new_callable=AsyncMock, + ) as mock_delete_verification_tokens: + mock_delete_verification_tokens.return_value = ( + { + "deleted_keys": ["expired-dashboard-token-1"], + "failed_tokens": ["expired-dashboard-token-2"], + }, + [expired_keys[0]], + ) + with patch( + "litellm.proxy.common_utils.expired_ui_session_key_cleanup_manager.KeyManagementEventHooks.async_key_deleted_hook", + new_callable=AsyncMock, + ): + deleted_count = await manager.cleanup_expired_keys() + + assert deleted_count == 1 + assert mock_delete_verification_tokens.call_args.kwargs["tokens"] == tokens + + @pytest.mark.asyncio + async def test_cleanup_expired_keys_counts_nested_delete_response(self): + mock_prisma_client = AsyncMock() + mock_cache = MagicMock() + manager = ExpiredUISessionKeyCleanupManager( + prisma_client=mock_prisma_client, + user_api_key_cache=mock_cache, + ) + expired_keys = [ + LiteLLM_VerificationToken( + token="expired-dashboard-token-1", + team_id=UI_SESSION_TOKEN_TEAM_ID, + expires=datetime.now(timezone.utc) - timedelta(seconds=1), + ), + LiteLLM_VerificationToken( + token="expired-dashboard-token-2", + team_id=UI_SESSION_TOKEN_TEAM_ID, + expires=datetime.now(timezone.utc) - timedelta(seconds=1), + ), + ] + tokens = [key.token for key in expired_keys] + manager._find_expired_ui_session_keys = AsyncMock(return_value=expired_keys) + + with patch( + "litellm.proxy.common_utils.expired_ui_session_key_cleanup_manager.delete_verification_tokens", + new_callable=AsyncMock, + ) as mock_delete_verification_tokens: + mock_delete_verification_tokens.return_value = ( + { + "deleted_keys": {"deleted_keys": 2}, + "failed_tokens": tokens, + }, + expired_keys, + ) + with patch( + "litellm.proxy.common_utils.expired_ui_session_key_cleanup_manager.KeyManagementEventHooks.async_key_deleted_hook", + new_callable=AsyncMock, + ): + deleted_count = await manager.cleanup_expired_keys() + + assert deleted_count == 2 + + @pytest.mark.asyncio + async def test_cleanup_expired_keys_treats_missing_keys_as_noop(self): + mock_prisma_client = AsyncMock() + mock_cache = MagicMock() + manager = ExpiredUISessionKeyCleanupManager( + prisma_client=mock_prisma_client, + user_api_key_cache=mock_cache, + ) + expired_key = LiteLLM_VerificationToken( + token="expired-dashboard-token", + team_id=UI_SESSION_TOKEN_TEAM_ID, + expires=datetime.now(timezone.utc) - timedelta(seconds=1), + ) + manager._find_expired_ui_session_keys = AsyncMock(return_value=[expired_key]) + + with patch( + "litellm.proxy.common_utils.expired_ui_session_key_cleanup_manager.delete_verification_tokens", + new_callable=AsyncMock, + ) as mock_delete_verification_tokens: + mock_delete_verification_tokens.side_effect = HTTPException( + status_code=status.HTTP_404_NOT_FOUND, + detail={"error": "No keys found"}, + ) + with patch( + "litellm.proxy.common_utils.expired_ui_session_key_cleanup_manager.KeyManagementEventHooks.async_key_deleted_hook", + new_callable=AsyncMock, + ) as mock_key_deleted_hook: + deleted_count = await manager.cleanup_expired_keys() + + assert deleted_count == 0 + mock_key_deleted_hook.assert_not_called() + @pytest.mark.asyncio async def test_cleanup_expired_keys_noops_when_no_keys_found(self): mock_prisma_client = AsyncMock() From 0beec45c138d3d92a956cc1e1dc75fa2bd2ae782 Mon Sep 17 00:00:00 2001 From: Mateo Wang <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 24 Apr 2026 18:10:21 -0700 Subject: [PATCH 08/57] [Feat] Add azure/gpt-5.5 + azure/gpt-5.5-pro entries (+ dated variants) (#26361) * feat(azure): add azure/gpt-5.5 + azure/gpt-5.5-pro entries (+ dated variants) Azure variants of OpenAI's GPT-5.5 family. Microsoft has not yet shipped GPT-5.5 on Azure OpenAI (latest GA on the Foundry models page is GPT-5.4 as of 2026-04-24), but adding the entries day-0 mirrors the established precedent for azure/gpt-5.4* (which were in the cost map before the Azure rollout) so cost tracking and capability flags work the moment customers deploy. Schema follows the existing azure/gpt-5.4* shape: - Same base/long-context pricing as openai/gpt-5.5*: $5/$30 chat, $60/$360 pro per 1M, with priority tier 2x base - Azure variants drop the flex/batches keys (Azure has no flex tier) but keep priority pricing, matching gpt-5.4* precedent - mode=chat for the thinking model, mode=responses for pro reasoning_effort capability flags mirror the OpenAI variants exactly since Azure proxies the same API contract: minimal rejection on both chat and pro, low/none rejection on pro. Once #26456 (which sets supports_low_reasoning_effort + minimal=false on openai/gpt-5.5*) lands, OpenAI and Azure flag profiles align. Tests pin entry presence + pricing for all four Azure variants and verify the live-API-derived reasoning_effort flags. * test: register supports_low_reasoning_effort in cost-map JSON schema azure/gpt-5.5-pro and azure/gpt-5.5-pro-2026-04-23 added in this branch carry supports_low_reasoning_effort=false. The strict 'additionalProperties: false' schema in test_aaamodel_prices_and_context_window_json_is_valid rejected the new key. Register it alongside the other supports_*_reasoning_effort entries. Note: the runtime side of this flag (code that reads it) lands in #26456. Until that PR merges the flag is inert for both Azure and OpenAI pro entries, but having the schema accept it lets cost-map tests pass on either merge order. --- ...odel_prices_and_context_window_backup.json | 163 ++++++++++++++++++ model_prices_and_context_window.json | 163 ++++++++++++++++++ .../llm_cost_calc/test_llm_cost_calc_utils.py | 57 ++++++ tests/test_litellm/test_utils.py | 1 + 4 files changed, 384 insertions(+) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 49ce5022c5..e11fab9077 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -4645,6 +4645,169 @@ "supports_vision": true, "supports_web_search": true }, + "azure/gpt-5.5": { + "cache_read_input_token_cost": 5e-07, + "cache_read_input_token_cost_above_272k_tokens": 1e-06, + "cache_read_input_token_cost_priority": 1e-06, + "cache_read_input_token_cost_above_272k_tokens_priority": 2e-06, + "input_cost_per_token": 5e-06, + "input_cost_per_token_above_272k_tokens": 1e-05, + "input_cost_per_token_priority": 1e-05, + "input_cost_per_token_above_272k_tokens_priority": 2e-05, + "litellm_provider": "azure", + "max_input_tokens": 1050000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 3e-05, + "output_cost_per_token_above_272k_tokens": 4.5e-05, + "output_cost_per_token_priority": 6e-05, + "output_cost_per_token_above_272k_tokens_priority": 9e-05, + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/batch", + "/v1/responses" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_native_streaming": true, + "supports_parallel_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_service_tier": true, + "supports_vision": true, + "supports_web_search": true, + "supports_none_reasoning_effort": true, + "supports_xhigh_reasoning_effort": true, + "supports_minimal_reasoning_effort": false + }, + "azure/gpt-5.5-2026-04-23": { + "cache_read_input_token_cost": 5e-07, + "cache_read_input_token_cost_above_272k_tokens": 1e-06, + "cache_read_input_token_cost_priority": 1e-06, + "cache_read_input_token_cost_above_272k_tokens_priority": 2e-06, + "input_cost_per_token": 5e-06, + "input_cost_per_token_above_272k_tokens": 1e-05, + "input_cost_per_token_priority": 1e-05, + "input_cost_per_token_above_272k_tokens_priority": 2e-05, + "litellm_provider": "azure", + "max_input_tokens": 1050000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 3e-05, + "output_cost_per_token_above_272k_tokens": 4.5e-05, + "output_cost_per_token_priority": 6e-05, + "output_cost_per_token_above_272k_tokens_priority": 9e-05, + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/batch", + "/v1/responses" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_native_streaming": true, + "supports_parallel_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_service_tier": true, + "supports_vision": true, + "supports_web_search": true + }, + "azure/gpt-5.5-pro": { + "cache_read_input_token_cost": 6e-06, + "cache_read_input_token_cost_above_272k_tokens": 1.2e-05, + "input_cost_per_token": 6e-05, + "input_cost_per_token_above_272k_tokens": 0.00012, + "litellm_provider": "azure", + "max_input_tokens": 1050000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "responses", + "output_cost_per_token": 0.00036, + "output_cost_per_token_above_272k_tokens": 0.00054, + "supported_endpoints": [ + "/v1/batch", + "/v1/responses" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_native_streaming": true, + "supports_parallel_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": false, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_vision": true, + "supports_web_search": true, + "supports_none_reasoning_effort": false, + "supports_xhigh_reasoning_effort": true, + "supports_minimal_reasoning_effort": false, + "supports_low_reasoning_effort": false + }, + "azure/gpt-5.5-pro-2026-04-23": { + "cache_read_input_token_cost": 6e-06, + "cache_read_input_token_cost_above_272k_tokens": 1.2e-05, + "input_cost_per_token": 6e-05, + "input_cost_per_token_above_272k_tokens": 0.00012, + "litellm_provider": "azure", + "max_input_tokens": 1050000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "responses", + "output_cost_per_token": 0.00036, + "output_cost_per_token_above_272k_tokens": 0.00054, + "supported_endpoints": [ + "/v1/batch", + "/v1/responses" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_native_streaming": true, + "supports_parallel_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": false, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_vision": true, + "supports_web_search": true + }, "azure/gpt-5.4-mini": { "cache_read_input_token_cost": 7.5e-08, "input_cost_per_token": 7.5e-07, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 3733f07a30..6eb6f1a9a9 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -4659,6 +4659,169 @@ "supports_vision": true, "supports_web_search": true }, + "azure/gpt-5.5": { + "cache_read_input_token_cost": 5e-07, + "cache_read_input_token_cost_above_272k_tokens": 1e-06, + "cache_read_input_token_cost_priority": 1e-06, + "cache_read_input_token_cost_above_272k_tokens_priority": 2e-06, + "input_cost_per_token": 5e-06, + "input_cost_per_token_above_272k_tokens": 1e-05, + "input_cost_per_token_priority": 1e-05, + "input_cost_per_token_above_272k_tokens_priority": 2e-05, + "litellm_provider": "azure", + "max_input_tokens": 1050000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 3e-05, + "output_cost_per_token_above_272k_tokens": 4.5e-05, + "output_cost_per_token_priority": 6e-05, + "output_cost_per_token_above_272k_tokens_priority": 9e-05, + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/batch", + "/v1/responses" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_native_streaming": true, + "supports_parallel_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_service_tier": true, + "supports_vision": true, + "supports_web_search": true, + "supports_none_reasoning_effort": true, + "supports_xhigh_reasoning_effort": true, + "supports_minimal_reasoning_effort": false + }, + "azure/gpt-5.5-2026-04-23": { + "cache_read_input_token_cost": 5e-07, + "cache_read_input_token_cost_above_272k_tokens": 1e-06, + "cache_read_input_token_cost_priority": 1e-06, + "cache_read_input_token_cost_above_272k_tokens_priority": 2e-06, + "input_cost_per_token": 5e-06, + "input_cost_per_token_above_272k_tokens": 1e-05, + "input_cost_per_token_priority": 1e-05, + "input_cost_per_token_above_272k_tokens_priority": 2e-05, + "litellm_provider": "azure", + "max_input_tokens": 1050000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 3e-05, + "output_cost_per_token_above_272k_tokens": 4.5e-05, + "output_cost_per_token_priority": 6e-05, + "output_cost_per_token_above_272k_tokens_priority": 9e-05, + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/batch", + "/v1/responses" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_native_streaming": true, + "supports_parallel_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_service_tier": true, + "supports_vision": true, + "supports_web_search": true + }, + "azure/gpt-5.5-pro": { + "cache_read_input_token_cost": 6e-06, + "cache_read_input_token_cost_above_272k_tokens": 1.2e-05, + "input_cost_per_token": 6e-05, + "input_cost_per_token_above_272k_tokens": 0.00012, + "litellm_provider": "azure", + "max_input_tokens": 1050000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "responses", + "output_cost_per_token": 0.00036, + "output_cost_per_token_above_272k_tokens": 0.00054, + "supported_endpoints": [ + "/v1/batch", + "/v1/responses" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_native_streaming": true, + "supports_parallel_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": false, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_vision": true, + "supports_web_search": true, + "supports_none_reasoning_effort": false, + "supports_xhigh_reasoning_effort": true, + "supports_minimal_reasoning_effort": false, + "supports_low_reasoning_effort": false + }, + "azure/gpt-5.5-pro-2026-04-23": { + "cache_read_input_token_cost": 6e-06, + "cache_read_input_token_cost_above_272k_tokens": 1.2e-05, + "input_cost_per_token": 6e-05, + "input_cost_per_token_above_272k_tokens": 0.00012, + "litellm_provider": "azure", + "max_input_tokens": 1050000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "responses", + "output_cost_per_token": 0.00036, + "output_cost_per_token_above_272k_tokens": 0.00054, + "supported_endpoints": [ + "/v1/batch", + "/v1/responses" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_native_streaming": true, + "supports_parallel_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": false, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_vision": true, + "supports_web_search": true + }, "azure/gpt-5.4-mini": { "cache_read_input_token_cost": 7.5e-08, "input_cost_per_token": 7.5e-07, diff --git a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py index 5e37e2a342..3016762043 100644 --- a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py +++ b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py @@ -449,6 +449,63 @@ def test_gpt55_dated_variants_match_base_reasoning_effort_capabilities( ) +@pytest.mark.parametrize( + "model,expected_mode,expected_input,expected_output,expected_cache_read", + [ + ("azure/gpt-5.5", "chat", 5e-6, 3e-5, 5e-7), + ("azure/gpt-5.5-2026-04-23", "chat", 5e-6, 3e-5, 5e-7), + ("azure/gpt-5.5-pro", "responses", 6e-5, 3.6e-4, 6e-6), + ("azure/gpt-5.5-pro-2026-04-23", "responses", 6e-5, 3.6e-4, 6e-6), + ], +) +def test_azure_gpt55_entries_present_with_correct_pricing( + model, expected_mode, expected_input, expected_output, expected_cache_read +): + """Day-0 Azure entries for GPT-5.5 mirror the OpenAI pricing structure. + + Pricing parity with openai/gpt-5.5* (verified against OpenAI's pricing page + on 2026-04-24): $5/$30 input/output per 1M for chat, $60/$360 for pro. + Cache discount is 10% of input. + """ + os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" + litellm.model_cost = litellm.get_model_cost_map(url="") + + m = litellm.model_cost[model] + assert m["litellm_provider"] == "azure" + assert m["mode"] == expected_mode + assert m["input_cost_per_token"] == expected_input + assert m["output_cost_per_token"] == expected_output + assert m["cache_read_input_token_cost"] == expected_cache_read + # Long-context window inherited from gpt-5.4 / openai gpt-5.5. + assert m["max_input_tokens"] == 1050000 + assert m["max_output_tokens"] == 128000 + + +@pytest.mark.parametrize( + "model,expected_none,expected_minimal,expected_xhigh", + [ + # Mirror live OpenAI API contract (verified via openai/gpt-5.5* on + # 2026-04-24): chat accepts {none, low, medium, high, xhigh} but NOT + # minimal; pro accepts {medium, high, xhigh} only. + # NOTE: openai/gpt-5.5* entries currently set supports_minimal=true on + # main (pre #26456). Once that PR lands, OpenAI + Azure flags align. + ("azure/gpt-5.5", True, False, True), + ("azure/gpt-5.5-pro", False, False, True), + ], +) +def test_azure_gpt55_reasoning_effort_flags_match_live_openai_api( + model, expected_none, expected_minimal, expected_xhigh +): + """Azure entries pin reasoning_effort flags to OpenAI's actual API contract.""" + os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" + litellm.model_cost = litellm.get_model_cost_map(url="") + + m = litellm.model_cost[model] + assert m.get("supports_none_reasoning_effort") is expected_none + assert m.get("supports_minimal_reasoning_effort") is expected_minimal + assert m.get("supports_xhigh_reasoning_effort") is expected_xhigh + + def test_generic_cost_per_token_anthropic_prompt_caching(): model = "claude-sonnet-4@20250514" usage = Usage( diff --git a/tests/test_litellm/test_utils.py b/tests/test_litellm/test_utils.py index 67b6269619..c6b0f49ff6 100644 --- a/tests/test_litellm/test_utils.py +++ b/tests/test_litellm/test_utils.py @@ -769,6 +769,7 @@ def test_aaamodel_prices_and_context_window_json_is_valid(): "uses_embed_content": {"type": "boolean"}, "supports_reasoning": {"type": "boolean"}, "supports_minimal_reasoning_effort": {"type": "boolean"}, + "supports_low_reasoning_effort": {"type": "boolean"}, "supports_none_reasoning_effort": {"type": "boolean"}, "supports_xhigh_reasoning_effort": {"type": "boolean"}, "supports_max_reasoning_effort": {"type": "boolean"}, From 3e4f9af9555df6fef78e22624778e655d3067173 Mon Sep 17 00:00:00 2001 From: Sameer Kankute Date: Mon, 27 Apr 2026 12:05:49 +0530 Subject: [PATCH 09/57] Add support for azure entra discovery endpoint --- .../mcp_server/mcp_server_manager.py | 41 +++++- .../mcp_server/test_mcp_server_manager.py | 130 ++++++++++++++++++ 2 files changed, 166 insertions(+), 5 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index 251f271903..6aaf37c91e 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -1488,11 +1488,18 @@ class MCPServerManager: client = get_async_httpx_client(llm_provider=httpxSpecialProvider.MCP) response = await client.get(server_url) response.raise_for_status() - verbose_logger.warning( - "MCP OAuth discovery unexpectedly succeeded for %s; server did not challenge", - server_url, + ( + authorization_servers, + resource_scopes, + ) = await self._attempt_well_known_discovery(server_url) + metadata = await self._fetch_authorization_server_metadata( + authorization_servers ) - raise RuntimeError("OAuth discovery must not succeed without a challenge") + if metadata is None and resource_scopes: + return MCPOAuthMetadata(scopes=resource_scopes) + if metadata is not None and resource_scopes: + metadata.scopes = resource_scopes + return metadata except HTTPStatusError as exc: verbose_logger.debug( "MCP OAuth discovery for %s received status error: %s", @@ -1674,6 +1681,9 @@ class MCPServerManager: f"{base}/.well-known/oauth-authorization-server/{path}" ) candidate_urls.append(f"{base}/.well-known/openid-configuration/{path}") + candidate_urls.append( + f"{issuer_url.rstrip('/')}/.well-known/openid-configuration" + ) candidate_urls.append(f"{base}/.well-known/oauth-authorization-server") candidate_urls.append(f"{base}/.well-known/openid-configuration") candidate_urls.append(issuer_url.rstrip("/")) @@ -1713,7 +1723,28 @@ class MCPServerManager: ): return metadata - return None + return self._build_azure_authorization_server_metadata(parsed) + + @staticmethod + def _build_azure_authorization_server_metadata( + parsed_issuer_url: Any, + ) -> Optional[MCPOAuthMetadata]: + path_parts = [ + part for part in (parsed_issuer_url.path or "").split("/") if part + ] + if ( + parsed_issuer_url.netloc != "login.microsoftonline.com" + or len(path_parts) != 2 + or path_parts[1] != "v2.0" + ): + return None + + tenant = path_parts[0] + base = f"{parsed_issuer_url.scheme}://{parsed_issuer_url.netloc}/{tenant}" + return MCPOAuthMetadata( + authorization_url=f"{base}/oauth2/v2.0/authorize", + token_url=f"{base}/oauth2/v2.0/token", + ) @staticmethod def _decrypt_credential_field( diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py index 447c28078e..8614cb3ba7 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py @@ -728,6 +728,136 @@ class TestMCPServerManager: ] assert scopes == ["read", "write"] + @pytest.mark.asyncio + async def test_descovery_metadata_probes_well_known_when_server_does_not_challenge( + self, + ): + manager = MCPServerManager() + + mock_response = MagicMock() + mock_response.raise_for_status = MagicMock() + + mock_client = MagicMock() + mock_client.get = AsyncMock(return_value=mock_response) + + mock_metadata = MCPOAuthMetadata( + scopes=None, + authorization_url="https://login.microsoftonline.com/tenant/oauth2/v2.0/authorize", + token_url="https://login.microsoftonline.com/tenant/oauth2/v2.0/token", + registration_url=None, + ) + + with ( + patch( + "litellm.proxy._experimental.mcp_server.mcp_server_manager.get_async_httpx_client", + return_value=mock_client, + ), + patch.object( + manager, + "_attempt_well_known_discovery", + AsyncMock( + return_value=( + ["https://login.microsoftonline.com/test-tenant-id/v2.0"], + ["api://some-scope/.default"], + ) + ), + ) as mock_well_known, + patch.object( + manager, + "_fetch_authorization_server_metadata", + AsyncMock(return_value=mock_metadata), + ) as mock_fetch_auth, + ): + result = await manager._descovery_metadata("http://localhost:8001/mcp") + + mock_well_known.assert_awaited_once_with("http://localhost:8001/mcp") + mock_fetch_auth.assert_awaited_once_with( + ["https://login.microsoftonline.com/test-tenant-id/v2.0"] + ) + assert result is mock_metadata + assert result.scopes == ["api://some-scope/.default"] + + @pytest.mark.asyncio + async def test_fetch_single_authorization_server_metadata_supports_azure_issuer_path( + self, + ): + manager = MCPServerManager() + issuer = "https://login.microsoftonline.com/test-tenant-id/v2.0" + + def build_response(url: str): + mock_response = MagicMock() + if url == f"{issuer}/.well-known/openid-configuration": + mock_response.json.return_value = { + "authorization_endpoint": "https://login.microsoftonline.com/test-tenant-id/oauth2/v2.0/authorize", + "token_endpoint": "https://login.microsoftonline.com/test-tenant-id/oauth2/v2.0/token", + "scopes_supported": ["api://some-scope/.default"], + } + mock_response.raise_for_status = MagicMock() + else: + request = httpx.Request("GET", url) + response_obj = httpx.Response(status_code=404, request=request) + mock_response.raise_for_status = MagicMock( + side_effect=httpx.HTTPStatusError( + "not found", request=request, response=response_obj + ) + ) + return mock_response + + mock_client = MagicMock() + mock_client.get = AsyncMock(side_effect=build_response) + + with patch( + "litellm.proxy._experimental.mcp_server.mcp_server_manager.get_async_httpx_client", + return_value=mock_client, + ): + result = await manager._fetch_single_authorization_server_metadata(issuer) + + assert result is not None + assert ( + result.authorization_url + == "https://login.microsoftonline.com/test-tenant-id/oauth2/v2.0/authorize" + ) + assert ( + result.token_url + == "https://login.microsoftonline.com/test-tenant-id/oauth2/v2.0/token" + ) + assert result.scopes == ["api://some-scope/.default"] + + @pytest.mark.asyncio + async def test_fetch_single_authorization_server_metadata_derives_azure_metadata( + self, + ): + manager = MCPServerManager() + issuer = "https://login.microsoftonline.com/test-tenant-id/v2.0" + + request = httpx.Request("GET", issuer) + response_obj = httpx.Response(status_code=404, request=request) + mock_response = MagicMock() + mock_response.raise_for_status = MagicMock( + side_effect=httpx.HTTPStatusError( + "not found", request=request, response=response_obj + ) + ) + + mock_client = MagicMock() + mock_client.get = AsyncMock(return_value=mock_response) + + with patch( + "litellm.proxy._experimental.mcp_server.mcp_server_manager.get_async_httpx_client", + return_value=mock_client, + ): + result = await manager._fetch_single_authorization_server_metadata(issuer) + + assert result is not None + assert ( + result.authorization_url + == "https://login.microsoftonline.com/test-tenant-id/oauth2/v2.0/authorize" + ) + assert ( + result.token_url + == "https://login.microsoftonline.com/test-tenant-id/oauth2/v2.0/token" + ) + @pytest.mark.asyncio async def test_descovery_metadata_falls_back_to_origin_when_no_auth_servers(self): manager = MCPServerManager() From 6c5b50135e15ac8db244e6eb63b61688e00b3d37 Mon Sep 17 00:00:00 2001 From: Sameer Kankute Date: Mon, 27 Apr 2026 18:57:22 +0530 Subject: [PATCH 10/57] Fix greptile review --- .../mcp_server/mcp_server_manager.py | 18 +++++++++++++++++- 1 file changed, 17 insertions(+), 1 deletion(-) diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index 6aaf37c91e..edc9c87cda 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -106,6 +106,12 @@ if not _separator_probe.is_valid: SEP_986_URL, ) +_AZURE_ENTRA_HOSTS = { + "login.microsoftonline.com", # Global + "login.microsoftonline.us", # US Government + "login.chinacloudapi.cn", # China +} + def _warn_on_server_name_fields( *, @@ -1495,6 +1501,16 @@ class MCPServerManager: metadata = await self._fetch_authorization_server_metadata( authorization_servers ) + if ( + metadata is None + and not resource_scopes + and authorization_servers + and response.status_code == 200 + ): + verbose_logger.warning( + "MCP OAuth discovery for %s received 200 OK without RFC 9728 challenge and no discoverable authorization metadata.", + server_url, + ) if metadata is None and resource_scopes: return MCPOAuthMetadata(scopes=resource_scopes) if metadata is not None and resource_scopes: @@ -1733,7 +1749,7 @@ class MCPServerManager: part for part in (parsed_issuer_url.path or "").split("/") if part ] if ( - parsed_issuer_url.netloc != "login.microsoftonline.com" + parsed_issuer_url.netloc not in _AZURE_ENTRA_HOSTS or len(path_parts) != 2 or path_parts[1] != "v2.0" ): From adff1c93d0a987b95eec0cf5ab28dad2fc76c324 Mon Sep 17 00:00:00 2001 From: Ryan Crabbe Date: Mon, 27 Apr 2026 14:27:11 -0700 Subject: [PATCH 11/57] refactor(ui): simplify LoggingSettings save flow via React Query callbacks Switch the spend-logs save flow from mutateAsync + try/catch to mutate + callbacks. Errors now surface through a single onError path (no more double toast on failure), and the delete-then-update sequencing runs through onSettled instead of awaited promises. handleFormSubmit is no longer async. Tighten the corresponding test to assert exactly one error toast fires. --- .../LoggingSettings/LoggingSettings.test.tsx | 79 +++++++------------ .../LoggingSettings/LoggingSettings.tsx | 67 ++++++++-------- 2 files changed, 61 insertions(+), 85 deletions(-) diff --git a/ui/litellm-dashboard/src/components/Settings/AdminSettings/LoggingSettings/LoggingSettings.test.tsx b/ui/litellm-dashboard/src/components/Settings/AdminSettings/LoggingSettings/LoggingSettings.test.tsx index 228e899a3a..74413333ec 100644 --- a/ui/litellm-dashboard/src/components/Settings/AdminSettings/LoggingSettings/LoggingSettings.test.tsx +++ b/ui/litellm-dashboard/src/components/Settings/AdminSettings/LoggingSettings/LoggingSettings.test.tsx @@ -36,18 +36,18 @@ const mockNotificationsManager = vi.mocked(NotificationsManager); const mockParseErrorMessage = vi.mocked(parseErrorMessage); describe("LoggingSettings", () => { - const mockMutateAsync = vi.fn(); + const mockMutate = vi.fn(); const mockDeleteField = vi.fn(); const mockRefetch = vi.fn(); beforeEach(() => { vi.clearAllMocks(); mockUseStoreRequestInSpendLogs.mockReturnValue({ - mutateAsync: mockMutateAsync, + mutate: mockMutate, isPending: false, } as any); mockUseDeleteProxyConfigField.mockReturnValue({ - mutateAsync: mockDeleteField, + mutate: mockDeleteField, isPending: false, } as any); mockUseProxyConfig.mockReturnValue({ @@ -94,10 +94,8 @@ describe("LoggingSettings", () => { it("should submit form with store prompts enabled and retention period", async () => { const user = userEvent.setup(); - mockMutateAsync.mockImplementation(async (_params, options) => { - await Promise.resolve(); + mockMutate.mockImplementation((_params, options) => { options?.onSuccess?.(); - return { message: "Success" }; }); renderWithProviders(); @@ -113,7 +111,7 @@ describe("LoggingSettings", () => { await waitFor(() => { expect(mockDeleteField).not.toHaveBeenCalled(); - expect(mockMutateAsync).toHaveBeenCalledWith( + expect(mockMutate).toHaveBeenCalledWith( { store_prompts_in_spend_logs: true, maximum_spend_logs_retention_period: "30d", @@ -125,11 +123,11 @@ describe("LoggingSettings", () => { it("should delete retention period field when left empty on submit", async () => { const user = userEvent.setup(); - mockDeleteField.mockResolvedValue({ message: "Field deleted successfully" }); - mockMutateAsync.mockImplementation(async (_params, options) => { - await Promise.resolve(); + mockDeleteField.mockImplementation((_params, options) => { + options?.onSettled?.(); + }); + mockMutate.mockImplementation((_params, options) => { options?.onSuccess?.(); - return { message: "Success" }; }); renderWithProviders(); @@ -139,7 +137,7 @@ describe("LoggingSettings", () => { await waitFor(() => { expect(mockDeleteField).toHaveBeenCalled(); - expect(mockMutateAsync).toHaveBeenCalledWith( + expect(mockMutate).toHaveBeenCalledWith( { store_prompts_in_spend_logs: false, }, @@ -150,11 +148,11 @@ describe("LoggingSettings", () => { it("should show success notification on successful submission", async () => { const user = userEvent.setup(); - mockDeleteField.mockResolvedValue({ message: "Field deleted successfully" }); - mockMutateAsync.mockImplementation(async (_params, options) => { - await Promise.resolve(); + mockDeleteField.mockImplementation((_params, options) => { + options?.onSettled?.(); + }); + mockMutate.mockImplementation((_params, options) => { options?.onSuccess?.(); - return { message: "Success" }; }); renderWithProviders(); @@ -167,30 +165,11 @@ describe("LoggingSettings", () => { }); }); - it("should show error notification when submission throws", async () => { - const user = userEvent.setup(); - const error = new Error("Network error"); - mockMutateAsync.mockRejectedValue(error); - mockParseErrorMessage.mockReturnValue("Network error"); - - renderWithProviders(); - - const saveButton = screen.getByRole("button", { name: "Save Settings" }); - await user.click(saveButton); - - await waitFor(() => { - expect(mockNotificationsManager.fromBackend).toHaveBeenCalledWith( - "Failed to save spend logs settings: Network error", - ); - }); - }); - - it("should show error notification via onError callback", async () => { + it("should show a single error notification via onError callback", async () => { const user = userEvent.setup(); const error = new Error("Backend error"); - mockMutateAsync.mockImplementation((_params, options) => { + mockMutate.mockImplementation((_params, options) => { options?.onError?.(error); - return Promise.reject(error); }); mockParseErrorMessage.mockReturnValue("Backend error"); @@ -204,11 +183,12 @@ describe("LoggingSettings", () => { "Failed to save spend logs settings: Backend error", ); }); + expect(mockNotificationsManager.fromBackend).toHaveBeenCalledTimes(1); }); it("should show loading state on save button when update pending", () => { mockUseStoreRequestInSpendLogs.mockReturnValue({ - mutateAsync: mockMutateAsync, + mutate: mockMutate, isPending: true, } as any); @@ -221,7 +201,7 @@ describe("LoggingSettings", () => { it("should show loading state on save button when delete pending", () => { mockUseDeleteProxyConfigField.mockReturnValue({ - mutateAsync: mockDeleteField, + mutate: mockDeleteField, isPending: true, } as any); @@ -297,11 +277,12 @@ describe("LoggingSettings", () => { it("should continue with update even if deleteField fails", async () => { const user = userEvent.setup(); const deleteError = new Error("Field does not exist"); - mockDeleteField.mockRejectedValue(deleteError); - mockMutateAsync.mockImplementation(async (_params, options) => { - await Promise.resolve(); + mockDeleteField.mockImplementation((_params, options) => { + options?.onError?.(deleteError); + options?.onSettled?.(); + }); + mockMutate.mockImplementation((_params, options) => { options?.onSuccess?.(); - return { message: "Success" }; }); renderWithProviders(); @@ -311,7 +292,7 @@ describe("LoggingSettings", () => { await waitFor(() => { expect(mockDeleteField).toHaveBeenCalled(); - expect(mockMutateAsync).toHaveBeenCalledWith( + expect(mockMutate).toHaveBeenCalledWith( { store_prompts_in_spend_logs: false, }, @@ -323,11 +304,11 @@ describe("LoggingSettings", () => { it("should submit with only store prompts enabled when retention is empty", async () => { const user = userEvent.setup(); - mockDeleteField.mockResolvedValue({ message: "Field deleted successfully" }); - mockMutateAsync.mockImplementation(async (_params, options) => { - await Promise.resolve(); + mockDeleteField.mockImplementation((_params, options) => { + options?.onSettled?.(); + }); + mockMutate.mockImplementation((_params, options) => { options?.onSuccess?.(); - return { message: "Success" }; }); renderWithProviders(); @@ -340,7 +321,7 @@ describe("LoggingSettings", () => { await waitFor(() => { expect(mockDeleteField).toHaveBeenCalled(); - expect(mockMutateAsync).toHaveBeenCalledWith( + expect(mockMutate).toHaveBeenCalledWith( { store_prompts_in_spend_logs: true, }, diff --git a/ui/litellm-dashboard/src/components/Settings/AdminSettings/LoggingSettings/LoggingSettings.tsx b/ui/litellm-dashboard/src/components/Settings/AdminSettings/LoggingSettings/LoggingSettings.tsx index c3aaebd3bf..456240ec5d 100644 --- a/ui/litellm-dashboard/src/components/Settings/AdminSettings/LoggingSettings/LoggingSettings.tsx +++ b/ui/litellm-dashboard/src/components/Settings/AdminSettings/LoggingSettings/LoggingSettings.tsx @@ -18,8 +18,8 @@ import React, { useMemo } from "react"; const LoggingSettings: React.FC = () => { const [form] = Form.useForm(); - const { mutateAsync, isPending } = useStoreRequestInSpendLogs(); - const { mutateAsync: deleteField, isPending: isDeletingField } = useDeleteProxyConfigField(); + const { mutate, isPending } = useStoreRequestInSpendLogs(); + const { mutate: deleteField, isPending: isDeletingField } = useDeleteProxyConfigField(); const { data: proxyConfigData, isLoading: isLoadingConfig } = useProxyConfig(ConfigType.GENERAL_SETTINGS); const storePromptsValue = Form.useWatch("store_prompts_in_spend_logs", form); @@ -42,44 +42,39 @@ const LoggingSettings: React.FC = () => { }; }, [proxyConfigData]); - const handleFormSubmit = async (formValues: StoreRequestInSpendLogsParams) => { - try { - const retentionPeriodValue = formValues.maximum_spend_logs_retention_period; - const shouldDeleteRetentionPeriod = - !retentionPeriodValue || - (typeof retentionPeriodValue === "string" && retentionPeriodValue.trim() === ""); + const handleFormSubmit = (formValues: StoreRequestInSpendLogsParams) => { + const retentionPeriodValue = formValues.maximum_spend_logs_retention_period; + const hasRetentionPeriod = + typeof retentionPeriodValue === "string" && retentionPeriodValue.trim() !== ""; - if (shouldDeleteRetentionPeriod) { - try { - await deleteField({ - config_type: ConfigType.GENERAL_SETTINGS, - field_name: GeneralSettingsFieldName.MAXIMUM_SPEND_LOGS_RETENTION_PERIOD, - }); - } catch (deleteError) { - console.warn("Failed to delete retention period field (may not exist):", deleteError); - } - } + const updateParams: StoreRequestInSpendLogsParams = { + store_prompts_in_spend_logs: formValues.store_prompts_in_spend_logs, + ...(hasRetentionPeriod && { maximum_spend_logs_retention_period: retentionPeriodValue }), + }; - const updateParams: StoreRequestInSpendLogsParams = { - store_prompts_in_spend_logs: formValues.store_prompts_in_spend_logs, - ...(retentionPeriodValue && - typeof retentionPeriodValue === "string" && - retentionPeriodValue.trim() !== "" && { - maximum_spend_logs_retention_period: retentionPeriodValue, - }), - }; - - await mutateAsync(updateParams, { - onSuccess: () => { - NotificationsManager.success("Spend logs settings updated successfully"); - }, - onError: (error) => { - NotificationsManager.fromBackend("Failed to save spend logs settings: " + parseErrorMessage(error)); - }, + const submitUpdate = () => + mutate(updateParams, { + onSuccess: () => NotificationsManager.success("Spend logs settings updated successfully"), + onError: (error) => + NotificationsManager.fromBackend("Failed to save spend logs settings: " + parseErrorMessage(error)), }); - } catch (error) { - NotificationsManager.fromBackend("Failed to save spend logs settings: " + parseErrorMessage(error)); + + if (hasRetentionPeriod) { + submitUpdate(); + return; } + + deleteField( + { + config_type: ConfigType.GENERAL_SETTINGS, + field_name: GeneralSettingsFieldName.MAXIMUM_SPEND_LOGS_RETENTION_PERIOD, + }, + { + onError: (deleteError) => + console.warn("Failed to delete retention period field (may not exist):", deleteError), + onSettled: submitUpdate, + }, + ); }; return ( From 503c3921c8ccda34862f955357c7bb8059c5e88c Mon Sep 17 00:00:00 2001 From: Liam McDonald Date: Mon, 27 Apr 2026 15:33:59 -0700 Subject: [PATCH 12/57] Fix gpt-5.5-pro pricing --- ...odel_prices_and_context_window_backup.json | 64 +++++++++---------- model_prices_and_context_window.json | 64 +++++++++---------- 2 files changed, 64 insertions(+), 64 deletions(-) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 5cccd5f00a..8511d785fb 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -4735,17 +4735,17 @@ "supports_web_search": true }, "azure/gpt-5.5-pro": { - "cache_read_input_token_cost": 6e-06, - "cache_read_input_token_cost_above_272k_tokens": 1.2e-05, - "input_cost_per_token": 6e-05, - "input_cost_per_token_above_272k_tokens": 0.00012, + "cache_read_input_token_cost": 3e-06, + "cache_read_input_token_cost_above_272k_tokens": 6e-06, + "input_cost_per_token": 3e-05, + "input_cost_per_token_above_272k_tokens": 6e-05, "litellm_provider": "azure", "max_input_tokens": 1050000, "max_output_tokens": 128000, "max_tokens": 128000, "mode": "responses", - "output_cost_per_token": 0.00036, - "output_cost_per_token_above_272k_tokens": 0.00054, + "output_cost_per_token": 0.00018, + "output_cost_per_token_above_272k_tokens": 0.00027, "supported_endpoints": [ "/v1/batch", "/v1/responses" @@ -4774,17 +4774,17 @@ "supports_low_reasoning_effort": false }, "azure/gpt-5.5-pro-2026-04-23": { - "cache_read_input_token_cost": 6e-06, - "cache_read_input_token_cost_above_272k_tokens": 1.2e-05, - "input_cost_per_token": 6e-05, - "input_cost_per_token_above_272k_tokens": 0.00012, + "cache_read_input_token_cost": 3e-06, + "cache_read_input_token_cost_above_272k_tokens": 6e-06, + "input_cost_per_token": 3e-05, + "input_cost_per_token_above_272k_tokens": 6e-05, "litellm_provider": "azure", "max_input_tokens": 1050000, "max_output_tokens": 128000, "max_tokens": 128000, "mode": "responses", - "output_cost_per_token": 0.00036, - "output_cost_per_token_above_272k_tokens": 0.00054, + "output_cost_per_token": 0.00018, + "output_cost_per_token_above_272k_tokens": 0.00027, "supported_endpoints": [ "/v1/batch", "/v1/responses" @@ -19898,21 +19898,21 @@ "supports_minimal_reasoning_effort": true }, "gpt-5.5-pro": { - "cache_read_input_token_cost": 6e-06, - "cache_read_input_token_cost_above_272k_tokens": 1.2e-05, - "input_cost_per_token": 6e-05, - "input_cost_per_token_above_272k_tokens": 0.00012, - "input_cost_per_token_flex": 3e-05, - "input_cost_per_token_batches": 3e-05, + "cache_read_input_token_cost": 3e-06, + "cache_read_input_token_cost_above_272k_tokens": 6e-06, + "input_cost_per_token": 3e-05, + "input_cost_per_token_above_272k_tokens": 6e-05, + "input_cost_per_token_flex": 1.5e-05, + "input_cost_per_token_batches": 1.5e-05, "litellm_provider": "openai", "max_input_tokens": 1050000, "max_output_tokens": 128000, "max_tokens": 128000, "mode": "responses", - "output_cost_per_token": 0.00036, - "output_cost_per_token_above_272k_tokens": 0.00054, - "output_cost_per_token_flex": 0.00018, - "output_cost_per_token_batches": 0.00018, + "output_cost_per_token": 0.00018, + "output_cost_per_token_above_272k_tokens": 0.00027, + "output_cost_per_token_flex": 9e-05, + "output_cost_per_token_batches": 9e-05, "supported_endpoints": [ "/v1/responses", "/v1/batch" @@ -19941,21 +19941,21 @@ "supports_minimal_reasoning_effort": true }, "gpt-5.5-pro-2026-04-23": { - "cache_read_input_token_cost": 6e-06, - "cache_read_input_token_cost_above_272k_tokens": 1.2e-05, - "input_cost_per_token": 6e-05, - "input_cost_per_token_above_272k_tokens": 0.00012, - "input_cost_per_token_flex": 3e-05, - "input_cost_per_token_batches": 3e-05, + "cache_read_input_token_cost": 3e-06, + "cache_read_input_token_cost_above_272k_tokens": 6e-06, + "input_cost_per_token": 3e-05, + "input_cost_per_token_above_272k_tokens": 6e-05, + "input_cost_per_token_flex": 1.5e-05, + "input_cost_per_token_batches": 1.5e-05, "litellm_provider": "openai", "max_input_tokens": 1050000, "max_output_tokens": 128000, "max_tokens": 128000, "mode": "responses", - "output_cost_per_token": 0.00036, - "output_cost_per_token_above_272k_tokens": 0.00054, - "output_cost_per_token_flex": 0.00018, - "output_cost_per_token_batches": 0.00018, + "output_cost_per_token": 0.00018, + "output_cost_per_token_above_272k_tokens": 0.00027, + "output_cost_per_token_flex": 9e-05, + "output_cost_per_token_batches": 9e-05, "supported_endpoints": [ "/v1/responses", "/v1/batch" diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 4d8f3a984f..114883f530 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -4749,17 +4749,17 @@ "supports_web_search": true }, "azure/gpt-5.5-pro": { - "cache_read_input_token_cost": 6e-06, - "cache_read_input_token_cost_above_272k_tokens": 1.2e-05, - "input_cost_per_token": 6e-05, - "input_cost_per_token_above_272k_tokens": 0.00012, + "cache_read_input_token_cost": 3e-06, + "cache_read_input_token_cost_above_272k_tokens": 6e-06, + "input_cost_per_token": 3e-05, + "input_cost_per_token_above_272k_tokens": 6e-05, "litellm_provider": "azure", "max_input_tokens": 1050000, "max_output_tokens": 128000, "max_tokens": 128000, "mode": "responses", - "output_cost_per_token": 0.00036, - "output_cost_per_token_above_272k_tokens": 0.00054, + "output_cost_per_token": 0.00018, + "output_cost_per_token_above_272k_tokens": 0.00027, "supported_endpoints": [ "/v1/batch", "/v1/responses" @@ -4788,17 +4788,17 @@ "supports_low_reasoning_effort": false }, "azure/gpt-5.5-pro-2026-04-23": { - "cache_read_input_token_cost": 6e-06, - "cache_read_input_token_cost_above_272k_tokens": 1.2e-05, - "input_cost_per_token": 6e-05, - "input_cost_per_token_above_272k_tokens": 0.00012, + "cache_read_input_token_cost": 3e-06, + "cache_read_input_token_cost_above_272k_tokens": 6e-06, + "input_cost_per_token": 3e-05, + "input_cost_per_token_above_272k_tokens": 6e-05, "litellm_provider": "azure", "max_input_tokens": 1050000, "max_output_tokens": 128000, "max_tokens": 128000, "mode": "responses", - "output_cost_per_token": 0.00036, - "output_cost_per_token_above_272k_tokens": 0.00054, + "output_cost_per_token": 0.00018, + "output_cost_per_token_above_272k_tokens": 0.00027, "supported_endpoints": [ "/v1/batch", "/v1/responses" @@ -19912,21 +19912,21 @@ "supports_minimal_reasoning_effort": true }, "gpt-5.5-pro": { - "cache_read_input_token_cost": 6e-06, - "cache_read_input_token_cost_above_272k_tokens": 1.2e-05, - "input_cost_per_token": 6e-05, - "input_cost_per_token_above_272k_tokens": 0.00012, - "input_cost_per_token_flex": 3e-05, - "input_cost_per_token_batches": 3e-05, + "cache_read_input_token_cost": 3e-06, + "cache_read_input_token_cost_above_272k_tokens": 6e-06, + "input_cost_per_token": 3e-05, + "input_cost_per_token_above_272k_tokens": 6e-05, + "input_cost_per_token_flex": 1.5e-05, + "input_cost_per_token_batches": 1.5e-05, "litellm_provider": "openai", "max_input_tokens": 1050000, "max_output_tokens": 128000, "max_tokens": 128000, "mode": "responses", - "output_cost_per_token": 0.00036, - "output_cost_per_token_above_272k_tokens": 0.00054, - "output_cost_per_token_flex": 0.00018, - "output_cost_per_token_batches": 0.00018, + "output_cost_per_token": 0.00018, + "output_cost_per_token_above_272k_tokens": 0.00027, + "output_cost_per_token_flex": 9e-05, + "output_cost_per_token_batches": 9e-05, "supported_endpoints": [ "/v1/responses", "/v1/batch" @@ -19955,21 +19955,21 @@ "supports_minimal_reasoning_effort": true }, "gpt-5.5-pro-2026-04-23": { - "cache_read_input_token_cost": 6e-06, - "cache_read_input_token_cost_above_272k_tokens": 1.2e-05, - "input_cost_per_token": 6e-05, - "input_cost_per_token_above_272k_tokens": 0.00012, - "input_cost_per_token_flex": 3e-05, - "input_cost_per_token_batches": 3e-05, + "cache_read_input_token_cost": 3e-06, + "cache_read_input_token_cost_above_272k_tokens": 6e-06, + "input_cost_per_token": 3e-05, + "input_cost_per_token_above_272k_tokens": 6e-05, + "input_cost_per_token_flex": 1.5e-05, + "input_cost_per_token_batches": 1.5e-05, "litellm_provider": "openai", "max_input_tokens": 1050000, "max_output_tokens": 128000, "max_tokens": 128000, "mode": "responses", - "output_cost_per_token": 0.00036, - "output_cost_per_token_above_272k_tokens": 0.00054, - "output_cost_per_token_flex": 0.00018, - "output_cost_per_token_batches": 0.00018, + "output_cost_per_token": 0.00018, + "output_cost_per_token_above_272k_tokens": 0.00027, + "output_cost_per_token_flex": 9e-05, + "output_cost_per_token_batches": 9e-05, "supported_endpoints": [ "/v1/responses", "/v1/batch" From 321575a29d1bacc4149d7ad9c7fb584c947c4047 Mon Sep 17 00:00:00 2001 From: Liam McDonald Date: Mon, 27 Apr 2026 15:37:51 -0700 Subject: [PATCH 13/57] Fix gpt-5.5-pro pricing tests --- .../llm_cost_calc/test_llm_cost_calc_utils.py | 21 +++++++++---------- 1 file changed, 10 insertions(+), 11 deletions(-) diff --git a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py index 3016762043..2771aae6b9 100644 --- a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py +++ b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py @@ -369,7 +369,7 @@ def test_generic_cost_per_token_gpt55(): def test_generic_cost_per_token_gpt55_pro(): - """gpt-5.5-pro: responses-only model — $60/1M input, $360/1M output, $6/1M cached input.""" + """gpt-5.5-pro: responses-only model — $30/1M input, $180/1M output, $3/1M cached input.""" model = "gpt-5.5-pro" custom_llm_provider = "openai" os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" @@ -378,18 +378,17 @@ def test_generic_cost_per_token_gpt55_pro(): model_cost_map = litellm.model_cost[model] # Sanity-check the map values match OpenAI's published pricing. - assert model_cost_map["input_cost_per_token"] == 6e-5 - assert model_cost_map["output_cost_per_token"] == 3.6e-4 - assert model_cost_map["cache_read_input_token_cost"] == 6e-6 + assert model_cost_map["input_cost_per_token"] == 3e-5 + assert model_cost_map["output_cost_per_token"] == 1.8e-4 + assert model_cost_map["cache_read_input_token_cost"] == 3e-6 assert model_cost_map["litellm_provider"] == "openai" # gpt-5.5-pro is a responses-only model (no /v1/chat/completions endpoint). assert model_cost_map["mode"] == "responses" assert "/v1/chat/completions" not in model_cost_map["supported_endpoints"] assert "/v1/responses" in model_cost_map["supported_endpoints"] - # Inherits GPT-5.4-pro's long-context window + tiered pricing (scaled 2x). - assert model_cost_map["max_input_tokens"] == 1050000 - assert model_cost_map["input_cost_per_token_above_272k_tokens"] == 1.2e-4 - assert model_cost_map["output_cost_per_token_above_272k_tokens"] == 5.4e-4 + # Inherits GPT-5.4-pro's long-context window + tiered pricing. + assert model_cost_map["input_cost_per_token_above_272k_tokens"] == 6e-5 + assert model_cost_map["output_cost_per_token_above_272k_tokens"] == 2.7e-4 prompt_tokens = 1000 completion_tokens = 500 @@ -454,8 +453,8 @@ def test_gpt55_dated_variants_match_base_reasoning_effort_capabilities( [ ("azure/gpt-5.5", "chat", 5e-6, 3e-5, 5e-7), ("azure/gpt-5.5-2026-04-23", "chat", 5e-6, 3e-5, 5e-7), - ("azure/gpt-5.5-pro", "responses", 6e-5, 3.6e-4, 6e-6), - ("azure/gpt-5.5-pro-2026-04-23", "responses", 6e-5, 3.6e-4, 6e-6), + ("azure/gpt-5.5-pro", "responses", 3e-5, 1.8e-4, 3e-6), + ("azure/gpt-5.5-pro-2026-04-23", "responses", 3e-5, 1.8e-4, 3e-6), ], ) def test_azure_gpt55_entries_present_with_correct_pricing( @@ -464,7 +463,7 @@ def test_azure_gpt55_entries_present_with_correct_pricing( """Day-0 Azure entries for GPT-5.5 mirror the OpenAI pricing structure. Pricing parity with openai/gpt-5.5* (verified against OpenAI's pricing page - on 2026-04-24): $5/$30 input/output per 1M for chat, $60/$360 for pro. + on 2026-04-24): $5/$30 input/output per 1M for chat, $30/$180 for pro. Cache discount is 10% of input. """ os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" From 325c74548df644dde1450fc65eb8f3cae63e796b Mon Sep 17 00:00:00 2001 From: Ryan Crabbe Date: Mon, 27 Apr 2026 15:48:27 -0700 Subject: [PATCH 14/57] refactor(ui): invalidate proxyConfig query after spend-logs mutations Previously, useStoreRequestInSpendLogs and useDeleteProxyConfigField did not refresh the proxyConfig cache on success, so the Logging Settings form continued to render the pre-save values until React Query refetched on its own. Wire both hooks to invalidate proxyConfigKeys on success so any active observer (currently the Logging Settings page) repulls fresh data. Export proxyConfigKeys for cross-hook reuse. --- .../hooks/proxyConfig/useProxyConfig.test.ts | 23 +++++++++++++++++++ .../hooks/proxyConfig/useProxyConfig.ts | 8 +++++-- .../useStoreRequestInSpendLogs.ts | 7 +++++- 3 files changed, 35 insertions(+), 3 deletions(-) diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/proxyConfig/useProxyConfig.test.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/proxyConfig/useProxyConfig.test.ts index a8ce55d274..4ba80df5c6 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/hooks/proxyConfig/useProxyConfig.test.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/proxyConfig/useProxyConfig.test.ts @@ -7,6 +7,7 @@ import { useDeleteProxyConfigField, getProxyConfigCall, deleteProxyConfigFieldCall, + proxyConfigKeys, ConfigType, GeneralSettingsFieldName, type ProxyConfigResponse, @@ -426,6 +427,28 @@ describe("useDeleteProxyConfigField", () => { expect(result.current.error).toBeDefined(); }); + + it("should invalidate proxyConfig queries after a successful delete", async () => { + (fetchSpy as any).mockResolvedValue({ + ok: true, + json: async () => mockDeleteResponse, + }); + + const invalidateSpy = vi.spyOn(queryClient, "invalidateQueries"); + + const { result } = renderHook(() => useDeleteProxyConfigField(), { wrapper }); + + result.current.mutate({ + config_type: ConfigType.GENERAL_SETTINGS, + field_name: GeneralSettingsFieldName.MAXIMUM_SPEND_LOGS_RETENTION_PERIOD, + }); + + await waitFor(() => { + expect(result.current.isSuccess).toBe(true); + }); + + expect(invalidateSpy).toHaveBeenCalledWith({ queryKey: proxyConfigKeys.all }); + }); }); describe("getProxyConfigCall", () => { diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/proxyConfig/useProxyConfig.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/proxyConfig/useProxyConfig.ts index b823ce4ffd..485c7cc1f9 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/hooks/proxyConfig/useProxyConfig.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/proxyConfig/useProxyConfig.ts @@ -1,4 +1,4 @@ -import { useQuery, useMutation, UseMutationResult } from "@tanstack/react-query"; +import { useQuery, useMutation, UseMutationResult, useQueryClient } from "@tanstack/react-query"; import { createQueryKeys } from "../common/queryKeysFactory"; import useAuthorized from "../useAuthorized"; import { proxyBaseUrl, getGlobalLitellmHeaderName, deriveErrorMessage, handleError } from "@/components/networking"; @@ -101,7 +101,7 @@ export const getProxyConfigCall = async (accessToken: string, configType: Config } }; -const proxyConfigKeys = createQueryKeys("proxyConfig"); +export const proxyConfigKeys = createQueryKeys("proxyConfig"); /** * Network call function to delete a proxy config field @@ -168,6 +168,7 @@ export const useDeleteProxyConfigField = (): UseMutationResult< DeleteProxyConfigFieldRequest > => { const { accessToken } = useAuthorized(); + const queryClient = useQueryClient(); return useMutation({ mutationFn: async (request: DeleteProxyConfigFieldRequest) => { @@ -176,5 +177,8 @@ export const useDeleteProxyConfigField = (): UseMutationResult< } return await deleteProxyConfigFieldCall(accessToken, request); }, + onSuccess: () => { + queryClient.invalidateQueries({ queryKey: proxyConfigKeys.all }); + }, }); }; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/storeRequestInSpendLogs/useStoreRequestInSpendLogs.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/storeRequestInSpendLogs/useStoreRequestInSpendLogs.ts index 9c6211c308..67b52997a0 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/hooks/storeRequestInSpendLogs/useStoreRequestInSpendLogs.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/storeRequestInSpendLogs/useStoreRequestInSpendLogs.ts @@ -1,6 +1,7 @@ -import { useMutation, UseMutationResult } from "@tanstack/react-query"; +import { useMutation, UseMutationResult, useQueryClient } from "@tanstack/react-query"; import { getProxyBaseUrl, getGlobalLitellmHeaderName } from "@/components/networking"; import useAuthorized from "../useAuthorized"; +import { proxyConfigKeys } from "../proxyConfig/useProxyConfig"; export interface StoreRequestInSpendLogsParams { store_prompts_in_spend_logs: boolean; @@ -51,6 +52,7 @@ export const useStoreRequestInSpendLogs = (): UseMutationResult< StoreRequestInSpendLogsParams > => { const { accessToken } = useAuthorized(); + const queryClient = useQueryClient(); return useMutation({ mutationFn: async (params: StoreRequestInSpendLogsParams) => { @@ -59,5 +61,8 @@ export const useStoreRequestInSpendLogs = (): UseMutationResult< } return await performStoreRequestInSpendLogs(accessToken, params); }, + onSuccess: () => { + queryClient.invalidateQueries({ queryKey: proxyConfigKeys.all }); + }, }); }; From ea0ce944cd37552c6a37cb148c8a9b5c2a5937a5 Mon Sep 17 00:00:00 2001 From: Liam McDonald Date: Mon, 27 Apr 2026 15:58:46 -0700 Subject: [PATCH 15/57] correct gpt-5.5-pro token pricing to match OpenAI --- .../litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py index 2771aae6b9..77284a64cf 100644 --- a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py +++ b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py @@ -387,6 +387,7 @@ def test_generic_cost_per_token_gpt55_pro(): assert "/v1/chat/completions" not in model_cost_map["supported_endpoints"] assert "/v1/responses" in model_cost_map["supported_endpoints"] # Inherits GPT-5.4-pro's long-context window + tiered pricing. + assert model_cost_map["max_input_tokens"] == 1050000 assert model_cost_map["input_cost_per_token_above_272k_tokens"] == 6e-5 assert model_cost_map["output_cost_per_token_above_272k_tokens"] == 2.7e-4 From 7f48284decb155c257445e27f77f98266b74737c Mon Sep 17 00:00:00 2001 From: Ryan Crabbe Date: Mon, 27 Apr 2026 16:28:48 -0700 Subject: [PATCH 16/57] test(ui): reset mocks between LoggingSettings tests to prevent bleed-through vi.clearAllMocks does not reset mockImplementation, so the error-notification test was inadvertently relying on a deleteField stub set up in earlier tests and would time out when run in isolation. --- .../AdminSettings/LoggingSettings/LoggingSettings.test.tsx | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/ui/litellm-dashboard/src/components/Settings/AdminSettings/LoggingSettings/LoggingSettings.test.tsx b/ui/litellm-dashboard/src/components/Settings/AdminSettings/LoggingSettings/LoggingSettings.test.tsx index 74413333ec..1adf8039e6 100644 --- a/ui/litellm-dashboard/src/components/Settings/AdminSettings/LoggingSettings/LoggingSettings.test.tsx +++ b/ui/litellm-dashboard/src/components/Settings/AdminSettings/LoggingSettings/LoggingSettings.test.tsx @@ -41,7 +41,7 @@ describe("LoggingSettings", () => { const mockRefetch = vi.fn(); beforeEach(() => { - vi.clearAllMocks(); + vi.resetAllMocks(); mockUseStoreRequestInSpendLogs.mockReturnValue({ mutate: mockMutate, isPending: false, @@ -168,6 +168,9 @@ describe("LoggingSettings", () => { it("should show a single error notification via onError callback", async () => { const user = userEvent.setup(); const error = new Error("Backend error"); + mockDeleteField.mockImplementation((_params, options) => { + options?.onSettled?.(); + }); mockMutate.mockImplementation((_params, options) => { options?.onError?.(error); }); From 7d69621b592321e04b6559b5853dd4b6e6d3ff36 Mon Sep 17 00:00:00 2001 From: Mateo Wang <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 27 Apr 2026 17:31:03 -0700 Subject: [PATCH 17/57] docs: update pull_request_template to add Linear ticket mentioning We are replacing daily updates with Linear tickets instead of GitHub PRs directly so linking the two is essential --- .github/pull_request_template.md | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/.github/pull_request_template.md b/.github/pull_request_template.md index 210f232b17..f9ce9e5dcb 100644 --- a/.github/pull_request_template.md +++ b/.github/pull_request_template.md @@ -2,6 +2,10 @@ +## Linear ticket + + + ## Pre-Submission checklist **Please complete all items before asking a LiteLLM maintainer to review your PR** From 3ca985451e5a416b33595a0031187cc7aa629fd0 Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Mon, 27 Apr 2026 23:37:09 -0700 Subject: [PATCH 18/57] fix(vertex): preserve items on array branches inside anyOf with null convert_anyof_null_to_nullable was stripping the items field from array branches inside anyOf when a sibling null branch was present, leaving {"type": "array"} without items. Vertex requires items whenever type == "array" (even inside anyOf) and rejects the call with INVALID_ARGUMENT. Leave the (possibly empty) items in place so the downstream process_items step can convert {} to {"type": "object"}, which is what Vertex wants. Also: - Update test_build_vertex_schema expected output, which was codifying the broken shape. - Convert test_gemini_tool_calling_not_working to a hermetic mock test that asserts the request body sent to Vertex includes items inside the callbacks anyOf array branch. The previous form made a real network call and was flaky in CI. Co-Authored-By: Claude Opus 4.7 (1M context) --- litellm/llms/vertex_ai/common_utils.py | 10 +-- .../test_amazing_vertex_completion.py | 70 +++++++++++++++++-- .../vertex_ai/test_vertex_ai_common_utils.py | 6 +- 3 files changed, 74 insertions(+), 12 deletions(-) diff --git a/litellm/llms/vertex_ai/common_utils.py b/litellm/llms/vertex_ai/common_utils.py index ccd4d4f293..9b23520dcd 100644 --- a/litellm/llms/vertex_ai/common_utils.py +++ b/litellm/llms/vertex_ai/common_utils.py @@ -710,14 +710,10 @@ def convert_anyof_null_to_nullable(schema, depth=0): if contains_null: # set all types to nullable following guidance found here: https://cloud.google.com/vertex-ai/generative-ai/docs/samples/generativeaionvertexai-gemini-controlled-generation-response-schema-3#generativeaionvertexai_gemini_controlled_generation_response_schema_3-python + # Empty `items: {}` on array branches is left in place; downstream + # process_items() converts it to {"type": "object"}, which Vertex + # requires whenever type == "array" (even inside anyOf). for atype in anyof: - # Remove items field if type is array and items is empty - if ( - atype.get("type") == "array" - and "items" in atype - and not atype["items"] - ): - atype.pop("items") atype["nullable"] = True properties = schema.get("properties", None) diff --git a/tests/local_testing/test_amazing_vertex_completion.py b/tests/local_testing/test_amazing_vertex_completion.py index 9070a9feab..3b4ecb82b1 100644 --- a/tests/local_testing/test_amazing_vertex_completion.py +++ b/tests/local_testing/test_amazing_vertex_completion.py @@ -3569,8 +3569,14 @@ def test_gemini_tool_calling_working_demo(): def test_gemini_tool_calling_not_working(): - load_vertex_ai_credentials() - litellm._turn_on_debug() + """ + Regression test: tool params with anyOf containing both an empty-items + array branch and a null branch must serialize with items present on the + array branch (Vertex rejects array types missing `items`). + """ + from litellm.llms.custom_httpx.http_handler import HTTPHandler + from litellm.llms.vertex_ai.vertex_llm_base import VertexBase + args = { "messages": [ { @@ -3637,8 +3643,64 @@ def test_gemini_tool_calling_not_working(): ], "vertex_location": "global", } - response = completion(model="vertex_ai/gemini-3-flash-preview", **args) - print(response) + + client = HTTPHandler() + mock_response = MagicMock() + mock_response.status_code = 200 + mock_response.headers = {"Content-Type": "application/json"} + mock_response.json.return_value = { + "candidates": [ + { + "content": { + "role": "model", + "parts": [{"text": "Hello!"}], + }, + "finishReason": "STOP", + } + ], + "usageMetadata": { + "promptTokenCount": 10, + "candidatesTokenCount": 5, + "totalTokenCount": 15, + }, + } + + with ( + patch.object(client, "post", return_value=mock_response) as mock_post, + patch.object( + VertexBase, + "_ensure_access_token", + return_value=("fake-token", "fake-project"), + ), + ): + completion( + model="vertex_ai/gemini-3-flash-preview", + client=client, + **args, + ) + + sent_body = mock_post.call_args.kwargs.get( + "json" + ) or mock_post.call_args.kwargs.get("data") + assert sent_body is not None, "expected request body to be sent" + if isinstance(sent_body, str): + sent_body = json.loads(sent_body) + + function_decl = sent_body["tools"][0]["function_declarations"][0] + callbacks_schema = function_decl["parameters"]["properties"]["config"][ + "properties" + ]["callbacks"] + array_branches = [ + branch + for branch in callbacks_schema["anyOf"] + if branch.get("type", "").lower() == "array" + ] + assert array_branches, "expected an array branch in callbacks anyOf" + for branch in array_branches: + assert "items" in branch and branch["items"], ( + f"array branch in callbacks.anyOf must include non-empty items " + f"(Vertex rejects array types missing items). Got: {branch}" + ) def test_vertex_ai_llama_tool_calling(): diff --git a/tests/test_litellm/llms/vertex_ai/test_vertex_ai_common_utils.py b/tests/test_litellm/llms/vertex_ai/test_vertex_ai_common_utils.py index ef93375c3c..dc3be7114f 100644 --- a/tests/test_litellm/llms/vertex_ai/test_vertex_ai_common_utils.py +++ b/tests/test_litellm/llms/vertex_ai/test_vertex_ai_common_utils.py @@ -225,7 +225,11 @@ def test_build_vertex_schema(): "metadata": {"type": "object"}, "callbacks": { "anyOf": [ - {"type": "array", "nullable": True}, + { + "type": "array", + "items": {"type": "object"}, + "nullable": True, + }, {"type": "object", "nullable": True}, ] }, From cfe4bc678e612bcc37831b69582e886bf05e03da Mon Sep 17 00:00:00 2001 From: Sameer Kankute Date: Tue, 28 Apr 2026 15:37:23 +0530 Subject: [PATCH 19/57] feat(vector-stores): support provider-specific Bedrock retrieval config Route vector store search `extra_body` into provider transformers and handle Bedrock `retrievalConfiguration` explicitly so only intended provider-specific fields are forwarded. Made-with: Cursor --- .../azure_ai/vector_stores/transformation.py | 1 + .../base_llm/vector_store/transformation.py | 3 ++ .../bedrock/vector_stores/transformation.py | 23 +++++++----- litellm/llms/custom_httpx/llm_http_handler.py | 3 ++ .../gemini/vector_stores/transformation.py | 1 + .../milvus/vector_stores/transformation.py | 1 + .../openai/vector_stores/transformation.py | 1 + .../pg_vector/vector_stores/transformation.py | 2 ++ .../ragflow/vector_stores/transformation.py | 1 + .../vector_stores/transformation.py | 2 ++ .../vector_stores/rag_api/transformation.py | 1 + .../search_api/transformation.py | 1 + ...est_bedrock_vector_store_transformation.py | 35 +++++++++++++++++++ 13 files changed, 66 insertions(+), 9 deletions(-) diff --git a/litellm/llms/azure_ai/vector_stores/transformation.py b/litellm/llms/azure_ai/vector_stores/transformation.py index b62acb6516..d2c8206ca9 100644 --- a/litellm/llms/azure_ai/vector_stores/transformation.py +++ b/litellm/llms/azure_ai/vector_stores/transformation.py @@ -92,6 +92,7 @@ class AzureAIVectorStoreConfig(BaseVectorStoreConfig, BaseAzureLLM): vector_store_id: str, query: Union[str, List[str]], vector_store_search_optional_params: VectorStoreSearchOptionalRequestParams, + extra_body: Optional[Dict[str, Any]], api_base: str, litellm_logging_obj: LiteLLMLoggingObj, litellm_params: dict, diff --git a/litellm/llms/base_llm/vector_store/transformation.py b/litellm/llms/base_llm/vector_store/transformation.py index 5fbf0a4b19..49d2f72db7 100644 --- a/litellm/llms/base_llm/vector_store/transformation.py +++ b/litellm/llms/base_llm/vector_store/transformation.py @@ -56,6 +56,7 @@ class BaseVectorStoreConfig: vector_store_id: str, query: Union[str, List[str]], vector_store_search_optional_params: VectorStoreSearchOptionalRequestParams, + extra_body: Optional[Dict[str, Any]], api_base: str, litellm_logging_obj: LiteLLMLoggingObj, litellm_params: dict, @@ -67,6 +68,7 @@ class BaseVectorStoreConfig: vector_store_id: str, query: Union[str, List[str]], vector_store_search_optional_params: VectorStoreSearchOptionalRequestParams, + extra_body: Optional[Dict[str, Any]], api_base: str, litellm_logging_obj: LiteLLMLoggingObj, litellm_params: dict, @@ -81,6 +83,7 @@ class BaseVectorStoreConfig: vector_store_id=vector_store_id, query=query, vector_store_search_optional_params=vector_store_search_optional_params, + extra_body=extra_body, api_base=api_base, litellm_logging_obj=litellm_logging_obj, litellm_params=litellm_params, diff --git a/litellm/llms/bedrock/vector_stores/transformation.py b/litellm/llms/bedrock/vector_stores/transformation.py index 4da0a7c779..5b0b64f242 100644 --- a/litellm/llms/bedrock/vector_stores/transformation.py +++ b/litellm/llms/bedrock/vector_stores/transformation.py @@ -1,4 +1,4 @@ -from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Union +from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Union, cast from urllib.parse import urlparse import httpx @@ -7,8 +7,8 @@ from litellm.llms.base_llm.vector_store.transformation import BaseVectorStoreCon from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM from litellm.types.integrations.rag.bedrock_knowledgebase import ( BedrockKBContent, - BedrockKBResponse, BedrockKBRetrievalConfiguration, + BedrockKBResponse, BedrockKBRetrievalQuery, ) from litellm.types.router import GenericLiteLLMParams @@ -199,6 +199,7 @@ class BedrockVectorStoreConfig(BaseVectorStoreConfig, BaseAWSLLM): vector_store_id: str, query: Union[str, List[str]], vector_store_search_optional_params: VectorStoreSearchOptionalRequestParams, + extra_body: Optional[Dict[str, Any]], api_base: str, litellm_logging_obj: LiteLLMLoggingObj, litellm_params: dict, @@ -213,6 +214,14 @@ class BedrockVectorStoreConfig(BaseVectorStoreConfig, BaseAWSLLM): } retrieval_config: Dict[str, Any] = {} + from litellm import verbose_logger + + if isinstance(extra_body, dict): + retrieval_config = dict( + extra_body.get("retrievalConfiguration") + or extra_body.get("retrieval_configuration") + or {} + ) max_results = vector_store_search_optional_params.get("max_num_results") if max_results is not None: retrieval_config.setdefault("vectorSearchConfiguration", {})[ @@ -224,13 +233,9 @@ class BedrockVectorStoreConfig(BaseVectorStoreConfig, BaseAWSLLM): "filter" ] = filters if retrieval_config: - # Create a properly typed retrieval configuration - typed_retrieval_config: BedrockKBRetrievalConfiguration = {} - if "vectorSearchConfiguration" in retrieval_config: - typed_retrieval_config["vectorSearchConfiguration"] = retrieval_config[ - "vectorSearchConfiguration" - ] - request_body["retrievalConfiguration"] = typed_retrieval_config + request_body["retrievalConfiguration"] = cast( + BedrockKBRetrievalConfiguration, retrieval_config + ) litellm_logging_obj.model_call_details["query"] = query return url, request_body diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index b9ada079f6..99f748c0c1 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -8582,6 +8582,7 @@ class BaseLLMHTTPHandler: vector_store_id=vector_store_id, query=query, vector_store_search_optional_params=vector_store_search_optional_params, + extra_body=extra_body, api_base=api_base, litellm_logging_obj=logging_obj, litellm_params=dict(litellm_params), @@ -8594,6 +8595,7 @@ class BaseLLMHTTPHandler: vector_store_id=vector_store_id, query=query, vector_store_search_optional_params=vector_store_search_optional_params, + extra_body=extra_body, api_base=api_base, litellm_logging_obj=logging_obj, litellm_params=dict(litellm_params), @@ -8694,6 +8696,7 @@ class BaseLLMHTTPHandler: vector_store_id=vector_store_id, query=query, vector_store_search_optional_params=vector_store_search_optional_params, + extra_body=extra_body, api_base=api_base, litellm_logging_obj=logging_obj, litellm_params=dict(litellm_params), diff --git a/litellm/llms/gemini/vector_stores/transformation.py b/litellm/llms/gemini/vector_stores/transformation.py index e6e8369643..b6cace066c 100644 --- a/litellm/llms/gemini/vector_stores/transformation.py +++ b/litellm/llms/gemini/vector_stores/transformation.py @@ -115,6 +115,7 @@ class GeminiVectorStoreConfig(BaseVectorStoreConfig): vector_store_id: str, query: Union[str, List[str]], vector_store_search_optional_params: VectorStoreSearchOptionalRequestParams, + extra_body: Optional[Dict[str, Any]], api_base: str, litellm_logging_obj: LiteLLMLoggingObj, litellm_params: dict, diff --git a/litellm/llms/milvus/vector_stores/transformation.py b/litellm/llms/milvus/vector_stores/transformation.py index fcf5d14db7..8c08b78338 100644 --- a/litellm/llms/milvus/vector_stores/transformation.py +++ b/litellm/llms/milvus/vector_stores/transformation.py @@ -127,6 +127,7 @@ class MilvusVectorStoreConfig(BaseVectorStoreConfig): vector_store_id: str, query: Union[str, List[str]], vector_store_search_optional_params: VectorStoreSearchOptionalRequestParams, + extra_body: Optional[Dict[str, Any]], api_base: str, litellm_logging_obj: LiteLLMLoggingObj, litellm_params: dict, diff --git a/litellm/llms/openai/vector_stores/transformation.py b/litellm/llms/openai/vector_stores/transformation.py index c763ed1c8d..b6eae390d9 100644 --- a/litellm/llms/openai/vector_stores/transformation.py +++ b/litellm/llms/openai/vector_stores/transformation.py @@ -103,6 +103,7 @@ class OpenAIVectorStoreConfig(BaseVectorStoreConfig): vector_store_id: str, query: Union[str, List[str]], vector_store_search_optional_params: VectorStoreSearchOptionalRequestParams, + extra_body: Optional[Dict[str, Any]], api_base: str, litellm_logging_obj: LiteLLMLoggingObj, litellm_params: dict, diff --git a/litellm/llms/pg_vector/vector_stores/transformation.py b/litellm/llms/pg_vector/vector_stores/transformation.py index ba87a8f2b0..8261036cae 100644 --- a/litellm/llms/pg_vector/vector_stores/transformation.py +++ b/litellm/llms/pg_vector/vector_stores/transformation.py @@ -77,6 +77,7 @@ class PGVectorStoreConfig(OpenAIVectorStoreConfig): vector_store_id: str, query: Union[str, List[str]], vector_store_search_optional_params: VectorStoreSearchOptionalRequestParams, + extra_body: Optional[Dict[str, Any]], api_base: str, litellm_logging_obj: LiteLLMLoggingObj, litellm_params: dict, @@ -86,6 +87,7 @@ class PGVectorStoreConfig(OpenAIVectorStoreConfig): vector_store_id=vector_store_id, query=query, vector_store_search_optional_params=vector_store_search_optional_params, + extra_body=extra_body, api_base=api_base, litellm_logging_obj=litellm_logging_obj, litellm_params=litellm_params, diff --git a/litellm/llms/ragflow/vector_stores/transformation.py b/litellm/llms/ragflow/vector_stores/transformation.py index ed5397eef0..ae28222c3c 100644 --- a/litellm/llms/ragflow/vector_stores/transformation.py +++ b/litellm/llms/ragflow/vector_stores/transformation.py @@ -99,6 +99,7 @@ class RAGFlowVectorStoreConfig(BaseVectorStoreConfig): vector_store_id: str, query: Union[str, List[str]], vector_store_search_optional_params: VectorStoreSearchOptionalRequestParams, + extra_body: Optional[Dict[str, Any]], api_base: str, litellm_logging_obj: LiteLLMLoggingObj, litellm_params: dict, diff --git a/litellm/llms/s3_vectors/vector_stores/transformation.py b/litellm/llms/s3_vectors/vector_stores/transformation.py index 19b5976986..0cf8635887 100644 --- a/litellm/llms/s3_vectors/vector_stores/transformation.py +++ b/litellm/llms/s3_vectors/vector_stores/transformation.py @@ -76,6 +76,7 @@ class S3VectorsVectorStoreConfig(BaseVectorStoreConfig, BaseAWSLLM): vector_store_id: str, query: Union[str, List[str]], vector_store_search_optional_params: VectorStoreSearchOptionalRequestParams, + extra_body: Optional[Dict[str, Any]], api_base: str, litellm_logging_obj: LiteLLMLoggingObj, litellm_params: dict, @@ -137,6 +138,7 @@ class S3VectorsVectorStoreConfig(BaseVectorStoreConfig, BaseAWSLLM): vector_store_id: str, query: Union[str, List[str]], vector_store_search_optional_params: VectorStoreSearchOptionalRequestParams, + extra_body: Optional[Dict[str, Any]], api_base: str, litellm_logging_obj: LiteLLMLoggingObj, litellm_params: dict, diff --git a/litellm/llms/vertex_ai/vector_stores/rag_api/transformation.py b/litellm/llms/vertex_ai/vector_stores/rag_api/transformation.py index 4baa5774c4..b3fcf4b394 100644 --- a/litellm/llms/vertex_ai/vector_stores/rag_api/transformation.py +++ b/litellm/llms/vertex_ai/vector_stores/rag_api/transformation.py @@ -97,6 +97,7 @@ class VertexVectorStoreConfig(BaseVectorStoreConfig, VertexBase): vector_store_id: str, query: Union[str, List[str]], vector_store_search_optional_params: VectorStoreSearchOptionalRequestParams, + extra_body: Optional[Dict[str, Any]], api_base: str, litellm_logging_obj: LiteLLMLoggingObj, litellm_params: dict, diff --git a/litellm/llms/vertex_ai/vector_stores/search_api/transformation.py b/litellm/llms/vertex_ai/vector_stores/search_api/transformation.py index 179bd7aeff..47dac8dca3 100644 --- a/litellm/llms/vertex_ai/vector_stores/search_api/transformation.py +++ b/litellm/llms/vertex_ai/vector_stores/search_api/transformation.py @@ -104,6 +104,7 @@ class VertexSearchAPIVectorStoreConfig(BaseVectorStoreConfig, VertexBase): vector_store_id: str, query: Union[str, List[str]], vector_store_search_optional_params: VectorStoreSearchOptionalRequestParams, + extra_body: Optional[Dict[str, Any]], api_base: str, litellm_logging_obj: LiteLLMLoggingObj, litellm_params: dict, diff --git a/tests/test_litellm/llms/bedrock/vector_stores/test_bedrock_vector_store_transformation.py b/tests/test_litellm/llms/bedrock/vector_stores/test_bedrock_vector_store_transformation.py index bcd4c62005..5fa48703b1 100644 --- a/tests/test_litellm/llms/bedrock/vector_stores/test_bedrock_vector_store_transformation.py +++ b/tests/test_litellm/llms/bedrock/vector_stores/test_bedrock_vector_store_transformation.py @@ -18,6 +18,7 @@ def test_transform_search_request(): vector_store_id="kb123", query="hello", vector_store_search_optional_params={}, + extra_body=None, api_base="https://bedrock-agent-runtime.us-west-2.amazonaws.com/knowledgebases", litellm_logging_obj=mock_log, litellm_params={}, @@ -25,3 +26,37 @@ def test_transform_search_request(): assert url.endswith("/kb123/retrieve") assert body["retrievalQuery"].get("text") == "hello" + + +def test_transform_search_request_uses_only_retrieval_config_from_extra_body(): + config = BedrockVectorStoreConfig() + mock_log = MagicMock() + mock_log.model_call_details = {} + + url, body = config.transform_search_vector_store_request( + vector_store_id="kb123", + query="hello", + vector_store_search_optional_params={}, + extra_body={ + "retrievalConfiguration": { + "vectorSearchConfiguration": { + "overrideSearchType": "HYBRID", + "numberOfResults": 8, + } + }, + "unrelatedField": {"should_not": "be_forwarded"}, + }, + api_base="https://bedrock-agent-runtime.us-west-2.amazonaws.com/knowledgebases", + litellm_logging_obj=mock_log, + litellm_params={}, + ) + + assert url.endswith("/kb123/retrieve") + assert body["retrievalQuery"].get("text") == "hello" + assert ( + body["retrievalConfiguration"]["vectorSearchConfiguration"][ + "overrideSearchType" + ] + == "HYBRID" + ) + assert "unrelatedField" not in body From 6b86e544e889f6a1aa464b559643e95a26b2e899 Mon Sep 17 00:00:00 2001 From: Sameer Kankute Date: Tue, 28 Apr 2026 15:51:48 +0530 Subject: [PATCH 20/57] Fix greptile reviews --- .../bedrock/vector_stores/transformation.py | 24 ++++++- ...est_bedrock_vector_store_transformation.py | 70 +++++++++++++++++++ 2 files changed, 92 insertions(+), 2 deletions(-) diff --git a/litellm/llms/bedrock/vector_stores/transformation.py b/litellm/llms/bedrock/vector_stores/transformation.py index 5b0b64f242..4e81fa4e66 100644 --- a/litellm/llms/bedrock/vector_stores/transformation.py +++ b/litellm/llms/bedrock/vector_stores/transformation.py @@ -1,8 +1,10 @@ +from copy import deepcopy from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Union, cast from urllib.parse import urlparse import httpx +from litellm._logging import verbose_logger from litellm.llms.base_llm.vector_store.transformation import BaseVectorStoreConfig from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM from litellm.types.integrations.rag.bedrock_knowledgebase import ( @@ -214,21 +216,39 @@ class BedrockVectorStoreConfig(BaseVectorStoreConfig, BaseAWSLLM): } retrieval_config: Dict[str, Any] = {} - from litellm import verbose_logger if isinstance(extra_body, dict): - retrieval_config = dict( + retrieval_config = deepcopy( extra_body.get("retrievalConfiguration") or extra_body.get("retrieval_configuration") or {} ) max_results = vector_store_search_optional_params.get("max_num_results") if max_results is not None: + existing_number_of_results = retrieval_config.get( + "vectorSearchConfiguration", {} + ).get("numberOfResults") + if ( + existing_number_of_results is not None + and existing_number_of_results != max_results + ): + verbose_logger.debug( + "Overriding extra_body retrievalConfiguration.vectorSearchConfiguration.numberOfResults (%s) with max_num_results=%s", + existing_number_of_results, + max_results, + ) retrieval_config.setdefault("vectorSearchConfiguration", {})[ "numberOfResults" ] = max_results filters = vector_store_search_optional_params.get("filters") if filters is not None: + existing_filter = retrieval_config.get("vectorSearchConfiguration", {}).get( + "filter" + ) + if existing_filter is not None and existing_filter != filters: + verbose_logger.debug( + "Overriding extra_body retrievalConfiguration.vectorSearchConfiguration.filter with filters from vector_store_search_optional_params" + ) retrieval_config.setdefault("vectorSearchConfiguration", {})[ "filter" ] = filters diff --git a/tests/test_litellm/llms/bedrock/vector_stores/test_bedrock_vector_store_transformation.py b/tests/test_litellm/llms/bedrock/vector_stores/test_bedrock_vector_store_transformation.py index 5fa48703b1..c211a3536e 100644 --- a/tests/test_litellm/llms/bedrock/vector_stores/test_bedrock_vector_store_transformation.py +++ b/tests/test_litellm/llms/bedrock/vector_stores/test_bedrock_vector_store_transformation.py @@ -60,3 +60,73 @@ def test_transform_search_request_uses_only_retrieval_config_from_extra_body(): == "HYBRID" ) assert "unrelatedField" not in body + + +def test_transform_search_request_does_not_mutate_extra_body_and_overrides_number_of_results(): + config = BedrockVectorStoreConfig() + mock_log = MagicMock() + mock_log.model_call_details = {} + extra_body = { + "retrievalConfiguration": { + "vectorSearchConfiguration": { + "overrideSearchType": "HYBRID", + "numberOfResults": 8, + } + } + } + + _, body = config.transform_search_vector_store_request( + vector_store_id="kb123", + query="hello", + vector_store_search_optional_params={"max_num_results": 10}, + extra_body=extra_body, + api_base="https://bedrock-agent-runtime.us-west-2.amazonaws.com/knowledgebases", + litellm_logging_obj=mock_log, + litellm_params={}, + ) + + assert ( + body["retrievalConfiguration"]["vectorSearchConfiguration"]["numberOfResults"] + == 10 + ) + assert ( + extra_body["retrievalConfiguration"]["vectorSearchConfiguration"][ + "numberOfResults" + ] + == 8 + ) + + +def test_transform_search_request_overrides_filter_without_mutating_extra_body(): + config = BedrockVectorStoreConfig() + mock_log = MagicMock() + mock_log.model_call_details = {} + extra_body = { + "retrievalConfiguration": { + "vectorSearchConfiguration": { + "filter": {"equals": {"key": "tenant", "value": "a"}} + } + } + } + new_filter = {"equals": {"key": "tenant", "value": "b"}} + + _, body = config.transform_search_vector_store_request( + vector_store_id="kb123", + query="hello", + vector_store_search_optional_params={"filters": new_filter}, + extra_body=extra_body, + api_base="https://bedrock-agent-runtime.us-west-2.amazonaws.com/knowledgebases", + litellm_logging_obj=mock_log, + litellm_params={}, + ) + + assert ( + body["retrievalConfiguration"]["vectorSearchConfiguration"]["filter"] + == new_filter + ) + assert ( + extra_body["retrievalConfiguration"]["vectorSearchConfiguration"]["filter"][ + "equals" + ]["value"] + == "a" + ) From 0dd64baa669aef52738f1d628982537707d29e95 Mon Sep 17 00:00:00 2001 From: michelligabriele Date: Tue, 28 Apr 2026 17:25:11 +0200 Subject: [PATCH 21/57] fix(caching): preserve prompt_tokens_details through embedding cache round-trip (#26653) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * fix(caching): preserve prompt_tokens_details through embedding cache round-trip The embedding caching layer was dropping prompt_tokens_details (including image_count) because CachedEmbedding had no field for usage metadata and the cache retrieval code reconstructed Usage without it. This caused inconsistent responses where the first call returned image_count but cached responses did not, breaking cost tracking for multimodal embeddings. Add prompt_tokens_details to CachedEmbedding, persist per-item details during cache storage, aggregate them on retrieval, and merge them in combine_usage() for partial cache hits. * style: apply Black formatting to caching files * fix(caching): address Greptile review — cyclic import, guarded construction, nested dict merge Move PromptTokensDetailsWrapper to inline import to resolve CodeQL cyclic import warning. Guard PromptTokensDetailsWrapper construction with try/except to handle unexpected cached keys. Add recursive dict merging in _merge_prompt_tokens_details for nested fields like cache_creation_token_details. --- litellm/caching/caching.py | 61 +++++- litellm/caching/caching_handler.py | 88 +++++++++ litellm/types/caching.py | 1 + .../caching/test_caching_handler.py | 180 ++++++++++++++++++ 4 files changed, 328 insertions(+), 2 deletions(-) diff --git a/litellm/caching/caching.py b/litellm/caching/caching.py index 6a68ba8c4d..ce1bc26c5e 100644 --- a/litellm/caching/caching.py +++ b/litellm/caching/caching.py @@ -650,7 +650,10 @@ class Cache: verbose_logger.exception(f"LiteLLM Cache: Excepton add_cache: {str(e)}") def _convert_to_cached_embedding( - self, embedding_response: Any, model: Optional[str] + self, + embedding_response: Any, + model: Optional[str], + prompt_tokens_details: Optional[dict] = None, ) -> CachedEmbedding: """ Convert any embedding response into the standardized CachedEmbedding TypedDict format. @@ -662,6 +665,7 @@ class Cache: "index": embedding_response.get("index"), "object": embedding_response.get("object"), "model": model, + "prompt_tokens_details": prompt_tokens_details, } elif hasattr(embedding_response, "model_dump"): data = embedding_response.model_dump() @@ -670,6 +674,7 @@ class Cache: "index": data.get("index"), "object": data.get("object"), "model": model, + "prompt_tokens_details": prompt_tokens_details, } else: data = vars(embedding_response) @@ -678,10 +683,54 @@ class Cache: "index": data.get("index"), "object": data.get("object"), "model": model, + "prompt_tokens_details": prompt_tokens_details, } except KeyError as e: raise ValueError(f"Missing expected key in embedding response: {e}") + def _get_per_item_prompt_tokens_details( + self, + result: EmbeddingResponse, + idx_in_result_data: int, + ) -> Optional[dict]: + """ + Extract per-item prompt_tokens_details from a response for caching. + + For single-item responses (common for multimodal providers like Bedrock Titan, + Nova, Vertex AI), returns the full prompt_tokens_details. + For multi-item responses, distributes integer fields evenly across items + so that summing all per-item details reconstructs the original totals. + """ + if result.usage is None or result.usage.prompt_tokens_details is None: + return None + + details = result.usage.prompt_tokens_details + if hasattr(details, "model_dump"): + details_dict = details.model_dump(exclude_none=True) + elif isinstance(details, dict): + details_dict = {k: v for k, v in details.items() if v is not None} + else: + return None + + if not details_dict: + return None + + num_items = len(result.data) + if num_items <= 1: + return details_dict + + # Distribute integer/float fields evenly across items + per_item: dict = {} + for key, value in details_dict.items(): + if isinstance(value, int): + quotient, remainder = divmod(value, num_items) + per_item[key] = quotient + (1 if idx_in_result_data < remainder else 0) + elif isinstance(value, float): + per_item[key] = value / num_items + else: + per_item[key] = value + return per_item if per_item else None + def add_embedding_response_to_cache( self, result: EmbeddingResponse, @@ -693,10 +742,18 @@ class Cache: kwargs["cache_key"] = preset_cache_key embedding_response = result.data[idx_in_result_data] + # Extract per-item prompt_tokens_details from response usage + prompt_tokens_details = self._get_per_item_prompt_tokens_details( + result=result, + idx_in_result_data=idx_in_result_data, + ) + # Always convert to properly typed CachedEmbedding model_name = result.model embedding_dict: CachedEmbedding = self._convert_to_cached_embedding( - embedding_response, model_name + embedding_response, + model_name, + prompt_tokens_details=prompt_tokens_details, ) cache_key, cached_data, kwargs = self._add_cache_logic( diff --git a/litellm/caching/caching_handler.py b/litellm/caching/caching_handler.py index 2bec705946..7d514e648f 100644 --- a/litellm/caching/caching_handler.py +++ b/litellm/caching/caching_handler.py @@ -59,6 +59,7 @@ from litellm.types.utils import ( if TYPE_CHECKING: from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj + from litellm.types.utils import PromptTokensDetailsWrapper else: LiteLLMLoggingObj = Any @@ -415,6 +416,7 @@ class LLMCachingHandler: final_embedding_cached_response._hidden_params["cache_hit"] = True prompt_tokens = 0 + aggregated_details: Optional[dict] = None for val in non_null_list: idx, cr = val # (idx, cr) tuple if cr is not None: @@ -431,11 +433,35 @@ class LLMCachingHandler: prompt_tokens += token_counter( text=kwargs_input_as_list[idx], count_response_tokens=True ) + # Aggregate prompt_tokens_details from cached items + item_details = cr.get("prompt_tokens_details") + if item_details: + if aggregated_details is None: + aggregated_details = {} + for key, value in item_details.items(): + if isinstance(value, (int, float)): + aggregated_details[key] = ( + aggregated_details.get(key, 0) + value + ) + else: + aggregated_details[key] = value + ## USAGE + prompt_tokens_details: Optional["PromptTokensDetailsWrapper"] = None + if aggregated_details: + from litellm.types.utils import PromptTokensDetailsWrapper + + try: + prompt_tokens_details = PromptTokensDetailsWrapper( + **aggregated_details + ) + except Exception: + prompt_tokens_details = None usage = Usage( prompt_tokens=prompt_tokens, completion_tokens=0, total_tokens=prompt_tokens, + prompt_tokens_details=prompt_tokens_details, ) final_embedding_cached_response.usage = usage if len(remaining_list) == 0: @@ -478,8 +504,70 @@ class LLMCachingHandler: prompt_tokens=usage1.prompt_tokens + usage2.prompt_tokens, completion_tokens=usage1.completion_tokens + usage2.completion_tokens, total_tokens=usage1.total_tokens + usage2.total_tokens, + prompt_tokens_details=self._merge_prompt_tokens_details( + usage1.prompt_tokens_details, + usage2.prompt_tokens_details, + ), ) + def _merge_prompt_tokens_details( + self, + details1: Optional["PromptTokensDetailsWrapper"], + details2: Optional["PromptTokensDetailsWrapper"], + ) -> Optional["PromptTokensDetailsWrapper"]: + """Merge two PromptTokensDetailsWrapper objects by summing numeric fields.""" + if details1 is None and details2 is None: + return None + if details1 is None: + return details2 + if details2 is None: + return details1 + + dict1 = ( + details1.model_dump(exclude_none=True) + if hasattr(details1, "model_dump") + else {} + ) + dict2 = ( + details2.model_dump(exclude_none=True) + if hasattr(details2, "model_dump") + else {} + ) + + merged: dict = {} + for key in set(dict1.keys()) | set(dict2.keys()): + v1 = dict1.get(key, 0) + v2 = dict2.get(key, 0) + if isinstance(v1, (int, float)) and isinstance(v2, (int, float)): + merged[key] = v1 + v2 + elif isinstance(v1, dict) and isinstance(v2, dict): + # Recursively merge nested dicts (e.g. cache_creation_token_details) + nested: dict = {} + for nk in set(v1.keys()) | set(v2.keys()): + nv1 = v1.get(nk, 0) + nv2 = v2.get(nk, 0) + if isinstance(nv1, (int, float)) and isinstance(nv2, (int, float)): + nested[nk] = nv1 + nv2 + elif nv1: + nested[nk] = nv1 + else: + nested[nk] = nv2 + merged[key] = nested + elif v1: + merged[key] = v1 + else: + merged[key] = v2 + + if not merged: + return None + + from litellm.types.utils import PromptTokensDetailsWrapper + + try: + return PromptTokensDetailsWrapper(**merged) + except Exception: + return None + def _combine_cached_embedding_response_with_api_result( self, _caching_handler_response: CachingHandlerResponse, diff --git a/litellm/types/caching.py b/litellm/types/caching.py index c8194ce2e7..f8050b292c 100644 --- a/litellm/types/caching.py +++ b/litellm/types/caching.py @@ -118,3 +118,4 @@ class CachedEmbedding(TypedDict): index: Optional[int] object: Optional[str] model: Optional[str] + prompt_tokens_details: Optional[dict] diff --git a/tests/test_litellm/caching/test_caching_handler.py b/tests/test_litellm/caching/test_caching_handler.py index 837ce7d405..742a4f410d 100644 --- a/tests/test_litellm/caching/test_caching_handler.py +++ b/tests/test_litellm/caching/test_caching_handler.py @@ -52,3 +52,183 @@ async def test_process_async_embedding_cached_response(): print(f"response: {response}") assert len(response.data) == 1 + + +@pytest.mark.asyncio +async def test_embedding_cache_preserves_prompt_tokens_details(): + """Test that prompt_tokens_details (including image_count) survives a full cache hit.""" + llm_caching_handler = LLMCachingHandler( + original_function=MagicMock(), + request_kwargs={}, + start_time=datetime.now(), + ) + + cached_result = [ + { + "embedding": [-0.025, -0.019], + "index": 0, + "object": "embedding", + "model": "amazon.titan-embed-image-v1", + "prompt_tokens_details": {"image_count": 1}, + } + ] + + mock_logging_obj = MagicMock() + mock_logging_obj.async_success_handler = AsyncMock() + response, cache_hit = llm_caching_handler._process_async_embedding_cached_response( + final_embedding_cached_response=None, + cached_result=cached_result, + kwargs={"model": "amazon.titan-embed-image-v1", "input": "base64imagedata"}, + logging_obj=mock_logging_obj, + start_time=datetime.now(), + model="amazon.titan-embed-image-v1", + ) + + assert cache_hit + assert response.usage is not None + assert response.usage.prompt_tokens_details is not None + assert response.usage.prompt_tokens_details.image_count == 1 + + +@pytest.mark.asyncio +async def test_embedding_cache_backward_compat_no_prompt_tokens_details(): + """Test that old cached items without prompt_tokens_details still work.""" + llm_caching_handler = LLMCachingHandler( + original_function=MagicMock(), + request_kwargs={}, + start_time=datetime.now(), + ) + + # Old-format cached item — no prompt_tokens_details field + cached_result = [ + { + "embedding": [-0.025, -0.019], + "index": 0, + "object": "embedding", + "model": "text-embedding-ada-002", + } + ] + + mock_logging_obj = MagicMock() + mock_logging_obj.async_success_handler = AsyncMock() + response, cache_hit = llm_caching_handler._process_async_embedding_cached_response( + final_embedding_cached_response=None, + cached_result=cached_result, + kwargs={"model": "text-embedding-ada-002", "input": "test"}, + logging_obj=mock_logging_obj, + start_time=datetime.now(), + model="text-embedding-ada-002", + ) + + assert cache_hit + assert response.usage is not None + assert response.usage.prompt_tokens_details is None + + +@pytest.mark.asyncio +async def test_embedding_cache_aggregates_multiple_image_counts(): + """Test that image_count is summed correctly across multiple cached items.""" + llm_caching_handler = LLMCachingHandler( + original_function=MagicMock(), + request_kwargs={}, + start_time=datetime.now(), + ) + + cached_result = [ + { + "embedding": [-0.025, -0.019], + "index": 0, + "object": "embedding", + "model": "amazon.titan-embed-image-v1", + "prompt_tokens_details": {"image_count": 1}, + }, + { + "embedding": [0.031, 0.042], + "index": 1, + "object": "embedding", + "model": "amazon.titan-embed-image-v1", + "prompt_tokens_details": {"image_count": 1}, + }, + ] + + mock_logging_obj = MagicMock() + mock_logging_obj.async_success_handler = AsyncMock() + response, cache_hit = llm_caching_handler._process_async_embedding_cached_response( + final_embedding_cached_response=None, + cached_result=cached_result, + kwargs={ + "model": "amazon.titan-embed-image-v1", + "input": ["img1", "img2"], + }, + logging_obj=mock_logging_obj, + start_time=datetime.now(), + model="amazon.titan-embed-image-v1", + ) + + assert cache_hit + assert response.usage.prompt_tokens_details is not None + assert response.usage.prompt_tokens_details.image_count == 2 + + +def test_combine_usage_merges_prompt_tokens_details(): + """Test that combine_usage merges prompt_tokens_details from both Usage objects.""" + from litellm.types.utils import PromptTokensDetailsWrapper, Usage + + llm_caching_handler = LLMCachingHandler( + original_function=MagicMock(), + request_kwargs={}, + start_time=datetime.now(), + ) + + usage1 = Usage( + prompt_tokens=10, + completion_tokens=0, + total_tokens=10, + prompt_tokens_details=PromptTokensDetailsWrapper(image_count=1), + ) + usage2 = Usage( + prompt_tokens=20, + completion_tokens=0, + total_tokens=20, + prompt_tokens_details=PromptTokensDetailsWrapper(image_count=2), + ) + + combined = llm_caching_handler.combine_usage(usage1, usage2) + + assert combined.prompt_tokens == 30 + assert combined.total_tokens == 30 + assert combined.prompt_tokens_details is not None + assert combined.prompt_tokens_details.image_count == 3 + + +def test_combine_usage_handles_none_details(): + """Test that combine_usage works when one or both sides have null prompt_tokens_details.""" + from litellm.types.utils import PromptTokensDetailsWrapper, Usage + + llm_caching_handler = LLMCachingHandler( + original_function=MagicMock(), + request_kwargs={}, + start_time=datetime.now(), + ) + + # Both null + usage_a = Usage(prompt_tokens=10, completion_tokens=0, total_tokens=10) + usage_b = Usage(prompt_tokens=20, completion_tokens=0, total_tokens=20) + combined = llm_caching_handler.combine_usage(usage_a, usage_b) + assert combined.prompt_tokens_details is None + + # Only first has details + usage_c = Usage( + prompt_tokens=10, + completion_tokens=0, + total_tokens=10, + prompt_tokens_details=PromptTokensDetailsWrapper(image_count=1), + ) + combined = llm_caching_handler.combine_usage(usage_c, usage_b) + assert combined.prompt_tokens_details is not None + assert combined.prompt_tokens_details.image_count == 1 + + # Only second has details + combined = llm_caching_handler.combine_usage(usage_a, usage_c) + assert combined.prompt_tokens_details is not None + assert combined.prompt_tokens_details.image_count == 1 From 10aed9e9816c61600765766428c1c167327e2c64 Mon Sep 17 00:00:00 2001 From: milan-berri Date: Tue, 28 Apr 2026 18:38:17 +0300 Subject: [PATCH 22/57] feat(logging): add retry settings for generic API logger (#26645) * Add retry settings for generic API logger Made-with: Cursor * Refine generic API retry behavior Made-with: Cursor --- .../generic_api/generic_api_callback.py | 72 +++++++++++--- .../logging_callback_manager.py | 13 +++ .../test_logging_callback_manager.py | 37 ++++++++ .../test_generic_api_callback.py | 94 +++++++++++++++++++ 4 files changed, 205 insertions(+), 11 deletions(-) diff --git a/litellm/integrations/generic_api/generic_api_callback.py b/litellm/integrations/generic_api/generic_api_callback.py index 9a8060520d..2982df8fda 100644 --- a/litellm/integrations/generic_api/generic_api_callback.py +++ b/litellm/integrations/generic_api/generic_api_callback.py @@ -11,8 +11,9 @@ import json import os import re import traceback -from typing import Dict, List, Literal, Optional, Union +from typing import Any, Dict, List, Literal, Optional, Union +import httpx import litellm from litellm._logging import verbose_logger from litellm._uuid import uuid @@ -103,6 +104,9 @@ class GenericAPILogger(CustomBatchLogger): event_types: Optional[List[API_EVENT_TYPES]] = None, callback_name: Optional[str] = None, log_format: Optional[LOG_FORMAT_TYPES] = None, + max_retries: int = 0, + retry_delay: float = 1.0, + timeout: Optional[Union[float, httpx.Timeout]] = None, **kwargs, ): """ @@ -114,6 +118,9 @@ class GenericAPILogger(CustomBatchLogger): event_types: Optional[List[API_EVENT_TYPES]] = None, callback_name: Optional[str] = None - If provided, loads config from generic_api_compatible_callbacks.json log_format: Optional[LOG_FORMAT_TYPES] = None - Format for log output: "json_array" (default), "ndjson", or "single" + max_retries: Number of retry attempts after the initial request fails. Defaults to 0. + retry_delay: Initial retry delay in seconds. Retries use exponential backoff. + timeout: Optional timeout to use for Generic API callback requests. """ ######################################################### # Check if callback_name is provided and load config @@ -162,6 +169,10 @@ class GenericAPILogger(CustomBatchLogger): self.endpoint: str = endpoint self.event_types: Optional[List[API_EVENT_TYPES]] = event_types self.callback_name: Optional[str] = callback_name + self.max_retries = max(0, int(max_retries or 0)) + retry_delay_value = 0.0 if retry_delay is None else retry_delay + self.retry_delay = max(0.0, float(retry_delay_value)) + self.timeout = timeout # Validate and store log_format if log_format is not None and log_format not in [ @@ -226,6 +237,53 @@ class GenericAPILogger(CustomBatchLogger): return headers_dict + def _should_retry_exception(self, exception: Exception) -> bool: + if isinstance(exception, (litellm.Timeout, httpx.TransportError)): + return True + + if isinstance(exception, httpx.HTTPStatusError): + return exception.response.status_code >= 500 + + return False + + async def _sleep_before_retry(self, attempt: int) -> None: + if self.retry_delay <= 0: + return + + delay = self.retry_delay * (2**attempt) + await asyncio.sleep(delay) + + async def _post_with_retries(self, data: str) -> httpx.Response: + post_kwargs: Dict[str, Any] = { + "url": self.endpoint, + "headers": self.headers, + "data": data, + } + if self.timeout is not None: + post_kwargs["timeout"] = self.timeout + + total_attempts = self.max_retries + 1 + for attempt in range(total_attempts): + try: + return await self.async_httpx_client.post(**post_kwargs) + except Exception as e: + is_last_attempt = attempt == self.max_retries + should_retry = self._should_retry_exception(e) + if is_last_attempt or not should_retry: + raise + + verbose_logger.warning( + "Generic API Logger - retrying request to %s after error: %s " + "(attempt %s/%s)", + self.endpoint, + str(e), + attempt + 1, + total_attempts, + ) + await self._sleep_before_retry(attempt) + + raise RuntimeError("Generic API Logger retry loop exited unexpectedly") + async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): """ Async Log success events to Generic API Endpoint @@ -325,11 +383,7 @@ class GenericAPILogger(CustomBatchLogger): # Send each log as individual HTTP request in parallel tasks = [] for log_entry in self.log_queue: - task = self.async_httpx_client.post( - url=self.endpoint, - headers=self.headers, - data=safe_dumps(log_entry), - ) + task = self._post_with_retries(data=safe_dumps(log_entry)) tasks.append(task) # Execute all requests in parallel @@ -356,11 +410,7 @@ class GenericAPILogger(CustomBatchLogger): raise ValueError(f"Unknown log_format: {self.log_format}") # Make POST request - response = await self.async_httpx_client.post( - url=self.endpoint, - headers=self.headers, - data=data, - ) + response = await self._post_with_retries(data=data) verbose_logger.debug( f"Generic API Logger - sent batch to {self.endpoint}, " diff --git a/litellm/litellm_core_utils/logging_callback_manager.py b/litellm/litellm_core_utils/logging_callback_manager.py index c5c150274c..6c749118de 100644 --- a/litellm/litellm_core_utils/logging_callback_manager.py +++ b/litellm/litellm_core_utils/logging_callback_manager.py @@ -221,6 +221,13 @@ class LoggingCallbackManager: headers = callback_config.get("headers") event_types = callback_config.get("event_types") log_format = callback_config.get("log_format") + max_retries = max(0, int(callback_config.get("max_retries", 0) or 0)) + retry_delay_value = callback_config.get("retry_delay") + retry_delay = max( + 0.0, + float(0.0 if retry_delay_value is None else retry_delay_value), + ) + timeout = callback_config.get("timeout") if endpoint is None or headers is None: verbose_logger.warning( @@ -236,6 +243,9 @@ class LoggingCallbackManager: and cached_logger.headers == headers and cached_logger.event_types == event_types and cached_logger.log_format == log_format + and cached_logger.max_retries == max_retries + and cached_logger.retry_delay == retry_delay + and cached_logger.timeout == timeout ): return cached_logger @@ -244,6 +254,9 @@ class LoggingCallbackManager: headers=headers, event_types=event_types, log_format=log_format, + max_retries=max_retries, + retry_delay=retry_delay, + timeout=timeout, ) _generic_api_logger_cache[callback] = new_logger return new_logger diff --git a/tests/litellm_utils_tests/test_logging_callback_manager.py b/tests/litellm_utils_tests/test_logging_callback_manager.py index 88ae07fd81..d9540f8f85 100644 --- a/tests/litellm_utils_tests/test_logging_callback_manager.py +++ b/tests/litellm_utils_tests/test_logging_callback_manager.py @@ -366,3 +366,40 @@ def test_generic_api_compatible_callbacks_json_unknown_callback(): # Should return the string unchanged assert result == "unknown_callback", "Unknown callback should be returned as-is" assert isinstance(result, str), "Unknown callback should remain a string" + + +@pytest.mark.asyncio +async def test_generic_api_callback_settings_retry_config(): + """ + Test that generic_api callback_settings are passed to GenericAPILogger. + """ + from litellm.integrations.generic_api.generic_api_callback import GenericAPILogger + from litellm.litellm_core_utils.logging_callback_manager import ( + _generic_api_logger_cache, + ) + + callback_name = "test_generic_api_retry_config" + _generic_api_logger_cache.pop(callback_name, None) + litellm.callback_settings[callback_name] = { + "callback_type": "generic_api", + "endpoint": "https://example.com/api/logs", + "headers": {"Content-Type": "application/json"}, + "max_retries": 2, + "retry_delay": 0.5, + "timeout": 3, + } + + try: + result = LoggingCallbackManager._add_custom_callback_generic_api_str( + callback_name + ) + + assert isinstance(result, GenericAPILogger) + assert result.endpoint == "https://example.com/api/logs" + assert result.headers == {"Content-Type": "application/json"} + assert result.max_retries == 2 + assert result.retry_delay == 0.5 + assert result.timeout == 3 + finally: + litellm.callback_settings.pop(callback_name, None) + _generic_api_logger_cache.pop(callback_name, None) diff --git a/tests/logging_callback_tests/test_generic_api_callback.py b/tests/logging_callback_tests/test_generic_api_callback.py index 528a5101df..6984b6fa00 100644 --- a/tests/logging_callback_tests/test_generic_api_callback.py +++ b/tests/logging_callback_tests/test_generic_api_callback.py @@ -8,6 +8,7 @@ sys.path.insert(0, os.path.abspath("../..")) import asyncio import litellm import gzip +import httpx import json import logging import time @@ -470,3 +471,96 @@ async def test_generic_api_callback_invalid_log_format(): endpoint=test_endpoint, log_format="invalid_format", # type: ignore # Intentionally invalid for testing ) + + +@pytest.mark.asyncio +async def test_generic_api_callback_retries_timeout_then_succeeds(): + """ + Test that GenericAPILogger retries LiteLLM timeout errors when configured. + """ + test_endpoint = "https://example.com/api/logs" + generic_logger = GenericAPILogger( + endpoint=test_endpoint, + max_retries=1, + retry_delay=0, + timeout=0.2, + ) + + mock_post = AsyncMock() + mock_post.side_effect = [ + litellm.Timeout( + message="Connection timed out", + model="default-model-name", + llm_provider="litellm-httpx-handler", + ), + type("Response", (), {"status_code": 200})(), + ] + generic_logger.async_httpx_client.post = mock_post + generic_logger.log_queue = [{"event": "timeout-retry"}] + + await generic_logger.async_send_batch() + + assert mock_post.call_count == 2 + first_call = mock_post.call_args_list[0][1] + assert first_call["url"] == test_endpoint + assert first_call["timeout"] == 0.2 + assert json.loads(first_call["data"]) == [{"event": "timeout-retry"}] + + +@pytest.mark.asyncio +async def test_generic_api_callback_retries_5xx_then_succeeds(): + """ + Test that GenericAPILogger retries transient HTTP 5xx errors when configured. + """ + test_endpoint = "https://example.com/api/logs" + generic_logger = GenericAPILogger( + endpoint=test_endpoint, + max_retries=1, + retry_delay=0, + ) + + request = httpx.Request("POST", test_endpoint) + response = httpx.Response(status_code=503, request=request) + mock_post = AsyncMock() + mock_post.side_effect = [ + httpx.HTTPStatusError( + "Server error", + request=request, + response=response, + ), + type("Response", (), {"status_code": 200})(), + ] + generic_logger.async_httpx_client.post = mock_post + generic_logger.log_queue = [{"event": "5xx-retry"}] + + await generic_logger.async_send_batch() + + assert mock_post.call_count == 2 + + +@pytest.mark.asyncio +async def test_generic_api_callback_does_not_retry_4xx(): + """ + Test that GenericAPILogger does not retry non-transient HTTP 4xx errors. + """ + test_endpoint = "https://example.com/api/logs" + generic_logger = GenericAPILogger( + endpoint=test_endpoint, + max_retries=2, + retry_delay=0, + ) + + request = httpx.Request("POST", test_endpoint) + response = httpx.Response(status_code=401, request=request) + mock_post = AsyncMock() + mock_post.side_effect = httpx.HTTPStatusError( + "Unauthorized", + request=request, + response=response, + ) + generic_logger.async_httpx_client.post = mock_post + generic_logger.log_queue = [{"event": "4xx-no-retry"}] + + await generic_logger.async_send_batch() + + mock_post.assert_called_once() From 52fb23a512894cc283c1a94a88eebea3745b05b5 Mon Sep 17 00:00:00 2001 From: milan-berri Date: Tue, 28 Apr 2026 18:41:20 +0300 Subject: [PATCH 23/57] fix(logging): backfill streaming hidden response cost (#26606) * fix(logging): backfill streaming hidden response cost Made-with: Cursor * fix(logging): avoid mutating streaming hidden params Backfill calculated streaming response cost into logging payload copies so OTEL spans expose hidden_params.response_cost without mutating the response object. Made-with: Cursor * fix black formatting Apply the repo-pinned Black 24.10.0 formatting expected by CI. Made-with: Cursor * fix(types): allow numeric hidden response cost Allow standard logging hidden params to carry numeric response_cost values, matching LiteLLM's calculated cost payloads. Made-with: Cursor * refactor(logging): simplify hidden response cost backfill Clean up metadata initialization and reuse the raw response cost when deciding whether to backfill hidden params. Made-with: Cursor --- litellm/litellm_core_utils/litellm_logging.py | 36 ++++--- litellm/types/utils.py | 2 +- .../test_litellm_logging.py | 98 +++++++++++++++++++ 3 files changed, 123 insertions(+), 13 deletions(-) diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 57341472b4..fb103afea0 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -1725,12 +1725,18 @@ class Logging(LiteLLMLoggingBaseClass): return if self.model_call_details.get("litellm_params") is None: return - self.model_call_details["litellm_params"].setdefault("metadata", {}) - if self.model_call_details["litellm_params"]["metadata"] is None: - self.model_call_details["litellm_params"]["metadata"] = {} - self.model_call_details["litellm_params"]["metadata"]["hidden_params"] = ( - getattr(logging_result, "_hidden_params", {}) - ) + metadata_hidden_params = hidden_params.copy() + response_cost = self.model_call_details.get("response_cost") + if ( + metadata_hidden_params.get("response_cost") is None + and response_cost is not None + ): + metadata_hidden_params["response_cost"] = response_cost + + litellm_params = self.model_call_details["litellm_params"] + metadata = litellm_params.get("metadata") or {} + litellm_params["metadata"] = metadata + metadata["hidden_params"] = metadata_hidden_params def _process_hidden_params_and_response_cost( self, @@ -5438,11 +5444,6 @@ def get_standard_logging_object_payload( completion_start_time_float=completion_start_time_float, stream=kwargs.get("stream", False), ) - # clean up litellm hidden params - clean_hidden_params = StandardLoggingPayloadSetup.get_hidden_params( - hidden_params - ) - # clean up litellm metadata clean_metadata = StandardLoggingPayloadSetup.get_standard_logging_metadata( metadata=metadata, @@ -5476,6 +5477,18 @@ def get_standard_logging_object_payload( ## Get model cost information ## base_model = _get_base_model_from_metadata(model_call_details=kwargs) custom_pricing = use_custom_pricing_for_model(litellm_params=litellm_params) + raw_response_cost = kwargs.get("response_cost") + response_cost: float = raw_response_cost or 0.0 + + # clean up litellm hidden params + clean_hidden_params = StandardLoggingPayloadSetup.get_hidden_params( + hidden_params + ) + if ( + clean_hidden_params["response_cost"] is None + and raw_response_cost is not None + ): + clean_hidden_params["response_cost"] = response_cost model_cost_information = StandardLoggingPayloadSetup.get_model_cost_information( base_model=base_model, @@ -5484,7 +5497,6 @@ def get_standard_logging_object_payload( init_response_obj=init_response_obj, api_base=litellm_params.get("api_base"), ) - response_cost: float = kwargs.get("response_cost", 0) or 0.0 error_information = StandardLoggingPayloadSetup.get_error_information( original_exception=original_exception, diff --git a/litellm/types/utils.py b/litellm/types/utils.py index a212d56c1a..ed29d49fc2 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -2659,7 +2659,7 @@ class StandardLoggingHiddenParams(TypedDict): ] # id of the model in the router, separates multiple models with the same name but different credentials cache_key: Optional[str] api_base: Optional[str] - response_cost: Optional[str] + response_cost: Optional[Union[str, float]] litellm_overhead_time_ms: Optional[float] additional_headers: Optional[StandardLoggingAdditionalHeaders] batch_models: Optional[List[str]] diff --git a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py index cf7be6bf1c..3348118a02 100644 --- a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py +++ b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py @@ -2337,6 +2337,104 @@ def test_merge_hidden_params_from_response_into_metadata_populates_metadata(): assert meta["hidden_params"]["model_id"] == "mid-test" +def test_merge_hidden_params_from_response_into_metadata_backfills_response_cost(): + """Streaming metadata should include the already-calculated response cost.""" + from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj + + logging_obj = LiteLLMLoggingObj( + model="gpt-4o-mini", + messages=[{"role": "user", "content": "hi"}], + stream=True, + call_type="acompletion", + start_time=time.time(), + litellm_call_id="merge-hp-cost-test", + function_id="merge-hp-cost-fn", + ) + logging_obj.model_call_details = { + "litellm_params": {"metadata": {}}, + "response_cost": 0.002, + } + + class _Resp: + _hidden_params = {"response_cost": None, "model_id": "mid-test"} + + response = _Resp() + logging_obj._merge_hidden_params_from_response_into_metadata(response) + meta = logging_obj.model_call_details["litellm_params"]["metadata"] + assert meta["hidden_params"]["response_cost"] == 0.002 + assert meta["hidden_params"]["model_id"] == "mid-test" + assert response._hidden_params["response_cost"] is None + + +def test_standard_logging_hidden_params_backfills_response_cost_without_mutating_response(): + """Streaming standard logging payload should expose the calculated response cost.""" + from datetime import datetime + + from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj + from litellm.types.utils import Usage + + logging_obj = LiteLLMLoggingObj( + model="gpt-4o-mini", + messages=[{"role": "user", "content": "hi"}], + stream=True, + call_type="acompletion", + start_time=time.time(), + litellm_call_id="standard-hp-cost-test", + function_id="standard-hp-cost-fn", + ) + logging_obj.model_call_details = { + "litellm_params": {"metadata": {}, "proxy_server_request": {}}, + "litellm_call_id": "standard-hp-cost-test", + "call_type": "acompletion", + "stream": True, + "model": "gpt-4o-mini", + "custom_llm_provider": "openai", + "optional_params": {"stream": True}, + "response_cost": 0.002, + } + response = ModelResponse( + id="standard-hp-cost-response", + model="gpt-4o-mini", + usage=Usage(prompt_tokens=10, completion_tokens=5, total_tokens=15), + ) + response._hidden_params = {"response_cost": None, "model_id": "mid-test"} + + payload = logging_obj._build_standard_logging_payload( + response, datetime.now(), datetime.now() + ) + + assert payload is not None + assert payload["hidden_params"]["response_cost"] == 0.002 + assert response._hidden_params["response_cost"] is None + + +def test_merge_hidden_params_from_response_into_metadata_preserves_response_cost(): + """Do not overwrite provider-supplied response cost when it already exists.""" + from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj + + logging_obj = LiteLLMLoggingObj( + model="gpt-4o-mini", + messages=[{"role": "user", "content": "hi"}], + stream=True, + call_type="acompletion", + start_time=time.time(), + litellm_call_id="merge-hp-preserve-cost-test", + function_id="merge-hp-preserve-cost-fn", + ) + logging_obj.model_call_details = { + "litellm_params": {"metadata": {}}, + "response_cost": 0.002, + } + + class _Resp: + _hidden_params = {"response_cost": 0.001, "model_id": "mid-test"} + + logging_obj._merge_hidden_params_from_response_into_metadata(_Resp()) + meta = logging_obj.model_call_details["litellm_params"]["metadata"] + assert meta["hidden_params"]["response_cost"] == 0.001 + assert meta["hidden_params"]["model_id"] == "mid-test" + + def test_merge_hidden_params_from_response_into_metadata_no_op_when_empty(): from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj From 1d56e732e835e9ad12fa63e92400e7b61b6c4440 Mon Sep 17 00:00:00 2001 From: Sameer Kankute Date: Tue, 28 Apr 2026 21:14:40 +0530 Subject: [PATCH 24/57] fix(vertex-ai): reuse anthropic messages config instances (#26099) Cache provider config lookups for Vertex Anthropic messages so repeated requests reuse the same config object and preserve credential cache state. Add a regression test to catch any future loss of config reuse. Made-with: Cursor --- litellm/utils.py | 15 +++++++++-- ...artner_models_anthropic_messages_config.py | 26 +++++++++++++++++++ 2 files changed, 39 insertions(+), 2 deletions(-) diff --git a/litellm/utils.py b/litellm/utils.py index e1ad1db63e..e63bf402bf 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -8410,6 +8410,17 @@ class ProviderConfigManager: model: str, provider: LlmProviders, ) -> Optional[BaseAnthropicMessagesConfig]: + return ProviderConfigManager._get_provider_anthropic_messages_config_cached( + model=model, provider=provider + ) + + @staticmethod + @lru_cache(maxsize=DEFAULT_MAX_LRU_CACHE_SIZE) + def _get_provider_anthropic_messages_config_cached( + model: str, + provider: LlmProviders, + ) -> Optional[BaseAnthropicMessagesConfig]: + model_lower = model.lower() if litellm.LlmProviders.ANTHROPIC == provider: return litellm.AnthropicMessagesConfig() # The 'BEDROCK' provider corresponds to Amazon's implementation of Anthropic Claude v3. @@ -8419,14 +8430,14 @@ class ProviderConfigManager: return BedrockModelInfo.get_bedrock_provider_config_for_messages_api(model) elif litellm.LlmProviders.VERTEX_AI == provider: - if "claude" in model.lower(): + if "claude" in model_lower: from litellm.llms.vertex_ai.vertex_ai_partner_models.anthropic.experimental_pass_through.transformation import ( VertexAIPartnerModelsAnthropicMessagesConfig, ) return VertexAIPartnerModelsAnthropicMessagesConfig() elif litellm.LlmProviders.AZURE_AI == provider: - if "claude" in model.lower(): + if "claude" in model_lower: from litellm.llms.azure_ai.anthropic.messages_transformation import ( AzureAnthropicMessagesConfig, ) diff --git a/tests/test_litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/test_vertex_ai_partner_models_anthropic_messages_config.py b/tests/test_litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/test_vertex_ai_partner_models_anthropic_messages_config.py index 214e5f0797..b8cd65d3c9 100644 --- a/tests/test_litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/test_vertex_ai_partner_models_anthropic_messages_config.py +++ b/tests/test_litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/test_vertex_ai_partner_models_anthropic_messages_config.py @@ -311,3 +311,29 @@ def test_transform_anthropic_messages_request_removes_scope_from_cache_control() # scope removed from message content assert "scope" not in result["messages"][0]["content"][0]["cache_control"] assert result["messages"][0]["content"][0]["cache_control"]["type"] == "ephemeral" + + +def test_provider_config_manager_reuses_vertex_anthropic_messages_config_instance(): + """ + Regression test: repeated provider config lookups for the same Vertex Claude model + should return the same config instance (which preserves auth cache state). + """ + import litellm + from litellm.utils import ProviderConfigManager + + ProviderConfigManager._get_provider_anthropic_messages_config_cached.cache_clear() + try: + first_config = ProviderConfigManager.get_provider_anthropic_messages_config( + model="claude-opus-4-6", + provider=litellm.LlmProviders.VERTEX_AI, + ) + second_config = ProviderConfigManager.get_provider_anthropic_messages_config( + model="claude-opus-4-6", + provider=litellm.LlmProviders.VERTEX_AI, + ) + + assert isinstance(first_config, VertexAIPartnerModelsAnthropicMessagesConfig) + assert isinstance(second_config, VertexAIPartnerModelsAnthropicMessagesConfig) + assert first_config is second_config + finally: + ProviderConfigManager._get_provider_anthropic_messages_config_cached.cache_clear() From 1af11d4371ac5aed4c0263a2d34061c28d9e3ba3 Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Tue, 28 Apr 2026 09:23:55 -0700 Subject: [PATCH 25/57] fix(vertex): synthesize items for array types missing items entirely MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Companion to the prior commit. process_items only converted empty `items: {}` to `{"type": "object"}`. But anyOf branches like `{"type": "array"}` (no items field at all) were untouched, so after convert_anyof_null_to_nullable stripped the null branch and added nullable, the array branch was sent to Vertex as `{"type": "array", "nullable": true}` — which Vertex rejects with INVALID_ARGUMENT (`any_of[0].items: missing field`). Make process_items synthesize `items: {"type": "object"}` for any `type == "array"` schema where items is missing or empty. Also: - Convert test_gemini_tool_calling_working_demo to a hermetic mock test asserting items is present on the array branch in the sent body. Was previously a real-network call to Vertex and was the test the user reported still failing in CI. - Add unit test test_build_vertex_schema_array_branch_missing_items_in_anyof covering the missing-items shape directly. Co-Authored-By: Claude Opus 4.7 (1M context) --- litellm/llms/vertex_ai/common_utils.py | 9 ++- .../test_amazing_vertex_completion.py | 70 +++++++++++++++++-- .../vertex_ai/test_vertex_ai_common_utils.py | 37 ++++++++++ 3 files changed, 111 insertions(+), 5 deletions(-) diff --git a/litellm/llms/vertex_ai/common_utils.py b/litellm/llms/vertex_ai/common_utils.py index 9b23520dcd..b4bfde5f54 100644 --- a/litellm/llms/vertex_ai/common_utils.py +++ b/litellm/llms/vertex_ai/common_utils.py @@ -597,7 +597,14 @@ def process_items(schema, depth=0): f"Max depth of {DEFAULT_MAX_RECURSE_DEPTH} exceeded while processing schema. Please check the schema for excessive nesting." ) if isinstance(schema, dict): - if "items" in schema and schema["items"] == {}: + # Vertex requires `items` whenever `type == "array"` (even inside anyOf). + # Normalize: empty `items: {}` and missing-items both become {"type": "object"}. + type_val = schema.get("type") + if ( + isinstance(type_val, str) + and type_val.lower() == "array" + and ("items" not in schema or schema.get("items") == {}) + ): schema["items"] = {"type": "object"} for key, value in schema.items(): if isinstance(value, dict): diff --git a/tests/local_testing/test_amazing_vertex_completion.py b/tests/local_testing/test_amazing_vertex_completion.py index 3b4ecb82b1..9782bf3c2a 100644 --- a/tests/local_testing/test_amazing_vertex_completion.py +++ b/tests/local_testing/test_amazing_vertex_completion.py @@ -3493,8 +3493,14 @@ def test_litellm_api_base(monkeypatch, provider, route): def test_gemini_tool_calling_working_demo(): - load_vertex_ai_credentials() - litellm._turn_on_debug() + """ + Regression test: tool params with anyOf containing a `{"type": "array"}` + branch (no items field at all) must synthesize items before the request + is sent to Vertex (Vertex rejects array types missing items). + """ + from litellm.llms.custom_httpx.http_handler import HTTPHandler + from litellm.llms.vertex_ai.vertex_llm_base import VertexBase + args = { "messages": [ { @@ -3564,8 +3570,64 @@ def test_gemini_tool_calling_working_demo(): ], "vertex_location": "global", } - response = completion(model="vertex_ai/gemini-3-flash-preview", **args) - print(response) + + client = HTTPHandler() + mock_response = MagicMock() + mock_response.status_code = 200 + mock_response.headers = {"Content-Type": "application/json"} + mock_response.json.return_value = { + "candidates": [ + { + "content": { + "role": "model", + "parts": [{"text": "Hello!"}], + }, + "finishReason": "STOP", + } + ], + "usageMetadata": { + "promptTokenCount": 10, + "candidatesTokenCount": 5, + "totalTokenCount": 15, + }, + } + + with ( + patch.object(client, "post", return_value=mock_response) as mock_post, + patch.object( + VertexBase, + "_ensure_access_token", + return_value=("fake-token", "fake-project"), + ), + ): + completion( + model="vertex_ai/gemini-3-flash-preview", + client=client, + **args, + ) + + sent_body = mock_post.call_args.kwargs.get( + "json" + ) or mock_post.call_args.kwargs.get("data") + assert sent_body is not None, "expected request body to be sent" + if isinstance(sent_body, str): + sent_body = json.loads(sent_body) + + function_decl = sent_body["tools"][0]["function_declarations"][0] + callbacks_schema = function_decl["parameters"]["properties"]["config"][ + "properties" + ]["callbacks"] + array_branches = [ + branch + for branch in callbacks_schema["anyOf"] + if branch.get("type", "").lower() == "array" + ] + assert array_branches, "expected an array branch in callbacks anyOf" + for branch in array_branches: + assert "items" in branch and branch["items"], ( + f"array branch in callbacks.anyOf must include non-empty items " + f"(Vertex rejects array types missing items). Got: {branch}" + ) def test_gemini_tool_calling_not_working(): diff --git a/tests/test_litellm/llms/vertex_ai/test_vertex_ai_common_utils.py b/tests/test_litellm/llms/vertex_ai/test_vertex_ai_common_utils.py index dc3be7114f..95507390df 100644 --- a/tests/test_litellm/llms/vertex_ai/test_vertex_ai_common_utils.py +++ b/tests/test_litellm/llms/vertex_ai/test_vertex_ai_common_utils.py @@ -292,6 +292,43 @@ def test_process_items_basic(): process_items(schema) assert schema["properties"]["nested"]["items"] == {"type": "object"} + # Vertex rejects array types missing `items` entirely (not just empty). + # Synthesize {"type": "object"} so the request validates. + schema = {"type": "array"} + process_items(schema) + assert schema["items"] == {"type": "object"} + + +def test_build_vertex_schema_array_branch_missing_items_in_anyof(): + """ + Regression: an `anyOf` branch with `{"type": "array"}` (no items) must + end up with synthesized `items: {"type": "object"}` after the schema + transform — Vertex returns INVALID_ARGUMENT otherwise. + """ + from litellm.llms.vertex_ai.common_utils import _build_vertex_schema + + parameters = { + "properties": { + "callbacks": { + "anyOf": [ + {"type": "array"}, + {"type": "object"}, + {"type": "null"}, + ] + } + }, + "type": "object", + } + + result = _build_vertex_schema(parameters) + callbacks_anyof = result["properties"]["callbacks"]["anyOf"] + array_branches = [b for b in callbacks_anyof if b.get("type") == "array"] + assert array_branches, "expected an array branch to remain after transform" + for branch in array_branches: + assert branch.get("items") == { + "type": "object" + }, f"array branch must have items synthesized; got {branch}" + def test_vertex_ai_complex_response_schema(): import json From 898040fcdda0f3193862ab495905198e26f5c45f Mon Sep 17 00:00:00 2001 From: Sameer Kankute Date: Tue, 28 Apr 2026 22:34:14 +0530 Subject: [PATCH 26/57] Fix tests --- .../s3_vectors/vector_stores/test_s3_vectors_transformation.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/test_litellm/llms/s3_vectors/vector_stores/test_s3_vectors_transformation.py b/tests/test_litellm/llms/s3_vectors/vector_stores/test_s3_vectors_transformation.py index 7c06823d16..9507ff401a 100644 --- a/tests/test_litellm/llms/s3_vectors/vector_stores/test_s3_vectors_transformation.py +++ b/tests/test_litellm/llms/s3_vectors/vector_stores/test_s3_vectors_transformation.py @@ -56,6 +56,7 @@ class TestS3VectorsVectorStoreConfig: vector_store_id="invalid-format", query="test query", vector_store_search_optional_params={}, + extra_body=None, api_base="https://s3vectors.us-west-2.api.aws", litellm_logging_obj=mock_logging_obj, litellm_params={}, From dc46467235fa498d3d84482b9942604d5d694b4f Mon Sep 17 00:00:00 2001 From: Ryan Crabbe Date: Tue, 28 Apr 2026 14:24:19 -0700 Subject: [PATCH 27/57] fix(tests): replace deprecated Bedrock Claude 3.7 Sonnet model ID AWS Bedrock has reached end-of-life for `claude-3-7-sonnet-20250219-v1:0`, returning 404s with "This model version has reached the end of its life." Update test references to `claude-sonnet-4-5-20250929-v1:0` (same capability surface: thinking, tools, prompt caching, PDF input, vision, computer use). The bedrock/invoke pass-through tests stay on Sonnet 3.5 since Sonnet 4.5 is converse-only on Bedrock. --- .../litellm_utils_tests/test_health_check.py | 4 +-- tests/litellm_utils_tests/test_utils.py | 6 ++-- .../test_bedrock_anthropic_regression.py | 12 ++++---- .../test_bedrock_completion.py | 10 +++---- .../test_litellm_proxy_provider.py | 2 +- tests/llm_translation/test_optional_params.py | 2 +- tests/local_testing/test_function_calling.py | 2 +- ..._anthropic_messages_prompt_caching_test.py | 4 +-- .../test_anthropic_messages_prompt_caching.py | 4 +-- .../open_telemetry/data/captured_kwargs.json | 2 +- .../data/captured_response.json | 2 +- .../test_anthropic_cache_control_hook.py | 22 +++++++-------- .../integrations/test_opentelemetry.py | 2 +- ...llm_core_utils_prompt_templates_factory.py | 2 +- .../chat/test_converse_transformation.py | 12 ++++---- tests/test_litellm/test_utils.py | 28 +++++++++---------- 16 files changed, 58 insertions(+), 58 deletions(-) diff --git a/tests/litellm_utils_tests/test_health_check.py b/tests/litellm_utils_tests/test_health_check.py index a41907722e..45c6a04ad5 100644 --- a/tests/litellm_utils_tests/test_health_check.py +++ b/tests/litellm_utils_tests/test_health_check.py @@ -314,11 +314,11 @@ def test_update_litellm_params_for_health_check(): # Issue #15807: Fixes health checks sending "region/model" as model ID to AWS model_info = {} litellm_params = { - "model": "bedrock/us-gov-west-1/anthropic.claude-3-7-sonnet-20250219-v1:0", + "model": "bedrock/us-gov-west-1/anthropic.claude-sonnet-4-5-20250929-v1:0", "api_key": "fake_key", } updated_params = _update_litellm_params_for_health_check(model_info, litellm_params) - assert updated_params["model"] == "anthropic.claude-3-7-sonnet-20250219-v1:0" + assert updated_params["model"] == "anthropic.claude-sonnet-4-5-20250929-v1:0" # Test with Bedrock cross-region inference profile - should preserve the inference profile prefix # AWS requires inference profile IDs like "us.anthropic.claude..." for cross-region routing diff --git a/tests/litellm_utils_tests/test_utils.py b/tests/litellm_utils_tests/test_utils.py index d5df4ef75a..20af6e1023 100644 --- a/tests/litellm_utils_tests/test_utils.py +++ b/tests/litellm_utils_tests/test_utils.py @@ -2309,11 +2309,11 @@ def test_get_provider_audio_transcription_config(): @pytest.mark.parametrize( "model, expected_bool", [ - ("anthropic.claude-3-7-sonnet-20250219-v1:0", True), - ("us.anthropic.claude-3-7-sonnet-20250219-v1:0", True), + ("anthropic.claude-sonnet-4-5-20250929-v1:0", True), + ("us.anthropic.claude-sonnet-4-5-20250929-v1:0", True), ], ) -def test_claude_3_7_sonnet_supports_pdf_input(model, expected_bool): +def test_claude_sonnet_4_5_supports_pdf_input(model, expected_bool): from litellm.utils import supports_pdf_input assert supports_pdf_input(model) == expected_bool diff --git a/tests/llm_translation/test_bedrock_anthropic_regression.py b/tests/llm_translation/test_bedrock_anthropic_regression.py index 5928ca0223..8b8ce0a6cc 100644 --- a/tests/llm_translation/test_bedrock_anthropic_regression.py +++ b/tests/llm_translation/test_bedrock_anthropic_regression.py @@ -134,7 +134,7 @@ class TestBedrockAnthropicPromptCachingRegression: if "converse" in model_prefix: config = AmazonConverseConfig() result = config.transform_request( - model="us.anthropic.claude-3-7-sonnet-20250219-v1:0", + model="us.anthropic.claude-sonnet-4-5-20250929-v1:0", messages=messages, optional_params={}, litellm_params={}, @@ -162,7 +162,7 @@ class TestBedrockAnthropicPromptCachingRegression: else: config = AmazonAnthropicClaudeConfig() result = config.transform_request( - model="us.anthropic.claude-3-7-sonnet-20250219-v1:0", + model="us.anthropic.claude-sonnet-4-5-20250929-v1:0", messages=messages, optional_params={}, litellm_params={}, @@ -227,7 +227,7 @@ class TestBedrockAnthropicPromptCachingRegression: if "converse" in model_prefix: config = AmazonConverseConfig() result = config._transform_request_helper( - model="us.anthropic.claude-3-7-sonnet-20250219-v1:0", + model="us.anthropic.claude-sonnet-4-5-20250929-v1:0", system_content_blocks=[], optional_params={}, messages=messages, @@ -236,7 +236,7 @@ class TestBedrockAnthropicPromptCachingRegression: else: config = AmazonAnthropicClaudeConfig() result = config.transform_request( - model="us.anthropic.claude-3-7-sonnet-20250219-v1:0", + model="us.anthropic.claude-sonnet-4-5-20250929-v1:0", messages=messages, optional_params={}, litellm_params={}, @@ -498,7 +498,7 @@ class TestBedrockAnthropicCombinedRegressions: if "converse" in model_prefix: config = AmazonConverseConfig() result = config._transform_request_helper( - model="us.anthropic.claude-3-7-sonnet-20250219-v1:0", + model="us.anthropic.claude-sonnet-4-5-20250929-v1:0", system_content_blocks=[], optional_params={}, messages=messages, @@ -518,7 +518,7 @@ class TestBedrockAnthropicCombinedRegressions: else: config = AmazonAnthropicClaudeConfig() result = config.transform_request( - model="us.anthropic.claude-3-7-sonnet-20250219-v1:0", + model="us.anthropic.claude-sonnet-4-5-20250929-v1:0", messages=messages, optional_params={}, litellm_params={}, diff --git a/tests/llm_translation/test_bedrock_completion.py b/tests/llm_translation/test_bedrock_completion.py index ddfe383f2a..15f950224d 100644 --- a/tests/llm_translation/test_bedrock_completion.py +++ b/tests/llm_translation/test_bedrock_completion.py @@ -1323,7 +1323,7 @@ def test_base_aws_llm_get_credentials(): def test_bedrock_completion_test_2(): litellm.set_verbose = True data = { - "model": "bedrock/anthropic.claude-3-7-sonnet-20250219-v1:0", + "model": "bedrock/anthropic.claude-sonnet-4-5-20250929-v1:0", "messages": [ { "role": "system", @@ -1630,7 +1630,7 @@ def test_bedrock_completion_test_4(modify_params): litellm.modify_params = modify_params data = { - "model": "anthropic.claude-3-7-sonnet-20250219-v1:0", + "model": "anthropic.claude-sonnet-4-5-20250929-v1:0", "messages": [ { "role": "user", @@ -2115,7 +2115,7 @@ class TestBedrockConverseAnthropicUnitTests(BaseAnthropicChatTest): def get_base_completion_call_args_with_thinking(self) -> dict: return { - "model": "bedrock/us.anthropic.claude-3-7-sonnet-20250219-v1:0", + "model": "bedrock/us.anthropic.claude-sonnet-4-5-20250929-v1:0", "thinking": {"type": "enabled", "budget_tokens": 16000}, } @@ -2828,7 +2828,7 @@ async def test_bedrock_thinking_in_assistant_message(sync_mode): client = AsyncHTTPHandler() params = { - "model": "bedrock/us.anthropic.claude-3-7-sonnet-20250219-v1:0", + "model": "bedrock/us.anthropic.claude-sonnet-4-5-20250929-v1:0", "messages": [ { "role": "assistant", @@ -2887,7 +2887,7 @@ async def test_bedrock_stream_thinking_content_openwebui(): ``` """ response = await litellm.acompletion( - model="bedrock/us.anthropic.claude-3-7-sonnet-20250219-v1:0", + model="bedrock/us.anthropic.claude-sonnet-4-5-20250929-v1:0", messages=[{"role": "user", "content": "Hello who is this?"}], stream=True, max_tokens=1080, diff --git a/tests/llm_translation/test_litellm_proxy_provider.py b/tests/llm_translation/test_litellm_proxy_provider.py index 8fc961d12d..8b6f37bfbc 100644 --- a/tests/llm_translation/test_litellm_proxy_provider.py +++ b/tests/llm_translation/test_litellm_proxy_provider.py @@ -580,7 +580,7 @@ def test_litellm_gateway_from_sdk_with_response_cost_in_additional_headers(): def test_litellm_gateway_from_sdk_with_thinking_param(): try: response = litellm.completion( - model="litellm_proxy/anthropic.claude-3-7-sonnet-20250219-v1:0", + model="litellm_proxy/anthropic.claude-sonnet-4-5-20250929-v1:0", messages=[{"role": "user", "content": "Hello world"}], api_base="http://0.0.0.0:4000", api_key="sk-PIp1h0RekR", diff --git a/tests/llm_translation/test_optional_params.py b/tests/llm_translation/test_optional_params.py index 82a3d96b02..b40ce11bb9 100644 --- a/tests/llm_translation/test_optional_params.py +++ b/tests/llm_translation/test_optional_params.py @@ -1828,7 +1828,7 @@ def test_azure_response_format_param(): "model, provider", [ ("claude-3-7-sonnet-20240620-v1:0", "anthropic"), - ("anthropic.claude-3-7-sonnet-20250219-v1:0", "bedrock"), + ("anthropic.claude-sonnet-4-5-20250929-v1:0", "bedrock"), ("invoke/anthropic.claude-3-7-sonnet-20240620-v1:0", "bedrock"), ("claude-3-7-sonnet@20250219", "vertex_ai"), ], diff --git a/tests/local_testing/test_function_calling.py b/tests/local_testing/test_function_calling.py index b52805c066..02affa1d57 100644 --- a/tests/local_testing/test_function_calling.py +++ b/tests/local_testing/test_function_calling.py @@ -159,7 +159,7 @@ def test_aaparallel_function_call(model): "model", [ "anthropic/claude-4-sonnet-20250514", - "bedrock/us.anthropic.claude-3-7-sonnet-20250219-v1:0", + "bedrock/us.anthropic.claude-sonnet-4-5-20250929-v1:0", ], ) @pytest.mark.flaky(retries=3, delay=1) diff --git a/tests/pass_through_unit_tests/base_anthropic_messages_prompt_caching_test.py b/tests/pass_through_unit_tests/base_anthropic_messages_prompt_caching_test.py index 3c71af97c9..d6502afbe7 100644 --- a/tests/pass_through_unit_tests/base_anthropic_messages_prompt_caching_test.py +++ b/tests/pass_through_unit_tests/base_anthropic_messages_prompt_caching_test.py @@ -96,8 +96,8 @@ class BaseAnthropicMessagesPromptCachingTest(ABC): Returns the model string to use for tests. Examples: - - "bedrock/converse/anthropic.claude-3-7-sonnet-20250219-v1:0" - - "bedrock/invoke/anthropic.claude-3-7-sonnet-20250219-v1:0" + - "bedrock/converse/anthropic.claude-sonnet-4-5-20250929-v1:0" + - "bedrock/invoke/anthropic.claude-3-5-sonnet-20241022-v2:0" """ pass diff --git a/tests/pass_through_unit_tests/test_anthropic_messages_prompt_caching.py b/tests/pass_through_unit_tests/test_anthropic_messages_prompt_caching.py index 83a47a0149..bfdbf75351 100644 --- a/tests/pass_through_unit_tests/test_anthropic_messages_prompt_caching.py +++ b/tests/pass_through_unit_tests/test_anthropic_messages_prompt_caching.py @@ -31,7 +31,7 @@ class TestBedrockConversePromptCaching(BaseAnthropicMessagesPromptCachingTest): """ def get_model(self) -> str: - return "bedrock/converse/us.anthropic.claude-3-7-sonnet-20250219-v1:0" + return "bedrock/converse/us.anthropic.claude-sonnet-4-5-20250929-v1:0" class TestBedrockInvokePromptCaching(BaseAnthropicMessagesPromptCachingTest): @@ -43,4 +43,4 @@ class TestBedrockInvokePromptCaching(BaseAnthropicMessagesPromptCachingTest): """ def get_model(self) -> str: - return "bedrock/invoke/us.anthropic.claude-3-7-sonnet-20250219-v1:0" + return "bedrock/invoke/us.anthropic.claude-3-5-sonnet-20241022-v2:0" diff --git a/tests/test_litellm/integrations/open_telemetry/data/captured_kwargs.json b/tests/test_litellm/integrations/open_telemetry/data/captured_kwargs.json index 913e3bfeda..818e4fa3ea 100644 --- a/tests/test_litellm/integrations/open_telemetry/data/captured_kwargs.json +++ b/tests/test_litellm/integrations/open_telemetry/data/captured_kwargs.json @@ -1 +1 @@ -{"litellm_trace_id": null, "litellm_call_id": "dbecd23a-e71a-49cf-90d4-712a8a8e29c5", "input": [{"role": "user", "content": "What is the capital of France?"}], "litellm_params": {"acompletion": true, "api_key": null, "force_timeout": 600, "logger_fn": null, "verbose": false, "custom_llm_provider": "bedrock", "api_base": "https://bedrock-runtime.us-west-2.amazonaws.com/model/arn%3Aaws%3Abedrock%3Aus-west-2%3A1234567890123%3Ainference-profile%2Fus.anthropic.claude-3-7-sonnet-20250219-v1%3A0/converse", "litellm_call_id": "dbecd23a-e71a-49cf-90d4-712a8a8e29c5", "model_alias_map": {}, "completion_call_id": null, "aembedding": null, "metadata": {"requester_metadata": {}, "user_api_key_hash": "unused-for-aws-bedrock", "user_api_key_alias": null, "user_api_key_team_id": null, "user_api_key_user_id": null, "user_api_key_org_id": null, "user_api_key_team_alias": null, "user_api_key_end_user_id": null, "user_api_key_user_email": null, "user_api_key": "unused-for-aws-bedrock", "user_api_end_user_max_budget": null, "litellm_api_version": "1.72.3", "global_max_parallel_requests": null, "user_api_key_team_max_budget": null, "user_api_key_team_spend": null, "user_api_key_spend": 0.0, "user_api_key_max_budget": null, "user_api_key_model_max_budget": {}, "user_api_key_metadata": {}, "headers": {"host": "0.0.0.0:44444", "accept-encoding": "gzip, deflate, zstd", "connection": "keep-alive", "accept": "application/json", "content-type": "application/json", "user-agent": "AsyncOpenAI/Python 1.84.0", "x-stainless-lang": "python", "x-stainless-package-version": "1.84.0", "x-stainless-os": "MacOS", "x-stainless-arch": "arm64", "x-stainless-runtime": "CPython", "x-stainless-runtime-version": "3.12.10", "x-stainless-async": "async:asyncio", "x-stainless-retry-count": "0", "x-stainless-read-timeout": "600", "content-length": "116"}, "endpoint": "http://0.0.0.0:44444/chat/completions", "litellm_parent_otel_span": null, "requester_ip_address": "", "model_group": "claude-3-7-sonnet", "model_group_size": 1, "deployment": "bedrock/arn:aws:bedrock:us-west-2:1234567890123:inference-profile/us.anthropic.claude-3-7-sonnet-20250219-v1:0", "model_info": {"id": "6bace4d6db0105943b3b0bfe7eb1a62c06e6f16f008cc4673fdf918eb3e9e62a", "db_model": false}, "api_base": null, "caching_groups": null, "hidden_params": {"custom_llm_provider": "bedrock", "region_name": null, "optional_params": {"stream": false, "max_retries": 0, "provider": "aws", "region": "us-west-2"}, "litellm_call_id": "dbecd23a-e71a-49cf-90d4-712a8a8e29c5", "api_base": null, "model_id": "6bace4d6db0105943b3b0bfe7eb1a62c06e6f16f008cc4673fdf918eb3e9e62a", "response_cost": 0.001047, "additional_headers": {"x-litellm-model-group": "claude-3-7-sonnet", "x-litellm-attempted-retries": 0, "x-litellm-attempted-fallbacks": 0}, "litellm_model_name": "bedrock/arn:aws:bedrock:us-west-2:1234567890123:inference-profile/us.anthropic.claude-3-7-sonnet-20250219-v1:0", "litellm_overhead_time_ms": 231.156, "_response_ms": 236.798}}, "model_info": {"id": "6bace4d6db0105943b3b0bfe7eb1a62c06e6f16f008cc4673fdf918eb3e9e62a", "db_model": false}, "proxy_server_request": {"url": "http://0.0.0.0:44444/chat/completions", "method": "POST", "headers": {"host": "0.0.0.0:44444", "accept-encoding": "gzip, deflate, zstd", "connection": "keep-alive", "accept": "application/json", "content-type": "application/json", "user-agent": "AsyncOpenAI/Python 1.84.0", "x-stainless-lang": "python", "x-stainless-package-version": "1.84.0", "x-stainless-os": "MacOS", "x-stainless-arch": "arm64", "x-stainless-runtime": "CPython", "x-stainless-runtime-version": "3.12.10", "x-stainless-async": "async:asyncio", "x-stainless-retry-count": "0", "x-stainless-read-timeout": "600", "content-length": "116"}, "body": {"messages": [{"role": "user", "content": "What is the capital of France?"}], "model": "claude-3-7-sonnet", "stream": false}}, "preset_cache_key": null, "no-log": null, "stream_response": {}, "input_cost_per_token": null, "input_cost_per_second": null, "output_cost_per_token": null, "output_cost_per_second": null, "cooldown_time": null, "text_completion": null, "azure_ad_token_provider": null, "user_continue_message": null, "base_model": null, "litellm_trace_id": "4c97150b-b1a3-4dec-bd7a-734786b1b3bc", "litellm_session_id": null, "hf_model_name": null, "custom_prompt_dict": {}, "litellm_metadata": null, "disable_add_transform_inline_image_block": null, "drop_params": null, "prompt_id": null, "prompt_variables": null, "async_call": null, "ssl_verify": null, "merge_reasoning_content_in_choices": false, "api_version": null, "azure_ad_token": null, "tenant_id": null, "client_id": null, "client_secret": null, "azure_username": null, "azure_password": null, "max_retries": 0, "timeout": 6000.0, "bucket_name": null, "vertex_credentials": null, "vertex_project": null, "use_litellm_proxy": false}, "applied_guardrails": [], "model": "arn:aws:bedrock:us-west-2:1234567890123:inference-profile/us.anthropic.claude-3-7-sonnet-20250219-v1:0", "messages": [{"role": "user", "content": "What is the capital of France?"}], "optional_params": {"stream": false, "max_retries": 0, "provider": "aws", "region": "us-west-2"}, "start_time": "2025-06-22 10:59:08.159939", "stream": false, "user": null, "call_type": "acompletion", "completion_start_time": "2025-06-22 10:59:08.399523", "standard_callback_dynamic_params": {}, "stream_options": null, "max_retries": 0, "provider": "aws", "region": "us-west-2", "custom_llm_provider": "bedrock", "api_key": "", "additional_args": {"complete_input_dict": "{\"messages\": [{\"role\": \"user\", \"content\": [{\"text\": \"What is the capital of France?\"}]}], \"additionalModelRequestFields\": {\"provider\": \"aws\", \"region\": \"us-west-2\"}, \"system\": [], \"inferenceConfig\": {}}"}, "log_event_type": "post_api_call", "api_call_start_time": "2025-06-22 10:59:08.387641", "llm_api_duration_ms": 5.642, "original_response": "{\"metrics\":{\"latencyMs\":1513},\"output\":{\"message\":{\"content\":[{\"text\":\"The capital of France is Paris. Paris has been the capital city of France since 987 CE when Hugh Capet, the first king of the Capetian dynasty, made the city his seat of government. Today, Paris is not only the political capital but also the cultural and economic center of France.\"}],\"role\":\"assistant\"}},\"stopReason\":\"end_turn\",\"usage\":{\"cacheReadInputTokenCount\":0,\"cacheReadInputTokens\":0,\"cacheWriteInputTokenCount\":0,\"cacheWriteInputTokens\":0,\"inputTokens\":14,\"outputTokens\":67,\"totalTokens\":81}}", "end_time": "2025-06-22 10:59:08.399523", "cache_hit": null, "response_cost": 0.001047, "standard_logging_object": {"id": "chatcmpl-fa9be5b7-9487-46ab-86de-6462d578fea1", "trace_id": "4c97150b-b1a3-4dec-bd7a-734786b1b3bc", "call_type": "acompletion", "cache_hit": null, "stream": true, "status": "success", "custom_llm_provider": "bedrock", "saved_cache_cost": 0.0, "startTime": 1750615148.162725, "endTime": 1750615148.399523, "completionStartTime": 1750615148.399523, "response_time": 0.23679804801940918, "model": "arn:aws:bedrock:us-west-2:1234567890123:inference-profile/us.anthropic.claude-3-7-sonnet-20250219-v1:0", "metadata": {"user_api_key_hash": "unused-for-aws-bedrock", "user_api_key_alias": null, "user_api_key_team_id": null, "user_api_key_org_id": null, "user_api_key_user_id": null, "user_api_key_team_alias": null, "user_api_key_user_email": null, "spend_logs_metadata": null, "requester_ip_address": "", "requester_metadata": {}, "user_api_key_end_user_id": null, "prompt_management_metadata": null, "applied_guardrails": [], "mcp_tool_call_metadata": null, "vector_store_request_metadata": null, "usage_object": {"completion_tokens": 67, "prompt_tokens": 14, "total_tokens": 81, "completion_tokens_details": null, "prompt_tokens_details": {"audio_tokens": null, "cached_tokens": 0}, "cache_creation_input_tokens": 0, "cache_read_input_tokens": 0}, "requester_custom_headers": {"x-stainless-lang": "python", "x-stainless-package-version": "1.84.0", "x-stainless-os": "MacOS", "x-stainless-arch": "arm64", "x-stainless-runtime": "CPython", "x-stainless-runtime-version": "3.12.10", "x-stainless-async": "async:asyncio", "x-stainless-retry-count": "0", "x-stainless-read-timeout": "600"}}, "cache_key": null, "response_cost": 0.001047, "total_tokens": 81, "prompt_tokens": 14, "completion_tokens": 67, "request_tags": [], "end_user": "", "api_base": "https://bedrock-runtime.us-west-2.amazonaws.com/model/arn%3Aaws%3Abedrock%3Aus-west-2%3A1234567890123%3Ainference-profile%2Fus.anthropic.claude-3-7-sonnet-20250219-v1%3A0/converse", "model_group": "claude-3-7-sonnet", "model_id": "6bace4d6db0105943b3b0bfe7eb1a62c06e6f16f008cc4673fdf918eb3e9e62a", "requester_ip_address": "", "messages": [{"role": "user", "content": "What is the capital of France?"}], "response": {"id": "chatcmpl-fa9be5b7-9487-46ab-86de-6462d578fea1", "created": 1750615148, "model": "arn:aws:bedrock:us-west-2:1234567890123:inference-profile/us.anthropic.claude-3-7-sonnet-20250219-v1:0", "object": "chat.completion", "system_fingerprint": null, "choices": [{"finish_reason": "stop", "index": 0, "message": {"content": "The capital of France is Paris. Paris has been the capital city of France since 987 CE when Hugh Capet, the first king of the Capetian dynasty, made the city his seat of government. Today, Paris is not only the political capital but also the cultural and economic center of France.", "role": "assistant", "tool_calls": null, "function_call": null}}], "usage": {"completion_tokens": 67, "prompt_tokens": 14, "total_tokens": 81, "completion_tokens_details": null, "prompt_tokens_details": {"audio_tokens": null, "cached_tokens": 0}, "cache_creation_input_tokens": 0, "cache_read_input_tokens": 0}}, "model_parameters": {"stream": false}, "hidden_params": {"model_id": "6bace4d6db0105943b3b0bfe7eb1a62c06e6f16f008cc4673fdf918eb3e9e62a", "cache_key": null, "api_base": null, "response_cost": 0.001047, "additional_headers": {}, "litellm_overhead_time_ms": 231.156, "batch_models": null, "litellm_model_name": "bedrock/arn:aws:bedrock:us-west-2:1234567890123:inference-profile/us.anthropic.claude-3-7-sonnet-20250219-v1:0", "usage_object": null}, "model_map_information": {"model_map_key": "arn:aws:bedrock:us-west-2:1234567890123:inference-profile/us.anthropic.claude-3-7-sonnet-20250219-v1:0", "model_map_value": {"key": "anthropic.claude-3-7-sonnet-20250219-v1:0", "max_tokens": 8192, "max_input_tokens": 200000, "max_output_tokens": 8192, "input_cost_per_token": 3e-06, "cache_creation_input_token_cost": 3.75e-06, "cache_read_input_token_cost": 3e-07, "input_cost_per_character": null, "input_cost_per_token_above_128k_tokens": null, "input_cost_per_token_above_200k_tokens": null, "input_cost_per_query": null, "input_cost_per_second": null, "input_cost_per_audio_token": null, "input_cost_per_token_batches": null, "output_cost_per_token_batches": null, "output_cost_per_token": 1.5e-05, "output_cost_per_audio_token": null, "output_cost_per_character": null, "output_cost_per_reasoning_token": null, "output_cost_per_token_above_128k_tokens": null, "output_cost_per_character_above_128k_tokens": null, "output_cost_per_token_above_200k_tokens": null, "output_cost_per_second": null, "output_cost_per_image": null, "output_vector_size": null, "litellm_provider": "bedrock_converse", "mode": "chat", "supports_system_messages": null, "supports_response_schema": true, "supports_vision": true, "supports_function_calling": true, "supports_tool_choice": true, "supports_assistant_prefill": true, "supports_prompt_caching": true, "supports_audio_input": null, "supports_audio_output": null, "supports_pdf_input": true, "supports_embedding_image_input": null, "supports_native_streaming": null, "supports_web_search": null, "supports_url_context": null, "supports_reasoning": true, "supports_computer_use": true, "search_context_cost_per_query": null, "tpm": null, "rpm": null, "supported_openai_params": ["max_tokens", "max_completion_tokens", "stream", "stream_options", "stop", "temperature", "top_p", "extra_headers", "response_format", "tools", "tool_choice", "thinking", "reasoning_effort"]}}, "error_str": null, "error_information": {"error_code": "", "error_class": "", "llm_provider": "", "traceback": "", "error_message": ""}, "response_cost_failure_debug_info": null, "guardrail_information": null, "standard_built_in_tools_params": {"web_search_options": null, "file_search": null}}, "async_complete_streaming_response": "ModelResponse(id='chatcmpl-fa9be5b7-9487-46ab-86de-6462d578fea1', created=1750615148, model='arn:aws:bedrock:us-west-2:1234567890123:inference-profile/us.anthropic.claude-3-7-sonnet-20250219-v1:0', object='chat.completion', system_fingerprint=None, choices=[Choices(finish_reason='stop', index=0, message=Message(content='The capital of France is Paris. Paris has been the capital city of France since 987 CE when Hugh Capet, the first king of the Capetian dynasty, made the city his seat of government. Today, Paris is not only the political capital but also the cultural and economic center of France.', role='assistant', tool_calls=None, function_call=None, provider_specific_fields=None))], usage=Usage(completion_tokens=67, prompt_tokens=14, total_tokens=81, completion_tokens_details=None, prompt_tokens_details=PromptTokensDetailsWrapper(audio_tokens=None, cached_tokens=0, text_tokens=None, image_tokens=None), cache_creation_input_tokens=0, cache_read_input_tokens=0))"} \ No newline at end of file +{"litellm_trace_id": null, "litellm_call_id": "dbecd23a-e71a-49cf-90d4-712a8a8e29c5", "input": [{"role": "user", "content": "What is the capital of France?"}], "litellm_params": {"acompletion": true, "api_key": null, "force_timeout": 600, "logger_fn": null, "verbose": false, "custom_llm_provider": "bedrock", "api_base": "https://bedrock-runtime.us-west-2.amazonaws.com/model/arn%3Aaws%3Abedrock%3Aus-west-2%3A1234567890123%3Ainference-profile%2Fus.anthropic.claude-sonnet-4-5-20250929-v1%3A0/converse", "litellm_call_id": "dbecd23a-e71a-49cf-90d4-712a8a8e29c5", "model_alias_map": {}, "completion_call_id": null, "aembedding": null, "metadata": {"requester_metadata": {}, "user_api_key_hash": "unused-for-aws-bedrock", "user_api_key_alias": null, "user_api_key_team_id": null, "user_api_key_user_id": null, "user_api_key_org_id": null, "user_api_key_team_alias": null, "user_api_key_end_user_id": null, "user_api_key_user_email": null, "user_api_key": "unused-for-aws-bedrock", "user_api_end_user_max_budget": null, "litellm_api_version": "1.72.3", "global_max_parallel_requests": null, "user_api_key_team_max_budget": null, "user_api_key_team_spend": null, "user_api_key_spend": 0.0, "user_api_key_max_budget": null, "user_api_key_model_max_budget": {}, "user_api_key_metadata": {}, "headers": {"host": "0.0.0.0:44444", "accept-encoding": "gzip, deflate, zstd", "connection": "keep-alive", "accept": "application/json", "content-type": "application/json", "user-agent": "AsyncOpenAI/Python 1.84.0", "x-stainless-lang": "python", "x-stainless-package-version": "1.84.0", "x-stainless-os": "MacOS", "x-stainless-arch": "arm64", "x-stainless-runtime": "CPython", "x-stainless-runtime-version": "3.12.10", "x-stainless-async": "async:asyncio", "x-stainless-retry-count": "0", "x-stainless-read-timeout": "600", "content-length": "116"}, "endpoint": "http://0.0.0.0:44444/chat/completions", "litellm_parent_otel_span": null, "requester_ip_address": "", "model_group": "claude-3-7-sonnet", "model_group_size": 1, "deployment": "bedrock/arn:aws:bedrock:us-west-2:1234567890123:inference-profile/us.anthropic.claude-sonnet-4-5-20250929-v1:0", "model_info": {"id": "6bace4d6db0105943b3b0bfe7eb1a62c06e6f16f008cc4673fdf918eb3e9e62a", "db_model": false}, "api_base": null, "caching_groups": null, "hidden_params": {"custom_llm_provider": "bedrock", "region_name": null, "optional_params": {"stream": false, "max_retries": 0, "provider": "aws", "region": "us-west-2"}, "litellm_call_id": "dbecd23a-e71a-49cf-90d4-712a8a8e29c5", "api_base": null, "model_id": "6bace4d6db0105943b3b0bfe7eb1a62c06e6f16f008cc4673fdf918eb3e9e62a", "response_cost": 0.001047, "additional_headers": {"x-litellm-model-group": "claude-3-7-sonnet", "x-litellm-attempted-retries": 0, "x-litellm-attempted-fallbacks": 0}, "litellm_model_name": "bedrock/arn:aws:bedrock:us-west-2:1234567890123:inference-profile/us.anthropic.claude-sonnet-4-5-20250929-v1:0", "litellm_overhead_time_ms": 231.156, "_response_ms": 236.798}}, "model_info": {"id": "6bace4d6db0105943b3b0bfe7eb1a62c06e6f16f008cc4673fdf918eb3e9e62a", "db_model": false}, "proxy_server_request": {"url": "http://0.0.0.0:44444/chat/completions", "method": "POST", "headers": {"host": "0.0.0.0:44444", "accept-encoding": "gzip, deflate, zstd", "connection": "keep-alive", "accept": "application/json", "content-type": "application/json", "user-agent": "AsyncOpenAI/Python 1.84.0", "x-stainless-lang": "python", "x-stainless-package-version": "1.84.0", "x-stainless-os": "MacOS", "x-stainless-arch": "arm64", "x-stainless-runtime": "CPython", "x-stainless-runtime-version": "3.12.10", "x-stainless-async": "async:asyncio", "x-stainless-retry-count": "0", "x-stainless-read-timeout": "600", "content-length": "116"}, "body": {"messages": [{"role": "user", "content": "What is the capital of France?"}], "model": "claude-3-7-sonnet", "stream": false}}, "preset_cache_key": null, "no-log": null, "stream_response": {}, "input_cost_per_token": null, "input_cost_per_second": null, "output_cost_per_token": null, "output_cost_per_second": null, "cooldown_time": null, "text_completion": null, "azure_ad_token_provider": null, "user_continue_message": null, "base_model": null, "litellm_trace_id": "4c97150b-b1a3-4dec-bd7a-734786b1b3bc", "litellm_session_id": null, "hf_model_name": null, "custom_prompt_dict": {}, "litellm_metadata": null, "disable_add_transform_inline_image_block": null, "drop_params": null, "prompt_id": null, "prompt_variables": null, "async_call": null, "ssl_verify": null, "merge_reasoning_content_in_choices": false, "api_version": null, "azure_ad_token": null, "tenant_id": null, "client_id": null, "client_secret": null, "azure_username": null, "azure_password": null, "max_retries": 0, "timeout": 6000.0, "bucket_name": null, "vertex_credentials": null, "vertex_project": null, "use_litellm_proxy": false}, "applied_guardrails": [], "model": "arn:aws:bedrock:us-west-2:1234567890123:inference-profile/us.anthropic.claude-sonnet-4-5-20250929-v1:0", "messages": [{"role": "user", "content": "What is the capital of France?"}], "optional_params": {"stream": false, "max_retries": 0, "provider": "aws", "region": "us-west-2"}, "start_time": "2025-06-22 10:59:08.159939", "stream": false, "user": null, "call_type": "acompletion", "completion_start_time": "2025-06-22 10:59:08.399523", "standard_callback_dynamic_params": {}, "stream_options": null, "max_retries": 0, "provider": "aws", "region": "us-west-2", "custom_llm_provider": "bedrock", "api_key": "", "additional_args": {"complete_input_dict": "{\"messages\": [{\"role\": \"user\", \"content\": [{\"text\": \"What is the capital of France?\"}]}], \"additionalModelRequestFields\": {\"provider\": \"aws\", \"region\": \"us-west-2\"}, \"system\": [], \"inferenceConfig\": {}}"}, "log_event_type": "post_api_call", "api_call_start_time": "2025-06-22 10:59:08.387641", "llm_api_duration_ms": 5.642, "original_response": "{\"metrics\":{\"latencyMs\":1513},\"output\":{\"message\":{\"content\":[{\"text\":\"The capital of France is Paris. Paris has been the capital city of France since 987 CE when Hugh Capet, the first king of the Capetian dynasty, made the city his seat of government. Today, Paris is not only the political capital but also the cultural and economic center of France.\"}],\"role\":\"assistant\"}},\"stopReason\":\"end_turn\",\"usage\":{\"cacheReadInputTokenCount\":0,\"cacheReadInputTokens\":0,\"cacheWriteInputTokenCount\":0,\"cacheWriteInputTokens\":0,\"inputTokens\":14,\"outputTokens\":67,\"totalTokens\":81}}", "end_time": "2025-06-22 10:59:08.399523", "cache_hit": null, "response_cost": 0.001047, "standard_logging_object": {"id": "chatcmpl-fa9be5b7-9487-46ab-86de-6462d578fea1", "trace_id": "4c97150b-b1a3-4dec-bd7a-734786b1b3bc", "call_type": "acompletion", "cache_hit": null, "stream": true, "status": "success", "custom_llm_provider": "bedrock", "saved_cache_cost": 0.0, "startTime": 1750615148.162725, "endTime": 1750615148.399523, "completionStartTime": 1750615148.399523, "response_time": 0.23679804801940918, "model": "arn:aws:bedrock:us-west-2:1234567890123:inference-profile/us.anthropic.claude-sonnet-4-5-20250929-v1:0", "metadata": {"user_api_key_hash": "unused-for-aws-bedrock", "user_api_key_alias": null, "user_api_key_team_id": null, "user_api_key_org_id": null, "user_api_key_user_id": null, "user_api_key_team_alias": null, "user_api_key_user_email": null, "spend_logs_metadata": null, "requester_ip_address": "", "requester_metadata": {}, "user_api_key_end_user_id": null, "prompt_management_metadata": null, "applied_guardrails": [], "mcp_tool_call_metadata": null, "vector_store_request_metadata": null, "usage_object": {"completion_tokens": 67, "prompt_tokens": 14, "total_tokens": 81, "completion_tokens_details": null, "prompt_tokens_details": {"audio_tokens": null, "cached_tokens": 0}, "cache_creation_input_tokens": 0, "cache_read_input_tokens": 0}, "requester_custom_headers": {"x-stainless-lang": "python", "x-stainless-package-version": "1.84.0", "x-stainless-os": "MacOS", "x-stainless-arch": "arm64", "x-stainless-runtime": "CPython", "x-stainless-runtime-version": "3.12.10", "x-stainless-async": "async:asyncio", "x-stainless-retry-count": "0", "x-stainless-read-timeout": "600"}}, "cache_key": null, "response_cost": 0.001047, "total_tokens": 81, "prompt_tokens": 14, "completion_tokens": 67, "request_tags": [], "end_user": "", "api_base": "https://bedrock-runtime.us-west-2.amazonaws.com/model/arn%3Aaws%3Abedrock%3Aus-west-2%3A1234567890123%3Ainference-profile%2Fus.anthropic.claude-sonnet-4-5-20250929-v1%3A0/converse", "model_group": "claude-3-7-sonnet", "model_id": "6bace4d6db0105943b3b0bfe7eb1a62c06e6f16f008cc4673fdf918eb3e9e62a", "requester_ip_address": "", "messages": [{"role": "user", "content": "What is the capital of France?"}], "response": {"id": "chatcmpl-fa9be5b7-9487-46ab-86de-6462d578fea1", "created": 1750615148, "model": "arn:aws:bedrock:us-west-2:1234567890123:inference-profile/us.anthropic.claude-sonnet-4-5-20250929-v1:0", "object": "chat.completion", "system_fingerprint": null, "choices": [{"finish_reason": "stop", "index": 0, "message": {"content": "The capital of France is Paris. Paris has been the capital city of France since 987 CE when Hugh Capet, the first king of the Capetian dynasty, made the city his seat of government. Today, Paris is not only the political capital but also the cultural and economic center of France.", "role": "assistant", "tool_calls": null, "function_call": null}}], "usage": {"completion_tokens": 67, "prompt_tokens": 14, "total_tokens": 81, "completion_tokens_details": null, "prompt_tokens_details": {"audio_tokens": null, "cached_tokens": 0}, "cache_creation_input_tokens": 0, "cache_read_input_tokens": 0}}, "model_parameters": {"stream": false}, "hidden_params": {"model_id": "6bace4d6db0105943b3b0bfe7eb1a62c06e6f16f008cc4673fdf918eb3e9e62a", "cache_key": null, "api_base": null, "response_cost": 0.001047, "additional_headers": {}, "litellm_overhead_time_ms": 231.156, "batch_models": null, "litellm_model_name": "bedrock/arn:aws:bedrock:us-west-2:1234567890123:inference-profile/us.anthropic.claude-sonnet-4-5-20250929-v1:0", "usage_object": null}, "model_map_information": {"model_map_key": "arn:aws:bedrock:us-west-2:1234567890123:inference-profile/us.anthropic.claude-sonnet-4-5-20250929-v1:0", "model_map_value": {"key": "anthropic.claude-sonnet-4-5-20250929-v1:0", "max_tokens": 8192, "max_input_tokens": 200000, "max_output_tokens": 8192, "input_cost_per_token": 3e-06, "cache_creation_input_token_cost": 3.75e-06, "cache_read_input_token_cost": 3e-07, "input_cost_per_character": null, "input_cost_per_token_above_128k_tokens": null, "input_cost_per_token_above_200k_tokens": null, "input_cost_per_query": null, "input_cost_per_second": null, "input_cost_per_audio_token": null, "input_cost_per_token_batches": null, "output_cost_per_token_batches": null, "output_cost_per_token": 1.5e-05, "output_cost_per_audio_token": null, "output_cost_per_character": null, "output_cost_per_reasoning_token": null, "output_cost_per_token_above_128k_tokens": null, "output_cost_per_character_above_128k_tokens": null, "output_cost_per_token_above_200k_tokens": null, "output_cost_per_second": null, "output_cost_per_image": null, "output_vector_size": null, "litellm_provider": "bedrock_converse", "mode": "chat", "supports_system_messages": null, "supports_response_schema": true, "supports_vision": true, "supports_function_calling": true, "supports_tool_choice": true, "supports_assistant_prefill": true, "supports_prompt_caching": true, "supports_audio_input": null, "supports_audio_output": null, "supports_pdf_input": true, "supports_embedding_image_input": null, "supports_native_streaming": null, "supports_web_search": null, "supports_url_context": null, "supports_reasoning": true, "supports_computer_use": true, "search_context_cost_per_query": null, "tpm": null, "rpm": null, "supported_openai_params": ["max_tokens", "max_completion_tokens", "stream", "stream_options", "stop", "temperature", "top_p", "extra_headers", "response_format", "tools", "tool_choice", "thinking", "reasoning_effort"]}}, "error_str": null, "error_information": {"error_code": "", "error_class": "", "llm_provider": "", "traceback": "", "error_message": ""}, "response_cost_failure_debug_info": null, "guardrail_information": null, "standard_built_in_tools_params": {"web_search_options": null, "file_search": null}}, "async_complete_streaming_response": "ModelResponse(id='chatcmpl-fa9be5b7-9487-46ab-86de-6462d578fea1', created=1750615148, model='arn:aws:bedrock:us-west-2:1234567890123:inference-profile/us.anthropic.claude-sonnet-4-5-20250929-v1:0', object='chat.completion', system_fingerprint=None, choices=[Choices(finish_reason='stop', index=0, message=Message(content='The capital of France is Paris. Paris has been the capital city of France since 987 CE when Hugh Capet, the first king of the Capetian dynasty, made the city his seat of government. Today, Paris is not only the political capital but also the cultural and economic center of France.', role='assistant', tool_calls=None, function_call=None, provider_specific_fields=None))], usage=Usage(completion_tokens=67, prompt_tokens=14, total_tokens=81, completion_tokens_details=None, prompt_tokens_details=PromptTokensDetailsWrapper(audio_tokens=None, cached_tokens=0, text_tokens=None, image_tokens=None), cache_creation_input_tokens=0, cache_read_input_tokens=0))"} \ No newline at end of file diff --git a/tests/test_litellm/integrations/open_telemetry/data/captured_response.json b/tests/test_litellm/integrations/open_telemetry/data/captured_response.json index 3cf77781cc..1fa1889909 100644 --- a/tests/test_litellm/integrations/open_telemetry/data/captured_response.json +++ b/tests/test_litellm/integrations/open_telemetry/data/captured_response.json @@ -1 +1 @@ -{"id": "chatcmpl-fa9be5b7-9487-46ab-86de-6462d578fea1", "created": 1750615148, "model": "arn:aws:bedrock:us-west-2:1234567890123:inference-profile/us.anthropic.claude-3-7-sonnet-20250219-v1:0", "object": "chat.completion", "system_fingerprint": null, "choices": [{"finish_reason": "stop", "index": 0, "message": {"content": "The capital of France is Paris. Paris has been the capital city of France since 987 CE when Hugh Capet, the first king of the Capetian dynasty, made the city his seat of government. Today, Paris is not only the political capital but also the cultural and economic center of France.", "role": "assistant", "tool_calls": null, "function_call": null}}], "usage": {"completion_tokens": 67, "prompt_tokens": 14, "total_tokens": 81, "completion_tokens_details": null, "prompt_tokens_details": {"audio_tokens": null, "cached_tokens": 0}, "cache_creation_input_tokens": 0, "cache_read_input_tokens": 0}} \ No newline at end of file +{"id": "chatcmpl-fa9be5b7-9487-46ab-86de-6462d578fea1", "created": 1750615148, "model": "arn:aws:bedrock:us-west-2:1234567890123:inference-profile/us.anthropic.claude-sonnet-4-5-20250929-v1:0", "object": "chat.completion", "system_fingerprint": null, "choices": [{"finish_reason": "stop", "index": 0, "message": {"content": "The capital of France is Paris. Paris has been the capital city of France since 987 CE when Hugh Capet, the first king of the Capetian dynasty, made the city his seat of government. Today, Paris is not only the political capital but also the cultural and economic center of France.", "role": "assistant", "tool_calls": null, "function_call": null}}], "usage": {"completion_tokens": 67, "prompt_tokens": 14, "total_tokens": 81, "completion_tokens_details": null, "prompt_tokens_details": {"audio_tokens": null, "cached_tokens": 0}, "cache_creation_input_tokens": 0, "cache_read_input_tokens": 0}} \ No newline at end of file diff --git a/tests/test_litellm/integrations/test_anthropic_cache_control_hook.py b/tests/test_litellm/integrations/test_anthropic_cache_control_hook.py index 271d58061a..1a4d03528e 100644 --- a/tests/test_litellm/integrations/test_anthropic_cache_control_hook.py +++ b/tests/test_litellm/integrations/test_anthropic_cache_control_hook.py @@ -220,7 +220,7 @@ async def test_anthropic_cache_control_hook_negative_indices(): with patch.object(client, "post", return_value=mock_response) as mock_post: # Test with multiple messages and negative indices response = await litellm.acompletion( - model="bedrock/us.anthropic.claude-3-7-sonnet-20250219-v1:0", + model="bedrock/us.anthropic.claude-sonnet-4-5-20250929-v1:0", messages=[ { "role": "system", @@ -352,7 +352,7 @@ async def test_anthropic_cache_control_hook_out_of_bounds_logging(): ] await litellm.acompletion( - model="bedrock/us.anthropic.claude-3-7-sonnet-20250219-v1:0", + model="bedrock/us.anthropic.claude-sonnet-4-5-20250929-v1:0", messages=messages, cache_control_injection_points=[ {"location": "message", "index": 10} @@ -420,7 +420,7 @@ async def test_anthropic_cache_control_hook_negative_out_of_bounds_logging(): ] await litellm.acompletion( - model="bedrock/us.anthropic.claude-3-7-sonnet-20250219-v1:0", + model="bedrock/us.anthropic.claude-sonnet-4-5-20250929-v1:0", messages=messages, cache_control_injection_points=[ { @@ -486,7 +486,7 @@ async def test_anthropic_cache_control_hook_multiple_user_messages(): with patch.object(client, "post", return_value=mock_response) as mock_post: # Test with multiple user messages and negative indices response = await litellm.acompletion( - model="bedrock/us.anthropic.claude-3-7-sonnet-20250219-v1:0", + model="bedrock/us.anthropic.claude-sonnet-4-5-20250929-v1:0", messages=[ { "role": "user", @@ -586,7 +586,7 @@ async def test_anthropic_cache_control_hook_out_of_bounds(bad_index): ] await litellm.acompletion( - model="bedrock/us.anthropic.claude-3-7-sonnet-20250219-v1:0", + model="bedrock/us.anthropic.claude-sonnet-4-5-20250929-v1:0", messages=messages, cache_control_injection_points=[ {"location": "message", "index": bad_index} @@ -651,7 +651,7 @@ async def test_anthropic_cache_control_hook_single_message(message_list): client = AsyncHTTPHandler() with patch.object(client, "post", return_value=mock_response) as mock_post: await litellm.acompletion( - model="bedrock/us.anthropic.claude-3-7-sonnet-20250219-v1:0", + model="bedrock/us.anthropic.claude-sonnet-4-5-20250929-v1:0", messages=message_list, cache_control_injection_points=[{"location": "message", "index": -1}], client=client, @@ -691,7 +691,7 @@ async def test_anthropic_cache_control_hook_empty_message_list(): match="bedrock requires at least one non-system message", ): await litellm.acompletion( - model="bedrock/us.anthropic.claude-3-7-sonnet-20250219-v1:0", + model="bedrock/us.anthropic.claude-sonnet-4-5-20250929-v1:0", messages=[], cache_control_injection_points=[ {"location": "message", "index": -1} @@ -742,7 +742,7 @@ async def test_anthropic_cache_control_hook_no_op(): ] await litellm.acompletion( - model="bedrock/us.anthropic.claude-3-7-sonnet-20250219-v1:0", + model="bedrock/us.anthropic.claude-sonnet-4-5-20250929-v1:0", messages=messages, # No cache_control_injection_points parameter client=client, @@ -799,7 +799,7 @@ async def test_anthropic_cache_control_hook_multiple_content_items_last_only(): client = AsyncHTTPHandler() with patch.object(client, "post", return_value=mock_response) as mock_post: response = await litellm.acompletion( - model="bedrock/us.anthropic.claude-3-7-sonnet-20250219-v1:0", + model="bedrock/us.anthropic.claude-sonnet-4-5-20250929-v1:0", messages=[ { "role": "user", @@ -874,7 +874,7 @@ async def test_anthropic_cache_control_hook_document_analysis_multiple_pages(): client = AsyncHTTPHandler() with patch.object(client, "post", return_value=mock_response) as mock_post: response = await litellm.acompletion( - model="bedrock/us.anthropic.claude-3-7-sonnet-20250219-v1:0", + model="bedrock/us.anthropic.claude-sonnet-4-5-20250929-v1:0", messages=[ { "role": "user", @@ -1057,7 +1057,7 @@ async def test_anthropic_cache_control_hook_string_negative_index(): client = AsyncHTTPHandler() with patch.object(client, "post", return_value=mock_response) as mock_post: await litellm.acompletion( - model="bedrock/us.anthropic.claude-3-7-sonnet-20250219-v1:0", + model="bedrock/us.anthropic.claude-sonnet-4-5-20250929-v1:0", messages=[ {"role": "user", "content": "First message"}, {"role": "assistant", "content": "First response"}, diff --git a/tests/test_litellm/integrations/test_opentelemetry.py b/tests/test_litellm/integrations/test_opentelemetry.py index f710647189..b31bbca889 100644 --- a/tests/test_litellm/integrations/test_opentelemetry.py +++ b/tests/test_litellm/integrations/test_opentelemetry.py @@ -262,7 +262,7 @@ class TestOpenTelemetryProviderInitialization(unittest.TestCase): class TestOpenTelemetry(unittest.TestCase): POLL_INTERVAL = 0.05 POLL_TIMEOUT = 2.0 - MODEL = "arn:aws:bedrock:us-west-2:1234567890123:inference-profile/us.anthropic.claude-3-7-sonnet-20250219-v1:0" + MODEL = "arn:aws:bedrock:us-west-2:1234567890123:inference-profile/us.anthropic.claude-sonnet-4-5-20250929-v1:0" HERE = os.path.dirname(__file__) @patch.dict(os.environ, {}, clear=True) diff --git a/tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py b/tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py index 72cfd89408..f8708dd2f7 100644 --- a/tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py +++ b/tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py @@ -77,7 +77,7 @@ async def test_anthropic_bedrock_thinking_blocks_with_none_content(): # test _bedrock_converse_messages_pt_async result = await BedrockConverseMessagesProcessor._bedrock_converse_messages_pt_async( messages=messages, - model="us.anthropic.claude-3-7-sonnet-20250219-v1:0", + model="us.anthropic.claude-sonnet-4-5-20250929-v1:0", llm_provider="bedrock", ) diff --git a/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py b/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py index 38a59c694e..8e53e57f1e 100644 --- a/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py +++ b/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py @@ -279,7 +279,7 @@ def test_reasoning_with_forced_tool_choice_switches_to_auto(): } optional_params = config.map_openai_params( - model="bedrock/us.anthropic.claude-3-7-sonnet-20250219-v1:0", + model="bedrock/us.anthropic.claude-sonnet-4-5-20250929-v1:0", non_default_params=non_default_params, optional_params={}, drop_params=False, @@ -2797,7 +2797,7 @@ def test_thinking_with_max_completion_tokens(): result = config.map_openai_params( non_default_params=non_default_params_with_max_completion, optional_params=optional_params, - model="us.anthropic.claude-3-7-sonnet-20250219-v1:0", + model="us.anthropic.claude-sonnet-4-5-20250929-v1:0", drop_params=False, ) @@ -2819,7 +2819,7 @@ def test_thinking_with_max_completion_tokens(): result = config.map_openai_params( non_default_params=non_default_params_with_max_tokens, optional_params=optional_params, - model="us.anthropic.claude-3-7-sonnet-20250219-v1:0", + model="us.anthropic.claude-sonnet-4-5-20250929-v1:0", drop_params=False, ) @@ -2842,7 +2842,7 @@ def test_thinking_with_max_completion_tokens(): result = config.map_openai_params( non_default_params=non_default_params_without_max, optional_params=optional_params, - model="us.anthropic.claude-3-7-sonnet-20250219-v1:0", + model="us.anthropic.claude-sonnet-4-5-20250929-v1:0", drop_params=False, ) @@ -3617,7 +3617,7 @@ class TestBedrockMinThinkingBudgetTokens: """Test that thinking.budget_tokens is clamped to the Bedrock minimum (1024).""" def _map_params( - self, thinking_value, model="anthropic.claude-3-7-sonnet-20250219-v1:0" + self, thinking_value, model="anthropic.claude-sonnet-4-5-20250929-v1:0" ): """Helper to call map_openai_params with the given thinking value.""" config = AmazonConverseConfig() @@ -3651,7 +3651,7 @@ class TestBedrockMinThinkingBudgetTokens: result = config.map_openai_params( non_default_params={}, optional_params={}, - model="anthropic.claude-3-7-sonnet-20250219-v1:0", + model="anthropic.claude-sonnet-4-5-20250929-v1:0", drop_params=False, ) assert "thinking" not in result or result.get("thinking") is None diff --git a/tests/test_litellm/test_utils.py b/tests/test_litellm/test_utils.py index 93c61e003d..b8a4220c67 100644 --- a/tests/test_litellm/test_utils.py +++ b/tests/test_litellm/test_utils.py @@ -1179,7 +1179,7 @@ def test_get_model_info_shows_supports_computer_use(): "model, custom_llm_provider", [ ("gpt-3.5-turbo", "openai"), - ("anthropic.claude-3-7-sonnet-20250219-v1:0", "bedrock"), + ("anthropic.claude-sonnet-4-5-20250929-v1:0", "bedrock"), ("gemini-2.5-pro", "vertex_ai"), ], ) @@ -1325,7 +1325,7 @@ class TestProxyFunctionCalling: ), ( "litellm_proxy/bedrock-claude-3-opus", - "bedrock/anthropic.claude-3-7-sonnet-20250219-v1:0", + "bedrock/anthropic.claude-sonnet-4-5-20250929-v1:0", False, ), ( @@ -1623,7 +1623,7 @@ class TestProxyFunctionCalling: ), ( "litellm_proxy/bedrock-claude-3-opus", - "bedrock/converse/anthropic.claude-3-7-sonnet-20250219-v1:0", + "bedrock/converse/anthropic.claude-sonnet-4-5-20250929-v1:0", False, "Bedrock Claude 3 Opus via Converse API", ), @@ -1710,7 +1710,7 @@ class TestProxyFunctionCalling: ), ( "litellm_proxy/staging-claude-opus", - "bedrock/converse/anthropic.claude-3-7-sonnet-20250219-v1:0", + "bedrock/converse/anthropic.claude-sonnet-4-5-20250929-v1:0", False, "Staging Claude Opus", ), @@ -1722,7 +1722,7 @@ class TestProxyFunctionCalling: ), ( "litellm_proxy/high-performance-claude", - "bedrock/converse/anthropic.claude-3-7-sonnet-20250219-v1:0", + "bedrock/converse/anthropic.claude-sonnet-4-5-20250929-v1:0", False, "High-performance Claude deployment", ), @@ -1860,7 +1860,7 @@ class TestProxyFunctionCalling: bedrock_models = [ "bedrock/converse/anthropic.claude-3-haiku-20240307-v1:0", "bedrock/converse/anthropic.claude-3-sonnet-20240229-v1:0", - "bedrock/converse/anthropic.claude-3-7-sonnet-20250219-v1:0", + "bedrock/converse/anthropic.claude-sonnet-4-5-20250929-v1:0", ] for model in bedrock_models: @@ -1892,7 +1892,7 @@ class TestProxyFunctionCalling: ), ( "litellm_proxy/bedrock-claude-3-opus", - "bedrock/converse/anthropic.claude-3-7-sonnet-20250219-v1:0", + "bedrock/converse/anthropic.claude-sonnet-4-5-20250929-v1:0", False, "Bedrock Claude 3 Opus via Converse API", ), @@ -1979,7 +1979,7 @@ class TestProxyFunctionCalling: ), ( "litellm_proxy/staging-claude-opus", - "bedrock/converse/anthropic.claude-3-7-sonnet-20250219-v1:0", + "bedrock/converse/anthropic.claude-sonnet-4-5-20250929-v1:0", False, "Staging Claude Opus", ), @@ -1991,7 +1991,7 @@ class TestProxyFunctionCalling: ), ( "litellm_proxy/high-performance-claude", - "bedrock/converse/anthropic.claude-3-7-sonnet-20250219-v1:0", + "bedrock/converse/anthropic.claude-sonnet-4-5-20250929-v1:0", False, "High-performance Claude deployment", ), @@ -2129,7 +2129,7 @@ class TestProxyFunctionCalling: bedrock_models = [ "bedrock/converse/anthropic.claude-3-haiku-20240307-v1:0", "bedrock/converse/anthropic.claude-3-sonnet-20240229-v1:0", - "bedrock/converse/anthropic.claude-3-7-sonnet-20250219-v1:0", + "bedrock/converse/anthropic.claude-sonnet-4-5-20250929-v1:0", ] for model in bedrock_models: @@ -2161,7 +2161,7 @@ class TestProxyFunctionCalling: ), ( "litellm_proxy/bedrock-claude-3-opus", - "bedrock/converse/anthropic.claude-3-7-sonnet-20250219-v1:0", + "bedrock/converse/anthropic.claude-sonnet-4-5-20250929-v1:0", False, "Bedrock Claude 3 Opus via Converse API", ), @@ -2248,7 +2248,7 @@ class TestProxyFunctionCalling: ), ( "litellm_proxy/staging-claude-opus", - "bedrock/converse/anthropic.claude-3-7-sonnet-20250219-v1:0", + "bedrock/converse/anthropic.claude-sonnet-4-5-20250929-v1:0", False, "Staging Claude Opus", ), @@ -2260,7 +2260,7 @@ class TestProxyFunctionCalling: ), ( "litellm_proxy/high-performance-claude", - "bedrock/converse/anthropic.claude-3-7-sonnet-20250219-v1:0", + "bedrock/converse/anthropic.claude-sonnet-4-5-20250929-v1:0", False, "High-performance Claude deployment", ), @@ -2398,7 +2398,7 @@ class TestProxyFunctionCalling: bedrock_models = [ "bedrock/converse/anthropic.claude-3-haiku-20240307-v1:0", "bedrock/converse/anthropic.claude-3-sonnet-20240229-v1:0", - "bedrock/converse/anthropic.claude-3-7-sonnet-20250219-v1:0", + "bedrock/converse/anthropic.claude-sonnet-4-5-20250929-v1:0", ] for model in bedrock_models: From b1a0a3fc17d616ad2993a4c73f4eecbf31ef005c Mon Sep 17 00:00:00 2001 From: Ryan Crabbe Date: Tue, 28 Apr 2026 14:51:47 -0700 Subject: [PATCH 28/57] fix(tests): use Sonnet 4.5 for Bedrock invoke prompt-caching tests Claude 3.5 Sonnet v2 reached EOL on Bedrock 2026-03-01, returning the same 404 EOL error as 3.7 Sonnet. Sonnet 4.5 supports both InvokeModel and Converse APIs on Bedrock, so use the same model for both routes. --- .../base_anthropic_messages_prompt_caching_test.py | 2 +- .../test_anthropic_messages_prompt_caching.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/pass_through_unit_tests/base_anthropic_messages_prompt_caching_test.py b/tests/pass_through_unit_tests/base_anthropic_messages_prompt_caching_test.py index d6502afbe7..5fc4ecefb3 100644 --- a/tests/pass_through_unit_tests/base_anthropic_messages_prompt_caching_test.py +++ b/tests/pass_through_unit_tests/base_anthropic_messages_prompt_caching_test.py @@ -97,7 +97,7 @@ class BaseAnthropicMessagesPromptCachingTest(ABC): Examples: - "bedrock/converse/anthropic.claude-sonnet-4-5-20250929-v1:0" - - "bedrock/invoke/anthropic.claude-3-5-sonnet-20241022-v2:0" + - "bedrock/invoke/anthropic.claude-sonnet-4-5-20250929-v1:0" """ pass diff --git a/tests/pass_through_unit_tests/test_anthropic_messages_prompt_caching.py b/tests/pass_through_unit_tests/test_anthropic_messages_prompt_caching.py index bfdbf75351..a194ded12f 100644 --- a/tests/pass_through_unit_tests/test_anthropic_messages_prompt_caching.py +++ b/tests/pass_through_unit_tests/test_anthropic_messages_prompt_caching.py @@ -43,4 +43,4 @@ class TestBedrockInvokePromptCaching(BaseAnthropicMessagesPromptCachingTest): """ def get_model(self) -> str: - return "bedrock/invoke/us.anthropic.claude-3-5-sonnet-20241022-v2:0" + return "bedrock/invoke/us.anthropic.claude-sonnet-4-5-20250929-v1:0" From 6052ce1017aa27e7692da2d0664bfe91f659acfc Mon Sep 17 00:00:00 2001 From: Michael Riad Zaky Date: Fri, 24 Apr 2026 16:44:50 -0700 Subject: [PATCH 29/57] cache LiteLLM_Config param reads in DualCache + batch scheduler-tick fetch --- litellm/proxy/proxy_server.py | 55 +++++++++--- litellm/proxy/utils.py | 89 +++++++++++++++++++ tests/test_litellm/proxy/test_proxy_server.py | 20 +++++ 3 files changed, 150 insertions(+), 14 deletions(-) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 007dbe5fa7..8f676df04c 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -497,14 +497,18 @@ from litellm.proxy.utils import ( _get_redoc_url, _is_projected_spend_over_limit, _is_valid_team_configs, + get_config_param, get_custom_url, get_error_message_str, get_server_root_path, handle_exception_on_proxy, hash_password, hash_token, + invalidate_config_param, + litellm_config_cache, migrate_passwords_to_scrypt_async, model_dump_with_preserved_fields, + prefetch_config_params, update_spend, ) from litellm.proxy.vector_store_endpoints.endpoints import router as vector_store_router @@ -2929,8 +2933,13 @@ class ProxyConfig: ## INIT PROXY REDIS USAGE CLIENT ## redis_usage_cache = litellm.cache.cache spend_counter_cache.redis_cache = redis_usage_cache + litellm_config_cache.redis_cache = redis_usage_cache # Note: PKCE verifier storage uses redis_usage_cache directly (not # user_api_key_cache) to avoid routing all API-key lookups through Redis. + elif litellm_config_cache.redis_cache is None: + verbose_proxy_logger.info( + "litellm_config_cache: no Redis configured; cluster-wide cache sharing disabled." + ) def switch_on_llm_response_caching(self): """ @@ -4846,10 +4855,7 @@ class ProxyConfig: "environment_variables", ] for k in keys: - response = prisma_client.get_generic_data( - key="param_name", value=k, table_name="config" - ) - _tasks.append(response) + _tasks.append(get_config_param(prisma_client, k)) responses = await asyncio.gather(*_tasks) for response in responses: @@ -4931,6 +4937,19 @@ class ProxyConfig: global llm_router, llm_model_list, master_key, general_settings try: + # warm the config cache so the per-param reads below all hit + await prefetch_config_params( + prisma_client, + [ + "general_settings", + "router_settings", + "litellm_settings", + "environment_variables", + "model_cost_map_reload_config", + "anthropic_beta_headers_reload_config", + ], + ) + # Only load models from DB if "models" is in supported_db_objects (or if supported_db_objects is not set) if self._should_load_db_object(object_type="models"): new_models = await self._get_models_from_db(prisma_client=prisma_client) @@ -4940,8 +4959,8 @@ class ProxyConfig: new_models=new_models, proxy_logging_obj=proxy_logging_obj ) - db_general_settings = await prisma_client.db.litellm_config.find_first( - where={"param_name": "general_settings"} + db_general_settings = await get_config_param( + prisma_client, "general_settings" ) # update general settings @@ -5034,10 +5053,7 @@ class ProxyConfig: from litellm.proxy.hooks.mcp_semantic_filter import SemanticToolFilterHook try: - # Load litellm_settings from DB - config_record = await prisma_client.db.litellm_config.find_unique( - where={"param_name": "litellm_settings"} - ) + config_record = await get_config_param(prisma_client, "litellm_settings") if config_record is None or config_record.param_value is None: return @@ -5192,8 +5208,8 @@ class ProxyConfig: """ try: # Get model cost map reload configuration from database - config_record = await prisma_client.db.litellm_config.find_unique( - where={"param_name": "model_cost_map_reload_config"} + config_record = await get_config_param( + prisma_client, "model_cost_map_reload_config" ) if config_record is None or config_record.param_value is None: @@ -5288,6 +5304,7 @@ class ProxyConfig: }, }, ) + await invalidate_config_param("model_cost_map_reload_config") verbose_proxy_logger.info( f"Model cost map reloaded successfully. Models count: {len(new_model_cost_map) if new_model_cost_map else 0}" @@ -5307,8 +5324,8 @@ class ProxyConfig: """ try: # Get anthropic beta headers reload configuration from database - config_record = await prisma_client.db.litellm_config.find_unique( - where={"param_name": "anthropic_beta_headers_reload_config"} + config_record = await get_config_param( + prisma_client, "anthropic_beta_headers_reload_config" ) if config_record is None or config_record.param_value is None: @@ -5396,6 +5413,7 @@ class ProxyConfig: }, }, ) + await invalidate_config_param("anthropic_beta_headers_reload_config") # Count providers in config provider_count = sum( @@ -12674,6 +12692,7 @@ async def update_config( # noqa: PLR0915 "update": {"param_value": v}, }, ) + await invalidate_config_param(k) ### OLD LOGIC [TODO] MOVE TO DB ### @@ -12861,6 +12880,7 @@ async def update_config_general_settings( "update": {"param_value": json.dumps(general_settings)}, # type: ignore }, ) + await invalidate_config_param("general_settings") return response @@ -13144,6 +13164,7 @@ async def delete_config_general_settings( "update": {"param_value": json.dumps(general_settings)}, # type: ignore }, ) + await invalidate_config_param("general_settings") return response @@ -13509,6 +13530,7 @@ async def reload_model_cost_map( }, }, ) + await invalidate_config_param("model_cost_map_reload_config") models_count = len(new_model_cost_map) if new_model_cost_map else 0 verbose_proxy_logger.info( @@ -13578,6 +13600,7 @@ async def schedule_model_cost_map_reload( }, }, ) + await invalidate_config_param("model_cost_map_reload_config") verbose_proxy_logger.info( f"Model cost map reload scheduled for every {hours} hours" @@ -13631,6 +13654,7 @@ async def cancel_model_cost_map_reload( await prisma_client.db.litellm_config.delete( where={"param_name": "model_cost_map_reload_config"} ) + await invalidate_config_param("model_cost_map_reload_config") verbose_proxy_logger.info("Model cost map reload schedule cancelled") @@ -13861,6 +13885,7 @@ async def reload_anthropic_beta_headers( }, }, ) + await invalidate_config_param("anthropic_beta_headers_reload_config") provider_count = sum( 1 for k in new_config.keys() if k not in ["provider_aliases", "description"] @@ -13934,6 +13959,7 @@ async def schedule_anthropic_beta_headers_reload( }, }, ) + await invalidate_config_param("anthropic_beta_headers_reload_config") verbose_proxy_logger.info( f"Anthropic beta headers reload scheduled for every {hours} hours" @@ -13987,6 +14013,7 @@ async def cancel_anthropic_beta_headers_reload( await prisma_client.db.litellm_config.delete( where={"param_name": "anthropic_beta_headers_reload_config"} ) + await invalidate_config_param("anthropic_beta_headers_reload_config") verbose_proxy_logger.info("Anthropic beta headers reload schedule cancelled") diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 712853a33c..3a1184c434 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -2442,6 +2442,92 @@ async def _lookup_deprecated_key( return None +# DualCache for LiteLLM_Config param_name reads. +# Redis layer is attached in proxy_server._init_cache. +LITELLM_CONFIG_CACHE_TTL_SECONDS: int = int( + os.environ.get("LITELLM_CONFIG_PARAM_CACHE_TTL_SECONDS", "60") +) +_CONFIG_CACHE_MISS: str = "__litellm_config_param_miss__" + +litellm_config_cache: DualCache = DualCache( + default_in_memory_ttl=LITELLM_CONFIG_CACHE_TTL_SECONDS, + default_redis_ttl=LITELLM_CONFIG_CACHE_TTL_SECONDS, +) + + +class _ConfigRow: + """Mimics the Prisma litellm_config row shape for cached entries.""" + + __slots__ = ("param_name", "param_value") + + def __init__(self, param_name: str, param_value: Any) -> None: + self.param_name = param_name + self.param_value = param_value + + +def _config_cache_key(param_name: str) -> str: + return f"litellm_config:param:{param_name}" + + +def _pack_config_row(row: Any) -> Dict[str, Any]: + return {"param_name": row.param_name, "param_value": row.param_value} + + +def _unpack_config_row(cached: Any) -> Optional[_ConfigRow]: + if cached is None or cached == _CONFIG_CACHE_MISS: + return None + if isinstance(cached, dict): + return _ConfigRow(cached["param_name"], cached["param_value"]) + return None + + +async def get_config_param(prisma_client: Any, param_name: str) -> Optional[Any]: + """Cached read of a LiteLLM_Config row; returns row, _ConfigRow shim, or None.""" + cache_key = _config_cache_key(param_name) + cached = await litellm_config_cache.async_get_cache(cache_key) + if cached is not None: + return _unpack_config_row(cached) + + row = await prisma_client.get_generic_data( + key="param_name", value=param_name, table_name="config" + ) + cache_value: Any = _pack_config_row(row) if row is not None else _CONFIG_CACHE_MISS + await litellm_config_cache.async_set_cache( + cache_key, cache_value, ttl=LITELLM_CONFIG_CACHE_TTL_SECONDS + ) + return row + + +async def invalidate_config_param(param_name: str) -> None: + """Evict from both cache layers; call after every LiteLLM_Config write.""" + await litellm_config_cache.async_delete_cache(_config_cache_key(param_name)) + + +async def prefetch_config_params(prisma_client: Any, param_names: List[str]) -> None: + """Batch-load LiteLLM_Config rows into the cache with one find_many.""" + if not param_names: + return + try: + rows = await prisma_client.db.litellm_config.find_many( + where={"param_name": {"in": param_names}} # type: ignore + ) + except Exception as e: + verbose_proxy_logger.debug( + "prefetch_config_params failed, falling through to per-param queries: %s", + e, + ) + return + by_name = {row.param_name: row for row in rows} + for name in param_names: + row = by_name.get(name) + cache_value: Any = ( + _pack_config_row(row) if row is not None else _CONFIG_CACHE_MISS + ) + await litellm_config_cache.async_set_cache( + _config_cache_key(name), cache_value, ttl=LITELLM_CONFIG_CACHE_TTL_SECONDS + ) + + class PrismaClient: spend_log_transactions: List = [] _spend_log_transactions_lock = asyncio.Lock() @@ -3310,6 +3396,9 @@ class PrismaClient: tasks.append(updated_table_row) await asyncio.gather(*tasks) + # invalidate cache so other pods see writes from save_config + for k in data.keys(): + await invalidate_config_param(k) verbose_proxy_logger.info("Data Inserted into Config Table") elif table_name == "spend": db_data = self.jsonify_object(data=data) diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index 3349a138ee..1f4f82a64e 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -2544,6 +2544,14 @@ class TestPriceDataReloadAPI: class TestPriceDataReloadIntegration: """Integration tests for the complete price data reload feature""" + @pytest.fixture(autouse=True) + def _flush_litellm_config_cache(self): + from litellm.proxy.utils import litellm_config_cache + + litellm_config_cache.flush_cache() + yield + litellm_config_cache.flush_cache() + @pytest.fixture def client_with_auth(self): """Create a test client with authentication""" @@ -2601,6 +2609,7 @@ class TestPriceDataReloadIntegration: def test_distributed_reload_check_function(self): """Test the _check_and_reload_model_cost_map function""" from litellm.proxy.proxy_server import ProxyConfig + from litellm.proxy.utils import litellm_config_cache proxy_config = ProxyConfig() @@ -2609,14 +2618,19 @@ class TestPriceDataReloadIntegration: # Test case 1: No config in database mock_prisma.db.litellm_config.find_unique = AsyncMock(return_value=None) + # _check_and_reload_model_cost_map routes through get_config_param, + # which calls prisma.get_generic_data on a cache miss. + mock_prisma.get_generic_data = AsyncMock(return_value=None) # Should return early without reloading asyncio.run(proxy_config._check_and_reload_model_cost_map(mock_prisma)) # Test case 2: Config with interval but not time to reload + litellm_config_cache.flush_cache() mock_config = MagicMock() mock_config.param_value = {"interval_hours": 6, "force_reload": False} mock_prisma.db.litellm_config.find_unique = AsyncMock(return_value=mock_config) + mock_prisma.get_generic_data = AsyncMock(return_value=mock_config) # Mock current time and last reload time with patch( @@ -2632,8 +2646,10 @@ class TestPriceDataReloadIntegration: asyncio.run(proxy_config._check_and_reload_model_cost_map(mock_prisma)) # Test case 3: Config with force reload + litellm_config_cache.flush_cache() mock_config.param_value = {"interval_hours": 6, "force_reload": True} mock_prisma.db.litellm_config.find_unique = AsyncMock(return_value=mock_config) + mock_prisma.get_generic_data = AsyncMock(return_value=mock_config) mock_prisma.db.litellm_config.upsert = AsyncMock(return_value=None) original_model_cost = litellm.model_cost.copy() @@ -2675,6 +2691,8 @@ class TestPriceDataReloadIntegration: mock_config = MagicMock() mock_config.param_value = {"interval_hours": 24, "force_reload": True} mock_prisma.db.litellm_config.find_unique = AsyncMock(return_value=mock_config) + # _check_and_reload_model_cost_map now reads through get_generic_data. + mock_prisma.get_generic_data = AsyncMock(return_value=mock_config) mock_prisma.db.litellm_config.upsert = AsyncMock(return_value=None) original_model_cost = litellm.model_cost.copy() @@ -2770,6 +2788,8 @@ class TestPriceDataReloadIntegration: mock_config = MagicMock() mock_config.param_value = {"interval_hours": 12, "force_reload": True} mock_prisma.db.litellm_config.find_unique = AsyncMock(return_value=mock_config) + # _check_and_reload_anthropic_beta_headers now reads through get_generic_data. + mock_prisma.get_generic_data = AsyncMock(return_value=mock_config) mock_prisma.db.litellm_config.upsert = AsyncMock(return_value=None) with patch( From 21ed38971d244c0a034604f6439c0584d55b4d20 Mon Sep 17 00:00:00 2001 From: Michael-RZ-Berri Date: Tue, 28 Apr 2026 17:04:40 -0700 Subject: [PATCH 30/57] lazy-load optional feature routers on first request (#26534) Co-authored-by: Michael Riad Zaky --- litellm/proxy/_lazy_features.py | 307 ++++++++++++++++++ litellm/proxy/proxy_server.py | 126 ++----- tests/proxy_unit_tests/test_proxy_routes.py | 14 + tests/test_litellm/proxy/test_proxy_server.py | 252 ++++++++++++++ .../test_vector_store_endpoints.py | 15 + 5 files changed, 609 insertions(+), 105 deletions(-) create mode 100644 litellm/proxy/_lazy_features.py diff --git a/litellm/proxy/_lazy_features.py b/litellm/proxy/_lazy_features.py new file mode 100644 index 0000000000..450c9483f3 --- /dev/null +++ b/litellm/proxy/_lazy_features.py @@ -0,0 +1,307 @@ +""" +Lazy registration for optional feature routers. Each LAZY_FEATURES entry +imports its module only on the first request matching its path prefix, +saving ~700 MB at idle for deployments that don't use these features. +First hit pays the import cost (1-3 s for heavy modules); /openapi.json +omits each feature's routes until the feature is warmed. +""" + +import asyncio +import importlib +from dataclasses import dataclass, field +from typing import TYPE_CHECKING, Callable, Tuple + +from starlette.types import Receive, Scope, Send + +from litellm._logging import verbose_proxy_logger + +if TYPE_CHECKING: + from fastapi import FastAPI + + +def _include_router(attr_name: str = "router") -> Callable[["FastAPI", object], None]: + def _register(app: "FastAPI", module: object) -> None: + app.include_router(getattr(module, attr_name)) + + return _register + + +def _mount_app( + prefix: str, attr_name: str = "app" +) -> Callable[["FastAPI", object], None]: + def _register(app: "FastAPI", module: object) -> None: + app.mount(path=prefix, app=getattr(module, attr_name)) + + return _register + + +@dataclass(frozen=True) +class LazyFeature: + name: str + module_path: str + path_prefixes: Tuple[str, ...] + register_fn: Callable[["FastAPI", object], None] = field( + default_factory=lambda: _include_router("router") + ) + # For routes whose path has a leading parameter (e.g. /{server}/authorize) + # — startswith can't match those, so the matcher also checks endswith. + path_suffixes: Tuple[str, ...] = () + + +LAZY_FEATURES: Tuple[LazyFeature, ...] = ( + LazyFeature( + name="guardrails", + module_path="litellm.proxy.guardrails.guardrail_endpoints", + path_prefixes=( + "/guardrails", + "/v2/guardrails", + "/apply_guardrail", + "/policies/usage", + ), + ), + LazyFeature( + name="policies", + module_path="litellm.proxy.management_endpoints.policy_endpoints", + # Trailing slash to avoid matching /policies/... (policy_engine). + path_prefixes=("/policy/", "/utils/test_policies_and_guardrails"), + ), + LazyFeature( + name="policy_engine", + module_path="litellm.proxy.policy_engine.policy_endpoints", + path_prefixes=("/policies",), + ), + LazyFeature( + name="policy_resolve", + module_path="litellm.proxy.policy_engine.policy_resolve_endpoints", + path_prefixes=("/policies/resolve", "/policies/attachments/estimate-impact"), + ), + LazyFeature( + name="agents", + module_path="litellm.proxy.agent_endpoints.endpoints", + path_prefixes=("/v1/agents", "/agents", "/agent/"), + ), + LazyFeature( + name="a2a", + module_path="litellm.proxy.agent_endpoints.a2a_endpoints", + path_prefixes=("/a2a", "/v1/a2a"), + ), + LazyFeature( + name="vector_stores", + module_path="litellm.proxy.vector_store_endpoints.endpoints", + path_prefixes=("/v1/vector_stores", "/vector_stores", "/v1/indexes"), + ), + LazyFeature( + name="vector_store_management", + module_path="litellm.proxy.vector_store_endpoints.management_endpoints", + # Trailing slash to avoid matching /vector_stores/... (vector_stores). + path_prefixes=("/vector_store/", "/v1/vector_store/"), + ), + LazyFeature( + name="vector_store_files", + # Routes appear under both /v1/vector_stores/{id}/files and the + # un-versioned form, so both prefixes must trigger the load. + module_path="litellm.proxy.vector_store_files_endpoints.endpoints", + path_prefixes=("/v1/vector_stores", "/vector_stores"), + ), + LazyFeature( + name="tools", + module_path="litellm.proxy.management_endpoints.tool_management_endpoints", + path_prefixes=("/v1/tool", "/tool"), + ), + LazyFeature( + name="search_tools", + module_path="litellm.proxy.search_endpoints.search_tool_management", + path_prefixes=("/search_tools",), + ), + # mcp_management owns most /v1/mcp/* admin routes; mcp_app is the mounted + # streaming sub-app at /mcp. + LazyFeature( + name="mcp_management", + module_path="litellm.proxy.management_endpoints.mcp_management_endpoints", + path_prefixes=("/v1/mcp/",), + ), + LazyFeature( + # Also serves /.well-known/oauth-* (OAuth metadata discovery). + # No /mcp/oauth prefix here: the mounted /mcp sub-app would + # shadow it, and there are no actual routes there anyway. + name="mcp_byok_oauth", + module_path="litellm.proxy._experimental.mcp_server.byok_oauth_endpoints", + path_prefixes=("/v1/mcp/oauth", "/.well-known/oauth-"), + ), + LazyFeature( + # Serves OAuth dance endpoints (/authorize, /token, /callback, + # /register) plus several /.well-known/ discovery URLs at the proxy + # root — needed for MCP-over-OAuth flows even before /mcp is hit. + name="mcp_discoverable", + module_path="litellm.proxy._experimental.mcp_server.discoverable_endpoints", + path_prefixes=( + "/.well-known/oauth-", + "/.well-known/openid-configuration", + "/.well-known/jwks.json", + "/authorize", + "/token", + "/callback", + "/register", + ), + # Catches the /{mcp_server_name}/authorize|token|register variants. + path_suffixes=("/authorize", "/token", "/register"), + ), + LazyFeature( + name="mcp_rest", + module_path="litellm.proxy._experimental.mcp_server.rest_endpoints", + path_prefixes=("/mcp-rest",), + ), + LazyFeature( + # Hardcoded /mcp matches BASE_MCP_ROUTE; importing the constant + # here would defeat lazy loading. + name="mcp_app", + module_path="litellm.proxy._experimental.mcp_server.server", + path_prefixes=("/mcp",), + register_fn=_mount_app("/mcp", attr_name="app"), + ), + LazyFeature( + name="config_overrides", + module_path="litellm.proxy.management_endpoints.config_override_endpoints", + path_prefixes=("/config_overrides",), + ), + LazyFeature( + name="realtime", + module_path="litellm.proxy.realtime_endpoints.endpoints", + path_prefixes=("/openai/v1/realtime", "/v1/realtime", "/realtime"), + ), + LazyFeature( + name="anthropic_passthrough", + module_path="litellm.proxy.anthropic_endpoints.endpoints", + path_prefixes=("/v1/messages", "/anthropic", "/api/event_logging"), + ), + LazyFeature( + name="anthropic_skills", + module_path="litellm.proxy.anthropic_endpoints.skills_endpoints", + path_prefixes=("/v1/skills", "/skills"), + ), + LazyFeature( + name="langfuse_passthrough", + module_path="litellm.proxy.vertex_ai_endpoints.langfuse_endpoints", + path_prefixes=("/langfuse",), + ), + LazyFeature( + name="evals", + module_path="litellm.proxy.openai_evals_endpoints.endpoints", + path_prefixes=("/v1/evals", "/evals"), + ), + LazyFeature( + name="claude_code_marketplace", + module_path="litellm.proxy.anthropic_endpoints.claude_code_endpoints", + path_prefixes=("/claude-code",), + register_fn=_include_router("claude_code_marketplace_router"), + ), + LazyFeature( + name="scim", + module_path="litellm.proxy.management_endpoints.scim.scim_v2", + path_prefixes=("/scim",), + register_fn=_include_router("scim_router"), + ), + LazyFeature( + name="cloudzero", + module_path="litellm.proxy.spend_tracking.cloudzero_endpoints", + path_prefixes=("/cloudzero",), + ), + LazyFeature( + name="vantage", + module_path="litellm.proxy.spend_tracking.vantage_endpoints", + path_prefixes=("/vantage",), + ), + LazyFeature( + name="usage_ai", + module_path="litellm.proxy.management_endpoints.usage_endpoints", + path_prefixes=("/usage/ai",), + ), + LazyFeature( + name="prompts", + module_path="litellm.proxy.prompts.prompt_endpoints", + path_prefixes=("/prompts", "/utils/dotprompt_json_converter"), + ), + LazyFeature( + name="jwt_mappings", + module_path="litellm.proxy.management_endpoints.jwt_key_mapping_endpoints", + path_prefixes=("/jwt/key/mapping",), + ), + LazyFeature( + name="compliance", + module_path="litellm.proxy.management_endpoints.compliance_endpoints", + path_prefixes=("/compliance",), + ), + LazyFeature( + name="access_groups", + module_path="litellm.proxy.management_endpoints.access_group_endpoints", + path_prefixes=("/access_group", "/v1/access_group", "/v1/unified_access_group"), + ), +) + + +class LazyFeatureMiddleware: + """ASGI middleware that imports + registers a feature router on first + matching request. Idempotent; once loaded, subsequent requests skip.""" + + def __init__( + self, + app, + fastapi_app: "FastAPI", + features: Tuple[LazyFeature, ...] = LAZY_FEATURES, + ): + self.app = app + self._fastapi_app = fastapi_app + self._features = features + self._loaded: set = set() + # Per-feature locks so independent features can load in parallel. + self._locks: dict = {} + + async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None: + # Short-circuit once every feature has loaded. + if scope["type"] in ("http", "websocket") and len(self._loaded) < len( + self._features + ): + path = scope.get("path", "") + for feat in self._features: + if feat.module_path in self._loaded: + continue + if any(path.startswith(p) for p in feat.path_prefixes) or any( + path.endswith(s) for s in feat.path_suffixes + ): + await self._load(feat) + await self.app(scope, receive, send) + + async def _load(self, feat: LazyFeature) -> None: + lock = self._locks.setdefault(feat.module_path, asyncio.Lock()) + async with lock: + if feat.module_path in self._loaded: + return + try: + # Import on a thread (heavy modules take 1-3 s). register_fn + # mutates app.router.routes, so it stays on the loop thread. + loop = asyncio.get_running_loop() + module = await loop.run_in_executor( + None, importlib.import_module, feat.module_path + ) + feat.register_fn(self._fastapi_app, module) + self._loaded.add(feat.module_path) + self._fastapi_app.openapi_schema = None + verbose_proxy_logger.info( + "Lazy-loaded optional feature %r (module: %s)", + feat.name, + feat.module_path, + ) + except Exception as exc: + # Mark loaded anyway so we don't retry on every request. + self._loaded.add(feat.module_path) + verbose_proxy_logger.warning( + "Failed to lazy-load optional feature %r (module: %s): %s. " + "This feature's endpoints will return 404 until restart.", + feat.name, + feat.module_path, + exc, + ) + + +def attach_lazy_features(app: "FastAPI") -> None: + app.add_middleware(LazyFeatureMiddleware, fastapi_app=app) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 8f676df04c..c03a63f211 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -235,37 +235,11 @@ from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging from litellm.litellm_core_utils.sensitive_data_masker import SensitiveDataMasker from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler from litellm.llms.vertex_ai.vertex_llm_base import VertexBase -from litellm.proxy._experimental.mcp_server.byok_oauth_endpoints import ( - router as mcp_byok_oauth_router, -) -from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( - router as mcp_discoverable_endpoints_router, -) -from litellm.proxy._experimental.mcp_server.rest_endpoints import ( - router as mcp_rest_endpoints_router, -) -from litellm.proxy._experimental.mcp_server.server import app as mcp_app -from litellm.proxy._experimental.mcp_server.tool_registry import ( - global_mcp_tool_registry, -) from litellm.proxy._types import * -from litellm.proxy.agent_endpoints.a2a_endpoints import router as a2a_router -from litellm.proxy.agent_endpoints.agent_registry import global_agent_registry -from litellm.proxy.agent_endpoints.endpoints import router as agent_endpoints_router -from litellm.proxy.agent_endpoints.model_list_helpers import ( - append_agents_to_model_group, - append_agents_to_model_info, -) +from litellm.proxy._lazy_features import attach_lazy_features from litellm.proxy.analytics_endpoints.analytics_endpoints import ( router as analytics_router, ) -from litellm.proxy.anthropic_endpoints.claude_code_endpoints import ( - claude_code_marketplace_router, -) -from litellm.proxy.anthropic_endpoints.endpoints import router as anthropic_router -from litellm.proxy.anthropic_endpoints.skills_endpoints import ( - router as anthropic_skills_router, -) from litellm.proxy.auth.auth_checks import ( ExperimentalUIJWTToken, get_team_object, @@ -328,7 +302,6 @@ from litellm.proxy.discovery_endpoints import ui_discovery_endpoints_router from litellm.proxy.fine_tuning_endpoints.endpoints import router as fine_tuning_router from litellm.proxy.fine_tuning_endpoints.endpoints import set_fine_tuning_config from litellm.proxy.google_endpoints.endpoints import router as google_router -from litellm.proxy.guardrails.guardrail_endpoints import router as guardrails_router from litellm.proxy.guardrails.init_guardrails import ( init_guardrails_v2, initialize_guardrails, @@ -344,9 +317,6 @@ from litellm.proxy.hooks.prompt_injection_detection import ( from litellm.proxy.hooks.proxy_track_cost_callback import _ProxyDBLogger from litellm.proxy.image_endpoints.endpoints import router as image_router from litellm.proxy.litellm_pre_call_utils import add_litellm_data_to_request -from litellm.proxy.management_endpoints.access_group_endpoints import ( - router as access_group_router, -) from litellm.proxy.management_endpoints.budget_management_endpoints import ( router as budget_management_router, ) @@ -360,12 +330,6 @@ from litellm.proxy.management_endpoints.common_utils import ( _user_has_admin_privileges, admin_can_invite_user, ) -from litellm.proxy.management_endpoints.compliance_endpoints import ( - router as compliance_router, -) -from litellm.proxy.management_endpoints.config_override_endpoints import ( - router as config_override_router, -) from litellm.proxy.management_endpoints.cost_tracking_settings import ( router as cost_tracking_settings_router, ) @@ -379,9 +343,6 @@ from litellm.proxy.management_endpoints.internal_user_endpoints import ( router as internal_user_router, ) from litellm.proxy.management_endpoints.internal_user_endpoints import user_update -from litellm.proxy.management_endpoints.jwt_key_mapping_endpoints import ( - router as jwt_key_mapping_router, -) from litellm.proxy.management_endpoints.key_management_endpoints import ( delete_verification_tokens, duration_in_seconds, @@ -390,9 +351,6 @@ from litellm.proxy.management_endpoints.key_management_endpoints import ( from litellm.proxy.management_endpoints.key_management_endpoints import ( router as key_management_router, ) -from litellm.proxy.management_endpoints.mcp_management_endpoints import ( - router as mcp_management_router, -) from litellm.proxy.management_endpoints.model_access_group_management_endpoints import ( router as model_access_group_management_router, ) @@ -407,11 +365,9 @@ from litellm.proxy.management_endpoints.model_management_endpoints import ( from litellm.proxy.management_endpoints.organization_endpoints import ( router as organization_router, ) -from litellm.proxy.management_endpoints.policy_endpoints import router as policy_router from litellm.proxy.management_endpoints.router_settings_endpoints import ( router as router_settings_router, ) -from litellm.proxy.management_endpoints.scim.scim_v2 import scim_router from litellm.proxy.management_endpoints.tag_management_endpoints import ( router as tag_management_router, ) @@ -423,15 +379,11 @@ from litellm.proxy.management_endpoints.team_endpoints import ( update_team, validate_membership, ) -from litellm.proxy.management_endpoints.tool_management_endpoints import ( - router as tool_management_router, -) from litellm.proxy.memory.memory_endpoints import router as memory_router from litellm.proxy.management_endpoints.ui_sso import ( get_disabled_non_admin_personal_key_creation, ) from litellm.proxy.management_endpoints.ui_sso import router as ui_sso_router -from litellm.proxy.management_endpoints.usage_endpoints import router as usage_ai_router from litellm.proxy.management_endpoints.user_agent_analytics_endpoints import ( router as user_agent_analytics_router, ) @@ -441,7 +393,6 @@ from litellm.proxy.middleware.in_flight_requests_middleware import ( ) from litellm.proxy.middleware.prometheus_auth_middleware import PrometheusAuthMiddleware from litellm.proxy.ocr_endpoints.endpoints import router as ocr_router -from litellm.proxy.openai_evals_endpoints.endpoints import router as evals_router from litellm.proxy.openai_files_endpoints.files_endpoints import ( router as openai_files_router, ) @@ -461,27 +412,16 @@ from litellm.proxy.pass_through_endpoints.pass_through_endpoints import ( from litellm.proxy.pass_through_endpoints.pass_through_endpoints import ( router as pass_through_router, ) -from litellm.proxy.policy_engine.policy_endpoints import router as policy_crud_router -from litellm.proxy.policy_engine.policy_resolve_endpoints import ( - router as policy_resolve_router, -) -from litellm.proxy.prompts.prompt_endpoints import router as prompts_router from litellm.proxy.public_endpoints import router as public_endpoints_router from litellm.proxy.rag_endpoints.endpoints import router as rag_router -from litellm.proxy.realtime_endpoints.endpoints import router as webrtc_router from litellm.proxy.rerank_endpoints.endpoints import router as rerank_router from litellm.proxy.response_api_endpoints.endpoints import router as response_router from litellm.proxy.route_llm_request import route_request from litellm.proxy.search_endpoints.endpoints import router as search_router -from litellm.proxy.search_endpoints.search_tool_management import ( - router as search_tool_management_router, -) -from litellm.proxy.spend_tracking.cloudzero_endpoints import router as cloudzero_router from litellm.proxy.spend_tracking.spend_management_endpoints import ( router as spend_management_router, ) from litellm.proxy.spend_tracking.spend_tracking_utils import get_logging_payload -from litellm.proxy.spend_tracking.vantage_endpoints import router as vantage_router from litellm.proxy.types_utils.utils import get_instance_fn from litellm.proxy.ui_crud_endpoints.proxy_setting_endpoints import ( router as ui_crud_endpoints_router, @@ -511,16 +451,6 @@ from litellm.proxy.utils import ( prefetch_config_params, update_spend, ) -from litellm.proxy.vector_store_endpoints.endpoints import router as vector_store_router -from litellm.proxy.vector_store_endpoints.management_endpoints import ( - router as vector_store_management_router, -) -from litellm.proxy.vector_store_files_endpoints.endpoints import ( - router as vector_store_files_router, -) -from litellm.proxy.vertex_ai_endpoints.langfuse_endpoints import ( - router as langfuse_router, -) from litellm.proxy.video_endpoints.endpoints import router as video_router from litellm.router import ( AssistantsTypedDict, @@ -3854,11 +3784,19 @@ class ProxyConfig: ## MCP TOOLS mcp_tools_config = config.get("mcp_tools", None) if mcp_tools_config: + from litellm.proxy._experimental.mcp_server.tool_registry import ( + global_mcp_tool_registry, + ) + global_mcp_tool_registry.load_tools_from_config(mcp_tools_config) ## AGENTS agent_config = config.get("agent_list", None) if agent_config: + from litellm.proxy.agent_endpoints.agent_registry import ( + global_agent_registry, + ) + global_agent_registry.load_agents_from_config(agent_config) # type: ignore mcp_servers_config = config.get("mcp_servers", None) @@ -10576,6 +10514,10 @@ async def model_info_v2( verbose_proxy_logger.debug("all_models: %s", all_models) # Append A2A agents to models list + from litellm.proxy.agent_endpoints.model_list_helpers import ( + append_agents_to_model_info, + ) + all_models = await append_agents_to_model_info( models=all_models, user_api_key_dict=user_api_key_dict, @@ -11425,6 +11367,10 @@ async def model_group_info( ) # Append A2A agents to model groups + from litellm.proxy.agent_endpoints.model_list_helpers import ( + append_agents_to_model_group, + ) + model_groups = await append_agents_to_model_group( model_groups=model_groups, user_api_key_dict=user_api_key_dict, @@ -14230,65 +14176,40 @@ app.include_router(container_router) app.include_router(search_router) app.include_router(image_router) app.include_router(fine_tuning_router) -app.include_router(vector_store_router) -app.include_router(vector_store_management_router) -app.include_router(vector_store_files_router) app.include_router(credential_router) app.include_router(llm_passthrough_router) -app.include_router(webrtc_router) -app.include_router(mcp_management_router) -app.include_router(mcp_byok_oauth_router) -app.include_router(anthropic_router) -app.include_router(anthropic_skills_router) -app.include_router(evals_router) -app.include_router(claude_code_marketplace_router) -app.include_router(google_router) -app.include_router(langfuse_router) app.include_router(pass_through_router) app.include_router(health_router) app.include_router(key_management_router) app.include_router(internal_user_router) app.include_router(team_router) app.include_router(ui_sso_router) -app.include_router(scim_router) app.include_router(organization_router) app.include_router(customer_router) app.include_router(spend_management_router) -app.include_router(cloudzero_router) -app.include_router(vantage_router) app.include_router(caching_router) app.include_router(analytics_router) -app.include_router(guardrails_router) -app.include_router(policy_router) -app.include_router(usage_ai_router) -app.include_router(policy_crud_router) -app.include_router(policy_resolve_router) -app.include_router(search_tool_management_router) -app.include_router(prompts_router) app.include_router(callback_management_endpoints_router) app.include_router(debugging_endpoints_router) app.include_router(ui_crud_endpoints_router) app.include_router(openai_files_router) app.include_router(team_callback_router) -app.include_router(jwt_key_mapping_router) app.include_router(budget_management_router) app.include_router(model_management_router) app.include_router(model_access_group_management_router) app.include_router(tag_management_router) -app.include_router(tool_management_router) app.include_router(memory_router) app.include_router(cost_tracking_settings_router) app.include_router(router_settings_router) app.include_router(fallback_management_router) app.include_router(cache_settings_router) -app.include_router(config_override_router) app.include_router(user_agent_analytics_router) app.include_router(enterprise_router) app.include_router(ui_discovery_endpoints_router) -app.include_router(agent_endpoints_router) -app.include_router(compliance_router) -app.include_router(a2a_router) -app.include_router(access_group_router) +# Eager: /models/{name}:method overlaps with the OpenAI /models endpoint. +app.include_router(google_router) + +attach_lazy_features(app) async def _stream_mcp_asgi_response( @@ -14521,8 +14442,3 @@ async def dynamic_mcp_route(mcp_server_name: str, request: Request): f"Error handling dynamic MCP route for {mcp_server_name}: {str(e)}" ) raise HTTPException(status_code=500, detail=f"Internal server error: {str(e)}") - - -app.mount(path=BASE_MCP_ROUTE, app=mcp_app) -app.include_router(mcp_rest_endpoints_router) -app.include_router(mcp_discoverable_endpoints_router) diff --git a/tests/proxy_unit_tests/test_proxy_routes.py b/tests/proxy_unit_tests/test_proxy_routes.py index 812e4e1ac4..67eca5206d 100644 --- a/tests/proxy_unit_tests/test_proxy_routes.py +++ b/tests/proxy_unit_tests/test_proxy_routes.py @@ -39,6 +39,20 @@ def test_routes_on_litellm_proxy(): this prevents accidentelly deleting /threads, or /batches etc """ + # Force-load lazy features so the test sees the full route set. Continue + # on per-feature import failure — the assertion below still catches + # missing-route regressions. + import importlib + + from litellm.proxy._lazy_features import LAZY_FEATURES + + for feat in LAZY_FEATURES: + try: + module = importlib.import_module(feat.module_path) + feat.register_fn(app, module) + except Exception as exc: + print(f"warning: failed to force-load {feat.name}: {exc}") + _all_routes = [] for route in app.routes: diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index 1f4f82a64e..7a96f6cbd1 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -5471,3 +5471,255 @@ async def test_reseed_warms_cache_even_on_zero_db_spend(): finally: ps.spend_counter_cache = orig_counter ps.prisma_client = orig_prisma + + +# --------------------------------------------------------------------------- +# Lazy feature loading (LazyFeatureMiddleware) — verifies that optional +# routers are NOT imported at module load and ARE imported on first request +# to a matching path prefix. The same module isn't re-imported on subsequent +# requests. +# --------------------------------------------------------------------------- + + +import sys + + +class TestLazyFeatureRegistry: + """Sanity checks on the registry shape — guards against accidental edits.""" + + def test_registry_entries_have_required_fields(self): + from litellm.proxy._lazy_features import LAZY_FEATURES, LazyFeature + + assert len(LAZY_FEATURES) > 0 + for feat in LAZY_FEATURES: + assert isinstance(feat, LazyFeature) + assert feat.name + assert feat.module_path + assert feat.path_prefixes + assert all(p.startswith("/") for p in feat.path_prefixes) + assert callable(feat.register_fn) + + def test_registry_names_unique(self): + from litellm.proxy._lazy_features import LAZY_FEATURES + + names = [f.name for f in LAZY_FEATURES] + assert len(names) == len(set(names)), "duplicate feature names" + + +class TestLazyFeaturesNotImportedAtStartup: + """ + The whole point of the refactor: gated feature modules must NOT be + present in `sys.modules` immediately after `proxy_server` imports. + """ + + def test_heavy_modules_absent_at_startup(self): + # Force a fresh `proxy_server` import in a subprocess so other tests + # in this run (which may have triggered lazy loads via the TestClient) + # don't pollute the result. + import subprocess + + check = ( + "import sys; " + "from litellm.proxy.proxy_server import app; " # noqa: F401 + "heavy = [" + "'litellm.proxy._experimental.mcp_server.rest_endpoints'," + "'litellm.proxy._experimental.mcp_server.server'," + "'litellm.proxy.management_endpoints.config_override_endpoints'," + "'litellm.proxy.guardrails.guardrail_endpoints'," + "'litellm.proxy.openai_evals_endpoints.endpoints'," + "]; " + "still_present = [m for m in heavy if m in sys.modules]; " + "print('PRESENT_AT_STARTUP:', still_present)" + ) + result = subprocess.run( + [sys.executable, "-c", check], + capture_output=True, + text=True, + timeout=120, + ) + # Last non-empty line of stdout (skip warnings printed before) + out_lines = [ + line for line in result.stdout.strip().splitlines() if line.strip() + ] + report = next((line for line in out_lines if "PRESENT_AT_STARTUP" in line), "") + assert report, f"no report emitted (stderr: {result.stderr[-500:]})" + assert ( + "PRESENT_AT_STARTUP: []" in report + ), f"expected no heavy modules at startup, got: {report}" + + +class TestLazyFeatureMiddleware: + """Behavior of the middleware itself, exercised in isolation.""" + + @pytest.mark.asyncio + async def test_first_request_triggers_load_subsequent_does_not(self): + from fastapi import FastAPI + + from litellm.proxy._lazy_features import ( + LazyFeature, + LazyFeatureMiddleware, + ) + + loads = [] + + def fake_register(app, module): + loads.append(getattr(module, "__name__", "?")) + + feat = LazyFeature( + name="dummy", + module_path="json", # any always-importable stdlib module + path_prefixes=("/dummy",), + register_fn=fake_register, + ) + + # Build a minimal ASGI receiver to satisfy the middleware contract + async def downstream(scope, receive, send): + # echo back; no-op handler + await send({"type": "http.response.start", "status": 200, "headers": []}) + await send({"type": "http.response.body", "body": b""}) + + target_app = FastAPI() + mw = LazyFeatureMiddleware(downstream, fastapi_app=target_app, features=(feat,)) + + async def receive(): + return {"type": "http.request", "body": b"", "more_body": False} + + sent: list = [] + + async def send(message): + sent.append(message) + + # First request matching the prefix triggers register + await mw( + {"type": "http", "path": "/dummy/x", "method": "GET", "headers": []}, + receive, + send, + ) + assert loads == ["json"] + + # Second matching request must NOT re-register + sent.clear() + await mw( + {"type": "http", "path": "/dummy/y", "method": "GET", "headers": []}, + receive, + send, + ) + assert loads == ["json"], "register_fn called twice for the same feature" + + # Non-matching path must not trigger anything + await mw( + {"type": "http", "path": "/unrelated", "method": "GET", "headers": []}, + receive, + send, + ) + assert loads == ["json"] + + @pytest.mark.asyncio + async def test_concurrent_first_requests_only_register_once(self): + """ + Two requests to the same prefix arriving in parallel must result in + exactly one `register_fn` invocation — the lock prevents the import + + register from racing with itself. + """ + from fastapi import FastAPI + + from litellm.proxy._lazy_features import ( + LazyFeature, + LazyFeatureMiddleware, + ) + + loads = [] + + def slow_register(app, module): + loads.append(getattr(module, "__name__", "?")) + + feat = LazyFeature( + name="dummy_concurrent", + module_path="json", + path_prefixes=("/dummy_c",), + register_fn=slow_register, + ) + + async def downstream(scope, receive, send): + await send({"type": "http.response.start", "status": 200, "headers": []}) + await send({"type": "http.response.body", "body": b""}) + + target_app = FastAPI() + mw = LazyFeatureMiddleware(downstream, fastapi_app=target_app, features=(feat,)) + + async def receive(): + return {"type": "http.request", "body": b"", "more_body": False} + + sent: list = [] + + async def send(message): + sent.append(message) + + async def hit(): + await mw( + { + "type": "http", + "path": "/dummy_c/x", + "method": "GET", + "headers": [], + }, + receive, + send, + ) + + await asyncio.gather(hit(), hit(), hit(), hit(), hit()) + assert loads == [ + "json" + ], f"expected one registration despite concurrent first hits, got {loads}" + + @pytest.mark.asyncio + async def test_failing_import_does_not_loop(self): + """ + If a feature's module can't be imported, the middleware should mark it + loaded anyway so subsequent requests don't repeatedly retry the failing + import (which would amplify the cost on every request). + """ + from fastapi import FastAPI + + from litellm.proxy._lazy_features import ( + LazyFeature, + LazyFeatureMiddleware, + ) + + attempts = [] + + def fail_register(app, module): + attempts.append("called") + raise RuntimeError("boom") + + feat = LazyFeature( + name="failing", + module_path="json", + path_prefixes=("/fail",), + register_fn=fail_register, + ) + + async def downstream(scope, receive, send): + await send({"type": "http.response.start", "status": 200, "headers": []}) + await send({"type": "http.response.body", "body": b""}) + + target_app = FastAPI() + mw = LazyFeatureMiddleware(downstream, fastapi_app=target_app, features=(feat,)) + + async def receive(): + return {"type": "http.request", "body": b"", "more_body": False} + + sent: list = [] + + async def send(message): + sent.append(message) + + for _ in range(3): + await mw( + {"type": "http", "path": "/fail/x", "method": "GET", "headers": []}, + receive, + send, + ) + assert attempts == [ + "called" + ], f"failing register_fn should be invoked once, not on every request; got {attempts}" diff --git a/tests/test_litellm/proxy/vector_store_endpoints/test_vector_store_endpoints.py b/tests/test_litellm/proxy/vector_store_endpoints/test_vector_store_endpoints.py index 44cc5cc445..1e596aa567 100644 --- a/tests/test_litellm/proxy/vector_store_endpoints/test_vector_store_endpoints.py +++ b/tests/test_litellm/proxy/vector_store_endpoints/test_vector_store_endpoints.py @@ -786,8 +786,23 @@ class TestVectorStoreManagementEndpointsExist: - POST /vector_store/info - POST /vector_store/update """ + import importlib + + from litellm.proxy._lazy_features import LAZY_FEATURES from litellm.proxy.proxy_server import app + # Force-register the lazy vector_store_management routes so the + # assertions can find them. + already_registered = any( + getattr(r, "path", None) == "/vector_store/new" for r in app.routes + ) + if not already_registered: + for feat in LAZY_FEATURES: + if feat.name == "vector_store_management": + module = importlib.import_module(feat.module_path) + feat.register_fn(app, module) + break + # Define expected endpoints expected_endpoints = [ ("POST", "/vector_store/new"), From 0520d5ce117a51994a862b8df6384fa6b1a52d74 Mon Sep 17 00:00:00 2001 From: Michael-RZ-Berri Date: Tue, 28 Apr 2026 17:05:36 -0700 Subject: [PATCH 31/57] [Fix] Unify cost calc in success_handler dict and typed branches (#26629) * Unify cost calc in success_handler dict and typed branches * Trim verbose comments and docstrings --------- Co-authored-by: Michael Riad Zaky Co-authored-by: Michael Riad Zaky --- litellm/litellm_core_utils/litellm_logging.py | 17 +-- .../test_litellm_logging.py | 138 ++++++++++++++++++ 2 files changed, 142 insertions(+), 13 deletions(-) diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index fb103afea0..829c1c9ca0 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -1467,6 +1467,8 @@ class Logging(LiteLLMLoggingBaseClass): LiteLLMRealtimeStreamLoggingObject, OpenAIModerationResponse, "SearchResponse", + dict, + list, ], cache_hit: Optional[bool] = None, litellm_model_name: Optional[str] = None, @@ -1744,6 +1746,7 @@ class Logging(LiteLLMLoggingBaseClass): start_time, end_time, ): + """Resolve hidden params, compute response cost, and emit the standard logging payload.""" hidden_params = getattr(logging_result, "_hidden_params", {}) if hidden_params: if self.model_call_details.get("litellm_params") is not None: @@ -1877,24 +1880,12 @@ class Logging(LiteLLMLoggingBaseClass): ): if self._is_recognized_call_type_for_logging( logging_result=logging_result - ): + ) or isinstance(logging_result, (dict, list)): self._process_hidden_params_and_response_cost( logging_result=logging_result, start_time=start_time, end_time=end_time, ) - elif isinstance(result, dict) or isinstance(result, list): - self.model_call_details["standard_logging_object"] = ( - self._build_standard_logging_payload( - result, start_time, end_time - ) - ) - if ( - standard_logging_payload := self.model_call_details.get( - "standard_logging_object" - ) - ) is not None: - emit_standard_logging_payload(standard_logging_payload) elif standard_logging_object is not None: self.model_call_details["standard_logging_object"] = ( standard_logging_object diff --git a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py index 3348118a02..1764d9c609 100644 --- a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py +++ b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py @@ -2534,3 +2534,141 @@ def test_get_standard_logging_object_payload_includes_litellm_call_id(logging_ob assert payload is not None assert payload["litellm_call_id"] == call_id + + +def _make_dict_logging_obj(): + """Build a Logging instance configured for a non-streaming dict result.""" + obj = LitellmLogging( + model="claude-haiku-4-5@20251001", + messages=[{"role": "user", "content": "hi"}], + stream=False, + call_type="acompletion", + litellm_call_id="test-call-id", + start_time=time.time(), + function_id="test-fn", + ) + obj.model_call_details = { + "model": "claude-haiku-4-5@20251001", + "custom_llm_provider": "vertex_ai", + "litellm_params": {"metadata": {}}, + "response_cost": None, + } + return obj + + +def test_success_handler_computes_cost_for_dict_response(): + """Non-streaming dict responses run through the cost calculator.""" + logging_obj = _make_dict_logging_obj() + expected_cost = 0.42 + with ( + patch.object( + logging_obj, + "_response_cost_calculator", + return_value=expected_cost, + ) as mock_calc, + patch.object( + logging_obj, + "_build_standard_logging_payload", + return_value={"response_cost": expected_cost}, + ), + patch( + "litellm.litellm_core_utils.litellm_logging.emit_standard_logging_payload" + ), + patch.object( + logging_obj, + "_is_recognized_call_type_for_logging", + return_value=False, + ), + patch.object( + logging_obj, + "_transform_usage_objects", + side_effect=lambda result: result, + ), + ): + logging_obj.success_handler( + result={"id": "msg_1"}, + start_time=time.time(), + end_time=time.time(), + ) + mock_calc.assert_called_once() + assert logging_obj.model_call_details["response_cost"] == expected_cost + + +def test_success_handler_preserves_precomputed_cost_for_dict_response(): + """Precomputed response_cost on model_call_details must not be overwritten.""" + logging_obj = _make_dict_logging_obj() + precomputed_cost = 1.23 + logging_obj.model_call_details["response_cost"] = precomputed_cost + with ( + patch.object( + logging_obj, + "_response_cost_calculator", + return_value=9.99, + ) as mock_calc, + patch.object( + logging_obj, + "_build_standard_logging_payload", + return_value={"response_cost": precomputed_cost}, + ), + patch( + "litellm.litellm_core_utils.litellm_logging.emit_standard_logging_payload" + ), + patch.object( + logging_obj, + "_is_recognized_call_type_for_logging", + return_value=False, + ), + patch.object( + logging_obj, + "_transform_usage_objects", + side_effect=lambda result: result, + ), + ): + logging_obj.success_handler( + result={"id": "msg_2"}, + start_time=time.time(), + end_time=time.time(), + ) + mock_calc.assert_not_called() + assert logging_obj.model_call_details["response_cost"] == precomputed_cost + + +def test_success_handler_unified_helper_runs_for_typed_results(): + """Recognized typed responses still flow through the unified helper.""" + logging_obj = _make_dict_logging_obj() + expected_cost = 0.10 + typed_result = MagicMock() + typed_result._hidden_params = {} + + with ( + patch.object( + logging_obj, + "_response_cost_calculator", + return_value=expected_cost, + ) as mock_calc, + patch.object( + logging_obj, + "_build_standard_logging_payload", + return_value={"response_cost": expected_cost}, + ), + patch( + "litellm.litellm_core_utils.litellm_logging.emit_standard_logging_payload" + ), + patch.object( + logging_obj, + "_is_recognized_call_type_for_logging", + return_value=True, + ), + patch.object( + logging_obj, + "_transform_usage_objects", + side_effect=lambda result: result, + ), + ): + logging_obj.success_handler( + result=typed_result, + start_time=time.time(), + end_time=time.time(), + ) + mock_calc.assert_called_once() + assert logging_obj.model_call_details["response_cost"] == expected_cost From fd32f29e39ad54aa058779dbb2c5f91f2946a39f Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Tue, 28 Apr 2026 17:21:41 -0700 Subject: [PATCH 32/57] Revert "lazy-load optional feature routers on first request (#26534)" (#26727) This reverts commit 21ed38971d244c0a034604f6439c0584d55b4d20. --- litellm/proxy/_lazy_features.py | 307 ------------------ litellm/proxy/proxy_server.py | 126 +++++-- tests/proxy_unit_tests/test_proxy_routes.py | 14 - tests/test_litellm/proxy/test_proxy_server.py | 252 -------------- .../test_vector_store_endpoints.py | 15 - 5 files changed, 105 insertions(+), 609 deletions(-) delete mode 100644 litellm/proxy/_lazy_features.py diff --git a/litellm/proxy/_lazy_features.py b/litellm/proxy/_lazy_features.py deleted file mode 100644 index 450c9483f3..0000000000 --- a/litellm/proxy/_lazy_features.py +++ /dev/null @@ -1,307 +0,0 @@ -""" -Lazy registration for optional feature routers. Each LAZY_FEATURES entry -imports its module only on the first request matching its path prefix, -saving ~700 MB at idle for deployments that don't use these features. -First hit pays the import cost (1-3 s for heavy modules); /openapi.json -omits each feature's routes until the feature is warmed. -""" - -import asyncio -import importlib -from dataclasses import dataclass, field -from typing import TYPE_CHECKING, Callable, Tuple - -from starlette.types import Receive, Scope, Send - -from litellm._logging import verbose_proxy_logger - -if TYPE_CHECKING: - from fastapi import FastAPI - - -def _include_router(attr_name: str = "router") -> Callable[["FastAPI", object], None]: - def _register(app: "FastAPI", module: object) -> None: - app.include_router(getattr(module, attr_name)) - - return _register - - -def _mount_app( - prefix: str, attr_name: str = "app" -) -> Callable[["FastAPI", object], None]: - def _register(app: "FastAPI", module: object) -> None: - app.mount(path=prefix, app=getattr(module, attr_name)) - - return _register - - -@dataclass(frozen=True) -class LazyFeature: - name: str - module_path: str - path_prefixes: Tuple[str, ...] - register_fn: Callable[["FastAPI", object], None] = field( - default_factory=lambda: _include_router("router") - ) - # For routes whose path has a leading parameter (e.g. /{server}/authorize) - # — startswith can't match those, so the matcher also checks endswith. - path_suffixes: Tuple[str, ...] = () - - -LAZY_FEATURES: Tuple[LazyFeature, ...] = ( - LazyFeature( - name="guardrails", - module_path="litellm.proxy.guardrails.guardrail_endpoints", - path_prefixes=( - "/guardrails", - "/v2/guardrails", - "/apply_guardrail", - "/policies/usage", - ), - ), - LazyFeature( - name="policies", - module_path="litellm.proxy.management_endpoints.policy_endpoints", - # Trailing slash to avoid matching /policies/... (policy_engine). - path_prefixes=("/policy/", "/utils/test_policies_and_guardrails"), - ), - LazyFeature( - name="policy_engine", - module_path="litellm.proxy.policy_engine.policy_endpoints", - path_prefixes=("/policies",), - ), - LazyFeature( - name="policy_resolve", - module_path="litellm.proxy.policy_engine.policy_resolve_endpoints", - path_prefixes=("/policies/resolve", "/policies/attachments/estimate-impact"), - ), - LazyFeature( - name="agents", - module_path="litellm.proxy.agent_endpoints.endpoints", - path_prefixes=("/v1/agents", "/agents", "/agent/"), - ), - LazyFeature( - name="a2a", - module_path="litellm.proxy.agent_endpoints.a2a_endpoints", - path_prefixes=("/a2a", "/v1/a2a"), - ), - LazyFeature( - name="vector_stores", - module_path="litellm.proxy.vector_store_endpoints.endpoints", - path_prefixes=("/v1/vector_stores", "/vector_stores", "/v1/indexes"), - ), - LazyFeature( - name="vector_store_management", - module_path="litellm.proxy.vector_store_endpoints.management_endpoints", - # Trailing slash to avoid matching /vector_stores/... (vector_stores). - path_prefixes=("/vector_store/", "/v1/vector_store/"), - ), - LazyFeature( - name="vector_store_files", - # Routes appear under both /v1/vector_stores/{id}/files and the - # un-versioned form, so both prefixes must trigger the load. - module_path="litellm.proxy.vector_store_files_endpoints.endpoints", - path_prefixes=("/v1/vector_stores", "/vector_stores"), - ), - LazyFeature( - name="tools", - module_path="litellm.proxy.management_endpoints.tool_management_endpoints", - path_prefixes=("/v1/tool", "/tool"), - ), - LazyFeature( - name="search_tools", - module_path="litellm.proxy.search_endpoints.search_tool_management", - path_prefixes=("/search_tools",), - ), - # mcp_management owns most /v1/mcp/* admin routes; mcp_app is the mounted - # streaming sub-app at /mcp. - LazyFeature( - name="mcp_management", - module_path="litellm.proxy.management_endpoints.mcp_management_endpoints", - path_prefixes=("/v1/mcp/",), - ), - LazyFeature( - # Also serves /.well-known/oauth-* (OAuth metadata discovery). - # No /mcp/oauth prefix here: the mounted /mcp sub-app would - # shadow it, and there are no actual routes there anyway. - name="mcp_byok_oauth", - module_path="litellm.proxy._experimental.mcp_server.byok_oauth_endpoints", - path_prefixes=("/v1/mcp/oauth", "/.well-known/oauth-"), - ), - LazyFeature( - # Serves OAuth dance endpoints (/authorize, /token, /callback, - # /register) plus several /.well-known/ discovery URLs at the proxy - # root — needed for MCP-over-OAuth flows even before /mcp is hit. - name="mcp_discoverable", - module_path="litellm.proxy._experimental.mcp_server.discoverable_endpoints", - path_prefixes=( - "/.well-known/oauth-", - "/.well-known/openid-configuration", - "/.well-known/jwks.json", - "/authorize", - "/token", - "/callback", - "/register", - ), - # Catches the /{mcp_server_name}/authorize|token|register variants. - path_suffixes=("/authorize", "/token", "/register"), - ), - LazyFeature( - name="mcp_rest", - module_path="litellm.proxy._experimental.mcp_server.rest_endpoints", - path_prefixes=("/mcp-rest",), - ), - LazyFeature( - # Hardcoded /mcp matches BASE_MCP_ROUTE; importing the constant - # here would defeat lazy loading. - name="mcp_app", - module_path="litellm.proxy._experimental.mcp_server.server", - path_prefixes=("/mcp",), - register_fn=_mount_app("/mcp", attr_name="app"), - ), - LazyFeature( - name="config_overrides", - module_path="litellm.proxy.management_endpoints.config_override_endpoints", - path_prefixes=("/config_overrides",), - ), - LazyFeature( - name="realtime", - module_path="litellm.proxy.realtime_endpoints.endpoints", - path_prefixes=("/openai/v1/realtime", "/v1/realtime", "/realtime"), - ), - LazyFeature( - name="anthropic_passthrough", - module_path="litellm.proxy.anthropic_endpoints.endpoints", - path_prefixes=("/v1/messages", "/anthropic", "/api/event_logging"), - ), - LazyFeature( - name="anthropic_skills", - module_path="litellm.proxy.anthropic_endpoints.skills_endpoints", - path_prefixes=("/v1/skills", "/skills"), - ), - LazyFeature( - name="langfuse_passthrough", - module_path="litellm.proxy.vertex_ai_endpoints.langfuse_endpoints", - path_prefixes=("/langfuse",), - ), - LazyFeature( - name="evals", - module_path="litellm.proxy.openai_evals_endpoints.endpoints", - path_prefixes=("/v1/evals", "/evals"), - ), - LazyFeature( - name="claude_code_marketplace", - module_path="litellm.proxy.anthropic_endpoints.claude_code_endpoints", - path_prefixes=("/claude-code",), - register_fn=_include_router("claude_code_marketplace_router"), - ), - LazyFeature( - name="scim", - module_path="litellm.proxy.management_endpoints.scim.scim_v2", - path_prefixes=("/scim",), - register_fn=_include_router("scim_router"), - ), - LazyFeature( - name="cloudzero", - module_path="litellm.proxy.spend_tracking.cloudzero_endpoints", - path_prefixes=("/cloudzero",), - ), - LazyFeature( - name="vantage", - module_path="litellm.proxy.spend_tracking.vantage_endpoints", - path_prefixes=("/vantage",), - ), - LazyFeature( - name="usage_ai", - module_path="litellm.proxy.management_endpoints.usage_endpoints", - path_prefixes=("/usage/ai",), - ), - LazyFeature( - name="prompts", - module_path="litellm.proxy.prompts.prompt_endpoints", - path_prefixes=("/prompts", "/utils/dotprompt_json_converter"), - ), - LazyFeature( - name="jwt_mappings", - module_path="litellm.proxy.management_endpoints.jwt_key_mapping_endpoints", - path_prefixes=("/jwt/key/mapping",), - ), - LazyFeature( - name="compliance", - module_path="litellm.proxy.management_endpoints.compliance_endpoints", - path_prefixes=("/compliance",), - ), - LazyFeature( - name="access_groups", - module_path="litellm.proxy.management_endpoints.access_group_endpoints", - path_prefixes=("/access_group", "/v1/access_group", "/v1/unified_access_group"), - ), -) - - -class LazyFeatureMiddleware: - """ASGI middleware that imports + registers a feature router on first - matching request. Idempotent; once loaded, subsequent requests skip.""" - - def __init__( - self, - app, - fastapi_app: "FastAPI", - features: Tuple[LazyFeature, ...] = LAZY_FEATURES, - ): - self.app = app - self._fastapi_app = fastapi_app - self._features = features - self._loaded: set = set() - # Per-feature locks so independent features can load in parallel. - self._locks: dict = {} - - async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None: - # Short-circuit once every feature has loaded. - if scope["type"] in ("http", "websocket") and len(self._loaded) < len( - self._features - ): - path = scope.get("path", "") - for feat in self._features: - if feat.module_path in self._loaded: - continue - if any(path.startswith(p) for p in feat.path_prefixes) or any( - path.endswith(s) for s in feat.path_suffixes - ): - await self._load(feat) - await self.app(scope, receive, send) - - async def _load(self, feat: LazyFeature) -> None: - lock = self._locks.setdefault(feat.module_path, asyncio.Lock()) - async with lock: - if feat.module_path in self._loaded: - return - try: - # Import on a thread (heavy modules take 1-3 s). register_fn - # mutates app.router.routes, so it stays on the loop thread. - loop = asyncio.get_running_loop() - module = await loop.run_in_executor( - None, importlib.import_module, feat.module_path - ) - feat.register_fn(self._fastapi_app, module) - self._loaded.add(feat.module_path) - self._fastapi_app.openapi_schema = None - verbose_proxy_logger.info( - "Lazy-loaded optional feature %r (module: %s)", - feat.name, - feat.module_path, - ) - except Exception as exc: - # Mark loaded anyway so we don't retry on every request. - self._loaded.add(feat.module_path) - verbose_proxy_logger.warning( - "Failed to lazy-load optional feature %r (module: %s): %s. " - "This feature's endpoints will return 404 until restart.", - feat.name, - feat.module_path, - exc, - ) - - -def attach_lazy_features(app: "FastAPI") -> None: - app.add_middleware(LazyFeatureMiddleware, fastapi_app=app) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index c03a63f211..8f676df04c 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -235,11 +235,37 @@ from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging from litellm.litellm_core_utils.sensitive_data_masker import SensitiveDataMasker from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler from litellm.llms.vertex_ai.vertex_llm_base import VertexBase +from litellm.proxy._experimental.mcp_server.byok_oauth_endpoints import ( + router as mcp_byok_oauth_router, +) +from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( + router as mcp_discoverable_endpoints_router, +) +from litellm.proxy._experimental.mcp_server.rest_endpoints import ( + router as mcp_rest_endpoints_router, +) +from litellm.proxy._experimental.mcp_server.server import app as mcp_app +from litellm.proxy._experimental.mcp_server.tool_registry import ( + global_mcp_tool_registry, +) from litellm.proxy._types import * -from litellm.proxy._lazy_features import attach_lazy_features +from litellm.proxy.agent_endpoints.a2a_endpoints import router as a2a_router +from litellm.proxy.agent_endpoints.agent_registry import global_agent_registry +from litellm.proxy.agent_endpoints.endpoints import router as agent_endpoints_router +from litellm.proxy.agent_endpoints.model_list_helpers import ( + append_agents_to_model_group, + append_agents_to_model_info, +) from litellm.proxy.analytics_endpoints.analytics_endpoints import ( router as analytics_router, ) +from litellm.proxy.anthropic_endpoints.claude_code_endpoints import ( + claude_code_marketplace_router, +) +from litellm.proxy.anthropic_endpoints.endpoints import router as anthropic_router +from litellm.proxy.anthropic_endpoints.skills_endpoints import ( + router as anthropic_skills_router, +) from litellm.proxy.auth.auth_checks import ( ExperimentalUIJWTToken, get_team_object, @@ -302,6 +328,7 @@ from litellm.proxy.discovery_endpoints import ui_discovery_endpoints_router from litellm.proxy.fine_tuning_endpoints.endpoints import router as fine_tuning_router from litellm.proxy.fine_tuning_endpoints.endpoints import set_fine_tuning_config from litellm.proxy.google_endpoints.endpoints import router as google_router +from litellm.proxy.guardrails.guardrail_endpoints import router as guardrails_router from litellm.proxy.guardrails.init_guardrails import ( init_guardrails_v2, initialize_guardrails, @@ -317,6 +344,9 @@ from litellm.proxy.hooks.prompt_injection_detection import ( from litellm.proxy.hooks.proxy_track_cost_callback import _ProxyDBLogger from litellm.proxy.image_endpoints.endpoints import router as image_router from litellm.proxy.litellm_pre_call_utils import add_litellm_data_to_request +from litellm.proxy.management_endpoints.access_group_endpoints import ( + router as access_group_router, +) from litellm.proxy.management_endpoints.budget_management_endpoints import ( router as budget_management_router, ) @@ -330,6 +360,12 @@ from litellm.proxy.management_endpoints.common_utils import ( _user_has_admin_privileges, admin_can_invite_user, ) +from litellm.proxy.management_endpoints.compliance_endpoints import ( + router as compliance_router, +) +from litellm.proxy.management_endpoints.config_override_endpoints import ( + router as config_override_router, +) from litellm.proxy.management_endpoints.cost_tracking_settings import ( router as cost_tracking_settings_router, ) @@ -343,6 +379,9 @@ from litellm.proxy.management_endpoints.internal_user_endpoints import ( router as internal_user_router, ) from litellm.proxy.management_endpoints.internal_user_endpoints import user_update +from litellm.proxy.management_endpoints.jwt_key_mapping_endpoints import ( + router as jwt_key_mapping_router, +) from litellm.proxy.management_endpoints.key_management_endpoints import ( delete_verification_tokens, duration_in_seconds, @@ -351,6 +390,9 @@ from litellm.proxy.management_endpoints.key_management_endpoints import ( from litellm.proxy.management_endpoints.key_management_endpoints import ( router as key_management_router, ) +from litellm.proxy.management_endpoints.mcp_management_endpoints import ( + router as mcp_management_router, +) from litellm.proxy.management_endpoints.model_access_group_management_endpoints import ( router as model_access_group_management_router, ) @@ -365,9 +407,11 @@ from litellm.proxy.management_endpoints.model_management_endpoints import ( from litellm.proxy.management_endpoints.organization_endpoints import ( router as organization_router, ) +from litellm.proxy.management_endpoints.policy_endpoints import router as policy_router from litellm.proxy.management_endpoints.router_settings_endpoints import ( router as router_settings_router, ) +from litellm.proxy.management_endpoints.scim.scim_v2 import scim_router from litellm.proxy.management_endpoints.tag_management_endpoints import ( router as tag_management_router, ) @@ -379,11 +423,15 @@ from litellm.proxy.management_endpoints.team_endpoints import ( update_team, validate_membership, ) +from litellm.proxy.management_endpoints.tool_management_endpoints import ( + router as tool_management_router, +) from litellm.proxy.memory.memory_endpoints import router as memory_router from litellm.proxy.management_endpoints.ui_sso import ( get_disabled_non_admin_personal_key_creation, ) from litellm.proxy.management_endpoints.ui_sso import router as ui_sso_router +from litellm.proxy.management_endpoints.usage_endpoints import router as usage_ai_router from litellm.proxy.management_endpoints.user_agent_analytics_endpoints import ( router as user_agent_analytics_router, ) @@ -393,6 +441,7 @@ from litellm.proxy.middleware.in_flight_requests_middleware import ( ) from litellm.proxy.middleware.prometheus_auth_middleware import PrometheusAuthMiddleware from litellm.proxy.ocr_endpoints.endpoints import router as ocr_router +from litellm.proxy.openai_evals_endpoints.endpoints import router as evals_router from litellm.proxy.openai_files_endpoints.files_endpoints import ( router as openai_files_router, ) @@ -412,16 +461,27 @@ from litellm.proxy.pass_through_endpoints.pass_through_endpoints import ( from litellm.proxy.pass_through_endpoints.pass_through_endpoints import ( router as pass_through_router, ) +from litellm.proxy.policy_engine.policy_endpoints import router as policy_crud_router +from litellm.proxy.policy_engine.policy_resolve_endpoints import ( + router as policy_resolve_router, +) +from litellm.proxy.prompts.prompt_endpoints import router as prompts_router from litellm.proxy.public_endpoints import router as public_endpoints_router from litellm.proxy.rag_endpoints.endpoints import router as rag_router +from litellm.proxy.realtime_endpoints.endpoints import router as webrtc_router from litellm.proxy.rerank_endpoints.endpoints import router as rerank_router from litellm.proxy.response_api_endpoints.endpoints import router as response_router from litellm.proxy.route_llm_request import route_request from litellm.proxy.search_endpoints.endpoints import router as search_router +from litellm.proxy.search_endpoints.search_tool_management import ( + router as search_tool_management_router, +) +from litellm.proxy.spend_tracking.cloudzero_endpoints import router as cloudzero_router from litellm.proxy.spend_tracking.spend_management_endpoints import ( router as spend_management_router, ) from litellm.proxy.spend_tracking.spend_tracking_utils import get_logging_payload +from litellm.proxy.spend_tracking.vantage_endpoints import router as vantage_router from litellm.proxy.types_utils.utils import get_instance_fn from litellm.proxy.ui_crud_endpoints.proxy_setting_endpoints import ( router as ui_crud_endpoints_router, @@ -451,6 +511,16 @@ from litellm.proxy.utils import ( prefetch_config_params, update_spend, ) +from litellm.proxy.vector_store_endpoints.endpoints import router as vector_store_router +from litellm.proxy.vector_store_endpoints.management_endpoints import ( + router as vector_store_management_router, +) +from litellm.proxy.vector_store_files_endpoints.endpoints import ( + router as vector_store_files_router, +) +from litellm.proxy.vertex_ai_endpoints.langfuse_endpoints import ( + router as langfuse_router, +) from litellm.proxy.video_endpoints.endpoints import router as video_router from litellm.router import ( AssistantsTypedDict, @@ -3784,19 +3854,11 @@ class ProxyConfig: ## MCP TOOLS mcp_tools_config = config.get("mcp_tools", None) if mcp_tools_config: - from litellm.proxy._experimental.mcp_server.tool_registry import ( - global_mcp_tool_registry, - ) - global_mcp_tool_registry.load_tools_from_config(mcp_tools_config) ## AGENTS agent_config = config.get("agent_list", None) if agent_config: - from litellm.proxy.agent_endpoints.agent_registry import ( - global_agent_registry, - ) - global_agent_registry.load_agents_from_config(agent_config) # type: ignore mcp_servers_config = config.get("mcp_servers", None) @@ -10514,10 +10576,6 @@ async def model_info_v2( verbose_proxy_logger.debug("all_models: %s", all_models) # Append A2A agents to models list - from litellm.proxy.agent_endpoints.model_list_helpers import ( - append_agents_to_model_info, - ) - all_models = await append_agents_to_model_info( models=all_models, user_api_key_dict=user_api_key_dict, @@ -11367,10 +11425,6 @@ async def model_group_info( ) # Append A2A agents to model groups - from litellm.proxy.agent_endpoints.model_list_helpers import ( - append_agents_to_model_group, - ) - model_groups = await append_agents_to_model_group( model_groups=model_groups, user_api_key_dict=user_api_key_dict, @@ -14176,40 +14230,65 @@ app.include_router(container_router) app.include_router(search_router) app.include_router(image_router) app.include_router(fine_tuning_router) +app.include_router(vector_store_router) +app.include_router(vector_store_management_router) +app.include_router(vector_store_files_router) app.include_router(credential_router) app.include_router(llm_passthrough_router) +app.include_router(webrtc_router) +app.include_router(mcp_management_router) +app.include_router(mcp_byok_oauth_router) +app.include_router(anthropic_router) +app.include_router(anthropic_skills_router) +app.include_router(evals_router) +app.include_router(claude_code_marketplace_router) +app.include_router(google_router) +app.include_router(langfuse_router) app.include_router(pass_through_router) app.include_router(health_router) app.include_router(key_management_router) app.include_router(internal_user_router) app.include_router(team_router) app.include_router(ui_sso_router) +app.include_router(scim_router) app.include_router(organization_router) app.include_router(customer_router) app.include_router(spend_management_router) +app.include_router(cloudzero_router) +app.include_router(vantage_router) app.include_router(caching_router) app.include_router(analytics_router) +app.include_router(guardrails_router) +app.include_router(policy_router) +app.include_router(usage_ai_router) +app.include_router(policy_crud_router) +app.include_router(policy_resolve_router) +app.include_router(search_tool_management_router) +app.include_router(prompts_router) app.include_router(callback_management_endpoints_router) app.include_router(debugging_endpoints_router) app.include_router(ui_crud_endpoints_router) app.include_router(openai_files_router) app.include_router(team_callback_router) +app.include_router(jwt_key_mapping_router) app.include_router(budget_management_router) app.include_router(model_management_router) app.include_router(model_access_group_management_router) app.include_router(tag_management_router) +app.include_router(tool_management_router) app.include_router(memory_router) app.include_router(cost_tracking_settings_router) app.include_router(router_settings_router) app.include_router(fallback_management_router) app.include_router(cache_settings_router) +app.include_router(config_override_router) app.include_router(user_agent_analytics_router) app.include_router(enterprise_router) app.include_router(ui_discovery_endpoints_router) -# Eager: /models/{name}:method overlaps with the OpenAI /models endpoint. -app.include_router(google_router) - -attach_lazy_features(app) +app.include_router(agent_endpoints_router) +app.include_router(compliance_router) +app.include_router(a2a_router) +app.include_router(access_group_router) async def _stream_mcp_asgi_response( @@ -14442,3 +14521,8 @@ async def dynamic_mcp_route(mcp_server_name: str, request: Request): f"Error handling dynamic MCP route for {mcp_server_name}: {str(e)}" ) raise HTTPException(status_code=500, detail=f"Internal server error: {str(e)}") + + +app.mount(path=BASE_MCP_ROUTE, app=mcp_app) +app.include_router(mcp_rest_endpoints_router) +app.include_router(mcp_discoverable_endpoints_router) diff --git a/tests/proxy_unit_tests/test_proxy_routes.py b/tests/proxy_unit_tests/test_proxy_routes.py index 67eca5206d..812e4e1ac4 100644 --- a/tests/proxy_unit_tests/test_proxy_routes.py +++ b/tests/proxy_unit_tests/test_proxy_routes.py @@ -39,20 +39,6 @@ def test_routes_on_litellm_proxy(): this prevents accidentelly deleting /threads, or /batches etc """ - # Force-load lazy features so the test sees the full route set. Continue - # on per-feature import failure — the assertion below still catches - # missing-route regressions. - import importlib - - from litellm.proxy._lazy_features import LAZY_FEATURES - - for feat in LAZY_FEATURES: - try: - module = importlib.import_module(feat.module_path) - feat.register_fn(app, module) - except Exception as exc: - print(f"warning: failed to force-load {feat.name}: {exc}") - _all_routes = [] for route in app.routes: diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index 7a96f6cbd1..1f4f82a64e 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -5471,255 +5471,3 @@ async def test_reseed_warms_cache_even_on_zero_db_spend(): finally: ps.spend_counter_cache = orig_counter ps.prisma_client = orig_prisma - - -# --------------------------------------------------------------------------- -# Lazy feature loading (LazyFeatureMiddleware) — verifies that optional -# routers are NOT imported at module load and ARE imported on first request -# to a matching path prefix. The same module isn't re-imported on subsequent -# requests. -# --------------------------------------------------------------------------- - - -import sys - - -class TestLazyFeatureRegistry: - """Sanity checks on the registry shape — guards against accidental edits.""" - - def test_registry_entries_have_required_fields(self): - from litellm.proxy._lazy_features import LAZY_FEATURES, LazyFeature - - assert len(LAZY_FEATURES) > 0 - for feat in LAZY_FEATURES: - assert isinstance(feat, LazyFeature) - assert feat.name - assert feat.module_path - assert feat.path_prefixes - assert all(p.startswith("/") for p in feat.path_prefixes) - assert callable(feat.register_fn) - - def test_registry_names_unique(self): - from litellm.proxy._lazy_features import LAZY_FEATURES - - names = [f.name for f in LAZY_FEATURES] - assert len(names) == len(set(names)), "duplicate feature names" - - -class TestLazyFeaturesNotImportedAtStartup: - """ - The whole point of the refactor: gated feature modules must NOT be - present in `sys.modules` immediately after `proxy_server` imports. - """ - - def test_heavy_modules_absent_at_startup(self): - # Force a fresh `proxy_server` import in a subprocess so other tests - # in this run (which may have triggered lazy loads via the TestClient) - # don't pollute the result. - import subprocess - - check = ( - "import sys; " - "from litellm.proxy.proxy_server import app; " # noqa: F401 - "heavy = [" - "'litellm.proxy._experimental.mcp_server.rest_endpoints'," - "'litellm.proxy._experimental.mcp_server.server'," - "'litellm.proxy.management_endpoints.config_override_endpoints'," - "'litellm.proxy.guardrails.guardrail_endpoints'," - "'litellm.proxy.openai_evals_endpoints.endpoints'," - "]; " - "still_present = [m for m in heavy if m in sys.modules]; " - "print('PRESENT_AT_STARTUP:', still_present)" - ) - result = subprocess.run( - [sys.executable, "-c", check], - capture_output=True, - text=True, - timeout=120, - ) - # Last non-empty line of stdout (skip warnings printed before) - out_lines = [ - line for line in result.stdout.strip().splitlines() if line.strip() - ] - report = next((line for line in out_lines if "PRESENT_AT_STARTUP" in line), "") - assert report, f"no report emitted (stderr: {result.stderr[-500:]})" - assert ( - "PRESENT_AT_STARTUP: []" in report - ), f"expected no heavy modules at startup, got: {report}" - - -class TestLazyFeatureMiddleware: - """Behavior of the middleware itself, exercised in isolation.""" - - @pytest.mark.asyncio - async def test_first_request_triggers_load_subsequent_does_not(self): - from fastapi import FastAPI - - from litellm.proxy._lazy_features import ( - LazyFeature, - LazyFeatureMiddleware, - ) - - loads = [] - - def fake_register(app, module): - loads.append(getattr(module, "__name__", "?")) - - feat = LazyFeature( - name="dummy", - module_path="json", # any always-importable stdlib module - path_prefixes=("/dummy",), - register_fn=fake_register, - ) - - # Build a minimal ASGI receiver to satisfy the middleware contract - async def downstream(scope, receive, send): - # echo back; no-op handler - await send({"type": "http.response.start", "status": 200, "headers": []}) - await send({"type": "http.response.body", "body": b""}) - - target_app = FastAPI() - mw = LazyFeatureMiddleware(downstream, fastapi_app=target_app, features=(feat,)) - - async def receive(): - return {"type": "http.request", "body": b"", "more_body": False} - - sent: list = [] - - async def send(message): - sent.append(message) - - # First request matching the prefix triggers register - await mw( - {"type": "http", "path": "/dummy/x", "method": "GET", "headers": []}, - receive, - send, - ) - assert loads == ["json"] - - # Second matching request must NOT re-register - sent.clear() - await mw( - {"type": "http", "path": "/dummy/y", "method": "GET", "headers": []}, - receive, - send, - ) - assert loads == ["json"], "register_fn called twice for the same feature" - - # Non-matching path must not trigger anything - await mw( - {"type": "http", "path": "/unrelated", "method": "GET", "headers": []}, - receive, - send, - ) - assert loads == ["json"] - - @pytest.mark.asyncio - async def test_concurrent_first_requests_only_register_once(self): - """ - Two requests to the same prefix arriving in parallel must result in - exactly one `register_fn` invocation — the lock prevents the import + - register from racing with itself. - """ - from fastapi import FastAPI - - from litellm.proxy._lazy_features import ( - LazyFeature, - LazyFeatureMiddleware, - ) - - loads = [] - - def slow_register(app, module): - loads.append(getattr(module, "__name__", "?")) - - feat = LazyFeature( - name="dummy_concurrent", - module_path="json", - path_prefixes=("/dummy_c",), - register_fn=slow_register, - ) - - async def downstream(scope, receive, send): - await send({"type": "http.response.start", "status": 200, "headers": []}) - await send({"type": "http.response.body", "body": b""}) - - target_app = FastAPI() - mw = LazyFeatureMiddleware(downstream, fastapi_app=target_app, features=(feat,)) - - async def receive(): - return {"type": "http.request", "body": b"", "more_body": False} - - sent: list = [] - - async def send(message): - sent.append(message) - - async def hit(): - await mw( - { - "type": "http", - "path": "/dummy_c/x", - "method": "GET", - "headers": [], - }, - receive, - send, - ) - - await asyncio.gather(hit(), hit(), hit(), hit(), hit()) - assert loads == [ - "json" - ], f"expected one registration despite concurrent first hits, got {loads}" - - @pytest.mark.asyncio - async def test_failing_import_does_not_loop(self): - """ - If a feature's module can't be imported, the middleware should mark it - loaded anyway so subsequent requests don't repeatedly retry the failing - import (which would amplify the cost on every request). - """ - from fastapi import FastAPI - - from litellm.proxy._lazy_features import ( - LazyFeature, - LazyFeatureMiddleware, - ) - - attempts = [] - - def fail_register(app, module): - attempts.append("called") - raise RuntimeError("boom") - - feat = LazyFeature( - name="failing", - module_path="json", - path_prefixes=("/fail",), - register_fn=fail_register, - ) - - async def downstream(scope, receive, send): - await send({"type": "http.response.start", "status": 200, "headers": []}) - await send({"type": "http.response.body", "body": b""}) - - target_app = FastAPI() - mw = LazyFeatureMiddleware(downstream, fastapi_app=target_app, features=(feat,)) - - async def receive(): - return {"type": "http.request", "body": b"", "more_body": False} - - sent: list = [] - - async def send(message): - sent.append(message) - - for _ in range(3): - await mw( - {"type": "http", "path": "/fail/x", "method": "GET", "headers": []}, - receive, - send, - ) - assert attempts == [ - "called" - ], f"failing register_fn should be invoked once, not on every request; got {attempts}" diff --git a/tests/test_litellm/proxy/vector_store_endpoints/test_vector_store_endpoints.py b/tests/test_litellm/proxy/vector_store_endpoints/test_vector_store_endpoints.py index 1e596aa567..44cc5cc445 100644 --- a/tests/test_litellm/proxy/vector_store_endpoints/test_vector_store_endpoints.py +++ b/tests/test_litellm/proxy/vector_store_endpoints/test_vector_store_endpoints.py @@ -786,23 +786,8 @@ class TestVectorStoreManagementEndpointsExist: - POST /vector_store/info - POST /vector_store/update """ - import importlib - - from litellm.proxy._lazy_features import LAZY_FEATURES from litellm.proxy.proxy_server import app - # Force-register the lazy vector_store_management routes so the - # assertions can find them. - already_registered = any( - getattr(r, "path", None) == "/vector_store/new" for r in app.routes - ) - if not already_registered: - for feat in LAZY_FEATURES: - if feat.name == "vector_store_management": - module = importlib.import_module(feat.module_path) - feat.register_fn(app, module) - break - # Define expected endpoints expected_endpoints = [ ("POST", "/vector_store/new"), From b07e1c03418312c4bbc617f611a75f64d64a64ec Mon Sep 17 00:00:00 2001 From: Michael Riad Zaky Date: Tue, 28 Apr 2026 17:38:42 -0700 Subject: [PATCH 33/57] drop response body from vertex/bedrock transformation errors --- .../bedrock/chat/converse_transformation.py | 4 +-- .../vertex_and_google_ai_studio_gemini.py | 8 ++--- .../chat/test_converse_transformation.py | 36 +++++++++++++++++++ ...test_vertex_and_google_ai_studio_gemini.py | 29 +++++++++++++++ 4 files changed, 71 insertions(+), 6 deletions(-) diff --git a/litellm/llms/bedrock/chat/converse_transformation.py b/litellm/llms/bedrock/chat/converse_transformation.py index a27153365d..61a7d4c08d 100644 --- a/litellm/llms/bedrock/chat/converse_transformation.py +++ b/litellm/llms/bedrock/chat/converse_transformation.py @@ -1942,8 +1942,8 @@ class AmazonConverseConfig(BaseConfig): completion_response = ConverseResponseBlock(**response.json()) # type: ignore except Exception as e: raise BedrockError( - message="Received={}, Error converting to valid response block={}. File an issue if litellm error - https://github.com/BerriAI/litellm/issues".format( - response.text, str(e) + message="Error converting to valid response block={}. File an issue if litellm error - https://github.com/BerriAI/litellm/issues".format( + str(e) ), status_code=422, ) diff --git a/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py b/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py index a90e05919f..474ddb402a 100644 --- a/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py +++ b/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py @@ -2395,8 +2395,8 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): completion_response = GenerateContentResponseBody(**raw_response.json()) # type: ignore except Exception as e: raise VertexAIError( - message="Received={}, Error converting to valid response block={}. File an issue if litellm error - https://github.com/BerriAI/litellm/issues".format( - raw_response.text, str(e) + message="Error converting to valid response block={}. File an issue if litellm error - https://github.com/BerriAI/litellm/issues".format( + str(e) ), status_code=422, headers=raw_response.headers, @@ -2530,8 +2530,8 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): except Exception as e: raise VertexAIError( - message="Received={}, Error converting to valid response block={}. File an issue if litellm error - https://github.com/BerriAI/litellm/issues".format( - completion_response, str(e) + message="Error converting to valid response block={}. File an issue if litellm error - https://github.com/BerriAI/litellm/issues".format( + str(e) ), status_code=422, headers=raw_response.headers, diff --git a/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py b/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py index 8e53e57f1e..633aed38a0 100644 --- a/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py +++ b/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py @@ -4146,3 +4146,39 @@ def test_transform_response_finish_reason_stop_when_json_mode_filters_all_tools( # finish_reason must be "stop", not "tool_calls" assert result.choices[0].finish_reason == "stop" + + +def test_transform_response_does_not_leak_body_on_parse_failure(): + from litellm.llms.bedrock.common_utils import BedrockError + + leaky_body = {"output": {"message": {"content": [{"text": "secret content"}]}}} + + class MockResponse: + def json(self): + return leaky_body + + @property + def text(self): + return json.dumps(leaky_body) + + with patch( + "litellm.llms.bedrock.chat.converse_transformation.ConverseResponseBlock", + side_effect=KeyError("missing required field"), + ): + with pytest.raises(BedrockError) as exc_info: + AmazonConverseConfig()._transform_response( + model="bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0", + response=MockResponse(), + model_response=ModelResponse(), + stream=False, + logging_obj=None, + optional_params={}, + api_key=None, + data=None, + messages=[], + encoding=None, + ) + + msg = str(exc_info.value) + assert "secret content" not in msg + assert "Error converting to valid response block" in msg diff --git a/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py b/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py index 0118b39ddd..353d19b019 100644 --- a/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py +++ b/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py @@ -4261,3 +4261,32 @@ def test_sync_streaming_uses_custom_client(): # Verify that gemini_client is in the partial's keywords assert "gemini_client" in partial_make_sync_call.keywords assert partial_make_sync_call.keywords["gemini_client"] is mock_client + + +def test_transform_response_does_not_leak_body_on_parse_failure(): + leaky_body = {"candidates": [{"content": {"parts": [{"text": "secret content"}]}}]} + raw_response = MagicMock() + raw_response.json.return_value = leaky_body + raw_response.text = json.dumps(leaky_body) + raw_response.headers = {} + + with patch( + "litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini.GenerateContentResponseBody", + side_effect=KeyError("missing required field"), + ): + with pytest.raises(VertexAIError) as exc_info: + VertexGeminiConfig().transform_response( + model="gemini-pro", + raw_response=raw_response, + model_response=ModelResponse(), + logging_obj=MagicMock(), + request_data={}, + messages=[], + optional_params={}, + litellm_params={}, + encoding=None, + ) + + msg = str(exc_info.value) + assert "secret content" not in msg + assert "Error converting to valid response block" in msg From f8bb29aebfb4530a66120f55157f5dc144e136b9 Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Tue, 28 Apr 2026 17:43:17 -0700 Subject: [PATCH 34/57] =?UTF-8?q?bump:=20version=201.83.14=20=E2=86=92=201?= =?UTF-8?q?.84.0?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- pyproject.toml | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index a15fa5a06a..657632d69e 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "litellm" -version = "1.83.14" +version = "1.84.0" description = "Library to easily interface with LLM API providers" readme = "README.md" requires-python = ">=3.10, <3.14" @@ -236,7 +236,7 @@ source-exclude = [ profile = "black" [tool.commitizen] -version = "1.83.14" +version = "1.84.0" version_files = [ "pyproject.toml:^version", ] From b4d9006f92c14b6fa7161b286b303561211ce04d Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Tue, 28 Apr 2026 17:43:36 -0700 Subject: [PATCH 35/57] uv lock --- uv.lock | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/uv.lock b/uv.lock index 53f032cfba..f837e2b5ef 100644 --- a/uv.lock +++ b/uv.lock @@ -9,7 +9,7 @@ resolution-markers = [ ] [options] -exclude-newer = "2026-04-23T02:32:27.506663Z" +exclude-newer = "0001-01-01T00:00:00Z" # This has no effect and is included for backwards compatibility when using relative exclude-newer values. exclude-newer-span = "P3D" [manifest] @@ -3085,7 +3085,7 @@ wheels = [ [[package]] name = "litellm" -version = "1.83.14" +version = "1.84.0" source = { editable = "." } dependencies = [ { name = "aiohttp" }, From fc49c181bca73d63d3b861d7fc2a46a834a3ba8e Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Wed, 29 Apr 2026 01:41:24 +0000 Subject: [PATCH 36/57] feat(mcp): opt-in short-ID tool prefix to stay under 60-char tool name limit MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Adds LITELLM_USE_SHORT_MCP_TOOL_PREFIX. When enabled, tool / prompt / resource / resource-template names emitted from MCP servers are prefixed with a deterministic three-character base62 ID derived from the server's server_id (SHA-256 → base62) instead of the (potentially long) alias / server_name. This keeps namespaced tool names well under the 60-character upper bound enforced by some model APIs while still letting us distinguish MCP-routed tools from local tools. Behavioural notes: - Default off — when the env var is unset, the long-prefix behaviour is unchanged. The plan is to flip the default in a future release and remove the gate after a deprecation window. - Prefix derivation is deterministic, so it is stable across processes, workers and restarts without any persistence layer. - Reverse-lookup is tolerant: _create_prefixed_tools registers every known prefix form (alias / server_name / server_id / short ID) in the routing map and _get_mcp_server_from_tool_name resolves any of them. Old clients holding cached long-prefixed names continue to route correctly even after the flag is enabled. - _get_allowed_mcp_servers_from_mcp_server_names accepts the short prefix in /mcp/{server_name}-style URLs. - The OpenAPI tool-listing path now filters by the active server prefix instead of server.name so spec-backed servers benefit too. Co-authored-by: Mateo Wang --- .../mcp_server/mcp_server_manager.py | 65 +++--- .../proxy/_experimental/mcp_server/server.py | 26 ++- .../proxy/_experimental/mcp_server/utils.py | 104 ++++++++- .../mcp_server/test_short_mcp_tool_prefix.py | 207 ++++++++++++++++++ 4 files changed, 362 insertions(+), 40 deletions(-) create mode 100644 tests/test_litellm/proxy/_experimental/mcp_server/test_short_mcp_tool_prefix.py diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index 251f271903..a2f0517f81 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -52,6 +52,7 @@ from litellm.proxy._experimental.mcp_server.utils import ( add_server_prefix_to_name, get_server_prefix, is_tool_name_prefixed, + iter_known_server_prefixes, merge_mcp_headers, normalize_server_name, split_server_prefix_from_name, @@ -1236,7 +1237,11 @@ class MCPServerManager: ## HANDLE OPENAPI TOOLS if server.spec_path: - _tools = global_mcp_tool_registry.list_tools(tool_prefix=server.name) + # OpenAPI tools were stored in the registry under the prefix + # active at registration time — fetch by that same prefix. + _tools = global_mcp_tool_registry.list_tools( + tool_prefix=get_server_prefix(server) + ) tools = global_mcp_tool_registry.convert_tools_to_mcp_sdk_tool_type( _tools ) @@ -1838,9 +1843,13 @@ class MCPServerManager: tool_copy.name = name_to_use prefixed_tools.append(tool_copy) - # Update tool to server mapping for resolution (support both forms) + # Register every known prefix form (alias, server_name, server_id, + # short ID) so call_tool can resolve regardless of which form a + # caller / cached client is using. self.tool_name_to_mcp_server_name_mapping[original_name] = prefix - self.tool_name_to_mcp_server_name_mapping[prefixed_name] = prefix + for known_prefix in iter_known_server_prefixes(server): + qualified = add_server_prefix_to_name(original_name, known_prefix) + self.tool_name_to_mcp_server_name_mapping[qualified] = prefix verbose_logger.info( f"Successfully fetched {len(prefixed_tools)} tools from server {server.name}" @@ -2601,37 +2610,43 @@ class MCPServerManager: Returns: MCPServer if found, None otherwise """ + registry_servers = list(self.get_registry().values()) + + # Build prefix → server lookup covering every known form a tool name + # may take (alias / server_name / server_id / short ID). This is what + # makes the short-prefix mode work without breaking historical names. + prefix_to_server: Dict[str, MCPServer] = {} + for server in registry_servers: + for known_prefix in iter_known_server_prefixes(server): + normalised = normalize_server_name(known_prefix) + prefix_to_server.setdefault(normalised, server) + # First try with the original tool name if tool_name in self.tool_name_to_mcp_server_name_mapping: server_name = self.tool_name_to_mcp_server_name_mapping[tool_name] - for server in self.get_registry().values(): - if normalize_server_name(server.name) == normalize_server_name( - server_name - ): + normalised_lookup = normalize_server_name(server_name) + if normalised_lookup in prefix_to_server: + return prefix_to_server[normalised_lookup] + for server in registry_servers: + if normalize_server_name(server.name) == normalised_lookup: return server - # If not found and tool name is prefixed, try extracting server name from prefix - known_prefixes = { - normalize_server_name(get_server_prefix(s)) - for s in self.get_registry().values() - if get_server_prefix(s) - } - if is_tool_name_prefixed(tool_name, known_server_prefixes=known_prefixes): + # If not found and tool name is prefixed, extract the prefix and + # match against any known form. + if is_tool_name_prefixed( + tool_name, known_server_prefixes=set(prefix_to_server.keys()) + ): ( original_tool_name, server_name_from_prefix, ) = split_server_prefix_from_name(tool_name) - if original_tool_name in self.tool_name_to_mcp_server_name_mapping: - for server in self.get_registry().values(): - if server.server_name is None: - if normalize_server_name(server.name) == normalize_server_name( - server_name_from_prefix - ): - return server - elif normalize_server_name( - server.server_name - ) == normalize_server_name(server_name_from_prefix): - return server + normalised_prefix = normalize_server_name(server_name_from_prefix) + server = prefix_to_server.get(normalised_prefix) + if server is not None and ( + original_tool_name in self.tool_name_to_mcp_server_name_mapping + or tool_name in self.tool_name_to_mcp_server_name_mapping + ): + return server return None diff --git a/litellm/proxy/_experimental/mcp_server/server.py b/litellm/proxy/_experimental/mcp_server/server.py index c2e998f01e..3924687a0b 100644 --- a/litellm/proxy/_experimental/mcp_server/server.py +++ b/litellm/proxy/_experimental/mcp_server/server.py @@ -49,6 +49,7 @@ from litellm.proxy._experimental.mcp_server.utils import ( LITELLM_MCP_SERVER_VERSION, add_server_prefix_to_name, get_server_prefix, + iter_known_server_prefixes, ) from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.auth.ip_address_utils import IPAddressUtils @@ -711,14 +712,13 @@ if MCP_AVAILABLE: for server in allowed_mcp_servers: if server: match_list = [ - s.lower() - for s in [ - server.alias, - server.server_name, - server.server_id, - ] - if s is not None + s.lower() for s in iter_known_server_prefixes(server) if s ] + # Always accept server_id even if it isn't part of the + # current prefix form (iter_known_server_prefixes only + # yields it when no other identifier exists). + if server.server_id: + match_list.append(server.server_id.lower()) if server_or_group.lower() in match_list: filtered_server[server.server_id] = server @@ -2031,11 +2031,13 @@ if MCP_AVAILABLE: # Remove prefix from tool name for logging and processing original_tool_name, server_name = split_server_prefix_from_name(name) - # If tool name is unprefixed, resolve its server so we can enforce permissions - if not server_name: - mcp_server = global_mcp_server_manager._get_mcp_server_from_tool_name(name) - if mcp_server: - server_name = mcp_server.name + # Resolve the actual MCP server up-front so the permission check uses + # the canonical server.name even when the tool name is prefixed with a + # short ID (LITELLM_USE_SHORT_MCP_TOOL_PREFIX) that doesn't match the + # server's display name directly. + mcp_server = global_mcp_server_manager._get_mcp_server_from_tool_name(name) + if mcp_server is not None: + server_name = mcp_server.name # Only enforce server-level permissions when we can resolve a server if server_name: diff --git a/litellm/proxy/_experimental/mcp_server/utils.py b/litellm/proxy/_experimental/mcp_server/utils.py index 146ba10bb7..42910cc790 100644 --- a/litellm/proxy/_experimental/mcp_server/utils.py +++ b/litellm/proxy/_experimental/mcp_server/utils.py @@ -2,10 +2,11 @@ MCP Server Utilities """ -from typing import Any, Dict, Mapping, Optional, Tuple +from typing import Any, Dict, Iterator, Mapping, Optional, Tuple -import os +import hashlib import importlib +import os # Constants LITELLM_MCP_SERVER_NAME = "litellm-mcp-server" @@ -14,6 +15,63 @@ LITELLM_MCP_SERVER_DESCRIPTION = "MCP Server for LiteLLM" MCP_TOOL_PREFIX_SEPARATOR = os.environ.get("MCP_TOOL_PREFIX_SEPARATOR", "-") MCP_TOOL_PREFIX_FORMAT = "{server_name}{separator}{tool_name}" +# --------------------------------------------------------------------------- +# Short-ID tool prefix (opt-in) +# --------------------------------------------------------------------------- +# When LITELLM_USE_SHORT_MCP_TOOL_PREFIX is truthy the prefix attached to MCP +# tool / prompt / resource / resource-template names switches from the +# (potentially long) human-readable server name to a deterministic three +# character base62 ID derived from the server's ``server_id``. +# +# Why three characters and base62 ([0-9A-Za-z])? +# * 62**3 = 238_328 distinct IDs — the chance of a real local tool name +# happening to begin with the exact prefix LiteLLM assigned to a given +# MCP server is negligible in practice. +# * The IDs are short enough that prefixed tool names stay well under the +# 60-character upper bound enforced by some model APIs (Anthropic etc.) +# even for long upstream tool names. +# * The mapping is deterministic (SHA-256 of ``server_id`` → first three +# base62 chars), which means the prefix is stable across processes, +# workers and restarts without any persistence layer. Two servers with +# different ``server_id`` values can in principle hash to the same +# three chars, but for the reverse-lookup path we register every known +# form of the prefix anyway, so a collision only affects the cosmetic +# emitted name, not routing correctness. +# +# This flag is intentionally opt-in for the first release so customers can +# migrate. It will become the default in a future release. +SHORT_MCP_TOOL_PREFIX_LENGTH = 3 +_BASE62_ALPHABET = "0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz" + + +def is_short_mcp_tool_prefix_enabled() -> bool: + """Return True when the short-ID tool prefix mode is enabled. + + Read at call time (not import time) so tests and runtime config changes + take effect without reimporting the module. + """ + raw = os.environ.get("LITELLM_USE_SHORT_MCP_TOOL_PREFIX", "") + return raw.strip().lower() in ("1", "true", "yes", "on") + + +def compute_short_server_prefix(server_id: str) -> str: + """Derive the deterministic three-character base62 prefix for a server. + + Uses SHA-256 of the server_id and folds the first eight bytes into a + base62 string. An empty server_id raises ValueError — short prefixes + require a stable identifier to be deterministic. + """ + if not server_id: + raise ValueError("compute_short_server_prefix requires a non-empty server_id") + + digest = hashlib.sha256(server_id.encode("utf-8")).digest() + value = int.from_bytes(digest[:8], "big") + chars = [] + for _ in range(SHORT_MCP_TOOL_PREFIX_LENGTH): + value, idx = divmod(value, len(_BASE62_ALPHABET)) + chars.append(_BASE62_ALPHABET[idx]) + return "".join(reversed(chars)) + def is_mcp_available() -> bool: """ @@ -82,7 +140,18 @@ def add_server_prefix_to_name(name: str, server_name: str) -> str: def get_server_prefix(server: Any) -> str: - """Return the prefix for a server: alias if present, else server_name, else server_id""" + """Return the prefix for a server. + + When the short-prefix mode is enabled (``LITELLM_USE_SHORT_MCP_TOOL_PREFIX``) + a deterministic three-character base62 ID derived from ``server_id`` is + returned. Otherwise we fall back to the historical behaviour: alias if + present, else server_name, else server_id. + """ + if is_short_mcp_tool_prefix_enabled(): + server_id = getattr(server, "server_id", None) + if server_id: + return compute_short_server_prefix(server_id) + if hasattr(server, "alias") and server.alias: return server.alias if hasattr(server, "server_name") and server.server_name: @@ -92,6 +161,35 @@ def get_server_prefix(server: Any) -> str: return "" +def iter_known_server_prefixes(server: Any) -> Iterator[str]: + """Yield every prefix form that may appear in tool names for ``server``. + + Always includes the *current* prefix returned by ``get_server_prefix``. + Additionally yields the historical (alias / server_name / server_id) and + short-ID forms so the routing layer can resolve tool names regardless of + which prefix mode was active when the client first observed them. + """ + seen = set() + + def _emit(value: Optional[str]) -> Iterator[str]: + if value and value not in seen: + seen.add(value) + yield value + + yield from _emit(get_server_prefix(server)) + + server_id = getattr(server, "server_id", None) + if server_id: + try: + yield from _emit(compute_short_server_prefix(server_id)) + except ValueError: + pass + + yield from _emit(getattr(server, "alias", None)) + yield from _emit(getattr(server, "server_name", None)) + yield from _emit(server_id) + + def split_server_prefix_from_name(prefixed_name: str) -> Tuple[str, str]: """Return the unprefixed name plus the server name used as prefix.""" if MCP_TOOL_PREFIX_SEPARATOR in prefixed_name: diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_short_mcp_tool_prefix.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_short_mcp_tool_prefix.py new file mode 100644 index 0000000000..aef6546c81 --- /dev/null +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_short_mcp_tool_prefix.py @@ -0,0 +1,207 @@ +""" +Tests for the short-ID MCP tool prefix (LITELLM_USE_SHORT_MCP_TOOL_PREFIX). + +The short-prefix mode swaps the historical alias/server_name prefix on +tool names for a deterministic three-character base62 ID derived from the +server's ``server_id``. This keeps tool names well below the 60-char +upper bound enforced by some model APIs while remaining stable across +processes/restarts and tolerant of mixed-version clients. +""" + +from typing import List + +import pytest +from mcp.types import Tool as MCPTool + +from litellm.proxy._experimental.mcp_server.mcp_server_manager import MCPServerManager +from litellm.proxy._experimental.mcp_server.utils import ( + SHORT_MCP_TOOL_PREFIX_LENGTH, + add_server_prefix_to_name, + compute_short_server_prefix, + get_server_prefix, + is_short_mcp_tool_prefix_enabled, + iter_known_server_prefixes, +) +from litellm.types.mcp_server.mcp_server_manager import MCPServer + + +def _make_server( + *, + server_id: str = "abcdef-1234", + server_name: str = "github_onprem", + alias: str = "github_onprem", +) -> MCPServer: + return MCPServer( + server_id=server_id, + name=alias or server_name, + alias=alias, + server_name=server_name, + transport="http", + ) + + +@pytest.fixture(autouse=True) +def _reset_env(monkeypatch): + monkeypatch.delenv("LITELLM_USE_SHORT_MCP_TOOL_PREFIX", raising=False) + yield + + +# --------------------------------------------------------------------------- +# Pure helpers +# --------------------------------------------------------------------------- + + +class TestShortPrefixHelpers: + def test_short_prefix_is_three_base62_chars(self): + prefix = compute_short_server_prefix("any-server-id") + assert len(prefix) == SHORT_MCP_TOOL_PREFIX_LENGTH + assert prefix.isalnum() and prefix.isascii() + + def test_short_prefix_is_deterministic(self): + assert compute_short_server_prefix("abc") == compute_short_server_prefix("abc") + assert compute_short_server_prefix("abc") != compute_short_server_prefix("abd") + + def test_short_prefix_requires_server_id(self): + with pytest.raises(ValueError): + compute_short_server_prefix("") + + def test_flag_defaults_to_false(self): + assert is_short_mcp_tool_prefix_enabled() is False + + @pytest.mark.parametrize("value", ["1", "true", "TRUE", "yes", "On"]) + def test_flag_truthy_values(self, monkeypatch, value): + monkeypatch.setenv("LITELLM_USE_SHORT_MCP_TOOL_PREFIX", value) + assert is_short_mcp_tool_prefix_enabled() is True + + @pytest.mark.parametrize("value", ["0", "false", "no", "off", ""]) + def test_flag_falsey_values(self, monkeypatch, value): + monkeypatch.setenv("LITELLM_USE_SHORT_MCP_TOOL_PREFIX", value) + assert is_short_mcp_tool_prefix_enabled() is False + + +# --------------------------------------------------------------------------- +# get_server_prefix behaviour +# --------------------------------------------------------------------------- + + +class TestGetServerPrefix: + def test_default_mode_uses_alias(self): + server = _make_server(alias="github_onprem", server_name="github_onprem") + assert get_server_prefix(server) == "github_onprem" + + def test_short_mode_uses_short_id(self, monkeypatch): + monkeypatch.setenv("LITELLM_USE_SHORT_MCP_TOOL_PREFIX", "true") + server = _make_server(server_id="abcdef-1234") + prefix = get_server_prefix(server) + assert prefix == compute_short_server_prefix("abcdef-1234") + assert len(prefix) == SHORT_MCP_TOOL_PREFIX_LENGTH + + def test_short_mode_falls_back_when_no_server_id(self, monkeypatch): + monkeypatch.setenv("LITELLM_USE_SHORT_MCP_TOOL_PREFIX", "true") + + class _Bare: + alias = "fallback_alias" + server_name = None + server_id = None + + assert get_server_prefix(_Bare()) == "fallback_alias" + + +# --------------------------------------------------------------------------- +# iter_known_server_prefixes — covers reverse-lookup tolerance +# --------------------------------------------------------------------------- + + +class TestIterKnownServerPrefixes: + def test_default_mode_includes_short_id_too(self): + server = _make_server() + prefixes = list(iter_known_server_prefixes(server)) + # Contains the live prefix and every known form so that mixed-mode + # clients can be resolved. + assert "github_onprem" in prefixes + assert compute_short_server_prefix(server.server_id) in prefixes + + def test_short_mode_still_yields_long_forms(self, monkeypatch): + monkeypatch.setenv("LITELLM_USE_SHORT_MCP_TOOL_PREFIX", "true") + server = _make_server() + prefixes = list(iter_known_server_prefixes(server)) + assert "github_onprem" in prefixes + assert compute_short_server_prefix(server.server_id) in prefixes + + +# --------------------------------------------------------------------------- +# Manager-level behaviour: list + reverse-lookup +# --------------------------------------------------------------------------- + + +def _stub_tools() -> List[MCPTool]: + return [ + MCPTool(name="get_repo", description="", inputSchema={"type": "object"}), + MCPTool(name="list_issues", description="", inputSchema={"type": "object"}), + ] + + +class TestManagerShortPrefix: + def test_list_tools_uses_short_prefix_when_flag_on(self, monkeypatch): + monkeypatch.setenv("LITELLM_USE_SHORT_MCP_TOOL_PREFIX", "true") + manager = MCPServerManager() + server = _make_server() + + out = manager._create_prefixed_tools(_stub_tools(), server) + + short = compute_short_server_prefix(server.server_id) + assert {t.name for t in out} == {f"{short}-get_repo", f"{short}-list_issues"} + + def test_call_tool_lookup_resolves_short_prefix(self, monkeypatch): + monkeypatch.setenv("LITELLM_USE_SHORT_MCP_TOOL_PREFIX", "true") + manager = MCPServerManager() + server = _make_server() + manager.registry[server.server_id] = server + manager._create_prefixed_tools(_stub_tools(), server) + + short = compute_short_server_prefix(server.server_id) + resolved = manager._get_mcp_server_from_tool_name(f"{short}-get_repo") + assert resolved is server + + def test_call_tool_lookup_resolves_long_prefix_in_short_mode(self, monkeypatch): + """Old clients that cached the long-prefix name must still route.""" + monkeypatch.setenv("LITELLM_USE_SHORT_MCP_TOOL_PREFIX", "true") + manager = MCPServerManager() + server = _make_server() + manager.registry[server.server_id] = server + manager._create_prefixed_tools(_stub_tools(), server) + + resolved = manager._get_mcp_server_from_tool_name("github_onprem-get_repo") + assert resolved is server + + def test_default_mode_unchanged(self): + manager = MCPServerManager() + server = _make_server() + + out = manager._create_prefixed_tools(_stub_tools(), server) + + assert {t.name for t in out} == { + "github_onprem-get_repo", + "github_onprem-list_issues", + } + assert ( + manager._get_mcp_server_from_tool_name("github_onprem-get_repo") is None + ) # registry empty + manager.registry[server.server_id] = server + assert ( + manager._get_mcp_server_from_tool_name("github_onprem-get_repo") is server + ) + + def test_total_tool_name_length_short_enough(self, monkeypatch): + """The short prefix keeps tool names under the 60-char limit even + when the upstream tool name is itself reasonably long.""" + monkeypatch.setenv("LITELLM_USE_SHORT_MCP_TOOL_PREFIX", "true") + long_server_name = "a" * 50 + server = _make_server( + server_id="server-id-1", + server_name=long_server_name, + alias=long_server_name, + ) + prefix = get_server_prefix(server) + full = add_server_prefix_to_name("get_repo", prefix) + assert len(full) < 60 From 1da1eb661b3aafd39d8705da66c915d330a258b8 Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Tue, 28 Apr 2026 19:33:18 -0700 Subject: [PATCH 37/57] ci(release): accept PEP 440 tag forms in create-release workflow The tag validator required a leading `v`, so dispatching create-release with `1.84.0` (or `1.84.0rc1`, `1.84.0.dev42`, `1.84.0.post1`) failed even though those are the new naming convention. Make the leading `v` optional in both create-release.yml and create-release-branch.yml so both legacy (`v1.83.10-stable`, `v1.83.14.rc.1`, `v1.82.3.dev.9`, `v1.82.3-stable.patch.4`, `v1.83.13-nightly`) and new PEP 440 forms are accepted during the transition. Refresh the input descriptions to show the new examples. --- .github/workflows/create-release-branch.yml | 8 ++++---- .github/workflows/create-release.yml | 6 +++--- 2 files changed, 7 insertions(+), 7 deletions(-) diff --git a/.github/workflows/create-release-branch.yml b/.github/workflows/create-release-branch.yml index 13b76c94df..ec2651306f 100644 --- a/.github/workflows/create-release-branch.yml +++ b/.github/workflows/create-release-branch.yml @@ -4,7 +4,7 @@ on: workflow_dispatch: inputs: tag: - description: "Release tag (e.g. v1.83.0-stable) — branch will be named release/" + description: "Release tag (e.g. 1.84.0, 1.84.0rc1, 1.84.0.dev42, 1.84.0.post1; legacy v1.83.10-stable still accepted) — branch will be named release/" required: true type: string commit_hash: @@ -14,7 +14,7 @@ on: workflow_call: inputs: tag: - description: "Release tag" + description: "Release tag (e.g. 1.84.0, 1.84.0rc1, 1.84.0.dev42, 1.84.0.post1; legacy v1.83.10-stable still accepted)" required: true type: string commit_hash: @@ -40,8 +40,8 @@ jobs: echo "::error::commit_hash must be a full 40-character commit SHA" exit 1 fi - if ! echo "${TAG}" | grep -qE '^v[0-9]+\.[0-9]+\.[0-9]+'; then - echo "::error::tag must start with vX.Y.Z" + if ! echo "${TAG}" | grep -qE '^v?[0-9]+\.[0-9]+\.[0-9]+'; then + echo "::error::tag must start with X.Y.Z (optional leading v), e.g. 1.84.0, 1.84.0rc1, 1.84.0.dev42, or v1.83.10-stable" exit 1 fi diff --git a/.github/workflows/create-release.yml b/.github/workflows/create-release.yml index 68ab397d82..c0aec1687e 100644 --- a/.github/workflows/create-release.yml +++ b/.github/workflows/create-release.yml @@ -4,7 +4,7 @@ on: workflow_dispatch: inputs: tag: - description: "Release tag (e.g. v1.83.0-stable)" + description: "Release tag (e.g. 1.84.0, 1.84.0rc1, 1.84.0.dev42, 1.84.0.post1; legacy v1.83.10-stable still accepted)" required: true type: string commit_hash: @@ -30,8 +30,8 @@ jobs: echo "::error::commit_hash must be a full 40-character commit SHA" exit 1 fi - if ! echo "${TAG}" | grep -qE '^v[0-9]+\.[0-9]+\.[0-9]+'; then - echo "::error::tag must start with vX.Y.Z" + if ! echo "${TAG}" | grep -qE '^v?[0-9]+\.[0-9]+\.[0-9]+'; then + echo "::error::tag must start with X.Y.Z (optional leading v), e.g. 1.84.0, 1.84.0rc1, 1.84.0.dev42, or v1.83.10-stable" exit 1 fi From 3a5980804c2aef672ef1f324e101e2c6694285f7 Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Tue, 28 Apr 2026 19:38:13 -0700 Subject: [PATCH 38/57] ci(release): mark rc / dev / nightly tags as GitHub pre-releases `prerelease: false` was hardcoded, so dispatching create-release with `1.84.0rc1`, `1.84.0.dev42`, or legacy `v1.83.13-nightly` would publish them as stable releases on the GitHub Releases page. Derive the flag from the tag instead. The detector matches `rc`, `.dev`, `nightly`, `alpha`, `beta`. PEP 440 post-releases (`1.84.0.post1`) and legacy `-stable[.patch.N]` are stable maintenance releases per PEP 440, so they intentionally do not match. --- .github/workflows/create-release.yml | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/.github/workflows/create-release.yml b/.github/workflows/create-release.yml index c0aec1687e..39d078267f 100644 --- a/.github/workflows/create-release.yml +++ b/.github/workflows/create-release.yml @@ -45,6 +45,11 @@ jobs: const tag = process.env.TAG; const commitHash = process.env.COMMIT_HASH; + // Mark RC / dev / nightly / alpha / beta tags as GitHub pre-releases. + // PEP 440 post-releases (e.g. `1.84.0.post1`) and legacy `-stable[.patch.N]` + // are stable maintenance releases, not pre-releases. + const isPrerelease = /(?:rc|nightly|alpha|beta|\.dev)/i.test(tag); + const cosignSection = [ `## Verify Docker Image Signature`, ``, @@ -89,7 +94,7 @@ jobs: target_commitish: commitHash, name: tag, owner: context.repo.owner, - prerelease: false, + prerelease: isPrerelease, repo: context.repo.repo, tag_name: tag, }); From 4e827446d263a228eda7ce8365196574b162322a Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 28 Apr 2026 19:58:56 -0700 Subject: [PATCH 39/57] fix: type error --- .../proxy/_experimental/mcp_server/mcp_server_manager.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index a2f0517f81..8e473c1cb2 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -2641,12 +2641,12 @@ class MCPServerManager: server_name_from_prefix, ) = split_server_prefix_from_name(tool_name) normalised_prefix = normalize_server_name(server_name_from_prefix) - server = prefix_to_server.get(normalised_prefix) - if server is not None and ( + matched_server = prefix_to_server.get(normalised_prefix) + if matched_server is not None and ( original_tool_name in self.tool_name_to_mcp_server_name_mapping or tool_name in self.tool_name_to_mcp_server_name_mapping ): - return server + return matched_server return None From cf74f55b7983e6e0fc58c96d9077803afb37c6ae Mon Sep 17 00:00:00 2001 From: Sameer Kankute Date: Wed, 29 Apr 2026 08:34:31 +0530 Subject: [PATCH 40/57] Fix extra body error --- .../llms/azure_ai/vector_stores/transformation.py | 2 +- litellm/llms/base_llm/vector_store/transformation.py | 6 +++--- litellm/llms/bedrock/vector_stores/transformation.py | 2 +- litellm/llms/custom_httpx/llm_http_handler.py | 6 +++--- litellm/llms/gemini/vector_stores/transformation.py | 2 +- litellm/llms/milvus/vector_stores/transformation.py | 2 +- litellm/llms/openai/vector_stores/transformation.py | 2 +- .../llms/pg_vector/vector_stores/transformation.py | 4 ++-- litellm/llms/ragflow/vector_stores/transformation.py | 2 +- .../llms/s3_vectors/vector_stores/transformation.py | 4 ++-- .../vector_stores/rag_api/transformation.py | 2 +- .../vector_stores/search_api/transformation.py | 2 +- .../test_bedrock_knowledgebase_hook.py | 1 + .../test_bedrock_vector_store_transformation.py | 12 ++++++------ .../vector_stores/test_s3_vectors_transformation.py | 2 +- .../vector_store_tests/test_ragflow_vector_store.py | 1 + 16 files changed, 27 insertions(+), 25 deletions(-) diff --git a/litellm/llms/azure_ai/vector_stores/transformation.py b/litellm/llms/azure_ai/vector_stores/transformation.py index d2c8206ca9..d1b93c9e7a 100644 --- a/litellm/llms/azure_ai/vector_stores/transformation.py +++ b/litellm/llms/azure_ai/vector_stores/transformation.py @@ -92,10 +92,10 @@ class AzureAIVectorStoreConfig(BaseVectorStoreConfig, BaseAzureLLM): vector_store_id: str, query: Union[str, List[str]], vector_store_search_optional_params: VectorStoreSearchOptionalRequestParams, - extra_body: Optional[Dict[str, Any]], api_base: str, litellm_logging_obj: LiteLLMLoggingObj, litellm_params: dict, + extra_body: Optional[Dict[str, Any]] = None, ) -> Tuple[str, Dict[str, Any]]: """ Transform search request for Azure AI Search API diff --git a/litellm/llms/base_llm/vector_store/transformation.py b/litellm/llms/base_llm/vector_store/transformation.py index 49d2f72db7..85a9c83826 100644 --- a/litellm/llms/base_llm/vector_store/transformation.py +++ b/litellm/llms/base_llm/vector_store/transformation.py @@ -56,10 +56,10 @@ class BaseVectorStoreConfig: vector_store_id: str, query: Union[str, List[str]], vector_store_search_optional_params: VectorStoreSearchOptionalRequestParams, - extra_body: Optional[Dict[str, Any]], api_base: str, litellm_logging_obj: LiteLLMLoggingObj, litellm_params: dict, + extra_body: Optional[Dict[str, Any]] = None, ) -> Tuple[str, Dict]: pass @@ -68,10 +68,10 @@ class BaseVectorStoreConfig: vector_store_id: str, query: Union[str, List[str]], vector_store_search_optional_params: VectorStoreSearchOptionalRequestParams, - extra_body: Optional[Dict[str, Any]], api_base: str, litellm_logging_obj: LiteLLMLoggingObj, litellm_params: dict, + extra_body: Optional[Dict[str, Any]] = None, ) -> Tuple[str, Dict]: """ Optional async version of transform_search_vector_store_request. @@ -83,10 +83,10 @@ class BaseVectorStoreConfig: vector_store_id=vector_store_id, query=query, vector_store_search_optional_params=vector_store_search_optional_params, - extra_body=extra_body, api_base=api_base, litellm_logging_obj=litellm_logging_obj, litellm_params=litellm_params, + extra_body=extra_body, ) @abstractmethod diff --git a/litellm/llms/bedrock/vector_stores/transformation.py b/litellm/llms/bedrock/vector_stores/transformation.py index 4e81fa4e66..f028503c6a 100644 --- a/litellm/llms/bedrock/vector_stores/transformation.py +++ b/litellm/llms/bedrock/vector_stores/transformation.py @@ -201,10 +201,10 @@ class BedrockVectorStoreConfig(BaseVectorStoreConfig, BaseAWSLLM): vector_store_id: str, query: Union[str, List[str]], vector_store_search_optional_params: VectorStoreSearchOptionalRequestParams, - extra_body: Optional[Dict[str, Any]], api_base: str, litellm_logging_obj: LiteLLMLoggingObj, litellm_params: dict, + extra_body: Optional[Dict[str, Any]] = None, ) -> Tuple[str, Dict]: if isinstance(query, list): query = " ".join(query) diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index 99f748c0c1..d8515d2c4d 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -8582,10 +8582,10 @@ class BaseLLMHTTPHandler: vector_store_id=vector_store_id, query=query, vector_store_search_optional_params=vector_store_search_optional_params, - extra_body=extra_body, api_base=api_base, litellm_logging_obj=logging_obj, litellm_params=dict(litellm_params), + extra_body=extra_body, ) else: ( @@ -8595,10 +8595,10 @@ class BaseLLMHTTPHandler: vector_store_id=vector_store_id, query=query, vector_store_search_optional_params=vector_store_search_optional_params, - extra_body=extra_body, api_base=api_base, litellm_logging_obj=logging_obj, litellm_params=dict(litellm_params), + extra_body=extra_body, ) all_optional_params: Dict[str, Any] = dict(litellm_params) all_optional_params.update(vector_store_search_optional_params or {}) @@ -8696,10 +8696,10 @@ class BaseLLMHTTPHandler: vector_store_id=vector_store_id, query=query, vector_store_search_optional_params=vector_store_search_optional_params, - extra_body=extra_body, api_base=api_base, litellm_logging_obj=logging_obj, litellm_params=dict(litellm_params), + extra_body=extra_body, ) all_optional_params: Dict[str, Any] = dict(litellm_params) diff --git a/litellm/llms/gemini/vector_stores/transformation.py b/litellm/llms/gemini/vector_stores/transformation.py index b6cace066c..35d83bd2ad 100644 --- a/litellm/llms/gemini/vector_stores/transformation.py +++ b/litellm/llms/gemini/vector_stores/transformation.py @@ -115,10 +115,10 @@ class GeminiVectorStoreConfig(BaseVectorStoreConfig): vector_store_id: str, query: Union[str, List[str]], vector_store_search_optional_params: VectorStoreSearchOptionalRequestParams, - extra_body: Optional[Dict[str, Any]], api_base: str, litellm_logging_obj: LiteLLMLoggingObj, litellm_params: dict, + extra_body: Optional[Dict[str, Any]] = None, ) -> Tuple[str, Dict]: """ Transform search request to Gemini's generateContent format. diff --git a/litellm/llms/milvus/vector_stores/transformation.py b/litellm/llms/milvus/vector_stores/transformation.py index 8c08b78338..af78cd8dbd 100644 --- a/litellm/llms/milvus/vector_stores/transformation.py +++ b/litellm/llms/milvus/vector_stores/transformation.py @@ -127,10 +127,10 @@ class MilvusVectorStoreConfig(BaseVectorStoreConfig): vector_store_id: str, query: Union[str, List[str]], vector_store_search_optional_params: VectorStoreSearchOptionalRequestParams, - extra_body: Optional[Dict[str, Any]], api_base: str, litellm_logging_obj: LiteLLMLoggingObj, litellm_params: dict, + extra_body: Optional[Dict[str, Any]] = None, ) -> Tuple[str, Dict[str, Any]]: """ Transform search request for Azure AI Search API diff --git a/litellm/llms/openai/vector_stores/transformation.py b/litellm/llms/openai/vector_stores/transformation.py index b6eae390d9..2c11d13748 100644 --- a/litellm/llms/openai/vector_stores/transformation.py +++ b/litellm/llms/openai/vector_stores/transformation.py @@ -103,10 +103,10 @@ class OpenAIVectorStoreConfig(BaseVectorStoreConfig): vector_store_id: str, query: Union[str, List[str]], vector_store_search_optional_params: VectorStoreSearchOptionalRequestParams, - extra_body: Optional[Dict[str, Any]], api_base: str, litellm_logging_obj: LiteLLMLoggingObj, litellm_params: dict, + extra_body: Optional[Dict[str, Any]] = None, ) -> Tuple[str, Dict]: url = f"{api_base}/{vector_store_id}/search" typed_request_body = VectorStoreSearchRequest( diff --git a/litellm/llms/pg_vector/vector_stores/transformation.py b/litellm/llms/pg_vector/vector_stores/transformation.py index 8261036cae..7b22edd867 100644 --- a/litellm/llms/pg_vector/vector_stores/transformation.py +++ b/litellm/llms/pg_vector/vector_stores/transformation.py @@ -77,19 +77,19 @@ class PGVectorStoreConfig(OpenAIVectorStoreConfig): vector_store_id: str, query: Union[str, List[str]], vector_store_search_optional_params: VectorStoreSearchOptionalRequestParams, - extra_body: Optional[Dict[str, Any]], api_base: str, litellm_logging_obj: LiteLLMLoggingObj, litellm_params: dict, + extra_body: Optional[Dict[str, Any]] = None, ) -> Tuple[str, Dict]: url = f"{api_base}/{vector_store_id}/search" _, request_body = super().transform_search_vector_store_request( vector_store_id=vector_store_id, query=query, vector_store_search_optional_params=vector_store_search_optional_params, - extra_body=extra_body, api_base=api_base, litellm_logging_obj=litellm_logging_obj, litellm_params=litellm_params, + extra_body=extra_body, ) return url, request_body diff --git a/litellm/llms/ragflow/vector_stores/transformation.py b/litellm/llms/ragflow/vector_stores/transformation.py index ae28222c3c..3238d3e9c1 100644 --- a/litellm/llms/ragflow/vector_stores/transformation.py +++ b/litellm/llms/ragflow/vector_stores/transformation.py @@ -99,10 +99,10 @@ class RAGFlowVectorStoreConfig(BaseVectorStoreConfig): vector_store_id: str, query: Union[str, List[str]], vector_store_search_optional_params: VectorStoreSearchOptionalRequestParams, - extra_body: Optional[Dict[str, Any]], api_base: str, litellm_logging_obj: LiteLLMLoggingObj, litellm_params: dict, + extra_body: Optional[Dict[str, Any]] = None, ) -> Tuple[str, Dict]: """RAGFlow vector stores are management-only, search is not supported.""" raise NotImplementedError( diff --git a/litellm/llms/s3_vectors/vector_stores/transformation.py b/litellm/llms/s3_vectors/vector_stores/transformation.py index 0cf8635887..8270e99d45 100644 --- a/litellm/llms/s3_vectors/vector_stores/transformation.py +++ b/litellm/llms/s3_vectors/vector_stores/transformation.py @@ -76,10 +76,10 @@ class S3VectorsVectorStoreConfig(BaseVectorStoreConfig, BaseAWSLLM): vector_store_id: str, query: Union[str, List[str]], vector_store_search_optional_params: VectorStoreSearchOptionalRequestParams, - extra_body: Optional[Dict[str, Any]], api_base: str, litellm_logging_obj: LiteLLMLoggingObj, litellm_params: dict, + extra_body: Optional[Dict[str, Any]] = None, ) -> Tuple[str, Dict]: """Sync version - generates embedding synchronously.""" # For S3 Vectors, vector_store_id should be in format: bucket_name:index_name @@ -138,10 +138,10 @@ class S3VectorsVectorStoreConfig(BaseVectorStoreConfig, BaseAWSLLM): vector_store_id: str, query: Union[str, List[str]], vector_store_search_optional_params: VectorStoreSearchOptionalRequestParams, - extra_body: Optional[Dict[str, Any]], api_base: str, litellm_logging_obj: LiteLLMLoggingObj, litellm_params: dict, + extra_body: Optional[Dict[str, Any]] = None, ) -> Tuple[str, Dict]: """Async version - generates embedding asynchronously.""" # For S3 Vectors, vector_store_id should be in format: bucket_name:index_name diff --git a/litellm/llms/vertex_ai/vector_stores/rag_api/transformation.py b/litellm/llms/vertex_ai/vector_stores/rag_api/transformation.py index b3fcf4b394..d31e1f6c8f 100644 --- a/litellm/llms/vertex_ai/vector_stores/rag_api/transformation.py +++ b/litellm/llms/vertex_ai/vector_stores/rag_api/transformation.py @@ -97,10 +97,10 @@ class VertexVectorStoreConfig(BaseVectorStoreConfig, VertexBase): vector_store_id: str, query: Union[str, List[str]], vector_store_search_optional_params: VectorStoreSearchOptionalRequestParams, - extra_body: Optional[Dict[str, Any]], api_base: str, litellm_logging_obj: LiteLLMLoggingObj, litellm_params: dict, + extra_body: Optional[Dict[str, Any]] = None, ) -> Tuple[str, Dict[str, Any]]: """ Transform search request for Vertex AI RAG API diff --git a/litellm/llms/vertex_ai/vector_stores/search_api/transformation.py b/litellm/llms/vertex_ai/vector_stores/search_api/transformation.py index 47dac8dca3..6cb7a86bea 100644 --- a/litellm/llms/vertex_ai/vector_stores/search_api/transformation.py +++ b/litellm/llms/vertex_ai/vector_stores/search_api/transformation.py @@ -104,10 +104,10 @@ class VertexSearchAPIVectorStoreConfig(BaseVectorStoreConfig, VertexBase): vector_store_id: str, query: Union[str, List[str]], vector_store_search_optional_params: VectorStoreSearchOptionalRequestParams, - extra_body: Optional[Dict[str, Any]], api_base: str, litellm_logging_obj: LiteLLMLoggingObj, litellm_params: dict, + extra_body: Optional[Dict[str, Any]] = None, ) -> Tuple[str, Dict[str, Any]]: """ Transform search request for Vertex AI RAG API diff --git a/tests/logging_callback_tests/test_bedrock_knowledgebase_hook.py b/tests/logging_callback_tests/test_bedrock_knowledgebase_hook.py index abe96b2ea2..3e8d59b299 100644 --- a/tests/logging_callback_tests/test_bedrock_knowledgebase_hook.py +++ b/tests/logging_callback_tests/test_bedrock_knowledgebase_hook.py @@ -354,6 +354,7 @@ async def test_bedrock_kb_request_body_has_transformed_filters( api_base=api_base, litellm_logging_obj=logging_obj, litellm_params=litellm_params_dict, + extra_body=None, ) ) captured_request_body["url"] = url diff --git a/tests/test_litellm/llms/bedrock/vector_stores/test_bedrock_vector_store_transformation.py b/tests/test_litellm/llms/bedrock/vector_stores/test_bedrock_vector_store_transformation.py index c211a3536e..d60d0487d0 100644 --- a/tests/test_litellm/llms/bedrock/vector_stores/test_bedrock_vector_store_transformation.py +++ b/tests/test_litellm/llms/bedrock/vector_stores/test_bedrock_vector_store_transformation.py @@ -18,10 +18,10 @@ def test_transform_search_request(): vector_store_id="kb123", query="hello", vector_store_search_optional_params={}, - extra_body=None, api_base="https://bedrock-agent-runtime.us-west-2.amazonaws.com/knowledgebases", litellm_logging_obj=mock_log, litellm_params={}, + extra_body=None, ) assert url.endswith("/kb123/retrieve") @@ -37,6 +37,9 @@ def test_transform_search_request_uses_only_retrieval_config_from_extra_body(): vector_store_id="kb123", query="hello", vector_store_search_optional_params={}, + api_base="https://bedrock-agent-runtime.us-west-2.amazonaws.com/knowledgebases", + litellm_logging_obj=mock_log, + litellm_params={}, extra_body={ "retrievalConfiguration": { "vectorSearchConfiguration": { @@ -46,9 +49,6 @@ def test_transform_search_request_uses_only_retrieval_config_from_extra_body(): }, "unrelatedField": {"should_not": "be_forwarded"}, }, - api_base="https://bedrock-agent-runtime.us-west-2.amazonaws.com/knowledgebases", - litellm_logging_obj=mock_log, - litellm_params={}, ) assert url.endswith("/kb123/retrieve") @@ -79,10 +79,10 @@ def test_transform_search_request_does_not_mutate_extra_body_and_overrides_numbe vector_store_id="kb123", query="hello", vector_store_search_optional_params={"max_num_results": 10}, - extra_body=extra_body, api_base="https://bedrock-agent-runtime.us-west-2.amazonaws.com/knowledgebases", litellm_logging_obj=mock_log, litellm_params={}, + extra_body=extra_body, ) assert ( @@ -114,10 +114,10 @@ def test_transform_search_request_overrides_filter_without_mutating_extra_body() vector_store_id="kb123", query="hello", vector_store_search_optional_params={"filters": new_filter}, - extra_body=extra_body, api_base="https://bedrock-agent-runtime.us-west-2.amazonaws.com/knowledgebases", litellm_logging_obj=mock_log, litellm_params={}, + extra_body=extra_body, ) assert ( diff --git a/tests/test_litellm/llms/s3_vectors/vector_stores/test_s3_vectors_transformation.py b/tests/test_litellm/llms/s3_vectors/vector_stores/test_s3_vectors_transformation.py index 9507ff401a..7085e45cdc 100644 --- a/tests/test_litellm/llms/s3_vectors/vector_stores/test_s3_vectors_transformation.py +++ b/tests/test_litellm/llms/s3_vectors/vector_stores/test_s3_vectors_transformation.py @@ -56,10 +56,10 @@ class TestS3VectorsVectorStoreConfig: vector_store_id="invalid-format", query="test query", vector_store_search_optional_params={}, - extra_body=None, api_base="https://s3vectors.us-west-2.api.aws", litellm_logging_obj=mock_logging_obj, litellm_params={}, + extra_body=None, ) def test_transform_search_response(self): diff --git a/tests/vector_store_tests/test_ragflow_vector_store.py b/tests/vector_store_tests/test_ragflow_vector_store.py index 4ca3823312..cb4cfd75c1 100644 --- a/tests/vector_store_tests/test_ragflow_vector_store.py +++ b/tests/vector_store_tests/test_ragflow_vector_store.py @@ -267,6 +267,7 @@ class TestRAGFlowVectorStore(BaseVectorStoreTest): api_base="http://localhost:9380", litellm_logging_obj=logging_obj, litellm_params={}, + extra_body=None, ) def test_transform_search_vector_store_response_not_implemented(self): From 4ae2996f08398bc4fd35c5e940fd94ec8fa0bbe6 Mon Sep 17 00:00:00 2001 From: ishaan-berri <155045088+ishaan-berri@users.noreply.github.com> Date: Tue, 28 Apr 2026 20:10:42 -0700 Subject: [PATCH 41/57] Add gpt-image-2 support (#26644) (#26705) * Add gpt-image-2 support * Address gpt-image-2 PR feedback Co-authored-by: Emerson Gomes --- .../get_llm_provider_logic.py | 1 + .../litellm_core_utils/llm_cost_calc/utils.py | 8 +- .../llms/azure/image_generation/__init__.py | 2 +- .../image_generation/gpt_transformation.py | 2 +- .../image_generation/cost_calculator.py | 6 +- .../image_generation/gpt_transformation.py | 2 +- ...odel_prices_and_context_window_backup.json | 64 +++++++++++++++ litellm/utils.py | 1 + model_prices_and_context_window.json | 64 +++++++++++++++ .../test_gpt_image_cost_calculator.py | 80 ++++++++++++++++++- tests/test_litellm/test_utils.py | 79 ++++++++++++++++++ 11 files changed, 298 insertions(+), 11 deletions(-) diff --git a/litellm/litellm_core_utils/get_llm_provider_logic.py b/litellm/litellm_core_utils/get_llm_provider_logic.py index 95bcd4d718..4ff077efe7 100644 --- a/litellm/litellm_core_utils/get_llm_provider_logic.py +++ b/litellm/litellm_core_utils/get_llm_provider_logic.py @@ -348,6 +348,7 @@ def get_llm_provider( # noqa: PLR0915 or "ft:gpt-3.5-turbo" in model or "ft:gpt-4" in model # catches ft:gpt-4-0613, ft:gpt-4o or model in litellm.openai_image_generation_models + or model.startswith("gpt-image") or model in litellm.openai_video_generation_models ): custom_llm_provider = "openai" diff --git a/litellm/litellm_core_utils/llm_cost_calc/utils.py b/litellm/litellm_core_utils/llm_cost_calc/utils.py index 888999504f..59d0465e6d 100644 --- a/litellm/litellm_core_utils/llm_cost_calc/utils.py +++ b/litellm/litellm_core_utils/llm_cost_calc/utils.py @@ -982,9 +982,9 @@ class CostCalculatorUtils: image_response=completion_response, ) elif custom_llm_provider == litellm.LlmProviders.OPENAI.value: - # Check if this is a gpt-image model (token-based pricing) + # gpt-image models use token-based pricing. model_lower = model.lower() - if "gpt-image-1" in model_lower: + if "gpt-image" in model_lower: from litellm.llms.openai.image_generation.cost_calculator import ( cost_calculator as openai_gpt_image_cost_calculator, ) @@ -1004,9 +1004,9 @@ class CostCalculatorUtils: optional_params=optional_params, ) elif custom_llm_provider == litellm.LlmProviders.AZURE.value: - # Check if this is a gpt-image model (token-based pricing) + # gpt-image models use token-based pricing. model_lower = model.lower() - if "gpt-image-1" in model_lower: + if "gpt-image" in model_lower: from litellm.llms.openai.image_generation.cost_calculator import ( cost_calculator as openai_gpt_image_cost_calculator, ) diff --git a/litellm/llms/azure/image_generation/__init__.py b/litellm/llms/azure/image_generation/__init__.py index fcdf49f291..a9cf151464 100644 --- a/litellm/llms/azure/image_generation/__init__.py +++ b/litellm/llms/azure/image_generation/__init__.py @@ -24,6 +24,6 @@ def get_azure_image_generation_config(model: str) -> BaseImageGenerationConfig: return AzureDallE3ImageGenerationConfig() else: verbose_logger.debug( - f"Using AzureGPTImageGenerationConfig for model: {model}. This follows the gpt-image-1 model format." + f"Using AzureGPTImageGenerationConfig for model: {model}. This follows the gpt-image model format." ) return AzureGPTImageGenerationConfig() diff --git a/litellm/llms/azure/image_generation/gpt_transformation.py b/litellm/llms/azure/image_generation/gpt_transformation.py index 1f5f65f693..2d46592e3f 100644 --- a/litellm/llms/azure/image_generation/gpt_transformation.py +++ b/litellm/llms/azure/image_generation/gpt_transformation.py @@ -3,7 +3,7 @@ from litellm.llms.openai.image_generation import GPTImageGenerationConfig class AzureGPTImageGenerationConfig(GPTImageGenerationConfig): """ - Azure gpt-image-1 image generation config + Azure gpt-image image generation config """ pass diff --git a/litellm/llms/openai/image_generation/cost_calculator.py b/litellm/llms/openai/image_generation/cost_calculator.py index 8bca75172f..d009a085fa 100644 --- a/litellm/llms/openai/image_generation/cost_calculator.py +++ b/litellm/llms/openai/image_generation/cost_calculator.py @@ -1,5 +1,5 @@ """ -Cost calculator for OpenAI image generation models (gpt-image-1, gpt-image-1-mini) +Cost calculator for OpenAI image generation models (gpt-image family) These models use token-based pricing instead of pixel-based pricing like DALL-E. """ @@ -17,13 +17,13 @@ def cost_calculator( custom_llm_provider: Optional[str] = None, ) -> float: """ - Calculate cost for OpenAI gpt-image-1 and gpt-image-1-mini models. + Calculate cost for OpenAI gpt-image models. Uses the same usage format as Responses API, so we reuse the helper to transform to chat completion format and use generic_cost_per_token. Args: - model: The model name (e.g., "gpt-image-1", "gpt-image-1-mini") + model: The model name (e.g., "gpt-image-1", "gpt-image-2") image_response: The ImageResponse containing usage data custom_llm_provider: Optional provider name diff --git a/litellm/llms/openai/image_generation/gpt_transformation.py b/litellm/llms/openai/image_generation/gpt_transformation.py index c106d7f17b..68f799e574 100644 --- a/litellm/llms/openai/image_generation/gpt_transformation.py +++ b/litellm/llms/openai/image_generation/gpt_transformation.py @@ -15,7 +15,7 @@ if TYPE_CHECKING: class GPTImageGenerationConfig(BaseImageGenerationConfig): """ - OpenAI gpt-image-1 image generation config + OpenAI gpt-image image generation config """ def get_supported_openai_params( diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 8511d785fb..e4268fac81 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -5103,6 +5103,38 @@ "/v1/images/edits" ] }, + "azure/gpt-image-2": { + "cache_read_input_image_token_cost": 2e-06, + "cache_read_input_token_cost": 1.25e-06, + "input_cost_per_token": 5e-06, + "input_cost_per_image_token": 8e-06, + "litellm_provider": "azure", + "mode": "image_generation", + "output_cost_per_token": 1e-05, + "output_cost_per_image_token": 3e-05, + "supported_endpoints": [ + "/v1/images/generations", + "/v1/images/edits" + ], + "supports_vision": true, + "supports_pdf_input": true + }, + "azure/gpt-image-2-2026-04-21": { + "cache_read_input_image_token_cost": 2e-06, + "cache_read_input_token_cost": 1.25e-06, + "input_cost_per_token": 5e-06, + "input_cost_per_image_token": 8e-06, + "litellm_provider": "azure", + "mode": "image_generation", + "output_cost_per_token": 1e-05, + "output_cost_per_image_token": 3e-05, + "supported_endpoints": [ + "/v1/images/generations", + "/v1/images/edits" + ], + "supports_vision": true, + "supports_pdf_input": true + }, "azure/low/1024-x-1024/gpt-image-1-mini": { "input_cost_per_pixel": 2.0751953125e-09, "litellm_provider": "azure", @@ -19083,6 +19115,38 @@ "supports_vision": true, "supports_pdf_input": true }, + "gpt-image-2": { + "cache_read_input_image_token_cost": 2e-06, + "cache_read_input_token_cost": 1.25e-06, + "input_cost_per_token": 5e-06, + "litellm_provider": "openai", + "mode": "image_generation", + "output_cost_per_token": 1e-05, + "input_cost_per_image_token": 8e-06, + "output_cost_per_image_token": 3e-05, + "supported_endpoints": [ + "/v1/images/generations", + "/v1/images/edits" + ], + "supports_vision": true, + "supports_pdf_input": true + }, + "gpt-image-2-2026-04-21": { + "cache_read_input_image_token_cost": 2e-06, + "cache_read_input_token_cost": 1.25e-06, + "input_cost_per_token": 5e-06, + "litellm_provider": "openai", + "mode": "image_generation", + "output_cost_per_token": 1e-05, + "input_cost_per_image_token": 8e-06, + "output_cost_per_image_token": 3e-05, + "supported_endpoints": [ + "/v1/images/generations", + "/v1/images/edits" + ], + "supports_vision": true, + "supports_pdf_input": true + }, "low/1024-x-1024/gpt-image-1.5": { "input_cost_per_image": 0.009, "litellm_provider": "openai", diff --git a/litellm/utils.py b/litellm/utils.py index e63bf402bf..027c9fedce 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -6526,6 +6526,7 @@ def validate_environment( # noqa: PLR0915 or model in litellm.open_ai_text_completion_models or model in litellm.open_ai_embedding_models or model in litellm.openai_image_generation_models + or model.startswith("gpt-image") ): if "OPENAI_API_KEY" in os.environ: keys_in_environment = True diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 114883f530..ca7d323ad6 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -5117,6 +5117,38 @@ "/v1/images/edits" ] }, + "azure/gpt-image-2": { + "cache_read_input_image_token_cost": 2e-06, + "cache_read_input_token_cost": 1.25e-06, + "input_cost_per_token": 5e-06, + "input_cost_per_image_token": 8e-06, + "litellm_provider": "azure", + "mode": "image_generation", + "output_cost_per_token": 1e-05, + "output_cost_per_image_token": 3e-05, + "supported_endpoints": [ + "/v1/images/generations", + "/v1/images/edits" + ], + "supports_vision": true, + "supports_pdf_input": true + }, + "azure/gpt-image-2-2026-04-21": { + "cache_read_input_image_token_cost": 2e-06, + "cache_read_input_token_cost": 1.25e-06, + "input_cost_per_token": 5e-06, + "input_cost_per_image_token": 8e-06, + "litellm_provider": "azure", + "mode": "image_generation", + "output_cost_per_token": 1e-05, + "output_cost_per_image_token": 3e-05, + "supported_endpoints": [ + "/v1/images/generations", + "/v1/images/edits" + ], + "supports_vision": true, + "supports_pdf_input": true + }, "azure/low/1024-x-1024/gpt-image-1-mini": { "input_cost_per_pixel": 2.0751953125e-09, "litellm_provider": "azure", @@ -19097,6 +19129,38 @@ "supports_vision": true, "supports_pdf_input": true }, + "gpt-image-2": { + "cache_read_input_image_token_cost": 2e-06, + "cache_read_input_token_cost": 1.25e-06, + "input_cost_per_token": 5e-06, + "litellm_provider": "openai", + "mode": "image_generation", + "output_cost_per_token": 1e-05, + "input_cost_per_image_token": 8e-06, + "output_cost_per_image_token": 3e-05, + "supported_endpoints": [ + "/v1/images/generations", + "/v1/images/edits" + ], + "supports_vision": true, + "supports_pdf_input": true + }, + "gpt-image-2-2026-04-21": { + "cache_read_input_image_token_cost": 2e-06, + "cache_read_input_token_cost": 1.25e-06, + "input_cost_per_token": 5e-06, + "litellm_provider": "openai", + "mode": "image_generation", + "output_cost_per_token": 1e-05, + "input_cost_per_image_token": 8e-06, + "output_cost_per_image_token": 3e-05, + "supported_endpoints": [ + "/v1/images/generations", + "/v1/images/edits" + ], + "supports_vision": true, + "supports_pdf_input": true + }, "low/1024-x-1024/gpt-image-1.5": { "input_cost_per_image": 0.009, "litellm_provider": "openai", diff --git a/tests/test_litellm/test_gpt_image_cost_calculator.py b/tests/test_litellm/test_gpt_image_cost_calculator.py index 620c073498..6644b1389c 100644 --- a/tests/test_litellm/test_gpt_image_cost_calculator.py +++ b/tests/test_litellm/test_gpt_image_cost_calculator.py @@ -29,8 +29,21 @@ from litellm.types.utils import ( ) +@pytest.fixture(autouse=True) +def _use_local_model_cost_map(monkeypatch): + original_model_cost = litellm.model_cost + monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") + litellm.model_cost = litellm.get_model_cost_map(url="") + litellm.get_model_info.cache_clear() + try: + yield + finally: + litellm.model_cost = original_model_cost + litellm.get_model_info.cache_clear() + + class TestGPTImageCostCalculator: - """Test the OpenAI gpt-image-1 cost calculator""" + """Test the OpenAI gpt-image cost calculator""" def test_gpt_image_1_cost_with_text_only(self): """Test cost calculation with only text input tokens""" @@ -149,6 +162,44 @@ class TestGPTImageCostCalculator: assert cost == 0.0 + def test_gpt_image_2_cost_with_text_and_image_tokens(self): + """Test cost calculation for gpt-image-2 token pricing""" + from litellm.llms.openai.image_generation.cost_calculator import cost_calculator + + usage = Usage( + prompt_tokens=600, + completion_tokens=5000, + total_tokens=5600, + prompt_tokens_details=PromptTokensDetailsWrapper( + text_tokens=100, + image_tokens=500, + ), + completion_tokens_details=CompletionTokensDetailsWrapper( + text_tokens=1000, + image_tokens=4000, + ), + ) + + image_response = ImageResponse( + created=1234567890, + data=[ImageObject(url="http://example.com/image.jpg")], + ) + image_response.usage = usage + + cost = cost_calculator( + model="gpt-image-2", + image_response=image_response, + custom_llm_provider="openai", + ) + + # GPT Image 2 pricing: + # Text input: 100 * $5/1M = 0.0005 + # Image input: 500 * $8/1M = 0.004 + # Text output: 1000 * $10/1M = 0.01 + # Image output: 4000 * $30/1M = 0.12 + expected_cost = 0.0005 + 0.004 + 0.01 + 0.12 + assert abs(cost - expected_cost) < 1e-6, f"Expected {expected_cost}, got {cost}" + class TestGPTImageCostRouting: """Test that gpt-image models are properly routed to the token-based calculator""" @@ -182,6 +233,33 @@ class TestGPTImageCostRouting: expected_cost = 0.0005 + 0.2 assert abs(cost - expected_cost) < 1e-6, f"Expected {expected_cost}, got {cost}" + def test_openai_gpt_image_2_routes_to_token_calculator(self): + """Test that OpenAI gpt-image-2 routes to token-based calculator""" + from litellm.litellm_core_utils.llm_cost_calc.utils import CostCalculatorUtils + + usage = Usage( + prompt_tokens=100, + completion_tokens=5000, + total_tokens=5100, + prompt_tokens_details=PromptTokensDetailsWrapper(text_tokens=100), + completion_tokens_details=CompletionTokensDetailsWrapper(image_tokens=5000), + ) + + image_response = ImageResponse( + created=1234567890, + data=[ImageObject(url="http://example.com/image.jpg")], + ) + image_response.usage = usage + + cost = CostCalculatorUtils.route_image_generation_cost_calculator( + model="gpt-image-2", + completion_response=image_response, + custom_llm_provider="openai", + ) + + expected_cost = 0.0005 + 0.15 + assert abs(cost - expected_cost) < 1e-6, f"Expected {expected_cost}, got {cost}" + def test_openai_dalle_routes_to_pixel_calculator(self): """Test that OpenAI DALL-E still routes to pixel-based calculator""" from litellm.litellm_core_utils.llm_cost_calc.utils import CostCalculatorUtils diff --git a/tests/test_litellm/test_utils.py b/tests/test_litellm/test_utils.py index b8a4220c67..f28fe3ed25 100644 --- a/tests/test_litellm/test_utils.py +++ b/tests/test_litellm/test_utils.py @@ -32,6 +32,19 @@ from litellm.utils import ( # Adds the parent directory to the system path +@pytest.fixture +def local_model_cost_map(monkeypatch): + original_model_cost = litellm.model_cost + monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") + litellm.model_cost = litellm.get_model_cost_map(url="") + litellm.get_model_info.cache_clear() + try: + yield + finally: + litellm.model_cost = original_model_cost + litellm.get_model_info.cache_clear() + + def test_check_provider_match_azure_ai_allows_openai_and_azure(): """ Test that azure_ai provider can match openai and azure models. @@ -198,6 +211,72 @@ def test_get_optional_params_image_gen_filters_empty_values(): assert optional_params == {} +def test_gpt_image_provider_detection_covers_existing_family(): + for image_model in ("gpt-image-1", "gpt-image-1-mini", "gpt-image-1.5"): + model, custom_llm_provider, _, _ = litellm.get_llm_provider(model=image_model) + + assert model == image_model + assert custom_llm_provider == "openai" + + +def test_gpt_image_2_provider_and_model_info(local_model_cost_map): + + model, custom_llm_provider, _, _ = litellm.get_llm_provider(model="gpt-image-2") + + assert model == "gpt-image-2" + assert custom_llm_provider == "openai" + + model_info = litellm.get_model_info(model="gpt-image-2") + assert model_info["litellm_provider"] == "openai" + assert model_info["mode"] == "image_generation" + assert model_info["input_cost_per_token"] == 5e-06 + assert model_info["input_cost_per_image_token"] == 8e-06 + assert model_info["output_cost_per_token"] == 1e-05 + assert model_info["output_cost_per_image_token"] == 3e-05 + assert ( + "/v1/images/generations" + in litellm.model_cost["gpt-image-2"]["supported_endpoints"] + ) + assert ( + "/v1/images/edits" in litellm.model_cost["gpt-image-2"]["supported_endpoints"] + ) + assert model_info["supports_vision"] is True + assert model_info["supports_pdf_input"] is True + + +def test_gpt_image_2_snapshot_model_info(local_model_cost_map): + model, custom_llm_provider, _, _ = litellm.get_llm_provider( + model="gpt-image-2-2026-04-21" + ) + + assert model == "gpt-image-2-2026-04-21" + assert custom_llm_provider == "openai" + + model_info = litellm.get_model_info(model="gpt-image-2-2026-04-21") + assert model_info["litellm_provider"] == "openai" + assert model_info["mode"] == "image_generation" + assert model_info["output_cost_per_image_token"] == 3e-05 + + +def test_azure_gpt_image_2_model_info(local_model_cost_map): + model, custom_llm_provider, _, _ = litellm.get_llm_provider( + model="azure/gpt-image-2" + ) + + assert model == "gpt-image-2" + assert custom_llm_provider == "azure" + + model_info = litellm.get_model_info( + model="gpt-image-2", custom_llm_provider="azure" + ) + assert model_info["litellm_provider"] == "azure" + assert model_info["mode"] == "image_generation" + assert model_info["input_cost_per_token"] == 5e-06 + assert model_info["input_cost_per_image_token"] == 8e-06 + assert model_info["output_cost_per_token"] == 1e-05 + assert model_info["output_cost_per_image_token"] == 3e-05 + + def test_all_model_configs(): from litellm.llms.vertex_ai.vertex_ai_partner_models.ai21.transformation import ( VertexAIAi21Config, From 44ab016743c9b59f2dcc4c17f4f6b6431d1108d2 Mon Sep 17 00:00:00 2001 From: xinrui <94846330+xinrui-z@users.noreply.github.com> Date: Wed, 29 Apr 2026 11:18:30 +0800 Subject: [PATCH 42/57] feat(provider): add AIHubMix as an OpenAI-compatible provider (#24294) * feat: add AIHubMix provider to providers.json * fix: add aihubmix to provider_endpoints_support.json for CI check --------- Co-authored-by: yuneng-jiang --- litellm/llms/openai_like/providers.json | 5 +++++ provider_endpoints_support.json | 17 +++++++++++++++++ 2 files changed, 22 insertions(+) diff --git a/litellm/llms/openai_like/providers.json b/litellm/llms/openai_like/providers.json index 275c352b39..5dd1247001 100644 --- a/litellm/llms/openai_like/providers.json +++ b/litellm/llms/openai_like/providers.json @@ -101,5 +101,10 @@ "param_mappings": { "max_completion_tokens": "max_tokens" } + }, + "aihubmix": { + "base_url": "https://aihubmix.com/v1", + "api_key_env": "AIHUBMIX_API_KEY", + "api_base_env": "AIHUBMIX_API_BASE" } } diff --git a/provider_endpoints_support.json b/provider_endpoints_support.json index 6f23c87f91..ed49c14621 100644 --- a/provider_endpoints_support.json +++ b/provider_endpoints_support.json @@ -193,6 +193,23 @@ "a2a": false } }, + "aihubmix": { + "display_name": "AIHubMix (`aihubmix`)", + "url": "https://docs.litellm.ai/docs/providers/aihubmix", + "endpoints": { + "chat_completions": true, + "messages": true, + "responses": true, + "embeddings": true, + "image_generations": true, + "audio_transcriptions": true, + "audio_speech": true, + "moderations": true, + "batches": false, + "rerank": true, + "a2a": false + } + }, "assemblyai": { "display_name": "AssemblyAI (`assemblyai`)", "url": "https://docs.litellm.ai/docs/pass_through/assembly_ai", From e0cd536eaa04880b1d142822d6cbb6320cc8ea7f Mon Sep 17 00:00:00 2001 From: Sameer Kankute Date: Wed, 29 Apr 2026 09:06:34 +0530 Subject: [PATCH 43/57] Fix lint --- litellm/proxy/_experimental/mcp_server/mcp_server_manager.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index edc9c87cda..eb0bda3fec 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -1533,8 +1533,8 @@ class MCPServerManager: header_value ) - authorization_servers: List[str] = [] - resource_scopes: Optional[List[str]] = None + authorization_servers = [] + resource_scopes = None if resource_metadata_url: ( authorization_servers, From df3dbd18d6fc022267416fbf93993e912029fe9c Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Wed, 29 Apr 2026 03:43:13 +0000 Subject: [PATCH 44/57] feat(mcp): rehash short tool prefix on collision and cache per server MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Two MCP servers can natural-hash to the same three-character base62 prefix. With 62**3 = 238_328 slots the birthday bound is ~488 servers for 50% collision probability, so a single proxy hosting more than ~100 MCP servers has a non-trivial chance of seeing a collision in practice — and a collision means tool names from two different servers share a routing key, causing silent mis-routing. Mitigation: - compute_short_server_prefix(server_id, attempt=N) folds an attempt counter into the SHA-256 seed, so rehashes are deterministic and produce a fresh three-char prefix space per attempt. - New MCPServer.short_prefix field caches the resolved (post-dedup) prefix on the model so it stays stable across the process lifetime. - MCPServerManager._assign_unique_short_prefix walks attempts 0..N until it finds a prefix not already used by another server in the combined registry. Logs an INFO line when a rehash happens so operators have a breadcrumb if it ever does. - Wired into every registration path: load_servers_from_config, add_server, update_server, reload_servers_from_database. The database reload path also carries the previously-resolved prefix forward so reloads don't churn it. - get_server_prefix prefers the cached short_prefix when set, so the resolved value (not the raw natural hash) is used everywhere. - iter_known_server_prefixes yields the cached short_prefix too, so reverse-lookup tolerance covers the rehashed form. No-op when LITELLM_USE_SHORT_MCP_TOOL_PREFIX is disabled — the field stays None and behaviour is unchanged. Co-authored-by: Mateo Wang --- .../mcp_server/mcp_server_manager.py | 78 ++++++++++++++++++ .../proxy/_experimental/mcp_server/utils.py | 29 +++++-- .../types/mcp_server/mcp_server_manager.py | 6 ++ .../mcp_server/test_short_mcp_tool_prefix.py | 81 +++++++++++++++++++ 4 files changed, 186 insertions(+), 8 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index 8e473c1cb2..853208a538 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -50,7 +50,9 @@ from litellm.proxy._experimental.mcp_server.oauth2_token_cache import resolve_mc from litellm.proxy._experimental.mcp_server.utils import ( MCP_TOOL_PREFIX_SEPARATOR, add_server_prefix_to_name, + compute_short_server_prefix, get_server_prefix, + is_short_mcp_tool_prefix_enabled, is_tool_name_prefixed, iter_known_server_prefixes, merge_mcp_headers, @@ -365,6 +367,7 @@ class MCPServerManager: aws_session_name=server_config.get("aws_session_name", None), instructions=server_config.get("instructions", None), ) + self._assign_unique_short_prefix(new_server) self.config_mcp_servers[server_id] = new_server # Check if this is an OpenAPI-based server @@ -727,6 +730,7 @@ class MCPServerManager: try: if mcp_server.server_id not in self.registry: new_server = await self.build_mcp_server_from_table(mcp_server) + self._assign_unique_short_prefix(new_server) self.registry[mcp_server.server_id] = new_server await self._maybe_register_openapi_tools(new_server) verbose_logger.debug(f"Added MCP Server: {new_server.name}") @@ -739,6 +743,12 @@ class MCPServerManager: try: if mcp_server.server_id in self.registry: new_server = await self.build_mcp_server_from_table(mcp_server) + # Carry the previously-resolved short prefix across so the + # tool names stay stable for clients holding cached lists. + existing_prefix = self.registry[mcp_server.server_id].short_prefix + if existing_prefix and not new_server.short_prefix: + new_server.short_prefix = existing_prefix + self._assign_unique_short_prefix(new_server) self.registry[mcp_server.server_id] = new_server await self._maybe_register_openapi_tools(new_server) verbose_logger.debug(f"Updated MCP Server: {new_server.name}") @@ -1815,6 +1825,63 @@ class MCPServerManager: verbose_logger.warning(f"Error listing tools from {server_name}: {str(e)}") return [] + _SHORT_PREFIX_MAX_REHASH_ATTEMPTS = 1024 + + def _assign_unique_short_prefix(self, server: MCPServer) -> None: + """Resolve and cache a collision-free short tool prefix on ``server``. + + Called at registration time for every MCP server entering the + registry. Mutates ``server.short_prefix`` in place. No-ops when + ``LITELLM_USE_SHORT_MCP_TOOL_PREFIX`` is disabled, when the server + has no ``server_id`` (synthetic temp-server objects), or when a + prefix is already cached. + + Collision strategy: take the natural hash; if it's already used by + a *different* server in the combined registry, rehash with an + incrementing attempt counter until we find an unused slot. The + attempt counter is folded into the hash so the resulting prefix is + still deterministic for a given (server_id, set-of-other-server-ids) + pair within one process. + """ + if not is_short_mcp_tool_prefix_enabled(): + return + if server.short_prefix: + return + if not server.server_id: + return + + used: Dict[str, str] = {} + for other in self.get_registry().values(): + if other.server_id == server.server_id: + continue + if other.short_prefix: + used[other.short_prefix] = other.server_id + + for attempt in range(self._SHORT_PREFIX_MAX_REHASH_ATTEMPTS): + candidate = compute_short_server_prefix(server.server_id, attempt=attempt) + if candidate not in used: + server.short_prefix = candidate + if attempt > 0: + verbose_logger.info( + "MCP short-prefix collision resolved for server %s: " + "natural hash collided with %s, using rehashed prefix " + "%s (attempt=%d).", + server.server_id, + used.get( + compute_short_server_prefix(server.server_id, attempt=0), + "", + ), + candidate, + attempt, + ) + return + + raise RuntimeError( + f"Unable to assign a unique short MCP tool prefix for server " + f"{server.server_id} after {self._SHORT_PREFIX_MAX_REHASH_ATTEMPTS} " + "attempts; the 3-character prefix space is too crowded." + ) + def _create_prefixed_tools( self, tools: List[MCPTool], server: MCPServer, add_prefix: bool = True ) -> List[MCPTool]: @@ -2681,6 +2748,9 @@ class MCPServerManager: previous_registry = self.registry new_registry: Dict[str, MCPServer] = {} + # Stage one: build every server. Stage two assigns short prefixes + # against the *full* set so dedup is deterministic regardless of + # iteration order. for server in db_mcp_servers: existing_server = previous_registry.get(server.server_id) @@ -2704,10 +2774,18 @@ class MCPServerManager: f"Building server from DB: {server.server_id} ({server.server_name})" ) new_server = await self.build_mcp_server_from_table(server) + # Carry the cached short_prefix from the previous registry entry + # (if any) so the prefix is stable across reloads. + if existing_server is not None and existing_server.short_prefix: + new_server.short_prefix = existing_server.short_prefix new_registry[server.server_id] = new_server await self._maybe_register_openapi_tools(new_server) + # Swap in the new registry first so _assign_unique_short_prefix + # sees the complete set when checking for collisions. self.registry = new_registry + for new_server in new_registry.values(): + self._assign_unique_short_prefix(new_server) verbose_logger.debug( "MCP registry refreshed (%s servers in registry)", len(new_registry) diff --git a/litellm/proxy/_experimental/mcp_server/utils.py b/litellm/proxy/_experimental/mcp_server/utils.py index 42910cc790..a0c278b906 100644 --- a/litellm/proxy/_experimental/mcp_server/utils.py +++ b/litellm/proxy/_experimental/mcp_server/utils.py @@ -54,17 +54,22 @@ def is_short_mcp_tool_prefix_enabled() -> bool: return raw.strip().lower() in ("1", "true", "yes", "on") -def compute_short_server_prefix(server_id: str) -> str: +def compute_short_server_prefix(server_id: str, attempt: int = 0) -> str: """Derive the deterministic three-character base62 prefix for a server. - Uses SHA-256 of the server_id and folds the first eight bytes into a - base62 string. An empty server_id raises ValueError — short prefixes - require a stable identifier to be deterministic. + Uses SHA-256 of ``f"{server_id}#{attempt}"`` and folds the first eight + bytes into a base62 string. Pass ``attempt > 0`` to rehash to a + different prefix when the natural hash collides with a prefix already + assigned to another server (see + ``MCPServerManager._assign_unique_short_prefix``). An empty server_id + raises ValueError — short prefixes require a stable identifier to be + deterministic. """ if not server_id: raise ValueError("compute_short_server_prefix requires a non-empty server_id") - digest = hashlib.sha256(server_id.encode("utf-8")).digest() + seed = server_id if attempt == 0 else f"{server_id}#{attempt}" + digest = hashlib.sha256(seed.encode("utf-8")).digest() value = int.from_bytes(digest[:8], "big") chars = [] for _ in range(SHORT_MCP_TOOL_PREFIX_LENGTH): @@ -143,11 +148,18 @@ def get_server_prefix(server: Any) -> str: """Return the prefix for a server. When the short-prefix mode is enabled (``LITELLM_USE_SHORT_MCP_TOOL_PREFIX``) - a deterministic three-character base62 ID derived from ``server_id`` is - returned. Otherwise we fall back to the historical behaviour: alias if - present, else server_name, else server_id. + a three-character base62 ID is returned. We prefer the cached + ``server.short_prefix`` value when set — that field is populated at + registration time by ``MCPServerManager._assign_unique_short_prefix`` + and resolves natural-hash collisions deterministically — and only fall + back to the natural hash for ad-hoc / temp-server objects without a + cached value. In default mode the historical behaviour is preserved: + alias if present, else server_name, else server_id. """ if is_short_mcp_tool_prefix_enabled(): + cached = getattr(server, "short_prefix", None) + if cached: + return cached server_id = getattr(server, "server_id", None) if server_id: return compute_short_server_prefix(server_id) @@ -177,6 +189,7 @@ def iter_known_server_prefixes(server: Any) -> Iterator[str]: yield value yield from _emit(get_server_prefix(server)) + yield from _emit(getattr(server, "short_prefix", None)) server_id = getattr(server, "server_id", None) if server_id: diff --git a/litellm/types/mcp_server/mcp_server_manager.py b/litellm/types/mcp_server/mcp_server_manager.py index ace8c8a418..8f8673b0a7 100644 --- a/litellm/types/mcp_server/mcp_server_manager.py +++ b/litellm/types/mcp_server/mcp_server_manager.py @@ -81,6 +81,12 @@ class MCPServer(BaseModel): # Defaults to the token's expires_in minus the expiry buffer, or # MCP_PER_USER_TOKEN_DEFAULT_TTL when expires_in is absent. token_storage_ttl_seconds: Optional[int] = None + # Resolved short-ID tool prefix when LITELLM_USE_SHORT_MCP_TOOL_PREFIX is + # enabled. Set by ``MCPServerManager._assign_unique_short_prefix`` at + # registration time so that natural-hash collisions between two + # different ``server_id`` values are bumped deterministically. Left + # ``None`` in default-prefix mode. + short_prefix: Optional[str] = None model_config = ConfigDict(arbitrary_types_allowed=True) @property diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_short_mcp_tool_prefix.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_short_mcp_tool_prefix.py index aef6546c81..cdaecfb227 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_short_mcp_tool_prefix.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_short_mcp_tool_prefix.py @@ -205,3 +205,84 @@ class TestManagerShortPrefix: prefix = get_server_prefix(server) full = add_server_prefix_to_name("get_repo", prefix) assert len(full) < 60 + + +# --------------------------------------------------------------------------- +# Collision-resolution at registration time +# --------------------------------------------------------------------------- + + +class TestShortPrefixCollisionResolution: + """``_assign_unique_short_prefix`` must rehash on collision. + + The dedup path is exercised by forcing two distinct ``server_id`` + values to both hash to the same natural prefix via a monkeypatched + ``compute_short_server_prefix``. + """ + + def test_no_op_when_flag_off(self): + manager = MCPServerManager() + server = _make_server(server_id="abc") + manager._assign_unique_short_prefix(server) + assert server.short_prefix is None + + def test_assigns_natural_hash_when_no_collision(self, monkeypatch): + from litellm.proxy._experimental.mcp_server import utils as mcp_utils + + monkeypatch.setenv("LITELLM_USE_SHORT_MCP_TOOL_PREFIX", "true") + manager = MCPServerManager() + server = _make_server(server_id="abc") + manager._assign_unique_short_prefix(server) + + assert server.short_prefix == mcp_utils.compute_short_server_prefix("abc") + + def test_rehashes_when_natural_hash_collides(self, monkeypatch): + """Two server_ids that natural-hash to the same prefix get + deterministic, distinct short prefixes.""" + monkeypatch.setenv("LITELLM_USE_SHORT_MCP_TOOL_PREFIX", "true") + + # Force every attempt=0 hash to "AAA" and attempt=1 to "AAB". + # That way the second server registered must rehash to "AAB". + from litellm.proxy._experimental.mcp_server import utils as mcp_utils + + def _fake_hash(server_id: str, attempt: int = 0) -> str: + return "AAA" if attempt == 0 else f"AA{chr(ord('A') + attempt)}" + + monkeypatch.setattr(mcp_utils, "compute_short_server_prefix", _fake_hash) + # Also patch the symbol that the manager imported at module load. + from litellm.proxy._experimental.mcp_server import ( + mcp_server_manager as mgr_module, + ) + + monkeypatch.setattr(mgr_module, "compute_short_server_prefix", _fake_hash) + + manager = MCPServerManager() + first = _make_server(server_id="server-1", alias="srv1") + second = _make_server(server_id="server-2", alias="srv2") + + # Pretend both are already in the registry so dedup sees both. + manager.registry[first.server_id] = first + manager._assign_unique_short_prefix(first) + manager.registry[second.server_id] = second + manager._assign_unique_short_prefix(second) + + assert first.short_prefix == "AAA" + assert second.short_prefix == "AAB" + assert first.short_prefix != second.short_prefix + + def test_cached_prefix_is_reused(self, monkeypatch): + monkeypatch.setenv("LITELLM_USE_SHORT_MCP_TOOL_PREFIX", "true") + manager = MCPServerManager() + server = _make_server(server_id="abc") + server.short_prefix = "ZZZ" # pretend a previous registration set this + + manager._assign_unique_short_prefix(server) + + assert server.short_prefix == "ZZZ" + + def test_get_server_prefix_prefers_cached(self, monkeypatch): + monkeypatch.setenv("LITELLM_USE_SHORT_MCP_TOOL_PREFIX", "true") + server = _make_server(server_id="abc") + server.short_prefix = "Q9q" + + assert get_server_prefix(server) == "Q9q" From 3215874e400de2989db7b15ef85c27dd703381af Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Wed, 29 Apr 2026 03:48:41 +0000 Subject: [PATCH 45/57] fix(test): scope ERROR log assertion to LiteLLM logger in test_model_alias_map The test was flaking on unrelated asyncio ERROR records (e.g. "Unclosed client session" from background tasks in other tests). Restrict the assertion to records emitted by LiteLLM loggers so the test only fails on errors actually produced by the code under test. Co-authored-by: Mateo Wang --- tests/local_testing/test_model_alias_map.py | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/tests/local_testing/test_model_alias_map.py b/tests/local_testing/test_model_alias_map.py index cf731d6628..14c1de2f6a 100644 --- a/tests/local_testing/test_model_alias_map.py +++ b/tests/local_testing/test_model_alias_map.py @@ -30,10 +30,9 @@ def test_model_alias_map(caplog): ) print(response.model) - captured_logs = [rec.levelname for rec in caplog.records] - - for log in captured_logs: - assert "ERROR" not in log + for rec in caplog.records: + if rec.levelname == "ERROR" and rec.name.startswith("LiteLLM"): + pytest.fail(f"Unexpected litellm ERROR log: {rec.getMessage()}") assert "llama-3.1-8b-instant" in response.model except litellm.ServiceUnavailableError: From 6b3f07ba25a375bd3949b6d0fd321a204de4d116 Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Wed, 29 Apr 2026 03:53:03 +0000 Subject: [PATCH 46/57] fix(mcp): register OpenAPI tools after short prefix collision resolution in reload --- litellm/proxy/_experimental/mcp_server/mcp_server_manager.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index 853208a538..8bb7a5d50d 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -2779,13 +2779,16 @@ class MCPServerManager: if existing_server is not None and existing_server.short_prefix: new_server.short_prefix = existing_server.short_prefix new_registry[server.server_id] = new_server - await self._maybe_register_openapi_tools(new_server) # Swap in the new registry first so _assign_unique_short_prefix # sees the complete set when checking for collisions. self.registry = new_registry for new_server in new_registry.values(): self._assign_unique_short_prefix(new_server) + # Register OpenAPI tools *after* the final short prefix is assigned + # so the tools are stored in the global registry under the same + # prefix that lookups will use. + await self._maybe_register_openapi_tools(new_server) verbose_logger.debug( "MCP registry refreshed (%s servers in registry)", len(new_registry) From 3fb50563057f88a5291d15e37ada705d942d5320 Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Wed, 29 Apr 2026 03:59:35 +0000 Subject: [PATCH 47/57] fix(mcp): address greptile review on short tool prefix - server.py: drop the redundant server_id append in _get_filtered_mcp_servers_from_mcp_server_names. iter_known_server_prefixes already yields server_id unconditionally, so the manual append (and its misleading comment) was a no-op duplicate. - utils.py: rewrite the SHORT_MCP_TOOL_PREFIX docstring to accurately describe the collision behaviour. The previous wording said collisions were 'cosmetic only', but a natural-hash collision IS a routing-correctness issue, which is precisely why we already added _assign_unique_short_prefix to rehash deterministically. The new comment cross-references that path. - utils.py: restrict the first character of the short prefix to [A-Za-z] via a 52-char alphabet for position 0 only. The remaining two positions still use the full base62 alphabet. This keeps prefixes valid identifiers on every backend and gives 52*62*62 = 199_888 distinct prefixes (still comfortably more than any realistic deployment). - tests: add coverage proving the first character of the prefix is always alphabetic across many server_ids and rehash attempts. Co-authored-by: Mateo Wang --- .../proxy/_experimental/mcp_server/server.py | 5 -- .../proxy/_experimental/mcp_server/utils.py | 71 ++++++++++++------- .../mcp_server/test_short_mcp_tool_prefix.py | 14 ++++ 3 files changed, 60 insertions(+), 30 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/server.py b/litellm/proxy/_experimental/mcp_server/server.py index 3924687a0b..ae6055217b 100644 --- a/litellm/proxy/_experimental/mcp_server/server.py +++ b/litellm/proxy/_experimental/mcp_server/server.py @@ -714,11 +714,6 @@ if MCP_AVAILABLE: match_list = [ s.lower() for s in iter_known_server_prefixes(server) if s ] - # Always accept server_id even if it isn't part of the - # current prefix form (iter_known_server_prefixes only - # yields it when no other identifier exists). - if server.server_id: - match_list.append(server.server_id.lower()) if server_or_group.lower() in match_list: filtered_server[server.server_id] = server diff --git a/litellm/proxy/_experimental/mcp_server/utils.py b/litellm/proxy/_experimental/mcp_server/utils.py index a0c278b906..df5705c342 100644 --- a/litellm/proxy/_experimental/mcp_server/utils.py +++ b/litellm/proxy/_experimental/mcp_server/utils.py @@ -21,27 +21,39 @@ MCP_TOOL_PREFIX_FORMAT = "{server_name}{separator}{tool_name}" # When LITELLM_USE_SHORT_MCP_TOOL_PREFIX is truthy the prefix attached to MCP # tool / prompt / resource / resource-template names switches from the # (potentially long) human-readable server name to a deterministic three -# character base62 ID derived from the server's ``server_id``. +# character ID derived from the server's ``server_id``. # -# Why three characters and base62 ([0-9A-Za-z])? -# * 62**3 = 238_328 distinct IDs — the chance of a real local tool name -# happening to begin with the exact prefix LiteLLM assigned to a given -# MCP server is negligible in practice. -# * The IDs are short enough that prefixed tool names stay well under the -# 60-character upper bound enforced by some model APIs (Anthropic etc.) -# even for long upstream tool names. -# * The mapping is deterministic (SHA-256 of ``server_id`` → first three -# base62 chars), which means the prefix is stable across processes, -# workers and restarts without any persistence layer. Two servers with -# different ``server_id`` values can in principle hash to the same -# three chars, but for the reverse-lookup path we register every known -# form of the prefix anyway, so a collision only affects the cosmetic -# emitted name, not routing correctness. +# Why three characters? +# * The first character is restricted to 52 alphabetic characters +# ([A-Za-z]) and the remaining two characters use the full base62 +# alphabet ([0-9A-Za-z]). That guarantees the prefix never starts +# with a digit so it remains a valid identifier for every model API +# (some providers historically required a leading alphabetic char). +# * 52 * 62 * 62 = 199_888 distinct IDs. The chance of a real local +# tool name happening to begin with the exact prefix LiteLLM assigned +# to a given MCP server is negligible in practice. +# * The IDs are short enough that prefixed tool names stay well under +# the 60-character upper bound enforced by some model APIs (Anthropic +# etc.) even for long upstream tool names. +# * The mapping is deterministic (SHA-256 of ``server_id`` → three +# characters drawn from the alphabets above), so the prefix is stable +# across processes, workers and restarts without any persistence +# layer. Two servers with different ``server_id`` values can in +# principle hash to the same three chars; that natural-hash collision +# IS a routing-correctness issue (the second registrant would otherwise +# have its tools misrouted to the first), so registration goes through +# ``MCPServerManager._assign_unique_short_prefix`` which rehashes with +# a deterministic attempt counter until it finds an unused prefix and +# caches the result on ``MCPServer.short_prefix``. A collision is +# logged at INFO when it happens. # # This flag is intentionally opt-in for the first release so customers can # migrate. It will become the default in a future release. SHORT_MCP_TOOL_PREFIX_LENGTH = 3 _BASE62_ALPHABET = "0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz" +# Subset of _BASE62_ALPHABET used for the *first* character only, to +# guarantee the prefix never starts with a digit. +_BASE52_ALPHA_ALPHABET = "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz" def is_short_mcp_tool_prefix_enabled() -> bool: @@ -55,15 +67,17 @@ def is_short_mcp_tool_prefix_enabled() -> bool: def compute_short_server_prefix(server_id: str, attempt: int = 0) -> str: - """Derive the deterministic three-character base62 prefix for a server. + """Derive the deterministic three-character prefix for a server. Uses SHA-256 of ``f"{server_id}#{attempt}"`` and folds the first eight - bytes into a base62 string. Pass ``attempt > 0`` to rehash to a - different prefix when the natural hash collides with a prefix already - assigned to another server (see - ``MCPServerManager._assign_unique_short_prefix``). An empty server_id - raises ValueError — short prefixes require a stable identifier to be - deterministic. + bytes into a fixed-length string whose first character is drawn from + ``_BASE52_ALPHA_ALPHABET`` (so the prefix never starts with a digit) + and whose remaining characters are drawn from the full base62 + alphabet. Pass ``attempt > 0`` to rehash to a different prefix when + the natural hash collides with a prefix already assigned to another + server (see ``MCPServerManager._assign_unique_short_prefix``). An + empty ``server_id`` raises ``ValueError`` — short prefixes require a + stable identifier to be deterministic. """ if not server_id: raise ValueError("compute_short_server_prefix requires a non-empty server_id") @@ -71,10 +85,17 @@ def compute_short_server_prefix(server_id: str, attempt: int = 0) -> str: seed = server_id if attempt == 0 else f"{server_id}#{attempt}" digest = hashlib.sha256(seed.encode("utf-8")).digest() value = int.from_bytes(digest[:8], "big") + + # Build chars from least-significant to most-significant; we reverse + # at the end so the first emitted char comes from the high-order + # bits of the digest (which is the position we constrain to be + # alphabetic). chars = [] - for _ in range(SHORT_MCP_TOOL_PREFIX_LENGTH): - value, idx = divmod(value, len(_BASE62_ALPHABET)) - chars.append(_BASE62_ALPHABET[idx]) + for position in range(SHORT_MCP_TOOL_PREFIX_LENGTH): + is_first_char = position == SHORT_MCP_TOOL_PREFIX_LENGTH - 1 + alphabet = _BASE52_ALPHA_ALPHABET if is_first_char else _BASE62_ALPHABET + value, idx = divmod(value, len(alphabet)) + chars.append(alphabet[idx]) return "".join(reversed(chars)) diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_short_mcp_tool_prefix.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_short_mcp_tool_prefix.py index cdaecfb227..bf90e9ebef 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_short_mcp_tool_prefix.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_short_mcp_tool_prefix.py @@ -57,6 +57,20 @@ class TestShortPrefixHelpers: assert len(prefix) == SHORT_MCP_TOOL_PREFIX_LENGTH assert prefix.isalnum() and prefix.isascii() + def test_short_prefix_first_char_is_alphabetic(self): + """The first char must be [A-Za-z] so the prefix is a valid identifier + on every model API (some providers historically required the first + character of a function name to be alphabetic).""" + # Sweep many server_ids and rehash attempts to give us coverage of + # every position the high-order bits can land on. + for i in range(200): + for attempt in range(4): + prefix = compute_short_server_prefix(f"server-{i}", attempt=attempt) + assert prefix[0].isalpha(), ( + f"prefix {prefix!r} for server-{i} (attempt={attempt}) " + f"starts with a non-alphabetic character" + ) + def test_short_prefix_is_deterministic(self): assert compute_short_server_prefix("abc") == compute_short_server_prefix("abc") assert compute_short_server_prefix("abc") != compute_short_server_prefix("abd") From 1c9c219a74e060271bb62402c8c05528c3a4055a Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Tue, 28 Apr 2026 23:44:34 -0700 Subject: [PATCH 48/57] fix(proxy): self-heal Prisma read paths + harden reconnect state machine MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Two related fixes layered on top of the existing reconnect plumbing: 1. Restore reconnect-and-retry on `PrismaClient.get_generic_data` (issue #25143). 1.83.x lost the transport-reconnect-and-retry-once branch that 1.82.6 had on this method, so transient `httpx.ReadError` flaps now surface immediately as `db_exceptions` alerts. `_update_config_from_db` fans out four concurrent `get_generic_data` reads, so a single transport blip used to mark four alerts and a stale config window. Adds `call_with_db_reconnect_retry` to `litellm/proxy/db/exception_handler.py` — a single canonical "try DB read, on transport error reconnect once and retry once" wrapper. Mirrors the inline pattern in `auth_checks._fetch_key_object_from_db_with_reconnect` so we have one implementation rather than three drifting copies, and gives future read paths a clean opt-in. 2. Fix the `_engine_confirmed_dead` flag-reset bug in `_run_reconnect_cycle`. The flag was cleared before `_do_heavy_reconnect()` ran, so any failure inside the heavy reconnect (timeout, missing DATABASE_URL, recreate failure) left the flag False — and the next attempt could silently demote to the lightweight path even though the engine was genuinely dead. Move the reset into the success branch so the flag stays True across heavy-reconnect failures and the next attempt re-enters the heavy branch. Tests: - `tests/test_litellm/proxy/db/test_exception_handler_reconnect_retry.py` (new) — 9 tests covering the helper's contract: happy path, retry on transport error, no retry on data-layer errors, propagation when reconnect fails, propagation after second transport error, `hasattr` guard for partial mocks, fresh-coroutine-per-call invariant, explicit timeout override, default timeouts read off the prisma_client. - `tests/test_litellm/proxy/db/test_prisma_self_heal.py` — adds: - `test_get_generic_data_retries_on_transport_error_for_config_table` - `test_get_generic_data_propagates_when_reconnect_fails` - `test_engine_confirmed_dead_persists_across_failed_heavy_reconnect` (regression test for the flag-reset bug). All 16 self-heal tests + 9 helper tests + 535 auth/exception-handler tests pass locally. --- litellm/proxy/db/exception_handler.py | 123 +++++++++- litellm/proxy/utils.py | 42 +++- .../test_exception_handler_reconnect_retry.py | 225 ++++++++++++++++++ .../proxy/db/test_prisma_self_heal.py | 123 ++++++++++ 4 files changed, 501 insertions(+), 12 deletions(-) create mode 100644 tests/test_litellm/proxy/db/test_exception_handler_reconnect_retry.py diff --git a/litellm/proxy/db/exception_handler.py b/litellm/proxy/db/exception_handler.py index cfa90a48ee..f096773200 100644 --- a/litellm/proxy/db/exception_handler.py +++ b/litellm/proxy/db/exception_handler.py @@ -1,5 +1,6 @@ -from typing import Union +from typing import Any, Awaitable, Callable, Optional, Union +from litellm._logging import verbose_proxy_logger from litellm.proxy._types import ( DB_CONNECTION_ERROR_TYPES, ProxyErrorTypes, @@ -123,3 +124,123 @@ class PrismaDBExceptionHandler: ): return None raise e + + +# Default fallback timeouts when neither the caller nor the prisma_client +# expose `_db_auth_reconnect_timeout_seconds` / `_db_auth_reconnect_lock_timeout_seconds`. +# Match the auth path's existing defaults so behavior is uniform across read paths. +_DEFAULT_RECONNECT_TIMEOUT_SECONDS = 2.0 +_DEFAULT_RECONNECT_LOCK_TIMEOUT_SECONDS = 0.1 + + +def _coerce_timeout(value: Any, fallback: float) -> float: + """Return `value` if it is a real int/float, else `fallback`. Guards + against tests that mock `prisma_client` and leave the timeout slots as + MagicMock instances.""" + if isinstance(value, (int, float)) and not isinstance(value, bool): + return float(value) + return fallback + + +async def call_with_db_reconnect_retry( + prisma_client: Any, + coro_factory: Callable[[], Awaitable[Any]], + *, + reason: str, + timeout_seconds: Optional[float] = None, + lock_timeout_seconds: Optional[float] = None, +) -> Any: + """Run a Prisma read coroutine with one transport-reconnect-and-retry. + + The canonical "self-heal a transient DB transport blip" wrapper used by + `PrismaClient.get_generic_data` and other read paths. Mirrors the inline + pattern in `auth_checks._fetch_key_object_from_db_with_reconnect` so we + have a single implementation rather than three drifting copies. + + Behavior: + 1. Await `coro_factory()`. On success, return its value. + 2. On exception, if it is NOT a transport error (per + `is_database_transport_error`), re-raise — data-layer errors like + `UniqueViolationError` mean the DB is reachable, reconnect would be + pointless. + 3. If `prisma_client` does not expose `attempt_db_reconnect`, re-raise. + This guards against partial stand-ins / older clients in tests. + 4. Call `prisma_client.attempt_db_reconnect(reason=...)`. If it returns + False (cooldown / lock contention / reconnect failure), re-raise. + 5. Otherwise await `coro_factory()` a second time and return / propagate + its result. At-most-one retry by construction — no infinite loop. + + `coro_factory` MUST be a zero-arg callable that returns a fresh awaitable + on each call. Passing an already-awaited coroutine would fail on retry + with `RuntimeError: cannot reuse already awaited coroutine`. + + `reason` should follow `___failure` so + telemetry distinguishes between fan-out callers (e.g. + `_update_config_from_db` issues four concurrent reads). + + Args: + prisma_client: The `PrismaClient` (or stand-in) that owns + `attempt_db_reconnect` and the `_db_auth_reconnect_*` defaults. + coro_factory: Zero-arg callable returning the read awaitable. + reason: Telemetry tag forwarded to `attempt_db_reconnect`. + timeout_seconds: Optional override for the reconnect cycle timeout. + Defaults to `prisma_client._db_auth_reconnect_timeout_seconds`, + then to 2.0s. + lock_timeout_seconds: Optional override for how long the helper will + wait to acquire the reconnect lock. Defaults to + `prisma_client._db_auth_reconnect_lock_timeout_seconds`, then to + 0.1s. + + Returns: + Whatever `coro_factory()` returns (on first or second attempt). + + Raises: + Whatever `coro_factory()` raises if the failure is not a transport + error, or if the reconnect attempt does not succeed, or if the retry + also fails. + """ + try: + return await coro_factory() + except Exception as first_exc: + if not PrismaDBExceptionHandler.is_database_transport_error(first_exc): + raise + if not hasattr(prisma_client, "attempt_db_reconnect"): + raise + + resolved_timeout = _coerce_timeout( + ( + timeout_seconds + if timeout_seconds is not None + else getattr(prisma_client, "_db_auth_reconnect_timeout_seconds", None) + ), + _DEFAULT_RECONNECT_TIMEOUT_SECONDS, + ) + resolved_lock_timeout = _coerce_timeout( + ( + lock_timeout_seconds + if lock_timeout_seconds is not None + else getattr( + prisma_client, "_db_auth_reconnect_lock_timeout_seconds", None + ) + ), + _DEFAULT_RECONNECT_LOCK_TIMEOUT_SECONDS, + ) + + verbose_proxy_logger.warning( + "DB transport error on read; attempting reconnect-and-retry. reason=%s error=%s", + reason, + first_exc, + ) + + did_reconnect = await prisma_client.attempt_db_reconnect( + reason=reason, + timeout_seconds=resolved_timeout, + lock_timeout_seconds=resolved_lock_timeout, + ) + if not did_reconnect: + raise + + # At most one retry. If the retry also raises a transport error, we + # propagate — repeated reconnect-loops are the watchdog's job, not + # this helper's. + return await coro_factory() diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 3a1184c434..93cc8bfd22 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -106,7 +106,10 @@ from litellm.proxy.db.create_views import ( should_create_missing_views, ) from litellm.proxy.db.db_spend_update_writer import DBSpendUpdateWriter -from litellm.proxy.db.exception_handler import PrismaDBExceptionHandler +from litellm.proxy.db.exception_handler import ( + PrismaDBExceptionHandler, + call_with_db_reconnect_retry, +) from litellm.proxy.db.log_db_metrics import log_db_metrics from litellm.proxy.db.prisma_client import PrismaWrapper from litellm.proxy.guardrails.guardrail_hooks.unified_guardrail.unified_guardrail import ( @@ -2779,30 +2782,42 @@ class PrismaClient: table_name: Literal["users", "keys", "config", "spend"], ): """ - Generic implementation of get data + Generic implementation of get data. + + Self-heals across a single transient transport blip via + `call_with_db_reconnect_retry`: on `httpx.ReadError` / + `ClientNotConnectedError` / similar, attempt one DB reconnect and + retry once before surfacing the failure. Restores the 1.82.6 behavior + that was lost in 1.83.x — see issue #25143. """ start_time = time.time() - try: + + async def _do_query(): if table_name == "users": - response = await self.db.litellm_usertable.find_first( + return await self.db.litellm_usertable.find_first( where={key: value} # type: ignore ) elif table_name == "keys": - response = await self.db.litellm_verificationtoken.find_first( # type: ignore + return await self.db.litellm_verificationtoken.find_first( # type: ignore where={key: value} # type: ignore ) elif table_name == "config": - response = await self.db.litellm_config.find_first( # type: ignore + return await self.db.litellm_config.find_first( # type: ignore where={key: value} # type: ignore ) elif table_name == "spend": - response = await self.db.l.find_first( # type: ignore + return await self.db.l.find_first( # type: ignore where={key: value} # type: ignore ) - return response - except Exception as e: - import traceback + return None + try: + return await call_with_db_reconnect_retry( + self, + _do_query, + reason=f"prisma_get_generic_data_{table_name}_lookup_failure", + ) + except Exception as e: error_msg = f"LiteLLM Prisma Client Exception get_generic_data: {str(e)}" verbose_proxy_logger.error(error_msg) error_msg = error_msg + "\nException Type: {}".format(type(e)) @@ -4204,7 +4219,6 @@ class PrismaClient: ) self._reap_all_zombies() self._cleanup_engine_watcher() - self._engine_confirmed_dead = False async def _do_heavy_reconnect() -> None: db_url = os.getenv("DATABASE_URL", "") @@ -4217,6 +4231,12 @@ class PrismaClient: await self._start_engine_watcher() await asyncio.wait_for(_do_heavy_reconnect(), timeout=effective_timeout) + # Only clear the "dead engine" flag after the heavy reconnect + # actually completed. If `_do_heavy_reconnect()` raises (timeout, + # missing DATABASE_URL, recreate failure), the flag stays True so + # the next attempt re-enters the heavy branch instead of silently + # demoting to the lightweight path. + self._engine_confirmed_dead = False else: verbose_proxy_logger.debug( "Performing Prisma DB reconnect (engine alive or unknown)." diff --git a/tests/test_litellm/proxy/db/test_exception_handler_reconnect_retry.py b/tests/test_litellm/proxy/db/test_exception_handler_reconnect_retry.py new file mode 100644 index 0000000000..22ccb02b7a --- /dev/null +++ b/tests/test_litellm/proxy/db/test_exception_handler_reconnect_retry.py @@ -0,0 +1,225 @@ +""" +Unit tests for `call_with_db_reconnect_retry` — the canonical "try DB read, +on transport error reconnect once and retry once" helper. + +Covers the regression in issue #25143 where read paths (e.g. +`PrismaClient.get_generic_data`) lost their reconnect-and-retry-once branch in +LiteLLM 1.83.x and started emitting `db_exceptions` alerts on transient +`httpx.ReadError` flaps that used to self-heal in 1.82.6. +""" + +import os +import sys +from unittest.mock import AsyncMock, MagicMock + +import httpx +import pytest +from prisma.errors import UniqueViolationError + +sys.path.insert(0, os.path.abspath("../../..")) + +from litellm.proxy.db.exception_handler import call_with_db_reconnect_retry + + +def _make_client( + *, + attempt_db_reconnect_return: bool = True, + has_attempt_db_reconnect: bool = True, +): + """Build a minimal stand-in for PrismaClient that exposes only the surface + `call_with_db_reconnect_retry` actually pokes at.""" + client = MagicMock() + if has_attempt_db_reconnect: + client.attempt_db_reconnect = AsyncMock( + return_value=attempt_db_reconnect_return + ) + else: + # `hasattr(client, "attempt_db_reconnect")` must return False — MagicMock + # auto-creates attributes, so we wipe it out via `spec`. + client = MagicMock(spec=[]) + client._db_auth_reconnect_timeout_seconds = 2.0 + client._db_auth_reconnect_lock_timeout_seconds = 0.1 + return client + + +@pytest.mark.asyncio +async def test_call_with_db_reconnect_retry_returns_value_on_first_success(): + """Happy path: factory succeeds first call, no reconnect attempted.""" + client = _make_client() + + async def _factory(): + return {"id": 1} + + result = await call_with_db_reconnect_retry(client, _factory, reason="happy_path") + + assert result == {"id": 1} + client.attempt_db_reconnect.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_call_with_db_reconnect_retry_retries_after_transport_error(): + """Transport error on first call → reconnect → second call succeeds.""" + client = _make_client(attempt_db_reconnect_return=True) + + invocations = [] + + async def _factory(): + invocations.append(None) + if len(invocations) == 1: + raise httpx.ReadError("transport blip") + return {"id": 1} + + result = await call_with_db_reconnect_retry( + client, _factory, reason="prisma_get_generic_data_config_lookup_failure" + ) + + assert result == {"id": 1} + assert len(invocations) == 2 + client.attempt_db_reconnect.assert_awaited_once() + call_kwargs = client.attempt_db_reconnect.await_args.kwargs + assert call_kwargs["reason"] == "prisma_get_generic_data_config_lookup_failure" + + +@pytest.mark.asyncio +async def test_call_with_db_reconnect_retry_does_not_retry_on_data_layer_error(): + """Data-layer errors (e.g. UniqueViolationError) are NOT transport errors — + propagate immediately, do not reconnect.""" + client = _make_client() + + async def _factory(): + raise UniqueViolationError( + data={"user_facing_error": {"meta": {}}}, + message="Unique constraint failed", + ) + + with pytest.raises(UniqueViolationError): + await call_with_db_reconnect_retry(client, _factory, reason="data_layer_test") + + client.attempt_db_reconnect.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_call_with_db_reconnect_retry_propagates_when_reconnect_fails(): + """Transport error, but reconnect returns False → propagate the original + exception. Do not call factory a second time.""" + client = _make_client(attempt_db_reconnect_return=False) + + invocations = [] + + async def _factory(): + invocations.append(None) + raise httpx.ReadError("transport blip") + + with pytest.raises(httpx.ReadError): + await call_with_db_reconnect_retry(client, _factory, reason="reconnect_fails") + + assert len(invocations) == 1 + client.attempt_db_reconnect.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_call_with_db_reconnect_retry_propagates_after_second_transport_error(): + """Transport error, reconnect succeeds, retry also raises transport error → + propagate. At most one retry by construction (no infinite loop).""" + client = _make_client(attempt_db_reconnect_return=True) + + invocations = [] + + async def _factory(): + invocations.append(None) + raise httpx.ReadError("still failing") + + with pytest.raises(httpx.ReadError): + await call_with_db_reconnect_retry( + client, _factory, reason="second_transport_error" + ) + + assert len(invocations) == 2 + client.attempt_db_reconnect.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_call_with_db_reconnect_retry_skips_when_no_attempt_db_reconnect_attr(): + """Older PrismaClient stand-ins / partial mocks may not expose + `attempt_db_reconnect`. The helper must not crash — just propagate the + original exception. Mirrors the `hasattr` guard from + `auth_checks._fetch_key_object_from_db_with_reconnect`.""" + client = _make_client(has_attempt_db_reconnect=False) + + async def _factory(): + raise httpx.ReadError("transport blip") + + with pytest.raises(httpx.ReadError): + await call_with_db_reconnect_retry(client, _factory, reason="no_reconnect_attr") + + +@pytest.mark.asyncio +async def test_call_with_db_reconnect_retry_invokes_factory_twice_not_same_coro(): + """Guard against the obvious bug of awaiting the same coroutine twice + (`RuntimeError: cannot reuse already awaited coroutine`). The helper must + call the factory a fresh time on retry, not cache an awaitable.""" + client = _make_client(attempt_db_reconnect_return=True) + + factory_call_count = 0 + + async def _factory(): + nonlocal factory_call_count + factory_call_count += 1 + if factory_call_count == 1: + raise httpx.ReadError("transport blip") + return "ok" + + result = await call_with_db_reconnect_retry( + client, _factory, reason="fresh_coro_on_retry" + ) + + assert result == "ok" + assert factory_call_count == 2 + + +@pytest.mark.asyncio +async def test_call_with_db_reconnect_retry_passes_explicit_timeouts(): + """Explicit timeout_seconds / lock_timeout_seconds override the auth + defaults read off the prisma_client object.""" + client = _make_client(attempt_db_reconnect_return=True) + + async def _factory(): + if not hasattr(_factory, "_called"): + _factory._called = True # type: ignore[attr-defined] + raise httpx.ReadError("transport blip") + return "ok" + + result = await call_with_db_reconnect_retry( + client, + _factory, + reason="explicit_timeouts", + timeout_seconds=5.5, + lock_timeout_seconds=0.25, + ) + + assert result == "ok" + call_kwargs = client.attempt_db_reconnect.await_args.kwargs + assert call_kwargs["timeout_seconds"] == 5.5 + assert call_kwargs["lock_timeout_seconds"] == 0.25 + + +@pytest.mark.asyncio +async def test_call_with_db_reconnect_retry_uses_auth_defaults_when_unset(): + """When timeouts are not provided, helper reads + `_db_auth_reconnect_timeout_seconds` / `_db_auth_reconnect_lock_timeout_seconds` + off the prisma_client (matching the auth path's existing convention).""" + client = _make_client(attempt_db_reconnect_return=True) + client._db_auth_reconnect_timeout_seconds = 3.0 + client._db_auth_reconnect_lock_timeout_seconds = 0.5 + + async def _factory(): + if not hasattr(_factory, "_called"): + _factory._called = True # type: ignore[attr-defined] + raise httpx.ReadError("transport blip") + return "ok" + + await call_with_db_reconnect_retry(client, _factory, reason="defaults") + + call_kwargs = client.attempt_db_reconnect.await_args.kwargs + assert call_kwargs["timeout_seconds"] == 3.0 + assert call_kwargs["lock_timeout_seconds"] == 0.5 diff --git a/tests/test_litellm/proxy/db/test_prisma_self_heal.py b/tests/test_litellm/proxy/db/test_prisma_self_heal.py index fb215e5477..af65cad09e 100644 --- a/tests/test_litellm/proxy/db/test_prisma_self_heal.py +++ b/tests/test_litellm/proxy/db/test_prisma_self_heal.py @@ -5,6 +5,7 @@ import sys import time from unittest.mock import AsyncMock, MagicMock, patch +import httpx import pytest sys.path.insert( @@ -358,3 +359,125 @@ async def test_lightweight_reconnect_skips_kill_on_successful_disconnect( await client._run_reconnect_cycle(timeout_seconds=5.0) mock_kill.assert_not_called() + + +# --------------------------------------------------------------------------- +# get_generic_data: transport-reconnect-and-retry coverage (issue #25143) +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_get_generic_data_retries_on_transport_error_for_config_table( + mock_proxy_logging, +): + """`get_generic_data(table_name="config")` self-heals on a transient + `httpx.ReadError`: reconnect once, retry once, return the row. + + Regression for issue #25143 — the 1.83.x line lost the reconnect-and-retry + branch that 1.82.6 had on this method. `_update_config_from_db` fans out + four concurrent `get_generic_data` calls, so a single transport flap used + to surface as four `db_exceptions` alerts and a stale config window. + """ + client = PrismaClient( + database_url="mock://test", proxy_logging_obj=mock_proxy_logging + ) + + expected_row = {"param_name": "general_settings", "param_value": {"foo": "bar"}} + invocations: list[None] = [] + + async def _flaky_find_first(**kwargs): + invocations.append(None) + if len(invocations) == 1: + raise httpx.ReadError("simulated transport blip") + return expected_row + + client.db.litellm_config.find_first = AsyncMock(side_effect=_flaky_find_first) + client.attempt_db_reconnect = AsyncMock(return_value=True) + + result = await client.get_generic_data( + key="param_name", + value="general_settings", + table_name="config", + ) + + assert result == expected_row + assert len(invocations) == 2 + client.attempt_db_reconnect.assert_awaited_once() + reconnect_kwargs = client.attempt_db_reconnect.await_args.kwargs + assert reconnect_kwargs["reason"] == "prisma_get_generic_data_config_lookup_failure" + + # The failure_handler telemetry side-effect must NOT fire on the first + # transport blip — only if the post-retry call also fails. Drain the + # event loop so any spuriously-spawned task would have run by now. + await asyncio.sleep(0) + mock_proxy_logging.failure_handler.assert_not_called() + + +@pytest.mark.asyncio +async def test_get_generic_data_propagates_when_reconnect_fails(mock_proxy_logging): + """If reconnect itself does not succeed, propagate the original transport + error and let the existing failure_handler / db_exceptions telemetry fire.""" + client = PrismaClient( + database_url="mock://test", proxy_logging_obj=mock_proxy_logging + ) + + client.db.litellm_config.find_first = AsyncMock( + side_effect=httpx.ReadError("simulated transport blip") + ) + client.attempt_db_reconnect = AsyncMock(return_value=False) + + with pytest.raises(httpx.ReadError): + await client.get_generic_data( + key="param_name", + value="general_settings", + table_name="config", + ) + + client.attempt_db_reconnect.assert_awaited_once() + # Failure telemetry IS expected here — the read genuinely failed. + await asyncio.sleep(0) + mock_proxy_logging.failure_handler.assert_called_once() + + +# --------------------------------------------------------------------------- +# _engine_confirmed_dead flag-reset bug (B2) +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_engine_confirmed_dead_persists_across_failed_heavy_reconnect( + mock_proxy_logging, +): + """Regression test for the flag-reset bug. + + Before the fix, `_run_reconnect_cycle` cleared + `self._engine_confirmed_dead = False` *before* awaiting + `_do_heavy_reconnect()`. If the heavy reconnect raised (e.g. timeout, + missing DATABASE_URL, recreate failure), the flag was left cleared and the + next attempt could demote to the lightweight path even though the engine + was genuinely dead. + + The fix moves the reset into the success branch — the flag must stay True + when heavy reconnect raises. + """ + client = PrismaClient( + database_url="mock://test", proxy_logging_obj=mock_proxy_logging + ) + client._engine_confirmed_dead = True + client._engine_pid = 0 # so `_is_engine_alive` is not consulted + + # Make the heavy reconnect path raise. + client.db.recreate_prisma_client = AsyncMock( + side_effect=RuntimeError("simulated heavy reconnect failure") + ) + client._start_engine_watcher = AsyncMock() + client._cleanup_engine_watcher = MagicMock() + client._reap_all_zombies = MagicMock() + + with patch.dict(os.environ, {"DATABASE_URL": "postgresql://test"}): + with pytest.raises(Exception): + await client._run_reconnect_cycle(timeout_seconds=5.0) + + # The flag must STILL be True so the next attempt re-enters the heavy + # branch instead of silently demoting to the lightweight path. + assert client._engine_confirmed_dead is True From aa2ef4120098981cb6160d94347743e2e77ffb61 Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Tue, 28 Apr 2026 23:55:46 -0700 Subject: [PATCH 49/57] fix(proxy): preserve original transport error if reconnect itself raises MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Greptile review on #26756 (P2): if `attempt_db_reconnect` itself raises (e.g. lock cancellation, timer error, unexpected internal failure), the original `httpx.ReadError` / transport error was lost — `failure_handler` and `db_exceptions` alerts then logged the reconnect exception instead of the actual DB transport problem, masking the root cause. Wrap the reconnect call in a try/except. On reconnect failure, re-raise the *original* `first_exc` and chain the reconnect error as `__cause__` so it remains visible for debuggability without becoming the primary exception observers see. Adds `test_call_with_db_reconnect_retry_preserves_original_error_when_reconnect_raises` asserting (a) the propagated exception is the original transport error and (b) the reconnect exception is attached as `__cause__`. --- litellm/proxy/db/exception_handler.py | 25 ++++++++++++---- .../test_exception_handler_reconnect_retry.py | 30 +++++++++++++++++++ 2 files changed, 50 insertions(+), 5 deletions(-) diff --git a/litellm/proxy/db/exception_handler.py b/litellm/proxy/db/exception_handler.py index f096773200..ab9d341aa5 100644 --- a/litellm/proxy/db/exception_handler.py +++ b/litellm/proxy/db/exception_handler.py @@ -232,11 +232,26 @@ async def call_with_db_reconnect_retry( first_exc, ) - did_reconnect = await prisma_client.attempt_db_reconnect( - reason=reason, - timeout_seconds=resolved_timeout, - lock_timeout_seconds=resolved_lock_timeout, - ) + # Preserve the original transport error in telemetry. If + # `attempt_db_reconnect` itself raises (e.g. lock cancellation, timer + # error, unexpected internal failure), surfacing that exception + # instead of `first_exc` would mask the actual DB transport problem + # in `failure_handler` / `db_exceptions` alerts. Chain the reconnect + # error as the cause for debuggability without losing the original. + try: + did_reconnect = await prisma_client.attempt_db_reconnect( + reason=reason, + timeout_seconds=resolved_timeout, + lock_timeout_seconds=resolved_lock_timeout, + ) + except Exception as reconnect_exc: + verbose_proxy_logger.warning( + "DB reconnect attempt raised; preserving original transport error. " + "reason=%s reconnect_error=%s", + reason, + reconnect_exc, + ) + raise first_exc from reconnect_exc if not did_reconnect: raise diff --git a/tests/test_litellm/proxy/db/test_exception_handler_reconnect_retry.py b/tests/test_litellm/proxy/db/test_exception_handler_reconnect_retry.py index 22ccb02b7a..ae0e1f845b 100644 --- a/tests/test_litellm/proxy/db/test_exception_handler_reconnect_retry.py +++ b/tests/test_litellm/proxy/db/test_exception_handler_reconnect_retry.py @@ -223,3 +223,33 @@ async def test_call_with_db_reconnect_retry_uses_auth_defaults_when_unset(): call_kwargs = client.attempt_db_reconnect.await_args.kwargs assert call_kwargs["timeout_seconds"] == 3.0 assert call_kwargs["lock_timeout_seconds"] == 0.5 + + +@pytest.mark.asyncio +async def test_call_with_db_reconnect_retry_preserves_original_error_when_reconnect_raises(): + """If `attempt_db_reconnect` itself raises (lock cancellation, timer + error, unexpected internal failure), the helper must surface the + *original* transport error to telemetry — not the reconnect exception. + Otherwise `failure_handler` / `db_exceptions` alerts log the wrong + error string and the actual DB transport problem becomes invisible. + + The reconnect error is chained as the `__cause__` for debuggability.""" + client = MagicMock() + reconnect_exc = RuntimeError("simulated reconnect lock cancellation") + client.attempt_db_reconnect = AsyncMock(side_effect=reconnect_exc) + client._db_auth_reconnect_timeout_seconds = 2.0 + client._db_auth_reconnect_lock_timeout_seconds = 0.1 + + original_exc = httpx.ReadError("transport blip") + + async def _factory(): + raise original_exc + + with pytest.raises(httpx.ReadError) as exc_info: + await call_with_db_reconnect_retry( + client, _factory, reason="reconnect_itself_raises" + ) + + assert exc_info.value is original_exc + assert exc_info.value.__cause__ is reconnect_exc + client.attempt_db_reconnect.assert_awaited_once() From 06d9a694441cdf399f86be26e441fbca42a03df6 Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Tue, 28 Apr 2026 23:57:28 -0700 Subject: [PATCH 50/57] docs(proxy): clarify _kill_engine_process is on the routine reconnect path MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Greptile review on #26225 (P2): the docstring said "Called when disconnect() fails", and the SIGTERM warning log read "after failed disconnect", but both were stale — `_kill_engine_process` is now invoked on every routine reconnect (via the unified `recreate_prisma_client` path), not as a disconnect-failure recovery branch. The misleading wording would have produced confusing log lines on every reconnect cycle in production. Update the docstring to explain the actual reason (avoiding the blocking `disconnect()` event-loop freeze) and reword the SIGTERM warning to "during reconnect" so it matches reality. No behavior change; logs only. --- litellm/proxy/db/prisma_client.py | 15 ++++++++++----- 1 file changed, 10 insertions(+), 5 deletions(-) diff --git a/litellm/proxy/db/prisma_client.py b/litellm/proxy/db/prisma_client.py index a8942f92a1..d112e22230 100644 --- a/litellm/proxy/db/prisma_client.py +++ b/litellm/proxy/db/prisma_client.py @@ -61,11 +61,16 @@ class PrismaWrapper: @staticmethod async def _kill_engine_process(pid: int) -> None: - """Force-kill an orphaned engine subprocess to prevent DB connection pool leaks. + """Force-kill the engine subprocess to prevent DB connection pool leaks. - Called when disconnect() fails and the old engine process may still be - holding open connections. Sends SIGTERM for graceful shutdown, waits - briefly, then SIGKILL as a backstop. + Called on every reconnect (in `recreate_prisma_client`) to retire the + old query-engine subprocess without invoking prisma-client-py's + synchronous `disconnect()` — which blocks the asyncio event loop on + `subprocess.Popen.wait()` for 30-120+ seconds when the engine is + stuck on TCP close. + + Sends SIGTERM for graceful shutdown, waits briefly, then SIGKILL as + a backstop. """ if pid <= 0: return @@ -74,7 +79,7 @@ class PrismaWrapper: except (ProcessLookupError, PermissionError, OSError): return # Already dead or inaccessible verbose_proxy_logger.warning( - "Sent SIGTERM to orphaned prisma-query-engine PID %s after failed disconnect.", + "Sent SIGTERM to prisma-query-engine PID %s during reconnect.", pid, ) # Brief wait for graceful shutdown, then force-kill From 848b79acb5e8f0d36e3b16321dc0fdabfaecd581 Mon Sep 17 00:00:00 2001 From: Yassin Kortam Date: Tue, 28 Apr 2026 18:04:29 -0700 Subject: [PATCH 51/57] fix: added keepalive args for aiohttp tcpconnector --- litellm/constants.py | 10 ++ litellm/llms/custom_httpx/http_handler.py | 62 +++++++ litellm/proxy/proxy_server.py | 7 + .../custom_httpx/test_aiohttp_so_keepalive.py | 161 ++++++++++++++++++ 4 files changed, 240 insertions(+) create mode 100644 tests/test_litellm/llms/custom_httpx/test_aiohttp_so_keepalive.py diff --git a/litellm/constants.py b/litellm/constants.py index fc9f5730cd..a0e99dd16b 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -224,6 +224,16 @@ AIOHTTP_CONNECTOR_LIMIT_PER_HOST = int( ) AIOHTTP_KEEPALIVE_TIMEOUT = int(os.getenv("AIOHTTP_KEEPALIVE_TIMEOUT", 120)) AIOHTTP_TTL_DNS_CACHE = int(os.getenv("AIOHTTP_TTL_DNS_CACHE", 300)) +# TCP keep-alive (SO_KEEPALIVE) — opt-in. Required when running behind NAT/LBs +# whose idle timeout is shorter than provider response timeouts (e.g. AWS NAT +# Gateway: 350s vs OpenAI/Azure: 600s). Without this, the kernel sends nothing +# during a long provider call and the NAT reaps the flow before the response +# arrives. Enabling SO_KEEPALIVE makes the kernel emit TCP probes that reset +# the NAT idle timer. +AIOHTTP_SO_KEEPALIVE = os.getenv("AIOHTTP_SO_KEEPALIVE", "False").lower() == "true" +AIOHTTP_TCP_KEEPIDLE = int(os.getenv("AIOHTTP_TCP_KEEPIDLE", 60)) +AIOHTTP_TCP_KEEPINTVL = int(os.getenv("AIOHTTP_TCP_KEEPINTVL", 30)) +AIOHTTP_TCP_KEEPCNT = int(os.getenv("AIOHTTP_TCP_KEEPCNT", 5)) # enable_cleanup_closed is only needed for Python versions with the SSL leak bug # Fixed in Python 3.12.7+ and 3.13.1+ (see https://github.com/python/cpython/pull/118960) # Reference: https://github.com/aio-libs/aiohttp/blob/master/aiohttp/connector.py#L74-L78 diff --git a/litellm/llms/custom_httpx/http_handler.py b/litellm/llms/custom_httpx/http_handler.py index 03d2af7232..dd955c23d2 100644 --- a/litellm/llms/custom_httpx/http_handler.py +++ b/litellm/llms/custom_httpx/http_handler.py @@ -1,5 +1,7 @@ import asyncio +import inspect import os +import socket import ssl import sys import time @@ -29,6 +31,10 @@ from litellm.constants import ( AIOHTTP_CONNECTOR_LIMIT_PER_HOST, AIOHTTP_KEEPALIVE_TIMEOUT, AIOHTTP_NEEDS_CLEANUP_CLOSED, + AIOHTTP_SO_KEEPALIVE, + AIOHTTP_TCP_KEEPCNT, + AIOHTTP_TCP_KEEPIDLE, + AIOHTTP_TCP_KEEPINTVL, AIOHTTP_TTL_DNS_CACHE, COMPLETION_HTTP_FALLBACK_SECONDS, DEFAULT_SSL_CIPHERS, @@ -54,6 +60,57 @@ except Exception: version = "0.0.0" +# aiohttp 3.10+ exposes a `socket_factory` kwarg on TCPConnector. Older +# versions don't — detect once and skip the keep-alive wiring there. +# https://docs.aiohttp.org/en/stable/client_reference.html#aiohttp.TCPConnector +_AIOHTTP_SUPPORTS_SOCKET_FACTORY = ( + "socket_factory" in inspect.signature(TCPConnector.__init__).parameters +) + + +def _build_aiohttp_keepalive_socket_factory() -> ( + Optional[Callable[[Tuple[Any, ...]], socket.socket]] +): + """ + Build a socket_factory that enables SO_KEEPALIVE on aiohttp TCP sockets. + + Why: by default, aiohttp creates sockets without SO_KEEPALIVE, so the kernel + sends nothing during a long idle TCP connection. NAT/LB hops (e.g. AWS NAT + Gateway, 350s idle timeout) reap the flow well before slow provider + responses (OpenAI/Azure: up to 600s) arrive. Enabling SO_KEEPALIVE makes + the kernel emit TCP probes that reset the NAT idle timer. + + Returns None when AIOHTTP_SO_KEEPALIVE is disabled or aiohttp is too old. + """ + if not AIOHTTP_SO_KEEPALIVE or not _AIOHTTP_SUPPORTS_SOCKET_FACTORY: + return None + + def factory(addr_info: Tuple[Any, ...]) -> socket.socket: + family, type_, proto = addr_info[0], addr_info[1], addr_info[2] + sock = socket.socket(family=family, type=type_, proto=proto) + sock.setblocking(False) + sock.setsockopt(socket.SOL_SOCKET, socket.SO_KEEPALIVE, 1) + # Linux: TCP_KEEPIDLE is idle-before-first-probe. + # macOS/Darwin: TCP_KEEPALIVE is the equivalent. + if hasattr(socket, "TCP_KEEPIDLE"): + sock.setsockopt( + socket.IPPROTO_TCP, socket.TCP_KEEPIDLE, AIOHTTP_TCP_KEEPIDLE + ) + elif hasattr(socket, "TCP_KEEPALIVE"): + sock.setsockopt( + socket.IPPROTO_TCP, socket.TCP_KEEPALIVE, AIOHTTP_TCP_KEEPIDLE + ) + if hasattr(socket, "TCP_KEEPINTVL"): + sock.setsockopt( + socket.IPPROTO_TCP, socket.TCP_KEEPINTVL, AIOHTTP_TCP_KEEPINTVL + ) + if hasattr(socket, "TCP_KEEPCNT"): + sock.setsockopt(socket.IPPROTO_TCP, socket.TCP_KEEPCNT, AIOHTTP_TCP_KEEPCNT) + return sock + + return factory + + def get_default_headers() -> dict: """ Get default headers for HTTP requests. @@ -935,6 +992,11 @@ class AsyncHTTPHandler: transport_connector_kwargs["limit_per_host"] = ( AIOHTTP_CONNECTOR_LIMIT_PER_HOST ) + # Returns None when SO_KEEPALIVE is disabled or aiohttp is too old to + # accept socket_factory — version detection lives inside the builder. + socket_factory = _build_aiohttp_keepalive_socket_factory() + if socket_factory is not None: + transport_connector_kwargs["socket_factory"] = socket_factory return LiteLLMAiohttpTransport( client=lambda: ClientSession( diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index c03a63f211..efa9cdb8e1 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -672,6 +672,10 @@ async def _initialize_shared_aiohttp_session(): try: from aiohttp import ClientSession, TCPConnector + from litellm.llms.custom_httpx.http_handler import ( + _build_aiohttp_keepalive_socket_factory, + ) + connector_kwargs: Dict[str, Any] = { "keepalive_timeout": AIOHTTP_KEEPALIVE_TIMEOUT, "ttl_dns_cache": AIOHTTP_TTL_DNS_CACHE, @@ -682,6 +686,9 @@ async def _initialize_shared_aiohttp_session(): connector_kwargs["limit"] = AIOHTTP_CONNECTOR_LIMIT if AIOHTTP_CONNECTOR_LIMIT_PER_HOST > 0: connector_kwargs["limit_per_host"] = AIOHTTP_CONNECTOR_LIMIT_PER_HOST + socket_factory = _build_aiohttp_keepalive_socket_factory() + if socket_factory is not None: + connector_kwargs["socket_factory"] = socket_factory connector = TCPConnector(**connector_kwargs) session = ClientSession(connector=connector) 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 new file mode 100644 index 0000000000..5515e6ce81 --- /dev/null +++ b/tests/test_litellm/llms/custom_httpx/test_aiohttp_so_keepalive.py @@ -0,0 +1,161 @@ +import socket +from unittest.mock import MagicMock, patch + + +def _invoke_connector_factory(http_handler_module): + """ + Drive the lambda factory installed on the transport so TCPConnector is + actually constructed. _create_aiohttp_transport returns a transport whose + _client_factory is the lambda that builds (TCPConnector → ClientSession); + invoking it directly avoids relying on _get_valid_client_session's internal + branching to trigger connector construction. + """ + transport = http_handler_module.AsyncHTTPHandler._create_aiohttp_transport( + shared_session=None + ) + transport._client_factory() + return transport + + +def test_socket_factory_omitted_when_disabled(monkeypatch): + from litellm.llms.custom_httpx import http_handler as http_handler_module + + monkeypatch.setattr(http_handler_module, "AIOHTTP_SO_KEEPALIVE", False) + monkeypatch.setattr(http_handler_module, "_AIOHTTP_SUPPORTS_SOCKET_FACTORY", True) + + connector_mock = MagicMock(name="connector") + session_mock = MagicMock(name="session") + + with patch.object( + http_handler_module, "TCPConnector", return_value=connector_mock + ) as mock_tcp_connector: + with patch.object( + http_handler_module, "ClientSession", return_value=session_mock + ): + _invoke_connector_factory(http_handler_module) + + assert mock_tcp_connector.call_count >= 1 + assert "socket_factory" not in mock_tcp_connector.call_args.kwargs + + +def test_socket_factory_attached_when_enabled(monkeypatch): + from litellm.llms.custom_httpx import http_handler as http_handler_module + + monkeypatch.setattr(http_handler_module, "AIOHTTP_SO_KEEPALIVE", True) + monkeypatch.setattr(http_handler_module, "_AIOHTTP_SUPPORTS_SOCKET_FACTORY", True) + + connector_mock = MagicMock(name="connector") + session_mock = MagicMock(name="session") + + with patch.object( + http_handler_module, "TCPConnector", return_value=connector_mock + ) as mock_tcp_connector: + with patch.object( + http_handler_module, "ClientSession", return_value=session_mock + ): + _invoke_connector_factory(http_handler_module) + + assert mock_tcp_connector.call_count >= 1 + factory = mock_tcp_connector.call_args.kwargs.get("socket_factory") + assert callable(factory) + + +def test_socket_factory_skipped_on_old_aiohttp(monkeypatch): + from litellm.llms.custom_httpx import http_handler as http_handler_module + + monkeypatch.setattr(http_handler_module, "AIOHTTP_SO_KEEPALIVE", True) + monkeypatch.setattr(http_handler_module, "_AIOHTTP_SUPPORTS_SOCKET_FACTORY", False) + + connector_mock = MagicMock(name="connector") + session_mock = MagicMock(name="session") + + with patch.object( + http_handler_module, "TCPConnector", return_value=connector_mock + ) as mock_tcp_connector: + with patch.object( + http_handler_module, "ClientSession", return_value=session_mock + ): + _invoke_connector_factory(http_handler_module) + + assert mock_tcp_connector.call_count >= 1 + assert "socket_factory" not in mock_tcp_connector.call_args.kwargs + + +def test_socket_factory_sets_keepalive_options(monkeypatch): + from litellm.llms.custom_httpx import http_handler as http_handler_module + + monkeypatch.setattr(http_handler_module, "AIOHTTP_SO_KEEPALIVE", True) + monkeypatch.setattr(http_handler_module, "_AIOHTTP_SUPPORTS_SOCKET_FACTORY", True) + monkeypatch.setattr(http_handler_module, "AIOHTTP_TCP_KEEPIDLE", 45) + monkeypatch.setattr(http_handler_module, "AIOHTTP_TCP_KEEPINTVL", 15) + monkeypatch.setattr(http_handler_module, "AIOHTTP_TCP_KEEPCNT", 4) + + factory = http_handler_module._build_aiohttp_keepalive_socket_factory() + assert factory is not None + + addr_info = (socket.AF_INET, socket.SOCK_STREAM, socket.IPPROTO_TCP, "", ("", 0)) + + fake_sock = MagicMock(spec=socket.socket) + with patch("socket.socket", return_value=fake_sock) as sock_ctor: + returned = factory(addr_info) + + sock_ctor.assert_called_once_with( + family=socket.AF_INET, type=socket.SOCK_STREAM, proto=socket.IPPROTO_TCP + ) + assert returned is fake_sock + fake_sock.setblocking.assert_called_once_with(False) + + setsockopt_calls = { + (call.args[0], call.args[1]): call.args[2] + for call in fake_sock.setsockopt.call_args_list + } + assert setsockopt_calls[(socket.SOL_SOCKET, socket.SO_KEEPALIVE)] == 1 + + if hasattr(socket, "TCP_KEEPIDLE"): + assert setsockopt_calls[(socket.IPPROTO_TCP, socket.TCP_KEEPIDLE)] == 45 + elif hasattr(socket, "TCP_KEEPALIVE"): + assert setsockopt_calls[(socket.IPPROTO_TCP, socket.TCP_KEEPALIVE)] == 45 + if hasattr(socket, "TCP_KEEPINTVL"): + assert setsockopt_calls[(socket.IPPROTO_TCP, socket.TCP_KEEPINTVL)] == 15 + if hasattr(socket, "TCP_KEEPCNT"): + assert setsockopt_calls[(socket.IPPROTO_TCP, socket.TCP_KEEPCNT)] == 4 + + +def test_socket_factory_uses_tcp_keepalive_when_keepidle_unavailable(monkeypatch): + """ + Cover the macOS/Darwin branch: when TCP_KEEPIDLE is missing but TCP_KEEPALIVE + is present, the factory should fall back to TCP_KEEPALIVE for the idle timer. + Linux CI runners always have TCP_KEEPIDLE, so we patch socket itself to + simulate the BSD-derived environment. + """ + from litellm.llms.custom_httpx import http_handler as http_handler_module + + monkeypatch.setattr(http_handler_module, "AIOHTTP_SO_KEEPALIVE", True) + monkeypatch.setattr(http_handler_module, "_AIOHTTP_SUPPORTS_SOCKET_FACTORY", True) + monkeypatch.setattr(http_handler_module, "AIOHTTP_TCP_KEEPIDLE", 60) + + factory = http_handler_module._build_aiohttp_keepalive_socket_factory() + assert factory is not None + + fake_socket_module = MagicMock(spec=[]) + fake_socket_module.SOL_SOCKET = socket.SOL_SOCKET + fake_socket_module.SO_KEEPALIVE = socket.SO_KEEPALIVE + fake_socket_module.IPPROTO_TCP = socket.IPPROTO_TCP + fake_socket_module.TCP_KEEPALIVE = getattr(socket, "TCP_KEEPALIVE", 0x10) + fake_sock = MagicMock(spec=socket.socket) + fake_socket_module.socket = MagicMock(return_value=fake_sock) + + addr_info = (socket.AF_INET, socket.SOCK_STREAM, socket.IPPROTO_TCP, "", ("", 0)) + + with patch.object(http_handler_module, "socket", fake_socket_module): + factory(addr_info) + + setsockopt_calls = { + (call.args[0], call.args[1]): call.args[2] + for call in fake_sock.setsockopt.call_args_list + } + assert setsockopt_calls[(socket.SOL_SOCKET, socket.SO_KEEPALIVE)] == 1 + assert ( + setsockopt_calls[(socket.IPPROTO_TCP, fake_socket_module.TCP_KEEPALIVE)] == 60 + ) + assert (socket.IPPROTO_TCP, getattr(socket, "TCP_KEEPIDLE", -1)) not in setsockopt_calls From ea275659ac80ba58aaaf4ab518d5fa2bfb47d4ae Mon Sep 17 00:00:00 2001 From: ishaan-berri <155045088+ishaan-berri@users.noreply.github.com> Date: Wed, 29 Apr 2026 09:28:57 -0700 Subject: [PATCH 52/57] remove /ui/chat page (#26739) * remove /ui/chat static page from dashboard build * add screenshot showing /ui/chat 404 * update screenshots: swagger working, /ui/chat broken * remove screenshots from repo * restore screenshots from previous PR --- litellm/proxy/_experimental/out/chat.html | 1 - litellm/proxy/_experimental/out/chat.txt | 22 ------------------- .../_experimental/out/chat/__next._full.txt | 22 ------------------- .../_experimental/out/chat/__next._head.txt | 6 ----- .../_experimental/out/chat/__next._index.txt | 8 ------- .../_experimental/out/chat/__next._tree.txt | 4 ---- .../out/chat/__next.chat.__PAGE__.txt | 9 -------- .../_experimental/out/chat/__next.chat.txt | 4 ---- 8 files changed, 76 deletions(-) delete mode 100644 litellm/proxy/_experimental/out/chat.html delete mode 100644 litellm/proxy/_experimental/out/chat.txt delete mode 100644 litellm/proxy/_experimental/out/chat/__next._full.txt delete mode 100644 litellm/proxy/_experimental/out/chat/__next._head.txt delete mode 100644 litellm/proxy/_experimental/out/chat/__next._index.txt delete mode 100644 litellm/proxy/_experimental/out/chat/__next._tree.txt delete mode 100644 litellm/proxy/_experimental/out/chat/__next.chat.__PAGE__.txt delete mode 100644 litellm/proxy/_experimental/out/chat/__next.chat.txt diff --git a/litellm/proxy/_experimental/out/chat.html b/litellm/proxy/_experimental/out/chat.html deleted file mode 100644 index 6ca94ee94f..0000000000 --- a/litellm/proxy/_experimental/out/chat.html +++ /dev/null @@ -1 +0,0 @@ -LiteLLM Dashboard \ No newline at end of file diff --git a/litellm/proxy/_experimental/out/chat.txt b/litellm/proxy/_experimental/out/chat.txt deleted file mode 100644 index 8ee24df074..0000000000 --- a/litellm/proxy/_experimental/out/chat.txt +++ /dev/null @@ -1,22 +0,0 @@ -1:"$Sreact.fragment" -2:I[867271,["/litellm-asset-prefix/_next/static/chunks/9e09de50158b3159.js","/litellm-asset-prefix/_next/static/chunks/7e5fe5584502da06.js"],"default"] -3:I[71195,["/litellm-asset-prefix/_next/static/chunks/9e09de50158b3159.js","/litellm-asset-prefix/_next/static/chunks/7e5fe5584502da06.js"],"default"] -4:I[339756,["/litellm-asset-prefix/_next/static/chunks/d96012bcfc98706a.js","/litellm-asset-prefix/_next/static/chunks/dbca964212122d58.js"],"default"] -5:I[837457,["/litellm-asset-prefix/_next/static/chunks/d96012bcfc98706a.js","/litellm-asset-prefix/_next/static/chunks/dbca964212122d58.js"],"default"] -6:I[347257,["/litellm-asset-prefix/_next/static/chunks/d96012bcfc98706a.js","/litellm-asset-prefix/_next/static/chunks/dbca964212122d58.js"],"ClientPageRoot"] -7:I[321443,["/litellm-asset-prefix/_next/static/chunks/9e09de50158b3159.js","/litellm-asset-prefix/_next/static/chunks/7e5fe5584502da06.js","/litellm-asset-prefix/_next/static/chunks/0a6c418370a8c183.js","/litellm-asset-prefix/_next/static/chunks/cf9c81fc7166f4d4.js","/litellm-asset-prefix/_next/static/chunks/542a1a209eb732c6.js","/litellm-asset-prefix/_next/static/chunks/00ff280cdb7d7ee5.js","/litellm-asset-prefix/_next/static/chunks/a1abfc2f35c701cc.js","/litellm-asset-prefix/_next/static/chunks/4980372eaa37b78b.js","/litellm-asset-prefix/_next/static/chunks/119ecee91f911bc8.js","/litellm-asset-prefix/_next/static/chunks/ac3cf77acb5bf234.js","/litellm-asset-prefix/_next/static/chunks/fcdf7322b0aa3e2e.js","/litellm-asset-prefix/_next/static/chunks/7e417dd24c8becd0.js","/litellm-asset-prefix/_next/static/chunks/89aa55578de861b7.js"],"default"] -a:I[897367,["/litellm-asset-prefix/_next/static/chunks/d96012bcfc98706a.js","/litellm-asset-prefix/_next/static/chunks/dbca964212122d58.js"],"OutletBoundary"] -b:"$Sreact.suspense" -d:I[897367,["/litellm-asset-prefix/_next/static/chunks/d96012bcfc98706a.js","/litellm-asset-prefix/_next/static/chunks/dbca964212122d58.js"],"ViewportBoundary"] -f:I[897367,["/litellm-asset-prefix/_next/static/chunks/d96012bcfc98706a.js","/litellm-asset-prefix/_next/static/chunks/dbca964212122d58.js"],"MetadataBoundary"] -11:I[168027,[],"default"] -:HL["/litellm-asset-prefix/_next/static/chunks/4e20891f2fd03463.css","style"] -:HL["/litellm-asset-prefix/_next/static/chunks/e0cb6755699177c1.css","style"] -:HL["/litellm-asset-prefix/_next/static/media/83afe278b6a6bb3c-s.p.3a6ba036.woff2","font",{"crossOrigin":"","type":"font/woff2"}] -0:{"P":null,"b":"zxkD4-EPlgfKHDTw8O869","c":["","chat"],"q":"","i":false,"f":[[["",{"children":["chat",{"children":["__PAGE__",{}]}]},"$undefined","$undefined",true],[["$","$1","c",{"children":[[["$","link","0",{"rel":"stylesheet","href":"/litellm-asset-prefix/_next/static/chunks/4e20891f2fd03463.css","precedence":"next","crossOrigin":"$undefined","nonce":"$undefined"}],["$","link","1",{"rel":"stylesheet","href":"/litellm-asset-prefix/_next/static/chunks/e0cb6755699177c1.css","precedence":"next","crossOrigin":"$undefined","nonce":"$undefined"}],["$","script","script-0",{"src":"/litellm-asset-prefix/_next/static/chunks/9e09de50158b3159.js","async":true,"nonce":"$undefined"}],["$","script","script-1",{"src":"/litellm-asset-prefix/_next/static/chunks/7e5fe5584502da06.js","async":true,"nonce":"$undefined"}]],["$","html",null,{"lang":"en","children":["$","body",null,{"className":"inter_5972bc34-module__OU16Qa__className","children":["$","$L2",null,{"children":["$","$L3",null,{"children":["$","$L4",null,{"parallelRouterKey":"children","error":"$undefined","errorStyles":"$undefined","errorScripts":"$undefined","template":["$","$L5",null,{}],"templateStyles":"$undefined","templateScripts":"$undefined","notFound":[[["$","title",null,{"children":"404: This page could not be found."}],["$","div",null,{"style":{"fontFamily":"system-ui,\"Segoe UI\",Roboto,Helvetica,Arial,sans-serif,\"Apple Color Emoji\",\"Segoe UI Emoji\"","height":"100vh","textAlign":"center","display":"flex","flexDirection":"column","alignItems":"center","justifyContent":"center"},"children":["$","div",null,{"children":[["$","style",null,{"dangerouslySetInnerHTML":{"__html":"body{color:#000;background:#fff;margin:0}.next-error-h1{border-right:1px solid rgba(0,0,0,.3)}@media (prefers-color-scheme:dark){body{color:#fff;background:#000}.next-error-h1{border-right:1px solid rgba(255,255,255,.3)}}"}}],["$","h1",null,{"className":"next-error-h1","style":{"display":"inline-block","margin":"0 20px 0 0","padding":"0 23px 0 0","fontSize":24,"fontWeight":500,"verticalAlign":"top","lineHeight":"49px"},"children":404}],["$","div",null,{"style":{"display":"inline-block"},"children":["$","h2",null,{"style":{"fontSize":14,"fontWeight":400,"lineHeight":"49px","margin":0},"children":"This page could not be found."}]}]]}]}]],[]],"forbidden":"$undefined","unauthorized":"$undefined"}]}]}]}]}]]}],{"children":[["$","$1","c",{"children":[null,["$","$L4",null,{"parallelRouterKey":"children","error":"$undefined","errorStyles":"$undefined","errorScripts":"$undefined","template":["$","$L5",null,{}],"templateStyles":"$undefined","templateScripts":"$undefined","notFound":"$undefined","forbidden":"$undefined","unauthorized":"$undefined"}]]}],{"children":[["$","$1","c",{"children":[["$","$L6",null,{"Component":"$7","serverProvidedParams":{"searchParams":{},"params":{},"promises":["$@8","$@9"]}}],[["$","script","script-0",{"src":"/litellm-asset-prefix/_next/static/chunks/0a6c418370a8c183.js","async":true,"nonce":"$undefined"}],["$","script","script-1",{"src":"/litellm-asset-prefix/_next/static/chunks/cf9c81fc7166f4d4.js","async":true,"nonce":"$undefined"}],["$","script","script-2",{"src":"/litellm-asset-prefix/_next/static/chunks/542a1a209eb732c6.js","async":true,"nonce":"$undefined"}],["$","script","script-3",{"src":"/litellm-asset-prefix/_next/static/chunks/00ff280cdb7d7ee5.js","async":true,"nonce":"$undefined"}],["$","script","script-4",{"src":"/litellm-asset-prefix/_next/static/chunks/a1abfc2f35c701cc.js","async":true,"nonce":"$undefined"}],["$","script","script-5",{"src":"/litellm-asset-prefix/_next/static/chunks/4980372eaa37b78b.js","async":true,"nonce":"$undefined"}],["$","script","script-6",{"src":"/litellm-asset-prefix/_next/static/chunks/119ecee91f911bc8.js","async":true,"nonce":"$undefined"}],["$","script","script-7",{"src":"/litellm-asset-prefix/_next/static/chunks/ac3cf77acb5bf234.js","async":true,"nonce":"$undefined"}],["$","script","script-8",{"src":"/litellm-asset-prefix/_next/static/chunks/fcdf7322b0aa3e2e.js","async":true,"nonce":"$undefined"}],["$","script","script-9",{"src":"/litellm-asset-prefix/_next/static/chunks/7e417dd24c8becd0.js","async":true,"nonce":"$undefined"}],["$","script","script-10",{"src":"/litellm-asset-prefix/_next/static/chunks/89aa55578de861b7.js","async":true,"nonce":"$undefined"}]],["$","$La",null,{"children":["$","$b",null,{"name":"Next.MetadataOutlet","children":"$@c"}]}]]}],{},null,false,false]},null,false,false]},null,false,false],["$","$1","h",{"children":[null,["$","$Ld",null,{"children":"$Le"}],["$","div",null,{"hidden":true,"children":["$","$Lf",null,{"children":["$","$b",null,{"name":"Next.Metadata","children":"$L10"}]}]}],["$","meta",null,{"name":"next-size-adjust","content":""}]]}],false]],"m":"$undefined","G":["$11",[]],"S":true} -8:{} -9:"$0:f:0:1:1:children:1:children:0:props:children:0:props:serverProvidedParams:params" -e:[["$","meta","0",{"charSet":"utf-8"}],["$","meta","1",{"name":"viewport","content":"width=device-width, initial-scale=1"}]] -12:I[27201,["/litellm-asset-prefix/_next/static/chunks/d96012bcfc98706a.js","/litellm-asset-prefix/_next/static/chunks/dbca964212122d58.js"],"IconMark"] -c:null -10:[["$","title","0",{"children":"LiteLLM Dashboard"}],["$","meta","1",{"name":"description","content":"LiteLLM Proxy Admin UI"}],["$","link","2",{"rel":"icon","href":"/favicon.ico?favicon.1d32c690.ico","sizes":"48x48","type":"image/x-icon"}],["$","link","3",{"rel":"icon","href":"./favicon.ico"}],["$","$L12","4",{}]] diff --git a/litellm/proxy/_experimental/out/chat/__next._full.txt b/litellm/proxy/_experimental/out/chat/__next._full.txt deleted file mode 100644 index 8ee24df074..0000000000 --- a/litellm/proxy/_experimental/out/chat/__next._full.txt +++ /dev/null @@ -1,22 +0,0 @@ -1:"$Sreact.fragment" -2:I[867271,["/litellm-asset-prefix/_next/static/chunks/9e09de50158b3159.js","/litellm-asset-prefix/_next/static/chunks/7e5fe5584502da06.js"],"default"] -3:I[71195,["/litellm-asset-prefix/_next/static/chunks/9e09de50158b3159.js","/litellm-asset-prefix/_next/static/chunks/7e5fe5584502da06.js"],"default"] -4:I[339756,["/litellm-asset-prefix/_next/static/chunks/d96012bcfc98706a.js","/litellm-asset-prefix/_next/static/chunks/dbca964212122d58.js"],"default"] -5:I[837457,["/litellm-asset-prefix/_next/static/chunks/d96012bcfc98706a.js","/litellm-asset-prefix/_next/static/chunks/dbca964212122d58.js"],"default"] -6:I[347257,["/litellm-asset-prefix/_next/static/chunks/d96012bcfc98706a.js","/litellm-asset-prefix/_next/static/chunks/dbca964212122d58.js"],"ClientPageRoot"] -7:I[321443,["/litellm-asset-prefix/_next/static/chunks/9e09de50158b3159.js","/litellm-asset-prefix/_next/static/chunks/7e5fe5584502da06.js","/litellm-asset-prefix/_next/static/chunks/0a6c418370a8c183.js","/litellm-asset-prefix/_next/static/chunks/cf9c81fc7166f4d4.js","/litellm-asset-prefix/_next/static/chunks/542a1a209eb732c6.js","/litellm-asset-prefix/_next/static/chunks/00ff280cdb7d7ee5.js","/litellm-asset-prefix/_next/static/chunks/a1abfc2f35c701cc.js","/litellm-asset-prefix/_next/static/chunks/4980372eaa37b78b.js","/litellm-asset-prefix/_next/static/chunks/119ecee91f911bc8.js","/litellm-asset-prefix/_next/static/chunks/ac3cf77acb5bf234.js","/litellm-asset-prefix/_next/static/chunks/fcdf7322b0aa3e2e.js","/litellm-asset-prefix/_next/static/chunks/7e417dd24c8becd0.js","/litellm-asset-prefix/_next/static/chunks/89aa55578de861b7.js"],"default"] -a:I[897367,["/litellm-asset-prefix/_next/static/chunks/d96012bcfc98706a.js","/litellm-asset-prefix/_next/static/chunks/dbca964212122d58.js"],"OutletBoundary"] -b:"$Sreact.suspense" -d:I[897367,["/litellm-asset-prefix/_next/static/chunks/d96012bcfc98706a.js","/litellm-asset-prefix/_next/static/chunks/dbca964212122d58.js"],"ViewportBoundary"] -f:I[897367,["/litellm-asset-prefix/_next/static/chunks/d96012bcfc98706a.js","/litellm-asset-prefix/_next/static/chunks/dbca964212122d58.js"],"MetadataBoundary"] -11:I[168027,[],"default"] -:HL["/litellm-asset-prefix/_next/static/chunks/4e20891f2fd03463.css","style"] -:HL["/litellm-asset-prefix/_next/static/chunks/e0cb6755699177c1.css","style"] -:HL["/litellm-asset-prefix/_next/static/media/83afe278b6a6bb3c-s.p.3a6ba036.woff2","font",{"crossOrigin":"","type":"font/woff2"}] -0:{"P":null,"b":"zxkD4-EPlgfKHDTw8O869","c":["","chat"],"q":"","i":false,"f":[[["",{"children":["chat",{"children":["__PAGE__",{}]}]},"$undefined","$undefined",true],[["$","$1","c",{"children":[[["$","link","0",{"rel":"stylesheet","href":"/litellm-asset-prefix/_next/static/chunks/4e20891f2fd03463.css","precedence":"next","crossOrigin":"$undefined","nonce":"$undefined"}],["$","link","1",{"rel":"stylesheet","href":"/litellm-asset-prefix/_next/static/chunks/e0cb6755699177c1.css","precedence":"next","crossOrigin":"$undefined","nonce":"$undefined"}],["$","script","script-0",{"src":"/litellm-asset-prefix/_next/static/chunks/9e09de50158b3159.js","async":true,"nonce":"$undefined"}],["$","script","script-1",{"src":"/litellm-asset-prefix/_next/static/chunks/7e5fe5584502da06.js","async":true,"nonce":"$undefined"}]],["$","html",null,{"lang":"en","children":["$","body",null,{"className":"inter_5972bc34-module__OU16Qa__className","children":["$","$L2",null,{"children":["$","$L3",null,{"children":["$","$L4",null,{"parallelRouterKey":"children","error":"$undefined","errorStyles":"$undefined","errorScripts":"$undefined","template":["$","$L5",null,{}],"templateStyles":"$undefined","templateScripts":"$undefined","notFound":[[["$","title",null,{"children":"404: This page could not be found."}],["$","div",null,{"style":{"fontFamily":"system-ui,\"Segoe UI\",Roboto,Helvetica,Arial,sans-serif,\"Apple Color Emoji\",\"Segoe UI Emoji\"","height":"100vh","textAlign":"center","display":"flex","flexDirection":"column","alignItems":"center","justifyContent":"center"},"children":["$","div",null,{"children":[["$","style",null,{"dangerouslySetInnerHTML":{"__html":"body{color:#000;background:#fff;margin:0}.next-error-h1{border-right:1px solid rgba(0,0,0,.3)}@media (prefers-color-scheme:dark){body{color:#fff;background:#000}.next-error-h1{border-right:1px solid rgba(255,255,255,.3)}}"}}],["$","h1",null,{"className":"next-error-h1","style":{"display":"inline-block","margin":"0 20px 0 0","padding":"0 23px 0 0","fontSize":24,"fontWeight":500,"verticalAlign":"top","lineHeight":"49px"},"children":404}],["$","div",null,{"style":{"display":"inline-block"},"children":["$","h2",null,{"style":{"fontSize":14,"fontWeight":400,"lineHeight":"49px","margin":0},"children":"This page could not be found."}]}]]}]}]],[]],"forbidden":"$undefined","unauthorized":"$undefined"}]}]}]}]}]]}],{"children":[["$","$1","c",{"children":[null,["$","$L4",null,{"parallelRouterKey":"children","error":"$undefined","errorStyles":"$undefined","errorScripts":"$undefined","template":["$","$L5",null,{}],"templateStyles":"$undefined","templateScripts":"$undefined","notFound":"$undefined","forbidden":"$undefined","unauthorized":"$undefined"}]]}],{"children":[["$","$1","c",{"children":[["$","$L6",null,{"Component":"$7","serverProvidedParams":{"searchParams":{},"params":{},"promises":["$@8","$@9"]}}],[["$","script","script-0",{"src":"/litellm-asset-prefix/_next/static/chunks/0a6c418370a8c183.js","async":true,"nonce":"$undefined"}],["$","script","script-1",{"src":"/litellm-asset-prefix/_next/static/chunks/cf9c81fc7166f4d4.js","async":true,"nonce":"$undefined"}],["$","script","script-2",{"src":"/litellm-asset-prefix/_next/static/chunks/542a1a209eb732c6.js","async":true,"nonce":"$undefined"}],["$","script","script-3",{"src":"/litellm-asset-prefix/_next/static/chunks/00ff280cdb7d7ee5.js","async":true,"nonce":"$undefined"}],["$","script","script-4",{"src":"/litellm-asset-prefix/_next/static/chunks/a1abfc2f35c701cc.js","async":true,"nonce":"$undefined"}],["$","script","script-5",{"src":"/litellm-asset-prefix/_next/static/chunks/4980372eaa37b78b.js","async":true,"nonce":"$undefined"}],["$","script","script-6",{"src":"/litellm-asset-prefix/_next/static/chunks/119ecee91f911bc8.js","async":true,"nonce":"$undefined"}],["$","script","script-7",{"src":"/litellm-asset-prefix/_next/static/chunks/ac3cf77acb5bf234.js","async":true,"nonce":"$undefined"}],["$","script","script-8",{"src":"/litellm-asset-prefix/_next/static/chunks/fcdf7322b0aa3e2e.js","async":true,"nonce":"$undefined"}],["$","script","script-9",{"src":"/litellm-asset-prefix/_next/static/chunks/7e417dd24c8becd0.js","async":true,"nonce":"$undefined"}],["$","script","script-10",{"src":"/litellm-asset-prefix/_next/static/chunks/89aa55578de861b7.js","async":true,"nonce":"$undefined"}]],["$","$La",null,{"children":["$","$b",null,{"name":"Next.MetadataOutlet","children":"$@c"}]}]]}],{},null,false,false]},null,false,false]},null,false,false],["$","$1","h",{"children":[null,["$","$Ld",null,{"children":"$Le"}],["$","div",null,{"hidden":true,"children":["$","$Lf",null,{"children":["$","$b",null,{"name":"Next.Metadata","children":"$L10"}]}]}],["$","meta",null,{"name":"next-size-adjust","content":""}]]}],false]],"m":"$undefined","G":["$11",[]],"S":true} -8:{} -9:"$0:f:0:1:1:children:1:children:0:props:children:0:props:serverProvidedParams:params" -e:[["$","meta","0",{"charSet":"utf-8"}],["$","meta","1",{"name":"viewport","content":"width=device-width, initial-scale=1"}]] -12:I[27201,["/litellm-asset-prefix/_next/static/chunks/d96012bcfc98706a.js","/litellm-asset-prefix/_next/static/chunks/dbca964212122d58.js"],"IconMark"] -c:null -10:[["$","title","0",{"children":"LiteLLM Dashboard"}],["$","meta","1",{"name":"description","content":"LiteLLM Proxy Admin UI"}],["$","link","2",{"rel":"icon","href":"/favicon.ico?favicon.1d32c690.ico","sizes":"48x48","type":"image/x-icon"}],["$","link","3",{"rel":"icon","href":"./favicon.ico"}],["$","$L12","4",{}]] diff --git a/litellm/proxy/_experimental/out/chat/__next._head.txt b/litellm/proxy/_experimental/out/chat/__next._head.txt deleted file mode 100644 index 635c8e398b..0000000000 --- a/litellm/proxy/_experimental/out/chat/__next._head.txt +++ /dev/null @@ -1,6 +0,0 @@ -1:"$Sreact.fragment" -2:I[897367,["/litellm-asset-prefix/_next/static/chunks/d96012bcfc98706a.js","/litellm-asset-prefix/_next/static/chunks/dbca964212122d58.js"],"ViewportBoundary"] -3:I[897367,["/litellm-asset-prefix/_next/static/chunks/d96012bcfc98706a.js","/litellm-asset-prefix/_next/static/chunks/dbca964212122d58.js"],"MetadataBoundary"] -4:"$Sreact.suspense" -5:I[27201,["/litellm-asset-prefix/_next/static/chunks/d96012bcfc98706a.js","/litellm-asset-prefix/_next/static/chunks/dbca964212122d58.js"],"IconMark"] -0:{"buildId":"zxkD4-EPlgfKHDTw8O869","rsc":["$","$1","h",{"children":[null,["$","$L2",null,{"children":[["$","meta","0",{"charSet":"utf-8"}],["$","meta","1",{"name":"viewport","content":"width=device-width, initial-scale=1"}]]}],["$","div",null,{"hidden":true,"children":["$","$L3",null,{"children":["$","$4",null,{"name":"Next.Metadata","children":[["$","title","0",{"children":"LiteLLM Dashboard"}],["$","meta","1",{"name":"description","content":"LiteLLM Proxy Admin UI"}],["$","link","2",{"rel":"icon","href":"/favicon.ico?favicon.1d32c690.ico","sizes":"48x48","type":"image/x-icon"}],["$","link","3",{"rel":"icon","href":"./favicon.ico"}],["$","$L5","4",{}]]}]}]}],["$","meta",null,{"name":"next-size-adjust","content":""}]]}],"loading":null,"isPartial":false} diff --git a/litellm/proxy/_experimental/out/chat/__next._index.txt b/litellm/proxy/_experimental/out/chat/__next._index.txt deleted file mode 100644 index d026a48036..0000000000 --- a/litellm/proxy/_experimental/out/chat/__next._index.txt +++ /dev/null @@ -1,8 +0,0 @@ -1:"$Sreact.fragment" -2:I[867271,["/litellm-asset-prefix/_next/static/chunks/9e09de50158b3159.js","/litellm-asset-prefix/_next/static/chunks/7e5fe5584502da06.js"],"default"] -3:I[71195,["/litellm-asset-prefix/_next/static/chunks/9e09de50158b3159.js","/litellm-asset-prefix/_next/static/chunks/7e5fe5584502da06.js"],"default"] -4:I[339756,["/litellm-asset-prefix/_next/static/chunks/d96012bcfc98706a.js","/litellm-asset-prefix/_next/static/chunks/dbca964212122d58.js"],"default"] -5:I[837457,["/litellm-asset-prefix/_next/static/chunks/d96012bcfc98706a.js","/litellm-asset-prefix/_next/static/chunks/dbca964212122d58.js"],"default"] -:HL["/litellm-asset-prefix/_next/static/chunks/4e20891f2fd03463.css","style"] -:HL["/litellm-asset-prefix/_next/static/chunks/e0cb6755699177c1.css","style"] -0:{"buildId":"zxkD4-EPlgfKHDTw8O869","rsc":["$","$1","c",{"children":[[["$","link","0",{"rel":"stylesheet","href":"/litellm-asset-prefix/_next/static/chunks/4e20891f2fd03463.css","precedence":"next"}],["$","link","1",{"rel":"stylesheet","href":"/litellm-asset-prefix/_next/static/chunks/e0cb6755699177c1.css","precedence":"next"}],["$","script","script-0",{"src":"/litellm-asset-prefix/_next/static/chunks/9e09de50158b3159.js","async":true}],["$","script","script-1",{"src":"/litellm-asset-prefix/_next/static/chunks/7e5fe5584502da06.js","async":true}]],["$","html",null,{"lang":"en","children":["$","body",null,{"className":"inter_5972bc34-module__OU16Qa__className","children":["$","$L2",null,{"children":["$","$L3",null,{"children":["$","$L4",null,{"parallelRouterKey":"children","template":["$","$L5",null,{}],"notFound":[[["$","title",null,{"children":"404: This page could not be found."}],["$","div",null,{"style":{"fontFamily":"system-ui,\"Segoe UI\",Roboto,Helvetica,Arial,sans-serif,\"Apple Color Emoji\",\"Segoe UI Emoji\"","height":"100vh","textAlign":"center","display":"flex","flexDirection":"column","alignItems":"center","justifyContent":"center"},"children":["$","div",null,{"children":[["$","style",null,{"dangerouslySetInnerHTML":{"__html":"body{color:#000;background:#fff;margin:0}.next-error-h1{border-right:1px solid rgba(0,0,0,.3)}@media (prefers-color-scheme:dark){body{color:#fff;background:#000}.next-error-h1{border-right:1px solid rgba(255,255,255,.3)}}"}}],["$","h1",null,{"className":"next-error-h1","style":{"display":"inline-block","margin":"0 20px 0 0","padding":"0 23px 0 0","fontSize":24,"fontWeight":500,"verticalAlign":"top","lineHeight":"49px"},"children":404}],["$","div",null,{"style":{"display":"inline-block"},"children":["$","h2",null,{"style":{"fontSize":14,"fontWeight":400,"lineHeight":"49px","margin":0},"children":"This page could not be found."}]}]]}]}]],[]]}]}]}]}]}]]}],"loading":null,"isPartial":false} diff --git a/litellm/proxy/_experimental/out/chat/__next._tree.txt b/litellm/proxy/_experimental/out/chat/__next._tree.txt deleted file mode 100644 index afb8782562..0000000000 --- a/litellm/proxy/_experimental/out/chat/__next._tree.txt +++ /dev/null @@ -1,4 +0,0 @@ -:HL["/litellm-asset-prefix/_next/static/chunks/4e20891f2fd03463.css","style"] -:HL["/litellm-asset-prefix/_next/static/chunks/e0cb6755699177c1.css","style"] -:HL["/litellm-asset-prefix/_next/static/media/83afe278b6a6bb3c-s.p.3a6ba036.woff2","font",{"crossOrigin":"","type":"font/woff2"}] -0:{"buildId":"zxkD4-EPlgfKHDTw8O869","tree":{"name":"","paramType":null,"paramKey":"","hasRuntimePrefetch":false,"slots":{"children":{"name":"chat","paramType":null,"paramKey":"chat","hasRuntimePrefetch":false,"slots":{"children":{"name":"__PAGE__","paramType":null,"paramKey":"__PAGE__","hasRuntimePrefetch":false,"slots":null,"isRootLayout":false}},"isRootLayout":false}},"isRootLayout":true},"staleTime":300} diff --git a/litellm/proxy/_experimental/out/chat/__next.chat.__PAGE__.txt b/litellm/proxy/_experimental/out/chat/__next.chat.__PAGE__.txt deleted file mode 100644 index 5d19f68046..0000000000 --- a/litellm/proxy/_experimental/out/chat/__next.chat.__PAGE__.txt +++ /dev/null @@ -1,9 +0,0 @@ -1:"$Sreact.fragment" -2:I[347257,["/litellm-asset-prefix/_next/static/chunks/d96012bcfc98706a.js","/litellm-asset-prefix/_next/static/chunks/dbca964212122d58.js"],"ClientPageRoot"] -3:I[321443,["/litellm-asset-prefix/_next/static/chunks/9e09de50158b3159.js","/litellm-asset-prefix/_next/static/chunks/7e5fe5584502da06.js","/litellm-asset-prefix/_next/static/chunks/0a6c418370a8c183.js","/litellm-asset-prefix/_next/static/chunks/cf9c81fc7166f4d4.js","/litellm-asset-prefix/_next/static/chunks/542a1a209eb732c6.js","/litellm-asset-prefix/_next/static/chunks/00ff280cdb7d7ee5.js","/litellm-asset-prefix/_next/static/chunks/a1abfc2f35c701cc.js","/litellm-asset-prefix/_next/static/chunks/4980372eaa37b78b.js","/litellm-asset-prefix/_next/static/chunks/119ecee91f911bc8.js","/litellm-asset-prefix/_next/static/chunks/ac3cf77acb5bf234.js","/litellm-asset-prefix/_next/static/chunks/fcdf7322b0aa3e2e.js","/litellm-asset-prefix/_next/static/chunks/7e417dd24c8becd0.js","/litellm-asset-prefix/_next/static/chunks/89aa55578de861b7.js"],"default"] -6:I[897367,["/litellm-asset-prefix/_next/static/chunks/d96012bcfc98706a.js","/litellm-asset-prefix/_next/static/chunks/dbca964212122d58.js"],"OutletBoundary"] -7:"$Sreact.suspense" -0:{"buildId":"zxkD4-EPlgfKHDTw8O869","rsc":["$","$1","c",{"children":[["$","$L2",null,{"Component":"$3","serverProvidedParams":{"searchParams":{},"params":{},"promises":["$@4","$@5"]}}],[["$","script","script-0",{"src":"/litellm-asset-prefix/_next/static/chunks/0a6c418370a8c183.js","async":true}],["$","script","script-1",{"src":"/litellm-asset-prefix/_next/static/chunks/cf9c81fc7166f4d4.js","async":true}],["$","script","script-2",{"src":"/litellm-asset-prefix/_next/static/chunks/542a1a209eb732c6.js","async":true}],["$","script","script-3",{"src":"/litellm-asset-prefix/_next/static/chunks/00ff280cdb7d7ee5.js","async":true}],["$","script","script-4",{"src":"/litellm-asset-prefix/_next/static/chunks/a1abfc2f35c701cc.js","async":true}],["$","script","script-5",{"src":"/litellm-asset-prefix/_next/static/chunks/4980372eaa37b78b.js","async":true}],["$","script","script-6",{"src":"/litellm-asset-prefix/_next/static/chunks/119ecee91f911bc8.js","async":true}],["$","script","script-7",{"src":"/litellm-asset-prefix/_next/static/chunks/ac3cf77acb5bf234.js","async":true}],["$","script","script-8",{"src":"/litellm-asset-prefix/_next/static/chunks/fcdf7322b0aa3e2e.js","async":true}],["$","script","script-9",{"src":"/litellm-asset-prefix/_next/static/chunks/7e417dd24c8becd0.js","async":true}],["$","script","script-10",{"src":"/litellm-asset-prefix/_next/static/chunks/89aa55578de861b7.js","async":true}]],["$","$L6",null,{"children":["$","$7",null,{"name":"Next.MetadataOutlet","children":"$@8"}]}]]}],"loading":null,"isPartial":false} -4:{} -5:"$0:rsc:props:children:0:props:serverProvidedParams:params" -8:null diff --git a/litellm/proxy/_experimental/out/chat/__next.chat.txt b/litellm/proxy/_experimental/out/chat/__next.chat.txt deleted file mode 100644 index 12f813af35..0000000000 --- a/litellm/proxy/_experimental/out/chat/__next.chat.txt +++ /dev/null @@ -1,4 +0,0 @@ -1:"$Sreact.fragment" -2:I[339756,["/litellm-asset-prefix/_next/static/chunks/d96012bcfc98706a.js","/litellm-asset-prefix/_next/static/chunks/dbca964212122d58.js"],"default"] -3:I[837457,["/litellm-asset-prefix/_next/static/chunks/d96012bcfc98706a.js","/litellm-asset-prefix/_next/static/chunks/dbca964212122d58.js"],"default"] -0:{"buildId":"zxkD4-EPlgfKHDTw8O869","rsc":["$","$1","c",{"children":[null,["$","$L2",null,{"parallelRouterKey":"children","template":["$","$L3",null,{}]}]]}],"loading":null,"isPartial":false} From 4cecfec9f9637a227a495a5519842e3b5e790b36 Mon Sep 17 00:00:00 2001 From: Sameer Kankute Date: Thu, 30 Apr 2026 01:04:14 +0530 Subject: [PATCH 53/57] feat(proxy): LiteLLM headers on Google native generateContent routes (#25500) * feat(proxy): return LiteLLM headers on Google native generateContent routes Wire build_litellm_proxy_success_headers_from_llm_response for :generateContent and :streamGenerateContent so x-litellm-*, rate limit, and provider headers match the OpenAI-style proxy path. Add unit test. Annotate httpx.HTTPStatusError branch so pyright accepts .response after optional exception transform. Remove unused variable in streaming tracer test (Ruff F841). Made-with: Cursor * fix(proxy): prefill Google GenAI stream _hidden_params for proxy headers - Pass model_id, api_base, and process_response_headers output into streaming iterators so streamGenerateContent gets the same x-litellm-* headers as non-streaming paths. - Drop request_data deployment mutation from build_litellm_proxy_success_headers_from_llm_response. - Avoid logging raw request key names in oversized debug payload (code scanning). - Extend tests for streaming iterator shape, metadata fallback, and helper. Made-with: Cursor * Update litellm/proxy/common_request_processing.py Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> * remove unused key count * Fix greptile review * Update litellm/proxy/common_request_processing.py Co-authored-by: Mateo Wang <277851410+mateo-berri@users.noreply.github.com> --------- Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> Co-authored-by: Mateo Wang <277851410+mateo-berri@users.noreply.github.com> --- litellm/google_genai/streaming_iterator.py | 8 +- litellm/llms/custom_httpx/llm_http_handler.py | 36 +++++ litellm/proxy/common_request_processing.py | 72 ++++++++- litellm/proxy/google_endpoints/endpoints.py | 29 +++- .../custom_httpx/test_llm_http_handler.py | 32 +++- .../proxy/test_common_request_processing.py | 142 +++++++++++++++++- 6 files changed, 303 insertions(+), 16 deletions(-) diff --git a/litellm/google_genai/streaming_iterator.py b/litellm/google_genai/streaming_iterator.py index 8cb2ee0937..3e97b48077 100644 --- a/litellm/google_genai/streaming_iterator.py +++ b/litellm/google_genai/streaming_iterator.py @@ -1,6 +1,6 @@ import asyncio from datetime import datetime -from typing import TYPE_CHECKING, Any, List, Optional +from typing import TYPE_CHECKING, Any, Dict, List, Optional from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.proxy.pass_through_endpoints.success_handler import ( @@ -29,12 +29,14 @@ class BaseGoogleGenAIGenerateContentStreamingIterator: litellm_logging_obj: LiteLLMLoggingObj, request_body: dict, model: str, + hidden_params: Optional[Dict[str, Any]] = None, ): self.litellm_logging_obj = litellm_logging_obj self.request_body = request_body self.start_time = datetime.now() self.collected_chunks: List[bytes] = [] self.model = model + self._hidden_params: Dict[str, Any] = hidden_params or {} async def _handle_async_streaming_logging( self, @@ -76,11 +78,13 @@ class GoogleGenAIGenerateContentStreamingIterator( litellm_metadata: dict, custom_llm_provider: str, request_body: Optional[dict] = None, + hidden_params: Optional[Dict[str, Any]] = None, ): super().__init__( litellm_logging_obj=logging_obj, request_body=request_body or {}, model=model, + hidden_params=hidden_params, ) self.response = response self.model = model @@ -130,11 +134,13 @@ class AsyncGoogleGenAIGenerateContentStreamingIterator( litellm_metadata: dict, custom_llm_provider: str, request_body: Optional[dict] = None, + hidden_params: Optional[Dict[str, Any]] = None, ): super().__init__( litellm_logging_obj=logging_obj, request_body=request_body or {}, model=model, + hidden_params=hidden_params, ) self.response = response self.model = model diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index b9ada079f6..6b9a2c6a5d 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -155,6 +155,30 @@ else: LiteLLMLoggingObj = Any +def _google_genai_streaming_hidden_params( + *, + api_base: str, + litellm_params: GenericLiteLLMParams, + logging_obj: LiteLLMLoggingObj, + response_headers: httpx.Headers, +) -> Dict[str, Any]: + """Pre-stream metadata for proxy response headers (mirrors CustomStreamWrapper._hidden_params).""" + from litellm.litellm_core_utils.core_helpers import process_response_headers + + _model_info: Dict[str, Any] = dict( + getattr(litellm_params, "model_info", None) or {} + ) + _raw_id = _model_info.get("id") or logging_obj.get_router_model_id() or "" + _model_id = _raw_id if isinstance(_raw_id, str) else str(_raw_id) + return { + "model_id": _model_id, + "api_base": api_base, + "cache_key": "", + "response_cost": "", + "additional_headers": process_response_headers(response_headers), + } + + class BaseLLMHTTPHandler: async def _make_common_async_call( self, @@ -10425,6 +10449,12 @@ class BaseLLMHTTPHandler: litellm_metadata=litellm_metadata or {}, custom_llm_provider=custom_llm_provider, request_body=data, + hidden_params=_google_genai_streaming_hidden_params( + api_base=api_base, + litellm_params=litellm_params, + logging_obj=logging_obj, + response_headers=response.headers, + ), ) else: response = sync_httpx_client.post( @@ -10534,6 +10564,12 @@ class BaseLLMHTTPHandler: litellm_metadata=litellm_metadata or {}, custom_llm_provider=custom_llm_provider, request_body=data, + hidden_params=_google_genai_streaming_hidden_params( + api_base=api_base, + litellm_params=litellm_params, + logging_obj=logging_obj, + response_headers=response.headers, + ), ) else: response = await async_httpx_client.post( diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index c8fab4be4a..76c52f83ee 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -619,6 +619,67 @@ class ProxyBaseLLMRequestProcessing: verbose_proxy_logger.error(f"Error setting custom headers: {e}") return {} + @staticmethod + async def build_litellm_proxy_success_headers_from_llm_response( + *, + response: Any, + request_data: dict, + request: Request, + user_api_key_dict: UserAPIKeyAuth, + logging_obj: LiteLLMLoggingObj, + version: Optional[str], + proxy_logging_obj: ProxyLogging, + ) -> Dict[str, str]: + """ + Build LiteLLM proxy response headers for routes that call the LLM directly + (e.g. Google native :generateContent) instead of base_process_llm_request. + """ + if isinstance(response, dict): + hidden_params = response.get("_hidden_params") or {} + else: + hidden_params = getattr(response, "_hidden_params", None) or {} + if not isinstance(hidden_params, dict): + hidden_params = {} + + model_id = ProxyBaseLLMRequestProcessing._get_model_id_from_response( + hidden_params, request_data + ) + + cache_key = hidden_params.get("cache_key", None) or "" + api_base = hidden_params.get("api_base", None) or "" + response_cost = hidden_params.get("response_cost", None) or "" + fastest_response_batch_completion = hidden_params.get( + "fastest_response_batch_completion", None + ) + additional_headers = hidden_params.get("additional_headers", {}) or {} + + custom_headers = ProxyBaseLLMRequestProcessing.get_custom_headers( + user_api_key_dict=user_api_key_dict, + call_id=logging_obj.litellm_call_id, + model_id=model_id, + cache_key=cache_key, + api_base=api_base, + version=version, + response_cost=response_cost, + model_region=getattr(user_api_key_dict, "allowed_model_region", ""), + fastest_response_batch_completion=fastest_response_batch_completion, + request_data=request_data, + hidden_params=hidden_params, + litellm_logging_obj=logging_obj, + **additional_headers, + ) + + callback_headers = await proxy_logging_obj.post_call_response_headers_hook( + data=request_data, + user_api_key_dict=user_api_key_dict, + response=response, + request_headers=dict(request.headers), + ) + if callback_headers: + custom_headers.update(callback_headers) + + return custom_headers + async def common_processing_pre_call_logic( self, request: Request, @@ -875,7 +936,7 @@ class ProxyBaseLLMRequestProcessing: else: verbose_proxy_logger.debug( "Request received by LiteLLM:\n%s", - json.dumps(self.data, indent=4, default=str), + _payload_str, ) async def base_process_llm_request( # noqa: PLR0915 @@ -1511,9 +1572,7 @@ class ProxyBaseLLMRequestProcessing: _response = assembled_response try: from litellm.proxy.proxy_server import llm_router as _global_llm_router - from litellm.proxy.utils import ( - _check_and_merge_model_level_guardrails, - ) + from litellm.proxy.utils import _check_and_merge_model_level_guardrails guardrail_data = _check_and_merge_model_level_guardrails( data=captured_data, llm_router=_global_llm_router @@ -1690,11 +1749,12 @@ class ProxyBaseLLMRequestProcessing: elif isinstance(e, httpx.HTTPStatusError): # Handle httpx.HTTPStatusError - extract actual error from response # This matches the original behavior before the refactor in commit 511d435f6f - error_body = await e.response.aread() + http_status_error: httpx.HTTPStatusError = e + error_body = await http_status_error.response.aread() error_text = error_body.decode("utf-8") raise HTTPException( - status_code=e.response.status_code, + status_code=http_status_error.response.status_code, detail={"error": error_text}, ) error_msg = f"{str(e)}" diff --git a/litellm/proxy/google_endpoints/endpoints.py b/litellm/proxy/google_endpoints/endpoints.py index 9768d93e92..6ada8f5878 100644 --- a/litellm/proxy/google_endpoints/endpoints.py +++ b/litellm/proxy/google_endpoints/endpoints.py @@ -35,6 +35,7 @@ async def google_generate_content( general_settings, llm_router, proxy_config, + proxy_logging_obj, version, ) @@ -73,6 +74,16 @@ async def google_generate_content( if llm_router is None: raise HTTPException(status_code=500, detail="Router not initialized") response = await llm_router.agenerate_content(**data) + success_headers = await ProxyBaseLLMRequestProcessing.build_litellm_proxy_success_headers_from_llm_response( + response=response, + request_data=data, + request=request, + user_api_key_dict=user_api_key_dict, + logging_obj=logging_obj, + version=version, + proxy_logging_obj=proxy_logging_obj, + ) + fastapi_response.headers.update(success_headers) return response @@ -95,6 +106,7 @@ async def google_stream_generate_content( general_settings, llm_router, proxy_config, + proxy_logging_obj, version, ) @@ -137,9 +149,24 @@ async def google_stream_generate_content( raise HTTPException(status_code=500, detail="Router not initialized") response = await llm_router.agenerate_content_stream(**data) + success_headers = await ProxyBaseLLMRequestProcessing.build_litellm_proxy_success_headers_from_llm_response( + response=response, + request_data=data, + request=request, + user_api_key_dict=user_api_key_dict, + logging_obj=logging_obj, + version=version, + proxy_logging_obj=proxy_logging_obj, + ) + # Check if response is an async iterator (streaming response) if response is not None and hasattr(response, "__aiter__"): - return StreamingResponse(content=response, media_type="text/event-stream") + return StreamingResponse( + content=response, + media_type="text/event-stream", + headers=success_headers, + ) + fastapi_response.headers.update(success_headers) return response diff --git a/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py b/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py index 6924eb8d3d..752b5ff090 100644 --- a/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py +++ b/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py @@ -2,12 +2,16 @@ import os import sys from unittest.mock import AsyncMock, Mock, patch +import httpx import pytest sys.path.insert( 0, os.path.abspath("../../../..") ) # Adds the parent directory to the system path -from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler +from litellm.llms.custom_httpx.llm_http_handler import ( + BaseLLMHTTPHandler, + _google_genai_streaming_hidden_params, +) from litellm.types.router import GenericLiteLLMParams @@ -320,3 +324,29 @@ async def test_async_anthropic_messages_handler_header_priority(): assert captured_headers["X-Forwarded-Only"] == "keep" assert captured_headers["X-Extra-Only"] == "also-keep" assert captured_headers["X-Provider-Only"] == "keep-this-too" + + +def test_google_genai_streaming_hidden_params_model_info_and_router_fallback(): + logging_obj = Mock() + logging_obj.get_router_model_id = Mock(return_value="router-model-id") + + from_model_info = _google_genai_streaming_hidden_params( + api_base="https://generativelanguage.googleapis.com/v1beta", + litellm_params=GenericLiteLLMParams(model_info={"id": "info-id"}), + logging_obj=logging_obj, + response_headers=httpx.Headers({"x-ratelimit-remaining": "10"}), + ) + assert from_model_info["model_id"] == "info-id" + assert ( + from_model_info["api_base"] + == "https://generativelanguage.googleapis.com/v1beta" + ) + assert isinstance(from_model_info["additional_headers"], dict) + + from_router = _google_genai_streaming_hidden_params( + api_base="https://x", + litellm_params=GenericLiteLLMParams(), + logging_obj=logging_obj, + response_headers=httpx.Headers({}), + ) + assert from_router["model_id"] == "router-model-id" diff --git a/tests/test_litellm/proxy/test_common_request_processing.py b/tests/test_litellm/proxy/test_common_request_processing.py index a635c16e98..d4b4617730 100644 --- a/tests/test_litellm/proxy/test_common_request_processing.py +++ b/tests/test_litellm/proxy/test_common_request_processing.py @@ -218,6 +218,141 @@ class TestProxyBaseLLMRequestProcessing: headers_with_invalid ) + @pytest.mark.asyncio + async def test_build_litellm_proxy_success_headers_from_llm_response(self): + """ + Google native :generateContent uses this helper instead of base_process_llm_request; + ensure x-litellm-* headers and callback hooks merge like the main proxy path. + """ + mock_request = MagicMock(spec=Request) + mock_request.headers = {} + + class _FakeGenaiResponse: + _hidden_params = { + "model_id": "deployment-model-id", + "cache_key": "ck-test", + "api_base": "https://generativelanguage.googleapis.com/v1beta", + "response_cost": 0.001, + "additional_headers": {"llm_provider-ratelimit-requests": "1000"}, + } + + logging_obj = MagicMock() + logging_obj.litellm_call_id = "call-id-test" + + mock_user = MagicMock() + mock_user.tpm_limit = None + mock_user.rpm_limit = None + mock_user.max_budget = None + mock_user.spend = 0.0 + mock_user.allowed_model_region = None + + proxy_logging_obj = MagicMock(spec=ProxyLogging) + proxy_logging_obj.post_call_response_headers_hook = AsyncMock( + return_value={"x-ratelimit-remaining-requests": "999"} + ) + + headers = await ProxyBaseLLMRequestProcessing.build_litellm_proxy_success_headers_from_llm_response( + response=_FakeGenaiResponse(), + request_data={"model": "gemini/gemini-1.5-flash"}, + request=mock_request, + user_api_key_dict=mock_user, + logging_obj=logging_obj, + version="9.9.9", + proxy_logging_obj=proxy_logging_obj, + ) + + assert headers["x-litellm-call-id"] == "call-id-test" + assert headers["x-litellm-model-id"] == "deployment-model-id" + assert headers["x-litellm-version"] == "9.9.9" + assert headers["llm_provider-ratelimit-requests"] == "1000" + assert headers["x-ratelimit-remaining-requests"] == "999" + proxy_logging_obj.post_call_response_headers_hook.assert_awaited_once() + + @pytest.mark.asyncio + async def test_build_litellm_proxy_success_headers_streaming_style_iterator(self): + """AsyncGoogleGenAIGenerateContentStreamingIterator sets _hidden_params at init; headers must propagate.""" + + class _FakeStreamLike: + def __aiter__(self): + return self + + async def __anext__(self): + raise StopAsyncIteration + + _hidden_params = { + "model_id": "stream-model-id", + "api_base": "https://generativelanguage.googleapis.com/v1beta", + "cache_key": "", + "response_cost": "", + "additional_headers": {"llm_provider-x": "y"}, + } + + mock_request = MagicMock(spec=Request) + mock_request.headers = {} + logging_obj = MagicMock() + logging_obj.litellm_call_id = "cid-stream" + mock_user = MagicMock() + mock_user.tpm_limit = None + mock_user.rpm_limit = None + mock_user.max_budget = None + mock_user.spend = 0.0 + mock_user.allowed_model_region = None + proxy_logging_obj = MagicMock(spec=ProxyLogging) + proxy_logging_obj.post_call_response_headers_hook = AsyncMock(return_value={}) + + headers = await ProxyBaseLLMRequestProcessing.build_litellm_proxy_success_headers_from_llm_response( + response=_FakeStreamLike(), + request_data={"model": "gemini/gemini-2.0-flash"}, + request=mock_request, + user_api_key_dict=mock_user, + logging_obj=logging_obj, + version="1.0.0", + proxy_logging_obj=proxy_logging_obj, + ) + + assert headers["x-litellm-model-id"] == "stream-model-id" + assert headers["x-litellm-model-api-base"] == ( + "https://generativelanguage.googleapis.com/v1beta" + ) + assert headers["llm_provider-x"] == "y" + + @pytest.mark.asyncio + async def test_build_litellm_proxy_success_headers_no_hidden_params_metadata_fallback( + self, + ): + """When response has no _hidden_params, model_id can still come from litellm_metadata.""" + + class _BareResponse: + pass + + mock_request = MagicMock(spec=Request) + mock_request.headers = {} + logging_obj = MagicMock() + logging_obj.litellm_call_id = "cid-meta" + mock_user = MagicMock() + mock_user.tpm_limit = None + mock_user.rpm_limit = None + mock_user.max_budget = None + mock_user.spend = 0.0 + mock_user.allowed_model_region = None + proxy_logging_obj = MagicMock(spec=ProxyLogging) + proxy_logging_obj.post_call_response_headers_hook = AsyncMock(return_value={}) + + headers = await ProxyBaseLLMRequestProcessing.build_litellm_proxy_success_headers_from_llm_response( + response=_BareResponse(), + request_data={ + "model": "gemini/gemini-1.5-flash", + "litellm_metadata": {"model_info": {"id": "meta-model-id"}}, + }, + request=mock_request, + user_api_key_dict=mock_user, + logging_obj=logging_obj, + version="1.0.0", + proxy_logging_obj=proxy_logging_obj, + ) + + assert headers["x-litellm-model-id"] == "meta-model-id" + @pytest.mark.asyncio async def test_add_litellm_data_to_request_with_stream_timeout_header(self): """ @@ -1158,13 +1293,6 @@ class TestCommonRequestProcessingHelpers: assert mock_tracer.trace.call_count == 4 # Verify that each call was made with the correct operation name - expected_calls = [ - (("streaming.chunk.yield",), {}), - (("streaming.chunk.yield",), {}), - (("streaming.chunk.yield",), {}), - (("streaming.chunk.yield",), {}), - ] - actual_calls = mock_tracer.trace.call_args_list assert len(actual_calls) == 4 From a291cc60cf9324e85a003641fa79e59a48834b6c Mon Sep 17 00:00:00 2001 From: Ryan Crabbe Date: Wed, 29 Apr 2026 15:11:29 -0700 Subject: [PATCH 54/57] fix: drop sensitive locals from re-raised error messages Remove parameters that may contain credentials from the messages built inside broad except handlers. These messages can surface in HTTP error responses, so caller-supplied secrets and integration tokens shouldn't be interpolated into them. --- litellm/integrations/prompt_management_base.py | 4 ++-- .../llm_response_utils/convert_dict_to_response.py | 2 -- 2 files changed, 2 insertions(+), 4 deletions(-) diff --git a/litellm/integrations/prompt_management_base.py b/litellm/integrations/prompt_management_base.py index 71da650dc4..ab6ef9e32d 100644 --- a/litellm/integrations/prompt_management_base.py +++ b/litellm/integrations/prompt_management_base.py @@ -88,7 +88,7 @@ class PromptManagementBase(ABC): messages = compiled_prompt_client["prompt_template"] + client_messages except Exception as e: raise ValueError( - f"Error compiling prompt: {e}. Prompt id={prompt_id}, prompt_variables={prompt_variables}, client_messages={client_messages}, dynamic_callback_params={dynamic_callback_params}" + f"Error compiling prompt: {e}. Prompt id={prompt_id}, prompt_variables={prompt_variables}, client_messages={client_messages}" ) compiled_prompt_client["completed_messages"] = messages @@ -117,7 +117,7 @@ class PromptManagementBase(ABC): messages = compiled_prompt_client["prompt_template"] + client_messages except Exception as e: raise ValueError( - f"Error compiling prompt: {e}. Prompt id={prompt_id}, prompt_variables={prompt_variables}, client_messages={client_messages}, dynamic_callback_params={dynamic_callback_params}" + f"Error compiling prompt: {e}. Prompt id={prompt_id}, prompt_variables={prompt_variables}, client_messages={client_messages}" ) compiled_prompt_client["completed_messages"] = messages diff --git a/litellm/litellm_core_utils/llm_response_utils/convert_dict_to_response.py b/litellm/litellm_core_utils/llm_response_utils/convert_dict_to_response.py index 78378faa26..5fd42fe0d3 100644 --- a/litellm/litellm_core_utils/llm_response_utils/convert_dict_to_response.py +++ b/litellm/litellm_core_utils/llm_response_utils/convert_dict_to_response.py @@ -824,8 +824,6 @@ def convert_to_model_response_object( # noqa: PLR0915 stream=stream, start_time=start_time, end_time=end_time, - hidden_params=hidden_params, - _response_headers=_response_headers, convert_tool_call_to_json_mode=convert_tool_call_to_json_mode, ) raise Exception( From f3fd79bf23b94dd628cf63fbf5d99a3802956e91 Mon Sep 17 00:00:00 2001 From: Ryan Crabbe Date: Wed, 29 Apr 2026 15:16:01 -0700 Subject: [PATCH 55/57] fix: trim caller-supplied dicts from compile_prompt error message Drop prompt_variables and client_messages from the re-raised error so callers cannot leak secrets, tokens, or PII embedded in those payloads through HTTP error responses. Both sync and async variants. --- litellm/integrations/prompt_management_base.py | 8 ++------ 1 file changed, 2 insertions(+), 6 deletions(-) diff --git a/litellm/integrations/prompt_management_base.py b/litellm/integrations/prompt_management_base.py index ab6ef9e32d..9c626aea84 100644 --- a/litellm/integrations/prompt_management_base.py +++ b/litellm/integrations/prompt_management_base.py @@ -87,9 +87,7 @@ class PromptManagementBase(ABC): try: messages = compiled_prompt_client["prompt_template"] + client_messages except Exception as e: - raise ValueError( - f"Error compiling prompt: {e}. Prompt id={prompt_id}, prompt_variables={prompt_variables}, client_messages={client_messages}" - ) + raise ValueError(f"Error compiling prompt: {e}. Prompt id={prompt_id}") compiled_prompt_client["completed_messages"] = messages return compiled_prompt_client @@ -116,9 +114,7 @@ class PromptManagementBase(ABC): try: messages = compiled_prompt_client["prompt_template"] + client_messages except Exception as e: - raise ValueError( - f"Error compiling prompt: {e}. Prompt id={prompt_id}, prompt_variables={prompt_variables}, client_messages={client_messages}" - ) + raise ValueError(f"Error compiling prompt: {e}. Prompt id={prompt_id}") compiled_prompt_client["completed_messages"] = messages return compiled_prompt_client From 4a7af1ff68d144b56927cb1a421b3c5538b60aab Mon Sep 17 00:00:00 2001 From: ishaan-berri <155045088+ishaan-berri@users.noreply.github.com> Date: Wed, 29 Apr 2026 17:12:18 -0700 Subject: [PATCH 56/57] feat(proxy): durable agent workflow run tracking via /v1/workflows/runs (#26793) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * feat(schema): add workflow run tracking tables (LiteLLM_WorkflowRun, LiteLLM_WorkflowEvent, LiteLLM_WorkflowMessage) * feat(proxy): add /v1/workflows/runs endpoints for durable agent workflow tracking * feat(proxy): register workflow management router in proxy_server * docs(workflows): add README for workflow run tracking API * test(workflows): add unit tests for /v1/workflows/runs endpoints * fix(workflows): atomic event+status update via tx(), run_id 404 guard, sequence retry on collision * test(workflows): add tx mock, 404 on unknown run_id, retry-on-collision tests * fix(workflows): constrain status to Literal enum, rename total→count in list responses * add tenant isolation and bounded limits to workflow endpoints * add created_by column and index to LiteLLM_WorkflowRun * add ownership and bounded-limit tests for workflow endpoints * Fix workflow run ownership for null owners * guard prisma import in workflow_management_endpoints * sync schema.prisma copies with workflow run models * black: format workflow_management_endpoints.py --------- Co-authored-by: Cursor Agent --- .../migration.sql | 75 ++ .../litellm_proxy_extras/schema.prisma | 77 ++ .../workflow_management_endpoints.py | 492 ++++++++++++ litellm/proxy/proxy_server.py | 4 + litellm/proxy/schema.prisma | 77 ++ litellm/proxy/workflows/README.md | 150 ++++ schema.prisma | 77 ++ .../test_workflow_management_endpoints.py | 611 ++++++++++++++ ui/litellm-dashboard/src/app/page.tsx | 3 + .../src/components/leftnav.test.tsx | 2 +- .../src/components/leftnav.tsx | 35 +- .../src/components/page_metadata.ts | 2 + .../src/components/workflow_runs/index.tsx | 751 ++++++++++++++++++ 13 files changed, 2345 insertions(+), 11 deletions(-) create mode 100644 litellm-proxy-extras/litellm_proxy_extras/migrations/20260429161855_workflow_runs_tables/migration.sql create mode 100644 litellm/proxy/management_endpoints/workflow_management_endpoints.py create mode 100644 litellm/proxy/workflows/README.md create mode 100644 tests/test_litellm/proxy/management_endpoints/test_workflow_management_endpoints.py create mode 100644 ui/litellm-dashboard/src/components/workflow_runs/index.tsx diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260429161855_workflow_runs_tables/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260429161855_workflow_runs_tables/migration.sql new file mode 100644 index 0000000000..6454f26765 --- /dev/null +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260429161855_workflow_runs_tables/migration.sql @@ -0,0 +1,75 @@ +-- CreateTable +CREATE TABLE "LiteLLM_WorkflowRun" ( + "run_id" TEXT NOT NULL, + "session_id" TEXT NOT NULL, + "workflow_type" TEXT NOT NULL, + "status" TEXT NOT NULL DEFAULT 'pending', + "created_by" TEXT, + "created_at" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP, + "updated_at" TIMESTAMP(3) NOT NULL, + "input" JSONB, + "output" JSONB, + "metadata" JSONB, + + CONSTRAINT "LiteLLM_WorkflowRun_pkey" PRIMARY KEY ("run_id") +); + +-- CreateTable +CREATE TABLE "LiteLLM_WorkflowEvent" ( + "event_id" TEXT NOT NULL, + "run_id" TEXT NOT NULL, + "event_type" TEXT NOT NULL, + "step_name" TEXT NOT NULL, + "sequence_number" INTEGER NOT NULL, + "data" JSONB, + "created_at" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP, + + CONSTRAINT "LiteLLM_WorkflowEvent_pkey" PRIMARY KEY ("event_id") +); + +-- CreateTable +CREATE TABLE "LiteLLM_WorkflowMessage" ( + "message_id" TEXT NOT NULL, + "run_id" TEXT NOT NULL, + "role" TEXT NOT NULL, + "content" TEXT NOT NULL, + "sequence_number" INTEGER NOT NULL, + "session_id" TEXT, + "created_at" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP, + + CONSTRAINT "LiteLLM_WorkflowMessage_pkey" PRIMARY KEY ("message_id") +); + +-- CreateIndex +CREATE UNIQUE INDEX "LiteLLM_WorkflowRun_session_id_key" ON "LiteLLM_WorkflowRun"("session_id"); + +-- CreateIndex +CREATE INDEX "LiteLLM_WorkflowRun_workflow_type_status_idx" ON "LiteLLM_WorkflowRun"("workflow_type", "status"); + +-- CreateIndex +CREATE INDEX "LiteLLM_WorkflowRun_session_id_idx" ON "LiteLLM_WorkflowRun"("session_id"); + +-- CreateIndex +CREATE INDEX "LiteLLM_WorkflowRun_created_at_idx" ON "LiteLLM_WorkflowRun"("created_at"); + +-- CreateIndex +CREATE INDEX "LiteLLM_WorkflowRun_created_by_idx" ON "LiteLLM_WorkflowRun"("created_by"); + +-- CreateIndex +CREATE INDEX "LiteLLM_WorkflowEvent_run_id_idx" ON "LiteLLM_WorkflowEvent"("run_id"); + +-- CreateIndex +CREATE UNIQUE INDEX "LiteLLM_WorkflowEvent_run_id_sequence_number_key" ON "LiteLLM_WorkflowEvent"("run_id", "sequence_number"); + +-- CreateIndex +CREATE INDEX "LiteLLM_WorkflowMessage_run_id_idx" ON "LiteLLM_WorkflowMessage"("run_id"); + +-- CreateIndex +CREATE UNIQUE INDEX "LiteLLM_WorkflowMessage_run_id_sequence_number_key" ON "LiteLLM_WorkflowMessage"("run_id", "sequence_number"); + +-- AddForeignKey +ALTER TABLE "LiteLLM_WorkflowEvent" ADD CONSTRAINT "LiteLLM_WorkflowEvent_run_id_fkey" FOREIGN KEY ("run_id") REFERENCES "LiteLLM_WorkflowRun"("run_id") ON DELETE RESTRICT ON UPDATE CASCADE; + +-- AddForeignKey +ALTER TABLE "LiteLLM_WorkflowMessage" ADD CONSTRAINT "LiteLLM_WorkflowMessage_run_id_fkey" FOREIGN KEY ("run_id") REFERENCES "LiteLLM_WorkflowRun"("run_id") ON DELETE RESTRICT ON UPDATE CASCADE; + diff --git a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma index 8f07c5afa3..fd54ed5243 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma +++ b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma @@ -1290,3 +1290,80 @@ model LiteLLM_AdaptiveRouterSession { @@id([session_id, router_name, model_name]) @@index([last_activity_at], map: "idx_adaptive_router_session_activity") } + +// --------------------------------------------------------------------------- +// Workflow Run Tracking +// +// Generic durable state tracking for any agent or automated workflow. +// Design: three tables — run (header + materialized status), event (append-only +// source of truth for state transitions), message (conversation inbox/outbox). +// +// Usage: +// - Set `workflow_type` to identify the owning system (e.g. "shin-builder"). +// - Store domain-specific fields in `metadata` (worktree_path, pr_url, etc.). +// - `session_id` on WorkflowRun matches `x-litellm-session-id` header sent to +// the proxy — all spend logs for this run are automatically tagged. +// --------------------------------------------------------------------------- + +// One instance of work being done. `status` is a materialized cache of the +// latest event; the event log is the authoritative source of truth. +model LiteLLM_WorkflowRun { + run_id String @id @default(uuid()) + session_id String @unique @default(uuid()) + workflow_type String + status String @default("pending") + created_by String? // user_id of the key that created this run; null = created by master key + created_at DateTime @default(now()) + updated_at DateTime @updatedAt + input Json? + output Json? + metadata Json? + + events LiteLLM_WorkflowEvent[] + messages LiteLLM_WorkflowMessage[] + + @@index([workflow_type, status]) + @@index([session_id]) + @@index([created_at]) + @@index([created_by]) +} + +// Append-only log of state transitions. Never mutate rows here. +// `step_name` and `event_type` are caller-defined strings — no hardcoded enums. +// Status auto-update rules (applied by the append endpoint): +// step.started → run.status = running +// step.failed → run.status = failed +// hook.waiting → run.status = paused +// hook.received → run.status = running +model LiteLLM_WorkflowEvent { + event_id String @id @default(uuid()) + run_id String + event_type String + step_name String + sequence_number Int + data Json? + created_at DateTime @default(now()) + + run LiteLLM_WorkflowRun @relation(fields: [run_id], references: [run_id]) + + @@unique([run_id, sequence_number]) + @@index([run_id]) +} + +// Conversation inbox/outbox — full message content, separate from the durable +// event log. Spend logs truncate messages; this table stores them in full. +// `session_id` here is the Claude --resume session ID (or similar). +model LiteLLM_WorkflowMessage { + message_id String @id @default(uuid()) + run_id String + role String + content String + sequence_number Int + session_id String? + created_at DateTime @default(now()) + + run LiteLLM_WorkflowRun @relation(fields: [run_id], references: [run_id]) + + @@unique([run_id, sequence_number]) + @@index([run_id]) +} diff --git a/litellm/proxy/management_endpoints/workflow_management_endpoints.py b/litellm/proxy/management_endpoints/workflow_management_endpoints.py new file mode 100644 index 0000000000..a19af4dd48 --- /dev/null +++ b/litellm/proxy/management_endpoints/workflow_management_endpoints.py @@ -0,0 +1,492 @@ +""" +WORKFLOW RUN MANAGEMENT + +Generic durable state tracking for agents and automated workflows. + +POST /v1/workflows/runs - Create a workflow run +GET /v1/workflows/runs - List runs (filter by type, status) +GET /v1/workflows/runs/{run_id} - Get run with latest event +PATCH /v1/workflows/runs/{run_id} - Update status, metadata, output +POST /v1/workflows/runs/{run_id}/events - Append event (updates run status) +GET /v1/workflows/runs/{run_id}/events - Full event log +POST /v1/workflows/runs/{run_id}/messages - Append conversation message +GET /v1/workflows/runs/{run_id}/messages - Fetch conversation history +""" + +import json +from typing import Any, Dict, Literal, Optional + +from fastapi import APIRouter, Depends, HTTPException, Query + +try: + from prisma.errors import UniqueViolationError +except ImportError: + UniqueViolationError = None # type: ignore +from pydantic import BaseModel + +from litellm._logging import verbose_proxy_logger +from litellm.proxy._types import CommonProxyErrors, LitellmUserRoles, UserAPIKeyAuth +from litellm.proxy.auth.user_api_key_auth import user_api_key_auth + +router = APIRouter() + +_MAX_SEQUENCE_RETRIES = 5 + + +def _json(value: Any) -> str: + """Serialize a Python value for prisma-client-py Json fields (must be a string).""" + return json.dumps(value) + + +def _is_admin(user_api_key_dict: UserAPIKeyAuth) -> bool: + return user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN.value + + +def _caller_key(user_api_key_dict: UserAPIKeyAuth) -> Optional[str]: + """Return the hashed key token that identifies this caller, or None for master key.""" + return user_api_key_dict.token + + +# Status transitions driven by event_type +_EVENT_STATUS_MAP: Dict[str, str] = { + "step.started": "running", + "step.failed": "failed", + "hook.waiting": "paused", + "hook.received": "running", +} + + +# --------------------------------------------------------------------------- +# Request / Response models +# --------------------------------------------------------------------------- + + +class WorkflowRunCreateRequest(BaseModel): + workflow_type: str + input: Optional[Dict[str, Any]] = None + metadata: Optional[Dict[str, Any]] = None + + +WorkflowRunStatus = Literal["pending", "running", "paused", "completed", "failed"] + + +class WorkflowRunUpdateRequest(BaseModel): + status: Optional[WorkflowRunStatus] = None + output: Optional[Dict[str, Any]] = None + metadata: Optional[Dict[str, Any]] = None + + +class WorkflowEventCreateRequest(BaseModel): + event_type: str + step_name: str + data: Optional[Dict[str, Any]] = None + + +class WorkflowMessageCreateRequest(BaseModel): + role: str + content: str + session_id: Optional[str] = None + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + + +async def _get_next_sequence_number(prisma_client: Any, run_id: str, table: str) -> int: + """Return MAX(sequence_number) + 1 for the given run, for either events or messages.""" + if table == "events": + rows = await prisma_client.db.litellm_workflowevent.find_many( + where={"run_id": run_id}, + order={"sequence_number": "desc"}, + take=1, + ) + else: + rows = await prisma_client.db.litellm_workflowmessage.find_many( + where={"run_id": run_id}, + order={"sequence_number": "desc"}, + take=1, + ) + return (rows[0].sequence_number + 1) if rows else 0 + + +async def _require_run( + prisma_client: Any, + run_id: str, + user_api_key_dict: Optional[UserAPIKeyAuth] = None, +) -> Any: + """Return the run or raise 404. For non-admin callers, also enforce key ownership.""" + run = await prisma_client.db.litellm_workflowrun.find_unique( + where={"run_id": run_id} + ) + if run is None: + raise HTTPException(status_code=404, detail=f"Run '{run_id}' not found") + if user_api_key_dict is not None and not _is_admin(user_api_key_dict): + caller = _caller_key(user_api_key_dict) + if not caller or run.created_by != caller: + raise HTTPException(status_code=404, detail=f"Run '{run_id}' not found") + return run + + +# --------------------------------------------------------------------------- +# Endpoints +# --------------------------------------------------------------------------- + + +@router.post( + "/v1/workflows/runs", + tags=["workflow management"], + dependencies=[Depends(user_api_key_auth)], +) +async def create_workflow_run( + data: WorkflowRunCreateRequest, + user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), +): + """Create a new workflow run. Returns run_id and session_id. + + The caller's API key token is stored as created_by so that non-admin keys + can only see and modify their own runs. + """ + from litellm.proxy.proxy_server import prisma_client + + if prisma_client is None: + raise HTTPException( + status_code=500, detail=CommonProxyErrors.db_not_connected_error.value + ) + + try: + create_data: Dict[str, Any] = { + "workflow_type": data.workflow_type, + "created_by": _caller_key(user_api_key_dict), + } + if data.input is not None: + create_data["input"] = _json(data.input) + if data.metadata is not None: + create_data["metadata"] = _json(data.metadata) + run = await prisma_client.db.litellm_workflowrun.create(data=create_data) + return run + except Exception as e: + verbose_proxy_logger.exception("Error creating workflow run: %s", e) + raise HTTPException(status_code=500, detail=str(e)) + + +@router.get( + "/v1/workflows/runs", + tags=["workflow management"], + dependencies=[Depends(user_api_key_auth)], +) +async def list_workflow_runs( + workflow_type: Optional[str] = Query(None), + status: Optional[str] = Query(None), + limit: int = Query(50, ge=1, le=250), + user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), +): + """List workflow runs. Filter by workflow_type and/or status. + + Non-admin callers only see runs created by their own API key. + """ + from litellm.proxy.proxy_server import prisma_client + + if prisma_client is None: + raise HTTPException( + status_code=500, detail=CommonProxyErrors.db_not_connected_error.value + ) + + where: Dict[str, Any] = {} + if workflow_type: + where["workflow_type"] = workflow_type + if status: + statuses = [s.strip() for s in status.split(",")] + where["status"] = {"in": statuses} if len(statuses) > 1 else statuses[0] + + # Non-admin callers are scoped to their own key. + if not _is_admin(user_api_key_dict): + caller = _caller_key(user_api_key_dict) + if caller: + where["created_by"] = caller + + try: + runs = await prisma_client.db.litellm_workflowrun.find_many( + where=where, + order={"created_at": "desc"}, + take=limit, + ) + return {"runs": runs, "count": len(runs)} + except Exception as e: + verbose_proxy_logger.exception("Error listing workflow runs: %s", e) + raise HTTPException(status_code=500, detail=str(e)) + + +@router.get( + "/v1/workflows/runs/{run_id}", + tags=["workflow management"], + dependencies=[Depends(user_api_key_auth)], +) +async def get_workflow_run( + run_id: str, + user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), +): + """Get a workflow run with its most recent event.""" + from litellm.proxy.proxy_server import prisma_client + + if prisma_client is None: + raise HTTPException( + status_code=500, detail=CommonProxyErrors.db_not_connected_error.value + ) + + try: + run = await prisma_client.db.litellm_workflowrun.find_unique( + where={"run_id": run_id}, + include={"events": {"order_by": {"sequence_number": "desc"}, "take": 1}}, + ) + if run is None: + raise HTTPException(status_code=404, detail=f"Run '{run_id}' not found") + if not _is_admin(user_api_key_dict): + caller = _caller_key(user_api_key_dict) + if not caller or run.created_by != caller: + raise HTTPException(status_code=404, detail=f"Run '{run_id}' not found") + return run + except HTTPException: + raise + except Exception as e: + verbose_proxy_logger.exception("Error getting workflow run: %s", e) + raise HTTPException(status_code=500, detail=str(e)) + + +@router.patch( + "/v1/workflows/runs/{run_id}", + tags=["workflow management"], + dependencies=[Depends(user_api_key_auth)], +) +async def update_workflow_run( + run_id: str, + data: WorkflowRunUpdateRequest, + user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), +): + """Update status, metadata, or output on a workflow run.""" + from litellm.proxy.proxy_server import prisma_client + + if prisma_client is None: + raise HTTPException( + status_code=500, detail=CommonProxyErrors.db_not_connected_error.value + ) + + update: Dict[str, Any] = {} + if data.status is not None: + update["status"] = data.status + if data.output is not None: + update["output"] = _json(data.output) + if data.metadata is not None: + update["metadata"] = _json(data.metadata) + + if not update: + raise HTTPException(status_code=400, detail="No fields to update") + + # Enforce ownership before writing. + await _require_run(prisma_client, run_id, user_api_key_dict) + + try: + run = await prisma_client.db.litellm_workflowrun.update( + where={"run_id": run_id}, + data=update, + ) + if run is None: + raise HTTPException(status_code=404, detail=f"Run '{run_id}' not found") + return run + except HTTPException: + raise + except Exception as e: + verbose_proxy_logger.exception("Error updating workflow run: %s", e) + raise HTTPException(status_code=500, detail=str(e)) + + +@router.post( + "/v1/workflows/runs/{run_id}/events", + tags=["workflow management"], + dependencies=[Depends(user_api_key_auth)], +) +async def append_workflow_event( + run_id: str, + data: WorkflowEventCreateRequest, + user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), +): + """Append an event to the run's event log. Also updates run.status if event_type maps to a status. + + Sequence numbers use optimistic concurrency: on a unique-constraint collision + (concurrent append), retries up to _MAX_SEQUENCE_RETRIES times with a fresh MAX+1. + The event+status update is atomic in a single DB transaction. + """ + from litellm.proxy.proxy_server import prisma_client + + if prisma_client is None: + raise HTTPException( + status_code=500, detail=CommonProxyErrors.db_not_connected_error.value + ) + + await _require_run(prisma_client, run_id, user_api_key_dict) + + new_status = _EVENT_STATUS_MAP.get(data.event_type) + + for attempt in range(_MAX_SEQUENCE_RETRIES): + try: + seq = await _get_next_sequence_number(prisma_client, run_id, "events") + event_data: Dict[str, Any] = { + "run_id": run_id, + "event_type": data.event_type, + "step_name": data.step_name, + "sequence_number": seq, + } + if data.data is not None: + event_data["data"] = _json(data.data) + + async with prisma_client.db.tx() as tx: + event = await tx.litellm_workflowevent.create(data=event_data) + if new_status: + await tx.litellm_workflowrun.update( + where={"run_id": run_id}, + data={"status": new_status}, + ) + + return event + + except Exception as e: + if UniqueViolationError is not None and isinstance(e, UniqueViolationError): + if attempt == _MAX_SEQUENCE_RETRIES - 1: + verbose_proxy_logger.exception( + "Sequence number collision after %d retries for run %s", + _MAX_SEQUENCE_RETRIES, + run_id, + ) + raise HTTPException( + status_code=409, + detail="Concurrent write conflict — please retry", + ) + continue + verbose_proxy_logger.exception("Error appending workflow event: %s", e) + raise HTTPException(status_code=500, detail=str(e)) + + raise HTTPException( + status_code=500, detail="Failed to append event" + ) # pragma: no cover + + +@router.get( + "/v1/workflows/runs/{run_id}/events", + tags=["workflow management"], + dependencies=[Depends(user_api_key_auth)], +) +async def list_workflow_events( + run_id: str, + limit: int = Query(100, ge=1, le=500), + user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), +): + """Fetch event log for a run, ordered by sequence_number. Default limit 100, max 500.""" + from litellm.proxy.proxy_server import prisma_client + + if prisma_client is None: + raise HTTPException( + status_code=500, detail=CommonProxyErrors.db_not_connected_error.value + ) + + await _require_run(prisma_client, run_id, user_api_key_dict) + + try: + events = await prisma_client.db.litellm_workflowevent.find_many( + where={"run_id": run_id}, + order={"sequence_number": "asc"}, + take=limit, + ) + return {"events": events, "count": len(events)} + except Exception as e: + verbose_proxy_logger.exception("Error listing workflow events: %s", e) + raise HTTPException(status_code=500, detail=str(e)) + + +@router.post( + "/v1/workflows/runs/{run_id}/messages", + tags=["workflow management"], + dependencies=[Depends(user_api_key_auth)], +) +async def append_workflow_message( + run_id: str, + data: WorkflowMessageCreateRequest, + user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), +): + """Append a conversation message. Stores full content (not truncated). + + Uses optimistic concurrency for sequence numbers. + """ + from litellm.proxy.proxy_server import prisma_client + + if prisma_client is None: + raise HTTPException( + status_code=500, detail=CommonProxyErrors.db_not_connected_error.value + ) + + await _require_run(prisma_client, run_id, user_api_key_dict) + + for attempt in range(_MAX_SEQUENCE_RETRIES): + try: + seq = await _get_next_sequence_number(prisma_client, run_id, "messages") + msg_data: Dict[str, Any] = { + "run_id": run_id, + "role": data.role, + "content": data.content, + "sequence_number": seq, + } + if data.session_id is not None: + msg_data["session_id"] = data.session_id + msg = await prisma_client.db.litellm_workflowmessage.create(data=msg_data) + return msg + + except Exception as e: + if UniqueViolationError is not None and isinstance(e, UniqueViolationError): + if attempt == _MAX_SEQUENCE_RETRIES - 1: + verbose_proxy_logger.exception( + "Sequence number collision after %d retries for run %s", + _MAX_SEQUENCE_RETRIES, + run_id, + ) + raise HTTPException( + status_code=409, + detail="Concurrent write conflict — please retry", + ) + continue + verbose_proxy_logger.exception("Error appending workflow message: %s", e) + raise HTTPException(status_code=500, detail=str(e)) + + raise HTTPException( + status_code=500, detail="Failed to append message" + ) # pragma: no cover + + +@router.get( + "/v1/workflows/runs/{run_id}/messages", + tags=["workflow management"], + dependencies=[Depends(user_api_key_auth)], +) +async def list_workflow_messages( + run_id: str, + limit: int = Query(100, ge=1, le=500), + user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), +): + """Fetch conversation history for a run, ordered by sequence_number. Default limit 100, max 500.""" + from litellm.proxy.proxy_server import prisma_client + + if prisma_client is None: + raise HTTPException( + status_code=500, detail=CommonProxyErrors.db_not_connected_error.value + ) + + await _require_run(prisma_client, run_id, user_api_key_dict) + + try: + messages = await prisma_client.db.litellm_workflowmessage.find_many( + where={"run_id": run_id}, + order={"sequence_number": "asc"}, + take=limit, + ) + return {"messages": messages, "count": len(messages)} + except Exception as e: + verbose_proxy_logger.exception("Error listing workflow messages: %s", e) + raise HTTPException(status_code=500, detail=str(e)) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 870ea78f17..1222995529 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -426,6 +426,9 @@ from litellm.proxy.management_endpoints.team_endpoints import ( from litellm.proxy.management_endpoints.tool_management_endpoints import ( router as tool_management_router, ) +from litellm.proxy.management_endpoints.workflow_management_endpoints import ( + router as workflow_management_router, +) from litellm.proxy.memory.memory_endpoints import router as memory_router from litellm.proxy.management_endpoints.ui_sso import ( get_disabled_non_admin_personal_key_creation, @@ -14283,6 +14286,7 @@ app.include_router(model_management_router) app.include_router(model_access_group_management_router) app.include_router(tag_management_router) app.include_router(tool_management_router) +app.include_router(workflow_management_router) app.include_router(memory_router) app.include_router(cost_tracking_settings_router) app.include_router(router_settings_router) diff --git a/litellm/proxy/schema.prisma b/litellm/proxy/schema.prisma index 8f07c5afa3..fd54ed5243 100644 --- a/litellm/proxy/schema.prisma +++ b/litellm/proxy/schema.prisma @@ -1290,3 +1290,80 @@ model LiteLLM_AdaptiveRouterSession { @@id([session_id, router_name, model_name]) @@index([last_activity_at], map: "idx_adaptive_router_session_activity") } + +// --------------------------------------------------------------------------- +// Workflow Run Tracking +// +// Generic durable state tracking for any agent or automated workflow. +// Design: three tables — run (header + materialized status), event (append-only +// source of truth for state transitions), message (conversation inbox/outbox). +// +// Usage: +// - Set `workflow_type` to identify the owning system (e.g. "shin-builder"). +// - Store domain-specific fields in `metadata` (worktree_path, pr_url, etc.). +// - `session_id` on WorkflowRun matches `x-litellm-session-id` header sent to +// the proxy — all spend logs for this run are automatically tagged. +// --------------------------------------------------------------------------- + +// One instance of work being done. `status` is a materialized cache of the +// latest event; the event log is the authoritative source of truth. +model LiteLLM_WorkflowRun { + run_id String @id @default(uuid()) + session_id String @unique @default(uuid()) + workflow_type String + status String @default("pending") + created_by String? // user_id of the key that created this run; null = created by master key + created_at DateTime @default(now()) + updated_at DateTime @updatedAt + input Json? + output Json? + metadata Json? + + events LiteLLM_WorkflowEvent[] + messages LiteLLM_WorkflowMessage[] + + @@index([workflow_type, status]) + @@index([session_id]) + @@index([created_at]) + @@index([created_by]) +} + +// Append-only log of state transitions. Never mutate rows here. +// `step_name` and `event_type` are caller-defined strings — no hardcoded enums. +// Status auto-update rules (applied by the append endpoint): +// step.started → run.status = running +// step.failed → run.status = failed +// hook.waiting → run.status = paused +// hook.received → run.status = running +model LiteLLM_WorkflowEvent { + event_id String @id @default(uuid()) + run_id String + event_type String + step_name String + sequence_number Int + data Json? + created_at DateTime @default(now()) + + run LiteLLM_WorkflowRun @relation(fields: [run_id], references: [run_id]) + + @@unique([run_id, sequence_number]) + @@index([run_id]) +} + +// Conversation inbox/outbox — full message content, separate from the durable +// event log. Spend logs truncate messages; this table stores them in full. +// `session_id` here is the Claude --resume session ID (or similar). +model LiteLLM_WorkflowMessage { + message_id String @id @default(uuid()) + run_id String + role String + content String + sequence_number Int + session_id String? + created_at DateTime @default(now()) + + run LiteLLM_WorkflowRun @relation(fields: [run_id], references: [run_id]) + + @@unique([run_id, sequence_number]) + @@index([run_id]) +} diff --git a/litellm/proxy/workflows/README.md b/litellm/proxy/workflows/README.md new file mode 100644 index 0000000000..f452066afb --- /dev/null +++ b/litellm/proxy/workflows/README.md @@ -0,0 +1,150 @@ +# Workflow Run Tracking + +Generic durable state tracking for agents and automated workflows built on the LiteLLM proxy. + +## The Problem + +Agents like [shin-builder](https://github.com/BerriAI/shin-builder) run multi-stage pipelines (triage → plan → implement → PR). Their task state and conversation history lived in memory — a process restart lost everything. + +## Three-Table Design + +``` +WorkflowRun one instance of work (header + materialized status) +WorkflowEvent append-only state transitions (source of truth for replay) +WorkflowMessage conversation inbox/outbox (full content, not truncated) +``` + +**WorkflowEvent is the source of truth.** `WorkflowRun.status` is a materialized cache updated automatically when events are appended. If you need to debug a run, replay its events. + +## API + +All endpoints require a valid LiteLLM API key (`Authorization: Bearer sk-...`). + +### Runs + +``` +POST /v1/workflows/runs Create a run +GET /v1/workflows/runs List runs (?workflow_type=&status=) +GET /v1/workflows/runs/{run_id} Get run + latest event +PATCH /v1/workflows/runs/{run_id} Update status / metadata / output +``` + +### Events + +``` +POST /v1/workflows/runs/{run_id}/events Append event (auto-updates run status) +GET /v1/workflows/runs/{run_id}/events Full event log (ordered by sequence) +``` + +### Messages + +``` +POST /v1/workflows/runs/{run_id}/messages Append message +GET /v1/workflows/runs/{run_id}/messages Conversation history (ordered by sequence) +``` + +## Quick Start + +```bash +# Create a run +curl -X POST http://localhost:4000/v1/workflows/runs \ + -H "Authorization: Bearer sk-1234" \ + -H "Content-Type: application/json" \ + -d '{"workflow_type": "shin-builder", "metadata": {"title": "Fix login bug"}}' + +# {"run_id": "abc-123", "session_id": "xyz-456", "status": "pending", ...} + +# Mark step started (sets status → running) +curl -X POST http://localhost:4000/v1/workflows/runs/abc-123/events \ + -H "Authorization: Bearer sk-1234" \ + -H "Content-Type: application/json" \ + -d '{"event_type": "step.started", "step_name": "grill", "data": {"claude_session_id": "sess-789"}}' + +# Store a conversation message +curl -X POST http://localhost:4000/v1/workflows/runs/abc-123/messages \ + -H "Authorization: Bearer sk-1234" \ + -H "Content-Type: application/json" \ + -d '{"role": "user", "content": "What is the expected behavior?", "session_id": "sess-789"}' + +# Restart recovery: fetch active runs and resume from last event's data.claude_session_id +curl "http://localhost:4000/v1/workflows/runs?status=running,paused&workflow_type=shin-builder" \ + -H "Authorization: Bearer sk-1234" +``` + +## Status Auto-Update Rules + +When you append an event, the run's status is updated automatically: + +| event_type | run.status | +|-----------------|------------| +| `step.started` | `running` | +| `step.failed` | `failed` | +| `hook.waiting` | `paused` | +| `hook.received` | `running` | + +Set `status = completed` explicitly via PATCH when the workflow finishes. + +## Linking to Spend Logs + +`WorkflowRun.session_id` is generated automatically (UUID). Pass it as the `x-litellm-session-id` header when making completions through the proxy: + +```python +headers = {"x-litellm-session-id": run.session_id} +``` + +All spend log entries for this run are then tagged automatically. Query cost per run: + +``` +POST /ui/spend_logs/view_session_spend_logs?session_id={run.session_id} +``` + +## Sequence Numbers + +Sequence numbers on events and messages are assigned server-side (`MAX + 1` per run). Callers never supply them. This guarantees ordering even under concurrent writes. + +## Using from shin-builder + +Replace the in-memory `tasks.py` dict with calls to these endpoints: + +```python +import httpx + +class WorkflowRunClient: + def __init__(self, base_url: str, api_key: str): + self._client = httpx.AsyncClient( + base_url=base_url, + headers={"Authorization": f"Bearer {api_key}"}, + ) + + async def create_task(self, title: str, **metadata) -> dict: + r = await self._client.post("/v1/workflows/runs", json={ + "workflow_type": "shin-builder", + "metadata": {"title": title, **metadata}, + }) + r.raise_for_status() + return r.json() + + async def list_active_tasks(self) -> list: + r = await self._client.get( + "/v1/workflows/runs", + params={"workflow_type": "shin-builder", "status": "running,paused"}, + ) + r.raise_for_status() + return r.json()["runs"] + + async def transition(self, run_id: str, step_name: str, event_type: str, data: dict = None): + r = await self._client.post(f"/v1/workflows/runs/{run_id}/events", json={ + "event_type": event_type, + "step_name": step_name, + "data": data or {}, + }) + r.raise_for_status() + + async def append_message(self, run_id: str, role: str, content: str, session_id: str = None): + r = await self._client.post(f"/v1/workflows/runs/{run_id}/messages", json={ + "role": role, "content": content, "session_id": session_id, + }) + r.raise_for_status() +``` + +On startup, call `list_active_tasks()` to restore in-flight runs. The last `step.started` event's `data.claude_session_id` gives you the `--resume` ID. diff --git a/schema.prisma b/schema.prisma index 8f07c5afa3..fd54ed5243 100644 --- a/schema.prisma +++ b/schema.prisma @@ -1290,3 +1290,80 @@ model LiteLLM_AdaptiveRouterSession { @@id([session_id, router_name, model_name]) @@index([last_activity_at], map: "idx_adaptive_router_session_activity") } + +// --------------------------------------------------------------------------- +// Workflow Run Tracking +// +// Generic durable state tracking for any agent or automated workflow. +// Design: three tables — run (header + materialized status), event (append-only +// source of truth for state transitions), message (conversation inbox/outbox). +// +// Usage: +// - Set `workflow_type` to identify the owning system (e.g. "shin-builder"). +// - Store domain-specific fields in `metadata` (worktree_path, pr_url, etc.). +// - `session_id` on WorkflowRun matches `x-litellm-session-id` header sent to +// the proxy — all spend logs for this run are automatically tagged. +// --------------------------------------------------------------------------- + +// One instance of work being done. `status` is a materialized cache of the +// latest event; the event log is the authoritative source of truth. +model LiteLLM_WorkflowRun { + run_id String @id @default(uuid()) + session_id String @unique @default(uuid()) + workflow_type String + status String @default("pending") + created_by String? // user_id of the key that created this run; null = created by master key + created_at DateTime @default(now()) + updated_at DateTime @updatedAt + input Json? + output Json? + metadata Json? + + events LiteLLM_WorkflowEvent[] + messages LiteLLM_WorkflowMessage[] + + @@index([workflow_type, status]) + @@index([session_id]) + @@index([created_at]) + @@index([created_by]) +} + +// Append-only log of state transitions. Never mutate rows here. +// `step_name` and `event_type` are caller-defined strings — no hardcoded enums. +// Status auto-update rules (applied by the append endpoint): +// step.started → run.status = running +// step.failed → run.status = failed +// hook.waiting → run.status = paused +// hook.received → run.status = running +model LiteLLM_WorkflowEvent { + event_id String @id @default(uuid()) + run_id String + event_type String + step_name String + sequence_number Int + data Json? + created_at DateTime @default(now()) + + run LiteLLM_WorkflowRun @relation(fields: [run_id], references: [run_id]) + + @@unique([run_id, sequence_number]) + @@index([run_id]) +} + +// Conversation inbox/outbox — full message content, separate from the durable +// event log. Spend logs truncate messages; this table stores them in full. +// `session_id` here is the Claude --resume session ID (or similar). +model LiteLLM_WorkflowMessage { + message_id String @id @default(uuid()) + run_id String + role String + content String + sequence_number Int + session_id String? + created_at DateTime @default(now()) + + run LiteLLM_WorkflowRun @relation(fields: [run_id], references: [run_id]) + + @@unique([run_id, sequence_number]) + @@index([run_id]) +} diff --git a/tests/test_litellm/proxy/management_endpoints/test_workflow_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_workflow_management_endpoints.py new file mode 100644 index 0000000000..a337ff6d88 --- /dev/null +++ b/tests/test_litellm/proxy/management_endpoints/test_workflow_management_endpoints.py @@ -0,0 +1,611 @@ +""" +Unit tests for workflow management endpoints (/v1/workflows/runs/*). +Uses FastAPI TestClient with a mocked prisma_client. +""" + +import os +import sys +from datetime import datetime, timezone +from typing import Any +from unittest.mock import AsyncMock, MagicMock, patch + +from fastapi import FastAPI +from fastapi.testclient import TestClient +from prisma.errors import UniqueViolationError + +sys.path.insert(0, os.path.abspath("../../..")) + +from litellm.proxy.management_endpoints.workflow_management_endpoints import router + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + + +def _make_run( + run_id: str = "run-1", + session_id: str = "sess-1", + workflow_type: str = "shin-builder", + status: str = "pending", + created_by: Any = "tok-test", +) -> MagicMock: + obj = MagicMock() + obj.run_id = run_id + obj.session_id = session_id + obj.workflow_type = workflow_type + obj.status = status + obj.created_by = created_by + obj.created_at = datetime.now(timezone.utc) + obj.updated_at = datetime.now(timezone.utc) + obj.input = None + obj.output = None + obj.metadata = None + return obj + + +def _make_event( + event_id: str = "evt-1", + run_id: str = "run-1", + event_type: str = "step.started", + step_name: str = "grill", + sequence_number: int = 0, +) -> MagicMock: + obj = MagicMock() + obj.event_id = event_id + obj.run_id = run_id + obj.event_type = event_type + obj.step_name = step_name + obj.sequence_number = sequence_number + obj.data = None + obj.created_at = datetime.now(timezone.utc) + return obj + + +def _make_message( + message_id: str = "msg-1", + run_id: str = "run-1", + role: str = "user", + content: str = "hello", + sequence_number: int = 0, +) -> MagicMock: + obj = MagicMock() + obj.message_id = message_id + obj.run_id = run_id + obj.role = role + obj.content = content + obj.sequence_number = sequence_number + obj.session_id = None + obj.created_at = datetime.now(timezone.utc) + return obj + + +def _make_tx(event_return=None, run_return=None, msg_return=None) -> MagicMock: + """Build an async context-manager mock for prisma_client.db.tx().""" + tx = MagicMock() + tx.litellm_workflowevent = MagicMock() + tx.litellm_workflowevent.create = AsyncMock( + return_value=event_return or _make_event() + ) + tx.litellm_workflowrun = MagicMock() + tx.litellm_workflowrun.update = AsyncMock(return_value=run_return or _make_run()) + tx.litellm_workflowmessage = MagicMock() + tx.litellm_workflowmessage.create = AsyncMock( + return_value=msg_return or _make_message() + ) + tx.__aenter__ = AsyncMock(return_value=tx) + tx.__aexit__ = AsyncMock(return_value=False) + return tx + + +def _make_prisma_client() -> MagicMock: + client = MagicMock() + client.db = MagicMock() + client.db.litellm_workflowrun = MagicMock() + client.db.litellm_workflowevent = MagicMock() + client.db.litellm_workflowmessage = MagicMock() + # default tx() returns a no-op transaction + client.db.tx = MagicMock(return_value=_make_tx()) + return client + + +def _make_app() -> FastAPI: + app = FastAPI() + app.include_router(router) + return app + + +def _override_auth() -> Any: + from litellm.proxy._types import UserAPIKeyAuth + + auth = UserAPIKeyAuth(api_key="sk-test", user_id="admin") + auth.token = "tok-test" + return auth + + +def _override_auth_admin() -> Any: + from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth + + auth = UserAPIKeyAuth(api_key="sk-master") + auth.user_role = LitellmUserRoles.PROXY_ADMIN # type: ignore[assignment] + return auth + + +def _override_auth_user_with_token(token: str = "tok-abc") -> Any: + """Return a non-admin caller whose hashed token equals `token`.""" + from litellm.proxy._types import UserAPIKeyAuth + + auth = UserAPIKeyAuth(api_key="sk-user", user_id="user-1") + auth.token = token # override the computed hash with a predictable value + return auth + + +# --------------------------------------------------------------------------- +# Tests +# --------------------------------------------------------------------------- + + +class TestCreateWorkflowRun: + def setup_method(self): + from litellm.proxy.auth.user_api_key_auth import user_api_key_auth + + self._prisma = _make_prisma_client() + app = _make_app() + app.dependency_overrides[user_api_key_auth] = _override_auth + self.client = TestClient(app, raise_server_exceptions=True) + + @patch("litellm.proxy.proxy_server.prisma_client") + def test_create_returns_run(self, mock_pc): + mock_pc.db = self._prisma.db + self._prisma.db.litellm_workflowrun.create = AsyncMock(return_value=_make_run()) + + resp = self.client.post( + "/v1/workflows/runs", + json={"workflow_type": "shin-builder"}, + ) + assert resp.status_code == 200 + self._prisma.db.litellm_workflowrun.create.assert_awaited_once() + + @patch("litellm.proxy.proxy_server.prisma_client", None) + def test_create_500_when_no_db(self): + resp = self.client.post( + "/v1/workflows/runs", + json={"workflow_type": "shin-builder"}, + ) + assert resp.status_code == 500 + + +class TestListWorkflowRuns: + def setup_method(self): + from litellm.proxy.auth.user_api_key_auth import user_api_key_auth + + self._prisma = _make_prisma_client() + app = _make_app() + app.dependency_overrides[user_api_key_auth] = _override_auth + self.client = TestClient(app, raise_server_exceptions=True) + + @patch("litellm.proxy.proxy_server.prisma_client") + def test_list_returns_runs(self, mock_pc): + mock_pc.db = self._prisma.db + self._prisma.db.litellm_workflowrun.find_many = AsyncMock( + return_value=[_make_run()] + ) + + resp = self.client.get("/v1/workflows/runs") + assert resp.status_code == 200 + data = resp.json() + assert data["count"] == 1 + + @patch("litellm.proxy.proxy_server.prisma_client") + def test_list_filters_by_status(self, mock_pc): + mock_pc.db = self._prisma.db + self._prisma.db.litellm_workflowrun.find_many = AsyncMock(return_value=[]) + + resp = self.client.get("/v1/workflows/runs?status=running") + assert resp.status_code == 200 + call_kwargs = self._prisma.db.litellm_workflowrun.find_many.call_args[1] + assert call_kwargs["where"]["status"] == "running" + + @patch("litellm.proxy.proxy_server.prisma_client") + def test_list_filters_by_multiple_statuses(self, mock_pc): + mock_pc.db = self._prisma.db + self._prisma.db.litellm_workflowrun.find_many = AsyncMock(return_value=[]) + + resp = self.client.get("/v1/workflows/runs?status=running,paused") + assert resp.status_code == 200 + call_kwargs = self._prisma.db.litellm_workflowrun.find_many.call_args[1] + assert call_kwargs["where"]["status"] == {"in": ["running", "paused"]} + + +class TestGetWorkflowRun: + def setup_method(self): + from litellm.proxy.auth.user_api_key_auth import user_api_key_auth + + self._prisma = _make_prisma_client() + app = _make_app() + app.dependency_overrides[user_api_key_auth] = _override_auth + self.client = TestClient(app, raise_server_exceptions=True) + + @patch("litellm.proxy.proxy_server.prisma_client") + def test_get_existing_run(self, mock_pc): + mock_pc.db = self._prisma.db + self._prisma.db.litellm_workflowrun.find_unique = AsyncMock( + return_value=_make_run() + ) + + resp = self.client.get("/v1/workflows/runs/run-1") + assert resp.status_code == 200 + + @patch("litellm.proxy.proxy_server.prisma_client") + def test_get_missing_run_returns_404(self, mock_pc): + mock_pc.db = self._prisma.db + self._prisma.db.litellm_workflowrun.find_unique = AsyncMock(return_value=None) + + resp = self.client.get("/v1/workflows/runs/nonexistent") + assert resp.status_code == 404 + + +class TestUpdateWorkflowRun: + def setup_method(self): + from litellm.proxy.auth.user_api_key_auth import user_api_key_auth + + self._prisma = _make_prisma_client() + app = _make_app() + app.dependency_overrides[user_api_key_auth] = _override_auth + self.client = TestClient(app, raise_server_exceptions=True) + + @patch("litellm.proxy.proxy_server.prisma_client") + def test_update_status(self, mock_pc): + mock_pc.db = self._prisma.db + self._prisma.db.litellm_workflowrun.find_unique = AsyncMock( + return_value=_make_run() + ) + updated = _make_run(status="completed") + self._prisma.db.litellm_workflowrun.update = AsyncMock(return_value=updated) + + resp = self.client.patch( + "/v1/workflows/runs/run-1", json={"status": "completed"} + ) + assert resp.status_code == 200 + self._prisma.db.litellm_workflowrun.update.assert_awaited_once() + + @patch("litellm.proxy.proxy_server.prisma_client") + def test_update_no_fields_returns_400(self, mock_pc): + mock_pc.db = self._prisma.db + resp = self.client.patch("/v1/workflows/runs/run-1", json={}) + assert resp.status_code == 400 + + +class TestAppendWorkflowEvent: + def setup_method(self): + from litellm.proxy.auth.user_api_key_auth import user_api_key_auth + + self._prisma = _make_prisma_client() + app = _make_app() + app.dependency_overrides[user_api_key_auth] = _override_auth + self.client = TestClient(app, raise_server_exceptions=True) + + @patch("litellm.proxy.proxy_server.prisma_client") + def test_append_event_updates_run_status(self, mock_pc): + mock_pc.db = self._prisma.db + # _require_run check + self._prisma.db.litellm_workflowrun.find_unique = AsyncMock( + return_value=_make_run() + ) + self._prisma.db.litellm_workflowevent.find_many = AsyncMock(return_value=[]) + tx = _make_tx( + event_return=_make_event(), run_return=_make_run(status="running") + ) + self._prisma.db.tx = MagicMock(return_value=tx) + + resp = self.client.post( + "/v1/workflows/runs/run-1/events", + json={"event_type": "step.started", "step_name": "grill"}, + ) + assert resp.status_code == 200 + # run status updated inside tx + tx.litellm_workflowrun.update.assert_awaited_once() + update_call = tx.litellm_workflowrun.update.call_args[1] + assert update_call["data"]["status"] == "running" + + @patch("litellm.proxy.proxy_server.prisma_client") + def test_append_event_no_status_update_for_unknown_type(self, mock_pc): + mock_pc.db = self._prisma.db + self._prisma.db.litellm_workflowrun.find_unique = AsyncMock( + return_value=_make_run() + ) + self._prisma.db.litellm_workflowevent.find_many = AsyncMock(return_value=[]) + tx = _make_tx(event_return=_make_event(event_type="custom.event")) + self._prisma.db.tx = MagicMock(return_value=tx) + + resp = self.client.post( + "/v1/workflows/runs/run-1/events", + json={"event_type": "custom.event", "step_name": "grill"}, + ) + assert resp.status_code == 200 + # no status update inside tx for unknown event_type + tx.litellm_workflowrun.update.assert_not_awaited() + + @patch("litellm.proxy.proxy_server.prisma_client") + def test_sequence_number_increments(self, mock_pc): + mock_pc.db = self._prisma.db + self._prisma.db.litellm_workflowrun.find_unique = AsyncMock( + return_value=_make_run() + ) + existing = _make_event(sequence_number=4) + self._prisma.db.litellm_workflowevent.find_many = AsyncMock( + return_value=[existing] + ) + tx = _make_tx(event_return=_make_event(sequence_number=5)) + self._prisma.db.tx = MagicMock(return_value=tx) + + self.client.post( + "/v1/workflows/runs/run-1/events", + json={"event_type": "step.started", "step_name": "plan"}, + ) + create_call = tx.litellm_workflowevent.create.call_args[1] + assert create_call["data"]["sequence_number"] == 5 + + @patch("litellm.proxy.proxy_server.prisma_client") + def test_unknown_run_id_returns_404(self, mock_pc): + mock_pc.db = self._prisma.db + self._prisma.db.litellm_workflowrun.find_unique = AsyncMock(return_value=None) + + resp = self.client.post( + "/v1/workflows/runs/nonexistent/events", + json={"event_type": "step.started", "step_name": "grill"}, + ) + assert resp.status_code == 404 + + @patch("litellm.proxy.proxy_server.prisma_client") + def test_sequence_collision_retries_and_succeeds(self, mock_pc): + """UniqueViolationError on first attempt triggers retry; second attempt succeeds.""" + mock_pc.db = self._prisma.db + self._prisma.db.litellm_workflowrun.find_unique = AsyncMock( + return_value=_make_run() + ) + self._prisma.db.litellm_workflowevent.find_many = AsyncMock(return_value=[]) + + # First tx raises UniqueViolationError; second succeeds. + tx_fail = _make_tx() + tx_fail.__aenter__ = AsyncMock(return_value=tx_fail) + tx_fail.litellm_workflowevent.create = AsyncMock( + side_effect=UniqueViolationError( + {"user_facing_error": {"message": "unique"}} + ) + ) + tx_fail.__aexit__ = AsyncMock(return_value=False) + + tx_ok = _make_tx(event_return=_make_event(sequence_number=1)) + + self._prisma.db.tx = MagicMock(side_effect=[tx_fail, tx_ok]) + + resp = self.client.post( + "/v1/workflows/runs/run-1/events", + json={"event_type": "step.started", "step_name": "grill"}, + ) + assert resp.status_code == 200 + + +class TestWorkflowMessages: + def setup_method(self): + from litellm.proxy.auth.user_api_key_auth import user_api_key_auth + + self._prisma = _make_prisma_client() + app = _make_app() + app.dependency_overrides[user_api_key_auth] = _override_auth + self.client = TestClient(app, raise_server_exceptions=True) + + @patch("litellm.proxy.proxy_server.prisma_client") + def test_append_message(self, mock_pc): + mock_pc.db = self._prisma.db + self._prisma.db.litellm_workflowrun.find_unique = AsyncMock( + return_value=_make_run() + ) + self._prisma.db.litellm_workflowmessage.find_many = AsyncMock(return_value=[]) + self._prisma.db.litellm_workflowmessage.create = AsyncMock( + return_value=_make_message() + ) + + resp = self.client.post( + "/v1/workflows/runs/run-1/messages", + json={"role": "user", "content": "fix the bug"}, + ) + assert resp.status_code == 200 + + @patch("litellm.proxy.proxy_server.prisma_client") + def test_append_message_unknown_run_returns_404(self, mock_pc): + mock_pc.db = self._prisma.db + self._prisma.db.litellm_workflowrun.find_unique = AsyncMock(return_value=None) + + resp = self.client.post( + "/v1/workflows/runs/nonexistent/messages", + json={"role": "user", "content": "hello"}, + ) + assert resp.status_code == 404 + + @patch("litellm.proxy.proxy_server.prisma_client") + def test_list_messages_ordered(self, mock_pc): + mock_pc.db = self._prisma.db + self._prisma.db.litellm_workflowrun.find_unique = AsyncMock( + return_value=_make_run() + ) + self._prisma.db.litellm_workflowmessage.find_many = AsyncMock( + return_value=[ + _make_message(sequence_number=0), + _make_message(sequence_number=1, role="assistant"), + ] + ) + + resp = self.client.get("/v1/workflows/runs/run-1/messages") + assert resp.status_code == 200 + data = resp.json() + assert data["count"] == 2 + call_kwargs = self._prisma.db.litellm_workflowmessage.find_many.call_args[1] + assert call_kwargs["order"] == {"sequence_number": "asc"} + + @patch("litellm.proxy.proxy_server.prisma_client") + def test_list_messages_respects_limit(self, mock_pc): + mock_pc.db = self._prisma.db + self._prisma.db.litellm_workflowrun.find_unique = AsyncMock( + return_value=_make_run() + ) + self._prisma.db.litellm_workflowmessage.find_many = AsyncMock(return_value=[]) + + resp = self.client.get("/v1/workflows/runs/run-1/messages?limit=25") + assert resp.status_code == 200 + call_kwargs = self._prisma.db.litellm_workflowmessage.find_many.call_args[1] + assert call_kwargs["take"] == 25 + + +class TestListWorkflowEvents: + def setup_method(self): + from litellm.proxy.auth.user_api_key_auth import user_api_key_auth + + self._prisma = _make_prisma_client() + app = _make_app() + app.dependency_overrides[user_api_key_auth] = _override_auth + self.client = TestClient(app, raise_server_exceptions=True) + + @patch("litellm.proxy.proxy_server.prisma_client") + def test_list_events_ordered(self, mock_pc): + mock_pc.db = self._prisma.db + self._prisma.db.litellm_workflowrun.find_unique = AsyncMock( + return_value=_make_run() + ) + self._prisma.db.litellm_workflowevent.find_many = AsyncMock( + return_value=[ + _make_event(sequence_number=0), + _make_event(sequence_number=1), + ] + ) + + resp = self.client.get("/v1/workflows/runs/run-1/events") + assert resp.status_code == 200 + data = resp.json() + assert data["count"] == 2 + call_kwargs = self._prisma.db.litellm_workflowevent.find_many.call_args[1] + assert call_kwargs["order"] == {"sequence_number": "asc"} + + @patch("litellm.proxy.proxy_server.prisma_client") + def test_list_events_respects_limit(self, mock_pc): + mock_pc.db = self._prisma.db + self._prisma.db.litellm_workflowrun.find_unique = AsyncMock( + return_value=_make_run() + ) + self._prisma.db.litellm_workflowevent.find_many = AsyncMock(return_value=[]) + + resp = self.client.get("/v1/workflows/runs/run-1/events?limit=10") + assert resp.status_code == 200 + call_kwargs = self._prisma.db.litellm_workflowevent.find_many.call_args[1] + assert call_kwargs["take"] == 10 + + @patch("litellm.proxy.proxy_server.prisma_client") + def test_list_events_unknown_run_returns_404(self, mock_pc): + mock_pc.db = self._prisma.db + self._prisma.db.litellm_workflowrun.find_unique = AsyncMock(return_value=None) + + resp = self.client.get("/v1/workflows/runs/nonexistent/events") + assert resp.status_code == 404 + + +class TestTenantIsolation: + """Ownership enforcement: non-admin callers only see their own runs.""" + + def _make_app_with_auth(self, auth_fn): + from litellm.proxy.auth.user_api_key_auth import user_api_key_auth + + self._prisma = _make_prisma_client() + app = _make_app() + app.dependency_overrides[user_api_key_auth] = auth_fn + return TestClient(app, raise_server_exceptions=True) + + @patch("litellm.proxy.proxy_server.prisma_client") + def test_create_stores_caller_token(self, mock_pc): + token = "tok-owner" + client = self._make_app_with_auth(lambda: _override_auth_user_with_token(token)) + mock_pc.db = self._prisma.db + self._prisma.db.litellm_workflowrun.create = AsyncMock( + return_value=_make_run(created_by=token) + ) + + resp = client.post("/v1/workflows/runs", json={"workflow_type": "test"}) + assert resp.status_code == 200 + create_call = self._prisma.db.litellm_workflowrun.create.call_args[1] + assert create_call["data"]["created_by"] == token + + @patch("litellm.proxy.proxy_server.prisma_client") + def test_non_admin_list_scoped_to_caller_token(self, mock_pc): + token = "tok-owner" + client = self._make_app_with_auth(lambda: _override_auth_user_with_token(token)) + mock_pc.db = self._prisma.db + self._prisma.db.litellm_workflowrun.find_many = AsyncMock(return_value=[]) + + resp = client.get("/v1/workflows/runs") + assert resp.status_code == 200 + call_kwargs = self._prisma.db.litellm_workflowrun.find_many.call_args[1] + assert call_kwargs["where"].get("created_by") == token + + @patch("litellm.proxy.proxy_server.prisma_client") + def test_admin_list_not_scoped(self, mock_pc): + client = self._make_app_with_auth(_override_auth_admin) + mock_pc.db = self._prisma.db + self._prisma.db.litellm_workflowrun.find_many = AsyncMock(return_value=[]) + + resp = client.get("/v1/workflows/runs") + assert resp.status_code == 200 + call_kwargs = self._prisma.db.litellm_workflowrun.find_many.call_args[1] + assert "created_by" not in call_kwargs["where"] + + @patch("litellm.proxy.proxy_server.prisma_client") + def test_non_admin_get_other_users_run_returns_404(self, mock_pc): + token = "tok-caller" + client = self._make_app_with_auth(lambda: _override_auth_user_with_token(token)) + mock_pc.db = self._prisma.db + # Run owned by a different key + self._prisma.db.litellm_workflowrun.find_unique = AsyncMock( + return_value=_make_run(created_by="tok-other-owner") + ) + + resp = client.get("/v1/workflows/runs/run-1") + assert resp.status_code == 404 + + @patch("litellm.proxy.proxy_server.prisma_client") + def test_non_admin_get_null_owner_run_returns_404(self, mock_pc): + token = "tok-caller" + client = self._make_app_with_auth(lambda: _override_auth_user_with_token(token)) + mock_pc.db = self._prisma.db + self._prisma.db.litellm_workflowrun.find_unique = AsyncMock( + return_value=_make_run(created_by=None) + ) + + resp = client.get("/v1/workflows/runs/run-1") + assert resp.status_code == 404 + + @patch("litellm.proxy.proxy_server.prisma_client") + def test_non_admin_update_null_owner_run_returns_404(self, mock_pc): + token = "tok-caller" + client = self._make_app_with_auth(lambda: _override_auth_user_with_token(token)) + mock_pc.db = self._prisma.db + self._prisma.db.litellm_workflowrun.find_unique = AsyncMock( + return_value=_make_run(created_by=None) + ) + self._prisma.db.litellm_workflowrun.update = AsyncMock( + return_value=_make_run(status="completed") + ) + + resp = client.patch("/v1/workflows/runs/run-1", json={"status": "completed"}) + assert resp.status_code == 404 + self._prisma.db.litellm_workflowrun.update.assert_not_awaited() + + @patch("litellm.proxy.proxy_server.prisma_client") + def test_non_admin_get_own_run_succeeds(self, mock_pc): + token = "tok-caller" + client = self._make_app_with_auth(lambda: _override_auth_user_with_token(token)) + mock_pc.db = self._prisma.db + self._prisma.db.litellm_workflowrun.find_unique = AsyncMock( + return_value=_make_run(created_by=token) + ) + + resp = client.get("/v1/workflows/runs/run-1") + assert resp.status_code == 200 diff --git a/ui/litellm-dashboard/src/app/page.tsx b/ui/litellm-dashboard/src/app/page.tsx index d32cf74d6a..a8553d5405 100644 --- a/ui/litellm-dashboard/src/app/page.tsx +++ b/ui/litellm-dashboard/src/app/page.tsx @@ -40,6 +40,7 @@ import { ProjectsPage } from "@/components/Projects/ProjectsPage"; import VectorStoreManagement from "@/components/vector_store_management"; import ToolPoliciesView from "@/components/ToolPoliciesView"; import { MemoryView } from "@/components/MemoryView"; +import WorkflowRuns from "@/components/workflow_runs"; import SpendLogsTable from "@/components/view_logs"; import ViewUserDashboard from "@/components/view_users"; import { ThemeProvider } from "@/contexts/ThemeContext"; @@ -631,6 +632,8 @@ function CreateKeyPageContent() { ) : page == "tool-policies" ? ( + ) : page == "workflows" ? ( + ) : page == "memory" ? ( { "Virtual Keys", "Playground", "Models + Endpoints", - "Agents", + "Agentic", "MCP Servers", "Guardrails", "Policies", diff --git a/ui/litellm-dashboard/src/components/leftnav.tsx b/ui/litellm-dashboard/src/components/leftnav.tsx index c340b65496..cecc99739b 100644 --- a/ui/litellm-dashboard/src/components/leftnav.tsx +++ b/ui/litellm-dashboard/src/components/leftnav.tsx @@ -3,6 +3,7 @@ import { useTeams } from "@/app/(dashboard)/hooks/teams/useTeams"; import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; import { ApiOutlined, + ApartmentOutlined, AppstoreOutlined, AuditOutlined, BankOutlined, @@ -120,11 +121,31 @@ const menuGroups: MenuGroup[] = [ roles: rolesWithWriteAccess, }, { - key: "agents", - page: "agents", - label: "Agents", + key: "agentic", + page: "agentic", + label: "Agentic", icon: , - roles: rolesWithWriteAccess, + children: [ + { + key: "agents", + page: "agents", + label: "Agents", + icon: , + roles: rolesWithWriteAccess, + }, + { + key: "workflows", + page: "workflows", + label: "Workflow Runs", + icon: , + }, + { + key: "memory", + page: "memory", + label: "Memory", + icon: , + }, + ], }, { key: "mcp-servers", @@ -139,12 +160,6 @@ const menuGroups: MenuGroup[] = [ icon: , roles: all_admin_roles, }, - { - key: "memory", - page: "memory", - label: "Memory", - icon: , - }, { key: "guardrails", page: "guardrails", diff --git a/ui/litellm-dashboard/src/components/page_metadata.ts b/ui/litellm-dashboard/src/components/page_metadata.ts index 63f03f4586..277bb0bd02 100644 --- a/ui/litellm-dashboard/src/components/page_metadata.ts +++ b/ui/litellm-dashboard/src/components/page_metadata.ts @@ -9,6 +9,8 @@ export const pageDescriptions: Record = { "llm-playground": "Interactive playground for testing LLM requests", models: "Configure and manage LLM models and endpoints", agents: "Create and manage AI agents", + agentic: "Manage agentic resources: agents, workflow runs, and memory", + workflows: "Track and inspect durable workflow run history", "mcp-servers": "Configure Model Context Protocol servers", memory: "Inspect and manage agent memory entries stored under /v1/memory", guardrails: "Set up content moderation and safety guardrails", diff --git a/ui/litellm-dashboard/src/components/workflow_runs/index.tsx b/ui/litellm-dashboard/src/components/workflow_runs/index.tsx new file mode 100644 index 0000000000..a69eac282b --- /dev/null +++ b/ui/litellm-dashboard/src/components/workflow_runs/index.tsx @@ -0,0 +1,751 @@ +import React, { useState, useEffect, useCallback } from "react"; +import { Button, Collapse, Drawer, Empty, Spin, Table, Tooltip, Typography } from "antd"; +import { ReloadOutlined } from "@ant-design/icons"; +import { proxyBaseUrl } from "@/components/networking"; + +const { Text } = Typography; + +interface WorkflowRunsProps { + accessToken: string | null; +} + +type RunStatus = "pending" | "running" | "paused" | "completed" | "failed"; + +interface RunMetadata { + title?: string; + state?: string; + pr_url?: string; + worktree_path?: string; + plan_text?: string; + grill_session_id?: string; + session_id?: string; + [key: string]: unknown; +} + +interface WorkflowRun { + run_id: string; + status: RunStatus; + workflow_type: string; + created_at: string; + metadata?: RunMetadata | null; +} + +interface WorkflowRunEvent { + event_id: string; + event_type: string; + step_name: string; + sequence_number: number; + created_at: string; + data?: Record | null; +} + +interface WorkflowRunMessage { + message_id: string; + role: string; + content: string; + sequence_number: number; + created_at: string; +} + +// ── design tokens ───────────────────────────────────────────────────────────── + +const STATUS_DOT: Record = { + pending: "#a1a1aa", + running: "#3b82f6", + paused: "#f59e0b", + completed: "#22c55e", + failed: "#ef4444", +}; + +const EVENT_COLOR: Record = { + "step.started": { bar: "#f0fdf4", border: "#86efac", text: "#16a34a" }, + "step.failed": { bar: "#fef2f2", border: "#fca5a5", text: "#dc2626" }, + "hook.waiting": { bar: "#fffbeb", border: "#fcd34d", text: "#d97706" }, + "hook.received": { bar: "#eff6ff", border: "#93c5fd", text: "#2563eb" }, +}; + +function eventStyle(type: string) { + return EVENT_COLOR[type] ?? { bar: "#f4f4f5", border: "#d4d4d8", text: "#52525b" }; +} + +// ── helpers ─────────────────────────────────────────────────────────────────── + +function timeAgo(iso: string): string { + const diff = Date.now() - new Date(iso).getTime(); + if (isNaN(diff)) return iso; + const s = Math.floor(diff / 1000); + if (s < 60) return `${s}s ago`; + const m = Math.floor(s / 60); + if (m < 60) return `${m}m ago`; + const h = Math.floor(m / 60); + if (h < 24) return `${h}h ago`; + return `${Math.floor(h / 24)}d ago`; +} + +function fmtDuration(ms: number): string { + if (ms < 0) return ""; + if (ms < 1000) return `${ms}ms`; + return `${(ms / 1000).toFixed(1)}s`; +} + +function runTitle(run: WorkflowRun): string { + const t = run.metadata?.title; + if (t) return String(t); + return run.workflow_type ?? run.run_id.slice(0, 8); +} + +function shortId(id: string): string { + return id.slice(0, 8); +} + +// ── status dot ──────────────────────────────────────────────────────────────── + +const StatusDot: React.FC<{ status: RunStatus; size?: number }> = ({ status, size = 8 }) => ( + +); + +// ── truncated text value ────────────────────────────────────────────────────── + +const TRUNCATE_AT = 120; + +const TruncatedValue: React.FC<{ value: string }> = ({ value }) => { + const [expanded, setExpanded] = useState(false); + if (value.length <= TRUNCATE_AT) { + return {value}; + } + return ( + + {expanded ? value : value.slice(0, TRUNCATE_AT) + "…"} + + + ); +}; + +// ── metadata card ───────────────────────────────────────────────────────────── + +const MetadataCard: React.FC<{ run: WorkflowRun }> = ({ run }) => { + const meta = run.metadata ?? {}; + + const primaryFields: { key: string; label: string }[] = [ + { key: "state", label: "state" }, + { key: "worktree_path", label: "worktree" }, + { key: "grill_session_id", label: "grill session" }, + { key: "session_id", label: "session" }, + ]; + + const primaryKeys = new Set(["title", ...primaryFields.map((f) => f.key)]); + const extraEntries = Object.entries(meta).filter( + ([k, v]) => !primaryKeys.has(k) && v !== null && v !== undefined && v !== "" + ); + + return ( +
+ {/* title bar */} +
+ + + {runTitle(run)} + + + {shortId(run.run_id)} + + + {run.workflow_type} + +
+ + {/* key fields grid */} +
+ + {run.status} + + + {timeAgo(run.created_at)} + + + {meta.pr_url && ( + + + {String(meta.pr_url)} + + + )} + + {primaryFields.map(({ key, label }) => { + const v = meta[key]; + if (v === null || v === undefined || v === "") return null; + const str = typeof v === "object" ? JSON.stringify(v) : String(v); + return ( + + + + ); + })} + + {extraEntries.map(([k, v]) => { + const str = typeof v === "object" ? JSON.stringify(v) : String(v); + return ( + + + + ); + })} +
+
+ ); +}; + +const FieldPair: React.FC<{ label: string; children: React.ReactNode }> = ({ + label, + children, +}) => ( +
+ + {label} + + {children} +
+); + +// ── gantt timeline ──────────────────────────────────────────────────────────── + +const GanttTimeline: React.FC<{ + run: WorkflowRun; + events: WorkflowRunEvent[]; +}> = ({ run, events }) => { + if (events.length === 0) { + return ( +
+ No events recorded +
+ ); + } + + const runStart = new Date(run.created_at).getTime(); + const eventTimes = events.map((e) => new Date(e.created_at).getTime()); + const lastTime = Math.max(...eventTimes); + const totalSpan = Math.max(lastTime - runStart, 1); + const totalDur = fmtDuration(lastTime - runStart); + + return ( +
+ {/* ruler */} +
+
+
+ {[0, 100].map((pct) => ( + + {pct === 0 ? "0" : totalDur} + + ))} +
+
+ + {/* outer run bar */} +
+
+ {runTitle(run)} +
+
+ {totalDur} +
+
+ + {/* event rows */} +
+ {events.map((ev) => { + const evTime = new Date(ev.created_at).getTime(); + const leftPct = ((evTime - runStart) / totalSpan) * 100; + + const nextIdx = events.findIndex((e) => e.sequence_number > ev.sequence_number); + const nextTime = + nextIdx >= 0 + ? new Date(events[nextIdx].created_at).getTime() + : lastTime + Math.max(totalSpan * 0.12, 500); + const widthPct = Math.max(8, ((nextTime - evTime) / totalSpan) * 100); + const style = eventStyle(ev.event_type); + const dur = fmtDuration(nextTime - evTime); + + return ( + +
+ {ev.step_name || ev.event_type} +
+
+ +
type: {ev.event_type}
+
step: {ev.step_name}
+
seq: {ev.sequence_number}
+
time: {timeAgo(ev.created_at)}
+ {ev.data && Object.keys(ev.data).length > 0 && ( +
data: {JSON.stringify(ev.data)}
+ )} +
+ } + > +
+ {ev.event_type} + {dur && {dur}} +
+ +
+ + ); + })} +
+
+ ); +}; + +// ── message row ─────────────────────────────────────────────────────────────── + +const MessageRow: React.FC<{ msg: WorkflowRunMessage }> = ({ msg }) => { + const roleColor: Record = { + user: "#2563eb", + assistant: "#16a34a", + system: "#7c3aed", + tool_result: "#d97706", + }; + const color = roleColor[msg.role] ?? "#52525b"; + + return ( +
+ [{msg.role}] +
+ + {msg.content} + + + {timeAgo(msg.created_at)} + +
+
+ ); +}; + +// ── main component ──────────────────────────────────────────────────────────── + +const WorkflowRuns: React.FC = ({ accessToken }) => { + const [runs, setRuns] = useState([]); + const [loadingRuns, setLoadingRuns] = useState(false); + const [selectedRun, setSelectedRun] = useState(null); + const [events, setEvents] = useState([]); + const [messages, setMessages] = useState([]); + const [loadingDetail, setLoadingDetail] = useState(false); + const [drawerOpen, setDrawerOpen] = useState(false); + + const fetchRuns = useCallback(async () => { + if (!accessToken) return; + setLoadingRuns(true); + try { + const res = await fetch(`${proxyBaseUrl ?? ""}/v1/workflows/runs?limit=100`, { + headers: { Authorization: `Bearer ${accessToken}` }, + }); + if (!res.ok) throw new Error(`HTTP ${res.status}`); + const data = await res.json(); + setRuns(data.runs ?? []); + } catch (err) { + console.error("workflow runs fetch failed:", err); + } finally { + setLoadingRuns(false); + } + }, [accessToken]); + + const fetchRunDetail = useCallback( + async (run: WorkflowRun) => { + if (!accessToken) return; + setSelectedRun(run); + setDrawerOpen(true); + setLoadingDetail(true); + setEvents([]); + setMessages([]); + try { + const base = proxyBaseUrl ?? ""; + const [evRes, msgRes] = await Promise.all([ + fetch(`${base}/v1/workflows/runs/${run.run_id}/events`, { + headers: { Authorization: `Bearer ${accessToken}` }, + }), + fetch(`${base}/v1/workflows/runs/${run.run_id}/messages`, { + headers: { Authorization: `Bearer ${accessToken}` }, + }), + ]); + const evData = evRes.ok ? await evRes.json() : { events: [] }; + const msgData = msgRes.ok ? await msgRes.json() : { messages: [] }; + setEvents( + [...(evData.events ?? [])].sort( + (a: WorkflowRunEvent, b: WorkflowRunEvent) => a.sequence_number - b.sequence_number + ) + ); + setMessages( + [...(msgData.messages ?? [])].sort( + (a: WorkflowRunMessage, b: WorkflowRunMessage) => a.sequence_number - b.sequence_number + ) + ); + } catch (err) { + console.error("workflow run detail fetch failed:", err); + } finally { + setLoadingDetail(false); + } + }, + [accessToken] + ); + + useEffect(() => { + fetchRuns(); + }, [fetchRuns]); + + const columns = [ + { + title: "Run", + dataIndex: "run_id", + key: "run", + render: (_: string, run: WorkflowRun) => ( +
+ +
+
+ {runTitle(run)} +
+
+ {shortId(run.run_id)} +
+
+
+ ), + }, + { + title: "Type", + dataIndex: "workflow_type", + key: "workflow_type", + render: (v: string) => ( + {v} + ), + }, + { + title: "Status", + dataIndex: "status", + key: "status", + render: (status: RunStatus, run: WorkflowRun) => { + const state = run.metadata?.state; + return ( +
+ + + {state ?? status} + +
+ ); + }, + }, + { + title: "Created", + dataIndex: "created_at", + key: "created_at", + render: (v: string) => ( + {timeAgo(v)} + ), + }, + ]; + + return ( +
+ {/* page header */} +
+
+
Workflow Runs
+
+ Durable state tracking for agents and automated workflows +
+
+ +
+ + {/* runs table — matches logs page density */} +
+
({ + onClick: () => fetchRunDetail(run), + style: { cursor: "pointer" }, + })} + locale={{ + emptyText: ( + No workflow runs yet} + image={Empty.PRESENTED_IMAGE_SIMPLE} + /> + ), + }} + className="[&_.ant-table-cell]:py-0.5 [&_.ant-table-thead_.ant-table-cell]:py-1" + style={{ border: "none" }} + /> + + + {/* detail drawer */} + setDrawerOpen(false)} + width={680} + title={null} + closable={false} + bodyStyle={{ padding: 0 }} + styles={{ body: { padding: 0 } }} + > + {!selectedRun ? null : loadingDetail ? ( +
+ +
+ ) : ( +
+ {/* drawer close + refresh */} +
+ + +
+ + {/* metadata card — top */} + + + {/* collapsible sections */} + + Timeline + + {events.length} {events.length === 1 ? "event" : "events"} + + + ), + children: ( +
+ +
+ ), + }, + { + key: "messages", + label: ( + + Messages + + {messages.length} + + + ), + children: messages.length === 0 ? ( +
+ No messages +
+ ) : ( +
+ {messages.map((msg) => ( + + ))} +
+ ), + }, + ]} + /> +
+ )} +
+ + ); +}; + +export default WorkflowRuns; From dedaf74a5ec16dcb98624ea9b4d6ceedd85b1ac1 Mon Sep 17 00:00:00 2001 From: stuxf <70670632+stuxf@users.noreply.github.com> Date: Wed, 29 Apr 2026 17:27:22 -0700 Subject: [PATCH 57/57] chore(auth): tighten clientside api_base handling (#26518) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * chore(auth): validate clientside api_base against SSRF guard; clear admin secrets on base override Two related issues with how the proxy handles client-supplied ``api_base`` / ``base_url`` overrides on chat-completion requests: 1. **SSRF gate bypass** — ``check_complete_credentials()`` returned ``True`` for any non-empty ``api_key``, allowing the ``is_request_body_safe`` ``banned_params`` loop to admit ``api_base`` / ``base_url`` values that point at private (RFC 1918), loopback, link-local, or cloud-metadata addresses. Now: when the gate sees a client-supplied ``api_base`` / ``base_url``, it runs the URL through ``litellm_core_utils.url_utils.validate_url`` (DNS-resolves, blocks internal/IMDS/LL networks, defends against rebinding). Rejection raises with a clear message. 2. **Admin-config leak on base override** — ``get_dynamic_litellm_params`` only carried the three clientside keys (``api_key``, ``api_base``, ``base_url``) from request to upstream call. Other admin-configured fields on ``litellm_params`` — ``organization``, ``extra_body``, ``extra_headers``, ``api_version``, ``azure_ad_token``, AWS / Vertex creds, etc. — flowed through unchanged. With base redirected to a client-controlled server, those admin secrets were sent to the attacker. Now: when ``api_base`` / ``base_url`` is in ``request_kwargs``, drop those admin-config fields from ``litellm_params`` unless the caller re-supplied them. Tests cover the SSRF-target rejection per URL field, the admin-secret clearing on base override, the don't-clear case when only ``api_key`` is overridden (BYOK pattern), and the don't-overwrite case when the caller resupplies fields like ``organization`` themselves. Co-Authored-By: Claude Opus 4.7 (1M context) * chore(vertex-batches): wrap api_base GET in safe_get for defense-in-depth The vertex batches status-poll fetches an attacker-influenceable ``api_base`` URL with a raw ``sync_handler.get()``. The proxy auth gate already validates clientside ``api_base`` before reaching this sink, so the proxy flow is covered. This adds the per-sink wrap so SDK callers and any future code path that bypasses the proxy gate pick up the same SSRF defense from ``url_utils.safe_get``. Operators with a legitimate private Vertex base can either allowlist the host via ``litellm.user_url_allowed_hosts`` or disable validation with ``litellm.user_url_validation = False``. Co-Authored-By: Claude Opus 4.7 (1M context) * refactor(auth): hoist url_utils import; derive admin-config field list from CredentialLiteLLMParams /simplify pass: - Move ``from litellm.litellm_core_utils.url_utils import SSRFError, validate_url`` to module top in ``proxy/auth/auth_utils.py``. CLAUDE.md prefers module-level imports unless avoiding a circular dependency, and there's no cycle here (``url_utils`` doesn't depend on ``proxy.auth``). - Replace the hardcoded ``_ADMIN_CONFIG_FIELDS_TO_CLEAR_ON_BASE_OVERRIDE`` literal with ``_admin_config_fields_to_clear_on_base_override()`` that derives the typed-field portion from ``CredentialLiteLLMParams.model_fields``. Adds three fields the hardcoded list missed (``aws_bedrock_runtime_endpoint``, ``watsonx_region_name``, ``region_name``) and stays in sync as new provider fields are declared on the model. The kwargs-only set (``organization``, ``extra_body``, ``azure_ad_token``, ``aws_session_token``, ``aws_sts_endpoint``, ``aws_web_identity_token``, ``aws_role_name``, …) remains explicit since those fields aren't on the typed model. Co-Authored-By: Claude Opus 4.7 (1M context) * fix(auth): close field-echo bypass; gate URL check on toggle; cover async batch path Three issues from review: 1. ``get_dynamic_litellm_params`` used ``if field not in request_kwargs: pop`` to clear admin-set provider config when the caller redirected ``api_base``. A caller could *echo* any clear-list field name (with any value, including an empty string) to skip the pop, leaving the admin's value in ``litellm_params`` to be forwarded to the redirected upstream. Fix: always pop, then write the caller's value back if they resupplied the field. 2. ``check_complete_credentials`` called ``validate_url`` directly. That helper doesn't itself consult ``litellm.user_url_validation``; the toggle is honoured by ``safe_get`` / ``async_safe_get``. Mirror that here so admins who explicitly disabled URL validation aren't blocked at the proxy boundary. 3. ``VertexAIBatchesHandler._async_retrieve_batch`` still used a bare ``await client.get(api_base, ...)`` while the sync sibling was wrapped in ``safe_get``. Wrap the async call in ``async_safe_get`` so SDK callers on the async path get the same DNS-rebind / private / cloud-metadata defenses as the sync path. Tests: - ``TestCheckCompleteCredentialsBlocksSSRF`` is now mock-only; an autouse fixture flips the toggle on, ``validate_url`` is patched in the parametrized blocking tests, and the positive path no longer makes a real DNS call to api.openai.com. - ``test_skips_url_validation_when_toggle_is_off`` documents the new toggle-off behaviour and asserts ``validate_url`` is not called. - ``test_caller_resupplied_value_overrides_admin_value_on_base_override`` replaces the prior test that asserted the buggy preserve-admin-value-on-echo behaviour. - ``test_field_echo_does_not_preserve_admin_value`` is a focused regression test for the empty-string echo vector. Co-Authored-By: Claude Opus 4.7 (1M context) * fix(auth): close provider-confusion credential exfil; expand banned-params; cover OCI Three additions on top of the entry-point URL gate so the cluster is fully closed against caller-supplied ``api_base`` redirection: 1. ``get_llm_provider_logic.py`` matched registered openai-compatible endpoints against ``api_base`` with an unanchored substring search (``if endpoint in api_base:``). A caller could pass an api_base like ``https://attacker.com/api.groq.com/openai/v1`` to coerce the proxy into reading ``GROQ_API_KEY`` from the environment and forwarding it as a Bearer credential to the attacker's host. Replaced with parsed- URL semantics (hostname exact-match plus segment-bounded path-prefix) in a new ``_endpoint_matches_api_base`` helper. 2. ``is_request_body_safe`` rejects ``api_base`` / ``base_url`` / ``user_config`` / a handful of AWS / vertex fields, but the list omitted three other endpoint-targeting fields: * ``aws_bedrock_runtime_endpoint`` — Bedrock endpoint redirect * ``langsmith_base_url`` / ``langfuse_host`` — observability callback hostnames; attacker-controlled values exfiltrate the entire request payload (incl. message content) via the logging hook. Added all three to the blocklist. 3. ``_admin_config_fields_to_clear_on_base_override`` derives its typed- field list from ``CredentialLiteLLMParams.model_fields``, which does not declare any of the OCI provider's auth fields. Added ``oci_signer``, ``oci_user``, ``oci_fingerprint``, ``oci_tenancy``, ``oci_key``, and ``oci_key_file`` to the kwargs-only fixed list so they are cleared on caller-redirected ``api_base`` like the AWS / Azure / Vertex equivalents. Tests: - ``TestEndpointMatchesApiBase`` — direct unit tests on the new matcher: legitimate provider URLs (5 shapes) match; attacker smuggling via path injection, suffix label, prefix label, userinfo ``@`` injection, and path-segment lookalikes (7 shapes) do not. - ``TestGetLlmProviderRejectsAttackerSmuggledApiBase`` — end-to-end invariant that ``GROQ_API_KEY`` is never read against an attacker- controlled host while the legitimate ``api.groq.com`` path still resolves the provider correctly. - ``TestIsRequestBodySafeBlocksEndpointTargetingFields`` — parametrized coverage that each of the three new banned-params raises a clear rejection naming the offending field. Co-Authored-By: Claude Opus 4.7 (1M context) * fix(auth): remove implicit api-key bypass + add posthog/braintrust/slack to blocklist The historical ``check_complete_credentials`` clause inside ``is_request_body_safe`` was a third, *implicit*, *caller-controlled* BYOK path: any caller that supplied a non-empty ``api_key`` caused the entire banned-params blocklist to be skipped. That turned every missing entry on the blocklist into an exploitable SSRF / credential-exfil hole and is the root cause of the chain of api_base advisories that have been re-discovered with each new integration: * GHSA-jh89-88fc-qrfp (critical, triage) — env-var exfil via api_base * GHSA-3frq-6r6h-7j64 (high, triage) — admin org / extra_body leak * veria-admin Dv_m860l, b_yRJeQ5, stN90yjP, LBlyOAc8, U2TD78kg — variations on "list X is missing field Y" Two explicit, admin-controlled BYOK paths already exist and remain: ``general_settings.allow_client_side_credentials = true`` (proxy-wide) and ``configurable_clientside_auth_params: [...]`` per deployment. Removing the implicit bypass converts the failure mode of a missing blocklist entry from "live credential leak" to "predictable 400 with a clear remediation message," which is the structural fix. Also adds the three remaining endpoint-targeting fields the dynamic callback layer reads from request body: ``posthog_host``, ``braintrust_host``, ``slack_webhook_url``. ``slack_webhook_url`` in particular was a direct exfil channel (caller-set webhook → proxy mirrors every request to attacker's Slack). Tests: - ``test_api_key_does_not_bypass_blocklist`` — parametrized regression asserting api_key=anything no longer skips the gate for any of the five highest-risk fields. - ``test_admin_opt_in_proxy_wide_still_allows`` — confirms the documented BYOK opt-in still works. - Extends ``test_endpoint_targeting_field_in_request_body_is_rejected`` to cover posthog / braintrust / slack. Co-Authored-By: Claude Opus 4.7 (1M context) * fix(auth): block sagemaker_base_url, s3_endpoint_url, deployment_url Provider-specific endpoint overrides surfaced by a wider audit of ``optional_params`` consumers in ``litellm/llms/``. Same threat as ``api_base``: a caller-supplied value redirects the outbound request to an attacker host. * ``s3_endpoint_url`` — read in ``litellm/llms/bedrock/files/transformation.py`` to build the S3 upload URL for Bedrock files. Caller redirects file uploads to attacker-controlled S3. * ``sagemaker_base_url`` — read in ``litellm/llms/sagemaker/{chat,completion}/*``. Caller redirects SageMaker traffic. This is the primary vector described in veria-admin mNqEBBtG. * ``deployment_url`` — popped in ``litellm/llms/sap/chat/transformation.py``. Caller redirects SAP deployment requests. Tests parametrize ``test_endpoint_targeting_field_in_request_body_is_rejected`` to cover the three new fields. Co-Authored-By: Claude Opus 4.7 (1M context) --------- Co-authored-by: Claude Opus 4.7 (1M context) --- .../get_llm_provider_logic.py | 40 ++- litellm/llms/vertex_ai/batches/handler.py | 20 +- litellm/proxy/auth/auth_utils.py | 75 ++++- .../clientside_credential_handler.py | 67 ++++ .../test_get_llm_provider_endpoint_match.py | 138 ++++++++ .../proxy/auth/test_auth_utils.py | 304 ++++++++++++++++++ 6 files changed, 629 insertions(+), 15 deletions(-) create mode 100644 tests/test_litellm/litellm_core_utils/test_get_llm_provider_endpoint_match.py diff --git a/litellm/litellm_core_utils/get_llm_provider_logic.py b/litellm/litellm_core_utils/get_llm_provider_logic.py index 4ff077efe7..c0ca6835ee 100644 --- a/litellm/litellm_core_utils/get_llm_provider_logic.py +++ b/litellm/litellm_core_utils/get_llm_provider_logic.py @@ -1,4 +1,5 @@ from typing import Optional, Tuple +from urllib.parse import urlparse import litellm from litellm.constants import REPLICATE_MODEL_NAME_WITH_ID_LENGTH @@ -8,6 +9,43 @@ from litellm.secret_managers.main import get_secret, get_secret_str from ..types.router import LiteLLM_Params +def _endpoint_matches_api_base(endpoint: str, api_base: str) -> bool: + """ + Match a registered openai-compatible endpoint against a caller-supplied + ``api_base`` using parsed-URL semantics, not unanchored substring search. + + Both inputs may be a bare hostname (``api.perplexity.ai``), host+path + (``api.deepinfra.com/v1/openai``), or a full URL + (``https://api.cerebras.ai/v1``). Hostnames must match exactly + (case-insensitive); if the registered endpoint has a non-trivial path, + the api_base path must start with it on a segment boundary. + + The naive ``endpoint in api_base`` shape lets a caller pass + ``https://attacker.com/api.groq.com/openai/v1`` to coerce the proxy + into reading the server's GROQ_API_KEY from the environment and + forwarding it to the attacker's host as a Bearer credential. + """ + + def _parse(value: str): + # Ensure urlparse sees a scheme so it populates hostname / path. + normalized = value if "://" in value else f"https://{value}" + return urlparse(normalized) + + parsed_endpoint = _parse(endpoint) + parsed_url = _parse(api_base) + + endpoint_host = (parsed_endpoint.hostname or "").lower() + url_host = (parsed_url.hostname or "").lower() + if not endpoint_host or endpoint_host != url_host: + return False + + endpoint_path = parsed_endpoint.path.rstrip("/") + if not endpoint_path: + return True + url_path = parsed_url.path.rstrip("/") + return url_path == endpoint_path or url_path.startswith(endpoint_path + "/") + + def _is_non_openai_azure_model(model: str) -> bool: try: model_name = model.split("/", 1)[1] @@ -210,7 +248,7 @@ def get_llm_provider( # noqa: PLR0915 # check if api base is a known openai compatible endpoint if api_base: for endpoint in litellm.openai_compatible_endpoints: - if endpoint in api_base: + if _endpoint_matches_api_base(endpoint, api_base): if endpoint == "api.perplexity.ai": custom_llm_provider = "perplexity" dynamic_api_key = get_secret_str("PERPLEXITYAI_API_KEY") diff --git a/litellm/llms/vertex_ai/batches/handler.py b/litellm/llms/vertex_ai/batches/handler.py index 028e02eb0c..7436bfef58 100644 --- a/litellm/llms/vertex_ai/batches/handler.py +++ b/litellm/llms/vertex_ai/batches/handler.py @@ -4,6 +4,7 @@ from typing import Any, Coroutine, Dict, Optional, Union import httpx import litellm +from litellm.litellm_core_utils.url_utils import async_safe_get, safe_get from litellm.llms.custom_httpx.http_handler import ( _get_httpx_client, get_async_httpx_client, @@ -224,8 +225,14 @@ class VertexAIBatchPrediction(VertexLLM): }, ) - response = sync_handler.get( - url=api_base, + # ``api_base`` here can come from caller-supplied request kwargs + # (clientside override). Wrap the fetch in ``safe_get`` so DNS + # rebind / private / cloud-metadata targets are rejected; the + # proxy auth gate already blocks malicious clientside ``api_base`` + # at the boundary — this is defense-in-depth for SDK callers. + response = safe_get( + sync_handler, + api_base, headers=headers, ) @@ -270,8 +277,13 @@ class VertexAIBatchPrediction(VertexLLM): }, ) - response = await client.get( - url=api_base, + # Mirror the sync path: ``api_base`` may come from caller-supplied + # request kwargs, so wrap the fetch in ``async_safe_get`` to reject + # DNS-rebind / private / cloud-metadata targets. Defense-in-depth + # behind the proxy auth gate's clientside ``api_base`` check. + response = await async_safe_get( + client, + api_base, headers=headers, ) if response.status_code != 200: diff --git a/litellm/proxy/auth/auth_utils.py b/litellm/proxy/auth/auth_utils.py index 448c975d12..91c8f2dd7c 100644 --- a/litellm/proxy/auth/auth_utils.py +++ b/litellm/proxy/auth/auth_utils.py @@ -6,9 +6,11 @@ from typing import Any, List, Optional, Tuple from fastapi import HTTPException, Request, status +import litellm from litellm import Router, provider_list from litellm._logging import verbose_proxy_logger from litellm.constants import STANDARD_CUSTOMER_ID_HEADERS +from litellm.litellm_core_utils.url_utils import SSRFError, validate_url from litellm.proxy._types import * from litellm.types.router import CONFIGURABLE_CLIENTSIDE_AUTH_PARAMS @@ -53,6 +55,12 @@ def _check_valid_ip( def check_complete_credentials(request_body: dict) -> bool: """ if 'api_base' in request body. Check if complete credentials given. Prevent malicious attacks. + + Supplying an ``api_key`` is necessary but not sufficient: even with + credentials supplied, an ``api_base`` / ``base_url`` that resolves to a + private/internal/cloud-metadata address would still allow the proxy to + be used as an SSRF pivot. Validate any URL fields here so the gate + can't be bypassed with ``api_key=anything`` plus a malicious target. """ given_model: Optional[str] = None @@ -70,10 +78,27 @@ def check_complete_credentials(request_body: dict) -> bool: return False api_key_value = request_body.get("api_key") - if api_key_value and isinstance(api_key_value, str) and api_key_value.strip(): - return True + if not (api_key_value and isinstance(api_key_value, str) and api_key_value.strip()): + return False - return False + # ``validate_url`` itself doesn't consult the toggle; ``safe_get`` / + # ``async_safe_get`` do. Mirror that here so admins who explicitly + # disabled URL validation (e.g. for an internal Ollama endpoint they + # accept the SSRF risk for) aren't blocked at the proxy boundary. + if getattr(litellm, "user_url_validation", False): + for url_field in ("api_base", "base_url"): + url_value = request_body.get(url_field) + if not url_value or not isinstance(url_value, str): + continue + try: + validate_url(url_value) + except SSRFError as e: + raise ValueError( + f"Rejected request: client-side {url_field}={url_value!r} " + f"is rejected by the SSRF guard ({e})." + ) + + return True def check_regex_or_str_match(request_body_value: Any, regex_str: str) -> bool: @@ -159,15 +184,42 @@ def is_request_body_safe( "aws_web_identity_token", "aws_role_name", "vertex_credentials", + # Endpoint-targeting fields that retarget the outbound request or + # an observability callback. An attacker-controlled value either + # exfiltrates the request payload (incl. messages + admin-set + # tokens) to the attacker's host, or coerces the proxy into + # authenticating against the attacker's host with admin secrets. + "aws_bedrock_runtime_endpoint", + "langsmith_base_url", + "langfuse_host", + "posthog_host", + "braintrust_host", + "slack_webhook_url", + # Provider-specific endpoint overrides that flow into the outbound + # request via ``optional_params``. Same threat as ``api_base``: + # ``s3_endpoint_url`` redirects Bedrock file uploads to attacker + # S3; ``sagemaker_base_url`` redirects all SageMaker traffic; + # ``deployment_url`` redirects SAP deployments. + "s3_endpoint_url", + "sagemaker_base_url", + "deployment_url", ] + # The blocklist is enforced unconditionally. Legitimate clientside + # credential / endpoint passthrough goes through one of the two + # explicit admin opt-ins (``general_settings.allow_client_side_credentials`` + # proxy-wide or ``configurable_clientside_auth_params`` per deployment). + # Historically there was a third, *implicit*, *caller-controlled* path: + # ``check_complete_credentials`` returned True when the caller supplied + # any non-empty ``api_key``, which made the entire blocklist a no-op. + # That bypass turned every missing entry on the blocklist into an + # exploitable SSRF / credential-exfil hole — see GHSA-jh89-88fc-qrfp, + # GHSA-3frq-6r6h-7j64, and the chain of veria-admin findings (Dv_m860l, + # b_yRJeQ5, stN90yjP, LBlyOAc8, U2TD78kg). Removed: the blocklist now + # has a single, predictable failure mode for missing entries (a 400), + # not a credential leak. for param in banned_params: - if ( - param in request_body - and not check_complete_credentials( # allow client-credentials to be passed to proxy - request_body=request_body - ) - ): + if param in request_body: if general_settings.get("allow_client_side_credentials") is True: return True elif ( @@ -182,7 +234,10 @@ def is_request_body_safe( return True raise ValueError( f"Rejected Request: {param} is not allowed in request body. " - "Enable with `general_settings::allow_client_side_credentials` on proxy config.yaml. " + "Clientside passthrough requires explicit admin opt-in via " + "either `general_settings.allow_client_side_credentials = true` " + "(proxy-wide) or `configurable_clientside_auth_params` on the " + "deployment in your proxy config.yaml. " "Relevant Issue: https://huntr.com/bounties/4001e1a2-7b7a-4776-a3ae-e6692ec3d997", ) diff --git a/litellm/router_utils/clientside_credential_handler.py b/litellm/router_utils/clientside_credential_handler.py index c98f614335..45ade81b2d 100644 --- a/litellm/router_utils/clientside_credential_handler.py +++ b/litellm/router_utils/clientside_credential_handler.py @@ -11,9 +11,60 @@ If given, generate a unique model_id for the deployment. Ensures cooldowns are applied correctly. """ +from typing import List + clientside_credential_keys = ["api_key", "api_base", "base_url"] +def _admin_config_fields_to_clear_on_base_override() -> List[str]: + """ + Provider-specific credential / endpoint-targeting fields that must NOT + flow through to a client-redirected upstream. + + Built dynamically from ``CredentialLiteLLMParams.model_fields`` so any + new provider field added there (Bedrock endpoint, Watsonx region, etc.) + is gated automatically — plus a fixed list of kwargs-only fields that + aren't declared on the typed model. + """ + from litellm.types.router import CredentialLiteLLMParams + + typed_fields = [ + f + for f in CredentialLiteLLMParams.model_fields + if f not in clientside_credential_keys + ] + kwargs_only_fields = [ + # Caller-supplied via **kwargs, not declared on CredentialLiteLLMParams. + "organization", + "extra_body", + "extra_headers", + "default_headers", + "api_type", + "azure_ad_token", + "azure_ad_token_provider", + "aws_session_token", + "aws_sts_endpoint", + "aws_web_identity_token", + "aws_role_name", + # OCI provider — consumed by litellm/llms/oci/* via optional_params + # and not declared on CredentialLiteLLMParams. Without these here, + # an admin's OCI signing key / tenancy / fingerprint would flow + # through to an attacker-redirected upstream. + "oci_signer", + "oci_user", + "oci_fingerprint", + "oci_tenancy", + "oci_key", + "oci_key_file", + ] + return typed_fields + kwargs_only_fields + + +_ADMIN_CONFIG_FIELDS_TO_CLEAR_ON_BASE_OVERRIDE = ( + _admin_config_fields_to_clear_on_base_override() +) + + def is_clientside_credential(request_kwargs: dict) -> bool: """ Check if the credential is a clientside credential. @@ -34,4 +85,20 @@ def get_dynamic_litellm_params(litellm_params: dict, request_kwargs: dict) -> di for key in clientside_credential_keys: if key in request_kwargs: litellm_params[key] = request_kwargs[key] + + # If the caller redirected api_base/base_url to a client-controlled value, + # don't forward the admin's organization / extra_body / region / token / + # vertex / aws fields — those were meant for the original upstream. + # Always drop the admin's value first, then write the caller's value back + # if they resupplied the field. The naive + # ``if field not in request_kwargs: pop`` shape lets a caller *echo* a + # field name (with any value, including an empty string) to keep the + # admin's value in ``litellm_params`` and have it forwarded to the + # redirected upstream. + if "api_base" in request_kwargs or "base_url" in request_kwargs: + for field in _ADMIN_CONFIG_FIELDS_TO_CLEAR_ON_BASE_OVERRIDE: + litellm_params.pop(field, None) + if field in request_kwargs: + litellm_params[field] = request_kwargs[field] + return litellm_params diff --git a/tests/test_litellm/litellm_core_utils/test_get_llm_provider_endpoint_match.py b/tests/test_litellm/litellm_core_utils/test_get_llm_provider_endpoint_match.py new file mode 100644 index 0000000000..fc5b39a2fd --- /dev/null +++ b/tests/test_litellm/litellm_core_utils/test_get_llm_provider_endpoint_match.py @@ -0,0 +1,138 @@ +""" +Regression tests for the parsed-URL hostname match used to identify a +caller-supplied ``api_base`` as a known openai-compatible provider. + +The previous shape (``if endpoint in api_base:``) used unanchored +substring search, which let a caller pass +``https://attacker.com/api.groq.com/openai/v1`` and have the proxy +return ``GROQ_API_KEY`` as the dynamic credential — exfiltrating the +server's real provider key to an attacker-controlled host on the +outbound request. +""" + +import os +import sys +from unittest.mock import patch + +import pytest + +sys.path.insert(0, os.path.abspath("../../..")) + +from litellm.litellm_core_utils.get_llm_provider_logic import ( + _endpoint_matches_api_base, + get_llm_provider, +) + + +class TestEndpointMatchesApiBase: + """Direct unit tests on the parsed-URL matcher.""" + + @pytest.mark.parametrize( + "endpoint, api_base", + [ + # Bare hostname endpoint, exact host match. + ("api.perplexity.ai", "https://api.perplexity.ai/v1"), + # Endpoint includes a path; api_base path starts with it. + ("api.groq.com/openai/v1", "https://api.groq.com/openai/v1"), + # Endpoint with full URL scheme. + ("https://api.cerebras.ai/v1", "https://api.cerebras.ai/v1/chat"), + # Trailing-slash on registered endpoint must not break match. + ("https://llm.chutes.ai/v1/", "https://llm.chutes.ai/v1/chat"), + # Case-insensitive on hostname. + ("api.groq.com/openai/v1", "https://API.GROQ.COM/openai/v1"), + ], + ) + def test_legitimate_provider_urls_match(self, endpoint, api_base): + assert _endpoint_matches_api_base(endpoint, api_base) is True + + @pytest.mark.parametrize( + "endpoint, api_base", + [ + # Attacker host, registered endpoint smuggled into path. + ( + "api.groq.com/openai/v1", + "https://attacker.com/api.groq.com/openai/v1", + ), + # Attacker host, registered endpoint smuggled into a path segment. + ( + "api.groq.com/openai/v1", + "https://attacker.com/foo/api.groq.com/openai/v1", + ), + # Lookalike host that contains the registered host as a suffix label. + ( + "api.groq.com/openai/v1", + "https://api.groq.com.attacker.com/openai/v1", + ), + # Lookalike host with the registered host as a prefix. + ( + "api.groq.com/openai/v1", + "https://api.groq.com.evil.example/openai/v1", + ), + # Right host, wrong path — endpoint requires ``/openai/v1`` prefix. + ("api.groq.com/openai/v1", "https://api.groq.com/v1"), + # Path-segment lookalike: ``/openai/v10`` must not match ``/openai/v1``. + ("api.groq.com/openai/v1", "https://api.groq.com/openai/v10"), + # Userinfo / @-injection trick — the ``hostname`` after ``@`` is + # what httpx connects to. + ( + "api.groq.com/openai/v1", + "https://api.groq.com@attacker.com/openai/v1", + ), + ], + ) + def test_attacker_smuggling_does_not_match(self, endpoint, api_base): + assert _endpoint_matches_api_base(endpoint, api_base) is False + + +class TestGetLlmProviderRejectsAttackerSmuggledApiBase: + """ + End-to-end: ``get_llm_provider`` must NOT return the server's stored + secret (e.g. ``GROQ_API_KEY``) for an api_base whose hostname is + attacker-controlled, even when the registered endpoint string appears + elsewhere in the URL. + """ + + def test_attacker_host_does_not_yield_groq_secret(self): + # The function may either fall through (different provider) or + # raise BadRequestError because the model can't be identified. + # The invariant under test is that ``GROQ_API_KEY`` is never + # looked up against an attacker-controlled hostname. + import litellm + + with patch( + "litellm.litellm_core_utils.get_llm_provider_logic.get_secret_str", + return_value="server-real-groq-key", + ) as mocked_secret: + try: + _, _, dynamic_api_key, _ = get_llm_provider( + model="some-model", + api_base="https://attacker.com/api.groq.com/openai/v1", + ) + # If it returned, the dynamic key must not be the secret. + assert dynamic_api_key != "server-real-groq-key" + except litellm.exceptions.BadRequestError: + # Acceptable outcome: provider unidentifiable, no secret + # was returned. + pass + + # Regardless of return / raise, the secret must never have been + # read against this attacker-controlled api_base. + groq_lookups = [ + call + for call in mocked_secret.call_args_list + if call.args and call.args[0] == "GROQ_API_KEY" + ] + assert groq_lookups == [] + + def test_legitimate_groq_api_base_still_resolves(self): + with patch( + "litellm.litellm_core_utils.get_llm_provider_logic.get_secret_str", + return_value="server-real-groq-key", + ): + _, provider, dynamic_api_key, _ = get_llm_provider( + model="some-model", + api_base="https://api.groq.com/openai/v1", + ) + + assert provider == "groq" + assert dynamic_api_key == "server-real-groq-key" diff --git a/tests/test_litellm/proxy/auth/test_auth_utils.py b/tests/test_litellm/proxy/auth/test_auth_utils.py index 15cfc84c23..91f300b88c 100644 --- a/tests/test_litellm/proxy/auth/test_auth_utils.py +++ b/tests/test_litellm/proxy/auth/test_auth_utils.py @@ -5,6 +5,8 @@ Unit tests for auth_utils functions related to rate limiting and customer ID ext from typing import Optional from unittest.mock import MagicMock, patch +import pytest + from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.auth.auth_utils import ( _get_customer_id_from_standard_headers, @@ -15,6 +17,7 @@ from litellm.proxy.auth.auth_utils import ( get_key_model_tpm_limit, get_project_model_rpm_limit, get_project_model_tpm_limit, + is_request_body_safe, ) @@ -660,3 +663,304 @@ class TestCheckCompleteCredentials: def test_returns_true_when_api_key_is_valid(self): result = check_complete_credentials({"model": "gpt-4", "api_key": "sk-valid"}) assert result is True + + +class TestCheckCompleteCredentialsBlocksSSRF: + """ + Even with credentials supplied, ``api_base`` / ``base_url`` must not + point at private / internal / cloud-metadata addresses. Without this + the gate accepts ``api_key=anything`` plus a malicious target and the + proxy is used as an SSRF pivot. + + The check only runs when ``litellm.user_url_validation`` is True, so + every test in this class flips the toggle. Tests stay mock-only — no + real DNS is performed. + """ + + @pytest.fixture(autouse=True) + def _enable_url_validation(self, monkeypatch): + import litellm + + monkeypatch.setattr(litellm, "user_url_validation", True, raising=False) + + @pytest.mark.parametrize( + "url_field", + ["api_base", "base_url"], + ) + @pytest.mark.parametrize( + "blocked_url", + [ + "http://169.254.169.254/latest/meta-data/iam/security-credentials/", + "http://metadata.google.internal/computeMetadata/v1/", + "http://127.0.0.1:8080/admin", + "http://10.0.0.1/", + "http://192.168.1.1/", + ], + ) + def test_rejects_private_or_metadata_targets(self, url_field, blocked_url): + from litellm.litellm_core_utils.url_utils import SSRFError + + with patch( + "litellm.proxy.auth.auth_utils.validate_url", + side_effect=SSRFError(f"blocked: {blocked_url}"), + ): + with pytest.raises(ValueError) as exc_info: + check_complete_credentials( + { + "model": "gpt-4", + "api_key": "sk-some-clientside-key", + url_field: blocked_url, + } + ) + assert url_field in str(exc_info.value) + assert "SSRF" in str(exc_info.value) + + def test_allows_public_target_when_validate_url_passes(self): + # ``validate_url`` is mocked so no real DNS is performed. + with patch( + "litellm.proxy.auth.auth_utils.validate_url", + return_value=("https://api.openai.com/v1", "api.openai.com"), + ): + result = check_complete_credentials( + { + "model": "gpt-4", + "api_key": "sk-some-clientside-key", + "api_base": "https://api.openai.com/v1", + } + ) + assert result is True + + def test_skips_url_validation_when_toggle_is_off(self, monkeypatch): + # Admins who disable ``user_url_validation`` (default) should not + # have requests rejected at the proxy boundary even if the URL + # would fail the SSRF guard. + import litellm + + monkeypatch.setattr(litellm, "user_url_validation", False, raising=False) + with patch( + "litellm.proxy.auth.auth_utils.validate_url", + ) as mocked: + result = check_complete_credentials( + { + "model": "gpt-4", + "api_key": "sk-some-clientside-key", + "api_base": "http://127.0.0.1:8080/admin", + } + ) + assert result is True + mocked.assert_not_called() + + +class TestGetDynamicLitellmParamsClearsAdminConfigOnBaseOverride: + """ + When the caller redirects ``api_base`` / ``base_url`` to their own + server, admin-set fields like ``OpenAI-Organization``, ``extra_body``, + AWS / Vertex / Azure tokens, and per-deployment ``api_version`` must + NOT flow through to that destination. + """ + + def test_clears_admin_organization_and_extra_body_on_base_override(self): + from litellm.router_utils.clientside_credential_handler import ( + get_dynamic_litellm_params, + ) + + admin_params = { + "model": "gpt-4", + "api_key": "sk-admin-key", + "api_base": "https://admin.upstream/v1", + "organization": "org-admin-corp", + "extra_body": {"x-admin-secret": "super-secret"}, + "api_version": "2026-04-01", + } + out = get_dynamic_litellm_params( + litellm_params=dict(admin_params), + request_kwargs={ + "api_key": "sk-attacker", + "api_base": "https://attacker.example", + }, + ) + assert out["api_base"] == "https://attacker.example" + assert out["api_key"] == "sk-attacker" + assert "organization" not in out + assert "extra_body" not in out + assert "api_version" not in out + + def test_clears_aws_and_vertex_secrets_on_base_override(self): + from litellm.router_utils.clientside_credential_handler import ( + get_dynamic_litellm_params, + ) + + admin_params = { + "model": "bedrock/claude-3", + "aws_access_key_id": "AKIA-EXAMPLE", + "aws_secret_access_key": "secret-example", + "aws_session_token": "session-example", + "vertex_credentials": '{"private_key":"-----BEGIN..."}', + "vertex_project": "admin-gcp-project", + } + out = get_dynamic_litellm_params( + litellm_params=dict(admin_params), + request_kwargs={"base_url": "https://attacker.example"}, + ) + assert "aws_access_key_id" not in out + assert "aws_secret_access_key" not in out + assert "aws_session_token" not in out + assert "vertex_credentials" not in out + assert "vertex_project" not in out + + def test_caller_resupplied_value_overrides_admin_value_on_base_override(self): + # When the caller redirects ``api_base`` and *also* supplies their + # own value for one of the admin fields (e.g. ``organization``), + # the caller's value must win — never the admin's. The naive + # ``if field not in request_kwargs: pop`` shape lets a caller echo + # the field name with any value (or empty string) to keep the + # admin's value forwarded, which is the exfiltration vector this + # test guards against. + from litellm.router_utils.clientside_credential_handler import ( + get_dynamic_litellm_params, + ) + + out = get_dynamic_litellm_params( + litellm_params={ + "api_base": "https://admin.upstream/v1", + "organization": "org-admin", + "extra_body": {"admin": "value"}, + }, + request_kwargs={ + "api_base": "https://attacker.example", + "organization": "org-attacker", + "extra_body": {"attacker": "value"}, + }, + ) + assert out["organization"] == "org-attacker" + assert out["extra_body"] == {"attacker": "value"} + + def test_field_echo_does_not_preserve_admin_value(self): + # Regression: a caller that echoes an admin-config field name with + # an *empty* value (or any value) must not be able to keep the + # admin's value in ``litellm_params``. + from litellm.router_utils.clientside_credential_handler import ( + get_dynamic_litellm_params, + ) + + out = get_dynamic_litellm_params( + litellm_params={ + "api_base": "https://admin.upstream/v1", + "organization": "org-admin-secret", + "extra_body": {"x-admin-only": "secret"}, + }, + request_kwargs={ + "api_base": "https://attacker.example", + "organization": "", + "extra_body": "", + }, + ) + assert out["organization"] == "" + assert out["extra_body"] == "" + assert "org-admin-secret" not in str(out) + + def test_no_clearing_when_only_api_key_overridden(self): + from litellm.router_utils.clientside_credential_handler import ( + get_dynamic_litellm_params, + ) + + # Caller only overrides api_key (BYOK pattern); admin's organization / + # extra_body / region still apply because the destination is unchanged. + out = get_dynamic_litellm_params( + litellm_params={ + "api_base": "https://admin.upstream/v1", + "organization": "org-admin", + "api_version": "2026-04-01", + }, + request_kwargs={"api_key": "sk-byok"}, + ) + assert out["organization"] == "org-admin" + assert out["api_version"] == "2026-04-01" + assert out["api_base"] == "https://admin.upstream/v1" + + +class TestIsRequestBodySafeBlocksEndpointTargetingFields: + """ + ``is_request_body_safe`` rejects request-body fields that retarget the + outbound request to a caller-controlled host. Beyond the original + ``api_base`` / ``base_url``, the same protection must apply to: + + * ``aws_bedrock_runtime_endpoint`` — Bedrock endpoint redirect; an + attacker-controlled value coerces the proxy to authenticate against + their host with the admin's AWS creds. + * ``langsmith_base_url`` — Langsmith callback host; attacker-controlled + values exfiltrate the entire request payload (incl. message content) + via the observability hook. + * ``langfuse_host`` — same exfil vector via the Langfuse hook. + """ + + @pytest.fixture(autouse=True) + def _disable_url_validation(self, monkeypatch): + # The new banned-params entries should be rejected even when + # ``user_url_validation`` is off — the gate isn't the URL guard, + # it's the banned-params list. + import litellm + + monkeypatch.setattr(litellm, "user_url_validation", False, raising=False) + + @pytest.mark.parametrize( + "field", + [ + "aws_bedrock_runtime_endpoint", + "langsmith_base_url", + "langfuse_host", + "posthog_host", + "braintrust_host", + "slack_webhook_url", + "s3_endpoint_url", + "sagemaker_base_url", + "deployment_url", + ], + ) + def test_endpoint_targeting_field_in_request_body_is_rejected(self, field): + with pytest.raises(ValueError) as exc: + is_request_body_safe( + request_body={"model": "gpt-4", field: "https://attacker.example"}, + general_settings={}, + llm_router=None, + model="gpt-4", + ) + # The function lists the offending param name in the error. + assert field in str(exc.value) + + @pytest.mark.parametrize( + "field", + ["api_base", "base_url", "user_config", "langfuse_host", "slack_webhook_url"], + ) + def test_api_key_does_not_bypass_blocklist(self, field): + # Regression: the historical ``check_complete_credentials`` clause + # made the entire blocklist a no-op for any caller that supplied + # a non-empty ``api_key``. That bypass turned every missing entry + # on the blocklist into an SSRF / credential-exfil hole. Verify + # that supplying an api_key (alongside the banned param) does NOT + # bypass the gate — it can only be opened by an admin opt-in. + with pytest.raises(ValueError) as exc: + is_request_body_safe( + request_body={ + "model": "gpt-4", + "api_key": "sk-anything", + field: "https://attacker.example", + }, + general_settings={}, + llm_router=None, + model="gpt-4", + ) + assert field in str(exc.value) + + def test_admin_opt_in_proxy_wide_still_allows(self): + # ``general_settings.allow_client_side_credentials = True`` remains + # the documented proxy-wide BYOK opt-in. + assert ( + is_request_body_safe( + request_body={"model": "gpt-4", "api_base": "https://my-byok.example"}, + general_settings={"allow_client_side_credentials": True}, + llm_router=None, + model="gpt-4", + ) + is True + )