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(