From 5705aaebbc3247cb1666dec0288ab7b5507b2c79 Mon Sep 17 00:00:00 2001 From: Rens Date: Thu, 18 Dec 2025 15:11:58 +0200 Subject: [PATCH 001/158] Fix Gemini 3 imgs in tool response --- .../prompt_templates/factory.py | 7 +-- ...llm_core_utils_prompt_templates_factory.py | 51 +++++++++++++++++++ 2 files changed, 55 insertions(+), 3 deletions(-) diff --git a/litellm/litellm_core_utils/prompt_templates/factory.py b/litellm/litellm_core_utils/prompt_templates/factory.py index 652692c7b8d..9afea83ef7b 100644 --- a/litellm/litellm_core_utils/prompt_templates/factory.py +++ b/litellm/litellm_core_utils/prompt_templates/factory.py @@ -1496,9 +1496,10 @@ def convert_to_gemini_tool_call_result( content_type = content.get("type", "") if content_type == "text": content_str += content.get("text", "") - elif content_type == "input_image": - # Extract image for inline_data (for Computer Use screenshots) - image_url = content.get("image_url", "") + elif content_type in ("input_image", "image_url"): + # Extract image for inline_data (for Computer Use screenshots and tool results) + image_url_data = content.get("image_url", "") + image_url = image_url_data.get("url", "") if isinstance(image_url_data, dict) else image_url_data if image_url: # Convert image to base64 blob format for Gemini diff --git a/tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py b/tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py index 41ac893b4d7..c8fe6efeaa1 100644 --- a/tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py +++ b/tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py @@ -497,6 +497,57 @@ def test_convert_gemini_messages(): ) +def test_convert_gemini_tool_call_result_with_image_url(): + """ + Test that image_url content type in tool results is handled correctly for Gemini. + Fixes: https://github.com/BerriAI/litellm/issues/18187 + """ + from litellm.litellm_core_utils.prompt_templates.factory import ( + convert_to_gemini_tool_call_result, + ) + from litellm.types.llms.openai import ChatCompletionToolMessage + + # Test with string image_url format + message_str_format = ChatCompletionToolMessage( + role="tool", + tool_call_id="call_123", + content=[{"type": "image_url", "image_url": "data:image/jpeg;base64,/9j/4AAQ"}], + ) + last_message_with_tool_calls = { + "role": "assistant", + "content": "", + "tool_calls": [ + { + "id": "call_123", + "type": "function", + "index": 0, + "function": {"name": "get_image", "arguments": "{}"}, + } + ], + } + + result = convert_to_gemini_tool_call_result( + message=message_str_format, + last_message_with_tool_calls=last_message_with_tool_calls, + ) + # Should have inline_data for the image + assert isinstance(result, list) and any("inline_data" in p for p in result) + + # Test with dict image_url format (OpenAI standard) + message_dict_format = ChatCompletionToolMessage( + role="tool", + tool_call_id="call_456", + content=[{"type": "image_url", "image_url": {"url": "data:image/jpeg;base64,/9j/4AAQ"}}], + ) + last_message_with_tool_calls["tool_calls"][0]["id"] = "call_456" + + result2 = convert_to_gemini_tool_call_result( + message=message_dict_format, + last_message_with_tool_calls=last_message_with_tool_calls, + ) + assert isinstance(result2, list) and any("inline_data" in p for p in result2) + + def test_bedrock_tools_unpack_defs(): """ Test that the unpack_defs method handles nested $ref inside anyOf items correctly From 9002f75277228c0723d0d503f71f1f1a5de4638f Mon Sep 17 00:00:00 2001 From: Emerson Gomes Date: Sat, 20 Dec 2025 11:01:41 -0600 Subject: [PATCH 002/158] Require auth for MCP connection test --- .../mcp_server/rest_endpoints.py | 8 +-- .../mcp_server/test_rest_endpoints.py | 55 +++++++++++++++++++ 2 files changed, 58 insertions(+), 5 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/rest_endpoints.py b/litellm/proxy/_experimental/mcp_server/rest_endpoints.py index 032331ece02..92c390f0a64 100644 --- a/litellm/proxy/_experimental/mcp_server/rest_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/rest_endpoints.py @@ -1,5 +1,4 @@ import importlib -import traceback from typing import Dict, List, Optional, Union from fastapi import APIRouter, Depends, Query, Request @@ -329,16 +328,15 @@ if MCP_AVAILABLE: except Exception as e: verbose_logger.error(f"Error in MCP operation: {e}", exc_info=True) - stack_trace = traceback.format_exc() return { "status": "error", - "message": f"An internal error has occurred: {str(e)}", - "stack_trace": stack_trace, + "message": "An internal error has occurred while testing the MCP server.", } - @router.post("/test/connection") + @router.post("/test/connection", dependencies=[Depends(user_api_key_auth)]) async def test_connection( request: NewMCPServerRequest, + user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), ): """ Test if we can connect to the provided MCP server before adding it diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py index a0c09663a88..85ec807b1ff 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py @@ -7,6 +7,7 @@ from litellm.proxy._experimental.mcp_server import rest_endpoints from litellm.proxy._experimental.mcp_server.auth import ( user_api_key_auth_mcp as auth_mcp, ) +from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy._types import NewMCPServerRequest, UserAPIKeyAuth from litellm.types.mcp import MCPAuth @@ -31,6 +32,60 @@ def _build_request(headers: Optional[Dict[str, str]] = None) -> Request: return Request(scope, receive=receive) +def _get_route(path: str, method: str): + for route in rest_endpoints.router.routes: + if getattr(route, "path", None) == path and method in getattr( + route, "methods", set() + ): + return route + raise AssertionError(f"Route {method} {path} not found") + + +def _route_has_dependency(route, dependency) -> bool: + if any( + getattr(dep, "dependency", None) == dependency + for dep in getattr(route, "dependencies", []) + ): + return True + dependant = getattr(route, "dependant", None) + if dependant is None: + return False + return any(getattr(dep, "call", None) == dependency for dep in dependant.dependencies) + + +@pytest.mark.asyncio +async def test_execute_with_mcp_client_redacts_stack_trace(monkeypatch): + def fake_create_client(*args, **kwargs): + return object() + + monkeypatch.setattr( + rest_endpoints.global_mcp_server_manager, + "_create_mcp_client", + fake_create_client, + ) + + async def failing_operation(client): + raise RuntimeError("boom") + + payload = NewMCPServerRequest( + server_name="example", + url="https://example.com", + auth_type=MCPAuth.none, + ) + + result = await rest_endpoints._execute_with_mcp_client( + payload, failing_operation + ) + + assert result["status"] == "error" + assert "stack_trace" not in result + + +def test_test_connection_requires_auth_dependency(): + route = _get_route("/mcp-rest/test/connection", "POST") + assert _route_has_dependency(route, user_api_key_auth) + + @pytest.mark.asyncio async def test_test_tools_list_forwards_mcp_auth_header(monkeypatch): """Ensure credential-based auth forwards the auth_value to the MCP client.""" From acce6b9c83f143116038b49be710398802cd540e Mon Sep 17 00:00:00 2001 From: Emerson Gomes Date: Sat, 20 Dec 2025 15:56:46 -0600 Subject: [PATCH 003/158] Apply suggestion from @Copilot Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> --- .../proxy/_experimental/mcp_server/test_rest_endpoints.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py index 85ec807b1ff..31ab4afb631 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py @@ -82,7 +82,7 @@ async def test_execute_with_mcp_client_redacts_stack_trace(monkeypatch): def test_test_connection_requires_auth_dependency(): - route = _get_route("/mcp-rest/test/connection", "POST") + route = _get_route("/test/connection", "POST") assert _route_has_dependency(route, user_api_key_auth) From 3a141e642a292fce73b24e2e58178cc7237afae2 Mon Sep 17 00:00:00 2001 From: Yuta Saito Date: Fri, 26 Dec 2025 15:37:41 +0900 Subject: [PATCH 004/158] fix: get_all_mcp_servers_with_health_and_teams --- .../mcp_server/mcp_server_manager.py | 146 ++++++++---------- 1 file changed, 63 insertions(+), 83 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index 2260649e8b2..a5790803677 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -2244,100 +2244,80 @@ class MCPServerManager: Returns: List of MCP server objects with health and team data """ - from litellm.proxy._experimental.mcp_server.db import ( - get_all_mcp_servers, - get_mcp_servers, - ) - from litellm.proxy.management_endpoints.common_utils import _user_has_admin_view - from litellm.proxy.proxy_server import prisma_client # Get allowed server IDs allowed_server_ids = await self.get_allowed_mcp_servers(user_api_key_auth) - # Get servers from database - list_mcp_servers: List[LiteLLM_MCPServerTable] = [] - if prisma_client is not None: - list_mcp_servers = await get_mcp_servers(prisma_client, allowed_server_ids) + async def _check_server_health(server_id: str) -> Optional[LiteLLM_MCPServerTable]: + """Helper function to check health of a single server""" + server = self.get_mcp_server_by_id(server_id) + if server is None: + verbose_logger.warning(f"MCP Server {server_id} not found") + return None - # If admin, also get all servers from database - if user_api_key_auth and _user_has_admin_view(user_api_key_auth): - all_mcp_servers = await get_all_mcp_servers(prisma_client) - for server in all_mcp_servers: - if server.server_id not in allowed_server_ids: - list_mcp_servers.append(server) + status = "unknown" + health_check_error = None - # Add config.yaml servers - for _server_id, _server_config in self.config_mcp_servers.items(): - if _server_id in allowed_server_ids: - list_mcp_servers.append( - LiteLLM_MCPServerTable( - **{ - **_server_config.model_dump(), - "created_at": datetime.datetime.now(), - "updated_at": datetime.datetime.now(), - "description": ( - _server_config.mcp_info.get("description") - if _server_config.mcp_info - else None - ), - "allowed_tools": _server_config.allowed_tools or [], - "mcp_info": _server_config.mcp_info, - "mcp_access_groups": _server_config.access_groups or [], - "extra_headers": _server_config.extra_headers or [], - "command": getattr(_server_config, "command", None), - "args": getattr(_server_config, "args", None) or [], - "env": getattr(_server_config, "env", None) or {}, - } - ) + # Check if we should skip health check based on auth configuration + should_skip_health_check = False + + # Skip if auth_type is oauth2 + if server.auth_type == MCPAuth.oauth2: + should_skip_health_check = True + # Skip if auth_type is not none and authentication_token is missing + elif server.auth_type and server.auth_type != MCPAuth.none and not server.authentication_token: + should_skip_health_check = True + + if not should_skip_health_check: + extra_headers = {} + if server.static_headers: + extra_headers.update(server.static_headers) + + client = self._create_mcp_client( + server=server, + mcp_auth_header=None, + extra_headers=extra_headers, + stdio_env=None, ) - # Get team information for non-admin users - server_to_teams_map: Dict[str, List[Dict[str, str]]] = {} - if ( - user_api_key_auth - and not _user_has_admin_view(user_api_key_auth) - and prisma_client is not None - ): - teams = await prisma_client.db.litellm_teamtable.find_many( - include={"object_permission": True} + try: + async def _noop(session): + return "ok" + + await client.run_with_session(_noop) + status = "healthy" + except Exception as e: + health_check_error = str(e) + status = "unhealthy" + + return LiteLLM_MCPServerTable( + **{ + **server.model_dump(), + "created_at": datetime.datetime.now(), + "updated_at": datetime.datetime.now(), + "description": ( + server.mcp_info.get("description") + if server.mcp_info + else None + ), + "allowed_tools": server.allowed_tools or [], + "mcp_info": server.mcp_info, + "mcp_access_groups": server.access_groups or [], + "extra_headers": server.extra_headers or [], + "command": getattr(server, "command", None), + "args": getattr(server, "args", None) or [], + "env": getattr(server, "env", None) or {}, + "status": status, + "health_check_error": health_check_error, + } ) - user_teams = [] - for team in teams: - if team.members_with_roles: - for member in team.members_with_roles: - if ( - "user_id" in member - and member["user_id"] is not None - and member["user_id"] == user_api_key_auth.user_id - ): - user_teams.append(team) + # Run health checks concurrently + tasks = [_check_server_health(server_id) for server_id in allowed_server_ids] + results = await asyncio.gather(*tasks) - # Create a mapping of server_id to teams that have access to it - for team in user_teams: - if team.object_permission and team.object_permission.mcp_servers: - for server_id in team.object_permission.mcp_servers: - if server_id not in server_to_teams_map: - server_to_teams_map[server_id] = [] - server_to_teams_map[server_id].append( - { - "team_id": team.team_id, - "team_alias": team.team_alias, - "organization_id": team.organization_id, - } - ) - - ## mark invalid servers w/ reason for being invalid - valid_server_ids = self.get_all_mcp_server_ids() - for server in list_mcp_servers: - if server.server_id not in valid_server_ids: - server.status = "unhealthy" - ## try adding server to registry to get error - try: - await self.add_update_server(server) - except Exception as e: - server.health_check_error = str(e) - server.health_check_error = "Server is not in in memory registry yet. This could be a temporary sync issue." + # Filter out None results (servers that were not found) + list_mcp_servers = [server for server in results if server is not None] return list_mcp_servers From c25ceb5596d7b95b201014d3073fa792cf8256f4 Mon Sep 17 00:00:00 2001 From: Yuta Saito Date: Fri, 26 Dec 2025 16:55:51 +0900 Subject: [PATCH 005/158] refactor: MCP health check --- .../mcp_server/mcp_server_manager.py | 210 ++++-------- .../mcp_management_endpoints.py | 122 +------ .../mcp_server/test_mcp_server_manager.py | 270 ++++++++------- .../test_mcp_management_endpoints.py | 317 +----------------- 4 files changed, 226 insertions(+), 693 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index a5790803677..1171eef75d7 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -2127,7 +2127,7 @@ class MCPServerManager: async def health_check_server( self, server_id: str, mcp_auth_header: Optional[str] = None - ) -> Dict[str, Any]: + ) -> LiteLLM_MCPServerTable: """ Perform a health check on a specific MCP server. @@ -2138,96 +2138,82 @@ class MCPServerManager: Returns: Dict containing health check results """ - import time from datetime import datetime server = self.get_mcp_server_by_id(server_id) if not server: - return { - "server_id": server_id, - "server_name": None, - "status": "unknown", - "error": "Server not found", - "last_health_check": datetime.now().isoformat(), - "response_time_ms": None, - } - - start_time = time.time() - try: - # Try to get tools from the server as a health check - tools = await self._get_tools_from_server(server, mcp_auth_header) - response_time = (time.time() - start_time) * 1000 - - return { - "server_id": server_id, - "server_name": server.name, - "status": "healthy", - "tools_count": len(tools), - "last_health_check": datetime.now().isoformat(), - "response_time_ms": round(response_time, 2), - "error": None, - } - except Exception as e: - response_time = (time.time() - start_time) * 1000 - error_message = str(e) - - return { - "server_id": server_id, - "server_name": server.name, - "status": "unhealthy", - "last_health_check": datetime.now().isoformat(), - "response_time_ms": round(response_time, 2), - "error": error_message, - } - - async def health_check_all_servers( - self, mcp_auth_header: Optional[str] = None - ) -> Dict[str, Any]: - """ - Perform health checks on all MCP servers. - - Args: - mcp_auth_header: Optional authentication header for the MCP servers - - Returns: - Dict containing health check results for all servers - """ - all_servers = self.get_registry() - results = {} - - for server_id, server in all_servers.items(): - results[server_id] = await self.health_check_server( - server_id, mcp_auth_header + verbose_logger.warning(f"MCP Server {server_id} not found") + return LiteLLM_MCPServerTable( + server_id=server_id, + server_name=None, + transport=MCPTransport.http, # Default transport for not found servers + status="unknown", + health_check_error="Server not found", + last_health_check=datetime.now(), ) - return results + status = "unknown" + health_check_error = None - async def health_check_allowed_servers( - self, - user_api_key_auth: Optional[UserAPIKeyAuth] = None, - mcp_auth_header: Optional[str] = None, - ) -> Dict[str, Any]: - """ - Perform health checks on all MCP servers that the user has access to. + # Check if we should skip health check based on auth configuration + should_skip_health_check = False - Args: - user_api_key_auth: User authentication info for access control - mcp_auth_header: Optional authentication header for the MCP servers + # Skip if auth_type is oauth2 + if server.auth_type == MCPAuth.oauth2: + should_skip_health_check = True + # Skip if auth_type is not none and authentication_token is missing + elif server.auth_type and server.auth_type != MCPAuth.none and not server.authentication_token: + should_skip_health_check = True - Returns: - Dict containing health check results for accessible servers - """ - # Get allowed servers for the user - allowed_server_ids = await self.get_allowed_mcp_servers(user_api_key_auth) + if not should_skip_health_check: + extra_headers = {} + if server.static_headers: + extra_headers.update(server.static_headers) - # Perform health checks on allowed servers - results = {} - for server_id in allowed_server_ids: - results[server_id] = await self.health_check_server( - server_id, mcp_auth_header + client = self._create_mcp_client( + server=server, + mcp_auth_header=None, + extra_headers=extra_headers, + stdio_env=None, ) - return results + try: + async def _noop(session): + return "ok" + + await client.run_with_session(_noop) + status = "healthy" + except Exception as e: + health_check_error = str(e) + status = "unhealthy" + + return LiteLLM_MCPServerTable( + server_id=server.server_id, + server_name=server.server_name, + alias=server.alias, + description=( + server.mcp_info.get("description") + if server.mcp_info + else None + ), + url=server.url, + transport=server.transport, + auth_type=server.auth_type, + created_at=datetime.now(), + updated_at=datetime.now(), + teams=[], + mcp_access_groups=server.access_groups or [], + allowed_tools=server.allowed_tools or [], + extra_headers=server.extra_headers or [], + mcp_info=server.mcp_info, + static_headers=server.static_headers, + status=status, + last_health_check=datetime.now(), + health_check_error=health_check_error, + command=getattr(server, "command", None), + args=getattr(server, "args", None) or [], + env=getattr(server, "env", None) or {}, + ) async def get_all_mcp_servers_with_health_and_teams( self, @@ -2248,72 +2234,8 @@ class MCPServerManager: # Get allowed server IDs allowed_server_ids = await self.get_allowed_mcp_servers(user_api_key_auth) - async def _check_server_health(server_id: str) -> Optional[LiteLLM_MCPServerTable]: - """Helper function to check health of a single server""" - server = self.get_mcp_server_by_id(server_id) - if server is None: - verbose_logger.warning(f"MCP Server {server_id} not found") - return None - - status = "unknown" - health_check_error = None - - # Check if we should skip health check based on auth configuration - should_skip_health_check = False - - # Skip if auth_type is oauth2 - if server.auth_type == MCPAuth.oauth2: - should_skip_health_check = True - # Skip if auth_type is not none and authentication_token is missing - elif server.auth_type and server.auth_type != MCPAuth.none and not server.authentication_token: - should_skip_health_check = True - - if not should_skip_health_check: - extra_headers = {} - if server.static_headers: - extra_headers.update(server.static_headers) - - client = self._create_mcp_client( - server=server, - mcp_auth_header=None, - extra_headers=extra_headers, - stdio_env=None, - ) - - try: - async def _noop(session): - return "ok" - - await client.run_with_session(_noop) - status = "healthy" - except Exception as e: - health_check_error = str(e) - status = "unhealthy" - - return LiteLLM_MCPServerTable( - **{ - **server.model_dump(), - "created_at": datetime.datetime.now(), - "updated_at": datetime.datetime.now(), - "description": ( - server.mcp_info.get("description") - if server.mcp_info - else None - ), - "allowed_tools": server.allowed_tools or [], - "mcp_info": server.mcp_info, - "mcp_access_groups": server.access_groups or [], - "extra_headers": server.extra_headers or [], - "command": getattr(server, "command", None), - "args": getattr(server, "args", None) or [], - "env": getattr(server, "env", None) or {}, - "status": status, - "health_check_error": health_check_error, - } - ) - # Run health checks concurrently - tasks = [_check_server_health(server_id) for server_id in allowed_server_ids] + tasks = [self.health_check_server(server_id) for server_id in allowed_server_ids] results = await asyncio.gather(*tasks) # Filter out None results (servers that were not found) diff --git a/litellm/proxy/management_endpoints/mcp_management_endpoints.py b/litellm/proxy/management_endpoints/mcp_management_endpoints.py index f0eddcc8683..a77c7c87542 100644 --- a/litellm/proxy/management_endpoints/mcp_management_endpoints.py +++ b/litellm/proxy/management_endpoints/mcp_management_endpoints.py @@ -296,116 +296,6 @@ if MCP_AVAILABLE: access_groups_list = sorted(list(access_groups)) return {"access_groups": access_groups_list} - @router.get( - "/server/{server_id}/health", - description="Perform health check on a specific MCP server", - dependencies=[Depends(user_api_key_auth)], - ) - async def health_check_mcp_server( - server_id: str, - user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), - ): - """ - Perform a health check on the MCP server specified by the `server_id` - Parameters: - - server_id: str - Required. The unique identifier of the mcp server to health check. - ``` - curl --location 'http://localhost:4000/v1/mcp/server/{server_id}/health' \ - --header 'Authorization: Bearer your_api_key_here' - ``` - """ - # Check if server exists and user has access - prisma_client = get_prisma_client_or_throw( - "Database not connected. Connect a database to your proxy" - ) - - # check to see if server exists for all users - mcp_server = await get_mcp_server(prisma_client, server_id) - if mcp_server is None: - raise HTTPException( - status_code=status.HTTP_404_NOT_FOUND, - detail={"error": f"MCP Server with id {server_id} not found"}, - ) - - # Implement authz restriction from requested user - if not _user_has_admin_view(user_api_key_dict): - # Perform authz check to filter the mcp servers user has access to - mcp_server_records = await get_all_mcp_servers_for_user( - prisma_client, user_api_key_dict - ) - exists = does_mcp_server_exist(mcp_server_records, server_id) - - if not exists: - raise HTTPException( - status_code=status.HTTP_403_FORBIDDEN, - detail={ - "error": f"User does not have permission to access mcp server with id {server_id}. You can only access mcp servers that you have access to." - }, - ) - - # Perform health check using server manager - try: - health_result = await global_mcp_server_manager.health_check_server( - server_id - ) - return health_result - except Exception as e: - verbose_proxy_logger.exception( - f"Error performing health check on MCP server {server_id}: {str(e)}" - ) - raise HTTPException( - status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, - detail={"error": f"Error performing health check: {str(e)}"}, - ) - - @router.get( - "/server/health", - description="Perform health check on all accessible MCP servers", - dependencies=[Depends(user_api_key_auth)], - ) - async def health_check_all_mcp_servers( - user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), - ): - """ - Perform health checks on all MCP servers accessible to the user - ``` - curl --location 'http://localhost:4000/v1/mcp/server/health' \ - --header 'Authorization: Bearer your_api_key_here' - ``` - """ - # Use server manager to get health checks for allowed servers - try: - all_health_results = ( - await global_mcp_server_manager.health_check_allowed_servers( - user_api_key_auth=user_api_key_dict - ) - ) - - return { - "total_servers": len(all_health_results), - "healthy_count": len( - [r for r in all_health_results.values() if r["status"] == "healthy"] - ), - "unhealthy_count": len( - [ - r - for r in all_health_results.values() - if r["status"] == "unhealthy" - ] - ), - "unknown_count": len( - [r for r in all_health_results.values() if r["status"] == "unknown"] - ), - "servers": all_health_results, - } - except Exception as e: - verbose_proxy_logger.exception( - f"Error performing health checks on MCP servers: {str(e)}" - ) - raise HTTPException( - status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, - detail={"error": f"Error performing health checks: {str(e)}"}, - ) ## FastAPI Routes @router.get( @@ -484,15 +374,9 @@ if MCP_AVAILABLE: server_id ) # Update the server object with health check results - mcp_server.status = health_result.get("status", "unknown") - mcp_server.last_health_check = ( - datetime.fromisoformat( - health_result.get("last_health_check", datetime.now().isoformat()) - ) - if health_result.get("last_health_check") - else None - ) - mcp_server.health_check_error = health_result.get("error") + mcp_server.status = health_result.status if health_result.status else "unknown" + mcp_server.last_health_check = health_result.last_health_check + mcp_server.health_check_error = health_result.health_check_error except Exception as e: verbose_proxy_logger.debug( f"Error performing health check on server {server_id}: {e}" diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py index ff016a1a130..37a93441513 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py @@ -641,37 +641,31 @@ class TestMCPServerManager: manager = MCPServerManager() # Mock server - server = MagicMock() - server.server_id = "test-server" - server.name = "test-server" + server = MCPServer( + server_id="test-server", + name="test-server", + transport=MCPTransport.http, + auth_type=None, + authentication_token="test-token", + url="http://test-server.com", + ) manager.get_mcp_server_by_id = MagicMock(return_value=server) - # Mock successful _get_tools_from_server - async def mock_get_tools_from_server( - server, - mcp_auth_header=None, - raw_headers=None, - ): - tool1 = MagicMock() - tool1.name = "tool1" - tool2 = MagicMock() - tool2.name = "tool2" - return [tool1, tool2] - - manager._get_tools_from_server = mock_get_tools_from_server + # Mock successful client.run_with_session + mock_client = AsyncMock() + mock_client.run_with_session = AsyncMock(return_value="ok") + manager._create_mcp_client = MagicMock(return_value=mock_client) # Perform health check result = await manager.health_check_server("test-server") - # Verify results - assert result["server_id"] == "test-server" - assert result["status"] == "healthy" - assert result["tools_count"] == 2 - assert result["error"] is None - assert "last_health_check" in result - assert "response_time_ms" in result - assert result["response_time_ms"] >= 0 # Allow 0 for very fast mocks + # Verify results - result is now LiteLLM_MCPServerTable + assert isinstance(result, LiteLLM_MCPServerTable) + assert result.server_id == "test-server" + assert result.status == "healthy" + assert result.health_check_error is None + assert result.last_health_check is not None @pytest.mark.asyncio async def test_health_check_server_unhealthy(self): @@ -679,32 +673,33 @@ class TestMCPServerManager: manager = MCPServerManager() # Mock server - server = MagicMock() - server.server_id = "test-server" - server.name = "test-server" + server = MCPServer( + server_id="test-server", + name="test-server", + transport=MCPTransport.http, + auth_type=None, + authentication_token="test-token", + url="http://test-server.com", + ) manager.get_mcp_server_by_id = MagicMock(return_value=server) - # Mock failed _get_tools_from_server - async def mock_get_tools_from_server( - server, - mcp_auth_header=None, - raw_headers=None, - ): - raise Exception("Connection timeout") - - manager._get_tools_from_server = mock_get_tools_from_server + # Mock failed client.run_with_session + mock_client = AsyncMock() + mock_client.run_with_session = AsyncMock( + side_effect=Exception("Connection timeout") + ) + manager._create_mcp_client = MagicMock(return_value=mock_client) # Perform health check result = await manager.health_check_server("test-server") # Verify results - assert result["server_id"] == "test-server" - assert result["status"] == "unhealthy" - assert result["error"] == "Connection timeout" - assert "last_health_check" in result - assert "response_time_ms" in result - assert result["response_time_ms"] >= 0 # Allow 0 for very fast mocks + assert isinstance(result, LiteLLM_MCPServerTable) + assert result.server_id == "test-server" + assert result.status == "unhealthy" + assert result.health_check_error == "Connection timeout" + assert result.last_health_check is not None @pytest.mark.asyncio async def test_health_check_server_not_found(self): @@ -718,104 +713,121 @@ class TestMCPServerManager: result = await manager.health_check_server("non-existent-server") # Verify results - assert result["server_id"] == "non-existent-server" - assert result["status"] == "unknown" - assert result["error"] == "Server not found" - assert result["response_time_ms"] is None - assert "last_health_check" in result + assert isinstance(result, LiteLLM_MCPServerTable) + assert result.server_id == "non-existent-server" + assert result.server_name is None + assert result.status == "unknown" + assert result.health_check_error == "Server not found" + assert result.last_health_check is not None @pytest.mark.asyncio - async def test_health_check_all_servers(self): - """Test health check for all servers""" + async def test_health_check_server_oauth2_skips_check(self): + """Test that health check is skipped for OAuth2 servers and returns unknown status""" manager = MCPServerManager() - # Mock servers - server1 = MagicMock() - server1.server_id = "server1" - server1.name = "server1" - - server2 = MagicMock() - server2.server_id = "server2" - server2.name = "server2" - - # Mock registry - manager.registry = {"server1": server1, "server2": server2} - - # Mock get_mcp_server_by_id - def mock_get_server_by_id(server_id): - if server_id == "server1": - return server1 - elif server_id == "server2": - return server2 - return None - - manager.get_mcp_server_by_id = mock_get_server_by_id - - # Mock _get_tools_from_server with different results - async def mock_get_tools_from_server( - server, - mcp_auth_header=None, - raw_headers=None, - ): - if server.server_id == "server1": - tool = MagicMock() - tool.name = "tool1" - return [tool] - elif server.server_id == "server2": - raise Exception("Connection failed") - return [] - - manager._get_tools_from_server = mock_get_tools_from_server - - # Perform health check for all servers - result = await manager.health_check_all_servers() - - # Verify results - assert len(result) == 2 - assert "server1" in result - assert "server2" in result - - # Check server1 (healthy) - assert result["server1"]["status"] == "healthy" - assert result["server1"]["tools_count"] == 1 - assert result["server1"]["error"] is None - - # Check server2 (unhealthy) - assert result["server2"]["status"] == "unhealthy" - assert result["server2"]["error"] == "Connection failed" - - @pytest.mark.asyncio - async def test_health_check_server_with_auth_header(self): - """Test health check with authentication header""" - manager = MCPServerManager() - - # Mock server - server = MagicMock() - server.server_id = "test-server" - server.name = "test-server" + # Mock OAuth2 server + server = MCPServer( + server_id="oauth2-server", + name="oauth2-server", + transport=MCPTransport.http, + auth_type=MCPAuth.oauth2, + url="http://oauth2-server.com", + ) manager.get_mcp_server_by_id = MagicMock(return_value=server) - # Mock _get_tools_from_server to verify auth header is passed - async def mock_get_tools_from_server( - server, - mcp_auth_header=None, - raw_headers=None, - ): - assert mcp_auth_header == "test-token" - tool = MagicMock() - tool.name = "tool1" - return [tool] + # _create_mcp_client should not be called for OAuth2 servers + manager._create_mcp_client = MagicMock() - manager._get_tools_from_server = mock_get_tools_from_server + # Perform health check + result = await manager.health_check_server("oauth2-server") - # Perform health check with auth header - result = await manager.health_check_server("test-server", "test-token") + # Verify that client was not created (health check was skipped) + manager._create_mcp_client.assert_not_called() # Verify results - assert result["server_id"] == "test-server" - assert result["status"] == "healthy" - assert result["tools_count"] == 1 + assert isinstance(result, LiteLLM_MCPServerTable) + assert result.server_id == "oauth2-server" + assert result.status == "unknown" + assert result.health_check_error is None + assert result.last_health_check is not None + + @pytest.mark.asyncio + async def test_health_check_server_no_token_skips_check(self): + """Test that health check is skipped when auth_type is set but authentication_token is missing""" + manager = MCPServerManager() + + # Mock server with auth_type but no authentication_token + server = MCPServer( + server_id="no-token-server", + name="no-token-server", + transport=MCPTransport.http, + auth_type=MCPAuth.bearer_token, + authentication_token=None, # No token + url="http://no-token-server.com", + ) + + manager.get_mcp_server_by_id = MagicMock(return_value=server) + + # _create_mcp_client should not be called + manager._create_mcp_client = MagicMock() + + # Perform health check + result = await manager.health_check_server("no-token-server") + + # Verify that client was not created (health check was skipped) + manager._create_mcp_client.assert_not_called() + + # Verify results + assert isinstance(result, LiteLLM_MCPServerTable) + assert result.server_id == "no-token-server" + assert result.status == "unknown" + assert result.health_check_error is None + assert result.last_health_check is not None + + @pytest.mark.asyncio + async def test_health_check_server_with_static_headers(self): + """Test health check with static headers configured""" + manager = MCPServerManager() + + # Mock server with static_headers + server = MCPServer( + server_id="test-server", + name="test-server", + transport=MCPTransport.http, + auth_type=None, + authentication_token="test-token", + url="http://test-server.com", + static_headers={"X-Custom-Header": "custom-value"}, + ) + + manager.get_mcp_server_by_id = MagicMock(return_value=server) + + # Mock successful client + mock_client = AsyncMock() + mock_client.run_with_session = AsyncMock(return_value="ok") + + # Capture the extra_headers passed to _create_mcp_client + captured_extra_headers = None + + def capture_create_mcp_client(server, mcp_auth_header, extra_headers, stdio_env): + nonlocal captured_extra_headers + captured_extra_headers = extra_headers + return mock_client + + manager._create_mcp_client = MagicMock(side_effect=capture_create_mcp_client) + + # Perform health check + result = await manager.health_check_server("test-server") + + # Verify static headers were passed + assert captured_extra_headers == {"X-Custom-Header": "custom-value"} + + # Verify results + assert isinstance(result, LiteLLM_MCPServerTable) + assert result.server_id == "test-server" + assert result.status == "healthy" + assert result.health_check_error is None @pytest.mark.asyncio async def test_pre_call_tool_check_allowed_tools_list_allows_tool(self): diff --git a/tests/test_litellm/proxy/management_endpoints/test_mcp_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_mcp_management_endpoints.py index 61342e8025b..8d5b3dcedc9 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_mcp_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_mcp_management_endpoints.py @@ -486,11 +486,14 @@ class TestListMCPServers: mock_server.credentials = {"auth_value": "top-secret"} mock_prisma_client = MagicMock() - mock_health_result = { - "status": "healthy", - "last_health_check": datetime.now().isoformat(), - "error": None, - } + + # Mock health check result as LiteLLM_MCPServerTable + mock_health_result = generate_mock_mcp_server_db_record( + server_id="server-1", alias="Server 1" + ) + mock_health_result.status = "healthy" + mock_health_result.last_health_check = datetime.now() + mock_health_result.health_check_error = None mock_user_auth = generate_mock_user_api_key_auth( user_role=LitellmUserRoles.PROXY_ADMIN @@ -531,11 +534,14 @@ class TestListMCPServers: delattr(mock_server, "credentials") mock_prisma_client = MagicMock() - mock_health_result = { - "status": "healthy", - "last_health_check": datetime.now().isoformat(), - "error": None, - } + + # Mock health check result as LiteLLM_MCPServerTable + mock_health_result = generate_mock_mcp_server_db_record( + server_id="server-2", alias="Server 2" + ) + mock_health_result.status = "healthy" + mock_health_result.last_health_check = datetime.now() + mock_health_result.health_check_error = None mock_user_auth = generate_mock_user_api_key_auth( user_role=LitellmUserRoles.PROXY_ADMIN @@ -568,296 +574,6 @@ class TestListMCPServers: assert result.status == "healthy" -class TestMCPHealthCheckEndpoints: - """Test MCP health check endpoints""" - - @pytest.mark.asyncio - async def test_health_check_mcp_server_success(self): - """Test successful health check for a specific MCP server""" - # Mock server - mock_server = generate_mock_mcp_server_db_record( - server_id="test-server", alias="Test Server" - ) - - # Mock dependencies - mock_prisma_client = MagicMock() - - # Mock global MCP server manager - mock_manager = MagicMock() - mock_manager.health_check_server = AsyncMock( - return_value={ - "server_id": "test-server", - "server_name": "Test Server", - "status": "healthy", - "tools_count": 3, - "last_health_check": "2024-01-01T12:00:00", - "response_time_ms": 150.5, - "error": None, - } - ) - - mock_user_auth = generate_mock_user_api_key_auth( - user_role=LitellmUserRoles.PROXY_ADMIN - ) - - with patch( - "litellm.proxy.management_endpoints.mcp_management_endpoints.get_prisma_client_or_throw", - return_value=mock_prisma_client, - ), patch( - "litellm.proxy.management_endpoints.mcp_management_endpoints._user_has_admin_view", - return_value=True, - ), patch( - "litellm.proxy.management_endpoints.mcp_management_endpoints.global_mcp_server_manager", - mock_manager, - ), patch( - "litellm.proxy.management_endpoints.mcp_management_endpoints.get_mcp_server", - AsyncMock(return_value=mock_server), - ): - # Import and call the function - from litellm.proxy.management_endpoints.mcp_management_endpoints import ( - health_check_mcp_server, - ) - - result = await health_check_mcp_server( - server_id="test-server", user_api_key_dict=mock_user_auth - ) - - # Verify results - assert result["server_id"] == "test-server" - assert result["server_name"] == "Test Server" - assert result["status"] == "healthy" - assert result["tools_count"] == 3 - assert result["response_time_ms"] == 150.5 - assert result["error"] is None - - @pytest.mark.asyncio - async def test_health_check_mcp_server_not_found(self): - """Test health check for a server that doesn't exist""" - # Mock dependencies - mock_prisma_client = MagicMock() - - mock_user_auth = generate_mock_user_api_key_auth( - user_role=LitellmUserRoles.PROXY_ADMIN - ) - - with patch( - "litellm.proxy.management_endpoints.mcp_management_endpoints.get_prisma_client_or_throw", - return_value=mock_prisma_client, - ), patch( - "litellm.proxy.management_endpoints.mcp_management_endpoints.get_mcp_server", - AsyncMock(return_value=None), - ): - # Import and call the function - from litellm.proxy.management_endpoints.mcp_management_endpoints import ( - health_check_mcp_server, - ) - - # Should raise HTTPException - with pytest.raises(Exception) as exc_info: - await health_check_mcp_server( - server_id="non-existent-server", user_api_key_dict=mock_user_auth - ) - - assert "not found" in str(exc_info.value) - - @pytest.mark.asyncio - async def test_health_check_mcp_server_unauthorized(self): - """Test health check for a server user doesn't have access to""" - # Mock server - mock_server = generate_mock_mcp_server_db_record( - server_id="test-server", alias="Test Server" - ) - - # Mock dependencies - mock_prisma_client = MagicMock() - - mock_user_auth = generate_mock_user_api_key_auth( - user_role=LitellmUserRoles.INTERNAL_USER # Non-admin user - ) - - # Mock user doesn't have access to this server - mock_user_servers = [] - - with patch( - "litellm.proxy.management_endpoints.mcp_management_endpoints.get_prisma_client_or_throw", - return_value=mock_prisma_client, - ), patch( - "litellm.proxy.management_endpoints.mcp_management_endpoints._user_has_admin_view", - return_value=False, - ), patch( - "litellm.proxy.management_endpoints.mcp_management_endpoints.get_all_mcp_servers_for_user", - return_value=mock_user_servers, - ), patch( - "litellm.proxy.management_endpoints.mcp_management_endpoints.get_mcp_server", - AsyncMock(return_value=mock_server), - ): - # Import and call the function - from litellm.proxy.management_endpoints.mcp_management_endpoints import ( - health_check_mcp_server, - ) - - # Should raise HTTPException - with pytest.raises(Exception) as exc_info: - await health_check_mcp_server( - server_id="test-server", user_api_key_dict=mock_user_auth - ) - - assert "permission" in str(exc_info.value) - - @pytest.mark.asyncio - async def test_health_check_all_mcp_servers(self): - """Test health check for all accessible MCP servers""" - # Mock team records - team_records = [ - generate_mock_team_record( - team_id="team1", - team_alias="Team 1", - organization_id="org1", - mcp_servers=["server1", "server2"], - ) - ] - - # Mock DB servers - db_servers = [ - generate_mock_mcp_server_db_record(server_id="server1"), - generate_mock_mcp_server_db_record(server_id="server2"), - ] - - # Mock dependencies - mock_prisma_client = MagicMock() - mock_prisma_client = setup_mock_prisma_client( - mock_prisma_client=mock_prisma_client, - team_records=team_records, - mcp_servers=db_servers, - ) - - # Mock global MCP server manager - mock_manager = MagicMock() - mock_manager.health_check_allowed_servers = AsyncMock( - return_value={ - "server1": { - "server_id": "server1", - "server_name": "Test DB Server", - "status": "healthy", - "tools_count": 2, - "last_health_check": "2024-01-01T12:00:00", - "response_time_ms": 100.0, - "error": None, - }, - "server2": { - "server_id": "server2", - "server_name": "Test DB Server", - "status": "unhealthy", - "last_health_check": "2024-01-01T12:00:00", - "response_time_ms": 5000.0, - "error": "Connection timeout", - }, - } - ) - mock_manager.get_allowed_mcp_servers = AsyncMock( - return_value=["server1", "server2"] - ) - - mock_user_auth = generate_mock_user_api_key_auth( - user_role=LitellmUserRoles.INTERNAL_USER - ) - - with patch( - "litellm.proxy.management_endpoints.mcp_management_endpoints.get_prisma_client_or_throw", - return_value=mock_prisma_client, - ), patch( - "litellm.proxy.management_endpoints.mcp_management_endpoints._user_has_admin_view", - return_value=False, - ), patch( - "litellm.proxy.management_endpoints.mcp_management_endpoints.global_mcp_server_manager", - mock_manager, - ): - # Import and call the function - from litellm.proxy.management_endpoints.mcp_management_endpoints import ( - health_check_all_mcp_servers, - ) - - result = await health_check_all_mcp_servers( - user_api_key_dict=mock_user_auth - ) - - # Verify results - assert result["total_servers"] == 2 - assert result["healthy_count"] == 1 - assert result["unhealthy_count"] == 1 - assert result["unknown_count"] == 0 - assert "server1" in result["servers"] - assert "server2" in result["servers"] - - # Check individual server results - assert result["servers"]["server1"]["status"] == "healthy" - assert result["servers"]["server1"]["tools_count"] == 2 - assert result["servers"]["server1"]["server_name"] == "Test DB Server" - assert result["servers"]["server2"]["status"] == "unhealthy" - assert result["servers"]["server2"]["error"] == "Connection timeout" - assert result["servers"]["server2"]["server_name"] == "Test DB Server" - - @pytest.mark.asyncio - async def test_fetch_all_mcp_servers_with_health_status(self): - """Test that fetch_all_mcp_servers includes health check status""" - # Mock server with health status - mock_server = generate_mock_mcp_server_db_record( - server_id="test-server", alias="Test Server" - ) - # Add health status to the mock server - mock_server.status = "healthy" - mock_server.last_health_check = datetime.now() - mock_server.health_check_error = None - - # Mock dependencies - mock_prisma_client = MagicMock() - mock_prisma_client = setup_mock_prisma_client( - mock_prisma_client=mock_prisma_client, - team_records=[], - mcp_servers=[], # Don't add servers here since we're mocking get_all_mcp_servers - ) - - # Mock global MCP server manager - mock_manager = MagicMock() - mock_manager.config_mcp_servers = {} - mock_manager.get_allowed_mcp_servers = AsyncMock(return_value=[]) - mock_manager.get_all_mcp_servers_with_health_and_teams = AsyncMock( - return_value=[mock_server] - ) - - mock_server.credentials = {"auth_value": "secret"} - - mock_user_auth = generate_mock_user_api_key_auth( - user_role=LitellmUserRoles.PROXY_ADMIN - ) - - with patch( - "litellm.proxy.management_endpoints.mcp_management_endpoints.get_prisma_client_or_throw", - return_value=mock_prisma_client, - ), patch( - "litellm.proxy.management_endpoints.mcp_management_endpoints._user_has_admin_view", - return_value=True, - ), patch( - "litellm.proxy.management_endpoints.mcp_management_endpoints.global_mcp_server_manager", - mock_manager, - ): - # Import and call the function - from litellm.proxy.management_endpoints.mcp_management_endpoints import ( - fetch_all_mcp_servers, - ) - - result = await fetch_all_mcp_servers(user_api_key_dict=mock_user_auth) - - # Verify health check status is included - assert len(result) == 1 - server = result[0] - assert server.server_id == "test-server" - assert server.status == "healthy" - assert server.last_health_check is not None - assert server.health_check_error is None - assert server.credentials is None - - class TestTemporaryMCPSessionEndpoints: def test_inherit_credentials_from_existing_server(self): payload = NewMCPServerRequest( @@ -1170,7 +886,6 @@ class TestTemporaryMCPSessionEndpoints: fallback_client_id="server-1", ) - class TestUpdateMCPServer: """Test suite for update MCP server functionality""" From aa0ff9131ea0be71cabbd0a3af3f14e8ee29bb6c Mon Sep 17 00:00:00 2001 From: Yuta Saito Date: Fri, 26 Dec 2025 16:59:31 +0900 Subject: [PATCH 006/158] fix: format --- .../mcp_server/mcp_server_manager.py | 19 ++++++++++++------- .../mcp_management_endpoints.py | 5 +++-- 2 files changed, 15 insertions(+), 9 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index 1171eef75d7..58e9a345dc7 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -11,7 +11,7 @@ import datetime import hashlib import json import re -from typing import Any, Dict, List, Optional, Set, Tuple, Union, cast +from typing import Any, Dict, List, Literal, Optional, Set, Tuple, Union, cast from urllib.parse import urlparse from fastapi import HTTPException @@ -2152,7 +2152,7 @@ class MCPServerManager: last_health_check=datetime.now(), ) - status = "unknown" + status: Literal["healthy", "unhealthy", "unknown"] = "unknown" health_check_error = None # Check if we should skip health check based on auth configuration @@ -2162,7 +2162,11 @@ class MCPServerManager: if server.auth_type == MCPAuth.oauth2: should_skip_health_check = True # Skip if auth_type is not none and authentication_token is missing - elif server.auth_type and server.auth_type != MCPAuth.none and not server.authentication_token: + elif ( + server.auth_type + and server.auth_type != MCPAuth.none + and not server.authentication_token + ): should_skip_health_check = True if not should_skip_health_check: @@ -2178,6 +2182,7 @@ class MCPServerManager: ) try: + async def _noop(session): return "ok" @@ -2192,9 +2197,7 @@ class MCPServerManager: server_name=server.server_name, alias=server.alias, description=( - server.mcp_info.get("description") - if server.mcp_info - else None + server.mcp_info.get("description") if server.mcp_info else None ), url=server.url, transport=server.transport, @@ -2235,7 +2238,9 @@ class MCPServerManager: allowed_server_ids = await self.get_allowed_mcp_servers(user_api_key_auth) # Run health checks concurrently - tasks = [self.health_check_server(server_id) for server_id in allowed_server_ids] + tasks = [ + self.health_check_server(server_id) for server_id in allowed_server_ids + ] results = await asyncio.gather(*tasks) # Filter out None results (servers that were not found) diff --git a/litellm/proxy/management_endpoints/mcp_management_endpoints.py b/litellm/proxy/management_endpoints/mcp_management_endpoints.py index a77c7c87542..59b9659dd88 100644 --- a/litellm/proxy/management_endpoints/mcp_management_endpoints.py +++ b/litellm/proxy/management_endpoints/mcp_management_endpoints.py @@ -296,7 +296,6 @@ if MCP_AVAILABLE: access_groups_list = sorted(list(access_groups)) return {"access_groups": access_groups_list} - ## FastAPI Routes @router.get( "/server", @@ -374,7 +373,9 @@ if MCP_AVAILABLE: server_id ) # Update the server object with health check results - mcp_server.status = health_result.status if health_result.status else "unknown" + mcp_server.status = ( + health_result.status if health_result.status else "unknown" + ) mcp_server.last_health_check = health_result.last_health_check mcp_server.health_check_error = health_result.health_check_error except Exception as e: From 6888d34ea5372bd17086e92ed30554baf3d7cdbd Mon Sep 17 00:00:00 2001 From: Yuta Saito Date: Sat, 27 Dec 2025 07:20:24 +0900 Subject: [PATCH 007/158] feat: Log actual executed event type in guardrail logging --- litellm/integrations/custom_guardrail.py | 46 +++- .../integrations/test_custom_guardrail.py | 238 +++++++++++++++++- 2 files changed, 273 insertions(+), 11 deletions(-) diff --git a/litellm/integrations/custom_guardrail.py b/litellm/integrations/custom_guardrail.py index fe0ce208ee6..7bc02bb3bb6 100644 --- a/litellm/integrations/custom_guardrail.py +++ b/litellm/integrations/custom_guardrail.py @@ -243,14 +243,14 @@ class CustomGuardrail(CustomLogger): def _is_valid_response_type(self, result: Any) -> bool: """ Check if result is a valid LLMResponseTypes instance. - + Safely handles TypedDict types which don't support isinstance checks. For non-LiteLLM responses (like passthrough httpx.Response), returns True to allow them through. """ if result is None: return False - + try: # Try isinstance check on valid types that support it response_types = get_args(LLMResponseTypes) @@ -506,6 +506,7 @@ class CustomGuardrail(CustomLogger): duration: Optional[float] = None, masked_entity_count: Optional[Dict[str, int]] = None, guardrail_provider: Optional[str] = None, + event_type: Optional[GuardrailEventHooks] = None, ) -> None: """ Builds `StandardLoggingGuardrailInformation` and adds it to the request metadata so it can be used for logging to DataDog, Langfuse, etc. @@ -514,14 +515,18 @@ class CustomGuardrail(CustomLogger): guardrail_json_response = str(guardrail_json_response) from litellm.types.utils import GuardrailMode + # Use event_type if provided, otherwise fall back to self.event_hook + if event_type is not None: + guardrail_mode = event_type + elif isinstance(self.event_hook, Mode): + guardrail_mode = GuardrailMode(**self.event_hook.model_dump()) + else: + guardrail_mode = self.event_hook + slg = StandardLoggingGuardrailInformation( guardrail_name=self.guardrail_name, guardrail_provider=guardrail_provider, - guardrail_mode=( - GuardrailMode(**self.event_hook.model_dump()) # type: ignore - if isinstance(self.event_hook, Mode) - else self.event_hook - ), + guardrail_mode=guardrail_mode, guardrail_response=guardrail_json_response, guardrail_status=guardrail_status, start_time=start_time, @@ -589,6 +594,7 @@ class CustomGuardrail(CustomLogger): start_time: Optional[float] = None, end_time: Optional[float] = None, duration: Optional[float] = None, + event_type: Optional[GuardrailEventHooks] = None, ): """ Add StandardLoggingGuardrailInformation to the request data @@ -605,6 +611,7 @@ class CustomGuardrail(CustomLogger): duration=duration, start_time=start_time, end_time=end_time, + event_type=event_type, ) return response @@ -615,6 +622,7 @@ class CustomGuardrail(CustomLogger): start_time: Optional[float] = None, end_time: Optional[float] = None, duration: Optional[float] = None, + event_type: Optional[GuardrailEventHooks] = None, ): """ Add StandardLoggingGuardrailInformation to the request data @@ -628,6 +636,7 @@ class CustomGuardrail(CustomLogger): duration=duration, start_time=start_time, end_time=end_time, + event_type=event_type, ) raise e @@ -712,16 +721,32 @@ def log_guardrail_information(func): Logs for: - pre_call - during_call - - TODO: log post_call. This is more involved since the logs are sent to DD, s3 before the guardrail is even run + - post_call """ import asyncio import functools + def _infer_event_type_from_function_name( + func_name: str, + ) -> Optional[GuardrailEventHooks]: + """Infer the actual event type from the function name""" + if func_name == "async_pre_call_hook": + return GuardrailEventHooks.pre_call + elif func_name == "async_moderation_hook": + return GuardrailEventHooks.during_call + elif func_name in ( + "async_post_call_success_hook", + "async_post_call_streaming_hook", + ): + return GuardrailEventHooks.post_call + return None + @functools.wraps(func) async def async_wrapper(*args, **kwargs): start_time = datetime.now() # Move start_time inside the wrapper self: CustomGuardrail = args[0] request_data: dict = kwargs.get("data") or kwargs.get("request_data") or {} + event_type = _infer_event_type_from_function_name(func.__name__) try: response = await func(*args, **kwargs) return self._process_response( @@ -730,6 +755,7 @@ def log_guardrail_information(func): start_time=start_time.timestamp(), end_time=datetime.now().timestamp(), duration=(datetime.now() - start_time).total_seconds(), + event_type=event_type, ) except Exception as e: return self._process_error( @@ -738,6 +764,7 @@ def log_guardrail_information(func): start_time=start_time.timestamp(), end_time=datetime.now().timestamp(), duration=(datetime.now() - start_time).total_seconds(), + event_type=event_type, ) @functools.wraps(func) @@ -745,18 +772,21 @@ def log_guardrail_information(func): start_time = datetime.now() # Move start_time inside the wrapper self: CustomGuardrail = args[0] request_data: dict = kwargs.get("data") or kwargs.get("request_data") or {} + event_type = _infer_event_type_from_function_name(func.__name__) try: response = func(*args, **kwargs) return self._process_response( response=response, request_data=request_data, duration=(datetime.now() - start_time).total_seconds(), + event_type=event_type, ) except Exception as e: return self._process_error( e=e, request_data=request_data, duration=(datetime.now() - start_time).total_seconds(), + event_type=event_type, ) @functools.wraps(func) diff --git a/tests/test_litellm/integrations/test_custom_guardrail.py b/tests/test_litellm/integrations/test_custom_guardrail.py index a719d102a7c..a322dfe9a2b 100644 --- a/tests/test_litellm/integrations/test_custom_guardrail.py +++ b/tests/test_litellm/integrations/test_custom_guardrail.py @@ -498,7 +498,7 @@ class TestPassthroughCallTypeHandling: def test_get_pre_call_type_with_allm_passthrough_route(self): """ Test that _get_pre_call_type correctly maps allm_passthrough_route. - + This tests Fix #1: allm_passthrough_route was not being handled, causing call_type to be None. """ from litellm.proxy.common_request_processing import ( @@ -509,14 +509,14 @@ class TestPassthroughCallTypeHandling: result = ProxyBaseLLMRequestProcessing._get_pre_call_type( route_type="allm_passthrough_route" ) - + # Should return allm_passthrough_route, not None assert result == "allm_passthrough_route" def test_get_pre_call_type_preserves_standard_mappings(self): """ Test that _get_pre_call_type still correctly maps standard route types. - + Ensures Fix #1 didn't break existing functionality. """ from litellm.proxy.common_request_processing import ( @@ -536,3 +536,235 @@ class TestPassthroughCallTypeHandling: ProxyBaseLLMRequestProcessing._get_pre_call_type(route_type="aresponses") == "responses" ) + + +class TestEventTypeLogging: + """Tests for event_type logging in guardrail information.""" + + @pytest.mark.asyncio + async def test_log_guardrail_information_infers_event_type_from_async_pre_call_hook( + self, + ): + """ + Test that log_guardrail_information decorator correctly infers GuardrailEventHooks.pre_call + from async_pre_call_hook function name. + """ + from litellm.integrations.custom_guardrail import log_guardrail_information + from litellm.types.guardrails import GuardrailEventHooks + + class TestGuardrail(CustomGuardrail): + def __init__(self): + super().__init__( + guardrail_name="test_event_type_guardrail", + event_hook=[ + GuardrailEventHooks.pre_call, + GuardrailEventHooks.post_call, + ], + ) + + @log_guardrail_information + async def async_pre_call_hook(self, data: dict, **kwargs): + return {"result": "pre_call_executed"} + + guardrail = TestGuardrail() + request_data = {"metadata": {}} + + await guardrail.async_pre_call_hook(data=request_data) + + # Check that the guardrail_mode was set to pre_call (not the full list) + logged_info = request_data["metadata"]["standard_logging_guardrail_information"] + assert len(logged_info) == 1 + assert logged_info[0]["guardrail_mode"] == GuardrailEventHooks.pre_call + + @pytest.mark.asyncio + async def test_log_guardrail_information_infers_event_type_from_async_post_call_success_hook( + self, + ): + """ + Test that log_guardrail_information decorator correctly infers GuardrailEventHooks.post_call + from async_post_call_success_hook function name. + """ + from litellm.integrations.custom_guardrail import log_guardrail_information + from litellm.types.guardrails import GuardrailEventHooks + + class TestGuardrail(CustomGuardrail): + def __init__(self): + super().__init__( + guardrail_name="test_event_type_guardrail", + event_hook=[ + GuardrailEventHooks.pre_call, + GuardrailEventHooks.post_call, + ], + ) + + @log_guardrail_information + async def async_post_call_success_hook(self, data: dict, **kwargs): + return {"result": "post_call_executed"} + + guardrail = TestGuardrail() + request_data = {"metadata": {}} + + await guardrail.async_post_call_success_hook(data=request_data) + + # Check that the guardrail_mode was set to post_call (not the full list) + logged_info = request_data["metadata"]["standard_logging_guardrail_information"] + assert len(logged_info) == 1 + assert logged_info[0]["guardrail_mode"] == GuardrailEventHooks.post_call + + @pytest.mark.asyncio + async def test_log_guardrail_information_infers_event_type_from_async_moderation_hook( + self, + ): + """ + Test that log_guardrail_information decorator correctly infers GuardrailEventHooks.during_call + from async_moderation_hook function name. + """ + from litellm.integrations.custom_guardrail import log_guardrail_information + from litellm.types.guardrails import GuardrailEventHooks + + class TestGuardrail(CustomGuardrail): + def __init__(self): + super().__init__( + guardrail_name="test_event_type_guardrail", + event_hook=[ + GuardrailEventHooks.during_call, + GuardrailEventHooks.post_call, + ], + ) + + @log_guardrail_information + async def async_moderation_hook(self, data: dict, **kwargs): + return {"result": "moderation_executed"} + + guardrail = TestGuardrail() + request_data = {"metadata": {}} + + await guardrail.async_moderation_hook(data=request_data) + + # Check that the guardrail_mode was set to during_call (not the full list) + logged_info = request_data["metadata"]["standard_logging_guardrail_information"] + assert len(logged_info) == 1 + assert logged_info[0]["guardrail_mode"] == GuardrailEventHooks.during_call + + @pytest.mark.asyncio + async def test_log_guardrail_information_infers_event_type_from_async_post_call_streaming_hook( + self, + ): + """ + Test that log_guardrail_information decorator correctly infers GuardrailEventHooks.post_call + from async_post_call_streaming_hook function name. + """ + from litellm.integrations.custom_guardrail import log_guardrail_information + from litellm.types.guardrails import GuardrailEventHooks + + class TestGuardrail(CustomGuardrail): + def __init__(self): + super().__init__( + guardrail_name="test_event_type_guardrail", + event_hook=[ + GuardrailEventHooks.pre_call, + GuardrailEventHooks.post_call, + ], + ) + + @log_guardrail_information + async def async_post_call_streaming_hook(self, data: dict, **kwargs): + return {"result": "streaming_executed"} + + guardrail = TestGuardrail() + request_data = {"metadata": {}} + + await guardrail.async_post_call_streaming_hook(data=request_data) + + # Check that the guardrail_mode was set to post_call (not the full list) + logged_info = request_data["metadata"]["standard_logging_guardrail_information"] + assert len(logged_info) == 1 + assert logged_info[0]["guardrail_mode"] == GuardrailEventHooks.post_call + + @pytest.mark.asyncio + async def test_log_guardrail_information_returns_none_for_unknown_function_name( + self, + ): + """ + Test that log_guardrail_information decorator returns None for event_type + when function name doesn't match known patterns, and falls back to self.event_hook. + """ + from litellm.integrations.custom_guardrail import log_guardrail_information + from litellm.types.guardrails import GuardrailEventHooks + + class TestGuardrail(CustomGuardrail): + def __init__(self): + super().__init__( + guardrail_name="test_event_type_guardrail", + event_hook=GuardrailEventHooks.pre_call, + ) + + @log_guardrail_information + async def some_other_hook(self, data: dict, **kwargs): + return {"result": "other_hook_executed"} + + guardrail = TestGuardrail() + request_data = {"metadata": {}} + + await guardrail.some_other_hook(data=request_data) + + # Check that the guardrail_mode falls back to self.event_hook + logged_info = request_data["metadata"]["standard_logging_guardrail_information"] + assert len(logged_info) == 1 + assert logged_info[0]["guardrail_mode"] == GuardrailEventHooks.pre_call + + def test_add_standard_logging_uses_event_type_over_event_hook(self): + """ + Test that add_standard_logging_guardrail_information_to_request_data + prioritizes event_type parameter over self.event_hook. + """ + from litellm.types.guardrails import GuardrailEventHooks + + guardrail = CustomGuardrail( + guardrail_name="test_guardrail", + event_hook=[GuardrailEventHooks.pre_call, GuardrailEventHooks.post_call], + ) + + request_data = {"metadata": {}} + + # Call with explicit event_type + guardrail.add_standard_logging_guardrail_information_to_request_data( + guardrail_json_response={"result": "ok"}, + request_data=request_data, + guardrail_status="success", + event_type=GuardrailEventHooks.post_call, + ) + + # Should use the provided event_type (post_call), not the full event_hook list + logged_info = request_data["metadata"]["standard_logging_guardrail_information"] + assert len(logged_info) == 1 + assert logged_info[0]["guardrail_mode"] == GuardrailEventHooks.post_call + + def test_add_standard_logging_falls_back_to_event_hook_when_event_type_is_none( + self, + ): + """ + Test that add_standard_logging_guardrail_information_to_request_data + falls back to self.event_hook when event_type is None. + """ + from litellm.types.guardrails import GuardrailEventHooks + + guardrail = CustomGuardrail( + guardrail_name="test_guardrail", + event_hook=GuardrailEventHooks.pre_call, + ) + + request_data = {"metadata": {}} + + # Call with event_type=None + guardrail.add_standard_logging_guardrail_information_to_request_data( + guardrail_json_response={"result": "ok"}, + request_data=request_data, + guardrail_status="success", + event_type=None, + ) + + # Should fall back to self.event_hook + logged_info = request_data["metadata"]["standard_logging_guardrail_information"] + assert len(logged_info) == 1 + assert logged_info[0]["guardrail_mode"] == GuardrailEventHooks.pre_call From 8bf0b9e19bb79c57b784cf74b88de5b26d1c28ad Mon Sep 17 00:00:00 2001 From: Yuta Saito Date: Sat, 27 Dec 2025 07:16:30 +0900 Subject: [PATCH 008/158] feat: add event_type parameter to add_standard_logging_guardrail_information_to_request_data --- .../guardrail_hooks/bedrock_guardrails.py | 8 +++++ .../guardrail_hooks/dynamoai/dynamoai.py | 6 ++++ .../ibm_guardrails/ibm_detector.py | 12 +++++++ .../guardrail_hooks/javelin/javelin.py | 6 +++- .../guardrail_hooks/lakera_ai_v2.py | 4 +++ .../model_armor/model_armor.py | 2 ++ .../guardrails/guardrail_hooks/noma/noma.py | 32 +++++++++++++++---- .../panw_prisma_airs/panw_prisma_airs.py | 24 ++++++++++++-- 8 files changed, 84 insertions(+), 10 deletions(-) diff --git a/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py b/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py index 62c997659bd..c89decfd01c 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py +++ b/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py @@ -449,6 +449,12 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): prepared_request.headers, ) + event_type = ( + GuardrailEventHooks.pre_call + if source == "INPUT" + else GuardrailEventHooks.post_call + ) + try: httpx_response = await self.async_handler.post( url=prepared_request.url, @@ -469,6 +475,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): start_time=start_time.timestamp(), end_time=datetime.now().timestamp(), duration=(datetime.now() - start_time).total_seconds(), + event_type=event_type, ) # Re-raise the exception to maintain existing behavior raise @@ -486,6 +493,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): start_time=start_time.timestamp(), end_time=datetime.now().timestamp(), duration=(datetime.now() - start_time).total_seconds(), + event_type=event_type, ) ######################################################### if httpx_response.status_code == 200: diff --git a/litellm/proxy/guardrails/guardrail_hooks/dynamoai/dynamoai.py b/litellm/proxy/guardrails/guardrail_hooks/dynamoai/dynamoai.py index 6915286a2d7..59381149809 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/dynamoai/dynamoai.py +++ b/litellm/proxy/guardrails/guardrail_hooks/dynamoai/dynamoai.py @@ -97,6 +97,7 @@ class DynamoAIGuardrails(CustomGuardrail): async def _call_dynamoai_guardrails( self, messages: List[Dict[str, Any]], + event_type: GuardrailEventHooks, text_type: str = "input", request_data: Optional[dict] = None, ) -> DynamoAIResponse: @@ -157,6 +158,7 @@ class DynamoAIGuardrails(CustomGuardrail): start_time=start_time.timestamp(), end_time=end_time.timestamp(), duration=duration, + event_type=event_type, ) return response_json @@ -177,6 +179,7 @@ class DynamoAIGuardrails(CustomGuardrail): start_time=start_time.timestamp(), end_time=end_time.timestamp(), duration=duration, + event_type=event_type, ) raise @@ -332,6 +335,7 @@ class DynamoAIGuardrails(CustomGuardrail): messages=_messages, text_type="input", request_data=data, + event_type=GuardrailEventHooks.pre_call, ) verbose_proxy_logger.debug( @@ -380,6 +384,7 @@ class DynamoAIGuardrails(CustomGuardrail): messages=_messages, text_type="input", request_data=data, + event_type=GuardrailEventHooks.during_call, ) verbose_proxy_logger.debug( @@ -460,6 +465,7 @@ class DynamoAIGuardrails(CustomGuardrail): messages=dynamoai_messages, text_type="output", request_data=data, + event_type=GuardrailEventHooks.post_call, ) verbose_proxy_logger.debug( diff --git a/litellm/proxy/guardrails/guardrail_hooks/ibm_guardrails/ibm_detector.py b/litellm/proxy/guardrails/guardrail_hooks/ibm_guardrails/ibm_detector.py index 55fa17c21e7..18228ae4b48 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/ibm_guardrails/ibm_detector.py +++ b/litellm/proxy/guardrails/guardrail_hooks/ibm_guardrails/ibm_detector.py @@ -108,6 +108,7 @@ class IBMGuardrailDetector(CustomGuardrail): async def _call_detector_server( self, contents: List[str], + event_type: GuardrailEventHooks, request_data: Optional[dict] = None, ) -> List[List[IBMDetectorDetection]]: """ @@ -172,6 +173,7 @@ class IBMGuardrailDetector(CustomGuardrail): start_time=start_time.timestamp(), end_time=end_time.timestamp(), duration=duration, + event_type=event_type, ) return response_json @@ -192,6 +194,7 @@ class IBMGuardrailDetector(CustomGuardrail): start_time=start_time.timestamp(), end_time=end_time.timestamp(), duration=duration, + event_type=event_type, ) raise @@ -199,6 +202,7 @@ class IBMGuardrailDetector(CustomGuardrail): async def _call_orchestrator( self, content: str, + event_type: GuardrailEventHooks, request_data: Optional[dict] = None, ) -> List[IBMDetectorDetection]: """ @@ -258,6 +262,7 @@ class IBMGuardrailDetector(CustomGuardrail): start_time=start_time.timestamp(), end_time=end_time.timestamp(), duration=duration, + event_type=event_type, ) return response_json.get("detections", []) @@ -278,6 +283,7 @@ class IBMGuardrailDetector(CustomGuardrail): start_time=start_time.timestamp(), end_time=end_time.timestamp(), duration=duration, + event_type=event_type, ) raise @@ -472,6 +478,7 @@ class IBMGuardrailDetector(CustomGuardrail): result = await self._call_detector_server( contents=contents_to_check, request_data=data, + event_type=GuardrailEventHooks.pre_call, ) verbose_proxy_logger.debug( @@ -500,6 +507,7 @@ class IBMGuardrailDetector(CustomGuardrail): orchestrator_result = await self._call_orchestrator( content=content, request_data=data, + event_type=GuardrailEventHooks.pre_call, ) verbose_proxy_logger.debug( @@ -557,6 +565,7 @@ class IBMGuardrailDetector(CustomGuardrail): result = await self._call_detector_server( contents=contents_to_check, request_data=data, + event_type=GuardrailEventHooks.during_call, ) verbose_proxy_logger.debug( @@ -585,6 +594,7 @@ class IBMGuardrailDetector(CustomGuardrail): orchestrator_result = await self._call_orchestrator( content=content, request_data=data, + event_type=GuardrailEventHooks.during_call, ) verbose_proxy_logger.debug( @@ -673,6 +683,7 @@ class IBMGuardrailDetector(CustomGuardrail): result = await self._call_detector_server( contents=contents_to_check, request_data=data, + event_type=GuardrailEventHooks.post_call, ) verbose_proxy_logger.debug( @@ -702,6 +713,7 @@ class IBMGuardrailDetector(CustomGuardrail): orchestrator_result = await self._call_orchestrator( content=content, request_data=data, + event_type=GuardrailEventHooks.post_call, ) verbose_proxy_logger.debug( diff --git a/litellm/proxy/guardrails/guardrail_hooks/javelin/javelin.py b/litellm/proxy/guardrails/guardrail_hooks/javelin/javelin.py index 6d4ed089818..953275acf14 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/javelin/javelin.py +++ b/litellm/proxy/guardrails/guardrail_hooks/javelin/javelin.py @@ -83,6 +83,7 @@ class JavelinGuardrail(CustomGuardrail): async def call_javelin_guard( self, request: JavelinGuardRequest, + event_type: GuardrailEventHooks, ) -> JavelinGuardResponse: """ Call the Javelin guard API. @@ -158,6 +159,7 @@ class JavelinGuardrail(CustomGuardrail): start_time=start_time.timestamp(), end_time=datetime.now().timestamp(), duration=(datetime.now() - start_time).total_seconds(), + event_type=event_type, ) async def async_pre_call_hook( @@ -208,7 +210,9 @@ class JavelinGuardrail(CustomGuardrail): config=self.config if self.config else {}, ) - javelin_response = await self.call_javelin_guard(request=javelin_guard_request) + javelin_response = await self.call_javelin_guard( + request=javelin_guard_request, event_type=GuardrailEventHooks.pre_call + ) assessments = javelin_response.get("assessments", []) reject_prompt = "" diff --git a/litellm/proxy/guardrails/guardrail_hooks/lakera_ai_v2.py b/litellm/proxy/guardrails/guardrail_hooks/lakera_ai_v2.py index 6d98866eadf..732331349e0 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/lakera_ai_v2.py +++ b/litellm/proxy/guardrails/guardrail_hooks/lakera_ai_v2.py @@ -70,6 +70,7 @@ class LakeraAIGuardrail(CustomGuardrail): self, messages: List[AllMessageValues], request_data: Dict, + event_type: GuardrailEventHooks, ) -> Tuple[LakeraAIResponse, Dict]: """ Call the Lakera AI v2 guard API. @@ -128,6 +129,7 @@ class LakeraAIGuardrail(CustomGuardrail): end_time=datetime.now().timestamp(), duration=(datetime.now() - start_time).total_seconds(), masked_entity_count=masked_entity_count, + event_type=event_type, ) def _mask_pii_in_messages( @@ -214,6 +216,7 @@ class LakeraAIGuardrail(CustomGuardrail): lakera_guardrail_response, masked_entity_count = await self.call_v2_guard( messages=new_messages, request_data=data, + event_type=GuardrailEventHooks.pre_call, ) ######################################################### @@ -279,6 +282,7 @@ class LakeraAIGuardrail(CustomGuardrail): lakera_guardrail_response, masked_entity_count = await self.call_v2_guard( messages=new_messages, request_data=data, + event_type=GuardrailEventHooks.during_call, ) ######################################################### diff --git a/litellm/proxy/guardrails/guardrail_hooks/model_armor/model_armor.py b/litellm/proxy/guardrails/guardrail_hooks/model_armor/model_armor.py index 51136c29eca..4b56fa5fc90 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/model_armor/model_armor.py +++ b/litellm/proxy/guardrails/guardrail_hooks/model_armor/model_armor.py @@ -327,6 +327,7 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase): start_time: Optional[float] = None, end_time: Optional[float] = None, duration: Optional[float] = None, + event_type: Optional[GuardrailEventHooks] = None, ): """ Override to store only the Model Armor API response, not the entire data dict. @@ -351,6 +352,7 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase): duration=duration, start_time=start_time, end_time=end_time, + event_type=event_type, ) return response diff --git a/litellm/proxy/guardrails/guardrail_hooks/noma/noma.py b/litellm/proxy/guardrails/guardrail_hooks/noma/noma.py index a0ea90ccf21..04f88e1a3da 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/noma/noma.py +++ b/litellm/proxy/guardrails/guardrail_hooks/noma/noma.py @@ -163,6 +163,7 @@ class NomaGuardrail(CustomGuardrail): self, request_data: dict, user_auth: UserAPIKeyAuth, + event_type: Optional[GuardrailEventHooks] = None, ) -> Optional[str]: """Shared logic for processing user message checks""" start_time = datetime.now() @@ -213,6 +214,7 @@ class NomaGuardrail(CustomGuardrail): start_time=start_time.timestamp(), end_time=end_time.timestamp(), duration=duration, + event_type=event_type, ) if self.monitor_mode: @@ -242,6 +244,7 @@ class NomaGuardrail(CustomGuardrail): request_data: dict, response: LLMResponse, user_auth: UserAPIKeyAuth, + event_type: Optional[GuardrailEventHooks] = None, ) -> Optional[str]: """Shared logic for processing LLM response checks""" @@ -293,6 +296,7 @@ class NomaGuardrail(CustomGuardrail): start_time=start_time.timestamp(), end_time=end_time.timestamp(), duration=duration, + event_type=event_type, ) if self.monitor_mode: @@ -602,7 +606,9 @@ class NomaGuardrail(CustomGuardrail): return data try: - return await self._check_user_message(data, user_api_key_dict) + return await self._check_user_message( + data, user_api_key_dict, GuardrailEventHooks.pre_call + ) except NomaBlockedMessage: # Blocked requests were already logged in _process_user_message_check with "blocked" status raise @@ -619,6 +625,7 @@ class NomaGuardrail(CustomGuardrail): start_time=start_time.timestamp(), end_time=start_time.timestamp(), duration=0.0, + event_type=GuardrailEventHooks.pre_call, ) verbose_proxy_logger.error(f"Noma pre-call hook failed: {str(e)}") @@ -650,7 +657,9 @@ class NomaGuardrail(CustomGuardrail): return data try: - return await self._check_user_message(data, user_api_key_dict) + return await self._check_user_message( + data, user_api_key_dict, GuardrailEventHooks.during_call + ) except NomaBlockedMessage: # Blocked requests were already logged in _process_user_message_check with "blocked" status raise @@ -667,6 +676,7 @@ class NomaGuardrail(CustomGuardrail): start_time=start_time.timestamp(), end_time=start_time.timestamp(), duration=0.0, + event_type=GuardrailEventHooks.during_call, ) verbose_proxy_logger.error(f"Noma moderation hook failed: {str(e)}") @@ -700,7 +710,9 @@ class NomaGuardrail(CustomGuardrail): return response try: - return await self._check_llm_response(data, response, user_api_key_dict) + return await self._check_llm_response( + data, response, user_api_key_dict, GuardrailEventHooks.post_call + ) except NomaBlockedMessage: # Blocked requests were already logged in _process_llm_response_check with "blocked" status raise @@ -717,6 +729,7 @@ class NomaGuardrail(CustomGuardrail): start_time=start_time.timestamp(), end_time=start_time.timestamp(), duration=0.0, + event_type=GuardrailEventHooks.post_call, ) verbose_proxy_logger.error(f"Noma post-call hook failed: {str(e)}") @@ -728,9 +741,12 @@ class NomaGuardrail(CustomGuardrail): self, request_data: dict, user_auth: UserAPIKeyAuth, + event_type: Optional[GuardrailEventHooks] = None, ) -> Union[Exception, str, dict, None]: """Check user message for policy violations""" - user_message = await self._process_user_message_check(request_data, user_auth) + user_message = await self._process_user_message_check( + request_data, user_auth, event_type + ) if not user_message: return request_data @@ -741,10 +757,11 @@ class NomaGuardrail(CustomGuardrail): request_data: dict, response: LLMResponse, user_auth: UserAPIKeyAuth, + event_type: Optional[GuardrailEventHooks] = None, ) -> Any: """Check LLM response for policy violations""" content = await self._process_llm_response_check( - request_data, response, user_auth + request_data, response, user_auth, event_type ) if not content: return response @@ -858,7 +875,10 @@ class NomaGuardrail(CustomGuardrail): if isinstance(assembled_model_response, ModelResponse): try: processed_response = await self._check_llm_response( - request_data, assembled_model_response, user_api_key_dict + request_data, + assembled_model_response, + user_api_key_dict, + GuardrailEventHooks.post_call, ) except NomaBlockedMessage: raise diff --git a/litellm/proxy/guardrails/guardrail_hooks/panw_prisma_airs/panw_prisma_airs.py b/litellm/proxy/guardrails/guardrail_hooks/panw_prisma_airs/panw_prisma_airs.py index 88145ae9e47..02e481acddd 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/panw_prisma_airs/panw_prisma_airs.py +++ b/litellm/proxy/guardrails/guardrail_hooks/panw_prisma_airs/panw_prisma_airs.py @@ -24,6 +24,7 @@ from litellm.llms.custom_httpx.http_handler import ( httpxSpecialProvider, ) from litellm.proxy._types import UserAPIKeyAuth +from litellm.types.guardrails import GuardrailEventHooks from litellm.types.utils import CallTypesLiteral, ModelResponse if TYPE_CHECKING: @@ -523,6 +524,7 @@ class PanwPrismaAirsHandler(CustomGuardrail): scan_result: Dict[str, Any], data: Dict[str, Any], start_time: datetime, + event_type: GuardrailEventHooks, is_response: bool = False, ) -> Optional[Dict[str, Any]]: """Handle API errors with fail-open/fail-closed logic.""" @@ -542,6 +544,7 @@ class PanwPrismaAirsHandler(CustomGuardrail): start_time=start_time.timestamp(), end_time=end_time.timestamp(), duration=duration, + event_type=event_type, ) if scan_result.get("_always_block"): @@ -735,7 +738,11 @@ class PanwPrismaAirsHandler(CustomGuardrail): if scan_result.get("_is_transient") or scan_result.get("_always_block"): return self._handle_api_error_with_logging( - scan_result, data, start_time, is_response=False + scan_result, + data, + start_time, + is_response=False, + event_type=GuardrailEventHooks.pre_call, ) end_time = datetime.now() @@ -749,6 +756,7 @@ class PanwPrismaAirsHandler(CustomGuardrail): start_time=start_time.timestamp(), end_time=end_time.timestamp(), duration=(end_time - start_time).total_seconds(), + event_type=GuardrailEventHooks.pre_call, ) action = scan_result.get("action", "block") @@ -872,7 +880,11 @@ class PanwPrismaAirsHandler(CustomGuardrail): if scan_result.get("_is_transient") or scan_result.get("_always_block"): self._handle_api_error_with_logging( - scan_result, data, start_time, is_response=True + scan_result, + data, + start_time, + is_response=True, + event_type=GuardrailEventHooks.post_call, ) return response @@ -887,6 +899,7 @@ class PanwPrismaAirsHandler(CustomGuardrail): start_time=start_time.timestamp(), end_time=end_time.timestamp(), duration=(end_time - start_time).total_seconds(), + event_type=GuardrailEventHooks.post_call, ) action = scan_result.get("action", "block") @@ -1066,7 +1079,11 @@ class PanwPrismaAirsHandler(CustomGuardrail): if scan_result.get("_is_transient") or scan_result.get("_always_block"): self._handle_api_error_with_logging( - scan_result, request_data, start_time, is_response=True + scan_result, + request_data, + start_time, + is_response=True, + event_type=EventHooks.post_call, ) for chunk in all_chunks: yield chunk @@ -1083,6 +1100,7 @@ class PanwPrismaAirsHandler(CustomGuardrail): start_time=start_time.timestamp(), end_time=end_time.timestamp(), duration=(end_time - start_time).total_seconds(), + event_type=EventHooks.post_call, ) # Add guardrail to applied guardrails header for observability From 7d673c308da5da3c3d8c889e71a30894a7d83a74 Mon Sep 17 00:00:00 2001 From: Yuta Saito Date: Sat, 27 Dec 2025 07:23:07 +0900 Subject: [PATCH 009/158] lint --- .../guardrail_hooks/bedrock_guardrails.py | 18 +++++++++--------- .../ibm_guardrails/ibm_detector.py | 1 - .../guardrail_hooks/model_armor/model_armor.py | 4 +++- .../guardrails/guardrail_hooks/noma/noma.py | 5 +---- 4 files changed, 13 insertions(+), 15 deletions(-) diff --git a/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py b/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py index c89decfd01c..c8cef6e2790 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py +++ b/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py @@ -613,10 +613,10 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): """ Only raise exception for "BLOCKED" actions, not for "ANONYMIZED" actions. - If `self.mask_request_content` or `self.mask_response_content` is set to `True`, + If `self.mask_request_content` or `self.mask_response_content` is set to `True`, then use the output from the guardrail to mask the request or response content. - - However, even with masking enabled, content with action="BLOCKED" should still + + However, even with masking enabled, content with action="BLOCKED" should still raise an exception, only content with action="ANONYMIZED" should be masked. """ @@ -739,9 +739,9 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): ######################################################### ########## 1. Make the Bedrock API request ########## ######################################################### - bedrock_guardrail_response: Optional[Union[BedrockGuardrailResponse, str]] = ( - None - ) + bedrock_guardrail_response: Optional[ + Union[BedrockGuardrailResponse, str] + ] = None try: bedrock_guardrail_response = await self.make_bedrock_api_request( source="INPUT", messages=filtered_messages, request_data=data @@ -811,9 +811,9 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): ######################################################### ########## 1. Make the Bedrock API request ########## ######################################################### - bedrock_guardrail_response: Optional[Union[BedrockGuardrailResponse, str]] = ( - None - ) + bedrock_guardrail_response: Optional[ + Union[BedrockGuardrailResponse, str] + ] = None try: bedrock_guardrail_response = await self.make_bedrock_api_request( source="INPUT", messages=filtered_messages, request_data=data diff --git a/litellm/proxy/guardrails/guardrail_hooks/ibm_guardrails/ibm_detector.py b/litellm/proxy/guardrails/guardrail_hooks/ibm_guardrails/ibm_detector.py index 18228ae4b48..2fc05213640 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/ibm_guardrails/ibm_detector.py +++ b/litellm/proxy/guardrails/guardrail_hooks/ibm_guardrails/ibm_detector.py @@ -143,7 +143,6 @@ class IBMGuardrailDetector(CustomGuardrail): ) try: - response = await self.async_handler.post( url=self.api_url, json=payload, diff --git a/litellm/proxy/guardrails/guardrail_hooks/model_armor/model_armor.py b/litellm/proxy/guardrails/guardrail_hooks/model_armor/model_armor.py index 4b56fa5fc90..a12eb2486d2 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/model_armor/model_armor.py +++ b/litellm/proxy/guardrails/guardrail_hooks/model_armor/model_armor.py @@ -295,7 +295,9 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase): filters = ( list(filter_results.values()) if isinstance(filter_results, dict) - else filter_results if isinstance(filter_results, list) else [] + else filter_results + if isinstance(filter_results, list) + else [] ) # Prefer sanitized text from deidentifyResult if present diff --git a/litellm/proxy/guardrails/guardrail_hooks/noma/noma.py b/litellm/proxy/guardrails/guardrail_hooks/noma/noma.py index 04f88e1a3da..1794751a08c 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/noma/noma.py +++ b/litellm/proxy/guardrails/guardrail_hooks/noma/noma.py @@ -119,9 +119,7 @@ class NomaGuardrail(CustomGuardrail): self.api_base = api_base or os.environ.get( "NOMA_API_BASE", NomaGuardrail._DEFAULT_API_BASE ) - self.application_id = application_id or os.environ.get( - "NOMA_APPLICATION_ID" - ) + self.application_id = application_id or os.environ.get("NOMA_APPLICATION_ID") self.default_application_id = "litellm" if monitor_mode is None: @@ -582,7 +580,6 @@ class NomaGuardrail(CustomGuardrail): data: dict, call_type: CallTypesLiteral, ) -> Optional[Union[Exception, str, dict]]: - verbose_proxy_logger.debug("Running Noma pre-call hook") if ( From c6d6ffd567055acde3247bcaa6d00a8adb59b9f8 Mon Sep 17 00:00:00 2001 From: Yuta Saito Date: Sat, 27 Dec 2025 07:28:16 +0900 Subject: [PATCH 010/158] fix: mypy error --- litellm/integrations/custom_guardrail.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/litellm/integrations/custom_guardrail.py b/litellm/integrations/custom_guardrail.py index 7bc02bb3bb6..6a76b57e7f7 100644 --- a/litellm/integrations/custom_guardrail.py +++ b/litellm/integrations/custom_guardrail.py @@ -516,12 +516,13 @@ class CustomGuardrail(CustomLogger): from litellm.types.utils import GuardrailMode # Use event_type if provided, otherwise fall back to self.event_hook + guardrail_mode: Union[GuardrailEventHooks, GuardrailMode, List[GuardrailEventHooks]] if event_type is not None: guardrail_mode = event_type elif isinstance(self.event_hook, Mode): - guardrail_mode = GuardrailMode(**self.event_hook.model_dump()) + guardrail_mode = GuardrailMode(**dict(self.event_hook.model_dump())) # type: ignore[typeddict-item] else: - guardrail_mode = self.event_hook + guardrail_mode = self.event_hook # type: ignore[assignment] slg = StandardLoggingGuardrailInformation( guardrail_name=self.guardrail_name, From f5024624d7b9fdc5fd888edffbb7b388f192b09f Mon Sep 17 00:00:00 2001 From: yurekami Date: Mon, 29 Dec 2025 03:44:03 +0900 Subject: [PATCH 011/158] fix: correct deepseek-v3p2 pricing for Fireworks AI MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Updated pricing for fireworks_ai/accounts/fireworks/models/deepseek-v3p2: - input_cost_per_token: 1.2e-06 -> 5.6e-07 ($0.56/1M tokens) - output_cost_per_token: 1.2e-06 -> 1.68e-06 ($1.68/1M tokens) Pricing verified from https://fireworks.ai/models/fireworks/deepseek-v3p2 Fixes #17998 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude Opus 4.5 --- litellm/model_prices_and_context_window_backup.json | 4 ++-- model_prices_and_context_window.json | 4 ++-- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 513a4a554e0..60ddb74533c 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -10885,13 +10885,13 @@ "supports_tool_choice": true }, "fireworks_ai/accounts/fireworks/models/deepseek-v3p2": { - "input_cost_per_token": 1.2e-06, + "input_cost_per_token": 5.6e-07, "litellm_provider": "fireworks_ai", "max_input_tokens": 163840, "max_output_tokens": 163840, "max_tokens": 163840, "mode": "chat", - "output_cost_per_token": 1.2e-06, + "output_cost_per_token": 1.68e-06, "source": "https://fireworks.ai/models/fireworks/deepseek-v3p2", "supports_function_calling": true, "supports_reasoning": true, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 4651107c5b8..9eae448da97 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -10943,13 +10943,13 @@ "supports_tool_choice": true }, "fireworks_ai/accounts/fireworks/models/deepseek-v3p2": { - "input_cost_per_token": 1.2e-06, + "input_cost_per_token": 5.6e-07, "litellm_provider": "fireworks_ai", "max_input_tokens": 163840, "max_output_tokens": 163840, "max_tokens": 163840, "mode": "chat", - "output_cost_per_token": 1.2e-06, + "output_cost_per_token": 1.68e-06, "source": "https://fireworks.ai/models/fireworks/deepseek-v3p2", "supports_function_calling": true, "supports_reasoning": true, From 93628c06eed2b2a5aa0da1bc85e355052f8c36b6 Mon Sep 17 00:00:00 2001 From: Chesars Date: Mon, 29 Dec 2025 00:19:20 -0300 Subject: [PATCH 012/158] feat: Add MiniMax provider support to UI dashboard - Add MiniMax to provider_create_fields.json with credential fields: - api_key (required, password field) - api_base (optional, defaults to https://api.minimax.io/v1) - Add MiniMax to UI provider enum and mappings - Includes tooltip for International vs China endpoints - Default model placeholder: minimax/MiniMax-M2 This enables users to configure MiniMax models directly through the proxy UI dashboard without needing to edit YAML config files. Fixes #18481 --- .../provider_create_fields.json | 28 +++++++++++++++++++ .../src/components/provider_info_helpers.tsx | 2 ++ 2 files changed, 30 insertions(+) diff --git a/litellm/proxy/public_endpoints/provider_create_fields.json b/litellm/proxy/public_endpoints/provider_create_fields.json index 9916bdf6923..0e0991dc214 100644 --- a/litellm/proxy/public_endpoints/provider_create_fields.json +++ b/litellm/proxy/public_endpoints/provider_create_fields.json @@ -1680,6 +1680,34 @@ ], "default_model_placeholder": "gpt-3.5-turbo" }, + { + "provider": "MINIMAX", + "provider_display_name": "MiniMax", + "litellm_provider": "minimax", + "credential_fields": [ + { + "key": "api_key", + "label": "API Key", + "placeholder": "your-minimax-api-key", + "tooltip": "MiniMax API Key from https://platform.minimaxi.com/", + "required": true, + "field_type": "password", + "options": null, + "default_value": null + }, + { + "key": "api_base", + "label": "API Base URL", + "placeholder": "https://api.minimax.io/v1", + "tooltip": "International: https://api.minimax.io/v1, China: https://api.minimaxi.com/v1", + "required": false, + "field_type": "text", + "options": null, + "default_value": "https://api.minimax.io/v1" + } + ], + "default_model_placeholder": "minimax/MiniMax-M2" + }, { "provider": "MOONSHOT", "provider_display_name": "Moonshot", diff --git a/ui/litellm-dashboard/src/components/provider_info_helpers.tsx b/ui/litellm-dashboard/src/components/provider_info_helpers.tsx index 277786aa8ce..d44fe9446b6 100644 --- a/ui/litellm-dashboard/src/components/provider_info_helpers.tsx +++ b/ui/litellm-dashboard/src/components/provider_info_helpers.tsx @@ -24,6 +24,7 @@ export enum Providers { Hosted_Vllm = "vllm", Infinity = "Infinity", JinaAI = "Jina AI", + MiniMax = "MiniMax", MistralAI = "Mistral AI", Ollama = "Ollama", OpenAI = "OpenAI", @@ -55,6 +56,7 @@ export const provider_map: Record = { Google_AI_Studio: "gemini", Bedrock: "bedrock", Groq: "groq", + MiniMax: "minimax", MistralAI: "mistral", Cohere: "cohere", OpenAI_Compatible: "openai", From e31ee9be51121c63080d16ef45b53c32fa7c0f98 Mon Sep 17 00:00:00 2001 From: Chesars Date: Mon, 29 Dec 2025 00:21:14 -0300 Subject: [PATCH 013/158] feat: Add MiniMax official logo to UI - Downloaded official MiniMax logo from HuggingFace repository - Added minimax.svg to assets/logos directory - Updated providerLogoMap to reference the logo - Logo source: https://huggingface.co/MiniMaxAI/MiniMax-VL-01 --- ui/litellm-dashboard/public/assets/logos/minimax.svg | 1 + ui/litellm-dashboard/src/components/provider_info_helpers.tsx | 1 + 2 files changed, 2 insertions(+) create mode 100644 ui/litellm-dashboard/public/assets/logos/minimax.svg diff --git a/ui/litellm-dashboard/public/assets/logos/minimax.svg b/ui/litellm-dashboard/public/assets/logos/minimax.svg new file mode 100644 index 00000000000..59b741bbcb7 --- /dev/null +++ b/ui/litellm-dashboard/public/assets/logos/minimax.svg @@ -0,0 +1 @@ +资源 2 \ No newline at end of file diff --git a/ui/litellm-dashboard/src/components/provider_info_helpers.tsx b/ui/litellm-dashboard/src/components/provider_info_helpers.tsx index d44fe9446b6..daea7abe877 100644 --- a/ui/litellm-dashboard/src/components/provider_info_helpers.tsx +++ b/ui/litellm-dashboard/src/components/provider_info_helpers.tsx @@ -112,6 +112,7 @@ export const providerLogoMap: Record = { [Providers.Google_AI_Studio]: `${asset_logos_folder}google.svg`, [Providers.Hosted_Vllm]: `${asset_logos_folder}vllm.png`, [Providers.Infinity]: `${asset_logos_folder}infinity.png`, + [Providers.MiniMax]: `${asset_logos_folder}minimax.svg`, [Providers.MistralAI]: `${asset_logos_folder}mistral.svg`, [Providers.Ollama]: `${asset_logos_folder}ollama.svg`, [Providers.OpenAI]: `${asset_logos_folder}openai_small.svg`, From e4c9b0bea2b92e932e2403bb8443c762beac4b2b Mon Sep 17 00:00:00 2001 From: Devaj Date: Mon, 29 Dec 2025 09:55:41 +0530 Subject: [PATCH 014/158] fix(vertex_ai): convert image URLs to base64 for Vertex AI Anthropic Fixes #18430 - Pass custom_llm_provider to anthropic_messages_pt instead of hardcoded 'anthropic' - Add check for vertex_ai provider to force base64 conversion for image URLs - Add tests to verify behavior for both Vertex AI and regular Anthropic --- .../prompt_templates/factory.py | 8 +- litellm/llms/anthropic/chat/transformation.py | 2 +- ..._vertex_ai_anthropic_image_url_handling.py | 179 ++++++++++++++++++ 3 files changed, 186 insertions(+), 3 deletions(-) create mode 100644 tests/test_litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/test_vertex_ai_anthropic_image_url_handling.py diff --git a/litellm/litellm_core_utils/prompt_templates/factory.py b/litellm/litellm_core_utils/prompt_templates/factory.py index 04c8c235557..1543c0f9f45 100644 --- a/litellm/litellm_core_utils/prompt_templates/factory.py +++ b/litellm/litellm_core_utils/prompt_templates/factory.py @@ -930,7 +930,8 @@ def create_anthropic_image_param( # Check if the image URL is an HTTP/HTTPS URL if image_url.startswith("http://") or image_url.startswith("https://"): - # For Bedrock invoke, always convert URLs to base64 (Bedrock invoke doesn't support URLs) + # For Bedrock invoke and Vertex AI Anthropic, always convert URLs to base64 + # as these providers don't support URL sources for images if is_bedrock_invoke or image_url.startswith("http://"): base64_url = convert_url_to_base64(url=image_url) image_chunk = convert_to_anthropic_image_obj( @@ -1914,9 +1915,12 @@ def anthropic_messages_pt( # noqa: PLR0915 "format": image_url_value.get("format"), } # Bedrock invoke models have format: invoke/... + # Vertex AI Anthropic also doesn't support URL sources for images is_bedrock_invoke = model.lower().startswith("invoke/") + is_vertex_ai = llm_provider.startswith("vertex_ai") if llm_provider else False + force_base64 = is_bedrock_invoke or is_vertex_ai _anthropic_content_element = create_anthropic_image_param( - image_url_input, format=format, is_bedrock_invoke=is_bedrock_invoke + image_url_input, format=format, is_bedrock_invoke=force_base64 ) _content_element = add_cache_control_to_content( anthropic_content_element=_anthropic_content_element, diff --git a/litellm/llms/anthropic/chat/transformation.py b/litellm/llms/anthropic/chat/transformation.py index b477dbd457e..1c108f1e94a 100644 --- a/litellm/llms/anthropic/chat/transformation.py +++ b/litellm/llms/anthropic/chat/transformation.py @@ -994,7 +994,7 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): anthropic_messages = anthropic_messages_pt( model=model, messages=messages, - llm_provider="anthropic", + llm_provider=self.custom_llm_provider or "anthropic", ) except Exception as e: raise AnthropicError( diff --git a/tests/test_litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/test_vertex_ai_anthropic_image_url_handling.py b/tests/test_litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/test_vertex_ai_anthropic_image_url_handling.py new file mode 100644 index 00000000000..fca784342d7 --- /dev/null +++ b/tests/test_litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/test_vertex_ai_anthropic_image_url_handling.py @@ -0,0 +1,179 @@ +""" +Tests for Vertex AI Anthropic image URL handling. + +Issue: https://github.com/BerriAI/litellm/issues/18430 +Vertex AI Anthropic models don't support URL sources for images. +LiteLLM should convert image URLs to base64 when using Vertex AI Anthropic. +""" +import os +import sys +from unittest.mock import patch, MagicMock + +import pytest + +sys.path.insert( + 0, os.path.abspath("../../../../../..") +) # Adds the parent directory to the system path + +from litellm.litellm_core_utils.prompt_templates.factory import ( + anthropic_messages_pt, + create_anthropic_image_param, +) + + +class TestVertexAIAnthropicImageURLHandling: + """Test that Vertex AI Anthropic converts image URLs to base64.""" + + @patch("litellm.litellm_core_utils.prompt_templates.factory.convert_url_to_base64") + def test_vertex_ai_anthropic_converts_https_url_to_base64( + self, mock_convert_url: MagicMock + ): + """ + Test that HTTPS image URLs are converted to base64 for Vertex AI Anthropic. + + For regular Anthropic, HTTPS URLs are passed through as URL type. + For Vertex AI Anthropic, HTTPS URLs should be converted to base64. + """ + mock_convert_url.return_value = "data:image/jpeg;base64,/9j/4AAQSkZJRgABAQAAAQ==" + + messages = [ + { + "role": "user", + "content": [ + {"type": "text", "text": "Describe this image"}, + { + "type": "image_url", + "image_url": {"url": "https://example.com/image.jpg"}, + }, + ], + } + ] + + # For Vertex AI, image URLs should be converted to base64 + result = anthropic_messages_pt( + messages=messages, + model="claude-sonnet-4", + llm_provider="vertex_ai", + ) + + # Verify convert_url_to_base64 was called + mock_convert_url.assert_called_once_with(url="https://example.com/image.jpg") + + # Check the result has base64 source type + user_message = result[0] + assert user_message["role"] == "user" + image_content = user_message["content"][1] + assert image_content["type"] == "image" + assert image_content["source"]["type"] == "base64" + + @patch("litellm.litellm_core_utils.prompt_templates.factory.convert_url_to_base64") + def test_regular_anthropic_uses_url_type_for_https( + self, mock_convert_url: MagicMock + ): + """ + Test that regular Anthropic API uses URL type for HTTPS images. + + This confirms the original behavior is preserved for non-Vertex AI. + """ + messages = [ + { + "role": "user", + "content": [ + {"type": "text", "text": "Describe this image"}, + { + "type": "image_url", + "image_url": {"url": "https://example.com/image.jpg"}, + }, + ], + } + ] + + # For regular Anthropic, HTTPS URLs should NOT be converted + result = anthropic_messages_pt( + messages=messages, + model="claude-sonnet-4", + llm_provider="anthropic", + ) + + # convert_url_to_base64 should NOT be called for regular Anthropic with HTTPS + mock_convert_url.assert_not_called() + + # Check the result has URL source type + user_message = result[0] + assert user_message["role"] == "user" + image_content = user_message["content"][1] + assert image_content["type"] == "image" + assert image_content["source"]["type"] == "url" + assert image_content["source"]["url"] == "https://example.com/image.jpg" + + @patch("litellm.litellm_core_utils.prompt_templates.factory.convert_url_to_base64") + def test_vertex_ai_beta_also_converts_to_base64( + self, mock_convert_url: MagicMock + ): + """ + Test that vertex_ai_beta provider also converts image URLs to base64. + """ + mock_convert_url.return_value = "data:image/png;base64,iVBORw0KGgo=" + + messages = [ + { + "role": "user", + "content": [ + {"type": "text", "text": "What is in this image?"}, + { + "type": "image_url", + "image_url": "https://example.com/photo.png", + }, + ], + } + ] + + result = anthropic_messages_pt( + messages=messages, + model="claude-3-opus", + llm_provider="vertex_ai_beta", + ) + + # Verify convert_url_to_base64 was called + mock_convert_url.assert_called_once() + + # Check the result has base64 source type + user_message = result[0] + image_content = user_message["content"][1] + assert image_content["source"]["type"] == "base64" + + +class TestCreateAnthropicImageParam: + """Test the create_anthropic_image_param function directly.""" + + @patch("litellm.litellm_core_utils.prompt_templates.factory.convert_url_to_base64") + def test_force_base64_converts_https_url(self, mock_convert_url: MagicMock): + """ + Test that is_bedrock_invoke=True (used for both Bedrock and Vertex AI) + forces conversion of HTTPS URLs to base64. + """ + mock_convert_url.return_value = "data:image/jpeg;base64,/9j/4AAQSkZJRg==" + + result = create_anthropic_image_param( + image_url_input="https://example.com/image.jpg", + format=None, + is_bedrock_invoke=True, # This flag is set for both Bedrock and Vertex AI + ) + + mock_convert_url.assert_called_once_with(url="https://example.com/image.jpg") + assert result["source"]["type"] == "base64" + + @patch("litellm.litellm_core_utils.prompt_templates.factory.convert_url_to_base64") + def test_no_force_uses_url_type(self, mock_convert_url: MagicMock): + """ + Test that without force, HTTPS URLs use URL type. + """ + result = create_anthropic_image_param( + image_url_input="https://example.com/image.jpg", + format=None, + is_bedrock_invoke=False, + ) + + mock_convert_url.assert_not_called() + assert result["source"]["type"] == "url" + assert result["source"]["url"] == "https://example.com/image.jpg" From b720fba142bfc52d961c8ca595c4042afcbba9e8 Mon Sep 17 00:00:00 2001 From: Prajeena Maharjan Date: Sun, 28 Dec 2025 23:48:43 -0600 Subject: [PATCH 015/158] Fix formatting in proxy configs documentation Got an error message: {"error":{"message":"Invalid JSON payload: trailing comma is not allowed: line 8 column 8 (char 141)","type":"invalid_request_error","param":"request_body","code":"400"}}% --- docs/my-website/docs/proxy/configs.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/my-website/docs/proxy/configs.md b/docs/my-website/docs/proxy/configs.md index ba4ca190aa9..bc2f6a13362 100644 --- a/docs/my-website/docs/proxy/configs.md +++ b/docs/my-website/docs/proxy/configs.md @@ -116,7 +116,7 @@ curl --location 'http://0.0.0.0:4000/chat/completions' \ "role": "user", "content": "what llm are you" } - ], + ] } ' ``` From d7468dab7e812054568c73e99392e954114674a8 Mon Sep 17 00:00:00 2001 From: Daniel Krueger Date: Mon, 29 Dec 2025 15:29:54 +0100 Subject: [PATCH 016/158] fix authentication errors at messages API via azure_ai Use x-api-key instead of api-key. This has been removed by commit 61e737e361d1b2649e9fd491580529ea894dd0d4 for unknown reason. --- .../anthropic/messages_transformation.py | 7 ++++++- ...azure_anthropic_messages_transformation.py | 19 ++++++++++--------- 2 files changed, 16 insertions(+), 10 deletions(-) diff --git a/litellm/llms/azure_ai/anthropic/messages_transformation.py b/litellm/llms/azure_ai/anthropic/messages_transformation.py index 73dc84167ab..55818cc07d6 100644 --- a/litellm/llms/azure_ai/anthropic/messages_transformation.py +++ b/litellm/llms/azure_ai/anthropic/messages_transformation.py @@ -48,7 +48,12 @@ class AzureAnthropicMessagesConfig(AnthropicMessagesConfig): headers = BaseAzureLLM._base_validate_azure_environment( headers=headers, litellm_params=litellm_params_obj ) - + + # Azure Anthropic uses x-api-key header (not api-key) + # Convert api-key to x-api-key if present + if "api-key" in headers and "x-api-key" not in headers: + headers["x-api-key"] = headers.pop("api-key") + # Set anthropic-version header if "anthropic-version" not in headers: headers["anthropic-version"] = "2023-06-01" diff --git a/tests/test_litellm/llms/azure_ai/claude/test_azure_anthropic_messages_transformation.py b/tests/test_litellm/llms/azure_ai/claude/test_azure_anthropic_messages_transformation.py index d78a638fd89..bdced849c7e 100644 --- a/tests/test_litellm/llms/azure_ai/claude/test_azure_anthropic_messages_transformation.py +++ b/tests/test_litellm/llms/azure_ai/claude/test_azure_anthropic_messages_transformation.py @@ -55,11 +55,12 @@ class TestAzureAnthropicMessagesConfig: assert isinstance(call_args[1]["litellm_params"], GenericLiteLLMParams) assert call_args[1]["litellm_params"].api_key == "test-api-key" assert "anthropic-version" in result - # api-key header is preserved as-is (no conversion to x-api-key) - assert "api-key" in result + assert "x-api-key" in result + assert result["x-api-key"] == "test-api-key" + assert "api-key" not in result - def test_validate_anthropic_messages_environment_preserves_api_key_header(self): - """Test that api-key header is preserved as-is (Azure handles the header internally)""" + def test_validate_anthropic_messages_environment_converts_api_key_to_x_api_key(self): + """Test that api-key header is converted to x-api-key""" config = AzureAnthropicMessagesConfig() headers = {} model = "claude-sonnet-4-5" @@ -79,9 +80,10 @@ class TestAzureAnthropicMessagesConfig: litellm_params=litellm_params, ) - # Verify api-key header is preserved as-is - assert "api-key" in result - assert result["api-key"] == "test-api-key" + # Verify api-key was converted to x-api-key + assert "x-api-key" in result + assert result["x-api-key"] == "test-api-key" + assert "api-key" not in result def test_validate_anthropic_messages_environment_sets_headers(self): """Test that required headers are set""" @@ -108,8 +110,7 @@ class TestAzureAnthropicMessagesConfig: assert result["anthropic-version"] == "2023-06-01" assert "content-type" in result assert result["content-type"] == "application/json" - # api-key header is preserved as-is - assert "api-key" in result + assert "x-api-key" in result def test_get_complete_url_with_base_url(self): """Test get_complete_url with base URL""" From a1849a152ce40eaefd86f00d406b73dbaae8b089 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Mon, 29 Dec 2025 11:27:22 -0800 Subject: [PATCH 017/158] Playwright setup in UI directory --- .gitignore | 2 + .../e2e_tests/playwright.config.ts | 52 +++++++++++++++ .../e2e_tests/tests/basic.spec.ts | 10 +++ ui/litellm-dashboard/package-lock.json | 64 +++++++++++++++++++ ui/litellm-dashboard/package.json | 5 +- ui/litellm-dashboard/vitest.config.ts | 2 + 6 files changed, 134 insertions(+), 1 deletion(-) create mode 100644 ui/litellm-dashboard/e2e_tests/playwright.config.ts create mode 100644 ui/litellm-dashboard/e2e_tests/tests/basic.spec.ts diff --git a/.gitignore b/.gitignore index 8196d1d9f24..b6d4fd44190 100644 --- a/.gitignore +++ b/.gitignore @@ -101,3 +101,5 @@ tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py litellm/proxy/_experimental/out/guardrails/index.html scripts/test_vertex_ai_search.py LAZY_LOADING_IMPROVEMENTS.md +**/test-results +**/playwright-report \ No newline at end of file diff --git a/ui/litellm-dashboard/e2e_tests/playwright.config.ts b/ui/litellm-dashboard/e2e_tests/playwright.config.ts new file mode 100644 index 00000000000..cce416075f4 --- /dev/null +++ b/ui/litellm-dashboard/e2e_tests/playwright.config.ts @@ -0,0 +1,52 @@ +import { defineConfig, devices } from "@playwright/test"; + +/** + * See https://playwright.dev/docs/test-configuration. + */ +export default defineConfig({ + testDir: "./tests", + testMatch: "**/*.spec.ts", + testIgnore: ["**/*.test.*"], + /* Run tests in files in parallel */ + fullyParallel: true, + /* Fail the build on CI if you accidentally left test.only in the source code. */ + forbidOnly: !!process.env.CI, + /* Retry on CI only */ + retries: process.env.CI ? 2 : 0, + /* Opt out of parallel tests on CI. */ + workers: process.env.CI ? 1 : undefined, + /* Reporter to use. See https://playwright.dev/docs/test-reporters */ + reporter: "html", + /* Shared settings for all the projects below. See https://playwright.dev/docs/api/class-testoptions. */ + use: { + /* Base URL to use in actions like `await page.goto('/')`. */ + baseURL: "http://localhost:3000", + + /* Collect trace when retrying the failed test. See https://playwright.dev/docs/trace-viewer */ + trace: "on-first-retry", + }, + + /* Configure projects for major browsers */ + projects: [ + { + name: "chromium", + use: { ...devices["Desktop Chrome"] }, + }, + + { + name: "firefox", + use: { ...devices["Desktop Firefox"] }, + }, + + { + name: "webkit", + use: { ...devices["Desktop Safari"] }, + }, + ], + + /* Timeout settings */ + timeout: 4 * 60 * 1000, + expect: { + timeout: 10 * 1000, + }, +}); diff --git a/ui/litellm-dashboard/e2e_tests/tests/basic.spec.ts b/ui/litellm-dashboard/e2e_tests/tests/basic.spec.ts new file mode 100644 index 00000000000..9e990f60866 --- /dev/null +++ b/ui/litellm-dashboard/e2e_tests/tests/basic.spec.ts @@ -0,0 +1,10 @@ +import { test, expect } from "@playwright/test"; + +test("basic test to verify playwright setup", async ({ page }) => { + // Navigate to the base URL + await page.goto("/"); + + // Wait for the page to load and check if we can find some basic content + // This is a very basic test just to verify the setup works + await expect(page).toHaveTitle(/.*LiteLLM.*/); +}); diff --git a/ui/litellm-dashboard/package-lock.json b/ui/litellm-dashboard/package-lock.json index b65c35e78d8..cc21b516e24 100644 --- a/ui/litellm-dashboard/package-lock.json +++ b/ui/litellm-dashboard/package-lock.json @@ -39,6 +39,7 @@ "uuid": "^11.1.0" }, "devDependencies": { + "@playwright/test": "^1.57.0", "@tailwindcss/forms": "^0.5.7", "@testing-library/dom": "^10.4.1", "@testing-library/jest-dom": "^6.8.0", @@ -4927,6 +4928,22 @@ "node": ">=12.4.0" } }, + "node_modules/@playwright/test": { + "version": "1.57.0", + "resolved": "https://registry.npmjs.org/@playwright/test/-/test-1.57.0.tgz", + "integrity": "sha512-6TyEnHgd6SArQO8UO2OMTxshln3QMWBtPGrOCgs3wVEmQmwyuNtB10IZMfmYDE0riwNR1cu4q+pPcxMVtaG3TA==", + "dev": true, + "license": "Apache-2.0", + "dependencies": { + "playwright": "1.57.0" + }, + "bin": { + "playwright": "cli.js" + }, + "engines": { + "node": ">=18" + } + }, "node_modules/@pnpm/config.env-replace": { "version": "1.1.0", "resolved": "https://registry.npmjs.org/@pnpm/config.env-replace/-/config.env-replace-1.1.0.tgz", @@ -19219,6 +19236,53 @@ "pathe": "^2.0.3" } }, + "node_modules/playwright": { + "version": "1.57.0", + "resolved": "https://registry.npmjs.org/playwright/-/playwright-1.57.0.tgz", + "integrity": "sha512-ilYQj1s8sr2ppEJ2YVadYBN0Mb3mdo9J0wQ+UuDhzYqURwSoW4n1Xs5vs7ORwgDGmyEh33tRMeS8KhdkMoLXQw==", + "dev": true, + "license": "Apache-2.0", + "dependencies": { + "playwright-core": "1.57.0" + }, + "bin": { + "playwright": "cli.js" + }, + "engines": { + "node": ">=18" + }, + "optionalDependencies": { + "fsevents": "2.3.2" + } + }, + "node_modules/playwright-core": { + "version": "1.57.0", + "resolved": "https://registry.npmjs.org/playwright-core/-/playwright-core-1.57.0.tgz", + "integrity": "sha512-agTcKlMw/mjBWOnD6kFZttAAGHgi/Nw0CZ2o6JqWSbMlI219lAFLZZCyqByTsvVAJq5XA5H8cA6PrvBRpBWEuQ==", + "dev": true, + "license": "Apache-2.0", + "bin": { + "playwright-core": "cli.js" + }, + "engines": { + "node": ">=18" + } + }, + "node_modules/playwright/node_modules/fsevents": { + "version": "2.3.2", + "resolved": "https://registry.npmjs.org/fsevents/-/fsevents-2.3.2.tgz", + "integrity": "sha512-xiqMQR4xAeHTuB9uWm+fFRcIOgKBMiOBP+eXiyT7jsgVCq1bkVygt00oASowB7EdtpOHaaPgKt812P9ab+DDKA==", + "dev": true, + "hasInstallScript": true, + "license": "MIT", + "optional": true, + "os": [ + "darwin" + ], + "engines": { + "node": "^8.16.0 || ^10.6.0 || >=11.0.0" + } + }, "node_modules/points-on-curve": { "version": "0.2.0", "resolved": "https://registry.npmjs.org/points-on-curve/-/points-on-curve-0.2.0.tgz", diff --git a/ui/litellm-dashboard/package.json b/ui/litellm-dashboard/package.json index ce42d0ba41a..a018a4e7a58 100644 --- a/ui/litellm-dashboard/package.json +++ b/ui/litellm-dashboard/package.json @@ -10,7 +10,9 @@ "test": "vitest", "test:watch": "vitest -w", "format": "prettier --write .", - "format:check": "prettier --check ." + "format:check": "prettier --check .", + "e2e": "playwright test --config e2e_tests/playwright.config.ts", + "e2e:ui": "playwright test --ui --config e2e_tests/playwright.config.ts" }, "dependencies": { "@anthropic-ai/sdk": "^0.54.0", @@ -44,6 +46,7 @@ "uuid": "^11.1.0" }, "devDependencies": { + "@playwright/test": "^1.57.0", "@tailwindcss/forms": "^0.5.7", "@testing-library/dom": "^10.4.1", "@testing-library/jest-dom": "^6.8.0", diff --git a/ui/litellm-dashboard/vitest.config.ts b/ui/litellm-dashboard/vitest.config.ts index 194b6fa691f..51e3df7a280 100644 --- a/ui/litellm-dashboard/vitest.config.ts +++ b/ui/litellm-dashboard/vitest.config.ts @@ -8,6 +8,8 @@ export default defineConfig({ globals: true, css: true, // lets you import CSS/modules without extra mocks coverage: { reporter: ["text", "lcov"] }, + exclude: ["e2e_tests/**,", "node_modules/**"], + include: ["src/**/*.test.ts", "src/**/*.test.tsx", "tests/**/*.test.ts", "tests/**/*.test.tsx"], }, resolve: { alias: { From 41fa2c0812d8aa143c0bf61d610850c537879b97 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Mon, 29 Dec 2025 12:00:03 -0800 Subject: [PATCH 018/158] migrate e2e tests to ui directory --- .circleci/config.yml | 4 ++-- .../e2e_tests/auth/login.setup.ts | 20 +++++++++++++++++++ .../e2e_tests/playwright.config.ts | 4 ++-- .../e2e_tests/tests/basic.spec.ts | 10 ---------- ui/litellm-dashboard/storageState.json | 15 ++++++++++++++ 5 files changed, 39 insertions(+), 14 deletions(-) create mode 100644 ui/litellm-dashboard/e2e_tests/auth/login.setup.ts delete mode 100644 ui/litellm-dashboard/e2e_tests/tests/basic.spec.ts create mode 100644 ui/litellm-dashboard/storageState.json diff --git a/.circleci/config.yml b/.circleci/config.yml index e30dc02b2ab..b1bee9675f3 100644 --- a/.circleci/config.yml +++ b/.circleci/config.yml @@ -3515,7 +3515,7 @@ jobs: - run: name: Run Playwright Tests command: | - npx playwright test e2e_ui_tests/ --reporter=html --output=test-results + cd ui/litellm-dashboard && npm run e2e no_output_timeout: 120m - store_artifacts: path: test-results @@ -3973,4 +3973,4 @@ workflows: - proxy_pass_through_endpoint_tests - check_code_and_doc_quality - publish_proxy_extras - - guardrails_testing \ No newline at end of file + - guardrails_testing diff --git a/ui/litellm-dashboard/e2e_tests/auth/login.setup.ts b/ui/litellm-dashboard/e2e_tests/auth/login.setup.ts new file mode 100644 index 00000000000..712191f6843 --- /dev/null +++ b/ui/litellm-dashboard/e2e_tests/auth/login.setup.ts @@ -0,0 +1,20 @@ +import { test, expect } from "@playwright/test"; + +test("login and save auth state", async ({ page }) => { + await page.goto("http://localhost:4000/ui"); + + await page.getByPlaceholder("Enter your username").fill("admin"); + await page.getByPlaceholder("Enter your password").fill("gm"); + + const loginButton = page.getByRole("button", { name: "Login" }); + await expect(loginButton).toBeEnabled(); + await loginButton.click(); + + // Assert successful login (important) + await expect(page.getByText("AI Gateway")).toBeVisible(); + + // 🔐 Save auth state + await page.context().storageState({ + path: "storageState.json", + }); +}); diff --git a/ui/litellm-dashboard/e2e_tests/playwright.config.ts b/ui/litellm-dashboard/e2e_tests/playwright.config.ts index cce416075f4..5ab68b3534f 100644 --- a/ui/litellm-dashboard/e2e_tests/playwright.config.ts +++ b/ui/litellm-dashboard/e2e_tests/playwright.config.ts @@ -4,8 +4,8 @@ import { defineConfig, devices } from "@playwright/test"; * See https://playwright.dev/docs/test-configuration. */ export default defineConfig({ - testDir: "./tests", - testMatch: "**/*.spec.ts", + testDir: ".", + testMatch: ["**/*.spec.ts", "**/*.setup.ts"], testIgnore: ["**/*.test.*"], /* Run tests in files in parallel */ fullyParallel: true, diff --git a/ui/litellm-dashboard/e2e_tests/tests/basic.spec.ts b/ui/litellm-dashboard/e2e_tests/tests/basic.spec.ts deleted file mode 100644 index 9e990f60866..00000000000 --- a/ui/litellm-dashboard/e2e_tests/tests/basic.spec.ts +++ /dev/null @@ -1,10 +0,0 @@ -import { test, expect } from "@playwright/test"; - -test("basic test to verify playwright setup", async ({ page }) => { - // Navigate to the base URL - await page.goto("/"); - - // Wait for the page to load and check if we can find some basic content - // This is a very basic test just to verify the setup works - await expect(page).toHaveTitle(/.*LiteLLM.*/); -}); diff --git a/ui/litellm-dashboard/storageState.json b/ui/litellm-dashboard/storageState.json new file mode 100644 index 00000000000..f2f012b34f8 --- /dev/null +++ b/ui/litellm-dashboard/storageState.json @@ -0,0 +1,15 @@ +{ + "cookies": [ + { + "name": "token", + "value": "eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJ1c2VyX2lkIjoiZGVmYXVsdF91c2VyX2lkIiwia2V5Ijoic2stUkJVZGFncVZuY0VWMDNVdnNDZkstdyIsInVzZXJfZW1haWwiOm51bGwsInVzZXJfcm9sZSI6InByb3h5X2FkbWluIiwibG9naW5fbWV0aG9kIjoidXNlcm5hbWVfcGFzc3dvcmQiLCJwcmVtaXVtX3VzZXIiOnRydWUsImF1dGhfaGVhZGVyX25hbWUiOiJBdXRob3JpemF0aW9uIiwiZGlzYWJsZWRfbm9uX2FkbWluX3BlcnNvbmFsX2tleV9jcmVhdGlvbiI6ZmFsc2UsInNlcnZlcl9yb290X3BhdGgiOiIvIn0.aW_1-qDI9HuZvDq53ufTX_HrE6EJ7XSCPBb-N136KjQ", + "domain": "localhost", + "path": "/", + "expires": -1, + "httpOnly": false, + "secure": false, + "sameSite": "Lax" + } + ], + "origins": [] +} \ No newline at end of file From ab89587acac69ea1108411a277778df9bfbd4f09 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Mon, 29 Dec 2025 12:49:28 -0800 Subject: [PATCH 019/158] circle ci config change for e2e --- .circleci/config.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.circleci/config.yml b/.circleci/config.yml index b1bee9675f3..723eba408c6 100644 --- a/.circleci/config.yml +++ b/.circleci/config.yml @@ -3515,7 +3515,7 @@ jobs: - run: name: Run Playwright Tests command: | - cd ui/litellm-dashboard && npm run e2e + cd ui/litellm-dashboard && npx playwright test --config=e2e_tests/playwright.config.ts no_output_timeout: 120m - store_artifacts: path: test-results From 657eb2b09b8b0042feecd182ac160fd3b6ae9827 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Mon, 29 Dec 2025 12:57:36 -0800 Subject: [PATCH 020/158] attempt fix 2 --- .circleci/config.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.circleci/config.yml b/.circleci/config.yml index 723eba408c6..f6100a0c432 100644 --- a/.circleci/config.yml +++ b/.circleci/config.yml @@ -3475,7 +3475,7 @@ jobs: - run: name: Install Playwright Browsers command: | - npx playwright install + cd ui/litellm-dashboard && npx playwright install - run: name: Build Docker image From 108fc4a59f991090755b55425c858b54b39b6c3d Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Mon, 29 Dec 2025 13:10:50 -0800 Subject: [PATCH 021/158] Using microsoft official docker for playwright --- .circleci/config.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.circleci/config.yml b/.circleci/config.yml index f6100a0c432..a777ec0255e 100644 --- a/.circleci/config.yml +++ b/.circleci/config.yml @@ -3423,7 +3423,7 @@ jobs: e2e_ui_testing: machine: - image: ubuntu-2204:2023.10.1 + image: mcr.microsoft.com/playwright:v1.57.0-noble resource_class: xlarge working_directory: ~/project steps: From 990b7ff93d8d00f5dcbde8bd774d7beadc3965a2 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Mon, 29 Dec 2025 13:18:48 -0800 Subject: [PATCH 022/158] Fixing CI/CD --- .circleci/config.yml | 10 ++-------- 1 file changed, 2 insertions(+), 8 deletions(-) diff --git a/.circleci/config.yml b/.circleci/config.yml index a777ec0255e..c18f5cd128a 100644 --- a/.circleci/config.yml +++ b/.circleci/config.yml @@ -3422,8 +3422,8 @@ jobs: --coverage.reportsDirectory=coverage/html e2e_ui_testing: - machine: - image: mcr.microsoft.com/playwright:v1.57.0-noble + docker: + - image: mcr.microsoft.com/playwright:v1.57.0-noble resource_class: xlarge working_directory: ~/project steps: @@ -3431,12 +3431,6 @@ jobs: - setup_google_dns - attach_workspace: at: ~/project - - run: - name: Upgrade Docker to v24.x (API 1.44+) - command: | - curl -fsSL https://get.docker.com | sh - sudo usermod -aG docker $USER - docker version - run: name: Install Python 3.9 command: | From ec29fe14720eae7d4dc853fc4777ca66bacc544c Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Mon, 29 Dec 2025 13:27:00 -0800 Subject: [PATCH 023/158] remove google dns for e2e --- .circleci/config.yml | 1 - 1 file changed, 1 deletion(-) diff --git a/.circleci/config.yml b/.circleci/config.yml index c18f5cd128a..4caa5189cb0 100644 --- a/.circleci/config.yml +++ b/.circleci/config.yml @@ -3428,7 +3428,6 @@ jobs: working_directory: ~/project steps: - checkout - - setup_google_dns - attach_workspace: at: ~/project - run: From f8d90939fe76662aae002302b14dd0d3ae2a49ec Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Mon, 29 Dec 2025 13:40:36 -0800 Subject: [PATCH 024/158] Attempt fix --- .circleci/config.yml | 8 ++------ 1 file changed, 2 insertions(+), 6 deletions(-) diff --git a/.circleci/config.yml b/.circleci/config.yml index 4caa5189cb0..6d9a9ff5bee 100644 --- a/.circleci/config.yml +++ b/.circleci/config.yml @@ -3433,14 +3433,10 @@ jobs: - run: name: Install Python 3.9 command: | - curl https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh --output miniconda.sh bash miniconda.sh -b -p $HOME/miniconda export PATH="$HOME/miniconda/bin:$PATH" - conda init bash - source ~/.bashrc - conda create -n myenv python=3.9 -y - conda activate myenv - python --version + $HOME/miniconda/bin/conda create -y -n myenv python=3.9 + conda run -n myenv python --version - run: name: Install Dependencies command: | From de551889f3459b376946ae46aa948dee2d65c42b Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Mon, 29 Dec 2025 14:20:04 -0800 Subject: [PATCH 025/158] Add optional field expand to /key/list --- litellm/proxy/_types.py | 1 + .../key_management_endpoints.py | 26 +++- .../test_key_management_endpoints.py | 112 ++++++++++++++++++ 3 files changed, 138 insertions(+), 1 deletion(-) diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 570a6bb9f3b..a94b4d4a077 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -2152,6 +2152,7 @@ class UserAPIKeyAuth( user_rpm_limit: Optional[int] = None user_email: Optional[str] = None request_route: Optional[str] = None + user: Optional[Any] = None # Expanded user object when expand=user is used model_config = ConfigDict(arbitrary_types_allowed=True) diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index cc2ac908149..65b4ec11f4f 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -3020,10 +3020,14 @@ async def list_keys( description="Column to sort by (e.g. 'user_id', 'created_at', 'spend')", ), sort_order: str = Query(default="desc", description="Sort order ('asc' or 'desc')"), + expand: Optional[List[str]] = Query(None, description="Expand related objects (e.g. 'user')"), ) -> KeyListResponseObject: """ List all keys for a given user / team / organization. + Parameters: + expand: Optional[List[str]] - Expand related objects (e.g. 'user' to include user information) + Returns: { "keys": List[str] or List[UserAPIKeyAuth], @@ -3031,6 +3035,9 @@ async def list_keys( "current_page": int, "total_pages": int, } + + When expand includes "user", each key object will include a "user" field with the associated user object. + Note: When expand=user is specified, full key objects are returned regardless of the return_full_object parameter. """ try: from litellm.proxy.proxy_server import prisma_client @@ -3080,6 +3087,7 @@ async def list_keys( include_created_by_keys=include_created_by_keys, sort_by=sort_by, sort_order=sort_order, + expand=expand, ) verbose_proxy_logger.debug("Successfully prepared response") @@ -3232,6 +3240,7 @@ async def _list_key_helper( include_created_by_keys: bool = False, sort_by: Optional[str] = None, sort_order: str = "desc", + expand: Optional[List[str]] = None, ) -> KeyListResponseObject: """ Helper function to list keys @@ -3334,13 +3343,28 @@ async def _list_key_helper( # Calculate total pages total_pages = -(-total_count // size) # Ceiling division + # Fetch user information if expand includes "user" + user_map = {} + if expand and "user" in expand: + user_ids = [key.user_id for key in keys if key.user_id] + if user_ids: + users = await prisma_client.db.litellm_usertable.find_many( + where={"user_id": {"in": list(set(user_ids))}} # Remove duplicates + ) + user_map = {user.user_id: user for user in users} + # Prepare response key_list: List[Union[str, UserAPIKeyAuth]] = [] for key in keys: key_dict = key.dict() # Attach object_permission if object_permission_id is set key_dict = await attach_object_permission_to_dict(key_dict, prisma_client) - if return_full_object is True: + + # Include user information if expand includes "user" + if expand and "user" in expand and key.user_id and key.user_id in user_map: + key_dict["user"] = user_map[key.user_id].dict() + + if return_full_object is True or (expand and "user" in expand): key_list.append(UserAPIKeyAuth(**key_dict)) # Return full key object else: _token = key_dict.get("token") diff --git a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py index 648045a7ea6..cba06e7fb3f 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py @@ -3405,3 +3405,115 @@ async def test_can_modify_verification_token_personal_key_no_user_id(monkeypatch ) assert result is False + + +@pytest.mark.asyncio +async def test_list_keys_with_expand_user(): + """ + Test that expand=user parameter correctly includes user information in the response. + """ + mock_prisma_client = AsyncMock() + + # Create mock keys with user_ids + mock_key1 = MagicMock() + mock_key1.token = "token1" + mock_key1.user_id = "user123" + mock_key1.dict.return_value = { + "token": "token1", + "user_id": "user123", + "key_alias": "key1", + "models": ["gpt-4"], + } + + mock_key2 = MagicMock() + mock_key2.token = "token2" + mock_key2.user_id = "user456" + mock_key2.dict.return_value = { + "token": "token2", + "user_id": "user456", + "key_alias": "key2", + "models": ["gpt-3.5-turbo"], + } + + mock_find_many_keys = AsyncMock(return_value=[mock_key1, mock_key2]) + mock_count_keys = AsyncMock(return_value=2) + + # Create mock users + mock_user1 = MagicMock() + mock_user1.user_id = "user123" + mock_user1.user_email = "user1@example.com" + mock_user1.dict.return_value = { + "user_id": "user123", + "user_email": "user1@example.com", + "user_alias": "User One", + } + + mock_user2 = MagicMock() + mock_user2.user_id = "user456" + mock_user2.user_email = "user2@example.com" + mock_user2.dict.return_value = { + "user_id": "user456", + "user_email": "user2@example.com", + "user_alias": "User Two", + } + + mock_find_many_users = AsyncMock(return_value=[mock_user1, mock_user2]) + + mock_prisma_client.db.litellm_verificationtoken.find_many = mock_find_many_keys + mock_prisma_client.db.litellm_verificationtoken.count = mock_count_keys + mock_prisma_client.db.litellm_usertable.find_many = mock_find_many_users + + args = { + "prisma_client": mock_prisma_client, + "page": 1, + "size": 50, + "user_id": None, + "team_id": None, + "organization_id": None, + "key_alias": None, + "key_hash": None, + "exclude_team_id": None, + "return_full_object": False, # This should be overridden by expand=user + "admin_team_ids": None, + "include_created_by_keys": False, + "expand": ["user"], # Test the expand parameter + } + + result = await _list_key_helper(**args) + + # Verify that keys were fetched + mock_find_many_keys.assert_called_once() + mock_count_keys.assert_called_once() + + # Verify that users were fetched + # Note: Order doesn't matter for the 'in' query, so we just check that both user_ids are present + call_args = mock_find_many_users.call_args + assert call_args is not None + where_clause = call_args.kwargs["where"] + assert "user_id" in where_clause + assert "in" in where_clause["user_id"] + user_ids_in_query = set(where_clause["user_id"]["in"]) + assert user_ids_in_query == {"user123", "user456"} + + # Verify response structure + assert len(result["keys"]) == 2 + assert result["total_count"] == 2 + assert result["current_page"] == 1 + assert result["total_pages"] == 1 + + # Verify that user data is included in the response + # Since expand=user is specified, keys should be full objects + assert isinstance(result["keys"][0], UserAPIKeyAuth) + assert isinstance(result["keys"][1], UserAPIKeyAuth) + + # Verify user data is attached to keys + assert result["keys"][0].user == { + "user_id": "user123", + "user_email": "user1@example.com", + "user_alias": "User One", + } + assert result["keys"][1].user == { + "user_id": "user456", + "user_email": "user2@example.com", + "user_alias": "User Two", + } From 03cee685881025be13584ca4729725b8985dde1c Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Mon, 29 Dec 2025 14:27:55 -0800 Subject: [PATCH 026/158] Ruff check --- .../key_management_endpoints.py | 94 ++++++++++++------- 1 file changed, 59 insertions(+), 35 deletions(-) diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index 65b4ec11f4f..4f9534521c3 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -3223,46 +3223,17 @@ def _validate_sort_params( return order_by -async def _list_key_helper( - prisma_client: PrismaClient, - page: int, - size: int, +def _build_key_filter_conditions( user_id: Optional[str], team_id: Optional[str], organization_id: Optional[str], key_alias: Optional[str], key_hash: Optional[str], - exclude_team_id: Optional[str] = None, - return_full_object: bool = False, - admin_team_ids: Optional[ - List[str] - ] = None, # New parameter for teams where user is admin - include_created_by_keys: bool = False, - sort_by: Optional[str] = None, - sort_order: str = "desc", - expand: Optional[List[str]] = None, -) -> KeyListResponseObject: - """ - Helper function to list keys - Args: - page: int - size: int - user_id: Optional[str] - team_id: Optional[str] - key_alias: Optional[str] - exclude_team_id: Optional[str] # exclude a specific team_id - return_full_object: bool # when true, will return UserAPIKeyAuth objects instead of just the token - admin_team_ids: Optional[List[str]] # list of team IDs where the user is an admin - - Returns: - KeyListResponseObject - { - "keys": List[str] or List[UserAPIKeyAuth], # Updated to reflect possible return types - "total_count": int, - "current_page": int, - "total_pages": int, - } - """ + exclude_team_id: Optional[str], + admin_team_ids: Optional[List[str]], + include_created_by_keys: bool, +) -> Dict[str, Union[str, Dict[str, Any], List[Dict[str, Any]]]]: + """Build filter conditions for key listing.""" # Prepare filter conditions where: Dict[str, Union[str, Dict[str, Any], List[Dict[str, Any]]]] = {} where.update(_get_condition_to_filter_out_ui_session_tokens()) @@ -3303,6 +3274,59 @@ async def _list_key_helper( where.update(or_conditions[0]) verbose_proxy_logger.debug(f"Filter conditions: {where}") + return where + + +async def _list_key_helper( + prisma_client: PrismaClient, + page: int, + size: int, + user_id: Optional[str], + team_id: Optional[str], + organization_id: Optional[str], + key_alias: Optional[str], + key_hash: Optional[str], + exclude_team_id: Optional[str] = None, + return_full_object: bool = False, + admin_team_ids: Optional[ + List[str] + ] = None, # New parameter for teams where user is admin + include_created_by_keys: bool = False, + sort_by: Optional[str] = None, + sort_order: str = "desc", + expand: Optional[List[str]] = None, +) -> KeyListResponseObject: + """ + Helper function to list keys + Args: + page: int + size: int + user_id: Optional[str] + team_id: Optional[str] + key_alias: Optional[str] + exclude_team_id: Optional[str] # exclude a specific team_id + return_full_object: bool # when true, will return UserAPIKeyAuth objects instead of just the token + admin_team_ids: Optional[List[str]] # list of team IDs where the user is an admin + + Returns: + KeyListResponseObject + { + "keys": List[str] or List[UserAPIKeyAuth], # Updated to reflect possible return types + "total_count": int, + "current_page": int, + "total_pages": int, + } + """ + where = _build_key_filter_conditions( + user_id=user_id, + team_id=team_id, + organization_id=organization_id, + key_alias=key_alias, + key_hash=key_hash, + exclude_team_id=exclude_team_id, + admin_team_ids=admin_team_ids, + include_created_by_keys=include_created_by_keys, + ) # Calculate skip for pagination skip = (page - 1) * size From b69c1dcee7e989574905ca1b4181e219ff0fe622 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Mon, 29 Dec 2025 15:09:33 -0800 Subject: [PATCH 027/158] Key List table uses expand from /key/list --- .../src/components/all_keys_table.tsx | 79 ++++++++++--------- .../components/key_team_helpers/key_list.tsx | 20 ++++- .../src/components/networking.tsx | 6 ++ .../components/templates/view_key_table.tsx | 22 +++--- 4 files changed, 75 insertions(+), 52 deletions(-) diff --git a/ui/litellm-dashboard/src/components/all_keys_table.tsx b/ui/litellm-dashboard/src/components/all_keys_table.tsx index 210ce09fa34..10f678be03f 100644 --- a/ui/litellm-dashboard/src/components/all_keys_table.tsx +++ b/ui/litellm-dashboard/src/components/all_keys_table.tsx @@ -1,23 +1,39 @@ "use client"; -import React, { useEffect, useState } from "react"; -import { ColumnDef, ColumnResizeMode, ColumnResizeDirection } from "@tanstack/react-table"; -import { Select, SelectItem } from "@tremor/react"; -import { Button } from "@tremor/react"; -import KeyInfoView from "./templates/key_info_view"; -import { Tooltip } from "antd"; -import { Team, KeyResponse } from "./key_team_helpers/key_list"; -import FilterComponent from "./molecules/filter"; -import { FilterOption } from "./molecules/filter"; -import { Organization, userListCall } from "./networking"; -import { useFilterLogic } from "./key_team_helpers/filter_logic"; import { Setter } from "@/types"; -import { updateExistingKeys } from "@/utils/dataUtils"; -import { flexRender, getCoreRowModel, getSortedRowModel, SortingState, useReactTable } from "@tanstack/react-table"; -import { Table, TableHead, TableHeaderCell, TableBody, TableRow, TableCell, Icon } from "@tremor/react"; -import { SwitchVerticalIcon, ChevronUpIcon, ChevronDownIcon, ChevronRightIcon } from "@heroicons/react/outline"; -import { Badge, Text } from "@tremor/react"; +import { formatNumberWithCommas, updateExistingKeys } from "@/utils/dataUtils"; +import { ChevronDownIcon, ChevronRightIcon, ChevronUpIcon, SwitchVerticalIcon } from "@heroicons/react/outline"; +import { + ColumnDef, + ColumnResizeDirection, + ColumnResizeMode, + flexRender, + getCoreRowModel, + getSortedRowModel, + SortingState, + useReactTable, +} from "@tanstack/react-table"; +import { + Badge, + Button, + Icon, + Select, + SelectItem, + Table, + TableBody, + TableCell, + TableHead, + TableHeaderCell, + TableRow, + Text, +} from "@tremor/react"; +import { Tooltip } from "antd"; +import React, { useEffect, useState } from "react"; import { getModelDisplayName } from "./key_team_helpers/fetch_available_models_team_key"; -import { formatNumberWithCommas } from "@/utils/dataUtils"; +import { useFilterLogic } from "./key_team_helpers/filter_logic"; +import { KeyResponse, Team } from "./key_team_helpers/key_list"; +import FilterComponent, { FilterOption } from "./molecules/filter"; +import { Organization } from "./networking"; +import KeyInfoView from "./templates/key_info_view"; interface AllKeysTableProps { keys: KeyResponse[]; @@ -124,7 +140,6 @@ export function AllKeysTable({ setAccessToken, }: AllKeysTableProps) { const [selectedKeyId, setSelectedKeyId] = useState(null); - const [userList, setUserList] = useState([]); const [columnResizeMode, setColumnResizeMode] = React.useState("onChange"); const [columnResizeDirection, setColumnResizeDirection] = React.useState("ltr"); const [sorting, setSorting] = React.useState(() => { @@ -155,17 +170,6 @@ export function AllKeysTable({ accessToken, }); - useEffect(() => { - if (accessToken) { - const user_IDs = keys.map((key) => key.user_id).filter((id) => id !== null); - const fetchUserList = async () => { - const userListData = await userListCall(accessToken, user_IDs, 1, 100); - setUserList(userListData.users); - }; - fetchUserList(); - } - }, [accessToken, keys]); - // Add a useEffect to call refresh when a key is created useEffect(() => { if (refresh) { @@ -269,18 +273,19 @@ export function AllKeysTable({ }, { id: "user_email", - accessorKey: "user_id", + accessorKey: "user", header: "User Email", size: 160, cell: (info) => { - const userId = info.getValue() as string; - const user = userList.find((u) => u.user_id === userId); - return user?.user_email ? ( - - {user?.user_email.slice(0, 20)}... + const user = info.getValue() as any; + const value = user?.user_email; + const width = info.cell.column.getSize(); + return ( + + + {value ?? "-"} + - ) : ( - "-" ); }, }, 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 6bb014e6187..9f95ba46cc1 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,6 @@ -import { useState, useEffect } from "react"; -import { keyListCall, Member, Organization } from "../networking"; import { Setter } from "@/types"; +import { useEffect, useState } from "react"; +import { keyListCall, Member, Organization } from "../networking"; export interface Team { team_id: string; @@ -106,6 +106,7 @@ interface UseKeyListProps { selectedKeyAlias: string | null; accessToken: string; createClicked: boolean; + expand?: string[]; } interface PaginationData { @@ -129,6 +130,7 @@ const useKeyList = ({ selectedKeyAlias, accessToken, createClicked, + expand = [], }: UseKeyListProps): UseKeyListReturn => { const [keyData, setKeyData] = useState({ keys: [], @@ -151,7 +153,19 @@ const useKeyList = ({ const page = typeof params.page === "number" ? params.page : 1; const pageSize = typeof params.pageSize === "number" ? params.pageSize : 100; - const data = await keyListCall(accessToken, null, null, null, null, null, page, pageSize); + const data = await keyListCall( + accessToken, + null, + null, + null, + null, + null, + page, + pageSize, + null, + null, + expand.join(","), + ); console.log("data", data); setKeyData(data); setError(null); diff --git a/ui/litellm-dashboard/src/components/networking.tsx b/ui/litellm-dashboard/src/components/networking.tsx index aeba207db26..987aaf96cb3 100644 --- a/ui/litellm-dashboard/src/components/networking.tsx +++ b/ui/litellm-dashboard/src/components/networking.tsx @@ -3250,6 +3250,7 @@ export const keyListCall = async ( pageSize: number, sortBy: string | null = null, sortOrder: string | null = null, + expand: string | null = null, ) => { /** * Get all available teams on proxy @@ -3294,6 +3295,11 @@ export const keyListCall = async ( if (sortOrder) { queryParams.append("sort_order", sortOrder); } + + if (expand) { + queryParams.append("expand", expand); + } + queryParams.append("return_full_object", "true"); queryParams.append("include_team_keys", "true"); queryParams.append("include_created_by_keys", "true"); diff --git a/ui/litellm-dashboard/src/components/templates/view_key_table.tsx b/ui/litellm-dashboard/src/components/templates/view_key_table.tsx index 5defc287a07..f8a86f0fad4 100644 --- a/ui/litellm-dashboard/src/components/templates/view_key_table.tsx +++ b/ui/litellm-dashboard/src/components/templates/view_key_table.tsx @@ -1,17 +1,14 @@ "use client"; -import React, { useEffect, useState } from "react"; -import { keyDeleteCall, Organization } from "../networking"; -import { add } from "date-fns"; -import { regenerateKeyCall } from "../networking"; -import { Grid, Col, Button, Text, Title, TextInput } from "@tremor/react"; -import { fetchAvailableModelsForTeamOrKey } from "../key_team_helpers/fetch_available_models_team_key"; -import { Modal, Form, InputNumber } from "antd"; -import { CopyToClipboard } from "react-copy-to-clipboard"; -import useKeyList from "../key_team_helpers/key_list"; -import { KeyResponse } from "../key_team_helpers/key_list"; -import { AllKeysTable } from "../all_keys_table"; -import { Team } from "../key_team_helpers/key_list"; import { Setter } from "@/types"; +import { Button, Col, Grid, Text, TextInput, Title } from "@tremor/react"; +import { Form, InputNumber, Modal } from "antd"; +import { add } from "date-fns"; +import React, { useEffect, useState } from "react"; +import { CopyToClipboard } from "react-copy-to-clipboard"; +import { AllKeysTable } from "../all_keys_table"; +import { fetchAvailableModelsForTeamOrKey } from "../key_team_helpers/fetch_available_models_team_key"; +import useKeyList, { KeyResponse, Team } from "../key_team_helpers/key_list"; +import { keyDeleteCall, Organization, regenerateKeyCall } from "../networking"; import NotificationManager from "../molecules/notifications_manager"; @@ -125,6 +122,7 @@ const ViewKeyTable: React.FC = ({ selectedKeyAlias, accessToken: accessToken || "", createClicked, + expand: ["user"], }); const handlePageChange = (newPage: number) => { From b21b156e5567ee2eb943f408b6f496e874397720 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Mon, 29 Dec 2025 15:21:19 -0800 Subject: [PATCH 028/158] adding tests --- .../src/components/all_keys_table.test.tsx | 36 +++++++++++++++++++ .../components/key_team_helpers/key_list.tsx | 4 +++ 2 files changed, 40 insertions(+) diff --git a/ui/litellm-dashboard/src/components/all_keys_table.test.tsx b/ui/litellm-dashboard/src/components/all_keys_table.test.tsx index 201a4be75bd..1458143d77e 100644 --- a/ui/litellm-dashboard/src/components/all_keys_table.test.tsx +++ b/ui/litellm-dashboard/src/components/all_keys_table.test.tsx @@ -93,6 +93,10 @@ const mockKey: KeyResponse = { user_tpm_limit: 1000, user_rpm_limit: 100, user_email: "user@example.com", + user: { + user_email: "user@example.com", + user_id: "user-1", + }, }; const mockTeam: Team = { @@ -190,3 +194,35 @@ it("should display key information correctly", async () => { expect(screen.getByText("5.5000")).toBeInTheDocument(); }); }); + +it("should display user email correctly", async () => { + const mockProps = { + keys: [mockKey], + setKeys: vi.fn(), + isLoading: false, + pagination: { + currentPage: 1, + totalPages: 1, + totalCount: 1, + }, + onPageChange: vi.fn(), + pageSize: 50, + teams: [mockTeam], + selectedTeam: null, + setSelectedTeam: vi.fn(), + selectedKeyAlias: null, + setSelectedKeyAlias: vi.fn(), + accessToken: "test-token", + userID: "user-1", + userRole: "admin", + organizations: [mockOrganization], + setCurrentOrg: vi.fn(), + premiumUser: false, + }; + + renderWithProviders(); + + await waitFor(() => { + expect(screen.getByText("user@example.com")).toBeInTheDocument(); + }); +}); 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 9f95ba46cc1..a04fbf3943d 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 @@ -91,6 +91,10 @@ export interface KeyResponse { last_rotation_at?: string; key_rotation_at?: string; next_rotation_at?: string; + user?: { + user_id: string; + user_email: string; + }; } interface KeyListResponse { From e0b96c4265ba4eb4640a947ac844bd0aff3d7f4d Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Mon, 29 Dec 2025 17:03:11 -0800 Subject: [PATCH 029/158] refactor keys table --- .../VirtualKeysTable.test.tsx} | 18 +-- .../VirtualKeysTable.tsx} | 102 +++------------ .../key_team_helpers/filter_logic.tsx | 4 +- .../organisms/create_key_button.tsx | 2 +- .../organisms/regenerate_key_modal.tsx | 33 ++--- .../KeyInfoView.handleKeyUpdate.test.tsx | 51 ++++---- .../templates/key_info_view.test.tsx | 119 +++++++++--------- .../components/templates/key_info_view.tsx | 16 +-- .../components/templates/view_key_table.tsx | 113 +---------------- 9 files changed, 122 insertions(+), 336 deletions(-) rename ui/litellm-dashboard/src/components/{all_keys_table.test.tsx => VirtualKeysPage/VirtualKeysTable.test.tsx} (90%) rename ui/litellm-dashboard/src/components/{all_keys_table.tsx => VirtualKeysPage/VirtualKeysTable.tsx} (89%) diff --git a/ui/litellm-dashboard/src/components/all_keys_table.test.tsx b/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.test.tsx similarity index 90% rename from ui/litellm-dashboard/src/components/all_keys_table.test.tsx rename to ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.test.tsx index 1458143d77e..4f40ed99a84 100644 --- a/ui/litellm-dashboard/src/components/all_keys_table.test.tsx +++ b/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.test.tsx @@ -1,13 +1,13 @@ import { screen, waitFor } from "@testing-library/react"; import { vi, it, expect } from "vitest"; -import { renderWithProviders } from "../../tests/test-utils"; -import { AllKeysTable } from "./all_keys_table"; -import { KeyResponse, Team } from "./key_team_helpers/key_list"; -import { Organization } from "./networking"; +import { renderWithProviders } from "../../../tests/test-utils"; +import { VirtualKeysTable } from "./VirtualKeysTable"; +import { KeyResponse, Team } from "../key_team_helpers/key_list"; +import { Organization } from "../networking"; // Mock network calls vi.mock("./networking", async (importOriginal) => { - const actual = await importOriginal(); + const actual = await importOriginal(); return { ...actual, userListCall: vi.fn().mockResolvedValue({ @@ -131,7 +131,7 @@ const mockOrganization: Organization = { members: [], }; -it("should render AllKeysTable component", () => { +it("should render VirtualKeysTable component", () => { const mockProps = { keys: [mockKey], setKeys: vi.fn(), @@ -156,7 +156,7 @@ it("should render AllKeysTable component", () => { premiumUser: false, }; - renderWithProviders(); + renderWithProviders(); expect(screen.getByText("Test Key Alias")).toBeInTheDocument(); }); @@ -186,7 +186,7 @@ it("should display key information correctly", async () => { premiumUser: false, }; - renderWithProviders(); + renderWithProviders(); await waitFor(() => { expect(screen.getByText("Test Key Alias")).toBeInTheDocument(); @@ -220,7 +220,7 @@ it("should display user email correctly", async () => { premiumUser: false, }; - renderWithProviders(); + renderWithProviders(); await waitFor(() => { expect(screen.getByText("user@example.com")).toBeInTheDocument(); diff --git a/ui/litellm-dashboard/src/components/all_keys_table.tsx b/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.tsx similarity index 89% rename from ui/litellm-dashboard/src/components/all_keys_table.tsx rename to ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.tsx index 10f678be03f..128129ce66a 100644 --- a/ui/litellm-dashboard/src/components/all_keys_table.tsx +++ b/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.tsx @@ -4,8 +4,6 @@ import { formatNumberWithCommas, updateExistingKeys } from "@/utils/dataUtils"; import { ChevronDownIcon, ChevronRightIcon, ChevronUpIcon, SwitchVerticalIcon } from "@heroicons/react/outline"; import { ColumnDef, - ColumnResizeDirection, - ColumnResizeMode, flexRender, getCoreRowModel, getSortedRowModel, @@ -16,8 +14,6 @@ import { Badge, Button, Icon, - Select, - SelectItem, Table, TableBody, TableCell, @@ -28,14 +24,15 @@ import { } from "@tremor/react"; import { Tooltip } from "antd"; import React, { useEffect, useState } from "react"; -import { getModelDisplayName } from "./key_team_helpers/fetch_available_models_team_key"; -import { useFilterLogic } from "./key_team_helpers/filter_logic"; -import { KeyResponse, Team } from "./key_team_helpers/key_list"; -import FilterComponent, { FilterOption } from "./molecules/filter"; -import { Organization } from "./networking"; -import KeyInfoView from "./templates/key_info_view"; +import { getModelDisplayName } from "../key_team_helpers/fetch_available_models_team_key"; +import { useFilterLogic } from "../key_team_helpers/filter_logic"; +import { KeyResponse, Team } from "../key_team_helpers/key_list"; +import FilterComponent, { FilterOption } from "../molecules/filter"; +import { Organization } from "../networking"; +import KeyInfoView from "../templates/key_info_view"; +import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; -interface AllKeysTableProps { +interface VirtualKeysTableProps { keys: KeyResponse[]; setKeys: (keys: KeyResponse[] | ((prev: KeyResponse[]) => KeyResponse[])) => void; isLoading?: boolean; @@ -51,9 +48,6 @@ interface AllKeysTableProps { setSelectedTeam: (team: Team | null) => void; selectedKeyAlias: string | null; setSelectedKeyAlias: Setter; - accessToken: string | null; - userID: string | null; - userRole: string | null; organizations: Organization[] | null; setCurrentOrg: React.Dispatch>; refresh?: () => void; @@ -62,61 +56,14 @@ interface AllKeysTableProps { sortBy: string; sortOrder: "asc" | "desc"; }; - premiumUser: boolean; - setAccessToken?: (token: string) => void; } -// Define columns similar to our logs table - -interface UserResponse { - user_id: string; - user_email: string; - user_role: string; -} - -const TeamFilter = ({ - teams, - selectedTeam, - setSelectedTeam, -}: { - teams: Team[] | null; - selectedTeam: Team | null; - setSelectedTeam: (team: Team | null) => void; -}) => { - const handleTeamChange = (value: string) => { - const team = teams?.find((t) => t.team_id === value); - setSelectedTeam(team || null); - }; - - return ( -
-
- Where Team is - -
-
- ); -}; - /** - * AllKeysTable – a new table for keys that mimics the table styling used in view_logs. + * VirtualKeysTable – a new table for keys that mimics the table styling used in view_logs. * The team selector and filtering have been removed so that all keys are shown. */ -export function AllKeysTable({ +export function VirtualKeysTable({ keys, setKeys, isLoading = false, @@ -124,24 +71,13 @@ export function AllKeysTable({ onPageChange, pageSize = 50, teams, - selectedTeam, - setSelectedTeam, - selectedKeyAlias, - setSelectedKeyAlias, - accessToken, - userID, - userRole, organizations, - setCurrentOrg, refresh, onSortChange, currentSort, - premiumUser, - setAccessToken, -}: AllKeysTableProps) { +}: VirtualKeysTableProps) { + const { accessToken, userId: userID, userRole, premiumUser } = useAuthorized(); const [selectedKeyId, setSelectedKeyId] = useState(null); - const [columnResizeMode, setColumnResizeMode] = React.useState("onChange"); - const [columnResizeDirection, setColumnResizeDirection] = React.useState("ltr"); const [sorting, setSorting] = React.useState(() => { if (currentSort) { return [ @@ -167,7 +103,6 @@ export function AllKeysTable({ keys, teams, organizations, - accessToken, }); // Add a useEffect to call refresh when a key is created @@ -557,8 +492,8 @@ export function AllKeysTable({ const table = useReactTable({ data: filteredKeys, columns: columns.filter((col) => col.id !== "expander"), - columnResizeMode, - columnResizeDirection, + columnResizeMode: "onChange", + columnResizeDirection: "ltr", state: { sorting, }, @@ -619,12 +554,7 @@ export function AllKeysTable({ setKeys((keys) => keys.filter((key) => key.token !== selectedKeyId)); if (refresh) refresh(); // Minimal fix: refresh the full key list after a delete }} - accessToken={accessToken} - userID={userID} - userRole={userRole} teams={allTeams} - premiumUser={premiumUser} - setAccessToken={setAccessToken} /> ) : (
@@ -736,10 +666,6 @@ export function AllKeysTable({ userSelect: "none", touchAction: "none", opacity: header.column.getIsResizing() ? 1 : 0, - transform: - columnResizeMode === "onEnd" && header.column.getIsResizing() - ? `translateX(${(table.options.columnResizeDirection === "rtl" ? -1 : 1) * (table.getState().columnSizingInfo.deltaOffset ?? 0)}px)` - : "", }} />
diff --git a/ui/litellm-dashboard/src/components/key_team_helpers/filter_logic.tsx b/ui/litellm-dashboard/src/components/key_team_helpers/filter_logic.tsx index a074a5484f6..5428efb28aa 100644 --- a/ui/litellm-dashboard/src/components/key_team_helpers/filter_logic.tsx +++ b/ui/litellm-dashboard/src/components/key_team_helpers/filter_logic.tsx @@ -6,6 +6,7 @@ import { useQuery } from "@tanstack/react-query"; import { fetchAllKeyAliases, fetchAllOrganizations, fetchAllTeams } from "./filter_helpers"; import { debounce } from "lodash"; import { defaultPageSize } from "../constants"; +import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; export interface FilterState { "Team ID": string; @@ -21,12 +22,10 @@ export function useFilterLogic({ keys, teams, organizations, - accessToken, }: { keys: KeyResponse[]; teams: Team[] | null; organizations: Organization[] | null; - accessToken: string | null; }) { const defaultFilters: FilterState = { "Team ID": "", @@ -36,6 +35,7 @@ export function useFilterLogic({ "Sort By": "created_at", "Sort Order": "desc", }; + const { accessToken } = useAuthorized(); const [filters, setFilters] = useState(defaultFilters); const [allTeams, setAllTeams] = useState(teams || []); const [allOrganizations, setAllOrganizations] = useState(organizations || []); 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 f3ad803e27c..83ef444d59e 100644 --- a/ui/litellm-dashboard/src/components/organisms/create_key_button.tsx +++ b/ui/litellm-dashboard/src/components/organisms/create_key_button.tsx @@ -401,7 +401,7 @@ const CreateKey: React.FC = ({ console.log("key create Response:", response); // Add the data to the state in the parent component - // Also directly update the keys list in AllKeysTable without an API call + // Also directly update the keys list in VirtualKeysTable without an API call addKey(response); setApiKey(response["key"]); diff --git a/ui/litellm-dashboard/src/components/organisms/regenerate_key_modal.tsx b/ui/litellm-dashboard/src/components/organisms/regenerate_key_modal.tsx index 583a097449f..a4339e11920 100644 --- a/ui/litellm-dashboard/src/components/organisms/regenerate_key_modal.tsx +++ b/ui/litellm-dashboard/src/components/organisms/regenerate_key_modal.tsx @@ -1,31 +1,22 @@ -import React, { useEffect, useState } from "react"; -import { Button, Text, TextInput, Title, Grid, Col } from "@tremor/react"; -import { Modal, Form, InputNumber } from "antd"; +import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; +import { Button, Col, Grid, Text, TextInput, Title } from "@tremor/react"; +import { Form, InputNumber, Modal } from "antd"; import { add } from "date-fns"; -import { regenerateKeyCall } from "../networking"; -import { KeyResponse } from "../key_team_helpers/key_list"; +import { useEffect, useState } from "react"; import { CopyToClipboard } from "react-copy-to-clipboard"; +import { KeyResponse } from "../key_team_helpers/key_list"; import NotificationManager from "../molecules/notifications_manager"; +import { regenerateKeyCall } from "../networking"; interface RegenerateKeyModalProps { selectedToken: KeyResponse | null; visible: boolean; onClose: () => void; - accessToken: string | null; - premiumUser: boolean; - setAccessToken?: (token: string) => void; onKeyUpdate?: (updatedKeyData: Partial) => void; } -export function RegenerateKeyModal({ - selectedToken, - visible, - onClose, - accessToken, - premiumUser, - setAccessToken, - onKeyUpdate, -}: RegenerateKeyModalProps) { +export function RegenerateKeyModal({ selectedToken, visible, onClose, onKeyUpdate }: RegenerateKeyModalProps) { + const { accessToken } = useAuthorized(); const [form] = Form.useForm(); const [regeneratedKey, setRegeneratedKey] = useState(null); const [regenerateFormData, setRegenerateFormData] = useState(null); @@ -132,14 +123,6 @@ export function RegenerateKeyModal({ console.log("Updated key data with new token:", updatedKeyData); // Debug log - // If user regenerated their own auth key, update both local and global access tokens - if (isOwnKey) { - setCurrentAccessToken(response.key); // Update local token immediately - if (setAccessToken) { - setAccessToken(response.key); // Update global token - } - } - // Update the parent component with new key data if (onKeyUpdate) { onKeyUpdate(updatedKeyData); 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 a4680d7992e..c2c32390cd8 100644 --- a/ui/litellm-dashboard/src/components/templates/KeyInfoView.handleKeyUpdate.test.tsx +++ b/ui/litellm-dashboard/src/components/templates/KeyInfoView.handleKeyUpdate.test.tsx @@ -2,15 +2,21 @@ import { fireEvent, render, screen, waitFor } from "@testing-library/react"; import { beforeEach, describe, expect, it, vi } from "vitest"; // ---- Hoisted shared mocks (safe to use inside vi.mock factories) ---- -const { keyUpdateCallMock, keyDeleteCallMock } = vi.hoisted(() => { +const { keyUpdateCallMock, keyDeleteCallMock, mockUseAuthorized } = vi.hoisted(() => { return { keyUpdateCallMock: vi.fn().mockResolvedValue({}), keyDeleteCallMock: vi.fn().mockResolvedValue({}), + mockUseAuthorized: vi.fn(), }; }); // ---- Module mocks ---- +// Mock useAuthorized hook FIRST (before component imports it) +vi.mock("@/app/(dashboard)/hooks/useAuthorized", () => ({ + default: mockUseAuthorized, +})); + // Networking: wire the hoisted fns so we can assert calls later vi.mock("../networking", () => { return { @@ -29,12 +35,12 @@ vi.mock("../molecules/notifications_manager", () => { return { default: Notifications }; }); -// Roles: ensure 'admin' has write access and include all role helper functions +// Roles: ensure 'Admin' has write access and include all role helper functions vi.mock("../../utils/roles", async (importOriginal) => { const actual = await importOriginal(); return { ...actual, - rolesWithWriteAccess: ["admin"], + rolesWithWriteAccess: ["Admin"], }; }); @@ -243,20 +249,6 @@ vi.mock("@/app/(dashboard)/hooks/useTeams", () => ({ })), })); -// Mock useAuthorized hook to avoid Next.js router dependency -vi.mock("@/app/(dashboard)/hooks/useAuthorized", () => ({ - default: vi.fn(() => ({ - accessToken: "access_abc", - userId: "user_1", - userRole: "admin", - premiumUser: true, - token: "token_123", - userEmail: "test@example.com", - disabledPersonalKeyCreation: false, - showSSOBanner: false, - })), -})); - // KeyEditView mock: triggers onSubmit with our injected form values vi.mock("./key_edit_view", async () => { const React = await import("react"); @@ -302,21 +294,29 @@ const baseKeyData = { next_rotation_at: null as any, }; -const renderView = (premiumUser: boolean) => - render( +const renderView = (premiumUser: boolean) => { + // Configure the mock for this test + mockUseAuthorized.mockReturnValue({ + accessToken: "access_abc", + userId: "user_1", + userRole: "Admin", + premiumUser, + token: "token_123", + userEmail: "test@example.com", + disabledPersonalKeyCreation: false, + showSSOBanner: false, + }); + + return render( {}} keyData={baseKeyData as any} onKeyDataUpdate={() => {}} - accessToken="access_abc" - userID="user_1" - userRole="admin" teams={[]} - premiumUser={premiumUser} - setAccessToken={() => {}} />, ); +}; beforeEach(() => { vi.clearAllMocks(); @@ -328,6 +328,7 @@ describe("KeyInfoView handleKeyUpdate premium guard", () => { it("removes guardrails & prompts for non-premium users and prevents metadata.guardrails", async () => { renderView(false); // premiumUser = false + fireEvent.click(screen.getByText("Settings")); fireEvent.click(screen.getByText("Edit Settings")); (globalThis as any).__TEST_FORM_VALUES = { token: "tok_123", @@ -352,6 +353,7 @@ describe("KeyInfoView handleKeyUpdate premium guard", () => { it("preserves guardrails & prompts for premium users and includes metadata.guardrails", async () => { renderView(true); // premiumUser = true + fireEvent.click(screen.getByText("Settings")); fireEvent.click(screen.getByText("Edit Settings")); (globalThis as any).__TEST_FORM_VALUES = { token: "tok_123", @@ -378,6 +380,7 @@ describe("KeyInfoView handleKeyUpdate empty strings", () => { it(`maps empty strings to null for ${limit}`, async () => { renderView(true); // premiumUser = true + fireEvent.click(screen.getByText("Settings")); fireEvent.click(screen.getByText("Edit Settings")); (globalThis as any).__TEST_FORM_VALUES = { token: "tok_123", diff --git a/ui/litellm-dashboard/src/components/templates/key_info_view.test.tsx b/ui/litellm-dashboard/src/components/templates/key_info_view.test.tsx index 008b4cbe6ca..96d4b574a18 100644 --- a/ui/litellm-dashboard/src/components/templates/key_info_view.test.tsx +++ b/ui/litellm-dashboard/src/components/templates/key_info_view.test.tsx @@ -1,4 +1,5 @@ import useTeams from "@/app/(dashboard)/hooks/useTeams"; +import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; import { render, screen, waitFor } from "@testing-library/react"; import { beforeEach, describe, expect, it, vi } from "vitest"; import { KeyResponse, Team } from "../key_team_helpers/key_list"; @@ -8,6 +9,10 @@ vi.mock("@/app/(dashboard)/hooks/useTeams", () => ({ default: vi.fn(), })); +vi.mock("@/app/(dashboard)/hooks/useAuthorized", () => ({ + default: vi.fn(), +})); + describe("KeyInfoView", () => { beforeEach(() => { vi.mocked(useTeams).mockReturnValue({ @@ -85,17 +90,27 @@ describe("KeyInfoView", () => { key_rotation_at: undefined, }; + // Base mock for useAuthorized hook + const baseUseAuthorizedMock = { + accessToken: "test-token", + userId: "test-user", + userRole: "admin", + premiumUser: true, + token: "test-token", + userEmail: null, + disabledPersonalKeyCreation: null, + showSSOBanner: false, + }; + it("should render tags", async () => { + vi.mocked(useAuthorized).mockReturnValue(baseUseAuthorizedMock); + const { getByText } = render( {}} keyId={"test-key-id"} onKeyDataUpdate={() => {}} - accessToken={"test-token"} - userID={"test-user"} - userRole={"admin"} - premiumUser={true} teams={[]} />, ); @@ -105,16 +120,14 @@ describe("KeyInfoView", () => { }); it("should not render tags in metadata textarea", async () => { + vi.mocked(useAuthorized).mockReturnValue(baseUseAuthorizedMock); + const { container, getByText } = render( {}} keyId={"test-key-id"} onKeyDataUpdate={() => {}} - accessToken={"test-token"} - userID={"test-user"} - userRole={"admin"} - premiumUser={true} teams={[]} />, ); @@ -132,19 +145,15 @@ describe("KeyInfoView", () => { setTeams: vi.fn(), }); + vi.mocked(useAuthorized).mockReturnValue({ + ...baseUseAuthorizedMock, + userId: "proxy-admin-user", + userRole: "proxy_admin", + }); + const keyData = { ...MOCK_KEY_DATA, user_id: "other-user-id" }; render( - {}} - keyId={"test-key-id"} - onKeyDataUpdate={() => {}} - accessToken={"test-token"} - userID={"proxy-admin-user"} - userRole={"proxy_admin"} - premiumUser={true} - teams={[]} - />, + {}} keyId={"test-key-id"} onKeyDataUpdate={() => {}} teams={[]} />, ); await waitFor(() => { @@ -180,19 +189,15 @@ describe("KeyInfoView", () => { setTeams: vi.fn(), }); + vi.mocked(useAuthorized).mockReturnValue({ + ...baseUseAuthorizedMock, + userId: teamAdminUserId, + userRole: "user", + }); + const keyData = { ...MOCK_KEY_DATA, team_id: teamId, user_id: "other-user-id" }; render( - {}} - keyId={"test-key-id"} - onKeyDataUpdate={() => {}} - accessToken={"test-token"} - userID={teamAdminUserId} - userRole={"user"} - premiumUser={true} - teams={[]} - />, + {}} keyId={"test-key-id"} onKeyDataUpdate={() => {}} teams={[]} />, ); await waitFor(() => { @@ -207,20 +212,16 @@ describe("KeyInfoView", () => { setTeams: vi.fn(), }); + vi.mocked(useAuthorized).mockReturnValue({ + ...baseUseAuthorizedMock, + userId: "owner-user-id", + userRole: "user", + }); + const ownerUserId = "owner-user-id"; const keyData = { ...MOCK_KEY_DATA, user_id: ownerUserId }; render( - {}} - keyId={"test-key-id"} - onKeyDataUpdate={() => {}} - accessToken={"test-token"} - userID={ownerUserId} - userRole={"user"} - premiumUser={true} - teams={[]} - />, + {}} keyId={"test-key-id"} onKeyDataUpdate={() => {}} teams={[]} />, ); await waitFor(() => { @@ -235,19 +236,15 @@ describe("KeyInfoView", () => { setTeams: vi.fn(), }); + vi.mocked(useAuthorized).mockReturnValue({ + ...baseUseAuthorizedMock, + userId: "other-user-id", + userRole: "user", + }); + const keyData = { ...MOCK_KEY_DATA, user_id: "owner-user-id" }; render( - {}} - keyId={"test-key-id"} - onKeyDataUpdate={() => {}} - accessToken={"test-token"} - userID={"other-user-id"} - userRole={"user"} - premiumUser={true} - teams={[]} - />, + {}} keyId={"test-key-id"} onKeyDataUpdate={() => {}} teams={[]} />, ); await waitFor(() => { @@ -262,20 +259,16 @@ describe("KeyInfoView", () => { setTeams: vi.fn(), }); + vi.mocked(useAuthorized).mockReturnValue({ + ...baseUseAuthorizedMock, + userId: "internal-viewer-user-id", + userRole: "Internal Viewer", + }); + const ownerUserId = "internal-viewer-user-id"; const keyData = { ...MOCK_KEY_DATA, user_id: ownerUserId }; render( - {}} - keyId={"test-key-id"} - onKeyDataUpdate={() => {}} - accessToken={"test-token"} - userID={ownerUserId} - userRole={"Internal Viewer"} - premiumUser={true} - teams={[]} - />, + {}} keyId={"test-key-id"} onKeyDataUpdate={() => {}} teams={[]} />, ); await waitFor(() => { diff --git a/ui/litellm-dashboard/src/components/templates/key_info_view.tsx b/ui/litellm-dashboard/src/components/templates/key_info_view.tsx index 47abda58c29..76361cca2fd 100644 --- a/ui/litellm-dashboard/src/components/templates/key_info_view.tsx +++ b/ui/litellm-dashboard/src/components/templates/key_info_view.tsx @@ -19,6 +19,7 @@ import ObjectPermissionsView from "../object_permissions_view"; import { RegenerateKeyModal } from "../organisms/regenerate_key_modal"; import { parseErrorMessage } from "../shared/errorUtils"; import { KeyEditView } from "./key_edit_view"; +import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; interface KeyInfoViewProps { keyId: string; @@ -26,12 +27,7 @@ interface KeyInfoViewProps { keyData: KeyResponse | undefined; onKeyDataUpdate?: (data: Partial) => void; onDelete?: () => void; - accessToken: string | null; - userID: string | null; - userRole: string | null; teams: any[] | null; - premiumUser: boolean; - setAccessToken?: (token: string) => void; backButtonText?: string; } @@ -43,19 +39,14 @@ interface KeyInfoViewProps { * ───────────────────────────────────────────────────────────────────────── */ export default function KeyInfoView({ - keyId, onClose, keyData, - accessToken, - userID, - userRole, teams, onKeyDataUpdate, onDelete, - premiumUser, - setAccessToken, backButtonText = "Back to Keys", }: KeyInfoViewProps) { + const { accessToken, userId: userID, userRole, premiumUser } = useAuthorized(); const { teams: teamsData } = useTeams(); const [isEditing, setIsEditing] = useState(false); const [form] = Form.useForm(); @@ -400,9 +391,6 @@ export default function KeyInfoView({ selectedToken={currentKeyData} visible={isRegenerateModalOpen} onClose={() => setIsRegenerateModalOpen(false)} - accessToken={accessToken} - premiumUser={premiumUser} - setAccessToken={setAccessToken} onKeyUpdate={handleRegenerateKeyUpdate} /> diff --git a/ui/litellm-dashboard/src/components/templates/view_key_table.tsx b/ui/litellm-dashboard/src/components/templates/view_key_table.tsx index f8a86f0fad4..de347fd3305 100644 --- a/ui/litellm-dashboard/src/components/templates/view_key_table.tsx +++ b/ui/litellm-dashboard/src/components/templates/view_key_table.tsx @@ -5,27 +5,10 @@ import { Form, InputNumber, Modal } from "antd"; import { add } from "date-fns"; import React, { useEffect, useState } from "react"; import { CopyToClipboard } from "react-copy-to-clipboard"; -import { AllKeysTable } from "../all_keys_table"; -import { fetchAvailableModelsForTeamOrKey } from "../key_team_helpers/fetch_available_models_team_key"; +import { VirtualKeysTable } from "../VirtualKeysPage/VirtualKeysTable"; import useKeyList, { KeyResponse, Team } from "../key_team_helpers/key_list"; -import { keyDeleteCall, Organization, regenerateKeyCall } from "../networking"; - import NotificationManager from "../molecules/notifications_manager"; - -interface EditKeyModalProps { - visible: boolean; - onCancel: () => void; - token: any; // Assuming TeamType is a type representing your team object - onSubmit: (data: FormData) => void; // Assuming FormData is the type of data to be submitted -} - -interface ModelLimitModalProps { - visible: boolean; - onCancel: () => void; - token: KeyResponse; - onSubmit: (updatedMetadata: any) => void; - accessToken: string; -} +import { keyDeleteCall, Organization, regenerateKeyCall } from "../networking"; // Define the props type interface ViewKeyTableProps { @@ -47,39 +30,6 @@ interface ViewKeyTableProps { setAccessToken?: (token: string) => void; } -interface ItemData { - key_alias: string | null; - key_name: string; - spend: string; - max_budget: string | null; - models: string[]; - tpm_limit: string | null; - rpm_limit: string | null; - token: string; - token_id: string | null; - id: number; - team_id: string; - metadata: any; - user_id: string | null; - expires: any; - budget_duration: string | null; - budget_reset_at: string | null; - // Add any other properties that exist in the item data -} - -interface ModelLimits { - [key: string]: number; // Index signature allowing string keys -} - -interface CombinedLimit { - tpm: number; - rpm: number; -} - -interface CombinedLimits { - [key: string]: CombinedLimit; // Index signature allowing string keys -} - const ViewKeyTable: React.FC = ({ userID, userRole, @@ -98,24 +48,10 @@ const ViewKeyTable: React.FC = ({ createClicked, setAccessToken, }) => { - const [isButtonClicked, setIsButtonClicked] = useState(false); const [isDeleteModalOpen, setIsDeleteModalOpen] = useState(false); const [keyToDelete, setKeyToDelete] = useState(null); const [deleteConfirmInput, setDeleteConfirmInput] = useState(""); - const [selectedItem, setSelectedItem] = useState(null); - const [spendData, setSpendData] = useState<{ day: string; spend: number }[] | null>(null); - // NEW: Declare filter states for team and key alias. - const [teamFilter, setTeamFilter] = useState(selectedTeam?.team_id || ""); - - // Keep the team filter in sync with the incoming prop. - useEffect(() => { - setTeamFilter(selectedTeam?.team_id || ""); - }, [selectedTeam]); - - // Build a memoized filters object for the backend call. - - // Pass filters into the hook so the API call includes these query parameters. const { keys, isLoading, error, pagination, refresh, setKeys } = useKeyList({ selectedTeam: selectedTeam || undefined, currentOrg, @@ -129,21 +65,13 @@ const ViewKeyTable: React.FC = ({ refresh({ page: newPage }); }; - const [editModalVisible, setEditModalVisible] = useState(false); - const [infoDialogVisible, setInfoDialogVisible] = useState(false); const [selectedToken, setSelectedToken] = useState(null); - const [userModels, setUserModels] = useState([]); - const initialKnownTeamIDs: Set = new Set(); - const [modelLimitModalVisible, setModelLimitModalVisible] = useState(false); const [regenerateDialogVisible, setRegenerateDialogVisible] = useState(false); const [regeneratedKey, setRegeneratedKey] = useState(null); const [regenerateFormData, setRegenerateFormData] = useState(null); const [regenerateForm] = Form.useForm(); const [newExpiryTime, setNewExpiryTime] = useState(null); - const [knownTeamIDs, setKnownTeamIDs] = useState(initialKnownTeamIDs); - const [guardrailsList, setGuardrailsList] = useState([]); - useEffect(() => { const calculateNewExpiryTime = (duration: string | undefined) => { if (!duration) { @@ -190,36 +118,6 @@ const ViewKeyTable: React.FC = ({ console.log("calculateNewExpiryTime:", newExpiryTime); }, [selectedToken, regenerateFormData?.duration]); - useEffect(() => { - const fetchUserModels = async () => { - try { - if (userID === null || userRole === null || accessToken === null) { - return; - } - - const models = await fetchAvailableModelsForTeamOrKey(userID, userRole, accessToken); - if (models) { - setUserModels(models); - } - } catch (error) { - NotificationManager.error({ description: "Error fetching user models" }); - } - }; - - fetchUserModels(); - }, [accessToken, userID, userRole]); - - useEffect(() => { - if (teams) { - const teamIDSet: Set = new Set(); - teams.forEach((team: any, index: number) => { - const team_obj: string = team.team_id; - teamIDSet.add(team_obj); - }); - setKnownTeamIDs(teamIDSet); - } - }, [teams]); - const confirmDelete = async () => { if (keyToDelete == null || keys == null) { return; @@ -305,7 +203,7 @@ const ViewKeyTable: React.FC = ({ return (
- = ({ teams={teams} selectedTeam={selectedTeam} setSelectedTeam={setSelectedTeam} - accessToken={accessToken} - userID={userID} - userRole={userRole} organizations={organizations} setCurrentOrg={setCurrentOrg} refresh={refresh} selectedKeyAlias={selectedKeyAlias} setSelectedKeyAlias={setSelectedKeyAlias} - premiumUser={premiumUser} - setAccessToken={setAccessToken} /> {isDeleteModalOpen && From 5cdd88bc5c716a95f4b2d9ed600afd6cad85a64a Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Mon, 29 Dec 2025 17:06:16 -0800 Subject: [PATCH 030/158] fixing build --- .../UsagePage/components/EntityUsage/TopKeyView.tsx | 11 +---------- .../src/components/view_logs/index.tsx | 9 +-------- 2 files changed, 2 insertions(+), 18 deletions(-) diff --git a/ui/litellm-dashboard/src/components/UsagePage/components/EntityUsage/TopKeyView.tsx b/ui/litellm-dashboard/src/components/UsagePage/components/EntityUsage/TopKeyView.tsx index 8f6bc411630..db268286007 100644 --- a/ui/litellm-dashboard/src/components/UsagePage/components/EntityUsage/TopKeyView.tsx +++ b/ui/litellm-dashboard/src/components/UsagePage/components/EntityUsage/TopKeyView.tsx @@ -268,16 +268,7 @@ const TopKeyView: React.FC = ({ topKeys, teams, showTags = fals {/* Content */}
- +
diff --git a/ui/litellm-dashboard/src/components/view_logs/index.tsx b/ui/litellm-dashboard/src/components/view_logs/index.tsx index 34aa0ef3be8..2a94284fca5 100644 --- a/ui/litellm-dashboard/src/components/view_logs/index.tsx +++ b/ui/litellm-dashboard/src/components/view_logs/index.tsx @@ -526,12 +526,8 @@ export default function SpendLogsTable({ setSelectedKeyIdInfoView(null)} - premiumUser={premiumUser} backButtonText="Back to Logs" /> ) : selectedSessionId ? ( @@ -968,10 +964,7 @@ export function RequestViewer({ row }: { row: Row }) { {/* Cost Breakdown - Show if cost breakdown data is available */} - + {/* Configuration Info Message - Show when data is missing */} From a08dae7da01bc2ef5512aab7e82e081e4b8f3308 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Mon, 29 Dec 2025 19:27:52 -0800 Subject: [PATCH 031/158] keys table refactor pt2 --- .../src/app/(dashboard)/hooks/keys/useKeys.ts | 36 ++ .../VirtualKeysPage/VirtualKeysTable.test.tsx | 24 +- .../VirtualKeysPage/VirtualKeysTable.tsx | 107 ++--- .../components/templates/view_key_table.tsx | 379 +----------------- .../src/components/user_dashboard.tsx | 56 +-- 5 files changed, 103 insertions(+), 499 deletions(-) create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/hooks/keys/useKeys.ts diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/keys/useKeys.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/keys/useKeys.ts new file mode 100644 index 00000000000..8ae4d76ff5d --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/keys/useKeys.ts @@ -0,0 +1,36 @@ +import { keepPreviousData, useQuery, UseQueryResult } from "@tanstack/react-query"; +import { createQueryKeys } from "../common/queryKeysFactory"; +import { keyListCall } from "@/components/networking"; +import { KeyResponse } from "@/components/key_team_helpers/key_list"; +import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; + +const keyKeys = createQueryKeys("keys"); + +export interface KeysResponse { + keys: KeyResponse[]; + total_count: number; + current_page: number; + total_pages: number; +} + +export const useKeys = (page: number, pageSize: number): UseQueryResult => { + const { accessToken } = useAuthorized(); + + return useQuery({ + queryKey: keyKeys.list({ page, limit: pageSize }), + queryFn: async () => + await keyListCall( + accessToken!, + null, // organizationID + null, // teamID + null, // selectedKeyAlias + null, // userID + null, // keyHash + page, + pageSize, + ), + enabled: Boolean(accessToken), + staleTime: 30000, // 30 seconds + placeholderData: keepPreviousData, + }); +}; diff --git a/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.test.tsx b/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.test.tsx index 4f40ed99a84..71573712cd8 100644 --- a/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.test.tsx +++ b/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.test.tsx @@ -136,13 +136,7 @@ it("should render VirtualKeysTable component", () => { keys: [mockKey], setKeys: vi.fn(), isLoading: false, - pagination: { - currentPage: 1, - totalPages: 1, - totalCount: 1, - }, - onPageChange: vi.fn(), - pageSize: 50, + totalCount: 1, teams: [mockTeam], selectedTeam: null, setSelectedTeam: vi.fn(), @@ -166,13 +160,7 @@ it("should display key information correctly", async () => { keys: [mockKey], setKeys: vi.fn(), isLoading: false, - pagination: { - currentPage: 1, - totalPages: 1, - totalCount: 1, - }, - onPageChange: vi.fn(), - pageSize: 50, + totalCount: 1, teams: [mockTeam], selectedTeam: null, setSelectedTeam: vi.fn(), @@ -200,13 +188,7 @@ it("should display user email correctly", async () => { keys: [mockKey], setKeys: vi.fn(), isLoading: false, - pagination: { - currentPage: 1, - totalPages: 1, - totalCount: 1, - }, - onPageChange: vi.fn(), - pageSize: 50, + totalCount: 1, teams: [mockTeam], selectedTeam: null, setSelectedTeam: vi.fn(), diff --git a/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.tsx b/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.tsx index 128129ce66a..8cfe5cc2f1a 100644 --- a/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.tsx +++ b/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.tsx @@ -1,12 +1,14 @@ "use client"; -import { Setter } from "@/types"; -import { formatNumberWithCommas, updateExistingKeys } from "@/utils/dataUtils"; +import { useKeys } from "@/app/(dashboard)/hooks/keys/useKeys"; +import { formatNumberWithCommas } from "@/utils/dataUtils"; import { ChevronDownIcon, ChevronRightIcon, ChevronUpIcon, SwitchVerticalIcon } from "@heroicons/react/outline"; import { ColumnDef, flexRender, getCoreRowModel, + getPaginationRowModel, getSortedRowModel, + PaginationState, SortingState, useReactTable, } from "@tanstack/react-table"; @@ -30,27 +32,10 @@ import { KeyResponse, Team } from "../key_team_helpers/key_list"; import FilterComponent, { FilterOption } from "../molecules/filter"; import { Organization } from "../networking"; import KeyInfoView from "../templates/key_info_view"; -import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; interface VirtualKeysTableProps { - keys: KeyResponse[]; - setKeys: (keys: KeyResponse[] | ((prev: KeyResponse[]) => KeyResponse[])) => void; - isLoading?: boolean; - pagination: { - currentPage: number; - totalPages: number; - totalCount: number; - }; - onPageChange: (page: number) => void; - pageSize?: number; teams: Team[] | null; - selectedTeam: Team | null; - setSelectedTeam: (team: Team | null) => void; - selectedKeyAlias: string | null; - setSelectedKeyAlias: Setter; organizations: Organization[] | null; - setCurrentOrg: React.Dispatch>; - refresh?: () => void; onSortChange?: (sortBy: string, sortOrder: "asc" | "desc") => void; currentSort?: { sortBy: string; @@ -63,21 +48,8 @@ interface VirtualKeysTableProps { * The team selector and filtering have been removed so that all keys are shown. */ -export function VirtualKeysTable({ - keys, - setKeys, - isLoading = false, - pagination, - onPageChange, - pageSize = 50, - teams, - organizations, - refresh, - onSortChange, - currentSort, -}: VirtualKeysTableProps) { - const { accessToken, userId: userID, userRole, premiumUser } = useAuthorized(); - const [selectedKeyId, setSelectedKeyId] = useState(null); +export function VirtualKeysTable({ teams, organizations, onSortChange, currentSort }: VirtualKeysTableProps) { + const [selectedKey, setSelectedKey] = useState(null); const [sorting, setSorting] = React.useState(() => { if (currentSort) { return [ @@ -94,22 +66,33 @@ export function VirtualKeysTable({ }, ]; }); + const [tablePagination, setTablePagination] = React.useState({ + pageIndex: 0, + pageSize: 100, + }); + + const { + data: keys, + isPending: isLoading, + refetch, + } = useKeys(tablePagination.pageIndex + 1, tablePagination.pageSize); + const totalCount = keys?.total_count || 0; const [expandedAccordions, setExpandedAccordions] = useState>({}); // Use the filter logic hook const { filters, filteredKeys, allKeyAliases, allTeams, allOrganizations, handleFilterChange, handleFilterReset } = useFilterLogic({ - keys, + keys: keys?.keys || [], teams, organizations, }); // Add a useEffect to call refresh when a key is created useEffect(() => { - if (refresh) { + if (refetch) { const handleStorageChange = () => { - refresh(); + refetch(); }; // Listen for storage events that might indicate a key was created @@ -119,7 +102,7 @@ export function VirtualKeysTable({ window.removeEventListener("storage", handleStorageChange); }; } - }, [refresh]); + }, [refetch]); const columns: ColumnDef[] = [ { @@ -145,7 +128,7 @@ export function VirtualKeysTable({ size="xs" variant="light" className="font-mono text-blue-500 bg-blue-50 hover:bg-blue-100 text-xs font-normal px-2 py-0.5 text-left overflow-hidden truncate max-w-[200px]" - onClick={() => setSelectedKeyId(info.getValue() as string)} + onClick={() => setSelectedKey(info.row.original)} > {info.getValue() ? `${(info.getValue() as string).slice(0, 7)}...` : "-"} @@ -496,6 +479,7 @@ export function VirtualKeysTable({ columnResizeDirection: "ltr", state: { sorting, + pagination: tablePagination, }, onSortingChange: (updaterOrValue) => { const newSorting = typeof updaterOrValue === "function" ? updaterOrValue(sorting) : updaterOrValue; @@ -514,10 +498,14 @@ export function VirtualKeysTable({ onSortChange?.(sortBy, sortOrder); } }, + onPaginationChange: setTablePagination, getCoreRowModel: getCoreRowModel(), getSortedRowModel: getSortedRowModel(), + getPaginationRowModel: getPaginationRowModel(), enableSorting: true, manualSorting: false, + manualPagination: true, + pageCount: Math.ceil(totalCount / tablePagination.pageSize), }); // Update local sorting state when currentSort prop changes @@ -534,26 +522,11 @@ export function VirtualKeysTable({ return (
- {selectedKeyId ? ( + {selectedKey ? ( setSelectedKeyId(null)} - keyData={filteredKeys.find((k) => k.token === selectedKeyId)} - onKeyDataUpdate={(updatedKeyData) => { - setKeys((keys) => - keys.map((key) => { - if (key.token === updatedKeyData.token) { - return updateExistingKeys(key, updatedKeyData); - } - return key; - }), - ); - if (refresh) refresh(); // Minimal fix: refresh the full key list after an update - }} - onDelete={() => { - setKeys((keys) => keys.filter((key) => key.token !== selectedKeyId)); - if (refresh) refresh(); // Minimal fix: refresh the full key list after a delete - }} + keyId={selectedKey.token} + onClose={() => setSelectedKey(null)} + keyData={selectedKey} teams={allTeams} /> ) : ( @@ -572,26 +545,30 @@ export function VirtualKeysTable({ Showing{" "} {isLoading ? "..." - : `${(pagination.currentPage - 1) * pageSize + 1} - ${Math.min(pagination.currentPage * pageSize, pagination.totalCount)}`}{" "} - of {isLoading ? "..." : pagination.totalCount} results + : `${table.getState().pagination.pageIndex * table.getState().pagination.pageSize + 1} - ${Math.min( + (table.getState().pagination.pageIndex + 1) * table.getState().pagination.pageSize, + totalCount, + )}`}{" "} + of {isLoading ? "..." : totalCount} results
- Page {isLoading ? "..." : pagination.currentPage} of {isLoading ? "..." : pagination.totalPages} + Page {isLoading ? "..." : table.getState().pagination.pageIndex + 1} of{" "} + {isLoading ? "..." : table.getPageCount()} -
-
-
-
- - - -
-
-

- Warning: You are about to delete this Virtual Key. -

-

- This action is irreversible and will immediately revoke access for any applications using this - key. -

-
-
-

Are you sure you want to delete this Virtual Key?

-
- - setDeleteConfirmInput(e.target.value)} - placeholder="Enter key name exactly" - className="w-full px-4 py-3 border border-gray-300 rounded-md focus:outline-none focus:ring-2 focus:ring-blue-500 focus:border-blue-500 text-base" - autoFocus - /> -
-
-
-
- - -
- - - ); - })()} - - {/* Regenerate Key Form Modal */} - { - setRegenerateDialogVisible(false); - regenerateForm.resetFields(); - }} - footer={[ - , - , - ]} - > - {premiumUser ? ( -
{ - if ("duration" in changedValues) { - handleRegenerateFormChange("duration", changedValues.duration); - } - }} - > - - - - - - - - - - - - - - - -
- Current expiry:{" "} - {selectedToken?.expires != null ? new Date(selectedToken.expires).toLocaleString() : "Never"} -
- {newExpiryTime &&
New expiry: {newExpiryTime}
} -
- ) : ( -
-

Upgrade to use this feature

- -
- )} - - - {/* Regenerated Key Display Modal */} - {regeneratedKey && ( - setRegeneratedKey(null)} - footer={[ - , - ]} - > - - Regenerated Key - -

- Please replace your old key with the new key generated. For security reasons,{" "} - you will not be able to view it again through your LiteLLM account. If you lose this secret key, - you will need to generate a new one. -

- - - Key Alias: -
- 
-                  {selectedToken?.key_alias || "No alias set"}
-
-
- New Virtual Key: -
- 
{regeneratedKey}
-
- NotificationManager.success({ description: "Virtual Key copied to clipboard" })} - > - - - - - - )} + ); }; diff --git a/ui/litellm-dashboard/src/components/user_dashboard.tsx b/ui/litellm-dashboard/src/components/user_dashboard.tsx index 4f910452bd8..c55cb09457e 100644 --- a/ui/litellm-dashboard/src/components/user_dashboard.tsx +++ b/ui/litellm-dashboard/src/components/user_dashboard.tsx @@ -1,23 +1,23 @@ "use client"; -import React, { useState, useEffect } from "react"; -import { - userInfoCall, - modelAvailableCall, - getProxyUISettings, - Organization, - keyInfoCall, - getProxyBaseUrl, -} from "./networking"; -import { fetchTeams } from "./common_components/fetch_teams"; -import { Grid, Col } from "@tremor/react"; -import CreateKey from "./organisms/create_key_button"; -import ViewKeyTable from "./templates/view_key_table"; -import Onboarding from "../app/onboarding/page"; -import { useSearchParams } from "next/navigation"; -import { KeyResponse, Team } from "./key_team_helpers/key_list"; -import { jwtDecode } from "jwt-decode"; -import { Typography } from "antd"; import { clearTokenCookies } from "@/utils/cookieUtils"; +import { Col, Grid } from "@tremor/react"; +import { Typography } from "antd"; +import { jwtDecode } from "jwt-decode"; +import { useSearchParams } from "next/navigation"; +import React, { useEffect, useState } from "react"; +import Onboarding from "../app/onboarding/page"; +import { fetchTeams } from "./common_components/fetch_teams"; +import { KeyResponse, Team } from "./key_team_helpers/key_list"; +import { + getProxyBaseUrl, + getProxyUISettings, + keyInfoCall, + modelAvailableCall, + Organization, + userInfoCall, +} from "./networking"; +import CreateKey from "./organisms/create_key_button"; +import { VirtualKeysTable } from "./VirtualKeysPage/VirtualKeysTable"; export interface ProxySettings { PROXY_BASE_URL: string | null; @@ -361,25 +361,7 @@ const UserDashboard: React.FC = ({ addKey={addKey} premiumUser={premiumUser} /> - - + From 2dad030462926f71dd022d4f889d5559c518e69b Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Mon, 29 Dec 2025 19:41:28 -0800 Subject: [PATCH 032/158] updating tests --- .../VirtualKeysPage/VirtualKeysTable.test.tsx | 108 +++++++++++------- 1 file changed, 68 insertions(+), 40 deletions(-) diff --git a/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.test.tsx b/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.test.tsx index 71573712cd8..e42c3d99c9a 100644 --- a/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.test.tsx +++ b/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.test.tsx @@ -1,9 +1,11 @@ import { screen, waitFor } from "@testing-library/react"; -import { vi, it, expect } from "vitest"; +import { vi, it, expect, beforeEach, MockedFunction } from "vitest"; import { renderWithProviders } from "../../../tests/test-utils"; import { VirtualKeysTable } from "./VirtualKeysTable"; import { KeyResponse, Team } from "../key_team_helpers/key_list"; import { Organization } from "../networking"; +import { KeysResponse, useKeys } from "@/app/(dashboard)/hooks/keys/useKeys"; +import { useFilterLogic } from "../key_team_helpers/filter_logic"; // Mock network calls vi.mock("./networking", async (importOriginal) => { @@ -39,6 +41,16 @@ vi.mock("./key_team_helpers/filter_helpers", () => ({ ]), })); +// Mock useKeys hook +vi.mock("@/app/(dashboard)/hooks/keys/useKeys", () => ({ + useKeys: vi.fn(), +})); + +// Mock useFilterLogic hook +vi.mock("../key_team_helpers/filter_logic", () => ({ + useFilterLogic: vi.fn(), +})); + const mockKey: KeyResponse = { token: "sk-1234567890abcdef", token_id: "key-1", @@ -131,23 +143,55 @@ const mockOrganization: Organization = { members: [], }; +// Mock hook implementations +const mockUseKeys = useKeys as MockedFunction; +const mockUseFilterLogic = useFilterLogic as MockedFunction; + +beforeEach(() => { + // Reset mocks before each test + vi.clearAllMocks(); + + // Setup default mock implementations + mockUseKeys.mockReturnValue({ + data: { + keys: [mockKey], + total_count: 1, + current_page: 1, + total_pages: 1, + } as KeysResponse, + isPending: false, + refetch: vi.fn(), + } as any); + + mockUseFilterLogic.mockReturnValue({ + filters: { + "Team ID": "team-1", + "Organization ID": "org-1", + "Key Alias": "Test Key Alias", + "User ID": "user-1", + "User Email": "user@example.com", + "User Role": "user", + "Sort By": "created_at", + "Sort Order": "desc", + }, + filteredKeys: [mockKey], + allKeyAliases: ["test-key-alias"], + allTeams: [mockTeam], + allOrganizations: [mockOrganization], + handleFilterChange: vi.fn(), + handleFilterReset: vi.fn(), + }); +}); + it("should render VirtualKeysTable component", () => { const mockProps = { - keys: [mockKey], - setKeys: vi.fn(), - isLoading: false, - totalCount: 1, teams: [mockTeam], - selectedTeam: null, - setSelectedTeam: vi.fn(), - selectedKeyAlias: null, - setSelectedKeyAlias: vi.fn(), - accessToken: "test-token", - userID: "user-1", - userRole: "admin", organizations: [mockOrganization], - setCurrentOrg: vi.fn(), - premiumUser: false, + onSortChange: vi.fn(), + currentSort: { + sortBy: "created_at", + sortOrder: "desc" as const, + }, }; renderWithProviders(); @@ -157,21 +201,13 @@ it("should render VirtualKeysTable component", () => { it("should display key information correctly", async () => { const mockProps = { - keys: [mockKey], - setKeys: vi.fn(), - isLoading: false, - totalCount: 1, teams: [mockTeam], - selectedTeam: null, - setSelectedTeam: vi.fn(), - selectedKeyAlias: null, - setSelectedKeyAlias: vi.fn(), - accessToken: "test-token", - userID: "user-1", - userRole: "admin", organizations: [mockOrganization], - setCurrentOrg: vi.fn(), - premiumUser: false, + onSortChange: vi.fn(), + currentSort: { + sortBy: "created_at", + sortOrder: "desc" as const, + }, }; renderWithProviders(); @@ -185,21 +221,13 @@ it("should display key information correctly", async () => { it("should display user email correctly", async () => { const mockProps = { - keys: [mockKey], - setKeys: vi.fn(), - isLoading: false, - totalCount: 1, teams: [mockTeam], - selectedTeam: null, - setSelectedTeam: vi.fn(), - selectedKeyAlias: null, - setSelectedKeyAlias: vi.fn(), - accessToken: "test-token", - userID: "user-1", - userRole: "admin", organizations: [mockOrganization], - setCurrentOrg: vi.fn(), - premiumUser: false, + onSortChange: vi.fn(), + currentSort: { + sortBy: "created_at", + sortOrder: "desc" as const, + }, }; renderWithProviders(); From c79db9121d38b33181e11887dbee7f17c54631c2 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Mon, 29 Dec 2025 19:42:38 -0800 Subject: [PATCH 033/158] Fixing build --- .../src/components/templates/view_key_table.tsx | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/ui/litellm-dashboard/src/components/templates/view_key_table.tsx b/ui/litellm-dashboard/src/components/templates/view_key_table.tsx index 7a71b8a0f26..0344edaa85d 100644 --- a/ui/litellm-dashboard/src/components/templates/view_key_table.tsx +++ b/ui/litellm-dashboard/src/components/templates/view_key_table.tsx @@ -1,8 +1,6 @@ "use client"; import { Setter } from "@/types"; -import { Form } from "antd"; -import { add } from "date-fns"; -import React, { useEffect, useState } from "react"; +import React from "react"; import { VirtualKeysTable } from "../VirtualKeysPage/VirtualKeysTable"; import useKeyList, { KeyResponse, Team } from "../key_team_helpers/key_list"; import { Organization } from "../networking"; From a4341c6e5b3df5e51d2c3ca4f2ab4d14a6e304c2 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Mon, 29 Dec 2025 19:56:44 -0800 Subject: [PATCH 034/158] test removing python install step --- .circleci/config.yml | 7 ------- 1 file changed, 7 deletions(-) diff --git a/.circleci/config.yml b/.circleci/config.yml index 6d9a9ff5bee..a010d2a1c8b 100644 --- a/.circleci/config.yml +++ b/.circleci/config.yml @@ -3430,13 +3430,6 @@ jobs: - checkout - attach_workspace: at: ~/project - - run: - name: Install Python 3.9 - command: | - bash miniconda.sh -b -p $HOME/miniconda - export PATH="$HOME/miniconda/bin:$PATH" - $HOME/miniconda/bin/conda create -y -n myenv python=3.9 - conda run -n myenv python --version - run: name: Install Dependencies command: | From 554830d2e2fb9d4089c2bb427e26d49bbe131e3a Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Mon, 29 Dec 2025 20:33:46 -0800 Subject: [PATCH 035/158] reverting --- .circleci/config.yml | 30 +++++++++++++++++++++++++----- 1 file changed, 25 insertions(+), 5 deletions(-) diff --git a/.circleci/config.yml b/.circleci/config.yml index a010d2a1c8b..34ba0c6c3a8 100644 --- a/.circleci/config.yml +++ b/.circleci/config.yml @@ -3422,14 +3422,32 @@ jobs: --coverage.reportsDirectory=coverage/html e2e_ui_testing: - docker: - - image: mcr.microsoft.com/playwright:v1.57.0-noble + machine: + image: ubuntu-2204:2023.10.1 resource_class: xlarge working_directory: ~/project steps: - checkout + - setup_google_dns - attach_workspace: at: ~/project + - run: + name: Upgrade Docker to v24.x (API 1.44+) + command: | + curl -fsSL https://get.docker.com | sh + sudo usermod -aG docker $USER + docker version + - run: + name: Install Python 3.9 + command: | + curl https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh --output miniconda.sh + bash miniconda.sh -b -p $HOME/miniconda + export PATH="$HOME/miniconda/bin:$PATH" + conda init bash + source ~/.bashrc + conda create -n myenv python=3.9 -y + conda activate myenv + python --version - run: name: Install Dependencies command: | @@ -3457,8 +3475,7 @@ jobs: - run: name: Install Playwright Browsers command: | - cd ui/litellm-dashboard && npx playwright install - + npx playwright install --with-deps - run: name: Build Docker image command: docker build -t my-app:latest -f ./docker/Dockerfile.database . @@ -3497,7 +3514,10 @@ jobs: - run: name: Run Playwright Tests command: | - cd ui/litellm-dashboard && npx playwright test --config=e2e_tests/playwright.config.ts + npx playwright test \ + --config ui/litellm-dashboard/e2e_tests/playwright.config.ts \ + --reporter=html \ + --output=test-results no_output_timeout: 120m - store_artifacts: path: test-results From 46ce68619c0762788c0c0fd4456c76e52729c737 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Mon, 29 Dec 2025 20:49:24 -0800 Subject: [PATCH 036/158] attempt fix 3 --- .circleci/config.yml | 2 +- ui/litellm-dashboard/e2e_tests/playwright.config.ts | 5 ----- 2 files changed, 1 insertion(+), 6 deletions(-) diff --git a/.circleci/config.yml b/.circleci/config.yml index 34ba0c6c3a8..d98fdaa4a4d 100644 --- a/.circleci/config.yml +++ b/.circleci/config.yml @@ -3475,7 +3475,7 @@ jobs: - run: name: Install Playwright Browsers command: | - npx playwright install --with-deps + npx playwright install - run: name: Build Docker image command: docker build -t my-app:latest -f ./docker/Dockerfile.database . diff --git a/ui/litellm-dashboard/e2e_tests/playwright.config.ts b/ui/litellm-dashboard/e2e_tests/playwright.config.ts index 5ab68b3534f..80d45a892c8 100644 --- a/ui/litellm-dashboard/e2e_tests/playwright.config.ts +++ b/ui/litellm-dashboard/e2e_tests/playwright.config.ts @@ -37,11 +37,6 @@ export default defineConfig({ name: "firefox", use: { ...devices["Desktop Firefox"] }, }, - - { - name: "webkit", - use: { ...devices["Desktop Safari"] }, - }, ], /* Timeout settings */ From 42e437a418a71cb2a82cd198c025f9429eb8a200 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Tue, 30 Dec 2025 10:35:25 -0800 Subject: [PATCH 037/158] Loading state for Keys table --- .../VirtualKeysPage/VirtualKeysTable.tsx | 70 +++++++++++-------- .../organisms/create_key_button.tsx | 27 ++----- .../src/components/user_dashboard.tsx | 10 --- 3 files changed, 45 insertions(+), 62 deletions(-) diff --git a/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.tsx b/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.tsx index 8cfe5cc2f1a..b95d675979c 100644 --- a/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.tsx +++ b/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.tsx @@ -24,7 +24,7 @@ import { TableRow, Text, } from "@tremor/react"; -import { Tooltip } from "antd"; +import { Skeleton, Tooltip } from "antd"; import React, { useEffect, useState } from "react"; import { getModelDisplayName } from "../key_team_helpers/fetch_available_models_team_key"; import { useFilterLogic } from "../key_team_helpers/filter_logic"; @@ -520,6 +520,10 @@ export function VirtualKeysTable({ teams, organizations, onSortChange, currentSo } }, [currentSort]); + const { pageIndex, pageSize } = table.getState().pagination; + const start = pageIndex * pageSize + 1; + const end = Math.min((pageIndex + 1) * pageSize, totalCount); + const rangeLabel = `${start} - ${end}`; return (
{selectedKey ? ( @@ -541,38 +545,46 @@ export function VirtualKeysTable({ teams, organizations, onSortChange, currentSo
- - Showing{" "} - {isLoading - ? "..." - : `${table.getState().pagination.pageIndex * table.getState().pagination.pageSize + 1} - ${Math.min( - (table.getState().pagination.pageIndex + 1) * table.getState().pagination.pageSize, - totalCount, - )}`}{" "} - of {isLoading ? "..." : totalCount} results - + {isLoading ? ( + + ) : ( + + Showing {rangeLabel} of {totalCount} results + + )}
- - Page {isLoading ? "..." : table.getState().pagination.pageIndex + 1} of{" "} - {isLoading ? "..." : table.getPageCount()} - + {isLoading ? ( + + ) : ( + + Page {pageIndex + 1} of {table.getPageCount()} + + )} - + {isLoading ? ( + + ) : ( + + )} - + {isLoading ? ( + + ) : ( + + )}
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 83ef444d59e..3cd8a04e067 100644 --- a/ui/litellm-dashboard/src/components/organisms/create_key_button.tsx +++ b/ui/litellm-dashboard/src/components/organisms/create_key_button.tsx @@ -1,4 +1,5 @@ "use client"; +import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; import { formatNumberWithCommas } from "@/utils/dataUtils"; import { InfoCircleOutlined } from "@ant-design/icons"; import { Accordion, AccordionBody, AccordionHeader, Button, Col, Grid, Text, TextInput, Title } from "@tremor/react"; @@ -39,14 +40,10 @@ import VectorStoreSelector from "../vector_store_management/VectorStoreSelector" const { Option } = Select; interface CreateKeyProps { - userID: string; team: Team | null; - userRole: string | null; - accessToken: string; data: any[] | null; teams: Team[] | null; addKey: (data: any) => void; - premiumUser?: boolean; } interface User { @@ -137,16 +134,8 @@ export const fetchUserModels = async ( * Please contribute to the new refactor. * ───────────────────────────────────────────────────────────────────────── */ -const CreateKey: React.FC = ({ - userID, - team, - teams, - userRole, - accessToken, - data, - addKey, - premiumUser = false, -}) => { +const CreateKey: React.FC = ({ team, teams, data, addKey }) => { + const { accessToken, userId: userID, userRole, premiumUser } = useAuthorized(); const [form] = Form.useForm(); const [isModalVisible, setIsModalVisible] = useState(false); const [apiKey, setApiKey] = useState(null); @@ -170,7 +159,6 @@ const CreateKey: React.FC = ({ const [modelAliases, setModelAliases] = useState<{ [key: string]: string }>({}); const [autoRotationEnabled, setAutoRotationEnabled] = useState(false); const [rotationInterval, setRotationInterval] = useState("30d"); - const handleOk = () => { setIsModalVisible(false); form.resetFields(); @@ -490,14 +478,7 @@ const CreateKey: React.FC = ({ + Create New Key )} - +
{/* Section 1: Key Ownership */}
diff --git a/ui/litellm-dashboard/src/components/user_dashboard.tsx b/ui/litellm-dashboard/src/components/user_dashboard.tsx index c55cb09457e..ec6de82fdf9 100644 --- a/ui/litellm-dashboard/src/components/user_dashboard.tsx +++ b/ui/litellm-dashboard/src/components/user_dashboard.tsx @@ -92,13 +92,7 @@ const UserDashboard: React.FC = ({ const [teamSpend, setTeamSpend] = useState(null); const [userModels, setUserModels] = useState([]); const [proxySettings, setProxySettings] = useState(null); - const defaultTeam: TeamInterface = { - models: [], - team_alias: "Default Team", - team_id: null, - }; const [selectedTeam, setSelectedTeam] = useState(null); - const [selectedKeyAlias, setSelectedKeyAlias] = useState(null); // check if window is not undefined if (typeof window !== "undefined") { window.addEventListener("beforeunload", function () { @@ -352,14 +346,10 @@ const UserDashboard: React.FC = ({ From b1dd8bf3e1a6030571fc3864bf5a99316f1fb4c0 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Tue, 30 Dec 2025 10:43:27 -0800 Subject: [PATCH 038/158] Adding tests --- .../VirtualKeysPage/VirtualKeysTable.test.tsx | 28 +++++++++++++++++++ 1 file changed, 28 insertions(+) diff --git a/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.test.tsx b/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.test.tsx index e42c3d99c9a..3f55b11769c 100644 --- a/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.test.tsx +++ b/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.test.tsx @@ -236,3 +236,31 @@ it("should display user email correctly", async () => { expect(screen.getByText("user@example.com")).toBeInTheDocument(); }); }); + +it("should show skeleton loaders when isLoading is true", () => { + // Mock loading state + mockUseKeys.mockReturnValue({ + data: null, + isPending: true, + refetch: vi.fn(), + } as any); + + const mockProps = { + teams: [mockTeam], + organizations: [mockOrganization], + onSortChange: vi.fn(), + currentSort: { + sortBy: "created_at", + sortOrder: "desc" as const, + }, + }; + + renderWithProviders(); + + // Check that loading message is shown + expect(screen.getByText("🚅 Loading keys...")).toBeInTheDocument(); + + // Check that actual key data is not shown + expect(screen.queryByText("Test Key Alias")).not.toBeInTheDocument(); + expect(screen.queryByText("Test Team")).not.toBeInTheDocument(); +}); From c049f6a0734aa2864907aee8cf082d99cec60980 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Tue, 30 Dec 2025 12:19:25 -0800 Subject: [PATCH 039/158] Adding e2e tests for sidebar --- .gitignore | 3 +- .../e2e_tests/auth/login.setup.ts | 20 ----------- .../e2e_tests/fixtures/roles.ts | 6 ++++ .../e2e_tests/fixtures/users.ts | 10 ++++++ ui/litellm-dashboard/e2e_tests/globalSetup.ts | 18 ++++++++++ .../e2e_tests/playwright.config.ts | 1 + .../e2e_tests/tests/login/login.spec.ts | 13 +++++++ .../tests/navigation/sidebar.spec.ts | 34 +++++++++++++++++++ 8 files changed, 84 insertions(+), 21 deletions(-) delete mode 100644 ui/litellm-dashboard/e2e_tests/auth/login.setup.ts create mode 100644 ui/litellm-dashboard/e2e_tests/fixtures/roles.ts create mode 100644 ui/litellm-dashboard/e2e_tests/fixtures/users.ts create mode 100644 ui/litellm-dashboard/e2e_tests/globalSetup.ts create mode 100644 ui/litellm-dashboard/e2e_tests/tests/login/login.spec.ts create mode 100644 ui/litellm-dashboard/e2e_tests/tests/navigation/sidebar.spec.ts diff --git a/.gitignore b/.gitignore index b6d4fd44190..9a4f666bf47 100644 --- a/.gitignore +++ b/.gitignore @@ -102,4 +102,5 @@ litellm/proxy/_experimental/out/guardrails/index.html scripts/test_vertex_ai_search.py LAZY_LOADING_IMPROVEMENTS.md **/test-results -**/playwright-report \ No newline at end of file +**/playwright-report +**/*.storageState.json \ No newline at end of file diff --git a/ui/litellm-dashboard/e2e_tests/auth/login.setup.ts b/ui/litellm-dashboard/e2e_tests/auth/login.setup.ts deleted file mode 100644 index 712191f6843..00000000000 --- a/ui/litellm-dashboard/e2e_tests/auth/login.setup.ts +++ /dev/null @@ -1,20 +0,0 @@ -import { test, expect } from "@playwright/test"; - -test("login and save auth state", async ({ page }) => { - await page.goto("http://localhost:4000/ui"); - - await page.getByPlaceholder("Enter your username").fill("admin"); - await page.getByPlaceholder("Enter your password").fill("gm"); - - const loginButton = page.getByRole("button", { name: "Login" }); - await expect(loginButton).toBeEnabled(); - await loginButton.click(); - - // Assert successful login (important) - await expect(page.getByText("AI Gateway")).toBeVisible(); - - // 🔐 Save auth state - await page.context().storageState({ - path: "storageState.json", - }); -}); diff --git a/ui/litellm-dashboard/e2e_tests/fixtures/roles.ts b/ui/litellm-dashboard/e2e_tests/fixtures/roles.ts new file mode 100644 index 00000000000..913230ad44b --- /dev/null +++ b/ui/litellm-dashboard/e2e_tests/fixtures/roles.ts @@ -0,0 +1,6 @@ +export enum Role { + ProxyAdmin = "proxy_admin", + ProxyAdminViewer = "proxy_admin_viewer", + InternalUser = "internal_user", + InternalUserViewer = "internal_user_viewer", +} diff --git a/ui/litellm-dashboard/e2e_tests/fixtures/users.ts b/ui/litellm-dashboard/e2e_tests/fixtures/users.ts new file mode 100644 index 00000000000..d1f1eab00e5 --- /dev/null +++ b/ui/litellm-dashboard/e2e_tests/fixtures/users.ts @@ -0,0 +1,10 @@ +import { Role } from "./roles"; + +const isCI = !!process.env.CI; + +export const users = { + [Role.ProxyAdmin]: { + email: "admin", + password: isCI ? "gm" : "sk-1234", + }, +}; diff --git a/ui/litellm-dashboard/e2e_tests/globalSetup.ts b/ui/litellm-dashboard/e2e_tests/globalSetup.ts new file mode 100644 index 00000000000..a725c58f35b --- /dev/null +++ b/ui/litellm-dashboard/e2e_tests/globalSetup.ts @@ -0,0 +1,18 @@ +import { chromium } from "@playwright/test"; +import { users } from "./fixtures/users"; +import { Role } from "./fixtures/roles"; + +async function globalSetup() { + const browser = await chromium.launch(); + const page = await browser.newPage(); + await page.goto("http://localhost:4000/ui/login"); + await page.getByPlaceholder("Enter your username").fill(users[Role.ProxyAdmin].email); + await page.getByPlaceholder("Enter your password").fill(users[Role.ProxyAdmin].password); + const loginButton = page.getByRole("button", { name: "Login" }); + await loginButton.click(); + await page.waitForSelector("text=AI Gateway"); + await page.context().storageState({ path: "admin.storageState.json" }); + await browser.close(); +} + +export default globalSetup; diff --git a/ui/litellm-dashboard/e2e_tests/playwright.config.ts b/ui/litellm-dashboard/e2e_tests/playwright.config.ts index 80d45a892c8..9df2281016c 100644 --- a/ui/litellm-dashboard/e2e_tests/playwright.config.ts +++ b/ui/litellm-dashboard/e2e_tests/playwright.config.ts @@ -44,4 +44,5 @@ export default defineConfig({ expect: { timeout: 10 * 1000, }, + globalSetup: require.resolve("./globalSetup"), }); diff --git a/ui/litellm-dashboard/e2e_tests/tests/login/login.spec.ts b/ui/litellm-dashboard/e2e_tests/tests/login/login.spec.ts new file mode 100644 index 00000000000..5ac977ff0c8 --- /dev/null +++ b/ui/litellm-dashboard/e2e_tests/tests/login/login.spec.ts @@ -0,0 +1,13 @@ +import { expect, test } from "@playwright/test"; +import { users } from "../../fixtures/users"; +import { Role } from "../../fixtures/roles"; + +test("user can log in", async ({ page }) => { + await page.goto("http://localhost:4000/ui/login"); + await page.getByPlaceholder("Enter your username").fill(users[Role.ProxyAdmin].email); + await page.getByPlaceholder("Enter your password").fill(users[Role.ProxyAdmin].password); + const loginButton = page.getByRole("button", { name: "Login" }); + await expect(loginButton).toBeEnabled(); + await loginButton.click(); + await expect(page.getByText("AI Gateway")).toBeVisible(); +}); diff --git a/ui/litellm-dashboard/e2e_tests/tests/navigation/sidebar.spec.ts b/ui/litellm-dashboard/e2e_tests/tests/navigation/sidebar.spec.ts new file mode 100644 index 00000000000..dafb03a7cbd --- /dev/null +++ b/ui/litellm-dashboard/e2e_tests/tests/navigation/sidebar.spec.ts @@ -0,0 +1,34 @@ +import test, { expect } from "@playwright/test"; +import { Role } from "../../fixtures/roles"; + +const sidebarButtons = { + [Role.ProxyAdmin]: [ + "Virtual Keys", + "Playground", + "Models", + "Usage", + "Teams", + "Internal User", + "Settings", + "Experimental", + "API Reference", + "AI Hub", + ], +}; + +const roles = [{ role: Role.ProxyAdmin, storage: "admin.storageState.json" }]; + +for (const { role, storage } of roles) { + test.describe(`${role} sidebar`, () => { + test.use({ storageState: storage }); + + test("can see and navigate all sidebar buttons", async ({ page }) => { + await page.goto("http://localhost:4000/ui"); + for (const button of sidebarButtons[role as keyof typeof sidebarButtons]) { + const tab = page.getByRole("menuitem", { name: button }); + await expect(tab).toBeVisible(); + await tab.click(); + } + }); + }); +} From 16f663778d267b0fe9bd7983fa403c99f431f79e Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Tue, 30 Dec 2025 13:30:54 -0800 Subject: [PATCH 040/158] Add Model Form Refactor --- .../hooks/guardrails/useGuardrails.ts | 18 + .../src/app/(dashboard)/hooks/tags/useTags.ts | 16 + .../ModelsAndEndpointsView.tsx | 1 - .../src/components/add_model/AddModelForm.tsx | 386 +++++++++++++++++ .../components/add_model/add_model_tab.tsx | 407 +----------------- 5 files changed, 441 insertions(+), 387 deletions(-) create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/hooks/guardrails/useGuardrails.ts create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/hooks/tags/useTags.ts create mode 100644 ui/litellm-dashboard/src/components/add_model/AddModelForm.tsx diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/guardrails/useGuardrails.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/guardrails/useGuardrails.ts new file mode 100644 index 00000000000..9786b7fa359 --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/guardrails/useGuardrails.ts @@ -0,0 +1,18 @@ +import { useQuery, UseQueryResult } from "@tanstack/react-query"; +import { createQueryKeys } from "../common/queryKeysFactory"; +import { getGuardrailsList } from "@/components/networking"; +import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; + +const guardrailKeys = createQueryKeys("guardrails"); + +export const useGuardrails = (): UseQueryResult => { + const { accessToken, userId, userRole } = useAuthorized(); + return useQuery({ + queryKey: guardrailKeys.list({}), + queryFn: async () => { + const response = await getGuardrailsList(accessToken!); + return response.guardrails.map((g: { guardrail_name: string }) => g.guardrail_name); + }, + enabled: Boolean(accessToken && userId && userRole), + }); +}; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/tags/useTags.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/tags/useTags.ts new file mode 100644 index 00000000000..8f82502a74c --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/tags/useTags.ts @@ -0,0 +1,16 @@ +import { useQuery, UseQueryResult } from "@tanstack/react-query"; +import { createQueryKeys } from "../common/queryKeysFactory"; +import { tagListCall } from "@/components/networking"; +import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; +import { TagListResponse } from "@/components/tag_management/types"; + +const tagKeys = createQueryKeys("tags"); + +export const useTags = (): UseQueryResult => { + const { accessToken, userId, userRole } = useAuthorized(); + return useQuery({ + queryKey: tagKeys.list({}), + queryFn: async () => await tagListCall(accessToken!), + enabled: Boolean(accessToken && userId && userRole), + }); +}; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/ModelsAndEndpointsView.tsx b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/ModelsAndEndpointsView.tsx index bf001c62126..315b41393d6 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/ModelsAndEndpointsView.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/ModelsAndEndpointsView.tsx @@ -679,7 +679,6 @@ const ModelsAndEndpointsView: React.FC = ({ credentials={credentialsList} accessToken={accessToken} userRole={userRole} - premiumUser={premiumUser} /> )} diff --git a/ui/litellm-dashboard/src/components/add_model/AddModelForm.tsx b/ui/litellm-dashboard/src/components/add_model/AddModelForm.tsx new file mode 100644 index 00000000000..7ddc7840464 --- /dev/null +++ b/ui/litellm-dashboard/src/components/add_model/AddModelForm.tsx @@ -0,0 +1,386 @@ +import { useProviderFields } from "@/app/(dashboard)/hooks/providers/useProviderFields"; +import { useGuardrails } from "@/app/(dashboard)/hooks/guardrails/useGuardrails"; +import { useTags } from "@/app/(dashboard)/hooks/tags/useTags"; +import { all_admin_roles } from "@/utils/roles"; +import { Switch, Text } from "@tremor/react"; +import type { FormInstance } from "antd"; +import { Select as AntdSelect, Button, Card, Col, Form, Modal, Row, Tooltip, Typography } from "antd"; +import type { UploadProps } from "antd/es/upload"; +import React, { useEffect, useMemo, useState } from "react"; +import TeamDropdown from "../common_components/team_dropdown"; +import type { Team } from "../key_team_helpers/key_list"; +import { type CredentialItem, type ProviderCreateInfo, modelAvailableCall } from "../networking"; +import { Providers, providerLogoMap } from "../provider_info_helpers"; +import { ProviderLogo } from "../molecules/models/ProviderLogo"; +import AdvancedSettings from "./advanced_settings"; +import ConditionalPublicModelName from "./conditional_public_model_name"; +import LiteLLMModelNameField from "./litellm_model_name"; +import ConnectionErrorDisplay from "./model_connection_test"; +import ProviderSpecificFields from "./provider_specific_fields"; +import { TEST_MODES } from "./add_model_modes"; +import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; + +interface AddModelFormProps { + form: FormInstance; // For the Add Model tab + handleOk: () => void; + selectedProvider: Providers; + setSelectedProvider: (provider: Providers) => void; + providerModels: string[]; + setProviderModelsFn: (provider: Providers) => void; + getPlaceholder: (provider: Providers) => string; + uploadProps: UploadProps; + showAdvancedSettings: boolean; + setShowAdvancedSettings: (show: boolean) => void; + teams: Team[] | null; + credentials: CredentialItem[]; +} + +const { Title, Link } = Typography; + +const AddModelForm: React.FC = ({ + form, + handleOk, + selectedProvider, + setSelectedProvider, + providerModels, + setProviderModelsFn, + getPlaceholder, + uploadProps, + showAdvancedSettings, + setShowAdvancedSettings, + teams, + credentials, +}) => { + const [testMode, setTestMode] = useState("chat"); + const [isResultModalVisible, setIsResultModalVisible] = useState(false); + const [isTestingConnection, setIsTestingConnection] = useState(false); + // Using a unique ID to force the ConnectionErrorDisplay to remount and run a fresh test + const [connectionTestId, setConnectionTestId] = useState(""); + + const { accessToken, userRole, premiumUser } = useAuthorized(); + const { + data: providerMetadata, + isLoading: isProviderMetadataLoading, + error: providerMetadataError, + } = useProviderFields(); + const { data: guardrailsList, isLoading: isGuardrailsLoading, error: guardrailsError } = useGuardrails(); + const { data: tagsList, isLoading: isTagsLoading, error: tagsError } = useTags(); + + const handleTestConnection = async () => { + setIsTestingConnection(true); + setConnectionTestId(`test-${Date.now()}`); + setIsResultModalVisible(true); + }; + + const [isTeamOnly, setIsTeamOnly] = useState(false); + + const [modelAccessGroups, setModelAccessGroups] = useState([]); + + useEffect(() => { + const fetchModelAccessGroups = async () => { + const response = await modelAvailableCall(accessToken, "", "", false, null, true, true); + setModelAccessGroups(response["data"].map((model: any) => model["id"])); + }; + fetchModelAccessGroups(); + }, [accessToken]); + + const sortedProviderMetadata: ProviderCreateInfo[] = useMemo(() => { + if (!providerMetadata) { + return []; + } + return [...providerMetadata].sort((a, b) => a.provider_display_name.localeCompare(b.provider_display_name)); + }, [providerMetadata]); + + const providerMetadataErrorText = providerMetadataError + ? providerMetadataError instanceof Error + ? providerMetadataError.message + : "Failed to load providers" + : null; + + const isAdmin = all_admin_roles.includes(userRole); + + return ( + <> + Add Model + + { + console.log("🔥 Form onFinish triggered with values:", values); + handleOk(); + }} + onFinishFailed={(errorInfo) => { + console.log("💥 Form onFinishFailed triggered:", errorInfo); + }} + labelCol={{ span: 10 }} + wrapperCol={{ span: 16 }} + labelAlign="left" + > + <> + {/* Provider Selection */} + + { + setSelectedProvider(value as Providers); + setProviderModelsFn(value as Providers); + form.setFieldsValue({ + custom_llm_provider: value, + }); + form.setFieldsValue({ + model: [], + model_name: undefined, + }); + }} + > + {providerMetadataErrorText && sortedProviderMetadata.length === 0 && ( + + {providerMetadataErrorText} + + )} + {sortedProviderMetadata.map((providerInfo) => { + const displayName = providerInfo.provider_display_name; + const providerKey = providerInfo.provider; + const logoSrc = providerLogoMap[displayName] ?? ""; + + return ( + +
+ + {displayName} +
+ + ); + })} + + + + + {/* Conditionally Render "Public Model Name" */} + + + {/* Select Mode */} + + setTestMode(value)} + options={TEST_MODES} + /> + + + + + + Optional - LiteLLM endpoint to use when health checking this model{" "} + + Learn more + + + + + + {/* Credentials */} +
+ + Either select existing credentials OR enter new provider credentials below + +
+ + + (option?.label ?? "").toLowerCase().includes(input.toLowerCase())} + options={[ + { value: null, label: "None" }, + ...credentials.map((credential) => ({ + value: credential.credential_name, + label: credential.credential_name, + })), + ]} + allowClear + /> + + + + prevValues.litellm_credential_name !== currentValues.litellm_credential_name || + prevValues.provider !== currentValues.provider + } + > + {({ getFieldValue }) => { + const credentialName = getFieldValue("litellm_credential_name"); + console.log("🔑 Credential Name Changed:", credentialName); + // Only show provider specific fields if no credentials selected + if (!credentialName) { + return ( + <> +
+
+ OR +
+
+ + + ); + } + return null; + }} + +
+
+ Additional Model Info Settings +
+
+ {/* Team-only Model Switch */} + + + { + setIsTeamOnly(checked); + if (!checked) { + form.setFieldValue("team_id", undefined); + } + }} + disabled={!premiumUser} + /> + + + + {/* Conditional Team Selection */} + {isTeamOnly && ( + + + + )} + {isAdmin && ( + <> + + ({ + value: group, + label: group, + }))} + maxTagCount="responsive" + allowClear + /> + + + )} + + +
+ + Need Help? + +
+ + +
+
+ + + + + {/* Test Connection Results Modal */} + { + setIsResultModalVisible(false); + setIsTestingConnection(false); + }} + footer={[ + , + ]} + width={700} + > + {/* Only render the ConnectionErrorDisplay when modal is visible and we have a test ID */} + {isResultModalVisible && ( + { + setIsResultModalVisible(false); + setIsTestingConnection(false); + }} + onTestComplete={() => setIsTestingConnection(false)} + /> + )} + + + ); +}; + +export default AddModelForm; diff --git a/ui/litellm-dashboard/src/components/add_model/add_model_tab.tsx b/ui/litellm-dashboard/src/components/add_model/add_model_tab.tsx index cfe00af5071..17ec1cbceee 100644 --- a/ui/litellm-dashboard/src/components/add_model/add_model_tab.tsx +++ b/ui/litellm-dashboard/src/components/add_model/add_model_tab.tsx @@ -1,34 +1,19 @@ -import { useProviderFields } from "@/app/(dashboard)/hooks/providers/useProviderFields"; -import { all_admin_roles } from "@/utils/roles"; -import { Switch, Tab, TabGroup, TabList, TabPanel, TabPanels, Text } from "@tremor/react"; +import { Tab, TabGroup, TabList, TabPanel, TabPanels } from "@tremor/react"; import type { FormInstance } from "antd"; -import { Select as AntdSelect, Button, Card, Col, Form, Modal, Row, Tooltip, Typography } from "antd"; +import { Form } from "antd"; import type { UploadProps } from "antd/es/upload"; -import React, { useEffect, useMemo, useState } from "react"; -import TeamDropdown from "../common_components/team_dropdown"; +import React from "react"; import type { Team } from "../key_team_helpers/key_list"; -import { - type CredentialItem, - type ProviderCreateInfo, - getGuardrailsList, - modelAvailableCall, - tagListCall, -} from "../networking"; -import { Providers, providerLogoMap } from "../provider_info_helpers"; -import { ProviderLogo } from "../molecules/models/ProviderLogo"; -import { Tag } from "../tag_management/types"; +import { type CredentialItem } from "../networking"; +import { Providers } from "../provider_info_helpers"; import AddAutoRouterTab from "./add_auto_router_tab"; -import { TEST_MODES } from "./add_model_modes"; +import AddModelForm from "./AddModelForm"; import AdvancedSettings from "./advanced_settings"; -import ConditionalPublicModelName from "./conditional_public_model_name"; import { handleAddAutoRouterSubmit } from "./handle_add_auto_router_submit"; -import LiteLLMModelNameField from "./litellm_model_name"; -import ConnectionErrorDisplay from "./model_connection_test"; -import ProviderSpecificFields from "./provider_specific_fields"; interface AddModelTabProps { form: FormInstance; // For the Add Model tab - handleOk: () => void; + handleOk: (values?: any) => void; selectedProvider: Providers; setSelectedProvider: (provider: Providers) => void; providerModels: string[]; @@ -41,11 +26,8 @@ interface AddModelTabProps { credentials: CredentialItem[]; accessToken: string; userRole: string; - premiumUser: boolean; } -const { Title, Link } = Typography; - const AddModelTab: React.FC = ({ form, handleOk, @@ -61,90 +43,9 @@ const AddModelTab: React.FC = ({ credentials, accessToken, userRole, - premiumUser, }) => { // Create separate form instance for auto router const [autoRouterForm] = Form.useForm(); - // State for test mode and connection testing - const [testMode, setTestMode] = useState("chat"); - const [isResultModalVisible, setIsResultModalVisible] = useState(false); - const [isTestingConnection, setIsTestingConnection] = useState(false); - const [guardrailsList, setGuardrailsList] = useState([]); - const [tagsList, setTagsList] = useState>({}); - // Using a unique ID to force the ConnectionErrorDisplay to remount and run a fresh test - const [connectionTestId, setConnectionTestId] = useState(""); - - // Provider metadata for driving the provider select from backend config - const { - data: providerMetadata, - isLoading: isProviderMetadataLoading, - error: providerMetadataError, - } = useProviderFields(); - - useEffect(() => { - const fetchGuardrails = async () => { - try { - const response = await getGuardrailsList(accessToken); - const guardrailNames = response.guardrails.map((g: { guardrail_name: string }) => g.guardrail_name); - setGuardrailsList(guardrailNames); - } catch (error) { - console.error("Failed to fetch guardrails:", error); - } - }; - - fetchGuardrails(); - }, [accessToken]); - - useEffect(() => { - const fetchTags = async () => { - try { - const response = await tagListCall(accessToken); - setTagsList(response); - } catch (error) { - console.error("Failed to fetch tags:", error); - } - }; - - fetchTags(); - }, [accessToken]); - - // Test connection when button is clicked - const handleTestConnection = async () => { - setIsTestingConnection(true); - // Generate a new test ID (using timestamp for uniqueness) - // This forces React to create a new instance of ConnectionErrorDisplay - setConnectionTestId(`test-${Date.now()}`); - // Show the modal with the fresh test - setIsResultModalVisible(true); - }; - - // State for team-only switch - const [isTeamOnly, setIsTeamOnly] = useState(false); - - const [modelAccessGroups, setModelAccessGroups] = useState([]); - - useEffect(() => { - const fetchModelAccessGroups = async () => { - const response = await modelAvailableCall(accessToken, "", "", false, null, true, true); - setModelAccessGroups(response["data"].map((model: any) => model["id"])); - }; - fetchModelAccessGroups(); - }, [accessToken]); - - const sortedProviderMetadata: ProviderCreateInfo[] = useMemo(() => { - if (!providerMetadata) { - return []; - } - return [...providerMetadata].sort((a, b) => a.provider_display_name.localeCompare(b.provider_display_name)); - }, [providerMetadata]); - - const providerMetadataErrorText = providerMetadataError - ? providerMetadataError instanceof Error - ? providerMetadataError.message - : "Failed to load providers" - : null; - - const isAdmin = all_admin_roles.includes(userRole); const handleAutoRouterOk = () => { autoRouterForm @@ -166,247 +67,20 @@ const AddModelTab: React.FC = ({ - Add Model - -
{ - console.log("🔥 Form onFinish triggered with values:", values); - handleOk(); - }} - onFinishFailed={(errorInfo) => { - console.log("💥 Form onFinishFailed triggered:", errorInfo); - }} - labelCol={{ span: 10 }} - wrapperCol={{ span: 16 }} - labelAlign="left" - > - <> - {/* Provider Selection */} - - { - setSelectedProvider(value as Providers); - setProviderModelsFn(value as Providers); - form.setFieldsValue({ - custom_llm_provider: value, - }); - form.setFieldsValue({ - model: [], - model_name: undefined, - }); - }} - > - {providerMetadataErrorText && sortedProviderMetadata.length === 0 && ( - - {providerMetadataErrorText} - - )} - {sortedProviderMetadata.map((providerInfo) => { - const displayName = providerInfo.provider_display_name; - const providerKey = providerInfo.provider; - const logoSrc = providerLogoMap[displayName] ?? ""; - - return ( - -
- - {displayName} -
- - ); - })} - - - - - {/* Conditionally Render "Public Model Name" */} - - - {/* Select Mode */} - - setTestMode(value)} - options={TEST_MODES} - /> - - - - - - Optional - LiteLLM endpoint to use when health checking this model{" "} - - Learn more - - - - - - {/* Credentials */} -
- - Either select existing credentials OR enter new provider credentials below - -
- - - - (option?.label ?? "").toLowerCase().includes(input.toLowerCase()) - } - options={[ - { value: null, label: "None" }, - ...credentials.map((credential) => ({ - value: credential.credential_name, - label: credential.credential_name, - })), - ]} - allowClear - /> - - - - prevValues.litellm_credential_name !== currentValues.litellm_credential_name || - prevValues.provider !== currentValues.provider - } - > - {({ getFieldValue }) => { - const credentialName = getFieldValue("litellm_credential_name"); - console.log("🔑 Credential Name Changed:", credentialName); - // Only show provider specific fields if no credentials selected - if (!credentialName) { - return ( - <> -
-
- OR -
-
- - - ); - } - return null; - }} - -
-
- Additional Model Info Settings -
-
- {/* Team-only Model Switch */} - - - { - setIsTeamOnly(checked); - if (!checked) { - form.setFieldValue("team_id", undefined); - } - }} - disabled={!premiumUser} - /> - - - - {/* Conditional Team Selection */} - {isTeamOnly && ( - - - - )} - {isAdmin && ( - <> - - ({ - value: group, - label: group, - }))} - maxTagCount="responsive" - allowClear - /> - - - )} - - -
- - Need Help? - -
- - -
-
- - -
+ = ({ - - {/* Test Connection Results Modal */} - { - setIsResultModalVisible(false); - setIsTestingConnection(false); - }} - footer={[ - , - ]} - width={700} - > - {/* Only render the ConnectionErrorDisplay when modal is visible and we have a test ID */} - {isResultModalVisible && ( - { - setIsResultModalVisible(false); - setIsTestingConnection(false); - }} - onTestComplete={() => setIsTestingConnection(false)} - /> - )} - ); }; From 7fdea85b5c6549ecc2c3a39bdf32c733745c0ae0 Mon Sep 17 00:00:00 2001 From: Yuta Saito Date: Wed, 31 Dec 2025 06:55:47 +0900 Subject: [PATCH 041/158] feat: optimize MCP server listing by separating health checks --- .../mcp_server/mcp_server_manager.py | 79 ++++++- .../mcp_management_endpoints.py | 55 ++++- .../test_mcp_management_endpoints.py | 199 ++++++++++++++++-- .../mcpServers/useMCPServerHealth.test.ts | 127 +++++++++++ .../hooks/mcpServers/useMCPServerHealth.ts | 22 ++ .../mcp_tools/mcp_server_columns.tsx | 14 ++ .../components/mcp_tools/mcp_servers.test.tsx | 105 +++++++++ .../src/components/mcp_tools/mcp_servers.tsx | 50 +++-- .../src/components/networking.tsx | 38 ++++ 9 files changed, 653 insertions(+), 36 deletions(-) create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/hooks/mcpServers/useMCPServerHealth.test.ts create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/hooks/mcpServers/useMCPServerHealth.ts diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index 58e9a345dc7..2d3ea1e827f 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -2186,8 +2186,12 @@ class MCPServerManager: async def _noop(session): return "ok" - await client.run_with_session(_noop) + # Add timeout wrapper to prevent hanging + await asyncio.wait_for(client.run_with_session(_noop), timeout=10.0) status = "healthy" + except asyncio.TimeoutError: + health_check_error = "Health check timed out after 10 seconds" + status = "unhealthy" except Exception as e: health_check_error = str(e) status = "unhealthy" @@ -2221,14 +2225,15 @@ class MCPServerManager: async def get_all_mcp_servers_with_health_and_teams( self, user_api_key_auth: Optional[UserAPIKeyAuth] = None, - include_health: bool = True, + server_ids: Optional[List[str]] = None, ) -> List[LiteLLM_MCPServerTable]: """ Get all MCP servers that the user has access to, with health status and team information. Args: user_api_key_auth: User authentication info for access control - include_health: Whether to include health check information + server_ids: Optional list of server IDs to filter. If provided, only these servers + will be checked (subject to access control). If None, all accessible servers are checked. Returns: List of MCP server objects with health and team data @@ -2237,10 +2242,16 @@ class MCPServerManager: # Get allowed server IDs allowed_server_ids = await self.get_allowed_mcp_servers(user_api_key_auth) + # Filter by requested server_ids if provided + if server_ids: + # Only check servers that are both requested AND accessible + target_server_ids = [sid for sid in server_ids if sid in allowed_server_ids] + else: + # Check all accessible servers + target_server_ids = allowed_server_ids + # Run health checks concurrently - tasks = [ - self.health_check_server(server_id) for server_id in allowed_server_ids - ] + tasks = [self.health_check_server(server_id) for server_id in target_server_ids] results = await asyncio.gather(*tasks) # Filter out None results (servers that were not found) @@ -2248,6 +2259,62 @@ class MCPServerManager: return list_mcp_servers + async def get_all_allowed_mcp_servers( + self, + user_api_key_auth: Optional[UserAPIKeyAuth] = None, + ) -> List[LiteLLM_MCPServerTable]: + """ + Get all MCP servers that the user has access to. + + Args: + user_api_key_auth: User authentication info for access control + + Returns: + List of MCP server objects without health status + """ + from datetime import datetime + + # Get allowed server IDs + allowed_server_ids = await self.get_allowed_mcp_servers(user_api_key_auth) + + list_mcp_servers: List[LiteLLM_MCPServerTable] = [] + + for server_id in allowed_server_ids: + server = self.get_mcp_server_by_id(server_id) + if not server: + verbose_logger.warning(f"MCP Server {server_id} not found in registry") + continue + + # Build LiteLLM_MCPServerTable without health check + mcp_server_table = LiteLLM_MCPServerTable( + server_id=server.server_id, + server_name=server.server_name, + alias=server.alias, + description=( + server.mcp_info.get("description") if server.mcp_info else None + ), + url=server.url, + transport=server.transport, + auth_type=server.auth_type, + created_at=datetime.now(), + updated_at=datetime.now(), + teams=[], + mcp_access_groups=server.access_groups or [], + allowed_tools=server.allowed_tools or [], + extra_headers=server.extra_headers or [], + mcp_info=server.mcp_info, + static_headers=server.static_headers, + status=None, # No health check performed + last_health_check=None, # No health check performed + health_check_error=None, + command=getattr(server, "command", None), + args=getattr(server, "args", None) or [], + env=getattr(server, "env", None) or {}, + ) + list_mcp_servers.append(mcp_server_table) + + return list_mcp_servers + async def reload_servers_from_database(self): """ Public method to reload all MCP servers from database into registry. diff --git a/litellm/proxy/management_endpoints/mcp_management_endpoints.py b/litellm/proxy/management_endpoints/mcp_management_endpoints.py index 59b9659dd88..500323d3beb 100644 --- a/litellm/proxy/management_endpoints/mcp_management_endpoints.py +++ b/litellm/proxy/management_endpoints/mcp_management_endpoints.py @@ -16,7 +16,7 @@ Endpoints here: import importlib from dataclasses import dataclass from datetime import datetime, timedelta -from typing import Any, Dict, Iterable, List, Optional +from typing import Any, Dict, Iterable, List, Literal, Optional from fastapi import ( APIRouter, @@ -24,6 +24,7 @@ from fastapi import ( Form, Header, HTTPException, + Query, Request, Response, status, @@ -318,7 +319,7 @@ if MCP_AVAILABLE: aggregated_servers: Dict[str, LiteLLM_MCPServerTable] = {} for auth_context in auth_contexts: - servers = await global_mcp_server_manager.get_all_mcp_servers_with_health_and_teams( + servers = await global_mcp_server_manager.get_all_allowed_mcp_servers( user_api_key_auth=auth_context ) for server in servers: @@ -336,6 +337,56 @@ if MCP_AVAILABLE: server.mcp_info["is_public"] = True return redacted_mcp_servers + @router.get( + "/server/health", + description="Health check for MCP servers", + dependencies=[Depends(user_api_key_auth)], + ) + async def health_check_servers( + server_ids: Optional[List[str]] = Query( + None, + description="Server IDs to check. If not provided, checks all accessible servers.", + ), + user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), + ): + """ + Perform health checks on one or more MCP servers. + + Parameters: + - server_ids: Optional list of server IDs. If not provided, checks all accessible servers. + + Returns: + - Health check results for requested servers + + ``` + # Check all accessible servers + curl --location 'http://localhost:4000/v1/mcp/server/health' \ + --header 'Authorization: Bearer your_api_key_here' + + # Check specific servers + curl --location 'http://localhost:4000/v1/mcp/server/health?server_ids=server-1&server_ids=server-2' \ + --header 'Authorization: Bearer your_api_key_here' + ``` + """ + auth_contexts = await build_effective_auth_contexts(user_api_key_dict) + + server_status_map: Dict[ + str, Optional[Literal["healthy", "unhealthy", "unknown"]] + ] = {} + for auth_context in auth_contexts: + servers = await global_mcp_server_manager.get_all_mcp_servers_with_health_and_teams( + user_api_key_auth=auth_context, + server_ids=server_ids, + ) + for server in servers: + if server.server_id not in server_status_map: + server_status_map[server.server_id] = server.status + + return [ + {"server_id": server_id, "status": status} + for server_id, status in server_status_map.items() + ] + @router.get( "/server/{server_id}", description="Returns the mcp server info", diff --git a/tests/test_litellm/proxy/management_endpoints/test_mcp_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_mcp_management_endpoints.py index 8d5b3dcedc9..cd1a1f5e10d 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_mcp_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_mcp_management_endpoints.py @@ -169,8 +169,8 @@ class TestListMCPServers: return_value=["config_server_1", "config_server_2"] ) - # Mock the new method that returns servers with health and team data - mock_servers_with_health = [ + # Mock the new method that returns servers without health check + mock_servers = [ generate_mock_mcp_server_db_record( server_id="config_server_1", alias="Zapier MCP", @@ -184,11 +184,11 @@ class TestListMCPServers: transport="http", ), ] - mock_manager.get_all_mcp_servers_with_health_and_teams = AsyncMock( - return_value=mock_servers_with_health + mock_manager.get_all_allowed_mcp_servers = AsyncMock( + return_value=mock_servers ) - for idx, server in enumerate(mock_servers_with_health): + for idx, server in enumerate(mock_servers): server.credentials = {"auth_value": f"secret_{idx}"} with patch( @@ -200,6 +200,9 @@ class TestListMCPServers: ), patch( "litellm.proxy.management_endpoints.mcp_management_endpoints.get_prisma_client_or_throw", return_value=mock_prisma_client, + ), patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.build_effective_auth_contexts", + AsyncMock(return_value=[mock_user_auth]), ): # Import and call the function from litellm.proxy.management_endpoints.mcp_management_endpoints import ( @@ -300,8 +303,8 @@ class TestListMCPServers: ] ) - # Mock the new method that returns servers with health and team data - mock_servers_with_health = [ + # Mock the new method that returns servers without health check + mock_servers = [ db_server_1, db_server_2, generate_mock_mcp_server_db_record( @@ -317,11 +320,11 @@ class TestListMCPServers: transport="http", ), ] - mock_manager.get_all_mcp_servers_with_health_and_teams = AsyncMock( - return_value=mock_servers_with_health + mock_manager.get_all_allowed_mcp_servers = AsyncMock( + return_value=mock_servers ) - for idx, server in enumerate(mock_servers_with_health): + for idx, server in enumerate(mock_servers): server.credentials = {"auth_value": f"secret_{idx}"} with patch( @@ -333,6 +336,9 @@ class TestListMCPServers: ), patch( "litellm.proxy.management_endpoints.mcp_management_endpoints.get_prisma_client_or_throw", return_value=mock_prisma_client, + ), patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.build_effective_auth_contexts", + AsyncMock(return_value=[mock_user_auth]), ): # Import and call the function from litellm.proxy.management_endpoints.mcp_management_endpoints import ( @@ -425,8 +431,8 @@ class TestListMCPServers: return_value=["db_server_allowed", "config_server_allowed"] ) - # Mock the new method that returns servers with health and team data - mock_servers_with_health = [ + # Mock the new method that returns servers without health check + mock_servers = [ db_server_allowed, generate_mock_mcp_server_db_record( server_id="config_server_allowed", @@ -434,11 +440,11 @@ class TestListMCPServers: url="https://actions.zapier.com/mcp/sse", ), ] - mock_manager.get_all_mcp_servers_with_health_and_teams = AsyncMock( - return_value=mock_servers_with_health + mock_manager.get_all_allowed_mcp_servers = AsyncMock( + return_value=mock_servers ) - for idx, server in enumerate(mock_servers_with_health): + for idx, server in enumerate(mock_servers): server.credentials = {"auth_value": f"secret_{idx}"} with patch( @@ -450,6 +456,9 @@ class TestListMCPServers: ), patch( "litellm.proxy.management_endpoints.mcp_management_endpoints.get_prisma_client_or_throw", return_value=mock_prisma_client, + ), patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.build_effective_auth_contexts", + AsyncMock(return_value=[mock_user_auth]), ): # Import and call the function from litellm.proxy.management_endpoints.mcp_management_endpoints import ( @@ -975,3 +984,163 @@ class TestUpdateMCPServer: # Verify the result includes extra_headers assert result.extra_headers == ["X-Custom-Header", "X-Another-Header"] assert result.alias == "Updated Test Server" + + +class TestHealthCheckServers: + """Test suite for health check servers endpoint""" + + @pytest.mark.asyncio + async def test_health_check_all_servers(self): + """ + Test health check for all accessible servers + + Scenario: User has access to 2 servers, checks all + Expected: Returns health status for both servers + """ + from litellm.proxy.management_endpoints.mcp_management_endpoints import ( + health_check_servers, + ) + + # Mock user auth + mock_user_auth = generate_mock_user_api_key_auth() + + # Mock health check results + mock_health_result_1 = generate_mock_mcp_server_db_record( + server_id="server-1", + alias="Server 1", + url="https://server1.example.com", + ) + mock_health_result_1.status = "healthy" + mock_health_result_1.last_health_check = datetime.now() + mock_health_result_1.health_check_error = None + + mock_health_result_2 = generate_mock_mcp_server_db_record( + server_id="server-2", + alias="Server 2", + url="https://server2.example.com", + ) + mock_health_result_2.status = "unhealthy" + mock_health_result_2.last_health_check = datetime.now() + mock_health_result_2.health_check_error = "Connection timeout" + + # Mock manager + mock_manager = MagicMock() + mock_manager.get_all_mcp_servers_with_health_and_teams = AsyncMock( + return_value=[mock_health_result_1, mock_health_result_2] + ) + + with patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.global_mcp_server_manager", + mock_manager, + ), patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.build_effective_auth_contexts", + AsyncMock(return_value=[mock_user_auth]), + ): + result = await health_check_servers( + server_ids=None, + user_api_key_dict=mock_user_auth, + ) + + # Verify results + assert len(result) == 2 + assert result[0]["server_id"] == "server-1" + assert result[0]["status"] == "healthy" + assert result[1]["server_id"] == "server-2" + assert result[1]["status"] == "unhealthy" + + @pytest.mark.asyncio + async def test_health_check_specific_servers(self): + """ + Test health check for specific servers + + Scenario: User requests health check for specific server IDs + Expected: Returns health status only for requested servers + """ + from litellm.proxy.management_endpoints.mcp_management_endpoints import ( + health_check_servers, + ) + + # Mock user auth + mock_user_auth = generate_mock_user_api_key_auth() + + # Mock health check result + mock_health_result = generate_mock_mcp_server_db_record( + server_id="server-1", + alias="Server 1", + url="https://server1.example.com", + ) + mock_health_result.status = "healthy" + mock_health_result.last_health_check = datetime.now() + mock_health_result.health_check_error = None + + # Mock manager + mock_manager = MagicMock() + mock_manager.get_all_mcp_servers_with_health_and_teams = AsyncMock( + return_value=[mock_health_result] + ) + + with patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.global_mcp_server_manager", + mock_manager, + ), patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.build_effective_auth_contexts", + AsyncMock(return_value=[mock_user_auth]), + ): + result = await health_check_servers( + server_ids=["server-1"], + user_api_key_dict=mock_user_auth, + ) + + # Verify results + assert len(result) == 1 + assert result[0]["server_id"] == "server-1" + assert result[0]["status"] == "healthy" + + @pytest.mark.asyncio + async def test_health_check_unauthorized_servers(self): + """ + Test health check with unauthorized servers + + Scenario: User requests health check for servers they don't have access to + Expected: Only checks accessible servers, unauthorized servers are filtered out + """ + from litellm.proxy.management_endpoints.mcp_management_endpoints import ( + health_check_servers, + ) + + # Mock user auth + mock_user_auth = generate_mock_user_api_key_auth() + + # Mock health check result for authorized server + mock_health_result = generate_mock_mcp_server_db_record( + server_id="server-1", + alias="Server 1", + url="https://server1.example.com", + ) + mock_health_result.status = "healthy" + mock_health_result.last_health_check = datetime.now() + mock_health_result.health_check_error = None + + # Mock manager - server_ids filter is applied inside get_all_mcp_servers_with_health_and_teams + # So it only returns servers the user has access to + mock_manager = MagicMock() + mock_manager.get_all_mcp_servers_with_health_and_teams = AsyncMock( + return_value=[mock_health_result] # Only server-1 is returned (accessible) + ) + + with patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.global_mcp_server_manager", + mock_manager, + ), patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.build_effective_auth_contexts", + AsyncMock(return_value=[mock_user_auth]), + ): + result = await health_check_servers( + server_ids=["server-1", "server-unauthorized"], + user_api_key_dict=mock_user_auth, + ) + + # Verify results - only accessible server is returned + assert len(result) == 1 + assert result[0]["server_id"] == "server-1" + assert result[0]["status"] == "healthy" diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/mcpServers/useMCPServerHealth.test.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/mcpServers/useMCPServerHealth.test.ts new file mode 100644 index 00000000000..be910acf7e4 --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/mcpServers/useMCPServerHealth.test.ts @@ -0,0 +1,127 @@ +/* @vitest-environment jsdom */ +import React from "react"; +import { renderHook, waitFor } from "@testing-library/react"; +import { describe, it, expect, vi, beforeEach } from "vitest"; +import { QueryClient, QueryClientProvider } from "@tanstack/react-query"; +import { useMCPServerHealth } from "./useMCPServerHealth"; +import * as networking from "@/components/networking"; + +// Mock the networking module +vi.mock("@/components/networking", () => ({ + fetchMCPServerHealth: vi.fn(), +})); + +// Mock useAuthorized hook +vi.mock("../useAuthorized", () => ({ + default: vi.fn(() => ({ + accessToken: "test-token-123", + })), +})); + +const createQueryClient = () => + new QueryClient({ + defaultOptions: { + queries: { + retry: false, + gcTime: 0, + }, + }, + }); + +const wrapper = ({ children }: { children: React.ReactNode }) => { + const queryClient = createQueryClient(); + return React.createElement(QueryClientProvider, { client: queryClient }, children); +}; + +describe("useMCPServerHealth", () => { + beforeEach(() => { + vi.clearAllMocks(); + }); + + it("should fetch health status for given server IDs", async () => { + const mockHealthStatuses = [ + { server_id: "server-1", status: "healthy" }, + { server_id: "server-2", status: "unhealthy" }, + ]; + + vi.mocked(networking.fetchMCPServerHealth).mockResolvedValue(mockHealthStatuses); + + const { result } = renderHook(() => useMCPServerHealth(["server-1", "server-2"]), { + wrapper, + }); + + await waitFor(() => { + expect(result.current.isSuccess).toBe(true); + }); + + expect(networking.fetchMCPServerHealth).toHaveBeenCalledWith("test-token-123", ["server-1", "server-2"]); + expect(result.current.data).toEqual(mockHealthStatuses); + }); + + it("should fetch health status for all servers when no server IDs provided", async () => { + const mockHealthStatuses = [ + { server_id: "server-1", status: "healthy" }, + { server_id: "server-2", status: "healthy" }, + { server_id: "server-3", status: "unhealthy" }, + ]; + + vi.mocked(networking.fetchMCPServerHealth).mockResolvedValue(mockHealthStatuses); + + const { result } = renderHook(() => useMCPServerHealth(), { + wrapper, + }); + + await waitFor(() => { + expect(result.current.isSuccess).toBe(true); + }); + + expect(networking.fetchMCPServerHealth).toHaveBeenCalledWith("test-token-123", undefined); + expect(result.current.data).toEqual(mockHealthStatuses); + }); + + it("should handle empty server list", async () => { + vi.mocked(networking.fetchMCPServerHealth).mockResolvedValue([]); + + const { result } = renderHook(() => useMCPServerHealth([]), { + wrapper, + }); + + await waitFor(() => { + expect(result.current.isSuccess).toBe(true); + }); + + expect(networking.fetchMCPServerHealth).toHaveBeenCalledWith("test-token-123", []); + expect(result.current.data).toEqual([]); + }); + + it("should handle errors when fetching health status", async () => { + const mockError = new Error("Failed to fetch health status"); + vi.mocked(networking.fetchMCPServerHealth).mockRejectedValue(mockError); + + const { result } = renderHook(() => useMCPServerHealth(["server-1"]), { + wrapper, + }); + + await waitFor(() => { + expect(result.current.isError).toBe(true); + }); + + expect(result.current.error).toEqual(mockError); + }); + + it("should not fetch when accessToken is not available", async () => { + // Mock useAuthorized to return no token + const useAuthorizedModule = await import("../useAuthorized"); + vi.mocked(useAuthorizedModule.default).mockReturnValue({ + accessToken: null, + } as any); + + const { result } = renderHook(() => useMCPServerHealth(["server-1"]), { + wrapper, + }); + + // Should remain in idle state since query is not enabled + expect(result.current.status).toBe("pending"); + expect(networking.fetchMCPServerHealth).not.toHaveBeenCalled(); + }); +}); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/mcpServers/useMCPServerHealth.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/mcpServers/useMCPServerHealth.ts new file mode 100644 index 00000000000..ad3c633eb91 --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/mcpServers/useMCPServerHealth.ts @@ -0,0 +1,22 @@ +import { useQuery } from "@tanstack/react-query"; +import { createQueryKeys } from "../common/queryKeysFactory"; +import { fetchMCPServerHealth } from "@/components/networking"; +import useAuthorized from "../useAuthorized"; + +const mcpServerHealthKeys = createQueryKeys("mcpServerHealth"); + +interface MCPServerHealth { + server_id: string; + status: string; +} + +export const useMCPServerHealth = (serverIds?: string[]) => { + const { accessToken } = useAuthorized(); + return useQuery({ + queryKey: mcpServerHealthKeys.list({ serverIds }), + queryFn: async () => await fetchMCPServerHealth(accessToken!, serverIds), + enabled: !!accessToken, + // Refetch health status every 30 seconds to keep it up to date + refetchInterval: 30000, + }); +}; diff --git a/ui/litellm-dashboard/src/components/mcp_tools/mcp_server_columns.tsx b/ui/litellm-dashboard/src/components/mcp_tools/mcp_server_columns.tsx index 1bf719ef904..f6a5d6622d4 100644 --- a/ui/litellm-dashboard/src/components/mcp_tools/mcp_server_columns.tsx +++ b/ui/litellm-dashboard/src/components/mcp_tools/mcp_server_columns.tsx @@ -10,6 +10,7 @@ export const mcpServerColumns = ( onView: (serverId: string) => void, onEdit: (serverId: string) => void, onDelete: (serverId: string) => void, + isLoadingHealth?: boolean, ): ColumnDef[] => [ { accessorKey: "server_id", @@ -58,6 +59,19 @@ export const mcpServerColumns = ( const lastCheck = server.last_health_check; const error = server.health_check_error; + // Show loading spinner if health check is in progress + if (isLoadingHealth) { + return ( +
+ + + + + Loading... +
+ ); + } + const getStatusColor = (status: string) => { switch (status) { case "healthy": diff --git a/ui/litellm-dashboard/src/components/mcp_tools/mcp_servers.test.tsx b/ui/litellm-dashboard/src/components/mcp_tools/mcp_servers.test.tsx index 776d579fc13..4b8698b9762 100644 --- a/ui/litellm-dashboard/src/components/mcp_tools/mcp_servers.test.tsx +++ b/ui/litellm-dashboard/src/components/mcp_tools/mcp_servers.test.tsx @@ -8,6 +8,7 @@ import * as networking from "../networking"; // Mock the networking module vi.mock("../networking", () => ({ fetchMCPServers: vi.fn(), + fetchMCPServerHealth: vi.fn(), deleteMCPServer: vi.fn(), getProxyBaseUrl: vi.fn().mockReturnValue("http://localhost:4000"), })); @@ -123,4 +124,108 @@ describe("MCPServers", () => { // Note: useMCPServers uses useAuthorized() internally, which returns "123" from global mock expect(networking.fetchMCPServers).toHaveBeenCalledWith("123"); }); + + it("should fetch and merge health status for servers", async () => { + // Mock MCP servers data without health status + const mockServers = [ + { + server_id: "server-1", + server_name: "Test Server 1", + alias: "test-server-1", + url: "https://example.com/mcp", + transport: "http", + auth_type: "none", + created_at: "2024-01-01T00:00:00Z", + created_by: "user-1", + updated_at: "2024-01-01T00:00:00Z", + updated_by: "user-1", + teams: [], + mcp_access_groups: [], + status: undefined, + }, + { + server_id: "server-2", + server_name: "Test Server 2", + alias: "test-server-2", + url: "https://example2.com/mcp", + transport: "sse", + auth_type: "api_key", + created_at: "2024-01-02T00:00:00Z", + created_by: "user-2", + updated_at: "2024-01-02T00:00:00Z", + updated_by: "user-2", + teams: [], + mcp_access_groups: ["group-1"], + status: undefined, + }, + ]; + + // Mock health status data + const mockHealthStatuses = [ + { server_id: "server-1", status: "healthy" }, + { server_id: "server-2", status: "unhealthy" }, + ]; + + vi.mocked(networking.fetchMCPServers).mockResolvedValue(mockServers); + vi.mocked(networking.fetchMCPServerHealth).mockResolvedValue(mockHealthStatuses); + + const queryClient = createQueryClient(); + const { getByText } = render( + + + , + ); + + // Wait for the component to load + await waitFor(() => { + expect(getByText("MCP Servers")).toBeInTheDocument(); + }); + + // Verify the health check API was called with server IDs + await waitFor(() => { + expect(networking.fetchMCPServerHealth).toHaveBeenCalledWith("123", ["server-1", "server-2"]); + }); + }); + + it("should display loading state while health check is in progress", async () => { + const mockServers = [ + { + server_id: "server-1", + server_name: "Test Server 1", + alias: "test-server-1", + url: "https://example.com/mcp", + transport: "http", + auth_type: "none", + created_at: "2024-01-01T00:00:00Z", + created_by: "user-1", + updated_at: "2024-01-01T00:00:00Z", + updated_by: "user-1", + teams: [], + mcp_access_groups: [], + }, + ]; + + vi.mocked(networking.fetchMCPServers).mockResolvedValue(mockServers); + // Mock health check to never resolve (to test loading state) + vi.mocked(networking.fetchMCPServerHealth).mockImplementation( + () => new Promise(() => {}), // Never resolves + ); + + const queryClient = createQueryClient(); + const { getByText } = render( + + + , + ); + + // Wait for the component to load + await waitFor(() => { + expect(getByText("MCP Servers")).toBeInTheDocument(); + }); + + // Verify that health check was initiated + await waitFor(() => { + expect(networking.fetchMCPServerHealth).toHaveBeenCalled(); + }); + }); }); diff --git a/ui/litellm-dashboard/src/components/mcp_tools/mcp_servers.tsx b/ui/litellm-dashboard/src/components/mcp_tools/mcp_servers.tsx index d5b147f85c7..f6669fb2829 100644 --- a/ui/litellm-dashboard/src/components/mcp_tools/mcp_servers.tsx +++ b/ui/litellm-dashboard/src/components/mcp_tools/mcp_servers.tsx @@ -2,8 +2,9 @@ import { isAdminRole } from "@/utils/roles"; import { QuestionCircleOutlined } from "@ant-design/icons"; import { Button, Tab, TabGroup, TabList, TabPanel, TabPanels, Text, Title } from "@tremor/react"; import { Descriptions, Modal, Select, Tooltip, Typography } from "antd"; -import React, { useEffect, useState } from "react"; +import React, { useEffect, useState, useMemo } from "react"; import { useMCPServers } from "../../app/(dashboard)/hooks/mcpServers/useMCPServers"; +import { useMCPServerHealth } from "../../app/(dashboard)/hooks/mcpServers/useMCPServerHealth"; import NotificationsManager from "../molecules/notifications_manager"; import { deleteMCPServer } from "../networking"; import { DataTable } from "../view_logs/table"; @@ -19,7 +20,29 @@ const EDIT_OAUTH_UI_STATE_KEY = "litellm-mcp-oauth-edit-state"; const { Option } = Select; const MCPServers: React.FC = ({ accessToken, userRole, userID }) => { - const { data: mcpServers, isLoading: isLoadingServers, refetch, dataUpdatedAt } = useMCPServers(); + const { data: mcpServers, isLoading: isLoadingServers, refetch } = useMCPServers(); + + // Fetch health status for all servers + const serverIds = useMemo(() => mcpServers?.map((server) => server.server_id), [mcpServers]); + const { data: healthStatuses, isLoading: isLoadingHealth } = useMCPServerHealth(serverIds); + + // Merge health status data into servers + const serversWithHealth = useMemo(() => { + if (!mcpServers) return []; + if (!healthStatuses) return mcpServers; + + const healthMap = new Map(healthStatuses.map((h) => [h.server_id, h.status])); + + return mcpServers.map((server) => { + const healthStatus = healthMap.get(server.server_id); + return { + ...server, + status: healthStatus + ? (healthStatus as "healthy" | "unhealthy" | "unknown") + : server.status, + }; + }); + }, [mcpServers, healthStatuses]); // Log allowed_tools from fetched servers React.useEffect(() => { @@ -65,10 +88,10 @@ const MCPServers: React.FC = ({ accessToken, userRole, userID }) // Get unique teams from all servers const uniqueTeams = React.useMemo(() => { - if (!mcpServers) return []; + if (!serversWithHealth) return []; const teamsSet = new Set(); const uniqueTeamsArray: Team[] = []; - mcpServers.forEach((server: MCPServer) => { + serversWithHealth.forEach((server: MCPServer) => { if (server.teams) { server.teams.forEach((team: Team) => { const teamKey = team.team_id; @@ -80,17 +103,17 @@ const MCPServers: React.FC = ({ accessToken, userRole, userID }) } }); return uniqueTeamsArray; - }, [mcpServers]); + }, [serversWithHealth]); // Get unique MCP access groups from all servers const uniqueMcpAccessGroups = React.useMemo(() => { - if (!mcpServers) return []; + if (!serversWithHealth) return []; return Array.from( new Set( - mcpServers.flatMap((server) => server.mcp_access_groups).filter((group): group is string => group != null), + serversWithHealth.flatMap((server) => server.mcp_access_groups).filter((group): group is string => group != null), ), ); - }, [mcpServers]); + }, [serversWithHealth]); // Handle team filter change const handleTeamChange = (teamId: string) => { @@ -106,8 +129,8 @@ const MCPServers: React.FC = ({ accessToken, userRole, userID }) // Filtering logic for both team and access group const filterServers = (teamId: string, group: string) => { - if (!mcpServers) return setFilteredServers([]); - let filtered = mcpServers; + if (!serversWithHealth) return setFilteredServers([]); + let filtered = serversWithHealth; if (teamId === "personal") { setFilteredServers([]); return; @@ -123,10 +146,10 @@ const MCPServers: React.FC = ({ accessToken, userRole, userID }) setFilteredServers(filtered); }; - // Initial and effect-based filtering (trigger on query data updates) + // Initial and effect-based filtering (trigger on query data updates and health data updates) useEffect(() => { filterServers(selectedTeam, selectedMcpAccessGroup); - }, [dataUpdatedAt]); + }, [serversWithHealth, selectedTeam, selectedMcpAccessGroup]); const columns = React.useMemo( () => @@ -141,8 +164,9 @@ const MCPServers: React.FC = ({ accessToken, userRole, userID }) setEditServer(true); }, handleDelete, + isLoadingHealth, ), - [userRole], + [userRole, isLoadingHealth], ); function handleDelete(server_id: string) { diff --git a/ui/litellm-dashboard/src/components/networking.tsx b/ui/litellm-dashboard/src/components/networking.tsx index aeba207db26..e244a096131 100644 --- a/ui/litellm-dashboard/src/components/networking.tsx +++ b/ui/litellm-dashboard/src/components/networking.tsx @@ -5687,6 +5687,44 @@ export const fetchMCPServers = async (accessToken: string) => { } }; +export const fetchMCPServerHealth = async (accessToken: string, serverIds?: string[]) => { + try { + // Construct base URL + let url = proxyBaseUrl ? `${proxyBaseUrl}/v1/mcp/server/health` : `/v1/mcp/server/health`; + + // Add server_ids query parameters if provided + if (serverIds && serverIds.length > 0) { + const params = new URLSearchParams(); + serverIds.forEach((id) => params.append("server_ids", id)); + url = `${url}?${params.toString()}`; + } + + console.log("Fetching MCP server health from:", url); + + const response = await fetch(url, { + method: HTTP_REQUEST.GET, + headers: { + [globalLitellmHeaderName]: `Bearer ${accessToken}`, + "Content-Type": "application/json", + }, + }); + + if (!response.ok) { + const errorData = await response.json(); + const errorMessage = deriveErrorMessage(errorData); + handleError(errorMessage); + throw new Error(errorMessage); + } + + const data = await response.json(); + console.log("Fetched MCP server health:", data); + return data; + } catch (error) { + console.error("Failed to fetch MCP server health:", error); + throw error; + } +}; + export const fetchMCPAccessGroups = async (accessToken: string) => { try { // Construct base URL From 0dc8c1441f91385d157ff4fe2782510c2b3c6ab7 Mon Sep 17 00:00:00 2001 From: Yuta Saito Date: Wed, 31 Dec 2025 07:26:35 +0900 Subject: [PATCH 042/158] fix: npm build error --- .../src/app/(dashboard)/hooks/mcpServers/useMCPServerHealth.ts | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/mcpServers/useMCPServerHealth.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/mcpServers/useMCPServerHealth.ts index ad3c633eb91..95d7f3bcee0 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/hooks/mcpServers/useMCPServerHealth.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/mcpServers/useMCPServerHealth.ts @@ -13,7 +13,7 @@ interface MCPServerHealth { export const useMCPServerHealth = (serverIds?: string[]) => { const { accessToken } = useAuthorized(); return useQuery({ - queryKey: mcpServerHealthKeys.list({ serverIds }), + queryKey: [...mcpServerHealthKeys.lists(), { serverIds }], queryFn: async () => await fetchMCPServerHealth(accessToken!, serverIds), enabled: !!accessToken, // Refetch health status every 30 seconds to keep it up to date From ac798b7c92a51ebd34a0c110594d1b07499944f2 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Tue, 30 Dec 2025 14:53:15 -0800 Subject: [PATCH 043/158] Team admin view for add models --- .../ModelsAndEndpointsView.tsx | 28 +- .../src/components/add_model/AddModelForm.tsx | 445 ++++++++++-------- .../add_model/add_model_tab.test.tsx | 14 +- .../components/add_model/add_model_tab.tsx | 3 +- .../components/templates/model_dashboard.tsx | 3 +- 5 files changed, 263 insertions(+), 230 deletions(-) diff --git a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/ModelsAndEndpointsView.tsx b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/ModelsAndEndpointsView.tsx index 315b41393d6..4b62ce8cf88 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/ModelsAndEndpointsView.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/ModelsAndEndpointsView.tsx @@ -536,21 +536,19 @@ const ModelsAndEndpointsView: React.FC = ({ ); }; - const handleOk = () => { - addModelForm - .validateFields() - .then((values: any) => { - handleAddModelSubmit(values, accessToken, addModelForm, handleRefreshClick); - }) - .catch((error: any) => { - const errorMessages = - error.errorFields - ?.map((field: any) => { - return `${field.name.join(".")}: ${field.errors.join(", ")}`; - }) - .join(" | ") || "Unknown validation error"; - NotificationsManager.fromBackend(`Please fill in the following required fields: ${errorMessages}`); - }); + const handleOk = async () => { + try { + const values = await addModelForm.validateFields(); + await handleAddModelSubmit(values, accessToken, addModelForm, handleRefreshClick); + } catch (error: any) { + const errorMessages = + error.errorFields + ?.map((field: any) => { + return `${field.name.join(".")}: ${field.errors.join(", ")}`; + }) + .join(" | ") || "Unknown validation error"; + NotificationsManager.fromBackend(`Please fill in the following required fields: ${errorMessages}`); + } }; Object.keys(Providers).find((key) => (Providers as { [index: string]: any })[key] === selectedProvider); diff --git a/ui/litellm-dashboard/src/components/add_model/AddModelForm.tsx b/ui/litellm-dashboard/src/components/add_model/AddModelForm.tsx index 7ddc7840464..59ac63cffe6 100644 --- a/ui/litellm-dashboard/src/components/add_model/AddModelForm.tsx +++ b/ui/litellm-dashboard/src/components/add_model/AddModelForm.tsx @@ -1,10 +1,10 @@ import { useProviderFields } from "@/app/(dashboard)/hooks/providers/useProviderFields"; import { useGuardrails } from "@/app/(dashboard)/hooks/guardrails/useGuardrails"; import { useTags } from "@/app/(dashboard)/hooks/tags/useTags"; -import { all_admin_roles } from "@/utils/roles"; +import { all_admin_roles, isUserTeamAdminForAnyTeam } from "@/utils/roles"; import { Switch, Text } from "@tremor/react"; import type { FormInstance } from "antd"; -import { Select as AntdSelect, Button, Card, Col, Form, Modal, Row, Tooltip, Typography } from "antd"; +import { Select as AntdSelect, Button, Card, Col, Form, Modal, Row, Tooltip, Typography, Alert } from "antd"; import type { UploadProps } from "antd/es/upload"; import React, { useEffect, useMemo, useState } from "react"; import TeamDropdown from "../common_components/team_dropdown"; @@ -22,7 +22,7 @@ import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; interface AddModelFormProps { form: FormInstance; // For the Add Model tab - handleOk: () => void; + handleOk: () => Promise; selectedProvider: Providers; setSelectedProvider: (provider: Providers) => void; providerModels: string[]; @@ -57,7 +57,7 @@ const AddModelForm: React.FC = ({ // Using a unique ID to force the ConnectionErrorDisplay to remount and run a fresh test const [connectionTestId, setConnectionTestId] = useState(""); - const { accessToken, userRole, premiumUser } = useAuthorized(); + const { accessToken, userRole, premiumUser, userId } = useAuthorized(); const { data: providerMetadata, isLoading: isProviderMetadataLoading, @@ -73,8 +73,9 @@ const AddModelForm: React.FC = ({ }; const [isTeamOnly, setIsTeamOnly] = useState(false); - const [modelAccessGroups, setModelAccessGroups] = useState([]); + // Team admin specific state + const [teamAdminSelectedTeam, setTeamAdminSelectedTeam] = useState(null); useEffect(() => { const fetchModelAccessGroups = async () => { @@ -98,16 +99,20 @@ const AddModelForm: React.FC = ({ : null; const isAdmin = all_admin_roles.includes(userRole); + const isTeamAdmin = isUserTeamAdminForAnyTeam(teams, userId); return ( <> Add Model +
{ + onFinish={async (values) => { console.log("🔥 Form onFinish triggered with values:", values); - handleOk(); + await handleOk().then(() => { + setTeamAdminSelectedTeam(null); + }); }} onFinishFailed={(errorInfo) => { console.log("💥 Form onFinishFailed triggered:", errorInfo); @@ -117,215 +122,245 @@ const AddModelForm: React.FC = ({ labelAlign="left" > <> - {/* Provider Selection */} - - { - setSelectedProvider(value as Providers); - setProviderModelsFn(value as Providers); - form.setFieldsValue({ - custom_llm_provider: value, - }); - form.setFieldsValue({ - model: [], - model_name: undefined, - }); - }} - > - {providerMetadataErrorText && sortedProviderMetadata.length === 0 && ( - - {providerMetadataErrorText} - - )} - {sortedProviderMetadata.map((providerInfo) => { - const displayName = providerInfo.provider_display_name; - const providerKey = providerInfo.provider; - const logoSrc = providerLogoMap[displayName] ?? ""; - - return ( - -
- - {displayName} -
- - ); - })} - - - - - {/* Conditionally Render "Public Model Name" */} - - - {/* Select Mode */} - - setTestMode(value)} - options={TEST_MODES} - /> - - - - - - Optional - LiteLLM endpoint to use when health checking this model{" "} - - Learn more - - - - - - {/* Credentials */} -
- - Either select existing credentials OR enter new provider credentials below - -
- - - (option?.label ?? "").toLowerCase().includes(input.toLowerCase())} - options={[ - { value: null, label: "None" }, - ...credentials.map((credential) => ({ - value: credential.credential_name, - label: credential.credential_name, - })), - ]} - allowClear - /> - - - - prevValues.litellm_credential_name !== currentValues.litellm_credential_name || - prevValues.provider !== currentValues.provider - } - > - {({ getFieldValue }) => { - const credentialName = getFieldValue("litellm_credential_name"); - console.log("🔑 Credential Name Changed:", credentialName); - // Only show provider specific fields if no credentials selected - if (!credentialName) { - return ( - <> -
-
- OR -
-
- - - ); - } - return null; - }} - -
-
- Additional Model Info Settings -
-
- {/* Team-only Model Switch */} - - - { - setIsTeamOnly(checked); - if (!checked) { - form.setFieldValue("team_id", undefined); - } - }} - disabled={!premiumUser} - /> - - - - {/* Conditional Team Selection */} - {isTeamOnly && ( - - - - )} - {isAdmin && ( + {isTeamAdmin && !isAdmin && ( <> + { + setTeamAdminSelectedTeam(value); + }} + /> + + {!teamAdminSelectedTeam && ( + + )} + + )} + {(isAdmin || (isTeamAdmin && teamAdminSelectedTeam)) && ( + <> + { + setSelectedProvider(value as Providers); + setProviderModelsFn(value as Providers); + form.setFieldsValue({ + custom_llm_provider: value, + }); + form.setFieldsValue({ + model: [], + model_name: undefined, + }); + }} + > + {providerMetadataErrorText && sortedProviderMetadata.length === 0 && ( + + {providerMetadataErrorText} + + )} + {sortedProviderMetadata.map((providerInfo) => { + const displayName = providerInfo.provider_display_name; + const providerKey = providerInfo.provider; + const logoSrc = providerLogoMap[displayName] ?? ""; + + return ( + +
+ + {displayName} +
+ + ); + })} + + + + + {/* Conditionally Render "Public Model Name" */} + + + {/* Select Mode */} + + setTestMode(value)} + options={TEST_MODES} + /> + + + + + + Optional - LiteLLM endpoint to use when health checking this model{" "} + + Learn more + + + + + + {/* Credentials */} +
+ + Either select existing credentials OR enter new provider credentials below + +
+ + + ({ - value: group, - label: group, - }))} - maxTagCount="responsive" + filterOption={(input, option) => (option?.label ?? "").toLowerCase().includes(input.toLowerCase())} + options={[ + { value: null, label: "None" }, + ...credentials.map((credential) => ({ + value: credential.credential_name, + label: credential.credential_name, + })), + ]} allowClear /> + + + prevValues.litellm_credential_name !== currentValues.litellm_credential_name || + prevValues.provider !== currentValues.provider + } + > + {({ getFieldValue }) => { + const credentialName = getFieldValue("litellm_credential_name"); + console.log("🔑 Credential Name Changed:", credentialName); + // Only show provider specific fields if no credentials selected + if (!credentialName) { + return ( + <> +
+
+ OR +
+
+ + + ); + } + return null; + }} + +
+
+ Additional Model Info Settings +
+
+ {/* Team-only Model Switch - Only show for proxy admins, not team admins */} + {(isAdmin || !isTeamAdmin) && ( + + + { + setIsTeamOnly(checked); + if (!checked) { + form.setFieldValue("team_id", undefined); + } + }} + disabled={!premiumUser} + /> + + + )} + + {/* Conditional Team Selection */} + {isTeamOnly && (isAdmin || !isTeamAdmin) && ( + + + + )} + {isAdmin && ( + <> + + ({ + value: group, + label: group, + }))} + maxTagCount="responsive" + allowClear + /> + + + )} + )} - -
Need Help? diff --git a/ui/litellm-dashboard/src/components/add_model/add_model_tab.test.tsx b/ui/litellm-dashboard/src/components/add_model/add_model_tab.test.tsx index 9362cbb3260..197bcd6569f 100644 --- a/ui/litellm-dashboard/src/components/add_model/add_model_tab.test.tsx +++ b/ui/litellm-dashboard/src/components/add_model/add_model_tab.test.tsx @@ -62,6 +62,14 @@ vi.mock("@/app/(dashboard)/hooks/providers/useProviderFields", () => ({ }), })); +vi.mock("@/app/(dashboard)/hooks/useAuthorized", () => ({ + default: vi.fn().mockReturnValue({ + accessToken: "test-access-token", + userRole: "Admin", + premiumUser: true, + }), +})); + const createQueryClient = () => new QueryClient({ defaultOptions: { @@ -137,7 +145,6 @@ const createTestProps = () => { uploadProps, accessToken: "test-access-token", userRole: "Admin", - premiumUser: true, }; }; @@ -163,7 +170,6 @@ describe("Add Model Tab", () => { credentials={props.credentials} accessToken={props.accessToken} userRole={props.userRole} - premiumUser={props.premiumUser} /> , ); @@ -192,7 +198,6 @@ describe("Add Model Tab", () => { credentials={props.credentials} accessToken={props.accessToken} userRole={props.userRole} - premiumUser={props.premiumUser} /> , ); @@ -222,7 +227,6 @@ describe("Add Model Tab", () => { credentials={props.credentials} accessToken={props.accessToken} userRole={props.userRole} - premiumUser={props.premiumUser} /> , ); @@ -251,7 +255,6 @@ describe("Add Model Tab", () => { credentials={props.credentials} accessToken={props.accessToken} userRole={props.userRole} - premiumUser={props.premiumUser} /> , ); @@ -289,7 +292,6 @@ describe("Add Model Tab", () => { credentials={props.credentials} accessToken={props.accessToken} userRole={props.userRole} - premiumUser={props.premiumUser} /> , ); diff --git a/ui/litellm-dashboard/src/components/add_model/add_model_tab.tsx b/ui/litellm-dashboard/src/components/add_model/add_model_tab.tsx index 17ec1cbceee..f9b6533ac60 100644 --- a/ui/litellm-dashboard/src/components/add_model/add_model_tab.tsx +++ b/ui/litellm-dashboard/src/components/add_model/add_model_tab.tsx @@ -8,12 +8,11 @@ import { type CredentialItem } from "../networking"; import { Providers } from "../provider_info_helpers"; import AddAutoRouterTab from "./add_auto_router_tab"; import AddModelForm from "./AddModelForm"; -import AdvancedSettings from "./advanced_settings"; import { handleAddAutoRouterSubmit } from "./handle_add_auto_router_submit"; interface AddModelTabProps { form: FormInstance; // For the Add Model tab - handleOk: (values?: any) => void; + handleOk: (values?: any) => Promise; selectedProvider: Providers; setSelectedProvider: (provider: Providers) => void; providerModels: string[]; diff --git a/ui/litellm-dashboard/src/components/templates/model_dashboard.tsx b/ui/litellm-dashboard/src/components/templates/model_dashboard.tsx index 9dbe04bffb1..d0db4204807 100644 --- a/ui/litellm-dashboard/src/components/templates/model_dashboard.tsx +++ b/ui/litellm-dashboard/src/components/templates/model_dashboard.tsx @@ -950,7 +950,7 @@ const OldModelDashboard: React.FC = ({ ); }; - const handleOk = () => { + const handleOk = async () => { console.log("🚀 handleOk called from model dashboard!"); console.log("Current form values:", addModelForm.getFieldsValue()); @@ -1354,7 +1354,6 @@ const OldModelDashboard: React.FC = ({ credentials={credentialsList} accessToken={accessToken} userRole={userRole} - premiumUser={premiumUser} /> From 3723a7146279593bf5cb6db0feba17071fb5f371 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Tue, 30 Dec 2025 15:07:34 -0800 Subject: [PATCH 044/158] Adding tests --- .../add_model/AddModelForm.test.tsx | 276 ++++++++++++++++++ 1 file changed, 276 insertions(+) create mode 100644 ui/litellm-dashboard/src/components/add_model/AddModelForm.test.tsx diff --git a/ui/litellm-dashboard/src/components/add_model/AddModelForm.test.tsx b/ui/litellm-dashboard/src/components/add_model/AddModelForm.test.tsx new file mode 100644 index 00000000000..6f63fcfa549 --- /dev/null +++ b/ui/litellm-dashboard/src/components/add_model/AddModelForm.test.tsx @@ -0,0 +1,276 @@ +import { renderHook, screen, waitFor, renderWithProviders } from "../../../tests/test-utils"; +import userEvent from "@testing-library/user-event"; +import { Form } from "antd"; +import type { UploadProps } from "antd/es/upload"; +import { describe, expect, it, vi } from "vitest"; +import type { Team } from "../key_team_helpers/key_list"; +import type { CredentialItem } from "../networking"; +import { Providers } from "../provider_info_helpers"; +import AddModelForm from "./AddModelForm"; + +vi.mock("../molecules/models/ProviderLogo", () => ({ + ProviderLogo: ({ provider, className }: { provider: string; className?: string }) => ( +
+ {provider} +
+ ), +})); + +vi.mock("../networking", async () => { + const actual = await vi.importActual("../networking"); + return { + ...actual, + getGuardrailsList: vi.fn().mockResolvedValue({ + guardrails: [{ guardrail_name: "test-guardrail-1" }, { guardrail_name: "test-guardrail-2" }], + }), + tagListCall: vi.fn().mockResolvedValue({}), + modelAvailableCall: vi.fn().mockResolvedValue({ + data: [{ id: "model-group-1" }, { id: "model-group-2" }], + }), + modelHubCall: vi.fn().mockResolvedValue({ + data: [ + { model_group: "gpt-4", mode: "chat" }, + { model_group: "gpt-3.5-turbo", mode: "chat" }, + ], + }), + getProviderCreateMetadata: vi.fn().mockResolvedValue([ + { + provider: "OpenAI", + provider_display_name: "OpenAI", + litellm_provider: "openai", + default_model_placeholder: "gpt-3.5-turbo", + credential_fields: [], + }, + ]), + }; +}); + +vi.mock("@/app/(dashboard)/hooks/providers/useProviderFields", () => ({ + useProviderFields: vi.fn().mockReturnValue({ + data: [ + { + provider: "OpenAI", + provider_display_name: "OpenAI", + litellm_provider: "openai", + default_model_placeholder: "gpt-3.5-turbo", + credential_fields: [], + }, + ], + isLoading: false, + error: null, + }), +})); + +vi.mock("@/app/(dashboard)/hooks/useAuthorized", () => ({ + default: vi.fn(), +})); + +vi.mock("@/app/(dashboard)/hooks/guardrails/useGuardrails", () => ({ + useGuardrails: vi.fn().mockReturnValue({ + data: [{ guardrail_name: "test-guardrail" }], + isLoading: false, + error: null, + }), +})); + +vi.mock("@/app/(dashboard)/hooks/tags/useTags", () => ({ + useTags: vi.fn().mockReturnValue({ + data: { tag1: ["model1", "model2"] }, + isLoading: false, + error: null, + }), +})); + +const mockAuthorizedUser = (userRole: string, userId: string, premiumUser: boolean) => ({ + token: "test-token", + accessToken: "test-access-token", + userId, + userEmail: "test@example.com", + userRole, + premiumUser, + disabledPersonalKeyCreation: false, + showSSOBanner: false, +}); + +const testTeam: Team = { + team_id: "team-1", + team_alias: "Test Team", + models: ["gpt-4"], + max_budget: 100, + budget_duration: "monthly", + tpm_limit: null, + rpm_limit: null, + organization_id: "org-1", + created_at: "2024-01-01T00:00:00Z", + keys: [], + members_with_roles: [], +}; + +const createTestProps = (userRole = "proxy_admin", userId = "user-1", isTeamAdmin = false) => { + const { result } = renderHook(() => Form.useForm()); + const [form] = result.current; + + const teams = [ + { + ...testTeam, + members_with_roles: isTeamAdmin ? [{ user_id: userId, role: "admin" }] : [], + }, + ]; + + const credentials: CredentialItem[] = [ + { + credential_name: "test-credential", + credential_values: {}, + credential_info: { + custom_llm_provider: "openai", + description: "Test credential", + }, + }, + ]; + + const uploadProps: UploadProps = { + beforeUpload: () => false, + showUploadList: false, + }; + + return { + form, + handleOk: vi.fn(), + setSelectedProvider: vi.fn(), + setProviderModelsFn: vi.fn(), + getPlaceholder: vi.fn((provider: Providers) => `Enter ${provider} model name`), + setShowAdvancedSettings: vi.fn(), + selectedProvider: Providers.OpenAI, + providerModels: ["gpt-4", "gpt-3.5-turbo"], + showAdvancedSettings: false, + teams, + credentials, + uploadProps, + userRole, + userId, + }; +}; + +describe("AddModelForm", () => { + it("should render", async () => { + const mockUseAuthorized = vi.mocked(await import("@/app/(dashboard)/hooks/useAuthorized")); + mockUseAuthorized.default.mockReturnValue(mockAuthorizedUser("proxy_admin", "user-1", true)); + + const props = createTestProps(); + + renderWithProviders(); + + expect(await screen.findByRole("heading", { name: "Add Model" })).toBeInTheDocument(); + }); + + it("should show proxy admin only (not team admin) - should not see Select Team dropdown unless switch is toggled", async () => { + const mockUseAuthorized = vi.mocked(await import("@/app/(dashboard)/hooks/useAuthorized")); + mockUseAuthorized.default.mockReturnValue(mockAuthorizedUser("proxy_admin", "user-1", true)); + + const props = createTestProps("proxy_admin", "user-1", false); + + renderWithProviders(); + + await screen.findByText("Provider"); + + expect(screen.queryByText("Team Selection Required")).not.toBeInTheDocument(); + expect(screen.queryByText("Select Team")).not.toBeInTheDocument(); + + const teamSwitch = screen.getByRole("switch"); + expect(teamSwitch).toBeInTheDocument(); + + expect(screen.queryByText("Select Team")).not.toBeInTheDocument(); + + await userEvent.click(teamSwitch); + + expect(await screen.findByText("Select Team")).toBeInTheDocument(); + }); + + it("should show proxy admin who is also team admin - should not see Select Team dropdown unless switch is toggled", async () => { + const mockUseAuthorized = vi.mocked(await import("@/app/(dashboard)/hooks/useAuthorized")); + mockUseAuthorized.default.mockReturnValue(mockAuthorizedUser("proxy_admin", "user-1", true)); + + const props = createTestProps("proxy_admin", "user-1", true); + + renderWithProviders(); + + await screen.findByText("Provider"); + + expect(screen.queryByText("Team Selection Required")).not.toBeInTheDocument(); + expect(screen.queryByText("Select Team")).not.toBeInTheDocument(); + + const teamSwitch = screen.getByRole("switch"); + expect(teamSwitch).toBeInTheDocument(); + + expect(screen.queryByText("Select Team")).not.toBeInTheDocument(); + + await userEvent.click(teamSwitch); + + expect(await screen.findByText("Select Team")).toBeInTheDocument(); + }); + + it("should show team admin (not proxy admin) - should see alert and team select, must select team before seeing remaining fields", async () => { + const mockUseAuthorized = vi.mocked(await import("@/app/(dashboard)/hooks/useAuthorized")); + mockUseAuthorized.default.mockReturnValue(mockAuthorizedUser("team_member", "user-1", true)); + + const props = createTestProps("team_member", "user-1", true); + + renderWithProviders(); + + await screen.findByRole("heading", { name: "Add Model" }); + + expect(screen.getByText("Team Selection Required")).toBeInTheDocument(); + + expect(screen.getByText("Select Team")).toBeInTheDocument(); + + expect(screen.queryByText("Provider")).not.toBeInTheDocument(); + + const teamSelect = screen.getByRole("combobox"); + await userEvent.click(teamSelect); + await userEvent.click(screen.getByText("Test Team")); + + await waitFor(() => { + expect(screen.getByText("Provider")).toBeInTheDocument(); + }); + }); + + it("should show team admin (not proxy admin) - should not see team-BYOK switch", async () => { + const mockUseAuthorized = vi.mocked(await import("@/app/(dashboard)/hooks/useAuthorized")); + mockUseAuthorized.default.mockReturnValue(mockAuthorizedUser("team_member", "user-1", true)); + + const props = createTestProps("team_member", "user-1", true); + + renderWithProviders(); + + await screen.findByText("Select Team"); + + const teamSelect = screen.getByRole("combobox"); + await userEvent.click(teamSelect); + await userEvent.click(screen.getByText("Test Team")); + + await waitFor(() => { + expect(screen.getByText("Provider")).toBeInTheDocument(); + }); + + expect(screen.queryByRole("switch")).not.toBeInTheDocument(); + }); + + it("should handle non-admin, non-team-admin users - should not see team selection or switch", async () => { + const mockUseAuthorized = vi.mocked(await import("@/app/(dashboard)/hooks/useAuthorized")); + mockUseAuthorized.default.mockReturnValue(mockAuthorizedUser("user", "user-1", false)); + + const props = createTestProps("user", "user-1", false); + + renderWithProviders(); + + await screen.findByRole("heading", { name: "Add Model" }); + + expect(screen.queryByText("Team Selection Required")).not.toBeInTheDocument(); + + expect(screen.queryByText("Select Team")).not.toBeInTheDocument(); + + expect(screen.queryByText("Provider")).not.toBeInTheDocument(); + + expect(screen.queryByRole("switch")).not.toBeInTheDocument(); + }); +}); From 2abd0941be981fab55c85de8b1767e25f6e88a00 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Tue, 30 Dec 2025 16:44:31 -0800 Subject: [PATCH 045/158] E2E Parity --- .../e2e_tests/playwright.config.ts | 2 +- .../auth/unauthenticatedRedirect.spec.ts | 11 +++ .../e2e_tests/tests/users/searchUsers.spec.ts | 90 +++++++++++++++++++ .../tests/users/viewInternalUsers.spec.ts | 52 +++++++++++ 4 files changed, 154 insertions(+), 1 deletion(-) create mode 100644 ui/litellm-dashboard/e2e_tests/tests/auth/unauthenticatedRedirect.spec.ts create mode 100644 ui/litellm-dashboard/e2e_tests/tests/users/searchUsers.spec.ts create mode 100644 ui/litellm-dashboard/e2e_tests/tests/users/viewInternalUsers.spec.ts diff --git a/ui/litellm-dashboard/e2e_tests/playwright.config.ts b/ui/litellm-dashboard/e2e_tests/playwright.config.ts index 9df2281016c..329bb7f7afc 100644 --- a/ui/litellm-dashboard/e2e_tests/playwright.config.ts +++ b/ui/litellm-dashboard/e2e_tests/playwright.config.ts @@ -20,7 +20,7 @@ export default defineConfig({ /* Shared settings for all the projects below. See https://playwright.dev/docs/api/class-testoptions. */ use: { /* Base URL to use in actions like `await page.goto('/')`. */ - baseURL: "http://localhost:3000", + baseURL: "http://localhost:4000", /* Collect trace when retrying the failed test. See https://playwright.dev/docs/trace-viewer */ trace: "on-first-retry", diff --git a/ui/litellm-dashboard/e2e_tests/tests/auth/unauthenticatedRedirect.spec.ts b/ui/litellm-dashboard/e2e_tests/tests/auth/unauthenticatedRedirect.spec.ts new file mode 100644 index 00000000000..d8cc26f8642 --- /dev/null +++ b/ui/litellm-dashboard/e2e_tests/tests/auth/unauthenticatedRedirect.spec.ts @@ -0,0 +1,11 @@ +import { test, expect } from "@playwright/test"; + +test.describe("Authentication Checks", () => { + test("should redirect unauthenticated user from a protected page", async ({ page }) => { + const protectedPageUrl = "http://localhost:4000/ui?page=llm-playground"; + const expectedRedirectUrl = "http://localhost:4000/ui/login/"; + await page.goto(protectedPageUrl, { waitUntil: "domcontentloaded" }); + await expect(page).toHaveURL(expectedRedirectUrl); + await expect(page.getByRole("heading", { name: "Login" })).toBeVisible(); + }); +}); diff --git a/ui/litellm-dashboard/e2e_tests/tests/users/searchUsers.spec.ts b/ui/litellm-dashboard/e2e_tests/tests/users/searchUsers.spec.ts new file mode 100644 index 00000000000..57fee8247dd --- /dev/null +++ b/ui/litellm-dashboard/e2e_tests/tests/users/searchUsers.spec.ts @@ -0,0 +1,90 @@ +import { test, expect, Page } from "@playwright/test"; +test.describe("Internal Users Search", () => { + test.use({ storageState: "admin.storageState.json" }); + + async function goToInternalUsers(page: Page) { + await page.goto("http://localhost:4000/ui"); + + const tab = page.getByRole("menuitem", { name: "Internal User" }); + await expect(tab).toBeVisible(); + await tab.click(); + + await expect(page.locator("tbody tr").first()).toBeVisible(); + await expect(page.locator(".ant-skeleton")).toHaveCount(0); + } + + test("can search users by email", async ({ page }) => { + await goToInternalUsers(page); + + const rows = page.locator("tbody tr"); + const searchInput = page.getByPlaceholder("Search by email..."); + + await expect(searchInput).toBeVisible(); + + // Ensure initial data is loaded + const initialCount = await rows.count(); + expect(initialCount).toBeGreaterThan(0); + + // 🔹 Apply filter + wait for backend response + await Promise.all([ + page.waitForResponse( + (res) => + res.url().includes("/user/list") && + res.url().includes("user_email=test%40") && // encoded "test@" + res.status() === 200, + ), + searchInput.fill("test@"), + ]); + + const filteredCount = await rows.count(); + await expect(filteredCount).toBeLessThan(initialCount); + + // 🔹 Clear filter + wait for unfiltered request + await Promise.all([ + page.waitForResponse( + (res) => res.url().includes("/user/list") && !res.url().includes("user_email=") && res.status() === 200, + ), + searchInput.clear(), + ]); + + const resetCount = await rows.count(); + await expect(resetCount).toBe(initialCount); + }); + + test("can filter users by user ID and SSO ID", async ({ page }) => { + await goToInternalUsers(page); + const rows = page.locator("tbody tr"); + + // Ensure initial data is loaded + const initialCount = await rows.count(); + expect(initialCount).toBeGreaterThan(0); + + const filtersButton = page.getByRole("button", { + name: "Filters", + exact: true, + }); + await filtersButton.click(); + + const userIdInput = page.getByPlaceholder("Filter by User ID"); + const ssoIdInput = page.getByPlaceholder("Filter by SSO ID"); + await Promise.all([ + page.waitForResponse( + (res) => res.url().includes("/user/list") && res.url().includes("user_ids=user") && res.status() === 200, + ), + userIdInput.fill("user"), + ]); + + await Promise.all([ + page.waitForResponse( + (res) => + res.url().includes("/user/list") && + res.url().includes("user_ids=user") && + res.url().includes("sso_user_ids=sso") && + res.status() === 200, + ), + ssoIdInput.fill("sso"), + ]); + const combinedFilteredCount = await rows.count(); + await expect(combinedFilteredCount).toBeLessThan(initialCount); + }); +}); diff --git a/ui/litellm-dashboard/e2e_tests/tests/users/viewInternalUsers.spec.ts b/ui/litellm-dashboard/e2e_tests/tests/users/viewInternalUsers.spec.ts new file mode 100644 index 00000000000..980c7233e42 --- /dev/null +++ b/ui/litellm-dashboard/e2e_tests/tests/users/viewInternalUsers.spec.ts @@ -0,0 +1,52 @@ +import { test, expect, Page } from "@playwright/test"; + +test.describe("Internal Users Page", () => { + test.use({ storageState: "admin.storageState.json" }); + + async function goToInternalUsers(page: Page) { + await page.goto("http://localhost:4000/ui"); + + const internalUserTab = page.getByRole("menuitem", { name: "Internal User" }); + await expect(internalUserTab).toBeVisible(); + await internalUserTab.click(); + + const firstRow = page.locator("tbody tr").first(); + await expect(firstRow).toBeVisible(); + await expect(page.locator(".ant-skeleton")).toHaveCount(0); + } + + test("renders internal users table correctly", async ({ page }) => { + await goToInternalUsers(page); + + const rows = page.locator("tbody tr"); + const rowCount = await rows.count(); + expect(rowCount).toBeGreaterThan(0); + + const userIdHeader = page.getByRole("columnheader", { name: "User ID" }); + await expect(userIdHeader).toBeVisible(); + + const virtualKeysHeader = page.getByRole("columnheader", { name: "Virtual Keys" }); + await expect(virtualKeysHeader).toBeVisible(); + }); + + test("pagination controls work correctly", async ({ page }) => { + await goToInternalUsers(page); + + const paginationInfo = page.locator(".text-sm.text-gray-700"); + const prevButton = page.getByRole("button", { name: "Previous" }); + const nextButton = page.getByRole("button", { name: "Next" }); + + const infoText = (await paginationInfo.textContent()) || ""; + + // On first page, Previous should be disabled + if (infoText.includes("1 -")) { + await expect(prevButton).toBeDisabled(); + } + + // Check if there are more pages + const hasMorePages = infoText.includes("of") && !infoText.endsWith("25 of 25"); + if (hasMorePages) { + await expect(nextButton).toBeEnabled(); + } + }); +}); From b35e0d5f83da98ff76444705bb1cec1fc91860d7 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Tue, 30 Dec 2025 17:18:06 -0800 Subject: [PATCH 046/158] Adding timeout for flaky test --- ui/litellm-dashboard/e2e_tests/tests/users/searchUsers.spec.ts | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/ui/litellm-dashboard/e2e_tests/tests/users/searchUsers.spec.ts b/ui/litellm-dashboard/e2e_tests/tests/users/searchUsers.spec.ts index 57fee8247dd..5873bb3125c 100644 --- a/ui/litellm-dashboard/e2e_tests/tests/users/searchUsers.spec.ts +++ b/ui/litellm-dashboard/e2e_tests/tests/users/searchUsers.spec.ts @@ -35,7 +35,7 @@ test.describe("Internal Users Search", () => { ), searchInput.fill("test@"), ]); - + await page.waitForTimeout(5000); const filteredCount = await rows.count(); await expect(filteredCount).toBeLessThan(initialCount); From f3f1acd338fa0e348bf90de3617c56a882fc496a Mon Sep 17 00:00:00 2001 From: xuan07t2 Date: Wed, 31 Dec 2025 08:48:35 +0700 Subject: [PATCH 047/158] fix(vertex_ai): separate Tool objects for each tool type per API spec MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Fix Vertex AI API error: "tools[0].tool_type: one_of 'tool_type' has more than one initialized field" The Vertex AI API requires each Tool object to contain exactly one type of tool (e.g., FunctionDeclaration, GoogleSearch, CodeExecution). Previously, all tool types were combined into a single Tool object, causing INVALID_ARGUMENT errors when using multiple tools simultaneously. This change creates separate Tool objects for each tool type: - Function declarations in one Tool - Google Search in its own Tool - Code Execution in its own Tool - etc. Ref: https://cloud.google.com/vertex-ai/generative-ai/docs/reference/rest/v1beta1/Tool 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude Opus 4.5 --- .../vertex_and_google_ai_studio_gemini.py | 44 +++- ...test_vertex_and_google_ai_studio_gemini.py | 191 ++++++++++++++++++ 2 files changed, 224 insertions(+), 11 deletions(-) diff --git a/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py b/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py index a5cc3dca8c1..ae4fabc6b64 100644 --- a/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py +++ b/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py @@ -552,24 +552,46 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): "Invalid tool={}. Use `litellm.set_verbose` or `litellm --detailed_debug` to see raw request." ) - # Only include function_declarations if there are actual functions - _tools = Tools() +# Build list of Tool objects - each Tool should contain exactly one type + # per Vertex AI API spec: "A Tool object should contain exactly one type of Tool" + _tools_list: List[Tools] = [] + + # Function declarations can be grouped together in one Tool if gtool_func_declarations: - _tools["function_declarations"] = gtool_func_declarations + func_tool = Tools() + func_tool["function_declarations"] = gtool_func_declarations + _tools_list.append(func_tool) + + # Each special tool type must be in its own Tool object if googleSearch is not None: - _tools[VertexToolName.GOOGLE_SEARCH.value] = googleSearch + search_tool = Tools() + search_tool[VertexToolName.GOOGLE_SEARCH.value] = googleSearch + _tools_list.append(search_tool) if googleSearchRetrieval is not None: - _tools[VertexToolName.GOOGLE_SEARCH_RETRIEVAL.value] = googleSearchRetrieval + retrieval_tool = Tools() + retrieval_tool[VertexToolName.GOOGLE_SEARCH_RETRIEVAL.value] = googleSearchRetrieval + _tools_list.append(retrieval_tool) if enterpriseWebSearch is not None: - _tools[VertexToolName.ENTERPRISE_WEB_SEARCH.value] = enterpriseWebSearch + enterprise_tool = Tools() + enterprise_tool[VertexToolName.ENTERPRISE_WEB_SEARCH.value] = enterpriseWebSearch + _tools_list.append(enterprise_tool) if code_execution is not None: - _tools[VertexToolName.CODE_EXECUTION.value] = code_execution + code_tool = Tools() + code_tool[VertexToolName.CODE_EXECUTION.value] = code_execution + _tools_list.append(code_tool) if urlContext is not None: - _tools[VertexToolName.URL_CONTEXT.value] = urlContext + url_tool = Tools() + url_tool[VertexToolName.URL_CONTEXT.value] = urlContext + _tools_list.append(url_tool) if googleMaps is not None: - _tools[VertexToolName.GOOGLE_MAPS.value] = googleMaps + maps_tool = Tools() + maps_tool[VertexToolName.GOOGLE_MAPS.value] = googleMaps + _tools_list.append(maps_tool) if computerUse is not None: - _tools[VertexToolName.COMPUTER_USE.value] = computerUse + computer_tool = Tools() + computer_tool[VertexToolName.COMPUTER_USE.value] = computerUse + _tools_list.append(computer_tool) + # Add retrieval config to toolConfig if googleMaps has location data if google_maps_retrieval_config is not None: @@ -579,7 +601,7 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): "retrievalConfig" ] = google_maps_retrieval_config - return [_tools] + return _tools_list def _map_response_schema(self, value: dict) -> dict: old_schema = deepcopy(value) diff --git a/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py b/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py index 783d85f471b..91a28ee6ec9 100644 --- a/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py +++ b/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py @@ -2279,3 +2279,194 @@ def test_partial_json_chunk_on_first_chunk(): assert result is None, "Partial first chunk should return None" assert iterator.chunk_type == "accumulated_json", "Should switch to accumulated_json mode" + +# ==================== Tool Type Separation Tests ==================== +# These tests verify that each Tool object contains exactly one type per Vertex AI API spec +# Ref: https://cloud.google.com/vertex-ai/generative-ai/docs/reference/rest/v1beta1/Tool + + +def test_vertex_ai_multiple_tool_types_separate_objects(): + """ + Test that multiple tool types are placed in separate Tool objects. + + This is required by Vertex AI API spec: + "A Tool object should contain exactly one type of Tool" + + Related error without this fix: + "tools[0].tool_type: one_of 'tool_type' has more than one initialized field: + enterprise_web_search, url_context" + + Input: + value=[ + {"enterpriseWebSearch": {}}, + {"url_context": {}}, + ] + + Expected Output: + tools=[ + {"enterpriseWebSearch": {}}, # First Tool object + {"url_context": {}}, # Second Tool object (separate!) + ] + + NOT (incorrect - causes API error): + tools=[ + {"enterpriseWebSearch": {}, "url_context": {}} # Multiple types in one object + ] + """ + v = VertexGeminiConfig() + optional_params = {} + + tools = v._map_function( + value=[ + {"enterpriseWebSearch": {}}, + {"url_context": {}}, + ], + optional_params=optional_params + ) + + # Should have 2 separate Tool objects + assert len(tools) == 2, f"Expected 2 separate Tool objects, got {len(tools)}" + + # Each Tool object should contain exactly ONE type + tool_types_in_first = [k for k in tools[0].keys()] + tool_types_in_second = [k for k in tools[1].keys()] + + assert len(tool_types_in_first) == 1, f"First Tool should have exactly 1 type, got {tool_types_in_first}" + assert len(tool_types_in_second) == 1, f"Second Tool should have exactly 1 type, got {tool_types_in_second}" + + # Verify the correct tool types are present + assert "enterpriseWebSearch" in tools[0], "First Tool should contain enterpriseWebSearch" + assert "url_context" in tools[1], "Second Tool should contain url_context" + + +def test_vertex_ai_function_declarations_with_other_tools_separate(): + """ + Test that function declarations and other tool types are in separate Tool objects. + + This ensures that when using both function calling AND special tools like + google_search or code_execution, they are properly separated per API spec. + + Input: + value=[ + {"type": "function", "function": {"name": "get_weather", "description": "Get weather"}}, + {"googleSearch": {}}, + {"code_execution": {}}, + ] + + Expected Output: + tools=[ + {"function_declarations": [{"name": "get_weather", "description": "Get weather"}]}, + {"googleSearch": {}}, + {"code_execution": {}}, + ] + """ + v = VertexGeminiConfig() + optional_params = {} + + tools = v._map_function( + value=[ + {"type": "function", "function": {"name": "get_weather", "description": "Get weather"}}, + {"googleSearch": {}}, + {"code_execution": {}}, + ], + optional_params=optional_params + ) + + # Should have 3 separate Tool objects + assert len(tools) == 3, f"Expected 3 separate Tool objects, got {len(tools)}" + + # Find each tool type + func_tool = None + search_tool = None + code_tool = None + + for tool in tools: + if "function_declarations" in tool: + func_tool = tool + elif "googleSearch" in tool: + search_tool = tool + elif "code_execution" in tool: + code_tool = tool + + # Verify all tools are present and separate + assert func_tool is not None, "function_declarations Tool should be present" + assert search_tool is not None, "googleSearch Tool should be present" + assert code_tool is not None, "code_execution Tool should be present" + + # Verify each Tool has exactly one type + assert len(func_tool.keys()) == 1, "function_declarations Tool should have only one key" + assert len(search_tool.keys()) == 1, "googleSearch Tool should have only one key" + assert len(code_tool.keys()) == 1, "code_execution Tool should have only one key" + + # Verify function declaration content + assert func_tool["function_declarations"][0]["name"] == "get_weather" + + +def test_vertex_ai_single_tool_type_still_works(): + """ + Test that single tool type usage still works correctly (backward compatibility). + + Input: + value=[{"code_execution": {}}] + + Expected Output: + tools=[{"code_execution": {}}] + """ + v = VertexGeminiConfig() + optional_params = {} + + tools = v._map_function( + value=[{"code_execution": {}}], + optional_params=optional_params + ) + + assert len(tools) == 1 + assert "code_execution" in tools[0] + assert tools[0]["code_execution"] == {} + + +def test_vertex_ai_multiple_function_declarations_grouped(): + """ + Test that multiple function declarations are grouped in ONE Tool object. + + Function declarations are the exception - they CAN be grouped together + in a single Tool object (up to 512 declarations). + + Input: + value=[ + {"type": "function", "function": {"name": "func1", "description": "First function"}}, + {"type": "function", "function": {"name": "func2", "description": "Second function"}}, + ] + + Expected Output: + tools=[ + { + "function_declarations": [ + {"name": "func1", "description": "First function"}, + {"name": "func2", "description": "Second function"}, + ] + } + ] + """ + v = VertexGeminiConfig() + optional_params = {} + + tools = v._map_function( + value=[ + {"type": "function", "function": {"name": "func1", "description": "First function"}}, + {"type": "function", "function": {"name": "func2", "description": "Second function"}}, + ], + optional_params=optional_params + ) + + # Should have only 1 Tool object (function declarations grouped) + assert len(tools) == 1, f"Expected 1 Tool object for grouped functions, got {len(tools)}" + + # Should contain function_declarations with 2 functions + assert "function_declarations" in tools[0] + assert len(tools[0]["function_declarations"]) == 2 + + # Verify function names + func_names = [f["name"] for f in tools[0]["function_declarations"]] + assert "func1" in func_names + assert "func2" in func_names From 316a498965df40c2b999b1eb59ef20ad6a9fb748 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Tue, 30 Dec 2025 18:36:26 -0800 Subject: [PATCH 048/158] New badge reusable component --- .../components/common_components/NewBadge.tsx | 5 +++++ .../src/components/leftnav.tsx | 21 ++++++++----------- .../src/components/settings.tsx | 5 ++++- 3 files changed, 18 insertions(+), 13 deletions(-) create mode 100644 ui/litellm-dashboard/src/components/common_components/NewBadge.tsx diff --git a/ui/litellm-dashboard/src/components/common_components/NewBadge.tsx b/ui/litellm-dashboard/src/components/common_components/NewBadge.tsx new file mode 100644 index 00000000000..2a3d1a51248 --- /dev/null +++ b/ui/litellm-dashboard/src/components/common_components/NewBadge.tsx @@ -0,0 +1,5 @@ +import { Badge } from "antd"; + +export default function NewBadge() { + return ; +} diff --git a/ui/litellm-dashboard/src/components/leftnav.tsx b/ui/litellm-dashboard/src/components/leftnav.tsx index 11e957ae6d4..fc248ee049a 100644 --- a/ui/litellm-dashboard/src/components/leftnav.tsx +++ b/ui/litellm-dashboard/src/components/leftnav.tsx @@ -24,11 +24,12 @@ import { UserOutlined, } from "@ant-design/icons"; import type { MenuProps } from "antd"; -import { Badge, ConfigProvider, Layout, Menu } from "antd"; +import { ConfigProvider, Layout, Menu } from "antd"; import { useMemo } from "react"; import { all_admin_roles, internalUserRoles, isAdminRole, rolesWithWriteAccess } from "../utils/roles"; import type { Organization } from "./networking"; import UsageIndicator from "./usage_indicator"; +import NewBadge from "./common_components/NewBadge"; const { Sider } = Layout; // Define the props type @@ -103,11 +104,7 @@ const Sidebar: React.FC = ({ setPage, defaultSelectedKey, collapse { key: "agents", page: "agents", - label: ( - - Agents - - ), + label: Agents, icon: , roles: rolesWithWriteAccess, }, @@ -155,11 +152,7 @@ const Sidebar: React.FC = ({ setPage, defaultSelectedKey, collapse page: "new_usage", icon: , roles: [...all_admin_roles, ...internalUserRoles], - label: ( - - Usage - - ), + label: Usage, }, { key: "logs", @@ -267,7 +260,11 @@ const Sidebar: React.FC = ({ setPage, defaultSelectedKey, collapse { key: "settings", page: "settings", - label: "Settings", + label: ( + + Settings + + ), icon: , roles: all_admin_roles, children: [ diff --git a/ui/litellm-dashboard/src/components/settings.tsx b/ui/litellm-dashboard/src/components/settings.tsx index dfe141b2a4b..735ee6ee75c 100644 --- a/ui/litellm-dashboard/src/components/settings.tsx +++ b/ui/litellm-dashboard/src/components/settings.tsx @@ -40,6 +40,7 @@ import { AlertingObject } from "./Settings/LoggingAndAlerts/LoggingCallbacks/typ import { parseErrorMessage } from "./shared/errorUtils"; import DeleteResourceModal from "./common_components/DeleteResourceModal"; import CloudZeroCostTracking from "./CloudZeroCostTracking/CloudZeroCostTracking"; +import NewBadge from "./common_components/NewBadge"; interface SettingsPageProps { accessToken: string | null; userRole: string | null; @@ -569,7 +570,9 @@ const Settings: React.FC = ({ accessToken, userRole, userID, Logging Callbacks - CloudZero Cost Tracking + + CloudZero Cost Tracking + Alerting Types Alerting Settings Email Alerts From a17980c744f9945ba0f48f1a26d8d63f5d05ee91 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Wed, 31 Dec 2025 10:42:34 -0800 Subject: [PATCH 049/158] feat: hide new badges --- .../hooks/useDisableShowNewBadge.ts | 35 ++++ .../UsagePage/components/UsagePageView.tsx | 14 +- .../common_components/NewBadge.test.tsx | 52 ++++++ .../components/common_components/NewBadge.tsx | 17 +- .../src/components/navbar.tsx | 40 ++++- .../src/utils/localStorageUtils.test.ts | 157 ++++++++++++++++++ .../src/utils/localStorageUtils.ts | 33 ++++ 7 files changed, 331 insertions(+), 17 deletions(-) create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/hooks/useDisableShowNewBadge.ts create mode 100644 ui/litellm-dashboard/src/components/common_components/NewBadge.test.tsx create mode 100644 ui/litellm-dashboard/src/utils/localStorageUtils.test.ts create mode 100644 ui/litellm-dashboard/src/utils/localStorageUtils.ts diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/useDisableShowNewBadge.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/useDisableShowNewBadge.ts new file mode 100644 index 00000000000..d0a618e27ba --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/useDisableShowNewBadge.ts @@ -0,0 +1,35 @@ +// hooks/useDisableShowNewBadge.ts +import { useSyncExternalStore } from "react"; +import { getLocalStorageItem } from "@/utils/localStorageUtils"; +import { LOCAL_STORAGE_EVENT } from "@/utils/localStorageUtils"; + +function subscribe(callback: () => void) { + const onStorage = (e: StorageEvent) => { + if (e.key === "disableShowNewBadge") { + callback(); + } + }; + + const onCustom = (e: Event) => { + const { key } = (e as CustomEvent).detail; + if (key === "disableShowNewBadge") { + callback(); + } + }; + + window.addEventListener("storage", onStorage); + window.addEventListener(LOCAL_STORAGE_EVENT, onCustom); + + return () => { + window.removeEventListener("storage", onStorage); + window.removeEventListener(LOCAL_STORAGE_EVENT, onCustom); + }; +} + +function getSnapshot() { + return getLocalStorageItem("disableShowNewBadge") === "true"; +} + +export function useDisableShowNewBadge() { + return useSyncExternalStore(subscribe, getSnapshot); +} diff --git a/ui/litellm-dashboard/src/components/UsagePage/components/UsagePageView.tsx b/ui/litellm-dashboard/src/components/UsagePage/components/UsagePageView.tsx index 920955138b9..a7107814113 100644 --- a/ui/litellm-dashboard/src/components/UsagePage/components/UsagePageView.tsx +++ b/ui/litellm-dashboard/src/components/UsagePage/components/UsagePageView.tsx @@ -27,7 +27,7 @@ import { Text, Title, } from "@tremor/react"; -import { Alert, Badge } from "antd"; +import { Alert } from "antd"; import React, { useCallback, useEffect, useMemo, useState } from "react"; import { useAgents } from "@/app/(dashboard)/hooks/agents/useAgents"; @@ -419,13 +419,11 @@ const UsagePage: React.FC = ({ teams, organizations }) => {
- - setUsageView(value)} - isAdmin={all_admin_roles.includes(userRole || "")} - /> - + setUsageView(value)} + isAdmin={all_admin_roles.includes(userRole || "")} + />
{/* Your Usage Panel */} diff --git a/ui/litellm-dashboard/src/components/common_components/NewBadge.test.tsx b/ui/litellm-dashboard/src/components/common_components/NewBadge.test.tsx new file mode 100644 index 00000000000..a24b52db6b8 --- /dev/null +++ b/ui/litellm-dashboard/src/components/common_components/NewBadge.test.tsx @@ -0,0 +1,52 @@ +import { render, screen } from "@testing-library/react"; +import { describe, it, expect, vi, beforeEach } from "vitest"; +import NewBadge from "./NewBadge"; + +// Mock the hook directly +vi.mock("@/app/(dashboard)/hooks/useDisableShowNewBadge", () => ({ + useDisableShowNewBadge: vi.fn(), +})); + +import { useDisableShowNewBadge } from "@/app/(dashboard)/hooks/useDisableShowNewBadge"; + +const mockUseDisableShowNewBadge = vi.mocked(useDisableShowNewBadge); + +describe("NewBadge", () => { + beforeEach(() => { + vi.clearAllMocks(); + }); + + it("should render the badge when disableShowNewBadge is false", () => { + mockUseDisableShowNewBadge.mockReturnValue(false); + + render(Test Content); + + expect(screen.getByText("New")).toBeInTheDocument(); + expect(screen.getByText("Test Content")).toBeInTheDocument(); + }); + + it("should render the badge when disableShowNewBadge is not set", () => { + mockUseDisableShowNewBadge.mockReturnValue(false); + + render(); + + expect(screen.getByText("New")).toBeInTheDocument(); + }); + + it("should render only children when disableShowNewBadge is true", () => { + mockUseDisableShowNewBadge.mockReturnValue(true); + + render(Test Content); + + expect(screen.queryByText("New")).not.toBeInTheDocument(); + expect(screen.getByText("Test Content")).toBeInTheDocument(); + }); + + it("should render nothing when disableShowNewBadge is true and no children", () => { + mockUseDisableShowNewBadge.mockReturnValue(true); + + const { container } = render(); + + expect(container.firstChild).toBeNull(); + }); +}); diff --git a/ui/litellm-dashboard/src/components/common_components/NewBadge.tsx b/ui/litellm-dashboard/src/components/common_components/NewBadge.tsx index 2a3d1a51248..97cdea8cfbb 100644 --- a/ui/litellm-dashboard/src/components/common_components/NewBadge.tsx +++ b/ui/litellm-dashboard/src/components/common_components/NewBadge.tsx @@ -1,5 +1,18 @@ import { Badge } from "antd"; +import { useDisableShowNewBadge } from "@/app/(dashboard)/hooks/useDisableShowNewBadge"; -export default function NewBadge() { - return ; +export default function NewBadge({ children }: { children?: React.ReactNode }) { + const disableShowNewBadge = useDisableShowNewBadge(); + + if (disableShowNewBadge) { + return children ? <>{children} : null; + } + + return children ? ( + + {children} + + ) : ( + + ); } diff --git a/ui/litellm-dashboard/src/components/navbar.tsx b/ui/litellm-dashboard/src/components/navbar.tsx index e032649ae29..23b31ec4821 100644 --- a/ui/litellm-dashboard/src/components/navbar.tsx +++ b/ui/litellm-dashboard/src/components/navbar.tsx @@ -1,7 +1,7 @@ import Link from "next/link"; import React, { useState, useEffect } from "react"; import type { MenuProps } from "antd"; -import { Dropdown, Tooltip } from "antd"; +import { Dropdown, Switch, Tooltip } from "antd"; import { getProxyBaseUrl } from "@/components/networking"; import { UserOutlined, @@ -13,8 +13,10 @@ import { MenuUnfoldOutlined, } from "@ant-design/icons"; import { clearTokenCookies } from "@/utils/cookieUtils"; +import { getLocalStorageItem, setLocalStorageItem, removeLocalStorageItem } from "@/utils/localStorageUtils"; import { fetchProxySettings } from "@/utils/proxyUtils"; import { useTheme } from "@/contexts/ThemeContext"; +import { emitLocalStorageChange } from "@/utils/localStorageUtils"; interface NavbarProps { userID: string | null; @@ -44,6 +46,7 @@ const Navbar: React.FC = ({ const baseUrl = getProxyBaseUrl(); const [logoutUrl, setLogoutUrl] = useState(""); const [version, setVersion] = useState(""); + const [disableShowNewBadge, setDisableShowNewBadge] = useState(false); const { logoUrl } = useTheme(); // Simple logo URL: use custom logo if available, otherwise default @@ -79,6 +82,11 @@ const Navbar: React.FC = ({ initializeProxySettings(); }, [accessToken]); + useEffect(() => { + const storedValue = getLocalStorageItem("disableShowNewBadge"); + setDisableShowNewBadge(storedValue === "true"); + }, []); + useEffect(() => { setLogoutUrl(proxySettings?.PROXY_LOGOUT_URL || ""); }, [proxySettings]); @@ -129,6 +137,28 @@ const Navbar: React.FC = ({ {userEmail || "Unknown"}
+
e.stopPropagation()} + > + Hide New Feature Indicators + { + setDisableShowNewBadge(checked); + if (checked) { + setLocalStorageItem("disableShowNewBadge", "true"); + emitLocalStorageChange("disableShowNewBadge"); + } else { + removeLocalStorageItem("disableShowNewBadge"); + emitLocalStorageChange("disableShowNewBadge"); + } + }} + aria-label="Toggle hide new feature indicators" + /> +
), @@ -148,11 +178,7 @@ const Navbar: React.FC = ({
From 4ea3c527c2f67b2839d594c0146abbabffc75b14 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Thu, 1 Jan 2026 17:23:28 -0800 Subject: [PATCH 081/158] Clicking on logo goes to correct url --- .../healthReadiness/useHealthReadiness.ts | 27 +++ .../src/components/navbar.test.tsx | 181 ++++++++++++++++++ .../src/components/navbar.tsx | 23 +-- 3 files changed, 213 insertions(+), 18 deletions(-) create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/hooks/healthReadiness/useHealthReadiness.ts create mode 100644 ui/litellm-dashboard/src/components/navbar.test.tsx diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/healthReadiness/useHealthReadiness.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/healthReadiness/useHealthReadiness.ts new file mode 100644 index 00000000000..e7e599d08be --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/healthReadiness/useHealthReadiness.ts @@ -0,0 +1,27 @@ +import { useQuery, UseQueryResult } from "@tanstack/react-query"; +import { createQueryKeys } from "../common/queryKeysFactory"; +import { getProxyBaseUrl } from "@/components/networking"; + +const healthReadinessKeys = createQueryKeys("healthReadiness"); + +interface HealthReadinessResponse { + litellm_version?: string; + [key: string]: any; +} + +const fetchHealthReadiness = async (): Promise => { + const baseUrl = getProxyBaseUrl(); + const response = await fetch(`${baseUrl}/health/readiness`); + if (!response.ok) { + throw new Error(`Failed to fetch health readiness: ${response.statusText}`); + } + return response.json(); +}; + +export const useHealthReadiness = (): UseQueryResult => { + return useQuery({ + queryKey: healthReadinessKeys.detail("readiness"), + queryFn: fetchHealthReadiness, + staleTime: 5 * 60 * 1000, // 5 minutes + }); +}; diff --git a/ui/litellm-dashboard/src/components/navbar.test.tsx b/ui/litellm-dashboard/src/components/navbar.test.tsx new file mode 100644 index 00000000000..711f17da32e --- /dev/null +++ b/ui/litellm-dashboard/src/components/navbar.test.tsx @@ -0,0 +1,181 @@ +import { describe, it, expect, vi } from "vitest"; +import { renderWithProviders, screen, waitFor } from "../../tests/test-utils"; +import userEvent from "@testing-library/user-event"; +import Navbar from "./navbar"; + +// Mock the hooks and utilities +vi.mock("@/components/networking", () => ({ + getProxyBaseUrl: vi.fn(() => "http://localhost:4000"), +})); + +vi.mock("@/utils/proxyUtils", () => ({ + fetchProxySettings: vi.fn(), +})); + +// Create mock functions that can be controlled in tests +let mockUseThemeImpl = () => ({ logoUrl: null as string | null }); +let mockUseHealthReadinessImpl = () => ({ data: null as any }); +let mockGetLocalStorageItemImpl = () => null as string | null; + +vi.mock("@/contexts/ThemeContext", () => ({ + useTheme: () => mockUseThemeImpl(), +})); + +vi.mock("@/app/(dashboard)/hooks/healthReadiness/useHealthReadiness", () => ({ + useHealthReadiness: () => mockUseHealthReadinessImpl(), +})); + +vi.mock("@/utils/localStorageUtils", () => ({ + getLocalStorageItem: () => mockGetLocalStorageItemImpl(), + setLocalStorageItem: vi.fn(), + removeLocalStorageItem: vi.fn(), + emitLocalStorageChange: vi.fn(), +})); + +vi.mock("@/utils/cookieUtils", () => ({ + clearTokenCookies: vi.fn(), +})); + +// Mock window.location.href for logout testing +Object.defineProperty(window, "location", { + value: { href: "" }, + writable: true, +}); + +describe("Navbar", () => { + const defaultProps = { + userID: "test-user", + userEmail: "test@example.com", + userRole: "Admin", + premiumUser: false, + proxySettings: {}, + setProxySettings: vi.fn(), + accessToken: "test-token", + isPublicPage: false, + }; + + it("should render without crashing", () => { + renderWithProviders(); + + expect(screen.getByText("Docs")).toBeInTheDocument(); + expect(screen.getByText("User")).toBeInTheDocument(); + }); + + it("should display user information in dropdown", async () => { + const user = userEvent.setup(); + renderWithProviders(); + + await user.click(screen.getByText("User")); + + await waitFor(() => { + expect(screen.getByText("test-user")).toBeInTheDocument(); + }); + expect(screen.getByText("Admin")).toBeInTheDocument(); + expect(screen.getByText("test@example.com")).toBeInTheDocument(); + }); + + it("should show sidebar toggle button when onToggleSidebar is provided", () => { + const mockToggle = vi.fn(); + renderWithProviders(); + + const toggleButton = screen.getByTitle("Collapse sidebar"); + expect(toggleButton).toBeInTheDocument(); + }); + + it("should call onToggleSidebar when sidebar button is clicked", async () => { + const mockToggle = vi.fn(); + const user = userEvent.setup(); + renderWithProviders(); + + const toggleButton = screen.getByTitle("Collapse sidebar"); + await user.click(toggleButton); + + expect(mockToggle).toHaveBeenCalledTimes(1); + }); + + it("should show premium user badge when premiumUser is true", async () => { + const user = userEvent.setup(); + const premiumProps = { ...defaultProps, premiumUser: true }; + renderWithProviders(); + + await user.click(screen.getByText("User")); + + await waitFor(() => { + expect(screen.getByText("Premium")).toBeInTheDocument(); + }); + }); + + it("should show version badge when health data contains version", () => { + mockUseHealthReadinessImpl = () => ({ data: { litellm_version: "1.0.0" } }); + + renderWithProviders(); + + expect(screen.getByText("v1.0.0")).toBeInTheDocument(); + + // Reset mock + mockUseHealthReadinessImpl = () => ({ data: null }); + }); + + it("should use custom logo from theme context", () => { + mockUseThemeImpl = () => ({ logoUrl: "https://example.com/custom-logo.png" }); + + renderWithProviders(); + + const logoImg = screen.getByAltText("LiteLLM Brand"); + expect(logoImg).toHaveAttribute("src", "https://example.com/custom-logo.png"); + + // Reset mock + mockUseThemeImpl = () => ({ logoUrl: null }); + }); + + it("should hide user dropdown on public pages", () => { + const publicPageProps = { ...defaultProps, isPublicPage: true }; + renderWithProviders(); + + expect(screen.queryByText("User")).not.toBeInTheDocument(); + }); + + it("should handle hide new features toggle", async () => { + const user = userEvent.setup(); + + // Initially disabled + mockGetLocalStorageItemImpl = () => "false"; + + renderWithProviders(); + + await user.click(screen.getByText("User")); + + await waitFor(() => { + expect(screen.getByText("test-user")).toBeInTheDocument(); + }); + + // Find and click the toggle switch + const toggleSwitch = screen.getByLabelText("Toggle hide new feature indicators"); + await user.click(toggleSwitch); + + // The functions are mocked globally, so we can check if they were called + // by accessing them through the mock registry + const localStorageUtils = vi.mocked(await import("@/utils/localStorageUtils")); + expect(localStorageUtils.setLocalStorageItem).toHaveBeenCalledWith("disableShowNewBadge", "true"); + expect(localStorageUtils.emitLocalStorageChange).toHaveBeenCalledWith("disableShowNewBadge"); + }); + + it("should handle logout functionality", async () => { + const user = userEvent.setup(); + + renderWithProviders(); + + await user.click(screen.getByText("User")); + + await waitFor(() => { + expect(screen.getByText("test-user")).toBeInTheDocument(); + }); + + // Click logout + await user.click(screen.getByText("Logout")); + + const cookieUtils = vi.mocked(await import("@/utils/cookieUtils")); + expect(cookieUtils.clearTokenCookies).toHaveBeenCalled(); + expect(window.location.href).toBe(""); + }); +}); diff --git a/ui/litellm-dashboard/src/components/navbar.tsx b/ui/litellm-dashboard/src/components/navbar.tsx index 23b31ec4821..7d5ee3dd6f4 100644 --- a/ui/litellm-dashboard/src/components/navbar.tsx +++ b/ui/litellm-dashboard/src/components/navbar.tsx @@ -17,6 +17,7 @@ import { getLocalStorageItem, setLocalStorageItem, removeLocalStorageItem } from import { fetchProxySettings } from "@/utils/proxyUtils"; import { useTheme } from "@/contexts/ThemeContext"; import { emitLocalStorageChange } from "@/utils/localStorageUtils"; +import { useHealthReadiness } from "@/app/(dashboard)/hooks/healthReadiness/useHealthReadiness"; interface NavbarProps { userID: string | null; @@ -44,30 +45,16 @@ const Navbar: React.FC = ({ onToggleSidebar, }) => { const baseUrl = getProxyBaseUrl(); + console.log("baseUrl", baseUrl); const [logoutUrl, setLogoutUrl] = useState(""); - const [version, setVersion] = useState(""); const [disableShowNewBadge, setDisableShowNewBadge] = useState(false); const { logoUrl } = useTheme(); + const { data: healthData } = useHealthReadiness(); + const version = healthData?.litellm_version; // Simple logo URL: use custom logo if available, otherwise default const imageUrl = logoUrl || `${baseUrl}/get_image`; - useEffect(() => { - const fetchVersion = async () => { - try { - const response = await fetch(`${baseUrl}/health/readiness`); - const data = await response.json(); - if (data.litellm_version) { - setVersion(data.litellm_version); - } - } catch (error) { - console.error("Failed to fetch version:", error); - } - }; - - fetchVersion(); - }, [baseUrl]); - useEffect(() => { const initializeProxySettings = async () => { if (accessToken) { @@ -190,7 +177,7 @@ const Navbar: React.FC = ({ )}
- +
LiteLLM Brand Date: Thu, 1 Jan 2026 17:24:05 -0800 Subject: [PATCH 082/158] fixing build --- .../healthReadiness/useHealthReadiness.ts | 2 +- .../src/components/navbar.test.tsx | 4 +-- .../src/components/navbar.tsx | 30 +++++++++++-------- 3 files changed, 20 insertions(+), 16 deletions(-) diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/healthReadiness/useHealthReadiness.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/healthReadiness/useHealthReadiness.ts index e7e599d08be..db394b9f7f8 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/hooks/healthReadiness/useHealthReadiness.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/healthReadiness/useHealthReadiness.ts @@ -1,6 +1,6 @@ +import { getProxyBaseUrl } from "@/components/networking"; import { useQuery, UseQueryResult } from "@tanstack/react-query"; import { createQueryKeys } from "../common/queryKeysFactory"; -import { getProxyBaseUrl } from "@/components/networking"; const healthReadinessKeys = createQueryKeys("healthReadiness"); diff --git a/ui/litellm-dashboard/src/components/navbar.test.tsx b/ui/litellm-dashboard/src/components/navbar.test.tsx index 711f17da32e..7b1c4451d7a 100644 --- a/ui/litellm-dashboard/src/components/navbar.test.tsx +++ b/ui/litellm-dashboard/src/components/navbar.test.tsx @@ -1,6 +1,6 @@ -import { describe, it, expect, vi } from "vitest"; -import { renderWithProviders, screen, waitFor } from "../../tests/test-utils"; import userEvent from "@testing-library/user-event"; +import { describe, expect, it, vi } from "vitest"; +import { renderWithProviders, screen, waitFor } from "../../tests/test-utils"; import Navbar from "./navbar"; // Mock the hooks and utilities diff --git a/ui/litellm-dashboard/src/components/navbar.tsx b/ui/litellm-dashboard/src/components/navbar.tsx index 7d5ee3dd6f4..0ef8a505257 100644 --- a/ui/litellm-dashboard/src/components/navbar.tsx +++ b/ui/litellm-dashboard/src/components/navbar.tsx @@ -1,23 +1,27 @@ -import Link from "next/link"; -import React, { useState, useEffect } from "react"; -import type { MenuProps } from "antd"; -import { Dropdown, Switch, Tooltip } from "antd"; +import { useHealthReadiness } from "@/app/(dashboard)/hooks/healthReadiness/useHealthReadiness"; import { getProxyBaseUrl } from "@/components/networking"; +import { useTheme } from "@/contexts/ThemeContext"; +import { clearTokenCookies } from "@/utils/cookieUtils"; +import { + emitLocalStorageChange, + getLocalStorageItem, + removeLocalStorageItem, + setLocalStorageItem, +} from "@/utils/localStorageUtils"; +import { fetchProxySettings } from "@/utils/proxyUtils"; import { - UserOutlined, - LogoutOutlined, CrownOutlined, + LogoutOutlined, MailOutlined, - SafetyOutlined, MenuFoldOutlined, MenuUnfoldOutlined, + SafetyOutlined, + UserOutlined, } from "@ant-design/icons"; -import { clearTokenCookies } from "@/utils/cookieUtils"; -import { getLocalStorageItem, setLocalStorageItem, removeLocalStorageItem } from "@/utils/localStorageUtils"; -import { fetchProxySettings } from "@/utils/proxyUtils"; -import { useTheme } from "@/contexts/ThemeContext"; -import { emitLocalStorageChange } from "@/utils/localStorageUtils"; -import { useHealthReadiness } from "@/app/(dashboard)/hooks/healthReadiness/useHealthReadiness"; +import type { MenuProps } from "antd"; +import { Dropdown, Switch, Tooltip } from "antd"; +import Link from "next/link"; +import React, { useEffect, useState } from "react"; interface NavbarProps { userID: string | null; From ccfc53e7c56a6866c0d8974cb66312afb0fc4f49 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Thu, 1 Jan 2026 17:27:28 -0800 Subject: [PATCH 083/158] Fixing test --- ui/litellm-dashboard/src/components/public_model_hub.test.tsx | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/ui/litellm-dashboard/src/components/public_model_hub.test.tsx b/ui/litellm-dashboard/src/components/public_model_hub.test.tsx index 47a44fa5147..5b669e5d63e 100644 --- a/ui/litellm-dashboard/src/components/public_model_hub.test.tsx +++ b/ui/litellm-dashboard/src/components/public_model_hub.test.tsx @@ -27,6 +27,10 @@ vi.mock("./networking", async (importOriginal) => { }; }); +vi.mock("./navbar", () => ({ + default: vi.fn(() =>
Navbar Component
), +})); + beforeAll(() => { Object.defineProperty(window, "matchMedia", { writable: true, From c3817d1606503a7a75cc63dc3127dc5b00eb6f9e Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Thu, 1 Jan 2026 17:59:56 -0800 Subject: [PATCH 084/158] Change budget panel to have tabs --- .../components/budgets/budget_panel.test.tsx | 104 +++++++- .../src/components/budgets/budget_panel.tsx | 237 ++++++++---------- .../src/components/budgets/constants.ts | 42 ++++ 3 files changed, 239 insertions(+), 144 deletions(-) create mode 100644 ui/litellm-dashboard/src/components/budgets/constants.ts diff --git a/ui/litellm-dashboard/src/components/budgets/budget_panel.test.tsx b/ui/litellm-dashboard/src/components/budgets/budget_panel.test.tsx index d6412025054..534693d3984 100644 --- a/ui/litellm-dashboard/src/components/budgets/budget_panel.test.tsx +++ b/ui/litellm-dashboard/src/components/budgets/budget_panel.test.tsx @@ -1,5 +1,6 @@ import * as networking from "../networking"; import { fireEvent, render, waitFor, screen } from "@testing-library/react"; +import { act } from "@testing-library/react"; import { afterEach, describe, expect, it, vi } from "vitest"; import BudgetPanel from "./budget_panel"; @@ -24,11 +25,11 @@ describe("Budget Panel", () => { }, ]); - const { getByText } = render(); + render(); await waitFor(() => { - expect(getByText("Create a budget to assign to customers.")).toBeInTheDocument(); - expect(getByText("budget-1")).toBeInTheDocument(); + expect(screen.getByText("Create a budget to assign to customers.")).toBeInTheDocument(); + expect(screen.getByText("budget-1")).toBeInTheDocument(); }); }); @@ -43,23 +44,102 @@ describe("Budget Panel", () => { }, ]); - const { getByText, container } = render(); + render(); await waitFor(() => { - expect(getByText("budget-to-delete")).toBeInTheDocument(); + expect(screen.getByText("budget-to-delete")).toBeInTheDocument(); }); - // Find the first table row in tbody and click the second icon (trash/delete) - const bodyRows = container.querySelectorAll("tbody tr"); - expect(bodyRows.length).toBeGreaterThan(0); - const firstRow = bodyRows[0]; - const rowClickableIcons = firstRow.querySelectorAll(".cursor-pointer"); - expect(rowClickableIcons.length).toBeGreaterThan(1); + const deleteButton = screen.getByTestId("delete-budget-button"); - fireEvent.click(rowClickableIcons[1]); + act(() => { + fireEvent.click(deleteButton); + }); await waitFor(() => { expect(screen.getByText("Delete Budget?")).toBeInTheDocument(); }); }); + + it("should successfully delete a budget", async () => { + vi.mocked(networking.getBudgetList).mockResolvedValue([ + { + budget_id: "budget-to-delete", + max_budget: "200", + rpm_limit: 20, + tpm_limit: 2000, + updated_at: "2024-01-02T00:00:00Z", + }, + ]); + vi.mocked(networking.budgetDeleteCall).mockResolvedValue(undefined); + + render(); + + await waitFor(() => { + expect(screen.getByText("budget-to-delete")).toBeInTheDocument(); + }); + + // Open delete modal + const deleteButton = screen.getByTestId("delete-budget-button"); + act(() => { + fireEvent.click(deleteButton); + }); + + await waitFor(() => { + expect(screen.getByText("Delete Budget?")).toBeInTheDocument(); + }); + + // Confirm delete + const confirmButton = screen.getByRole("button", { name: /delete/i }); + act(() => { + fireEvent.click(confirmButton); + }); + + await waitFor(() => { + expect(networking.budgetDeleteCall).toHaveBeenCalledWith("token-123", "budget-to-delete"); + expect(networking.getBudgetList).toHaveBeenCalledTimes(2); // Initial load + refresh after delete + }); + }); + + it("should handle delete error", async () => { + vi.mocked(networking.getBudgetList).mockResolvedValue([ + { + budget_id: "budget-to-delete", + max_budget: "200", + rpm_limit: 20, + tpm_limit: 2000, + updated_at: "2024-01-02T00:00:00Z", + }, + ]); + vi.mocked(networking.budgetDeleteCall).mockRejectedValue(new Error("Delete failed")); + + render(); + + await waitFor(() => { + expect(screen.getByText("budget-to-delete")).toBeInTheDocument(); + }); + + // Open delete modal + const deleteButton = screen.getByTestId("delete-budget-button"); + act(() => { + fireEvent.click(deleteButton); + }); + + await waitFor(() => { + expect(screen.getByText("Delete Budget?")).toBeInTheDocument(); + }); + + // Confirm delete + const confirmButton = screen.getByRole("button", { name: /delete/i }); + act(() => { + fireEvent.click(confirmButton); + }); + + await waitFor(() => { + expect(networking.budgetDeleteCall).toHaveBeenCalledWith("token-123", "budget-to-delete"); + }); + + // Modal should still be open (error handling) + expect(screen.getByText("Delete Budget?")).toBeInTheDocument(); + }); }); diff --git a/ui/litellm-dashboard/src/components/budgets/budget_panel.tsx b/ui/litellm-dashboard/src/components/budgets/budget_panel.tsx index 252287191b7..b52ef5ab947 100644 --- a/ui/litellm-dashboard/src/components/budgets/budget_panel.tsx +++ b/ui/litellm-dashboard/src/components/budgets/budget_panel.tsx @@ -27,6 +27,7 @@ import NotificationsManager from "../molecules/notifications_manager"; import { budgetDeleteCall, getBudgetList } from "../networking"; import BudgetModal from "./budget_modal"; import EditBudgetModal from "./edit_budget_modal"; +import { CREATE_END_USER_CURL_COMMAND, CHAT_COMPLETIONS_CURL_COMMAND, OPENAI_SDK_PYTHON_CODE } from "./constants"; interface BudgetSettingsPageProps { accessToken: string | null; @@ -110,139 +111,111 @@ const BudgetPanel: React.FC = ({ accessToken }) => { - - {selectedBudget && ( - - )} - - Create a budget to assign to customers. - - - - Budget ID - Max Budget - TPM - RPM - - + + + Budgets + Examples + + + +
+ + {selectedBudget && ( + + )} + + Create a budget to assign to customers. +
+ + + Budget ID + Max Budget + TPM + RPM + + - - {budgetList - .slice() // Creates a shallow copy to avoid mutating the original array - .sort((a, b) => new Date(b.updated_at).getTime() - new Date(a.updated_at).getTime()) // Sort by updated_at in descending order - .map((value: budgetItem, index: number) => ( - - {value.budget_id} - {value.max_budget ? value.max_budget : "n/a"} - {value.tpm_limit ? value.tpm_limit : "n/a"} - {value.rpm_limit ? value.rpm_limit : "n/a"} - handleEditCall(value)} - /> - handleDeleteClick(value)} - /> - - ))} - -
- - -
- How to use budget id - - - Assign Budget to Customer - Test it (Curl) - - Test it (OpenAI SDK) - - - - - {` -curl -X POST --location '/end_user/new' \ - --H 'Authorization: Bearer ' \ - --H 'Content-Type: application/json' \ - --d '{"user_id": "my-customer-id', "budget_id": ""}' # 👈 KEY CHANGE - - `} - - - - - {` -curl -X POST --location '/chat/completions' \ - --H 'Authorization: Bearer ' \ - --H 'Content-Type: application/json' \ - --d '{ - "model": "gpt-3.5-turbo', - "messages":[{"role": "user", "content": "Hey, how's it going?"}], - "user": "my-customer-id" -}' # 👈 KEY CHANGE - - `} - - - - - {`from openai import OpenAI -client = OpenAI( - base_url="", - api_key="" -) - -completion = client.chat.completions.create( - model="gpt-3.5-turbo", - messages=[ - {"role": "system", "content": "You are a helpful assistant."}, - {"role": "user", "content": "Hello!"} - ], - user="my-customer-id" -) - -print(completion.choices[0].message)`} - - - - -
+ + {budgetList + .slice() // Creates a shallow copy to avoid mutating the original array + .sort((a, b) => new Date(b.updated_at).getTime() - new Date(a.updated_at).getTime()) // Sort by updated_at in descending order + .map((value: budgetItem, index: number) => ( + + {value.budget_id} + {value.max_budget ? value.max_budget : "n/a"} + {value.tpm_limit ? value.tpm_limit : "n/a"} + {value.rpm_limit ? value.rpm_limit : "n/a"} + handleEditCall(value)} + dataTestId="edit-budget-button" + /> + handleDeleteClick(value)} + dataTestId="delete-budget-button" + /> + + ))} + + + + +
+ + +
+ How to use budget id + + + Assign Budget to Customer + Test it (Curl) + Test it (OpenAI SDK) + + + + {CREATE_END_USER_CURL_COMMAND} + + + {CHAT_COMPLETIONS_CURL_COMMAND} + + + {OPENAI_SDK_PYTHON_CODE} + + + +
+ + +
); }; diff --git a/ui/litellm-dashboard/src/components/budgets/constants.ts b/ui/litellm-dashboard/src/components/budgets/constants.ts new file mode 100644 index 00000000000..9d6736db1f1 --- /dev/null +++ b/ui/litellm-dashboard/src/components/budgets/constants.ts @@ -0,0 +1,42 @@ +export const CREATE_END_USER_CURL_COMMAND = ` +curl -X POST --location '/end_user/new' \\ + +-H 'Authorization: Bearer ' \\ + +-H 'Content-Type: application/json' \\ + +-d '{"user_id": "my-customer-id', "budget_id": ""}' # 👈 KEY CHANGE + +`; + +export const CHAT_COMPLETIONS_CURL_COMMAND = ` +curl -X POST --location '/chat/completions' \\ + +-H 'Authorization: Bearer ' \\ + +-H 'Content-Type: application/json' \\ + +-d '{ + "model": "gpt-3.5-turbo', + "messages":[{"role": "user", "content": "Hey, how's it going?"}], + "user": "my-customer-id" +}' # 👈 KEY CHANGE + +`; + +export const OPENAI_SDK_PYTHON_CODE = `from openai import OpenAI +client = OpenAI( + base_url="", + api_key="" +) + +completion = client.chat.completions.create( + model="gpt-3.5-turbo", + messages=[ + {"role": "system", "content": "You are a helpful assistant."}, + {"role": "user", "content": "Hello!"} + ], + user="my-customer-id" +) + +print(completion.choices[0].message)`; From a1200a84274d942a386fe34c9b66ecd94881cf10 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Thu, 1 Jan 2026 18:26:46 -0800 Subject: [PATCH 085/158] tests for react query hooks --- .../hooks/agents/useAgents.test.ts | 332 ++++++++++++++++ .../hooks/credentials/useCredentials.test.ts | 194 ++++++++++ .../hooks/customers/useCustomers.test.ts | 334 ++++++++++++++++ .../hooks/guardrails/useGuardrails.test.ts | 11 - .../(dashboard)/hooks/keys/useKeys.test.ts | 362 ++++++++++++++++++ .../hooks/models/useModelCostMap.test.ts | 11 - .../organizations/useOrganizations.test.ts | 11 - .../hooks/organizations/useOrganizations.ts | 3 +- .../hooks/providers/useProviderFields.test.ts | 11 - .../(dashboard)/hooks/tags/useTags.test.ts | 11 - .../(dashboard)/hooks/teams/useTeams.test.ts | 275 +++++++++++++ 11 files changed, 1498 insertions(+), 57 deletions(-) create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/hooks/agents/useAgents.test.ts create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/hooks/credentials/useCredentials.test.ts create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/hooks/customers/useCustomers.test.ts create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/hooks/keys/useKeys.test.ts create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/hooks/teams/useTeams.test.ts diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/agents/useAgents.test.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/agents/useAgents.test.ts new file mode 100644 index 00000000000..44fc6a96836 --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/agents/useAgents.test.ts @@ -0,0 +1,332 @@ +import { describe, it, expect, vi, beforeEach } from "vitest"; +import { renderHook, waitFor } from "@testing-library/react"; +import { QueryClient, QueryClientProvider } from "@tanstack/react-query"; +import React, { ReactNode } from "react"; +import { useAgents } from "./useAgents"; +import { getAgentsList } from "@/components/networking"; +import type { AgentsResponse, Agent } from "@/components/agents/types"; + +// Mock the networking function +vi.mock("@/components/networking", () => ({ + getAgentsList: vi.fn(), +})); + +// Mock useAuthorized hook - we can override this in individual tests +const mockUseAuthorized = vi.fn(); +vi.mock("@/app/(dashboard)/hooks/useAuthorized", () => ({ + default: () => mockUseAuthorized(), +})); + +// Import actual roles instead of mocking them + +// Mock data +const mockAgents: Agent[] = [ + { + agent_id: "agent-1", + agent_name: "Test Agent 1", + litellm_params: { + model: "gpt-3.5-turbo", + api_key: "test-key-1", + }, + agent_card_params: { + description: "A test agent for unit testing", + }, + created_at: "2024-01-01T00:00:00Z", + updated_at: "2024-01-01T00:00:00Z", + created_by: "user-1", + updated_by: "user-1", + }, + { + agent_id: "agent-2", + agent_name: "Test Agent 2", + litellm_params: { + model: "claude-3", + api_key: "test-key-2", + }, + agent_card_params: { + description: "Another test agent", + }, + created_at: "2024-01-01T00:00:00Z", + updated_at: "2024-01-01T00:00:00Z", + created_by: "user-2", + updated_by: "user-2", + }, +]; + +const mockAgentsResponse: AgentsResponse = { + agents: mockAgents, +}; + +describe("useAgents", () => { + let queryClient: QueryClient; + + beforeEach(() => { + queryClient = new QueryClient({ + defaultOptions: { + queries: { + retry: false, + }, + }, + }); + + // Reset all mocks + vi.clearAllMocks(); + + // Set default mock for useAuthorized (enabled state) + mockUseAuthorized.mockReturnValue({ + accessToken: "test-access-token", + userRole: "Admin", + userId: "test-user-id", + token: "test-token", + userEmail: "test@example.com", + premiumUser: false, + disabledPersonalKeyCreation: null, + showSSOBanner: false, + }); + }); + + const wrapper = ({ children }: { children: ReactNode }) => + React.createElement(QueryClientProvider, { client: queryClient }, children); + + it("should return agents data when query is successful", async () => { + // Mock successful API call + (getAgentsList as any).mockResolvedValue(mockAgentsResponse); + + const { result } = renderHook(() => useAgents(), { wrapper }); + + // Initially loading + expect(result.current.isLoading).toBe(true); + expect(result.current.data).toBeUndefined(); + + // Wait for success + await waitFor(() => { + expect(result.current.isLoading).toBe(false); + expect(result.current.isSuccess).toBe(true); + }); + + expect(result.current.data).toEqual(mockAgentsResponse); + expect(result.current.error).toBeNull(); + expect(getAgentsList).toHaveBeenCalledWith("test-access-token"); + expect(getAgentsList).toHaveBeenCalledTimes(1); + }); + + it("should handle error when getAgentsList fails", async () => { + const errorMessage = "Failed to fetch agents"; + const testError = new Error(errorMessage); + + // Mock failed API call + (getAgentsList as any).mockRejectedValue(testError); + + const { result } = renderHook(() => useAgents(), { wrapper }); + + // Initially loading + expect(result.current.isLoading).toBe(true); + + // Wait for error + await waitFor(() => { + expect(result.current.isLoading).toBe(false); + expect(result.current.isError).toBe(true); + }); + + expect(result.current.error).toEqual(testError); + expect(result.current.data).toBeUndefined(); + expect(getAgentsList).toHaveBeenCalledWith("test-access-token"); + expect(getAgentsList).toHaveBeenCalledTimes(1); + }); + + it("should not execute query when accessToken is missing", async () => { + // Mock missing accessToken + mockUseAuthorized.mockReturnValue({ + accessToken: null, + userRole: "Admin", + userId: "test-user-id", + token: null, + userEmail: "test@example.com", + premiumUser: false, + disabledPersonalKeyCreation: null, + showSSOBanner: false, + }); + + const { result } = renderHook(() => useAgents(), { wrapper }); + + // Query should not execute + expect(result.current.isLoading).toBe(false); + expect(result.current.data).toBeUndefined(); + expect(result.current.isFetched).toBe(false); + + // API should not be called + expect(getAgentsList).not.toHaveBeenCalled(); + }); + + it("should not execute query when userRole is not an admin role", async () => { + // Mock non-admin userRole + mockUseAuthorized.mockReturnValue({ + accessToken: "test-access-token", + userRole: "member", // Not in all_admin_roles + userId: "test-user-id", + token: "test-token", + userEmail: "test@example.com", + premiumUser: false, + disabledPersonalKeyCreation: null, + showSSOBanner: false, + }); + + const { result } = renderHook(() => useAgents(), { wrapper }); + + // Query should not execute + expect(result.current.isLoading).toBe(false); + expect(result.current.data).toBeUndefined(); + expect(result.current.isFetched).toBe(false); + + // API should not be called + expect(getAgentsList).not.toHaveBeenCalled(); + }); + + it("should not execute query when userRole is null", async () => { + // Mock null userRole + mockUseAuthorized.mockReturnValue({ + accessToken: "test-access-token", + userRole: null, + userId: "test-user-id", + token: "test-token", + userEmail: "test@example.com", + premiumUser: false, + disabledPersonalKeyCreation: null, + showSSOBanner: false, + }); + + const { result } = renderHook(() => useAgents(), { wrapper }); + + // Query should not execute + expect(result.current.isLoading).toBe(false); + expect(result.current.data).toBeUndefined(); + expect(result.current.isFetched).toBe(false); + + // API should not be called + expect(getAgentsList).not.toHaveBeenCalled(); + }); + + it("should not execute query when userRole is empty string", async () => { + // Mock empty string userRole + mockUseAuthorized.mockReturnValue({ + accessToken: "test-access-token", + userRole: "", + userId: "test-user-id", + token: "test-token", + userEmail: "test@example.com", + premiumUser: false, + disabledPersonalKeyCreation: null, + showSSOBanner: false, + }); + + const { result } = renderHook(() => useAgents(), { wrapper }); + + // Query should not execute + expect(result.current.isLoading).toBe(false); + expect(result.current.data).toBeUndefined(); + expect(result.current.isFetched).toBe(false); + + // API should not be called + expect(getAgentsList).not.toHaveBeenCalled(); + }); + + it("should not execute query when both accessToken and userRole are missing", async () => { + // Mock both auth values missing + mockUseAuthorized.mockReturnValue({ + accessToken: null, + userRole: null, + userId: "test-user-id", + token: null, + userEmail: "test@example.com", + premiumUser: false, + disabledPersonalKeyCreation: null, + showSSOBanner: false, + }); + + const { result } = renderHook(() => useAgents(), { wrapper }); + + // Query should not execute + expect(result.current.isLoading).toBe(false); + expect(result.current.data).toBeUndefined(); + expect(result.current.isFetched).toBe(false); + + // API should not be called + expect(getAgentsList).not.toHaveBeenCalled(); + }); + + it("should execute query when accessToken is present and userRole is Admin", async () => { + // Mock successful API call + (getAgentsList as any).mockResolvedValue(mockAgentsResponse); + + // Ensure auth values are set (already done in beforeEach) + const { result } = renderHook(() => useAgents(), { wrapper }); + + // Wait for query to execute + await waitFor(() => { + expect(result.current.isLoading).toBe(false); + }); + + expect(getAgentsList).toHaveBeenCalledWith("test-access-token"); + expect(getAgentsList).toHaveBeenCalledTimes(1); + }); + + it("should execute query when accessToken is present and userRole is proxy_admin", async () => { + // Mock successful API call + (getAgentsList as any).mockResolvedValue(mockAgentsResponse); + + // Mock proxy_admin role + mockUseAuthorized.mockReturnValue({ + accessToken: "test-access-token", + userRole: "proxy_admin", + userId: "test-user-id", + token: "test-token", + userEmail: "test@example.com", + premiumUser: false, + disabledPersonalKeyCreation: null, + showSSOBanner: false, + }); + + const { result } = renderHook(() => useAgents(), { wrapper }); + + // Wait for query to execute + await waitFor(() => { + expect(result.current.isLoading).toBe(false); + }); + + expect(getAgentsList).toHaveBeenCalledWith("test-access-token"); + expect(getAgentsList).toHaveBeenCalledTimes(1); + }); + + it("should return empty agents array when API returns empty data", async () => { + // Mock API returning empty agents array + (getAgentsList as any).mockResolvedValue({ agents: [] }); + + const { result } = renderHook(() => useAgents(), { wrapper }); + + // Wait for success + await waitFor(() => { + expect(result.current.isLoading).toBe(false); + expect(result.current.isSuccess).toBe(true); + }); + + expect(result.current.data).toEqual({ agents: [] }); + expect(getAgentsList).toHaveBeenCalledWith("test-access-token"); + }); + + it("should handle network timeout error", async () => { + const timeoutError = new Error("Network timeout"); + + // Mock network timeout + (getAgentsList as any).mockRejectedValue(timeoutError); + + const { result } = renderHook(() => useAgents(), { wrapper }); + + // Wait for error + await waitFor(() => { + expect(result.current.isError).toBe(true); + }); + + expect(result.current.error).toEqual(timeoutError); + expect(result.current.data).toBeUndefined(); + }); +}); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/credentials/useCredentials.test.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/credentials/useCredentials.test.ts new file mode 100644 index 00000000000..ee903628f08 --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/credentials/useCredentials.test.ts @@ -0,0 +1,194 @@ +import { describe, it, expect, vi, beforeEach } from "vitest"; +import { renderHook, waitFor } from "@testing-library/react"; +import { QueryClient, QueryClientProvider } from "@tanstack/react-query"; +import React, { ReactNode } from "react"; +import { useCredentials } from "./useCredentials"; +import { credentialListCall, CredentialsResponse, CredentialItem } from "@/components/networking"; + +// Mock the networking function +vi.mock("@/components/networking", () => ({ + credentialListCall: vi.fn(), +})); + +// Mock useAuthorized hook - we can override this in individual tests +const mockUseAuthorized = vi.fn(); +vi.mock("@/app/(dashboard)/hooks/useAuthorized", () => ({ + default: () => mockUseAuthorized(), +})); + +// Mock data +const mockCredentialItems: CredentialItem[] = [ + { + credential_name: "openai-api-key", + credential_values: { api_key: "sk-test123" }, + credential_info: { + custom_llm_provider: "openai", + description: "OpenAI API Key for GPT models", + required: true, + }, + }, + { + credential_name: "anthropic-api-key", + credential_values: { api_key: "sk-ant-test456" }, + credential_info: { + custom_llm_provider: "anthropic", + description: "Anthropic API Key for Claude models", + required: true, + }, + }, +]; + +const mockCredentialsResponse: CredentialsResponse = { + credentials: mockCredentialItems, +}; + +describe("useCredentials", () => { + let queryClient: QueryClient; + + beforeEach(() => { + queryClient = new QueryClient({ + defaultOptions: { + queries: { + retry: false, + }, + }, + }); + + // Reset all mocks + vi.clearAllMocks(); + + // Set default mock for useAuthorized (enabled state) + mockUseAuthorized.mockReturnValue({ + accessToken: "test-access-token", + userRole: "Admin", + userId: "test-user-id", + token: "test-token", + userEmail: "test@example.com", + premiumUser: false, + disabledPersonalKeyCreation: null, + showSSOBanner: false, + }); + }); + + const wrapper = ({ children }: { children: ReactNode }) => + React.createElement(QueryClientProvider, { client: queryClient }, children); + + it("should return credentials data when query is successful", async () => { + // Mock successful API call + (credentialListCall as any).mockResolvedValue(mockCredentialsResponse); + + const { result } = renderHook(() => useCredentials(), { wrapper }); + + // Initially loading + expect(result.current.isLoading).toBe(true); + expect(result.current.data).toBeUndefined(); + + // Wait for success + await waitFor(() => { + expect(result.current.isLoading).toBe(false); + expect(result.current.isSuccess).toBe(true); + }); + + expect(result.current.data).toEqual(mockCredentialsResponse); + expect(result.current.error).toBeNull(); + expect(credentialListCall).toHaveBeenCalledWith("test-access-token"); + expect(credentialListCall).toHaveBeenCalledTimes(1); + }); + + it("should handle error when credentialListCall fails", async () => { + const errorMessage = "Failed to fetch credentials"; + const testError = new Error(errorMessage); + + // Mock failed API call + (credentialListCall as any).mockRejectedValue(testError); + + const { result } = renderHook(() => useCredentials(), { wrapper }); + + // Initially loading + expect(result.current.isLoading).toBe(true); + + // Wait for error + await waitFor(() => { + expect(result.current.isLoading).toBe(false); + expect(result.current.isError).toBe(true); + }); + + expect(result.current.error).toEqual(testError); + expect(result.current.data).toBeUndefined(); + expect(credentialListCall).toHaveBeenCalledWith("test-access-token"); + expect(credentialListCall).toHaveBeenCalledTimes(1); + }); + + it("should not execute query when accessToken is missing", async () => { + // Mock missing accessToken + mockUseAuthorized.mockReturnValue({ + accessToken: null, + userRole: "Admin", + userId: "test-user-id", + token: null, + userEmail: "test@example.com", + premiumUser: false, + disabledPersonalKeyCreation: null, + showSSOBanner: false, + }); + + const { result } = renderHook(() => useCredentials(), { wrapper }); + + // Query should not execute + expect(result.current.isLoading).toBe(false); + expect(result.current.data).toBeUndefined(); + expect(result.current.isFetched).toBe(false); + + // API should not be called + expect(credentialListCall).not.toHaveBeenCalled(); + }); + + it("should return empty credentials array when API returns empty data", async () => { + // Mock API returning empty credentials array + (credentialListCall as any).mockResolvedValue({ credentials: [] }); + + const { result } = renderHook(() => useCredentials(), { wrapper }); + + // Wait for success + await waitFor(() => { + expect(result.current.isLoading).toBe(false); + expect(result.current.isSuccess).toBe(true); + }); + + expect(result.current.data).toEqual({ credentials: [] }); + expect(credentialListCall).toHaveBeenCalledWith("test-access-token"); + }); + + it("should handle network timeout error", async () => { + const timeoutError = new Error("Network timeout"); + + // Mock network timeout + (credentialListCall as any).mockRejectedValue(timeoutError); + + const { result } = renderHook(() => useCredentials(), { wrapper }); + + // Wait for error + await waitFor(() => { + expect(result.current.isError).toBe(true); + }); + + expect(result.current.error).toEqual(timeoutError); + expect(result.current.data).toBeUndefined(); + }); + + it("should execute query when accessToken is present", async () => { + // Mock successful API call + (credentialListCall as any).mockResolvedValue(mockCredentialsResponse); + + // Ensure auth values are set (already done in beforeEach) + const { result } = renderHook(() => useCredentials(), { wrapper }); + + // Wait for query to execute + await waitFor(() => { + expect(result.current.isLoading).toBe(false); + }); + + expect(credentialListCall).toHaveBeenCalledWith("test-access-token"); + expect(credentialListCall).toHaveBeenCalledTimes(1); + }); +}); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/customers/useCustomers.test.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/customers/useCustomers.test.ts new file mode 100644 index 00000000000..716d6f75399 --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/customers/useCustomers.test.ts @@ -0,0 +1,334 @@ +import { allEndUsersCall } from "@/components/networking"; +import { QueryClient, QueryClientProvider } from "@tanstack/react-query"; +import { renderHook, waitFor } from "@testing-library/react"; +import React, { ReactNode } from "react"; +import { beforeEach, describe, expect, it, vi } from "vitest"; +import type { Customer, CustomersResponse } from "./useCustomers"; +import { useCustomers } from "./useCustomers"; + +// Mock the networking function +vi.mock("@/components/networking", () => ({ + allEndUsersCall: vi.fn(), +})); + +// Mock useAuthorized hook - we can override this in individual tests +const mockUseAuthorized = vi.fn(); +vi.mock("@/app/(dashboard)/hooks/useAuthorized", () => ({ + default: () => mockUseAuthorized(), +})); + +// Import actual roles instead of mocking them + +// Mock data +const mockCustomers: Customer[] = [ + { + user_id: "customer-1", + alias: "Test Customer 1", + spend: 150.5, + blocked: false, + allowed_model_region: "us-east-1", + default_model: "gpt-3.5-turbo", + budget_id: "budget-1", + litellm_budget_table: { + budget_id: "budget-1", + max_budget: 1000, + soft_budget: 800, + max_parallel_requests: 10, + tpm_limit: 1000, + rpm_limit: 100, + model_max_budget: { "gpt-4": 500 }, + budget_duration: "monthly", + budget_reset_at: "2024-02-01T00:00:00Z", + created_at: "2024-01-01T00:00:00Z", + created_by: "admin-1", + updated_at: "2024-01-01T00:00:00Z", + updated_by: "admin-1", + }, + }, + { + user_id: "customer-2", + alias: null, + spend: 0, + blocked: true, + allowed_model_region: null, + default_model: null, + budget_id: null, + litellm_budget_table: null, + }, +]; + +const mockCustomersResponse: CustomersResponse = mockCustomers; + +describe("useCustomers", () => { + let queryClient: QueryClient; + + beforeEach(() => { + queryClient = new QueryClient({ + defaultOptions: { + queries: { + retry: false, + }, + }, + }); + + // Reset all mocks + vi.clearAllMocks(); + + // Set default mock for useAuthorized (enabled state) + mockUseAuthorized.mockReturnValue({ + accessToken: "test-access-token", + userRole: "Admin", + userId: "test-user-id", + token: "test-token", + userEmail: "test@example.com", + premiumUser: false, + disabledPersonalKeyCreation: null, + showSSOBanner: false, + }); + }); + + const wrapper = ({ children }: { children: ReactNode }) => + React.createElement(QueryClientProvider, { client: queryClient }, children); + + it("should return customers data when query is successful", async () => { + // Mock successful API call + (allEndUsersCall as any).mockResolvedValue(mockCustomersResponse); + + const { result } = renderHook(() => useCustomers(), { wrapper }); + + // Initially loading + expect(result.current.isLoading).toBe(true); + expect(result.current.data).toBeUndefined(); + + // Wait for success + await waitFor(() => { + expect(result.current.isLoading).toBe(false); + expect(result.current.isSuccess).toBe(true); + }); + + expect(result.current.data).toEqual(mockCustomersResponse); + expect(result.current.error).toBeNull(); + expect(allEndUsersCall).toHaveBeenCalledWith("test-access-token"); + expect(allEndUsersCall).toHaveBeenCalledTimes(1); + }); + + it("should handle error when allEndUsersCall fails", async () => { + const errorMessage = "Failed to fetch customers"; + const testError = new Error(errorMessage); + + // Mock failed API call + (allEndUsersCall as any).mockRejectedValue(testError); + + const { result } = renderHook(() => useCustomers(), { wrapper }); + + // Initially loading + expect(result.current.isLoading).toBe(true); + + // Wait for error + await waitFor(() => { + expect(result.current.isLoading).toBe(false); + expect(result.current.isError).toBe(true); + }); + + expect(result.current.error).toEqual(testError); + expect(result.current.data).toBeUndefined(); + expect(allEndUsersCall).toHaveBeenCalledWith("test-access-token"); + expect(allEndUsersCall).toHaveBeenCalledTimes(1); + }); + + it("should not execute query when accessToken is missing", async () => { + // Mock missing accessToken + mockUseAuthorized.mockReturnValue({ + accessToken: null, + userRole: "Admin", + userId: "test-user-id", + token: null, + userEmail: "test@example.com", + premiumUser: false, + disabledPersonalKeyCreation: null, + showSSOBanner: false, + }); + + const { result } = renderHook(() => useCustomers(), { wrapper }); + + // Query should not execute + expect(result.current.isLoading).toBe(false); + expect(result.current.data).toBeUndefined(); + expect(result.current.isFetched).toBe(false); + + // API should not be called + expect(allEndUsersCall).not.toHaveBeenCalled(); + }); + + it("should not execute query when userRole is not an admin role", async () => { + // Mock non-admin userRole + mockUseAuthorized.mockReturnValue({ + accessToken: "test-access-token", + userRole: "member", // Not in all_admin_roles + userId: "test-user-id", + token: "test-token", + userEmail: "test@example.com", + premiumUser: false, + disabledPersonalKeyCreation: null, + showSSOBanner: false, + }); + + const { result } = renderHook(() => useCustomers(), { wrapper }); + + // Query should not execute + expect(result.current.isLoading).toBe(false); + expect(result.current.data).toBeUndefined(); + expect(result.current.isFetched).toBe(false); + + // API should not be called + expect(allEndUsersCall).not.toHaveBeenCalled(); + }); + + it("should not execute query when userRole is null", async () => { + // Mock null userRole + mockUseAuthorized.mockReturnValue({ + accessToken: "test-access-token", + userRole: null, + userId: "test-user-id", + token: "test-token", + userEmail: "test@example.com", + premiumUser: false, + disabledPersonalKeyCreation: null, + showSSOBanner: false, + }); + + const { result } = renderHook(() => useCustomers(), { wrapper }); + + // Query should not execute + expect(result.current.isLoading).toBe(false); + expect(result.current.data).toBeUndefined(); + expect(result.current.isFetched).toBe(false); + + // API should not be called + expect(allEndUsersCall).not.toHaveBeenCalled(); + }); + + it("should not execute query when userRole is empty string", async () => { + // Mock empty string userRole + mockUseAuthorized.mockReturnValue({ + accessToken: "test-access-token", + userRole: "", + userId: "test-user-id", + token: "test-token", + userEmail: "test@example.com", + premiumUser: false, + disabledPersonalKeyCreation: null, + showSSOBanner: false, + }); + + const { result } = renderHook(() => useCustomers(), { wrapper }); + + // Query should not execute + expect(result.current.isLoading).toBe(false); + expect(result.current.data).toBeUndefined(); + expect(result.current.isFetched).toBe(false); + + // API should not be called + expect(allEndUsersCall).not.toHaveBeenCalled(); + }); + + it("should not execute query when both accessToken and userRole are missing", async () => { + // Mock both auth values missing + mockUseAuthorized.mockReturnValue({ + accessToken: null, + userRole: null, + userId: "test-user-id", + token: null, + userEmail: "test@example.com", + premiumUser: false, + disabledPersonalKeyCreation: null, + showSSOBanner: false, + }); + + const { result } = renderHook(() => useCustomers(), { wrapper }); + + // Query should not execute + expect(result.current.isLoading).toBe(false); + expect(result.current.data).toBeUndefined(); + expect(result.current.isFetched).toBe(false); + + // API should not be called + expect(allEndUsersCall).not.toHaveBeenCalled(); + }); + + it("should execute query when accessToken is present and userRole is Admin", async () => { + // Mock successful API call + (allEndUsersCall as any).mockResolvedValue(mockCustomersResponse); + + // Ensure auth values are set (already done in beforeEach) + const { result } = renderHook(() => useCustomers(), { wrapper }); + + // Wait for query to execute + await waitFor(() => { + expect(result.current.isLoading).toBe(false); + }); + + expect(allEndUsersCall).toHaveBeenCalledWith("test-access-token"); + expect(allEndUsersCall).toHaveBeenCalledTimes(1); + }); + + it("should execute query when accessToken is present and userRole is proxy_admin", async () => { + // Mock successful API call + (allEndUsersCall as any).mockResolvedValue(mockCustomersResponse); + + // Mock proxy_admin role + mockUseAuthorized.mockReturnValue({ + accessToken: "test-access-token", + userRole: "proxy_admin", + userId: "test-user-id", + token: "test-token", + userEmail: "test@example.com", + premiumUser: false, + disabledPersonalKeyCreation: null, + showSSOBanner: false, + }); + + const { result } = renderHook(() => useCustomers(), { wrapper }); + + // Wait for query to execute + await waitFor(() => { + expect(result.current.isLoading).toBe(false); + }); + + expect(allEndUsersCall).toHaveBeenCalledWith("test-access-token"); + expect(allEndUsersCall).toHaveBeenCalledTimes(1); + }); + + it("should return empty customers array when API returns empty data", async () => { + // Mock API returning empty customers array + (allEndUsersCall as any).mockResolvedValue([]); + + const { result } = renderHook(() => useCustomers(), { wrapper }); + + // Wait for success + await waitFor(() => { + expect(result.current.isLoading).toBe(false); + expect(result.current.isSuccess).toBe(true); + }); + + expect(result.current.data).toEqual([]); + expect(allEndUsersCall).toHaveBeenCalledWith("test-access-token"); + }); + + it("should handle network timeout error", async () => { + const timeoutError = new Error("Network timeout"); + + // Mock network timeout + (allEndUsersCall as any).mockRejectedValue(timeoutError); + + const { result } = renderHook(() => useCustomers(), { wrapper }); + + // Wait for error + await waitFor(() => { + expect(result.current.isError).toBe(true); + }); + + expect(result.current.error).toEqual(timeoutError); + expect(result.current.data).toBeUndefined(); + }); +}); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/guardrails/useGuardrails.test.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/guardrails/useGuardrails.test.ts index 48344a2167f..d9e96a5308c 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/hooks/guardrails/useGuardrails.test.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/guardrails/useGuardrails.test.ts @@ -10,17 +10,6 @@ vi.mock("@/components/networking", () => ({ getGuardrailsList: vi.fn(), })); -// Mock the queryKeysFactory - we'll mock the specific return value -vi.mock("../common/queryKeysFactory", () => ({ - createQueryKeys: vi.fn((resource: string) => ({ - all: [resource], - lists: () => [resource, "list"], - list: (params?: any) => [resource, "list", { params }], - details: () => [resource, "detail"], - detail: (uid: string) => [resource, "detail", uid], - })), -})); - // Mock useAuthorized hook - we can override this in individual tests const mockUseAuthorized = vi.fn(); vi.mock("@/app/(dashboard)/hooks/useAuthorized", () => ({ diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/keys/useKeys.test.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/keys/useKeys.test.ts new file mode 100644 index 00000000000..c4ffb7041aa --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/keys/useKeys.test.ts @@ -0,0 +1,362 @@ +import { describe, it, expect, vi, beforeEach } from "vitest"; +import { renderHook, waitFor } from "@testing-library/react"; +import { QueryClient, QueryClientProvider } from "@tanstack/react-query"; +import React, { ReactNode } from "react"; +import { useKeys } from "./useKeys"; +import { keyListCall } from "@/components/networking"; +import type { KeyResponse } from "@/components/key_team_helpers/key_list"; + +// Mock the networking function +vi.mock("@/components/networking", () => ({ + keyListCall: vi.fn(), +})); + +// Mock useAuthorized hook - we can override this in individual tests +const mockUseAuthorized = vi.fn(); +vi.mock("@/app/(dashboard)/hooks/useAuthorized", () => ({ + default: () => mockUseAuthorized(), +})); + +// Mock data +const mockKeys: KeyResponse[] = [ + { + token: "sk-test-key-1", + token_id: "key-1", + key_name: "Test Key 1", + key_alias: "test-key-1", + spend: 10.5, + max_budget: 100, + expires: "2024-12-31T23:59:59Z", + models: ["gpt-3.5-turbo"], + aliases: {}, + config: {}, + user_id: "user-1", + team_id: null, + max_parallel_requests: 10, + metadata: {}, + tpm_limit: 1000, + rpm_limit: 100, + duration: "30d", + budget_duration: "1mo", + budget_reset_at: "2024-02-01T00:00:00Z", + allowed_cache_controls: [], + allowed_routes: [], + permissions: {}, + model_spend: { "gpt-3.5-turbo": 10.5 }, + model_max_budget: { "gpt-3.5-turbo": 100 }, + soft_budget_cooldown: false, + blocked: false, + litellm_budget_table: {}, + organization_id: null, + created_at: "2024-01-01T00:00:00Z", + updated_at: "2024-01-01T00:00:00Z", + team_spend: 0, + team_alias: "", + team_tpm_limit: 0, + team_rpm_limit: 0, + team_max_budget: 0, + team_models: [], + team_blocked: false, + soft_budget: 0, + team_model_aliases: {}, + team_member_spend: 0, + team_metadata: {}, + end_user_id: "", + end_user_tpm_limit: 0, + end_user_rpm_limit: 0, + end_user_max_budget: 0, + last_refreshed_at: 0, + api_key: "", + user_role: "user", + rpm_limit_per_model: {}, + tpm_limit_per_model: {}, + user_tpm_limit: 0, + user_rpm_limit: 0, + user_email: "", + }, + { + token: "sk-test-key-2", + token_id: "key-2", + key_name: "Test Key 2", + key_alias: "test-key-2", + spend: 25.0, + max_budget: 200, + expires: "2024-12-31T23:59:59Z", + models: ["claude-3"], + aliases: {}, + config: {}, + user_id: "user-2", + team_id: "team-1", + max_parallel_requests: 5, + metadata: {}, + tpm_limit: 500, + rpm_limit: 50, + duration: "30d", + budget_duration: "1mo", + budget_reset_at: "2024-02-01T00:00:00Z", + allowed_cache_controls: [], + allowed_routes: [], + permissions: {}, + model_spend: { "claude-3": 25.0 }, + model_max_budget: { "claude-3": 200 }, + soft_budget_cooldown: false, + blocked: false, + litellm_budget_table: {}, + organization_id: null, + created_at: "2024-01-01T00:00:00Z", + updated_at: "2024-01-01T00:00:00Z", + team_spend: 0, + team_alias: "test-team", + team_tpm_limit: 1000, + team_rpm_limit: 100, + team_max_budget: 500, + team_models: ["claude-3"], + team_blocked: false, + soft_budget: 0, + team_model_aliases: {}, + team_member_spend: 0, + team_metadata: {}, + end_user_id: "", + end_user_tpm_limit: 0, + end_user_rpm_limit: 0, + end_user_max_budget: 0, + last_refreshed_at: 0, + api_key: "", + user_role: "user", + rpm_limit_per_model: {}, + tpm_limit_per_model: {}, + user_tpm_limit: 0, + user_rpm_limit: 0, + user_email: "", + }, +]; + +const mockKeysResponse = { + keys: mockKeys, + total_count: 2, + current_page: 1, + total_pages: 1, +}; + +describe("useKeys", () => { + let queryClient: QueryClient; + + beforeEach(() => { + queryClient = new QueryClient({ + defaultOptions: { + queries: { + retry: false, + }, + }, + }); + + // Reset all mocks + vi.clearAllMocks(); + + // Set default mock for useAuthorized (enabled state) + mockUseAuthorized.mockReturnValue({ + accessToken: "test-access-token", + userRole: "Admin", + userId: "test-user-id", + token: "test-token", + userEmail: "test@example.com", + premiumUser: false, + disabledPersonalKeyCreation: null, + showSSOBanner: false, + }); + }); + + const wrapper = ({ children }: { children: ReactNode }) => + React.createElement(QueryClientProvider, { client: queryClient }, children); + + it("should return keys data when query is successful", async () => { + // Mock successful API call + (keyListCall as any).mockResolvedValue(mockKeysResponse); + + const { result } = renderHook(() => useKeys(1, 10), { wrapper }); + + // Initially loading + expect(result.current.isLoading).toBe(true); + expect(result.current.data).toBeUndefined(); + + // Wait for success + await waitFor(() => { + expect(result.current.isLoading).toBe(false); + expect(result.current.isSuccess).toBe(true); + }); + + expect(result.current.data).toEqual(mockKeysResponse); + expect(result.current.error).toBeNull(); + expect(keyListCall).toHaveBeenCalledWith( + "test-access-token", + null, // organizationID + null, // teamID + null, // selectedKeyAlias + null, // userID + null, // keyHash + 1, // page + 10, // pageSize + ); + expect(keyListCall).toHaveBeenCalledTimes(1); + }); + + it("should handle error when keyListCall fails", async () => { + const errorMessage = "Failed to fetch keys"; + const testError = new Error(errorMessage); + + // Mock failed API call + (keyListCall as any).mockRejectedValue(testError); + + const { result } = renderHook(() => useKeys(1, 10), { wrapper }); + + // Initially loading + expect(result.current.isLoading).toBe(true); + + // Wait for error + await waitFor(() => { + expect(result.current.isLoading).toBe(false); + expect(result.current.isError).toBe(true); + }); + + expect(result.current.error).toEqual(testError); + expect(result.current.data).toBeUndefined(); + expect(keyListCall).toHaveBeenCalledWith( + "test-access-token", + null, // organizationID + null, // teamID + null, // selectedKeyAlias + null, // userID + null, // keyHash + 1, // page + 10, // pageSize + ); + expect(keyListCall).toHaveBeenCalledTimes(1); + }); + + it("should not execute query when accessToken is missing", async () => { + // Mock missing accessToken + mockUseAuthorized.mockReturnValue({ + accessToken: null, + userRole: "Admin", + userId: "test-user-id", + token: null, + userEmail: "test@example.com", + premiumUser: false, + disabledPersonalKeyCreation: null, + showSSOBanner: false, + }); + + const { result } = renderHook(() => useKeys(1, 10), { wrapper }); + + // Query should not execute + expect(result.current.isLoading).toBe(false); + expect(result.current.data).toBeUndefined(); + expect(result.current.isFetched).toBe(false); + + // API should not be called + expect(keyListCall).not.toHaveBeenCalled(); + }); + + it("should pass correct page and pageSize parameters to the API", async () => { + // Mock successful API call + (keyListCall as any).mockResolvedValue(mockKeysResponse); + + const page = 2; + const pageSize = 20; + + const { result } = renderHook(() => useKeys(page, pageSize), { wrapper }); + + // Wait for success + await waitFor(() => { + expect(result.current.isLoading).toBe(false); + }); + + expect(keyListCall).toHaveBeenCalledWith( + "test-access-token", + null, // organizationID + null, // teamID + null, // selectedKeyAlias + null, // userID + null, // keyHash + page, // page + pageSize, // pageSize + ); + }); + + it("should return empty keys array when API returns empty data", async () => { + // Mock API returning empty keys array + const emptyResponse = { + keys: [], + total_count: 0, + current_page: 1, + total_pages: 0, + }; + (keyListCall as any).mockResolvedValue(emptyResponse); + + const { result } = renderHook(() => useKeys(1, 10), { wrapper }); + + // Wait for success + await waitFor(() => { + expect(result.current.isLoading).toBe(false); + expect(result.current.isSuccess).toBe(true); + }); + + expect(result.current.data).toEqual(emptyResponse); + expect(keyListCall).toHaveBeenCalledWith( + "test-access-token", + null, // organizationID + null, // teamID + null, // selectedKeyAlias + null, // userID + null, // keyHash + 1, // page + 10, // pageSize + ); + }); + + it("should handle network timeout error", async () => { + const timeoutError = new Error("Network timeout"); + + // Mock network timeout + (keyListCall as any).mockRejectedValue(timeoutError); + + const { result } = renderHook(() => useKeys(1, 10), { wrapper }); + + // Wait for error + await waitFor(() => { + expect(result.current.isError).toBe(true); + }); + + expect(result.current.error).toEqual(timeoutError); + expect(result.current.data).toBeUndefined(); + }); + + it("should handle pagination correctly", async () => { + const paginatedResponse = { + keys: [mockKeys[0]], // Only first key + total_count: 15, + current_page: 2, + total_pages: 2, + }; + (keyListCall as any).mockResolvedValue(paginatedResponse); + + const { result } = renderHook(() => useKeys(2, 10), { wrapper }); + + // Wait for success + await waitFor(() => { + expect(result.current.isLoading).toBe(false); + }); + + expect(result.current.data).toEqual(paginatedResponse); + expect(keyListCall).toHaveBeenCalledWith( + "test-access-token", + null, // organizationID + null, // teamID + null, // selectedKeyAlias + null, // userID + null, // keyHash + 2, // page + 10, // pageSize + ); + }); +}); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/models/useModelCostMap.test.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/models/useModelCostMap.test.ts index 325503408c7..f79ca33bc5d 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/hooks/models/useModelCostMap.test.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/models/useModelCostMap.test.ts @@ -10,17 +10,6 @@ vi.mock("@/components/networking", () => ({ modelCostMap: vi.fn(), })); -// Mock the queryKeysFactory - we'll mock the specific return value -vi.mock("../common/queryKeysFactory", () => ({ - createQueryKeys: vi.fn((resource: string) => ({ - all: [resource], - lists: () => [resource, "list"], - list: (params?: any) => [resource, "list", { params }], - details: () => [resource, "detail"], - detail: (uid: string) => [resource, "detail", uid], - })), -})); - // Mock data const mockModelCostData: Record = { "gpt-3.5-turbo": { diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/organizations/useOrganizations.test.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/organizations/useOrganizations.test.ts index 2f13a68bab4..66c005f37c4 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/hooks/organizations/useOrganizations.test.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/organizations/useOrganizations.test.ts @@ -11,17 +11,6 @@ vi.mock("@/components/networking", () => ({ organizationListCall: vi.fn(), })); -// Mock the queryKeysFactory - we'll mock the specific return value -vi.mock("../common/queryKeysFactory", () => ({ - createQueryKeys: vi.fn((resource: string) => ({ - all: [resource], - lists: () => [resource, "list"], - list: (params?: any) => [resource, "list", { params }], - details: () => [resource, "detail"], - detail: (uid: string) => [resource, "detail", uid], - })), -})); - // Mock useAuthorized hook - we can override this in individual tests const mockUseAuthorized = vi.fn(); vi.mock("@/app/(dashboard)/hooks/useAuthorized", () => ({ diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/organizations/useOrganizations.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/organizations/useOrganizations.ts index 57c9c057652..27a946d112a 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/hooks/organizations/useOrganizations.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/organizations/useOrganizations.ts @@ -6,8 +6,7 @@ import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; const organizationKeys = createQueryKeys("organizations"); export const useOrganizations = (): UseQueryResult => { - const { accessToken } = useAuthorized(); - const { userId, userRole } = useAuthorized(); + const { accessToken, userId, userRole } = useAuthorized(); return useQuery({ queryKey: organizationKeys.list({}), queryFn: async () => await organizationListCall(accessToken!), diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/providers/useProviderFields.test.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/providers/useProviderFields.test.ts index 952b52bb409..33242e0452f 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/hooks/providers/useProviderFields.test.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/providers/useProviderFields.test.ts @@ -11,17 +11,6 @@ vi.mock("@/components/networking", () => ({ getProviderCreateMetadata: vi.fn(), })); -// Mock the queryKeysFactory - we'll mock the specific return value -vi.mock("../common/queryKeysFactory", () => ({ - createQueryKeys: vi.fn((resource: string) => ({ - all: [resource], - lists: () => [resource, "list"], - list: (params?: any) => [resource, "list", { params }], - details: () => [resource, "detail"], - detail: (uid: string) => [resource, "detail", uid], - })), -})); - // Mock data const mockProviderFields: ProviderCreateInfo[] = [ { diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/tags/useTags.test.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/tags/useTags.test.ts index 9ae96623597..a1751339568 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/hooks/tags/useTags.test.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/tags/useTags.test.ts @@ -11,17 +11,6 @@ vi.mock("@/components/networking", () => ({ tagListCall: vi.fn(), })); -// Mock the queryKeysFactory - we'll mock the specific return value -vi.mock("../common/queryKeysFactory", () => ({ - createQueryKeys: vi.fn((resource: string) => ({ - all: [resource], - lists: () => [resource, "list"], - list: (params?: any) => [resource, "list", { params }], - details: () => [resource, "detail"], - detail: (uid: string) => [resource, "detail", uid], - })), -})); - // Mock useAuthorized hook - we can override this in individual tests const mockUseAuthorized = vi.fn(); vi.mock("@/app/(dashboard)/hooks/useAuthorized", () => ({ diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/teams/useTeams.test.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/teams/useTeams.test.ts new file mode 100644 index 00000000000..91ffbcfafa2 --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/teams/useTeams.test.ts @@ -0,0 +1,275 @@ +import { describe, it, expect, vi, beforeEach } from "vitest"; +import { renderHook, waitFor } from "@testing-library/react"; +import { QueryClient, QueryClientProvider } from "@tanstack/react-query"; +import React, { ReactNode } from "react"; +import { useTeams } from "./useTeams"; +import { fetchTeams } from "@/app/(dashboard)/networking"; +import type { Team } from "@/components/key_team_helpers/key_list"; + +// Mock the networking function +vi.mock("@/app/(dashboard)/networking", () => ({ + fetchTeams: vi.fn(), +})); + +// Mock useAuthorized hook - we can override this in individual tests +const mockUseAuthorized = vi.fn(); +vi.mock("@/app/(dashboard)/hooks/useAuthorized", () => ({ + default: () => mockUseAuthorized(), +})); + +// Mock data +const mockTeams: Team[] = [ + { + team_id: "team-1", + team_alias: "Test Team 1", + models: ["gpt-3.5-turbo", "claude-3"], + max_budget: 100.0, + budget_duration: "monthly", + tpm_limit: 1000, + rpm_limit: 100, + organization_id: "org-1", + created_at: "2024-01-01T00:00:00Z", + keys: [], + members_with_roles: [], + }, + { + team_id: "team-2", + team_alias: "Test Team 2", + models: ["gpt-4"], + max_budget: 200.0, + budget_duration: "monthly", + tpm_limit: 2000, + rpm_limit: 200, + organization_id: "org-1", + created_at: "2024-01-02T00:00:00Z", + keys: [], + members_with_roles: [], + }, +]; + +describe("useTeams", () => { + let queryClient: QueryClient; + + beforeEach(() => { + queryClient = new QueryClient({ + defaultOptions: { + queries: { + retry: false, + }, + }, + }); + + // Reset all mocks + vi.clearAllMocks(); + + // Set default mock for useAuthorized (enabled state) + mockUseAuthorized.mockReturnValue({ + accessToken: "test-access-token", + userId: "test-user-id", + userRole: "Admin", + token: "test-token", + userEmail: "test@example.com", + premiumUser: false, + disabledPersonalKeyCreation: null, + showSSOBanner: false, + }); + }); + + const wrapper = ({ children }: { children: ReactNode }) => + React.createElement(QueryClientProvider, { client: queryClient }, children); + + it("should return teams data when query is successful", async () => { + // Mock successful API call + (fetchTeams as any).mockResolvedValue(mockTeams); + + const { result } = renderHook(() => useTeams(), { wrapper }); + + // Initially loading + expect(result.current.isLoading).toBe(true); + expect(result.current.data).toBeUndefined(); + + // Wait for success + await waitFor(() => { + expect(result.current.isLoading).toBe(false); + expect(result.current.isSuccess).toBe(true); + }); + + expect(result.current.data).toEqual(mockTeams); + expect(result.current.error).toBeNull(); + expect(fetchTeams).toHaveBeenCalledWith("test-access-token", "test-user-id", "Admin", null); + expect(fetchTeams).toHaveBeenCalledTimes(1); + }); + + it("should handle error when fetchTeams fails", async () => { + const errorMessage = "Failed to fetch teams"; + const testError = new Error(errorMessage); + + // Mock failed API call + (fetchTeams as any).mockRejectedValue(testError); + + const { result } = renderHook(() => useTeams(), { wrapper }); + + // Initially loading + expect(result.current.isLoading).toBe(true); + + // Wait for error + await waitFor(() => { + expect(result.current.isLoading).toBe(false); + expect(result.current.isError).toBe(true); + }); + + expect(result.current.error).toEqual(testError); + expect(result.current.data).toBeUndefined(); + expect(fetchTeams).toHaveBeenCalledWith("test-access-token", "test-user-id", "Admin", null); + expect(fetchTeams).toHaveBeenCalledTimes(1); + }); + + it("should not execute query when accessToken is missing", async () => { + // Mock missing accessToken + mockUseAuthorized.mockReturnValue({ + accessToken: null, + userId: "test-user-id", + userRole: "Admin", + token: null, + userEmail: "test@example.com", + premiumUser: false, + disabledPersonalKeyCreation: null, + showSSOBanner: false, + }); + + const { result } = renderHook(() => useTeams(), { wrapper }); + + // Query should not execute + expect(result.current.isLoading).toBe(false); + expect(result.current.data).toBeUndefined(); + expect(result.current.isFetched).toBe(false); + + // API should not be called + expect(fetchTeams).not.toHaveBeenCalled(); + }); + + it("should not execute query when accessToken is empty string", async () => { + // Mock empty string accessToken + mockUseAuthorized.mockReturnValue({ + accessToken: "", + userId: "test-user-id", + userRole: "Admin", + token: "", + userEmail: "test@example.com", + premiumUser: false, + disabledPersonalKeyCreation: null, + showSSOBanner: false, + }); + + const { result } = renderHook(() => useTeams(), { wrapper }); + + // Query should not execute + expect(result.current.isLoading).toBe(false); + expect(result.current.data).toBeUndefined(); + expect(result.current.isFetched).toBe(false); + + // API should not be called + expect(fetchTeams).not.toHaveBeenCalled(); + }); + + it("should execute query when accessToken is present", async () => { + // Mock successful API call + (fetchTeams as any).mockResolvedValue(mockTeams); + + // Ensure auth values are set (already done in beforeEach) + const { result } = renderHook(() => useTeams(), { wrapper }); + + // Wait for query to execute + await waitFor(() => { + expect(result.current.isLoading).toBe(false); + }); + + expect(fetchTeams).toHaveBeenCalledWith("test-access-token", "test-user-id", "Admin", null); + expect(fetchTeams).toHaveBeenCalledTimes(1); + }); + + it("should return empty teams array when API returns empty data", async () => { + // Mock API returning empty teams array + (fetchTeams as any).mockResolvedValue([]); + + const { result } = renderHook(() => useTeams(), { wrapper }); + + // Wait for success + await waitFor(() => { + expect(result.current.isLoading).toBe(false); + expect(result.current.isSuccess).toBe(true); + }); + + expect(result.current.data).toEqual([]); + expect(fetchTeams).toHaveBeenCalledWith("test-access-token", "test-user-id", "Admin", null); + }); + + it("should handle network timeout error", async () => { + const timeoutError = new Error("Network timeout"); + + // Mock network timeout + (fetchTeams as any).mockRejectedValue(timeoutError); + + const { result } = renderHook(() => useTeams(), { wrapper }); + + // Wait for error + await waitFor(() => { + expect(result.current.isError).toBe(true); + }); + + expect(result.current.error).toEqual(timeoutError); + expect(result.current.data).toBeUndefined(); + }); + + it("should pass userId and userRole to fetchTeams", async () => { + // Mock successful API call + (fetchTeams as any).mockResolvedValue(mockTeams); + + // Mock specific userId and userRole + mockUseAuthorized.mockReturnValue({ + accessToken: "test-access-token", + userId: "custom-user-id", + userRole: "member", + token: "test-token", + userEmail: "test@example.com", + premiumUser: false, + disabledPersonalKeyCreation: null, + showSSOBanner: false, + }); + + const { result } = renderHook(() => useTeams(), { wrapper }); + + // Wait for query to execute + await waitFor(() => { + expect(result.current.isLoading).toBe(false); + }); + + expect(fetchTeams).toHaveBeenCalledWith("test-access-token", "custom-user-id", "member", null); + }); + + it("should handle null userId", async () => { + // Mock successful API call + (fetchTeams as any).mockResolvedValue(mockTeams); + + // Mock null userId + mockUseAuthorized.mockReturnValue({ + accessToken: "test-access-token", + userId: null, + userRole: "Admin", + token: "test-token", + userEmail: "test@example.com", + premiumUser: false, + disabledPersonalKeyCreation: null, + showSSOBanner: false, + }); + + const { result } = renderHook(() => useTeams(), { wrapper }); + + // Wait for query to execute + await waitFor(() => { + expect(result.current.isLoading).toBe(false); + }); + + expect(fetchTeams).toHaveBeenCalledWith("test-access-token", null, "Admin", null); + }); +}); From 9815b00deb95ef118a018d67ea82ad10b6b23c4a Mon Sep 17 00:00:00 2001 From: Yuta Saito Date: Fri, 2 Jan 2026 11:47:17 +0900 Subject: [PATCH 086/158] feat: add selectable mcp servers to the playground --- .../mcp/litellm_proxy_mcp_handler.py | 10 +- .../components/playground/chat_ui/ChatUI.tsx | 253 +++++++++++------- .../playground/chat_ui/CodeSnippets.tsx | 9 +- .../llm_calls/anthropic_messages.tsx | 23 -- .../playground/llm_calls/chat_completion.tsx | 53 ++-- .../playground/llm_calls/fetch_mcp_tools.tsx | 29 -- .../playground/llm_calls/responses_api.tsx | 40 ++- 7 files changed, 236 insertions(+), 181 deletions(-) delete mode 100644 ui/litellm-dashboard/src/components/playground/llm_calls/fetch_mcp_tools.tsx diff --git a/litellm/responses/mcp/litellm_proxy_mcp_handler.py b/litellm/responses/mcp/litellm_proxy_mcp_handler.py index 2eea28f6cc1..9cdcd3894e0 100644 --- a/litellm/responses/mcp/litellm_proxy_mcp_handler.py +++ b/litellm/responses/mcp/litellm_proxy_mcp_handler.py @@ -205,7 +205,15 @@ class LiteLLM_Proxy_MCP_Handler: else: tool_name = getattr(mcp_tool, "name", None) - if tool_name and tool_name in allowed_tool_names: + if not tool_name: + continue + + if tool_name in allowed_tool_names: + filtered_tools.append(mcp_tool) + continue + + unprefixed_name, _ = split_server_prefix_from_name(tool_name) + if unprefixed_name in allowed_tool_names: filtered_tools.append(mcp_tool) return filtered_tools diff --git a/ui/litellm-dashboard/src/components/playground/chat_ui/ChatUI.tsx b/ui/litellm-dashboard/src/components/playground/chat_ui/ChatUI.tsx index 1cf9c7eb8df..7acc723dd56 100644 --- a/ui/litellm-dashboard/src/components/playground/chat_ui/ChatUI.tsx +++ b/ui/litellm-dashboard/src/components/playground/chat_ui/ChatUI.tsx @@ -45,8 +45,8 @@ import { makeOpenAIAudioSpeechRequest } from "../llm_calls/audio_speech"; import { makeOpenAIAudioTranscriptionRequest } from "../llm_calls/audio_transcriptions"; import { makeOpenAIChatCompletionRequest } from "../llm_calls/chat_completion"; import { makeOpenAIEmbeddingsRequest } from "../llm_calls/embeddings_api"; -import type { MCPTool } from "../llm_calls/fetch_mcp_tools"; -import { fetchAvailableMCPTools } from "../llm_calls/fetch_mcp_tools"; +import { listMCPTools, fetchMCPServers } from "../../networking"; +import { MCPServer } from "../../mcp_tools/types"; import { fetchAvailableModels, ModelGroup } from "../llm_calls/fetch_models"; import { makeOpenAIImageEditsRequest } from "../llm_calls/image_edits"; import { makeOpenAIImageGenerationRequest } from "../llm_calls/image_generation"; @@ -97,20 +97,27 @@ const ChatUI: React.FC = ({ disabledPersonalKeyCreation, proxySettings, }) => { - const [isMCPToolsModalVisible, setIsMCPToolsModalVisible] = useState(false); - const [mcpTools, setMCPTools] = useState([]); - const [selectedMCPTools, setSelectedMCPTools] = useState(() => { - const saved = sessionStorage.getItem("selectedMCPTools"); + const [mcpServers, setMCPServers] = useState([]); + const [selectedMCPServers, setSelectedMCPServers] = useState(() => { + const saved = sessionStorage.getItem("selectedMCPServers"); try { - const parsed = saved ? JSON.parse(saved) : []; - // Convert from single string to array if needed for backward compatibility - return Array.isArray(parsed) ? parsed : parsed ? [parsed] : []; + return saved ? JSON.parse(saved) : []; } catch (error) { - console.error("Error parsing selectedMCPTools from sessionStorage", error); + console.error("Error parsing selectedMCPServers from sessionStorage", error); return []; } }); - const [isLoadingMCPTools, setIsLoadingMCPTools] = useState(false); + const [isLoadingMCPServers, setIsLoadingMCPServers] = useState(false); + const [serverToolsMap, setServerToolsMap] = useState>({}); + const [mcpServerToolRestrictions, setMCPServerToolRestrictions] = useState>(() => { + const saved = sessionStorage.getItem("mcpServerToolRestrictions"); + try { + return saved ? JSON.parse(saved) : {}; + } catch (error) { + console.error("Error parsing mcpServerToolRestrictions from sessionStorage", error); + return {}; + } + }); const [apiKeySource, setApiKeySource] = useState<"session" | "custom">(() => { const saved = sessionStorage.getItem("apiKeySource"); if (saved) { @@ -211,27 +218,38 @@ const ChatUI: React.FC = ({ const chatEndRef = useRef(null); - // Fetch MCP tools - const loadMCPTools = async () => { + // Fetch MCP servers + const loadMCPServers = async () => { const userApiKey = apiKeySource === "session" ? accessToken : apiKey; if (!userApiKey) return; - setIsLoadingMCPTools(true); + setIsLoadingMCPServers(true); try { - const tools = await fetchAvailableMCPTools(userApiKey); - setMCPTools(tools); + const servers = await fetchMCPServers(userApiKey); + setMCPServers(Array.isArray(servers) ? servers : servers.data || []); } catch (error) { - console.error("Error fetching MCP tools:", error); + console.error("Error fetching MCP servers:", error); } finally { - setIsLoadingMCPTools(false); + setIsLoadingMCPServers(false); } }; - useEffect(() => { - if (isMCPToolsModalVisible) { - loadMCPTools(); + // Fetch tools for a specific server + const loadServerTools = async (serverId: string) => { + const userApiKey = apiKeySource === "session" ? accessToken : apiKey; + if (!userApiKey || serverToolsMap[serverId]) return; + + try { + const response = await listMCPTools(userApiKey, serverId); + setServerToolsMap(prev => ({ + ...prev, + [serverId]: response.tools || [] + })); + } catch (error) { + console.error(`Error fetching tools for server ${serverId}:`, error); } - }, [isMCPToolsModalVisible, accessToken, apiKey, apiKeySource]); + }; + useEffect(() => { if (isGetCodeModalVisible) { @@ -244,7 +262,9 @@ const ChatUI: React.FC = ({ selectedTags, selectedVectorStores, selectedGuardrails, - selectedMCPTools, + selectedMCPServers, + mcpServers, + mcpServerToolRestrictions, endpointType, selectedModel, selectedSdk, @@ -264,7 +284,9 @@ const ChatUI: React.FC = ({ selectedTags, selectedVectorStores, selectedGuardrails, - selectedMCPTools, + selectedMCPServers, + mcpServers, + mcpServerToolRestrictions, endpointType, selectedModel, proxySettings, @@ -287,8 +309,10 @@ const ChatUI: React.FC = ({ sessionStorage.setItem("selectedTags", JSON.stringify(selectedTags)); sessionStorage.setItem("selectedVectorStores", JSON.stringify(selectedVectorStores)); sessionStorage.setItem("selectedGuardrails", JSON.stringify(selectedGuardrails)); - sessionStorage.setItem("selectedMCPTools", JSON.stringify(selectedMCPTools)); + sessionStorage.setItem("selectedMCPServers", JSON.stringify(selectedMCPServers)); + sessionStorage.setItem("mcpServerToolRestrictions", JSON.stringify(mcpServerToolRestrictions)); sessionStorage.setItem("selectedVoice", selectedVoice); + sessionStorage.removeItem("selectedMCPTools"); // Clean up old key if (selectedModel) { sessionStorage.setItem("selectedModel", selectedModel); @@ -318,7 +342,8 @@ const ChatUI: React.FC = ({ messageTraceId, responsesSessionId, useApiSessionManagement, - selectedMCPTools, + selectedMCPServers, + mcpServerToolRestrictions, selectedVoice, ]); @@ -355,7 +380,7 @@ const ChatUI: React.FC = ({ }; loadModels(); - loadMCPTools(); + loadMCPServers(); }, [accessToken, userID, userRole, apiKeySource, apiKey, token]); // Fetch agents when A2A endpoint is selected @@ -869,12 +894,14 @@ const ChatUI: React.FC = ({ traceId, selectedVectorStores.length > 0 ? selectedVectorStores : undefined, selectedGuardrails.length > 0 ? selectedGuardrails : undefined, - selectedMCPTools, + selectedMCPServers, updateChatImageUI, updateSearchResults, useAdvancedParams ? temperature : undefined, useAdvancedParams ? maxTokens : undefined, updateTotalLatency, + mcpServers, + mcpServerToolRestrictions, ); } else if (endpointType === EndpointType.IMAGE) { // For image generation @@ -940,12 +967,14 @@ const ChatUI: React.FC = ({ traceId, selectedVectorStores.length > 0 ? selectedVectorStores : undefined, selectedGuardrails.length > 0 ? selectedGuardrails : undefined, - selectedMCPTools, // Pass the selected tools array + selectedMCPServers, // Pass the selected servers array useApiSessionManagement ? responsesSessionId : null, // Only pass session ID if API mode is enabled handleResponseId, // Pass callback to capture new response ID handleMCPEvent, // Pass MCP event handler codeInterpreter.enabled, // Enable Code Interpreter tool codeInterpreter.setResult, // Handle code interpreter output + mcpServers, + mcpServerToolRestrictions, ); } else if (endpointType === EndpointType.ANTHROPIC_MESSAGES) { const apiChatHistory = [ @@ -968,7 +997,9 @@ const ChatUI: React.FC = ({ traceId, selectedVectorStores.length > 0 ? selectedVectorStores : undefined, selectedGuardrails.length > 0 ? selectedGuardrails : undefined, - selectedMCPTools, // Pass the selected tools array + selectedMCPServers, // Pass the selected servers array + mcpServers, + mcpServerToolRestrictions, ); } else if (endpointType === EndpointType.EMBEDDINGS) { await makeOpenAIEmbeddingsRequest( @@ -1329,13 +1360,13 @@ const ChatUI: React.FC = ({ /> - {/* MCP Tool Selection */} + {/* MCP Server Selection */}
- MCP Tool + MCP Servers @@ -1343,30 +1374,103 @@ const ChatUI: React.FC = ({ + + {/* Tool restrictions UI (optional) */} + {selectedMCPServers.length > 0 && + !selectedMCPServers.includes("__all__") && + MCP_SUPPORTED_ENDPOINTS.has(endpointType as EndpointType) && ( +
+ {selectedMCPServers.map(serverId => { + const server = mcpServers.find(s => s.server_id === serverId); + const tools = serverToolsMap[serverId] || []; + if (tools.length === 0) return null; + + return ( +
+ + Limit tools for {server?.alias || server?.server_name || serverId}: + + setSelectedMCPTools(value)} - optionLabelProp="label" - allowClear - maxTagCount="responsive" - > - {mcpTools.map((tool) => ( - {tool.name}
} - > -
- {tool.name} - {tool.description} -
- - ))} - -
- )} - - )}
); }; diff --git a/ui/litellm-dashboard/src/components/playground/chat_ui/CodeSnippets.tsx b/ui/litellm-dashboard/src/components/playground/chat_ui/CodeSnippets.tsx index 60ce6879bea..84f7f6b9e1c 100644 --- a/ui/litellm-dashboard/src/components/playground/chat_ui/CodeSnippets.tsx +++ b/ui/litellm-dashboard/src/components/playground/chat_ui/CodeSnippets.tsx @@ -1,5 +1,6 @@ import { MessageType } from "./types"; import { EndpointType } from "./mode_endpoint_mapping"; +import { MCPServer } from "../../mcp_tools/types"; interface CodeGenMetadata { tags?: string[]; @@ -16,7 +17,9 @@ interface GenerateCodeParams { selectedTags: string[]; selectedVectorStores: string[]; selectedGuardrails: string[]; - selectedMCPTools: string[]; + selectedMCPServers: string[]; + mcpServers?: MCPServer[]; + mcpServerToolRestrictions?: Record; selectedVoice?: string; endpointType: string; selectedModel: string | undefined; @@ -37,7 +40,9 @@ export const generateCodeSnippet = (params: GenerateCodeParams): string => { selectedTags, selectedVectorStores, selectedGuardrails, - selectedMCPTools, + selectedMCPServers, + mcpServers, + mcpServerToolRestrictions, selectedVoice, endpointType, selectedModel, diff --git a/ui/litellm-dashboard/src/components/playground/llm_calls/anthropic_messages.tsx b/ui/litellm-dashboard/src/components/playground/llm_calls/anthropic_messages.tsx index 3f8c90424c6..3aef2f2d5ef 100644 --- a/ui/litellm-dashboard/src/components/playground/llm_calls/anthropic_messages.tsx +++ b/ui/litellm-dashboard/src/components/playground/llm_calls/anthropic_messages.tsx @@ -17,7 +17,6 @@ export async function makeAnthropicMessagesRequest( traceId?: string, vector_store_ids?: string[], guardrails?: string[], - selectedMCPTools?: string[], ) { if (!accessToken) { throw new Error("Virtual Key is required"); @@ -47,23 +46,6 @@ export async function makeAnthropicMessagesRequest( const startTime = Date.now(); let firstTokenReceived = false; - // Format MCP tools if selected - const tools = - selectedMCPTools && selectedMCPTools.length > 0 - ? [ - { - type: "mcp", - server_label: "litellm", - server_url: `${proxyBaseUrl}/mcp`, - require_approval: "never", - allowed_tools: selectedMCPTools, - headers: { - "x-litellm-api-key": `Bearer ${accessToken}`, - }, - }, - ] - : undefined; - const requestBody: any = { model: selectedModel, messages: messages.map((m) => ({ role: m.role, content: m.content })), @@ -75,11 +57,6 @@ export async function makeAnthropicMessagesRequest( if (vector_store_ids) requestBody.vector_store_ids = vector_store_ids; if (guardrails) requestBody.guardrails = guardrails; - if (tools) { - requestBody.tools = tools; - requestBody.tool_choice = "auto"; - } - // Use the streaming helper method for cleaner async iteration // @ts-ignore - The SDK types might not include all litellm-specific parameters const stream = client.messages.stream(requestBody, { signal }); diff --git a/ui/litellm-dashboard/src/components/playground/llm_calls/chat_completion.tsx b/ui/litellm-dashboard/src/components/playground/llm_calls/chat_completion.tsx index 8bdaa94ec65..906c72245fd 100644 --- a/ui/litellm-dashboard/src/components/playground/llm_calls/chat_completion.tsx +++ b/ui/litellm-dashboard/src/components/playground/llm_calls/chat_completion.tsx @@ -3,6 +3,7 @@ import { ChatCompletionMessageParam } from "openai/resources/chat/completions"; import { TokenUsage } from "../chat_ui/ResponseMetrics"; import { VectorStoreSearchResponse } from "../chat_ui/types"; import { getProxyBaseUrl } from "@/components/networking"; +import { MCPServer } from "../../mcp_tools/types"; export async function makeOpenAIChatCompletionRequest( chatHistory: { role: string; content: string | any[] }[], @@ -17,12 +18,14 @@ export async function makeOpenAIChatCompletionRequest( traceId?: string, vector_store_ids?: string[], guardrails?: string[], - selectedMCPTools?: string[], + selectedMCPServers?: string[], onImageGenerated?: (imageUrl: string, model?: string) => void, onSearchResults?: (searchResults: VectorStoreSearchResponse[]) => void, temperature?: number, max_tokens?: number, onTotalLatency?: (latency: number) => void, + mcpServers?: MCPServer[], + mcpServerToolRestrictions?: Record, ) { // base url should be the current base_url const isLocal = process.env.NODE_ENV === "development"; @@ -53,22 +56,36 @@ export async function makeOpenAIChatCompletionRequest( let fullResponseContent = ""; let fullReasoningContent = ""; - // Format MCP tools if selected - const tools = - selectedMCPTools && selectedMCPTools.length > 0 - ? [ - { - type: "mcp", - server_label: "litellm", - server_url: 'litellm_proxy/mcp', - require_approval: "never", - allowed_tools: selectedMCPTools, - headers: { - "x-litellm-api-key": `Bearer ${accessToken}`, - }, - }, - ] - : undefined; + // Build tools array + const tools: any[] = []; + + // Add MCP servers if selected + if (selectedMCPServers && selectedMCPServers.length > 0) { + if (selectedMCPServers.includes("__all__")) { + // All MCP Servers selected + tools.push({ + type: "mcp", + server_label: "litellm", + server_url: "litellm_proxy/mcp", + require_approval: "never", + }); + } else { + // Individual servers selected - create one entry per server + selectedMCPServers.forEach((serverId) => { + const server = mcpServers?.find((s) => s.server_id === serverId); + const serverName = server?.alias || server?.server_name || serverId; + const allowedTools = mcpServerToolRestrictions?.[serverId] || []; + + tools.push({ + type: "mcp", + server_label: "litellm", + server_url: `litellm_proxy/mcp/${serverName}`, + require_approval: "never", + ...(allowedTools.length > 0 ? { allowed_tools: allowedTools } : {}), + }); + }); + } + } // @ts-ignore const response = await client.chat.completions.create( @@ -82,7 +99,7 @@ export async function makeOpenAIChatCompletionRequest( messages: chatHistory as ChatCompletionMessageParam[], ...(vector_store_ids ? { vector_store_ids } : {}), ...(guardrails ? { guardrails } : {}), - ...(tools ? { tools, tool_choice: "auto" } : {}), + ...(tools.length > 0 ? { tools, tool_choice: "auto" } : {}), ...(temperature !== undefined ? { temperature } : {}), ...(max_tokens !== undefined ? { max_tokens } : {}), }, diff --git a/ui/litellm-dashboard/src/components/playground/llm_calls/fetch_mcp_tools.tsx b/ui/litellm-dashboard/src/components/playground/llm_calls/fetch_mcp_tools.tsx deleted file mode 100644 index d89ecbfb481..00000000000 --- a/ui/litellm-dashboard/src/components/playground/llm_calls/fetch_mcp_tools.tsx +++ /dev/null @@ -1,29 +0,0 @@ -import { mcpToolsCall } from "../../networking"; - -export interface MCPTool { - name: string; - description: string; - title: string | null; - inputSchema: { - type: string; - properties: Record; - required: string[]; - }; - outputSchema: any; - annotations: any; - _meta: any; -} - -interface MCPToolsResponse { - tools: MCPTool[]; -} - -export async function fetchAvailableMCPTools(accessToken: string): Promise { - try { - const data = (await mcpToolsCall(accessToken)) as MCPToolsResponse; - return data.tools || []; - } catch (error) { - console.error("Error fetching MCP tools:", error); - return []; - } -} diff --git a/ui/litellm-dashboard/src/components/playground/llm_calls/responses_api.tsx b/ui/litellm-dashboard/src/components/playground/llm_calls/responses_api.tsx index 8488f012e72..e30c14111fd 100644 --- a/ui/litellm-dashboard/src/components/playground/llm_calls/responses_api.tsx +++ b/ui/litellm-dashboard/src/components/playground/llm_calls/responses_api.tsx @@ -4,6 +4,7 @@ import { TokenUsage } from "../chat_ui/ResponseMetrics"; import { getProxyBaseUrl } from "@/components/networking"; import NotificationManager from "@/components/molecules/notifications_manager"; import { MCPEvent } from "../chat_ui/MCPEventsDisplay"; +import { MCPServer } from "../../mcp_tools/types"; import { CodeInterpreterResult, CodeInterpreterState, @@ -26,12 +27,14 @@ export async function makeOpenAIResponsesRequest( traceId?: string, vector_store_ids?: string[], guardrails?: string[], - selectedMCPTools?: string[], + selectedMCPServers?: string[], previousResponseId?: string | null, onResponseId?: (responseId: string) => void, onMCPEvent?: (event: MCPEvent) => void, codeInterpreterEnabled?: boolean, onCodeInterpreterResult?: (result: CodeInterpreterResult) => void, + mcpServers?: MCPServer[], + mcpServerToolRestrictions?: Record, ) { if (!accessToken) { throw new Error("Virtual Key is required"); @@ -86,15 +89,32 @@ export async function makeOpenAIResponsesRequest( // Build tools array const tools: any[] = []; - // Add MCP tools if selected - if (selectedMCPTools && selectedMCPTools.length > 0) { - tools.push({ - type: "mcp", - server_label: "litellm", - server_url: `litellm_proxy/mcp`, - require_approval: "never", - allowed_tools: selectedMCPTools, - }); + // Add MCP servers if selected + if (selectedMCPServers && selectedMCPServers.length > 0) { + if (selectedMCPServers.includes("__all__")) { + // All MCP Servers selected + tools.push({ + type: "mcp", + server_label: "litellm", + server_url: "litellm_proxy/mcp", + require_approval: "never", + }); + } else { + // Individual servers selected - create one entry per server + selectedMCPServers.forEach((serverId) => { + const server = mcpServers?.find((s) => s.server_id === serverId); + const serverName = server?.alias || server?.server_name || serverId; + const allowedTools = mcpServerToolRestrictions?.[serverId] || []; + + tools.push({ + type: "mcp", + server_label: "litellm", + server_url: `litellm_proxy/mcp/${serverName}`, + require_approval: "never", + ...(allowedTools.length > 0 ? { allowed_tools: allowedTools } : {}), + }); + }); + } } // Add code_interpreter tool if enabled (OpenAI auto-creates container) From df4eb844981e510f265f49b285e171c9e154bc3f Mon Sep 17 00:00:00 2001 From: Yuta Saito Date: Fri, 2 Jan 2026 11:55:55 +0900 Subject: [PATCH 087/158] chore: add test --- .../playground/chat_ui/ChatUI.test.tsx | 2 +- .../llm_calls/chat_completion.test.tsx | 54 ++++++++ .../llm_calls/responses_api.test.tsx | 125 ++++++++++++++++++ 3 files changed, 180 insertions(+), 1 deletion(-) create mode 100644 ui/litellm-dashboard/src/components/playground/llm_calls/responses_api.test.tsx diff --git a/ui/litellm-dashboard/src/components/playground/chat_ui/ChatUI.test.tsx b/ui/litellm-dashboard/src/components/playground/chat_ui/ChatUI.test.tsx index 8c85fb3bb2b..d992ff0ecb4 100644 --- a/ui/litellm-dashboard/src/components/playground/chat_ui/ChatUI.test.tsx +++ b/ui/litellm-dashboard/src/components/playground/chat_ui/ChatUI.test.tsx @@ -253,7 +253,7 @@ describe("ChatUI", () => { }; const getMcpSelect = () => - screen.getByText("MCP Tool").closest("div")?.querySelector(".ant-select") as HTMLElement | null; + screen.getByText("MCP Servers").closest("div")?.querySelector(".ant-select") as HTMLElement | null; await selectEndpointOption("/v1/embeddings"); diff --git a/ui/litellm-dashboard/src/components/playground/llm_calls/chat_completion.test.tsx b/ui/litellm-dashboard/src/components/playground/llm_calls/chat_completion.test.tsx index 9e4cd9abebb..f86c051951e 100644 --- a/ui/litellm-dashboard/src/components/playground/llm_calls/chat_completion.test.tsx +++ b/ui/litellm-dashboard/src/components/playground/llm_calls/chat_completion.test.tsx @@ -115,4 +115,58 @@ describe("chat_completion", () => { max_tokens: 100, }); }); + + it("should configure MCP tools per server with restrictions", async () => { + const selectedMCPServers = ["server-1", "server-2"]; + const mcpServers = [ + { server_id: "server-1", alias: "alpha", server_name: "Alpha" }, + { server_id: "server-2", server_name: "Beta" }, + ]; + const mcpServerToolRestrictions = { + "server-1": ["toolA", "toolB"], + "server-2": ["toolC"], + } as Record; + + await makeOpenAIChatCompletionRequest( + mockChatHistory, + mockUpdateUI, + "gpt-4", + "test-token", + undefined, + undefined, + undefined, + undefined, + undefined, + undefined, + undefined, + undefined, + selectedMCPServers, + undefined, + undefined, + undefined, + undefined, + undefined, + mcpServers, + mcpServerToolRestrictions, + ); + + const callArgs = mockCreate.mock.calls[0][0]; + expect(callArgs.tool_choice).toBe("auto"); + expect(callArgs.tools).toEqual([ + { + type: "mcp", + server_label: "litellm", + server_url: "litellm_proxy/mcp/alpha", + require_approval: "never", + allowed_tools: ["toolA", "toolB"], + }, + { + type: "mcp", + server_label: "litellm", + server_url: "litellm_proxy/mcp/Beta", + require_approval: "never", + allowed_tools: ["toolC"], + }, + ]); + }); }); diff --git a/ui/litellm-dashboard/src/components/playground/llm_calls/responses_api.test.tsx b/ui/litellm-dashboard/src/components/playground/llm_calls/responses_api.test.tsx new file mode 100644 index 00000000000..40c8e41ffb3 --- /dev/null +++ b/ui/litellm-dashboard/src/components/playground/llm_calls/responses_api.test.tsx @@ -0,0 +1,125 @@ +import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; +import { makeOpenAIResponsesRequest } from "./responses_api"; +import { MessageType } from "../chat_ui/types"; + +vi.mock("@/components/networking", () => ({ + getProxyBaseUrl: vi.fn(() => "https://example.com"), +})); + +const mockResponsesCreate = vi.fn(); +const mockClient = { + responses: { + create: mockResponsesCreate, + }, +}; + +vi.mock("openai", () => ({ + default: { + OpenAI: vi.fn(() => mockClient), + }, +})); + +describe("responses_api", () => { + const mockUpdateTextUI = vi.fn(); + const messages: MessageType[] = [{ role: "user", content: "Hello" }]; + + beforeEach(() => { + const mockEvents = [ + { type: "response.output_text.delta", delta: "Hi" }, + { + type: "response.completed", + response: { + id: "resp_123", + usage: { output_tokens: 2, input_tokens: 5, total_tokens: 7 }, + }, + }, + ]; + + async function* mockStream() { + for (const event of mockEvents) { + yield event; + } + } + + mockResponsesCreate.mockResolvedValue(mockStream()); + }); + + afterEach(() => { + vi.clearAllMocks(); + }); + + it("should send a basic responses request", async () => { + await makeOpenAIResponsesRequest(messages, mockUpdateTextUI, "gpt-4", "test-token"); + + expect(mockResponsesCreate).toHaveBeenCalledTimes(1); + expect(mockResponsesCreate).toHaveBeenCalledWith( + expect.objectContaining({ + model: "gpt-4", + input: [ + { + role: "user", + content: "Hello", + type: "message", + }, + ], + stream: true, + }), + { signal: undefined }, + ); + expect(mockUpdateTextUI).toHaveBeenCalledWith("assistant", "Hi", "gpt-4"); + }); + + it("should configure MCP tools per server with restrictions", async () => { + const selectedMCPServers = ["server-1", "server-2"]; + const mcpServers = [ + { server_id: "server-1", alias: "alpha", server_name: "Alpha" }, + { server_id: "server-2", server_name: "Beta" }, + ]; + const mcpServerToolRestrictions: Record = { + "server-1": ["toolA"], + "server-2": ["toolB", "toolC"], + }; + + await makeOpenAIResponsesRequest( + messages, + mockUpdateTextUI, + "gpt-4", + "test-token", + undefined, + undefined, + undefined, + undefined, + undefined, + undefined, + undefined, + undefined, + selectedMCPServers, + undefined, + undefined, + undefined, + undefined, + undefined, + mcpServers, + mcpServerToolRestrictions, + ); + + const callArgs = mockResponsesCreate.mock.calls[0][0]; + expect(callArgs.tool_choice).toBe("auto"); + expect(callArgs.tools).toEqual([ + { + type: "mcp", + server_label: "litellm", + server_url: "litellm_proxy/mcp/alpha", + require_approval: "never", + allowed_tools: ["toolA"], + }, + { + type: "mcp", + server_label: "litellm", + server_url: "litellm_proxy/mcp/Beta", + require_approval: "never", + allowed_tools: ["toolB", "toolC"], + }, + ]); + }); +}); From ca94990307393e48b6eef60481c8ac0c0a1463df Mon Sep 17 00:00:00 2001 From: Sameer Kankute Date: Fri, 2 Jan 2026 11:08:35 +0530 Subject: [PATCH 088/158] remove prompt caching headers as the support has been removed --- docs/my-website/docs/providers/anthropic.md | 4 +- litellm/llms/anthropic/chat/transformation.py | 9 ++- litellm/llms/anthropic/common_utils.py | 16 +++-- ...odel_prices_and_context_window_backup.json | 58 +++++++++++++++++++ .../test_anthropic_prompt_caching.py | 8 +-- .../tests/test_anthropic_context_caching.py | 1 - 6 files changed, 81 insertions(+), 15 deletions(-) diff --git a/docs/my-website/docs/providers/anthropic.md b/docs/my-website/docs/providers/anthropic.md index bcfb698a0f8..cae8657f1a0 100644 --- a/docs/my-website/docs/providers/anthropic.md +++ b/docs/my-website/docs/providers/anthropic.md @@ -444,7 +444,7 @@ Here's what a sample Raw Request from LiteLLM for Anthropic Context Caching look POST Request Sent from LiteLLM: curl -X POST \ https://api.anthropic.com/v1/messages \ --H 'accept: application/json' -H 'anthropic-version: 2023-06-01' -H 'content-type: application/json' -H 'x-api-key: sk-...' -H 'anthropic-beta: prompt-caching-2024-07-31' \ +-H 'accept: application/json' -H 'anthropic-version: 2023-06-01' -H 'content-type: application/json' -H 'x-api-key: sk-...' \ -d '{'model': 'claude-3-5-sonnet-20240620', [ { "role": "user", @@ -472,6 +472,8 @@ https://api.anthropic.com/v1/messages \ "max_tokens": 10 }' ``` + +**Note:** Anthropic no longer requires the `anthropic-beta: prompt-caching-2024-07-31` header. Prompt caching now works automatically when you use `cache_control` in your messages. ::: ### Caching - Large Context Caching diff --git a/litellm/llms/anthropic/chat/transformation.py b/litellm/llms/anthropic/chat/transformation.py index 6bdc17f7979..5f22ed0c56e 100644 --- a/litellm/llms/anthropic/chat/transformation.py +++ b/litellm/llms/anthropic/chat/transformation.py @@ -54,7 +54,10 @@ from litellm.types.utils import ( CompletionTokensDetailsWrapper, ) from litellm.types.utils import Message as LitellmMessage -from litellm.types.utils import PromptTokensDetailsWrapper, ServerToolUse +from litellm.types.utils import ( + PromptTokensDetailsWrapper, + ServerToolUse, +) from litellm.utils import ( ModelResponse, Usage, @@ -204,9 +207,11 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): ) # Relevant issue: https://github.com/BerriAI/litellm/issues/7755 def get_cache_control_headers(self) -> dict: + # Anthropic no longer requires the prompt-caching beta header + # Prompt caching now works automatically when cache_control is used in messages + # Reference: https://docs.anthropic.com/en/docs/build-with-claude/prompt-caching return { "anthropic-version": "2023-06-01", - "anthropic-beta": "prompt-caching-2024-07-31", } def _map_tool_choice( diff --git a/litellm/llms/anthropic/common_utils.py b/litellm/llms/anthropic/common_utils.py index 098694f15ae..fcbe9823ed4 100644 --- a/litellm/llms/anthropic/common_utils.py +++ b/litellm/llms/anthropic/common_utils.py @@ -12,7 +12,11 @@ from litellm.litellm_core_utils.prompt_templates.common_utils import ( ) from litellm.llms.base_llm.base_utils import BaseLLMModelInfo, BaseTokenCounter from litellm.llms.base_llm.chat.transformation import BaseLLMException -from litellm.types.llms.anthropic import AllAnthropicToolsValues, AnthropicMcpServerTool, ANTHROPIC_HOSTED_TOOLS +from litellm.types.llms.anthropic import ( + ANTHROPIC_HOSTED_TOOLS, + AllAnthropicToolsValues, + AnthropicMcpServerTool, +) from litellm.types.llms.openai import AllMessageValues from litellm.types.utils import TokenCountResponse @@ -273,8 +277,9 @@ class AnthropicModelInfo(BaseLLMModelInfo): beta_header = self.get_computer_tool_beta_header(computer_tool_used) betas.append(beta_header) - if prompt_caching_set: - betas.append("prompt-caching-2024-07-31") + # Anthropic no longer requires the prompt-caching beta header + # Prompt caching now works automatically when cache_control is used in messages + # Reference: https://docs.anthropic.com/en/docs/build-with-claude/prompt-caching if file_id_used: betas.append("files-api-2025-04-14") @@ -305,8 +310,9 @@ class AnthropicModelInfo(BaseLLMModelInfo): container_with_skills_used: bool = False, ) -> dict: betas = set() - if prompt_caching_set: - betas.add("prompt-caching-2024-07-31") + # Anthropic no longer requires the prompt-caching beta header + # Prompt caching now works automatically when cache_control is used in messages + # Reference: https://docs.anthropic.com/en/docs/build-with-claude/prompt-caching if computer_tool_used: beta_header = self.get_computer_tool_beta_header(computer_tool_used) betas.add(beta_header) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index d4869c902b4..5c668fc59df 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -249,6 +249,30 @@ "/v1/images/generations" ] }, + "aiml/google/imagen-4.0-ultra-generate-001": { + "litellm_provider": "aiml", + "metadata": { + "notes": "Imagen 4.0 Ultra Generate API - Photorealistic image generation with precise text rendering" + }, + "mode": "image_generation", + "output_cost_per_image": 0.063, + "source": "https://docs.aimlapi.com/api-references/image-models/google/imagen-4-ultra-generate", + "supported_endpoints": [ + "/v1/images/generations" + ] + }, + "aiml/google/nano-banana-pro": { + "litellm_provider": "aiml", + "metadata": { + "notes": "Gemini 3 Pro Image (Nano Banana Pro) - Advanced text-to-image generation with reasoning and 4K resolution support" + }, + "mode": "image_generation", + "output_cost_per_image": 0.1575, + "source": "https://docs.aimlapi.com/api-references/image-models/google/gemini-3-pro-image-preview", + "supported_endpoints": [ + "/v1/images/generations" + ] + }, "amazon.nova-canvas-v1:0": { "litellm_provider": "bedrock", "max_input_tokens": 2600, @@ -3508,6 +3532,40 @@ "supports_service_tier": true, "supports_vision": true }, + "azure/gpt-5.2-chat": { + "cache_read_input_token_cost": 1.75e-07, + "cache_read_input_token_cost_priority": 3.5e-07, + "input_cost_per_token": 1.75e-06, + "input_cost_per_token_priority": 3.5e-06, + "litellm_provider": "azure", + "max_input_tokens": 128000, + "max_output_tokens": 16384, + "max_tokens": 16384, + "mode": "chat", + "output_cost_per_token": 1.4e-05, + "output_cost_per_token_priority": 2.8e-05, + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/responses" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_native_streaming": true, + "supports_parallel_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_vision": true + }, "azure/gpt-5.2-chat-2025-12-11": { "cache_read_input_token_cost": 1.75e-07, "cache_read_input_token_cost_priority": 3.5e-07, diff --git a/tests/local_testing/test_anthropic_prompt_caching.py b/tests/local_testing/test_anthropic_prompt_caching.py index 0926bd17b70..c8589dd8844 100644 --- a/tests/local_testing/test_anthropic_prompt_caching.py +++ b/tests/local_testing/test_anthropic_prompt_caching.py @@ -104,7 +104,6 @@ async def test_litellm_anthropic_prompt_caching_tools(): ], extra_headers={ "anthropic-version": "2023-06-01", - "anthropic-beta": "prompt-caching-2024-07-31", }, ) @@ -112,11 +111,12 @@ async def test_litellm_anthropic_prompt_caching_tools(): print("call args=", mock_post.call_args) expected_url = "https://api.anthropic.com/v1/messages" + # Note: anthropic-beta header for prompt-caching is no longer required + # Anthropic now supports prompt caching automatically when cache_control is used expected_headers = { "accept": "application/json", "content-type": "application/json", "anthropic-version": "2023-06-01", - "anthropic-beta": "prompt-caching-2024-07-31", "x-api-key": "mock_api_key", } @@ -285,7 +285,6 @@ async def test_anthropic_api_prompt_caching_basic(): max_tokens=10, extra_headers={ "anthropic-version": "2023-06-01", - "anthropic-beta": "prompt-caching-2024-07-31", }, ) @@ -356,7 +355,6 @@ async def test_anthropic_api_prompt_caching_basic_with_cache_creation(): max_tokens=10, extra_headers={ "anthropic-version": "2023-06-01", - "anthropic-beta": "prompt-caching-2024-07-31", }, ) @@ -645,7 +643,6 @@ async def test_litellm_anthropic_prompt_caching_system(): ], extra_headers={ "anthropic-version": "2023-06-01", - "anthropic-beta": "prompt-caching-2024-07-31", }, ) @@ -657,7 +654,6 @@ async def test_litellm_anthropic_prompt_caching_system(): "accept": "application/json", "content-type": "application/json", "anthropic-version": "2023-06-01", - "anthropic-beta": "prompt-caching-2024-07-31", "x-api-key": "mock_api_key", } diff --git a/tests/old_proxy_tests/tests/test_anthropic_context_caching.py b/tests/old_proxy_tests/tests/test_anthropic_context_caching.py index 7a153295f35..6b37873df4e 100644 --- a/tests/old_proxy_tests/tests/test_anthropic_context_caching.py +++ b/tests/old_proxy_tests/tests/test_anthropic_context_caching.py @@ -30,7 +30,6 @@ response = client.chat.completions.create( ], extra_headers={ "anthropic-version": "2023-06-01", - "anthropic-beta": "prompt-caching-2024-07-31", }, ) From 35f9a75d5525bf25cc48d0acd56b70031d784f93 Mon Sep 17 00:00:00 2001 From: Yuta Saito Date: Fri, 2 Jan 2026 14:07:58 +0900 Subject: [PATCH 089/158] feat: add UI support for configuring meta URLs --- docs/my-website/docs/mcp.md | 16 +++ .../migration.sql | 5 + .../litellm_proxy_extras/schema.prisma | 3 + .../mcp_server/mcp_server_manager.py | 30 ++++- litellm/proxy/_types.py | 113 ++++++++++-------- .../mcp_management_endpoints.py | 9 +- litellm/proxy/schema.prisma | 3 + schema.prisma | 3 + tests/mcp_tests/test_mcp_server.py | 8 +- .../mcp_server/test_mcp_custom_fields.py | 43 +++---- .../mcp_server/test_mcp_server_manager.py | 4 +- .../mcp_tools/create_mcp_server.tsx | 51 ++++++++ .../components/mcp_tools/mcp_server_edit.tsx | 51 ++++++++ .../src/components/mcp_tools/types.tsx | 3 + .../src/hooks/useTestMCPConnection.tsx | 6 + 15 files changed, 263 insertions(+), 85 deletions(-) create mode 100644 litellm-proxy-extras/litellm_proxy_extras/migrations/20260102131258_add_metadata_urls_to_mcp_servers/migration.sql diff --git a/docs/my-website/docs/mcp.md b/docs/my-website/docs/mcp.md index a70e3d24188..8c11b8fd654 100644 --- a/docs/my-website/docs/mcp.md +++ b/docs/my-website/docs/mcp.md @@ -110,6 +110,22 @@ For stdio MCP servers, select "Standard Input/Output (stdio)" as the transport t

+### OAuth Configuration & Overrides + +LiteLLM attempts [OAuth 2.0 Authorization Server Discovery](https://datatracker.ietf.org/doc/html/rfc8414) by default. When you create an MCP server in the UI and set `Authentication: OAuth`, LiteLLM will locate the provider metadata, dynamically register a client, and perform PKCE-based authorization without you providing any additional details. + +**Customize the OAuth flow when needed:** + + + +- **Provide explicit client credentials** – If the MCP provider does not offer dynamic client registration or you prefer to manage the client yourself, fill in `client_id`, `client_secret`, and the desired `scopes`. +- **Override discovery URLs** – In some environments, LiteLLM might not be able to reach the provider's metadata endpoints. Use the optional `authorization_url`, `token_url`, and `registration_url` fields to point LiteLLM directly to the correct endpoints. + +
+ ### Static Headers Sometimes your MCP server needs specific headers on every request. Maybe it's an API key, maybe it's a custom header the server expects. Instead of configuring auth, you can just set them directly. diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260102131258_add_metadata_urls_to_mcp_servers/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260102131258_add_metadata_urls_to_mcp_servers/migration.sql new file mode 100644 index 00000000000..8eebb797e2c --- /dev/null +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260102131258_add_metadata_urls_to_mcp_servers/migration.sql @@ -0,0 +1,5 @@ +-- AlterTable +ALTER TABLE "LiteLLM_MCPServerTable" ADD COLUMN "authorization_url" TEXT, +ADD COLUMN "registration_url" TEXT, +ADD COLUMN "token_url" TEXT; + diff --git a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma index aac0b5b35de..ea47b6ed03b 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma +++ b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma @@ -208,6 +208,9 @@ model LiteLLM_MCPServerTable { command String? args String[] @default([]) env Json? @default("{}") + authorization_url String? + token_url String? + registration_url String? } // Generate Tokens for Proxy diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index 2d3ea1e827f..15c41ecbd5d 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -536,9 +536,12 @@ class MCPServerManager: client_secret=client_secret_value or getattr(mcp_server, "client_secret", None), scopes=resolved_scopes, - authorization_url=getattr(mcp_oauth_metadata, "authorization_url", None), - token_url=getattr(mcp_oauth_metadata, "token_url", None), - registration_url=getattr(mcp_oauth_metadata, "registration_url", None), + authorization_url=mcp_server.authorization_url + or getattr(mcp_oauth_metadata, "authorization_url", None), + token_url=mcp_server.token_url + or getattr(mcp_oauth_metadata, "token_url", None), + registration_url=mcp_server.registration_url + or getattr(mcp_oauth_metadata, "registration_url", None), command=getattr(mcp_server, "command", None), args=getattr(mcp_server, "args", None) or [], env=env_dict, @@ -548,7 +551,7 @@ class MCPServerManager: ) return new_server - async def add_update_server(self, mcp_server: LiteLLM_MCPServerTable): + async def add_server(self, mcp_server: LiteLLM_MCPServerTable): try: if mcp_server.server_id not in self.registry: new_server = await self.build_mcp_server_from_table(mcp_server) @@ -559,6 +562,17 @@ class MCPServerManager: verbose_logger.debug(f"Failed to add MCP server: {str(e)}") raise e + async def update_server(self, mcp_server: LiteLLM_MCPServerTable): + try: + if mcp_server.server_id in self.registry: + new_server = await self.build_mcp_server_from_table(mcp_server) + self.registry[mcp_server.server_id] = new_server + verbose_logger.debug(f"Updated MCP Server: {new_server.name}") + + except Exception as e: + verbose_logger.debug(f"Failed to udpate MCP server: {str(e)}") + raise e + def get_all_mcp_server_ids(self) -> Set[str]: """ Get all MCP server IDs @@ -2040,7 +2054,7 @@ class MCPServerManager: verbose_logger.debug( f"Adding server to registry: {server.server_id} ({server.server_name})" ) - await self.add_update_server(server) + await self.add_server(server) verbose_logger.debug( f"Registry now contains {len(self.get_registry())} servers" @@ -2220,6 +2234,9 @@ class MCPServerManager: command=getattr(server, "command", None), args=getattr(server, "args", None) or [], env=getattr(server, "env", None) or {}, + authorization_url=server.authorization_url, + token_url=server.token_url, + registration_url=server.registration_url, ) async def get_all_mcp_servers_with_health_and_teams( @@ -2310,6 +2327,9 @@ class MCPServerManager: command=getattr(server, "command", None), args=getattr(server, "args", None) or [], env=getattr(server, "env", None) or {}, + authorization_url=server.authorization_url, + token_url=server.token_url, + registration_url=server.registration_url, ) list_mcp_servers.append(mcp_server_table) diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index a94b4d4a077..6ffc53f6df9 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -412,7 +412,6 @@ class LiteLLMRoutes(enum.Enum): agent_routes = [ "/v1/agents", "/agents", - "/a2a/{agent_id}", "/a2a/{agent_id}/message/send", "/a2a/{agent_id}/message/stream", @@ -830,9 +829,9 @@ class GenerateRequestBase(LiteLLMPydanticObjectBase): allowed_cache_controls: Optional[list] = [] config: Optional[dict] = {} permissions: Optional[dict] = {} - model_max_budget: Optional[dict] = ( - {} - ) # {"gpt-4": 5.0, "gpt-3.5-turbo": 5.0}, defaults to {} + model_max_budget: Optional[ + dict + ] = {} # {"gpt-4": 5.0, "gpt-3.5-turbo": 5.0}, defaults to {} model_config = ConfigDict(protected_namespaces=()) model_rpm_limit: Optional[dict] = None @@ -1035,6 +1034,9 @@ class NewMCPServerRequest(LiteLLMPydanticObjectBase): command: Optional[str] = None args: List[str] = Field(default_factory=list) env: Dict[str, str] = Field(default_factory=dict) + authorization_url: Optional[str] = None + token_url: Optional[str] = None + registration_url: Optional[str] = None @model_validator(mode="before") @classmethod @@ -1092,6 +1094,9 @@ class UpdateMCPServerRequest(LiteLLMPydanticObjectBase): command: Optional[str] = None args: List[str] = Field(default_factory=list) env: Dict[str, str] = Field(default_factory=dict) + authorization_url: Optional[str] = None + token_url: Optional[str] = None + registration_url: Optional[str] = None @model_validator(mode="before") @classmethod @@ -1141,6 +1146,9 @@ class LiteLLM_MCPServerTable(LiteLLMPydanticObjectBase): command: Optional[str] = None args: List[str] = Field(default_factory=list) env: Dict[str, str] = Field(default_factory=dict) + authorization_url: Optional[str] = None + token_url: Optional[str] = None + registration_url: Optional[str] = None class MakeMCPServersPublicRequest(LiteLLMPydanticObjectBase): @@ -1160,6 +1168,9 @@ class NewSkillRequest(LiteLLMPydanticObjectBase): file_name: Optional[str] = None # Original filename file_type: Optional[str] = None # MIME type (e.g., "application/zip") metadata: Optional[Dict[str, Any]] = None + authorization_url: Optional[str] = None + token_url: Optional[str] = None + registration_url: Optional[str] = None class UpdateSkillRequest(LiteLLMPydanticObjectBase): @@ -1347,12 +1358,12 @@ class NewCustomerRequest(BudgetNewRequest): blocked: bool = False # allow/disallow requests for this end-user budget_id: Optional[str] = None # give either a budget_id or max_budget spend: Optional[float] = None - allowed_model_region: Optional[AllowedModelRegion] = ( - None # require all user requests to use models in this specific region - ) - default_model: Optional[str] = ( - None # if no equivalent model in allowed region - default all requests to this model - ) + allowed_model_region: Optional[ + AllowedModelRegion + ] = None # require all user requests to use models in this specific region + default_model: Optional[ + str + ] = None # if no equivalent model in allowed region - default all requests to this model @model_validator(mode="before") @classmethod @@ -1374,12 +1385,12 @@ class UpdateCustomerRequest(LiteLLMPydanticObjectBase): blocked: bool = False # allow/disallow requests for this end-user max_budget: Optional[float] = None budget_id: Optional[str] = None # give either a budget_id or max_budget - allowed_model_region: Optional[AllowedModelRegion] = ( - None # require all user requests to use models in this specific region - ) - default_model: Optional[str] = ( - None # if no equivalent model in allowed region - default all requests to this model - ) + allowed_model_region: Optional[ + AllowedModelRegion + ] = None # require all user requests to use models in this specific region + default_model: Optional[ + str + ] = None # if no equivalent model in allowed region - default all requests to this model class DeleteCustomerRequest(LiteLLMPydanticObjectBase): @@ -1464,15 +1475,15 @@ class NewTeamRequest(TeamBase): ] = None # raise an error if 'guaranteed_throughput' is set and we're overallocating tpm model_tpm_limit: Optional[Dict[str, int]] = None - team_member_budget: Optional[float] = ( - None # allow user to set a budget for all team members - ) - team_member_rpm_limit: Optional[int] = ( - None # allow user to set RPM limit for all team members - ) - team_member_tpm_limit: Optional[int] = ( - None # allow user to set TPM limit for all team members - ) + team_member_budget: Optional[ + float + ] = None # allow user to set a budget for all team members + team_member_rpm_limit: Optional[ + int + ] = None # allow user to set RPM limit for all team members + team_member_tpm_limit: Optional[ + int + ] = None # allow user to set TPM limit for all team members team_member_key_duration: Optional[str] = None # e.g. "1d", "1w", "1m" allowed_vector_store_indexes: Optional[List[AllowedVectorStoreIndexItem]] = None @@ -1558,9 +1569,9 @@ class BlockKeyRequest(LiteLLMPydanticObjectBase): class AddTeamCallback(LiteLLMPydanticObjectBase): callback_name: str - callback_type: Optional[Literal["success", "failure", "success_and_failure"]] = ( - "success_and_failure" - ) + callback_type: Optional[ + Literal["success", "failure", "success_and_failure"] + ] = "success_and_failure" callback_vars: Dict[str, str] @model_validator(mode="before") @@ -1788,9 +1799,10 @@ class DynamoDBArgs(LiteLLMPydanticObjectBase): class PassThroughGuardrailSettings(LiteLLMPydanticObjectBase): """ Settings for a specific guardrail on a passthrough endpoint. - + Allows field-level targeting for guardrail execution. """ + request_fields: Optional[List[str]] = Field( default=None, description="JSONPath expressions for input field targeting (pre_call). Examples: 'query', 'documents[*].text', 'messages[*].content'. If not specified, guardrail runs on entire request payload.", @@ -1871,9 +1883,9 @@ class ConfigList(LiteLLMPydanticObjectBase): stored_in_db: Optional[bool] field_default_value: Any premium_field: bool = False - nested_fields: Optional[List[FieldDetail]] = ( - None # For nested dictionary or Pydantic fields - ) + nested_fields: Optional[ + List[FieldDetail] + ] = None # For nested dictionary or Pydantic fields class UserHeaderMapping(LiteLLMPydanticObjectBase): @@ -2259,9 +2271,9 @@ class LiteLLM_OrganizationMembershipTable(LiteLLMPydanticObjectBase): budget_id: Optional[str] = None created_at: datetime updated_at: datetime - user: Optional[Any] = ( - None # You might want to replace 'Any' with a more specific type if available - ) + user: Optional[ + Any + ] = None # You might want to replace 'Any' with a more specific type if available litellm_budget_table: Optional[LiteLLM_BudgetTable] = None model_config = ConfigDict(protected_namespaces=()) @@ -2703,7 +2715,7 @@ class AllCallbacks(LiteLLMPydanticObjectBase): "TRACELOOP_API_KEY", ], ui_callback_name="Traceloop", - ) + ) class SpendLogsMetadata(TypedDict): @@ -2737,9 +2749,7 @@ class SpendLogsMetadata(TypedDict): cold_storage_object_key: Optional[ str ] # S3/GCS object key for cold storage retrieval - litellm_overhead_time_ms: Optional[ - float - ] # LiteLLM overhead time in milliseconds + litellm_overhead_time_ms: Optional[float] # LiteLLM overhead time in milliseconds cost_breakdown: Optional[ CostBreakdown ] # Detailed cost breakdown (input_cost, output_cost, margin, discount, etc.) @@ -3223,9 +3233,9 @@ class TeamModelDeleteRequest(BaseModel): # Organization Member Requests class OrganizationMemberAddRequest(OrgMemberAddRequest): organization_id: str - max_budget_in_organization: Optional[float] = ( - None # Users max budget within the organization - ) + max_budget_in_organization: Optional[ + float + ] = None # Users max budget within the organization class OrganizationMemberDeleteRequest(MemberDeleteRequest): @@ -3440,9 +3450,9 @@ class ProviderBudgetResponse(LiteLLMPydanticObjectBase): Maps provider names to their budget configs. """ - providers: Dict[str, ProviderBudgetResponseObject] = ( - {} - ) # Dictionary mapping provider names to their budget configurations + providers: Dict[ + str, ProviderBudgetResponseObject + ] = {} # Dictionary mapping provider names to their budget configurations class ProxyStateVariables(TypedDict): @@ -3577,9 +3587,9 @@ class LiteLLM_JWTAuth(LiteLLMPydanticObjectBase): enforce_rbac: bool = False roles_jwt_field: Optional[str] = None # v2 on role mappings role_mappings: Optional[List[RoleMapping]] = None - object_id_jwt_field: Optional[str] = ( - None # can be either user / team, inferred from the role mapping - ) + object_id_jwt_field: Optional[ + str + ] = None # can be either user / team, inferred from the role mapping scope_mappings: Optional[List[ScopeMapping]] = None enforce_scope_based_access: bool = False enforce_team_based_model_access: bool = False @@ -3730,13 +3740,16 @@ class DailyOrganizationSpendTransaction(BaseDailySpendTransaction): class DailyUserSpendTransaction(BaseDailySpendTransaction): user_id: str + class DailyEndUserSpendTransaction(BaseDailySpendTransaction): end_user_id: str + class DailyTagSpendTransaction(BaseDailySpendTransaction): request_id: Optional[str] tag: str + class DailyAgentSpendTransaction(BaseDailySpendTransaction): agent_id: str @@ -3768,8 +3781,8 @@ class LiteLLM_ManagedFileTable(LiteLLMPydanticObjectBase): flat_model_file_ids: List[str] created_by: Optional[str] updated_by: Optional[str] - storage_backend: Optional[str] = None - storage_url: Optional[str] = None + storage_backend: Optional[str] = None + storage_url: Optional[str] = None class LiteLLM_ManagedObjectTable(LiteLLMPydanticObjectBase): diff --git a/litellm/proxy/management_endpoints/mcp_management_endpoints.py b/litellm/proxy/management_endpoints/mcp_management_endpoints.py index 500323d3beb..9111f53a517 100644 --- a/litellm/proxy/management_endpoints/mcp_management_endpoints.py +++ b/litellm/proxy/management_endpoints/mcp_management_endpoints.py @@ -209,6 +209,9 @@ if MCP_AVAILABLE: command=payload.command, args=payload.args, env=payload.env, + authorization_url=payload.authorization_url, + token_url=payload.token_url, + registration_url=payload.registration_url, ) def get_prisma_client_or_throw(message: str): @@ -448,7 +451,7 @@ if MCP_AVAILABLE: exists = does_mcp_server_exist(mcp_server_records, server_id) if exists: - await global_mcp_server_manager.add_update_server(mcp_server) + await global_mcp_server_manager.add_server(mcp_server) return _redact_mcp_credentials(mcp_server) else: raise HTTPException( @@ -522,7 +525,7 @@ if MCP_AVAILABLE: payload, touched_by=user_api_key_dict.user_id or LITELLM_PROXY_ADMIN_NAME, ) - await global_mcp_server_manager.add_update_server(new_mcp_server) + await global_mcp_server_manager.add_server(new_mcp_server) # Ensure registry is up to date by reloading from database await global_mcp_server_manager.reload_servers_from_database() @@ -803,7 +806,7 @@ if MCP_AVAILABLE: "error": f"MCP Server not found, passed server_id={payload.server_id}" }, ) - await global_mcp_server_manager.add_update_server(mcp_server_record_updated) + await global_mcp_server_manager.update_server(mcp_server_record_updated) # Ensure registry is up to date by reloading from database await global_mcp_server_manager.reload_servers_from_database() diff --git a/litellm/proxy/schema.prisma b/litellm/proxy/schema.prisma index aac0b5b35de..ea47b6ed03b 100644 --- a/litellm/proxy/schema.prisma +++ b/litellm/proxy/schema.prisma @@ -208,6 +208,9 @@ model LiteLLM_MCPServerTable { command String? args String[] @default([]) env Json? @default("{}") + authorization_url String? + token_url String? + registration_url String? } // Generate Tokens for Proxy diff --git a/schema.prisma b/schema.prisma index aac0b5b35de..ea47b6ed03b 100644 --- a/schema.prisma +++ b/schema.prisma @@ -208,6 +208,9 @@ model LiteLLM_MCPServerTable { command String? args String[] @default([]) env Json? @default("{}") + authorization_url String? + token_url String? + registration_url String? } // Generate Tokens for Proxy diff --git a/tests/mcp_tests/test_mcp_server.py b/tests/mcp_tests/test_mcp_server.py index 9242dfc75f4..8fb0e80cc39 100644 --- a/tests/mcp_tests/test_mcp_server.py +++ b/tests/mcp_tests/test_mcp_server.py @@ -1057,7 +1057,7 @@ async def test_mcp_server_manager_config_integration_with_database(): ) # Test the add_update_server method (this tests our fix) - await test_manager.add_update_server(db_server) + await test_manager.add_server(db_server) # Verify the server was added with correct access_groups registry = test_manager.get_registry() @@ -1381,7 +1381,7 @@ async def test_add_update_server_with_alias(): mock_mcp_server.token_url = None # Add server to manager - await test_manager.add_update_server(mock_mcp_server) + await test_manager.add_server(mock_mcp_server) # Verify server was added with correct name (should use alias) assert "test-server-123" in test_manager.registry @@ -1421,7 +1421,7 @@ async def test_add_update_server_without_alias(): mock_mcp_server.token_url = None # Add server to manager - await test_manager.add_update_server(mock_mcp_server) + await test_manager.add_server(mock_mcp_server) # Verify server was added with correct name (should use server_name) assert "test-server-123" in test_manager.registry @@ -1461,7 +1461,7 @@ async def test_add_update_server_fallback_to_server_id(): mock_mcp_server.token_url = None # Add server to manager - await test_manager.add_update_server(mock_mcp_server) + await test_manager.add_server(mock_mcp_server) # Verify server was added with correct name (should use server_id) assert "test-server-123" in test_manager.registry diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_custom_fields.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_custom_fields.py index 5581070be71..a2425cc659a 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_custom_fields.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_custom_fields.py @@ -71,28 +71,29 @@ class TestMCPCustomFields: manager = MCPServerManager() # Mock database record with custom fields - mock_server = Mock(spec=LiteLLM_MCPServerTable) - mock_server.server_id = "test-server-id" - mock_server.server_name = "Test Server" - mock_server.description = "A test server" - mock_server.url = "http://localhost:3000" - mock_server.transport = "http" - mock_server.auth_type = MCPAuth.bearer_token - mock_server.alias = None - mock_server.mcp_info = { - "server_name": "Test Server", - "description": "A test server", - "custom_db_field": "database_value", - "metadata": {"source": "database"}, - "version": "1.0.0" - } - mock_server.command = None - mock_server.args = None - mock_server.env = None - mock_server.mcp_access_groups = None + mock_server = LiteLLM_MCPServerTable( + server_id="test-server-id", + server_name="Test Server", + alias=None, + description="A test server", + url="http://localhost:3000", + transport="http", + auth_type=MCPAuth.bearer_token, + mcp_info={ + "server_name": "Test Server", + "description": "A test server", + "custom_db_field": "database_value", + "metadata": {"source": "database"}, + "version": "1.0.0", + }, + command=None, + args=[], + env={}, + mcp_access_groups=[], + ) # Add server to manager - await manager.add_update_server(mock_server) + await manager.add_server(mock_server) # Get the added server server = manager.get_mcp_server_by_id("test-server-id") @@ -209,4 +210,4 @@ class TestMCPCustomFields: # Should use mcp_info description, not config level assert mcp_info["description"] == "MCP info description" - assert mcp_info["custom_field"] == "custom_value" \ No newline at end of file + assert mcp_info["custom_field"] == "custom_value" diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py index 37a93441513..6491e11024a 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py @@ -65,7 +65,7 @@ class TestMCPServerManager: updated_at=datetime.now(), ) - await manager.add_update_server(stdio_server) + await manager.add_server(stdio_server) # Verify server was added assert "stdio-server-1" in manager.registry @@ -1364,7 +1364,7 @@ class TestMCPServerManager: "env": {}, }, ) - await manager.add_update_server(server) + await manager.add_server(server) assert server.server_id in manager.get_registry() @pytest.mark.asyncio diff --git a/ui/litellm-dashboard/src/components/mcp_tools/create_mcp_server.tsx b/ui/litellm-dashboard/src/components/mcp_tools/create_mcp_server.tsx index ba6739d07ce..e9f9636bf2c 100644 --- a/ui/litellm-dashboard/src/components/mcp_tools/create_mcp_server.tsx +++ b/ui/litellm-dashboard/src/components/mcp_tools/create_mcp_server.tsx @@ -111,6 +111,9 @@ const CreateMCPServer: React.FC = ({ transport, auth_type: AUTH_TYPE.OAUTH2, credentials: values.credentials, + authorization_url: values.authorization_url, + token_url: values.token_url, + registration_url: values.registration_url, mcp_access_groups: values.mcp_access_groups, static_headers: staticHeaders, command: values.command, @@ -608,6 +611,54 @@ const CreateMCPServer: React.FC = ({ size="large" /> + + Authorization URL Override (optional) + + + + + } + name="authorization_url" + > + + + + Token URL Override (optional) + + + + + } + name="token_url" + > + + + + Registration URL Override (optional) + + + + + } + name="registration_url" + > + +

Complete the OAuth authorization flow to fetch an access token and store it as the authentication value. diff --git a/ui/litellm-dashboard/src/components/mcp_tools/mcp_server_edit.tsx b/ui/litellm-dashboard/src/components/mcp_tools/mcp_server_edit.tsx index f27e58e4a93..b4486bdfb23 100644 --- a/ui/litellm-dashboard/src/components/mcp_tools/mcp_server_edit.tsx +++ b/ui/litellm-dashboard/src/components/mcp_tools/mcp_server_edit.tsx @@ -229,6 +229,9 @@ const MCPServerEdit: React.FC = ({ transport: mcpServer.transport, auth_type: mcpServer.auth_type, mcp_info: mcpServer.mcp_info, + authorization_url: mcpServer.authorization_url, + token_url: mcpServer.token_url, + registration_url: mcpServer.registration_url, }; const toolsResponse = await testMCPToolsListRequest(accessToken, mcpServerConfig, oauthAccessToken); @@ -495,6 +498,54 @@ const MCPServerEdit: React.FC = ({ size="large" /> + + Authorization URL Override (optional) + + + + + } + name="authorization_url" + > + + + + Token URL Override (optional) + + + + + } + name="token_url" + > + + + + Registration URL Override (optional) + + + + + } + name="registration_url" + > + +

Use OAuth to fetch a fresh access token and save it as the authentication value.

+ )} +
+
+ + {isSSOConfigured ? ( + renderSSOSettings() + ) : ( + setIsAddModalVisible(true)} /> + )} + + + setIsDeleteModalVisible(false)} + onSuccess={() => refetch()} + accessToken={accessToken} + /> + + setIsAddModalVisible(false)} + onSuccess={() => { + setIsAddModalVisible(false); + refetch(); + }} + accessToken={accessToken} + /> + + ); +} diff --git a/ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/SSOSettingsEmptyPlaceholder.tsx b/ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/SSOSettingsEmptyPlaceholder.tsx new file mode 100644 index 00000000000..fc315493a54 --- /dev/null +++ b/ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/SSOSettingsEmptyPlaceholder.tsx @@ -0,0 +1,30 @@ +import { Empty, Typography, Button } from "antd"; + +const { Title, Paragraph } = Typography; + +interface SSOSettingsEmptyPlaceholderProps { + onAdd: () => void; +} + +export default function SSOSettingsEmptyPlaceholder({ onAdd }: SSOSettingsEmptyPlaceholderProps) { + return ( +
+ + No SSO Configuration Found + + Configure Single Sign-On (SSO) to enable seamless authentication for your team members using your identity + provider. + +
+ } + > + + + + ); +} diff --git a/ui/litellm-dashboard/src/components/admins.tsx b/ui/litellm-dashboard/src/components/admins.tsx index 4ddd5cd5d1f..6af5a226da6 100644 --- a/ui/litellm-dashboard/src/components/admins.tsx +++ b/ui/litellm-dashboard/src/components/admins.tsx @@ -55,6 +55,7 @@ import { getSSOSettings, } from "./networking"; import UISettings from "./Settings/AdminSettings/UISettings/UISettings"; +import SSOSettings from "./Settings/AdminSettings/SSOSettings/SSOSettings"; const AdminPanel: React.FC = ({ searchParams, @@ -496,11 +497,15 @@ const AdminPanel: React.FC = ({ Go to 'Internal Users' page to add other admins. + SSO Settings Security Settings SCIM UI Settings + + + ✨ Security Settings From 546fba98498d8019d8ba911e2a4ee62dd35df1a7 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Fri, 2 Jan 2026 18:28:10 -0800 Subject: [PATCH 125/158] tests --- .../Modals/AddSSOSettingsModal.test.tsx | 26 +++++++++++++ .../Modals/DeleteSSOSettingsModal.test.tsx | 20 ++++++++++ .../SSOSettings/SSOSettings.test.tsx | 37 +++++++++++++++++++ .../SSOSettingsEmptyPlaceholder.test.tsx | 14 +++++++ 4 files changed, 97 insertions(+) create mode 100644 ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/Modals/AddSSOSettingsModal.test.tsx create mode 100644 ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/Modals/DeleteSSOSettingsModal.test.tsx create mode 100644 ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/SSOSettings.test.tsx create mode 100644 ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/SSOSettingsEmptyPlaceholder.test.tsx diff --git a/ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/Modals/AddSSOSettingsModal.test.tsx b/ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/Modals/AddSSOSettingsModal.test.tsx new file mode 100644 index 00000000000..13363a11643 --- /dev/null +++ b/ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/Modals/AddSSOSettingsModal.test.tsx @@ -0,0 +1,26 @@ +import { render, screen } from "@testing-library/react"; +import { describe, expect, it, vi } from "vitest"; +import AddSSOSettingsModal from "./AddSSOSettingsModal"; + +// Mock networking functions +vi.mock("@/components/networking", () => ({ + updateSSOSettings: vi.fn(), +})); + +// Mock error utils +vi.mock("@/components/shared/errorUtils", () => ({ + parseErrorMessage: vi.fn((error) => error?.message || "Unknown error"), +})); + +describe("AddSSOSettingsModal", () => { + it("should render", () => { + const onCancel = vi.fn(); + const onSuccess = vi.fn(); + + render(); + + expect(screen.getByText("SSO Provider")).toBeInTheDocument(); + expect(screen.getByText("Cancel")).toBeInTheDocument(); + expect(screen.getAllByText("Add SSO")).toHaveLength(2); // Title and button + }); +}); diff --git a/ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/Modals/DeleteSSOSettingsModal.test.tsx b/ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/Modals/DeleteSSOSettingsModal.test.tsx new file mode 100644 index 00000000000..ef6ec6c7055 --- /dev/null +++ b/ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/Modals/DeleteSSOSettingsModal.test.tsx @@ -0,0 +1,20 @@ +import { render, screen } from "@testing-library/react"; +import { describe, expect, it, vi } from "vitest"; +import DeleteSSOSettingsModal from "./DeleteSSOSettingsModal"; + +describe("DeleteSSOSettingsModal", () => { + it("should render", () => { + const onCancel = vi.fn(); + const onSuccess = vi.fn(); + + render( + , + ); + + expect(screen.getByText("Confirm Clear SSO Settings")).toBeInTheDocument(); + expect( + screen.getByText("Are you sure you want to clear all SSO settings? This action cannot be undone."), + ).toBeInTheDocument(); + expect(screen.getByText("Users will no longer be able to login using SSO after this change.")).toBeInTheDocument(); + }); +}); diff --git a/ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/SSOSettings.test.tsx b/ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/SSOSettings.test.tsx new file mode 100644 index 00000000000..5e7908a872b --- /dev/null +++ b/ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/SSOSettings.test.tsx @@ -0,0 +1,37 @@ +import { QueryClient, QueryClientProvider } from "@tanstack/react-query"; +import { render, screen } from "@testing-library/react"; +import { describe, expect, it, vi } from "vitest"; +import SSOSettings from "./SSOSettings"; + +// Mock the useSSOSettings hook +vi.mock("@/app/(dashboard)/hooks/sso/useSSOSettings", () => ({ + useSSOSettings: () => ({ + data: null, + refetch: vi.fn(), + }), +})); + +const createQueryClient = () => + new QueryClient({ + defaultOptions: { + queries: { + retry: false, + gcTime: 0, + }, + }, + }); + +describe("SSOSettings", () => { + it("should render", () => { + const queryClient = createQueryClient(); + + render( + + + , + ); + + expect(screen.getByText("SSO Configuration")).toBeInTheDocument(); + expect(screen.getByText("Manage Single Sign-On authentication settings")).toBeInTheDocument(); + }); +}); diff --git a/ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/SSOSettingsEmptyPlaceholder.test.tsx b/ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/SSOSettingsEmptyPlaceholder.test.tsx new file mode 100644 index 00000000000..6676ba1c2c9 --- /dev/null +++ b/ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/SSOSettingsEmptyPlaceholder.test.tsx @@ -0,0 +1,14 @@ +import { render, screen } from "@testing-library/react"; +import { describe, expect, it, vi } from "vitest"; +import SSOSettingsEmptyPlaceholder from "./SSOSettingsEmptyPlaceholder"; + +describe("SSOSettingsEmptyPlaceholder", () => { + it("should render", () => { + const onAdd = vi.fn(); + + render(); + + expect(screen.getByText("No SSO Configuration Found")).toBeInTheDocument(); + expect(screen.getByText("Configure SSO")).toBeInTheDocument(); + }); +}); From f4c712506dc9b139b2e0b0e46865d3fd3a109121 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Fri, 2 Jan 2026 18:49:01 -0800 Subject: [PATCH 126/158] Unit tests to increase test coverage --- .../hooks/uiSettings/useUISettings.test.ts | 185 +++++++ .../playground/chat_ui/EndpointUtils.test.tsx | 219 ++++++++ .../prompt_editor_view/ToolsCard.test.tsx | 82 +++ .../VersionHistorySidePanel.test.tsx | 473 ++++++++++++++++++ .../prompts/prompt_editor_view/utils.test.ts | 444 ++++++++++++++++ 5 files changed, 1403 insertions(+) create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/hooks/uiSettings/useUISettings.test.ts create mode 100644 ui/litellm-dashboard/src/components/playground/chat_ui/EndpointUtils.test.tsx create mode 100644 ui/litellm-dashboard/src/components/prompts/prompt_editor_view/ToolsCard.test.tsx create mode 100644 ui/litellm-dashboard/src/components/prompts/prompt_editor_view/VersionHistorySidePanel.test.tsx create mode 100644 ui/litellm-dashboard/src/components/prompts/prompt_editor_view/utils.test.ts diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/uiSettings/useUISettings.test.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/uiSettings/useUISettings.test.ts new file mode 100644 index 00000000000..785f003d2f8 --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/uiSettings/useUISettings.test.ts @@ -0,0 +1,185 @@ +import { getUiSettings } from "@/components/networking"; +import { QueryClient, QueryClientProvider } from "@tanstack/react-query"; +import { renderHook, waitFor } from "@testing-library/react"; +import React, { ReactNode } from "react"; +import { beforeEach, describe, expect, it, vi } from "vitest"; +import { useUISettings } from "./useUISettings"; + +// Mock the networking function +vi.mock("@/components/networking", () => ({ + getUiSettings: vi.fn(), +})); + +// Mock useAuthorized hook - we can override this in individual tests +const mockUseAuthorized = vi.fn(); +vi.mock("../useAuthorized", () => ({ + default: () => mockUseAuthorized(), +})); + +// Mock data +const mockUISettings: Record = { + theme: "dark", + language: "en", + notifications: true, + dashboard_layout: "compact", + api_keys_visible: false, +}; + +describe("useUISettings", () => { + let queryClient: QueryClient; + + beforeEach(() => { + queryClient = new QueryClient({ + defaultOptions: { + queries: { + retry: false, + }, + }, + }); + + // Reset all mocks + vi.clearAllMocks(); + + // Set default mock for useAuthorized (enabled state) + mockUseAuthorized.mockReturnValue({ + accessToken: "test-access-token", + userRole: "Admin", + userId: "test-user-id", + token: "test-token", + userEmail: "test@example.com", + premiumUser: false, + disabledPersonalKeyCreation: null, + showSSOBanner: false, + }); + }); + + const wrapper = ({ children }: { children: ReactNode }) => + React.createElement(QueryClientProvider, { client: queryClient }, children); + + it("should return UI settings data when query is successful", async () => { + // Mock successful API call + (getUiSettings as any).mockResolvedValue(mockUISettings); + + const { result } = renderHook(() => useUISettings(), { wrapper }); + + // Initially loading + expect(result.current.isLoading).toBe(true); + expect(result.current.data).toBeUndefined(); + + // Wait for success + await waitFor(() => { + expect(result.current.isLoading).toBe(false); + expect(result.current.isSuccess).toBe(true); + }); + + expect(result.current.data).toEqual(mockUISettings); + expect(result.current.error).toBeNull(); + expect(getUiSettings).toHaveBeenCalledWith("test-access-token"); + expect(getUiSettings).toHaveBeenCalledTimes(1); + }); + + it("should handle error when getUiSettings fails", async () => { + const errorMessage = "Failed to fetch UI settings"; + const testError = new Error(errorMessage); + + // Mock failed API call + (getUiSettings as any).mockRejectedValue(testError); + + const { result } = renderHook(() => useUISettings(), { wrapper }); + + // Initially loading + expect(result.current.isLoading).toBe(true); + + // Wait for error + await waitFor(() => { + expect(result.current.isLoading).toBe(false); + expect(result.current.isError).toBe(true); + }); + + expect(result.current.error).toEqual(testError); + expect(result.current.data).toBeUndefined(); + expect(getUiSettings).toHaveBeenCalledWith("test-access-token"); + expect(getUiSettings).toHaveBeenCalledTimes(1); + }); + + it("should not execute query when accessToken is missing", async () => { + // Mock missing accessToken + mockUseAuthorized.mockReturnValue({ + accessToken: null, + userRole: "Admin", + userId: "test-user-id", + token: null, + userEmail: "test@example.com", + premiumUser: false, + disabledPersonalKeyCreation: null, + showSSOBanner: false, + }); + + const { result } = renderHook(() => useUISettings(), { wrapper }); + + // Query should not execute + expect(result.current.isLoading).toBe(false); + expect(result.current.data).toBeUndefined(); + expect(result.current.isFetched).toBe(false); + + // API should not be called + expect(getUiSettings).not.toHaveBeenCalled(); + }); + + it("should not execute query when accessToken is empty string", async () => { + // Mock empty accessToken + mockUseAuthorized.mockReturnValue({ + accessToken: "", + userRole: "Admin", + userId: "test-user-id", + token: "", + userEmail: "test@example.com", + premiumUser: false, + disabledPersonalKeyCreation: null, + showSSOBanner: false, + }); + + const { result } = renderHook(() => useUISettings(), { wrapper }); + + // Query should not execute + expect(result.current.isLoading).toBe(false); + expect(result.current.data).toBeUndefined(); + expect(result.current.isFetched).toBe(false); + + // API should not be called + expect(getUiSettings).not.toHaveBeenCalled(); + }); + + it("should return empty object when API returns empty settings", async () => { + // Mock API returning empty object + (getUiSettings as any).mockResolvedValue({}); + + const { result } = renderHook(() => useUISettings(), { wrapper }); + + // Wait for success + await waitFor(() => { + expect(result.current.isLoading).toBe(false); + expect(result.current.isSuccess).toBe(true); + }); + + expect(result.current.data).toEqual({}); + expect(getUiSettings).toHaveBeenCalledWith("test-access-token"); + }); + + it("should handle network timeout error", async () => { + const timeoutError = new Error("Network timeout"); + + // Mock network timeout + (getUiSettings as any).mockRejectedValue(timeoutError); + + const { result } = renderHook(() => useUISettings(), { wrapper }); + + // Wait for error + await waitFor(() => { + expect(result.current.isError).toBe(true); + }); + + expect(result.current.error).toEqual(timeoutError); + expect(result.current.data).toBeUndefined(); + }); +}); diff --git a/ui/litellm-dashboard/src/components/playground/chat_ui/EndpointUtils.test.tsx b/ui/litellm-dashboard/src/components/playground/chat_ui/EndpointUtils.test.tsx new file mode 100644 index 00000000000..6eb481381ac --- /dev/null +++ b/ui/litellm-dashboard/src/components/playground/chat_ui/EndpointUtils.test.tsx @@ -0,0 +1,219 @@ +import { beforeEach, describe, expect, it, vi } from "vitest"; +import type { ModelGroup } from "../llm_calls/fetch_models"; +import { determineEndpointType } from "./EndpointUtils"; +import { EndpointType } from "./mode_endpoint_mapping"; + +// Mock the getEndpointType function +vi.mock("./mode_endpoint_mapping", () => ({ + EndpointType: { + IMAGE: "image", + VIDEO: "video", + CHAT: "chat", + RESPONSES: "responses", + IMAGE_EDITS: "image_edits", + ANTHROPIC_MESSAGES: "anthropic_messages", + EMBEDDINGS: "embeddings", + SPEECH: "speech", + TRANSCRIPTION: "transcription", + A2A_AGENTS: "a2a_agents", + }, + getEndpointType: vi.fn(), + ModelMode: { + AUDIO_SPEECH: "audio_speech", + AUDIO_TRANSCRIPTION: "audio_transcription", + IMAGE_GENERATION: "image_generation", + VIDEO_GENERATION: "video_generation", + CHAT: "chat", + RESPONSES: "responses", + IMAGE_EDITS: "image_edits", + ANTHROPIC_MESSAGES: "anthropic_messages", + EMBEDDING: "embedding", + }, +})); + +// Import the mocked function +import { getEndpointType } from "./mode_endpoint_mapping"; + +describe("determineEndpointType", () => { + beforeEach(() => { + vi.clearAllMocks(); + }); + + it("should return the correct endpoint type when model is found and has a valid mode", () => { + const mockModelInfo: ModelGroup[] = [ + { + model_group: "gpt-3.5-turbo", + mode: "chat", + }, + { + model_group: "dall-e-3", + mode: "image_generation", + }, + ]; + + // Mock getEndpointType to return IMAGE for image_generation mode + vi.mocked(getEndpointType).mockReturnValue(EndpointType.IMAGE); + + const result = determineEndpointType("dall-e-3", mockModelInfo); + + expect(getEndpointType).toHaveBeenCalledWith("image_generation"); + expect(result).toBe(EndpointType.IMAGE); + }); + + it("should return CHAT endpoint type when model is found but has no mode", () => { + const mockModelInfo: ModelGroup[] = [ + { + model_group: "gpt-3.5-turbo", + // No mode property + }, + ]; + + const result = determineEndpointType("gpt-3.5-turbo", mockModelInfo); + + expect(getEndpointType).not.toHaveBeenCalled(); + expect(result).toBe(EndpointType.CHAT); + }); + + it("should return CHAT endpoint type when model is not found in modelInfo", () => { + const mockModelInfo: ModelGroup[] = [ + { + model_group: "gpt-3.5-turbo", + mode: "chat", + }, + ]; + + const result = determineEndpointType("non-existent-model", mockModelInfo); + + expect(getEndpointType).not.toHaveBeenCalled(); + expect(result).toBe(EndpointType.CHAT); + }); + + it("should return CHAT endpoint type when modelInfo array is empty", () => { + const mockModelInfo: ModelGroup[] = []; + + const result = determineEndpointType("any-model", mockModelInfo); + + expect(getEndpointType).not.toHaveBeenCalled(); + expect(result).toBe(EndpointType.CHAT); + }); + + it("should handle different mode types correctly", () => { + const mockModelInfo: ModelGroup[] = [ + { + model_group: "tts-model", + mode: "audio_speech", + }, + { + model_group: "whisper-model", + mode: "audio_transcription", + }, + { + model_group: "embedding-model", + mode: "embedding", + }, + { + model_group: "video-model", + mode: "video_generation", + }, + ]; + + // Test speech mode + vi.mocked(getEndpointType).mockReturnValueOnce(EndpointType.SPEECH); + const speechResult = determineEndpointType("tts-model", mockModelInfo); + expect(getEndpointType).toHaveBeenCalledWith("audio_speech"); + expect(speechResult).toBe(EndpointType.SPEECH); + + // Reset mock for next test + vi.clearAllMocks(); + + // Test transcription mode + vi.mocked(getEndpointType).mockReturnValueOnce(EndpointType.TRANSCRIPTION); + const transcriptionResult = determineEndpointType("whisper-model", mockModelInfo); + expect(getEndpointType).toHaveBeenCalledWith("audio_transcription"); + expect(transcriptionResult).toBe(EndpointType.TRANSCRIPTION); + + // Reset mock for next test + vi.clearAllMocks(); + + // Test embedding mode + vi.mocked(getEndpointType).mockReturnValueOnce(EndpointType.EMBEDDINGS); + const embeddingResult = determineEndpointType("embedding-model", mockModelInfo); + expect(getEndpointType).toHaveBeenCalledWith("embedding"); + expect(embeddingResult).toBe(EndpointType.EMBEDDINGS); + + // Reset mock for next test + vi.clearAllMocks(); + + // Test video mode + vi.mocked(getEndpointType).mockReturnValueOnce(EndpointType.VIDEO); + const videoResult = determineEndpointType("video-model", mockModelInfo); + expect(getEndpointType).toHaveBeenCalledWith("video_generation"); + expect(videoResult).toBe(EndpointType.VIDEO); + }); + + it("should prioritize the first matching model when there are duplicates", () => { + const mockModelInfo: ModelGroup[] = [ + { + model_group: "gpt-3.5-turbo", + mode: "chat", + }, + { + model_group: "gpt-3.5-turbo", + mode: "image_generation", // Different mode for same model name + }, + ]; + + vi.mocked(getEndpointType).mockReturnValue(EndpointType.CHAT); + + const result = determineEndpointType("gpt-3.5-turbo", mockModelInfo); + + expect(getEndpointType).toHaveBeenCalledWith("chat"); + expect(result).toBe(EndpointType.CHAT); + }); + + it("should handle models with undefined mode property explicitly set", () => { + const mockModelInfo: ModelGroup[] = [ + { + model_group: "test-model", + mode: undefined, + }, + ]; + + const result = determineEndpointType("test-model", mockModelInfo); + + expect(getEndpointType).not.toHaveBeenCalled(); + expect(result).toBe(EndpointType.CHAT); + }); + + it("should handle models with empty string mode", () => { + const mockModelInfo: ModelGroup[] = [ + { + model_group: "test-model", + mode: "", + }, + ]; + + const result = determineEndpointType("test-model", mockModelInfo); + + // Empty string is falsy, so getEndpointType should not be called + expect(getEndpointType).not.toHaveBeenCalled(); + expect(result).toBe(EndpointType.CHAT); + }); + + it("should handle case-sensitive model group matching", () => { + const mockModelInfo: ModelGroup[] = [ + { + model_group: "GPT-3.5-TURBO", + mode: "chat", + }, + ]; + + vi.mocked(getEndpointType).mockReturnValue(EndpointType.CHAT); + + // Test with different case - should not match + const result = determineEndpointType("gpt-3.5-turbo", mockModelInfo); + + expect(getEndpointType).not.toHaveBeenCalled(); + expect(result).toBe(EndpointType.CHAT); + }); +}); diff --git a/ui/litellm-dashboard/src/components/prompts/prompt_editor_view/ToolsCard.test.tsx b/ui/litellm-dashboard/src/components/prompts/prompt_editor_view/ToolsCard.test.tsx new file mode 100644 index 00000000000..742b8aa37fc --- /dev/null +++ b/ui/litellm-dashboard/src/components/prompts/prompt_editor_view/ToolsCard.test.tsx @@ -0,0 +1,82 @@ +import { act, fireEvent, render, screen } from "@testing-library/react"; +import { describe, expect, it, vi } from "vitest"; +import ToolsCard from "./ToolsCard"; +import { Tool } from "./types"; + +describe("ToolsCard", () => { + const mockTools: Tool[] = [ + { + name: "Calculator", + description: "Performs mathematical calculations", + json: '{"type": "function", "function": {"name": "calculate"}}', + }, + { + name: "Weather API", + description: "Gets current weather information", + json: '{"type": "function", "function": {"name": "get_weather"}}', + }, + ]; + + const defaultProps = { + tools: [] as Tool[], + onAddTool: vi.fn(), + onEditTool: vi.fn(), + onRemoveTool: vi.fn(), + }; + + it("should render the component", () => { + render(); + expect(screen.getByText("Tools")).toBeInTheDocument(); + }); + + it("should display no tools message when tools array is empty", () => { + render(); + expect(screen.getByText("No tools added")).toBeInTheDocument(); + }); + + it("should render tools when provided", () => { + render(); + + expect(screen.getByText("Calculator")).toBeInTheDocument(); + expect(screen.getByText("Performs mathematical calculations")).toBeInTheDocument(); + expect(screen.getByText("Weather API")).toBeInTheDocument(); + expect(screen.getByText("Gets current weather information")).toBeInTheDocument(); + }); + + it("should call onAddTool when Add button is clicked", () => { + const mockOnAddTool = vi.fn(); + render(); + + act(() => { + fireEvent.click(screen.getByRole("button", { name: /add/i })); + }); + + expect(mockOnAddTool).toHaveBeenCalledTimes(1); + }); + + it("should call onEditTool with correct index when Edit button is clicked", () => { + const mockOnEditTool = vi.fn(); + render(); + + const editButtons = screen.getAllByText("Edit"); + act(() => { + fireEvent.click(editButtons[0]); + }); + + expect(mockOnEditTool).toHaveBeenCalledWith(0); + expect(mockOnEditTool).toHaveBeenCalledTimes(1); + }); + + it("should call onRemoveTool with correct index when remove button is clicked", () => { + const mockOnRemoveTool = vi.fn(); + render(); + + const removeButtons = screen.getAllByRole("button", { name: "" }); + act(() => { + fireEvent.click(removeButtons[0]); + }); + + expect(mockOnRemoveTool).toHaveBeenCalledWith(0); + expect(mockOnRemoveTool).toHaveBeenCalledTimes(1); + }); +}); diff --git a/ui/litellm-dashboard/src/components/prompts/prompt_editor_view/VersionHistorySidePanel.test.tsx b/ui/litellm-dashboard/src/components/prompts/prompt_editor_view/VersionHistorySidePanel.test.tsx new file mode 100644 index 00000000000..b0346e03a2d --- /dev/null +++ b/ui/litellm-dashboard/src/components/prompts/prompt_editor_view/VersionHistorySidePanel.test.tsx @@ -0,0 +1,473 @@ +import { describe, it, expect, vi, beforeEach, afterEach, type Mock } from "vitest"; +import { render, screen, fireEvent, act, waitFor } from "@testing-library/react"; +import VersionHistorySidePanel from "./VersionHistorySidePanel"; +import { getPromptVersions } from "../../networking"; +import type { PromptSpec } from "../../networking"; + +// Mock the networking function +vi.mock("../../networking", () => ({ + getPromptVersions: vi.fn(), +})); + +const mockGetPromptVersions = getPromptVersions as Mock; + +// Mock Ant Design components that might need special handling +vi.mock("antd", async () => { + const actual = await vi.importActual("antd"); + return { + ...actual, + Drawer: ({ children, title, onClose, open, width, placement, mask, maskClosable }: any) => ( +
+
{title}
+ +
{children}
+
+ ), + List: ({ children, dataSource, renderItem }: any) => ( +
{dataSource?.map((item: any, index: number) => renderItem(item, index))}
+ ), + Skeleton: ({ active }: any) => ( +
+ Loading... +
+ ), + Tag: ({ children, color, className }: any) => ( + + {children} + + ), + Typography: { + Text: ({ children, type, className }: any) => ( + + {children} + + ), + }, + }; +}); + +describe("VersionHistorySidePanel", () => { + // Mock data + const mockPromptVersions: PromptSpec[] = [ + { + prompt_id: "test-prompt.v2", + litellm_params: { prompt_id: "test-prompt.v2" }, + prompt_info: { prompt_type: "db" }, + version: 2, + created_at: "2024-01-15T10:30:00Z", + }, + { + prompt_id: "test-prompt.v1", + litellm_params: { prompt_id: "test-prompt.v1" }, + prompt_info: { prompt_type: "db" }, + version: 1, + created_at: "2024-01-10T09:00:00Z", + }, + { + prompt_id: "test-prompt.v3", + litellm_params: { prompt_id: "test-prompt.v3" }, + prompt_info: { prompt_type: "config" }, + version: 3, + created_at: "2024-01-20T14:15:00Z", + }, + ]; + + const mockPromptVersionsWithoutExplicitVersion = [ + { + prompt_id: "test-prompt.v2", + litellm_params: { prompt_id: "test-prompt.v2" }, + prompt_info: { prompt_type: "db" }, + created_at: "2024-01-15T10:30:00Z", + }, + { + prompt_id: "test-prompt.v1", + litellm_params: { prompt_id: "test-prompt.v1" }, + prompt_info: { prompt_type: "db" }, + created_at: "2024-01-10T09:00:00Z", + }, + ]; + + const defaultProps = { + isOpen: true, + onClose: vi.fn(), + accessToken: "test-token", + promptId: "test-prompt.v2", + activeVersionId: "test-prompt.v2", + onSelectVersion: vi.fn(), + }; + + beforeEach(() => { + vi.clearAllMocks(); + // Mock successful response by default + mockGetPromptVersions.mockResolvedValue({ + prompts: mockPromptVersions, + }); + }); + + afterEach(() => { + vi.clearAllTimers(); + }); + + describe("Component Rendering", () => { + it("should render the component with drawer", async () => { + await act(async () => { + render(); + }); + expect(screen.getByTestId("drawer")).toBeInTheDocument(); + expect(screen.getByText("Version History")).toBeInTheDocument(); + }); + + it("should not render when isOpen is false", async () => { + await act(async () => { + render(); + }); + // The drawer should still be rendered but with open=false + const drawer = screen.getByTestId("drawer"); + expect(drawer).toHaveAttribute("data-open", "false"); + }); + + it("should show loading skeleton initially", async () => { + // Mock a delayed response to show loading state + mockGetPromptVersions.mockImplementationOnce( + () => new Promise((resolve) => setTimeout(() => resolve({ prompts: mockPromptVersions }), 100)), + ); + + render(); + expect(screen.getByTestId("skeleton")).toBeInTheDocument(); + + // Wait for loading to complete + await waitFor(() => { + expect(screen.queryByTestId("skeleton")).not.toBeInTheDocument(); + }); + }); + + it("should show empty state when no versions are available", async () => { + mockGetPromptVersions.mockResolvedValueOnce({ prompts: [] }); + + render(); + + await waitFor(() => { + expect(screen.getByText("No version history available.")).toBeInTheDocument(); + }); + }); + + it("should render version list when data is loaded", async () => { + render(); + + await waitFor(() => { + expect(screen.getByText("v2")).toBeInTheDocument(); + expect(screen.getByText("v1")).toBeInTheDocument(); + expect(screen.getByText("v3")).toBeInTheDocument(); + }); + + // Check that Latest tag is shown for the first item + const latestTags = screen.getAllByText("Latest"); + expect(latestTags.length).toBeGreaterThan(0); + + // Check Active tag is shown for the active version + expect(screen.getByText("Active")).toBeInTheDocument(); + }); + }); + + describe("Version Selection and Highlighting", () => { + it("should highlight the active version correctly", async () => { + render(); + + await waitFor(() => { + const versionItems = screen.getAllByTestId("tag"); + // Should have Active tag for the selected version + expect(screen.getByText("Active")).toBeInTheDocument(); + }); + }); + + it("should highlight the latest version when no activeVersionId is provided", async () => { + render(); + + await waitFor(() => { + const latestTags = screen.getAllByText("Latest"); + expect(latestTags.length).toBeGreaterThan(0); + }); + }); + + it("should call onSelectVersion when a version is clicked", async () => { + const mockOnSelectVersion = vi.fn(); + render(); + + await waitFor(() => { + expect(screen.getByText("v1")).toBeInTheDocument(); + }); + + const versionItem = screen.getByText("v1").closest("div"); + expect(versionItem).toBeInTheDocument(); + + act(() => { + fireEvent.click(versionItem!); + }); + + expect(mockOnSelectVersion).toHaveBeenCalledWith(mockPromptVersions[1]); + }); + }); + + describe("Version Number Extraction", () => { + it("should extract version from explicit version field", async () => { + render(); + + await waitFor(() => { + expect(screen.getByText("v2")).toBeInTheDocument(); + expect(screen.getByText("v3")).toBeInTheDocument(); + }); + }); + + it("should extract version from prompt_id with .v suffix", async () => { + mockGetPromptVersions.mockResolvedValueOnce({ + prompts: mockPromptVersionsWithoutExplicitVersion, + }); + + render(); + + await waitFor(() => { + expect(screen.getByText("v2")).toBeInTheDocument(); + expect(screen.getByText("v1")).toBeInTheDocument(); + }); + }); + + it("should extract version from prompt_id with _v suffix", async () => { + const versionsWithUnderscore = [ + { + prompt_id: "test-prompt_v2", + litellm_params: { prompt_id: "test-prompt_v2" }, + prompt_info: { prompt_type: "db" }, + created_at: "2024-01-15T10:30:00Z", + }, + ]; + + mockGetPromptVersions.mockResolvedValueOnce({ + prompts: versionsWithUnderscore, + }); + + render(); + + await waitFor(() => { + expect(screen.getByText("v2")).toBeInTheDocument(); + }); + }); + + it("should default to v1 when no version info is available", async () => { + const versionWithoutVersionInfo = [ + { + prompt_id: "test-prompt", + litellm_params: { prompt_id: "test-prompt" }, + prompt_info: { prompt_type: "db" }, + created_at: "2024-01-15T10:30:00Z", + }, + ]; + + mockGetPromptVersions.mockResolvedValueOnce({ + prompts: versionWithoutVersionInfo, + }); + + render(); + + await waitFor(() => { + expect(screen.getByText("v1")).toBeInTheDocument(); + }); + }); + }); + + describe("Date Formatting", () => { + it("should format dates correctly", async () => { + render(); + + await waitFor(() => { + // Check that dates are displayed (format: YYYY-MM-DD HH:MM:SS) + const dateElements = screen.getAllByTestId("text"); + const dateText = dateElements.find((el) => el.textContent?.match(/\d{4}-\d{2}-\d{2} \d{2}:\d{2}:\d{2}/)); + expect(dateText).toBeTruthy(); + }); + }); + + it("should show dash for missing dates", async () => { + const versionsWithoutDates = [ + { + prompt_id: "test-prompt.v1", + litellm_params: { prompt_id: "test-prompt.v1" }, + prompt_info: { prompt_type: "db" }, + version: 1, + }, + ]; + + mockGetPromptVersions.mockResolvedValueOnce({ + prompts: versionsWithoutDates, + }); + + render(); + + await waitFor(() => { + expect(screen.getByText("-")).toBeInTheDocument(); + }); + }); + }); + + describe("Prompt Type Display", () => { + it("should show 'Saved to Database' for db prompts", async () => { + render(); + + await waitFor(() => { + const dbTexts = screen.getAllByText("Saved to Database"); + expect(dbTexts.length).toBeGreaterThan(0); + }); + }); + + it("should show 'Config Prompt' for config prompts", async () => { + render(); + + await waitFor(() => { + expect(screen.getByText("Config Prompt")).toBeInTheDocument(); + }); + }); + }); + + describe("Network Calls and Data Fetching", () => { + it("should call getPromptVersions with correct parameters", async () => { + render(); + + await waitFor(() => { + expect(getPromptVersions).toHaveBeenCalledWith("test-token", "test-prompt"); + }); + }); + + it("should strip .v suffix from promptId when fetching versions", async () => { + render(); + + await waitFor(() => { + expect(getPromptVersions).toHaveBeenCalledWith("test-token", "test-prompt"); + }); + }); + + it("should not fetch versions when isOpen is false", () => { + render(); + + expect(getPromptVersions).not.toHaveBeenCalled(); + }); + + it("should not fetch versions when accessToken is null", () => { + render(); + + expect(getPromptVersions).not.toHaveBeenCalled(); + }); + + it("should not fetch versions when promptId is not provided", () => { + render(); + + expect(getPromptVersions).not.toHaveBeenCalled(); + }); + + it("should refetch versions when props change", async () => { + const { rerender } = render(); + + await waitFor(() => { + expect(getPromptVersions).toHaveBeenCalledTimes(1); + }); + + rerender(); + + await waitFor(() => { + expect(getPromptVersions).toHaveBeenCalledTimes(2); + expect(getPromptVersions).toHaveBeenCalledWith("test-token", "different-prompt"); + }); + }); + }); + + describe("Error Handling", () => { + it("should handle network errors gracefully", async () => { + const consoleSpy = vi.spyOn(console, "error").mockImplementation(() => {}); + mockGetPromptVersions.mockRejectedValueOnce(new Error("Network error")); + + render(); + + await waitFor(() => { + expect(consoleSpy).toHaveBeenCalledWith("Error fetching prompt versions:", expect.any(Error)); + }); + + // Should show empty state when there's an error + expect(screen.getByText("No version history available.")).toBeInTheDocument(); + + consoleSpy.mockRestore(); + }); + }); + + describe("User Interactions", () => { + it("should call onClose when close button is clicked", () => { + const mockOnClose = vi.fn(); + render(); + + const drawer = screen.getByTestId("drawer"); + act(() => { + fireEvent.click(drawer); // Simulate close action + }); + + // Note: This test assumes the drawer handles close events. + // In a real scenario, you'd test the actual close trigger. + }); + + it("should prevent interaction with main content when drawer is open", () => { + render(); + + const drawer = screen.getByTestId("drawer"); + // The mask and maskClosable props are passed as boolean false to disable them + expect(drawer).toHaveAttribute("data-mask", "false"); + expect(drawer).toHaveAttribute("data-maskclosable", "false"); + }); + }); + + describe("Edge Cases", () => { + it("should handle activeVersionId with .v suffix correctly", async () => { + render(); + + await waitFor(() => { + expect(screen.getByText("Active")).toBeInTheDocument(); + }); + }); + + it("should handle activeVersionId with _v suffix correctly", async () => { + const versionsWithUnderscore = [ + { + prompt_id: "test-prompt_v2", + litellm_params: { prompt_id: "test-prompt_v2" }, + prompt_info: { prompt_type: "db" }, + version: 2, + created_at: "2024-01-15T10:30:00Z", + }, + ]; + + mockGetPromptVersions.mockResolvedValueOnce({ + prompts: versionsWithUnderscore, + }); + + render(); + + await waitFor(() => { + expect(screen.getByText("Active")).toBeInTheDocument(); + }); + }); + + it("should sort versions correctly with version field", async () => { + // The component doesn't explicitly sort, but we can verify the order from the API response + render(); + + await waitFor(() => { + const versionElements = screen.getAllByTestId("tag"); + // Verify versions are displayed as they come from the API + expect(screen.getByText("v2")).toBeInTheDocument(); + }); + }); + }); +}); diff --git a/ui/litellm-dashboard/src/components/prompts/prompt_editor_view/utils.test.ts b/ui/litellm-dashboard/src/components/prompts/prompt_editor_view/utils.test.ts new file mode 100644 index 00000000000..59fcaccb639 --- /dev/null +++ b/ui/litellm-dashboard/src/components/prompts/prompt_editor_view/utils.test.ts @@ -0,0 +1,444 @@ +import { describe, expect, it } from "vitest"; +import { PromptType } from "./types"; +import { + convertToDotPrompt, + extractVariables, + getVersionNumber, + parseExistingPrompt, + stripVersionFromPromptId, +} from "./utils"; + +describe("extractVariables", () => { + it("should extract variables from messages", () => { + const prompt: PromptType = { + name: "test", + model: "gpt-4", + config: {}, + tools: [], + developerMessage: "", + messages: [ + { role: "user", content: "Hello {{name}}, how are you?" }, + { role: "assistant", content: "I am fine {{name}}" }, + ], + }; + + const result = extractVariables(prompt); + expect(result).toEqual(["name"]); + }); + + it("should extract variables from developer message", () => { + const prompt: PromptType = { + name: "test", + model: "gpt-4", + config: {}, + tools: [], + developerMessage: "You are {{role}} assistant", + messages: [{ role: "user", content: "Hello" }], + }; + + const result = extractVariables(prompt); + expect(result).toEqual(["role"]); + }); + + it("should extract variables from both messages and developer message", () => { + const prompt: PromptType = { + name: "test", + model: "gpt-4", + config: {}, + tools: [], + developerMessage: "You are {{role}} assistant", + messages: [ + { role: "user", content: "Hello {{name}}" }, + { role: "assistant", content: "Hi {{name}}, I am {{role}}" }, + ], + }; + + const result = extractVariables(prompt); + expect(result.sort()).toEqual(["name", "role"].sort()); + }); + + it("should return empty array when no variables present", () => { + const prompt: PromptType = { + name: "test", + model: "gpt-4", + config: {}, + tools: [], + developerMessage: "You are an assistant", + messages: [{ role: "user", content: "Hello world" }], + }; + + const result = extractVariables(prompt); + expect(result).toEqual([]); + }); + + it("should handle duplicate variables", () => { + const prompt: PromptType = { + name: "test", + model: "gpt-4", + config: {}, + tools: [], + developerMessage: "", + messages: [ + { role: "user", content: "Hello {{name}}" }, + { role: "assistant", content: "Hi {{name}} again" }, + ], + }; + + const result = extractVariables(prompt); + expect(result).toEqual(["name"]); + }); +}); + +describe("convertToDotPrompt", () => { + it("should convert basic prompt to dot prompt format", () => { + const prompt: PromptType = { + name: "test", + model: "gpt-4", + config: {}, + tools: [], + developerMessage: "", + messages: [{ role: "user", content: "Hello world" }], + }; + + const result = convertToDotPrompt(prompt); + expect(result).toContain("---"); + expect(result).toContain("model: gpt-4"); + expect(result).toContain("input:"); + expect(result).toContain("schema:"); + expect(result).toContain("output:"); + expect(result).toContain("format: text"); + expect(result).toContain("User: Hello world"); + }); + + it("should include config parameters when set", () => { + const prompt: PromptType = { + name: "test", + model: "gpt-4", + config: { + temperature: 0.7, + max_tokens: 100, + top_p: 0.9, + }, + tools: [], + developerMessage: "", + messages: [{ role: "user", content: "Hello" }], + }; + + const result = convertToDotPrompt(prompt); + expect(result).toContain("temperature: 0.7"); + expect(result).toContain("max_tokens: 100"); + expect(result).toContain("top_p: 0.9"); + }); + + it("should include input schema with variables", () => { + const prompt: PromptType = { + name: "test", + model: "gpt-4", + config: {}, + tools: [], + developerMessage: "", + messages: [{ role: "user", content: "Hello {{name}}" }], + }; + + const result = convertToDotPrompt(prompt); + expect(result).toContain("input:"); + expect(result).toContain("schema:"); + expect(result).toContain("name: string"); + }); + + it("should include developer message when present", () => { + const prompt: PromptType = { + name: "test", + model: "gpt-4", + config: {}, + tools: [], + developerMessage: "You are a helpful assistant", + messages: [{ role: "user", content: "Hello" }], + }; + + const result = convertToDotPrompt(prompt); + expect(result).toContain("Developer: You are a helpful assistant"); + }); + + it("should include tools when present", () => { + const prompt: PromptType = { + name: "test", + model: "gpt-4", + config: {}, + tools: [ + { + name: "get_weather", + description: "Get weather information", + json: '{"type": "function", "function": {"name": "get_weather"}}', + }, + ], + developerMessage: "", + messages: [{ role: "user", content: "Hello" }], + }; + + const result = convertToDotPrompt(prompt); + expect(result).toContain("tools:"); + expect(result).toContain('{"type":"function","function":{"name":"get_weather"}}'); + }); + + it("should handle multiple messages with different roles", () => { + const prompt: PromptType = { + name: "test", + model: "gpt-4", + config: {}, + tools: [], + developerMessage: "", + messages: [ + { role: "user", content: "Hello" }, + { role: "assistant", content: "Hi there" }, + { role: "user", content: "How are you?" }, + ], + }; + + const result = convertToDotPrompt(prompt); + expect(result).toContain("User: Hello"); + expect(result).toContain("Assistant: Hi there"); + expect(result).toContain("User: How are you?"); + }); +}); + +describe("parseExistingPrompt", () => { + it("should parse basic dotprompt content", () => { + const apiResponse = { + prompt_spec: { + litellm_params: { + dotprompt_content: `--- +model: gpt-4 +input: + schema: +output: + format: text +--- + +User: Hello world`, + }, + prompt_id: "test-prompt", + }, + }; + + const result = parseExistingPrompt(apiResponse); + expect(result.name).toBe("test-prompt"); + expect(result.model).toBe("gpt-4"); + expect(result.messages).toEqual([{ role: "user", content: "Hello world" }]); + }); + + it("should parse with config parameters", () => { + const apiResponse = { + prompt_spec: { + litellm_params: { + dotprompt_content: `--- +model: gpt-4 +temperature: 0.7 +max_tokens: 100 +top_p: 0.9 +input: + schema: +output: + format: text +--- + +User: Hello`, + }, + prompt_id: "test-prompt", + }, + }; + + const result = parseExistingPrompt(apiResponse); + expect(result.config.temperature).toBe(0.7); + expect(result.config.max_tokens).toBe(100); + expect(result.config.top_p).toBe(0.9); + }); + + it("should parse with developer message", () => { + const apiResponse = { + prompt_spec: { + litellm_params: { + dotprompt_content: `--- +model: gpt-4 +input: + schema: +output: + format: text +--- + +Developer: You are a helpful assistant + +User: Hello`, + }, + prompt_id: "test-prompt", + }, + }; + + const result = parseExistingPrompt(apiResponse); + expect(result.developerMessage).toBe("You are a helpful assistant"); + }); + + it("should parse multiple messages", () => { + const apiResponse = { + prompt_spec: { + litellm_params: { + dotprompt_content: `--- +model: gpt-4 +input: + schema: +output: + format: text +--- + +User: Hello +How are you? + +Assistant: I am fine +Thank you for asking + +User: Great!`, + }, + prompt_id: "test-prompt", + }, + }; + + const result = parseExistingPrompt(apiResponse); + expect(result.messages).toEqual([ + { role: "user", content: "Hello\nHow are you?" }, + { role: "assistant", content: "I am fine\nThank you for asking" }, + { role: "user", content: "Great!" }, + ]); + }); + + it("should handle prompt with version suffix", () => { + const apiResponse = { + prompt_spec: { + litellm_params: { + dotprompt_content: `--- +model: gpt-4 +input: + schema: +output: + format: text +--- + +User: Hello`, + }, + prompt_id: "test-prompt.v2", + }, + }; + + const result = parseExistingPrompt(apiResponse); + expect(result.name).toBe("test-prompt"); + }); + + it("should throw error when no dotprompt_content", () => { + const apiResponse = { + prompt_spec: { + litellm_params: {}, + }, + }; + + expect(() => parseExistingPrompt(apiResponse)).toThrow("No dotprompt_content found in API response"); + }); + + it("should throw error for invalid dotprompt format", () => { + const apiResponse = { + prompt_spec: { + litellm_params: { + dotprompt_content: "invalid format", + }, + }, + }; + + expect(() => parseExistingPrompt(apiResponse)).toThrow("Invalid dotprompt format"); + }); + + it("should provide default values when parsing fails", () => { + const apiResponse = { + prompt_spec: { + litellm_params: { + dotprompt_content: `--- +model: gpt-4 +input: + schema: +output: + format: text +--- + +`, + }, + prompt_id: "test-prompt", + }, + }; + + const result = parseExistingPrompt(apiResponse); + expect(result.messages).toEqual([ + { role: "user", content: "Enter task specifics. Use {{template_variables}} for dynamic inputs" }, + ]); + }); +}); + +describe("getVersionNumber", () => { + it("should return '1' for undefined promptId", () => { + const result = getVersionNumber(undefined); + expect(result).toBe("1"); + }); + + it("should return '1' for promptId without version", () => { + const result = getVersionNumber("test-prompt"); + expect(result).toBe("1"); + }); + + it("should extract version with dot separator", () => { + const result = getVersionNumber("test-prompt.v2"); + expect(result).toBe("2"); + }); + + it("should extract version with underscore separator", () => { + const result = getVersionNumber("test-prompt_v3"); + expect(result).toBe("3"); + }); + + it("should extract version with hyphen separator", () => { + const result = getVersionNumber("test-prompt-v4"); + expect(result).toBe("4"); + }); + + it("should extract multi-digit version", () => { + const result = getVersionNumber("test-prompt.v123"); + expect(result).toBe("123"); + }); +}); + +describe("stripVersionFromPromptId", () => { + it("should return empty string for undefined promptId", () => { + const result = stripVersionFromPromptId(undefined); + expect(result).toBe(""); + }); + + it("should return promptId unchanged when no version present", () => { + const result = stripVersionFromPromptId("test-prompt"); + expect(result).toBe("test-prompt"); + }); + + it("should strip version with dot separator", () => { + const result = stripVersionFromPromptId("test-prompt.v2"); + expect(result).toBe("test-prompt"); + }); + + it("should strip version with underscore separator", () => { + const result = stripVersionFromPromptId("test-prompt_v3"); + expect(result).toBe("test-prompt"); + }); + + it("should strip version with hyphen separator", () => { + const result = stripVersionFromPromptId("test-prompt-v4"); + expect(result).toBe("test-prompt"); + }); + + it("should strip multi-digit version", () => { + const result = stripVersionFromPromptId("test-prompt.v123"); + expect(result).toBe("test-prompt"); + }); +}); From e6da33dc4ac9f3c39c2b58cb113d2e4df62ec8f0 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Fri, 2 Jan 2026 19:22:56 -0800 Subject: [PATCH 127/158] Remove modal in useful links --- .../AIHub/UsefulLinksManagement.test.tsx | 158 +++++++++++++++++- .../AIHub/UsefulLinksManagement.tsx | 90 +++++----- 2 files changed, 189 insertions(+), 59 deletions(-) diff --git a/ui/litellm-dashboard/src/components/AIHub/UsefulLinksManagement.test.tsx b/ui/litellm-dashboard/src/components/AIHub/UsefulLinksManagement.test.tsx index 7b0651ad2d6..0a859ca95f8 100644 --- a/ui/litellm-dashboard/src/components/AIHub/UsefulLinksManagement.test.tsx +++ b/ui/litellm-dashboard/src/components/AIHub/UsefulLinksManagement.test.tsx @@ -1,10 +1,9 @@ +import NotificationsManager from "@/components/molecules/notifications_manager"; +import { getProxyBaseUrl, getPublicModelHubInfo, updateUsefulLinksCall } from "@/components/networking"; import { render, screen, waitFor } from "@testing-library/react"; import userEvent from "@testing-library/user-event"; -import { Modal } from "antd"; import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; import UsefulLinksManagement from "./UsefulLinksManagement"; -import NotificationsManager from "@/components/molecules/notifications_manager"; -import { getPublicModelHubInfo, updateUsefulLinksCall, getProxyBaseUrl } from "@/components/networking"; vi.mock("@/components/networking", () => ({ getPublicModelHubInfo: vi.fn(), @@ -25,8 +24,6 @@ const mockedUpdateUsefulLinksCall = vi.mocked(updateUsefulLinksCall); const mockedGetProxyBaseUrl = vi.mocked(getProxyBaseUrl); const mockedNotifications = vi.mocked(NotificationsManager); -let modalSuccessSpy: any; - describe("UsefulLinksManagement", () => { beforeEach(() => { mockedGetPublicModelHubInfo.mockResolvedValue({ @@ -37,11 +34,9 @@ describe("UsefulLinksManagement", () => { }); mockedUpdateUsefulLinksCall.mockResolvedValue({}); mockedGetProxyBaseUrl.mockReturnValue("https://proxy.example.com"); - modalSuccessSpy = vi.spyOn(Modal, "success").mockImplementation(() => ({ destroy: vi.fn() }) as any); }); afterEach(() => { - modalSuccessSpy.mockRestore(); vi.clearAllMocks(); }); @@ -108,4 +103,153 @@ describe("UsefulLinksManagement", () => { expect(mockedNotifications.success).toHaveBeenCalledWith("Link order saved successfully"); }); + + it("should display the Model Hub link", async () => { + render(); + + expect(await screen.findByRole("link", { name: /public model hub/i })).toBeInTheDocument(); + }); + + it("should edit a link when edit button is clicked", async () => { + const user = userEvent.setup(); + mockedGetPublicModelHubInfo.mockResolvedValue({ + docs_title: "Docs", + custom_docs_description: null, + litellm_version: "1.0.0", + useful_links: { + "Test Link": "https://test.example.com", + }, + }); + + render(); + + await waitFor(() => expect(screen.getByText("Test Link")).toBeInTheDocument()); + + // Click edit button + const editButton = screen.getByTestId("edit-link-0-Test Link"); + await user.click(editButton); + + // Should show input fields in edit mode + expect(screen.getByDisplayValue("Test Link")).toBeInTheDocument(); + expect(screen.getByDisplayValue("https://test.example.com")).toBeInTheDocument(); + }); + + it("should update a link when save is clicked in edit mode", async () => { + const user = userEvent.setup(); + mockedGetPublicModelHubInfo.mockResolvedValue({ + docs_title: "Docs", + custom_docs_description: null, + litellm_version: "1.0.0", + useful_links: { + "Test Link": "https://test.example.com", + }, + }); + + render(); + + await waitFor(() => expect(screen.getByText("Test Link")).toBeInTheDocument()); + + // Click edit button + const editButton = screen.getByTestId("edit-link-0-Test Link"); + await user.click(editButton); + + // Update the display name + const displayNameInput = screen.getByDisplayValue("Test Link"); + await user.clear(displayNameInput); + await user.type(displayNameInput, "Updated Link"); + + // Click save + await user.click(screen.getByRole("button", { name: /save/i })); + + await waitFor(() => + expect(mockedUpdateUsefulLinksCall).toHaveBeenCalledWith("token", { + "Updated Link": { url: "https://test.example.com", index: 0 }, + }), + ); + + expect(mockedNotifications.success).toHaveBeenCalledWith("Link updated successfully"); + }); + + it("should cancel editing when cancel button is clicked", async () => { + const user = userEvent.setup(); + mockedGetPublicModelHubInfo.mockResolvedValue({ + docs_title: "Docs", + custom_docs_description: null, + litellm_version: "1.0.0", + useful_links: { + "Test Link": "https://test.example.com", + }, + }); + + render(); + + await waitFor(() => expect(screen.getByText("Test Link")).toBeInTheDocument()); + + // Click edit button + const editButton = screen.getByTestId("edit-link-0-Test Link"); + await user.click(editButton); + + // Update the display name + const displayNameInput = screen.getByDisplayValue("Test Link"); + await user.clear(displayNameInput); + await user.type(displayNameInput, "Updated Link"); + + // Click cancel + await user.click(screen.getByRole("button", { name: /cancel/i })); + + // Should go back to normal view + expect(screen.getByText("Test Link")).toBeInTheDocument(); + expect(screen.queryByDisplayValue("Updated Link")).not.toBeInTheDocument(); + }); + + it("should not move down the last item in rearrange mode", async () => { + const user = userEvent.setup(); + mockedGetPublicModelHubInfo.mockResolvedValue({ + docs_title: "Docs", + custom_docs_description: null, + litellm_version: "1.0.0", + useful_links: { + "First Link": "https://first.example.com", + "Second Link": "https://second.example.com", + }, + }); + + render(); + + await waitFor(() => expect(screen.getByText("First Link")).toBeInTheDocument()); + + // Enter rearrange mode + await user.click(screen.getByRole("button", { name: /rearrange order/i })); + + // Try to move down the last item (should not do anything) + const secondLinkMoveDownButton = screen.getByTestId("move-down-1-Second Link"); + await user.click(secondLinkMoveDownButton); + + // Links should remain in same order + const linksAfter = screen.getAllByText(/First Link|Second Link/); + expect(linksAfter[0]).toHaveTextContent("First Link"); + expect(linksAfter[1]).toHaveTextContent("Second Link"); + }); + + it("should expand and collapse the component", async () => { + const user = userEvent.setup(); + render(); + + await waitFor(() => expect(screen.getByText("Link Management")).toBeInTheDocument()); + + // Initially expanded + expect(screen.getByText("Manage Existing Links")).toBeInTheDocument(); + + // Click to collapse + await user.click(screen.getByText("Link Management")); + + // Should be collapsed + expect(screen.queryByText("Manage Existing Links")).not.toBeInTheDocument(); + + // Click to expand again + await user.click(screen.getByText("Link Management")); + + // Should be expanded + expect(screen.getByText("Manage Existing Links")).toBeInTheDocument(); + }); }); diff --git a/ui/litellm-dashboard/src/components/AIHub/UsefulLinksManagement.tsx b/ui/litellm-dashboard/src/components/AIHub/UsefulLinksManagement.tsx index 220113ded13..c73eaf52384 100644 --- a/ui/litellm-dashboard/src/components/AIHub/UsefulLinksManagement.tsx +++ b/ui/litellm-dashboard/src/components/AIHub/UsefulLinksManagement.tsx @@ -1,11 +1,11 @@ -import React, { useState, useEffect } from "react"; -import { Modal } from "antd"; -import { PlusCircleIcon, ChevronDownIcon, ChevronRightIcon } from "@heroicons/react/outline"; -import { isAdminRole } from "@/utils/roles"; -import { getPublicModelHubInfo, updateUsefulLinksCall, getProxyBaseUrl } from "../networking"; -import { Card, Title, Text, Table, TableHead, TableHeaderCell, TableBody, TableRow, TableCell } from "@tremor/react"; -import NotificationsManager from "@/components/molecules/notifications_manager"; import TableIconActionButton from "@/components/common_components/IconActionButton/TableIconActionButtons/TableIconActionButton"; +import NotificationsManager from "@/components/molecules/notifications_manager"; +import { isAdminRole } from "@/utils/roles"; +import { ChevronDownIcon, ChevronRightIcon, ExternalLinkIcon, PlusCircleIcon } from "@heroicons/react/outline"; +import { Card, Table, TableBody, TableCell, TableHead, TableHeaderCell, TableRow, Text, Title } from "@tremor/react"; +import Link from "next/link"; +import React, { useEffect, useState } from "react"; +import { getProxyBaseUrl, getPublicModelHubInfo, updateUsefulLinksCall } from "../networking"; interface UsefulLinksManagementProps { accessToken: string | null; @@ -102,32 +102,6 @@ const UsefulLinksManagement: React.FC = ({ accessTok }); await updateUsefulLinksCall(accessToken, linksObject); - // show success modal with public model hub link - Modal.success({ - title: "Links Saved Successfully", - content: ( -
-

- Your useful links have been saved and are now visible on the public model hub. -

-
-

View your updated model hub:

- - Open Public Model Hub → - -
-
- ), - width: 500, - okText: "Close", - maskClosable: true, - keyboard: true, - }); return true; } catch (error) { @@ -319,29 +293,41 @@ const UsefulLinksManagement: React.FC = ({ accessTok
Manage Existing Links - {!isRearranging ? ( - - ) : ( -
+ Public Model Hub + + + {!isRearranging ? ( - -
- )} + ) : ( +
+ + +
+ )} +
From 0aae5153b6b59f9bbe6f479e70e4b97eaafa8365 Mon Sep 17 00:00:00 2001 From: Krish Dholakia Date: Sat, 3 Jan 2026 16:06:07 +0530 Subject: [PATCH 128/158] docs: Clarify Bedrock AgentCore documentation (#18603) Co-authored-by: Cursor Agent --- docs/my-website/docs/providers/bedrock_agentcore.md | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/docs/my-website/docs/providers/bedrock_agentcore.md b/docs/my-website/docs/providers/bedrock_agentcore.md index 43df7f82519..e3e352f7ab6 100644 --- a/docs/my-website/docs/providers/bedrock_agentcore.md +++ b/docs/my-website/docs/providers/bedrock_agentcore.md @@ -11,6 +11,12 @@ Call Bedrock AgentCore in the OpenAI Request/Response format. | Provider Route on LiteLLM | `bedrock/agentcore/{AGENT_RUNTIME_ARN}` | | Provider Doc | [AWS Bedrock AgentCore ↗](https://docs.aws.amazon.com/bedrock/latest/APIReference/API_agentcore_InvokeAgentRuntime.html) | +:::info + +This documentation is for **AgentCore Agents** (agent runtimes). If you want to use AgentCore MCP servers, add them as you would any other MCP server. See the [MCP documentation](https://docs.litellm.ai/docs/mcp) for details. + +::: + ## Quick Start ### Model Format to LiteLLM From 87fe62229f4b8b5dddefa7c22521eb5662928ca1 Mon Sep 17 00:00:00 2001 From: Krish Dholakia Date: Sat, 3 Jan 2026 21:51:19 +0530 Subject: [PATCH 129/158] feat: Add adopters page and data structure (#18605) Co-authored-by: Cursor Agent --- docs/my-website/src/data/adopters/README.md | 88 +++++++++++++++++++ .../src/data/adopters/adopters.json | 8 ++ docs/my-website/src/data/adopters/index.js | 23 +++++ .../img/adopters/placeholder-company.svg | 8 ++ 4 files changed, 127 insertions(+) create mode 100644 docs/my-website/src/data/adopters/README.md create mode 100644 docs/my-website/src/data/adopters/adopters.json create mode 100644 docs/my-website/src/data/adopters/index.js create mode 100644 docs/my-website/static/img/adopters/placeholder-company.svg diff --git a/docs/my-website/src/data/adopters/README.md b/docs/my-website/src/data/adopters/README.md new file mode 100644 index 00000000000..61a5215f802 --- /dev/null +++ b/docs/my-website/src/data/adopters/README.md @@ -0,0 +1,88 @@ +# LiteLLM Adopters + +This directory contains data for organizations that use LiteLLM in production. + +## Adding Your Organization + +We've made it super easy to add your organization! Just follow the steps below. + +### Quick Add (Recommended) + +**[Edit adopters.json on GitHub →](https://github.com/BerriAI/litellm/edit/main/docs/my-website/src/data/adopters/adopters.json)** + +This will open the GitHub editor in your browser where you can: + +1. Add your organization's entry to the JSON array +2. Commit your changes +3. GitHub will automatically create a pull request for you! + +No need to clone the repository or set up a development environment. + +### JSON Format + +Add your organization to the array in `adopters.json`: + +```json +{ + "name": "Your Organization Name", + "logoUrl": "https://yoursite.com/logo.svg", + "url": "https://yourcompany.com", + "description": "Brief description of how you use LiteLLM (shown on hover)" +} +``` + +### Fields + +- **`name`** (required): Your organization's display name +- **`logoUrl`** (required): URL to your logo - can be either: + - External URL: `https://yoursite.com/logo.svg` (easiest!) + - Local path: `/img/adopters/your-logo.svg` (requires uploading logo file) +- **`url`** (optional): Your organization's website (makes the logo clickable) +- **`description`** (optional): Brief description shown when users hover over your logo + +### Logo Options + +#### Option 1: External URL (Easiest) + +Simply provide a direct link to your logo hosted anywhere: + +```json +"logoUrl": "https://yourcompany.com/assets/logo.svg" +``` + +#### Option 2: Local Logo (Better Performance) + +If you prefer to host the logo locally: + +1. Add your logo to `docs/my-website/static/img/adopters/your-company.svg` +2. Reference it as: `"logoUrl": "/img/adopters/your-company.svg"` + +**Logo Specifications:** + +- **Format**: SVG preferred (PNG also acceptable) +- **Dimensions**: 240x160px or similar 3:2 ratio recommended +- **Background**: Transparent or white background works best + +### Example + +```json +{ + "name": "Acme Corporation", + "logoUrl": "https://acme.com/logo.svg", + "url": "https://acme.com", + "description": "Using LiteLLM to route requests across 50+ LLM providers" +} +``` + +### Display Order + +Adopters are displayed alphabetically by organization name, so your position will be determined automatically. + +### Need Help? + +If you have questions about adding your organization: + +- Ask in [GitHub Discussions](https://github.com/BerriAI/litellm/discussions) +- Join our [Discord community](https://discord.com/invite/wuPM9dRgDw) + +Thank you for supporting LiteLLM! 🚅 diff --git a/docs/my-website/src/data/adopters/adopters.json b/docs/my-website/src/data/adopters/adopters.json new file mode 100644 index 00000000000..52319c149e2 --- /dev/null +++ b/docs/my-website/src/data/adopters/adopters.json @@ -0,0 +1,8 @@ +[ + { + "name": "Your Logo Here", + "logoUrl": "/img/adopters/placeholder-company.svg", + "description": "Add your organization to show support for LiteLLM", + "url": "https://github.com/BerriAI/litellm/edit/main/docs/my-website/src/data/adopters/adopters.json" + } +] diff --git a/docs/my-website/src/data/adopters/index.js b/docs/my-website/src/data/adopters/index.js new file mode 100644 index 00000000000..b1a242dcc33 --- /dev/null +++ b/docs/my-website/src/data/adopters/index.js @@ -0,0 +1,23 @@ +import adoptersData from './adopters.json'; + +/** + * @typedef {Object} Adopter + * @property {string} name - The organization's display name + * @property {string} logoUrl - URL to the organization's logo + * @property {string} [url] - The organization's website URL + * @property {string} [description] - Brief description shown on hover + */ + +/** + * List of organizations using LiteLLM + * @type {Adopter[]} + */ +export const adopters = adoptersData; + +/** + * Adopters sorted alphabetically by name + * @type {Adopter[]} + */ +export const sortedAdopters = [...adopters].sort((a, b) => + a.name.localeCompare(b.name) +); diff --git a/docs/my-website/static/img/adopters/placeholder-company.svg b/docs/my-website/static/img/adopters/placeholder-company.svg new file mode 100644 index 00000000000..937dffc6eaf --- /dev/null +++ b/docs/my-website/static/img/adopters/placeholder-company.svg @@ -0,0 +1,8 @@ + + + + + + Add Your Logo + Click to contribute + From bdd05475bca872f944a7cf704d20cd3fdb72e0cc Mon Sep 17 00:00:00 2001 From: Cesar Garcia <128240629+Chesars@users.noreply.github.com> Date: Sat, 3 Jan 2026 15:39:00 -0300 Subject: [PATCH 130/158] fix: correct cost calculation when reasoning_tokens present without text_tokens (#18607) Fixes #18599 When OpenAI models (gpt-5-nano, o1-*, o3-*) and other providers return reasoning_tokens in completion_tokens_details but don't provide text_tokens, LiteLLM was incorrectly calculating costs using only reasoning_tokens, ignoring the remaining completion tokens. Changes: - Modified generic_cost_per_token() in llm_cost_calc/utils.py to calculate text_tokens as: completion_tokens - reasoning_tokens - audio_tokens - image_tokens when text_tokens is not explicitly provided - Added comprehensive test case test_reasoning_tokens_without_text_tokens_gpt5_nano() to verify all completion_tokens are billed correctly Example: - completion_tokens: 977 - reasoning_tokens: 768 - Before: only 768 tokens billed (21% less) - After: all 977 tokens billed correctly Affected models: - OpenAI: gpt-5-nano, o1-*, o3-* - Perplexity: sonar-reasoning* - Any model returning reasoning_tokens without text_tokens --- .../litellm_core_utils/llm_cost_calc/utils.py | 20 ++++++-- .../llm_cost_calc/test_llm_cost_calc_utils.py | 51 +++++++++++++++++++ 2 files changed, 66 insertions(+), 5 deletions(-) diff --git a/litellm/litellm_core_utils/llm_cost_calc/utils.py b/litellm/litellm_core_utils/llm_cost_calc/utils.py index e36d6d68367..cbc0763382c 100644 --- a/litellm/litellm_core_utils/llm_cost_calc/utils.py +++ b/litellm/litellm_core_utils/llm_cost_calc/utils.py @@ -604,12 +604,22 @@ def generic_cost_per_token( reasoning_tokens = completion_tokens_details["reasoning_tokens"] image_tokens = completion_tokens_details["image_tokens"] - # Only assume all tokens are text if there's NO breakdown at all - # If image_tokens, audio_tokens, or reasoning_tokens exist, respect text_tokens=0 + # Handle text_tokens calculation: + # 1. If text_tokens is explicitly provided and > 0, use it + # 2. If there's a breakdown (reasoning/audio/image tokens), calculate text_tokens as the remainder + # 3. If no breakdown at all, assume all completion_tokens are text_tokens has_token_breakdown = image_tokens > 0 or audio_tokens > 0 or reasoning_tokens > 0 - if text_tokens == 0 and not has_token_breakdown: - text_tokens = usage.completion_tokens - is_text_tokens_total = True + if text_tokens == 0: + if has_token_breakdown: + # Calculate text tokens as remainder when we have a breakdown + # This handles cases like OpenAI's reasoning models where text_tokens isn't provided + text_tokens = max( + 0, usage.completion_tokens - reasoning_tokens - audio_tokens - image_tokens + ) + else: + # No breakdown at all, all tokens are text tokens + text_tokens = usage.completion_tokens + is_text_tokens_total = True ## TEXT COST completion_cost = float(text_tokens) * completion_base_cost diff --git a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py index 65e3dbec8bd..5ba78d9eed1 100644 --- a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py +++ b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py @@ -809,3 +809,54 @@ def test_bedrock_anthropic_prompt_caching(): assert completion_cost >= 0 assert round(prompt_cost, 3) == 0.111 assert round(completion_cost, 5) == 0.00820 + + +def test_reasoning_tokens_without_text_tokens_gpt5_nano(): + """ + Test fix for GitHub issue #18599: + https://github.com/BerriAI/litellm/issues/18599 + + When OpenAI models (gpt-5-nano, o1, o3) return reasoning_tokens but don't provide + text_tokens, LiteLLM should calculate text_tokens as: + text_tokens = completion_tokens - reasoning_tokens - audio_tokens - image_tokens + + This ensures ALL completion tokens are billed, not just reasoning tokens. + """ + model = "gpt-5-nano" + custom_llm_provider = "openai" + + # Simulate OpenAI gpt-5-nano response where text_tokens is NOT provided + # completion_tokens: 977 total + # reasoning_tokens: 768 + # text_tokens: should be calculated as 977 - 768 = 209 + usage = Usage( + prompt_tokens=17, + completion_tokens=977, + total_tokens=994, + completion_tokens_details=CompletionTokensDetailsWrapper( + reasoning_tokens=768, + audio_tokens=0, + # text_tokens NOT provided - this is the key part of the bug + ), + ) + + prompt_cost, completion_cost = generic_cost_per_token( + model=model, + usage=usage, + custom_llm_provider=custom_llm_provider, + ) + + # gpt-5-nano pricing: $0.05/1M input, $0.40/1M output + expected_prompt_cost = 17 * 0.05 / 1_000_000 + expected_completion_cost = 977 * 0.40 / 1_000_000 # ALL tokens, not just reasoning + + assert abs(prompt_cost - expected_prompt_cost) < 1e-10, \ + f"Prompt cost incorrect: {prompt_cost} vs {expected_prompt_cost}" + + assert abs(completion_cost - expected_completion_cost) < 1e-10, \ + f"Completion cost incorrect: {completion_cost} vs {expected_completion_cost}" + + # Verify it's NOT using only reasoning_tokens (the bug) + wrong_cost = 768 * 0.40 / 1_000_000 # Only reasoning tokens + assert abs(completion_cost - wrong_cost) > 1e-6, \ + "Bug detected: Cost calculation is using only reasoning_tokens instead of all completion_tokens!" From 969790c4631f9efcbed0c55e0fb185ba98a0aa77 Mon Sep 17 00:00:00 2001 From: Krish Dholakia Date: Sun, 4 Jan 2026 00:10:07 +0530 Subject: [PATCH 131/158] Iam roles anywhere docs (#18559) * Add documentation for IAM Roles Anywhere Co-authored-by: krrishdholakia * Refactor Bedrock provider docs for IAM Roles Anywhere Co-authored-by: krrishdholakia --------- Co-authored-by: Cursor Agent --- docs/my-website/docs/providers/bedrock.md | 47 +++++++++++++++++++++++ 1 file changed, 47 insertions(+) diff --git a/docs/my-website/docs/providers/bedrock.md b/docs/my-website/docs/providers/bedrock.md index 122554fe8a4..f1eed4b4d52 100644 --- a/docs/my-website/docs/providers/bedrock.md +++ b/docs/my-website/docs/providers/bedrock.md @@ -2208,6 +2208,53 @@ response = completion( | `aws_role_name` | `RoleArn` | The Amazon Resource Name (ARN) of the role to assume | [AssumeRole API](https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/sts.html#STS.Client.assume_role) | | `aws_session_name` | `RoleSessionName` | An identifier for the assumed role session | [AssumeRole API](https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/sts.html#STS.Client.assume_role) | +### IAM Roles Anywhere (On-Premise / External Workloads) + +[IAM Roles Anywhere](https://docs.aws.amazon.com/rolesanywhere/latest/userguide/introduction.html) extends IAM roles to workloads **outside of AWS** (on-premise servers, edge devices, other clouds). It uses the same STS mechanism as regular IAM roles but authenticates via X.509 certificates instead of AWS credentials. + +**Setup**: Configure the [AWS Signing Helper](https://docs.aws.amazon.com/rolesanywhere/latest/userguide/credential-helper.html) as a credential process in `~/.aws/config`: + +```ini +[profile litellm-roles-anywhere] +credential_process = aws_signing_helper credential-process \ + --certificate /path/to/certificate.pem \ + --private-key /path/to/private-key.pem \ + --trust-anchor-arn arn:aws:rolesanywhere:us-east-1:123456789012:trust-anchor/abc123 \ + --profile-arn arn:aws:rolesanywhere:us-east-1:123456789012:profile/def456 \ + --role-arn arn:aws:iam::123456789012:role/MyBedrockRole +``` + +**Usage**: Reference the profile in LiteLLM: + + + + +```python +from litellm import completion + +response = completion( + model="bedrock/anthropic.claude-3-sonnet-20240229-v1:0", + messages=[{"role": "user", "content": "Hello!"}], + aws_profile_name="litellm-roles-anywhere", +) +``` + + + + +```yaml +model_list: + - model_name: bedrock-claude + litellm_params: + model: bedrock/anthropic.claude-3-sonnet-20240229-v1:0 + aws_profile_name: "litellm-roles-anywhere" +``` + + + + +See the [IAM Roles Anywhere Getting Started Guide](https://docs.aws.amazon.com/rolesanywhere/latest/userguide/getting-started.html) for trust anchor and profile setup. + Make the bedrock completion call From a3503e59c227f5f1dd15e9d8910f67a6b7dca2a4 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Mehmet=20Can=20=C5=9Eakiro=C4=9Flu?= <53798389+cansakiroglu@users.noreply.github.com> Date: Sat, 3 Jan 2026 21:52:50 +0300 Subject: [PATCH 132/158] Litellm feat helm lifecycle support (#18517) * feat(helm): add lifecycle hook support for helm * add tests --- .../litellm-helm/templates/deployment.yaml | 4 ++++ .../litellm-helm/tests/deployment_tests.yaml | 24 ++++++++++++++++++- 2 files changed, 27 insertions(+), 1 deletion(-) diff --git a/deploy/charts/litellm-helm/templates/deployment.yaml b/deploy/charts/litellm-helm/templates/deployment.yaml index 0dab2ec40e0..19fa0479091 100644 --- a/deploy/charts/litellm-helm/templates/deployment.yaml +++ b/deploy/charts/litellm-helm/templates/deployment.yaml @@ -182,6 +182,10 @@ spec: {{- with .Values.volumeMounts }} {{- toYaml . | nindent 12 }} {{- end }} + {{- with .Values.lifecycle }} + lifecycle: + {{- toYaml . | nindent 12 }} + {{- end }} {{- with .Values.extraContainers }} {{- toYaml . | nindent 8 }} {{- end }} diff --git a/deploy/charts/litellm-helm/tests/deployment_tests.yaml b/deploy/charts/litellm-helm/tests/deployment_tests.yaml index f9c83966696..182a2362392 100644 --- a/deploy/charts/litellm-helm/tests/deployment_tests.yaml +++ b/deploy/charts/litellm-helm/tests/deployment_tests.yaml @@ -136,4 +136,26 @@ tests: path: spec.template.spec.containers[0].volumeMounts content: name: litellm-config - mountPath: /etc/litellm/ \ No newline at end of file + mountPath: /etc/litellm/ + - it: should work with lifecycle hooks + template: deployment.yaml + set: + lifecycle: + preStop: + exec: + command: + - /bin/sh + - -c + - echo "Container stopping" + asserts: + - exists: + path: spec.template.spec.containers[0].lifecycle + - equal: + path: spec.template.spec.containers[0].lifecycle.preStop.exec.command[0] + value: /bin/sh + - equal: + path: spec.template.spec.containers[0].lifecycle.preStop.exec.command[1] + value: -c + - equal: + path: spec.template.spec.containers[0].lifecycle.preStop.exec.command[2] + value: echo "Container stopping" \ No newline at end of file From 3a4ebf173f637dc9aa1fdc037b27a3757e246851 Mon Sep 17 00:00:00 2001 From: Deepak Walia <58362408+dee-walia20@users.noreply.github.com> Date: Sun, 4 Jan 2026 00:35:53 +0530 Subject: [PATCH 133/158] fix(sap): honor allowed_openai_params in transform_request (#18432) --- litellm/llms/sap/chat/transformation.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/litellm/llms/sap/chat/transformation.py b/litellm/llms/sap/chat/transformation.py index 01ceb72c0de..e13abca59f8 100755 --- a/litellm/llms/sap/chat/transformation.py +++ b/litellm/llms/sap/chat/transformation.py @@ -203,9 +203,13 @@ class GenAIHubOrchestrationConfig(OpenAIGPTConfig): headers: dict, ) -> dict: supported_params = self.get_supported_openai_params(model) + # Include extra params that passed validation (e.g., thinking_config for Gemini models via allowed_openai_params) + extra_params = [k for k in optional_params if k not in supported_params and k not in {"tools", "model_version"}] + supported_params = supported_params + extra_params model_params = { k: v for k, v in optional_params.items() if k in supported_params } + model_version = optional_params.pop("model_version", "latest") template = [] for message in messages: From dc62cdb3009bff03a2282de1a981f29421e0383c Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E3=82=86=E3=82=8A?= Date: Sun, 4 Jan 2026 03:07:52 +0800 Subject: [PATCH 134/158] fix: handle empty error objects in response conversion (#18493) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Some OpenAI-compatible providers (e.g., Apertis) return empty error objects even on successful responses. The previous check only verified that error was not None, causing spurious APIErrors. Now the code checks if the error object contains meaningful data: - For dict errors: non-empty message OR non-null code - For string errors: non-empty string - Other truthy values are still treated as errors Fixes #18407 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-authored-by: yurekami Co-authored-by: Claude Opus 4.5 --- .../convert_dict_to_response.py | 46 ++++-- .../test_convert_dict_to_chat_completion.py | 156 ++++++++++++++++++ 2 files changed, 188 insertions(+), 14 deletions(-) diff --git a/litellm/litellm_core_utils/llm_response_utils/convert_dict_to_response.py b/litellm/litellm_core_utils/llm_response_utils/convert_dict_to_response.py index 59d2a8a8dd0..bbe28e3ec2c 100644 --- a/litellm/litellm_core_utils/llm_response_utils/convert_dict_to_response.py +++ b/litellm/litellm_core_utils/llm_response_utils/convert_dict_to_response.py @@ -445,25 +445,43 @@ def convert_to_model_response_object( # noqa: PLR0915 hidden_params["additional_headers"] = additional_headers ### CHECK IF ERROR IN RESPONSE ### - openrouter returns these in the dictionary + # Some OpenAI-compatible providers (e.g., Apertis) return empty error objects + # even on success. Only raise if the error contains meaningful data. if ( response_object is not None and "error" in response_object and response_object["error"] is not None ): - error_args = {"status_code": 422, "message": "Error in response object"} - if isinstance(response_object["error"], dict): - if "code" in response_object["error"]: - error_args["status_code"] = response_object["error"]["code"] - if "message" in response_object["error"]: - if isinstance(response_object["error"]["message"], dict): - message_str = json.dumps(response_object["error"]["message"]) - else: - message_str = str(response_object["error"]["message"]) - error_args["message"] = message_str - raised_exception = Exception() - setattr(raised_exception, "status_code", error_args["status_code"]) - setattr(raised_exception, "message", error_args["message"]) - raise raised_exception + error_obj = response_object["error"] + has_meaningful_error = False + + if isinstance(error_obj, dict): + # Check if error dict has non-empty message or non-null code + error_message = error_obj.get("message", "") + error_code = error_obj.get("code") + has_meaningful_error = bool(error_message) or error_code is not None + elif isinstance(error_obj, str): + # String error is meaningful if non-empty + has_meaningful_error = bool(error_obj) + else: + # Any other truthy value is considered meaningful + has_meaningful_error = True + + if has_meaningful_error: + error_args = {"status_code": 422, "message": "Error in response object"} + if isinstance(error_obj, dict): + if "code" in error_obj: + error_args["status_code"] = error_obj["code"] + if "message" in error_obj: + if isinstance(error_obj["message"], dict): + message_str = json.dumps(error_obj["message"]) + else: + message_str = str(error_obj["message"]) + error_args["message"] = message_str + raised_exception = Exception() + setattr(raised_exception, "status_code", error_args["status_code"]) + setattr(raised_exception, "message", error_args["message"]) + raise raised_exception try: if response_type == "completion" and ( diff --git a/tests/llm_translation/test_llm_response_utils/test_convert_dict_to_chat_completion.py b/tests/llm_translation/test_llm_response_utils/test_convert_dict_to_chat_completion.py index 7e269f21451..c151150f634 100644 --- a/tests/llm_translation/test_llm_response_utils/test_convert_dict_to_chat_completion.py +++ b/tests/llm_translation/test_llm_response_utils/test_convert_dict_to_chat_completion.py @@ -903,3 +903,159 @@ def test_convert_to_model_response_object_with_thinking_content(): resp: ModelResponse = convert_to_model_response_object(**args) assert resp is not None assert resp.choices[0].message.reasoning_content is not None + + +def test_convert_to_model_response_object_with_empty_error_object(): + """ + Test that convert_to_model_response_object handles empty error objects gracefully. + + This is a regression test for issue #18407 where providers like Apertis return + empty error objects even on successful responses, causing spurious APIErrors. + + The error object structure: + { + "error": { + "message": "", + "type": "", + "param": "", + "code": null + } + } + """ + response_object = { + "model": "minimax-m2.1", + "choices": [ + { + "index": 0, + "message": { + "role": "assistant", + "content": "Hey! I'm doing well, thanks for asking!", + }, + "finish_reason": "stop", + } + ], + "usage": { + "prompt_tokens": 49, + "completion_tokens": 87, + "total_tokens": 136, + }, + "error": { + "message": "", + "type": "", + "param": "", + "code": None, + }, + } + + # This should NOT raise an exception + result = convert_to_model_response_object( + model_response_object=ModelResponse(), + response_object=response_object, + stream=False, + start_time=datetime.now(), + end_time=datetime.now(), + hidden_params=None, + _response_headers=None, + convert_tool_call_to_json_mode=False, + ) + + assert isinstance(result, ModelResponse) + assert result.model == "minimax-m2.1" + assert len(result.choices) == 1 + assert result.choices[0].message.content == "Hey! I'm doing well, thanks for asking!" + + +def test_convert_to_model_response_object_with_real_error(): + """ + Test that convert_to_model_response_object still raises for real errors. + + Ensures the empty error fix doesn't break legitimate error handling. + """ + response_object = { + "error": { + "message": "Rate limit exceeded", + "type": "rate_limit_error", + "param": None, + "code": 429, + }, + } + + with pytest.raises(Exception) as exc_info: + convert_to_model_response_object( + model_response_object=ModelResponse(), + response_object=response_object, + stream=False, + start_time=datetime.now(), + end_time=datetime.now(), + hidden_params=None, + _response_headers=None, + convert_tool_call_to_json_mode=False, + ) + + # The exception should have the error message + assert hasattr(exc_info.value, "message") + assert "Rate limit exceeded" in str(exc_info.value.message) + + +def test_convert_to_model_response_object_with_empty_dict_error(): + """ + Test that convert_to_model_response_object handles completely empty error dict. + """ + response_object = { + "model": "test-model", + "choices": [ + { + "index": 0, + "message": { + "role": "assistant", + "content": "Hello!", + }, + "finish_reason": "stop", + } + ], + "usage": { + "prompt_tokens": 10, + "completion_tokens": 5, + "total_tokens": 15, + }, + "error": {}, # Completely empty error object + } + + # This should NOT raise an exception + result = convert_to_model_response_object( + model_response_object=ModelResponse(), + response_object=response_object, + stream=False, + start_time=datetime.now(), + end_time=datetime.now(), + hidden_params=None, + _response_headers=None, + convert_tool_call_to_json_mode=False, + ) + + assert isinstance(result, ModelResponse) + assert result.choices[0].message.content == "Hello!" + + +def test_convert_to_model_response_object_with_error_code_only(): + """ + Test that errors with only a code (no message) are still treated as real errors. + """ + response_object = { + "error": { + "message": "", + "code": 500, + }, + } + + with pytest.raises(Exception): + convert_to_model_response_object( + model_response_object=ModelResponse(), + response_object=response_object, + stream=False, + start_time=datetime.now(), + end_time=datetime.now(), + hidden_params=None, + _response_headers=None, + convert_tool_call_to_json_mode=False, + ) From 9ba27d85cee19e2cced4fbac4a2b68e7d7ae7dcd Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E3=82=86=E3=82=8A?= Date: Sun, 4 Jan 2026 03:08:32 +0800 Subject: [PATCH 135/158] feat(types): add output_text property to ResponsesAPIResponse (#18491) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Add the output_text convenience property to ResponsesAPIResponse that aggregates all output_text items from the output list, matching the OpenAI SDK's Response.output_text behavior. The property iterates through output items, collects text content from message-type outputs, and returns them concatenated into a single string. Returns empty string if no output_text content exists. Handles both dict and Pydantic model access patterns for compatibility with different output formats. Fixes #18470 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-authored-by: yurekami Co-authored-by: Claude Opus 4.5 --- litellm/types/llms/openai.py | 33 +++++ .../types/llms/test_types_llms_openai.py | 134 ++++++++++++++++++ 2 files changed, 167 insertions(+) diff --git a/litellm/types/llms/openai.py b/litellm/types/llms/openai.py index ceeae958a80..c2912558cab 100644 --- a/litellm/types/llms/openai.py +++ b/litellm/types/llms/openai.py @@ -1197,6 +1197,39 @@ class ResponsesAPIResponse(BaseLiteLLMOpenAIResponseObject): # Define private attributes using PrivateAttr _hidden_params: dict = PrivateAttr(default_factory=dict) + @property + def output_text(self) -> str: + """ + Convenience property that aggregates all `output_text` items from the `output` list. + + If no `output_text` content blocks exist, then an empty string is returned. + + This matches the OpenAI SDK's Response.output_text behavior. + """ + texts: List[str] = [] + for output_item in self.output: + # Handle both dict and object access patterns + if isinstance(output_item, dict): + item_type = output_item.get("type") + content = output_item.get("content", []) + else: + item_type = getattr(output_item, "type", None) + content = getattr(output_item, "content", []) + + if item_type == "message": + for content_item in content: + if isinstance(content_item, dict): + content_type = content_item.get("type") + text = content_item.get("text", "") + else: + content_type = getattr(content_item, "type", None) + text = getattr(content_item, "text", "") or "" + + if content_type == "output_text": + texts.append(text) + + return "".join(texts) + class ResponsesAPIStreamEvents(str, Enum): """ diff --git a/tests/test_litellm/types/llms/test_types_llms_openai.py b/tests/test_litellm/types/llms/test_types_llms_openai.py index 05dec06d469..87cc9586665 100644 --- a/tests/test_litellm/types/llms/test_types_llms_openai.py +++ b/tests/test_litellm/types/llms/test_types_llms_openai.py @@ -35,3 +35,137 @@ def test_output_item_added_event(): assert event.sequence_number == 4 assert event.output_index == 1 assert event.item is None + + +class TestResponsesAPIResponseOutputText: + """Tests for the output_text property on ResponsesAPIResponse""" + + def test_output_text_with_single_message(self): + """Test output_text with a single message containing text output""" + from litellm.types.llms.openai import ResponsesAPIResponse + + response = ResponsesAPIResponse( + id="resp_123", + created_at=1234567890, + output=[ + { + "type": "message", + "id": "msg_123", + "status": "completed", + "role": "assistant", + "content": [ + { + "type": "output_text", + "text": "Hello, world!", + } + ], + } + ], + ) + + assert response.output_text == "Hello, world!" + + def test_output_text_with_multiple_messages(self): + """Test output_text with multiple messages aggregates all text""" + from litellm.types.llms.openai import ResponsesAPIResponse + + response = ResponsesAPIResponse( + id="resp_123", + created_at=1234567890, + output=[ + { + "type": "message", + "id": "msg_1", + "status": "completed", + "role": "assistant", + "content": [ + { + "type": "output_text", + "text": "First part. ", + } + ], + }, + { + "type": "message", + "id": "msg_2", + "status": "completed", + "role": "assistant", + "content": [ + { + "type": "output_text", + "text": "Second part.", + } + ], + }, + ], + ) + + assert response.output_text == "First part. Second part." + + def test_output_text_with_no_text_content(self): + """Test output_text returns empty string when no output_text content exists""" + from litellm.types.llms.openai import ResponsesAPIResponse + + response = ResponsesAPIResponse( + id="resp_123", + created_at=1234567890, + output=[ + { + "type": "function_call", + "id": "call_123", + "status": "completed", + "name": "get_weather", + "arguments": "{}", + } + ], + ) + + assert response.output_text == "" + + def test_output_text_with_mixed_content(self): + """Test output_text only aggregates output_text type content""" + from litellm.types.llms.openai import ResponsesAPIResponse + + response = ResponsesAPIResponse( + id="resp_123", + created_at=1234567890, + output=[ + { + "type": "message", + "id": "msg_1", + "status": "completed", + "role": "assistant", + "content": [ + { + "type": "output_text", + "text": "The weather is sunny. ", + }, + { + "type": "refusal", + "refusal": "I cannot do that.", + }, + ], + }, + { + "type": "function_call", + "id": "call_123", + "status": "completed", + "name": "get_weather", + "arguments": "{}", + }, + ], + ) + + assert response.output_text == "The weather is sunny. " + + def test_output_text_with_empty_output(self): + """Test output_text returns empty string with empty output list""" + from litellm.types.llms.openai import ResponsesAPIResponse + + response = ResponsesAPIResponse( + id="resp_123", + created_at=1234567890, + output=[], + ) + + assert response.output_text == "" From 37c908caf9777a757f813596ec04897aa995170f Mon Sep 17 00:00:00 2001 From: Lu Date: Sun, 4 Jan 2026 03:13:22 +0800 Subject: [PATCH 136/158] google genai adapter inline data support (#18477) * support inline data * add test --- .../google_genai/adapters/transformation.py | 62 +++++++-- .../google_genai/test_google_genai_adapter.py | 127 +++++++++++++++++- 2 files changed, 179 insertions(+), 10 deletions(-) diff --git a/litellm/google_genai/adapters/transformation.py b/litellm/google_genai/adapters/transformation.py index 9d3f990b1aa..58a52666d38 100644 --- a/litellm/google_genai/adapters/transformation.py +++ b/litellm/google_genai/adapters/transformation.py @@ -8,8 +8,10 @@ from litellm.types.llms.openai import ( AllMessageValues, ChatCompletionAssistantMessage, ChatCompletionAssistantToolCall, + ChatCompletionImageObject, ChatCompletionRequest, ChatCompletionSystemMessage, + ChatCompletionTextObject, ChatCompletionToolCallFunctionChunk, ChatCompletionToolChoiceValues, ChatCompletionToolMessage, @@ -385,13 +387,36 @@ class GoogleGenAIAdapter: if role == "user": # Handle user messages with potential function responses - combined_text = "" + content_parts: List[ + Union[ChatCompletionTextObject, ChatCompletionImageObject] + ] = [] tool_messages: List[ChatCompletionToolMessage] = [] for part in parts: if isinstance(part, dict): if "text" in part: - combined_text += part["text"] + content_parts.append( + cast( + ChatCompletionTextObject, + {"type": "text", "text": part["text"]}, + ) + ) + elif "inline_data" in part: + # Handle Base64 image data + inline_data = part["inline_data"] + mime_type = inline_data.get("mime_type", "image/jpeg") + data = inline_data.get("data", "") + content_parts.append( + cast( + ChatCompletionImageObject, + { + "type": "image_url", + "image_url": { + "url": f"data:{mime_type};base64,{data}" + }, + }, + ) + ) elif "functionResponse" in part: # Transform function response to tool message func_response = part["functionResponse"] @@ -402,13 +427,33 @@ class GoogleGenAIAdapter: ) tool_messages.append(tool_message) elif isinstance(part, str): - combined_text += part + content_parts.append( + cast( + ChatCompletionTextObject, {"type": "text", "text": part} + ) + ) - # Add user message if there's text content - if combined_text: - messages.append( - ChatCompletionUserMessage(role="user", content=combined_text) - ) + # Add user message if there's content + if content_parts: + # If only one text part, use simple string format for backward compatibility + if ( + len(content_parts) == 1 + and isinstance(content_parts[0], dict) + and content_parts[0].get("type") == "text" + ): + text_part = cast(ChatCompletionTextObject, content_parts[0]) + messages.append( + ChatCompletionUserMessage( + role="user", content=text_part["text"] + ) + ) + else: + # Use multimodal format (array of content parts) + messages.append( + ChatCompletionUserMessage( + role="user", content=content_parts + ) + ) # Add tool messages messages.extend(tool_messages) @@ -468,7 +513,6 @@ class GoogleGenAIAdapter: Dict in Google GenAI generate_content response format """ - # Extract the main response content choice = response.choices[0] if response.choices else None if not choice: diff --git a/tests/test_litellm/google_genai/test_google_genai_adapter.py b/tests/test_litellm/google_genai/test_google_genai_adapter.py index e8882a1acb3..135881ad209 100644 --- a/tests/test_litellm/google_genai/test_google_genai_adapter.py +++ b/tests/test_litellm/google_genai/test_google_genai_adapter.py @@ -1197,6 +1197,131 @@ async def test_agenerate_content_x_goog_api_key_header(): # Verify other expected headers assert headers.get("Content-Type") == "application/json", f"Expected Content-Type application/json, got {headers.get('Content-Type')}" - + print(f"✓ Test passed: x-goog-api-key header correctly set to {api_key_value}") print(f"✓ All headers: {list(headers.keys())}") + + +def test_inline_data_base64_image_transformation(): + """Test transformation of Gemini inline_data (Base64 images) to OpenAI format""" + from litellm.google_genai.adapters.transformation import GoogleGenAIAdapter + + adapter = GoogleGenAIAdapter() + + # Test input with Base64 image + model = "gpt-4-vision-preview" + contents = { + "role": "user", + "parts": [ + {"text": "What's in this image?"}, + { + "inline_data": { + "mime_type": "image/jpeg", + "data": "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNk+M9QDwADhgGAWjR9awAAAABJRU5ErkJggg==" + } + } + ] + } + + # Transform to completion format + completion_request = adapter.translate_generate_content_to_completion( + model=model, + contents=contents + ) + + # Verify the transformation + assert completion_request["model"] == model + assert len(completion_request["messages"]) == 1 + assert completion_request["messages"][0]["role"] == "user" + + # Verify content is an array (multimodal format) + content = completion_request["messages"][0]["content"] + assert isinstance(content, list), "Content should be a list for multimodal messages" + assert len(content) == 2, "Should have 2 content parts (text + image)" + + # Verify text part + text_part = content[0] + assert text_part["type"] == "text" + assert text_part["text"] == "What's in this image?" + + # Verify image part + image_part = content[1] + assert image_part["type"] == "image_url" + assert "image_url" in image_part + assert "url" in image_part["image_url"] + assert image_part["image_url"]["url"].startswith("data:image/jpeg;base64,") + assert "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNk+M9QDwADhgGAWjR9awAAAABJRU5ErkJggg==" in image_part["image_url"]["url"] + + +def test_inline_data_image_only_transformation(): + """Test transformation of Gemini inline_data with only image (no text)""" + from litellm.google_genai.adapters.transformation import GoogleGenAIAdapter + + adapter = GoogleGenAIAdapter() + + # Test input with only Base64 image (no text) + model = "gpt-4-vision-preview" + contents = { + "role": "user", + "parts": [ + { + "inline_data": { + "mime_type": "image/png", + "data": "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNk+M9QDwADhgGAWjR9awAAAABJRU5ErkJggg==" + } + } + ] + } + + # Transform to completion format + completion_request = adapter.translate_generate_content_to_completion( + model=model, + contents=contents + ) + + # Verify the transformation + assert completion_request["model"] == model + assert len(completion_request["messages"]) == 1 + assert completion_request["messages"][0]["role"] == "user" + + # Verify content is an array (multimodal format) + content = completion_request["messages"][0]["content"] + assert isinstance(content, list), "Content should be a list for multimodal messages" + assert len(content) == 1, "Should have 1 content part (image only)" + + # Verify image part + image_part = content[0] + assert image_part["type"] == "image_url" + assert "image_url" in image_part + assert "url" in image_part["image_url"] + assert image_part["image_url"]["url"].startswith("data:image/png;base64,") + + +def test_inline_data_backward_compatibility_text_only(): + """Test that pure text messages still use simple string format (backward compatibility)""" + from litellm.google_genai.adapters.transformation import GoogleGenAIAdapter + + adapter = GoogleGenAIAdapter() + + # Test input with only text (no images) + model = "gpt-3.5-turbo" + contents = { + "role": "user", + "parts": [{"text": "Hello, how are you?"}] + } + + # Transform to completion format + completion_request = adapter.translate_generate_content_to_completion( + model=model, + contents=contents + ) + + # Verify the transformation + assert completion_request["model"] == model + assert len(completion_request["messages"]) == 1 + assert completion_request["messages"][0]["role"] == "user" + + # Verify content is a simple string (not an array) for backward compatibility + content = completion_request["messages"][0]["content"] + assert isinstance(content, str), "Content should be a string for text-only messages (backward compatibility)" + assert content == "Hello, how are you?" From 9b1c5f7e360e7b448d9f16500e4ababba6f47b42 Mon Sep 17 00:00:00 2001 From: cantalupo555 Date: Sat, 3 Jan 2026 16:14:19 -0300 Subject: [PATCH 137/158] feat(zai): Add GLM-4.7 model with reasoning support (#18476) Add support for Z.AI GLM-4.7, latest flagship model with enhanced reasoning capabilities. Changes: - Add zai/glm-4.7 to model pricing with /bin/bash.60/M input, .20/M output - Add cached input pricing (/bin/bash.11/M) for GLM-4.7 - Add supports_reasoning flag to enable thinking parameter - Update ZAIChatConfig to support thinking parameter for models with reasoning - Update documentation with GLM-4.7 as latest flagship model - Add cached input column to pricing table (GLM-4.7 only) - Add tests for GLM-4.7 reasoning support and cost calculation - Update all examples to use GLM-4.7 Model specifications: - Context: 200K input, 128K output - Supports: reasoning, function calling, tool choice, prompt caching - Pricing: Same as GLM-4.6 with cache support See: https://docs.z.ai/guides/llm/glm-4.7 --- docs/my-website/docs/providers/zai.md | 36 +++++++++--------- litellm/llms/zai/chat/transformation.py | 11 +++++- ...odel_prices_and_context_window_backup.json | 14 +++++++ model_prices_and_context_window.json | 14 +++++++ .../llms/zai/test_zai_provider.py | 37 +++++++++++++++++++ 5 files changed, 94 insertions(+), 18 deletions(-) diff --git a/docs/my-website/docs/providers/zai.md b/docs/my-website/docs/providers/zai.md index 5055d0c1cdd..937ccd67680 100644 --- a/docs/my-website/docs/providers/zai.md +++ b/docs/my-website/docs/providers/zai.md @@ -19,7 +19,7 @@ import os os.environ['ZAI_API_KEY'] = "" response = completion( - model="zai/glm-4.6", + model="zai/glm-4.7", messages=[ {"role": "user", "content": "hello from litellm"} ], @@ -34,7 +34,7 @@ import os os.environ['ZAI_API_KEY'] = "" response = completion( - model="zai/glm-4.6", + model="zai/glm-4.7", messages=[ {"role": "user", "content": "hello from litellm"} ], @@ -51,7 +51,8 @@ We support ALL Z.AI GLM models, just set `zai/` as a prefix when sending complet | Model Name | Function Call | Notes | |------------|---------------|-------| -| glm-4.6 | `completion(model="zai/glm-4.6", messages)` | Latest flagship model, 200K context | +| glm-4.7 | `completion(model="zai/glm-4.7", messages)` | **Latest flagship**, 200K context, **Reasoning** | +| glm-4.6 | `completion(model="zai/glm-4.6", messages)` | 200K context | | glm-4.5 | `completion(model="zai/glm-4.5", messages)` | 128K context | | glm-4.5v | `completion(model="zai/glm-4.5v", messages)` | Vision model | | glm-4.5-x | `completion(model="zai/glm-4.5-x", messages)` | Premium tier | @@ -62,16 +63,17 @@ We support ALL Z.AI GLM models, just set `zai/` as a prefix when sending complet ## Model Pricing -| Model | Input ($/1M tokens) | Output ($/1M tokens) | Context Window | -|-------|---------------------|----------------------|----------------| -| glm-4.6 | $0.60 | $2.20 | 200K | -| glm-4.5 | $0.60 | $2.20 | 128K | -| glm-4.5v | $0.60 | $1.80 | 128K | -| glm-4.5-x | $2.20 | $8.90 | 128K | -| glm-4.5-air | $0.20 | $1.10 | 128K | -| glm-4.5-airx | $1.10 | $4.50 | 128K | -| glm-4-32b-0414-128k | $0.10 | $0.10 | 128K | -| glm-4.5-flash | **FREE** | **FREE** | 128K | +| Model | Input ($/1M tokens) | Output ($/1M tokens) | Cached Input ($/1M tokens) | Context Window | +|-------|---------------------|----------------------|---------------------------|----------------| +| glm-4.7 | $0.60 | $2.20 | $0.11 | 200K | +| glm-4.6 | $0.60 | $2.20 | - | 200K | +| glm-4.5 | $0.60 | $2.20 | - | 128K | +| glm-4.5v | $0.60 | $1.80 | - | 128K | +| glm-4.5-x | $2.20 | $8.90 | - | 128K | +| glm-4.5-air | $0.20 | $1.10 | - | 128K | +| glm-4.5-airx | $1.10 | $4.50 | - | 128K | +| glm-4-32b-0414-128k | $0.10 | $0.10 | - | 128K | +| glm-4.5-flash | **FREE** | **FREE** | - | 128K | ## Using with LiteLLM Proxy @@ -84,7 +86,7 @@ import os os.environ['ZAI_API_KEY'] = "" response = completion( - model="zai/glm-4.6", + model="zai/glm-4.7", messages=[{"role": "user", "content": "Hello, how are you?"}], ) @@ -98,9 +100,9 @@ print(response.choices[0].message.content) ```yaml model_list: - - model_name: glm-4.6 + - model_name: glm-4.7 litellm_params: - model: zai/glm-4.6 + model: zai/glm-4.7 api_key: os.environ/ZAI_API_KEY - model_name: glm-4.5-flash # Free tier litellm_params: @@ -121,7 +123,7 @@ curl -L -X POST 'http://0.0.0.0:4000/v1/chat/completions' \ -H 'Content-Type: application/json' \ -H 'Authorization: Bearer sk-1234' \ -d '{ - "model": "glm-4.6", + "model": "glm-4.7", "messages": [ { "role": "user", diff --git a/litellm/llms/zai/chat/transformation.py b/litellm/llms/zai/chat/transformation.py index 47b314d4e0d..4380256f0a4 100644 --- a/litellm/llms/zai/chat/transformation.py +++ b/litellm/llms/zai/chat/transformation.py @@ -20,7 +20,7 @@ class ZAIChatConfig(OpenAIGPTConfig): return api_base, dynamic_api_key def get_supported_openai_params(self, model: str) -> list: - return [ + base_params = [ "max_tokens", "stream", "stream_options", @@ -31,3 +31,12 @@ class ZAIChatConfig(OpenAIGPTConfig): "tool_choice", ] + import litellm + + try: + if litellm.supports_reasoning(model=model, custom_llm_provider=self.custom_llm_provider): + base_params.append("thinking") + except Exception: + pass + + return base_params diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index c585dac9063..d32adf54b5e 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -29649,6 +29649,20 @@ "supports_vision": true, "supports_web_search": true }, + "zai/glm-4.7": { + "cache_creation_input_token_cost": 0, + "cache_read_input_token_cost": 1.1e-07, + "input_cost_per_token": 6e-07, + "output_cost_per_token": 2.2e-06, + "litellm_provider": "zai", + "max_input_tokens": 200000, + "max_output_tokens": 128000, + "mode": "chat", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "source": "https://docs.z.ai/guides/overview/pricing" + }, "zai/glm-4.6": { "input_cost_per_token": 6e-07, "output_cost_per_token": 2.2e-06, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index df286e6540a..81b4469f24c 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -29691,6 +29691,20 @@ "supports_vision": true, "supports_web_search": true }, + "zai/glm-4.7": { + "cache_creation_input_token_cost": 0, + "cache_read_input_token_cost": 1.1e-07, + "input_cost_per_token": 6e-07, + "output_cost_per_token": 2.2e-06, + "litellm_provider": "zai", + "max_input_tokens": 200000, + "max_output_tokens": 128000, + "mode": "chat", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "source": "https://docs.z.ai/guides/overview/pricing" + }, "zai/glm-4.6": { "input_cost_per_token": 6e-07, "output_cost_per_token": 2.2e-06, diff --git a/tests/test_litellm/llms/zai/test_zai_provider.py b/tests/test_litellm/llms/zai/test_zai_provider.py index a3d47d666bc..d1e4359d048 100644 --- a/tests/test_litellm/llms/zai/test_zai_provider.py +++ b/tests/test_litellm/llms/zai/test_zai_provider.py @@ -1,6 +1,7 @@ """ Tests for Z.AI (Zhipu AI) provider - GLM models """ + import json import math @@ -50,10 +51,12 @@ def test_zai_in_provider_lists(): def test_zai_models_in_model_cost(): """Test that ZAI models are in the model cost map""" import os + os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" litellm.model_cost = litellm.get_model_cost_map(url="") zai_models = [ + "zai/glm-4.7", "zai/glm-4.6", "zai/glm-4.5", "zai/glm-4.5v", @@ -72,6 +75,7 @@ def test_zai_models_in_model_cost(): def test_zai_glm46_cost_calculation(): """Test the cost calculation for glm-4.6""" import os + os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" litellm.model_cost = litellm.get_model_cost_map(url="") @@ -92,6 +96,7 @@ def test_zai_glm46_cost_calculation(): def test_zai_flash_model_is_free(): """Test that glm-4.5-flash has zero cost""" import os + os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" litellm.model_cost = litellm.get_model_cost_map(url="") @@ -102,6 +107,38 @@ def test_zai_flash_model_is_free(): assert info["output_cost_per_token"] == 0 +def test_glm47_supports_reasoning(): + """Test that GLM-4.7 supports reasoning""" + import os + + os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" + litellm.model_cost = litellm.get_model_cost_map(url="") + + key = "zai/glm-4.7" + assert key in litellm.model_cost, f"Model {key} not found in model_cost" + + info = litellm.model_cost[key] + assert info["supports_reasoning"] is True + + +def test_glm47_cost_calculation(): + """Test cost calculation for GLM-4.7""" + import os + + os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" + litellm.model_cost = litellm.get_model_cost_map(url="") + + prompt_cost, completion_cost = cost_per_token( + model="zai/glm-4.7", + prompt_tokens=1000000, # 1M tokens + completion_tokens=1000000, + ) + + # GLM-4.7: $0.6/M input, $2.2/M output (same as GLM-4.6) + assert math.isclose(prompt_cost, 0.6, rel_tol=1e-6) + assert math.isclose(completion_cost, 2.2, rel_tol=1e-6) + + @pytest.mark.asyncio async def test_zai_completion_call(respx_mock, zai_response, monkeypatch): """Test completion call with zai provider using mocked response""" From 89b4a6d67c2e603e294bcf2f35bf34bd7e9ba2ae Mon Sep 17 00:00:00 2001 From: Anders Kaseorg Date: Sat, 3 Jan 2026 11:15:15 -0800 Subject: [PATCH 138/158] Allow installation with current grpcio on old Python (#18473) Instead of limiting grpcio < 1.68.0, specifically exclude the versions affected by the reconnect bug, and allow installation with either older or newer versions. Signed-off-by: Anders Kaseorg --- poetry.lock | 72 +----------------------------------------------- pyproject.toml | 2 +- requirements.txt | 2 +- 3 files changed, 3 insertions(+), 73 deletions(-) diff --git a/poetry.lock b/poetry.lock index ee97c00594c..a0a0f8540e5 100644 --- a/poetry.lock +++ b/poetry.lock @@ -2273,75 +2273,6 @@ googleapis-common-protos = {version = ">=1.56.0,<2.0.0", extras = ["grpc"]} grpcio = ">=1.44.0,<2.0.0" protobuf = ">=3.20.2,<4.21.1 || >4.21.1,<4.21.2 || >4.21.2,<4.21.3 || >4.21.3,<4.21.4 || >4.21.4,<4.21.5 || >4.21.5,<7.0.0" -[[package]] -name = "grpcio" -version = "1.67.1" -description = "HTTP/2-based RPC framework" -optional = false -python-versions = ">=3.8" -groups = ["main", "dev", "proxy-dev"] -markers = "python_version < \"3.14\"" -files = [ - {file = "grpcio-1.67.1-cp310-cp310-linux_armv7l.whl", hash = "sha256:8b0341d66a57f8a3119b77ab32207072be60c9bf79760fa609c5609f2deb1f3f"}, - {file = "grpcio-1.67.1-cp310-cp310-macosx_12_0_universal2.whl", hash = "sha256:f5a27dddefe0e2357d3e617b9079b4bfdc91341a91565111a21ed6ebbc51b22d"}, - {file = "grpcio-1.67.1-cp310-cp310-manylinux_2_17_aarch64.whl", hash = "sha256:43112046864317498a33bdc4797ae6a268c36345a910de9b9c17159d8346602f"}, - {file = "grpcio-1.67.1-cp310-cp310-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:c9b929f13677b10f63124c1a410994a401cdd85214ad83ab67cc077fc7e480f0"}, - {file = "grpcio-1.67.1-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:e7d1797a8a3845437d327145959a2c0c47c05947c9eef5ff1a4c80e499dcc6fa"}, - {file = "grpcio-1.67.1-cp310-cp310-musllinux_1_1_i686.whl", hash = "sha256:0489063974d1452436139501bf6b180f63d4977223ee87488fe36858c5725292"}, - {file = "grpcio-1.67.1-cp310-cp310-musllinux_1_1_x86_64.whl", hash = "sha256:9fd042de4a82e3e7aca44008ee2fb5da01b3e5adb316348c21980f7f58adc311"}, - {file = "grpcio-1.67.1-cp310-cp310-win32.whl", hash = "sha256:638354e698fd0c6c76b04540a850bf1db27b4d2515a19fcd5cf645c48d3eb1ed"}, - {file = "grpcio-1.67.1-cp310-cp310-win_amd64.whl", hash = "sha256:608d87d1bdabf9e2868b12338cd38a79969eaf920c89d698ead08f48de9c0f9e"}, - {file = "grpcio-1.67.1-cp311-cp311-linux_armv7l.whl", hash = "sha256:7818c0454027ae3384235a65210bbf5464bd715450e30a3d40385453a85a70cb"}, - {file = "grpcio-1.67.1-cp311-cp311-macosx_10_9_universal2.whl", hash = "sha256:ea33986b70f83844cd00814cee4451055cd8cab36f00ac64a31f5bb09b31919e"}, - {file = "grpcio-1.67.1-cp311-cp311-manylinux_2_17_aarch64.whl", hash = "sha256:c7a01337407dd89005527623a4a72c5c8e2894d22bead0895306b23c6695698f"}, - {file = "grpcio-1.67.1-cp311-cp311-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:80b866f73224b0634f4312a4674c1be21b2b4afa73cb20953cbbb73a6b36c3cc"}, - {file = "grpcio-1.67.1-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:f9fff78ba10d4250bfc07a01bd6254a6d87dc67f9627adece85c0b2ed754fa96"}, - {file = "grpcio-1.67.1-cp311-cp311-musllinux_1_1_i686.whl", hash = "sha256:8a23cbcc5bb11ea7dc6163078be36c065db68d915c24f5faa4f872c573bb400f"}, - {file = "grpcio-1.67.1-cp311-cp311-musllinux_1_1_x86_64.whl", hash = "sha256:1a65b503d008f066e994f34f456e0647e5ceb34cfcec5ad180b1b44020ad4970"}, - {file = "grpcio-1.67.1-cp311-cp311-win32.whl", hash = "sha256:e29ca27bec8e163dca0c98084040edec3bc49afd10f18b412f483cc68c712744"}, - {file = "grpcio-1.67.1-cp311-cp311-win_amd64.whl", hash = "sha256:786a5b18544622bfb1e25cc08402bd44ea83edfb04b93798d85dca4d1a0b5be5"}, - {file = "grpcio-1.67.1-cp312-cp312-linux_armv7l.whl", hash = "sha256:267d1745894200e4c604958da5f856da6293f063327cb049a51fe67348e4f953"}, - {file = "grpcio-1.67.1-cp312-cp312-macosx_10_9_universal2.whl", hash = "sha256:85f69fdc1d28ce7cff8de3f9c67db2b0ca9ba4449644488c1e0303c146135ddb"}, - {file = "grpcio-1.67.1-cp312-cp312-manylinux_2_17_aarch64.whl", hash = "sha256:f26b0b547eb8d00e195274cdfc63ce64c8fc2d3e2d00b12bf468ece41a0423a0"}, - {file = "grpcio-1.67.1-cp312-cp312-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:4422581cdc628f77302270ff839a44f4c24fdc57887dc2a45b7e53d8fc2376af"}, - {file = "grpcio-1.67.1-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:1d7616d2ded471231c701489190379e0c311ee0a6c756f3c03e6a62b95a7146e"}, - {file = "grpcio-1.67.1-cp312-cp312-musllinux_1_1_i686.whl", hash = "sha256:8a00efecde9d6fcc3ab00c13f816313c040a28450e5e25739c24f432fc6d3c75"}, - {file = "grpcio-1.67.1-cp312-cp312-musllinux_1_1_x86_64.whl", hash = "sha256:699e964923b70f3101393710793289e42845791ea07565654ada0969522d0a38"}, - {file = "grpcio-1.67.1-cp312-cp312-win32.whl", hash = "sha256:4e7b904484a634a0fff132958dabdb10d63e0927398273917da3ee103e8d1f78"}, - {file = "grpcio-1.67.1-cp312-cp312-win_amd64.whl", hash = "sha256:5721e66a594a6c4204458004852719b38f3d5522082be9061d6510b455c90afc"}, - {file = "grpcio-1.67.1-cp313-cp313-linux_armv7l.whl", hash = "sha256:aa0162e56fd10a5547fac8774c4899fc3e18c1aa4a4759d0ce2cd00d3696ea6b"}, - {file = "grpcio-1.67.1-cp313-cp313-macosx_10_13_universal2.whl", hash = "sha256:beee96c8c0b1a75d556fe57b92b58b4347c77a65781ee2ac749d550f2a365dc1"}, - {file = "grpcio-1.67.1-cp313-cp313-manylinux_2_17_aarch64.whl", hash = "sha256:a93deda571a1bf94ec1f6fcda2872dad3ae538700d94dc283c672a3b508ba3af"}, - {file = "grpcio-1.67.1-cp313-cp313-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:0e6f255980afef598a9e64a24efce87b625e3e3c80a45162d111a461a9f92955"}, - {file = "grpcio-1.67.1-cp313-cp313-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:9e838cad2176ebd5d4a8bb03955138d6589ce9e2ce5d51c3ada34396dbd2dba8"}, - {file = "grpcio-1.67.1-cp313-cp313-musllinux_1_1_i686.whl", hash = "sha256:a6703916c43b1d468d0756c8077b12017a9fcb6a1ef13faf49e67d20d7ebda62"}, - {file = "grpcio-1.67.1-cp313-cp313-musllinux_1_1_x86_64.whl", hash = "sha256:917e8d8994eed1d86b907ba2a61b9f0aef27a2155bca6cbb322430fc7135b7bb"}, - {file = "grpcio-1.67.1-cp313-cp313-win32.whl", hash = "sha256:e279330bef1744040db8fc432becc8a727b84f456ab62b744d3fdb83f327e121"}, - {file = "grpcio-1.67.1-cp313-cp313-win_amd64.whl", hash = "sha256:fa0c739ad8b1996bd24823950e3cb5152ae91fca1c09cc791190bf1627ffefba"}, - {file = "grpcio-1.67.1-cp38-cp38-linux_armv7l.whl", hash = "sha256:178f5db771c4f9a9facb2ab37a434c46cb9be1a75e820f187ee3d1e7805c4f65"}, - {file = "grpcio-1.67.1-cp38-cp38-macosx_10_9_universal2.whl", hash = "sha256:0f3e49c738396e93b7ba9016e153eb09e0778e776df6090c1b8c91877cc1c426"}, - {file = "grpcio-1.67.1-cp38-cp38-manylinux_2_17_aarch64.whl", hash = "sha256:24e8a26dbfc5274d7474c27759b54486b8de23c709d76695237515bc8b5baeab"}, - {file = "grpcio-1.67.1-cp38-cp38-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:3b6c16489326d79ead41689c4b84bc40d522c9a7617219f4ad94bc7f448c5085"}, - {file = "grpcio-1.67.1-cp38-cp38-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:60e6a4dcf5af7bbc36fd9f81c9f372e8ae580870a9e4b6eafe948cd334b81cf3"}, - {file = "grpcio-1.67.1-cp38-cp38-musllinux_1_1_i686.whl", hash = "sha256:95b5f2b857856ed78d72da93cd7d09b6db8ef30102e5e7fe0961fe4d9f7d48e8"}, - {file = "grpcio-1.67.1-cp38-cp38-musllinux_1_1_x86_64.whl", hash = "sha256:b49359977c6ec9f5d0573ea4e0071ad278ef905aa74e420acc73fd28ce39e9ce"}, - {file = "grpcio-1.67.1-cp38-cp38-win32.whl", hash = "sha256:f5b76ff64aaac53fede0cc93abf57894ab2a7362986ba22243d06218b93efe46"}, - {file = "grpcio-1.67.1-cp38-cp38-win_amd64.whl", hash = "sha256:804c6457c3cd3ec04fe6006c739579b8d35c86ae3298ffca8de57b493524b771"}, - {file = "grpcio-1.67.1-cp39-cp39-linux_armv7l.whl", hash = "sha256:a25bdea92b13ff4d7790962190bf6bf5c4639876e01c0f3dda70fc2769616335"}, - {file = "grpcio-1.67.1-cp39-cp39-macosx_10_9_universal2.whl", hash = "sha256:cdc491ae35a13535fd9196acb5afe1af37c8237df2e54427be3eecda3653127e"}, - {file = "grpcio-1.67.1-cp39-cp39-manylinux_2_17_aarch64.whl", hash = "sha256:85f862069b86a305497e74d0dc43c02de3d1d184fc2c180993aa8aa86fbd19b8"}, - {file = "grpcio-1.67.1-cp39-cp39-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:ec74ef02010186185de82cc594058a3ccd8d86821842bbac9873fd4a2cf8be8d"}, - {file = "grpcio-1.67.1-cp39-cp39-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:01f616a964e540638af5130469451cf580ba8c7329f45ca998ab66e0c7dcdb04"}, - {file = "grpcio-1.67.1-cp39-cp39-musllinux_1_1_i686.whl", hash = "sha256:299b3d8c4f790c6bcca485f9963b4846dd92cf6f1b65d3697145d005c80f9fe8"}, - {file = "grpcio-1.67.1-cp39-cp39-musllinux_1_1_x86_64.whl", hash = "sha256:60336bff760fbb47d7e86165408126f1dded184448e9a4c892189eb7c9d3f90f"}, - {file = "grpcio-1.67.1-cp39-cp39-win32.whl", hash = "sha256:5ed601c4c6008429e3d247ddb367fe8c7259c355757448d7c1ef7bd4a6739e8e"}, - {file = "grpcio-1.67.1-cp39-cp39-win_amd64.whl", hash = "sha256:5db70d32d6703b89912af16d6d45d78406374a8b8ef0d28140351dd0ec610e98"}, - {file = "grpcio-1.67.1.tar.gz", hash = "sha256:3dc2ed4cabea4dc14d5e708c2b426205956077cc5de419b4d4079315017e9732"}, -] - -[package.extras] -protobuf = ["grpcio-tools (>=1.67.1)"] - [[package]] name = "grpcio" version = "1.76.0" @@ -2349,7 +2280,6 @@ description = "HTTP/2-based RPC framework" optional = false python-versions = ">=3.9" groups = ["main", "dev", "proxy-dev"] -markers = "python_version >= \"3.14\"" files = [ {file = "grpcio-1.76.0-cp310-cp310-linux_armv7l.whl", hash = "sha256:65a20de41e85648e00305c1bb09a3598f840422e522277641145a32d42dcefcc"}, {file = "grpcio-1.76.0-cp310-cp310-macosx_11_0_universal2.whl", hash = "sha256:40ad3afe81676fd9ec6d9d406eda00933f218038433980aa19d401490e46ecde"}, @@ -8051,4 +7981,4 @@ utils = ["numpydoc"] [metadata] lock-version = "2.1" python-versions = ">=3.9,<4.0" -content-hash = "b010d9da7f5a765670932b78d720aae4fcb819daba050683ee125b4367972419" +content-hash = "7eed2b2c25173a275ac83c55fd901b9b84663b1d7daa54f0e78b30bf1c8f0e3e" diff --git a/pyproject.toml b/pyproject.toml index f929fb94cb0..3b09119a748 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -72,7 +72,7 @@ soundfile = {version = "^0.12.1", optional = true} # - 1.68.0-1.68.1 has reconnect bug (https://github.com/grpc/grpc/issues/38290) # - 1.75.0+ has Python 3.14 wheels and bug fix grpcio = [ - {version = ">=1.62.3,<1.68.0", python = "<3.14"}, + {version = ">=1.62.3,!=1.68.*,!=1.69.*,!=1.70.*,!=1.71.0,!=1.71.1,!=1.72.0,!=1.72.1,!=1.73.0", python = "<3.14"}, {version = ">=1.75.0", python = ">=3.14"}, ] diff --git a/requirements.txt b/requirements.txt index 3bc968c8cb8..06a7c17336c 100644 --- a/requirements.txt +++ b/requirements.txt @@ -41,7 +41,7 @@ opentelemetry-api==1.25.0 opentelemetry-sdk==1.25.0 opentelemetry-exporter-otlp==1.25.0 # grpcio: 1.68.0-1.68.1 has reconnect bug (#38290), 1.75+ has Python 3.14 wheels + fix -grpcio>=1.62.3,<1.68.0; python_version < "3.14" +grpcio>=1.62.3,!=1.68.*,!=1.69.*,!=1.70.*,!=1.71.0,!=1.71.1,!=1.72.0,!=1.72.1,!=1.73.0; python_version < "3.14" grpcio>=1.75.0; python_version >= "3.14" sentry_sdk==2.21.0 # for sentry error handling detect-secrets==1.5.0 # Enterprise - secret detection / masking in LLM requests From 099e108b51df00d2c8e855b459264973b1bb03b0 Mon Sep 17 00:00:00 2001 From: lif <1835304752@qq.com> Date: Sun, 4 Jan 2026 03:15:45 +0800 Subject: [PATCH 139/158] fix: correctly route codestral chat and FIM endpoints (#18467) Fixed duplicate condition that made text-completion-codestral provider unreachable. Now: - codestral.mistral.ai/v1/chat/completions -> codestral - codestral.mistral.ai/v1/fim/completions -> text-completion-codestral Fixes #18464 Signed-off-by: majiayu000 <1835304752@qq.com> --- .../get_llm_provider_logic.py | 4 +- .../test_codestral_provider_routing.py | 69 +++++++++++++++++++ 2 files changed, 71 insertions(+), 2 deletions(-) create mode 100644 tests/test_litellm/litellm_core_utils/test_codestral_provider_routing.py diff --git a/litellm/litellm_core_utils/get_llm_provider_logic.py b/litellm/litellm_core_utils/get_llm_provider_logic.py index 164e2a73e65..b753e9fa8b5 100644 --- a/litellm/litellm_core_utils/get_llm_provider_logic.py +++ b/litellm/litellm_core_utils/get_llm_provider_logic.py @@ -229,10 +229,10 @@ def get_llm_provider( # noqa: PLR0915 elif endpoint == "https://api.ai21.com/studio/v1": custom_llm_provider = "ai21_chat" dynamic_api_key = get_secret_str("AI21_API_KEY") - elif endpoint == "https://codestral.mistral.ai/v1": + elif endpoint == "codestral.mistral.ai/v1/chat/completions": custom_llm_provider = "codestral" dynamic_api_key = get_secret_str("CODESTRAL_API_KEY") - elif endpoint == "https://codestral.mistral.ai/v1": + elif endpoint == "codestral.mistral.ai/v1/fim/completions": custom_llm_provider = "text-completion-codestral" dynamic_api_key = get_secret_str("CODESTRAL_API_KEY") elif endpoint == "app.empower.dev/api/v1": diff --git a/tests/test_litellm/litellm_core_utils/test_codestral_provider_routing.py b/tests/test_litellm/litellm_core_utils/test_codestral_provider_routing.py new file mode 100644 index 00000000000..1a6ed51afd0 --- /dev/null +++ b/tests/test_litellm/litellm_core_utils/test_codestral_provider_routing.py @@ -0,0 +1,69 @@ +""" +Unit tests for codestral provider routing. + +These tests verify that the chat and FIM endpoints for codestral +are correctly routed to different providers: +- Chat endpoint -> codestral provider +- FIM endpoint -> text-completion-codestral provider + +Related issue: https://github.com/BerriAI/litellm/issues/18464 +""" +import pytest + +import litellm + + +class TestCodestralProviderRouting: + """Tests for codestral endpoint routing in get_llm_provider""" + + def test_codestral_chat_endpoint_routes_to_codestral_provider(self): + """ + Test that the codestral chat endpoint routes to the 'codestral' provider. + + The chat/completions endpoint should be handled by the codestral provider. + """ + model, custom_llm_provider, _, api_base = litellm.get_llm_provider( + model="codestral-latest", + api_base="https://codestral.mistral.ai/v1/chat/completions", + ) + + assert custom_llm_provider == "codestral" + + def test_codestral_fim_endpoint_routes_to_text_completion_provider(self): + """ + Test that the codestral FIM endpoint routes to 'text-completion-codestral'. + + The fim/completions endpoint should be handled by the + text-completion-codestral provider for fill-in-the-middle completions. + """ + model, custom_llm_provider, _, api_base = litellm.get_llm_provider( + model="codestral-latest", + api_base="https://codestral.mistral.ai/v1/fim/completions", + ) + + assert custom_llm_provider == "text-completion-codestral" + + def test_codestral_endpoints_are_different_providers(self): + """ + Test that chat and FIM endpoints route to different providers. + + This is the core fix for issue #18464 - previously both endpoints + would route to 'codestral' due to duplicate conditions. + """ + _, chat_provider, _, _ = litellm.get_llm_provider( + model="codestral-latest", + api_base="https://codestral.mistral.ai/v1/chat/completions", + ) + + _, fim_provider, _, _ = litellm.get_llm_provider( + model="codestral-latest", + api_base="https://codestral.mistral.ai/v1/fim/completions", + ) + + assert chat_provider != fim_provider + assert chat_provider == "codestral" + assert fim_provider == "text-completion-codestral" + + +if __name__ == "__main__": + pytest.main([__file__, "-v"]) From 64cfe75bfd398aa856c117aedfb6b72c70f48d16 Mon Sep 17 00:00:00 2001 From: lif <1835304752@qq.com> Date: Sun, 4 Jan 2026 03:17:38 +0800 Subject: [PATCH 140/158] fix: extract pure base64 data from data URLs for Ollama (#18465) Fix Ollama_chatException "illegal base64 data at input byte 4" error when using images with ollama_chat provider. Ollama expects pure base64 data, not the full data URL format (data:image/png;base64,...). Fixes #18338 Signed-off-by: majiayu000 <1835304752@qq.com> --- .../prompt_templates/common_utils.py | 32 +++- .../test_extract_base64_image.py | 156 ++++++++++++++++++ 2 files changed, 185 insertions(+), 3 deletions(-) create mode 100644 tests/test_litellm/litellm_core_utils/test_extract_base64_image.py diff --git a/litellm/litellm_core_utils/prompt_templates/common_utils.py b/litellm/litellm_core_utils/prompt_templates/common_utils.py index ca2a092dbc8..b100b9b516b 100644 --- a/litellm/litellm_core_utils/prompt_templates/common_utils.py +++ b/litellm/litellm_core_utils/prompt_templates/common_utils.py @@ -1087,9 +1087,35 @@ def _parse_content_for_reasoning( return None, message_text +def _extract_base64_data(image_url: str) -> str: + """ + Extract pure base64 data from an image URL. + + If the URL is a data URL (e.g., "data:image/png;base64,iVBOR..."), + extract and return only the base64 data portion. + Otherwise, return the original URL unchanged. + + This is needed for providers like Ollama that expect pure base64 data + rather than full data URLs. + + Args: + image_url: The image URL or data URL to process + + Returns: + The base64 data if it's a data URL, otherwise the original URL + """ + if image_url.startswith("data:") and ";base64," in image_url: + return image_url.split(";base64,", 1)[1] + return image_url + + def extract_images_from_message(message: AllMessageValues) -> List[str]: """ - Extract images from a message + Extract images from a message. + + For data URLs (e.g., "data:image/png;base64,iVBOR..."), only the base64 + data portion is extracted. This is required for providers like Ollama + that expect pure base64 data rather than full data URLs. """ images = [] message_content = message.get("content") @@ -1098,7 +1124,7 @@ def extract_images_from_message(message: AllMessageValues) -> List[str]: image_url = m.get("image_url") if image_url: if isinstance(image_url, str): - images.append(image_url) + images.append(_extract_base64_data(image_url)) elif isinstance(image_url, dict) and "url" in image_url: - images.append(image_url["url"]) + images.append(_extract_base64_data(image_url["url"])) return images diff --git a/tests/test_litellm/litellm_core_utils/test_extract_base64_image.py b/tests/test_litellm/litellm_core_utils/test_extract_base64_image.py new file mode 100644 index 00000000000..b17c02d7006 --- /dev/null +++ b/tests/test_litellm/litellm_core_utils/test_extract_base64_image.py @@ -0,0 +1,156 @@ +""" +Unit tests for _extract_base64_data and extract_images_from_message functions. + +These tests verify that base64 image data is correctly extracted from data URLs, +which fixes the Ollama error "illegal base64 data at input byte 4". + +Related issue: https://github.com/BerriAI/litellm/issues/18338 +""" +import pytest + +from litellm.litellm_core_utils.prompt_templates.common_utils import ( + _extract_base64_data, + extract_images_from_message, +) + + +class TestExtractBase64Data: + """Tests for _extract_base64_data function""" + + def test_extract_base64_from_png_data_url(self): + """Test extracting base64 data from a PNG data URL""" + data_url = "data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNk" + expected = "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNk" + assert _extract_base64_data(data_url) == expected + + def test_extract_base64_from_jpeg_data_url(self): + """Test extracting base64 data from a JPEG data URL""" + data_url = "data:image/jpeg;base64,/9j/4AAQSkZJRgABAQAAAQABAAD" + expected = "/9j/4AAQSkZJRgABAQAAAQABAAD" + assert _extract_base64_data(data_url) == expected + + def test_extract_base64_from_gif_data_url(self): + """Test extracting base64 data from a GIF data URL""" + data_url = "data:image/gif;base64,R0lGODlhAQABAIAAAAAAAP" + expected = "R0lGODlhAQABAIAAAAAAAP" + assert _extract_base64_data(data_url) == expected + + def test_regular_url_unchanged(self): + """Test that regular HTTP URLs are returned unchanged""" + url = "https://example.com/image.png" + assert _extract_base64_data(url) == url + + def test_file_path_unchanged(self): + """Test that file paths are returned unchanged""" + path = "/path/to/image.png" + assert _extract_base64_data(path) == path + + def test_data_url_without_base64_unchanged(self): + """Test that data URLs without base64 encoding are returned unchanged""" + # This is a data URL with URL encoding, not base64 + url = "data:text/plain,Hello%20World" + assert _extract_base64_data(url) == url + + def test_base64_data_with_special_chars(self): + """Test extracting base64 data that contains valid special characters""" + # Base64 can contain +, /, and = characters + data_url = "data:image/png;base64,abc+def/ghi===" + expected = "abc+def/ghi===" + assert _extract_base64_data(data_url) == expected + + +class TestExtractImagesFromMessage: + """Tests for extract_images_from_message function""" + + def test_extract_from_message_with_data_url_string(self): + """Test extracting images when image_url is a string data URL""" + message = { + "role": "user", + "content": [ + { + "type": "image_url", + "image_url": "data:image/png;base64,iVBORw0KGgo", + } + ], + } + result = extract_images_from_message(message) + assert result == ["iVBORw0KGgo"] + + def test_extract_from_message_with_data_url_dict(self): + """Test extracting images when image_url is a dict with url key""" + message = { + "role": "user", + "content": [ + { + "type": "image_url", + "image_url": {"url": "data:image/png;base64,iVBORw0KGgo"}, + } + ], + } + result = extract_images_from_message(message) + assert result == ["iVBORw0KGgo"] + + def test_extract_from_message_with_regular_url(self): + """Test that regular URLs are preserved""" + message = { + "role": "user", + "content": [ + { + "type": "image_url", + "image_url": {"url": "https://example.com/image.png"}, + } + ], + } + result = extract_images_from_message(message) + assert result == ["https://example.com/image.png"] + + def test_extract_multiple_images(self): + """Test extracting multiple images from a single message""" + message = { + "role": "user", + "content": [ + { + "type": "image_url", + "image_url": "data:image/png;base64,image1base64", + }, + { + "type": "image_url", + "image_url": {"url": "data:image/jpeg;base64,image2base64"}, + }, + { + "type": "image_url", + "image_url": "https://example.com/image3.png", + }, + ], + } + result = extract_images_from_message(message) + assert result == [ + "image1base64", + "image2base64", + "https://example.com/image3.png", + ] + + def test_empty_content(self): + """Test message with empty content""" + message = {"role": "user", "content": []} + result = extract_images_from_message(message) + assert result == [] + + def test_no_images_in_content(self): + """Test message with content but no images""" + message = { + "role": "user", + "content": [{"type": "text", "text": "Hello world"}], + } + result = extract_images_from_message(message) + assert result == [] + + def test_string_content(self): + """Test message with string content (no images possible)""" + message = {"role": "user", "content": "Hello world"} + result = extract_images_from_message(message) + assert result == [] + + +if __name__ == "__main__": + pytest.main([__file__, "-v"]) From 7d81d245fb2665f279c69da1a2066532b07f3045 Mon Sep 17 00:00:00 2001 From: lif <1835304752@qq.com> Date: Sun, 4 Jan 2026 03:18:09 +0800 Subject: [PATCH 141/158] fix: align prometheus metric names with DEFINED_PROMETHEUS_METRICS (#18463) Fix metric name inconsistency for litellm_remaining_requests_metric and litellm_remaining_tokens_metric. The factory received names without the _metric suffix, causing _is_metric_enabled to fail when users configured these metrics in prometheus_metrics_config. Fixes #18221 Signed-off-by: majiayu000 <1835304752@qq.com> --- litellm/integrations/prometheus.py | 4 +- ...test_prometheus_metric_name_consistency.py | 106 ++++++++++++++++++ 2 files changed, 108 insertions(+), 2 deletions(-) create mode 100644 tests/test_litellm/integrations/test_prometheus_metric_name_consistency.py diff --git a/litellm/integrations/prometheus.py b/litellm/integrations/prometheus.py index 20f1357a1c8..c01f7481277 100644 --- a/litellm/integrations/prometheus.py +++ b/litellm/integrations/prometheus.py @@ -214,7 +214,7 @@ class PrometheusLogger(CustomLogger): # Remaining Rate Limit for model self.litellm_remaining_requests_metric = self._gauge_factory( - "litellm_remaining_requests", + "litellm_remaining_requests_metric", "LLM Deployment Analytics - remaining requests for model, returned from LLM API Provider", labelnames=self.get_labels_for_metric( "litellm_remaining_requests_metric" @@ -222,7 +222,7 @@ class PrometheusLogger(CustomLogger): ) self.litellm_remaining_tokens_metric = self._gauge_factory( - "litellm_remaining_tokens", + "litellm_remaining_tokens_metric", "remaining tokens for model, returned from LLM API Provider", labelnames=self.get_labels_for_metric( "litellm_remaining_tokens_metric" diff --git a/tests/test_litellm/integrations/test_prometheus_metric_name_consistency.py b/tests/test_litellm/integrations/test_prometheus_metric_name_consistency.py new file mode 100644 index 00000000000..9658eff3cc5 --- /dev/null +++ b/tests/test_litellm/integrations/test_prometheus_metric_name_consistency.py @@ -0,0 +1,106 @@ +""" +Unit tests for prometheus metric name consistency + +This test ensures that the metric names used when creating Prometheus metrics +match the names defined in DEFINED_PROMETHEUS_METRICS, so that metric filtering +configuration works correctly. + +Related issue: https://github.com/BerriAI/litellm/issues/18221 +""" +from typing import get_args + +import pytest + + +def test_remaining_requests_metric_name_in_defined_metrics(): + """ + Test that litellm_remaining_requests_metric is defined in DEFINED_PROMETHEUS_METRICS. + + The metric name should include the _metric suffix to be consistent with the + configuration format users specify in prometheus_metrics_config. + """ + from litellm.types.integrations.prometheus import DEFINED_PROMETHEUS_METRICS + + defined_metrics = get_args(DEFINED_PROMETHEUS_METRICS) + assert ( + "litellm_remaining_requests_metric" in defined_metrics + ), "litellm_remaining_requests_metric should be in DEFINED_PROMETHEUS_METRICS" + + +def test_remaining_tokens_metric_name_in_defined_metrics(): + """ + Test that litellm_remaining_tokens_metric is defined in DEFINED_PROMETHEUS_METRICS. + + The metric name should include the _metric suffix to be consistent with the + configuration format users specify in prometheus_metrics_config. + """ + from litellm.types.integrations.prometheus import DEFINED_PROMETHEUS_METRICS + + defined_metrics = get_args(DEFINED_PROMETHEUS_METRICS) + assert ( + "litellm_remaining_tokens_metric" in defined_metrics + ), "litellm_remaining_tokens_metric should be in DEFINED_PROMETHEUS_METRICS" + + +def test_prometheus_metric_labels_have_remaining_metrics(): + """ + Test that PrometheusMetricLabels has label definitions for remaining metrics. + + This ensures that the labels can be retrieved when creating the metrics. + """ + from litellm.types.integrations.prometheus import PrometheusMetricLabels + + # Test that labels can be retrieved for remaining metrics + remaining_requests_labels = PrometheusMetricLabels.get_labels( + "litellm_remaining_requests_metric" + ) + remaining_tokens_labels = PrometheusMetricLabels.get_labels( + "litellm_remaining_tokens_metric" + ) + + assert isinstance( + remaining_requests_labels, list + ), "Labels for litellm_remaining_requests_metric should be a list" + assert isinstance( + remaining_tokens_labels, list + ), "Labels for litellm_remaining_tokens_metric should be a list" + + # These metrics should have api_provider and api_base labels + assert ( + "api_provider" in remaining_requests_labels + ), "litellm_remaining_requests_metric should have api_provider label" + assert ( + "api_base" in remaining_requests_labels + ), "litellm_remaining_requests_metric should have api_base label" + assert ( + "api_provider" in remaining_tokens_labels + ), "litellm_remaining_tokens_metric should have api_provider label" + assert ( + "api_base" in remaining_tokens_labels + ), "litellm_remaining_tokens_metric should have api_base label" + + +def test_all_defined_metrics_have_consistent_naming(): + """ + Test that all metrics defined in DEFINED_PROMETHEUS_METRICS follow + a consistent naming convention. + + This helps prevent similar inconsistencies in the future. + """ + from litellm.types.integrations.prometheus import DEFINED_PROMETHEUS_METRICS + + defined_metrics = get_args(DEFINED_PROMETHEUS_METRICS) + + for metric_name in defined_metrics: + # All metrics should start with 'litellm_' + assert metric_name.startswith( + "litellm_" + ), f"Metric {metric_name} should start with 'litellm_'" + + +if __name__ == "__main__": + test_remaining_requests_metric_name_in_defined_metrics() + test_remaining_tokens_metric_name_in_defined_metrics() + test_prometheus_metric_labels_have_remaining_metrics() + test_all_defined_metrics_have_consistent_naming() + print("All prometheus metric name consistency tests passed!") From 1452f0150551193bd26227a348f2e0acd6d294fd Mon Sep 17 00:00:00 2001 From: Alexsander Hamir Date: Sat, 3 Jan 2026 12:12:24 -0800 Subject: [PATCH 142/158] refactor: lazy load get_llm_provider and remove_index_from_tool_calls (#18608) --- litellm/_lazy_imports_registry.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/litellm/_lazy_imports_registry.py b/litellm/_lazy_imports_registry.py index 1a54e77998a..93fa8b39af2 100644 --- a/litellm/_lazy_imports_registry.py +++ b/litellm/_lazy_imports_registry.py @@ -32,6 +32,7 @@ UTILS_NAMES = ( "ModelResponse", "ModelResponseStream", "EmbeddingResponse", "ImageResponse", "TranscriptionResponse", "TextCompletionResponse", "get_provider_fields", "ModelResponseListIterator", "get_valid_models", "timeout", + "get_llm_provider", "remove_index_from_tool_calls", ) # Token counter names that support lazy loading via _lazy_import_token_counter @@ -336,6 +337,8 @@ _UTILS_IMPORT_MAP = { "ModelResponseListIterator": (".utils", "ModelResponseListIterator"), "get_valid_models": (".utils", "get_valid_models"), "timeout": (".timeout", "timeout"), + "get_llm_provider": ("litellm.litellm_core_utils.get_llm_provider_logic", "get_llm_provider"), + "remove_index_from_tool_calls": ("litellm.litellm_core_utils.core_helpers", "remove_index_from_tool_calls"), } _COST_CALCULATOR_IMPORT_MAP = { From 4904ed394eeace50bf951c8a0c9b2fd49b77b2c1 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Sat, 3 Jan 2026 12:36:26 -0800 Subject: [PATCH 143/158] Edit path for SSO Settings --- .../hooks/sso/useEditSSOSettings.ts | 38 +++ .../(dashboard)/hooks/sso/useSSOSettings.ts | 9 +- .../Modals/AddSSOSettingsModal.test.tsx | 2 +- .../Modals/AddSSOSettingsModal.tsx | 245 +++------------- .../Modals/BaseSSOSettingsForm.tsx | 249 ++++++++++++++++ .../Modals/EditSSOSettingsModal.tsx | 131 +++++++++ .../SSOSettings/RedactableField.tsx | 38 +++ .../AdminSettings/SSOSettings/SSOSettings.tsx | 84 ++++-- .../AdminSettings/SSOSettings/constants.ts | 15 + .../AdminSettings/SSOSettings/utils.test.ts | 274 ++++++++++++++++++ .../AdminSettings/SSOSettings/utils.ts | 54 ++++ 11 files changed, 894 insertions(+), 245 deletions(-) create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/hooks/sso/useEditSSOSettings.ts create mode 100644 ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/Modals/BaseSSOSettingsForm.tsx create mode 100644 ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/Modals/EditSSOSettingsModal.tsx create mode 100644 ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/RedactableField.tsx create mode 100644 ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/constants.ts create mode 100644 ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/utils.test.ts create mode 100644 ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/utils.ts diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/sso/useEditSSOSettings.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/sso/useEditSSOSettings.ts new file mode 100644 index 00000000000..69e52d0ff25 --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/sso/useEditSSOSettings.ts @@ -0,0 +1,38 @@ +import { useMutation, UseMutationResult } from "@tanstack/react-query"; +import { updateSSOSettings } from "@/components/networking"; +import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; + +export interface EditSSOSettingsParams { + google_client_id?: string | null; + google_client_secret?: string | null; + microsoft_client_id?: string | null; + microsoft_client_secret?: string | null; + microsoft_tenant?: string | null; + generic_client_id?: string | null; + generic_client_secret?: string | null; + generic_authorization_endpoint?: string | null; + generic_token_endpoint?: string | null; + generic_userinfo_endpoint?: string | null; + proxy_base_url?: string | null; + user_email?: string | null; + sso_provider?: string | null; + role_mappings?: any; + [key: string]: any; +} + +export interface EditSSOSettingsResponse { + [key: string]: any; +} + +export const useEditSSOSettings = (): UseMutationResult => { + const { accessToken } = useAuthorized(); + + return useMutation({ + mutationFn: async (params: EditSSOSettingsParams) => { + if (!accessToken) { + throw new Error("Access token is required"); + } + return await updateSSOSettings(accessToken, params); + }, + }); +}; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/sso/useSSOSettings.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/sso/useSSOSettings.ts index 6453b35ff93..3e09c3c2ca8 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/hooks/sso/useSSOSettings.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/sso/useSSOSettings.ts @@ -27,7 +27,14 @@ export interface SSOSettingsValues { proxy_base_url: string | null; user_email: string | null; ui_access_mode: string | null; - role_mappings: string | null; + role_mappings: { + provider: string; + group_claim: string; + default_role: "internal_user" | "internal_user_viewer" | "proxy_admin" | "proxy_admin_viewer"; + roles: { + [key: string]: string[]; + }; + }; } export interface SSOSettingsResponse { diff --git a/ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/Modals/AddSSOSettingsModal.test.tsx b/ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/Modals/AddSSOSettingsModal.test.tsx index 13363a11643..aae28191031 100644 --- a/ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/Modals/AddSSOSettingsModal.test.tsx +++ b/ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/Modals/AddSSOSettingsModal.test.tsx @@ -17,7 +17,7 @@ describe("AddSSOSettingsModal", () => { const onCancel = vi.fn(); const onSuccess = vi.fn(); - render(); + render(); expect(screen.getByText("SSO Provider")).toBeInTheDocument(); expect(screen.getByText("Cancel")).toBeInTheDocument(); diff --git a/ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/Modals/AddSSOSettingsModal.tsx b/ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/Modals/AddSSOSettingsModal.tsx index a4ef2938f60..7af6240b19e 100644 --- a/ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/Modals/AddSSOSettingsModal.tsx +++ b/ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/Modals/AddSSOSettingsModal.tsx @@ -1,145 +1,36 @@ "use client"; import NotificationsManager from "@/components/molecules/notifications_manager"; -import { updateSSOSettings } from "@/components/networking"; import { parseErrorMessage } from "@/components/shared/errorUtils"; -import { TextInput } from "@tremor/react"; -import { Button as Button2, Form, Input, Modal, Select } from "antd"; -import React, { useState } from "react"; +import { Button, Form, Modal, Space } from "antd"; +import React from "react"; +import BaseSSOSettingsForm from "./BaseSSOSettingsForm"; +import { useEditSSOSettings } from "@/app/(dashboard)/hooks/sso/useEditSSOSettings"; +import { processSSOSettingsPayload } from "../utils"; interface AddSSOSettingsModalProps { isVisible: boolean; onCancel: () => void; onSuccess: () => void; - accessToken: string | null; } -const ssoProviderLogoMap: Record = { - google: "https://artificialanalysis.ai/img/logos/google_small.svg", - microsoft: "https://upload.wikimedia.org/wikipedia/commons/a/a8/Microsoft_Azure_Logo.svg", - okta: "https://www.okta.com/sites/default/files/Okta_Logo_BrightBlue_Medium.png", - generic: "", -}; - -// Define the SSO provider configuration type -interface SSOProviderConfig { - envVarMap: Record; - fields: Array<{ - label: string; - name: string; - placeholder?: string; - }>; -} - -// Define configurations for each SSO provider -const ssoProviderConfigs: Record = { - google: { - envVarMap: { - google_client_id: "GOOGLE_CLIENT_ID", - google_client_secret: "GOOGLE_CLIENT_SECRET", - }, - fields: [ - { label: "Google Client ID", name: "google_client_id" }, - { label: "Google Client Secret", name: "google_client_secret" }, - ], - }, - microsoft: { - envVarMap: { - microsoft_client_id: "MICROSOFT_CLIENT_ID", - microsoft_client_secret: "MICROSOFT_CLIENT_SECRET", - microsoft_tenant: "MICROSOFT_TENANT", - }, - fields: [ - { label: "Microsoft Client ID", name: "microsoft_client_id" }, - { label: "Microsoft Client Secret", name: "microsoft_client_secret" }, - { label: "Microsoft Tenant", name: "microsoft_tenant" }, - ], - }, - okta: { - envVarMap: { - generic_client_id: "GENERIC_CLIENT_ID", - generic_client_secret: "GENERIC_CLIENT_SECRET", - generic_authorization_endpoint: "GENERIC_AUTHORIZATION_ENDPOINT", - generic_token_endpoint: "GENERIC_TOKEN_ENDPOINT", - generic_userinfo_endpoint: "GENERIC_USERINFO_ENDPOINT", - }, - fields: [ - { label: "Generic Client ID", name: "generic_client_id" }, - { label: "Generic Client Secret", name: "generic_client_secret" }, - { - label: "Authorization Endpoint", - name: "generic_authorization_endpoint", - placeholder: "https://your-domain/authorize", - }, - { label: "Token Endpoint", name: "generic_token_endpoint", placeholder: "https://your-domain/token" }, - { - label: "Userinfo Endpoint", - name: "generic_userinfo_endpoint", - placeholder: "https://your-domain/userinfo", - }, - ], - }, - generic: { - envVarMap: { - generic_client_id: "GENERIC_CLIENT_ID", - generic_client_secret: "GENERIC_CLIENT_SECRET", - generic_authorization_endpoint: "GENERIC_AUTHORIZATION_ENDPOINT", - generic_token_endpoint: "GENERIC_TOKEN_ENDPOINT", - generic_userinfo_endpoint: "GENERIC_USERINFO_ENDPOINT", - }, - fields: [ - { label: "Generic Client ID", name: "generic_client_id" }, - { label: "Generic Client Secret", name: "generic_client_secret" }, - { label: "Authorization Endpoint", name: "generic_authorization_endpoint" }, - { label: "Token Endpoint", name: "generic_token_endpoint" }, - { label: "Userinfo Endpoint", name: "generic_userinfo_endpoint" }, - ], - }, -}; - -const AddSSOSettingsModal: React.FC = ({ isVisible, onCancel, onSuccess, accessToken }) => { +const AddSSOSettingsModal: React.FC = ({ isVisible, onCancel, onSuccess }) => { const [form] = Form.useForm(); - const [isSubmitting, setIsSubmitting] = useState(false); + const { mutateAsync, isPending } = useEditSSOSettings(); // Enhanced form submission handler const handleFormSubmit = async (formValues: Record) => { - if (!accessToken) { - NotificationsManager.fromBackend("No access token available"); - return; - } + const payload = processSSOSettingsPayload(formValues); - setIsSubmitting(true); - try { - // Save SSO settings using the new API - await updateSSOSettings(accessToken, formValues); - - NotificationsManager.success("SSO settings added successfully"); - - // Reset form and close modal - form.resetFields(); - onSuccess(); - } catch (error: unknown) { - NotificationsManager.fromBackend("Failed to save SSO settings: " + parseErrorMessage(error)); - } finally { - setIsSubmitting(false); - } - }; - - // Helper function to render provider fields - const renderProviderFields = (provider: string) => { - const config = ssoProviderConfigs[provider]; - if (!config) return null; - - return config.fields.map((field) => ( - - {field.name.includes("client") ? : } - - )); + await mutateAsync(payload, { + onSuccess: () => { + NotificationsManager.success("SSO settings added successfully"); + onSuccess(); + }, + onError: (error) => { + NotificationsManager.fromBackend("Failed to save SSO settings: " + parseErrorMessage(error)); + }, + }); }; const handleCancel = () => { @@ -148,91 +39,23 @@ const AddSSOSettingsModal: React.FC = ({ isVisible, on }; return ( - -
- - - - - prevValues.sso_provider !== currentValues.sso_provider} - > - {({ getFieldValue }) => { - const provider = getFieldValue("sso_provider"); - return provider ? renderProviderFields(provider) : null; - }} - - - - - - value?.trim()} - rules={[ - { required: true, message: "Please enter the proxy base url" }, - { - pattern: /^https?:\/\/.+/, - message: "URL must start with http:// or https://", - }, - { - validator: (_, value) => { - // Only check for trailing slash if the URL starts with http:// or https:// - if (value && /^https?:\/\/.+/.test(value) && value.endsWith("/")) { - return Promise.reject("URL must not end with a trailing slash"); - } - return Promise.resolve(); - }, - }, - ]} - > - - - -
- Cancel - - Add SSO - -
-
+ + + + + } + onCancel={handleCancel} + > + ); }; diff --git a/ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/Modals/BaseSSOSettingsForm.tsx b/ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/Modals/BaseSSOSettingsForm.tsx new file mode 100644 index 00000000000..a4b36e5190e --- /dev/null +++ b/ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/Modals/BaseSSOSettingsForm.tsx @@ -0,0 +1,249 @@ +"use client"; + +import { TextInput } from "@tremor/react"; +import { Checkbox, Form, Input, Select } from "antd"; +import React from "react"; +import { ssoProviderLogoMap, ssoProviderDisplayNames } from "../constants"; + +export interface BaseSSOSettingsFormProps { + form: any; // Replace with proper Form type if available + onFormSubmit: (formValues: Record) => Promise; +} + +// Define the SSO provider configuration type +export interface SSOProviderConfig { + envVarMap: Record; + fields: Array<{ + label: string; + name: string; + placeholder?: string; + }>; +} + +// Define configurations for each SSO provider +export const ssoProviderConfigs: Record = { + google: { + envVarMap: { + google_client_id: "GOOGLE_CLIENT_ID", + google_client_secret: "GOOGLE_CLIENT_SECRET", + }, + fields: [ + { label: "Google Client ID", name: "google_client_id" }, + { label: "Google Client Secret", name: "google_client_secret" }, + ], + }, + microsoft: { + envVarMap: { + microsoft_client_id: "MICROSOFT_CLIENT_ID", + microsoft_client_secret: "MICROSOFT_CLIENT_SECRET", + microsoft_tenant: "MICROSOFT_TENANT", + }, + fields: [ + { label: "Microsoft Client ID", name: "microsoft_client_id" }, + { label: "Microsoft Client Secret", name: "microsoft_client_secret" }, + { label: "Microsoft Tenant", name: "microsoft_tenant" }, + ], + }, + okta: { + envVarMap: { + generic_client_id: "GENERIC_CLIENT_ID", + generic_client_secret: "GENERIC_CLIENT_SECRET", + generic_authorization_endpoint: "GENERIC_AUTHORIZATION_ENDPOINT", + generic_token_endpoint: "GENERIC_TOKEN_ENDPOINT", + generic_userinfo_endpoint: "GENERIC_USERINFO_ENDPOINT", + }, + fields: [ + { label: "Generic Client ID", name: "generic_client_id" }, + { label: "Generic Client Secret", name: "generic_client_secret" }, + { + label: "Authorization Endpoint", + name: "generic_authorization_endpoint", + placeholder: "https://your-domain/authorize", + }, + { label: "Token Endpoint", name: "generic_token_endpoint", placeholder: "https://your-domain/token" }, + { + label: "Userinfo Endpoint", + name: "generic_userinfo_endpoint", + placeholder: "https://your-domain/userinfo", + }, + ], + }, + generic: { + envVarMap: { + generic_client_id: "GENERIC_CLIENT_ID", + generic_client_secret: "GENERIC_CLIENT_SECRET", + generic_authorization_endpoint: "GENERIC_AUTHORIZATION_ENDPOINT", + generic_token_endpoint: "GENERIC_TOKEN_ENDPOINT", + generic_userinfo_endpoint: "GENERIC_USERINFO_ENDPOINT", + }, + fields: [ + { label: "Generic Client ID", name: "generic_client_id" }, + { label: "Generic Client Secret", name: "generic_client_secret" }, + { label: "Authorization Endpoint", name: "generic_authorization_endpoint" }, + { label: "Token Endpoint", name: "generic_token_endpoint" }, + { label: "Userinfo Endpoint", name: "generic_userinfo_endpoint" }, + ], + }, +}; + +// Helper function to render provider fields +export const renderProviderFields = (provider: string) => { + const config = ssoProviderConfigs[provider]; + if (!config) return null; + + return config.fields.map((field) => ( + + {field.name.includes("client") ? : } + + )); +}; + +const BaseSSOSettingsForm: React.FC = ({ form, onFormSubmit }) => { + return ( +
+
+ + + + + prevValues.sso_provider !== currentValues.sso_provider} + > + {({ getFieldValue }) => { + const provider = getFieldValue("sso_provider"); + return provider ? renderProviderFields(provider) : null; + }} + + + + + + value?.trim()} + rules={[ + { required: true, message: "Please enter the proxy base url" }, + { + pattern: /^https?:\/\/.+/, + message: "URL must start with http:// or https://", + }, + { + validator: (_, value) => { + // Only check for trailing slash if the URL starts with http:// or https:// + if (value && /^https?:\/\/.+/.test(value) && value.endsWith("/")) { + return Promise.reject("URL must not end with a trailing slash"); + } + return Promise.resolve(); + }, + }, + ]} + > + + + + prevValues.sso_provider !== currentValues.sso_provider} + > + {({ getFieldValue }) => { + const provider = getFieldValue("sso_provider"); + return provider === "okta" || provider === "generic" ? ( + + + + ) : null; + }} + + + prevValues.use_role_mappings !== currentValues.use_role_mappings} + > + {({ getFieldValue }) => { + const useRoleMappings = getFieldValue("use_role_mappings"); + return useRoleMappings ? ( + + + + ) : null; + }} + + + prevValues.use_role_mappings !== currentValues.use_role_mappings} + > + {({ getFieldValue }) => { + const useRoleMappings = getFieldValue("use_role_mappings"); + return useRoleMappings ? ( + <> + + + + + + + + + + + + + + + + + + + + + ) : null; + }} + +
+
+ ); +}; + +export default BaseSSOSettingsForm; diff --git a/ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/Modals/EditSSOSettingsModal.tsx b/ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/Modals/EditSSOSettingsModal.tsx new file mode 100644 index 00000000000..297698a7ba0 --- /dev/null +++ b/ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/Modals/EditSSOSettingsModal.tsx @@ -0,0 +1,131 @@ +"use client"; + +import { Button, Form, Modal, Space } from "antd"; +import React, { useEffect } from "react"; +import BaseSSOSettingsForm from "./BaseSSOSettingsForm"; +import NotificationsManager from "@/components/molecules/notifications_manager"; +import { parseErrorMessage } from "@/components/shared/errorUtils"; +import { processSSOSettingsPayload } from "../utils"; +import { useSSOSettings } from "@/app/(dashboard)/hooks/sso/useSSOSettings"; +import { useEditSSOSettings } from "@/app/(dashboard)/hooks/sso/useEditSSOSettings"; + +interface EditSSOSettingsModalProps { + isVisible: boolean; + onCancel: () => void; + onSuccess: () => void; +} + +const EditSSOSettingsModal: React.FC = ({ isVisible, onCancel, onSuccess }) => { + const [form] = Form.useForm(); + + // Use react-query hooks for SSO settings + const ssoSettings = useSSOSettings(); + const { mutateAsync, isPending } = useEditSSOSettings(); + useEffect(() => { + if (isVisible && ssoSettings.data && ssoSettings.data.values) { + const ssoData = ssoSettings.data; + console.log("Raw SSO data received:", ssoData); // Debug log + console.log("SSO values:", ssoData.values); // Debug log + console.log("user_email from API:", ssoData.values.user_email); // Debug log + + // Determine which SSO provider is configured + let selectedProvider = null; + if (ssoData.values.google_client_id) { + selectedProvider = "google"; + } else if (ssoData.values.microsoft_client_id) { + selectedProvider = "microsoft"; + } else if (ssoData.values.generic_client_id) { + // Check if it looks like Okta based on endpoints + if ( + ssoData.values.generic_authorization_endpoint?.includes("okta") || + ssoData.values.generic_authorization_endpoint?.includes("auth0") + ) { + selectedProvider = "okta"; + } else { + selectedProvider = "generic"; + } + } + + // Extract role mappings if they exist + let roleMappingFields = {}; + if (ssoData.values.role_mappings) { + const roleMappings = ssoData.values.role_mappings; + + // Helper function to join arrays into comma-separated strings + const joinTeams = (teams: string[] | undefined): string => { + if (!teams || teams.length === 0) return ""; + return teams.join(", "); + }; + + roleMappingFields = { + use_role_mappings: true, + group_claim: roleMappings.group_claim, + default_role: roleMappings.default_role || "internal_user", + proxy_admin_teams: joinTeams(roleMappings.roles?.proxy_admin), + admin_viewer_teams: joinTeams(roleMappings.roles?.proxy_admin_viewer), + internal_user_teams: joinTeams(roleMappings.roles?.internal_user), + internal_viewer_teams: joinTeams(roleMappings.roles?.internal_user_viewer), + }; + } + + // Set form values with existing data (excluding UI access control fields) + const formValues = { + sso_provider: selectedProvider, + ...ssoData.values, + ...roleMappingFields, + }; + + console.log("Setting form values:", formValues); // Debug log + + // Clear form first, then set values with a small delay to ensure proper initialization + form.resetFields(); + setTimeout(() => { + form.setFieldsValue(formValues); + console.log("Form values set, current form values:", form.getFieldsValue()); // Debug log + }, 100); + } + }, [isVisible, ssoSettings.data, form]); + + // Enhanced form submission handler + const handleFormSubmit = async (formValues: Record) => { + const payload = processSSOSettingsPayload(formValues); + + await mutateAsync(payload, { + onSuccess: () => { + NotificationsManager.success("SSO settings updated successfully"); + onSuccess(); + }, + onError: (error) => { + NotificationsManager.fromBackend("Failed to save SSO settings: " + parseErrorMessage(error)); + }, + }); + }; + + const handleCancel = () => { + form.resetFields(); + onCancel(); + }; + + return ( + + + + + } + onCancel={handleCancel} + > + + + ); +}; + +export default EditSSOSettingsModal; diff --git a/ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/RedactableField.tsx b/ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/RedactableField.tsx new file mode 100644 index 00000000000..44fef5cc7f8 --- /dev/null +++ b/ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/RedactableField.tsx @@ -0,0 +1,38 @@ +import { useState } from "react"; +import { Button } from "antd"; +import { Eye, EyeOff } from "lucide-react"; + +export default function RedactableField({ + defaultHidden = true, + value, +}: { + defaultHidden?: boolean; + value: string | null; +}) { + const [isHidden, setIsHidden] = useState(defaultHidden); + + return ( +
+ + {value ? ( + isHidden ? ( + "•".repeat(value.length) + ) : ( + value + ) + ) : ( + Not configured + )} + + {value && ( +
+ ); +} diff --git a/ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/SSOSettings.tsx b/ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/SSOSettings.tsx index 0ece207662c..975a3bc7d78 100644 --- a/ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/SSOSettings.tsx +++ b/ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/SSOSettings.tsx @@ -3,11 +3,14 @@ import { useSSOSettings, type SSOSettingsValues } from "@/app/(dashboard)/hooks/sso/useSSOSettings"; import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; import { Badge, Button, Card, Descriptions, Space, Typography } from "antd"; -import { Shield, Trash2 } from "lucide-react"; +import { Shield, Trash2, Edit } from "lucide-react"; import { useState } from "react"; import AddSSOSettingsModal from "./Modals/AddSSOSettingsModal"; import DeleteSSOSettingsModal from "./Modals/DeleteSSOSettingsModal"; +import EditSSOSettingsModal from "./Modals/EditSSOSettingsModal"; import SSOSettingsEmptyPlaceholder from "./SSOSettingsEmptyPlaceholder"; +import RedactableField from "./RedactableField"; +import { ssoProviderLogoMap, ssoProviderDisplayNames } from "./constants"; const { Title, Text } = Typography; @@ -16,6 +19,7 @@ export default function SSOSettings() { const { accessToken } = useAuthorized(); const [isDeleteModalVisible, setIsDeleteModalVisible] = useState(false); const [isAddModalVisible, setIsAddModalVisible] = useState(false); + const [isEditModalVisible, setIsEditModalVisible] = useState(false); const isSSOConfigured = Boolean(ssoSettings?.values.google_client_id) || Boolean(ssoSettings?.values.microsoft_client_id) || @@ -39,12 +43,6 @@ export default function SSOSettings() { } } - const renderRedactedValue = (value?: string | null) => ( - - {value ? "••••••••••••••••••••••••••••••••" : Not configured} - - ); - const renderEndpointValue = (value?: string | null) => ( {value || Not configured} @@ -67,44 +65,44 @@ export default function SSOSettings() { const providerConfigs = { google: { - providerText: "Google OAuth", + providerText: ssoProviderDisplayNames.google, fields: [ { - label: "Client ID (Redacted)", - render: (values: SSOSettingsValues) => renderRedactedValue(values.google_client_id), + label: "Client ID", + render: (values: SSOSettingsValues) => , }, { - label: "Client Secret (Redacted)", - render: (values: SSOSettingsValues) => renderRedactedValue(values.google_client_secret), + label: "Client Secret", + render: (values: SSOSettingsValues) => , }, { label: "Proxy Base URL", render: (values: SSOSettingsValues) => renderSimpleValue(values.proxy_base_url) }, ], }, microsoft: { - providerText: "Microsoft OAuth", + providerText: ssoProviderDisplayNames.microsoft, fields: [ { - label: "Client ID (Redacted)", - render: (values: SSOSettingsValues) => renderRedactedValue(values.microsoft_client_id), + label: "Client ID", + render: (values: SSOSettingsValues) => , }, { - label: "Client Secret (Redacted)", - render: (values: SSOSettingsValues) => renderRedactedValue(values.microsoft_client_secret), + label: "Client Secret", + render: (values: SSOSettingsValues) => , }, { label: "Tenant", render: (values: any) => renderSimpleValue(values.microsoft_tenant) }, { label: "Proxy Base URL", render: (values: SSOSettingsValues) => renderSimpleValue(values.proxy_base_url) }, ], }, okta: { - providerText: "Okta/Auth0", + providerText: ssoProviderDisplayNames.okta, fields: [ { - label: "Client ID (Redacted)", - render: (values: SSOSettingsValues) => renderRedactedValue(values.generic_client_id), + label: "Client ID", + render: (values: SSOSettingsValues) => , }, { - label: "Client Secret (Redacted)", - render: (values: SSOSettingsValues) => renderRedactedValue(values.generic_client_secret), + label: "Client Secret", + render: (values: SSOSettingsValues) => , }, { label: "Authorization Endpoint", @@ -122,15 +120,15 @@ export default function SSOSettings() { ], }, generic: { - providerText: "Generic OAuth", + providerText: ssoProviderDisplayNames.generic, fields: [ { - label: "Client ID (Redacted)", - render: (values: SSOSettingsValues) => renderRedactedValue(values.generic_client_id), + label: "Client ID", + render: (values: SSOSettingsValues) => , }, { - label: "Client Secret (Redacted)", - render: (values: SSOSettingsValues) => renderRedactedValue(values.generic_client_secret), + label: "Client Secret", + render: (values: SSOSettingsValues) => , }, { label: "Authorization Endpoint", @@ -160,7 +158,16 @@ export default function SSOSettings() { return ( - +
+ {ssoProviderLogoMap[selectedProvider] && ( + {selectedProvider} + )} + {config.providerText} +
{config.fields.map((field, index) => ( @@ -186,9 +193,14 @@ export default function SSOSettings() {
{isSSOConfigured && ( - + <> + + + )}
@@ -214,7 +226,15 @@ export default function SSOSettings() { setIsAddModalVisible(false); refetch(); }} - accessToken={accessToken} + /> + + setIsEditModalVisible(false)} + onSuccess={() => { + setIsEditModalVisible(false); + refetch(); + }} /> ); diff --git a/ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/constants.ts b/ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/constants.ts new file mode 100644 index 00000000000..595a961401a --- /dev/null +++ b/ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/constants.ts @@ -0,0 +1,15 @@ +// SSO Provider logos +export const ssoProviderLogoMap: Record = { + google: "https://artificialanalysis.ai/img/logos/google_small.svg", + microsoft: "https://upload.wikimedia.org/wikipedia/commons/a/a8/Microsoft_Azure_Logo.svg", + okta: "https://www.okta.com/sites/default/files/Okta_Logo_BrightBlue_Medium.png", + generic: "", +}; + +// SSO Provider display names (consistent between select dropdown and table) +export const ssoProviderDisplayNames: Record = { + google: "Google SSO", + microsoft: "Microsoft SSO", + okta: "Okta / Auth0 SSO", + generic: "Generic SSO", +}; diff --git a/ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/utils.test.ts b/ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/utils.test.ts new file mode 100644 index 00000000000..1c878d7b7b3 --- /dev/null +++ b/ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/utils.test.ts @@ -0,0 +1,274 @@ +import { processSSOSettingsPayload } from "./utils"; +import { describe, it, expect } from "vitest"; + +describe("processSSOSettingsPayload", () => { + describe("without role mappings", () => { + it("should return all fields except role mapping fields when use_role_mappings is false", () => { + const formValues = { + proxy_admin_teams: "team1, team2", + admin_viewer_teams: "viewer1", + internal_user_teams: "user1", + internal_viewer_teams: "viewer1", + default_role: "proxy_admin", + group_claim: "groups", + use_role_mappings: false, + other_field: "value", + another_field: 123, + }; + + const result = processSSOSettingsPayload(formValues); + + expect(result).toEqual({ + other_field: "value", + another_field: 123, + }); + expect(result.role_mappings).toBeUndefined(); + }); + + it("should return all fields except role mapping fields when use_role_mappings is not present", () => { + const formValues = { + proxy_admin_teams: "team1", + admin_viewer_teams: "viewer1", + internal_user_teams: "user1", + internal_viewer_teams: "viewer1", + default_role: "proxy_admin", + group_claim: "groups", + other_field: "value", + }; + + const result = processSSOSettingsPayload(formValues); + + expect(result).toEqual({ + other_field: "value", + }); + expect(result.role_mappings).toBeUndefined(); + }); + }); + + describe("with role mappings enabled", () => { + it("should create role mappings with all team types populated", () => { + const formValues = { + proxy_admin_teams: "admin1, admin2", + admin_viewer_teams: "viewer1, viewer2, viewer3", + internal_user_teams: "user1", + internal_viewer_teams: "internal_viewer1, internal_viewer2", + default_role: "proxy_admin", + group_claim: "groups", + use_role_mappings: true, + other_field: "value", + }; + + const result = processSSOSettingsPayload(formValues); + + expect(result.other_field).toBe("value"); + expect(result.role_mappings).toEqual({ + provider: "generic", + group_claim: "groups", + default_role: "proxy_admin", + roles: { + proxy_admin: ["admin1", "admin2"], + proxy_admin_viewer: ["viewer1", "viewer2", "viewer3"], + internal_user: ["user1"], + internal_user_viewer: ["internal_viewer1", "internal_viewer2"], + }, + }); + }); + + it("should handle empty team strings", () => { + const formValues = { + proxy_admin_teams: "", + admin_viewer_teams: "", + internal_user_teams: "", + internal_viewer_teams: "", + default_role: "internal_user", + group_claim: "groups", + use_role_mappings: true, + }; + + const result = processSSOSettingsPayload(formValues); + + expect(result.role_mappings.roles).toEqual({ + proxy_admin: [], + proxy_admin_viewer: [], + internal_user: [], + internal_user_viewer: [], + }); + }); + + it("should handle undefined team fields", () => { + const formValues = { + default_role: "internal_user_viewer", + group_claim: "groups", + use_role_mappings: true, + }; + + const result = processSSOSettingsPayload(formValues); + + expect(result.role_mappings.roles).toEqual({ + proxy_admin: [], + proxy_admin_viewer: [], + internal_user: [], + internal_user_viewer: [], + }); + }); + + it("should handle whitespace-only team strings", () => { + const formValues = { + proxy_admin_teams: " ", + admin_viewer_teams: ", , ,", + internal_user_teams: "user1, , user2", + internal_viewer_teams: "viewer1, ,viewer2", + default_role: "proxy_admin_viewer", + group_claim: "groups", + use_role_mappings: true, + }; + + const result = processSSOSettingsPayload(formValues); + + expect(result.role_mappings.roles).toEqual({ + proxy_admin: [], + proxy_admin_viewer: [], + internal_user: ["user1", "user2"], + internal_user_viewer: ["viewer1", "viewer2"], + }); + }); + + it("should trim whitespace from team names", () => { + const formValues = { + proxy_admin_teams: " admin1 , admin2 ", + admin_viewer_teams: " viewer1 ", + internal_user_teams: " user1 , user2 ", + internal_viewer_teams: "viewer1,viewer2", + default_role: "internal_user", + group_claim: "groups", + use_role_mappings: true, + }; + + const result = processSSOSettingsPayload(formValues); + + expect(result.role_mappings.roles).toEqual({ + proxy_admin: ["admin1", "admin2"], + proxy_admin_viewer: ["viewer1"], + internal_user: ["user1", "user2"], + internal_user_viewer: ["viewer1", "viewer2"], + }); + }); + + it("should filter out empty strings after trimming", () => { + const formValues = { + proxy_admin_teams: "admin1,,admin2, , admin3", + default_role: "internal_user", + group_claim: "groups", + use_role_mappings: true, + }; + + const result = processSSOSettingsPayload(formValues); + + expect(result.role_mappings.roles.proxy_admin).toEqual(["admin1", "admin2", "admin3"]); + }); + }); + + describe("default role mapping", () => { + it("should map internal_user_viewer correctly", () => { + const formValues = { + default_role: "internal_user_viewer", + group_claim: "groups", + use_role_mappings: true, + }; + + const result = processSSOSettingsPayload(formValues); + + expect(result.role_mappings.default_role).toBe("internal_user_viewer"); + }); + + it("should map internal_user correctly", () => { + const formValues = { + default_role: "internal_user", + group_claim: "groups", + use_role_mappings: true, + }; + + const result = processSSOSettingsPayload(formValues); + + expect(result.role_mappings.default_role).toBe("internal_user"); + }); + + it("should map proxy_admin_viewer correctly", () => { + const formValues = { + default_role: "proxy_admin_viewer", + group_claim: "groups", + use_role_mappings: true, + }; + + const result = processSSOSettingsPayload(formValues); + + expect(result.role_mappings.default_role).toBe("proxy_admin_viewer"); + }); + + it("should map proxy_admin correctly", () => { + const formValues = { + default_role: "proxy_admin", + group_claim: "groups", + use_role_mappings: true, + }; + + const result = processSSOSettingsPayload(formValues); + + expect(result.role_mappings.default_role).toBe("proxy_admin"); + }); + + it("should default to internal_user for unknown roles", () => { + const formValues = { + default_role: "unknown_role", + group_claim: "groups", + use_role_mappings: true, + }; + + const result = processSSOSettingsPayload(formValues); + + expect(result.role_mappings.default_role).toBe("internal_user"); + }); + + it("should default to internal_user for undefined default_role", () => { + const formValues = { + group_claim: "groups", + use_role_mappings: true, + }; + + const result = processSSOSettingsPayload(formValues); + + expect(result.role_mappings.default_role).toBe("internal_user"); + }); + }); + + describe("edge cases", () => { + it("should handle empty form values", () => { + const result = processSSOSettingsPayload({}); + + expect(result).toEqual({}); + }); + + it("should preserve other fields in the payload", () => { + const formValues = { + use_role_mappings: false, + sso_provider: "google", + client_id: "123", + client_secret: "secret", + redirect_url: "http://example.com", + custom_field: { nested: "value" }, + array_field: [1, 2, 3], + }; + + const result = processSSOSettingsPayload(formValues); + + expect(result).toEqual({ + sso_provider: "google", + client_id: "123", + client_secret: "secret", + redirect_url: "http://example.com", + custom_field: { nested: "value" }, + array_field: [1, 2, 3], + }); + }); + }); +}); diff --git a/ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/utils.ts b/ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/utils.ts new file mode 100644 index 00000000000..3533e1226c2 --- /dev/null +++ b/ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/utils.ts @@ -0,0 +1,54 @@ +/** + * Processes SSO settings form values and transforms them into the payload format expected by the API + * Handles role mappings transformation and field extraction + */ +export const processSSOSettingsPayload = (formValues: Record): Record => { + const { + proxy_admin_teams, + admin_viewer_teams, + internal_user_teams, + internal_viewer_teams, + default_role, + group_claim, + use_role_mappings, + ...rest + } = formValues; + + const payload: any = { + ...rest, + }; + + // Add role mappings if use_role_mappings is checked + if (use_role_mappings) { + // Helper function to split comma-separated string into array + const splitTeams = (teams: string | undefined): string[] => { + if (!teams || teams.trim() === "") return []; + return teams + .split(",") + .map((team) => team.trim()) + .filter((team) => team.length > 0); + }; + + // Map default role display values to backend values + const defaultRoleMapping: Record = { + internal_user_viewer: "internal_user_viewer", + internal_user: "internal_user", + proxy_admin_viewer: "proxy_admin_viewer", + proxy_admin: "proxy_admin", + }; + + payload.role_mappings = { + provider: "generic", + group_claim, + default_role: defaultRoleMapping[default_role] || "internal_user", + roles: { + proxy_admin: splitTeams(proxy_admin_teams), + proxy_admin_viewer: splitTeams(admin_viewer_teams), + internal_user: splitTeams(internal_user_teams), + internal_user_viewer: splitTeams(internal_viewer_teams), + }, + }; + } + + return payload; +}; From 5c7523b11e3bf1786addd39ed71246f2c8c08b50 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Sat, 3 Jan 2026 12:43:13 -0800 Subject: [PATCH 144/158] Fixing tests --- .../Modals/AddSSOSettingsModal.test.tsx | 23 +++- .../SSOSettings/RedactableField.test.tsx | 108 ++++++++++++++++++ .../AdminSettings/SSOSettings/SSOSettings.tsx | 8 +- 3 files changed, 133 insertions(+), 6 deletions(-) create mode 100644 ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/RedactableField.test.tsx diff --git a/ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/Modals/AddSSOSettingsModal.test.tsx b/ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/Modals/AddSSOSettingsModal.test.tsx index aae28191031..e5a5af6cba6 100644 --- a/ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/Modals/AddSSOSettingsModal.test.tsx +++ b/ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/Modals/AddSSOSettingsModal.test.tsx @@ -1,5 +1,6 @@ -import { render, screen } from "@testing-library/react"; +import { screen } from "@testing-library/react"; import { describe, expect, it, vi } from "vitest"; +import { renderWithProviders } from "../../../../../../tests/test-utils"; import AddSSOSettingsModal from "./AddSSOSettingsModal"; // Mock networking functions @@ -12,12 +13,30 @@ vi.mock("@/components/shared/errorUtils", () => ({ parseErrorMessage: vi.fn((error) => error?.message || "Unknown error"), })); +// Mock the useAuthorized hook +vi.mock("@/app/(dashboard)/hooks/useAuthorized", () => ({ + default: () => ({ + accessToken: "test-access-token", + userId: "test-user-id", + userEmail: "test@example.com", + userRole: "admin", + }), +})); + +// Mock NotificationsManager +vi.mock("@/components/molecules/notifications_manager", () => ({ + default: { + success: vi.fn(), + fromBackend: vi.fn(), + }, +})); + describe("AddSSOSettingsModal", () => { it("should render", () => { const onCancel = vi.fn(); const onSuccess = vi.fn(); - render(); + renderWithProviders(); expect(screen.getByText("SSO Provider")).toBeInTheDocument(); expect(screen.getByText("Cancel")).toBeInTheDocument(); diff --git a/ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/RedactableField.test.tsx b/ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/RedactableField.test.tsx new file mode 100644 index 00000000000..a047d7aea4f --- /dev/null +++ b/ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/RedactableField.test.tsx @@ -0,0 +1,108 @@ +import { render, screen, fireEvent } from "@testing-library/react"; +import { describe, expect, it } from "vitest"; +import RedactableField from "./RedactableField"; + +describe("RedactableField", () => { + describe("when value is null", () => { + it("should display 'Not configured' text", () => { + render(); + + expect(screen.getByText("Not configured")).toBeInTheDocument(); + }); + + it("should not display toggle button", () => { + render(); + + // There should be no button elements + const buttons = screen.queryAllByRole("button"); + expect(buttons).toHaveLength(0); + }); + }); + + describe("when value is provided", () => { + const testValue = "secret-password"; + + it("should be hidden by default and show redacted dots", () => { + render(); + + // Should show dots equal to the length of the value + expect(screen.getByText("•".repeat(testValue.length))).toBeInTheDocument(); + expect(screen.queryByText(testValue)).not.toBeInTheDocument(); + }); + + it("should show actual value when defaultHidden is false", () => { + render(); + + expect(screen.getByText(testValue)).toBeInTheDocument(); + expect(screen.queryByText("•".repeat(testValue.length))).not.toBeInTheDocument(); + }); + + it("should display toggle button with eye icon when hidden", () => { + render(); + + const button = screen.getByRole("button"); + expect(button).toBeInTheDocument(); + + // Check that the Eye icon is rendered (we can check by title or by the presence of the icon) + // The button should contain the Eye icon when hidden + const eyeIcon = button.querySelector("svg"); + expect(eyeIcon).toBeInTheDocument(); + }); + + it("should display toggle button with eye-off icon when shown", () => { + render(); + + const button = screen.getByRole("button"); + expect(button).toBeInTheDocument(); + + // The button should contain the EyeOff icon when shown + const eyeOffIcon = button.querySelector("svg"); + expect(eyeOffIcon).toBeInTheDocument(); + }); + + it("should toggle visibility when button is clicked", () => { + render(); + + // Initially hidden + expect(screen.getByText("•".repeat(testValue.length))).toBeInTheDocument(); + expect(screen.queryByText(testValue)).not.toBeInTheDocument(); + + // Click to show + const button = screen.getByRole("button"); + fireEvent.click(button); + + // Should now show the actual value + expect(screen.getByText(testValue)).toBeInTheDocument(); + expect(screen.queryByText("•".repeat(testValue.length))).not.toBeInTheDocument(); + + // Click again to hide + fireEvent.click(button); + + // Should be hidden again + expect(screen.getByText("•".repeat(testValue.length))).toBeInTheDocument(); + expect(screen.queryByText(testValue)).not.toBeInTheDocument(); + }); + + it("should handle empty string value", () => { + render(); + + // Empty string should show "Not configured" since value is falsy + expect(screen.getByText("Not configured")).toBeInTheDocument(); + + // No toggle button for empty string + const buttons = screen.queryAllByRole("button"); + expect(buttons).toHaveLength(0); + }); + + it("should handle different value lengths correctly", () => { + const shortValue = "hi"; + const longValue = "this-is-a-very-long-secret-value"; + + const { rerender } = render(); + expect(screen.getByText("••")).toBeInTheDocument(); + + rerender(); + expect(screen.getByText("•".repeat(longValue.length))).toBeInTheDocument(); + }); + }); +}); diff --git a/ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/SSOSettings.tsx b/ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/SSOSettings.tsx index 975a3bc7d78..d339f6d0e36 100644 --- a/ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/SSOSettings.tsx +++ b/ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/SSOSettings.tsx @@ -2,15 +2,15 @@ import { useSSOSettings, type SSOSettingsValues } from "@/app/(dashboard)/hooks/sso/useSSOSettings"; import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; -import { Badge, Button, Card, Descriptions, Space, Typography } from "antd"; -import { Shield, Trash2, Edit } from "lucide-react"; +import { Button, Card, Descriptions, Space, Typography } from "antd"; +import { Edit, Shield, Trash2 } from "lucide-react"; import { useState } from "react"; import AddSSOSettingsModal from "./Modals/AddSSOSettingsModal"; import DeleteSSOSettingsModal from "./Modals/DeleteSSOSettingsModal"; import EditSSOSettingsModal from "./Modals/EditSSOSettingsModal"; -import SSOSettingsEmptyPlaceholder from "./SSOSettingsEmptyPlaceholder"; import RedactableField from "./RedactableField"; -import { ssoProviderLogoMap, ssoProviderDisplayNames } from "./constants"; +import SSOSettingsEmptyPlaceholder from "./SSOSettingsEmptyPlaceholder"; +import { ssoProviderDisplayNames, ssoProviderLogoMap } from "./constants"; const { Title, Text } = Typography; From 2cbcaf2abf463a137e9a7f164c305da15e6c9413 Mon Sep 17 00:00:00 2001 From: Alexsander Hamir Date: Sat, 3 Jan 2026 13:34:17 -0800 Subject: [PATCH 145/158] refactor(utils): lazy load heavy imports to improve import time and memory usage (#18610) * refactor(utils): lazy load heavy imports to improve import time - Move BaseVectorStore, CredentialAccessor, and exception_mapping_utils imports to lazy loading via __getattr__ - Add _get_utils_globals() helper function following pattern from _lazy_imports.py - Refactor __getattr__ to use consistent caching pattern matching __init__.py - Update load_credentials_from_list to use lazy-loaded CredentialAccessor This reduces import time and memory usage by only loading these modules when they're actually accessed, not during module import. * refactor(utils): lazy load additional heavy imports to improve import time - Move get_llm_provider, _is_non_openai_azure_model to lazy loading - Move get_supported_openai_params to lazy loading - Move convert_dict_to_response functions (LiteLLMResponseObjectHandler, convert_to_model_response_object, etc.) to lazy loading - Move get_api_base and ResponseMetadata to lazy loading - Move _parse_content_for_reasoning to lazy loading - Update all internal usages to access via getattr(sys.modules[__name__], ...) This reduces import time and memory usage by only loading these modules when they're actually accessed, not during module import. * fix(utils): suppress PLR0915 linter warning for __getattr__ function The __getattr__ function intentionally has many statements to handle multiple lazy-loaded imports. Add noqa comment to suppress the warning. * fix(utils): add type stubs for lazy-loaded functions in TYPE_CHECKING block Add type imports and declarations in TYPE_CHECKING block to help mypy understand the types of lazy-loaded functions accessed via __getattr__. This follows the same pattern used in __init__.py for lazy-loaded items. --- litellm/utils.py | 249 +++++++++++++++++++++++++++++++++++++++-------- 1 file changed, 209 insertions(+), 40 deletions(-) diff --git a/litellm/utils.py b/litellm/utils.py index 3dbeeb970a4..e5eb57b0712 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -71,9 +71,6 @@ from litellm.constants import ( OPENAI_EMBEDDING_PARAMS, TOOL_CHOICE_OBJECT_TOKEN_COUNT, ) -from litellm.integrations.vector_store_integrations.base_vector_store import ( - BaseVectorStore, -) # Import cached imports utilities from litellm.litellm_core_utils.cached_imports import ( @@ -86,48 +83,21 @@ from litellm.litellm_core_utils.core_helpers import ( map_finish_reason, process_response_headers, ) -from litellm.litellm_core_utils.credential_accessor import CredentialAccessor from litellm.litellm_core_utils.dot_notation_indexing import ( delete_nested_value, is_nested_path, ) -from litellm.litellm_core_utils.exception_mapping_utils import ( - _get_response_headers, - exception_type, - get_error_message, -) from litellm.litellm_core_utils.get_litellm_params import ( _get_base_model_from_litellm_call_metadata, get_litellm_params, ) -from litellm.litellm_core_utils.get_llm_provider_logic import ( - _is_non_openai_azure_model, - get_llm_provider, -) -from litellm.litellm_core_utils.get_supported_openai_params import ( - get_supported_openai_params, -) from litellm.litellm_core_utils.llm_request_utils import _ensure_extra_body_is_safe -from litellm.litellm_core_utils.llm_response_utils.convert_dict_to_response import ( - LiteLLMResponseObjectHandler, - _handle_invalid_parallel_tool_calls, - convert_to_model_response_object, - convert_to_streaming_response, - convert_to_streaming_response_async, -) -from litellm.litellm_core_utils.llm_response_utils.get_api_base import get_api_base from litellm.litellm_core_utils.llm_response_utils.get_formatted_prompt import ( get_formatted_prompt, ) from litellm.litellm_core_utils.llm_response_utils.get_headers import ( get_response_headers, ) -from litellm.litellm_core_utils.llm_response_utils.response_metadata import ( - ResponseMetadata, -) -from litellm.litellm_core_utils.prompt_templates.common_utils import ( - _parse_content_for_reasoning, -) from litellm.litellm_core_utils.redact_messages import ( LiteLLMLoggingObject, redact_message_input_output_from_logging, @@ -346,6 +316,30 @@ if TYPE_CHECKING: from litellm.integrations.custom_logger import CustomLogger from litellm.llms.base_llm.files.transformation import BaseFilesConfig from litellm.proxy._types import AllowedModelRegion + # Type stubs for lazy-loaded functions to help mypy understand their types + # These imports allow mypy to understand the types when these are accessed via __getattr__ + from litellm.litellm_core_utils.exception_mapping_utils import exception_type + from litellm.litellm_core_utils.get_llm_provider_logic import ( + _is_non_openai_azure_model, + get_llm_provider, + ) + from litellm.litellm_core_utils.get_supported_openai_params import ( + get_supported_openai_params, + ) + from litellm.litellm_core_utils.llm_response_utils.convert_dict_to_response import ( + LiteLLMResponseObjectHandler, + _handle_invalid_parallel_tool_calls, + convert_to_model_response_object, + convert_to_streaming_response, + convert_to_streaming_response_async, + ) + from litellm.litellm_core_utils.llm_response_utils.get_api_base import get_api_base + from litellm.litellm_core_utils.llm_response_utils.response_metadata import ( + ResponseMetadata, + ) + from litellm.litellm_core_utils.prompt_templates.common_utils import ( + _parse_content_for_reasoning, + ) from litellm.llms.base_llm.batches.transformation import BaseBatchesConfig from litellm.llms.base_llm.chat.transformation import BaseConfig @@ -618,10 +612,22 @@ def get_applied_guardrails(kwargs: Dict[str, Any]) -> List[str]: return applied_guardrails +def _get_utils_globals() -> dict: + """ + Get the globals dictionary of the utils module. + + This is where we cache imported attributes so we don't import them twice. + """ + return sys.modules[__name__].__dict__ + + def load_credentials_from_list(kwargs: dict): """ Updates kwargs with the credentials if credential_name in kwarg """ + # Access CredentialAccessor via module to trigger lazy loading if needed + CredentialAccessor = getattr(sys.modules[__name__], 'CredentialAccessor') + credential_name = kwargs.get("litellm_credential_name") if credential_name and litellm.credential_list: credential_accessor = CredentialAccessor.get_credential_values(credential_name) @@ -2259,6 +2265,7 @@ def supports_response_schema( """ ## GET LLM PROVIDER ## try: + get_llm_provider = getattr(sys.modules[__name__], 'get_llm_provider') model, custom_llm_provider, _, _ = get_llm_provider( model=model, custom_llm_provider=custom_llm_provider ) @@ -2956,6 +2963,9 @@ def get_optional_params_embeddings( # noqa: PLR0915 additional_drop_params: Optional[List[str]] = None, **kwargs, ): + # Lazy load get_supported_openai_params + get_supported_openai_params = getattr(sys.modules[__name__], 'get_supported_openai_params') + # retrieve all parameters passed to the function passed_params = locals() custom_llm_provider = passed_params.pop("custom_llm_provider", None) @@ -3758,6 +3768,7 @@ def get_optional_params( # noqa: PLR0915 message=f"{custom_llm_provider} does not support parameters: {list(unsupported_params.keys())}, for model={model}. To drop these, set `litellm.drop_params=True` or for proxy:\n\n`litellm_settings:\n drop_params: true`\n. \n If you want to use these params dynamically send allowed_openai_params={list(unsupported_params.keys())} in your request.", ) + get_supported_openai_params = getattr(sys.modules[__name__], 'get_supported_openai_params') supported_params = get_supported_openai_params( model=model, custom_llm_provider=custom_llm_provider ) @@ -4895,6 +4906,7 @@ def get_max_tokens(model: str) -> Optional[int]: return litellm.model_cost[model]["max_output_tokens"] elif "max_tokens" in litellm.model_cost[model]: return litellm.model_cost[model]["max_tokens"] + get_llm_provider = getattr(sys.modules[__name__], 'get_llm_provider') model, custom_llm_provider, _, _ = get_llm_provider(model=model) if custom_llm_provider == "huggingface": max_tokens = _get_max_position_embeddings(model_name=model) @@ -5015,6 +5027,7 @@ def _get_potential_model_names( if custom_llm_provider is None: # Get custom_llm_provider try: + get_llm_provider = getattr(sys.modules[__name__], 'get_llm_provider') split_model, custom_llm_provider, _, _ = get_llm_provider(model=model) except Exception: split_model = model @@ -5737,6 +5750,7 @@ def validate_environment( # noqa: PLR0915 } ## EXTRACT LLM PROVIDER - if model name provided try: + get_llm_provider = getattr(sys.modules[__name__], 'get_llm_provider') _, custom_llm_provider, _, _ = get_llm_provider(model=model) except Exception: custom_llm_provider = None @@ -6299,6 +6313,7 @@ def register_prompt_template( complete_model = model potential_models = [complete_model] try: + get_llm_provider = getattr(sys.modules[__name__], 'get_llm_provider') model = get_llm_provider(model=model)[0] potential_models.append(model) except Exception: @@ -6384,6 +6399,7 @@ class TextCompletionStreamWrapper: except StopIteration: raise StopIteration except Exception as e: + exception_type = getattr(sys.modules[__name__], 'exception_type') raise exception_type( model=self.model, custom_llm_provider=self.custom_llm_provider or "", @@ -8705,16 +8721,169 @@ def should_run_mock_completion( return False -# Re-export encoding from main.py for backward compatibility -# This allows tests to import: from litellm.utils import encoding -# We use a lazy import to avoid loading main.py at utils.py import time -def __getattr__(name: str) -> Any: +def __getattr__(name: str) -> Any: # noqa: PLR0915 """Lazy import handler for utils module""" + _globals = _get_utils_globals() + + # Lazy load encoding from main.py to avoid heavy tiktoken import if name == "encoding": - # Cache it in the module's __dict__ for subsequent accesses - import sys - - from litellm.main import encoding as _encoding - sys.modules[__name__].__dict__["encoding"] = _encoding - return _encoding + # Check if already cached + if "encoding" not in _globals: + from litellm.main import encoding as _encoding + _globals["encoding"] = _encoding + return _globals["encoding"] + + # Lazy load BaseVectorStore to avoid loading it at module import time + if name == "BaseVectorStore": + # Check if already cached + if "BaseVectorStore" not in _globals: + from litellm.integrations.vector_store_integrations.base_vector_store import ( + BaseVectorStore as _BaseVectorStore, + ) + _globals["BaseVectorStore"] = _BaseVectorStore + return _globals["BaseVectorStore"] + + # Lazy load CredentialAccessor to avoid loading it at module import time + if name == "CredentialAccessor": + # Check if already cached + if "CredentialAccessor" not in _globals: + from litellm.litellm_core_utils.credential_accessor import ( + CredentialAccessor as _CredentialAccessor, + ) + _globals["CredentialAccessor"] = _CredentialAccessor + return _globals["CredentialAccessor"] + + # Lazy load exception_mapping_utils functions to avoid loading at module import time + if name == "exception_type": + # Check if already cached + if "exception_type" not in _globals: + from litellm.litellm_core_utils.exception_mapping_utils import ( + exception_type as _exception_type, + ) + _globals["exception_type"] = _exception_type + return _globals["exception_type"] + + if name == "get_error_message": + # Check if already cached + if "get_error_message" not in _globals: + from litellm.litellm_core_utils.exception_mapping_utils import ( + get_error_message as _get_error_message, + ) + _globals["get_error_message"] = _get_error_message + return _globals["get_error_message"] + + if name == "_get_response_headers": + # Check if already cached + if "_get_response_headers" not in _globals: + from litellm.litellm_core_utils.exception_mapping_utils import ( + _get_response_headers as __get_response_headers, + ) + _globals["_get_response_headers"] = __get_response_headers + return _globals["_get_response_headers"] + + # Lazy load get_llm_provider_logic functions to avoid loading at module import time + if name == "get_llm_provider": + # Check if already cached + if "get_llm_provider" not in _globals: + from litellm.litellm_core_utils.get_llm_provider_logic import ( + get_llm_provider as _get_llm_provider, + ) + _globals["get_llm_provider"] = _get_llm_provider + return _globals["get_llm_provider"] + + if name == "_is_non_openai_azure_model": + # Check if already cached + if "_is_non_openai_azure_model" not in _globals: + from litellm.litellm_core_utils.get_llm_provider_logic import ( + _is_non_openai_azure_model as __is_non_openai_azure_model, + ) + _globals["_is_non_openai_azure_model"] = __is_non_openai_azure_model + return _globals["_is_non_openai_azure_model"] + + # Lazy load get_supported_openai_params to avoid loading at module import time + if name == "get_supported_openai_params": + # Check if already cached + if "get_supported_openai_params" not in _globals: + from litellm.litellm_core_utils.get_supported_openai_params import ( + get_supported_openai_params as _get_supported_openai_params, + ) + _globals["get_supported_openai_params"] = _get_supported_openai_params + return _globals["get_supported_openai_params"] + + # Lazy load convert_dict_to_response functions to avoid loading at module import time + if name == "LiteLLMResponseObjectHandler": + # Check if already cached + if "LiteLLMResponseObjectHandler" not in _globals: + from litellm.litellm_core_utils.llm_response_utils.convert_dict_to_response import ( + LiteLLMResponseObjectHandler as _LiteLLMResponseObjectHandler, + ) + _globals["LiteLLMResponseObjectHandler"] = _LiteLLMResponseObjectHandler + return _globals["LiteLLMResponseObjectHandler"] + + if name == "_handle_invalid_parallel_tool_calls": + # Check if already cached + if "_handle_invalid_parallel_tool_calls" not in _globals: + from litellm.litellm_core_utils.llm_response_utils.convert_dict_to_response import ( + _handle_invalid_parallel_tool_calls as __handle_invalid_parallel_tool_calls, + ) + _globals["_handle_invalid_parallel_tool_calls"] = __handle_invalid_parallel_tool_calls + return _globals["_handle_invalid_parallel_tool_calls"] + + if name == "convert_to_model_response_object": + # Check if already cached + if "convert_to_model_response_object" not in _globals: + from litellm.litellm_core_utils.llm_response_utils.convert_dict_to_response import ( + convert_to_model_response_object as _convert_to_model_response_object, + ) + _globals["convert_to_model_response_object"] = _convert_to_model_response_object + return _globals["convert_to_model_response_object"] + + if name == "convert_to_streaming_response": + # Check if already cached + if "convert_to_streaming_response" not in _globals: + from litellm.litellm_core_utils.llm_response_utils.convert_dict_to_response import ( + convert_to_streaming_response as _convert_to_streaming_response, + ) + _globals["convert_to_streaming_response"] = _convert_to_streaming_response + return _globals["convert_to_streaming_response"] + + if name == "convert_to_streaming_response_async": + # Check if already cached + if "convert_to_streaming_response_async" not in _globals: + from litellm.litellm_core_utils.llm_response_utils.convert_dict_to_response import ( + convert_to_streaming_response_async as _convert_to_streaming_response_async, + ) + _globals["convert_to_streaming_response_async"] = _convert_to_streaming_response_async + return _globals["convert_to_streaming_response_async"] + + # Lazy load get_api_base to avoid loading at module import time + if name == "get_api_base": + # Check if already cached + if "get_api_base" not in _globals: + from litellm.litellm_core_utils.llm_response_utils.get_api_base import ( + get_api_base as _get_api_base, + ) + _globals["get_api_base"] = _get_api_base + return _globals["get_api_base"] + + # Lazy load ResponseMetadata to avoid loading at module import time + if name == "ResponseMetadata": + # Check if already cached + if "ResponseMetadata" not in _globals: + from litellm.litellm_core_utils.llm_response_utils.response_metadata import ( + ResponseMetadata as _ResponseMetadata, + ) + _globals["ResponseMetadata"] = _ResponseMetadata + return _globals["ResponseMetadata"] + + # Lazy load _parse_content_for_reasoning to avoid loading at module import time + if name == "_parse_content_for_reasoning": + # Check if already cached + if "_parse_content_for_reasoning" not in _globals: + from litellm.litellm_core_utils.prompt_templates.common_utils import ( + _parse_content_for_reasoning as __parse_content_for_reasoning, + ) + _globals["_parse_content_for_reasoning"] = __parse_content_for_reasoning + return _globals["_parse_content_for_reasoning"] + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") From 1c5c303e986bb5c11756ab411713badd6c1a1362 Mon Sep 17 00:00:00 2001 From: Alexsander Hamir Date: Sat, 3 Jan 2026 14:07:57 -0800 Subject: [PATCH 146/158] refactor(utils): implement lazy loading for provider configs, model info classes, streaming handlers, and redact utilities (#18611) * refactor(utils): lazy load redact_messages imports to improve import time - Move LiteLLMLoggingObject and redact_message_input_output_from_logging to lazy loading via __getattr__ - Add type stubs in TYPE_CHECKING block for mypy type checking - These are only used in type annotations (with from __future__ import annotations), so lazy loading works correctly This reduces import time by deferring the redact_messages module import until these are actually accessed. * refactor(utils): lazy load CustomStreamWrapper to improve import time - Move CustomStreamWrapper from streaming_handler to lazy loading via __getattr__ - Add type stub in TYPE_CHECKING block for mypy type checking - CustomStreamWrapper is not used internally in utils.py, only exported for other modules This reduces import time by deferring the streaming_handler module import until CustomStreamWrapper is actually accessed. * refactor(utils): lazy load BaseGoogleGenAIGenerateContentConfig to improve import time - Move BaseGoogleGenAIGenerateContentConfig from google_genai.transformation to lazy loading via __getattr__ - Add type stub in TYPE_CHECKING block for mypy type checking - BaseGoogleGenAIGenerateContentConfig is only used in type annotations (with from __future__ import annotations), so lazy loading works correctly This reduces import time by deferring the google_genai.transformation module import until BaseGoogleGenAIGenerateContentConfig is actually accessed. * refactor(utils): lazy load BaseOCRConfig, BaseSearchConfig, and BaseTextToSpeechConfig - Move BaseOCRConfig, BaseSearchConfig, and BaseTextToSpeechConfig to lazy loading via __getattr__ - Add type stubs in TYPE_CHECKING block for mypy type checking - These config classes are only used in quoted type annotations (forward references), so lazy loading works correctly This reduces import time by deferring the transformation module imports until these config classes are actually accessed. * refactor(utils): lazy load BedrockModelInfo, CohereModelInfo, and MistralOCRConfig - Move BedrockModelInfo, CohereModelInfo, and MistralOCRConfig to lazy loading via __getattr__ - Add type stubs in TYPE_CHECKING block for mypy type checking - Update internal usages to use getattr pattern for accessing lazy-loaded classes - These provider-specific model info classes are only used in specific code paths, so lazy loading reduces initial import time This reduces import time by deferring the bedrock, cohere, and mistral module imports until these classes are actually accessed. * fix(utils): remove duplicate MistralOCRConfig import in TYPE_CHECKING block --- litellm/utils.py | 130 ++++++++++++++++++++++++++++++++++++++++++----- 1 file changed, 116 insertions(+), 14 deletions(-) diff --git a/litellm/utils.py b/litellm/utils.py index e5eb57b0712..d373b5102ef 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -98,22 +98,8 @@ from litellm.litellm_core_utils.llm_response_utils.get_formatted_prompt import ( from litellm.litellm_core_utils.llm_response_utils.get_headers import ( get_response_headers, ) -from litellm.litellm_core_utils.redact_messages import ( - LiteLLMLoggingObject, - redact_message_input_output_from_logging, -) from litellm.litellm_core_utils.rules import Rules -from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper -from litellm.llms.base_llm.google_genai.transformation import ( - BaseGoogleGenAIGenerateContentConfig, -) -from litellm.llms.base_llm.ocr.transformation import BaseOCRConfig -from litellm.llms.base_llm.search.transformation import BaseSearchConfig -from litellm.llms.base_llm.text_to_speech.transformation import BaseTextToSpeechConfig -from litellm.llms.bedrock.common_utils import BedrockModelInfo -from litellm.llms.cohere.common_utils import CohereModelInfo from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler -from litellm.llms.mistral.ocr.transformation import MistralOCRConfig from litellm.router_utils.get_retry_from_policy import ( get_num_retries_from_retry_policy, reset_retry_policy, @@ -340,6 +326,20 @@ if TYPE_CHECKING: from litellm.litellm_core_utils.prompt_templates.common_utils import ( _parse_content_for_reasoning, ) + from litellm.litellm_core_utils.redact_messages import ( + LiteLLMLoggingObject, + redact_message_input_output_from_logging, + ) + from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper + from litellm.llms.base_llm.google_genai.transformation import ( + BaseGoogleGenAIGenerateContentConfig, + ) + from litellm.llms.base_llm.ocr.transformation import BaseOCRConfig + from litellm.llms.base_llm.search.transformation import BaseSearchConfig + from litellm.llms.base_llm.text_to_speech.transformation import BaseTextToSpeechConfig + from litellm.llms.bedrock.common_utils import BedrockModelInfo + from litellm.llms.cohere.common_utils import CohereModelInfo + from litellm.llms.mistral.ocr.transformation import MistralOCRConfig from litellm.llms.base_llm.batches.transformation import BaseBatchesConfig from litellm.llms.base_llm.chat.transformation import BaseConfig @@ -4023,6 +4023,7 @@ def get_optional_params( # noqa: PLR0915 ), ) elif custom_llm_provider == "bedrock": + BedrockModelInfo = getattr(sys.modules[__name__], 'BedrockModelInfo') bedrock_route = BedrockModelInfo.get_bedrock_route(model) bedrock_base_model = BedrockModelInfo.get_base_model(model) if bedrock_route == "converse" or bedrock_route == "converse_like": @@ -7440,6 +7441,7 @@ class ProviderConfigManager: litellm.LlmProviders.COHERE_CHAT == provider or litellm.LlmProviders.COHERE == provider ): + CohereModelInfo = getattr(sys.modules[__name__], 'CohereModelInfo') route = CohereModelInfo.get_cohere_route(model) if route == "v2": return litellm.CohereV2ChatConfig() @@ -8290,6 +8292,7 @@ class ProviderConfigManager: return get_vertex_ai_ocr_config(model=model) + MistralOCRConfig = getattr(sys.modules[__name__], 'MistralOCRConfig') PROVIDER_TO_CONFIG_MAP = { litellm.LlmProviders.MISTRAL: MistralOCRConfig, } @@ -8886,4 +8889,103 @@ def __getattr__(name: str) -> Any: # noqa: PLR0915 _globals["_parse_content_for_reasoning"] = __parse_content_for_reasoning return _globals["_parse_content_for_reasoning"] + # Lazy load redact_messages to avoid loading at module import time + if name == "LiteLLMLoggingObject": + # Check if already cached + if "LiteLLMLoggingObject" not in _globals: + from litellm.litellm_core_utils.redact_messages import ( + LiteLLMLoggingObject as _LiteLLMLoggingObject, + ) + _globals["LiteLLMLoggingObject"] = _LiteLLMLoggingObject + return _globals["LiteLLMLoggingObject"] + + if name == "redact_message_input_output_from_logging": + # Check if already cached + if "redact_message_input_output_from_logging" not in _globals: + from litellm.litellm_core_utils.redact_messages import ( + redact_message_input_output_from_logging as _redact_message_input_output_from_logging, + ) + _globals["redact_message_input_output_from_logging"] = _redact_message_input_output_from_logging + return _globals["redact_message_input_output_from_logging"] + + # Lazy load CustomStreamWrapper to avoid loading at module import time + if name == "CustomStreamWrapper": + # Check if already cached + if "CustomStreamWrapper" not in _globals: + from litellm.litellm_core_utils.streaming_handler import ( + CustomStreamWrapper as _CustomStreamWrapper, + ) + _globals["CustomStreamWrapper"] = _CustomStreamWrapper + return _globals["CustomStreamWrapper"] + + # Lazy load BaseGoogleGenAIGenerateContentConfig to avoid loading at module import time + if name == "BaseGoogleGenAIGenerateContentConfig": + # Check if already cached + if "BaseGoogleGenAIGenerateContentConfig" not in _globals: + from litellm.llms.base_llm.google_genai.transformation import ( + BaseGoogleGenAIGenerateContentConfig as _BaseGoogleGenAIGenerateContentConfig, + ) + _globals["BaseGoogleGenAIGenerateContentConfig"] = _BaseGoogleGenAIGenerateContentConfig + return _globals["BaseGoogleGenAIGenerateContentConfig"] + + # Lazy load BaseOCRConfig to avoid loading at module import time + if name == "BaseOCRConfig": + # Check if already cached + if "BaseOCRConfig" not in _globals: + from litellm.llms.base_llm.ocr.transformation import ( + BaseOCRConfig as _BaseOCRConfig, + ) + _globals["BaseOCRConfig"] = _BaseOCRConfig + return _globals["BaseOCRConfig"] + + # Lazy load BaseSearchConfig to avoid loading at module import time + if name == "BaseSearchConfig": + # Check if already cached + if "BaseSearchConfig" not in _globals: + from litellm.llms.base_llm.search.transformation import ( + BaseSearchConfig as _BaseSearchConfig, + ) + _globals["BaseSearchConfig"] = _BaseSearchConfig + return _globals["BaseSearchConfig"] + + # Lazy load BaseTextToSpeechConfig to avoid loading at module import time + if name == "BaseTextToSpeechConfig": + # Check if already cached + if "BaseTextToSpeechConfig" not in _globals: + from litellm.llms.base_llm.text_to_speech.transformation import ( + BaseTextToSpeechConfig as _BaseTextToSpeechConfig, + ) + _globals["BaseTextToSpeechConfig"] = _BaseTextToSpeechConfig + return _globals["BaseTextToSpeechConfig"] + + # Lazy load BedrockModelInfo to avoid loading at module import time + if name == "BedrockModelInfo": + # Check if already cached + if "BedrockModelInfo" not in _globals: + from litellm.llms.bedrock.common_utils import ( + BedrockModelInfo as _BedrockModelInfo, + ) + _globals["BedrockModelInfo"] = _BedrockModelInfo + return _globals["BedrockModelInfo"] + + # Lazy load CohereModelInfo to avoid loading at module import time + if name == "CohereModelInfo": + # Check if already cached + if "CohereModelInfo" not in _globals: + from litellm.llms.cohere.common_utils import ( + CohereModelInfo as _CohereModelInfo, + ) + _globals["CohereModelInfo"] = _CohereModelInfo + return _globals["CohereModelInfo"] + + # Lazy load MistralOCRConfig to avoid loading at module import time + if name == "MistralOCRConfig": + # Check if already cached + if "MistralOCRConfig" not in _globals: + from litellm.llms.mistral.ocr.transformation import ( + MistralOCRConfig as _MistralOCRConfig, + ) + _globals["MistralOCRConfig"] = _MistralOCRConfig + return _globals["MistralOCRConfig"] + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") From a6c3fb1fb598422e251fa5bd13e2aa4d16c7d3a6 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Sat, 3 Jan 2026 15:10:52 -0800 Subject: [PATCH 147/158] E2E test see models for specific provider --- ui/litellm-dashboard/e2e_tests/constants.ts | 1 + .../tests/modelsPage/addModel.spec.ts | 23 +++++++++++++++++++ .../tests/navigation/sidebar.spec.ts | 3 ++- 3 files changed, 26 insertions(+), 1 deletion(-) create mode 100644 ui/litellm-dashboard/e2e_tests/constants.ts create mode 100644 ui/litellm-dashboard/e2e_tests/tests/modelsPage/addModel.spec.ts diff --git a/ui/litellm-dashboard/e2e_tests/constants.ts b/ui/litellm-dashboard/e2e_tests/constants.ts new file mode 100644 index 00000000000..b07bd68fcf1 --- /dev/null +++ b/ui/litellm-dashboard/e2e_tests/constants.ts @@ -0,0 +1 @@ +export const ADMIN_STORAGE_PATH = "admin.storageState.json"; diff --git a/ui/litellm-dashboard/e2e_tests/tests/modelsPage/addModel.spec.ts b/ui/litellm-dashboard/e2e_tests/tests/modelsPage/addModel.spec.ts new file mode 100644 index 00000000000..5fa11a98ef6 --- /dev/null +++ b/ui/litellm-dashboard/e2e_tests/tests/modelsPage/addModel.spec.ts @@ -0,0 +1,23 @@ +import { test, expect } from "@playwright/test"; +import { ADMIN_STORAGE_PATH } from "../../constants"; + +test.describe("Add Model", () => { + test.use({ storageState: ADMIN_STORAGE_PATH }); + + test("Able to see all models for a specific provider in the model dropdown", async ({ page }) => { + await page.goto("http://localhost:4000/ui"); + + await page.getByText("Models + Endpoints").click(); + await page.getByRole("tab", { name: "Add Model" }).click(); + + const providerInputDropdown = page.getByRole("combobox", { name: /Provider/i }); + await providerInputDropdown.fill("Anthropic"); + await page.waitForTimeout(1000); + await providerInputDropdown.press("Enter"); + await page.waitForTimeout(1000); + + const providerModelsDropdown = page.locator(".ant-select-selection-overflow").first(); + await providerModelsDropdown.click(); + await expect(page.getByTitle("claude-haiku-4-5", { exact: true })).toBeVisible(); + }); +}); diff --git a/ui/litellm-dashboard/e2e_tests/tests/navigation/sidebar.spec.ts b/ui/litellm-dashboard/e2e_tests/tests/navigation/sidebar.spec.ts index dafb03a7cbd..6801f891e87 100644 --- a/ui/litellm-dashboard/e2e_tests/tests/navigation/sidebar.spec.ts +++ b/ui/litellm-dashboard/e2e_tests/tests/navigation/sidebar.spec.ts @@ -1,5 +1,6 @@ import test, { expect } from "@playwright/test"; import { Role } from "../../fixtures/roles"; +import { ADMIN_STORAGE_PATH } from "../../constants"; const sidebarButtons = { [Role.ProxyAdmin]: [ @@ -16,7 +17,7 @@ const sidebarButtons = { ], }; -const roles = [{ role: Role.ProxyAdmin, storage: "admin.storageState.json" }]; +const roles = [{ role: Role.ProxyAdmin, storage: ADMIN_STORAGE_PATH }]; for (const { role, storage } of roles) { test.describe(`${role} sidebar`, () => { From dd1ccec7348b73f9a3d7d1ec8ac5fd29651de759 Mon Sep 17 00:00:00 2001 From: Alexsander Hamir Date: Sat, 3 Jan 2026 15:50:53 -0800 Subject: [PATCH 148/158] refactor(utils): lazy load 15 additional imports to improve import time (#18613) * refactor(utils): lazy load 15 additional imports to improve import time - Move Rules, AsyncHTTPHandler, HTTPHandler to lazy loading via __getattr__ - Move get_num_retries_from_retry_policy, reset_retry_policy to lazy loading - Move get_secret to lazy loading - Move cached_imports functions (get_coroutine_checker, get_litellm_logging_class, get_set_callbacks) to lazy loading - Move core_helpers functions (get_litellm_metadata_from_kwargs, map_finish_reason, process_response_headers) to lazy loading - Move dot_notation_indexing functions (delete_nested_value, is_nested_path) to lazy loading - Move get_litellm_params functions to lazy loading - Move _ensure_extra_body_is_safe, get_formatted_prompt, get_response_headers, update_response_metadata to lazy loading - Move executor to lazy loading - Move BaseAnthropicMessagesConfig, BaseAudioTranscriptionConfig to lazy loading - Add type stubs in TYPE_CHECKING block for mypy type checking - These functions/classes are exported for other modules but not used internally in utils.py, so lazy loading is safe and improves startup performance * fix(utils): use getattr for Rules and get_coroutine_checker in client decorator - Update Rules() instantiation in client decorator to use getattr for lazy loading - Update Rules.has_pre_call_rules() usage in function_setup to use getattr - Update get_coroutine_checker() usage in client decorator to use getattr - Fixes NameError: name 'Rules' is not defined error that occurs when Rules is lazy-loaded * fix(utils): use getattr for get_litellm_logging_class in function_setup - Update get_litellm_logging_class() usage in function_setup to use getattr for lazy loading - Fixes NameError: name 'get_litellm_logging_class' is not defined error that occurs when get_litellm_logging_class is lazy-loaded * fix(utils): use getattr for get_set_callbacks in function_setup - Update get_set_callbacks() usage in function_setup to use getattr for lazy loading - Fixes NameError: name 'get_set_callbacks' is not defined error that occurs when get_set_callbacks is lazy-loaded * fix(utils): use getattr for all lazy-loaded imports in utils.py - Update update_response_metadata (4 occurrences) to use getattr - Update executor.submit (1 occurrence) to use getattr - Update get_num_retries_from_retry_policy (2 occurrences) to use getattr - Update reset_retry_policy (2 occurrences) to use getattr - Update is_nested_path and delete_nested_value (1 occurrence each) to use getattr - Update _ensure_extra_body_is_safe (1 occurrence) to use getattr Fixes NameError errors that occur when these functions/classes are lazy-loaded but used directly in utils.py * fix(utils): use getattr for _get_base_model_from_litellm_call_metadata in _get_base_model_from_metadata - Update _get_base_model_from_litellm_call_metadata usage to use getattr for lazy loading - Fixes NameError: name '_get_base_model_from_litellm_call_metadata' is not defined * fix(utils): use getattr for second _get_base_model_from_litellm_call_metadata usage - Fix the second occurrence of _get_base_model_from_litellm_call_metadata on line 7052 - Both occurrences in _get_base_model_from_metadata now use getattr for lazy loading * fix(utils): fix indentation in _get_base_model_from_metadata function * fix(utils): use getattr for get_litellm_metadata_from_kwargs in _get_litellm_params - Update get_litellm_metadata_from_kwargs usage to use getattr for lazy loading - Fixes NameError: name 'get_litellm_metadata_from_kwargs' is not defined * fix(utils): fix syntax error in get_litellm_metadata_from_kwargs fix - Move getattr call before cast statement to fix syntax error --- litellm/utils.py | 330 ++++++++++++++++++++++++++++++++++++++++------- 1 file changed, 284 insertions(+), 46 deletions(-) diff --git a/litellm/utils.py b/litellm/utils.py index d373b5102ef..6f4652c8278 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -72,39 +72,7 @@ from litellm.constants import ( TOOL_CHOICE_OBJECT_TOKEN_COUNT, ) -# Import cached imports utilities -from litellm.litellm_core_utils.cached_imports import ( - get_coroutine_checker, - get_litellm_logging_class, - get_set_callbacks, -) -from litellm.litellm_core_utils.core_helpers import ( - get_litellm_metadata_from_kwargs, - map_finish_reason, - process_response_headers, -) -from litellm.litellm_core_utils.dot_notation_indexing import ( - delete_nested_value, - is_nested_path, -) -from litellm.litellm_core_utils.get_litellm_params import ( - _get_base_model_from_litellm_call_metadata, - get_litellm_params, -) -from litellm.litellm_core_utils.llm_request_utils import _ensure_extra_body_is_safe -from litellm.litellm_core_utils.llm_response_utils.get_formatted_prompt import ( - get_formatted_prompt, -) -from litellm.litellm_core_utils.llm_response_utils.get_headers import ( - get_response_headers, -) -from litellm.litellm_core_utils.rules import Rules -from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler -from litellm.router_utils.get_retry_from_policy import ( - get_num_retries_from_retry_policy, - reset_retry_policy, -) -from litellm.secret_managers.main import get_secret + _CachingHandlerResponse = None _LLMCachingHandler = None @@ -280,16 +248,7 @@ from typing import ( from openai import OpenAIError as OriginalError -from litellm.litellm_core_utils.llm_response_utils.response_metadata import ( - update_response_metadata, -) -from litellm.litellm_core_utils.thread_pool_executor import executor -from litellm.llms.base_llm.anthropic_messages.transformation import ( - BaseAnthropicMessagesConfig, -) -from litellm.llms.base_llm.audio_transcription.transformation import ( - BaseAudioTranscriptionConfig, -) +# These are lazy loaded via __getattr__ from litellm.llms.base_llm.base_utils import ( BaseLLMModelInfo, type_to_response_format_param, @@ -340,6 +299,49 @@ if TYPE_CHECKING: from litellm.llms.bedrock.common_utils import BedrockModelInfo from litellm.llms.cohere.common_utils import CohereModelInfo from litellm.llms.mistral.ocr.transformation import MistralOCRConfig + # Type stubs for lazy-loaded functions and classes + from litellm.litellm_core_utils.cached_imports import ( + get_coroutine_checker, + get_litellm_logging_class, + get_set_callbacks, + ) + from litellm.litellm_core_utils.core_helpers import ( + get_litellm_metadata_from_kwargs, + map_finish_reason, + process_response_headers, + ) + from litellm.litellm_core_utils.dot_notation_indexing import ( + delete_nested_value, + is_nested_path, + ) + from litellm.litellm_core_utils.get_litellm_params import ( + _get_base_model_from_litellm_call_metadata, + get_litellm_params, + ) + from litellm.litellm_core_utils.llm_request_utils import _ensure_extra_body_is_safe + from litellm.litellm_core_utils.llm_response_utils.get_formatted_prompt import ( + get_formatted_prompt, + ) + from litellm.litellm_core_utils.llm_response_utils.get_headers import ( + get_response_headers, + ) + from litellm.litellm_core_utils.llm_response_utils.response_metadata import ( + update_response_metadata, + ) + from litellm.litellm_core_utils.rules import Rules + from litellm.litellm_core_utils.thread_pool_executor import executor + from litellm.llms.base_llm.anthropic_messages.transformation import ( + BaseAnthropicMessagesConfig, + ) + from litellm.llms.base_llm.audio_transcription.transformation import ( + BaseAudioTranscriptionConfig, + ) + from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler + from litellm.router_utils.get_retry_from_policy import ( + get_num_retries_from_retry_policy, + reset_retry_policy, + ) + from litellm.secret_managers.main import get_secret from litellm.llms.base_llm.batches.transformation import BaseBatchesConfig from litellm.llms.base_llm.chat.transformation import BaseConfig @@ -816,6 +818,7 @@ def function_setup( # noqa: PLR0915 + litellm.failure_callback ) ) + get_set_callbacks = getattr(sys.modules[__name__], 'get_set_callbacks') get_set_callbacks()(callback_list=callback_list, function_id=function_id) ## ASYNC CALLBACKS if len(litellm.input_callback) > 0: @@ -943,6 +946,7 @@ def function_setup( # noqa: PLR0915 elif kwargs.get("messages", None): messages = kwargs["messages"] ### PRE-CALL RULES ### + Rules = getattr(sys.modules[__name__], 'Rules') if ( Rules.has_pre_call_rules() and isinstance(messages, list) @@ -1075,6 +1079,7 @@ def function_setup( # noqa: PLR0915 call_type=call_type, ): stream = True + get_litellm_logging_class = getattr(sys.modules[__name__], 'get_litellm_logging_class') logging_obj = get_litellm_logging_class()( # Victim for object pool model=model, # type: ignore messages=messages, @@ -1158,6 +1163,8 @@ def _get_wrapper_num_retries( if num_retries is None: num_retries = litellm.num_retries if kwargs.get("retry_policy", None): + get_num_retries_from_retry_policy = getattr(sys.modules[__name__], 'get_num_retries_from_retry_policy') + reset_retry_policy = getattr(sys.modules[__name__], 'reset_retry_policy') retry_policy_num_retries = get_num_retries_from_retry_policy( exception=exception, retry_policy=kwargs.get("retry_policy"), @@ -1343,6 +1350,7 @@ def post_call_processing( def client(original_function): # noqa: PLR0915 + Rules = getattr(sys.modules[__name__], 'Rules') rules_obj = Rules() @wraps(original_function) @@ -1528,6 +1536,7 @@ def client(original_function): # noqa: PLR0915 ) else: # RETURN RESULT + update_response_metadata = getattr(sys.modules[__name__], 'update_response_metadata') update_response_metadata( result=result, logging_obj=logging_obj, @@ -1571,6 +1580,7 @@ def client(original_function): # noqa: PLR0915 # Copy the current context to propagate it to the background thread # This is essential for OpenTelemetry span context propagation ctx = contextvars.copy_context() + executor = getattr(sys.modules[__name__], 'executor') executor.submit( ctx.run, logging_obj.success_handler, @@ -1579,6 +1589,7 @@ def client(original_function): # noqa: PLR0915 end_time, ) # RETURN RESULT + update_response_metadata = getattr(sys.modules[__name__], 'update_response_metadata') update_response_metadata( result=result, logging_obj=logging_obj, @@ -1595,6 +1606,8 @@ def client(original_function): # noqa: PLR0915 kwargs.get("num_retries", None) or litellm.num_retries or None ) if kwargs.get("retry_policy", None): + get_num_retries_from_retry_policy = getattr(sys.modules[__name__], 'get_num_retries_from_retry_policy') + reset_retry_policy = getattr(sys.modules[__name__], 'reset_retry_policy') num_retries = get_num_retries_from_retry_policy( exception=e, retry_policy=kwargs.get("retry_policy"), @@ -1766,6 +1779,7 @@ def client(original_function): # noqa: PLR0915 chunks, messages=kwargs.get("messages", None) ) else: + update_response_metadata = getattr(sys.modules[__name__], 'update_response_metadata') update_response_metadata( result=result, logging_obj=logging_obj, @@ -1830,6 +1844,7 @@ def client(original_function): # noqa: PLR0915 end_time=end_time, ) + update_response_metadata = getattr(sys.modules[__name__], 'update_response_metadata') update_response_metadata( result=result, logging_obj=logging_obj, @@ -1905,6 +1920,7 @@ def client(original_function): # noqa: PLR0915 setattr(e, "timeout", timeout) raise e + get_coroutine_checker = getattr(sys.modules[__name__], 'get_coroutine_checker') is_coroutine = get_coroutine_checker().is_async_callable(original_function) # Return the appropriate wrapper based on the original function type @@ -4448,6 +4464,8 @@ def get_optional_params( # noqa: PLR0915 # Apply nested drops from additional_drop_params if additional_drop_params: + is_nested_path = getattr(sys.modules[__name__], 'is_nested_path') + delete_nested_value = getattr(sys.modules[__name__], 'delete_nested_value') nested_paths = [p for p in additional_drop_params if is_nested_path(p)] for path in nested_paths: optional_params = delete_nested_value(optional_params, path) @@ -4497,6 +4515,7 @@ def add_provider_specific_params_to_optional_params( else: processed_extra_body = initial_extra_body + _ensure_extra_body_is_safe = getattr(sys.modules[__name__], '_ensure_extra_body_is_safe') optional_params["extra_body"] = _ensure_extra_body_is_safe( extra_body=processed_extra_body ) @@ -7022,14 +7041,14 @@ def _get_base_model_from_metadata(model_call_details=None): return _base_model metadata = litellm_params.get("metadata", {}) - base_model_from_metadata = _get_base_model_from_litellm_call_metadata( - metadata=metadata - ) + _get_base_model_from_litellm_call_metadata = getattr(sys.modules[__name__], '_get_base_model_from_litellm_call_metadata') + base_model_from_metadata = _get_base_model_from_litellm_call_metadata(metadata=metadata) if base_model_from_metadata is not None: return base_model_from_metadata # Also check litellm_metadata (used by Responses API and other generic API calls) litellm_metadata = litellm_params.get("litellm_metadata", {}) + _get_base_model_from_litellm_call_metadata = getattr(sys.modules[__name__], '_get_base_model_from_litellm_call_metadata') return _get_base_model_from_litellm_call_metadata(metadata=litellm_metadata) return None @@ -8433,6 +8452,7 @@ def get_end_user_id_for_cost_tracking( service_type: "litellm_logging" or "prometheus" - used to allow prometheus only disable cost tracking. """ + get_litellm_metadata_from_kwargs = getattr(sys.modules[__name__], 'get_litellm_metadata_from_kwargs') _metadata = cast( dict, get_litellm_metadata_from_kwargs(dict(litellm_params=litellm_params)) ) @@ -8988,4 +9008,222 @@ def __getattr__(name: str) -> Any: # noqa: PLR0915 _globals["MistralOCRConfig"] = _MistralOCRConfig return _globals["MistralOCRConfig"] + # Lazy load Rules to avoid loading at module import time + if name == "Rules": + # Check if already cached + if "Rules" not in _globals: + from litellm.litellm_core_utils.rules import Rules as _Rules + _globals["Rules"] = _Rules + return _globals["Rules"] + + # Lazy load AsyncHTTPHandler and HTTPHandler to avoid loading at module import time + if name == "AsyncHTTPHandler": + # Check if already cached + if "AsyncHTTPHandler" not in _globals: + from litellm.llms.custom_httpx.http_handler import ( + AsyncHTTPHandler as _AsyncHTTPHandler, + ) + _globals["AsyncHTTPHandler"] = _AsyncHTTPHandler + return _globals["AsyncHTTPHandler"] + + if name == "HTTPHandler": + # Check if already cached + if "HTTPHandler" not in _globals: + from litellm.llms.custom_httpx.http_handler import ( + HTTPHandler as _HTTPHandler, + ) + _globals["HTTPHandler"] = _HTTPHandler + return _globals["HTTPHandler"] + + # Lazy load get_num_retries_from_retry_policy and reset_retry_policy to avoid loading at module import time + if name == "get_num_retries_from_retry_policy": + # Check if already cached + if "get_num_retries_from_retry_policy" not in _globals: + from litellm.router_utils.get_retry_from_policy import ( + get_num_retries_from_retry_policy as _get_num_retries_from_retry_policy, + ) + _globals["get_num_retries_from_retry_policy"] = _get_num_retries_from_retry_policy + return _globals["get_num_retries_from_retry_policy"] + + if name == "reset_retry_policy": + # Check if already cached + if "reset_retry_policy" not in _globals: + from litellm.router_utils.get_retry_from_policy import ( + reset_retry_policy as _reset_retry_policy, + ) + _globals["reset_retry_policy"] = _reset_retry_policy + return _globals["reset_retry_policy"] + + # Lazy load get_secret to avoid loading at module import time + if name == "get_secret": + # Check if already cached + if "get_secret" not in _globals: + from litellm.secret_managers.main import get_secret as _get_secret + _globals["get_secret"] = _get_secret + return _globals["get_secret"] + + # Lazy load cached_imports functions to avoid loading at module import time + if name == "get_coroutine_checker": + # Check if already cached + if "get_coroutine_checker" not in _globals: + from litellm.litellm_core_utils.cached_imports import ( + get_coroutine_checker as _get_coroutine_checker, + ) + _globals["get_coroutine_checker"] = _get_coroutine_checker + return _globals["get_coroutine_checker"] + + if name == "get_litellm_logging_class": + # Check if already cached + if "get_litellm_logging_class" not in _globals: + from litellm.litellm_core_utils.cached_imports import ( + get_litellm_logging_class as _get_litellm_logging_class, + ) + _globals["get_litellm_logging_class"] = _get_litellm_logging_class + return _globals["get_litellm_logging_class"] + + if name == "get_set_callbacks": + # Check if already cached + if "get_set_callbacks" not in _globals: + from litellm.litellm_core_utils.cached_imports import ( + get_set_callbacks as _get_set_callbacks, + ) + _globals["get_set_callbacks"] = _get_set_callbacks + return _globals["get_set_callbacks"] + + # Lazy load core_helpers functions to avoid loading at module import time + if name == "get_litellm_metadata_from_kwargs": + # Check if already cached + if "get_litellm_metadata_from_kwargs" not in _globals: + from litellm.litellm_core_utils.core_helpers import ( + get_litellm_metadata_from_kwargs as _get_litellm_metadata_from_kwargs, + ) + _globals["get_litellm_metadata_from_kwargs"] = _get_litellm_metadata_from_kwargs + return _globals["get_litellm_metadata_from_kwargs"] + + if name == "map_finish_reason": + # Check if already cached + if "map_finish_reason" not in _globals: + from litellm.litellm_core_utils.core_helpers import ( + map_finish_reason as _map_finish_reason, + ) + _globals["map_finish_reason"] = _map_finish_reason + return _globals["map_finish_reason"] + + if name == "process_response_headers": + # Check if already cached + if "process_response_headers" not in _globals: + from litellm.litellm_core_utils.core_helpers import ( + process_response_headers as _process_response_headers, + ) + _globals["process_response_headers"] = _process_response_headers + return _globals["process_response_headers"] + + # Lazy load dot_notation_indexing functions to avoid loading at module import time + if name == "delete_nested_value": + # Check if already cached + if "delete_nested_value" not in _globals: + from litellm.litellm_core_utils.dot_notation_indexing import ( + delete_nested_value as _delete_nested_value, + ) + _globals["delete_nested_value"] = _delete_nested_value + return _globals["delete_nested_value"] + + if name == "is_nested_path": + # Check if already cached + if "is_nested_path" not in _globals: + from litellm.litellm_core_utils.dot_notation_indexing import ( + is_nested_path as _is_nested_path, + ) + _globals["is_nested_path"] = _is_nested_path + return _globals["is_nested_path"] + + # Lazy load get_litellm_params functions to avoid loading at module import time + if name == "_get_base_model_from_litellm_call_metadata": + # Check if already cached + if "_get_base_model_from_litellm_call_metadata" not in _globals: + from litellm.litellm_core_utils.get_litellm_params import ( + _get_base_model_from_litellm_call_metadata as __get_base_model_from_litellm_call_metadata, + ) + _globals["_get_base_model_from_litellm_call_metadata"] = __get_base_model_from_litellm_call_metadata + return _globals["_get_base_model_from_litellm_call_metadata"] + + if name == "get_litellm_params": + # Check if already cached + if "get_litellm_params" not in _globals: + from litellm.litellm_core_utils.get_litellm_params import ( + get_litellm_params as _get_litellm_params, + ) + _globals["get_litellm_params"] = _get_litellm_params + return _globals["get_litellm_params"] + + # Lazy load _ensure_extra_body_is_safe to avoid loading at module import time + if name == "_ensure_extra_body_is_safe": + # Check if already cached + if "_ensure_extra_body_is_safe" not in _globals: + from litellm.litellm_core_utils.llm_request_utils import ( + _ensure_extra_body_is_safe as __ensure_extra_body_is_safe, + ) + _globals["_ensure_extra_body_is_safe"] = __ensure_extra_body_is_safe + return _globals["_ensure_extra_body_is_safe"] + + # Lazy load get_formatted_prompt to avoid loading at module import time + if name == "get_formatted_prompt": + # Check if already cached + if "get_formatted_prompt" not in _globals: + from litellm.litellm_core_utils.llm_response_utils.get_formatted_prompt import ( + get_formatted_prompt as _get_formatted_prompt, + ) + _globals["get_formatted_prompt"] = _get_formatted_prompt + return _globals["get_formatted_prompt"] + + # Lazy load get_response_headers to avoid loading at module import time + if name == "get_response_headers": + # Check if already cached + if "get_response_headers" not in _globals: + from litellm.litellm_core_utils.llm_response_utils.get_headers import ( + get_response_headers as _get_response_headers, + ) + _globals["get_response_headers"] = _get_response_headers + return _globals["get_response_headers"] + + # Lazy load update_response_metadata to avoid loading at module import time + if name == "update_response_metadata": + # Check if already cached + if "update_response_metadata" not in _globals: + from litellm.litellm_core_utils.llm_response_utils.response_metadata import ( + update_response_metadata as _update_response_metadata, + ) + _globals["update_response_metadata"] = _update_response_metadata + return _globals["update_response_metadata"] + + # Lazy load executor to avoid loading at module import time + if name == "executor": + # Check if already cached + if "executor" not in _globals: + from litellm.litellm_core_utils.thread_pool_executor import ( + executor as _executor, + ) + _globals["executor"] = _executor + return _globals["executor"] + + # Lazy load BaseAnthropicMessagesConfig to avoid loading at module import time + if name == "BaseAnthropicMessagesConfig": + # Check if already cached + if "BaseAnthropicMessagesConfig" not in _globals: + from litellm.llms.base_llm.anthropic_messages.transformation import ( + BaseAnthropicMessagesConfig as _BaseAnthropicMessagesConfig, + ) + _globals["BaseAnthropicMessagesConfig"] = _BaseAnthropicMessagesConfig + return _globals["BaseAnthropicMessagesConfig"] + + # Lazy load BaseAudioTranscriptionConfig to avoid loading at module import time + if name == "BaseAudioTranscriptionConfig": + # Check if already cached + if "BaseAudioTranscriptionConfig" not in _globals: + from litellm.llms.base_llm.audio_transcription.transformation import ( + BaseAudioTranscriptionConfig as _BaseAudioTranscriptionConfig, + ) + _globals["BaseAudioTranscriptionConfig"] = _BaseAudioTranscriptionConfig + return _globals["BaseAudioTranscriptionConfig"] + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") From b6d601c2f02d17ae9409366f24320e6cef70f97b Mon Sep 17 00:00:00 2001 From: Alexsander Hamir Date: Sat, 3 Jan 2026 16:16:17 -0800 Subject: [PATCH 149/158] perf(utils): lazy load 15+ unused imports (#18616) - Move BaseBatchesConfig, BaseContainerConfig, BaseEmbeddingConfig, BaseImageEditConfig, BaseImageGenerationConfig, BaseImageVariationConfig, BasePassthroughConfig, BaseRealtimeConfig, BaseRerankConfig, BaseVectorStoreConfig, BaseVectorStoreFilesConfig, BaseVideoConfig to lazy loading - Move ANTHROPIC_API_ONLY_HEADERS, AnthropicThinkingParam, RerankResponse to lazy loading - Move ChatCompletionDeltaToolCallChunk, ChatCompletionToolCallChunk, ChatCompletionToolCallFunctionChunk, LiteLLM_Params to lazy loading - Add type stubs to TYPE_CHECKING block for mypy support - Add lazy loading handlers in __getattr__ method - These imports are not used in utils.py runtime code, only in type annotations (safe with from __future__ import annotations) --- litellm/utils.py | 245 +++++++++++++++++++++++++++++++++++++++++------ 1 file changed, 216 insertions(+), 29 deletions(-) diff --git a/litellm/utils.py b/litellm/utils.py index 6f4652c8278..df0b2317123 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -149,10 +149,6 @@ def _get_cached_audio_utils(): _audio_utils_module = litellm.litellm_core_utils.audio_utils.utils return _audio_utils_module -from litellm.types.llms.anthropic import ( - ANTHROPIC_API_ONLY_HEADERS, - AnthropicThinkingParam, -) from litellm.types.llms.openai import ( AllMessageValues, AllPromptValues, @@ -163,7 +159,6 @@ from litellm.types.llms.openai import ( OpenAITextCompletionUserMessage, OpenAIWebSearchOptions, ) -from litellm.types.rerank import RerankResponse from litellm.types.utils import FileTypes # type: ignore from litellm.types.utils import ( OPENAI_RESPONSE_HEADERS, @@ -342,29 +337,41 @@ if TYPE_CHECKING: reset_retry_policy, ) from litellm.secret_managers.main import get_secret + # Type stubs for lazy-loaded config classes and types + from litellm.llms.base_llm.batches.transformation import BaseBatchesConfig + from litellm.llms.base_llm.containers.transformation import BaseContainerConfig + from litellm.llms.base_llm.embedding.transformation import BaseEmbeddingConfig + from litellm.llms.base_llm.image_edit.transformation import BaseImageEditConfig + from litellm.llms.base_llm.image_generation.transformation import ( + BaseImageGenerationConfig, + ) + from litellm.llms.base_llm.image_variations.transformation import ( + BaseImageVariationConfig, + ) + from litellm.llms.base_llm.passthrough.transformation import BasePassthroughConfig + from litellm.llms.base_llm.realtime.transformation import BaseRealtimeConfig + from litellm.llms.base_llm.rerank.transformation import BaseRerankConfig + from litellm.llms.base_llm.vector_store.transformation import BaseVectorStoreConfig + from litellm.llms.base_llm.vector_store_files.transformation import ( + BaseVectorStoreFilesConfig, + ) + from litellm.llms.base_llm.videos.transformation import BaseVideoConfig + from litellm.types.llms.anthropic import ( + ANTHROPIC_API_ONLY_HEADERS, + AnthropicThinkingParam, + ) + from litellm.types.rerank import RerankResponse + from litellm.types.llms.openai import ( + ChatCompletionDeltaToolCallChunk, + ChatCompletionToolCallChunk, + ChatCompletionToolCallFunctionChunk, + ) + from litellm.types.router import LiteLLM_Params -from litellm.llms.base_llm.batches.transformation import BaseBatchesConfig from litellm.llms.base_llm.chat.transformation import BaseConfig from litellm.llms.base_llm.completion.transformation import BaseTextCompletionConfig -from litellm.llms.base_llm.containers.transformation import BaseContainerConfig -from litellm.llms.base_llm.embedding.transformation import BaseEmbeddingConfig -from litellm.llms.base_llm.image_edit.transformation import BaseImageEditConfig -from litellm.llms.base_llm.image_generation.transformation import ( - BaseImageGenerationConfig, -) -from litellm.llms.base_llm.image_variations.transformation import ( - BaseImageVariationConfig, -) -from litellm.llms.base_llm.passthrough.transformation import BasePassthroughConfig -from litellm.llms.base_llm.realtime.transformation import BaseRealtimeConfig -from litellm.llms.base_llm.rerank.transformation import BaseRerankConfig from litellm.llms.base_llm.responses.transformation import BaseResponsesAPIConfig from litellm.llms.base_llm.skills.transformation import BaseSkillsAPIConfig -from litellm.llms.base_llm.vector_store.transformation import BaseVectorStoreConfig -from litellm.llms.base_llm.vector_store_files.transformation import ( - BaseVectorStoreFilesConfig, -) -from litellm.llms.base_llm.videos.transformation import BaseVideoConfig from ._logging import _is_debugging_on, verbose_logger from .caching.caching import ( @@ -392,12 +399,6 @@ from .exceptions import ( UnprocessableEntityError, UnsupportedParamsError, ) -from .types.llms.openai import ( - ChatCompletionDeltaToolCallChunk, - ChatCompletionToolCallChunk, - ChatCompletionToolCallFunctionChunk, -) -from .types.router import LiteLLM_Params if TYPE_CHECKING: from litellm import MockException @@ -9226,4 +9227,190 @@ def __getattr__(name: str) -> Any: # noqa: PLR0915 _globals["BaseAudioTranscriptionConfig"] = _BaseAudioTranscriptionConfig return _globals["BaseAudioTranscriptionConfig"] + # Lazy load BaseBatchesConfig to avoid loading at module import time + if name == "BaseBatchesConfig": + # Check if already cached + if "BaseBatchesConfig" not in _globals: + from litellm.llms.base_llm.batches.transformation import ( + BaseBatchesConfig as _BaseBatchesConfig, + ) + _globals["BaseBatchesConfig"] = _BaseBatchesConfig + return _globals["BaseBatchesConfig"] + + # Lazy load BaseContainerConfig to avoid loading at module import time + if name == "BaseContainerConfig": + # Check if already cached + if "BaseContainerConfig" not in _globals: + from litellm.llms.base_llm.containers.transformation import ( + BaseContainerConfig as _BaseContainerConfig, + ) + _globals["BaseContainerConfig"] = _BaseContainerConfig + return _globals["BaseContainerConfig"] + + # Lazy load BaseEmbeddingConfig to avoid loading at module import time + if name == "BaseEmbeddingConfig": + # Check if already cached + if "BaseEmbeddingConfig" not in _globals: + from litellm.llms.base_llm.embedding.transformation import ( + BaseEmbeddingConfig as _BaseEmbeddingConfig, + ) + _globals["BaseEmbeddingConfig"] = _BaseEmbeddingConfig + return _globals["BaseEmbeddingConfig"] + + # Lazy load BaseImageEditConfig to avoid loading at module import time + if name == "BaseImageEditConfig": + # Check if already cached + if "BaseImageEditConfig" not in _globals: + from litellm.llms.base_llm.image_edit.transformation import ( + BaseImageEditConfig as _BaseImageEditConfig, + ) + _globals["BaseImageEditConfig"] = _BaseImageEditConfig + return _globals["BaseImageEditConfig"] + + # Lazy load BaseImageGenerationConfig to avoid loading at module import time + if name == "BaseImageGenerationConfig": + # Check if already cached + if "BaseImageGenerationConfig" not in _globals: + from litellm.llms.base_llm.image_generation.transformation import ( + BaseImageGenerationConfig as _BaseImageGenerationConfig, + ) + _globals["BaseImageGenerationConfig"] = _BaseImageGenerationConfig + return _globals["BaseImageGenerationConfig"] + + # Lazy load BaseImageVariationConfig to avoid loading at module import time + if name == "BaseImageVariationConfig": + # Check if already cached + if "BaseImageVariationConfig" not in _globals: + from litellm.llms.base_llm.image_variations.transformation import ( + BaseImageVariationConfig as _BaseImageVariationConfig, + ) + _globals["BaseImageVariationConfig"] = _BaseImageVariationConfig + return _globals["BaseImageVariationConfig"] + + # Lazy load BasePassthroughConfig to avoid loading at module import time + if name == "BasePassthroughConfig": + # Check if already cached + if "BasePassthroughConfig" not in _globals: + from litellm.llms.base_llm.passthrough.transformation import ( + BasePassthroughConfig as _BasePassthroughConfig, + ) + _globals["BasePassthroughConfig"] = _BasePassthroughConfig + return _globals["BasePassthroughConfig"] + + # Lazy load BaseRealtimeConfig to avoid loading at module import time + if name == "BaseRealtimeConfig": + # Check if already cached + if "BaseRealtimeConfig" not in _globals: + from litellm.llms.base_llm.realtime.transformation import ( + BaseRealtimeConfig as _BaseRealtimeConfig, + ) + _globals["BaseRealtimeConfig"] = _BaseRealtimeConfig + return _globals["BaseRealtimeConfig"] + + # Lazy load BaseRerankConfig to avoid loading at module import time + if name == "BaseRerankConfig": + # Check if already cached + if "BaseRerankConfig" not in _globals: + from litellm.llms.base_llm.rerank.transformation import ( + BaseRerankConfig as _BaseRerankConfig, + ) + _globals["BaseRerankConfig"] = _BaseRerankConfig + return _globals["BaseRerankConfig"] + + # Lazy load BaseVectorStoreConfig to avoid loading at module import time + if name == "BaseVectorStoreConfig": + # Check if already cached + if "BaseVectorStoreConfig" not in _globals: + from litellm.llms.base_llm.vector_store.transformation import ( + BaseVectorStoreConfig as _BaseVectorStoreConfig, + ) + _globals["BaseVectorStoreConfig"] = _BaseVectorStoreConfig + return _globals["BaseVectorStoreConfig"] + + # Lazy load BaseVectorStoreFilesConfig to avoid loading at module import time + if name == "BaseVectorStoreFilesConfig": + # Check if already cached + if "BaseVectorStoreFilesConfig" not in _globals: + from litellm.llms.base_llm.vector_store_files.transformation import ( + BaseVectorStoreFilesConfig as _BaseVectorStoreFilesConfig, + ) + _globals["BaseVectorStoreFilesConfig"] = _BaseVectorStoreFilesConfig + return _globals["BaseVectorStoreFilesConfig"] + + # Lazy load BaseVideoConfig to avoid loading at module import time + if name == "BaseVideoConfig": + # Check if already cached + if "BaseVideoConfig" not in _globals: + from litellm.llms.base_llm.videos.transformation import ( + BaseVideoConfig as _BaseVideoConfig, + ) + _globals["BaseVideoConfig"] = _BaseVideoConfig + return _globals["BaseVideoConfig"] + + # Lazy load ANTHROPIC_API_ONLY_HEADERS to avoid loading at module import time + if name == "ANTHROPIC_API_ONLY_HEADERS": + # Check if already cached + if "ANTHROPIC_API_ONLY_HEADERS" not in _globals: + from litellm.types.llms.anthropic import ( + ANTHROPIC_API_ONLY_HEADERS as _ANTHROPIC_API_ONLY_HEADERS, + ) + _globals["ANTHROPIC_API_ONLY_HEADERS"] = _ANTHROPIC_API_ONLY_HEADERS + return _globals["ANTHROPIC_API_ONLY_HEADERS"] + + # Lazy load AnthropicThinkingParam to avoid loading at module import time + if name == "AnthropicThinkingParam": + # Check if already cached + if "AnthropicThinkingParam" not in _globals: + from litellm.types.llms.anthropic import ( + AnthropicThinkingParam as _AnthropicThinkingParam, + ) + _globals["AnthropicThinkingParam"] = _AnthropicThinkingParam + return _globals["AnthropicThinkingParam"] + + # Lazy load RerankResponse to avoid loading at module import time + if name == "RerankResponse": + # Check if already cached + if "RerankResponse" not in _globals: + from litellm.types.rerank import RerankResponse as _RerankResponse + _globals["RerankResponse"] = _RerankResponse + return _globals["RerankResponse"] + + # Lazy load ChatCompletionDeltaToolCallChunk to avoid loading at module import time + if name == "ChatCompletionDeltaToolCallChunk": + # Check if already cached + if "ChatCompletionDeltaToolCallChunk" not in _globals: + from litellm.types.llms.openai import ( + ChatCompletionDeltaToolCallChunk as _ChatCompletionDeltaToolCallChunk, + ) + _globals["ChatCompletionDeltaToolCallChunk"] = _ChatCompletionDeltaToolCallChunk + return _globals["ChatCompletionDeltaToolCallChunk"] + + # Lazy load ChatCompletionToolCallChunk to avoid loading at module import time + if name == "ChatCompletionToolCallChunk": + # Check if already cached + if "ChatCompletionToolCallChunk" not in _globals: + from litellm.types.llms.openai import ( + ChatCompletionToolCallChunk as _ChatCompletionToolCallChunk, + ) + _globals["ChatCompletionToolCallChunk"] = _ChatCompletionToolCallChunk + return _globals["ChatCompletionToolCallChunk"] + + # Lazy load ChatCompletionToolCallFunctionChunk to avoid loading at module import time + if name == "ChatCompletionToolCallFunctionChunk": + # Check if already cached + if "ChatCompletionToolCallFunctionChunk" not in _globals: + from litellm.types.llms.openai import ( + ChatCompletionToolCallFunctionChunk as _ChatCompletionToolCallFunctionChunk, + ) + _globals["ChatCompletionToolCallFunctionChunk"] = _ChatCompletionToolCallFunctionChunk + return _globals["ChatCompletionToolCallFunctionChunk"] + + # Lazy load LiteLLM_Params to avoid loading at module import time + if name == "LiteLLM_Params": + # Check if already cached + if "LiteLLM_Params" not in _globals: + from litellm.types.router import LiteLLM_Params as _LiteLLM_Params + _globals["LiteLLM_Params"] = _LiteLLM_Params + return _globals["LiteLLM_Params"] + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") From d6296411ec569516a0dabf92ca79b495546146a4 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Sat, 3 Jan 2026 16:23:10 -0800 Subject: [PATCH 150/158] SSO Settings Loading, deprecate previous flow --- .../AdminSettings/SSOSettings/SSOSettings.tsx | 106 ++++++++++-------- .../SSOSettingsLoadingSkeleton.tsx | 66 +++++++++++ .../src/components/admins.tsx | 8 +- 3 files changed, 130 insertions(+), 50 deletions(-) create mode 100644 ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/SSOSettingsLoadingSkeleton.tsx diff --git a/ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/SSOSettings.tsx b/ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/SSOSettings.tsx index d339f6d0e36..27ff96af05f 100644 --- a/ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/SSOSettings.tsx +++ b/ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/SSOSettings.tsx @@ -10,12 +10,13 @@ import DeleteSSOSettingsModal from "./Modals/DeleteSSOSettingsModal"; import EditSSOSettingsModal from "./Modals/EditSSOSettingsModal"; import RedactableField from "./RedactableField"; import SSOSettingsEmptyPlaceholder from "./SSOSettingsEmptyPlaceholder"; +import SSOSettingsLoadingSkeleton from "./SSOSettingsLoadingSkeleton"; import { ssoProviderDisplayNames, ssoProviderLogoMap } from "./constants"; const { Title, Text } = Typography; export default function SSOSettings() { - const { data: ssoSettings, refetch } = useSSOSettings(); + const { data: ssoSettings, refetch, isLoading } = useSSOSettings(); const { accessToken } = useAuthorized(); const [isDeleteModalVisible, setIsDeleteModalVisible] = useState(false); const [isAddModalVisible, setIsAddModalVisible] = useState(false); @@ -26,27 +27,28 @@ export default function SSOSettings() { Boolean(ssoSettings?.values.generic_client_id); // Determine the SSO provider based on the configuration - let selectedProvider: string | null = null; - if (ssoSettings?.values.google_client_id) { - selectedProvider = "google"; - } else if (ssoSettings?.values.microsoft_client_id) { - selectedProvider = "microsoft"; - } else if (ssoSettings?.values.generic_client_id) { - // Check if it looks like Okta based on endpoints - if ( - ssoSettings.values.generic_authorization_endpoint?.includes("okta") || - ssoSettings.values.generic_authorization_endpoint?.includes("auth0") - ) { - selectedProvider = "okta"; - } else { - selectedProvider = "generic"; + const detectSSOProvider = (values: SSOSettingsValues): string | null => { + if (values.google_client_id) return "google"; + if (values.microsoft_client_id) return "microsoft"; + if (values.generic_client_id) { + // Check if it looks like Okta/Auth0 based on endpoints + if ( + values.generic_authorization_endpoint?.includes("okta") || + values.generic_authorization_endpoint?.includes("auth0") + ) { + return "okta"; + } + return "generic"; } - } + return null; + }; + + const selectedProvider = ssoSettings?.values ? detectSSOProvider(ssoSettings.values) : null; const renderEndpointValue = (value?: string | null) => ( - - {value || Not configured} - + + {value || "-"} + ); const renderSimpleValue = (value?: string | null) => @@ -179,38 +181,44 @@ export default function SSOSettings() { }; return ( - - - {/* Header Section */} -
-
- -
- SSO Configuration - Manage Single Sign-On authentication settings + <> + {isLoading ? ( + + ) : ( + + + {/* Header Section */} +
+
+ +
+ SSO Configuration + Manage Single Sign-On authentication settings +
+
+ +
+ {isSSOConfigured && ( + <> + + + + )} +
-
-
- {isSSOConfigured && ( - <> - - - + {isSSOConfigured ? ( + renderSSOSettings() + ) : ( + setIsAddModalVisible(true)} /> )} -
-
- - {isSSOConfigured ? ( - renderSSOSettings() - ) : ( - setIsAddModalVisible(true)} /> - )} - + + + )} - + ); } diff --git a/ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/SSOSettingsLoadingSkeleton.tsx b/ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/SSOSettingsLoadingSkeleton.tsx new file mode 100644 index 00000000000..59e34f255e3 --- /dev/null +++ b/ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/SSOSettingsLoadingSkeleton.tsx @@ -0,0 +1,66 @@ +"use client"; + +import { Card, Descriptions, Skeleton, Space, Typography } from "antd"; +import { Shield } from "lucide-react"; + +const { Title, Text } = Typography; +export default function SSOSettingsLoadingSkeleton() { + const descriptionsConfig = { + column: { + xxl: 1, + xl: 1, + lg: 1, + md: 1, + sm: 1, + xs: 1, + }, + }; + + return ( + + + {/* Header Section */} +
+
+ +
+ SSO Configuration + Manage Single Sign-On authentication settings +
+
+ +
+ + +
+
+ + {/* Descriptions Table Skeleton */} + + {/* Provider Row */} + }> +
+ +
+ + + }> + + + + }> + + + + }> + + + + }> + + + + + + ); +} diff --git a/ui/litellm-dashboard/src/components/admins.tsx b/ui/litellm-dashboard/src/components/admins.tsx index 6af5a226da6..9de971bcd62 100644 --- a/ui/litellm-dashboard/src/components/admins.tsx +++ b/ui/litellm-dashboard/src/components/admins.tsx @@ -3,7 +3,7 @@ * Use this to avoid sharing master key with others */ import React, { useState, useEffect } from "react"; -import { Typography } from "antd"; +import { Alert, Typography } from "antd"; import { useRouter } from "next/navigation"; import { Button as Button2, Modal, Form, Input } from "antd"; import { Select, SelectItem } from "@tremor/react"; @@ -509,6 +509,12 @@ const AdminPanel: React.FC = ({ ✨ Security Settings +
Date: Sat, 3 Jan 2026 17:25:46 -0800 Subject: [PATCH 151/158] Adding unit testing coverage --- .../Modals/EditSSOSettingsModal.test.tsx | 620 ++++++++++++++++++ .../SSOSettingsLoadingSkeleton.test.tsx | 222 +++++++ .../VectorStoreSelector.test.tsx | 524 +++++++++++++++ .../VectorStoreTable.test.tsx | 415 ++++++++++++ .../src/utils/cookieUtils.test.ts | 72 ++ .../src/utils/proxyUtils.test.ts | 78 +++ .../src/utils/textUtils.test.ts | 21 + 7 files changed, 1952 insertions(+) create mode 100644 ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/Modals/EditSSOSettingsModal.test.tsx create mode 100644 ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/SSOSettingsLoadingSkeleton.test.tsx create mode 100644 ui/litellm-dashboard/src/components/vector_store_management/VectorStoreSelector.test.tsx create mode 100644 ui/litellm-dashboard/src/components/vector_store_management/VectorStoreTable.test.tsx create mode 100644 ui/litellm-dashboard/src/utils/proxyUtils.test.ts diff --git a/ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/Modals/EditSSOSettingsModal.test.tsx b/ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/Modals/EditSSOSettingsModal.test.tsx new file mode 100644 index 00000000000..559d837b409 --- /dev/null +++ b/ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/Modals/EditSSOSettingsModal.test.tsx @@ -0,0 +1,620 @@ +import { render, screen, fireEvent, waitFor } from "@testing-library/react"; +import { describe, it, expect, vi, beforeEach, Mock } from "vitest"; +import EditSSOSettingsModal from "./EditSSOSettingsModal"; +import { useSSOSettings } from "@/app/(dashboard)/hooks/sso/useSSOSettings"; +import { useEditSSOSettings } from "@/app/(dashboard)/hooks/sso/useEditSSOSettings"; +import NotificationsManager from "@/components/molecules/notifications_manager"; +import { parseErrorMessage } from "@/components/shared/errorUtils"; +import { processSSOSettingsPayload } from "../utils"; + +// Constants +const SSO_PROVIDERS = { + GOOGLE: "google", + MICROSOFT: "microsoft", + OKTA: "okta", + AUTH0: "auth0", + GENERIC: "generic", +} as const; + +const TEST_DATA = { + MODAL_TITLE: "Edit SSO Settings", + MODAL_WIDTH: "800", + SUCCESS_MESSAGE: "SSO settings updated successfully", + ERROR_MESSAGE_PREFIX: "Failed to save SSO settings:", + BUTTON_TEXT: { + CANCEL: "Cancel", + SAVE: "Save", + SAVING: "Saving...", + }, +} as const; + +const TEST_IDS = { + MODAL: "modal", + BUTTON: "button", + BASE_SSO_FORM: "base-sso-form", + TRIGGER_FORM_SUBMIT: "trigger-form-submit", +} as const; + +// Mock form instance +const mockForm = { + resetFields: vi.fn(), + setFieldsValue: vi.fn(), + getFieldsValue: vi.fn(), + submit: vi.fn(), +}; + +// Types +type SSOData = { + values: Record; +} & Record; + +type SSOSettingsHookReturn = { + data: SSOData | null; + isLoading: boolean; + error: any; +}; + +type EditSSOSettingsHookReturn = { + mutateAsync: ReturnType; + isPending: boolean; +}; + +// Test data factories +const createSSOData = (overrides: Record = {}): SSOData => ({ + values: { + user_email: "test@example.com", + ...overrides, + }, +}); + +const createGoogleSSOData = (overrides: Record = {}) => + createSSOData({ + google_client_id: "test-google-id", + google_client_secret: "test-google-secret", + ...overrides, + }); + +const createMicrosoftSSOData = (overrides: Record = {}) => + createSSOData({ + microsoft_client_id: "test-microsoft-id", + microsoft_client_secret: "test-microsoft-secret", + microsoft_tenant: "test-tenant", + ...overrides, + }); + +const createGenericSSOData = (overrides: Record = {}) => + createSSOData({ + generic_client_id: "test-generic-id", + generic_client_secret: "test-generic-secret", + generic_authorization_endpoint: overrides.authorization_endpoint || "https://custom.example.com/oauth", + ...overrides, + }); + +const createRoleMappingsSSOData = (overrides: Record = {}) => + createGoogleSSOData({ + role_mappings: { + group_claim: "groups", + default_role: "internal_user", + roles: { + proxy_admin: overrides.proxy_admin || ["admin-group"], + proxy_admin_viewer: overrides.proxy_admin_viewer || ["viewer-group"], + internal_user: overrides.internal_user || ["user-group"], + internal_user_viewer: overrides.internal_user_viewer || ["readonly-group"], + }, + }, + ...overrides, + }); + +// Mock utilities +const createMockHooks = (): { + useSSOSettings: SSOSettingsHookReturn; + useEditSSOSettings: EditSSOSettingsHookReturn; +} => ({ + useSSOSettings: { + data: null, + isLoading: false, + error: null, + }, + useEditSSOSettings: { + mutateAsync: vi.fn(), + isPending: false, + }, +}); + +vi.mock("antd", () => ({ + Modal: ({ children, open, title, footer, onCancel, width, ...props }: any) => ( +
+
{children}
+
{footer}
+
+ ), + Button: ({ children, onClick, loading, disabled, ...props }: any) => ( + + ), + Form: { + useForm: () => [mockForm], + }, + Space: ({ children, ...props }: any) => ( +
+ {children} +
+ ), +})); + +vi.mock("./BaseSSOSettingsForm", () => ({ + default: ({ form, onFormSubmit }: any) => ( +
+ +
+ ), +})); + +vi.mock("@/app/(dashboard)/hooks/sso/useSSOSettings", () => ({ + useSSOSettings: vi.fn(), +})); + +vi.mock("@/app/(dashboard)/hooks/sso/useEditSSOSettings", () => ({ + useEditSSOSettings: vi.fn(), +})); + +vi.mock("@/components/molecules/notifications_manager", () => ({ + default: { + success: vi.fn(), + fromBackend: vi.fn(), + }, +})); + +vi.mock("@/components/shared/errorUtils", () => ({ + parseErrorMessage: vi.fn(), +})); + +vi.mock("../utils", () => ({ + processSSOSettingsPayload: vi.fn(), +})); + +// Test helpers +const setupMocks = ( + overrides: Partial<{ + useSSOSettings: Partial; + useEditSSOSettings: Partial; + }> = {}, +) => { + const defaultMocks = createMockHooks(); + const mocks = { + useSSOSettings: { ...defaultMocks.useSSOSettings, ...overrides.useSSOSettings }, + useEditSSOSettings: { ...defaultMocks.useEditSSOSettings, ...overrides.useEditSSOSettings }, + }; + + (useSSOSettings as Mock).mockReturnValue(mocks.useSSOSettings); + (useEditSSOSettings as Mock).mockReturnValue(mocks.useEditSSOSettings); + + return mocks; +}; + +const renderComponent = (props: Partial> = {}) => { + const defaultProps = { + isVisible: true, + onCancel: vi.fn(), + onSuccess: vi.fn(), + }; + + return { + ...render(), + mockOnCancel: defaultProps.onCancel, + mockOnSuccess: defaultProps.onSuccess, + }; +}; + +const getButtons = () => screen.getAllByTestId(TEST_IDS.BUTTON); +const getCancelButton = () => getButtons()[0]; +const getSaveButton = () => getButtons()[1]; + +describe("EditSSOSettingsModal", () => { + beforeEach(() => { + vi.clearAllMocks(); + setupMocks(); + }); + + describe("Rendering", () => { + it("renders without crashing", () => { + expect(() => renderComponent()).not.toThrow(); + }); + + it("displays modal with correct configuration", () => { + renderComponent(); + + const modal = screen.getByTestId(TEST_IDS.MODAL); + expect(modal).toHaveAttribute("data-open", "true"); + expect(modal).toHaveAttribute("data-title", TEST_DATA.MODAL_TITLE); + expect(modal).toHaveAttribute("data-width", TEST_DATA.MODAL_WIDTH); + }); + + it("displays modal as closed when not visible", () => { + renderComponent({ isVisible: false }); + + const modal = screen.getByTestId(TEST_IDS.MODAL); + expect(modal).toHaveAttribute("data-open", "false"); + }); + }); + + describe("Footer Actions", () => { + it("renders cancel and save buttons", () => { + renderComponent(); + + const buttons = getButtons(); + expect(buttons).toHaveLength(2); + expect(buttons[0]).toHaveTextContent(TEST_DATA.BUTTON_TEXT.CANCEL); + expect(buttons[1]).toHaveTextContent(TEST_DATA.BUTTON_TEXT.SAVE); + }); + + it("calls onCancel and resets form when cancel button is clicked", () => { + const { mockOnCancel } = renderComponent(); + + fireEvent.click(getCancelButton()); + + expect(mockForm.resetFields).toHaveBeenCalled(); + expect(mockOnCancel).toHaveBeenCalled(); + }); + + it("calls form.submit when save button is clicked", () => { + renderComponent(); + + fireEvent.click(getSaveButton()); + + expect(mockForm.submit).toHaveBeenCalled(); + }); + + describe("Loading States", () => { + it("disables cancel button during submission", () => { + setupMocks({ + useEditSSOSettings: { mutateAsync: vi.fn(), isPending: true }, + }); + + renderComponent(); + + expect(getCancelButton()).toBeDisabled(); + }); + + it("shows loading state on save button during submission", () => { + setupMocks({ + useEditSSOSettings: { mutateAsync: vi.fn(), isPending: true }, + }); + + renderComponent(); + + expect(getSaveButton()).toHaveAttribute("data-loading", "true"); + expect(getSaveButton()).toHaveTextContent(TEST_DATA.BUTTON_TEXT.SAVING); + }); + }); + }); + + describe("Form Submission", () => { + const formValues = { testField: "testValue" }; + const processedPayload = { processed: "payload" }; + + beforeEach(() => { + (processSSOSettingsPayload as any).mockReturnValue(processedPayload); + }); + + it("processes form values and submits successfully", async () => { + const mockMutateAsync = vi.fn().mockImplementation((payload, options) => { + options.onSuccess(); + return Promise.resolve({ success: true }); + }); + + setupMocks({ + useEditSSOSettings: { mutateAsync: mockMutateAsync, isPending: false }, + }); + + const { mockOnSuccess } = renderComponent(); + + fireEvent.click(screen.getByTestId(TEST_IDS.TRIGGER_FORM_SUBMIT)); + + expect(processSSOSettingsPayload).toHaveBeenCalledWith(formValues); + expect(mockMutateAsync).toHaveBeenCalledWith( + processedPayload, + expect.objectContaining({ + onSuccess: expect.any(Function), + onError: expect.any(Function), + }), + ); + }); + + it("shows success notification and calls onSuccess callback", async () => { + const mockMutateAsync = vi.fn().mockImplementation((payload, options) => { + options.onSuccess(); + return Promise.resolve({ success: true }); + }); + + setupMocks({ + useEditSSOSettings: { mutateAsync: mockMutateAsync, isPending: false }, + }); + + const { mockOnSuccess } = renderComponent(); + + fireEvent.click(screen.getByTestId(TEST_IDS.TRIGGER_FORM_SUBMIT)); + + expect(NotificationsManager.success).toHaveBeenCalledWith(TEST_DATA.SUCCESS_MESSAGE); + expect(mockOnSuccess).toHaveBeenCalled(); + }); + + it("handles submission errors gracefully", async () => { + const error = new Error("Submission failed"); + const mockMutateAsync = vi.fn().mockImplementation((payload, options) => { + options.onError(error); + return Promise.reject(error); + }); + + setupMocks({ + useEditSSOSettings: { mutateAsync: mockMutateAsync, isPending: false }, + }); + + (parseErrorMessage as any).mockReturnValue("Parsed error message"); + + renderComponent(); + + fireEvent.click(screen.getByTestId(TEST_IDS.TRIGGER_FORM_SUBMIT)); + + expect(parseErrorMessage).toHaveBeenCalledWith(error); + expect(NotificationsManager.fromBackend).toHaveBeenCalledWith( + `${TEST_DATA.ERROR_MESSAGE_PREFIX} Parsed error message`, + ); + }); + }); + + describe("Form Initialization", () => { + describe("Provider Detection", () => { + const testProviderDetection = (testName: string, ssoData: SSOData, expectedProvider: string) => { + it(`detects ${testName} provider`, async () => { + setupMocks({ + useSSOSettings: { data: ssoData, isLoading: false, error: null }, + }); + + renderComponent(); + + await waitFor(() => { + expect(mockForm.setFieldsValue).toHaveBeenCalledWith({ + sso_provider: expectedProvider, + ...ssoData.values, + }); + }); + }); + }; + + testProviderDetection("Google", createGoogleSSOData(), SSO_PROVIDERS.GOOGLE); + + testProviderDetection("Microsoft", createMicrosoftSSOData(), SSO_PROVIDERS.MICROSOFT); + + testProviderDetection( + "Okta", + createGenericSSOData({ + authorization_endpoint: "https://okta.example.com/oauth2/authorize", + }), + SSO_PROVIDERS.OKTA, + ); + + testProviderDetection( + "Auth0 (detected as Okta)", + createGenericSSOData({ + authorization_endpoint: "https://auth0.example.com/authorize", + }), + SSO_PROVIDERS.OKTA, // Auth0 URLs are detected as Okta provider + ); + + testProviderDetection("generic", createGenericSSOData(), SSO_PROVIDERS.GENERIC); + }); + + describe("Role Mappings", () => { + it("processes role mappings with all roles assigned", async () => { + const ssoData = createRoleMappingsSSOData(); + + setupMocks({ + useSSOSettings: { data: ssoData, isLoading: false, error: null }, + }); + + renderComponent(); + + await waitFor(() => { + expect(mockForm.setFieldsValue).toHaveBeenCalledWith({ + sso_provider: SSO_PROVIDERS.GOOGLE, + ...ssoData.values, + use_role_mappings: true, + group_claim: "groups", + default_role: "internal_user", + proxy_admin_teams: "admin-group", + admin_viewer_teams: "viewer-group", + internal_user_teams: "user-group", + internal_viewer_teams: "readonly-group", + }); + }); + }); + + it("handles empty role mapping arrays", async () => { + const ssoData = createRoleMappingsSSOData({ + proxy_admin: [], + proxy_admin_viewer: [], + internal_user_viewer: [], + }); + + setupMocks({ + useSSOSettings: { data: ssoData, isLoading: false, error: null }, + }); + + renderComponent(); + + await waitFor(() => { + expect(mockForm.setFieldsValue).toHaveBeenCalledWith({ + sso_provider: SSO_PROVIDERS.GOOGLE, + ...ssoData.values, + use_role_mappings: true, + group_claim: "groups", + default_role: "internal_user", + proxy_admin_teams: "", + admin_viewer_teams: "", + internal_user_teams: "user-group", + internal_viewer_teams: "", + }); + }); + }); + }); + + describe("Initialization Guards", () => { + it("resets form before setting values", async () => { + const ssoData = createGoogleSSOData(); + + setupMocks({ + useSSOSettings: { data: ssoData, isLoading: false, error: null }, + }); + + renderComponent(); + + await waitFor(() => { + expect(mockForm.resetFields).toHaveBeenCalled(); + expect(mockForm.setFieldsValue).toHaveBeenCalled(); + }); + }); + + it("skips initialization when modal is not visible", () => { + const ssoData = createGoogleSSOData(); + + setupMocks({ + useSSOSettings: { data: ssoData, isLoading: false, error: null }, + }); + + renderComponent({ isVisible: false }); + + expect(mockForm.setFieldsValue).not.toHaveBeenCalled(); + }); + + it("skips initialization when SSO data is unavailable", () => { + setupMocks({ + useSSOSettings: { data: null, isLoading: false, error: null }, + }); + + renderComponent(); + + expect(mockForm.setFieldsValue).not.toHaveBeenCalled(); + }); + }); + }); + + describe("Error Handling", () => { + it("handles form submission errors with undefined error message", async () => { + const error = new Error("Network error"); + const mockMutateAsync = vi.fn().mockImplementation((payload, options) => { + options.onError(error); + return Promise.reject(error); + }); + + setupMocks({ + useEditSSOSettings: { mutateAsync: mockMutateAsync, isPending: false }, + }); + + (parseErrorMessage as any).mockReturnValue(undefined); + + renderComponent(); + + fireEvent.click(screen.getByTestId(TEST_IDS.TRIGGER_FORM_SUBMIT)); + + expect(NotificationsManager.fromBackend).toHaveBeenCalledWith(`${TEST_DATA.ERROR_MESSAGE_PREFIX} undefined`); + }); + + it("handles form submission with malformed data", async () => { + const mockMutateAsync = vi.fn().mockImplementation((payload, options) => { + options.onError(new Error("Invalid data")); + return Promise.reject(new Error("Invalid data")); + }); + + setupMocks({ + useEditSSOSettings: { mutateAsync: mockMutateAsync, isPending: false }, + }); + + (processSSOSettingsPayload as any).mockImplementation(() => { + throw new Error("Processing failed"); + }); + + renderComponent(); + + fireEvent.click(screen.getByTestId(TEST_IDS.TRIGGER_FORM_SUBMIT)); + + expect(processSSOSettingsPayload).toHaveBeenCalled(); + expect(mockMutateAsync).not.toHaveBeenCalled(); + }); + }); + + describe("Edge Cases", () => { + it("handles role mappings with undefined roles object", async () => { + const ssoData = createGoogleSSOData({ + role_mappings: { + group_claim: "groups", + default_role: "internal_user", + // roles is undefined + }, + }); + + setupMocks({ + useSSOSettings: { data: ssoData, isLoading: false, error: null }, + }); + + renderComponent(); + + await waitFor(() => { + expect(mockForm.setFieldsValue).toHaveBeenCalledWith({ + sso_provider: SSO_PROVIDERS.GOOGLE, + ...ssoData.values, + use_role_mappings: true, + group_claim: "groups", + default_role: "internal_user", + proxy_admin_teams: "", + admin_viewer_teams: "", + internal_user_teams: "", + internal_viewer_teams: "", + }); + }); + }); + + it("handles provider detection with partial SSO data", async () => { + const ssoData = createSSOData({ + // Only has generic fields, no specific provider identifiers + generic_client_id: "test-id", + generic_authorization_endpoint: "https://unknown.provider.com/auth", + }); + + setupMocks({ + useSSOSettings: { data: ssoData, isLoading: false, error: null }, + }); + + renderComponent(); + + await waitFor(() => { + expect(mockForm.setFieldsValue).toHaveBeenCalledWith({ + sso_provider: SSO_PROVIDERS.GENERIC, + ...ssoData.values, + }); + }); + }); + + it("handles form submission when processing throws error", async () => { + setupMocks({ + useEditSSOSettings: { mutateAsync: vi.fn(), isPending: false }, + }); + + (processSSOSettingsPayload as any).mockImplementation(() => { + throw new Error("Processing error"); + }); + + renderComponent(); + + expect(() => { + fireEvent.click(screen.getByTestId(TEST_IDS.TRIGGER_FORM_SUBMIT)); + }).not.toThrow(); + + expect(processSSOSettingsPayload).toHaveBeenCalled(); + }); + }); +}); diff --git a/ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/SSOSettingsLoadingSkeleton.test.tsx b/ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/SSOSettingsLoadingSkeleton.test.tsx new file mode 100644 index 00000000000..fd4fde69588 --- /dev/null +++ b/ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/SSOSettingsLoadingSkeleton.test.tsx @@ -0,0 +1,222 @@ +import { render, screen } from "@testing-library/react"; +import { describe, it, expect, vi } from "vitest"; +import SSOSettingsLoadingSkeleton from "./SSOSettingsLoadingSkeleton"; + +// Mock lucide-react icons +vi.mock("lucide-react", () => ({ + Shield: ({ className }: any) =>
, +})); + +// Mock Ant Design components +vi.mock("antd", () => ({ + Card: ({ children, ...props }: any) => ( +
+ {children} +
+ ), + Descriptions: Object.assign( + ({ children, bordered, column, ...props }: any) => ( +
+ {children} +
+ ), + { + Item: ({ children, label, ...props }: any) => ( +
+
{label}
+
{children}
+
+ ), + }, + ), + Typography: { + Title: ({ children, level, ...props }: any) => ( +
+ {children} +
+ ), + Text: ({ children, type, ...props }: any) => ( +
+ {children} +
+ ), + }, + Space: ({ children, direction, size, className, ...props }: any) => ( +
+ {children} +
+ ), + Skeleton: { + Button: ({ active, size, style, ...props }: any) => ( +
+ Button Skeleton +
+ ), + Node: ({ active, style, ...props }: any) => ( +
+ Node Skeleton +
+ ), + }, +})); + +describe("SSOSettingsLoadingSkeleton", () => { + it("should render without crashing", () => { + expect(() => render()).not.toThrow(); + }); + + it("should render Card component", () => { + render(); + expect(screen.getByTestId("card")).toBeInTheDocument(); + }); + + it("should render Space component with correct props", () => { + render(); + const space = screen.getByTestId("space"); + expect(space).toBeInTheDocument(); + expect(space).toHaveAttribute("data-direction", "vertical"); + expect(space).toHaveAttribute("data-size", "large"); + expect(space).toHaveClass("w-full"); + }); + + describe("Header Section", () => { + it("should render Shield icon", () => { + render(); + const shieldIcon = screen.getByTestId("shield-icon"); + expect(shieldIcon).toBeInTheDocument(); + expect(shieldIcon).toHaveClass("w-6 h-6 text-gray-400"); + }); + + it("should render title with correct text and level", () => { + render(); + const title = screen.getByTestId("typography-title"); + expect(title).toBeInTheDocument(); + expect(title).toHaveAttribute("data-level", "3"); + expect(title).toHaveTextContent("SSO Configuration"); + }); + + it("should render subtitle text", () => { + render(); + const text = screen.getByTestId("typography-text"); + expect(text).toBeInTheDocument(); + expect(text).toHaveAttribute("data-type", "secondary"); + expect(text).toHaveTextContent("Manage Single Sign-On authentication settings"); + }); + + it("should render two skeleton buttons with correct styles", () => { + render(); + const buttons = screen.getAllByTestId("skeleton-button"); + expect(buttons).toHaveLength(2); + + // First button + expect(buttons[0]).toHaveAttribute("data-active", "true"); + expect(buttons[0]).toHaveAttribute("data-size", "default"); + expect(buttons[0]).toHaveAttribute("data-style", JSON.stringify({ width: 170, height: 32 })); + + // Second button + expect(buttons[1]).toHaveAttribute("data-active", "true"); + expect(buttons[1]).toHaveAttribute("data-size", "default"); + expect(buttons[1]).toHaveAttribute("data-style", JSON.stringify({ width: 190, height: 32 })); + }); + }); + + describe("Descriptions Table", () => { + it("should render Descriptions component with bordered prop", () => { + render(); + const descriptions = screen.getByTestId("descriptions"); + expect(descriptions).toBeInTheDocument(); + expect(descriptions).toHaveAttribute("data-bordered", "true"); + }); + + it("should apply correct column configuration", () => { + render(); + const descriptions = screen.getByTestId("descriptions"); + const expectedColumn = { + xxl: 1, + xl: 1, + lg: 1, + md: 1, + sm: 1, + xs: 1, + }; + expect(descriptions).toHaveAttribute("data-column", JSON.stringify(expectedColumn)); + }); + + it("should render exactly 5 description items", () => { + render(); + const items = screen.getAllByTestId("descriptions-item"); + expect(items).toHaveLength(5); + }); + + describe("Description Items Structure", () => { + it("should render exactly 10 skeleton nodes total", () => { + render(); + const skeletonNodes = screen.getAllByTestId("skeleton-node"); + expect(skeletonNodes).toHaveLength(10); + }); + + it("should render 5 skeleton nodes for labels with width 80", () => { + render(); + const skeletonNodes = screen.getAllByTestId("skeleton-node"); + + const labelNodes = skeletonNodes.filter( + (node) => node.getAttribute("data-style") === JSON.stringify({ width: 80, height: 16 }), + ); + expect(labelNodes).toHaveLength(5); + + labelNodes.forEach((node) => { + expect(node).toHaveAttribute("data-active", "true"); + }); + }); + + it("should render skeleton nodes for content with correct widths", () => { + render(); + const skeletonNodes = screen.getAllByTestId("skeleton-node"); + + // Expected content widths: [100, 200, 250, 180, 220] + const expectedWidths = [100, 200, 250, 180, 220]; + expectedWidths.forEach((width) => { + const contentNode = skeletonNodes.find( + (node) => node.getAttribute("data-style") === JSON.stringify({ width, height: 16 }), + ); + expect(contentNode).toBeInTheDocument(); + expect(contentNode).toHaveAttribute("data-active", "true"); + }); + }); + }); + }); + + describe("Accessibility and Structure", () => { + it("should have proper semantic structure", () => { + render(); + // Card contains Space + const card = screen.getByTestId("card"); + const space = screen.getByTestId("space"); + expect(card).toContainElement(space); + + // Space contains header section and descriptions + const descriptions = screen.getByTestId("descriptions"); + expect(space).toContainElement(descriptions); + }); + + it("should render all skeleton elements as active", () => { + render(); + const skeletonNodes = screen.getAllByTestId("skeleton-node"); + const skeletonButtons = screen.getAllByTestId("skeleton-button"); + + skeletonNodes.forEach((node) => { + expect(node).toHaveAttribute("data-active", "true"); + }); + + skeletonButtons.forEach((button) => { + expect(button).toHaveAttribute("data-active", "true"); + }); + }); + }); +}); diff --git a/ui/litellm-dashboard/src/components/vector_store_management/VectorStoreSelector.test.tsx b/ui/litellm-dashboard/src/components/vector_store_management/VectorStoreSelector.test.tsx new file mode 100644 index 00000000000..8c6b85a53de --- /dev/null +++ b/ui/litellm-dashboard/src/components/vector_store_management/VectorStoreSelector.test.tsx @@ -0,0 +1,524 @@ +import { render, screen, waitFor, fireEvent } from "@testing-library/react"; +import userEvent from "@testing-library/user-event"; +import { beforeEach, describe, expect, it, vi } from "vitest"; +import VectorStoreSelector from "./VectorStoreSelector"; +import { vectorStoreListCall } from "../networking"; +import { VectorStore } from "./types"; + +// Mock dependencies +const mockVectorStoreListCall = vi.fn(); + +vi.mock("../networking", () => ({ + vectorStoreListCall: (...args: any[]) => mockVectorStoreListCall(...args), +})); + +// Mock antd Select component +vi.mock("antd", () => ({ + Select: vi.fn(), +})); + +// Import the mocked Select +import { Select as MockedSelect } from "antd"; + +// Configure the mock to render a simple div with data attributes +(MockedSelect as any).mockImplementation((props: any) => { + const { + onChange, + value, + placeholder, + loading, + className, + disabled, + options, + mode, + showSearch, + optionFilterProp, + style, + } = props; + + return ( +
{ + // For testing purposes, allow simulating different selection behaviors + // The test can control this by setting data attributes on the element + const testSelection = e.target.getAttribute("data-test-selection"); + if (testSelection && onChange) { + onChange(JSON.parse(testSelection)); + } else if (onChange && options?.length > 0) { + // Default behavior: select first option + onChange([options[0].value]); + } + }} + > + {options?.map((opt: any) => ( +
+ {opt.label} +
+ ))} +
+ ); +}); + +// Test helpers +const mockOnChange = vi.fn(); +const mockAccessToken = "test-token"; + +const mockVectorStores: VectorStore[] = [ + { + vector_store_id: "store-1", + custom_llm_provider: "openai", + vector_store_name: "My Store", + vector_store_description: "A test store", + created_at: "2024-01-01T00:00:00Z", + updated_at: "2024-01-01T00:00:00Z", + }, + { + vector_store_id: "store-2", + custom_llm_provider: "azure", + vector_store_name: "Another Store", + vector_store_description: "Another test store", + created_at: "2024-01-02T00:00:00Z", + updated_at: "2024-01-02T00:00:00Z", + }, + { + vector_store_id: "store-3", + custom_llm_provider: "pg_vector", + // No vector_store_name to test fallback to vector_store_id + vector_store_description: "Store without name", + created_at: "2024-01-03T00:00:00Z", + updated_at: "2024-01-03T00:00:00Z", + }, +]; + +const defaultProps = { + onChange: mockOnChange, + accessToken: mockAccessToken, +}; + +// Helper functions +const renderComponent = (props = {}) => { + return render(); +}; + +const waitForDataFetch = async () => { + await waitFor(() => { + expect(mockVectorStoreListCall).toHaveBeenCalled(); + }); +}; + +const getSelectElement = () => screen.getByTestId("vector-store-select"); + +const getOptionElements = () => + screen.getAllByTestId(/^vector-store-select/).filter((el) => el.hasAttribute("data-option-value")); + +describe("VectorStoreSelector", () => { + beforeEach(() => { + vi.clearAllMocks(); + mockVectorStoreListCall.mockResolvedValue({ + data: mockVectorStores, + }); + }); + + describe("Rendering", () => { + it("should render the select component", () => { + renderComponent(); + expect(getSelectElement()).toBeInTheDocument(); + }); + + it("should render with default placeholder", () => { + renderComponent(); + const select = getSelectElement(); + expect(select).toHaveAttribute("data-placeholder", "Select vector stores"); + }); + + it("should render with custom placeholder", () => { + renderComponent({ placeholder: "Choose stores" }); + const select = getSelectElement(); + expect(select).toHaveAttribute("data-placeholder", "Choose stores"); + }); + + it("should apply custom className", () => { + renderComponent({ className: "custom-class" }); + const select = getSelectElement(); + expect(select).toHaveClass("custom-class"); + }); + + it("should render with disabled state", () => { + renderComponent({ disabled: true }); + const select = getSelectElement(); + expect(select).toHaveAttribute("data-disabled", "true"); + }); + + it("should render with enabled state by default", () => { + renderComponent(); + const select = getSelectElement(); + expect(select).toHaveAttribute("data-disabled", "false"); + }); + + it("should render with multiple mode", () => { + renderComponent(); + const select = getSelectElement(); + expect(select).toHaveAttribute("data-mode", "multiple"); + }); + + it("should render with showSearch enabled", () => { + renderComponent(); + const select = getSelectElement(); + expect(select).toHaveAttribute("data-show-search", "true"); + }); + + it("should render with optionFilterProp set to label", () => { + renderComponent(); + const select = getSelectElement(); + expect(select).toHaveAttribute("data-option-filter-prop", "label"); + }); + + it("should render with full width style", () => { + renderComponent(); + const select = getSelectElement(); + expect(select).toHaveStyle({ width: "100%" }); + }); + }); + + describe("Data fetching", () => { + it("should fetch vector stores on mount when accessToken is provided", async () => { + renderComponent(); + await waitFor(() => { + expect(mockVectorStoreListCall).toHaveBeenCalledWith(mockAccessToken); + }); + }); + + it("should not fetch vector stores when accessToken is falsy", () => { + const { rerender } = render(); + expect(mockVectorStoreListCall).not.toHaveBeenCalled(); + + rerender(); + expect(mockVectorStoreListCall).not.toHaveBeenCalled(); + + rerender(); + expect(mockVectorStoreListCall).not.toHaveBeenCalled(); + }); + + it("should fetch vector stores again when accessToken changes", async () => { + const { rerender } = render(); + await waitFor(() => { + expect(mockVectorStoreListCall).toHaveBeenCalledWith("token-1"); + }); + + vi.clearAllMocks(); + rerender(); + await waitFor(() => { + expect(mockVectorStoreListCall).toHaveBeenCalledWith("token-2"); + }); + }); + + it("should set loading state while fetching", async () => { + let resolvePromise: (value: any) => void; + const promise = new Promise((resolve) => { + resolvePromise = resolve; + }); + mockVectorStoreListCall.mockReturnValue(promise); + + renderComponent(); + const select = getSelectElement(); + expect(select).toHaveAttribute("data-loading", "true"); + + resolvePromise!({ data: mockVectorStores }); + await waitFor(() => { + expect(select).toHaveAttribute("data-loading", "false"); + }); + }); + + it("should clear loading state after successful fetch", async () => { + renderComponent(); + await waitForDataFetch(); + const select = getSelectElement(); + expect(select).toHaveAttribute("data-loading", "false"); + }); + + it("should clear loading state after failed fetch", async () => { + const consoleErrorSpy = vi.spyOn(console, "error").mockImplementation(() => {}); + mockVectorStoreListCall.mockRejectedValueOnce(new Error("Network error")); + + renderComponent(); + await waitForDataFetch(); + + const select = getSelectElement(); + expect(select).toHaveAttribute("data-loading", "false"); + consoleErrorSpy.mockRestore(); + }); + }); + + describe("Options rendering", () => { + it("should render vector store options after successful fetch", async () => { + renderComponent(); + await waitForDataFetch(); + + expect(screen.getByText("My Store (store-1)")).toBeInTheDocument(); + expect(screen.getByText("Another Store (store-2)")).toBeInTheDocument(); + expect(screen.getByText("store-3 (store-3)")).toBeInTheDocument(); + }); + + it("should use vector_store_name when available for label", async () => { + renderComponent(); + await waitForDataFetch(); + + const option1 = screen.getByText("My Store (store-1)"); + expect(option1).toBeInTheDocument(); + expect(option1).toHaveAttribute("data-option-title", "A test store"); + }); + + it("should fallback to vector_store_id when vector_store_name is missing", async () => { + renderComponent(); + await waitForDataFetch(); + + const option3 = screen.getByText("store-3 (store-3)"); + expect(option3).toBeInTheDocument(); + // When vector_store_name is missing, title uses vector_store_description if available, otherwise vector_store_id + expect(option3).toHaveAttribute("data-option-title", "Store without name"); + }); + + it("should use vector_store_description as title when available", async () => { + renderComponent(); + await waitForDataFetch(); + + const option1 = screen.getByText("My Store (store-1)"); + expect(option1).toHaveAttribute("data-option-title", "A test store"); + }); + + it("should fallback to vector_store_id as title when vector_store_description is missing", async () => { + const storesWithoutDescription: VectorStore[] = [ + { + vector_store_id: "store-no-desc", + custom_llm_provider: "openai", + created_at: "2024-01-01T00:00:00Z", + updated_at: "2024-01-01T00:00:00Z", + }, + ]; + mockVectorStoreListCall.mockResolvedValueOnce({ + data: storesWithoutDescription, + }); + + renderComponent(); + await waitForDataFetch(); + + const option = screen.getByText("store-no-desc (store-no-desc)"); + expect(option).toHaveAttribute("data-option-title", "store-no-desc"); + }); + + it("should use vector_store_id as option value", async () => { + renderComponent(); + await waitForDataFetch(); + + const option1 = screen.getByText("My Store (store-1)"); + expect(option1).toHaveAttribute("data-option-value", "store-1"); + }); + + it("should handle empty vector stores array", async () => { + mockVectorStoreListCall.mockResolvedValueOnce({ + data: [], + }); + + renderComponent(); + await waitForDataFetch(); + + const options = getOptionElements(); + expect(options.length).toBe(0); + }); + + it("should handle response without data property", async () => { + mockVectorStoreListCall.mockResolvedValueOnce({}); + + renderComponent(); + await waitForDataFetch(); + + const options = getOptionElements(); + expect(options.length).toBe(0); + }); + }); + + describe("Value prop", () => { + it("should set initial value when value prop is provided", async () => { + renderComponent({ value: ["store-1", "store-2"] }); + await waitForDataFetch(); + + const select = getSelectElement(); + const dataValue = select.getAttribute("data-value"); + expect(dataValue).toBe(JSON.stringify(["store-1", "store-2"])); + }); + + it("should handle empty value array", async () => { + renderComponent({ value: [] }); + await waitForDataFetch(); + + const select = getSelectElement(); + const dataValue = select.getAttribute("data-value"); + expect(dataValue).toBe(JSON.stringify([])); + }); + + it("should handle undefined value", async () => { + renderComponent({ value: undefined }); + await waitForDataFetch(); + + const select = getSelectElement(); + const dataValue = select.getAttribute("data-value"); + expect(dataValue).toBeNull(); // undefined value results in no data-value attribute + }); + }); + + describe("onChange callback", () => { + it("should call onChange when selection changes", async () => { + renderComponent(); + await waitForDataFetch(); + + const select = getSelectElement(); + // Simulate selecting store-1 by setting test data attribute + select.setAttribute("data-test-selection", '["store-1"]'); + fireEvent.click(select); + + expect(mockOnChange).toHaveBeenCalledWith(["store-1"]); + }); + + it("should call onChange with multiple selected values", async () => { + renderComponent(); + await waitForDataFetch(); + + const select = getSelectElement(); + // Simulate selecting multiple values + select.setAttribute("data-test-selection", '["store-1", "store-2"]'); + fireEvent.click(select); + + expect(mockOnChange).toHaveBeenCalledWith(["store-1", "store-2"]); + }); + + it("should call onChange when deselecting options", async () => { + renderComponent({ value: ["store-1", "store-2"] }); + await waitForDataFetch(); + + const select = getSelectElement(); + // Simulate deselecting store-1 + select.setAttribute("data-test-selection", '["store-2"]'); + fireEvent.click(select); + + expect(mockOnChange).toHaveBeenCalledWith(["store-2"]); + }); + }); + + describe("Error handling", () => { + it("should handle fetch errors gracefully", async () => { + const consoleErrorSpy = vi.spyOn(console, "error").mockImplementation(() => {}); + const error = new Error("Network error"); + mockVectorStoreListCall.mockRejectedValueOnce(error); + + renderComponent(); + await waitForDataFetch(); + + expect(consoleErrorSpy).toHaveBeenCalledWith("Error fetching vector stores:", error); + consoleErrorSpy.mockRestore(); + }); + + it("should not crash when fetch throws non-Error", async () => { + const consoleErrorSpy = vi.spyOn(console, "error").mockImplementation(() => {}); + mockVectorStoreListCall.mockRejectedValueOnce("String error"); + + renderComponent(); + await waitForDataFetch(); + + expect(consoleErrorSpy).toHaveBeenCalledWith("Error fetching vector stores:", "String error"); + consoleErrorSpy.mockRestore(); + }); + + it("should continue to work after error", async () => { + const consoleErrorSpy = vi.spyOn(console, "error").mockImplementation(() => {}); + mockVectorStoreListCall.mockRejectedValueOnce(new Error("Network error")); + + renderComponent(); + await waitForDataFetch(); + + // Component should still render + expect(getSelectElement()).toBeInTheDocument(); + consoleErrorSpy.mockRestore(); + }); + }); + + describe("Edge cases", () => { + it("should handle vector stores with all optional fields missing", async () => { + const minimalStores: VectorStore[] = [ + { + vector_store_id: "minimal-store", + custom_llm_provider: "openai", + created_at: "2024-01-01T00:00:00Z", + updated_at: "2024-01-01T00:00:00Z", + }, + ]; + mockVectorStoreListCall.mockResolvedValueOnce({ + data: minimalStores, + }); + + renderComponent(); + await waitForDataFetch(); + + expect(screen.getByText("minimal-store (minimal-store)")).toBeInTheDocument(); + const option = screen.getByText("minimal-store (minimal-store)"); + expect(option).toHaveAttribute("data-option-title", "minimal-store"); + }); + + it("should handle very long vector store names", async () => { + const longNameStores: VectorStore[] = [ + { + vector_store_id: "store-long", + custom_llm_provider: "openai", + vector_store_name: "A".repeat(200), + created_at: "2024-01-01T00:00:00Z", + updated_at: "2024-01-01T00:00:00Z", + }, + ]; + mockVectorStoreListCall.mockResolvedValueOnce({ + data: longNameStores, + }); + + renderComponent(); + await waitForDataFetch(); + + const expectedLabel = `${"A".repeat(200)} (store-long)`; + expect(screen.getByText(expectedLabel)).toBeInTheDocument(); + }); + + it("should handle special characters in vector store names", async () => { + const specialCharStores: VectorStore[] = [ + { + vector_store_id: "store-special", + custom_llm_provider: "openai", + vector_store_name: 'Store & Co. "Quotes"', + created_at: "2024-01-01T00:00:00Z", + updated_at: "2024-01-01T00:00:00Z", + }, + ]; + mockVectorStoreListCall.mockResolvedValueOnce({ + data: specialCharStores, + }); + + renderComponent(); + await waitForDataFetch(); + + expect(screen.getByText(/Store & Co\. "Quotes"/)).toBeInTheDocument(); + }); + }); +}); diff --git a/ui/litellm-dashboard/src/components/vector_store_management/VectorStoreTable.test.tsx b/ui/litellm-dashboard/src/components/vector_store_management/VectorStoreTable.test.tsx new file mode 100644 index 00000000000..16d5d3623eb --- /dev/null +++ b/ui/litellm-dashboard/src/components/vector_store_management/VectorStoreTable.test.tsx @@ -0,0 +1,415 @@ +import { render, screen, fireEvent, waitFor } from "@testing-library/react"; +import userEvent from "@testing-library/user-event"; +import { beforeEach, describe, expect, it, vi } from "vitest"; +import VectorStoreTable from "./VectorStoreTable"; +import { VectorStore } from "./types"; + +// Mock dependencies +const mockGetProviderLogoAndName = vi.fn(); +const mockTableIconActionButton = vi.fn(); + +vi.mock("../provider_info_helpers", () => ({ + getProviderLogoAndName: (...args: any[]) => mockGetProviderLogoAndName(...args), +})); + +vi.mock("../common_components/IconActionButton/TableIconActionButtons/TableIconActionButton", () => ({ + default: (props: any) => { + mockTableIconActionButton(props); + return ( + + ); + }, +})); + +// Mock Tremor components to avoid complex styling issues +vi.mock("@tremor/react", () => ({ + Table: ({ children, ...props }: any) => {children}
, + TableHead: ({ children, ...props }: any) => {children}, + TableBody: ({ children, ...props }: any) => {children}, + TableRow: ({ children, ...props }: any) => {children}, + TableHeaderCell: ({ children, ...props }: any) => {children}, + TableCell: ({ children, ...props }: any) => {children}, +})); + +// Mock antd Tooltip +vi.mock("antd", () => ({ + Tooltip: ({ children, title }: any) => ( +
+ {children} +
+ ), +})); + +// Mock Heroicons +vi.mock("@heroicons/react/outline", () => ({ + ChevronDownIcon: (props: any) =>
, + ChevronUpIcon: (props: any) =>
, + SwitchVerticalIcon: (props: any) =>
, +})); + +// Test data +const mockVectorStores: VectorStore[] = [ + { + vector_store_id: "short-id", + custom_llm_provider: "openai", + vector_store_name: "My OpenAI Store", + vector_store_description: "A store for OpenAI vectors", + created_at: "2024-01-15T10:30:00Z", + updated_at: "2024-01-15T11:00:00Z", + created_by: "user-1", + updated_by: "user-1", + }, + { + vector_store_id: "very-long-vector-store-id-that-should-be-truncated", + custom_llm_provider: "azure", + vector_store_name: undefined, // Test missing name + vector_store_description: "A store for Azure vectors with a very long description that should show a tooltip", + created_at: "2024-01-10T09:15:00Z", + updated_at: "2024-01-12T14:20:00Z", + }, + { + vector_store_id: "store-3", + custom_llm_provider: "pg_vector", + vector_store_name: "PostgreSQL Store", + vector_store_description: undefined, // Test missing description + created_at: "2024-01-05T08:00:00Z", + updated_at: "2024-01-08T16:45:00Z", + }, +]; + +// Mock functions +const mockOnView = vi.fn(); +const mockOnEdit = vi.fn(); +const mockOnDelete = vi.fn(); + +const defaultProps = { + data: mockVectorStores, + onView: mockOnView, + onEdit: mockOnEdit, + onDelete: mockOnDelete, +}; + +// Helper function to render component +const renderComponent = (props = {}) => { + return render(); +}; + +describe("VectorStoreTable", () => { + beforeEach(() => { + vi.clearAllMocks(); + + // Setup default mock returns for getProviderLogoAndName + mockGetProviderLogoAndName.mockImplementation((provider: string) => { + const providerMap: Record = { + openai: { displayName: "OpenAI", logo: "/openai-logo.png" }, + azure: { displayName: "Azure", logo: "/azure-logo.png" }, + pg_vector: { displayName: "PostgreSQL Vector", logo: "/pg-logo.png" }, + }; + return providerMap[provider] || { displayName: provider, logo: "" }; + }); + }); + + describe("Rendering", () => { + it("should render the table with data", () => { + renderComponent(); + expect(screen.getByRole("table")).toBeInTheDocument(); + }); + + it("should render table headers", () => { + renderComponent(); + expect(screen.getByText("Vector Store ID")).toBeInTheDocument(); + expect(screen.getByText("Name")).toBeInTheDocument(); + expect(screen.getByText("Description")).toBeInTheDocument(); + expect(screen.getByText("Provider")).toBeInTheDocument(); + expect(screen.getByText("Created At")).toBeInTheDocument(); + expect(screen.getByText("Updated At")).toBeInTheDocument(); + // Check that we have the expected number of header cells (6 data + 1 actions) + const headers = screen.getAllByRole("columnheader"); + expect(headers).toHaveLength(7); + }); + + it("should render all vector store rows", () => { + renderComponent(); + expect(screen.getAllByRole("row")).toHaveLength(mockVectorStores.length + 1); // +1 for header row + }); + + it("should render empty state when no data", () => { + renderComponent({ data: [] }); + expect(screen.getByText("No vector stores found")).toBeInTheDocument(); + }); + }); + + describe("Vector Store ID Column", () => { + it("should render short vector store IDs fully", () => { + renderComponent(); + expect(screen.getByText("short-id")).toBeInTheDocument(); + }); + + it("should truncate long vector store IDs", () => { + renderComponent(); + // Check that the truncated text is rendered (first 15 chars + ...) + const truncatedText = "very-long-vecto..."; + expect(screen.getByText(truncatedText)).toBeInTheDocument(); + }); + + it("should make vector store ID clickable", async () => { + const user = userEvent.setup(); + renderComponent(); + const idButton = screen.getByText("short-id"); + await user.click(idButton); + expect(mockOnView).toHaveBeenCalledWith("short-id"); + }); + + it("should have correct styling for vector store ID button", () => { + renderComponent(); + const idButton = screen.getByText("short-id").closest("button"); + expect(idButton).toHaveClass("font-mono", "text-blue-500", "bg-blue-50", "hover:bg-blue-100"); + }); + }); + + describe("Name Column", () => { + it("should render vector store name", () => { + renderComponent(); + expect(screen.getByText("My OpenAI Store")).toBeInTheDocument(); + }); + + it("should render fallback for missing name", () => { + renderComponent(); + const fallbackElements = screen.getAllByText("-"); + expect(fallbackElements.length).toBe(2); // One for missing name, one for missing description + }); + + it("should wrap name in tooltip", () => { + renderComponent(); + const tooltips = screen.getAllByTestId("tooltip"); + const nameTooltip = tooltips.find((t) => t.getAttribute("data-title") === "My OpenAI Store"); + expect(nameTooltip).toBeInTheDocument(); + }); + }); + + describe("Description Column", () => { + it("should render vector store description", () => { + renderComponent(); + expect(screen.getByText("A store for OpenAI vectors")).toBeInTheDocument(); + }); + + it("should render fallback for missing description", () => { + renderComponent(); + const fallbackElements = screen.getAllByText("-"); + expect(fallbackElements.length).toBe(2); // One for missing name, one for missing description + }); + + it("should wrap description in tooltip", () => { + renderComponent(); + const tooltips = screen.getAllByTestId("tooltip"); + const descTooltip = tooltips.find( + (t) => + t.getAttribute("data-title") === + "A store for Azure vectors with a very long description that should show a tooltip", + ); + expect(descTooltip).toBeInTheDocument(); + }); + }); + + describe("Provider Column", () => { + it("should render provider display name", () => { + renderComponent(); + expect(screen.getByText("OpenAI")).toBeInTheDocument(); + expect(screen.getByText("Azure")).toBeInTheDocument(); + expect(screen.getByText("PostgreSQL Vector")).toBeInTheDocument(); + }); + + it("should render provider logo when available", () => { + renderComponent(); + const logos = screen.getAllByRole("img"); + expect(logos).toHaveLength(3); // All providers have logos in our mock + expect(logos[0]).toHaveAttribute("src", "/openai-logo.png"); + expect(logos[0]).toHaveAttribute("alt", "OpenAI"); + }); + + it("should call getProviderLogoAndName for each provider", () => { + renderComponent(); + expect(mockGetProviderLogoAndName).toHaveBeenCalledWith("openai"); + expect(mockGetProviderLogoAndName).toHaveBeenCalledWith("azure"); + expect(mockGetProviderLogoAndName).toHaveBeenCalledWith("pg_vector"); + }); + }); + + describe("Date Columns", () => { + it("should render created at dates", () => { + renderComponent(); + const dateElements = screen.getAllByText(/1\/\d+\/2024/); + expect(dateElements.length).toBe(6); // 3 created_at + 3 updated_at dates + }); + + it("should render updated at dates", () => { + renderComponent(); + const dateElements = screen.getAllByText(/1\/\d+\/2024/); + expect(dateElements.length).toBe(6); // 3 created_at + 3 updated_at dates + }); + }); + + describe("Actions Column", () => { + it("should render edit and delete action buttons for each row", () => { + renderComponent(); + expect(screen.getAllByTestId("action-button-edit")).toHaveLength(mockVectorStores.length); + expect(screen.getAllByTestId("action-button-delete")).toHaveLength(mockVectorStores.length); + }); + + it("should call onEdit when edit button is clicked", async () => { + const user = userEvent.setup(); + renderComponent(); + const editButtons = screen.getAllByTestId("action-button-edit"); + await user.click(editButtons[0]); + expect(mockOnEdit).toHaveBeenCalledWith("short-id"); + }); + + it("should call onDelete when delete button is clicked", async () => { + const user = userEvent.setup(); + renderComponent(); + const deleteButtons = screen.getAllByTestId("action-button-delete"); + await user.click(deleteButtons[0]); + expect(mockOnDelete).toHaveBeenCalledWith("short-id"); + }); + + it("should pass correct props to TableIconActionButton", () => { + renderComponent(); + expect(mockTableIconActionButton).toHaveBeenCalledWith( + expect.objectContaining({ + variant: "Edit", + tooltipText: "Edit vector store", + onClick: expect.any(Function), + }), + ); + expect(mockTableIconActionButton).toHaveBeenCalledWith( + expect.objectContaining({ + variant: "Delete", + tooltipText: "Delete vector store", + onClick: expect.any(Function), + }), + ); + }); + }); + + describe("Sorting", () => { + it("should initialize with created_at descending sort", () => { + renderComponent(); + // The table should initialize with sorting state + expect(screen.getByTestId("chevron-down")).toBeInTheDocument(); + }); + + it("should render sort icons for sortable columns", () => { + renderComponent(); + // Should have sort icons for Created At and Updated At columns + const sortIcons = screen.getAllByTestId(/^chevron-(up|down)$|^switch-vertical$/); + expect(sortIcons.length).toBeGreaterThan(0); + }); + + it("should make header cells clickable for sorting", () => { + renderComponent(); + const headerCells = screen.getAllByRole("columnheader"); + const sortableHeaders = headerCells.filter((cell) => cell.textContent !== ""); + expect(sortableHeaders.length).toBeGreaterThan(0); + }); + + it("should show ascending icon when sorted ascending", () => { + renderComponent(); + // Initially shows descending, but we can test the logic by checking the icons are present + expect(screen.getByTestId("chevron-down")).toBeInTheDocument(); + }); + }); + + describe("Styling and Layout", () => { + it("should apply correct CSS classes to table container", () => { + renderComponent(); + const tableContainer = screen.getByRole("table").parentElement?.parentElement; + expect(tableContainer).toHaveClass("rounded-lg", "custom-border", "relative"); + }); + + it("should apply overflow styling to table wrapper", () => { + renderComponent(); + const tableWrapper = screen.getByRole("table").parentElement; + expect(tableWrapper).toHaveClass("overflow-x-auto"); + }); + + it("should apply sticky styling to actions column", () => { + renderComponent(); + const headerCells = screen.getAllByRole("columnheader"); + const actionsHeader = headerCells[headerCells.length - 1]; + expect(actionsHeader).toHaveClass("sticky", "right-0", "bg-white"); + }); + + it("should apply sticky styling to action cells", () => { + renderComponent(); + const rows = screen.getAllByRole("row").slice(1); // Skip header row + rows.forEach((row) => { + const cells = row.querySelectorAll("td"); + const lastCell = cells[cells.length - 1]; + expect(lastCell).toHaveClass("sticky", "right-0", "bg-white"); + }); + }); + }); + + describe("Table Row Styling", () => { + it("should apply correct height to table rows", () => { + renderComponent(); + const rows = screen.getAllByRole("row").slice(1); // Skip header row + rows.forEach((row) => { + expect(row).toHaveClass("h-8"); + }); + }); + + it("should apply correct cell padding and styling", () => { + renderComponent(); + const cells = screen.getAllByRole("cell"); + cells.forEach((cell) => { + expect(cell).toHaveClass("py-0.5", "max-h-8", "overflow-hidden", "text-ellipsis", "whitespace-nowrap"); + }); + }); + }); + + describe("Empty State", () => { + it("should render single row with centered message when no data", () => { + renderComponent({ data: [] }); + const rows = screen.getAllByRole("row"); + expect(rows).toHaveLength(2); // Header + empty state row + expect(screen.getByText("No vector stores found")).toBeInTheDocument(); + }); + + it("should span all columns in empty state", () => { + renderComponent({ data: [] }); + const emptyCell = screen.getByText("No vector stores found").closest("td"); + expect(emptyCell).toHaveAttribute("colSpan", "7"); // 6 data columns + 1 actions column + }); + }); + + describe("Data Edge Cases", () => { + it("should handle vector stores with minimal data", () => { + const minimalData: VectorStore[] = [ + { + vector_store_id: "minimal", + custom_llm_provider: "test", + created_at: "2024-01-01T00:00:00Z", + updated_at: "2024-01-01T00:00:00Z", + }, + ]; + + renderComponent({ data: minimalData }); + expect(screen.getByText("minimal")).toBeInTheDocument(); + expect(screen.getAllByText("-")).toHaveLength(2); // Name and description fallbacks + }); + + it("should handle single vector store", () => { + const singleData = [mockVectorStores[0]]; + renderComponent({ data: singleData }); + expect(screen.getAllByRole("row")).toHaveLength(2); // Header + 1 data row + }); + }); +}); diff --git a/ui/litellm-dashboard/src/utils/cookieUtils.test.ts b/ui/litellm-dashboard/src/utils/cookieUtils.test.ts index 1993819be88..8b066e6a8ea 100644 --- a/ui/litellm-dashboard/src/utils/cookieUtils.test.ts +++ b/ui/litellm-dashboard/src/utils/cookieUtils.test.ts @@ -44,6 +44,78 @@ describe("cookieUtils", () => { expect(getCookie("token")).toBeNull(); }); + + it("should return early when window is undefined (server-side rendering)", () => { + const originalWindow = global.window; + const originalDocument = global.document; + + // Mock server-side environment + delete (global as any).window; + delete (global as any).document; + + // This should not throw an error and should return early + expect(() => clearTokenCookies()).not.toThrow(); + + // Restore globals + global.window = originalWindow; + global.document = originalDocument; + }); + + it("should return early when document is undefined (server-side rendering)", () => { + const originalDocument = global.document; + + // Mock server-side environment where document is undefined + delete (global as any).document; + + // This should not throw an error and should return early + expect(() => clearTokenCookies()).not.toThrow(); + + // Restore globals + global.document = originalDocument; + }); + + it("should add current path directory to paths array when different from root and /ui", () => { + // Mock window.location.pathname using vi.stubGlobal + const originalLocation = window.location; + vi.stubGlobal('location', { ...originalLocation, pathname: '/custom/path/page.html' }); + + // Spy on document.cookie to verify the paths being used + const cookieSpy = vi.spyOn(document, 'cookie', 'set'); + + clearTokenCookies(); + + // Verify that cookies were cleared for /custom/path/ path + expect(cookieSpy).toHaveBeenCalledWith( + expect.stringContaining('path=/custom/path/') + ); + + vi.restoreAllMocks(); + }); + + it("should not add current path directory when it's already in paths array", () => { + // Mock window.location.pathname using vi.stubGlobal + const originalLocation = window.location; + vi.stubGlobal('location', { ...originalLocation, pathname: '/' }); + + // Spy on document.cookie to count calls + const cookieSpy = vi.spyOn(document, 'cookie', 'set'); + + clearTokenCookies(); + + // Count how many times each path was used + const rootPathCalls = cookieSpy.mock.calls.filter(call => + call[0].includes('path=/;') || call[0].includes('path=/ ') + ); + const uiPathCalls = cookieSpy.mock.calls.filter(call => + call[0].includes('path=/ui;') || call[0].includes('path=/ui ') + ); + + // Should have calls for root and /ui paths, but not duplicate root + expect(rootPathCalls.length).toBeGreaterThan(0); + expect(uiPathCalls.length).toBeGreaterThan(0); + + vi.restoreAllMocks(); + }); }); describe("getCookie", () => { diff --git a/ui/litellm-dashboard/src/utils/proxyUtils.test.ts b/ui/litellm-dashboard/src/utils/proxyUtils.test.ts new file mode 100644 index 00000000000..37bbd429db9 --- /dev/null +++ b/ui/litellm-dashboard/src/utils/proxyUtils.test.ts @@ -0,0 +1,78 @@ +import { describe, it, expect, vi, beforeEach, afterEach } from "vitest"; +import { fetchProxySettings } from "./proxyUtils"; +import { getProxyUISettings } from "@/components/networking"; + +vi.mock("@/components/networking", () => ({ + getProxyUISettings: vi.fn(), +})); + +describe("fetchProxySettings", () => { + beforeEach(() => { + vi.clearAllMocks(); + }); + + afterEach(() => { + vi.restoreAllMocks(); + }); + + it("should return null when accessToken is null", async () => { + const result = await fetchProxySettings(null); + + expect(result).toBeNull(); + expect(getProxyUISettings).not.toHaveBeenCalled(); + }); + + it("should return null when accessToken is undefined", async () => { + const result = await fetchProxySettings(undefined as any); + + expect(result).toBeNull(); + expect(getProxyUISettings).not.toHaveBeenCalled(); + }); + + it("should return proxy settings when getProxyUISettings succeeds", async () => { + const mockProxySettings = { someSetting: "value", anotherSetting: 123 }; + const accessToken = "test-token"; + + vi.mocked(getProxyUISettings).mockResolvedValue(mockProxySettings); + + const result = await fetchProxySettings(accessToken); + + expect(result).toEqual(mockProxySettings); + expect(getProxyUISettings).toHaveBeenCalledOnce(); + expect(getProxyUISettings).toHaveBeenCalledWith(accessToken); + }); + + it("should return null and log error when getProxyUISettings throws", async () => { + const accessToken = "test-token"; + const mockError = new Error("Network error"); + const consoleSpy = vi.spyOn(console, "error").mockImplementation(() => {}); + + vi.mocked(getProxyUISettings).mockRejectedValue(mockError); + + const result = await fetchProxySettings(accessToken); + + expect(result).toBeNull(); + expect(getProxyUISettings).toHaveBeenCalledOnce(); + expect(getProxyUISettings).toHaveBeenCalledWith(accessToken); + expect(consoleSpy).toHaveBeenCalledWith("Error fetching proxy settings:", mockError); + + consoleSpy.mockRestore(); + }); + + it("should return null and log error when getProxyUISettings throws a string", async () => { + const accessToken = "test-token"; + const mockError = "String error"; + const consoleSpy = vi.spyOn(console, "error").mockImplementation(() => {}); + + vi.mocked(getProxyUISettings).mockRejectedValue(mockError); + + const result = await fetchProxySettings(accessToken); + + expect(result).toBeNull(); + expect(getProxyUISettings).toHaveBeenCalledOnce(); + expect(getProxyUISettings).toHaveBeenCalledWith(accessToken); + expect(consoleSpy).toHaveBeenCalledWith("Error fetching proxy settings:", mockError); + + consoleSpy.mockRestore(); + }); +}); \ No newline at end of file diff --git a/ui/litellm-dashboard/src/utils/textUtils.test.ts b/ui/litellm-dashboard/src/utils/textUtils.test.ts index b7c91b9dd33..dfb37ad63b4 100644 --- a/ui/litellm-dashboard/src/utils/textUtils.test.ts +++ b/ui/litellm-dashboard/src/utils/textUtils.test.ts @@ -5,6 +5,15 @@ describe("formatLabel", () => { it("should format label", () => { expect(formatLabel("test_label")).toBe("Test Label"); }); + + it("should return empty string when text is empty string", () => { + expect(formatLabel("")).toBe(""); + }); + + it("should return the same value when text is falsy", () => { + expect(formatLabel(null as any)).toBe(null); + expect(formatLabel(undefined as any)).toBe(undefined); + }); }); describe("truncateString", () => { @@ -26,4 +35,16 @@ describe("formItemValidateJSON", () => { it("should reject with an error message for invalid JSON", async () => { await expect(formItemValidateJSON({}, "invalid JSON")).rejects.toBe("Please enter valid JSON"); }); + + it("should resolve when value is empty string", async () => { + await expect(formItemValidateJSON({}, "")).resolves.toBeUndefined(); + }); + + it("should resolve when value is null", async () => { + await expect(formItemValidateJSON({}, null as any)).resolves.toBeUndefined(); + }); + + it("should resolve when value is undefined", async () => { + await expect(formItemValidateJSON({}, undefined as any)).resolves.toBeUndefined(); + }); }); From 1184db079ec96bee585b08f5e62971534774048f Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Sat, 3 Jan 2026 17:30:52 -0800 Subject: [PATCH 152/158] fixing tests --- .../Modals/EditSSOSettingsModal.tsx | 25 +++++++++++-------- 1 file changed, 15 insertions(+), 10 deletions(-) diff --git a/ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/Modals/EditSSOSettingsModal.tsx b/ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/Modals/EditSSOSettingsModal.tsx index 297698a7ba0..a731af68ff1 100644 --- a/ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/Modals/EditSSOSettingsModal.tsx +++ b/ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/Modals/EditSSOSettingsModal.tsx @@ -88,17 +88,22 @@ const EditSSOSettingsModal: React.FC = ({ isVisible, // Enhanced form submission handler const handleFormSubmit = async (formValues: Record) => { - const payload = processSSOSettingsPayload(formValues); + try { + const payload = processSSOSettingsPayload(formValues); - await mutateAsync(payload, { - onSuccess: () => { - NotificationsManager.success("SSO settings updated successfully"); - onSuccess(); - }, - onError: (error) => { - NotificationsManager.fromBackend("Failed to save SSO settings: " + parseErrorMessage(error)); - }, - }); + await mutateAsync(payload, { + onSuccess: () => { + NotificationsManager.success("SSO settings updated successfully"); + onSuccess(); + }, + onError: (error) => { + NotificationsManager.fromBackend("Failed to save SSO settings: " + parseErrorMessage(error)); + }, + }); + } catch (error) { + // Handle processing errors gracefully + NotificationsManager.fromBackend("Failed to process SSO settings: " + parseErrorMessage(error)); + } }; const handleCancel = () => { From 816124a40bc8e9d0614e11be10becd5a33352d6c Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Sat, 3 Jan 2026 17:39:59 -0800 Subject: [PATCH 153/158] Fixign build --- .../vector_store_management/VectorStoreSelector.test.tsx | 6 ++---- .../vector_store_management/VectorStoreTable.test.tsx | 2 +- 2 files changed, 3 insertions(+), 5 deletions(-) diff --git a/ui/litellm-dashboard/src/components/vector_store_management/VectorStoreSelector.test.tsx b/ui/litellm-dashboard/src/components/vector_store_management/VectorStoreSelector.test.tsx index 8c6b85a53de..67476b5559d 100644 --- a/ui/litellm-dashboard/src/components/vector_store_management/VectorStoreSelector.test.tsx +++ b/ui/litellm-dashboard/src/components/vector_store_management/VectorStoreSelector.test.tsx @@ -1,9 +1,7 @@ -import { render, screen, waitFor, fireEvent } from "@testing-library/react"; -import userEvent from "@testing-library/user-event"; +import { fireEvent, render, screen, waitFor } from "@testing-library/react"; import { beforeEach, describe, expect, it, vi } from "vitest"; -import VectorStoreSelector from "./VectorStoreSelector"; -import { vectorStoreListCall } from "../networking"; import { VectorStore } from "./types"; +import VectorStoreSelector from "./VectorStoreSelector"; // Mock dependencies const mockVectorStoreListCall = vi.fn(); diff --git a/ui/litellm-dashboard/src/components/vector_store_management/VectorStoreTable.test.tsx b/ui/litellm-dashboard/src/components/vector_store_management/VectorStoreTable.test.tsx index 16d5d3623eb..65d15260c4c 100644 --- a/ui/litellm-dashboard/src/components/vector_store_management/VectorStoreTable.test.tsx +++ b/ui/litellm-dashboard/src/components/vector_store_management/VectorStoreTable.test.tsx @@ -1,4 +1,4 @@ -import { render, screen, fireEvent, waitFor } from "@testing-library/react"; +import { render, screen } from "@testing-library/react"; import userEvent from "@testing-library/user-event"; import { beforeEach, describe, expect, it, vi } from "vitest"; import VectorStoreTable from "./VectorStoreTable"; From 1112974112ecc4aacf7ca1ab811c7e3cbac48b8f Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Sat, 3 Jan 2026 19:22:38 -0800 Subject: [PATCH 154/158] Virtual Keys Table Loading State --- .../VirtualKeysPage/VirtualKeysTable.test.tsx | 132 +++++++++++++++++- .../VirtualKeysPage/VirtualKeysTable.tsx | 21 +-- 2 files changed, 142 insertions(+), 11 deletions(-) diff --git a/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.test.tsx b/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.test.tsx index 3f55b11769c..cbd3d2c7320 100644 --- a/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.test.tsx +++ b/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.test.tsx @@ -1,4 +1,4 @@ -import { screen, waitFor } from "@testing-library/react"; +import { screen, waitFor, fireEvent } from "@testing-library/react"; import { vi, it, expect, beforeEach, MockedFunction } from "vitest"; import { renderWithProviders } from "../../../tests/test-utils"; import { VirtualKeysTable } from "./VirtualKeysTable"; @@ -264,3 +264,133 @@ it("should show skeleton loaders when isLoading is true", () => { expect(screen.queryByText("Test Key Alias")).not.toBeInTheDocument(); expect(screen.queryByText("Test Team")).not.toBeInTheDocument(); }); + +it("should show 'No keys found' message when filteredKeys is empty", () => { + // Mock empty filteredKeys + mockUseFilterLogic.mockReturnValue({ + filters: { + "Team ID": "", + "Organization ID": "", + "Key Alias": "", + "User ID": "", + "Sort By": "created_at", + "Sort Order": "desc", + }, + filteredKeys: [], + allKeyAliases: [], + allTeams: [mockTeam], + allOrganizations: [mockOrganization], + handleFilterChange: vi.fn(), + handleFilterReset: vi.fn(), + }); + + const mockProps = { + teams: [mockTeam], + organizations: [mockOrganization], + onSortChange: vi.fn(), + currentSort: { + sortBy: "created_at", + sortOrder: "desc" as const, + }, + }; + + renderWithProviders(); + + expect(screen.getByText("No keys found")).toBeInTheDocument(); +}); + +it("should handle models with more than 3 entries to trigger expansion UI", () => { + const keyWithManyModels = { + ...mockKey, + models: ["gpt-3.5-turbo", "gpt-4", "gpt-4-turbo", "claude-3", "claude-3-5-sonnet"], + }; + + mockUseFilterLogic.mockReturnValue({ + filters: { + "Team ID": "", + "Organization ID": "", + "Key Alias": "", + "User ID": "", + "Sort By": "created_at", + "Sort Order": "desc", + }, + filteredKeys: [keyWithManyModels], + allKeyAliases: ["test-key-alias"], + allTeams: [mockTeam], + allOrganizations: [mockOrganization], + handleFilterChange: vi.fn(), + handleFilterReset: vi.fn(), + }); + + const mockProps = { + teams: [mockTeam], + organizations: [mockOrganization], + onSortChange: vi.fn(), + currentSort: { + sortBy: "created_at", + sortOrder: "desc" as const, + }, + }; + + renderWithProviders(); + + // This test ensures the ChevronDownIcon import (line 6) is used + // by having a key with > 3 models which triggers the expansion logic + // that uses ChevronDownIcon and ChevronRightIcon + expect(screen.getByText("Test Key Alias")).toBeInTheDocument(); +}); + +it("should render table headers correctly", () => { + const mockProps = { + teams: [mockTeam], + organizations: [mockOrganization], + onSortChange: vi.fn(), + currentSort: { + sortBy: "created_at", + sortOrder: "desc" as const, + }, + }; + + renderWithProviders(); + + // Check that main headers are rendered (testing the header.isPlaceholder condition path) + expect(screen.getByText("Key ID")).toBeInTheDocument(); + expect(screen.getByText("Key Alias")).toBeInTheDocument(); + expect(screen.getByText("Team Alias")).toBeInTheDocument(); + expect(screen.getByText("Models")).toBeInTheDocument(); + expect(screen.getByText("Spend (USD)")).toBeInTheDocument(); +}); + +it("should handle column resizing hover events", () => { + const mockProps = { + teams: [mockTeam], + organizations: [mockOrganization], + onSortChange: vi.fn(), + currentSort: { + sortBy: "created_at", + sortOrder: "desc" as const, + }, + }; + + renderWithProviders(); + + // Find a header cell with data-header-id attribute + const headerCell = document.querySelector("[data-header-id]") as HTMLElement; + + expect(headerCell).toBeInTheDocument(); + + // Check that the resizer element exists within the header + const resizer = headerCell?.querySelector(".resizer") as HTMLElement; + expect(resizer).toBeInTheDocument(); + + // Initially, resizer should have opacity 0 + expect(resizer.style.opacity).toBe("0"); + + // Simulate mouse enter using fireEvent - should set opacity to 0.5 (lines 612-616) + fireEvent.mouseEnter(headerCell); + expect(resizer.style.opacity).toBe("0.5"); + + // Simulate mouse leave using fireEvent - should set opacity back to 0 (lines 618-622) + fireEvent.mouseLeave(headerCell); + expect(resizer.style.opacity).toBe("0"); +}); diff --git a/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.tsx b/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.tsx index b95d675979c..3bda8ee2f02 100644 --- a/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.tsx +++ b/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.tsx @@ -68,12 +68,13 @@ export function VirtualKeysTable({ teams, organizations, onSortChange, currentSo }); const [tablePagination, setTablePagination] = React.useState({ pageIndex: 0, - pageSize: 100, + pageSize: 50, }); const { data: keys, isPending: isLoading, + isFetching, refetch, } = useKeys(tablePagination.pageIndex + 1, tablePagination.pageSize); const totalCount = keys?.total_count || 0; @@ -545,8 +546,8 @@ export function VirtualKeysTable({ teams, organizations, onSortChange, currentSo
- {isLoading ? ( - + {isLoading || isFetching ? ( + ) : ( Showing {rangeLabel} of {totalCount} results @@ -554,32 +555,32 @@ export function VirtualKeysTable({ teams, organizations, onSortChange, currentSo )}
- {isLoading ? ( - + {isLoading || isFetching ? ( + ) : ( Page {pageIndex + 1} of {table.getPageCount()} )} - {isLoading ? ( + {isLoading || isFetching ? ( ) : ( )} - {isLoading ? ( + {isLoading || isFetching ? ( ) : (