From d3d2ecb44e9009fd9a24c4bb99c13af6bc61a6b4 Mon Sep 17 00:00:00 2001 From: Ryan Crabbe Date: Sat, 14 Mar 2026 10:47:19 -0700 Subject: [PATCH] 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(