From ee9f0db1cff3d3694a9e246351be785366c0058a Mon Sep 17 00:00:00 2001 From: shivam Date: Thu, 23 Jul 2026 23:37:40 +0000 Subject: [PATCH 01/29] fix(bedrock-mantle): backfill usage on non-streaming /v1/messages responses Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../bedrock/messages/mantle_transformation.py | 23 +++++ .../test_litellm/llms/bedrock/test_mantle.py | 97 +++++++++++++++++++ 2 files changed, 120 insertions(+) diff --git a/litellm/llms/bedrock/messages/mantle_transformation.py b/litellm/llms/bedrock/messages/mantle_transformation.py index da7b8697a6b..65a2ab3b9a7 100644 --- a/litellm/llms/bedrock/messages/mantle_transformation.py +++ b/litellm/llms/bedrock/messages/mantle_transformation.py @@ -17,6 +17,10 @@ from litellm.llms.bedrock.common_utils import build_mantle_messages_url from litellm.llms.bedrock.messages.invoke_transformations.anthropic_claude3_transformation import ( AmazonAnthropicClaudeMessagesConfig, ) +from litellm.types.llms.anthropic_messages.anthropic_response import ( + AnthropicMessagesResponse, + AnthropicUsage, +) from litellm.types.router import GenericLiteLLMParams if TYPE_CHECKING: @@ -103,6 +107,25 @@ class AmazonMantleMessagesConfig(AmazonAnthropicClaudeMessagesConfig): ) return {**request, "model": model_id, **stream_fields} + def transform_anthropic_messages_response( + self, + model: str, + raw_response: httpx.Response, + logging_obj: LiteLLMLoggingObj, + ) -> AnthropicMessagesResponse: + response = super().transform_anthropic_messages_response( + model=model, + raw_response=raw_response, + logging_obj=logging_obj, + ) + existing_usage: AnthropicUsage = response.get("usage") or AnthropicUsage() + normalized_usage: AnthropicUsage = { + "input_tokens": 0, + "output_tokens": 0, + **existing_usage, + } + return {**response, "usage": normalized_usage} + def get_async_streaming_response_iterator( self, model: str, diff --git a/tests/test_litellm/llms/bedrock/test_mantle.py b/tests/test_litellm/llms/bedrock/test_mantle.py index c5ce0d5aba7..d34517f61f6 100644 --- a/tests/test_litellm/llms/bedrock/test_mantle.py +++ b/tests/test_litellm/llms/bedrock/test_mantle.py @@ -399,6 +399,103 @@ async def test_mantle_anthropic_messages_sends_workspace_header_and_clean_body() assert "aws_bedrock_project_id" not in requests[0]["body"] +def _usageless_anthropic_response(url: str) -> httpx.Response: + return httpx.Response( + status_code=200, + json={ + "id": "msg_classifier", + "type": "message", + "role": "assistant", + "model": "anthropic.claude-opus-4-8", + "content": [{"type": "text", "text": "safe"}], + "stop_reason": "end_turn", + "stop_sequence": None, + }, + request=httpx.Request("POST", url), + ) + + +@pytest.mark.asyncio +async def test_mantle_anthropic_messages_backfills_missing_usage(): + """ + Regression for LIT-4758: a Mantle non-streaming response with no `usage` + object must not reach the client usage-less, or Claude Code's auto-mode + classifier crashes on `usage.input_tokens`. + """ + import litellm + + async def mock_post(self, url, data=None, headers=None, **kwargs): + return _usageless_anthropic_response(str(url)) + + try: + with patch( + "litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post", + new=mock_post, + ): + response = await litellm.anthropic_messages( + model="bedrock/mantle/anthropic.claude-opus-4-8", + messages=[{"role": "user", "content": "is `Bash(ls)` safe?"}], + max_tokens=10, + aws_access_key_id="fake-key", + aws_secret_access_key="fake-secret", + aws_region_name="us-east-1", + ) + finally: + await litellm.close_litellm_async_clients() + + assert response["usage"]["input_tokens"] == 0 + assert response["usage"]["output_tokens"] == 0 + + +@pytest.mark.asyncio +async def test_mantle_anthropic_messages_preserves_upstream_usage(): + """Backfill must not clobber a usage object the upstream did return.""" + import litellm + + def _response_with_usage(url: str) -> httpx.Response: + return httpx.Response( + status_code=200, + json={ + "id": "msg_test", + "type": "message", + "role": "assistant", + "model": "anthropic.claude-opus-4-8", + "content": [{"type": "text", "text": "ok"}], + "stop_reason": "end_turn", + "stop_sequence": None, + "usage": { + "input_tokens": 42, + "output_tokens": 7, + "cache_read_input_tokens": 5, + }, + }, + request=httpx.Request("POST", url), + ) + + async def mock_post(self, url, data=None, headers=None, **kwargs): + return _response_with_usage(str(url)) + + try: + with patch( + "litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post", + new=mock_post, + ): + response = await litellm.anthropic_messages( + model="bedrock/mantle/anthropic.claude-opus-4-8", + messages=[{"role": "user", "content": "hello"}], + max_tokens=10, + aws_access_key_id="fake-key", + aws_secret_access_key="fake-secret", + aws_region_name="us-east-1", + ) + finally: + await litellm.close_litellm_async_clients() + + assert response["usage"]["input_tokens"] == 42 + assert response["usage"]["output_tokens"] == 7 + assert response["usage"]["cache_read_input_tokens"] == 5 + + @pytest.mark.asyncio async def test_mantle_anthropic_messages_routes_to_vpc_api_base(): import litellm From a469bb79243524120fb32e75f4c103d83b72167d Mon Sep 17 00:00:00 2001 From: Tin Chi Lo Date: Thu, 23 Jul 2026 18:06:57 -0700 Subject: [PATCH 02/29] fix(ui): keep a key's MCP toolsets when saving an edit The key edit form seeded mcp_servers_and_groups from the key with only servers and accessGroups, but handleKeyUpdate writes mcp_toolsets from that same value, so every save posted an empty list and the backend merge applied it literally. A key granted a toolset lost the grant on any edit, including a budget change, and then got a 403 from /toolset//mcp Read toolsets in both places the form initializes from keyData, declare mcp_toolsets on KeyResponse.object_permission so a write-without-read is a type error, and carry toolsets through the create flow, which only looked at servers and accessGroups --- .../components/key_team_helpers/key_list.tsx | 1 + .../organisms/create_key_button.test.tsx | 30 ++++++++++++++++ .../organisms/create_key_button.tsx | 8 +++-- .../KeyInfoView.handleKeyUpdate.test.tsx | 22 ++++++++++++ .../templates/key_edit_view.test.tsx | 34 +++++++++++++++++++ .../components/templates/key_edit_view.tsx | 2 ++ 6 files changed, 95 insertions(+), 2 deletions(-) diff --git a/ui/litellm-dashboard/src/components/key_team_helpers/key_list.tsx b/ui/litellm-dashboard/src/components/key_team_helpers/key_list.tsx index e1c1fcb232c..d4bd6756df2 100644 --- a/ui/litellm-dashboard/src/components/key_team_helpers/key_list.tsx +++ b/ui/litellm-dashboard/src/components/key_team_helpers/key_list.tsx @@ -94,6 +94,7 @@ export interface KeyResponse { object_permission_id: string; mcp_servers: string[]; mcp_access_groups?: string[]; + mcp_toolsets?: string[]; mcp_tool_permissions?: Record; vector_stores: string[]; agents?: string[]; diff --git a/ui/litellm-dashboard/src/components/organisms/create_key_button.test.tsx b/ui/litellm-dashboard/src/components/organisms/create_key_button.test.tsx index f84b95b8d4d..2f72fa37d86 100644 --- a/ui/litellm-dashboard/src/components/organisms/create_key_button.test.tsx +++ b/ui/litellm-dashboard/src/components/organisms/create_key_button.test.tsx @@ -443,6 +443,36 @@ describe("CreateKey", () => { }); }); + it("should include mcp_toolsets in keyCreateCall payload when only toolsets are selected", async () => { + renderWithProviders(); + + act(() => { + fireEvent.click(screen.getByRole("button", { name: /create new key/i })); + }); + + await waitFor(() => { + expect(screen.getByRole("button", { name: /create key/i })).toBeInTheDocument(); + }); + + act(() => { + formMock.setFieldValue("key_alias", "Test Key"); + formMock.setFieldValue("allowed_mcp_servers_and_groups", { + servers: [], + accessGroups: [], + toolsets: ["ts-1"], + }); + }); + + act(() => { + fireEvent.click(screen.getByRole("button", { name: /create key/i })); + }); + + await waitFor(() => { + expect(mockKeyCreateCall).toHaveBeenCalled(); + }); + expect(mockKeyCreateCall.mock.calls[0][2].object_permission?.mcp_toolsets).toEqual(["ts-1"]); + }); + it("should prefill models when provided without team_id", async () => { renderWithProviders( = ({ team, teams, data, addKey, autoOp if ( formValues.allowed_mcp_servers_and_groups && (formValues.allowed_mcp_servers_and_groups.servers?.length > 0 || - formValues.allowed_mcp_servers_and_groups.accessGroups?.length > 0) + formValues.allowed_mcp_servers_and_groups.accessGroups?.length > 0 || + formValues.allowed_mcp_servers_and_groups.toolsets?.length > 0) ) { if (!formValues.object_permission) { formValues.object_permission = {}; } - const { servers, accessGroups } = formValues.allowed_mcp_servers_and_groups; + const { servers, accessGroups, toolsets } = formValues.allowed_mcp_servers_and_groups; if (servers && servers.length > 0) { formValues.object_permission.mcp_servers = servers; } if (accessGroups && accessGroups.length > 0) { formValues.object_permission.mcp_access_groups = accessGroups; } + if (toolsets && toolsets.length > 0) { + formValues.object_permission.mcp_toolsets = toolsets; + } // Remove the original field as it's now part of object_permission delete formValues.allowed_mcp_servers_and_groups; } diff --git a/ui/litellm-dashboard/src/components/templates/KeyInfoView.handleKeyUpdate.test.tsx b/ui/litellm-dashboard/src/components/templates/KeyInfoView.handleKeyUpdate.test.tsx index e58e24d0bf7..fa8ffcd7366 100644 --- a/ui/litellm-dashboard/src/components/templates/KeyInfoView.handleKeyUpdate.test.tsx +++ b/ui/litellm-dashboard/src/components/templates/KeyInfoView.handleKeyUpdate.test.tsx @@ -453,6 +453,28 @@ describe("KeyInfoView handleKeyUpdate guardrails guard", () => { }); }); +describe("KeyInfoView handleKeyUpdate mcp_toolsets", () => { + it("should forward the toolsets the edit form supplies into object_permission", async () => { + renderView(true); + + fireEvent.click(screen.getByText("Settings")); + fireEvent.click(screen.getByText("Edit Settings")); + (globalThis as any).__TEST_FORM_VALUES = { + token: "tok_123", + max_budget: 40000, + mcp_servers_and_groups: { servers: [], accessGroups: [], toolsets: ["ts-1"] }, + }; + + fireEvent.click(screen.getByText("Mock Submit")); + + await waitFor(() => expect(keyUpdateCallMock).toHaveBeenCalled()); + + const [, sentPayload] = keyUpdateCallMock.mock.calls[0]; + expect(sentPayload.object_permission.mcp_toolsets).toEqual(["ts-1"]); + expect(sentPayload.max_budget).toBe(40000); + }); +}); + describe("KeyInfoView handleKeyUpdate budget_duration", () => { it("should send a canonical budget_duration through unchanged", async () => { renderView(true); diff --git a/ui/litellm-dashboard/src/components/templates/key_edit_view.test.tsx b/ui/litellm-dashboard/src/components/templates/key_edit_view.test.tsx index 693c50374c1..756481487e7 100644 --- a/ui/litellm-dashboard/src/components/templates/key_edit_view.test.tsx +++ b/ui/litellm-dashboard/src/components/templates/key_edit_view.test.tsx @@ -631,6 +631,40 @@ describe("KeyEditView", () => { }); }); + it("should keep mcp_toolsets when saving an edit that does not touch the MCP selector", async () => { + const onSubmitMock = vi.fn().mockResolvedValue(undefined); + const keyDataWithToolset = { + ...MOCK_KEY_DATA, + object_permission: { + ...MOCK_KEY_DATA.object_permission!, + mcp_toolsets: ["ts-1"], + }, + }; + + renderWithProviders( + {}} + onSubmit={onSubmitMock} + accessToken="test-token" + userID="test-user" + userRole="admin" + premiumUser={false} + />, + ); + + await waitFor(() => { + expect(screen.getByText("Save Changes")).toBeInTheDocument(); + }); + + await userEvent.click(screen.getByRole("button", { name: /save changes/i })); + + await waitFor(() => { + expect(onSubmitMock).toHaveBeenCalled(); + }); + expect(onSubmitMock.mock.calls[0][0].mcp_servers_and_groups.toolsets).toEqual(["ts-1"]); + }); + it("should submit budget_limits: [] when the last budget window is deleted", async () => { const onSubmitMock = vi.fn().mockResolvedValue(undefined); const keyDataWithWindow = { diff --git a/ui/litellm-dashboard/src/components/templates/key_edit_view.tsx b/ui/litellm-dashboard/src/components/templates/key_edit_view.tsx index 23b868d2b87..a9fa05d817d 100644 --- a/ui/litellm-dashboard/src/components/templates/key_edit_view.tsx +++ b/ui/litellm-dashboard/src/components/templates/key_edit_view.tsx @@ -200,6 +200,7 @@ export function KeyEditView({ mcp_servers_and_groups: { servers: keyData.object_permission?.mcp_servers || [], accessGroups: keyData.object_permission?.mcp_access_groups || [], + toolsets: keyData.object_permission?.mcp_toolsets || [], }, mcp_tool_permissions: keyData.object_permission?.mcp_tool_permissions || {}, agents_and_groups: { @@ -233,6 +234,7 @@ export function KeyEditView({ mcp_servers_and_groups: { servers: keyData.object_permission?.mcp_servers || [], accessGroups: keyData.object_permission?.mcp_access_groups || [], + toolsets: keyData.object_permission?.mcp_toolsets || [], }, mcp_tool_permissions: keyData.object_permission?.mcp_tool_permissions || {}, throttle_on_budget_exceeded: keyData.metadata?.throttle_on_budget_exceeded || false, From dce1b0d1fd0e5483efb112c4da8dff0f11f07c04 Mon Sep 17 00:00:00 2001 From: Tin Chi Lo Date: Thu, 23 Jul 2026 18:33:08 -0700 Subject: [PATCH 03/29] fix(ui): allow null for mcp_toolsets in the dashboard key response type The generated schema declares object_permission.mcp_toolsets as string[] | null; the handwritten KeyResponse shape omitted the null. ObjectPermissionsView consumes the same value, so its prop type widens with it --- .../src/components/key_team_helpers/key_list.tsx | 2 +- ui/litellm-dashboard/src/components/object_permissions_view.tsx | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/ui/litellm-dashboard/src/components/key_team_helpers/key_list.tsx b/ui/litellm-dashboard/src/components/key_team_helpers/key_list.tsx index d4bd6756df2..ceff1809b7b 100644 --- a/ui/litellm-dashboard/src/components/key_team_helpers/key_list.tsx +++ b/ui/litellm-dashboard/src/components/key_team_helpers/key_list.tsx @@ -94,7 +94,7 @@ export interface KeyResponse { object_permission_id: string; mcp_servers: string[]; mcp_access_groups?: string[]; - mcp_toolsets?: string[]; + mcp_toolsets?: string[] | null; mcp_tool_permissions?: Record; vector_stores: string[]; agents?: string[]; diff --git a/ui/litellm-dashboard/src/components/object_permissions_view.tsx b/ui/litellm-dashboard/src/components/object_permissions_view.tsx index be021a5b59d..b0ee38bd834 100644 --- a/ui/litellm-dashboard/src/components/object_permissions_view.tsx +++ b/ui/litellm-dashboard/src/components/object_permissions_view.tsx @@ -9,7 +9,7 @@ interface ObjectPermission { mcp_servers: string[]; mcp_access_groups?: string[]; mcp_tool_permissions?: Record; - mcp_toolsets?: string[]; + mcp_toolsets?: string[] | null; vector_stores: string[]; agents?: string[]; agent_access_groups?: string[]; From 58e87985e538aa3a91fdf1431f3651169fd73df8 Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Thu, 23 Jul 2026 23:52:09 -0700 Subject: [PATCH 04/29] test: remove tests that mutation analysis proved assert nothing 25 test functions across three files pass unchanged when every function they execute is mutated; the owning file killed zero of their scored mutants. Four zero-kill tests tied to the fix in #31288 are kept for rewrite instead of removal. --- .../test_litellm/caching/test_redis_cache.py | 425 ------------------ .../guardrail_translation/test_handler.py | 189 +------- tests/test_litellm/proxy/test_proxy_server.py | 44 -- 3 files changed, 1 insertion(+), 657 deletions(-) diff --git a/tests/test_litellm/caching/test_redis_cache.py b/tests/test_litellm/caching/test_redis_cache.py index de5b3b32105..a2e18a62638 100644 --- a/tests/test_litellm/caching/test_redis_cache.py +++ b/tests/test_litellm/caching/test_redis_cache.py @@ -22,78 +22,6 @@ def redis_no_ping(): yield -@pytest.mark.parametrize("namespace", [None, "test"]) -@pytest.mark.asyncio -async def test_redis_cache_async_increment(namespace, monkeypatch, redis_no_ping): - monkeypatch.setenv("REDIS_HOST", "https://my-test-host") - redis_cache = RedisCache(namespace=namespace) - # Create an AsyncMock for the Redis client - mock_redis_instance = AsyncMock() - - # Make sure the mock can be used as an async context manager - mock_redis_instance.__aenter__.return_value = mock_redis_instance - mock_redis_instance.__aexit__.return_value = None - - assert redis_cache is not None - - expected_key = "test:test" if namespace else "test" - - with patch.object( - redis_cache, "init_async_client", return_value=mock_redis_instance - ): - # Call async_set_cache - await redis_cache.async_increment(key=expected_key, value=1) - - # Verify that the set method was called on the mock Redis instance - mock_redis_instance.incrbyfloat.assert_called_once_with( - name=expected_key, amount=1 - ) - - -@pytest.mark.asyncio -async def test_redis_cache_async_increment_refresh_ttl_true_bumps_existing_ttl( - monkeypatch, redis_no_ping -): - """With refresh_ttl=True, every increment should call expire() to bump - the TTL, even when the key already has a TTL (counter-style use).""" - monkeypatch.setenv("REDIS_HOST", "https://my-test-host") - redis_cache = RedisCache() - mock_redis_instance = AsyncMock() - mock_redis_instance.__aenter__.return_value = mock_redis_instance - mock_redis_instance.__aexit__.return_value = None - mock_redis_instance.ttl.return_value = 42 # key already has ~42s left - - with patch.object( - redis_cache, "init_async_client", return_value=mock_redis_instance - ): - await redis_cache.async_increment( - key="spend:team_member:u:t", value=0.05, refresh_ttl=True - ) - - mock_redis_instance.expire.assert_awaited_once_with("spend:team_member:u:t", 60) - - -@pytest.mark.asyncio -async def test_redis_cache_async_increment_default_does_not_bump_existing_ttl( - monkeypatch, redis_no_ping -): - """Default (refresh_ttl=False) preserves window-style semantics: TTL is - set only on first creation, never refreshed (used by rate-limit windows).""" - monkeypatch.setenv("REDIS_HOST", "https://my-test-host") - redis_cache = RedisCache() - mock_redis_instance = AsyncMock() - mock_redis_instance.__aenter__.return_value = mock_redis_instance - mock_redis_instance.__aexit__.return_value = None - mock_redis_instance.ttl.return_value = 42 # key already has ~42s left - - with patch.object( - redis_cache, "init_async_client", return_value=mock_redis_instance - ): - await redis_cache.async_increment(key="rate_limit:window", value=1) - - mock_redis_instance.expire.assert_not_awaited() - - @pytest.mark.parametrize("namespace", [None, "litellm"]) @pytest.mark.asyncio async def test_async_delete_cache_applies_namespace( @@ -140,42 +68,6 @@ async def test_redis_client_init_with_socket_timeout(monkeypatch, redis_no_ping) assert client.connection_pool.connection_kwargs["socket_timeout"] == 1.0 -@pytest.mark.asyncio -async def test_redis_cache_async_batch_get_cache(monkeypatch, redis_no_ping): - monkeypatch.setenv("REDIS_HOST", "https://my-test-host") - redis_cache = RedisCache() - - # Create an AsyncMock for the Redis client - mock_redis_instance = AsyncMock() - - # Make sure the mock can be used as an async context manager - mock_redis_instance.__aenter__.return_value = mock_redis_instance - mock_redis_instance.__aexit__.return_value = None - - # Setup the return value for mget - mock_redis_instance.mget.return_value = [ - b'{"key1": "value1"}', - None, - b'{"key3": "value3"}', - ] - - test_keys = ["key1", "key2", "key3"] - - with patch.object( - redis_cache, "init_async_client", return_value=mock_redis_instance - ): - # Call async_batch_get_cache - result = await redis_cache.async_batch_get_cache(key_list=test_keys) - - # Verify mget was called with the correct keys - mock_redis_instance.mget.assert_called_once() - - # Check that results were properly decoded - assert result["key1"] == {"key1": "value1"} - assert result["key2"] is None - assert result["key3"] == {"key3": "value3"} - - @pytest.mark.asyncio async def test_handle_lpop_count_for_older_redis_versions(monkeypatch): """Test the helper method that handles LPOP with count for Redis versions < 7.0""" @@ -202,41 +94,6 @@ async def test_handle_lpop_count_for_older_redis_versions(monkeypatch): assert mock_pipeline.execute.call_count == 2 -@pytest.mark.asyncio -async def test_async_rpush_pipeline_executes_all_operations(monkeypatch, redis_no_ping): - """Verify that multiple rpush ops are batched into a single pipeline execute""" - monkeypatch.setenv("REDIS_HOST", "https://my-test-host") - redis_cache = RedisCache() - - mock_redis_instance = AsyncMock() - mock_pipeline = MagicMock() - mock_pipeline.__aenter__ = AsyncMock(return_value=mock_pipeline) - mock_pipeline.__aexit__ = AsyncMock(return_value=None) - mock_pipeline.rpush = MagicMock() - mock_pipeline.execute = AsyncMock(return_value=[3, 5, 1]) - mock_redis_instance.pipeline = MagicMock(return_value=mock_pipeline) - - from litellm.types.caching import RedisPipelineRpushOperation - - rpush_list = [ - RedisPipelineRpushOperation(key="key1", values=["a", "b"]), - RedisPipelineRpushOperation(key="key2", values=["c"]), - RedisPipelineRpushOperation(key="key3", values=["d", "e", "f"]), - ] - - with patch.object( - redis_cache, "init_async_client", return_value=mock_redis_instance - ): - result = await redis_cache.async_rpush_pipeline(rpush_list=rpush_list) - - assert result == [3, 5, 1] - assert mock_pipeline.rpush.call_count == 3 - mock_pipeline.rpush.assert_any_call("key1", "a", "b") - mock_pipeline.rpush.assert_any_call("key2", "c") - mock_pipeline.rpush.assert_any_call("key3", "d", "e", "f") - mock_pipeline.execute.assert_called_once() - - @pytest.mark.asyncio async def test_async_rpush_pipeline_empty_list_returns_empty( monkeypatch, redis_no_ping @@ -256,183 +113,6 @@ async def test_async_rpush_pipeline_empty_list_returns_empty( mock_redis_instance.pipeline.assert_not_called() -@pytest.mark.asyncio -async def test_async_rpush_pipeline_raises_on_redis_error(monkeypatch, redis_no_ping): - """Pipeline errors should propagate""" - monkeypatch.setenv("REDIS_HOST", "https://my-test-host") - redis_cache = RedisCache() - - mock_redis_instance = AsyncMock() - mock_pipeline = MagicMock() - mock_pipeline.__aenter__ = AsyncMock(return_value=mock_pipeline) - mock_pipeline.__aexit__ = AsyncMock(return_value=None) - mock_pipeline.rpush = MagicMock() - mock_pipeline.execute = AsyncMock(side_effect=ConnectionError("Redis down")) - mock_redis_instance.pipeline = MagicMock(return_value=mock_pipeline) - - from litellm.types.caching import RedisPipelineRpushOperation - - rpush_list = [RedisPipelineRpushOperation(key="key1", values=["a"])] - - with patch.object( - redis_cache, "init_async_client", return_value=mock_redis_instance - ): - with pytest.raises(ConnectionError, match="Redis down"): - await redis_cache.async_rpush_pipeline(rpush_list=rpush_list) - - -@pytest.mark.asyncio -async def test_async_lpop_pipeline_single_round_trip(monkeypatch, redis_no_ping): - """Verify that multiple lpop ops are batched into a single pipeline execute""" - monkeypatch.setenv("REDIS_HOST", "https://my-test-host") - redis_cache = RedisCache() - redis_cache.redis_version = "7.0.0" - - mock_redis_instance = AsyncMock() - mock_pipeline = MagicMock() - mock_pipeline.__aenter__ = AsyncMock(return_value=mock_pipeline) - mock_pipeline.__aexit__ = AsyncMock(return_value=None) - mock_pipeline.lpop = MagicMock() - mock_pipeline.execute = AsyncMock( - return_value=[ - [b"val1", b"val2"], # key1 results - None, # key2 empty - [b"val3"], # key3 results - ] - ) - mock_redis_instance.pipeline = MagicMock(return_value=mock_pipeline) - - from litellm.types.caching import RedisPipelineLpopOperation - - lpop_list = [ - RedisPipelineLpopOperation(key="key1", count=10), - RedisPipelineLpopOperation(key="key2", count=10), - RedisPipelineLpopOperation(key="key3", count=5), - ] - - with patch.object( - redis_cache, "init_async_client", return_value=mock_redis_instance - ): - results = await redis_cache.async_lpop_pipeline(lpop_list=lpop_list) - - assert len(results) == 3 - assert results[0] == ["val1", "val2"] - assert results[1] is None - assert results[2] == ["val3"] - mock_pipeline.execute.assert_called_once() - - -@pytest.mark.asyncio -async def test_async_lpop_pipeline_redis_lt7_regroups_flat_results( - monkeypatch, redis_no_ping -): - """Verify Redis < 7 fallback issues individual LPOPs and regroups correctly""" - monkeypatch.setenv("REDIS_HOST", "https://my-test-host") - redis_cache = RedisCache() - redis_cache.redis_version = "6.2.0" - - mock_redis_instance = AsyncMock() - mock_pipeline = MagicMock() - mock_pipeline.__aenter__ = AsyncMock(return_value=mock_pipeline) - mock_pipeline.__aexit__ = AsyncMock(return_value=None) - mock_pipeline.lpop = MagicMock() - - # With count=3 for key1 and count=2 for key2, we get 5 individual LPOP commands - # Simulate: key1 has 2 values then None, key2 has 1 value then None - mock_pipeline.execute = AsyncMock( - return_value=[ - b"val1", - b"val2", - None, # 3 LPOPs for key1 - b"val3", - None, # 2 LPOPs for key2 - ] - ) - mock_redis_instance.pipeline = MagicMock(return_value=mock_pipeline) - - from litellm.types.caching import RedisPipelineLpopOperation - - lpop_list = [ - RedisPipelineLpopOperation(key="key1", count=3), - RedisPipelineLpopOperation(key="key2", count=2), - ] - - with patch.object( - redis_cache, "init_async_client", return_value=mock_redis_instance - ): - results = await redis_cache.async_lpop_pipeline(lpop_list=lpop_list) - - assert len(results) == 2 - assert results[0] == ["val1", "val2"] # 2 values, None filtered out - assert results[1] == ["val3"] # 1 value, None filtered out - # All 5 individual LPOPs should be queued, but only 1 execute() call - assert mock_pipeline.lpop.call_count == 5 - mock_pipeline.execute.assert_called_once() - - -@pytest.mark.asyncio -async def test_async_rpush_pipeline_raises_on_per_command_error( - monkeypatch, redis_no_ping -): - """Verify that per-command errors in pipeline results are raised, not silently dropped""" - monkeypatch.setenv("REDIS_HOST", "https://my-test-host") - redis_cache = RedisCache() - - mock_redis_instance = AsyncMock() - mock_pipeline = MagicMock() - mock_pipeline.__aenter__ = AsyncMock(return_value=mock_pipeline) - mock_pipeline.__aexit__ = AsyncMock(return_value=None) - mock_pipeline.rpush = MagicMock() - # Simulate: first RPUSH succeeds, second returns a per-command error - mock_pipeline.execute = AsyncMock(return_value=[3, Exception("WRONGTYPE")]) - mock_redis_instance.pipeline = MagicMock(return_value=mock_pipeline) - - from litellm.types.caching import RedisPipelineRpushOperation - - rpush_list = [ - RedisPipelineRpushOperation(key="key1", values=["a"]), - RedisPipelineRpushOperation(key="key2", values=["b"]), - ] - - with patch.object( - redis_cache, "init_async_client", return_value=mock_redis_instance - ): - with pytest.raises(Exception, match="WRONGTYPE"): - await redis_cache.async_rpush_pipeline(rpush_list=rpush_list) - - -@pytest.mark.asyncio -async def test_async_lpop_pipeline_raises_on_per_command_error( - monkeypatch, redis_no_ping -): - """Verify that per-command errors in LPOP pipeline results are raised, not silently dropped""" - monkeypatch.setenv("REDIS_HOST", "https://my-test-host") - redis_cache = RedisCache() - redis_cache.redis_version = "7.0.0" - - mock_redis_instance = AsyncMock() - mock_pipeline = MagicMock() - mock_pipeline.__aenter__ = AsyncMock(return_value=mock_pipeline) - mock_pipeline.__aexit__ = AsyncMock(return_value=None) - mock_pipeline.lpop = MagicMock() - # Simulate: first LPOP succeeds, second returns a per-command error - mock_pipeline.execute = AsyncMock(return_value=[[b"val1"], Exception("WRONGTYPE")]) - mock_redis_instance.pipeline = MagicMock(return_value=mock_pipeline) - - from litellm.types.caching import RedisPipelineLpopOperation - - lpop_list = [ - RedisPipelineLpopOperation(key="key1", count=10), - RedisPipelineLpopOperation(key="key2", count=10), - ] - - with patch.object( - redis_cache, "init_async_client", return_value=mock_redis_instance - ): - with pytest.raises(Exception, match="WRONGTYPE"): - await redis_cache.async_lpop_pipeline(lpop_list=lpop_list) - - @pytest.mark.asyncio async def test_async_lpop_pipeline_empty_list(monkeypatch, redis_no_ping): """Empty lpop_list should return empty list without touching Redis""" @@ -450,111 +130,6 @@ async def test_async_lpop_pipeline_empty_list(monkeypatch, redis_no_ping): mock_redis_instance.pipeline.assert_not_called() -@pytest.mark.asyncio -async def test_async_lpop_pipeline_propagates_redis_exception( - monkeypatch, redis_no_ping -): - """Pipeline errors should propagate""" - monkeypatch.setenv("REDIS_HOST", "https://my-test-host") - redis_cache = RedisCache() - redis_cache.redis_version = "7.0.0" - - mock_redis_instance = AsyncMock() - mock_pipeline = MagicMock() - mock_pipeline.__aenter__ = AsyncMock(return_value=mock_pipeline) - mock_pipeline.__aexit__ = AsyncMock(return_value=None) - mock_pipeline.lpop = MagicMock() - mock_pipeline.execute = AsyncMock(side_effect=ConnectionError("Redis down")) - mock_redis_instance.pipeline = MagicMock(return_value=mock_pipeline) - - from litellm.types.caching import RedisPipelineLpopOperation - - lpop_list = [RedisPipelineLpopOperation(key="key1", count=10)] - - with patch.object( - redis_cache, "init_async_client", return_value=mock_redis_instance - ): - with pytest.raises(ConnectionError, match="Redis down"): - await redis_cache.async_lpop_pipeline(lpop_list=lpop_list) - - -@pytest.mark.asyncio -@pytest.mark.parametrize( - "redis_version", - [ - # Standard cases - "7.0.0", # Standard Redis string version - 7.0, # Valkey/ElastiCache float version (THE BUG this fix addresses) - 7, # Integer version (e.g., from some Redis forks) - # Version < 7 - "6", # String without dots, version < 7 - # Malformed versions (fallback to 7) - "latest", # Non-numeric version - "", # Empty string - -7.0, # Negative float - # Format variations - " 7.0.0 ", # Whitespace (should be stripped) - "7.0.0-rc1", # Version with suffix - "10.0.0", # Double digit major version - ], -) -async def test_async_lpop_with_float_redis_version( - monkeypatch, redis_no_ping, redis_version -): - """ - Test async_lpop with various Redis version formats (especially float). - - This test specifically addresses the issue where AWS ElastiCache Valkey - returns redis_version as a float (e.g., 7.0) instead of a string (e.g., "7.0.0"), - which caused a 'float' object has no attribute 'split' error when trying to - use the Redis transaction buffer feature. - - The fix converts the version to a string and handles edge cases like: - - Floats (7.0) and integers (7) - - Strings with/without dots ("7" vs "7.0.0") - - Malformed versions ("v7.0.0", "latest") - fallback to version 7 - - Whitespace (" 7.0.0 ") - - Negative versions (fallback to version 7) - - Related: Database deadlock issues when use_redis_transaction_buffer is enabled. - """ - monkeypatch.setenv("REDIS_HOST", "https://my-test-host") - - # Create RedisCache instance - redis_cache = RedisCache() - redis_cache.redis_version = redis_version # Set the version to test - - # Create an AsyncMock for the Redis client - mock_redis_instance = AsyncMock() - mock_redis_instance.__aenter__.return_value = mock_redis_instance - mock_redis_instance.__aexit__.return_value = None - - # Mock lpop to return a test value (Redis >= 7.0 behavior) - mock_redis_instance.lpop.return_value = [b"value1", b"value2"] - - # Mock pipeline for Redis < 7.0 (used when major_version < 7) - mock_pipeline = MagicMock() - mock_pipeline.__aenter__ = AsyncMock(return_value=mock_pipeline) - mock_pipeline.__aexit__ = AsyncMock(return_value=None) - # Make pipeline() a regular method (not async) that returns the mock - mock_redis_instance.pipeline = MagicMock(return_value=mock_pipeline) - - # Mock handle_lpop_count_for_older_redis_versions for Redis < 7 - with patch.object( - redis_cache, - "handle_lpop_count_for_older_redis_versions", - return_value=[b"value1", b"value2"], - ): - with patch.object( - redis_cache, "init_async_client", return_value=mock_redis_instance - ): - # Call async_lpop with count - this should not raise AttributeError - result = await redis_cache.async_lpop(key="test_key", count=2) - - # Verify the method completed without error - assert result is not None - - # LIT-3374: the namespace must be applied uniformly across every key-taking # Redis operation, not just get/set/increment. Before the fix these paths wrote # or read raw keys, so with a namespace configured the prefixed keys other diff --git a/tests/test_litellm/llms/pass_through/guardrail_translation/test_handler.py b/tests/test_litellm/llms/pass_through/guardrail_translation/test_handler.py index f8bd83fc7df..1043c26c6ec 100644 --- a/tests/test_litellm/llms/pass_through/guardrail_translation/test_handler.py +++ b/tests/test_litellm/llms/pass_through/guardrail_translation/test_handler.py @@ -1,15 +1,10 @@ """ -Tests for LlmPassthroughRouteHandler and the guardrail_translation_mappings registry. +Tests for the guardrail_translation_mappings registry. Validates: - allm_passthrough_route is registered in the mappings (regression: this was the bug) -- Bedrock provider is dispatched to BedrockPassthroughGuardrailHandler -- Unknown provider skips apply_guardrail """ -import pytest -from unittest.mock import AsyncMock, MagicMock, patch - from litellm.llms.pass_through.guardrail_translation import ( guardrail_translation_mappings, ) @@ -40,185 +35,3 @@ class TestRegistry: is PassThroughEndpointHandler ) - -def _make_guardrail() -> MagicMock: - g = MagicMock() - g.guardrail_name = "test-guard" - g.apply_guardrail = AsyncMock(return_value={"texts": []}) - g.skip_system_message_in_guardrail = False - g.skip_tool_message_in_guardrail = False - return g - - -class TestLlmPassthroughRouteHandlerInput: - @pytest.mark.asyncio - async def test_bedrock_provider_delegates_to_bedrock_handler(self): - handler = LlmPassthroughRouteHandler() - data = { - "custom_llm_provider": "bedrock", - "endpoint": "model/anthropic.claude-3-sonnet/converse", - "data": {"messages": [{"role": "user", "content": [{"text": "hi"}]}]}, - } - guardrail = _make_guardrail() - - await handler.process_input_messages(data=data, guardrail_to_apply=guardrail) - - guardrail.apply_guardrail.assert_called_once() - - @pytest.mark.asyncio - async def test_unknown_provider_skips_apply_guardrail(self): - handler = LlmPassthroughRouteHandler() - data = { - "custom_llm_provider": "some_unknown_provider", - "endpoint": "v1/chat/completions", - "data": {"messages": [{"role": "user", "content": "hi"}]}, - } - guardrail = _make_guardrail() - - result = await handler.process_input_messages( - data=data, guardrail_to_apply=guardrail - ) - - guardrail.apply_guardrail.assert_not_called() - assert result is data - - @pytest.mark.asyncio - async def test_missing_provider_skips(self): - handler = LlmPassthroughRouteHandler() - data = {"endpoint": "foo/bar", "data": {}} - guardrail = _make_guardrail() - - result = await handler.process_input_messages( - data=data, guardrail_to_apply=guardrail - ) - - guardrail.apply_guardrail.assert_not_called() - assert result is data - - -class TestLlmPassthroughRouteHandlerOutput: - @pytest.mark.asyncio - async def test_bedrock_provider_delegates_output_to_bedrock_handler(self): - handler = LlmPassthroughRouteHandler() - response = { - "output": { - "message": { - "role": "assistant", - "content": [{"text": "hello"}], - } - } - } - request_data = { - "custom_llm_provider": "bedrock", - "endpoint": "model/anthropic.claude-3-sonnet/converse", - } - guardrail = _make_guardrail() - - await handler.process_output_response( - response=response, - guardrail_to_apply=guardrail, - request_data=request_data, - ) - - guardrail.apply_guardrail.assert_called_once() - - @pytest.mark.asyncio - async def test_unknown_provider_skips_output(self): - handler = LlmPassthroughRouteHandler() - response = {"some": "response"} - request_data = {"custom_llm_provider": "unknown"} - guardrail = _make_guardrail() - - result = await handler.process_output_response( - response=response, - guardrail_to_apply=guardrail, - request_data=request_data, - ) - - guardrail.apply_guardrail.assert_not_called() - assert result is response - - -class TestDeAnonymizeEventStream: - @pytest.mark.asyncio - async def test_bedrock_provider_dispatches_to_handler(self): - body = b"original-stream-bytes" - expected = b"de-anonymized-bytes" - proxy_logging_obj = MagicMock() - user_api_key_dict = MagicMock() - - with patch( - "litellm.llms.bedrock.passthrough.guardrail_translation.handler." - "BedrockPassthroughGuardrailHandler.de_anonymize_event_stream", - new=AsyncMock(return_value=expected), - ) as mock_handler: - result = await LlmPassthroughRouteHandler.de_anonymize_event_stream( - body_bytes=body, - proxy_logging_obj=proxy_logging_obj, - user_api_key_dict=user_api_key_dict, - data={"custom_llm_provider": "bedrock"}, - ) - - mock_handler.assert_awaited_once() - assert result == expected - - @pytest.mark.asyncio - async def test_unknown_provider_returns_original_bytes(self): - body = b"original-stream-bytes" - - result = await LlmPassthroughRouteHandler.de_anonymize_event_stream( - body_bytes=body, - proxy_logging_obj=MagicMock(), - user_api_key_dict=MagicMock(), - data={"custom_llm_provider": "anthropic"}, - ) - - assert result is body - - @pytest.mark.asyncio - async def test_missing_provider_returns_original_bytes(self): - body = b"original-stream-bytes" - - result = await LlmPassthroughRouteHandler.de_anonymize_event_stream( - body_bytes=body, - proxy_logging_obj=MagicMock(), - user_api_key_dict=MagicMock(), - data={}, - ) - - assert result is body - - -class TestSupportsEventStreamDeAnonymization: - def test_bedrock_converse_stream_is_supported(self): - assert ( - LlmPassthroughRouteHandler.supports_event_stream_de_anonymization( - "bedrock", "model/us.amazon.nova-lite-v1:0/converse-stream" - ) - is True - ) - - def test_bedrock_invoke_stream_is_not_supported(self): - assert ( - LlmPassthroughRouteHandler.supports_event_stream_de_anonymization( - "bedrock", - "model/us.amazon.nova-lite-v1:0/invoke-with-response-stream", - ) - is False - ) - - def test_unknown_provider_is_not_supported(self): - assert ( - LlmPassthroughRouteHandler.supports_event_stream_de_anonymization( - "anthropic", "model/foo/converse-stream" - ) - is False - ) - - def test_missing_provider_is_not_supported(self): - assert ( - LlmPassthroughRouteHandler.supports_event_stream_de_anonymization( - None, "model/foo/converse-stream" - ) - is False - ) diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index 54db0c0fd4f..040fba7c53a 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -1206,50 +1206,6 @@ def test_team_info_masking(): assert "public-test-key" not in str(exc_info.value) -def test_embedding_input_array_of_tokens(client_no_auth): - """ - Test to bypass decoding input as array of tokens for selected providers - - Ref: https://github.com/BerriAI/litellm/issues/10113 - """ - from litellm.proxy import proxy_server - - # The client_no_auth fixture should initialize the router - # Assert this to catch any router initialization regressions - assert proxy_server.llm_router is not None, ( - "llm_router is None after client_no_auth fixture initialized. " - "This indicates a router initialization issue that should be investigated." - ) - - try: - with mock.patch.object( - proxy_server.llm_router, - "aembedding", - return_value=example_embedding_result, - ) as mock_aembedding: - test_data = { - "model": "vllm_embed_model", - "input": [[2046, 13269, 158208]], - } - - response = client_no_auth.post("/v1/embeddings", json=test_data) - - # Assert that aembedding was called, and that input was not modified - mock_aembedding.assert_called_once() - call_args, call_kwargs = mock_aembedding.call_args - assert call_kwargs["model"] == "vllm_embed_model" - assert call_kwargs["input"] == [[2046, 13269, 158208]] - - assert response.status_code == 200 - result = response.json() - print(len(result["data"][0]["embedding"])) - assert ( - len(result["data"][0]["embedding"]) > 10 - ) # this usually has len==1536 so - except Exception as e: - pytest.fail(f"LiteLLM Proxy test failed. Exception - {str(e)}") - - @pytest.mark.asyncio async def test_get_all_team_models(): """ From 8177230a29529aba4e9c70fb78f48c1055ac1fac Mon Sep 17 00:00:00 2001 From: Noah Nistler <60981020+noahnistler@users.noreply.github.com> Date: Fri, 24 Jul 2026 15:25:58 -0500 Subject: [PATCH 05/29] feat(guardrails): add run_in_parallel opt-in for concurrent pre_call and post_call guardrails (#33770) * feat(guardrails): add run_in_parallel opt-in for concurrent pre_call guardrails Pre-call guardrails run sequentially because each may mutate the request payload and later guardrails depend on earlier mutations. Deployments with several slow block-only pre_call guardrails (external moderation, Bedrock, LLM-judge) therefore pay the sum of their latencies. during_call guardrails run concurrently but alongside the LLM call, so a violating payload has already been sent, which is unacceptable when the request must never reach the model. This adds a per-guardrail run_in_parallel flag (default off). Guardrails that opt in are pulled out of the sequential loop and run concurrently via asyncio.gather after every sequential (payload-mutating) guardrail has run, so they observe the mutated payload and still form a hard barrier before the LLM call; the first to raise blocks the request. Their returned data is discarded since they are declared block-only. The flag is wired from LitellmParams onto the guardrail instance at the same generic choke point in initialize_guardrail that already sets skip_system_message_in_guardrail, so no per-provider initializer needs to change. * feat(guardrails): extend run_in_parallel opt-in to post_call guardrails post_call_success_hook ran guardrails sequentially for the same reason pre_call did: response-modifying guardrails thread the response forward. But block-only output scanners (which read the response and reject on violation without changing it) serialize for no benefit and add latency. This reuses the existing run_in_parallel flag for the post_call hook. Opted-in post_call guardrails are pulled out of the sequential loop and run concurrently via asyncio.gather after the sequential (response-modifying) guardrails and before the non-guardrail CustomLogger callbacks, so they inspect the final response and still block it from reaching the client if any raises. Their returned response is discarded since they are block-only. The apply_guardrail path sets data["guardrail_to_apply"] immediately before awaiting, and unified_guardrail pops it before its first suspension point, so concurrent guardrails never race on that key under asyncio's cooperative scheduling. * fix(guardrails): await all parallel guardrails and prioritize blocks over reroutes Addresses review feedback on the run_in_parallel opt-in. asyncio.gather propagated the first exception without cancelling or awaiting the siblings, so a block at t=0 left the other guardrails running as unobserved background tasks (wasted external calls plus event-loop warnings), and a fast SensitiveDataRouteException/ModifyResponseException could return a reroute or passthrough before a slower block finished, letting crafted input bypass the block. Both the pre_call and post_call parallel batches now gather with return_exceptions=True so every guardrail runs to completion, then raise any blocking exception ahead of a flow-changing one. The registry choke point wrote bool(None)==False onto every instance when the config omitted run_in_parallel, silently disabling a constructor-set default; it now only writes when the config provides an explicit value. * fix(guardrails): record lifecycle logs for every concurrently-run guardrail The log_guardrail_information decorator skipped its auto-record when it saw that the count of standard_logging_guardrail_information entries in the shared request_data had grown during the wrapped call, taking that as proof the wrapped function had recorded its own richer entry. That heuristic breaks the moment guardrails run concurrently (parallel pre_call/post_call, during_call): a sibling guardrail's append inflates the shared count, so a guardrail that did not self-record wrongly concludes it already did and drops its own entry. The result is that enabling run_in_parallel silently loses per-guardrail lifecycle logs, so the Admin UI Request Lifecycle timeline and downstream loggers (Datadog, Langfuse, OTEL, spend logs) show only one of the concurrent guardrails. Replace the shared-count heuristic with a ContextVar flag set when a guardrail records its own entry. asyncio copies the context into each gathered task, so the flag is isolated per concurrent guardrail while still catching the self-record-then-skip-auto-record case within a single invocation. * test(guardrails): declare run_in_parallel on post_call guardrail mocks The post_call partition reads run_in_parallel on every CustomGuardrail callback. A MagicMock(spec=CustomGuardrail) has no run_in_parallel (it is set in __init__, not on the class) so the attribute access raised, and even a class-level default would return a truthy child mock that wrongly routes the double into the parallel batch. Declare the flag False on the shared mock factories so these pre-existing hook tests exercise the sequential path they assert on. * fix(guardrails): harden run_in_parallel reads and address review feedback Read run_in_parallel via getattr(..., False) in the pre_call and post_call partitions so a third-party CustomGuardrail subclass that overrides __init__ without chaining super().__init__() no longer raises AttributeError on a path that previously worked. Drop the redundant in-function GuardrailEventHooks import in _run_parallel_post_call_guardrails (already imported module-level). Remove the flaky wall-clock upper-bound assertions from the two concurrency tests; the all-start-before-any-end overlap assertion is the timing-independent signal that actually proves concurrency. --- litellm/integrations/custom_guardrail.py | 44 +- .../proxy/guardrails/guardrail_registry.py | 3 + litellm/proxy/utils.py | 135 +++++ litellm/types/guardrails.py | 11 + ruff-strict-budget.json | 4 +- tests/proxy_unit_tests/test_proxy_utils.py | 461 +++++++++++++++++- .../integrations/test_custom_guardrail.py | 33 ++ .../guardrails/test_guardrail_registry.py | 40 ++ .../proxy/guardrails/test_init_guardrails.py | 24 + .../proxy_logging/test_guardrail_pipeline.py | 1 + .../test_post_call_success_hook.py | 1 + 11 files changed, 736 insertions(+), 21 deletions(-) diff --git a/litellm/integrations/custom_guardrail.py b/litellm/integrations/custom_guardrail.py index cf9dafcb222..f639ad49d5e 100644 --- a/litellm/integrations/custom_guardrail.py +++ b/litellm/integrations/custom_guardrail.py @@ -1,3 +1,4 @@ +import contextvars import hashlib import os import secrets @@ -64,6 +65,10 @@ from litellm.exceptions import ( # proxy's metadata sanitizer. _PRE_CALL_EXECUTED_TOKEN = secrets.token_hex(16) +_guardrail_self_recorded: contextvars.ContextVar[bool] = contextvars.ContextVar( + "litellm_guardrail_self_recorded", default=False +) + def _strict_guardrail_modes_enabled() -> bool: """Whether guardrail-mode validation raises (default) or logs a warning. @@ -117,6 +122,7 @@ class CustomGuardrail(CustomLogger): on_sensitive_data: Optional[str] = None, sensitive_data_route_to_model: Optional[str] = None, sticky_session_routing: bool = True, + run_in_parallel: bool = False, only_scan_new_messages: bool = False, **kwargs, ): @@ -136,6 +142,9 @@ class CustomGuardrail(CustomLogger): on_sensitive_data: Action when sensitive data is detected. 'block' (default) or 'route' sensitive_data_route_to_model: Model to route to when on_sensitive_data='route' sticky_session_routing: When True, all subsequent requests in the session use the same model + run_in_parallel: When True, this pre_call or post_call guardrail runs concurrently with + other opted-in guardrails of the same hook. Only safe for block-only guardrails that + do not mutate the request or response. """ self.guardrail_name = guardrail_name self.supported_event_hooks = supported_event_hooks @@ -150,6 +159,7 @@ class CustomGuardrail(CustomLogger): self.on_sensitive_data: Optional[str] = on_sensitive_data self.sensitive_data_route_to_model: Optional[str] = sensitive_data_route_to_model self.sticky_session_routing: bool = sticky_session_routing + self.run_in_parallel: bool = run_in_parallel self.only_scan_new_messages: bool = only_scan_new_messages if supported_event_hooks: @@ -956,6 +966,8 @@ class CustomGuardrail(CustomLogger): request_data["metadata"] = {} _append_guardrail_info(request_data["metadata"]) + _guardrail_self_recorded.set(True) + # Emit the otel guardrail span here, where every guardrail execution lands, # rather than relying on a post-call hook that does not fire on every path # (e.g. a pass-through request that passes its guardrails). @@ -1238,8 +1250,12 @@ def log_guardrail_information(func): (structured detections, tracing detail) than this decorator's "allow"/"mask"/raw-response default. To avoid double-recording in that case (which would emit two spans, two Datadog records, two spend-log - entries, etc.), snapshot the entry count before invocation: if the - wrapped function already appended its own entry, skip the auto-record. + entries, etc.), a context-local flag records whether the wrapped function + appended its own entry; if so, the auto-record is skipped. The flag is a + ``ContextVar`` rather than a count of entries in the shared ``request_data`` + so it stays correct when guardrails run concurrently (asyncio copies the + context into each gathered task): counting shared entries would let one + guardrail's append hide another guardrail's missing record. """ import functools import inspect @@ -1259,16 +1275,6 @@ def log_guardrail_information(func): return GuardrailEventHooks.post_call return None - def _count_recorded_guardrail_entries(request_data: dict) -> int: - total = 0 - for container_key in ("metadata", "litellm_metadata"): - container = request_data.get(container_key) - if isinstance(container, dict): - entries = container.get("standard_logging_guardrail_information") - if isinstance(entries, list): - total += len(entries) - return total - @functools.wraps(func) async def async_wrapper(*args, **kwargs): start_time = datetime.now() # Move start_time inside the wrapper @@ -1282,10 +1288,10 @@ def log_guardrail_information(func): original_inputs = kwargs.get("inputs") logging_obj = kwargs.get("logging_obj") - entries_before = _count_recorded_guardrail_entries(request_data) + self_recorded_token = _guardrail_self_recorded.set(False) try: response = await func(*args, **kwargs) - if _count_recorded_guardrail_entries(request_data) > entries_before: + if _guardrail_self_recorded.get(): return response return self._process_response( response=response, @@ -1297,7 +1303,7 @@ def log_guardrail_information(func): original_inputs=original_inputs, ) except Exception as e: - if _count_recorded_guardrail_entries(request_data) > entries_before: + if _guardrail_self_recorded.get(): raise return self._process_error( e=e, @@ -1308,6 +1314,7 @@ def log_guardrail_information(func): event_type=event_type, ) finally: + _guardrail_self_recorded.reset(self_recorded_token) _sync_guardrail_info_to_logging_obj(request_data, logging_obj) @functools.wraps(func) @@ -1323,10 +1330,10 @@ def log_guardrail_information(func): original_inputs = kwargs.get("inputs") logging_obj = kwargs.get("logging_obj") - entries_before = _count_recorded_guardrail_entries(request_data) + self_recorded_token = _guardrail_self_recorded.set(False) try: response = func(*args, **kwargs) - if _count_recorded_guardrail_entries(request_data) > entries_before: + if _guardrail_self_recorded.get(): return response return self._process_response( response=response, @@ -1336,7 +1343,7 @@ def log_guardrail_information(func): original_inputs=original_inputs, ) except Exception as e: - if _count_recorded_guardrail_entries(request_data) > entries_before: + if _guardrail_self_recorded.get(): raise return self._process_error( e=e, @@ -1345,6 +1352,7 @@ def log_guardrail_information(func): event_type=event_type, ) finally: + _guardrail_self_recorded.reset(self_recorded_token) _sync_guardrail_info_to_logging_obj(request_data, logging_obj) @functools.wraps(func) diff --git a/litellm/proxy/guardrails/guardrail_registry.py b/litellm/proxy/guardrails/guardrail_registry.py index 1f2c9e0c182..bd00e9815a8 100644 --- a/litellm/proxy/guardrails/guardrail_registry.py +++ b/litellm/proxy/guardrails/guardrail_registry.py @@ -489,6 +489,9 @@ class InMemoryGuardrailHandler: "skip_tool_message_in_guardrail", getattr(litellm_params, "skip_tool_message_in_guardrail", None), ) + configured_run_in_parallel = getattr(litellm_params, "run_in_parallel", None) + if configured_run_in_parallel is not None: + custom_guardrail_callback.run_in_parallel = bool(configured_run_in_parallel) parsed_guardrail = Guardrail( guardrail_id=guardrail.get("guardrail_id"), diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 5b81d1f2da3..171e17ce650 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -1400,6 +1400,14 @@ class ProxyLogging: self._process_guardrail_metadata(data) return data + parallel_guardrails: tuple[CustomGuardrail, ...] = tuple( + cb + for cb in caps.resolved_callbacks + if isinstance(cb, CustomGuardrail) + and getattr(cb, "run_in_parallel", False) + and not (cb.guardrail_name and cb.guardrail_name in pipeline_managed) + ) + deferred_route_exc: Optional[SensitiveDataRouteException] = None for _callback in caps.resolved_callbacks: start_time = time.time() @@ -1409,6 +1417,9 @@ class ProxyLogging: if _callback.guardrail_name and _callback.guardrail_name in pipeline_managed: continue + if getattr(_callback, "run_in_parallel", False): + continue + result = await self._process_guardrail_callback( callback=_callback, data=data, # type: ignore @@ -1465,6 +1476,14 @@ class ProxyLogging: if deferred_route_exc is not None and data is not None: data = await self._handle_sensitive_data_route_exception(deferred_route_exc, data, user_api_key_dict) + if parallel_guardrails and data is not None: + await self._run_parallel_pre_call_guardrails( + guardrails=parallel_guardrails, + data=data, + user_api_key_dict=user_api_key_dict, + call_type=call_type, + ) + if data is not None: self._process_guardrail_metadata(data) @@ -1477,6 +1496,47 @@ class ProxyLogging: except Exception as e: raise e + async def _run_parallel_pre_call_guardrails( + self, + guardrails: tuple[CustomGuardrail, ...], + data: dict, + user_api_key_dict: UserAPIKeyAuth, + call_type: CallTypesLiteral, + ) -> None: + """ + Run opted-in pre_call guardrails concurrently against one shared payload + snapshot. These guardrails are declared block-only, so any modified data + they return is discarded; they run for their blocking side effect (raising + to reject the request before it reaches the LLM). Every guardrail is + awaited to completion (``return_exceptions=True``) so a raise by one never + leaves the others running as unobserved background tasks. A guardrail that + blocks (any exception other than a reroute or passthrough) takes precedence + over one that only changes the request flow, so a fast reroute can never + let a slower block be bypassed; the request is rejected before it reaches + the LLM, preserving the pre-call barrier that ``during_call`` guardrails + cannot provide. Per-guardrail latency is recorded by + ``_process_guardrail_callback``'s own metrics. + """ + results = await asyncio.gather( + *( + self._process_guardrail_callback( + callback=callback, + data=data, + user_api_key_dict=user_api_key_dict, + call_type=call_type, + event_type=GuardrailEventHooks.pre_call, + ) + for callback in guardrails + ), + return_exceptions=True, + ) + raised = tuple(result for result in results if isinstance(result, BaseException)) + blocking = next((exc for exc in raised if not _exception_changes_request_flow(exc)), None) + if blocking is not None: + raise blocking + if raised: + raise raised[0] + async def _handle_sensitive_data_route_exception( self, exc: SensitiveDataRouteException, @@ -2277,9 +2337,16 @@ class ProxyLogging: # Merge model-level guardrails before checking which guardrails to run guardrail_data = _check_and_merge_model_level_guardrails(data=data, llm_router=llm_router) + parallel_guardrails: tuple[CustomGuardrail, ...] = tuple( + callback for callback in guardrail_callbacks if getattr(callback, "run_in_parallel", False) + ) + for callback in guardrail_callbacks: # Main - V2 Guardrails implementation + if getattr(callback, "run_in_parallel", False): + continue + if ( callback.should_run_guardrail( data=guardrail_data, @@ -2316,6 +2383,15 @@ class ProxyLogging: if guardrail_response is not None: response = guardrail_response + if parallel_guardrails: + await self._run_parallel_post_call_guardrails( + guardrails=parallel_guardrails, + data=data, + guardrail_data=guardrail_data, + response=response, + user_api_key_dict=user_api_key_dict, + ) + ############ Handle CustomLogger ############################### ################################################################# @@ -2329,6 +2405,65 @@ class ProxyLogging: raise e return response + async def _run_parallel_post_call_guardrails( + self, + guardrails: tuple[CustomGuardrail, ...], + data: dict, + guardrail_data: dict, + response: LLMResponseTypes, + user_api_key_dict: UserAPIKeyAuth, + ) -> None: + """ + Run opted-in post_call guardrails concurrently against the response + produced by the sequential guardrails. These guardrails are declared + block-only, so any modified response they return is discarded; they run + for their blocking side effect (raising to reject the response before it + reaches the client). Every guardrail is awaited to completion + (``return_exceptions=True``) so a raise by one never leaves the others + running as unobserved background tasks. A guardrail that blocks (any + exception other than a passthrough) takes precedence over one that only + changes the response flow, so a fast passthrough can never let a slower + block be bypassed. Each per-guardrail coroutine sets ``guardrail_to_apply`` + immediately before awaiting, and the unified hook pops it before its first + suspension point, so concurrent guardrails never race on that key. + """ + + async def _run_one(callback: CustomGuardrail) -> None: + if callback.should_run_guardrail(data=guardrail_data, event_type=GuardrailEventHooks.post_call) is not True: + return + if "apply_guardrail" in type(callback).__dict__: + data["guardrail_to_apply"] = callback + await self._run_guardrail_with_metrics( + callback, + unified_guardrail.async_post_call_success_hook( + user_api_key_dict=user_api_key_dict, + data=data, + response=response, + ), + "post_call", + ) + else: + await self._run_guardrail_with_metrics( + callback, + callback.async_post_call_success_hook( + user_api_key_dict=user_api_key_dict, + data=data, + response=response, + ), + "post_call", + ) + + results = await asyncio.gather( + *(_run_one(callback) for callback in guardrails), + return_exceptions=True, + ) + raised = tuple(result for result in results if isinstance(result, BaseException)) + blocking = next((exc for exc in raised if not _exception_changes_request_flow(exc)), None) + if blocking is not None: + raise blocking + if raised: + raise raised[0] + async def post_call_response_headers_hook( self, data: dict, diff --git a/litellm/types/guardrails.py b/litellm/types/guardrails.py index c86794b90f8..a324e71e289 100644 --- a/litellm/types/guardrails.py +++ b/litellm/types/guardrails.py @@ -910,6 +910,17 @@ class BaseLitellmParams(ContentFilterConfigModel): # works for new and patch up ), ) + run_in_parallel: Optional[bool] = Field( + default=None, + description=( + "When True, this pre_call or post_call guardrail runs concurrently with other opted-in " + "guardrails of the same hook, after the sequential guardrails have run. Use only for " + "block-only guardrails that inspect and reject; do not enable it for guardrails that " + "modify the request or response (e.g. PII masking or sensitive-data routing), since " + "parallel runs share one snapshot and their mutations would race." + ), + ) + @field_validator( "mode", "default_action", diff --git a/ruff-strict-budget.json b/ruff-strict-budget.json index d3d70ff5ff4..39e9bf4773d 100644 --- a/ruff-strict-budget.json +++ b/ruff-strict-budget.json @@ -222,7 +222,7 @@ "limit": 38 }, "RET504": { - "limit": 721 + "limit": 719 }, "RUF010": { "limit": 874 @@ -324,7 +324,7 @@ "limit": 883 }, "UP006": { - "limit": 12792 + "limit": 12789 }, "UP007": { "limit": 2570 diff --git a/tests/proxy_unit_tests/test_proxy_utils.py b/tests/proxy_unit_tests/test_proxy_utils.py index ee18c96c393..d2218b08386 100644 --- a/tests/proxy_unit_tests/test_proxy_utils.py +++ b/tests/proxy_unit_tests/test_proxy_utils.py @@ -7,10 +7,12 @@ from typing import Any, Dict, List, Optional, Union from unittest.mock import Mock import pytest -from fastapi import Request +from fastapi import HTTPException, Request from starlette.datastructures import State +from litellm.integrations.custom_guardrail import CustomGuardrail from litellm.proxy.utils import _get_docs_url, _get_openapi_url, _get_redoc_url +from litellm.types.guardrails import GuardrailEventHooks sys.path.insert( 0, os.path.abspath("../..") @@ -2638,6 +2640,463 @@ async def test_during_call_hook_parallel_execution_with_error(): litellm.callbacks = original_callbacks +class _PreCallGuardrail(CustomGuardrail): + """Test double for pre_call guardrails; records timing and observed payload.""" + + def __init__(self, name, run_in_parallel, execution_order, sleep=0.1, default_on=True): + super().__init__( + guardrail_name=name, + event_hook=GuardrailEventHooks.pre_call, + default_on=default_on, + run_in_parallel=run_in_parallel, + ) + self.name = name + self.sleep = sleep + self.execution_order = execution_order + self.observed_content = None + self.was_called = False + + async def async_pre_call_hook(self, user_api_key_dict, cache, data, call_type): + self.was_called = True + self.observed_content = data["messages"][0]["content"] + self.execution_order.append(f"{self.name}_start") + await asyncio.sleep(self.sleep) + self.execution_order.append(f"{self.name}_end") + return None + + +@pytest.mark.asyncio +async def test_pre_call_hook_runs_opted_in_guardrails_in_parallel(): + """run_in_parallel pre_call guardrails execute concurrently (all start before any ends).""" + from litellm.caching.caching import DualCache + from litellm.proxy.utils import ProxyLogging + + proxy_logging = ProxyLogging(user_api_key_cache=DualCache()) + execution_order = [] + original_callbacks = litellm.callbacks.copy() if litellm.callbacks else [] + + try: + litellm.callbacks = [ + _PreCallGuardrail(f"g{i}", run_in_parallel=True, execution_order=execution_order) for i in range(3) + ] + + result = await proxy_logging.pre_call_hook( + user_api_key_dict=UserAPIKeyAuth(api_key="test_key", user_id="test_user"), + data={"model": "gpt-4", "messages": [{"role": "user", "content": "hi"}]}, + call_type="completion", + ) + + first_end_idx = next(i for i, item in enumerate(execution_order) if "end" in item) + starts_before_first_end = sum(1 for item in execution_order[:first_end_idx] if "start" in item) + assert starts_before_first_end == 3, f"expected 3 concurrent starts, got {starts_before_first_end}" + assert result["model"] == "gpt-4" + finally: + litellm.callbacks = original_callbacks + + +@pytest.mark.asyncio +async def test_pre_call_hook_runs_default_guardrails_sequentially(): + """Guardrails without run_in_parallel keep the sequential, one-at-a-time behavior.""" + from litellm.caching.caching import DualCache + from litellm.proxy.utils import ProxyLogging + + proxy_logging = ProxyLogging(user_api_key_cache=DualCache()) + execution_order = [] + original_callbacks = litellm.callbacks.copy() if litellm.callbacks else [] + + try: + litellm.callbacks = [ + _PreCallGuardrail(f"g{i}", run_in_parallel=False, execution_order=execution_order) for i in range(2) + ] + + start = asyncio.get_event_loop().time() + await proxy_logging.pre_call_hook( + user_api_key_dict=UserAPIKeyAuth(api_key="test_key", user_id="test_user"), + data={"model": "gpt-4", "messages": [{"role": "user", "content": "hi"}]}, + call_type="completion", + ) + elapsed = asyncio.get_event_loop().time() - start + + assert execution_order == ["g0_start", "g0_end", "g1_start", "g1_end"] + assert elapsed >= 0.18, f"sequential run took {elapsed}s, expected >= 0.18s" + finally: + litellm.callbacks = original_callbacks + + +@pytest.mark.asyncio +async def test_pre_call_hook_sequential_mutations_precede_parallel_batch(): + """Sequential (mutating) guardrails run before the parallel batch, which sees their changes.""" + from litellm.caching.caching import DualCache + from litellm.proxy.utils import ProxyLogging + + proxy_logging = ProxyLogging(user_api_key_cache=DualCache()) + execution_order = [] + original_callbacks = litellm.callbacks.copy() if litellm.callbacks else [] + + class MaskingGuardrail(CustomGuardrail): + def __init__(self): + super().__init__( + guardrail_name="masker", + event_hook=GuardrailEventHooks.pre_call, + default_on=True, + run_in_parallel=False, + ) + + async def async_pre_call_hook(self, user_api_key_dict, cache, data, call_type): + data["messages"][0]["content"] = "MASKED" + return data + + parallel_observer = _PreCallGuardrail("observer", run_in_parallel=True, execution_order=execution_order) + + try: + litellm.callbacks = [parallel_observer, MaskingGuardrail()] + + await proxy_logging.pre_call_hook( + user_api_key_dict=UserAPIKeyAuth(api_key="test_key", user_id="test_user"), + data={"model": "gpt-4", "messages": [{"role": "user", "content": "secret"}]}, + call_type="completion", + ) + + assert parallel_observer.observed_content == "MASKED" + finally: + litellm.callbacks = original_callbacks + + +@pytest.mark.asyncio +async def test_pre_call_hook_parallel_guardrail_blocks_request(): + """A raising parallel guardrail blocks the request before it reaches the LLM.""" + from litellm.caching.caching import DualCache + from litellm.proxy.utils import ProxyLogging + + proxy_logging = ProxyLogging(user_api_key_cache=DualCache()) + original_callbacks = litellm.callbacks.copy() if litellm.callbacks else [] + + class BlockingGuardrail(CustomGuardrail): + def __init__(self): + super().__init__( + guardrail_name="blocker", + event_hook=GuardrailEventHooks.pre_call, + default_on=True, + run_in_parallel=True, + ) + + async def async_pre_call_hook(self, user_api_key_dict, cache, data, call_type): + raise HTTPException(status_code=400, detail="blocked by guardrail") + + try: + litellm.callbacks = [BlockingGuardrail()] + + with pytest.raises(HTTPException) as exc_info: + await proxy_logging.pre_call_hook( + user_api_key_dict=UserAPIKeyAuth(api_key="test_key", user_id="test_user"), + data={"model": "gpt-4", "messages": [{"role": "user", "content": "hi"}]}, + call_type="completion", + ) + + assert exc_info.value.status_code == 400 + assert "blocked by guardrail" in str(exc_info.value.detail) + finally: + litellm.callbacks = original_callbacks + + +@pytest.mark.asyncio +async def test_pre_call_hook_parallel_guardrail_skipped_when_should_not_run(): + """A parallel guardrail that should_run_guardrail rejects is never invoked.""" + from litellm.caching.caching import DualCache + from litellm.proxy.utils import ProxyLogging + + proxy_logging = ProxyLogging(user_api_key_cache=DualCache()) + execution_order = [] + original_callbacks = litellm.callbacks.copy() if litellm.callbacks else [] + + try: + guardrail = _PreCallGuardrail( + "off_by_default", run_in_parallel=True, execution_order=execution_order, default_on=False + ) + litellm.callbacks = [guardrail] + + result = await proxy_logging.pre_call_hook( + user_api_key_dict=UserAPIKeyAuth(api_key="test_key", user_id="test_user"), + data={"model": "gpt-4", "messages": [{"role": "user", "content": "hi"}]}, + call_type="completion", + ) + + assert guardrail.was_called is False + assert result["model"] == "gpt-4" + finally: + litellm.callbacks = original_callbacks + + +@pytest.mark.asyncio +async def test_pre_call_hook_parallel_block_wins_over_reroute(): + """A slower block must win over a faster reroute so crafted input cannot bypass a block.""" + from litellm.caching.caching import DualCache + from litellm.exceptions import SensitiveDataRouteException + from litellm.proxy.utils import ProxyLogging + + proxy_logging = ProxyLogging(user_api_key_cache=DualCache()) + original_callbacks = litellm.callbacks.copy() if litellm.callbacks else [] + + class FastRerouteGuardrail(CustomGuardrail): + def __init__(self): + super().__init__( + guardrail_name="rerouter", + event_hook=GuardrailEventHooks.pre_call, + default_on=True, + run_in_parallel=True, + ) + + async def async_pre_call_hook(self, user_api_key_dict, cache, data, call_type): + raise SensitiveDataRouteException(route_to_model="on-prem", session_id="s1", guardrail_name="rerouter") + + class SlowBlockingGuardrail(CustomGuardrail): + def __init__(self): + super().__init__( + guardrail_name="blocker", + event_hook=GuardrailEventHooks.pre_call, + default_on=True, + run_in_parallel=True, + ) + + async def async_pre_call_hook(self, user_api_key_dict, cache, data, call_type): + await asyncio.sleep(0.1) + raise HTTPException(status_code=400, detail="blocked by guardrail") + + try: + litellm.callbacks = [FastRerouteGuardrail(), SlowBlockingGuardrail()] + + with pytest.raises(HTTPException) as exc_info: + await proxy_logging.pre_call_hook( + user_api_key_dict=UserAPIKeyAuth(api_key="test_key", user_id="test_user"), + data={"model": "gpt-4", "messages": [{"role": "user", "content": "hi"}]}, + call_type="completion", + ) + + assert exc_info.value.status_code == 400 + assert "blocked by guardrail" in str(exc_info.value.detail) + finally: + litellm.callbacks = original_callbacks + + +@pytest.mark.asyncio +async def test_pre_call_hook_parallel_awaits_all_when_one_blocks(): + """A block must not orphan sibling guardrails; every parallel guardrail runs to completion.""" + from litellm.caching.caching import DualCache + from litellm.proxy.utils import ProxyLogging + + proxy_logging = ProxyLogging(user_api_key_cache=DualCache()) + original_callbacks = litellm.callbacks.copy() if litellm.callbacks else [] + completed = [] + + class FastBlockingGuardrail(CustomGuardrail): + def __init__(self): + super().__init__( + guardrail_name="fast_blocker", + event_hook=GuardrailEventHooks.pre_call, + default_on=True, + run_in_parallel=True, + ) + + async def async_pre_call_hook(self, user_api_key_dict, cache, data, call_type): + raise HTTPException(status_code=400, detail="blocked") + + class SlowGuardrail(CustomGuardrail): + def __init__(self): + super().__init__( + guardrail_name="slow", + event_hook=GuardrailEventHooks.pre_call, + default_on=True, + run_in_parallel=True, + ) + + async def async_pre_call_hook(self, user_api_key_dict, cache, data, call_type): + await asyncio.sleep(0.1) + completed.append("slow") + return None + + try: + litellm.callbacks = [FastBlockingGuardrail(), SlowGuardrail()] + + with pytest.raises(HTTPException): + await proxy_logging.pre_call_hook( + user_api_key_dict=UserAPIKeyAuth(api_key="test_key", user_id="test_user"), + data={"model": "gpt-4", "messages": [{"role": "user", "content": "hi"}]}, + call_type="completion", + ) + + assert completed == ["slow"], "slow guardrail was orphaned instead of awaited to completion" + finally: + litellm.callbacks = original_callbacks + + +class _PostCallGuardrail(CustomGuardrail): + """Test double for post_call guardrails; records timing and invocation.""" + + def __init__(self, name, run_in_parallel, execution_order, sleep=0.1, default_on=True): + super().__init__( + guardrail_name=name, + event_hook=GuardrailEventHooks.post_call, + default_on=default_on, + run_in_parallel=run_in_parallel, + ) + self.name = name + self.sleep = sleep + self.execution_order = execution_order + self.was_called = False + + async def async_post_call_success_hook(self, data, user_api_key_dict, response): + self.was_called = True + self.execution_order.append(f"{self.name}_start") + await asyncio.sleep(self.sleep) + self.execution_order.append(f"{self.name}_end") + return None + + +@pytest.mark.asyncio +async def test_post_call_hook_runs_opted_in_guardrails_in_parallel(): + """run_in_parallel post_call guardrails execute concurrently (all start before any ends).""" + from litellm.caching.caching import DualCache + from litellm.proxy.utils import ProxyLogging + + proxy_logging = ProxyLogging(user_api_key_cache=DualCache()) + execution_order = [] + original_callbacks = litellm.callbacks.copy() if litellm.callbacks else [] + + try: + litellm.callbacks = [ + _PostCallGuardrail(f"g{i}", run_in_parallel=True, execution_order=execution_order) for i in range(3) + ] + + await proxy_logging.post_call_success_hook( + data={"model": "gpt-4", "messages": [{"role": "user", "content": "hi"}]}, + response=litellm.ModelResponse(), + user_api_key_dict=UserAPIKeyAuth(api_key="test_key", user_id="test_user"), + ) + + first_end_idx = next(i for i, item in enumerate(execution_order) if "end" in item) + starts_before_first_end = sum(1 for item in execution_order[:first_end_idx] if "start" in item) + assert starts_before_first_end == 3, f"expected 3 concurrent starts, got {starts_before_first_end}" + finally: + litellm.callbacks = original_callbacks + + +@pytest.mark.asyncio +async def test_post_call_hook_runs_default_guardrails_sequentially(): + """post_call guardrails without run_in_parallel keep the sequential behavior.""" + from litellm.caching.caching import DualCache + from litellm.proxy.utils import ProxyLogging + + proxy_logging = ProxyLogging(user_api_key_cache=DualCache()) + execution_order = [] + original_callbacks = litellm.callbacks.copy() if litellm.callbacks else [] + + try: + litellm.callbacks = [ + _PostCallGuardrail(f"g{i}", run_in_parallel=False, execution_order=execution_order) for i in range(2) + ] + + start = asyncio.get_event_loop().time() + await proxy_logging.post_call_success_hook( + data={"model": "gpt-4", "messages": [{"role": "user", "content": "hi"}]}, + response=litellm.ModelResponse(), + user_api_key_dict=UserAPIKeyAuth(api_key="test_key", user_id="test_user"), + ) + elapsed = asyncio.get_event_loop().time() - start + + assert execution_order == ["g0_start", "g0_end", "g1_start", "g1_end"] + assert elapsed >= 0.18, f"sequential run took {elapsed}s, expected >= 0.18s" + finally: + litellm.callbacks = original_callbacks + + +@pytest.mark.asyncio +async def test_post_call_hook_parallel_guardrail_blocks_response(): + """A raising parallel post_call guardrail blocks the response before it reaches the client.""" + from litellm.caching.caching import DualCache + from litellm.proxy.utils import ProxyLogging + + proxy_logging = ProxyLogging(user_api_key_cache=DualCache()) + original_callbacks = litellm.callbacks.copy() if litellm.callbacks else [] + + class BlockingPostCallGuardrail(CustomGuardrail): + def __init__(self): + super().__init__( + guardrail_name="post_blocker", + event_hook=GuardrailEventHooks.post_call, + default_on=True, + run_in_parallel=True, + ) + + async def async_post_call_success_hook(self, data, user_api_key_dict, response): + raise HTTPException(status_code=400, detail="blocked response by guardrail") + + try: + litellm.callbacks = [BlockingPostCallGuardrail()] + + with pytest.raises(HTTPException) as exc_info: + await proxy_logging.post_call_success_hook( + data={"model": "gpt-4", "messages": [{"role": "user", "content": "hi"}]}, + response=litellm.ModelResponse(), + user_api_key_dict=UserAPIKeyAuth(api_key="test_key", user_id="test_user"), + ) + + assert exc_info.value.status_code == 400 + assert "blocked response by guardrail" in str(exc_info.value.detail) + finally: + litellm.callbacks = original_callbacks + + +@pytest.mark.asyncio +async def test_post_call_hook_parallel_awaits_all_when_one_blocks(): + """A blocking post_call guardrail must not orphan its siblings; all run to completion.""" + from litellm.caching.caching import DualCache + from litellm.proxy.utils import ProxyLogging + + proxy_logging = ProxyLogging(user_api_key_cache=DualCache()) + original_callbacks = litellm.callbacks.copy() if litellm.callbacks else [] + completed = [] + + class FastBlockingPostCall(CustomGuardrail): + def __init__(self): + super().__init__( + guardrail_name="fast_post_blocker", + event_hook=GuardrailEventHooks.post_call, + default_on=True, + run_in_parallel=True, + ) + + async def async_post_call_success_hook(self, data, user_api_key_dict, response): + raise HTTPException(status_code=400, detail="blocked") + + class SlowPostCall(CustomGuardrail): + def __init__(self): + super().__init__( + guardrail_name="slow_post", + event_hook=GuardrailEventHooks.post_call, + default_on=True, + run_in_parallel=True, + ) + + async def async_post_call_success_hook(self, data, user_api_key_dict, response): + await asyncio.sleep(0.1) + completed.append("slow") + return None + + try: + litellm.callbacks = [FastBlockingPostCall(), SlowPostCall()] + + with pytest.raises(HTTPException): + await proxy_logging.post_call_success_hook( + data={"model": "gpt-4", "messages": [{"role": "user", "content": "hi"}]}, + response=litellm.ModelResponse(), + user_api_key_dict=UserAPIKeyAuth(api_key="test_key", user_id="test_user"), + ) + + assert completed == ["slow"], "slow post_call guardrail was orphaned instead of awaited to completion" + finally: + litellm.callbacks = original_callbacks + + @pytest.mark.asyncio async def test_handle_logging_proxy_only_error_preserves_pass_through_call_type(): """Ensure _handle_logging_proxy_only_error does not overwrite call_type diff --git a/tests/test_litellm/integrations/test_custom_guardrail.py b/tests/test_litellm/integrations/test_custom_guardrail.py index 64813c1eda7..4ea79f9e2a4 100644 --- a/tests/test_litellm/integrations/test_custom_guardrail.py +++ b/tests/test_litellm/integrations/test_custom_guardrail.py @@ -1,3 +1,4 @@ +import asyncio from unittest.mock import AsyncMock import pytest @@ -1394,6 +1395,38 @@ class TestEventTypeLogging: assert len(logged_info) == 1 assert logged_info[0]["guardrail_status"] == "guardrail_intervened" + @pytest.mark.asyncio + async def test_log_guardrail_information_records_every_concurrent_guardrail(self): + """Guardrails run concurrently (parallel pre_call/post_call, during_call) share one + request_data dict. Each must still record its own entry. The previous guard counted + entries in that shared dict, so a sibling's append made a guardrail think it had already + recorded and skip its own auto-record — silently dropping lifecycle logs the UI shows.""" + from litellm.integrations.custom_guardrail import log_guardrail_information + from litellm.types.guardrails import GuardrailEventHooks + + class SleeperGuardrail(CustomGuardrail): + def __init__(self, name, sleep): + super().__init__(guardrail_name=name, event_hook=GuardrailEventHooks.pre_call) + self._sleep = sleep + + @log_guardrail_information + async def async_pre_call_hook(self, data: dict, **kwargs): + await asyncio.sleep(self._sleep) + return data + + request_data = {"metadata": {}} + # Different sleeps guarantee overlapping execution windows: the faster guardrail + # records while the slower one is still awaiting, which is exactly what tripped the + # old shared-count guard. + await asyncio.gather( + SleeperGuardrail("guardrail-a", 0.05).async_pre_call_hook(data=request_data), + SleeperGuardrail("guardrail-b", 0.15).async_pre_call_hook(data=request_data), + ) + + logged = request_data["metadata"]["standard_logging_guardrail_information"] + assert {entry["guardrail_name"] for entry in logged} == {"guardrail-a", "guardrail-b"} + assert len(logged) == 2 + def test_add_standard_logging_falls_back_to_event_hook_when_event_type_is_none( self, ): diff --git a/tests/test_litellm/proxy/guardrails/test_guardrail_registry.py b/tests/test_litellm/proxy/guardrails/test_guardrail_registry.py index 26feddadf79..14cab50f441 100644 --- a/tests/test_litellm/proxy/guardrails/test_guardrail_registry.py +++ b/tests/test_litellm/proxy/guardrails/test_guardrail_registry.py @@ -1,3 +1,5 @@ +import pytest + from litellm.integrations.custom_guardrail import CustomGuardrail from litellm.proxy.guardrails.guardrail_registry import ( get_guardrail_initializer_from_hooks, @@ -32,6 +34,44 @@ def test_noma_registry_resolution(): assert "noma_v2" in guardrail_initializer_registry +@pytest.mark.parametrize( + "configured, expected", + [(None, True), (False, False), (True, True)], +) +def test_initialize_guardrail_run_in_parallel_preserves_constructor_default(configured, expected): + """ + A guardrail whose constructor sets run_in_parallel=True must keep that default when + the config omits the key; only an explicit config value may override it. The + previous code wrote bool(None)==False on every instance, silently disabling the + opt-in for such guardrails. + """ + from litellm.proxy.guardrails import guardrail_registry as registry_module + + def _initializer(litellm_params, guardrail): + return CustomGuardrail( + guardrail_name=guardrail["guardrail_name"], + event_hook=GuardrailEventHooks.pre_call, + default_on=True, + run_in_parallel=True, + ) + + registry_module.guardrail_initializer_registry["parallel_default_test"] = _initializer + try: + params = {"guardrail": "parallel_default_test", "mode": "pre_call"} + if configured is not None: + params["run_in_parallel"] = configured + + handler = InMemoryGuardrailHandler() + result = handler.initialize_guardrail( + guardrail={"guardrail_name": "cf-parallel-default", "litellm_params": params}, + ) + + stored = handler.guardrail_id_to_custom_guardrail[result["guardrail_id"]] + assert stored.run_in_parallel is expected + finally: + registry_module.guardrail_initializer_registry.pop("parallel_default_test", None) + + def test_update_in_memory_guardrail(): handler = InMemoryGuardrailHandler() handler.guardrail_id_to_custom_guardrail["123"] = CustomGuardrail( diff --git a/tests/test_litellm/proxy/guardrails/test_init_guardrails.py b/tests/test_litellm/proxy/guardrails/test_init_guardrails.py index 83593c20110..71e775842e3 100644 --- a/tests/test_litellm/proxy/guardrails/test_init_guardrails.py +++ b/tests/test_litellm/proxy/guardrails/test_init_guardrails.py @@ -62,3 +62,27 @@ def test_initialize_guardrail_preserves_guardrail_info(): assert result["guardrail_info"] == {"type": "PII", "description": "masks PII"} stored = guardrail_handler.IN_MEMORY_GUARDRAILS[result["guardrail_id"]] assert stored["guardrail_info"] == {"type": "PII", "description": "masks PII"} + + +@pytest.mark.parametrize( + "config_value, expected", + [(True, True), (False, False), (None, False)], +) +def test_initialize_guardrail_sets_run_in_parallel(config_value, expected): + """run_in_parallel from litellm_params must reach the built guardrail instance.""" + litellm_params = { + "guardrail": SupportedGuardrailIntegrations.PRESIDIO.value, + "mode": "pre_call", + "presidio_analyzer_api_base": "https://fakelink.com/v1/presidio/analyze", + "presidio_anonymizer_api_base": "https://fakelink.com/v1/presidio/anonymize", + } + if config_value is not None: + litellm_params["run_in_parallel"] = config_value + + guardrail_handler = InMemoryGuardrailHandler() + result = guardrail_handler.initialize_guardrail( + guardrail={"guardrail_name": "test_parallel_flag", "litellm_params": litellm_params}, + ) + + custom_guardrail = guardrail_handler.guardrail_id_to_custom_guardrail[result["guardrail_id"]] + assert custom_guardrail.run_in_parallel is expected diff --git a/tests/test_litellm/proxy/utils/proxy_logging/test_guardrail_pipeline.py b/tests/test_litellm/proxy/utils/proxy_logging/test_guardrail_pipeline.py index 64c14abfd83..5c711fc6c34 100644 --- a/tests/test_litellm/proxy/utils/proxy_logging/test_guardrail_pipeline.py +++ b/tests/test_litellm/proxy/utils/proxy_logging/test_guardrail_pipeline.py @@ -626,6 +626,7 @@ def _moderation_guardrail() -> MagicMock: cb.should_run_guardrail = MagicMock(return_value=True) cb.async_moderation_hook = AsyncMock(return_value=None) cb.async_post_call_success_hook = AsyncMock(return_value=None) + cb.run_in_parallel = False return cb diff --git a/tests/test_litellm/proxy/utils/proxy_logging/test_post_call_success_hook.py b/tests/test_litellm/proxy/utils/proxy_logging/test_post_call_success_hook.py index 6a339b37a80..715d66db181 100644 --- a/tests/test_litellm/proxy/utils/proxy_logging/test_post_call_success_hook.py +++ b/tests/test_litellm/proxy/utils/proxy_logging/test_post_call_success_hook.py @@ -28,6 +28,7 @@ def _make_guardrail(name="g", should_run=True, override=None): cb.event_hook = GuardrailEventHooks.post_call cb.should_run_guardrail = MagicMock(return_value=should_run) cb.async_post_call_success_hook = AsyncMock(return_value=override) + cb.run_in_parallel = False return cb From 7d9eec623081f72abeb3770001828ab24b490503 Mon Sep 17 00:00:00 2001 From: shivam Date: Fri, 24 Jul 2026 20:32:09 +0000 Subject: [PATCH 06/29] fix(proxy): return 400 instead of 500 for chat completions without messages Router.acompletion() takes messages positionally, so splatting a body that omits it raised a TypeError that the generic handler mapped to a 500. Validate the required body param at the routing boundary and raise the existing 400 contract instead. Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/route_llm_request.py | 28 ++++++- .../proxy/test_route_llm_request.py | 77 ++++++++++++++++--- 2 files changed, 94 insertions(+), 11 deletions(-) diff --git a/litellm/proxy/route_llm_request.py b/litellm/proxy/route_llm_request.py index 25fa0819930..1f5aacc2115 100644 --- a/litellm/proxy/route_llm_request.py +++ b/litellm/proxy/route_llm_request.py @@ -1,5 +1,5 @@ import asyncio -from typing import TYPE_CHECKING, Any, Literal, Optional +from typing import TYPE_CHECKING, Any, Literal, Mapping, Optional import httpx from fastapi import HTTPException, status @@ -145,6 +145,30 @@ class ProxyModelNotFoundError(HTTPException): super().__init__(status_code=status.HTTP_400_BAD_REQUEST, detail=detail) +REQUIRED_BODY_PARAM_BY_ROUTE: Mapping[str, str] = { + "acompletion": "messages", + "aembedding": "input", +} + + +class ProxyMissingRequiredParamError(HTTPException): + def __init__(self, route: str, param: str): + detail = {"error": f"{route}: Missing required parameter: '{param}'."} + super().__init__(status_code=status.HTTP_400_BAD_REQUEST, detail=detail) + self.type = "invalid_request_error" + self.param = param + + +def raise_if_required_body_param_missing(route_type: str, data: Mapping[str, object]) -> None: + required_param = REQUIRED_BODY_PARAM_BY_ROUTE.get(route_type) + if required_param is None or data.get(required_param) is not None: + return + raise ProxyMissingRequiredParamError( + route=ROUTE_ENDPOINT_MAPPING.get(route_type, route_type), + param=required_param, + ) + + def get_team_id_from_data(data: dict) -> Optional[str]: """ Get the team id from the data's metadata or litellm_metadata params. @@ -353,6 +377,8 @@ async def route_request( """ Common helper to route the request """ + raise_if_required_body_param_missing(route_type=route_type, data=data) + await add_shared_session_to_data(data) # Strip router-internal mock_testing_* flags. Combined with an diff --git a/tests/test_litellm/proxy/test_route_llm_request.py b/tests/test_litellm/proxy/test_route_llm_request.py index f506b9665a6..93b3ef1cce8 100644 --- a/tests/test_litellm/proxy/test_route_llm_request.py +++ b/tests/test_litellm/proxy/test_route_llm_request.py @@ -12,24 +12,25 @@ from litellm.proxy.route_llm_request import ProxyModelNotFoundError, route_reque @pytest.mark.parametrize( - "route_type", + "route_type, required_body_params", [ - "atext_completion", - "acompletion", - "aembedding", - "aimage_generation", - "aspeech", - "atranscription", - "amoderation", - "arerank", + ("atext_completion", {}), + ("acompletion", {"messages": [{"role": "user", "content": "Hello"}]}), + ("aembedding", {"input": "Hello"}), + ("aimage_generation", {}), + ("aspeech", {}), + ("atranscription", {}), + ("amoderation", {}), + ("arerank", {}), ], ) @pytest.mark.asyncio -async def test_route_request_dynamic_credentials(route_type): +async def test_route_request_dynamic_credentials(route_type, required_body_params): data = { "model": "openai/gpt-4o-mini-2024-07-18", "api_key": "my-bad-key", "api_base": "https://api.openai.com/v1 ", + **required_body_params, } llm_router = MagicMock() # Ensure that the dynamic method exists on the llm_router mock. @@ -887,3 +888,59 @@ async def test_route_request_override_enable_tag_filtering_beats_body_value(): call_kwargs = llm_router.acompletion.call_args[1] assert call_kwargs["enable_tag_filtering"] is True + + +@pytest.mark.parametrize( + "route_type, param, route", + [ + ("acompletion", "messages", "/chat/completions"), + ("aembedding", "input", "/embeddings"), + ], +) +@pytest.mark.parametrize("data_extra", [{}, {"messages": None, "input": None}]) +def test_raise_if_required_body_param_missing_rejects_missing_param(route_type, param, route, data_extra): + from litellm.proxy.route_llm_request import ( + ProxyMissingRequiredParamError, + raise_if_required_body_param_missing, + ) + + with pytest.raises(ProxyMissingRequiredParamError) as exc_info: + raise_if_required_body_param_missing(route_type=route_type, data={"model": "gpt-4o", **data_extra}) + + assert exc_info.value.status_code == 400 + assert exc_info.value.param == param + assert exc_info.value.type == "invalid_request_error" + assert exc_info.value.detail == {"error": f"{route}: Missing required parameter: '{param}'."} + + +@pytest.mark.parametrize( + "route_type, data", + [ + ("acompletion", {"model": "gpt-4o", "messages": [{"role": "user", "content": "hi"}]}), + ("acompletion", {"model": "gpt-4o", "messages": []}), + ("atext_completion", {"model": "gpt-4o"}), + ("aembedding", {"model": "text-embedding-3-small", "input": "hi"}), + ("arerank", {"model": "rerank-model"}), + ("aimage_generation", {"model": "dall-e-3"}), + ], +) +def test_raise_if_required_body_param_missing_allows_valid_requests(route_type, data): + from litellm.proxy.route_llm_request import raise_if_required_body_param_missing + + raise_if_required_body_param_missing(route_type=route_type, data=data) + + +@pytest.mark.asyncio +async def test_route_request_rejects_chat_completion_without_messages(): + """A /chat/completions body without `messages` used to splat into + Router.acompletion() and surface the resulting TypeError as a 500.""" + from litellm.proxy.route_llm_request import ProxyMissingRequiredParamError + + llm_router = MagicMock() + + with pytest.raises(ProxyMissingRequiredParamError) as exc_info: + await route_request({"model": "gpt-4o"}, llm_router, None, "acompletion") + + assert exc_info.value.status_code == 400 + assert exc_info.value.param == "messages" + llm_router.acompletion.assert_not_called() From a376f724002185917340d7d74252b5ef1894c5cb Mon Sep 17 00:00:00 2001 From: shivam Date: Fri, 24 Jul 2026 20:53:05 +0000 Subject: [PATCH 07/29] fix(responses): stop treating stream_options as a Responses API param Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/types/llms/openai.py | 1 - ...erimental_pass_through_messages_handler.py | 64 +++++++++++++++++++ .../test_responses_api_request_body.py | 27 ++++++++ 3 files changed, 91 insertions(+), 1 deletion(-) diff --git a/litellm/types/llms/openai.py b/litellm/types/llms/openai.py index 9f689a2dd31..2263d53182a 100644 --- a/litellm/types/llms/openai.py +++ b/litellm/types/llms/openai.py @@ -1171,7 +1171,6 @@ class ResponsesAPIOptionalRequestParams(TypedDict, total=False): max_tool_calls: Optional[int] prompt_cache_key: Optional[str] prompt_cache_retention: Optional[str] - stream_options: Optional[dict] top_logprobs: Optional[int] partial_images: Optional[int] # Number of partial images to generate (1-3) for streaming image generation context_management: Optional[List[ContextManagementEntry]] diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_anthropic_experimental_pass_through_messages_handler.py b/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_anthropic_experimental_pass_through_messages_handler.py index 3327fc39f73..8875a75e86f 100644 --- a/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_anthropic_experimental_pass_through_messages_handler.py +++ b/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_anthropic_experimental_pass_through_messages_handler.py @@ -2,6 +2,7 @@ import json import os import sys +import httpx import pytest from fastapi.testclient import TestClient @@ -9,6 +10,7 @@ sys.path.insert(0, os.path.abspath("../../../../..")) from unittest.mock import AsyncMock, MagicMock, patch +import litellm from litellm.anthropic_interface import messages from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler from litellm.types.utils import Delta, ModelResponse, StreamingChoices @@ -37,6 +39,68 @@ def test_anthropic_experimental_pass_through_messages_handler(): assert mock_responses.call_args.kwargs["api_key"] == "test-api-key" +@pytest.mark.asyncio +async def test_openai_model_does_not_forward_stream_options_to_responses_api(): + """ + Regression test for LIT-4779. `always_include_stream_usage` injects + stream_options={'include_usage': True} into every streaming request, but OpenAI + models on /v1/messages go to the Responses API, which 400s on that param. + """ + responses_payload = { + "id": "resp_stream_options", + "object": "response", + "created_at": 1734366691, + "status": "completed", + "model": "gpt-5.5", + "output": [ + { + "type": "message", + "id": "msg_1", + "status": "completed", + "role": "assistant", + "content": [{"type": "output_text", "text": "hi", "annotations": []}], + } + ], + "parallel_tool_calls": True, + "usage": {"input_tokens": 1, "output_tokens": 1, "total_tokens": 2}, + "error": None, + "incomplete_details": None, + "instructions": None, + "metadata": None, + "temperature": None, + "tool_choice": "auto", + "tools": [], + "top_p": None, + "max_output_tokens": None, + "previous_response_id": None, + "reasoning": None, + "truncation": None, + "user": None, + } + + mock_response = MagicMock() + mock_response.status_code = 200 + mock_response.text = json.dumps(responses_payload) + mock_response.headers = httpx.Headers({}) + mock_response.json.return_value = responses_payload + + with patch.object(AsyncHTTPHandler, "post", new_callable=AsyncMock) as mock_post: + mock_post.return_value = mock_response + + await litellm.anthropic.messages.acreate( + max_tokens=100, + messages=[{"role": "user", "content": "Hello, how are you?"}], + model="openai/gpt-5.5", + api_key="test-api-key", + stream_options={"include_usage": True}, + ) + + mock_post.assert_called_once() + post_kwargs = mock_post.call_args.kwargs + request_body = post_kwargs["json"] if "json" in post_kwargs else json.loads(post_kwargs["data"]) + assert "stream_options" not in request_body + + def test_anthropic_experimental_pass_through_messages_handler_dynamic_api_key_and_api_base_and_custom_values(): """ Test that api key, api base, and extra kwargs are forwarded to litellm.completion for Azure models. diff --git a/tests/test_litellm/responses/test_responses_api_request_body.py b/tests/test_litellm/responses/test_responses_api_request_body.py index 44dfa240d42..2922c9738aa 100644 --- a/tests/test_litellm/responses/test_responses_api_request_body.py +++ b/tests/test_litellm/responses/test_responses_api_request_body.py @@ -198,6 +198,33 @@ async def test_aresponses_azure_shell_tool_400_maps_to_bad_request_error(): assert "not supported" in str(excinfo.value).lower() +@pytest.mark.asyncio +async def test_aresponses_drops_stream_options(): + """ + stream_options is a Chat Completions param; the Responses API rejects it with + "Unknown parameter: 'stream_options.include_usage'". It must never reach the wire. + """ + with patch( + "litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post", + new_callable=AsyncMock, + ) as mock_post: + mock_post.return_value = MockResponse( + _minimal_responses_api_payload("resp_stream_options_test", "gpt-5.5"), 200 + ) + + await litellm.aresponses( + model="openai/gpt-5.5", + api_key="fake-api-key", + input="hi", + stream_options={"include_usage": True}, + ) + + mock_post.assert_called_once() + post_kwargs = mock_post.call_args.kwargs + request_body = post_kwargs["json"] if "json" in post_kwargs else json.loads(post_kwargs["data"]) + assert "stream_options" not in request_body + + @pytest.mark.asyncio async def test_aresponses_request_level_drop_params_drops_bedrock_mantle_service_tier( monkeypatch, From 61d32c9aac791918b2fd262241930149f4d1fd1a Mon Sep 17 00:00:00 2001 From: yucheng-berri Date: Fri, 24 Jul 2026 15:10:16 -0700 Subject: [PATCH 08/29] fix: handle explicit outputInfo: null in Vertex AI batch response (#34473) * fix: handle explicit outputInfo: null in Vertex AI batch response Vertex AI can return HTTP 200 for a create_batch/get_batch call with an explicit "outputInfo": null body (the output directory is assigned asynchronously and may not be populated yet at response time). _get_output_file_id_from_vertex_ai_batch_response did: response.get("outputInfo", OutputInfo()).get("gcsOutputDirectory", "") dict.get(key, default) only substitutes default when the key is absent, not when it is present but explicitly None, so this crashed with: AttributeError: 'NoneType' object has no attribute 'get' surfaced to callers as an opaque openai.InternalServerError 500 from litellm.create_batch()/retrieve_batch() for any Vertex AI batch job, regardless of whether the job ultimately succeeds. Fixed by guarding with `response.get("outputInfo") or OutputInfo()`, matching the existing null-safe pattern already used by the sibling _get_input_file_id_from_vertex_ai_batch_response for inputConfig. The existing outputConfig fallback branch (a few lines below) already handles this case correctly once it's reachable - it just never was. Added 2 regression tests covering outputInfo: null with and without an outputConfig fallback available. * test: drop explanatory comment from regression test --------- Co-authored-by: htourinho-clgx --- litellm/llms/vertex_ai/batches/transformation.py | 3 ++- .../llms/vertex_ai/batches/test_transformation.py | 12 ++++++++++++ 2 files changed, 14 insertions(+), 1 deletion(-) diff --git a/litellm/llms/vertex_ai/batches/transformation.py b/litellm/llms/vertex_ai/batches/transformation.py index 6bbe8f75701..df903ba7ef0 100644 --- a/litellm/llms/vertex_ai/batches/transformation.py +++ b/litellm/llms/vertex_ai/batches/transformation.py @@ -123,7 +123,8 @@ class VertexAIBatchTransformation: Gets the output file id from the Vertex AI Batch response """ - output_file_id: str = response.get("outputInfo", OutputInfo()).get("gcsOutputDirectory", "") + output_info = response.get("outputInfo") or OutputInfo() + output_file_id: str = output_info.get("gcsOutputDirectory", "") if output_file_id: output_file_id = output_file_id.rstrip("/") + "/predictions.jsonl" if output_file_id and output_file_id != "/predictions.jsonl": diff --git a/tests/test_litellm/llms/vertex_ai/batches/test_transformation.py b/tests/test_litellm/llms/vertex_ai/batches/test_transformation.py index 71da1d39876..1b37ade6b30 100644 --- a/tests/test_litellm/llms/vertex_ai/batches/test_transformation.py +++ b/tests/test_litellm/llms/vertex_ai/batches/test_transformation.py @@ -226,6 +226,18 @@ def test_get_output_file_id_empty_output_info_falls_through_to_output_config(): assert T._get_output_file_id_from_vertex_ai_batch_response(resp) == "gs://b/cfg/predictions.jsonl" +def test_get_output_file_id_output_info_explicit_none_falls_through_to_output_config(): + resp = { + "outputInfo": None, + "outputConfig": {"gcsDestination": {"outputUriPrefix": "gs://b/cfg"}}, + } + assert T._get_output_file_id_from_vertex_ai_batch_response(resp) == "gs://b/cfg/predictions.jsonl" + + +def test_get_output_file_id_output_info_explicit_none_and_no_output_config(): + assert T._get_output_file_id_from_vertex_ai_batch_response({"outputInfo": None}) == "" + + def test_get_output_file_id_no_output_info_and_no_output_config(): assert T._get_output_file_id_from_vertex_ai_batch_response({}) == "" From 7263aa00281ee6303a86993e8f1aa05afa0cf034 Mon Sep 17 00:00:00 2001 From: Yassin Kortam Date: Fri, 24 Jul 2026 15:11:00 -0700 Subject: [PATCH 09/29] fix(otel): keep an MCP tool call in one trace, anchored to its own request (#34537) Under otel_v2 a single MCP tool call surfaced in APM as two disconnected traces joined only by a span link: the HTTP transport transaction POST /{mcp_server_name}/mcp and the tools/call span carrying error.type=MCPToolResultError. resolve_mcp_span_context parented the MCP span to the W3C trace context the client propagates in params._meta (SEP-414) and recorded the transport as a link, so with no traceparent propagated (the common case today, including MCP Inspector) the span started its own root trace. Nest the MCP span under the transport span when nothing is propagated, so the call stays in one trace; the propagated-context path is unchanged and still parents to the remote context and links the transport per the OTel GenAI MCP semconv. The transport has to be resolved per message rather than read from the request-root ContextVar. A stateful streamable-HTTP session runs every message on the single task the session's initialize POST spawned, so that ContextVar is frozen at initialize inside the handler: live capture on staging showed the tools/call span linking the initialize POST rather than the POST that carried it, and nesting on that anchor would hang every tool call of a session off the first request's already-ended span. The gateway now resolves the current request's span on the ASGI task and carries it to the handler on the authenticated-user object, the same way per-request auth already crosses that boundary. --- litellm/integrations/otel/model/spans.py | 27 ++-- litellm/integrations/otel/plumbing/context.py | 110 +++++++++++--- .../mcp_server/auth/litellm_auth_handler.py | 9 +- .../proxy/_experimental/mcp_server/server.py | 92 +++++++++++- .../integrations/otel/test_otel_v2_logger.py | 136 +++++++++++++++--- 5 files changed, 323 insertions(+), 51 deletions(-) diff --git a/litellm/integrations/otel/model/spans.py b/litellm/integrations/otel/model/spans.py index c93f95ec97d..fa41070a8de 100644 --- a/litellm/integrations/otel/model/spans.py +++ b/litellm/integrations/otel/model/spans.py @@ -18,12 +18,14 @@ before the LLM call even starts), so a guardrail is a sibling of the LLM call, not a child of it. The emitter parents every span to the ambient OTel context (the active server span), which matches this. -MCP spans (``MCP_TOOL_CALL``, ``MCP_LIST_TOOLS``) are intentionally NOT in this -tree. Per the OTel GenAI MCP semconv, MCP and the HTTP transport are independent -contexts, so an MCP span parents to the trace context the client propagated in -``params._meta`` (or starts its own root when none is propagated) and records the -``PROXY_REQUEST`` transport span as a span *link*, never a parent. The registry -encodes this as ``parent=None, links=PROXY_REQUEST``. +MCP spans (``MCP_TOOL_CALL``, ``MCP_LIST_TOOLS``) have two shapes, chosen at emit +time by :func:`resolve_mcp_span_context`. When the client propagates trace context +in ``params._meta`` MCP and the HTTP transport are independent contexts per the +OTel GenAI MCP semconv, so the span parents to that propagated context and records +the ``PROXY_REQUEST`` transport span as a span *link*, never a parent — the shape +this registry's ``parent=None, links=PROXY_REQUEST`` entry encodes. When nothing is +propagated (the common case) the span nests under the transport span of the request +carrying that message, so the tool call stays in one trace. Not every service call becomes a span — :func:`span_role_for_service` decides: @@ -89,12 +91,13 @@ class SpanSpec: SPAN_REGISTRY: dict[SpanRole, SpanSpec] = { SpanRole.PROXY_REQUEST: SpanSpec(SpanRole.PROXY_REQUEST, LiteLLMSpanKind.SERVER, parent=None), SpanRole.LLM_CALL: SpanSpec(SpanRole.LLM_CALL, LiteLLMSpanKind.CLIENT, parent=SpanRole.PROXY_REQUEST), - # MCP and the HTTP transport are independent contexts (OTel GenAI MCP semconv), - # so an MCP span does not nest under the transport span. The proxy is an MCP - # client to the upstream server, so it's a CLIENT span; it parents to the trace - # context the client propagated in ``params._meta`` (or starts its own root when - # none is propagated) and records the PROXY_REQUEST transport span as a span - # *link*, never a parent — hence ``parent=None, links=PROXY_REQUEST``. + # The proxy is an MCP client to the upstream server, so MCP spans are CLIENT + # spans. With trace context propagated in ``params._meta``, MCP and the HTTP + # transport are independent contexts (OTel GenAI MCP semconv): the span parents + # to the propagated context and records the PROXY_REQUEST transport span as a + # span *link*, never a parent — the shape ``parent=None, links=PROXY_REQUEST`` + # encodes. With nothing propagated, ``resolve_mcp_span_context`` nests the span + # under that message's transport span instead, keeping the call in one trace. SpanRole.MCP_TOOL_CALL: SpanSpec( SpanRole.MCP_TOOL_CALL, LiteLLMSpanKind.CLIENT, parent=None, links=SpanRole.PROXY_REQUEST ), diff --git a/litellm/integrations/otel/plumbing/context.py b/litellm/integrations/otel/plumbing/context.py index 8acac112c3d..939559347b1 100644 --- a/litellm/integrations/otel/plumbing/context.py +++ b/litellm/integrations/otel/plumbing/context.py @@ -5,7 +5,14 @@ from typing import Mapping from opentelemetry import baggage from opentelemetry.context import Context, get_current -from opentelemetry.trace import Link, Span, get_current_span, set_span_in_context +from opentelemetry.trace import ( + Link, + NonRecordingSpan, + Span, + SpanContext, + get_current_span, + set_span_in_context, +) from opentelemetry.trace.propagation.tracecontext import ( TraceContextTextMapPropagator, ) @@ -72,6 +79,62 @@ def reset_mcp_message_trace_carrier(token: "Token[Mapping[str, str] | None]") -> _mcp_message_trace_carrier.reset(token) +# The transport span of the HTTP request carrying the CURRENT MCP message, as a +# plain ``SpanContext`` so it can cross a task boundary. +# +# ``_request_root_span`` above cannot be used for MCP: a *stateful* streamable-HTTP +# session runs every message on the single task spawned by that session's +# ``initialize`` POST, so the ContextVar the ASGI request task writes at auth time +# is frozen at ``initialize`` there and never sees the later ``tools/call`` POSTs. +# Reading it from the message handler would parent every tool call in the session +# to the first request's (already ended) server span. The gateway instead resolves +# the current message's transport span on the request task and hands it over the +# same way it hands over per-request auth, and the handler publishes it here for +# the span emitter to pick up. +_mcp_message_transport_span_context: "ContextVar[SpanContext | None]" = ContextVar( + "litellm_otel_mcp_message_transport_span_context", default=None +) + + +def set_mcp_message_transport_span_context( + span_context: "SpanContext | None", +) -> "Token[SpanContext | None]": + """Publish the transport span of the request carrying the current MCP message. + + Returns the reset token; the caller must reset it once the message is handled + so the transport never leaks to the next message on the same session task. + """ + return _mcp_message_transport_span_context.set(span_context) + + +def reset_mcp_message_transport_span_context(token: "Token[SpanContext | None]") -> None: + _mcp_message_transport_span_context.reset(token) + + +def request_root_span_context() -> "SpanContext | None": + """The anchored request root span's context, safe to hand to another task. + + A ``SpanContext`` is an immutable value, unlike the live ``Span``, so passing it + across the MCP session-task boundary cannot keep a finished span alive or invite + writes to it from the wrong request. + """ + span = request_root_span() + return span.get_span_context() if span is not None else None + + +def _mcp_transport_span_context() -> "SpanContext | None": + """The transport span an MCP message span should attach to. + + Prefers the transport the gateway published for this specific message; falls + back to the ambient request anchor for paths that emit an MCP span on the + request task itself (the REST MCP endpoints, the SDK). + """ + published = _mcp_message_transport_span_context.get() + if published is not None and published.is_valid: + return published + return request_root_span_context() + + def set_request_baggage(values: Mapping[str, str], context: Context | None = None) -> Context: """Return a context with ``values`` written into Baggage.""" ctx = context @@ -132,33 +195,44 @@ def resolve_request_span_context() -> Context: def resolve_mcp_span_context( carrier: "Mapping[str, str] | None" = None, ) -> "tuple[Context, tuple[Link, ...]]": - """Parent context + links for an MCP message span, per the OTel GenAI MCP semconv. + """Parent context + links for an MCP message span. - MCP and the underlying transport (HTTP) are independent lifecycles — one - streamable-HTTP session multiplexes many messages, so nesting the message span - under the HTTP/session span is wrong (it renders the message at the session's - start, skewed by however long the session has been open). Instead: + When the client propagates W3C trace context in the request's ``params._meta`` + (SEP-414), MCP and the underlying transport are independent lifecycles — one + streamable-HTTP session multiplexes many messages, and the client's own span is + the truthful parent. So, per the OTel GenAI MCP semconv: - * parent to the trace context the client propagated in the request's - ``params._meta`` (a *remote* parent), and - * record the transport/session span as a *link*, never the parent. + * parent to the trace context the client propagated (a *remote* parent), and + * record the transport span as a *link*, never the parent. + + Almost no client implements SEP-414 yet, so in practice nothing is propagated. + Rooting the span there splits a single tool call into two disconnected traces + joined only by a link, which is how it surfaces in APM: the ``POST`` transaction + and the ``tools/call`` span share no trace. With no remote parent to honor, + parent to the transport span of the request carrying this message instead, so + the call stays in one trace; no link is added since the transport is now the + real parent. The transport comes from :func:`_mcp_transport_span_context`, which + is the *current message's* POST rather than whatever request happened to open + the session, so a long-lived session does not glue every message under its + first request. With neither a remote parent nor a transport the returned context + carries no span and the span legitimately starts its own root trace. Only trace context (``traceparent``/``tracestate``) is extracted, never the client's W3C Baggage: ``params._meta`` is caller-controlled, and the otel baggage processor stamps allowlisted baggage keys (``litellm.team.id``, ``litellm.metadata.*``, ...) onto the span as attributes, so honoring remote - baggage would let a client spoof a span's identity attribution. - - With no propagated context the returned context carries no span, so the span - starts its own root trace (still linked to the transport). The base context is - explicitly empty so an absent ``traceparent`` can never fall through to the - ambient (stale session) span. + baggage would let a client spoof a span's identity attribution. The base context + for extraction is explicitly empty so an absent or malformed ``traceparent`` can + never fall through to the ambient (stale session) span. """ source = carrier if carrier is not None else _mcp_message_trace_carrier.get() parent = _PROPAGATOR.extract(dict(source or {}), context=Context()) - transport = request_root_span() - links = (Link(transport.get_span_context()),) if transport is not None else () - return parent, links + transport = _mcp_transport_span_context() + if is_recordable_span(get_current_span(parent)): + return parent, (Link(transport),) if transport is not None else () + if transport is not None: + return context_from_span(NonRecordingSpan(transport)), () + return parent, () def is_recordable_span(obj: object) -> bool: diff --git a/litellm/proxy/_experimental/mcp_server/auth/litellm_auth_handler.py b/litellm/proxy/_experimental/mcp_server/auth/litellm_auth_handler.py index 7122c64ec64..f7bc14575c7 100644 --- a/litellm/proxy/_experimental/mcp_server/auth/litellm_auth_handler.py +++ b/litellm/proxy/_experimental/mcp_server/auth/litellm_auth_handler.py @@ -1,9 +1,12 @@ -from typing import Dict, List, Optional +from typing import TYPE_CHECKING, Dict, List, Optional from mcp.server.auth.middleware.bearer_auth import AuthenticatedUser from litellm.proxy._types import UserAPIKeyAuth +if TYPE_CHECKING: + from opentelemetry.trace import SpanContext + class MCPAuthenticatedUser(AuthenticatedUser): """ @@ -16,6 +19,8 @@ class MCPAuthenticatedUser(AuthenticatedUser): 4. Server-specific authentication headers 5. OAuth2 headers 6. Raw headers - allows forwarding specific headers to the MCP server, specified by the admin. + 7. Transport span context - the tracing span of the HTTP request carrying the current + message, which a stateful session's message handler cannot read from its own task. """ def __init__( @@ -28,6 +33,7 @@ class MCPAuthenticatedUser(AuthenticatedUser): mcp_protocol_version: Optional[str] = None, raw_headers: Optional[Dict[str, str]] = None, client_ip: Optional[str] = None, + transport_span_context: Optional["SpanContext"] = None, ): self.user_api_key_auth = user_api_key_auth self.mcp_auth_header = mcp_auth_header @@ -37,3 +43,4 @@ class MCPAuthenticatedUser(AuthenticatedUser): self.oauth2_headers = oauth2_headers self.raw_headers = raw_headers self.client_ip = client_ip + self.transport_span_context = transport_span_context diff --git a/litellm/proxy/_experimental/mcp_server/server.py b/litellm/proxy/_experimental/mcp_server/server.py index 4fca4406a6f..483f57f9139 100644 --- a/litellm/proxy/_experimental/mcp_server/server.py +++ b/litellm/proxy/_experimental/mcp_server/server.py @@ -15,6 +15,7 @@ import types import uuid from datetime import datetime from typing import ( + TYPE_CHECKING, Any, AsyncIterator, Callable, @@ -107,6 +108,9 @@ _MAX_STATEFUL_SESSIONS_PER_OWNER = 100 # arbitrarily large body just to make a routing decision. _MCP_ROUTING_PEEK_MAX_BYTES = 4096 +if TYPE_CHECKING: + from opentelemetry.trace import SpanContext + def _invalidate_byok_cred_cache(user_id: str, server_id: str) -> None: """Remove a (user_id, server_id) entry from the BYOK credential cache. @@ -242,10 +246,12 @@ def _mcp_meta_trace_carrier(req_ctx: object) -> Optional[dict[str, str]]: """The W3C trace context (``traceparent``/``tracestate``) the MCP client propagated in the request's ``params._meta`` (SEP-414), or ``None``. - Per the OTel MCP semconv the MCP span parents to this propagated context rather - than to the HTTP/session transport (which is recorded as a link instead), so a - streamable-HTTP session that multiplexes many messages does not glue every - message under the session's first request. The client's W3C Baggage is + When present, per the OTel MCP semconv the MCP span parents to this propagated + context rather than to the HTTP transport (which is recorded as a link instead). + When absent, the span nests under the transport span of the request carrying + this specific message, so a streamable-HTTP session that multiplexes many + messages still does not glue every message under the session's first request; + see ``resolve_mcp_span_context``. The client's W3C Baggage is deliberately excluded: it is caller-controlled, and the otel baggage processor stamps allowlisted baggage keys (``litellm.team.id``, ``litellm.metadata.*``, ...) onto the span, so honoring remote baggage would let a client spoof a @@ -288,6 +294,56 @@ def _otel_reset_mcp_trace_carrier(token: object) -> None: return +def _otel_request_transport_span_context() -> Optional["SpanContext"]: + """The tracing span of the HTTP request being handled, as a portable value. + + Resolved on the ASGI request task, where the proxy's server span is anchored, + and carried to the MCP message handler on the authenticated-user object. A + stateful streamable-HTTP session handles every message on the task spawned by + its ``initialize`` POST, so the handler's own task cannot see later requests' + spans; this is the same reason per-request auth is carried across rather than + read from a ContextVar. Lazily imported so opentelemetry stays an optional + dependency; returns ``None`` when otel_v2 is unavailable or no request span is + anchored.""" + try: + from litellm.integrations.otel.plumbing.context import ( + request_root_span_context, + ) + + return request_root_span_context() + except ImportError: + return None + + +def _otel_set_mcp_transport_span_context(span_context: Optional["SpanContext"]) -> object: + """Publish the current message's transport span for the otel_v2 MCP span and + return a reset token, or ``None`` when otel_v2 is unavailable.""" + if span_context is None: + return None + try: + from litellm.integrations.otel.plumbing.context import ( + set_mcp_message_transport_span_context, + ) + + return set_mcp_message_transport_span_context(span_context) + except ImportError: + return None + + +def _otel_reset_mcp_transport_span_context(token: object) -> None: + """Paired with ``_otel_set_mcp_transport_span_context``.""" + if token is None: + return + try: + from litellm.integrations.otel.plumbing.context import ( + reset_mcp_message_transport_span_context, + ) + + reset_mcp_message_transport_span_context(token) + except ImportError: + return + + def _proxy_exception_to_http_exception(exc: ProxyException) -> HTTPException: """Map a ``ProxyException`` to an ``HTTPException`` that preserves its real status code and headers. @@ -654,6 +710,18 @@ if MCP_AVAILABLE: ############### MCP Server Routes ####################### ######################################################## + def _current_transport_span_context() -> Optional["SpanContext"]: + """The transport span of the HTTP request carrying the message being handled. + + Published by the ASGI request task onto the authenticated-user object, because + a stateful session's message handler runs on the task spawned by that session's + ``initialize`` POST and so cannot read later requests' spans from its own task. + """ + auth_user = auth_context_var.get() + if not isinstance(auth_user, MCPAuthenticatedUser): + auth_user = _recover_auth_from_session() + return auth_user.transport_span_context if auth_user is not None else None + @server.list_tools() async def handle_list_tools() -> "ListToolsResult | List[Tool]": """ @@ -670,9 +738,11 @@ if MCP_AVAILABLE: if req_ctx: _session_reset_token = active_mcp_session_var.set(req_ctx.session) _trace_token = None + _transport_token = None try: _trace_token = _otel_set_mcp_trace_carrier(_mcp_meta_trace_carrier(req_ctx)) + _transport_token = _otel_set_mcp_transport_span_context(_current_transport_span_context()) # Get user authentication from context variable ( user_api_key_auth, @@ -728,6 +798,7 @@ if MCP_AVAILABLE: # This prevents the HTTP stream from failing and allows the client to get a response return [] finally: + _otel_reset_mcp_transport_span_context(_transport_token) _otel_reset_mcp_trace_carrier(_trace_token) if _session_reset_token is not None: active_mcp_session_var.reset(_session_reset_token) @@ -901,9 +972,11 @@ if MCP_AVAILABLE: if req_ctx: _session_reset_token = active_mcp_session_var.set(req_ctx.session) _trace_token = None + _transport_token = None try: _trace_token = _otel_set_mcp_trace_carrier(_mcp_meta_trace_carrier(req_ctx)) + _transport_token = _otel_set_mcp_transport_span_context(_current_transport_span_context()) # Validate arguments ( user_api_key_auth, @@ -1042,6 +1115,7 @@ if MCP_AVAILABLE: return response finally: + _otel_reset_mcp_transport_span_context(_transport_token) _otel_reset_mcp_trace_carrier(_trace_token) if _session_reset_token is not None: active_mcp_session_var.reset(_session_reset_token) @@ -4197,6 +4271,7 @@ if MCP_AVAILABLE: session_id=session_id if use_stateful else None, touch_last_seen=(scope.get("method") or "").upper() != "DELETE", copy_existing_session_auth_context=is_initialize, + transport_span_context=_otel_request_transport_span_context(), ) local_send = send if use_stateful and is_initialize: @@ -4421,6 +4496,7 @@ if MCP_AVAILABLE: oauth2_headers: Optional[Dict[str, str]] = None, raw_headers: Optional[Dict[str, str]] = None, client_ip: Optional[str] = None, + transport_span_context: Optional["SpanContext"] = None, ) -> None: auth_user.user_api_key_auth = user_api_key_auth auth_user.mcp_auth_header = mcp_auth_header @@ -4429,6 +4505,7 @@ if MCP_AVAILABLE: auth_user.oauth2_headers = oauth2_headers auth_user.raw_headers = raw_headers auth_user.client_ip = client_ip + auth_user.transport_span_context = transport_span_context def set_auth_context( user_api_key_auth: Optional[UserAPIKeyAuth], @@ -4438,6 +4515,7 @@ if MCP_AVAILABLE: oauth2_headers: Optional[Dict[str, str]] = None, raw_headers: Optional[Dict[str, str]] = None, client_ip: Optional[str] = None, + transport_span_context: Optional["SpanContext"] = None, ) -> MCPAuthenticatedUser: """ Set the UserAPIKeyAuth in the auth context variable. @@ -4448,6 +4526,7 @@ if MCP_AVAILABLE: mcp_servers: Optional list of server names and access groups to filter by mcp_server_auth_headers: Optional dict of server-specific auth headers {server_alias: auth_value} client_ip: Client IP address for MCP access control + transport_span_context: Tracing span of the HTTP request carrying this message """ auth_user = MCPAuthenticatedUser( user_api_key_auth=user_api_key_auth, @@ -4457,6 +4536,7 @@ if MCP_AVAILABLE: oauth2_headers=oauth2_headers, raw_headers=raw_headers, client_ip=client_ip, + transport_span_context=transport_span_context, ) auth_context_var.set(auth_user) return auth_user @@ -4472,6 +4552,7 @@ if MCP_AVAILABLE: session_id: Optional[str] = None, touch_last_seen: bool = True, copy_existing_session_auth_context: bool = False, + transport_span_context: Optional["SpanContext"] = None, ) -> MCPAuthenticatedUser: auth_user = _stateful_session_auth_contexts.get(session_id) if session_id else None if auth_user is not None and session_id is not None: @@ -4486,6 +4567,7 @@ if MCP_AVAILABLE: oauth2_headers=oauth2_headers, raw_headers=raw_headers, client_ip=client_ip, + transport_span_context=transport_span_context, ) _update_auth_context( auth_user=auth_user, @@ -4496,6 +4578,7 @@ if MCP_AVAILABLE: oauth2_headers=oauth2_headers, raw_headers=raw_headers, client_ip=client_ip, + transport_span_context=transport_span_context, ) auth_context_var.set(auth_user) return auth_user @@ -4507,6 +4590,7 @@ if MCP_AVAILABLE: oauth2_headers=oauth2_headers, raw_headers=raw_headers, client_ip=client_ip, + transport_span_context=transport_span_context, ) def _wrap_send_with_stateful_session_auth_context( diff --git a/tests/test_litellm/integrations/otel/test_otel_v2_logger.py b/tests/test_litellm/integrations/otel/test_otel_v2_logger.py index 5f6002f4cdf..d3593f2c06b 100644 --- a/tests/test_litellm/integrations/otel/test_otel_v2_logger.py +++ b/tests/test_litellm/integrations/otel/test_otel_v2_logger.py @@ -29,7 +29,9 @@ from litellm.integrations.otel import ( # noqa: E402 from litellm.integrations.otel.plumbing import providers # noqa: E402 from litellm.integrations.otel.plumbing.context import ( # noqa: E402 reset_mcp_message_trace_carrier, + reset_mcp_message_transport_span_context, set_mcp_message_trace_carrier, + set_mcp_message_transport_span_context, set_request_root_span, ) from litellm.integrations.otel.logger import OpenTelemetryV2 # noqa: E402 @@ -56,9 +58,11 @@ def _reset_request_root_span(): _otel_context._request_root_span.set(None) _otel_context._mcp_message_trace_carrier.set(None) + _otel_context._mcp_message_transport_span_context.set(None) yield _otel_context._request_root_span.set(None) _otel_context._mcp_message_trace_carrier.set(None) + _otel_context._mcp_message_transport_span_context.set(None) def _payload(**overrides): @@ -528,15 +532,15 @@ _MCP_SPAN_CASES = [ @pytest.mark.parametrize("make_payload, span_name", _MCP_SPAN_CASES) -def test_mcp_span_roots_and_links_transport_without_propagated_context( +def test_mcp_span_nests_under_transport_without_propagated_context( make_payload, span_name ): - """MCP and the HTTP transport are independent lifecycles (one streamable-HTTP - session multiplexes many messages), so per the MCP semconv the message span - must NOT nest under the session/transport span — that is what made it render - skewed at the session's start. With no propagated ``params._meta`` context it - starts its own root trace and records the transport span as a *link*, never - the parent.""" + """Almost no MCP client implements SEP-414, so ``params._meta`` normally carries + no trace context. Rooting the span there split one tool call into two traces + joined only by a link, which is how it surfaced in APM: the ``POST`` transaction + and the ``tools/call`` span shared no ``trace_id``. With no remote parent to + honor the span nests under the transport span instead, and records no link since + the transport is now the real parent.""" logger, exporter = _logger() transport = logger._emitter.start_span( SpanRole.PROXY_REQUEST, LITELLM_PROXY_REQUEST_SPAN_NAME @@ -549,11 +553,75 @@ def test_mcp_span_roots_and_links_transport_without_propagated_context( ) transport.end() span = next(s for s in exporter.get_finished_spans() if s.name == span_name) + assert span.parent is not None + assert span.parent.span_id == transport.get_span_context().span_id + assert span.context.trace_id == transport.get_span_context().trace_id + assert span.links == () + + +@pytest.mark.parametrize("make_payload, span_name", _MCP_SPAN_CASES) +def test_mcp_span_nests_under_this_messages_transport_not_the_session_opener( + make_payload, span_name +): + """A *stateful* streamable-HTTP session runs every message on the single task + spawned by that session's ``initialize`` POST, so the ``_request_root_span`` + ContextVar the ASGI request task writes is frozen at ``initialize`` inside the + handler and never sees the later ``tools/call`` POST. Nesting on that anchor + would hang every tool call of the session off the first request's (already + ended) span, rendering skewed at the session's start. The gateway resolves the + current message's transport on the request task and publishes it, so the span + parents to the POST that actually carried this message.""" + logger, exporter = _logger() + session_opener = logger._emitter.start_span( + SpanRole.PROXY_REQUEST, LITELLM_PROXY_REQUEST_SPAN_NAME + ) + this_message = logger._emitter.start_span( + SpanRole.PROXY_REQUEST, LITELLM_PROXY_REQUEST_SPAN_NAME + ) + + async def session_task(): + token = set_mcp_message_transport_span_context( + this_message.get_span_context() + ) + try: + await logger.async_log_success_event( + {"standard_logging_object": make_payload()}, None, None, None + ) + finally: + reset_mcp_message_transport_span_context(token) + + async def initialize_request(): + # The anchor the session task inherits is the one ``initialize`` left behind; + # spawning here reproduces the SDK's session task, which outlives this request. + set_request_root_span(session_opener) + await asyncio.create_task(session_task()) + + asyncio.run(initialize_request()) + session_opener.end() + this_message.end() + span = next(s for s in exporter.get_finished_spans() if s.name == span_name) + assert span.parent is not None + assert span.parent.span_id == this_message.get_span_context().span_id + assert span.context.trace_id == this_message.get_span_context().trace_id + assert span.parent.span_id != session_opener.get_span_context().span_id + assert span.context.trace_id != session_opener.get_span_context().trace_id + + +@pytest.mark.parametrize("make_payload, span_name", _MCP_SPAN_CASES) +def test_mcp_span_roots_without_transport_or_propagated_context( + make_payload, span_name +): + """With neither a remote parent nor a transport span there is nothing to nest + under, so the span legitimately starts its own root trace with no links.""" + logger, exporter = _logger() + asyncio.run( + logger.async_log_success_event( + {"standard_logging_object": make_payload()}, None, None, None + ) + ) + span = next(s for s in exporter.get_finished_spans() if s.name == span_name) assert span.parent is None - assert span.context.trace_id != transport.get_span_context().trace_id - assert [link.context.span_id for link in span.links] == [ - transport.get_span_context().span_id - ] + assert span.links == () @pytest.mark.parametrize("make_payload, span_name", _MCP_SPAN_CASES) @@ -641,10 +709,11 @@ def test_mcp_span_carries_authenticated_identity(make_payload, span_name): assert span.attributes[LiteLLM.TEAM_ID] == "t1" -def test_mcp_span_malformed_traceparent_starts_root(): +def test_mcp_span_malformed_traceparent_nests_under_transport(): """A malformed traceparent in ``params._meta`` must not crash or parent to a - bogus span: the propagator ignores it, so the span starts its own root trace and - still links the transport span.""" + bogus span: the propagator ignores it, leaving no remote parent, so the span + falls back to nesting under the transport span rather than starting a + disconnected root trace.""" logger, exporter = _logger() transport = logger._emitter.start_span( SpanRole.PROXY_REQUEST, LITELLM_PROXY_REQUEST_SPAN_NAME @@ -661,9 +730,44 @@ def test_mcp_span_malformed_traceparent_starts_root(): reset_mcp_message_trace_carrier(token) transport.end() span = next(s for s in exporter.get_finished_spans() if s.name == "tools/list") - assert span.parent is None + assert span.parent is not None + assert span.parent.span_id == transport.get_span_context().span_id + assert span.links == () + + +def test_mcp_span_links_this_messages_transport_when_context_is_propagated(): + """On the semconv path the transport is recorded as a link, and that link must + point at the POST carrying this message too. Reading the stale session anchor + would attribute the tool call to whichever request opened the session.""" + logger, exporter = _logger() + session_opener = logger._emitter.start_span( + SpanRole.PROXY_REQUEST, LITELLM_PROXY_REQUEST_SPAN_NAME + ) + this_message = logger._emitter.start_span( + SpanRole.PROXY_REQUEST, LITELLM_PROXY_REQUEST_SPAN_NAME + ) + set_request_root_span(session_opener) + trace_token = set_mcp_message_trace_carrier( + {"traceparent": "00-11111111111111111111111111111111-2222222222222222-01"} + ) + transport_token = set_mcp_message_transport_span_context( + this_message.get_span_context() + ) + try: + asyncio.run( + logger.async_log_success_event( + {"standard_logging_object": _mcp_list_payload()}, None, None, None + ) + ) + finally: + reset_mcp_message_transport_span_context(transport_token) + reset_mcp_message_trace_carrier(trace_token) + session_opener.end() + this_message.end() + span = next(s for s in exporter.get_finished_spans() if s.name == "tools/list") + assert span.parent is not None and span.parent.span_id == 0x2222222222222222 assert [link.context.span_id for link in span.links] == [ - transport.get_span_context().span_id + this_message.get_span_context().span_id ] From 9f9714c209b3a00083d5bb3f93c1894aa93841dc Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Fri, 24 Jul 2026 16:05:18 -0700 Subject: [PATCH 10/29] test: remove five more zero-kill tests from test_http_handler The http_handler pair only received its full mutation verdict after the first removal batch landed; these five tests pass unchanged when every function they execute is mutated and the owning file killed none of their scored mutants. The ssl tests excluded from mutation scoring are untouched. --- .../llms/custom_httpx/test_http_handler.py | 111 ------------------ 1 file changed, 111 deletions(-) diff --git a/tests/test_litellm/llms/custom_httpx/test_http_handler.py b/tests/test_litellm/llms/custom_httpx/test_http_handler.py index 7bd1d7a6031..87d67e0e8b7 100644 --- a/tests/test_litellm/llms/custom_httpx/test_http_handler.py +++ b/tests/test_litellm/llms/custom_httpx/test_http_handler.py @@ -181,28 +181,6 @@ async def test_force_ipv4_transport(): litellm.disable_aiohttp_transport = original_disable -@pytest.mark.asyncio -async def test_ssl_context_transport(): - """Test transport creation with SSL context""" - # Create a test SSL context - ssl_context = ssl.create_default_context() - - transport = AsyncHTTPHandler._create_async_transport(ssl_context=ssl_context) - assert transport is not None - - try: - if isinstance(transport, LiteLLMAiohttpTransport): - # Get the client session and verify SSL context is passed through - client_session = transport._get_valid_client_session() - assert isinstance(client_session, ClientSession) - assert isinstance(client_session.connector, TCPConnector) - # Verify the connector has SSL context set by checking if it's using SSL - assert client_session.connector._ssl is not None - finally: - if isinstance(transport, LiteLLMAiohttpTransport): - await transport.aclose() - - @pytest.mark.asyncio async def test_aiohttp_disabled_transport(): """Test transport creation with aiohttp disabled""" @@ -339,44 +317,6 @@ async def test_ssl_context_with_shared_session(): litellm.disable_aiohttp_transport = original_disable -@pytest.mark.asyncio -async def test_aiohttp_transport_trust_env_setting(monkeypatch): - """Test that trust_env setting is properly configured in aiohttp transport""" - transports = [] - try: - # Test 1: Default trust_env behavior - transport = AsyncHTTPHandler._create_aiohttp_transport() - transports.append(transport) - client_session = transport._get_valid_client_session() - - # Default should be False (litellm.aiohttp_trust_env default) - default_trust_env = getattr(litellm, "aiohttp_trust_env", False) - assert client_session._trust_env == default_trust_env - - # Test 2: Environment variable override - monkeypatch.setenv("AIOHTTP_TRUST_ENV", "True") - transport_with_env = AsyncHTTPHandler._create_aiohttp_transport() - transports.append(transport_with_env) - client_session_with_env = transport_with_env._get_valid_client_session() - - # Should be True when environment variable is set - assert client_session_with_env._trust_env is True - - # Test 3: Verify environment variable with False value - monkeypatch.setenv("AIOHTTP_TRUST_ENV", "False") - transport_with_false_env = AsyncHTTPHandler._create_aiohttp_transport() - transports.append(transport_with_false_env) - client_session_with_false_env = ( - transport_with_false_env._get_valid_client_session() - ) - - # Should respect the litellm.aiohttp_trust_env setting when env var is False - assert client_session_with_false_env._trust_env == default_trust_env - finally: - for t in transports: - await t.aclose() - - def test_get_ssl_configuration(): """Test that get_ssl_configuration() returns a proper SSL context with certifi CA bundle when no environment variables are set.""" @@ -443,36 +383,6 @@ async def test_create_aiohttp_transport_with_shared_session(): assert not callable(transport.client) # Should not be callable -@pytest.mark.asyncio -async def test_create_aiohttp_transport_without_shared_session(): - """Test that _create_aiohttp_transport creates new session when none provided""" - from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler - - # Test without shared session - transport = AsyncHTTPHandler._create_aiohttp_transport(shared_session=None) - - # Verify the transport uses a lambda function (for backward compatibility) - assert callable(transport.client) # Should be a lambda function - - -@pytest.mark.asyncio -async def test_create_aiohttp_transport_with_closed_session(): - """Test that _create_aiohttp_transport creates new session when shared session is closed""" - from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler - - # Create a mock closed session - mock_session = MockClientSession() - mock_session.closed = True - - # Test with closed session - transport = AsyncHTTPHandler._create_aiohttp_transport( - shared_session=mock_session # type: ignore - ) - - # Verify the transport creates a new session (lambda function) - assert callable(transport.client) # Should be a lambda function - - @pytest.mark.asyncio async def test_async_handler_with_shared_session(): """Test AsyncHTTPHandler initialization with shared session""" @@ -622,27 +532,6 @@ async def test_session_reuse_integration(): await client2.close() -@pytest.mark.asyncio -async def test_session_validation(): - """Test that session validation works correctly""" - from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler - - # Test with None session - transport1 = AsyncHTTPHandler._create_aiohttp_transport(shared_session=None) - assert callable(transport1.client) # Should create lambda - - # Test with closed session - mock_closed_session = MockClientSession() - mock_closed_session.closed = True - transport2 = AsyncHTTPHandler._create_aiohttp_transport(shared_session=mock_closed_session) # type: ignore - assert callable(transport2.client) # Should create lambda - - # Test with valid session - mock_valid_session = MockClientSession() - transport3 = AsyncHTTPHandler._create_aiohttp_transport(shared_session=mock_valid_session) # type: ignore - assert transport3.client is mock_valid_session # Should reuse session - - @pytest.mark.parametrize( "env_curve,litellm_curve,expected_curve,should_call", [ From ba9f6d75d8b57df480a32a9ba8209e32886aa9e8 Mon Sep 17 00:00:00 2001 From: Tin Chi Lo Date: Thu, 23 Jul 2026 18:44:36 -0700 Subject: [PATCH 11/29] refactor(ui): derive the dashboard object_permission type from the generated schema The dashboard declared the server-owned object_permission shape by hand in five places, each with a different subset of fields and none matching the OpenAPI schema. That is what hid LIT-4766: KeyResponse.object_permission never declared mcp_toolsets, so a form that wrote the field without reading it compiled cleanly and silently wiped the grant Replace four of those copies with one alias over the generated LiteLLM_ObjectPermissionTable. The agent shape stays separate because the agent endpoint really does return a narrower type, so it points at its own generated AgentObjectPermission --- .../src/components/agents/types.ts | 8 +++----- .../src/components/key_team_helpers/key_list.tsx | 12 ++---------- .../src/components/networking.tsx | 9 ++------- .../src/components/object_permission_types.ts | 3 +++ .../src/components/object_permissions_view.tsx | 15 ++------------- .../src/components/team/TeamInfo.tsx | 13 ++----------- 6 files changed, 14 insertions(+), 46 deletions(-) create mode 100644 ui/litellm-dashboard/src/components/object_permission_types.ts diff --git a/ui/litellm-dashboard/src/components/agents/types.ts b/ui/litellm-dashboard/src/components/agents/types.ts index c29c566a5fe..24ff0c0e12c 100644 --- a/ui/litellm-dashboard/src/components/agents/types.ts +++ b/ui/litellm-dashboard/src/components/agents/types.ts @@ -1,14 +1,12 @@ +import type { components } from "@/lib/http/schema"; + export interface AgentAttachedKey { token: string; key_alias?: string | null; key_name?: string | null; } -export interface AgentObjectPermission { - mcp_servers?: string[]; - mcp_access_groups?: string[]; - mcp_tool_permissions?: Record; -} +export type AgentObjectPermission = components["schemas"]["AgentObjectPermission"]; export interface Agent { agent_id: string; diff --git a/ui/litellm-dashboard/src/components/key_team_helpers/key_list.tsx b/ui/litellm-dashboard/src/components/key_team_helpers/key_list.tsx index ceff1809b7b..4b446b0c283 100644 --- a/ui/litellm-dashboard/src/components/key_team_helpers/key_list.tsx +++ b/ui/litellm-dashboard/src/components/key_team_helpers/key_list.tsx @@ -1,6 +1,7 @@ import { Setter } from "@/types"; import { useEffect, useState } from "react"; import { keyListCall, Member, Organization } from "../networking"; +import type { ObjectPermission } from "../object_permission_types"; export interface Team { team_id: string; @@ -90,16 +91,7 @@ export interface KeyResponse { user_tpm_limit: number; user_rpm_limit: number; user_email: string; - object_permission?: { - object_permission_id: string; - mcp_servers: string[]; - mcp_access_groups?: string[]; - mcp_toolsets?: string[] | null; - mcp_tool_permissions?: Record; - vector_stores: string[]; - agents?: string[]; - agent_access_groups?: string[]; - }; + object_permission?: ObjectPermission | null; access_group_ids?: string[]; budget_fallbacks?: Record; budget_limits?: Array<{ budget_duration: string; max_budget: number; reset_at?: string }>; diff --git a/ui/litellm-dashboard/src/components/networking.tsx b/ui/litellm-dashboard/src/components/networking.tsx index 051b83f4e27..5a467826f72 100644 --- a/ui/litellm-dashboard/src/components/networking.tsx +++ b/ui/litellm-dashboard/src/components/networking.tsx @@ -28,6 +28,7 @@ import { TagNewRequest, TagUpdateRequest, TagListResponse, TagInfoResponse } fro import { Team } from "./key_team_helpers/key_list"; import { EmailEventSettingsResponse, EmailEventSettingsUpdateRequest } from "./email_events/types"; import type { SkillRegisterRequest } from "./claude_code_plugins/types"; +import type { ObjectPermission } from "./object_permission_types"; import { jsonFields } from "./common_components/check_openapi_schema"; import NotificationsManager from "./molecules/notifications_manager"; import type { MCPUserEnvVarsStatus } from "./mcp_tools/types"; @@ -208,13 +209,7 @@ export interface Organization { teams: any[] | null; users: any[] | null; members: any[] | null; - object_permission?: { - object_permission_id: string; - mcp_servers: string[]; - mcp_access_groups?: string[]; - mcp_toolsets?: string[]; - vector_stores: string[]; - }; + object_permission?: ObjectPermission | null; } export interface CredentialItem { diff --git a/ui/litellm-dashboard/src/components/object_permission_types.ts b/ui/litellm-dashboard/src/components/object_permission_types.ts new file mode 100644 index 00000000000..bde7281faec --- /dev/null +++ b/ui/litellm-dashboard/src/components/object_permission_types.ts @@ -0,0 +1,3 @@ +import type { components } from "@/lib/http/schema"; + +export type ObjectPermission = Partial; diff --git a/ui/litellm-dashboard/src/components/object_permissions_view.tsx b/ui/litellm-dashboard/src/components/object_permissions_view.tsx index b0ee38bd834..687d1a5a846 100644 --- a/ui/litellm-dashboard/src/components/object_permissions_view.tsx +++ b/ui/litellm-dashboard/src/components/object_permissions_view.tsx @@ -3,21 +3,10 @@ import { Text } from "@tremor/react"; import VectorStorePermissions from "./permissions/VectorStorePermissions"; import MCPServerPermissions from "./permissions/MCPServerPermissions"; import AgentPermissions from "./permissions/AgentPermissions"; - -interface ObjectPermission { - object_permission_id: string; - mcp_servers: string[]; - mcp_access_groups?: string[]; - mcp_tool_permissions?: Record; - mcp_toolsets?: string[] | null; - vector_stores: string[]; - agents?: string[]; - agent_access_groups?: string[]; - search_tools?: string[]; -} +import type { ObjectPermission } from "./object_permission_types"; interface ObjectPermissionsViewProps { - objectPermission?: ObjectPermission; + objectPermission?: ObjectPermission | null; variant?: "card" | "inline"; className?: string; accessToken?: string | null; diff --git a/ui/litellm-dashboard/src/components/team/TeamInfo.tsx b/ui/litellm-dashboard/src/components/team/TeamInfo.tsx index eaf2faa08ee..ff881a68938 100644 --- a/ui/litellm-dashboard/src/components/team/TeamInfo.tsx +++ b/ui/litellm-dashboard/src/components/team/TeamInfo.tsx @@ -17,6 +17,7 @@ import { import { useGuardrails, GuardrailListItem } from "@/app/(dashboard)/hooks/guardrails/useGuardrails"; import { formatNumberWithCommas } from "@/utils/dataUtils"; import { mapEmptyStringToNull } from "@/utils/keyUpdateUtils"; +import type { ObjectPermission } from "@/components/object_permission_types"; import { isProxyAdminRole } from "@/utils/roles"; import { EditOutlined, @@ -118,17 +119,7 @@ export interface TeamData { router_settings?: Record; guardrails?: string[]; policies?: string[]; - object_permission?: { - object_permission_id: string; - mcp_servers: string[]; - mcp_access_groups?: string[]; - mcp_tool_permissions?: Record; - mcp_toolsets?: string[]; - vector_stores: string[]; - agents?: string[]; - agent_access_groups?: string[]; - search_tools?: string[]; - }; + object_permission?: ObjectPermission | null; team_member_budget_table: { max_budget: number; budget_duration: string; From 3476240f11bbe288d3e79f47b4e91973ca57dc6c Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Fri, 24 Jul 2026 16:07:23 -0700 Subject: [PATCH 12/29] fix(ui): keep entity usage tabs aligned with their panels Tremor's TabPanels hands each child an index via React.Children.map, while the selected index comes from HeadlessUI counting only real Tab elements. An empty fragment, false, or null still consumes a panel index but contributes no tab, so the team-only Agent Activity conditional made the two lists drift for every non-team entity type: Key Activity resolved to the empty slot and rendered nothing at all, and Endpoint Activity rendered the key metrics Drive both lists from a single tab array so adding or removing a conditional tab touches one place and the indices cannot diverge --- .../EntityUsage/EntityUsage.test.tsx | 61 +- .../components/EntityUsage/EntityUsage.tsx | 618 +++++++++--------- 2 files changed, 364 insertions(+), 315 deletions(-) diff --git a/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/EntityUsage/EntityUsage.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/EntityUsage/EntityUsage.test.tsx index cbf3a2cc1f6..89c38c6274f 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/EntityUsage/EntityUsage.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/EntityUsage/EntityUsage.test.tsx @@ -25,8 +25,17 @@ vi.mock("@/components/networking", () => ({ // Mock the child components to simplify testing vi.mock("@/components/activity_metrics", () => ({ - ActivityMetrics: () =>
Activity Metrics
, - processActivityData: () => ({ data: [], metadata: {} }), + ActivityMetrics: ({ modelMetrics }: { modelMetrics?: { __source?: string } }) => ( +
+ Activity Metrics + {`metrics-source:${modelMetrics?.__source ?? "none"}`} +
+ ), + processActivityData: (_data: unknown, key: string) => ({ __source: key }), +})); + +vi.mock("../EndpointUsage/EndpointUsage", () => ({ + default: () =>
Endpoint Usage Panel
, })); vi.mock("@/components/UsagePage/components/EntityUsage/TopKeyView", () => ({ @@ -481,6 +490,54 @@ describe("EntityUsage", () => { expect(screen.getAllByText("Activity Metrics")[1]).toBeInTheDocument(); }); + const selectedPanels = (container: HTMLElement) => + Array.from(container.querySelectorAll("div.tremor-TabPanel-root")).filter( + (panel) => panel.getAttribute("aria-selected") === "true", + ); + + it.each([ + ["Cost", "Tag Spend Overview"], + ["Model Activity", "metrics-source:models"], + ["Key Activity", "metrics-source:api_keys"], + ["Endpoint Activity", "Endpoint Usage Panel"], + ])("shows only the %s panel for a non-team entity type", async (tabLabel, marker) => { + const { container } = render(); + + await waitFor(() => { + expect(mockTagDailyActivityCall).toHaveBeenCalled(); + }); + + act(() => { + fireEvent.click(screen.getByText(tabLabel)); + }); + + const selected = selectedPanels(container); + expect(selected).toHaveLength(1); + expect(selected[0].textContent).toContain(marker); + }); + + it.each([ + ["Cost", "Team Spend Overview"], + ["Model Activity", "metrics-source:models"], + ["Agent Activity", "metrics-source:entities"], + ["Key Activity", "metrics-source:api_keys"], + ["Endpoint Activity", "Endpoint Usage Panel"], + ])("shows only the %s panel for the team entity type", async (tabLabel, marker) => { + const { container } = render(); + + await waitFor(() => { + expect(mockTeamDailyActivityCall).toHaveBeenCalled(); + }); + + act(() => { + fireEvent.click(screen.getByText(tabLabel)); + }); + + const selected = selectedPanels(container); + expect(selected).toHaveLength(1); + expect(selected[0].textContent).toContain(marker); + }); + it("should handle empty data gracefully", async () => { const emptyData = { results: [], diff --git a/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/EntityUsage/EntityUsage.tsx b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/EntityUsage/EntityUsage.tsx index 534e2be7fe8..e330983b6f9 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/EntityUsage/EntityUsage.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/EntityUsage/EntityUsage.tsx @@ -25,7 +25,7 @@ import { } from "@tremor/react"; import { ExportOutlined, LoadingOutlined } from "@ant-design/icons"; import { Alert, Button } from "antd"; -import React, { useMemo, useState } from "react"; +import React, { type ReactNode, useMemo, useState } from "react"; import TeamMultiSelect from "@/components/common_components/team_multi_select"; import { ActivityMetrics, processActivityData } from "@/components/activity_metrics"; import { UsageExportHeader } from "@/components/EntityUsageExport"; @@ -406,6 +406,304 @@ const EntityUsage: React.FC = ({ accessToken, entityType, enti const capitalizedEntityLabel = entityType.charAt(0).toUpperCase() + entityType.slice(1); + const costPanel = ( + + {/* Total Spend Card */} + + + {capitalizedEntityLabel} Spend Overview + + + Total Spend + + ${formatNumberWithCommas(spendData.metadata.total_spend, 2)} + + + + Total Requests + {spendData.metadata.total_api_requests.toLocaleString()} + + + Successful Requests + + {spendData.metadata.total_successful_requests.toLocaleString()} + + + + Failed Requests + + {spendData.metadata.total_failed_requests.toLocaleString()} + + + + Total Tokens + {spendData.metadata.total_tokens.toLocaleString()} + + + + + + {/* Daily Spend Chart */} + + + + Daily Spend + + + new Date(a.date).getTime() - new Date(b.date).getTime())} + index="date" + categories={["metrics.spend"]} + colors={["cyan"]} + valueFormatter={valueFormatterSpend} + yAxisWidth={100} + showLegend={false} + customTooltip={({ payload, active }) => { + if (!active || !payload?.[0]) return null; + const data = payload[0].payload; + const entityCount = Object.keys(data.breakdown.entities || {}).length; + return ( +
+

{data.date}

+

Total Spend: ${formatNumberWithCommas(data.metrics.spend, 2)}

+

Total Requests: {data.metrics.api_requests}

+

Successful: {data.metrics.successful_requests}

+

Failed: {data.metrics.failed_requests}

+

Total Tokens: {data.metrics.total_tokens}

+

+ Total {capitalizedEntityLabel}s: {entityCount} +

+
+

Spend by {capitalizedEntityLabel}:

+ {Object.entries(data.breakdown.entities || {}) + .sort(([, a], [, b]) => { + const spendA = (a as EntityMetrics).metrics.spend; + const spendB = (b as EntityMetrics).metrics.spend; + return spendB - spendA; + }) + .slice(0, 5) + .map(([entity, entityData]) => { + const metrics = entityData as EntityMetrics; + return ( +

+ {getEntityLabel(entity, metrics.metadata)}: $ + {formatNumberWithCommas(metrics.metrics.spend, 2)} +

+ ); + })} + {entityCount > 5 &&

...and {entityCount - 5} more

} +
+
+ ); + }} + /> + + + + + {/* Entity Breakdown Section */} + + +
+
+ Spend Per {capitalizedEntityLabel} + Showing Top 5 by Spend +
+ Get Started by Tracking cost per {capitalizedEntityLabel} + + here + +
+
+ + + { + if (!active || !payload?.[0]) return null; + const data = payload[0].payload; + return ( +
+

{data.metadata.alias}

+

Spend: ${formatNumberWithCommas(data.metrics.spend, 4)}

+

Requests: {data.metrics.api_requests.toLocaleString()}

+

+ Successful: {data.metrics.successful_requests.toLocaleString()} +

+

Failed: {data.metrics.failed_requests.toLocaleString()}

+

Tokens: {data.metrics.total_tokens.toLocaleString()}

+
+ ); + }} + /> + + +
+ + + + {capitalizedEntityLabel} + Spend + Successful + Failed + Tokens + + + + {getEntityBreakdown() + .filter((entity) => entity.metrics.spend > 0) + .map((entity) => ( + + {entity.metadata.alias} + + + + + {entity.metrics.successful_requests.toLocaleString()} + + + {entity.metrics.failed_requests.toLocaleString()} + + {entity.metrics.total_tokens.toLocaleString()} + + ))} + +
+
+ + +
+ + + + {/* Top API Keys */} + + + Top Virtual Keys + + + + + {/* Top Models */} + + + {entityType === "agent" ? "Top Agents" : "Top Models"} + + + + + {/* Top Agents - only for team entity type */} + {entityType === "team" && ( + + + Top Agents Driving Spend + + + + )} + + {/* Spend by Provider */} + + +
+ Provider Usage + + + `$${formatNumberWithCommas(value, 2)}`} + colors={["cyan", "blue", "indigo", "violet", "purple"]} + showLabel + startAngle={90} + endAngle={-270} + /> + + + + + + Provider + Spend + Successful + Failed + Tokens + + + + {getProviderSpend().map((provider) => ( + + +
+ {provider.provider && } + {provider.provider} +
+ + + + + + {provider.successful_requests.toLocaleString()} + + {provider.failed_requests.toLocaleString()} + {provider.tokens.toLocaleString()} + + ))} + +
+ + +
+ + + + ); + + const tabs: readonly { key: string; label: string; content: ReactNode }[] = [ + { key: "cost", label: "Cost", content: costPanel }, + { + key: "models", + label: entityType === "agent" ? "Request / Token Consumption" : "Model Activity", + content: , + }, + ...(entityType === "team" + ? [{ key: "agents", label: "Agent Activity", content: }] + : []), + { + key: "keys", + label: "Key Activity", + content: , + }, + { key: "endpoints", label: "Endpoint Activity", content: }, + ]; + return (
{isFetchingMore && ( @@ -501,320 +799,14 @@ const EntityUsage: React.FC = ({ accessToken, entityType, enti /> - Cost - {entityType === "agent" ? "Request / Token Consumption" : "Model Activity"} - {entityType === "team" ? Agent Activity : <>} - Key Activity - Endpoint Activity + {tabs.map(({ key, label }) => ( + {label} + ))} - - - {/* Total Spend Card */} - - - {capitalizedEntityLabel} Spend Overview - - - Total Spend - - ${formatNumberWithCommas(spendData.metadata.total_spend, 2)} - - - - Total Requests - - {spendData.metadata.total_api_requests.toLocaleString()} - - - - Successful Requests - - {spendData.metadata.total_successful_requests.toLocaleString()} - - - - Failed Requests - - {spendData.metadata.total_failed_requests.toLocaleString()} - - - - Total Tokens - - {spendData.metadata.total_tokens.toLocaleString()} - - - - - - - {/* Daily Spend Chart */} - - - - Daily Spend - - - new Date(a.date).getTime() - new Date(b.date).getTime(), - )} - index="date" - categories={["metrics.spend"]} - colors={["cyan"]} - valueFormatter={valueFormatterSpend} - yAxisWidth={100} - showLegend={false} - customTooltip={({ payload, active }) => { - if (!active || !payload?.[0]) return null; - const data = payload[0].payload; - const entityCount = Object.keys(data.breakdown.entities || {}).length; - return ( -
-

{data.date}

-

- Total Spend: ${formatNumberWithCommas(data.metrics.spend, 2)} -

-

Total Requests: {data.metrics.api_requests}

-

Successful: {data.metrics.successful_requests}

-

Failed: {data.metrics.failed_requests}

-

Total Tokens: {data.metrics.total_tokens}

-

- Total {capitalizedEntityLabel}s: {entityCount} -

-
-

Spend by {capitalizedEntityLabel}:

- {Object.entries(data.breakdown.entities || {}) - .sort(([, a], [, b]) => { - const spendA = (a as EntityMetrics).metrics.spend; - const spendB = (b as EntityMetrics).metrics.spend; - return spendB - spendA; - }) - .slice(0, 5) - .map(([entity, entityData]) => { - const metrics = entityData as EntityMetrics; - return ( -

- {getEntityLabel(entity, metrics.metadata)}: $ - {formatNumberWithCommas(metrics.metrics.spend, 2)} -

- ); - })} - {entityCount > 5 && ( -

...and {entityCount - 5} more

- )} -
-
- ); - }} - /> - - - - - {/* Entity Breakdown Section */} - - -
-
- Spend Per {capitalizedEntityLabel} - Showing Top 5 by Spend -
- Get Started by Tracking cost per {capitalizedEntityLabel} - - here - -
-
- - - { - if (!active || !payload?.[0]) return null; - const data = payload[0].payload; - return ( -
-

{data.metadata.alias}

-

Spend: ${formatNumberWithCommas(data.metrics.spend, 4)}

-

Requests: {data.metrics.api_requests.toLocaleString()}

-

- Successful: {data.metrics.successful_requests.toLocaleString()} -

-

Failed: {data.metrics.failed_requests.toLocaleString()}

-

Tokens: {data.metrics.total_tokens.toLocaleString()}

-
- ); - }} - /> - - -
- - - - {capitalizedEntityLabel} - Spend - Successful - Failed - Tokens - - - - {getEntityBreakdown() - .filter((entity) => entity.metrics.spend > 0) - .map((entity) => ( - - {entity.metadata.alias} - - - - - {entity.metrics.successful_requests.toLocaleString()} - - - {entity.metrics.failed_requests.toLocaleString()} - - {entity.metrics.total_tokens.toLocaleString()} - - ))} - -
-
- - -
- - - - {/* Top API Keys */} - - - Top Virtual Keys - - - - - {/* Top Models */} - - - {entityType === "agent" ? "Top Agents" : "Top Models"} - - - - - {/* Top Agents - only for team entity type */} - {entityType === "team" && ( - - - Top Agents Driving Spend - - - - )} - - {/* Spend by Provider */} - - -
- Provider Usage - - - `$${formatNumberWithCommas(value, 2)}`} - colors={["cyan", "blue", "indigo", "violet", "purple"]} - showLabel - startAngle={90} - endAngle={-270} - /> - - - - - - Provider - Spend - Successful - Failed - Tokens - - - - {getProviderSpend().map((provider) => ( - - -
- {provider.provider && } - {provider.provider} -
- - - - - - {provider.successful_requests.toLocaleString()} - - - {provider.failed_requests.toLocaleString()} - - {provider.tokens.toLocaleString()} - - ))} - -
- - -
- - - - - - - - {entityType === "team" ? ( - - - - ) : ( - <> - )} - - - - - - + {tabs.map(({ key, content }) => ( + {content} + ))}
From 9e56630347d881b370a1e82a07a04eb17ec6ddcd Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Fri, 24 Jul 2026 16:08:29 -0700 Subject: [PATCH 13/29] refactor(ui): migrate routing groups table onto the shared DataTable Rebuilds the Router Settings > Routing Groups table on the shared DataTable and cell library, the last antd entity grid in the dashboard. The table splits into a thin RoutingGroupsTable container plus RoutingGroupsTableColumns, with the usage snippets moving to their own RoutingGroupUsagePanel on ui/tabs and the shared CodeBlock instead of antd Tabs and Paragraph copyable. Models render through the shared ModelsCell so long lists collapse behind "+N more" rather than wrapping the row, and the two inline icon buttons become a single overflow menu with Edit and Delete. antd gave the snippet panel its own chevron column; under the shared pattern a row has two click targets, the name cell and the overflow menu, so clicking the group name now opens the panel. Column set, order, actions, and the backend row order are otherwise unchanged. --- ui/litellm-dashboard/eslint-suppressions.json | 10 - .../routing_groups/RoutingGroupUsagePanel.tsx | 91 +++++++ .../RoutingGroupsTable.test.tsx | 145 ++++++++++ .../routing_groups/RoutingGroupsTable.tsx | 257 ++++-------------- .../RoutingGroupsTableColumns.tsx | 111 ++++++++ .../src/components/routing_groups/index.tsx | 2 +- .../src/components/routing_groups/strategy.ts | 8 + 7 files changed, 411 insertions(+), 213 deletions(-) create mode 100644 ui/litellm-dashboard/src/components/routing_groups/RoutingGroupUsagePanel.tsx create mode 100644 ui/litellm-dashboard/src/components/routing_groups/RoutingGroupsTable.test.tsx create mode 100644 ui/litellm-dashboard/src/components/routing_groups/RoutingGroupsTableColumns.tsx create mode 100644 ui/litellm-dashboard/src/components/routing_groups/strategy.ts diff --git a/ui/litellm-dashboard/eslint-suppressions.json b/ui/litellm-dashboard/eslint-suppressions.json index ce743063309..90ba4dc7f59 100644 --- a/ui/litellm-dashboard/eslint-suppressions.json +++ b/ui/litellm-dashboard/eslint-suppressions.json @@ -2964,11 +2964,6 @@ "count": 1 } }, - "src/components/common_components/TableHeaderSortDropdown/TableHeaderSortDropdown.tsx": { - "no-restricted-imports": { - "count": 1 - } - }, "src/components/common_components/budget_duration_dropdown.tsx": { "local/filename-pascal-case": { "count": 1 @@ -3554,11 +3549,6 @@ "count": 1 } }, - "src/components/routing_groups/RoutingGroupsTable.tsx": { - "no-restricted-imports": { - "count": 2 - } - }, "src/components/routing_groups/index.tsx": { "local/filename-pascal-case": { "count": 1 diff --git a/ui/litellm-dashboard/src/components/routing_groups/RoutingGroupUsagePanel.tsx b/ui/litellm-dashboard/src/components/routing_groups/RoutingGroupUsagePanel.tsx new file mode 100644 index 00000000000..fafea569012 --- /dev/null +++ b/ui/litellm-dashboard/src/components/routing_groups/RoutingGroupUsagePanel.tsx @@ -0,0 +1,91 @@ +"use client"; + +import { Code2 } from "lucide-react"; +import React from "react"; + +import CodeBlock from "@/components/CodeBlock"; +import { Tabs, TabsContent, TabsList, TabsTrigger } from "@/components/ui/tabs"; + +import { formatStrategyLabel } from "./strategy"; +import type { RoutingGroup } from "./types"; + +interface RoutingGroupUsagePanelProps { + group: RoutingGroup; + baseUrl: string; +} + +const exampleModel = (group: RoutingGroup): string => group.models[0] ?? ""; + +const buildCurlSnippet = (group: RoutingGroup, baseUrl: string): string => + `curl -X POST '${baseUrl}/v1/chat/completions' \\ + -H 'Content-Type: application/json' \\ + -H 'Authorization: Bearer $LITELLM_API_KEY' \\ + -d '{ + "model": "${exampleModel(group)}", + "messages": [{"role": "user", "content": "Hello!"}] + }'`; + +const buildPythonSnippet = (group: RoutingGroup, baseUrl: string): string => + `from openai import OpenAI + +client = OpenAI( + api_key="$LITELLM_API_KEY", + base_url="${baseUrl}", +) + +response = client.chat.completions.create( + model="${exampleModel(group)}", + messages=[{"role": "user", "content": "Hello!"}], +) + +print(response)`; + +const buildJsSnippet = (group: RoutingGroup, baseUrl: string): string => + `import OpenAI from "openai"; + +const client = new OpenAI({ + apiKey: process.env.LITELLM_API_KEY, + baseURL: "${baseUrl}", +}); + +const response = await client.chat.completions.create({ + model: "${exampleModel(group)}", + messages: [{ role: "user", content: "Hello!" }], +}); + +console.log(response);`; + +const SNIPPET_TABS = [ + { value: "curl", label: "cURL", language: "bash", build: buildCurlSnippet }, + { value: "python", label: "Python (OpenAI SDK)", language: "python", build: buildPythonSnippet }, + { value: "javascript", label: "JavaScript (OpenAI SDK)", language: "javascript", build: buildJsSnippet }, +] as const; + +export function RoutingGroupUsagePanel({ group, baseUrl }: RoutingGroupUsagePanelProps) { + return ( +
+
+ + How routing works for this group +
+

+ Callers request any model in the group by name; LiteLLM picks a deployment behind the scenes using the{" "} + {formatStrategyLabel(group.routing_strategy)} strategy. +

+ + + {SNIPPET_TABS.map((tab) => ( + + {tab.label} + + ))} + + {SNIPPET_TABS.map((tab) => ( + + + + ))} + +
+ ); +} diff --git a/ui/litellm-dashboard/src/components/routing_groups/RoutingGroupsTable.test.tsx b/ui/litellm-dashboard/src/components/routing_groups/RoutingGroupsTable.test.tsx new file mode 100644 index 00000000000..6f14b76e2fd --- /dev/null +++ b/ui/litellm-dashboard/src/components/routing_groups/RoutingGroupsTable.test.tsx @@ -0,0 +1,145 @@ +import { render, screen, within } from "@testing-library/react"; +import userEvent from "@testing-library/user-event"; +import { beforeEach, describe, expect, it, vi } from "vitest"; + +import RoutingGroupsTable from "./RoutingGroupsTable"; +import type { RoutingGroup } from "./types"; + +describe("RoutingGroupsTable", () => { + const onEdit = vi.fn(); + const onDelete = vi.fn(); + + const prodGroup: RoutingGroup = { + group_name: "prod-group", + models: ["gpt-4o", "claude-sonnet-4-5"], + routing_strategy: "usage-based-routing", + }; + + const devGroup: RoutingGroup = { + group_name: "dev-group", + models: ["gpt-4o-mini"], + routing_strategy: "simple-shuffle", + }; + + const defaultProps = { + groups: [] as RoutingGroup[], + onEdit, + onDelete, + proxyBaseUrl: "https://proxy.example.com", + }; + + const rowFor = (groupName: string): HTMLElement => { + const row = document.querySelector(`[data-row-id="${groupName}"]`); + if (!(row instanceof HTMLElement)) { + throw new Error(`No row rendered for ${groupName}`); + } + return row; + }; + + beforeEach(() => { + vi.clearAllMocks(); + }); + + it("should render every column header", () => { + render(); + for (const header of ["Group Name", "Models", "Strategy"]) { + expect(screen.getByText(header)).toBeInTheDocument(); + } + }); + + it("should show the empty state when there are no groups", () => { + render(); + expect(screen.getByText("No routing groups yet")).toBeInTheDocument(); + }); + + it("should render the group name, its models, and a human-readable strategy label", () => { + render(); + const row = rowFor("prod-group"); + expect(within(row).getByText("prod-group")).toBeInTheDocument(); + expect(within(row).getByText("gpt-4o")).toBeInTheDocument(); + expect(within(row).getByText("claude-sonnet-4-5")).toBeInTheDocument(); + expect(within(row).getByText("Usage Based")).toBeInTheDocument(); + }); + + it("should fall back to the raw strategy value when it has no friendly label", () => { + render(); + expect(within(rowFor("prod-group")).getByText("custom-strategy")).toBeInTheDocument(); + }); + + it("should collapse models beyond the first three behind a +N more badge", () => { + const wideGroup: RoutingGroup = { ...prodGroup, models: ["a", "b", "c", "d", "e"] }; + render(); + const row = rowFor("prod-group"); + expect(within(row).getByText("+2 more")).toBeInTheDocument(); + expect(within(row).queryByText("d")).not.toBeInTheDocument(); + }); + + it("should keep the incoming order until a column is sorted", async () => { + const user = userEvent.setup(); + render(); + + const namesInOrder = () => + screen + .getAllByRole("row") + .slice(1) + .map((row) => row.getAttribute("data-row-id")); + + expect(namesInOrder()).toEqual(["prod-group", "dev-group"]); + + await user.click(screen.getByTestId("sort-header-group_name")); + expect(namesInOrder()).toEqual(["dev-group", "prod-group"]); + }); + + it("should toggle the usage panel when the group name is clicked", async () => { + const user = userEvent.setup(); + render(); + + expect(screen.queryByText("How routing works for this group")).not.toBeInTheDocument(); + + await user.click(screen.getByRole("button", { name: "prod-group" })); + expect(await screen.findByText("How routing works for this group")).toBeInTheDocument(); + + await user.click(screen.getByRole("button", { name: "prod-group" })); + expect(screen.queryByText("How routing works for this group")).not.toBeInTheDocument(); + }); + + it("should build the usage snippet from the proxy base url and the group's first model", async () => { + const user = userEvent.setup(); + render(); + await user.click(screen.getByRole("button", { name: "prod-group" })); + + const panel = (await screen.findByText("How routing works for this group")).closest("div")?.parentElement; + expect(panel?.textContent).toContain("https://proxy.example.com"); + expect(panel?.textContent).toContain("gpt-4o"); + }); + + it("should expand only the clicked group", async () => { + const user = userEvent.setup(); + render(); + + await user.click(screen.getByRole("button", { name: "dev-group" })); + expect(await screen.findAllByText("How routing works for this group")).toHaveLength(1); + expect(within(rowFor("prod-group")).queryByText("How routing works for this group")).not.toBeInTheDocument(); + }); + + it("should edit a group through the actions menu", async () => { + const user = userEvent.setup(); + render(); + await user.click(screen.getByTestId("routing-group-actions-prod-group")); + await user.click(await screen.findByTestId("routing-group-action-edit")); + expect(onEdit).toHaveBeenCalledWith(prodGroup); + }); + + it("should delete a group through the actions menu", async () => { + const user = userEvent.setup(); + render(); + await user.click(screen.getByTestId("routing-group-actions-prod-group")); + await user.click(await screen.findByTestId("routing-group-action-delete")); + expect(onDelete).toHaveBeenCalledWith(prodGroup); + }); + + it("should show skeleton rows instead of the empty state while loading", () => { + render(); + expect(screen.queryByText("No routing groups yet")).not.toBeInTheDocument(); + }); +}); diff --git a/ui/litellm-dashboard/src/components/routing_groups/RoutingGroupsTable.tsx b/ui/litellm-dashboard/src/components/routing_groups/RoutingGroupsTable.tsx index 96f6e578c83..fce887fc63b 100644 --- a/ui/litellm-dashboard/src/components/routing_groups/RoutingGroupsTable.tsx +++ b/ui/litellm-dashboard/src/components/routing_groups/RoutingGroupsTable.tsx @@ -1,229 +1,82 @@ "use client"; -import React, { useState } from "react"; -import { Flex, Table, Tabs, Tag, Tooltip, Typography, Button } from "antd"; -import type { ColumnsType } from "antd/es/table"; -import { BranchesOutlined, DeleteOutlined, EditOutlined, CodeOutlined } from "@ant-design/icons"; -import type { RoutingGroup } from "./types"; +import type { ExpandedState, SortingState } from "@tanstack/react-table"; +import { Inbox } from "lucide-react"; +import React, { useCallback, useMemo, useState } from "react"; -const { Text, Paragraph } = Typography; +import { DataTable } from "@/components/shared/DataTable"; + +import { RoutingGroupUsagePanel } from "./RoutingGroupUsagePanel"; +import { getRoutingGroupsTableColumns } from "./RoutingGroupsTableColumns"; +import type { RoutingGroup } from "./types"; interface RoutingGroupsTableProps { groups: RoutingGroup[]; - loading?: boolean; + isLoading?: boolean; onEdit: (group: RoutingGroup) => void; onDelete: (group: RoutingGroup) => void; proxyBaseUrl?: string; } -const formatStrategyLabel = (strategy: string): string => { - switch (strategy) { - case "simple-shuffle": - return "Simple Shuffle"; - case "least-busy": - return "Least Busy"; - case "usage-based-routing": - return "Usage Based"; - case "latency-based-routing": - return "Latency Based"; - default: - return strategy; - } -}; - const resolveBaseUrl = (proxyBaseUrl?: string): string => { if (proxyBaseUrl && proxyBaseUrl.trim()) return proxyBaseUrl; if (typeof window !== "undefined" && window.location?.origin) return window.location.origin; return ""; }; -const exampleModel = (group: RoutingGroup): string => group.models[0] ?? ""; - -const buildCurlSnippet = (group: RoutingGroup, baseUrl: string): string => - `curl -X POST '${baseUrl}/v1/chat/completions' \\ - -H 'Content-Type: application/json' \\ - -H 'Authorization: Bearer $LITELLM_API_KEY' \\ - -d '{ - "model": "${exampleModel(group)}", - "messages": [{"role": "user", "content": "Hello!"}] - }'`; - -const buildPythonSnippet = (group: RoutingGroup, baseUrl: string): string => - `from openai import OpenAI - -client = OpenAI( - api_key="$LITELLM_API_KEY", - base_url="${baseUrl}", -) - -response = client.chat.completions.create( - model="${exampleModel(group)}", - messages=[{"role": "user", "content": "Hello!"}], -) - -print(response)`; - -const buildJsSnippet = (group: RoutingGroup, baseUrl: string): string => - `import OpenAI from "openai"; - -const client = new OpenAI({ - apiKey: process.env.LITELLM_API_KEY, - baseURL: "${baseUrl}", -}); - -const response = await client.chat.completions.create({ - model: "${exampleModel(group)}", - messages: [{ role: "user", content: "Hello!" }], -}); - -console.log(response);`; - -interface RoutingGroupSnippetProps { - group: RoutingGroup; - baseUrl: string; +function EmptyState() { + return ( +
+
+ +
+
No routing groups yet
+
+ Create a group to load-balance a set of models behind one name. +
+
+ ); } -const SNIPPET_BLOCK_STYLE: React.CSSProperties = { - backgroundColor: "#111827", - color: "#f3f4f6", - borderRadius: 6, - padding: 16, - fontSize: 12, - whiteSpace: "pre", - overflowX: "auto", -}; - -const RoutingGroupSnippet: React.FC = ({ group, baseUrl }) => { - const snippets = { - curl: buildCurlSnippet(group, baseUrl), - python: buildPythonSnippet(group, baseUrl), - javascript: buildJsSnippet(group, baseUrl), - } as const; - type SnippetKey = keyof typeof snippets; - const [activeKey, setActiveKey] = useState("curl"); - - const items = [ - { key: "curl", label: "cURL" }, - { key: "python", label: "Python (OpenAI SDK)" }, - { key: "javascript", label: "JavaScript (OpenAI SDK)" }, - ].map(({ key, label }) => ({ - key, - label, - children: ( - - {snippets[key as SnippetKey]} - - ), - })); - - return ( - setActiveKey(k as SnippetKey)} - items={items} - tabBarExtraContent={ - - } - /> - ); -}; - -const RoutingGroupsTable: React.FC = ({ groups, loading, onEdit, onDelete, proxyBaseUrl }) => { - const [expandedRowKeys, setExpandedRowKeys] = useState([]); +const RoutingGroupsTable: React.FC = ({ + groups, + isLoading, + onEdit, + onDelete, + proxyBaseUrl, +}) => { + const [sorting, setSorting] = useState([]); + const [expanded, setExpanded] = useState({}); const baseUrl = resolveBaseUrl(proxyBaseUrl); - const columns: ColumnsType = [ - { - title: "GROUP NAME", - dataIndex: "group_name", - key: "group_name", - render: (name: string) => ( - - {name} - - ), - }, - { - title: "MODELS", - dataIndex: "models", - key: "models", - render: (models: string[]) => ( - - {models.map((m) => ( - {m} - ))} - - ), - }, - { - title: "STRATEGY", - dataIndex: "routing_strategy", - key: "routing_strategy", - render: (strategy: string) => ( - - - {formatStrategyLabel(strategy)} - - ), - }, - { - title: "ACTIONS", - key: "actions", - width: 120, - align: "right", - render: (_, group) => ( - - - + + {rotationInterval} + + ); +}; + +const getDurationInput = (isCreateMode = true) => + screen.getByPlaceholderText(isCreateMode ? CREATE_PLACEHOLDER : EDIT_PLACEHOLDER) as HTMLInputElement; describe("KeyLifecycleSettings", () => { - const mockForm = { - getFieldValue: vi.fn(), - setFieldValue: vi.fn(), - setFieldsValue: vi.fn(), - }; - - const defaultProps = { - form: mockForm, - autoRotationEnabled: false, - onAutoRotationChange: vi.fn(), - rotationInterval: "", - onRotationIntervalChange: vi.fn(), - isCreateMode: false, - }; - beforeEach(() => { vi.clearAllMocks(); - mockForm.getFieldValue.mockReturnValue(""); }); - it("should render without crashing", () => { - renderWithProviders(); - + it("renders the expiry and auto-rotation sections", () => { + renderWithProviders(); expect(screen.getByText("Key Expiry Settings")).toBeInTheDocument(); expect(screen.getByText("Auto-Rotation Settings")).toBeInTheDocument(); + expect(getDurationInput()).toBeInTheDocument(); }); - describe("Key Expiry Settings", () => { - it("should render expiry input field", () => { - renderWithProviders(); + it("uses the create-mode placeholder in create mode", () => { + renderWithProviders(); + expect(screen.getByPlaceholderText(CREATE_PLACEHOLDER)).toBeInTheDocument(); + }); - expect(screen.getByText("Expire Key")).toBeInTheDocument(); - expect(screen.getByTestId("duration-input")).toBeInTheDocument(); - }); + it("uses the edit-mode placeholder in edit mode", () => { + renderWithProviders(); + expect(screen.getByPlaceholderText(EDIT_PLACEHOLDER)).toBeInTheDocument(); + }); - it("should show correct placeholder in create mode", () => { - renderWithProviders(); - - const input = screen.getByTestId("duration-input"); - expect(input).toHaveAttribute("placeholder", "e.g., 30d or leave empty to never expire"); - }); - - it("should show correct placeholder in edit mode", () => { - renderWithProviders(); - - const input = screen.getByTestId("duration-input"); - expect(input).toHaveAttribute("placeholder", "e.g., 30d"); - }); - - it("should show correct tooltip in create mode", () => { - renderWithProviders(); - - const tooltips = screen.getAllByTestId("tooltip"); - const expiryTooltip = tooltips.find((tooltip) => - tooltip.getAttribute("title")?.includes("Leave empty to keep the current expiry unchanged"), - ); - expect(expiryTooltip).toBeInTheDocument(); - expect(expiryTooltip).toHaveAttribute( - "title", - "Set when this key should expire. Format: 30s (seconds), 30m (minutes), 30h (hours), 30d (days). Leave empty to keep the current expiry unchanged.", - ); - }); - - it("should show correct tooltip in edit mode", () => { - renderWithProviders(); - - const tooltips = screen.getAllByTestId("tooltip"); - const expiryTooltip = tooltips.find((tooltip) => - tooltip.getAttribute("title")?.includes("Leave empty to keep the current expiry unchanged"), - ); - expect(expiryTooltip).toBeInTheDocument(); - expect(expiryTooltip).toHaveAttribute( - "title", - "Set when this key should expire. Format: 30s (seconds), 30m (minutes), 30h (hours), 30d (days). Leave empty to keep the current expiry unchanged.", - ); - }); - - it("should initialize with form value if present", () => { - mockForm.getFieldValue.mockReturnValue("30d"); - renderWithProviders(); - - const input = screen.getByTestId("duration-input") as HTMLInputElement; - expect(input.value).toBe("30d"); - }); - - it("should update form using setFieldValue when duration changes", async () => { + describe("duration is a single source of truth (regression for pre-filled value dropped on submit)", () => { + it("submits the duration the user typed", async () => { const user = userEvent.setup(); - renderWithProviders(); + const onFinish = vi.fn(); + renderWithProviders(); - const input = screen.getByTestId("duration-input"); - await user.type(input, "60d"); + await user.type(getDurationInput(), "1d"); + await user.click(screen.getByRole("button", { name: "submit" })); - expect(mockForm.setFieldValue).toHaveBeenCalledWith("duration", "60d"); + await waitFor(() => expect(onFinish).toHaveBeenCalledTimes(1)); + expect(onFinish.mock.calls[0][0]).toMatchObject({ duration: "1d" }); }); - it("should update form using setFieldsValue when setFieldValue is not available", async () => { + it("clears the displayed value when the form is reset, so no stale value lingers", async () => { const user = userEvent.setup(); - const formWithoutSetFieldValue = { - getFieldValue: vi.fn().mockReturnValue(""), - setFieldsValue: vi.fn(), - }; - renderWithProviders(); + renderWithProviders(); - const input = screen.getByTestId("duration-input"); - await user.type(input, "90d"); + await user.type(getDurationInput(), "1d"); + expect(getDurationInput().value).toBe("1d"); - expect(formWithoutSetFieldValue.setFieldsValue).toHaveBeenCalledWith({ duration: "90d" }); + await user.click(screen.getByRole("button", { name: "reset" })); + + await waitFor(() => expect(getDurationInput().value).toBe("")); + }); + + it("never submits a value that differs from what is displayed after a reset", async () => { + const user = userEvent.setup(); + const onFinish = vi.fn(); + renderWithProviders(); + + // First create: type "1d" and submit -> "1d" is sent. + await user.type(getDurationInput(), "1d"); + await user.click(screen.getByRole("button", { name: "submit" })); + await waitFor(() => expect(onFinish).toHaveBeenCalledTimes(1)); + expect(onFinish.mock.calls[0][0]).toMatchObject({ duration: "1d" }); + + // Second create: form resets, so the field must show empty AND submit empty. + // The old bug showed a stale "1d" while submitting null/empty. + await user.click(screen.getByRole("button", { name: "reset" })); + await waitFor(() => expect(getDurationInput().value).toBe("")); + + await user.click(screen.getByRole("button", { name: "submit" })); + await waitFor(() => expect(onFinish).toHaveBeenCalledTimes(2)); + expect(onFinish.mock.calls[1][0].duration).not.toBe("1d"); + expect(getDurationInput().value).toBe(onFinish.mock.calls[1][0].duration ?? ""); }); }); - describe("Auto-Rotation Settings", () => { - it("should render auto-rotation switch", () => { - renderWithProviders(); - - expect(screen.getByText("Enable Auto-Rotation")).toBeInTheDocument(); - expect(screen.getByTestId("switch")).toBeInTheDocument(); - }); - - it("should show switch as unchecked when autoRotationEnabled is false", () => { - renderWithProviders(); - - const switchElement = screen.getByTestId("switch") as HTMLInputElement; - expect(switchElement.checked).toBe(false); - }); - - it("should show switch as checked when autoRotationEnabled is true", () => { - renderWithProviders(); - - const switchElement = screen.getByTestId("switch") as HTMLInputElement; - expect(switchElement.checked).toBe(true); - }); - - it("should call onAutoRotationChange when switch is toggled", async () => { + describe("Never Expire", () => { + it("clears and disables the duration input, then submits an empty duration", async () => { const user = userEvent.setup(); - const onAutoRotationChange = vi.fn(); - renderWithProviders(); + const onFinish = vi.fn(); + renderWithProviders(); - const switchElement = screen.getByTestId("switch"); - await user.click(switchElement); + await user.type(getDurationInput(false), "30d"); + expect(getDurationInput(false).value).toBe("30d"); - expect(onAutoRotationChange).toHaveBeenCalledWith(true); + await user.click(screen.getByRole("checkbox", { name: /never expire/i })); + + await waitFor(() => expect(getDurationInput(false).value).toBe("")); + expect(getDurationInput(false)).toBeDisabled(); + + await user.click(screen.getByRole("button", { name: "submit" })); + await waitFor(() => expect(onFinish).toHaveBeenCalledTimes(1)); + expect(onFinish.mock.calls[0][0]).toMatchObject({ duration: "" }); }); + }); - it("should not show rotation interval section when auto-rotation is disabled", () => { - renderWithProviders(); + describe("Auto-Rotation", () => { + it("reveals the rotation interval controls when enabled", async () => { + const user = userEvent.setup(); + renderWithProviders(); expect(screen.queryByText("Rotation Interval")).not.toBeInTheDocument(); - expect(screen.queryByTestId("select")).not.toBeInTheDocument(); + await user.click(screen.getByRole("switch")); + + await waitFor(() => expect(screen.getByText("Rotation Interval")).toBeInTheDocument()); }); - it("should show rotation interval section when auto-rotation is enabled", () => { - renderWithProviders(); - - expect(screen.getByText("Rotation Interval")).toBeInTheDocument(); - expect(screen.getByTestId("select")).toBeInTheDocument(); - }); - - it("should show all predefined interval options", () => { - renderWithProviders(); - - expect(screen.getByText("7 days")).toBeInTheDocument(); - expect(screen.getByText("30 days")).toBeInTheDocument(); - expect(screen.getByText("90 days")).toBeInTheDocument(); - expect(screen.getByText("180 days")).toBeInTheDocument(); - expect(screen.getByText("365 days")).toBeInTheDocument(); - expect(screen.getByText("Custom interval")).toBeInTheDocument(); - }); - - it("should display current rotation interval in select", () => { - renderWithProviders(); - - const select = screen.getByTestId("select") as HTMLSelectElement; - expect(select.value).toBe("90d"); - }); - - it("should call onRotationIntervalChange when predefined interval is selected", async () => { + it("propagates a selected predefined interval", async () => { const user = userEvent.setup(); - const onRotationIntervalChange = vi.fn(); - renderWithProviders( - , - ); + renderWithProviders(); - const select = screen.getByTestId("select"); - await user.selectOptions(select, "30d"); + await user.click(screen.getByRole("switch")); + await waitFor(() => expect(screen.getByText("Rotation Interval")).toBeInTheDocument()); - expect(onRotationIntervalChange).toHaveBeenCalledWith("30d"); + await user.click(screen.getByRole("combobox")); + await user.click(await screen.findByText("90 days")); + + await waitFor(() => expect(document.querySelector(".ant-select-selection-item")?.textContent).toBe("90 days")); + expect(screen.getByTestId("rotation-interval-value")).toHaveTextContent("90d"); }); - it("should show custom input when custom option is selected", async () => { + it("shows the custom interval input when Custom interval is selected, without propagating yet", async () => { const user = userEvent.setup(); - renderWithProviders(); + renderWithProviders(); - const select = screen.getByTestId("select"); - await user.selectOptions(select, "custom"); + await user.click(screen.getByRole("switch")); + await waitFor(() => expect(screen.getByText("Rotation Interval")).toBeInTheDocument()); - expect(screen.getByTestId("custom-interval-input")).toBeInTheDocument(); + await user.click(screen.getByRole("combobox")); + await user.click(await screen.findByText("Custom interval")); + + expect(await screen.findByPlaceholderText("e.g., 1s, 5m, 2h, 14d")).toBeInTheDocument(); expect(screen.getByText("Supported formats: seconds (s), minutes (m), hours (h), days (d)")).toBeInTheDocument(); + expect(screen.getByTestId("rotation-interval-value")).toHaveTextContent(""); }); - it("should hide custom input when predefined interval is selected after custom", async () => { + it("propagates a typed custom interval to the parent", async () => { const user = userEvent.setup(); - const onRotationIntervalChange = vi.fn(); - renderWithProviders( - , - ); + renderWithProviders(); - const select = screen.getByTestId("select"); - await user.selectOptions(select, "7d"); + await user.click(screen.getByRole("switch")); + await waitFor(() => expect(screen.getByText("Rotation Interval")).toBeInTheDocument()); - expect(screen.queryByTestId("custom-interval-input")).not.toBeInTheDocument(); - expect(onRotationIntervalChange).toHaveBeenCalledWith("7d"); - }); + await user.click(screen.getByRole("combobox")); + await user.click(await screen.findByText("Custom interval")); - it("should call onRotationIntervalChange when custom interval is entered", async () => { - const user = userEvent.setup(); - const onRotationIntervalChange = vi.fn(); - renderWithProviders( - , - ); - - const select = screen.getByTestId("select"); - await user.selectOptions(select, "custom"); - - const customInput = screen.getByTestId("custom-interval-input"); + const customInput = await screen.findByPlaceholderText("e.g., 1s, 5m, 2h, 14d"); await user.type(customInput, "14d"); - expect(onRotationIntervalChange).toHaveBeenCalledWith("14d"); + await waitFor(() => expect(screen.getByTestId("rotation-interval-value")).toHaveTextContent("14d")); + expect((customInput as HTMLInputElement).value).toBe("14d"); }); - it("should show info message when auto-rotation is enabled", () => { - renderWithProviders(); - - expect( - screen.getByText( - "When rotation occurs, you'll receive a notification with the new key. The old key will be deactivated after a brief grace period.", - ), - ).toBeInTheDocument(); - }); - - it("should not show info message when auto-rotation is disabled", () => { - renderWithProviders(); - - expect( - screen.queryByText( - "When rotation occurs, you'll receive a notification with the new key. The old key will be deactivated after a brief grace period.", - ), - ).not.toBeInTheDocument(); - }); - - it("should initialize with custom interval input visible when custom interval is provided", () => { - renderWithProviders(); - - expect(screen.getByTestId("custom-interval-input")).toBeInTheDocument(); - const customInput = screen.getByTestId("custom-interval-input") as HTMLInputElement; - expect(customInput.value).toBe("14d"); - }); - - it("should show custom option selected when custom interval is provided", () => { - renderWithProviders(); - - const select = screen.getByTestId("select") as HTMLSelectElement; - expect(select.value).toBe("custom"); - }); - - it("should not call onRotationIntervalChange when selecting custom option", async () => { + it("hides the custom input and propagates the value when switching back to a predefined interval", async () => { const user = userEvent.setup(); - const onRotationIntervalChange = vi.fn(); - renderWithProviders( - , - ); + renderWithProviders(); - const select = screen.getByTestId("select"); - await user.selectOptions(select, "custom"); + await user.click(screen.getByRole("switch")); + await waitFor(() => expect(screen.getByText("Rotation Interval")).toBeInTheDocument()); - expect(onRotationIntervalChange).not.toHaveBeenCalled(); + await user.click(screen.getByRole("combobox")); + await user.click(await screen.findByText("Custom interval")); + const customInput = await screen.findByPlaceholderText("e.g., 1s, 5m, 2h, 14d"); + await user.type(customInput, "14d"); + await waitFor(() => expect(screen.getByTestId("rotation-interval-value")).toHaveTextContent("14d")); + + await user.click(screen.getByRole("combobox")); + await user.click(await screen.findByText("7 days")); + + await waitFor(() => expect(screen.getByTestId("rotation-interval-value")).toHaveTextContent("7d")); + expect(screen.queryByPlaceholderText("e.g., 1s, 5m, 2h, 14d")).not.toBeInTheDocument(); }); }); }); diff --git a/ui/litellm-dashboard/src/components/common_components/KeyLifecycleSettings.tsx b/ui/litellm-dashboard/src/components/common_components/KeyLifecycleSettings.tsx index 7c4738f9ede..8e88fab1095 100644 --- a/ui/litellm-dashboard/src/components/common_components/KeyLifecycleSettings.tsx +++ b/ui/litellm-dashboard/src/components/common_components/KeyLifecycleSettings.tsx @@ -1,5 +1,5 @@ import React, { useState } from "react"; -import { Select, Tooltip, Divider, Switch, Checkbox } from "antd"; +import { Select, Tooltip, Divider, Switch, Checkbox, Form } from "antd"; import { InfoCircleOutlined } from "@ant-design/icons"; import { TextInput } from "@tremor/react"; @@ -34,7 +34,6 @@ const KeyLifecycleSettings: React.FC = ({ const [showCustomInput, setShowCustomInput] = useState(isCustomInterval); const [customInterval, setCustomInterval] = useState(isCustomInterval ? rotationInterval : ""); - const [durationValue, setDurationValue] = useState(form?.getFieldValue?.("duration") || ""); const handleIntervalChange = (value: string) => { if (value === "custom") { @@ -53,14 +52,6 @@ const KeyLifecycleSettings: React.FC = ({ onRotationIntervalChange(value); }; - const handleDurationChange = (value: string) => { - setDurationValue(value); - if (form && typeof form.setFieldValue === "function") { - form.setFieldValue("duration", value); - } else if (form && typeof form.setFieldsValue === "function") { - form.setFieldsValue({ duration: value }); - } - }; return (
{/* Key Expiry Section */} @@ -80,7 +71,6 @@ const KeyLifecycleSettings: React.FC = ({ const checked = e.target.checked; onNeverExpireChange(checked); if (checked) { - setDurationValue(""); if (form && typeof form.setFieldValue === "function") { form.setFieldValue("duration", ""); } else if (form && typeof form.setFieldsValue === "function") { @@ -94,14 +84,13 @@ const KeyLifecycleSettings: React.FC = ({ )} - + + +
diff --git a/ui/litellm-dashboard/src/components/organisms/create_key_button.tsx b/ui/litellm-dashboard/src/components/organisms/create_key_button.tsx index 2a04bad6aa8..711652eb783 100644 --- a/ui/litellm-dashboard/src/components/organisms/create_key_button.tsx +++ b/ui/litellm-dashboard/src/components/organisms/create_key_button.tsx @@ -1644,9 +1644,6 @@ const CreateKey: React.FC = ({ team, teams, data, addKey, autoOp /> - diff --git a/ui/litellm-dashboard/src/components/templates/key_edit_view.tsx b/ui/litellm-dashboard/src/components/templates/key_edit_view.tsx index a9fa05d817d..962c6bc3568 100644 --- a/ui/litellm-dashboard/src/components/templates/key_edit_view.tsx +++ b/ui/litellm-dashboard/src/components/templates/key_edit_view.tsx @@ -860,9 +860,6 @@ export function KeyEditView({ neverExpire={neverExpire} onNeverExpireChange={setNeverExpire} /> - {/* Hidden form field for token */} From 3966fbf5ec9f87e680daed03d751b69d36a5e207 Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Fri, 24 Jul 2026 16:08:35 -0700 Subject: [PATCH 15/29] chore(ui): drop the unused antd table header sort dropdown TableHeaderSortDropdown had no importers left; the shared DataTable's DataTableSortHeader covers the same ascending/descending/reset menu on Base UI. knip did not flag it because its own test file counted as a usage. --- .../TableHeaderSortDropdown.test.tsx | 148 ------------------ .../TableHeaderSortDropdown.tsx | 82 ---------- 2 files changed, 230 deletions(-) delete mode 100644 ui/litellm-dashboard/src/components/common_components/TableHeaderSortDropdown/TableHeaderSortDropdown.test.tsx delete mode 100644 ui/litellm-dashboard/src/components/common_components/TableHeaderSortDropdown/TableHeaderSortDropdown.tsx diff --git a/ui/litellm-dashboard/src/components/common_components/TableHeaderSortDropdown/TableHeaderSortDropdown.test.tsx b/ui/litellm-dashboard/src/components/common_components/TableHeaderSortDropdown/TableHeaderSortDropdown.test.tsx deleted file mode 100644 index 58395371bbe..00000000000 --- a/ui/litellm-dashboard/src/components/common_components/TableHeaderSortDropdown/TableHeaderSortDropdown.test.tsx +++ /dev/null @@ -1,148 +0,0 @@ -import { render, screen, waitFor } from "@testing-library/react"; -import userEvent from "@testing-library/user-event"; -import { describe, expect, it, vi } from "vitest"; -import { TableHeaderSortDropdown } from "./TableHeaderSortDropdown"; - -describe("TableHeaderSortDropdown", () => { - it("should render", () => { - const onSortChange = vi.fn(); - render(); - expect(screen.getByRole("button")).toBeInTheDocument(); - }); - - it("should open dropdown menu when button is clicked", async () => { - const user = userEvent.setup(); - const onSortChange = vi.fn(); - render(); - - const button = screen.getByRole("button"); - await user.click(button); - - await waitFor(() => { - expect(screen.getByText("Ascending")).toBeInTheDocument(); - expect(screen.getByText("Descending")).toBeInTheDocument(); - expect(screen.getByText("Reset")).toBeInTheDocument(); - }); - }); - - it("should call onSortChange with asc when ascending option is clicked", async () => { - const user = userEvent.setup(); - const onSortChange = vi.fn(); - render(); - - const button = screen.getByRole("button"); - await user.click(button); - - await waitFor(() => { - expect(screen.getByText("Ascending")).toBeInTheDocument(); - }); - - const ascendingOption = screen.getByText("Ascending"); - await user.click(ascendingOption); - - expect(onSortChange).toHaveBeenCalledTimes(1); - expect(onSortChange).toHaveBeenCalledWith("asc"); - }); - - it("should call onSortChange with desc when descending option is clicked", async () => { - const user = userEvent.setup(); - const onSortChange = vi.fn(); - render(); - - const button = screen.getByRole("button"); - await user.click(button); - - await waitFor(() => { - expect(screen.getByText("Descending")).toBeInTheDocument(); - }); - - const descendingOption = screen.getByText("Descending"); - await user.click(descendingOption); - - expect(onSortChange).toHaveBeenCalledTimes(1); - expect(onSortChange).toHaveBeenCalledWith("desc"); - }); - - it("should call onSortChange with false when reset option is clicked", async () => { - const user = userEvent.setup(); - const onSortChange = vi.fn(); - render(); - - const button = screen.getByRole("button"); - await user.click(button); - - await waitFor(() => { - expect(screen.getByText("Reset")).toBeInTheDocument(); - }); - - const resetOption = screen.getByText("Reset"); - await user.click(resetOption); - - expect(onSortChange).toHaveBeenCalledTimes(1); - expect(onSortChange).toHaveBeenCalledWith(false); - }); - - it("should highlight ascending option when sort state is asc", async () => { - const user = userEvent.setup(); - const onSortChange = vi.fn(); - render(); - - const button = screen.getByRole("button"); - await user.click(button); - - await waitFor(() => { - const ascendingOption = screen.getByText("Ascending"); - const menuItem = ascendingOption.closest(".ant-dropdown-menu-item"); - expect(menuItem).toHaveClass("ant-dropdown-menu-item-selected"); - }); - }); - - it("should highlight descending option when sort state is desc", async () => { - const user = userEvent.setup(); - const onSortChange = vi.fn(); - render(); - - const button = screen.getByRole("button"); - await user.click(button); - - await waitFor(() => { - const descendingOption = screen.getByText("Descending"); - const menuItem = descendingOption.closest(".ant-dropdown-menu-item"); - expect(menuItem).toHaveClass("ant-dropdown-menu-item-selected"); - }); - }); - - it("should not highlight any option when sort state is false", async () => { - const user = userEvent.setup(); - const onSortChange = vi.fn(); - render(); - - const button = screen.getByRole("button"); - await user.click(button); - - await waitFor(() => { - expect(screen.getByText("Ascending")).toBeInTheDocument(); - }); - - const ascendingOption = screen.getByText("Ascending"); - const menuItem = ascendingOption.closest(".ant-dropdown-menu-item"); - expect(menuItem).not.toHaveClass("ant-dropdown-menu-item-selected"); - }); - - it("should stop event propagation when button is clicked", async () => { - const user = userEvent.setup(); - const onSortChange = vi.fn(); - const onParentClick = vi.fn(); - - render( -
- -
, - ); - - const button = screen.getByRole("button"); - await user.click(button); - - expect(onParentClick).not.toHaveBeenCalled(); - }); -}); diff --git a/ui/litellm-dashboard/src/components/common_components/TableHeaderSortDropdown/TableHeaderSortDropdown.tsx b/ui/litellm-dashboard/src/components/common_components/TableHeaderSortDropdown/TableHeaderSortDropdown.tsx deleted file mode 100644 index c83257c5c83..00000000000 --- a/ui/litellm-dashboard/src/components/common_components/TableHeaderSortDropdown/TableHeaderSortDropdown.tsx +++ /dev/null @@ -1,82 +0,0 @@ -import React from "react"; -import { Button, Dropdown, MenuProps } from "antd"; -import { SwitchVerticalIcon, ChevronUpIcon, ChevronDownIcon, XIcon } from "@heroicons/react/outline"; - -export type SortState = "asc" | "desc" | false; - -interface TableHeaderSortDropdownProps { - /** - * Current sort state: "asc", "desc", or false for neutral - */ - sortState: SortState; - /** - * Callback when sort state changes - * @param newState - The new sort state: "asc", "desc", or false - */ - onSortChange: (newState: SortState) => void; - /** - * Optional column ID for identification - */ - columnId?: string; -} - -export const TableHeaderSortDropdown: React.FC = ({ sortState, onSortChange }) => { - const handleMenuClick: MenuProps["onClick"] = ({ key }) => { - if (key === "asc") { - onSortChange("asc"); - } else if (key === "desc") { - onSortChange("desc"); - } else if (key === "reset") { - onSortChange(false); - } - }; - - const menuItems: MenuProps["items"] = [ - { - key: "asc", - label: "Ascending", - icon: , - }, - { - key: "desc", - label: "Descending", - icon: , - }, - { - key: "reset", - label: "Reset", - icon: , - }, - ]; - - // Determine which icon to display based on current sort state - const renderIcon = () => { - if (sortState === "asc") { - return ; - } else if (sortState === "desc") { - return ; - } else { - return ; - } - }; - - return ( - - - - - + , // Assuming formValues is an object -) => { - try { - if (formValues.metadata) { - // if there's an exception JSON.parse, show it in the message - try { - formValues.metadata = JSON.parse(formValues.metadata); - } catch (error) { - console.error("Failed to parse metadata:", error); - throw new Error("Failed to parse metadata: " + error); - } - } - - const data = await apiClient.post(`/organization/new`, { - accessToken, - body: { - ...formValues, // Include formValues in the request body - }, - }); - return data; - // Handle success - you might want to update some state or UI based on the created key - } catch (error) { - console.error("Failed to create key:", error); - throw error; - } -}; - export const organizationUpdateCall = async ( accessToken: string, formValues: Record, // Assuming formValues is an object diff --git a/ui/litellm-dashboard/src/components/organization/org-create/OrgCreateDialog.test.tsx b/ui/litellm-dashboard/src/components/organization/org-create/OrgCreateDialog.test.tsx new file mode 100644 index 00000000000..5ec9bb1e633 --- /dev/null +++ b/ui/litellm-dashboard/src/components/organization/org-create/OrgCreateDialog.test.tsx @@ -0,0 +1,203 @@ +import { QueryClient, QueryClientProvider } from "@tanstack/react-query"; +import { render, screen, waitFor } from "@testing-library/react"; +import userEvent from "@testing-library/user-event"; +import React from "react"; +import { describe, expect, it, vi } from "vitest"; + +vi.mock("@/components/molecules/notifications_manager", () => ({ + __esModule: true, + default: { success: vi.fn(), fromBackend: vi.fn() }, +})); +vi.mock("@/components/ModelSelect/ModelSelect", () => ({ + ModelSelect: ({ onChange }: { onChange: (values: string[]) => void }) => ( + + ), +})); +vi.mock("@/components/vector_store_management/VectorStoreSelector", () => ({ + __esModule: true, + default: ({ onChange }: { onChange: (values: string[]) => void }) => ( + + ), +})); +vi.mock("@/components/mcp_server_management/MCPServerSelector", () => ({ + __esModule: true, + default: ({ + onChange, + }: { + onChange: (values: { servers: string[]; accessGroups: string[]; toolsets: string[] }) => void; + }) => ( + + ), +})); + +import { OrgCreateDialog } from "./OrgCreateDialog"; + +const Harness = ({ createOrganization }: { createOrganization: (body: unknown) => Promise }) => { + const [open, setOpen] = React.useState(true); + return ( + <> + + + + ); +}; + +const renderDialog = (overrides?: { createOrganization?: ReturnType }) => { + const createOrganization = overrides?.createOrganization ?? vi.fn().mockResolvedValue({}); + const queryClient = new QueryClient({ defaultOptions: { queries: { retry: false } } }); + render( + + + , + ); + return { createOrganization }; +}; + +describe("OrgCreateDialog", () => { + it("blocks submit and shows an error when the name is missing", async () => { + const user = userEvent.setup(); + const { createOrganization } = renderDialog(); + + await user.click(screen.getByRole("button", { name: "Create Organization" })); + + expect(await screen.findByRole("alert")).toHaveTextContent("Please input an organization name"); + expect(createOrganization).not.toHaveBeenCalled(); + }); + + it("sends only alias and models for a minimal create and closes the dialog", async () => { + const user = userEvent.setup(); + const { createOrganization } = renderDialog(); + + await user.type(screen.getByLabelText("Organization Name"), "new-org"); + await user.click(screen.getByRole("button", { name: "Create Organization" })); + + await waitFor(() => expect(createOrganization).toHaveBeenCalledTimes(1)); + expect(createOrganization.mock.calls[0][0]).toStrictEqual({ organization_alias: "new-org", models: [] }); + await waitFor(() => expect(screen.queryByLabelText("Organization Name")).not.toBeInTheDocument()); + }); + + it("maps selectors and limits into the create body", async () => { + const user = userEvent.setup(); + const { createOrganization } = renderDialog(); + + await user.type(screen.getByLabelText("Organization Name"), "new-org"); + await user.click(screen.getByRole("button", { name: "set-models" })); + await user.type(screen.getByLabelText("Tokens per minute Limit (TPM)"), "1000"); + await user.click(screen.getByRole("button", { name: "set-vector-stores" })); + await user.click(screen.getByRole("button", { name: "set-mcp" })); + await user.click(screen.getByRole("button", { name: "Create Organization" })); + + await waitFor(() => expect(createOrganization).toHaveBeenCalledTimes(1)); + const expectedBody = { + organization_alias: "new-org", + models: ["gpt-5.2"], + tpm_limit: 1000, + object_permission: { + vector_stores: ["vs-1"], + mcp_servers: ["srv-1"], + mcp_toolsets: ["ts-1"], + }, + }; + expect(createOrganization.mock.calls[0][0]).toStrictEqual(expectedBody); + }); + + it("blocks submit and shows an error for invalid metadata JSON", async () => { + const user = userEvent.setup(); + const { createOrganization } = renderDialog(); + + await user.type(screen.getByLabelText("Organization Name"), "new-org"); + await user.type(screen.getByLabelText("Metadata"), "not json"); + await user.click(screen.getByRole("button", { name: "Create Organization" })); + + expect(await screen.findByRole("alert")).toHaveTextContent("Metadata must be a valid JSON object"); + expect(createOrganization).not.toHaveBeenCalled(); + }); + + it("keeps the dialog open with the entered values when the create fails", async () => { + const user = userEvent.setup(); + const { createOrganization } = renderDialog({ + createOrganization: vi.fn().mockRejectedValue(new Error("boom")), + }); + + await user.type(screen.getByLabelText("Organization Name"), "new-org"); + await user.click(screen.getByRole("button", { name: "Create Organization" })); + + await waitFor(() => expect(createOrganization).toHaveBeenCalledTimes(1)); + expect(screen.getByLabelText("Organization Name")).toHaveValue("new-org"); + }); + + it("resets the form when the dialog is cancelled and reopened", async () => { + const user = userEvent.setup(); + renderDialog(); + + await user.type(screen.getByLabelText("Organization Name"), "abandoned"); + await user.click(screen.getByRole("button", { name: "Cancel" })); + await waitFor(() => expect(screen.queryByLabelText("Organization Name")).not.toBeInTheDocument()); + + await user.click(screen.getByRole("button", { name: "reopen" })); + expect(screen.getByLabelText("Organization Name")).toHaveValue(""); + }); + + it("resets the form when the dialog is dismissed with Escape and reopened", async () => { + const user = userEvent.setup(); + renderDialog(); + + await user.type(screen.getByLabelText("Organization Name"), "abandoned"); + await user.keyboard("{Escape}"); + await waitFor(() => expect(screen.queryByLabelText("Organization Name")).not.toBeInTheDocument()); + + await user.click(screen.getByRole("button", { name: "reopen" })); + expect(screen.getByLabelText("Organization Name")).toHaveValue(""); + }); + + it("cannot be dismissed while a create is pending, then closes once on success", async () => { + const user = userEvent.setup(); + let resolveCreate: (value: unknown) => void = () => {}; + const createOrganization = vi.fn().mockImplementation( + () => + new Promise((resolve) => { + resolveCreate = resolve; + }), + ); + renderDialog({ createOrganization }); + + await user.type(screen.getByLabelText("Organization Name"), "new-org"); + await user.keyboard("{Enter}"); + await waitFor(() => expect(createOrganization).toHaveBeenCalledTimes(1)); + + await user.keyboard("{Escape}"); + expect(screen.getByLabelText("Organization Name")).toHaveValue("new-org"); + + resolveCreate({}); + await waitFor(() => expect(screen.queryByLabelText("Organization Name")).not.toBeInTheDocument()); + }); + + it("does not fire a second create while one is pending", async () => { + const user = userEvent.setup(); + let resolveCreate: (value: unknown) => void = () => {}; + const createOrganization = vi.fn().mockImplementation( + () => + new Promise((resolve) => { + resolveCreate = resolve; + }), + ); + renderDialog({ createOrganization }); + + await user.type(screen.getByLabelText("Organization Name"), "new-org"); + await user.keyboard("{Enter}"); + await waitFor(() => expect(createOrganization).toHaveBeenCalledTimes(1)); + await user.keyboard("{Enter}"); + + expect(createOrganization).toHaveBeenCalledTimes(1); + resolveCreate({}); + await waitFor(() => expect(screen.queryByLabelText("Organization Name")).not.toBeInTheDocument()); + }); +}); diff --git a/ui/litellm-dashboard/src/components/organization/org-create/OrgCreateDialog.tsx b/ui/litellm-dashboard/src/components/organization/org-create/OrgCreateDialog.tsx new file mode 100644 index 00000000000..998d9446365 --- /dev/null +++ b/ui/litellm-dashboard/src/components/organization/org-create/OrgCreateDialog.tsx @@ -0,0 +1,186 @@ +"use client"; + +import { useMutation, useQueryClient } from "@tanstack/react-query"; +import * as React from "react"; + +import { organizationKeys } from "@/app/(dashboard)/hooks/organizations/useOrganizations"; +import { ModelSelect } from "@/components/ModelSelect/ModelSelect"; +import MCPServerSelector from "@/components/mcp_server_management/MCPServerSelector"; +import NotificationsManager from "@/components/molecules/notifications_manager"; +import { FieldGroup } from "@/components/shared/form/field"; +import { FormField } from "@/components/shared/form/FormField"; +import { Button } from "@/components/ui/button"; +import { Dialog, DialogContent, DialogFooter, DialogHeader, DialogTitle } from "@/components/ui/dialog"; +import { Input } from "@/components/ui/input"; +import { Select, SelectContent, SelectItem, SelectTrigger, SelectValue } from "@/components/ui/select"; +import { Textarea } from "@/components/ui/textarea"; +import VectorStoreSelector from "@/components/vector_store_management/VectorStoreSelector"; +import { useZodForm } from "@/lib/forms/useZodForm"; +import { fetchClient } from "@/lib/http/api"; + +import { BUDGET_DURATION_OPTIONS, NO_RESET } from "../org-settings/OrgSettingsForm"; +import { orgSettingsSchema } from "../org-settings/schema"; +import { buildOrgCreateBody, emptyOrgFormValues, type OrgCreateBody } from "./mapper"; + +const defaultCreateOrganization = async (body: OrgCreateBody): Promise => { + const { data } = await fetchClient.POST("/organization/new", { body }); + return data; +}; + +interface OrgCreateDialogProps { + open: boolean; + onOpenChange: (open: boolean) => void; + accessToken: string; + createOrganization?: (body: OrgCreateBody) => Promise; +} + +export const OrgCreateDialog = ({ + open, + onOpenChange, + accessToken, + createOrganization = defaultCreateOrganization, +}: OrgCreateDialogProps) => { + const queryClient = useQueryClient(); + const form = useZodForm(orgSettingsSchema, { defaultValues: emptyOrgFormValues }); + + const closeAndReset = () => { + form.reset(emptyOrgFormValues); + onOpenChange(false); + }; + + const mutation = useMutation({ + mutationFn: (body: OrgCreateBody) => createOrganization(body), + onSuccess: () => { + NotificationsManager.success("Organization created successfully"); + queryClient.invalidateQueries({ queryKey: organizationKeys.all }); + closeAndReset(); + }, + onError: (error: unknown) => + NotificationsManager.fromBackend(error instanceof Error ? error.message : "Failed to create organization"), + }); + + const handleOpenChange = (nextOpen: boolean) => { + if (!nextOpen && mutation.isPending) return; + if (!nextOpen) { + form.reset(emptyOrgFormValues); + } + onOpenChange(nextOpen); + }; + + const onSubmit = form.handleSubmit((values) => { + if (mutation.isPending) return; + mutation.mutate(buildOrgCreateBody(values)); + }); + + return ( + + + + Create Organization + + +
+ + + {({ ref, ...field }) => } + + + + {(field) => ( + + )} + + + + {({ ref, ...field }) => } + + + + {({ id, value, onChange, "aria-invalid": ariaInvalid, "aria-describedby": ariaDescribedBy }) => ( + + )} + + + + {({ ref, ...field }) => } + + + + {({ ref, ...field }) => } + + + + {(field) => ( + + )} + + + + {(field) => ( + + )} + + + + {({ ref, ...field }) =>