From fbcdacc44659813b05dd2f9d145ec63575e2843d Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Tue, 21 Apr 2026 23:29:17 -0700 Subject: [PATCH 01/26] [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 73735796eb3..a8942f92a1c 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 f21a729f551..48931a0c284 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 786167b9486..0d241f75749 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 fb215e54777..57693a1a175 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 3e4f9af9555df6fef78e22624778e655d3067173 Mon Sep 17 00:00:00 2001 From: Sameer Kankute Date: Mon, 27 Apr 2026 12:05:49 +0530 Subject: [PATCH 02/26] 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 251f271903b..6aaf37c91ec 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 447c28078ec..8614cb3ba72 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 03/26] 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 6aaf37c91ec..edc9c87cdaf 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 cfe4bc678e612bcc37831b69582e886bf05e03da Mon Sep 17 00:00:00 2001 From: Sameer Kankute Date: Tue, 28 Apr 2026 15:37:23 +0530 Subject: [PATCH 04/26] 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 b62acb65166..d2c8206ca9a 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 5fbf0a4b19f..49d2f72db7c 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 4da0a7c7791..5b0b64f2429 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 b9ada079f6b..99f748c0c1a 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 e6e8369643e..b6cace066c8 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 fcf5d14db7c..8c08b783387 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 c763ed1c8da..b6eae390d93 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 ba87a8f2b01..8261036cae0 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 ed5397eef0c..ae28222c3ce 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 19b59769863..0cf86358873 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 4baa5774c48..b3fcf4b394c 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 179bd7aeff1..47dac8dca32 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 bcd4c620055..5fa48703b12 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 05/26] 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 5b0b64f2429..4e81fa4e66e 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 5fa48703b12..c211a3536e3 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 898040fcdda0f3193862ab495905198e26f5c45f Mon Sep 17 00:00:00 2001 From: Sameer Kankute Date: Tue, 28 Apr 2026 22:34:14 +0530 Subject: [PATCH 06/26] 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 7c06823d167..9507ff401aa 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 b07e1c03418312c4bbc617f611a75f64d64a64ec Mon Sep 17 00:00:00 2001 From: Michael Riad Zaky Date: Tue, 28 Apr 2026 17:38:42 -0700 Subject: [PATCH 07/26] 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 a27153365d2..61a7d4c08db 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 a90e05919fb..474ddb402a1 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 8e53e57f1e0..633aed38a08 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 0118b39ddd1..353d19b0198 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 fc49c181bca73d63d3b861d7fc2a46a834a3ba8e Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Wed, 29 Apr 2026 01:41:24 +0000 Subject: [PATCH 08/26] 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 251f271903b..a2f0517f81a 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 c2e998f01e5..3924687a0b1 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 146ba10bb76..42910cc790d 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 00000000000..aef6546c81a --- /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 09/26] 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 13b76c94dfa..ec2651306f2 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 68ab397d827..c0aec1687ef 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 10/26] 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 c0aec1687ef..39d078267f6 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 11/26] 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 a2f0517f81a..8e473c1cb21 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 12/26] 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 d2c8206ca9a..d1b93c9e7a3 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 49d2f72db7c..85a9c838264 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 4e81fa4e66e..f028503c6a2 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 99f748c0c1a..d8515d2c4d6 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 b6cace066c8..35d83bd2adc 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 8c08b783387..af78cd8dbda 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 b6eae390d93..2c11d137480 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 8261036cae0..7b22edd8676 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 ae28222c3ce..3238d3e9c14 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 0cf86358873..8270e99d456 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 b3fcf4b394c..d31e1f6c8f2 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 47dac8dca32..6cb7a86bea2 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 abe96b2ea2e..3e8d59b2992 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 c211a3536e3..d60d0487d06 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 9507ff401aa..7085e45cdc3 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 4ca38233129..cb4cfd75c1f 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 13/26] 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 95bcd4d7186..4ff077efe7c 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 888999504fe..59d0465e6d4 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 fcdf49f2916..a9cf151464b 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 1f5f65f693a..2d46592e3fb 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 8bca75172fa..d009a085fab 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 c106d7f17b6..68f799e5747 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 8511d785fb7..e4268fac81a 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 e63bf402bf8..027c9fedced 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 114883f530b..ca7d323ad6c 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 620c0734980..6644b1389cf 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 b8a4220c679..f28fe3ed258 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 14/26] 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 275c352b39e..5dd1247001e 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 6f23c87f911..ed49c146210 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 15/26] 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 edc9c87cdaf..eb0bda3fec2 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 16/26] 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 8e473c1cb21..853208a5382 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 42910cc790d..a0c278b906a 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 ace8c8a4188..8f8673b0a7d 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 aef6546c81a..cdaecfb2271 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 17/26] 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 cf731d66283..14c1de2f6a7 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 18/26] 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 853208a5382..8bb7a5d50d5 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 19/26] 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 3924687a0b1..ae6055217b8 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 a0c278b906a..df5705c3425 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 cdaecfb2271..bf90e9ebef6 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 20/26] 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 cfa90a48eee..f0967732006 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 3a1184c434e..93cc8bfd22b 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 00000000000..22ccb02b7a2 --- /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 fb215e54777..af65cad09e3 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 21/26] 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 f0967732006..ab9d341aa51 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 22ccb02b7a2..ae0e1f845b0 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 22/26] 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 a8942f92a1c..d112e222307 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 23/26] 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 fc9f5730cdf..a0e99dd16b7 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 03d2af72329..dd955c23d23 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 c03a63f211a..efa9cdb8e1b 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 00000000000..5515e6ce815 --- /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 24/26] 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 6ca94ee94f5..00000000000 --- 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 8ee24df0744..00000000000 --- 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 8ee24df0744..00000000000 --- 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 635c8e398b9..00000000000 --- 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 d026a480367..00000000000 --- 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 afb8782562e..00000000000 --- 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 5d19f680466..00000000000 --- 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 12f813af359..00000000000 --- 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 25/26] 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 8cb2ee09370..3e97b480779 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 b9ada079f6b..6b9a2c6a5df 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 c8fab4be4ae..76c52f83ee4 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 9768d93e922..6ada8f58783 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 6924eb8d3d9..752b5ff0905 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 a635c16e98d..d4b4617730b 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 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 26/26] 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 00000000000..6454f267656 --- /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 8f07c5afa3f..fd54ed5243a 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 00000000000..a19af4dd484 --- /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 870ea78f17a..12229955299 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 8f07c5afa3f..fd54ed5243a 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 00000000000..f452066afb0 --- /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 8f07c5afa3f..fd54ed5243a 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 00000000000..a337ff6d888 --- /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 d32cf74d6ad..a8553d5405b 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 c340b65496c..cecc99739b1 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 63f03f4586c..277bb0bd021 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 00000000000..a69eac282bf --- /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;