From ec59eef43bc8ead54b9b713531a3328329f7b2d2 Mon Sep 17 00:00:00 2001 From: Ryan Crabbe Date: Fri, 13 Mar 2026 17:20:34 -0700 Subject: [PATCH 1/9] feat(proxy): add control plane worker registry backend Add backend support for a control plane mode where one LiteLLM instance serves as a worker directory. Parses `worker_registry` from YAML config, exposes `is_control_plane` flag and worker list via the existing `/.well-known/litellm-ui-config` endpoint, and returns JWT in /v2/login JSON body for cross-origin auth. - WorkerRegistryEntry pydantic model with HTTP URL validation - Config parsing in ProxyConfig._init_non_llm_configs - is_control_plane + workers fields on UiDiscoveryEndpoints response - Example control_plane_config.yaml - Tests for discovery endpoint and /v2/login response shape --- .../ui_discovery_endpoints.py | 7 +++ litellm/proxy/proxy_server.py | 14 ++++- .../types/proxy/control_plane_endpoints.py | 14 +++++ .../ui_discovery_endpoints.py | 6 ++- .../test_ui_discovery_endpoints.py | 52 +++++++++++++++++-- tests/test_litellm/proxy/test_proxy_server.py | 8 +-- 6 files changed, 91 insertions(+), 10 deletions(-) create mode 100644 litellm/types/proxy/control_plane_endpoints.py diff --git a/litellm/proxy/discovery_endpoints/ui_discovery_endpoints.py b/litellm/proxy/discovery_endpoints/ui_discovery_endpoints.py index 2a38ceffba5..8154303a7ad 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, ) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 2a9be0a67c9..b37d5ea956c 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -322,6 +322,7 @@ from litellm.proxy.credential_endpoints.endpoints import router as credential_ro from litellm.proxy.db.db_transaction_queue.spend_log_cleanup import SpendLogCleanup from litellm.proxy.db.exception_handler import PrismaDBExceptionHandler from litellm.proxy.discovery_endpoints import ui_discovery_endpoints_router +from litellm.types.proxy.control_plane_endpoints import WorkerRegistryEntry from litellm.proxy.fine_tuning_endpoints.endpoints import router as fine_tuning_router from litellm.proxy.fine_tuning_endpoints.endpoints import set_fine_tuning_config from litellm.proxy.google_endpoints.endpoints import router as google_router @@ -2294,6 +2295,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): @@ -3384,7 +3386,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, @@ -11031,7 +11041,7 @@ async def login_v2(request: Request): # noqa: PLR0915 litellm_dashboard_ui += "?login=success" json_response = JSONResponse( - content={"redirect_url": litellm_dashboard_ui}, + content={"redirect_url": litellm_dashboard_ui, "token": jwt_token}, status_code=status.HTTP_200_OK, ) json_response.set_cookie(key="token", value=jwt_token) 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..189105aadd1 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,48 @@ 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 + + +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 diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index 112a06b1731..de3ffb9f839 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -104,10 +104,10 @@ def test_login_v2_returns_redirect_url_and_sets_cookie(monkeypatch): ) assert response.status_code == 200 - assert ( - response.json() - == {"redirect_url": "http://testserver/ui/?login=success"} - ) + assert response.json() == { + "redirect_url": "http://testserver/ui/?login=success", + "token": "signed-token", + } assert response.cookies.get("token") == "signed-token" mock_authenticate_user.assert_awaited_once_with( From d3d2ecb44e9009fd9a24c4bb99c13af6bc61a6b4 Mon Sep 17 00:00:00 2001 From: Ryan Crabbe Date: Sat, 14 Mar 2026 10:47:19 -0700 Subject: [PATCH 2/9] Gate control-plane-only fields behind worker_registry presence Only return JWT token in /v2/login response body when workers are configured (control plane mode). Non-CP instances continue to set the cookie only, avoiding unnecessary token exposure. Also make workers list in discovery endpoint explicitly conditional on is_control_plane flag. --- litellm/proxy/discovery_endpoints/ui_discovery_endpoints.py | 2 +- litellm/proxy/proxy_server.py | 6 +++++- .../discovery_endpoints/test_ui_discovery_endpoints.py | 5 +++++ tests/test_litellm/proxy/test_proxy_server.py | 3 ++- 4 files changed, 13 insertions(+), 3 deletions(-) diff --git a/litellm/proxy/discovery_endpoints/ui_discovery_endpoints.py b/litellm/proxy/discovery_endpoints/ui_discovery_endpoints.py index 8154303a7ad..233df5c6c57 100644 --- a/litellm/proxy/discovery_endpoints/ui_discovery_endpoints.py +++ b/litellm/proxy/discovery_endpoints/ui_discovery_endpoints.py @@ -39,5 +39,5 @@ async def get_ui_config(): admin_ui_disabled=admin_ui_disabled, sso_configured=sso_configured, is_control_plane=is_control_plane, - workers=proxy_config.worker_registry, + workers=proxy_config.worker_registry if is_control_plane else [], ) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index b37d5ea956c..0c52e30e99c 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -11040,8 +11040,12 @@ async def login_v2(request: Request): # noqa: PLR0915 litellm_dashboard_ui += "/ui/" litellm_dashboard_ui += "?login=success" + response_content: dict = {"redirect_url": litellm_dashboard_ui} + if len(proxy_config.worker_registry) > 0: + response_content["token"] = jwt_token + json_response = JSONResponse( - content={"redirect_url": litellm_dashboard_ui, "token": jwt_token}, + content=response_content, status_code=status.HTTP_200_OK, ) json_response.set_cookie(key="token", value=jwt_token) 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 189105aadd1..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 @@ -281,6 +281,10 @@ def test_ui_discovery_endpoints_is_control_plane_true_when_workers_configured(): 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(): @@ -302,3 +306,4 @@ def test_ui_discovery_endpoints_is_control_plane_false_when_no_workers(): 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/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index de3ffb9f839..864383f79ab 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -25,6 +25,7 @@ sys.path.insert( import litellm from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.proxy_server import app, initialize +from litellm.types.proxy.control_plane_endpoints import WorkerRegistryEntry from litellm.utils import _invalidate_model_cost_lowercase_map example_embedding_result = { @@ -106,8 +107,8 @@ def test_login_v2_returns_redirect_url_and_sets_cookie(monkeypatch): assert response.status_code == 200 assert response.json() == { "redirect_url": "http://testserver/ui/?login=success", - "token": "signed-token", } + assert "token" not in response.json() assert response.cookies.get("token") == "signed-token" mock_authenticate_user.assert_awaited_once_with( From 1786f5ad972ff305cec56c532a378851934028a6 Mon Sep 17 00:00:00 2001 From: Ryan Crabbe Date: Sat, 14 Mar 2026 11:06:38 -0700 Subject: [PATCH 3/9] Move WorkerRegistryEntry import to types import block --- litellm/proxy/proxy_server.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 0c52e30e99c..3c79c48689b 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -322,7 +322,6 @@ from litellm.proxy.credential_endpoints.endpoints import router as credential_ro from litellm.proxy.db.db_transaction_queue.spend_log_cleanup import SpendLogCleanup from litellm.proxy.db.exception_handler import PrismaDBExceptionHandler from litellm.proxy.discovery_endpoints import ui_discovery_endpoints_router -from litellm.types.proxy.control_plane_endpoints import WorkerRegistryEntry from litellm.proxy.fine_tuning_endpoints.endpoints import router as fine_tuning_router from litellm.proxy.fine_tuning_endpoints.endpoints import set_fine_tuning_config from litellm.proxy.google_endpoints.endpoints import router as google_router @@ -541,6 +540,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, ) From 6686a1213fdb09d97d0ae3b1931587f561fa0e0d Mon Sep 17 00:00:00 2001 From: Ryan Crabbe Date: Sat, 14 Mar 2026 11:24:06 -0700 Subject: [PATCH 4/9] Fix login test fragility and add control-plane token test Explicitly mock proxy_config.worker_registry in the existing login test to isolate it from global state. Add positive test verifying the token is returned in the response body when workers are configured. --- tests/test_litellm/proxy/test_proxy_server.py | 40 +++++++++++++++++++ 1 file changed, 40 insertions(+) diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index 864383f79ab..219a0ca7825 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -95,6 +95,9 @@ def test_login_v2_returns_redirect_url_and_sets_cookie(monkeypatch): monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", {}) 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) @@ -129,6 +132,43 @@ def test_login_v2_returns_redirect_url_and_sets_cookie(monkeypatch): ) +def test_login_v2_includes_token_in_body_when_control_plane(monkeypatch): + """Token should appear in the response body when workers are configured.""" + 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", {}) + 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 = [ + WorkerRegistryEntry( + worker_id="team-a", name="Team A", url="https://worker-1:4001" + ), + ] + 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( + "/v2/login", + json={"username": "alice", "password": "secret"}, + ) + + assert response.status_code == 200 + assert response.json()["token"] == "signed-token" + assert response.cookies.get("token") == "signed-token" + + def test_login_v2_returns_json_on_proxy_exception(monkeypatch): """Test that /v2/login returns JSON error when ProxyException is raised""" from litellm.proxy._types import ProxyErrorTypes, ProxyException From 96cd74440c23ea928dd84ddec87faac67250bf7e Mon Sep 17 00:00:00 2001 From: Ryan Crabbe Date: Sat, 14 Mar 2026 12:58:40 -0700 Subject: [PATCH 5/9] Add /v3/login endpoint for cross-origin control-plane login Workers don't have worker_registry, so /v2/login omits the token from the response body. The control plane UI needs the token in the body to set the cookie cross-origin (document.cookie) when authenticating against a worker. /v3/login always includes it. --- litellm/proxy/proxy_server.py | 72 +++++++++++++++++++ tests/test_litellm/proxy/test_proxy_server.py | 69 ++++++++++++++++++ 2 files changed, 141 insertions(+) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 3c79c48689b..5ce22c366f8 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -11075,6 +11075,78 @@ 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: + 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" + + json_response = JSONResponse( + content={"redirect_url": litellm_dashboard_ui, "token": jwt_token}, + status_code=status.HTTP_200_OK, + ) + json_response.set_cookie(key="token", value=jwt_token) + return json_response + 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, + ) + + @app.get("/onboarding/get_token", include_in_schema=False) async def onboarding(invite_link: str, request: Request): """ diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index 219a0ca7825..ab6396a43e8 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -277,6 +277,75 @@ def test_login_v2_returns_json_on_invalid_json_body(monkeypatch): assert isinstance(data["error"], dict) +def test_login_v3_always_includes_token_in_body(monkeypatch): + """v3/login always returns token in body, even without worker_registry.""" + 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", {}) + 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 + assert response.json()["token"] == "signed-token" + assert response.json()["redirect_url"] == "http://testserver/ui/?login=success" + assert response.cookies.get("token") == "signed-token" + + +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.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") From d04f8bcc457361484d28ffae596035cd2d91372b Mon Sep 17 00:00:00 2001 From: Ryan Crabbe Date: Tue, 17 Mar 2026 11:57:09 -0700 Subject: [PATCH 6/9] =?UTF-8?q?Keep=20/v2/login=20untouched=20=E2=80=94=20?= =?UTF-8?q?use=20/v3/login=20exclusively=20for=20control-plane=20cross-ori?= =?UTF-8?q?gin=20auth?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit /v2/login is battle-tested; control-plane token-in-body behavior belongs only in /v3/login. Removes the conditional worker_registry check from v2 and its associated test. --- litellm/proxy/proxy_server.py | 6 +-- tests/test_litellm/proxy/test_proxy_server.py | 49 ++----------------- 2 files changed, 5 insertions(+), 50 deletions(-) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 5ce22c366f8..48d9d18fe46 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -11040,12 +11040,8 @@ async def login_v2(request: Request): # noqa: PLR0915 litellm_dashboard_ui += "/ui/" litellm_dashboard_ui += "?login=success" - response_content: dict = {"redirect_url": litellm_dashboard_ui} - if len(proxy_config.worker_registry) > 0: - response_content["token"] = jwt_token - json_response = JSONResponse( - content=response_content, + content={"redirect_url": litellm_dashboard_ui}, status_code=status.HTTP_200_OK, ) json_response.set_cookie(key="token", value=jwt_token) diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index ab6396a43e8..e6269936b5c 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -25,7 +25,6 @@ sys.path.insert( import litellm from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.proxy_server import app, initialize -from litellm.types.proxy.control_plane_endpoints import WorkerRegistryEntry from litellm.utils import _invalidate_model_cost_lowercase_map example_embedding_result = { @@ -95,9 +94,6 @@ def test_login_v2_returns_redirect_url_and_sets_cookie(monkeypatch): monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", {}) 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) @@ -108,10 +104,10 @@ def test_login_v2_returns_redirect_url_and_sets_cookie(monkeypatch): ) assert response.status_code == 200 - assert response.json() == { - "redirect_url": "http://testserver/ui/?login=success", - } - assert "token" not in response.json() + assert ( + response.json() + == {"redirect_url": "http://testserver/ui/?login=success"} + ) assert response.cookies.get("token") == "signed-token" mock_authenticate_user.assert_awaited_once_with( @@ -132,43 +128,6 @@ def test_login_v2_returns_redirect_url_and_sets_cookie(monkeypatch): ) -def test_login_v2_includes_token_in_body_when_control_plane(monkeypatch): - """Token should appear in the response body when workers are configured.""" - 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", {}) - 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 = [ - WorkerRegistryEntry( - worker_id="team-a", name="Team A", url="https://worker-1:4001" - ), - ] - 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( - "/v2/login", - json={"username": "alice", "password": "secret"}, - ) - - assert response.status_code == 200 - assert response.json()["token"] == "signed-token" - assert response.cookies.get("token") == "signed-token" - - def test_login_v2_returns_json_on_proxy_exception(monkeypatch): """Test that /v2/login returns JSON error when ProxyException is raised""" from litellm.proxy._types import ProxyErrorTypes, ProxyException From 5543718caf1ddb35308a44d0123706cce1faa0fe Mon Sep 17 00:00:00 2001 From: Ryan Crabbe Date: Thu, 19 Mar 2026 14:02:49 -0700 Subject: [PATCH 7/9] feat: add cross-origin SSO cookie handoff for control plane - Add return_to param to /sso/key/generate for cross-origin redirect - Store return_to in httpOnly cookie during SSO redirect - Read cookie in /sso/callback and pass to OpenID handler - Redirect to control plane with token in URL, delete cookie after use --- litellm/proxy/management_endpoints/ui_sso.py | 45 +++++++++++++++++++- 1 file changed, 44 insertions(+), 1 deletion(-) diff --git a/litellm/proxy/management_endpoints/ui_sso.py b/litellm/proxy/management_endpoints/ui_sso.py index d55aa85a9b8..611a0e89bf4 100644 --- a/litellm/proxy/management_endpoints/ui_sso.py +++ b/litellm/proxy/management_endpoints/ui_sso.py @@ -301,6 +301,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 +395,22 @@ 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: + 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 @@ -1306,12 +1316,19 @@ 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 + cp_return_to: Optional[str] = None + cookie_return_to = request.cookies.get("litellm_cp_return_to") + if cookie_return_to: + cp_return_to = cookie_return_to.strip('"') + 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, ) @@ -2352,6 +2369,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 @@ -2528,6 +2546,31 @@ class SSOAuthenticationHandler: master_key or "", algorithm="HS256", ) + + # Control-plane cross-origin: redirect back to the control plane UI + # with the token in the URL (cookie won't work cross-origin) + if return_to is not None: + from urllib.parse import urlencode, urlparse + + parsed = urlparse(return_to) + if parsed.scheme not in ("http", "https"): + raise HTTPException( + status_code=400, + detail="return_to must be an HTTP(S) URL", + ) + separator = "&" if "?" in return_to else "?" + redirect_url = ( + return_to + + separator + + urlencode({"login": "success", "token": jwt_token}) + ) + verbose_proxy_logger.info( + f"Cross-origin SSO: redirecting to control plane at {parsed.netloc}" + ) + 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}") From 719974f0b39d9688bc42cff55a2c74ee33485bdd Mon Sep 17 00:00:00 2001 From: Ryan Crabbe Date: Thu, 19 Mar 2026 15:35:13 -0700 Subject: [PATCH 8/9] fix: validate return_to against control_plane_url to prevent open redirect Add origin validation for the return_to parameter in SSO flow to prevent JWT theft via crafted redirect URLs. Workers must configure control_plane_url in general_settings. --- litellm/proxy/management_endpoints/ui_sso.py | 41 ++++++++++---- .../proxy/management_endpoints/test_ui_sso.py | 53 ++++++++++++++++++- 2 files changed, 84 insertions(+), 10 deletions(-) diff --git a/litellm/proxy/management_endpoints/ui_sso.py b/litellm/proxy/management_endpoints/ui_sso.py index 611a0e89bf4..5e6aabd70ac 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 @@ -403,6 +404,7 @@ async def google_login( 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, @@ -1771,6 +1773,34 @@ 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", + ) + + parsed_cp = urlparse(control_plane_url) + parsed_return = urlparse(return_to) + cp_origin = f"{parsed_cp.scheme}://{parsed_cp.netloc}" + return_origin = f"{parsed_return.scheme}://{parsed_return.netloc}" + if cp_origin != return_origin: + 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, @@ -2550,14 +2580,7 @@ class SSOAuthenticationHandler: # Control-plane cross-origin: redirect back to the control plane UI # with the token in the URL (cookie won't work cross-origin) if return_to is not None: - from urllib.parse import urlencode, urlparse - - parsed = urlparse(return_to) - if parsed.scheme not in ("http", "https"): - raise HTTPException( - status_code=400, - detail="return_to must be an HTTP(S) URL", - ) + SSOAuthenticationHandler._validate_return_to(return_to) separator = "&" if "?" in return_to else "?" redirect_url = ( return_to @@ -2565,7 +2588,7 @@ class SSOAuthenticationHandler: + urlencode({"login": "success", "token": jwt_token}) ) verbose_proxy_logger.info( - f"Cross-origin SSO: redirecting to control plane at {parsed.netloc}" + "Cross-origin SSO: redirecting to control plane" ) redirect_response = RedirectResponse(url=redirect_url, status_code=303) redirect_response.delete_cookie("litellm_cp_return_to") 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..d48dcc52f59 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,54 @@ 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 + From f8d1be2e9b4af21a31138d4dba7408d5f6c7612d Mon Sep 17 00:00:00 2001 From: Ryan Crabbe Date: Thu, 19 Mar 2026 16:00:19 -0700 Subject: [PATCH 9/9] fix: gate /v3/login behind control_plane_url and fix RFC 3986 hostname comparison - /v3/login now returns 404 unless control_plane_url is configured - _validate_return_to uses urlparse().hostname for case-insensitive comparison - Add rejection test for /v3/login without control_plane_url - Add case-insensitive hostname test for _validate_return_to --- litellm/proxy/management_endpoints/ui_sso.py | 6 +--- litellm/proxy/proxy_server.py | 8 +++++ .../proxy/management_endpoints/test_ui_sso.py | 9 ++++++ tests/test_litellm/proxy/test_proxy_server.py | 30 +++++++++++++++++-- 4 files changed, 45 insertions(+), 8 deletions(-) diff --git a/litellm/proxy/management_endpoints/ui_sso.py b/litellm/proxy/management_endpoints/ui_sso.py index 5e6aabd70ac..39e015f3928 100644 --- a/litellm/proxy/management_endpoints/ui_sso.py +++ b/litellm/proxy/management_endpoints/ui_sso.py @@ -1791,11 +1791,7 @@ class SSOAuthenticationHandler: detail="return_to is not allowed: control_plane_url is not configured", ) - parsed_cp = urlparse(control_plane_url) - parsed_return = urlparse(return_to) - cp_origin = f"{parsed_cp.scheme}://{parsed_cp.netloc}" - return_origin = f"{parsed_return.scheme}://{parsed_return.netloc}" - if cp_origin != return_origin: + if urlparse(return_to).hostname != urlparse(control_plane_url).hostname: raise HTTPException( status_code=400, detail="return_to does not match the configured control_plane_url", diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 48d9d18fe46..73b224d4bba 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -11080,6 +11080,14 @@ async def login_v3(request: Request): # noqa: PLR0915 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")) 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 d48dcc52f59..4d531da67ca 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py +++ b/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py @@ -5211,3 +5211,12 @@ class TestValidateReturnTo: 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") + diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index e6269936b5c..f23f5a39353 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -236,8 +236,25 @@ def test_login_v2_returns_json_on_invalid_json_body(monkeypatch): assert isinstance(data["error"], dict) -def test_login_v3_always_includes_token_in_body(monkeypatch): - """v3/login always returns token in body, even without worker_registry.""" +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_includes_token_in_body(monkeypatch): + """v3/login returns token in body when control_plane_url is configured.""" mock_prisma_client = MagicMock() monkeypatch.setattr( "litellm.proxy.auth.login_utils.authenticate_user", @@ -249,7 +266,10 @@ def test_login_v3_always_includes_token_in_body(monkeypatch): ) 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", {}) + 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() @@ -289,6 +309,10 @@ def test_login_v3_returns_json_on_proxy_exception(monkeypatch): 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)