mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-09 03:18:44 +00:00
feat: add control plane with worker registry, cross-origin SSO, and worker switching
- Add worker_registry config for control plane to manage multiple worker proxies - Add /.well-known/litellm-ui-config discovery endpoint exposing workers and CP status - Add /v3/login endpoint returning token in response body for cross-origin auth - Encode return_to URL in OAuth state for cross-origin SSO callback redirect - Add worker selection dropdown on login page with SSO support - Add searchable WorkerDropdown in navbar for switching between workers - Clear localStorage worker selection and return URL cookie on worker switch - Redirect to login with worker pre-selected when switching workers
This commit is contained in:
Ryan Crabbe
parent
0e44c460b5
commit
243e6cf5b4
16 changed files with 580 additions and 30 deletions
|
|
@ -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 [],
|
||||
)
|
||||
|
|
|
|||
|
|
@ -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}")
|
||||
|
|
|
|||
|
|
@ -540,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,
|
||||
)
|
||||
|
|
@ -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,
|
||||
|
|
@ -11061,6 +11071,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):
|
||||
"""
|
||||
|
|
|
|||
14
litellm/types/proxy/control_plane_endpoints.py
Normal file
14
litellm/types/proxy/control_plane_endpoints.py
Normal file
|
|
@ -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
|
||||
|
|
@ -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] = []
|
||||
|
|
|
|||
|
|
@ -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"] == []
|
||||
|
|
|
|||
|
|
@ -236,6 +236,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")
|
||||
|
||||
|
|
|
|||
|
|
@ -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;
|
||||
},
|
||||
});
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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();
|
||||
|
|
|
|||
|
|
@ -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 { 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<string | null>(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,31 @@ function LoginPageContent() {
|
|||
return;
|
||||
}
|
||||
|
||||
// Cross-origin SSO: worker redirected back with token in URL
|
||||
const params = new URLSearchParams(window.location.search);
|
||||
const urlToken = params.get("token");
|
||||
if (urlToken && !isJwtExpired(urlToken)) {
|
||||
document.cookie = `token=${urlToken}; path=/; SameSite=Lax`;
|
||||
// Strip token from URL to keep it out of browser history
|
||||
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 +75,7 @@ function LoginPageContent() {
|
|||
if (returnUrl) {
|
||||
router.replace(returnUrl);
|
||||
} else {
|
||||
router.replace(`${getProxyBaseUrl()}/ui`);
|
||||
router.replace("/ui");
|
||||
}
|
||||
return;
|
||||
}
|
||||
|
|
@ -58,16 +95,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 +210,22 @@ function LoginPageContent() {
|
|||
{error && <Alert message={error} type="error" showIcon />}
|
||||
|
||||
<Form onFinish={handleSubmit} layout="vertical" requiredMark={true}>
|
||||
{uiConfig?.is_control_plane && workers.length > 0 && (
|
||||
<Form.Item label="Worker" style={{ marginBottom: 16 }}>
|
||||
<Select
|
||||
value={selectedWorkerId || undefined}
|
||||
onChange={(value) => setSelectedWorkerId(value)}
|
||||
placeholder="Choose a worker to connect to"
|
||||
size="large"
|
||||
suffixIcon={<CloudServerOutlined />}
|
||||
options={workers.map((w) => ({
|
||||
label: w.name,
|
||||
value: w.worker_id,
|
||||
}))}
|
||||
/>
|
||||
</Form.Item>
|
||||
)}
|
||||
|
||||
<Form.Item
|
||||
label="Username"
|
||||
name="username"
|
||||
|
|
@ -209,10 +281,20 @@ function LoginPageContent() {
|
|||
</Popover>
|
||||
) : (
|
||||
<Button
|
||||
disabled={isLoginLoading}
|
||||
onClick={() =>
|
||||
router.push(`${getProxyBaseUrl()}/sso/key/generate`)
|
||||
}
|
||||
disabled={isLoginLoading || (!!selectedWorkerId && workers.length === 0)}
|
||||
onClick={() => {
|
||||
const selectedWorker = workers.find((w) => w.worker_id === selectedWorkerId);
|
||||
if (selectedWorker) {
|
||||
// Store worker selection so useWorker hook restores it after redirect
|
||||
localStorage.setItem("litellm_selected_worker_id", selectedWorkerId!);
|
||||
switchToWorkerUrl(selectedWorker.url);
|
||||
}
|
||||
// SSO on the worker (or this instance if no worker), always
|
||||
// include return_to so the callback redirects back here
|
||||
const ssoBase = selectedWorker?.url ?? getProxyBaseUrl();
|
||||
const returnTo = encodeURIComponent(window.location.origin + "/ui/login");
|
||||
router.push(`${ssoBase}/sso/key/generate?return_to=${returnTo}`);
|
||||
}}
|
||||
block
|
||||
size="large"
|
||||
>
|
||||
|
|
|
|||
|
|
@ -0,0 +1,38 @@
|
|||
"use client";
|
||||
|
||||
import React from "react";
|
||||
import { Select } from "antd";
|
||||
import { CloudServerOutlined } from "@ant-design/icons";
|
||||
import { useWorker } from "@/hooks/useWorker";
|
||||
|
||||
interface WorkerDropdownProps {
|
||||
onWorkerSwitch: (workerId: string) => void;
|
||||
}
|
||||
|
||||
const WorkerDropdown: React.FC<WorkerDropdownProps> = ({ onWorkerSwitch }) => {
|
||||
const { isControlPlane, selectedWorker, workers } = useWorker();
|
||||
|
||||
if (!isControlPlane || !selectedWorker) return null;
|
||||
|
||||
return (
|
||||
<Select
|
||||
showSearch
|
||||
filterOption={(input, option) =>
|
||||
(option?.label as string ?? "").toLowerCase().includes(input.toLowerCase())
|
||||
}
|
||||
value={selectedWorker.worker_id}
|
||||
style={{ minWidth: 180 }}
|
||||
suffixIcon={<CloudServerOutlined />}
|
||||
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;
|
||||
|
|
@ -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,17 @@ const Navbar: React.FC<NavbarProps> = ({
|
|||
|
||||
const handleLogout = () => {
|
||||
clearTokenCookies();
|
||||
localStorage.removeItem("litellm_selected_worker_id");
|
||||
window.location.href = logoutUrl;
|
||||
};
|
||||
|
||||
const handleWorkerSwitch = (workerId: string) => {
|
||||
clearTokenCookies();
|
||||
clearStoredReturnUrl();
|
||||
localStorage.removeItem("litellm_selected_worker_id");
|
||||
window.location.href = `/ui/login?worker=${encodeURIComponent(workerId)}`;
|
||||
};
|
||||
|
||||
return (
|
||||
<nav className="bg-white border-b border-gray-200 sticky top-0 z-10">
|
||||
<div className="w-full">
|
||||
|
|
@ -169,6 +179,7 @@ const Navbar: React.FC<NavbarProps> = ({
|
|||
NEW
|
||||
</span>
|
||||
</a>
|
||||
<WorkerDropdown onWorkerSwitch={handleWorkerSwitch} />
|
||||
<CommunityEngagementButtons />
|
||||
{/* Dark mode is currently a work in progress. To test, you can change 'false' to 'true' below.
|
||||
Do not set this to true by default until all components are confirmed to support dark mode styles. */}
|
||||
|
|
|
|||
|
|
@ -86,7 +86,23 @@ const defaultProxyBaseUrl =
|
|||
: null;
|
||||
const defaultServerRootPath = "/";
|
||||
export let serverRootPath = defaultServerRootPath;
|
||||
export let proxyBaseUrl = defaultProxyBaseUrl;
|
||||
const WORKER_URL_KEY = "litellm_worker_url";
|
||||
// If a worker URL is in localStorage, use it as the initial proxyBaseUrl.
|
||||
// This survives page navigation and the sessionStorage.clear() in user_dashboard.
|
||||
const _rawWorkerUrl =
|
||||
typeof window !== "undefined" ? window.localStorage.getItem(WORKER_URL_KEY) : null;
|
||||
// Validate stored worker URL — reject non-HTTP schemes to prevent exfiltration
|
||||
const _initialWorkerUrl = (() => {
|
||||
if (!_rawWorkerUrl) return null;
|
||||
try {
|
||||
const parsed = new URL(_rawWorkerUrl);
|
||||
if (parsed.protocol === "http:" || parsed.protocol === "https:") return _rawWorkerUrl;
|
||||
} catch { /* invalid URL */ }
|
||||
// Invalid URL in storage — clear it
|
||||
if (typeof window !== "undefined") window.localStorage.removeItem(WORKER_URL_KEY);
|
||||
return null;
|
||||
})();
|
||||
export let proxyBaseUrl: string | null = _initialWorkerUrl ?? defaultProxyBaseUrl;
|
||||
if (isLocal != true) {
|
||||
console.log = function () { };
|
||||
}
|
||||
|
|
@ -102,6 +118,10 @@ const updateProxyBaseUrl = (serverRootPath: string, receivedProxyBaseUrl: string
|
|||
/**
|
||||
* Special function for updating the proxy base url. Should only be called by getUiConfig.
|
||||
*/
|
||||
// If a worker URL is in localStorage, don't let getUiConfig overwrite it
|
||||
if (typeof window !== "undefined" && window.localStorage.getItem(WORKER_URL_KEY)) {
|
||||
return;
|
||||
}
|
||||
const browserLocation = getWindowLocation();
|
||||
const resolvedDefaultProxyBaseUrl =
|
||||
isLocal && process.env.NEXT_PUBLIC_USE_REWRITES !== "true"
|
||||
|
|
@ -137,6 +157,36 @@ export const getProxyBaseUrl = (): string => {
|
|||
return browserLocation?.origin ?? "";
|
||||
};
|
||||
|
||||
/**
|
||||
* Switch API calls to point at a worker (or back to the control plane).
|
||||
* Persists to localStorage so it survives page navigation and the
|
||||
* sessionStorage.clear() in user_dashboard. Also updates the module-level
|
||||
* proxyBaseUrl so in-flight code in this JS execution sees the new value
|
||||
* immediately.
|
||||
*/
|
||||
function isValidHttpUrl(url: string): boolean {
|
||||
try {
|
||||
const parsed = new URL(url);
|
||||
return parsed.protocol === "http:" || parsed.protocol === "https:";
|
||||
} catch {
|
||||
return false;
|
||||
}
|
||||
}
|
||||
|
||||
export function switchToWorkerUrl(workerUrl: string | null): void {
|
||||
if (workerUrl && !isValidHttpUrl(workerUrl)) {
|
||||
return;
|
||||
}
|
||||
if (typeof window !== "undefined") {
|
||||
if (workerUrl) {
|
||||
window.localStorage.setItem(WORKER_URL_KEY, workerUrl);
|
||||
} else {
|
||||
window.localStorage.removeItem(WORKER_URL_KEY);
|
||||
}
|
||||
}
|
||||
proxyBaseUrl = workerUrl ?? defaultProxyBaseUrl;
|
||||
}
|
||||
|
||||
const HTTP_REQUEST = {
|
||||
GET: "GET",
|
||||
POST: "POST",
|
||||
|
|
@ -262,12 +312,20 @@ interface PublicModelHubInfo {
|
|||
useful_links: Record<string, string | { url: string; index: number }>;
|
||||
}
|
||||
|
||||
export interface WorkerInfo {
|
||||
worker_id: string;
|
||||
name: string;
|
||||
url: string;
|
||||
}
|
||||
|
||||
export interface LiteLLMWellKnownUiConfig {
|
||||
server_root_path: string;
|
||||
proxy_base_url: string | null;
|
||||
auto_redirect_to_sso: boolean;
|
||||
admin_ui_disabled: boolean;
|
||||
sso_configured: boolean;
|
||||
is_control_plane?: boolean;
|
||||
workers?: WorkerInfo[];
|
||||
}
|
||||
|
||||
export interface CredentialsResponse {
|
||||
|
|
@ -8972,15 +9030,18 @@ export const deriveErrorMessage = (errorData: any): string => {
|
|||
export interface LoginRequest {
|
||||
username: string;
|
||||
password: string;
|
||||
useV3?: boolean;
|
||||
}
|
||||
|
||||
interface LoginResponse {
|
||||
redirect_url: string;
|
||||
token?: string;
|
||||
}
|
||||
|
||||
export const loginCall = async (username: string, password: string): Promise<LoginResponse> => {
|
||||
export const loginCall = async (username: string, password: string, useV3?: boolean): Promise<LoginResponse> => {
|
||||
const proxyBaseUrl = getProxyBaseUrl();
|
||||
const loginUrl = proxyBaseUrl ? `${proxyBaseUrl}/v2/login` : "/v2/login";
|
||||
const loginPath = useV3 ? "/v3/login" : "/v2/login";
|
||||
const loginUrl = proxyBaseUrl ? `${proxyBaseUrl}${loginPath}` : loginPath;
|
||||
|
||||
const body = JSON.stringify({
|
||||
username,
|
||||
|
|
@ -9002,7 +9063,15 @@ export const loginCall = async (username: string, password: string): Promise<Log
|
|||
throw new Error(errorMessage);
|
||||
}
|
||||
|
||||
const data = await response.json();
|
||||
const data: LoginResponse = await response.json();
|
||||
|
||||
// Write the token cookie from the response body so it works cross-origin
|
||||
// (control plane UI → worker). The Set-Cookie header is still sent for
|
||||
// same-origin deployments as a fallback.
|
||||
if (data.token) {
|
||||
document.cookie = `token=${data.token}; path=/; SameSite=Lax`;
|
||||
}
|
||||
|
||||
return data;
|
||||
};
|
||||
|
||||
|
|
|
|||
65
ui/litellm-dashboard/src/hooks/useWorker.ts
Normal file
65
ui/litellm-dashboard/src/hooks/useWorker.ts
Normal file
|
|
@ -0,0 +1,65 @@
|
|||
"use client";
|
||||
|
||||
import { useState, useCallback, useEffect } from "react";
|
||||
import { switchToWorkerUrl, WorkerInfo } from "@/components/networking";
|
||||
import { useUIConfig } from "@/app/(dashboard)/hooks/uiConfig/useUIConfig";
|
||||
|
||||
const SELECTED_WORKER_KEY = "litellm_selected_worker_id";
|
||||
|
||||
interface UseWorkerReturn {
|
||||
isControlPlane: boolean;
|
||||
workers: WorkerInfo[];
|
||||
selectedWorkerId: string | null;
|
||||
selectedWorker: WorkerInfo | null;
|
||||
selectWorker: (workerId: string) => void;
|
||||
disconnectFromWorker: () => void;
|
||||
}
|
||||
|
||||
export const useWorker = (): UseWorkerReturn => {
|
||||
const { data: uiConfig } = useUIConfig();
|
||||
const isControlPlane = uiConfig?.is_control_plane ?? false;
|
||||
const workers: WorkerInfo[] = uiConfig?.workers ?? [];
|
||||
|
||||
const [selectedWorkerId, setSelectedWorkerId] = useState<string | null>(() => {
|
||||
if (typeof window === "undefined") return null;
|
||||
return localStorage.getItem(SELECTED_WORKER_KEY);
|
||||
});
|
||||
|
||||
// Once workers are loaded, restore proxyBaseUrl from the persisted selection
|
||||
useEffect(() => {
|
||||
if (!selectedWorkerId || workers.length === 0) return;
|
||||
const worker = workers.find((w) => w.worker_id === selectedWorkerId);
|
||||
if (worker) {
|
||||
switchToWorkerUrl(worker.url);
|
||||
}
|
||||
}, [selectedWorkerId, workers]);
|
||||
|
||||
const selectedWorker =
|
||||
workers.find((w) => w.worker_id === selectedWorkerId) ?? null;
|
||||
|
||||
const selectWorker = useCallback(
|
||||
(workerId: string) => {
|
||||
const worker = workers.find((w) => w.worker_id === workerId);
|
||||
if (!worker) return;
|
||||
setSelectedWorkerId(workerId);
|
||||
localStorage.setItem(SELECTED_WORKER_KEY, workerId);
|
||||
switchToWorkerUrl(worker.url);
|
||||
},
|
||||
[workers],
|
||||
);
|
||||
|
||||
const disconnectFromWorker = useCallback(() => {
|
||||
setSelectedWorkerId(null);
|
||||
localStorage.removeItem(SELECTED_WORKER_KEY);
|
||||
switchToWorkerUrl(null);
|
||||
}, []);
|
||||
|
||||
return {
|
||||
isControlPlane,
|
||||
workers,
|
||||
selectedWorkerId,
|
||||
selectedWorker,
|
||||
selectWorker,
|
||||
disconnectFromWorker,
|
||||
};
|
||||
};
|
||||
|
|
@ -14,7 +14,7 @@
|
|||
"moduleResolution": "bundler",
|
||||
"resolveJsonModule": true,
|
||||
"isolatedModules": true,
|
||||
"jsx": "preserve",
|
||||
"jsx": "react-jsx",
|
||||
"incremental": true,
|
||||
"plugins": [
|
||||
{
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue