diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index d62bfbb7d5..91a953c217 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -1672,6 +1672,7 @@ class NewTeamRequest(TeamBase): int ] = None # allow user to set TPM limit for all team members team_member_key_duration: Optional[str] = None # e.g. "1d", "1w", "1m" + team_member_budget_duration: Optional[str] = None # e.g. "30d", "1mo" allowed_vector_store_indexes: Optional[List[AllowedVectorStoreIndexItem]] = None enforced_batch_output_expires_after: Optional[dict] = None enforced_file_expires_after: Optional[dict] = None diff --git a/litellm/proxy/discovery_endpoints/ui_discovery_endpoints.py b/litellm/proxy/discovery_endpoints/ui_discovery_endpoints.py index 2a38ceffba..233df5c6c5 100644 --- a/litellm/proxy/discovery_endpoints/ui_discovery_endpoints.py +++ b/litellm/proxy/discovery_endpoints/ui_discovery_endpoints.py @@ -27,10 +27,17 @@ async def get_ui_config(): admin_ui_disabled = os.getenv("DISABLE_ADMIN_UI", "false").lower() == "true" sso_configured = _has_user_setup_sso() + + from litellm.proxy.proxy_server import proxy_config + + is_control_plane = len(proxy_config.worker_registry) > 0 + return UiDiscoveryEndpoints( server_root_path=get_server_root_path(), proxy_base_url=get_proxy_base_url(), auto_redirect_to_sso=sso_configured and auto_redirect_ui_login_to_sso, admin_ui_disabled=admin_ui_disabled, sso_configured=sso_configured, + is_control_plane=is_control_plane, + workers=proxy_config.worker_registry if is_control_plane else [], ) diff --git a/litellm/proxy/management_endpoints/mcp_management_endpoints.py b/litellm/proxy/management_endpoints/mcp_management_endpoints.py index f29a721ede..e4bb288cda 100644 --- a/litellm/proxy/management_endpoints/mcp_management_endpoints.py +++ b/litellm/proxy/management_endpoints/mcp_management_endpoints.py @@ -1214,17 +1214,9 @@ if MCP_AVAILABLE: "error": "User does not have permission to create mcp servers. You can only create mcp servers if you are a PROXY_ADMIN." }, ) - elif payload.server_id is not None: - # fail if the mcp server with id already exists - mcp_server = await get_mcp_server(prisma_client, payload.server_id) - if mcp_server is not None: - raise HTTPException( - status_code=status.HTTP_400_BAD_REQUEST, - detail={ - "error": f"MCP Server with id {payload.server_id} already exists. Cannot create another." - }, - ) - elif ( + + # Block reserved special server IDs + if ( SpecialMCPServerName.all_team_servers == payload.server_id or SpecialMCPServerName.all_proxy_servers == payload.server_id ): @@ -1235,6 +1227,17 @@ if MCP_AVAILABLE: }, ) + if payload.server_id is not None: + # fail if the mcp server with id already exists + mcp_server = await get_mcp_server(prisma_client, payload.server_id) + if mcp_server is not None: + raise HTTPException( + status_code=status.HTTP_400_BAD_REQUEST, + detail={ + "error": f"MCP Server with id {payload.server_id} already exists. Cannot create another." + }, + ) + # TODO: audit log for create # Admin-created servers are always active — clear any submission lifecycle diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index 3d4488b8a7..3643373be6 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -724,6 +724,7 @@ async def new_team( # noqa: PLR0915 - disable_global_guardrails: Optional[bool] - Whether to disable global guardrails for the key. - object_permission: Optional[LiteLLM_ObjectPermissionBase] - team-specific object permission. Example - {"vector_stores": ["vector_store_1", "vector_store_2"], "agents": ["agent_1", "agent_2"], "agent_access_groups": ["dev_group"]}. IF null or {} then no object permission. - team_member_budget: Optional[float] - The maximum budget allocated to an individual team member. + - team_member_budget_duration: Optional[str] - The duration of the budget for the team member. Doc [here](https://docs.litellm.ai/docs/proxy/team_budgets) - team_member_rpm_limit: Optional[int] - The RPM (Requests Per Minute) limit for individual team members. - team_member_tpm_limit: Optional[int] - The TPM (Tokens Per Minute) limit for individual team members. - team_member_key_duration: Optional[str] - The duration for a team member's key. e.g. "1d", "1w", "1mo" @@ -934,6 +935,7 @@ async def new_team( # noqa: PLR0915 team_member_budget=data.team_member_budget, team_member_rpm_limit=data.team_member_rpm_limit, team_member_tpm_limit=data.team_member_tpm_limit, + team_member_budget_duration=data.team_member_budget_duration, ): data_json = await TeamMemberBudgetHandler.create_team_member_budget_table( data=data, @@ -942,6 +944,7 @@ async def new_team( # noqa: PLR0915 team_member_budget=data.team_member_budget, team_member_rpm_limit=data.team_member_rpm_limit, team_member_tpm_limit=data.team_member_tpm_limit, + team_member_budget_duration=data.team_member_budget_duration, ) ## ADD TO TEAM TABLE diff --git a/litellm/proxy/management_endpoints/ui_sso.py b/litellm/proxy/management_endpoints/ui_sso.py index daf1d6f131..d06ce56f81 100644 --- a/litellm/proxy/management_endpoints/ui_sso.py +++ b/litellm/proxy/management_endpoints/ui_sso.py @@ -16,6 +16,7 @@ import os import secrets from copy import deepcopy from typing import TYPE_CHECKING, Any, Dict, List, Literal, Optional, Tuple, Union, cast +from urllib.parse import urlencode, urlparse if TYPE_CHECKING: import httpx @@ -301,6 +302,7 @@ async def google_login( source: Optional[str] = None, key: Optional[str] = None, existing_key: Optional[str] = None, + return_to: Optional[str] = None, ): # noqa: PLR0915 """ Create Proxy API Keys using Google Workspace SSO. Requires setting PROXY_BASE_URL in .env @@ -394,13 +396,23 @@ async def google_login( is True ): verbose_proxy_logger.info(f"Redirecting to SSO login for {redirect_url}") - return await SSOAuthenticationHandler.get_sso_login_redirect( + sso_redirect = await SSOAuthenticationHandler.get_sso_login_redirect( redirect_url=redirect_url, microsoft_client_id=microsoft_client_id, google_client_id=google_client_id, generic_client_id=generic_client_id, state=cli_state, ) + if return_to is not None and sso_redirect is not None: + SSOAuthenticationHandler._validate_return_to(return_to) + sso_redirect.set_cookie( + key="litellm_cp_return_to", + value=return_to, + max_age=600, + httponly=True, + samesite="lax", + ) + return sso_redirect elif ui_username is not None: # No Google, Microsoft SSO # Use UI Credentials set in .env @@ -1312,12 +1324,17 @@ async def auth_callback(request: Request, state: Optional[str] = None): # noqa: request=request, key=key_id, existing_key=existing_key, result=result ) + # Control-plane cross-origin: read return_to from cookie. + # Starlette's cookie_parser already handles RFC 2109 unquoting. + cp_return_to: Optional[str] = request.cookies.get("litellm_cp_return_to") + return await SSOAuthenticationHandler.get_redirect_response_from_openid( result=result, request=request, received_response=received_response, generic_client_id=generic_client_id, ui_access_mode=ui_access_mode, + return_to=cp_return_to, ) @@ -1760,6 +1777,38 @@ class SSOAuthenticationHandler: Handler for SSO Authentication across all SSO providers """ + @staticmethod + def _validate_return_to(return_to: str) -> None: + """ + Validate that return_to matches the configured control_plane_url origin. + + Raises HTTPException(400) if: + - control_plane_url is not configured in general_settings + - return_to origin does not match control_plane_url origin + """ + from litellm.proxy.proxy_server import general_settings + + control_plane_url = general_settings.get("control_plane_url") + if control_plane_url is None: + raise HTTPException( + status_code=400, + detail="return_to is not allowed: control_plane_url is not configured", + ) + + def _origin(url: str) -> tuple: + parsed = urlparse(url) + scheme = (parsed.scheme or "").lower() + hostname = (parsed.hostname or "").lower() + default_port = 443 if scheme == "https" else 80 + port = parsed.port if parsed.port is not None else default_port + return (scheme, hostname, port) + + if _origin(return_to) != _origin(control_plane_url): + raise HTTPException( + status_code=400, + detail="return_to does not match the configured control_plane_url", + ) + @staticmethod async def get_sso_login_redirect( redirect_url: str, @@ -2358,6 +2407,7 @@ class SSOAuthenticationHandler: received_response: Optional[dict] = None, generic_client_id: Optional[str] = None, ui_access_mode: Optional[Dict] = None, + return_to: Optional[str] = None, ) -> RedirectResponse: import jwt @@ -2367,6 +2417,7 @@ class SSOAuthenticationHandler: master_key, premium_user, proxy_logging_obj, + redis_usage_cache, user_api_key_cache, user_custom_sso, ) @@ -2534,6 +2585,36 @@ class SSOAuthenticationHandler: master_key or "", algorithm="HS256", ) + + # Control-plane cross-origin: store JWT behind a single-use opaque + # code (60s TTL) so the token never appears in browser history / logs. + # The control plane redeems it via POST /v3/login/exchange. + if return_to is not None: + SSOAuthenticationHandler._validate_return_to(return_to) + + code = secrets.token_urlsafe(32) + cache_key = f"login_code:{code}" + cache_value = {"token": jwt_token, "redirect_url": return_to} + if redis_usage_cache is not None: + await redis_usage_cache.async_set_cache( + key=cache_key, value=cache_value, ttl=60 + ) + else: + await user_api_key_cache.async_set_cache( + key=cache_key, value=cache_value, ttl=60 + ) + + separator = "&" if "?" in return_to else "?" + redirect_url = ( + return_to + separator + urlencode({"login": "success", "code": code}) + ) + verbose_proxy_logger.info( + "Cross-origin SSO: redirecting to control plane with login code" + ) + redirect_response = RedirectResponse(url=redirect_url, status_code=303) + redirect_response.delete_cookie("litellm_cp_return_to") + return redirect_response + if user_id is not None and isinstance(user_id, str): litellm_dashboard_ui += "?login=success" verbose_proxy_logger.info(f"Redirecting to {litellm_dashboard_ui}") diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 9c29927c5c..e982c934aa 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -541,6 +541,7 @@ from litellm.types.llms.anthropic import ( AnthropicResponseUsageBlock, ) from litellm.types.llms.openai import HttpxBinaryResponseContent +from litellm.types.proxy.control_plane_endpoints import WorkerRegistryEntry from litellm.types.proxy.management_endpoints.model_management_endpoints import ( ModelGroupInfoProxy, ) @@ -1546,6 +1547,7 @@ user_custom_key_generate = None # Sentinel: prevents PKCE-no-Redis advisory from re-logging on config hot-reload. # Tests that need to reset it can patch 'litellm.proxy.proxy_server._pkce_no_redis_warning_emitted'. _pkce_no_redis_warning_emitted: bool = False +_cp_no_redis_warning_emitted: bool = False user_custom_sso = None user_custom_ui_sso_sign_in_handler = None use_background_health_checks = None @@ -2295,6 +2297,7 @@ class ProxyConfig: self.config: Dict[str, Any] = {} self._last_semantic_filter_config: Optional[Dict[str, Any]] = None self._last_hashicorp_vault_config: Optional[Dict[str, Any]] = None + self.worker_registry: List["WorkerRegistryEntry"] = [] def is_yaml(self, config_file_path: str) -> bool: if not os.path.isfile(config_file_path): @@ -3095,6 +3098,21 @@ class ProxyConfig: "Set PKCE_STRICT_CACHE_MISS=true to fail fast with a 401 on cache misses " "instead of continuing without a code_verifier." ) + + ### CONTROL PLANE CODE-EXCHANGE PREREQUISITE CHECK ### + cp_url = general_settings.get("control_plane_url") + if cp_url and redis_usage_cache is None: + global _cp_no_redis_warning_emitted + if not _cp_no_redis_warning_emitted: + _cp_no_redis_warning_emitted = True + verbose_proxy_logger.warning( + "control_plane_url is configured but Redis is not configured for LiteLLM caching. " + "Login codes (SSO and /v3/login) will not be shared across instances — " + "the /v3/login/exchange call may land on a different pod and fail with 401. " + "Configure Redis via the 'cache' section in your proxy config, " + "or ensure sticky sessions for single-instance deployments." + ) + ### STORE MODEL IN DB ### feature flag for `/model/new` store_model_in_db = general_settings.get("store_model_in_db", False) if store_model_in_db is None: @@ -3385,7 +3403,15 @@ class ProxyConfig: litellm.vector_store_registry.load_vector_stores_from_config( vector_store_registry_config ) - pass + + ## WORKER REGISTRY (Control Plane) + worker_registry_config = config.get("worker_registry", None) + if worker_registry_config: + self.worker_registry = [ + WorkerRegistryEntry(**e) for e in worker_registry_config + ] + else: + self.worker_registry = [] async def _init_policy_engine( self, @@ -11095,6 +11121,165 @@ async def login_v2(request: Request): # noqa: PLR0915 ) +@router.post( + "/v3/login", include_in_schema=False +) # control-plane login — always returns token in body for cross-origin use +async def login_v3(request: Request): # noqa: PLR0915 + global premium_user, general_settings, master_key + from litellm.proxy.auth.login_utils import authenticate_user, create_ui_token_object + from litellm.proxy.utils import get_custom_url + + try: + if not general_settings.get("control_plane_url"): + raise ProxyException( + message="/v3/login is only available on workers with control_plane_url configured", + type=ProxyErrorTypes.not_found_error, + param="control_plane_url", + code=status.HTTP_404_NOT_FOUND, + ) + + body = await request.json() + username = str(body.get("username")) + password = str(body.get("password")) + + login_result = await authenticate_user( + username=username, + password=password, + master_key=master_key, + prisma_client=prisma_client, + ) + + returned_ui_token_object = create_ui_token_object( + login_result=login_result, + general_settings=general_settings, + premium_user=premium_user, + ) + + import jwt + + jwt_token = jwt.encode( + cast(dict, returned_ui_token_object), + cast(str, master_key), + algorithm="HS256", + ) + + litellm_dashboard_ui = get_custom_url(str(request.base_url)) + if litellm_dashboard_ui.endswith("/"): + litellm_dashboard_ui += "ui/" + else: + litellm_dashboard_ui += "/ui/" + litellm_dashboard_ui += "?login=success" + + # Store JWT behind a single-use opaque code (60s TTL) + code = secrets.token_urlsafe(32) + cache_key = f"login_code:{code}" + cache_value = {"token": jwt_token, "redirect_url": litellm_dashboard_ui} + if redis_usage_cache is not None: + await redis_usage_cache.async_set_cache( + key=cache_key, value=cache_value, ttl=60 + ) + else: + await user_api_key_cache.async_set_cache( + key=cache_key, value=cache_value, ttl=60 + ) + + return JSONResponse( + content={"code": code, "expires_in": 60}, + status_code=status.HTTP_200_OK, + ) + except Exception as e: + verbose_proxy_logger.exception( + "litellm.proxy.proxy_server.login_v3(): Exception occurred - {}".format( + str(e) + ) + ) + if isinstance(e, ProxyException): + raise e + elif isinstance(e, HTTPException): + raise ProxyException( + message=getattr(e, "detail", str(e)), + type=ProxyErrorTypes.auth_error, + param=getattr(e, "param", "None"), + code=getattr(e, "status_code", status.HTTP_500_INTERNAL_SERVER_ERROR), + ) + else: + error_msg = f"{str(e)}" + raise ProxyException( + message=error_msg, + type=ProxyErrorTypes.auth_error, + param="None", + code=status.HTTP_500_INTERNAL_SERVER_ERROR, + ) + + +@router.post( + "/v3/login/exchange", include_in_schema=False +) # exchange single-use opaque code for JWT +async def login_v3_exchange(request: Request): + try: + if not general_settings.get("control_plane_url"): + raise ProxyException( + message="/v3/login/exchange is only available on workers with control_plane_url configured", + type=ProxyErrorTypes.not_found_error, + param="control_plane_url", + code=status.HTTP_404_NOT_FOUND, + ) + + body = await request.json() + code = body.get("code") + if not code: + raise ProxyException( + message="Missing 'code' parameter", + type=ProxyErrorTypes.auth_error, + param="code", + code=status.HTTP_400_BAD_REQUEST, + ) + + cache_key = f"login_code:{code}" + if redis_usage_cache is not None: + cached_data = await redis_usage_cache.async_get_cache(key=cache_key) + else: + cached_data = await user_api_key_cache.async_get_cache(key=cache_key) + + if not cached_data or not isinstance(cached_data, dict): + raise ProxyException( + message="Invalid or expired login code", + type=ProxyErrorTypes.auth_error, + param="code", + code=status.HTTP_401_UNAUTHORIZED, + ) + + # Single-use: delete immediately + if redis_usage_cache is not None: + await redis_usage_cache.async_delete_cache(key=cache_key) + else: + await user_api_key_cache.async_delete_cache(key=cache_key) + + json_response = JSONResponse( + content={ + "token": cached_data["token"], + "redirect_url": cached_data["redirect_url"], + }, + status_code=status.HTTP_200_OK, + ) + json_response.set_cookie(key="token", value=cached_data["token"]) + return json_response + except ProxyException: + raise + except Exception as e: + verbose_proxy_logger.exception( + "litellm.proxy.proxy_server.login_v3_exchange(): Exception occurred - {}".format( + str(e) + ) + ) + raise ProxyException( + message=str(e), + type=ProxyErrorTypes.auth_error, + param="None", + code=status.HTTP_500_INTERNAL_SERVER_ERROR, + ) + + @app.get("/onboarding/get_token", include_in_schema=False) async def onboarding(invite_link: str, request: Request): """ diff --git a/litellm/types/proxy/control_plane_endpoints.py b/litellm/types/proxy/control_plane_endpoints.py new file mode 100644 index 0000000000..8bf4c44b20 --- /dev/null +++ b/litellm/types/proxy/control_plane_endpoints.py @@ -0,0 +1,14 @@ +from pydantic import BaseModel, field_validator + + +class WorkerRegistryEntry(BaseModel): + worker_id: str + name: str + url: str + + @field_validator("url") + @classmethod + def url_must_be_http(cls, v: str) -> str: + if not v.startswith(("http://", "https://")): + raise ValueError("Worker URL must start with http:// or https://") + return v diff --git a/litellm/types/proxy/discovery_endpoints/ui_discovery_endpoints.py b/litellm/types/proxy/discovery_endpoints/ui_discovery_endpoints.py index 4a4cdaa2ba..46cd3f49f1 100644 --- a/litellm/types/proxy/discovery_endpoints/ui_discovery_endpoints.py +++ b/litellm/types/proxy/discovery_endpoints/ui_discovery_endpoints.py @@ -1,7 +1,9 @@ -from typing import Optional +from typing import List, Optional from pydantic import BaseModel +from litellm.types.proxy.control_plane_endpoints import WorkerRegistryEntry + class UiDiscoveryEndpoints(BaseModel): server_root_path: str @@ -9,3 +11,5 @@ class UiDiscoveryEndpoints(BaseModel): auto_redirect_to_sso: bool admin_ui_disabled: bool sso_configured: bool + is_control_plane: bool = False + workers: List[WorkerRegistryEntry] = [] diff --git a/tests/test_litellm/llms/anthropic/test_anthropic_files_and_batches.py b/tests/test_litellm/llms/anthropic/test_anthropic_files_and_batches.py index 4757122017..0755c189af 100644 --- a/tests/test_litellm/llms/anthropic/test_anthropic_files_and_batches.py +++ b/tests/test_litellm/llms/anthropic/test_anthropic_files_and_batches.py @@ -336,7 +336,7 @@ class TestAnthropicFilesHandler: "extra_body": None } - with patch.object(handler.anthropic_model_info, "get_api_key", return_value=None): + with patch.object(handler.anthropic_model_info, "get_auth_header", return_value=None): with pytest.raises(ValueError, match="Missing Anthropic API Key"): await handler.afile_content( file_content_request=file_content_request, diff --git a/tests/test_litellm/proxy/discovery_endpoints/test_ui_discovery_endpoints.py b/tests/test_litellm/proxy/discovery_endpoints/test_ui_discovery_endpoints.py index f15960a607..54a127f435 100644 --- a/tests/test_litellm/proxy/discovery_endpoints/test_ui_discovery_endpoints.py +++ b/tests/test_litellm/proxy/discovery_endpoints/test_ui_discovery_endpoints.py @@ -1,6 +1,6 @@ import os import sys -from unittest.mock import patch +from unittest.mock import MagicMock, patch import pytest from fastapi import FastAPI @@ -11,6 +11,7 @@ sys.path.insert( ) from litellm.proxy.discovery_endpoints.ui_discovery_endpoints import router +from litellm.types.proxy.control_plane_endpoints import WorkerRegistryEntry def test_ui_discovery_endpoints_with_defaults(): @@ -245,9 +246,9 @@ def test_ui_discovery_endpoints_with_admin_ui_enabled(): patch("litellm.proxy.utils.get_proxy_base_url", return_value=None), \ patch("litellm.proxy.auth.auth_utils._has_user_setup_sso", return_value=False), \ patch.dict(os.environ, {"DISABLE_ADMIN_UI": "false"}, clear=False): - + response = client.get("/.well-known/litellm-ui-config") - + assert response.status_code == 200 data = response.json() assert data["server_root_path"] == "/" @@ -256,3 +257,53 @@ def test_ui_discovery_endpoints_with_admin_ui_enabled(): assert data["admin_ui_disabled"] is False assert data["sso_configured"] is False + +def test_ui_discovery_endpoints_is_control_plane_true_when_workers_configured(): + app = FastAPI() + app.include_router(router) + client = TestClient(app) + + mock_config = MagicMock() + mock_config.worker_registry = [ + WorkerRegistryEntry( + worker_id="team-a", name="Team A", url="https://worker-1:4001" + ), + ] + + with patch("litellm.proxy.utils.get_server_root_path", return_value="/"), \ + patch("litellm.proxy.utils.get_proxy_base_url", return_value=None), \ + patch("litellm.proxy.auth.auth_utils._has_user_setup_sso", return_value=False), \ + patch("litellm.proxy.proxy_server.proxy_config", mock_config), \ + patch.dict(os.environ, {"DISABLE_ADMIN_UI": "false"}, clear=False): + + response = client.get("/.well-known/litellm-ui-config") + + assert response.status_code == 200 + data = response.json() + assert data["is_control_plane"] is True + assert len(data["workers"]) == 1 + assert data["workers"][0]["worker_id"] == "team-a" + assert data["workers"][0]["name"] == "Team A" + assert data["workers"][0]["url"] == "https://worker-1:4001" + + +def test_ui_discovery_endpoints_is_control_plane_false_when_no_workers(): + app = FastAPI() + app.include_router(router) + client = TestClient(app) + + mock_config = MagicMock() + mock_config.worker_registry = [] + + with patch("litellm.proxy.utils.get_server_root_path", return_value="/"), \ + patch("litellm.proxy.utils.get_proxy_base_url", return_value=None), \ + patch("litellm.proxy.auth.auth_utils._has_user_setup_sso", return_value=False), \ + patch("litellm.proxy.proxy_server.proxy_config", mock_config), \ + patch.dict(os.environ, {"DISABLE_ADMIN_UI": "false"}, clear=False): + + response = client.get("/.well-known/litellm-ui-config") + + assert response.status_code == 200 + data = response.json() + assert data["is_control_plane"] is False + assert data["workers"] == [] diff --git a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py index c325b0b6fc..366f659bda 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py @@ -6441,3 +6441,53 @@ async def test_list_team_v1_batches_key_queries(): assert result[0].keys == [key1, key2] assert result[1].team_id == "team-2" assert result[1].keys == [key3] + + +def test_new_team_request_accepts_team_member_budget_duration(): + """Test that NewTeamRequest does not silently drop team_member_budget_duration.""" + from litellm.proxy._types import NewTeamRequest + + request = NewTeamRequest( + team_member_budget=20.0, + team_member_budget_duration="30d", + ) + assert request.team_member_budget == 20.0 + assert request.team_member_budget_duration == "30d" + + +@pytest.mark.asyncio +async def test_create_team_member_budget_table_with_duration(): + """Verify that create_team_member_budget_table passes budget_duration + through to the new_budget call when team_member_budget_duration is provided.""" + from litellm.proxy._types import NewTeamRequest, UserAPIKeyAuth, LitellmUserRoles + from litellm.proxy.management_endpoints.team_endpoints import TeamMemberBudgetHandler + + mock_budget_response = MagicMock(budget_id="budget-abc") + mock_admin = UserAPIKeyAuth( + user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN + ) + + data = NewTeamRequest( + team_alias="test-team", + team_member_budget=20.0, + team_member_budget_duration="30d", + ) + + with patch( + "litellm.proxy.management_endpoints.budget_management_endpoints.new_budget", + new_callable=AsyncMock, + return_value=mock_budget_response, + ) as mock_new_budget: + result = await TeamMemberBudgetHandler.create_team_member_budget_table( + data=data, + new_team_data_json={"metadata": None}, + user_api_key_dict=mock_admin, + team_member_budget=20.0, + team_member_budget_duration="30d", + ) + + mock_new_budget.assert_awaited_once() + budget_request = mock_new_budget.call_args.kwargs["budget_obj"] + assert budget_request.budget_duration == "30d" + assert budget_request.max_budget == 20.0 + assert result["metadata"]["team_member_budget_id"] == "budget-abc" diff --git a/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py b/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py index d43b2c4ba0..fc9c37b7f8 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py +++ b/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py @@ -6,7 +6,7 @@ from unittest.mock import AsyncMock, MagicMock, patch import httpx import pytest -from fastapi import Request +from fastapi import HTTPException, Request from litellm._uuid import uuid @@ -5160,3 +5160,99 @@ def test_generic_response_convertor_extra_attributes_missing_field(monkeypatch): assert result.extra_fields["missing_field"] is None assert result.extra_fields["another_missing"] is None + +class TestValidateReturnTo: + """Tests for SSOAuthenticationHandler._validate_return_to""" + + def test_rejects_when_no_control_plane_url_configured(self, monkeypatch): + """return_to should be rejected if control_plane_url is not in general_settings.""" + monkeypatch.setattr( + "litellm.proxy.proxy_server.general_settings", {} + ) + with pytest.raises(HTTPException) as exc_info: + SSOAuthenticationHandler._validate_return_to("https://cp.example.com/ui") + assert exc_info.value.status_code == 400 + assert "not configured" in exc_info.value.detail + + def test_allows_matching_origin(self, monkeypatch): + """return_to matching the configured control_plane_url origin should pass.""" + monkeypatch.setattr( + "litellm.proxy.proxy_server.general_settings", + {"control_plane_url": "https://cp.example.com"}, + ) + # Should not raise + SSOAuthenticationHandler._validate_return_to("https://cp.example.com/ui?page=models") + + def test_allows_matching_origin_with_trailing_slash(self, monkeypatch): + """Trailing slash on control_plane_url should not affect origin comparison.""" + monkeypatch.setattr( + "litellm.proxy.proxy_server.general_settings", + {"control_plane_url": "https://cp.example.com/"}, + ) + SSOAuthenticationHandler._validate_return_to("https://cp.example.com/ui") + + def test_rejects_prefix_attack(self, monkeypatch): + """return_to like cp.example.com.evil.com must be rejected (not just prefix match).""" + monkeypatch.setattr( + "litellm.proxy.proxy_server.general_settings", + {"control_plane_url": "https://cp.example.com"}, + ) + with pytest.raises(HTTPException) as exc_info: + SSOAuthenticationHandler._validate_return_to("https://cp.example.com.evil.com/steal") + assert exc_info.value.status_code == 400 + + def test_rejects_different_origin(self, monkeypatch): + """return_to pointing to a completely different domain should be rejected.""" + monkeypatch.setattr( + "litellm.proxy.proxy_server.general_settings", + {"control_plane_url": "https://cp.example.com"}, + ) + with pytest.raises(HTTPException) as exc_info: + SSOAuthenticationHandler._validate_return_to("https://evil.com/phish") + assert exc_info.value.status_code == 400 + + def test_case_insensitive_hostname(self, monkeypatch): + """Hostname comparison should be case-insensitive per RFC 3986.""" + monkeypatch.setattr( + "litellm.proxy.proxy_server.general_settings", + {"control_plane_url": "https://CP.Example.COM"}, + ) + # Should not raise + SSOAuthenticationHandler._validate_return_to("https://cp.example.com/ui") + + def test_rejects_scheme_mismatch(self, monkeypatch): + """http:// must be rejected when control_plane_url uses https://.""" + monkeypatch.setattr( + "litellm.proxy.proxy_server.general_settings", + {"control_plane_url": "https://cp.example.com"}, + ) + with pytest.raises(HTTPException) as exc_info: + SSOAuthenticationHandler._validate_return_to("http://cp.example.com/ui") + assert exc_info.value.status_code == 400 + + def test_rejects_port_mismatch(self, monkeypatch): + """Non-default port must be rejected.""" + monkeypatch.setattr( + "litellm.proxy.proxy_server.general_settings", + {"control_plane_url": "https://cp.example.com"}, + ) + with pytest.raises(HTTPException) as exc_info: + SSOAuthenticationHandler._validate_return_to("https://cp.example.com:8443/ui") + assert exc_info.value.status_code == 400 + + def test_allows_explicit_default_port(self, monkeypatch): + """https://host:443 should match https://host (default port normalisation).""" + monkeypatch.setattr( + "litellm.proxy.proxy_server.general_settings", + {"control_plane_url": "https://cp.example.com"}, + ) + SSOAuthenticationHandler._validate_return_to("https://cp.example.com:443/ui") + + def test_allows_matching_custom_port(self, monkeypatch): + """Both sides on the same custom port should match.""" + monkeypatch.setattr( + "litellm.proxy.proxy_server.general_settings", + {"control_plane_url": "https://cp.example.com:3000"}, + ) + SSOAuthenticationHandler._validate_return_to("https://cp.example.com:3000/ui") + diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index 112a06b173..bd6162f225 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -236,6 +236,217 @@ def test_login_v2_returns_json_on_invalid_json_body(monkeypatch): assert isinstance(data["error"], dict) +def test_login_v3_rejected_without_control_plane_url(monkeypatch): + """v3/login returns 404 when control_plane_url is not configured.""" + mock_prisma_client = MagicMock() + monkeypatch.setattr("litellm.proxy.proxy_server.master_key", "test-master-key") + monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", {}) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) + + client = TestClient(app) + response = client.post( + "/v3/login", + json={"username": "alice", "password": "secret"}, + ) + + assert response.status_code == 404 + assert "control_plane_url" in response.json()["error"]["message"] + + +def test_login_v3_returns_code(monkeypatch): + """v3/login returns an opaque code, not the JWT directly.""" + mock_prisma_client = MagicMock() + monkeypatch.setattr( + "litellm.proxy.auth.login_utils.authenticate_user", + AsyncMock(return_value={"user_id": "test-user"}), + ) + monkeypatch.setattr( + "litellm.proxy.auth.login_utils.create_ui_token_object", + MagicMock(return_value={"user_id": "test-user"}), + ) + monkeypatch.setattr("jwt.encode", MagicMock(return_value="signed-token")) + monkeypatch.setattr("litellm.proxy.proxy_server.master_key", "test-master-key") + monkeypatch.setattr( + "litellm.proxy.proxy_server.general_settings", + {"control_plane_url": "https://cp.example.com"}, + ) + monkeypatch.setattr("litellm.proxy.proxy_server.premium_user", False) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) + mock_config = MagicMock() + mock_config.worker_registry = [] + monkeypatch.setattr("litellm.proxy.proxy_server.proxy_config", mock_config) + monkeypatch.setattr("litellm.proxy.utils.get_server_root_path", lambda: "") + monkeypatch.setattr("litellm.proxy.utils.get_proxy_base_url", lambda: None) + + client = TestClient(app) + response = client.post( + "/v3/login", + json={"username": "alice", "password": "secret"}, + ) + + assert response.status_code == 200 + data = response.json() + assert "code" in data + assert data["expires_in"] == 60 + assert "token" not in data + + +def test_login_v3_exchange_happy_path(monkeypatch): + """Full flow: v3/login returns code, v3/login/exchange redeems it for JWT.""" + mock_prisma_client = MagicMock() + monkeypatch.setattr( + "litellm.proxy.auth.login_utils.authenticate_user", + AsyncMock(return_value={"user_id": "test-user"}), + ) + monkeypatch.setattr( + "litellm.proxy.auth.login_utils.create_ui_token_object", + MagicMock(return_value={"user_id": "test-user"}), + ) + monkeypatch.setattr("jwt.encode", MagicMock(return_value="signed-token")) + monkeypatch.setattr("litellm.proxy.proxy_server.master_key", "test-master-key") + monkeypatch.setattr( + "litellm.proxy.proxy_server.general_settings", + {"control_plane_url": "https://cp.example.com"}, + ) + monkeypatch.setattr("litellm.proxy.proxy_server.premium_user", False) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) + mock_config = MagicMock() + mock_config.worker_registry = [] + monkeypatch.setattr("litellm.proxy.proxy_server.proxy_config", mock_config) + monkeypatch.setattr("litellm.proxy.utils.get_server_root_path", lambda: "") + monkeypatch.setattr("litellm.proxy.utils.get_proxy_base_url", lambda: None) + + client = TestClient(app) + + # Step 1: login — get code + login_response = client.post( + "/v3/login", + json={"username": "alice", "password": "secret"}, + ) + assert login_response.status_code == 200 + code = login_response.json()["code"] + + # Step 2: exchange — get JWT + exchange_response = client.post( + "/v3/login/exchange", + json={"code": code}, + ) + assert exchange_response.status_code == 200 + exchange_data = exchange_response.json() + assert exchange_data["token"] == "signed-token" + assert "redirect_url" in exchange_data + assert exchange_response.cookies.get("token") == "signed-token" + + +def test_login_v3_exchange_single_use(monkeypatch): + """Code can only be redeemed once.""" + mock_prisma_client = MagicMock() + monkeypatch.setattr( + "litellm.proxy.auth.login_utils.authenticate_user", + AsyncMock(return_value={"user_id": "test-user"}), + ) + monkeypatch.setattr( + "litellm.proxy.auth.login_utils.create_ui_token_object", + MagicMock(return_value={"user_id": "test-user"}), + ) + monkeypatch.setattr("jwt.encode", MagicMock(return_value="signed-token")) + monkeypatch.setattr("litellm.proxy.proxy_server.master_key", "test-master-key") + monkeypatch.setattr( + "litellm.proxy.proxy_server.general_settings", + {"control_plane_url": "https://cp.example.com"}, + ) + monkeypatch.setattr("litellm.proxy.proxy_server.premium_user", False) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) + mock_config = MagicMock() + mock_config.worker_registry = [] + monkeypatch.setattr("litellm.proxy.proxy_server.proxy_config", mock_config) + monkeypatch.setattr("litellm.proxy.utils.get_server_root_path", lambda: "") + monkeypatch.setattr("litellm.proxy.utils.get_proxy_base_url", lambda: None) + + client = TestClient(app) + + login_response = client.post( + "/v3/login", + json={"username": "alice", "password": "secret"}, + ) + code = login_response.json()["code"] + + # First exchange succeeds + first = client.post("/v3/login/exchange", json={"code": code}) + assert first.status_code == 200 + + # Second exchange fails + second = client.post("/v3/login/exchange", json={"code": code}) + assert second.status_code == 401 + + +def test_login_v3_exchange_invalid_code(monkeypatch): + """Random code returns 401.""" + monkeypatch.setattr( + "litellm.proxy.proxy_server.general_settings", + {"control_plane_url": "https://cp.example.com"}, + ) + client = TestClient(app) + response = client.post( + "/v3/login/exchange", + json={"code": "nonexistent-code"}, + ) + assert response.status_code == 401 + + +def test_login_v3_exchange_rejected_without_control_plane_url(monkeypatch): + """v3/login/exchange returns 404 when control_plane_url is not configured.""" + monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", {}) + + client = TestClient(app) + response = client.post( + "/v3/login/exchange", + json={"code": "some-code"}, + ) + + assert response.status_code == 404 + assert "control_plane_url" in response.json()["error"]["message"] + + +def test_login_v3_returns_json_on_proxy_exception(monkeypatch): + """Test that /v3/login returns JSON error when ProxyException is raised""" + from litellm.proxy._types import ProxyErrorTypes, ProxyException + + mock_prisma_client = MagicMock() + mock_authenticate_user = AsyncMock( + side_effect=ProxyException( + message="Invalid credentials", + type=ProxyErrorTypes.auth_error, + param="password", + code=401, + ) + ) + + monkeypatch.setattr( + "litellm.proxy.auth.login_utils.authenticate_user", + mock_authenticate_user, + ) + monkeypatch.setattr("litellm.proxy.proxy_server.master_key", "test-master-key") + monkeypatch.setattr( + "litellm.proxy.proxy_server.general_settings", + {"control_plane_url": "https://cp.example.com"}, + ) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) + + client = TestClient(app) + response = client.post( + "/v3/login", + json={"username": "alice", "password": "wrong"}, + ) + + assert response.status_code == 401 + assert response.headers["content-type"] == "application/json" + data = response.json() + assert "error" in data + assert data["error"]["message"] == "Invalid credentials" + assert data["error"]["type"] == "auth_error" + + def test_fallback_login_has_no_deprecation_banner(client_no_auth): response = client_no_auth.get("/fallback/login") diff --git a/ui/litellm-dashboard/src/app/(dashboard)/api-reference/page.tsx b/ui/litellm-dashboard/src/app/(dashboard)/api-reference/page.tsx index f9f26c0eac..02bed1adbe 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/api-reference/page.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/api-reference/page.tsx @@ -1,16 +1,10 @@ "use client"; import APIReferenceView from "@/app/(dashboard)/api-reference/APIReferenceView"; -import { useState } from "react"; - -interface ProxySettings { - PROXY_BASE_URL: string; - PROXY_LOGOUT_URL: string; - LITELLM_UI_API_DOC_BASE_URL?: string | null; -} +import useProxySettings from "@/app/(dashboard)/hooks/proxySettings/useProxySettings"; const APIReferencePage = () => { - const [proxySettings, setProxySettings] = useState({ PROXY_BASE_URL: "", PROXY_LOGOUT_URL: "" }); + const proxySettings = useProxySettings(); return ; }; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/components/Sidebar2.tsx b/ui/litellm-dashboard/src/app/(dashboard)/components/Sidebar2.tsx index 18ab475f22..27a6e6c13b 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/components/Sidebar2.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/components/Sidebar2.tsx @@ -195,7 +195,7 @@ const menuItems: MenuItemCfg[] = [ icon: , roles: all_admin_roles, }, - { key: "14", page: "api_ref", label: "API Reference", icon: }, + { key: "14", page: "api-reference", label: "API Reference", icon: }, { key: "16", page: "model-hub-table", diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/common/queryKeysFactory.test.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/common/queryKeysFactory.test.ts new file mode 100644 index 0000000000..39afd04409 --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/common/queryKeysFactory.test.ts @@ -0,0 +1,34 @@ +import { describe, it, expect } from "vitest"; +import { createQueryKeys } from "./queryKeysFactory"; + +describe("createQueryKeys", () => { + const keys = createQueryKeys("books"); + + it("should return the resource name as the base key", () => { + expect(keys.all).toEqual(["books"]); + }); + + it("should generate a lists key", () => { + expect(keys.lists()).toEqual(["books", "list"]); + }); + + it("should generate a list key with params", () => { + expect(keys.list({ page: 1, limit: 10 })).toEqual([ + "books", + "list", + { params: { page: 1, limit: 10 } }, + ]); + }); + + it("should generate a list key with undefined params when none provided", () => { + expect(keys.list()).toEqual(["books", "list", { params: undefined }]); + }); + + it("should generate a details key", () => { + expect(keys.details()).toEqual(["books", "detail"]); + }); + + it("should generate a detail key for a specific ID", () => { + expect(keys.detail("123")).toEqual(["books", "detail", "123"]); + }); +}); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/login/useLogin.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/login/useLogin.ts index a15b4a06d1..be53b1c80a 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/hooks/login/useLogin.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/login/useLogin.ts @@ -3,8 +3,8 @@ import { loginCall, LoginRequest } from "@/components/networking"; export const useLogin = () => { return useMutation({ - mutationFn: async ({ username, password }: LoginRequest) => { - const result = await loginCall(username, password); + mutationFn: async ({ username, password, useV3 }: LoginRequest) => { + const result = await loginCall(username, password, useV3); return result; }, }); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/proxySettings/useProxySettings.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/proxySettings/useProxySettings.ts new file mode 100644 index 0000000000..d4fb307385 --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/proxySettings/useProxySettings.ts @@ -0,0 +1,21 @@ +import { useState, useEffect } from "react"; +import { fetchProxySettings } from "@/utils/proxyUtils"; +import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; + +export default function useProxySettings() { + const { accessToken } = useAuthorized(); + const [proxySettings, setProxySettings] = useState({ + PROXY_BASE_URL: "", + PROXY_LOGOUT_URL: "", + LITELLM_UI_API_DOC_BASE_URL: null as string | null, + }); + + useEffect(() => { + if (!accessToken) return; + fetchProxySettings(accessToken).then((settings) => { + if (settings) setProxySettings(settings); + }); + }, [accessToken]); + + return proxySettings; +} diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/uiConfig/useUIConfig.test.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/uiConfig/useUIConfig.test.ts index aba5dddf13..b05bae1e81 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/hooks/uiConfig/useUIConfig.test.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/uiConfig/useUIConfig.test.ts @@ -28,6 +28,8 @@ const mockUIConfig: LiteLLMWellKnownUiConfig = { proxy_base_url: "https://proxy.example.com", auto_redirect_to_sso: true, admin_ui_disabled: false, + is_control_plane: false, + workers: [], }; describe("useUIConfig", () => { @@ -102,6 +104,8 @@ describe("useUIConfig", () => { auto_redirect_to_sso: false, sso_configured: false, admin_ui_disabled: true, + is_control_plane: false, + workers: [], }; // Mock successful API call with different data diff --git a/ui/litellm-dashboard/src/app/(dashboard)/teams/components/TeamsHeaderTabs.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/teams/components/TeamsHeaderTabs.test.tsx new file mode 100644 index 0000000000..50a7f10f04 --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/teams/components/TeamsHeaderTabs.test.tsx @@ -0,0 +1,54 @@ +import { render, screen } from "@testing-library/react"; +import React from "react"; +import { describe, expect, it, vi } from "vitest"; +import TeamsHeaderTabs from "./TeamsHeaderTabs"; + +vi.mock("@tremor/react", () => ({ + TabGroup: ({ children, ...props }: any) =>
{children}
, + TabList: ({ children, ...props }: any) =>
{children}
, + Tab: ({ children, ...props }: any) => , + TabPanels: ({ children, ...props }: any) =>
{children}
, + Text: ({ children, ...props }: any) => {children}, + Icon: ({ onClick, ...props }: any) => + + ); + } + + return ( + + columns={teamColumns} + dataSource={displayTeams} + rowKey="team_id" + pagination={false} + onChange={handleTableSort} + locale={{ + emptyText: ( +
+ +
+ No teams yet +
+
+ + Create your first team to organize members and manage access to models. + +
+ {canCreateOrManageTeams(userRole, userID, organizations) && ( + + )} +
+ ), + }} + scroll={{ x: 1000 }} + size="middle" + /> + ); + }; + + const tabItems = [ + { + key: "your-teams", + label: "Your Teams", + children: ( + <> + + + + } + suffix={isSearching ? : null} + placeholder="Search teams by name..." + onChange={(e) => handleSearchChange(e.target.value)} + allowClear + style={{ maxWidth: 400 }} + /> + handleFilterChange("organization_id", value || "")} + loading={isLoading} + /> + + { + setCurrentPage(page); + setPageSize(size); + fetchTeamsV2({ page, size }); + }} + size="small" + showTotal={(total) => `${total} teams`} + showSizeChanger + pageSizeOptions={["10", "20", "50"]} + /> + + + {renderTeamsContent()} + + + + + ), + }, + { + key: "available-teams", + label: "Available Teams", + children: , + }, + ...(isProxyAdminRole(userRole || "") + ? [ + { + key: "default-settings", + label: "Default Team Settings", + children: , + }, + ] + : []), + ]; + return ( -
- - - {canCreateOrManageTeams(userRole, userID, organizations) && ( - - )} - {selectedTeamId ? ( - { - setTeams((teams) => { - if (teams == null) { - return teams; - } - const updated = teams.map((team) => { - if (data.team_id === team.team_id) { - return updateExistingKeys(team, data); - } - return team; - }); - // Minimal fix: refresh the full team list after an update - if (accessToken) { - fetchTeams(accessToken, userID, userRole, currentOrg, setTeams); - } - return updated; - }); - }} - onClose={() => { - setSelectedTeamId(null); - setEditTeam(false); - }} - accessToken={accessToken} - is_team_admin={is_team_admin(teams?.find((team) => team.team_id === selectedTeamId))} - is_proxy_admin={userRole == "Admin"} - userModels={userModels} - editTeam={editTeam} - premiumUser={premiumUser} - /> - ) : ( - - -
- Your Teams - Available Teams - {isProxyAdminRole(userRole || "") && Default Team Settings} -
-
- {lastRefreshed && Last Refreshed: {lastRefreshed}} - -
- - - - - Click on “Team ID” to view team details and manage team members. - - - - -
-
- {/* Search and Filter Controls */} -
- {/* Team Alias Search */} - handleFilterChange("team_alias", value)} - icon={Search} - /> + + {selectedTeamId ? ( + { + setTeams((teams) => { + if (teams == null) { + return teams; + } + return teams.map((team) => { + if (data.team_id === team.team_id) { + return updateExistingKeys(team, data); + } + return team; + }); + }); + fetchTeamsV2(); + }} + onClose={() => { + setSelectedTeamId(null); + setEditTeam(false); + }} + accessToken={accessToken} + is_team_admin={is_team_admin(teams?.find((team) => team.team_id === selectedTeamId))} + is_proxy_admin={userRole == "Admin"} + userModels={userModels} + editTeam={editTeam} + premiumUser={premiumUser} + /> + ) : ( + <> + + + + <TeamOutlined style={{ marginRight: 8 }} /> + Teams + + + Manage teams, members, and their access to models and budgets + + + {canCreateOrManageTeams(userRole, userID, organizations) && ( + + )} + - {/* Filter Button */} - setShowFilters(!showFilters)} - active={showFilters} - hasActiveFilters={!!(filters.team_id || filters.team_alias || filters.organization_id)} - /> + + + )} - {/* Reset Filters Button */} - -
- - {/* Additional Filters */} - {showFilters && ( -
- {/* Team ID Search */} - handleFilterChange("team_id", value)} - icon={User} - /> - - {/* Organization Dropdown */} -
- -
-
- )} -
-
- - - - Team Name - Team ID - Created - Spend (USD) - Budget (USD) - Models - Organization - Info - Actions - - - - - {teams && teams.length > 0 ? ( - teams - .filter((team) => { - if (!currentOrg) return true; - return team.organization_id === currentOrg.organization_id; - }) - .sort((a, b) => new Date(b.created_at).getTime() - new Date(a.created_at).getTime()) - .map((team: any) => ( - - - {team["team_alias"]} - - -
- - - -
- - - {team.created_at ? new Date(team.created_at).toLocaleDateString() : "N/A"} - - - {formatNumberWithCommas(team["spend"], 4)} - - - {team["max_budget"] !== null && team["max_budget"] !== undefined - ? team["max_budget"] - : "No limit"} - - 3 ? "px-0" : ""} - > -
- {Array.isArray(team.models) ? ( -
- {team.models.length === 0 ? ( - - All Proxy Models - - ) : ( - <> -
- {team.models.length > 3 && ( -
- { - setExpandedAccordions((prev) => ({ - ...prev, - [team.team_id]: !prev[team.team_id], - })); - }} - /> -
- )} -
- {team.models.slice(0, 3).map((model: string, index: number) => - model === "all-proxy-models" ? ( - - All Proxy Models - - ) : ( - - - {model.length > 30 - ? `${getModelDisplayName(model).slice(0, 30)}...` - : getModelDisplayName(model)} - - - ), - )} - {team.models.length > 3 && !expandedAccordions[team.team_id] && ( - - - +{team.models.length - 3}{" "} - {team.models.length - 3 === 1 ? "more model" : "more models"} - - - )} - {expandedAccordions[team.team_id] && ( -
- {team.models.slice(3).map((model: string, index: number) => - model === "all-proxy-models" ? ( - - All Proxy Models - - ) : ( - - - {model.length > 30 - ? `${getModelDisplayName(model).slice(0, 30)}...` - : getModelDisplayName(model)} - - - ), - )} -
- )} -
-
- - )} -
- ) : null} -
- - - - {getOrganizationAlias(team.organization_id, organizationsData || organizations)} - - - - {perTeamInfo && - team.team_id && - perTeamInfo[team.team_id] && - perTeamInfo[team.team_id].keys && - perTeamInfo[team.team_id].keys.length}{" "} - Keys - - - {perTeamInfo && - team.team_id && - perTeamInfo[team.team_id] && - perTeamInfo[team.team_id].team_info && - perTeamInfo[team.team_id].team_info.members_with_roles && - perTeamInfo[team.team_id].team_info.members_with_roles.length}{" "} - Members - - - - {userRole == "Admin" ? ( - <> - { - setSelectedTeamId(team.team_id); - setEditTeam(true); - }} - dataTestId="edit-team-button" - tooltipText="Edit team" - /> - handleDelete(team)} - dataTestId="delete-team-button" - tooltipText="Delete team" - /> - - ) : null} - - - )) - ) : ( - - -
- No teams found - Adjust your filters or create a new team -
- - - )} - -
- - - - - - - - - {isProxyAdminRole(userRole || "") && ( - - - - )} - - - )} - {canCreateOrManageTeams(userRole, userID, organizations) && ( + {canCreateOrManageTeams(userRole, userID, organizations) && ( = ({ : "" } > - = ({ optionFilterProp="children" > {adminOrgs?.map((org) => ( - + {org.organization_alias}{" "} ({org.organization_id}) - + ))} - + {/* Show message when org admin needs to select organization */} {isOrgAdmin && !isSingleOrg && adminOrgs.length > 1 && (
- + Please select an organization to create a team for. You can only create teams within organizations where you are an admin. @@ -1190,11 +1211,11 @@ const Teams: React.FC = ({ - - daily - weekly - monthly - + @@ -1313,7 +1334,7 @@ const Teams: React.FC = ({ className="mt-8" help="Select existing guardrails or enter new ones" > - = ({ className="mt-8" help="Select existing policies or enter new ones" > - = ({
- + Create custom aliases for models that can be used by team members in API calls. This allows you to create shortcuts for specific models. @@ -1548,14 +1569,12 @@ const Teams: React.FC = ({
- Create Team +
)} - - -
+ ); }; diff --git a/ui/litellm-dashboard/src/components/Projects/ProjectModals/CreateProjectModal.tsx b/ui/litellm-dashboard/src/components/Projects/ProjectModals/CreateProjectModal.tsx index e490f89303..bbf56e4930 100644 --- a/ui/litellm-dashboard/src/components/Projects/ProjectModals/CreateProjectModal.tsx +++ b/ui/litellm-dashboard/src/components/Projects/ProjectModals/CreateProjectModal.tsx @@ -1,5 +1,6 @@ -import { Modal, Form, Button, Typography, message } from "antd"; +import { Modal, Form, Button, Typography } from "antd"; import { FolderAddOutlined } from "@ant-design/icons"; +import MessageManager from "@/components/molecules/message_manager"; import { useCreateProject, ProjectCreateParams, @@ -32,12 +33,12 @@ export function CreateProjectModal({ createMutation.mutate(params, { onSuccess: () => { - message.success("Project created successfully"); + MessageManager.success("Project created successfully"); form.resetFields(); onClose(); }, onError: (error) => { - message.error(error.message || "Failed to create project"); + MessageManager.error(error.message || "Failed to create project"); }, }); } catch (error) { diff --git a/ui/litellm-dashboard/src/components/Projects/ProjectModals/EditProjectModal.tsx b/ui/litellm-dashboard/src/components/Projects/ProjectModals/EditProjectModal.tsx index 75f56b1373..dc3b43ef73 100644 --- a/ui/litellm-dashboard/src/components/Projects/ProjectModals/EditProjectModal.tsx +++ b/ui/litellm-dashboard/src/components/Projects/ProjectModals/EditProjectModal.tsx @@ -1,6 +1,7 @@ import { useEffect } from "react"; -import { Modal, Form, Button, Typography, message } from "antd"; +import { Modal, Form, Button, Typography } from "antd"; import { SaveOutlined } from "@ant-design/icons"; +import MessageManager from "@/components/molecules/message_manager"; import { ProjectResponse } from "@/app/(dashboard)/hooks/projects/useProjects"; import { useUpdateProject, @@ -80,12 +81,12 @@ export function EditProjectModal({ { projectId: project.project_id, params }, { onSuccess: () => { - message.success("Project updated successfully"); + MessageManager.success("Project updated successfully"); onSuccess?.(); onClose(); }, onError: (error) => { - message.error(error.message || "Failed to update project"); + MessageManager.error(error.message || "Failed to update project"); }, }, ); diff --git a/ui/litellm-dashboard/src/components/SearchTools/SearchToolTester.tsx b/ui/litellm-dashboard/src/components/SearchTools/SearchToolTester.tsx index 28b34b0608..d1cb5077b1 100644 --- a/ui/litellm-dashboard/src/components/SearchTools/SearchToolTester.tsx +++ b/ui/litellm-dashboard/src/components/SearchTools/SearchToolTester.tsx @@ -1,5 +1,6 @@ import React, { useState } from "react"; -import { Button, Input, Typography, Spin, message } from "antd"; +import { Button, Input, Typography, Spin } from "antd"; +import MessageManager from "@/components/molecules/message_manager"; import { SearchOutlined, LoadingOutlined } from "@ant-design/icons"; import { searchToolQueryCall } from "../networking"; import NotificationsManager from "../molecules/notifications_manager"; @@ -39,7 +40,7 @@ export const SearchToolTester: React.FC = ({ searchToolNa const handleSearch = async () => { if (!query.trim()) { - message.warning("Please enter a search query"); + MessageManager.warning("Please enter a search query"); return; } diff --git a/ui/litellm-dashboard/src/components/Settings/RouterSettings/Fallbacks/AddFallbacks.tsx b/ui/litellm-dashboard/src/components/Settings/RouterSettings/Fallbacks/AddFallbacks.tsx index 34f059516a..e55089d27d 100644 --- a/ui/litellm-dashboard/src/components/Settings/RouterSettings/Fallbacks/AddFallbacks.tsx +++ b/ui/litellm-dashboard/src/components/Settings/RouterSettings/Fallbacks/AddFallbacks.tsx @@ -5,8 +5,9 @@ */ import { Button as TremorButton } from "@tremor/react"; -import { Button, message } from "antd"; +import { Button } from "antd"; import React, { useEffect, useState } from "react"; +import MessageManager from "@/components/molecules/message_manager"; import NotificationManager from "../../../molecules/notifications_manager"; import { fetchAvailableModels, ModelGroup } from "../../../playground/llm_calls/fetch_models"; import { AddFallbacksModal } from "./AddFallbacksModal"; @@ -90,7 +91,7 @@ export default function AddFallbacks({ (g) => !g.primaryModel || g.fallbackModels.length === 0, ); if (invalidGroups.length > 0) { - message.error( + MessageManager.error( `Please complete configuration for all groups. ${invalidGroups.length} group(s) incomplete.`, ); return; diff --git a/ui/litellm-dashboard/src/components/Settings/RouterSettings/Fallbacks/FallbackSelectionForm.tsx b/ui/litellm-dashboard/src/components/Settings/RouterSettings/Fallbacks/FallbackSelectionForm.tsx index 08b031c683..0482161bc6 100644 --- a/ui/litellm-dashboard/src/components/Settings/RouterSettings/Fallbacks/FallbackSelectionForm.tsx +++ b/ui/litellm-dashboard/src/components/Settings/RouterSettings/Fallbacks/FallbackSelectionForm.tsx @@ -5,9 +5,10 @@ */ import { Button } from "@tremor/react"; -import { message, Tabs } from "antd"; +import { Tabs } from "antd"; import { Plus } from "lucide-react"; import React, { useEffect, useState } from "react"; +import MessageManager from "@/components/molecules/message_manager"; import { FallbackGroup, FallbackGroupConfig } from "./FallbackGroupConfig"; interface FallbackSelectionFormProps { @@ -60,7 +61,7 @@ export function FallbackSelectionForm({ const handleRemoveGroup = (targetId: string) => { if (groups.length === 1) { - message.warning("At least one group is required"); + MessageManager.warning("At least one group is required"); return; } const newGroups = groups.filter((g) => g.id !== targetId); diff --git a/ui/litellm-dashboard/src/components/agents/add_agent_form.tsx b/ui/litellm-dashboard/src/components/agents/add_agent_form.tsx index c5518596b8..b6d96f0445 100644 --- a/ui/litellm-dashboard/src/components/agents/add_agent_form.tsx +++ b/ui/litellm-dashboard/src/components/agents/add_agent_form.tsx @@ -1,5 +1,6 @@ import React, { useState, useEffect } from "react"; -import { Modal, Form, message, Select, Input, Steps, Radio, Tag, Divider, Switch, InputNumber, Collapse } from "antd"; +import { Modal, Form, Select, Input, Steps, Radio, Tag, Divider, Switch, InputNumber, Collapse } from "antd"; +import MessageManager from "@/components/molecules/message_manager"; import { Button } from "@tremor/react"; import { CheckCircleFilled, KeyOutlined, RobotOutlined, AppstoreOutlined, InfoCircleOutlined } from "@ant-design/icons"; import CreatedKeyDisplay from "../shared/CreatedKeyDisplay"; @@ -216,7 +217,7 @@ const AddAgentForm: React.FC = ({ const handleCreateAgent = async () => { if (!accessToken) { - message.error("No access token available"); + MessageManager.error("No access token available"); return; } @@ -226,7 +227,7 @@ const AddAgentForm: React.FC = ({ const values = { ...form.getFieldsValue(true) }; const agentData = buildAgentData(values); if (!agentData) { - message.error("Failed to build agent data"); + MessageManager.error("Failed to build agent data"); setIsSubmitting(false); return; } @@ -301,7 +302,7 @@ const AddAgentForm: React.FC = ({ setCreatedKeyValue(keyResponse.key || null); } else if (keyAssignOption === "existing_key") { if (!selectedExistingKey) { - message.error("Please select an existing key to assign"); + MessageManager.error("Please select an existing key to assign"); setIsSubmitting(false); return; } @@ -318,7 +319,7 @@ const AddAgentForm: React.FC = ({ } catch (error) { console.error("Error creating agent:", error); const errorMessage = error instanceof Error ? error.message : String(error); - message.error(errorMessage ? `Failed to create agent: ${errorMessage}` : "Failed to create agent"); + MessageManager.error(errorMessage ? `Failed to create agent: ${errorMessage}` : "Failed to create agent"); } finally { setIsSubmitting(false); } diff --git a/ui/litellm-dashboard/src/components/agents/agent_info.tsx b/ui/litellm-dashboard/src/components/agents/agent_info.tsx index b41e318a76..d543be8356 100644 --- a/ui/litellm-dashboard/src/components/agents/agent_info.tsx +++ b/ui/litellm-dashboard/src/components/agents/agent_info.tsx @@ -1,6 +1,7 @@ import React, { useState, useEffect } from "react"; import { Card, Title, Text, Button as TremorButton, Tab, TabGroup, TabList, TabPanel, TabPanels} from "@tremor/react"; -import { Form, Input, InputNumber, Button as AntButton, message, Spin, Descriptions, Divider } from "antd"; +import { Form, Input, InputNumber, Button as AntButton, Spin, Descriptions, Divider } from "antd"; +import MessageManager from "@/components/molecules/message_manager"; import { ArrowLeftIcon } from "@heroicons/react/outline"; import { getAgentInfo, patchAgentCall, getAgentCreateMetadata, AgentCreateInfo } from "../networking"; import { Agent } from "./types"; @@ -72,7 +73,7 @@ const AgentInfoView: React.FC = ({ } } catch (error) { console.error("Error fetching agent info:", error); - message.error("Failed to load agent information"); + MessageManager.error("Failed to load agent information"); } finally { setIsLoading(false); } @@ -111,12 +112,12 @@ const AgentInfoView: React.FC = ({ } await patchAgentCall(accessToken, agentId, updateData); - message.success("Agent updated successfully"); + MessageManager.success("Agent updated successfully"); setIsEditing(false); fetchAgentInfo(); } catch (error) { console.error("Error updating agent:", error); - message.error("Failed to update agent"); + MessageManager.error("Failed to update agent"); } finally { setIsSaving(false); } diff --git a/ui/litellm-dashboard/src/components/chat/ChatPage.tsx b/ui/litellm-dashboard/src/components/chat/ChatPage.tsx index ccf39d2147..b547f74917 100644 --- a/ui/litellm-dashboard/src/components/chat/ChatPage.tsx +++ b/ui/litellm-dashboard/src/components/chat/ChatPage.tsx @@ -1,7 +1,8 @@ "use client"; import React, { useCallback, useEffect, useRef, useState, useLayoutEffect } from "react"; -import { Tooltip, Skeleton, Popover, message } from "antd"; +import { Tooltip, Skeleton, Popover } from "antd"; +import MessageManager from "@/components/molecules/message_manager"; import { SettingOutlined, PlusOutlined, @@ -212,7 +213,7 @@ const ChatPage: React.FC = ({ accessToken, userRole, userId, user localStorage.setItem(LOCALSTORAGE_MODEL_KEY, JSON.stringify([names[0]])); } }) - .catch(() => message.error("Could not load models")) + .catch(() => MessageManager.error("Could not load models")) .finally(() => setIsLoadingModels(false)); }, [accessToken]); diff --git a/ui/litellm-dashboard/src/components/chat/MCPAppsPanel.tsx b/ui/litellm-dashboard/src/components/chat/MCPAppsPanel.tsx index b805cab71e..d25db73ae4 100644 --- a/ui/litellm-dashboard/src/components/chat/MCPAppsPanel.tsx +++ b/ui/litellm-dashboard/src/components/chat/MCPAppsPanel.tsx @@ -5,7 +5,7 @@ import { Spin, Input, Button, Skeleton } from "antd"; import { SearchOutlined, ArrowLeftOutlined, RightOutlined, ToolOutlined, CheckCircleOutlined } from "@ant-design/icons"; import { deleteMCPOAuthUserCredential, fetchMCPServers, getMCPOAuthUserCredentialStatus, listMCPTools } from "../networking"; import { AUTH_TYPE, MCPServer, MCPTool, handleTransport } from "../mcp_tools/types"; -import { message } from "antd"; +import MessageManager from "@/components/molecules/message_manager"; import { useUserMcpOAuthFlow } from "@/hooks/useUserMcpOAuthFlow"; // ── OAuth2 connect button ───────────────────────────────────────────────────── @@ -198,7 +198,7 @@ const MCPAppsPanel: React.FC = ({ accessToken, selectedServers, onChange const idToFetch = serverId ?? serverName; const result = await listMCPTools(accessToken, idToFetch); if (result?.error) { - message.warning(`Could not load tools for ${serverName}`); + MessageManager.warning(`Could not load tools for ${serverName}`); return; } // Use the ref so we read the most up-to-date list; guard against duplicates @@ -207,7 +207,7 @@ const MCPAppsPanel: React.FC = ({ accessToken, selectedServers, onChange onChange([...selectedServersRef.current, serverName]); } } catch { - message.warning(`Could not load tools for ${serverName}`); + MessageManager.warning(`Could not load tools for ${serverName}`); } finally { setTogglingOn((prev) => { const next = new Set(prev); diff --git a/ui/litellm-dashboard/src/components/chat/MCPConnectPicker.tsx b/ui/litellm-dashboard/src/components/chat/MCPConnectPicker.tsx index a52ce18156..6ef3aecc46 100644 --- a/ui/litellm-dashboard/src/components/chat/MCPConnectPicker.tsx +++ b/ui/litellm-dashboard/src/components/chat/MCPConnectPicker.tsx @@ -1,5 +1,6 @@ import React, { useEffect, useState } from "react"; -import { Switch, Spin, message } from "antd"; +import { Switch, Spin } from "antd"; +import MessageManager from "@/components/molecules/message_manager"; import { fetchMCPServers, listMCPTools } from "../networking"; import { MCPServer } from "../mcp_tools/types"; @@ -57,7 +58,7 @@ const MCPConnectPicker: React.FC = ({ accessToken, selectedServers, onCha const result = await listMCPTools(accessToken, serverName); // listMCPTools never throws; it returns { tools, error, message } on failure if (result?.error) { - message.warning( + MessageManager.warning( `Could not load tools for ${serverName} — it will be excluded from this message.` ); // Do not add to selectedServers @@ -65,7 +66,7 @@ const MCPConnectPicker: React.FC = ({ accessToken, selectedServers, onCha } onChange([...selectedServers, serverName]); } catch { - message.warning( + MessageManager.warning( `Could not load tools for ${serverName} — it will be excluded from this message.` ); // Do not add to selectedServers diff --git a/ui/litellm-dashboard/src/components/chat/MCPCredentialsTab.tsx b/ui/litellm-dashboard/src/components/chat/MCPCredentialsTab.tsx index 363ec8c0e4..e2d879b22a 100644 --- a/ui/litellm-dashboard/src/components/chat/MCPCredentialsTab.tsx +++ b/ui/litellm-dashboard/src/components/chat/MCPCredentialsTab.tsx @@ -8,7 +8,8 @@ */ import React, { useCallback, useEffect, useState } from "react"; -import { Spin, message } from "antd"; +import { Spin } from "antd"; +import MessageManager from "@/components/molecules/message_manager"; import { DeleteOutlined, LinkOutlined } from "@ant-design/icons"; import { Badge, Table, TableBody, TableCell, TableHead, TableHeaderCell, TableRow } from "@tremor/react"; import { @@ -77,7 +78,7 @@ const MCPCredentialsTab: React.FC = ({ accessToken }) => { await deleteMCPOAuthUserCredential(accessToken, serverId); setCredentials((prev) => prev.filter((c) => c.server_id !== serverId)); } catch { - message.error("Failed to revoke connection. Please try again."); + MessageManager.error("Failed to revoke connection. Please try again."); } finally { setRevoking((prev) => { const n = new Set(prev); n.delete(serverId); return n; }); } diff --git a/ui/litellm-dashboard/src/components/claude_code_plugins/add_plugin_form.tsx b/ui/litellm-dashboard/src/components/claude_code_plugins/add_plugin_form.tsx index 217851f128..d5e417a9a8 100644 --- a/ui/litellm-dashboard/src/components/claude_code_plugins/add_plugin_form.tsx +++ b/ui/litellm-dashboard/src/components/claude_code_plugins/add_plugin_form.tsx @@ -1,5 +1,6 @@ import React, { useState } from "react"; -import { Modal, Form, Input, Select, message } from "antd"; +import { Modal, Form, Input, Select } from "antd"; +import MessageManager from "@/components/molecules/message_manager"; import { Button } from "@tremor/react"; import { registerClaudeCodePlugin } from "../networking"; import { @@ -43,13 +44,13 @@ const AddPluginForm: React.FC = ({ const handleSubmit = async (values: any) => { if (!accessToken) { - message.error("No access token available"); + MessageManager.error("No access token available"); return; } // Validate plugin name if (!validatePluginName(values.name)) { - message.error( + MessageManager.error( "Plugin name must be kebab-case (lowercase letters, numbers, and hyphens only)" ); return; @@ -57,7 +58,7 @@ const AddPluginForm: React.FC = ({ // Validate semantic version if provided if (values.version && !isValidSemanticVersion(values.version)) { - message.error( + MessageManager.error( "Version must be in semantic versioning format (e.g., 1.0.0)" ); return; @@ -65,13 +66,13 @@ const AddPluginForm: React.FC = ({ // Validate email if provided if (values.authorEmail && !isValidEmail(values.authorEmail)) { - message.error("Invalid email format"); + MessageManager.error("Invalid email format"); return; } // Validate homepage URL if provided if (values.homepage && !isValidUrl(values.homepage)) { - message.error("Invalid homepage URL format"); + MessageManager.error("Invalid homepage URL format"); return; } @@ -119,14 +120,14 @@ const AddPluginForm: React.FC = ({ } await registerClaudeCodePlugin(accessToken, pluginData); - message.success("Plugin registered successfully"); + MessageManager.success("Plugin registered successfully"); form.resetFields(); setSourceType("github"); onSuccess(); onClose(); } catch (error) { console.error("Error registering plugin:", error); - message.error("Failed to register plugin"); + MessageManager.error("Failed to register plugin"); } finally { setIsSubmitting(false); } diff --git a/ui/litellm-dashboard/src/components/common_components/IconActionButton/TableIconActionButtons/TableIconActionButton.tsx b/ui/litellm-dashboard/src/components/common_components/IconActionButton/TableIconActionButtons/TableIconActionButton.tsx index 488913a734..2f146aab72 100644 --- a/ui/litellm-dashboard/src/components/common_components/IconActionButton/TableIconActionButtons/TableIconActionButton.tsx +++ b/ui/litellm-dashboard/src/components/common_components/IconActionButton/TableIconActionButtons/TableIconActionButton.tsx @@ -6,6 +6,7 @@ import { ChevronUpIcon, ChevronDownIcon, ExternalLinkIcon, + ClipboardCopyIcon, } from "@heroicons/react/outline"; import { Tooltip } from "antd"; import BaseActionButton from "../BaseActionButton"; @@ -32,6 +33,7 @@ export const TableIconActionButtonMap: Record void; disabled?: boolean; loading?: boolean; + style?: React.CSSProperties; } const OrganizationDropdown: React.FC = ({ @@ -16,16 +19,18 @@ const OrganizationDropdown: React.FC = ({ onChange, disabled, loading, + style, }) => { return ( diff --git a/ui/litellm-dashboard/src/components/key_team_helpers/filter_helpers.test.ts b/ui/litellm-dashboard/src/components/key_team_helpers/filter_helpers.test.ts new file mode 100644 index 0000000000..e3f6a5989f --- /dev/null +++ b/ui/litellm-dashboard/src/components/key_team_helpers/filter_helpers.test.ts @@ -0,0 +1,90 @@ +import { describe, expect, it, vi } from "vitest"; +import { fetchTeamFilterOptions } from "./filter_helpers"; + +const mockKeyListCall = vi.fn(); + +vi.mock("@/components/networking", () => ({ + keyListCall: (...args: unknown[]) => mockKeyListCall(...args), + teamListCall: vi.fn(), + organizationListCall: vi.fn(), +})); + +describe("fetchTeamFilterOptions", () => { + it("should return empty arrays when accessToken is null", async () => { + const result = await fetchTeamFilterOptions(null, "team-1"); + + expect(result).toEqual({ keyAliases: [], organizationIds: [], userIds: [] }); + expect(mockKeyListCall).not.toHaveBeenCalled(); + }); + + it("should return empty arrays when teamId is empty", async () => { + const result = await fetchTeamFilterOptions("tok-123", ""); + + expect(result).toEqual({ keyAliases: [], organizationIds: [], userIds: [] }); + expect(mockKeyListCall).not.toHaveBeenCalled(); + }); + + it("should return sorted key aliases from fetched keys", async () => { + mockKeyListCall.mockResolvedValue({ + keys: [ + { key_alias: "zeta-key" }, + { key_alias: "alpha-key" }, + { key_alias: "mid-key" }, + ], + total_pages: 1, + }); + + const result = await fetchTeamFilterOptions("tok-123", "team-1"); + + expect(result.keyAliases).toEqual(["alpha-key", "mid-key", "zeta-key"]); + }); + + it("should deduplicate organization IDs across pages", async () => { + mockKeyListCall + .mockResolvedValueOnce({ + keys: [ + { organization_id: "org-b" }, + { organization_id: "org-a" }, + ], + total_pages: 2, + }) + .mockResolvedValueOnce({ + keys: [ + { organization_id: "org-a" }, + { organization_id: "org-c" }, + ], + total_pages: 2, + }); + + const result = await fetchTeamFilterOptions("tok-123", "team-1"); + + expect(result.organizationIds).toEqual(["org-a", "org-b", "org-c"]); + }); + + it("should map user IDs with email addresses", async () => { + mockKeyListCall.mockResolvedValue({ + keys: [ + { user_id: "u1", user: { user_email: "alice@example.com" } }, + { user_id: "u2", user: { user_email: "bob@example.com" } }, + ], + total_pages: 1, + }); + + const result = await fetchTeamFilterOptions("tok-123", "team-1"); + + expect(result.userIds).toEqual( + expect.arrayContaining([ + { id: "u1", email: "alice@example.com" }, + { id: "u2", email: "bob@example.com" }, + ]), + ); + }); + + it("should handle API errors gracefully and return empty arrays", async () => { + mockKeyListCall.mockRejectedValue(new Error("Network error")); + + const result = await fetchTeamFilterOptions("tok-123", "team-1"); + + expect(result).toEqual({ keyAliases: [], organizationIds: [], userIds: [] }); + }); +}); diff --git a/ui/litellm-dashboard/src/components/key_team_helpers/transform_key_info.test.ts b/ui/litellm-dashboard/src/components/key_team_helpers/transform_key_info.test.ts new file mode 100644 index 0000000000..a1139addbf --- /dev/null +++ b/ui/litellm-dashboard/src/components/key_team_helpers/transform_key_info.test.ts @@ -0,0 +1,62 @@ +import { describe, it, expect } from "vitest"; +import { transformKeyInfo } from "./transform_key_info"; + +describe("transformKeyInfo", () => { + it("should combine key and info fields into a single object", () => { + const apiResponse = { + key: "sk-abc123", + info: { + token_id: "tok_1", + key_name: "my-key", + spend: 10.5, + }, + }; + const result = transformKeyInfo(apiResponse); + expect(result).toEqual({ + token: "sk-abc123", + token_id: "tok_1", + key_name: "my-key", + spend: 10.5, + }); + }); + + it("should set the token field from the key property", () => { + const apiResponse = { + key: "sk-xyz789", + info: { key_name: "test" }, + }; + const result = transformKeyInfo(apiResponse); + expect(result.token).toBe("sk-xyz789"); + }); + + it("should preserve all info fields in the result", () => { + const apiResponse = { + key: "sk-abc", + info: { + token_id: "tok_2", + key_name: "prod-key", + spend: 42, + models: ["gpt-4"], + team_id: "team-1", + metadata: { env: "production" }, + }, + }; + const result = transformKeyInfo(apiResponse); + expect(result.token_id).toBe("tok_2"); + expect(result.key_name).toBe("prod-key"); + expect(result.spend).toBe(42); + expect(result.models).toEqual(["gpt-4"]); + expect(result.team_id).toBe("team-1"); + expect(result.metadata).toEqual({ env: "production" }); + }); + + it("should handle empty info object", () => { + const apiResponse = { + key: "sk-empty", + info: {}, + }; + const result = transformKeyInfo(apiResponse); + expect(result.token).toBe("sk-empty"); + expect(Object.keys(result)).toContain("token"); + }); +}); diff --git a/ui/litellm-dashboard/src/components/leftnav.tsx b/ui/litellm-dashboard/src/components/leftnav.tsx index d3789fcffa..09ab380942 100644 --- a/ui/litellm-dashboard/src/components/leftnav.tsx +++ b/ui/litellm-dashboard/src/components/leftnav.tsx @@ -232,8 +232,8 @@ const menuGroups: MenuGroup[] = [ groupLabel: "DEVELOPER TOOLS", items: [ { - key: "api_ref", - page: "api_ref", + key: "api-reference", + page: "api-reference", label: "API Reference", icon: , }, diff --git a/ui/litellm-dashboard/src/components/mcp_tools/ByokCredentialModal.tsx b/ui/litellm-dashboard/src/components/mcp_tools/ByokCredentialModal.tsx index 473918c126..58c2a965e0 100644 --- a/ui/litellm-dashboard/src/components/mcp_tools/ByokCredentialModal.tsx +++ b/ui/litellm-dashboard/src/components/mcp_tools/ByokCredentialModal.tsx @@ -1,7 +1,8 @@ "use client"; import React, { useState } from "react"; -import { Modal, Input, Switch, message } from "antd"; +import { Modal, Input, Switch } from "antd"; +import MessageManager from "@/components/molecules/message_manager"; import { KeyOutlined, LockOutlined, @@ -46,7 +47,7 @@ export const ByokCredentialModal: React.FC = ({ const handleAuthorize = async () => { if (!apiKey.trim()) { - message.error("Please enter your API key"); + MessageManager.error("Please enter your API key"); return; } setLoading(true); @@ -63,11 +64,11 @@ export const ByokCredentialModal: React.FC = ({ const err = await response.json(); throw new Error(err?.detail?.error || "Failed to save credential"); } - message.success(`Connected to ${serverDisplayName}`); + MessageManager.success(`Connected to ${serverDisplayName}`); onSuccess(server.server_id); handleClose(); } catch (e: any) { - message.error(e.message || "Failed to connect"); + MessageManager.error(e.message || "Failed to connect"); } finally { setLoading(false); } diff --git a/ui/litellm-dashboard/src/components/molecules/message_manager.tsx b/ui/litellm-dashboard/src/components/molecules/message_manager.tsx new file mode 100644 index 0000000000..e4c1552d5e --- /dev/null +++ b/ui/litellm-dashboard/src/components/molecules/message_manager.tsx @@ -0,0 +1,38 @@ +import { message as staticMessage } from "antd"; +import type { MessageInstance } from "antd/es/message/interface"; + +let messageInstance: MessageInstance | null = null; + +export const setMessageInstance = (instance: MessageInstance) => { + messageInstance = instance; +}; + +const getMessageApi = () => messageInstance || staticMessage; + +const MessageManager = { + success(content: string, duration?: number) { + getMessageApi().success(content, duration); + }, + + error(content: string, duration?: number) { + getMessageApi().error(content, duration); + }, + + warning(content: string, duration?: number) { + getMessageApi().warning(content, duration); + }, + + info(content: string, duration?: number) { + getMessageApi().info(content, duration); + }, + + loading(content: string, duration?: number) { + return getMessageApi().loading(content, duration); + }, + + destroy() { + getMessageApi().destroy(); + }, +}; + +export default MessageManager; diff --git a/ui/litellm-dashboard/src/components/navbar.tsx b/ui/litellm-dashboard/src/components/navbar.tsx index c46a3af5a6..96d6ce613b 100644 --- a/ui/litellm-dashboard/src/components/navbar.tsx +++ b/ui/litellm-dashboard/src/components/navbar.tsx @@ -4,6 +4,7 @@ import { getProxyBaseUrl } from "@/components/networking"; import { useUIConfig } from "@/app/(dashboard)/hooks/uiConfig/useUIConfig"; import { useTheme } from "@/contexts/ThemeContext"; import { clearTokenCookies } from "@/utils/cookieUtils"; +import { clearStoredReturnUrl } from "@/utils/returnUrlUtils"; import { fetchProxySettings } from "@/utils/proxyUtils"; import { MenuFoldOutlined, MenuUnfoldOutlined, MessageOutlined, MoonOutlined, SunOutlined } from "@ant-design/icons"; import { Button, Switch, Tag } from "antd"; @@ -12,6 +13,7 @@ import React, { useEffect, useState } from "react"; import { BlogDropdown } from "./Navbar/BlogDropdown/BlogDropdown"; import { CommunityEngagementButtons } from "./Navbar/CommunityEngagementButtons/CommunityEngagementButtons"; import UserDropdown from "./Navbar/UserDropdown/UserDropdown"; +import WorkerDropdown from "./Navbar/WorkerDropdown/WorkerDropdown"; interface NavbarProps { userID: string | null; @@ -77,9 +79,19 @@ const Navbar: React.FC = ({ const handleLogout = () => { clearTokenCookies(); + localStorage.removeItem("litellm_selected_worker_id"); + localStorage.removeItem("litellm_worker_url"); window.location.href = logoutUrl; }; + const handleWorkerSwitch = (workerId: string) => { + clearTokenCookies(); + clearStoredReturnUrl(); + localStorage.removeItem("litellm_selected_worker_id"); + localStorage.removeItem("litellm_worker_url"); + window.location.href = `/ui/login?worker=${encodeURIComponent(workerId)}`; + }; + return (