From ad43a35d762d62cd2cb9a18a0f1fbdef4f3fb67c Mon Sep 17 00:00:00 2001 From: Ryan Crabbe Date: Thu, 19 Mar 2026 22:50:19 -0700 Subject: [PATCH 1/2] feat: add control plane for multi-proxy worker management MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Adds a control plane capability that enables a central admin instance to manage multiple regional worker proxies from a single UI. Backend: - Worker registry loaded from YAML config (worker_id, name, url) - /.well-known/litellm-ui-config exposes is_control_plane and workers list - /v3/login + /v3/login/exchange: opaque code exchange for cross-origin username/password auth (JWT never in URL/logs, single-use 60s TTL) - SSO cookie handoff with return_to → opaque code → exchange - _validate_return_to: full origin validation (scheme+hostname+port) - Startup warning when control_plane_url set without Redis - Both /v3 endpoints gated behind control_plane_url config Frontend: - Worker selector dropdown on login page (gated behind is_control_plane) - Cross-origin SSO code exchange handling on callback - switchToWorkerUrl: localStorage-persisted worker URL for API calls - useWorker hook: shared worker state management - WorkerDropdown in navbar for switching workers - Logout/switch clears worker state from localStorage Tests: - 7 tests for /v3/login + /v3/login/exchange - 10 tests for _validate_return_to - 2 tests for control plane discovery endpoint --- .../ui_discovery_endpoints.py | 7 + litellm/proxy/management_endpoints/ui_sso.py | 85 ++++++- litellm/proxy/proxy_server.py | 187 +++++++++++++++- .../types/proxy/control_plane_endpoints.py | 14 ++ .../ui_discovery_endpoints.py | 6 +- .../test_ui_discovery_endpoints.py | 57 ++++- .../proxy/management_endpoints/test_ui_sso.py | 98 +++++++- tests/test_litellm/proxy/test_proxy_server.py | 211 ++++++++++++++++++ .../app/(dashboard)/hooks/login/useLogin.ts | 4 +- .../hooks/uiConfig/useUIConfig.test.ts | 4 + .../src/app/login/LoginPage.test.tsx | 15 +- .../src/app/login/LoginPage.tsx | 125 +++++++++-- .../Navbar/WorkerDropdown/WorkerDropdown.tsx | 38 ++++ .../src/components/navbar.tsx | 13 ++ .../src/components/networking.tsx | 126 ++++++++++- ui/litellm-dashboard/src/hooks/useWorker.ts | 65 ++++++ 16 files changed, 1025 insertions(+), 30 deletions(-) create mode 100644 litellm/types/proxy/control_plane_endpoints.py create mode 100644 ui/litellm-dashboard/src/components/Navbar/WorkerDropdown/WorkerDropdown.tsx create mode 100644 ui/litellm-dashboard/src/hooks/useWorker.ts diff --git a/litellm/proxy/discovery_endpoints/ui_discovery_endpoints.py b/litellm/proxy/discovery_endpoints/ui_discovery_endpoints.py index 2a38ceffba5..233df5c6c57 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/ui_sso.py b/litellm/proxy/management_endpoints/ui_sso.py index daf1d6f1316..8634a8bb33d 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,38 @@ 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 9c29927c5cb..e982c934aa6 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 00000000000..8bf4c44b205 --- /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 4a4cdaa2bae..46cd3f49f1a 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/proxy/discovery_endpoints/test_ui_discovery_endpoints.py b/tests/test_litellm/proxy/discovery_endpoints/test_ui_discovery_endpoints.py index f15960a6072..54a127f435b 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_ui_sso.py b/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py index d43b2c4ba05..fc9c37b7f84 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 112a06b1731..bd6162f225a 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)/hooks/login/useLogin.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/login/useLogin.ts index a15b4a06d13..be53b1c80a8 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/uiConfig/useUIConfig.test.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/uiConfig/useUIConfig.test.ts index aba5dddf13d..b05bae1e818 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/login/LoginPage.test.tsx b/ui/litellm-dashboard/src/app/login/LoginPage.test.tsx index ad2dde2da83..866b7d0f172 100644 --- a/ui/litellm-dashboard/src/app/login/LoginPage.test.tsx +++ b/ui/litellm-dashboard/src/app/login/LoginPage.test.tsx @@ -41,6 +41,17 @@ vi.mock("@/app/(dashboard)/hooks/login/useLogin", () => ({ })), })); +vi.mock("@/hooks/useWorker", () => ({ + useWorker: vi.fn(() => ({ + isControlPlane: false, + workers: [], + selectedWorkerId: null, + selectedWorker: null, + selectWorker: vi.fn(), + disconnectFromWorker: vi.fn(), + })), +})); + import { useUIConfig } from "@/app/(dashboard)/hooks/uiConfig/useUIConfig"; import { getCookie } from "@/utils/cookieUtils"; import { isJwtExpired } from "@/utils/jwtUtils"; @@ -108,7 +119,7 @@ describe("LoginPage", () => { ); await waitFor(() => { - expect(mockReplace).toHaveBeenCalledWith("http://localhost:4000/ui"); + expect(mockReplace).toHaveBeenCalledWith("/ui"); }); }); @@ -189,7 +200,7 @@ describe("LoginPage", () => { ); await waitFor(() => { - expect(mockReplace).toHaveBeenCalledWith("http://localhost:4000/ui"); + expect(mockReplace).toHaveBeenCalledWith("/ui"); }); expect(mockPush).not.toHaveBeenCalled(); diff --git a/ui/litellm-dashboard/src/app/login/LoginPage.tsx b/ui/litellm-dashboard/src/app/login/LoginPage.tsx index d54140c5a29..5a9d420456c 100644 --- a/ui/litellm-dashboard/src/app/login/LoginPage.tsx +++ b/ui/litellm-dashboard/src/app/login/LoginPage.tsx @@ -3,14 +3,15 @@ import { useLogin } from "@/app/(dashboard)/hooks/login/useLogin"; import { useUIConfig } from "@/app/(dashboard)/hooks/uiConfig/useUIConfig"; import LoadingScreen from "@/components/common_components/LoadingScreen"; -import { getProxyBaseUrl } from "@/components/networking"; -import { getCookie } from "@/utils/cookieUtils"; +import { exchangeLoginCode, getProxyBaseUrl, switchToWorkerUrl } from "@/components/networking"; +import { clearTokenCookies, getCookie } from "@/utils/cookieUtils"; import { isJwtExpired } from "@/utils/jwtUtils"; import { consumeReturnUrl, getReturnUrl, isValidReturnUrl } from "@/utils/returnUrlUtils"; -import { InfoCircleOutlined } from "@ant-design/icons"; -import { Alert, Button, Card, Form, Input, Popover, Space, Typography } from "antd"; +import { InfoCircleOutlined, CloudServerOutlined } from "@ant-design/icons"; +import { Alert, Button, Card, Form, Input, Popover, Select, Space, Typography } from "antd"; import { useRouter } from "next/navigation"; import { useEffect, useState } from "react"; +import { useWorker } from "@/hooks/useWorker"; function LoginPageContent() { const [username, setUsername] = useState(""); @@ -19,6 +20,17 @@ function LoginPageContent() { const { data: uiConfig, isLoading: isConfigLoading } = useUIConfig(); const loginMutation = useLogin(); const router = useRouter(); + const { workers, selectWorker } = useWorker(); + const [selectedWorkerId, setSelectedWorkerId] = useState(null); + + // Pre-select worker from URL param (e.g. /ui/login?worker=team-b) + useEffect(() => { + const params = new URLSearchParams(window.location.search); + const workerParam = params.get("worker"); + if (workerParam) { + setSelectedWorkerId(workerParam); + } + }, []); useEffect(() => { if (isConfigLoading) { @@ -31,6 +43,44 @@ function LoginPageContent() { return; } + // Cross-origin SSO: worker redirected back with a single-use code. + // Exchange it for the JWT via the worker's /v3/login/exchange endpoint. + const params = new URLSearchParams(window.location.search); + const ssoCode = params.get("code"); + if (ssoCode) { + const workerUrl = localStorage.getItem("litellm_worker_url"); + exchangeLoginCode(ssoCode, workerUrl).then(() => { + params.delete("code"); + const cleanSearch = params.toString(); + window.history.replaceState(null, "", window.location.pathname + (cleanSearch ? `?${cleanSearch}` : "")); + router.replace("/ui/?login=success"); + }); + return; + } + + // Backwards compat: handle direct token in URL (legacy flow) + const urlToken = params.get("token"); + if (urlToken && !isJwtExpired(urlToken)) { + document.cookie = `token=${urlToken}; path=/; SameSite=Lax`; + params.delete("token"); + const cleanSearch = params.toString(); + window.history.replaceState( + null, + "", + window.location.pathname + (cleanSearch ? `?${cleanSearch}` : ""), + ); + router.replace("/ui/?login=success"); + return; + } + + // If switching workers on a control plane, clear the old token and show login + const switchingWorker = params.has("worker"); + if (switchingWorker && uiConfig?.is_control_plane) { + clearTokenCookies(); + setIsLoading(false); + return; + } + const rawToken = getCookie("token"); if (rawToken && !isJwtExpired(rawToken)) { // User already logged in - redirect to return URL or default @@ -38,7 +88,7 @@ function LoginPageContent() { if (returnUrl) { router.replace(returnUrl); } else { - router.replace(`${getProxyBaseUrl()}/ui`); + router.replace("/ui"); } return; } @@ -58,16 +108,35 @@ function LoginPageContent() { }, [isConfigLoading, router, uiConfig]); const handleSubmit = () => { + // If a worker is selected, point proxyBaseUrl at it before login + const selectedWorker = workers.find((w) => w.worker_id === selectedWorkerId); + if (selectedWorker) { + switchToWorkerUrl(selectedWorker.url); + } + loginMutation.mutate( - { username, password }, + { username, password, useV3: !!selectedWorker }, { onSuccess: (data) => { - // Check if we have a return URL to use instead of the default redirect - const returnUrl = consumeReturnUrl(); - if (returnUrl) { - router.push(returnUrl); + // Update the worker context with the selected worker + if (selectedWorker) { + selectWorker(selectedWorker.worker_id); + // Stay on the CP's UI — proxyBaseUrl already points at the worker + router.push("/ui/?login=success"); } else { - router.push(data.redirect_url); + // Normal (non-control-plane) login — follow the server's redirect + const returnUrl = consumeReturnUrl(); + if (returnUrl) { + router.push(returnUrl); + } else { + router.push(data.redirect_url); + } + } + }, + onError: () => { + // Reset proxyBaseUrl on login failure + if (selectedWorker) { + switchToWorkerUrl(null); } }, }, @@ -154,6 +223,22 @@ function LoginPageContent() { {error && }
+ {uiConfig?.is_control_plane && workers.length > 0 && ( + + + (option?.label as string ?? "").toLowerCase().includes(input.toLowerCase()) + } + value={selectedWorker.worker_id} + style={{ minWidth: 180 }} + suffixIcon={} + options={workers.map((w) => ({ + label: w.name, + value: w.worker_id, + disabled: w.worker_id === selectedWorker.worker_id, + }))} + onChange={(newWorkerId) => { + onWorkerSwitch(newWorkerId); + }} + /> + ); +}; + +export default WorkerDropdown; diff --git a/ui/litellm-dashboard/src/components/navbar.tsx b/ui/litellm-dashboard/src/components/navbar.tsx index c46a3af5a6e..96d6ce613bf 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 (