mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-10 22:41:41 +00:00
Merge branch 'main' into main
This commit is contained in:
Dávid Schneidhoffer
GitHub
commit
145acb7947
19 changed files with 1077 additions and 30 deletions
|
|
@ -1672,6 +1672,7 @@ class NewTeamRequest(TeamBase):
|
|||
int
|
||||
] = None # allow user to set TPM limit for all team members
|
||||
team_member_key_duration: Optional[str] = None # e.g. "1d", "1w", "1m"
|
||||
team_member_budget_duration: Optional[str] = None # e.g. "30d", "1mo"
|
||||
allowed_vector_store_indexes: Optional[List[AllowedVectorStoreIndexItem]] = None
|
||||
enforced_batch_output_expires_after: Optional[dict] = None
|
||||
enforced_file_expires_after: Optional[dict] = None
|
||||
|
|
|
|||
|
|
@ -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 [],
|
||||
)
|
||||
|
|
|
|||
|
|
@ -724,6 +724,7 @@ async def new_team( # noqa: PLR0915
|
|||
- disable_global_guardrails: Optional[bool] - Whether to disable global guardrails for the key.
|
||||
- object_permission: Optional[LiteLLM_ObjectPermissionBase] - team-specific object permission. Example - {"vector_stores": ["vector_store_1", "vector_store_2"], "agents": ["agent_1", "agent_2"], "agent_access_groups": ["dev_group"]}. IF null or {} then no object permission.
|
||||
- team_member_budget: Optional[float] - The maximum budget allocated to an individual team member.
|
||||
- team_member_budget_duration: Optional[str] - The duration of the budget for the team member. Doc [here](https://docs.litellm.ai/docs/proxy/team_budgets)
|
||||
- team_member_rpm_limit: Optional[int] - The RPM (Requests Per Minute) limit for individual team members.
|
||||
- team_member_tpm_limit: Optional[int] - The TPM (Tokens Per Minute) limit for individual team members.
|
||||
- team_member_key_duration: Optional[str] - The duration for a team member's key. e.g. "1d", "1w", "1mo"
|
||||
|
|
@ -934,6 +935,7 @@ async def new_team( # noqa: PLR0915
|
|||
team_member_budget=data.team_member_budget,
|
||||
team_member_rpm_limit=data.team_member_rpm_limit,
|
||||
team_member_tpm_limit=data.team_member_tpm_limit,
|
||||
team_member_budget_duration=data.team_member_budget_duration,
|
||||
):
|
||||
data_json = await TeamMemberBudgetHandler.create_team_member_budget_table(
|
||||
data=data,
|
||||
|
|
@ -942,6 +944,7 @@ async def new_team( # noqa: PLR0915
|
|||
team_member_budget=data.team_member_budget,
|
||||
team_member_rpm_limit=data.team_member_rpm_limit,
|
||||
team_member_tpm_limit=data.team_member_tpm_limit,
|
||||
team_member_budget_duration=data.team_member_budget_duration,
|
||||
)
|
||||
|
||||
## ADD TO TEAM TABLE
|
||||
|
|
|
|||
|
|
@ -16,6 +16,7 @@ import os
|
|||
import secrets
|
||||
from copy import deepcopy
|
||||
from typing import TYPE_CHECKING, Any, Dict, List, Literal, Optional, Tuple, Union, cast
|
||||
from urllib.parse import urlencode, urlparse
|
||||
|
||||
if TYPE_CHECKING:
|
||||
import httpx
|
||||
|
|
@ -301,6 +302,7 @@ async def google_login(
|
|||
source: Optional[str] = None,
|
||||
key: Optional[str] = None,
|
||||
existing_key: Optional[str] = None,
|
||||
return_to: Optional[str] = None,
|
||||
): # noqa: PLR0915
|
||||
"""
|
||||
Create Proxy API Keys using Google Workspace SSO. Requires setting PROXY_BASE_URL in .env
|
||||
|
|
@ -394,13 +396,23 @@ async def google_login(
|
|||
is True
|
||||
):
|
||||
verbose_proxy_logger.info(f"Redirecting to SSO login for {redirect_url}")
|
||||
return await SSOAuthenticationHandler.get_sso_login_redirect(
|
||||
sso_redirect = await SSOAuthenticationHandler.get_sso_login_redirect(
|
||||
redirect_url=redirect_url,
|
||||
microsoft_client_id=microsoft_client_id,
|
||||
google_client_id=google_client_id,
|
||||
generic_client_id=generic_client_id,
|
||||
state=cli_state,
|
||||
)
|
||||
if return_to is not None and sso_redirect is not None:
|
||||
SSOAuthenticationHandler._validate_return_to(return_to)
|
||||
sso_redirect.set_cookie(
|
||||
key="litellm_cp_return_to",
|
||||
value=return_to,
|
||||
max_age=600,
|
||||
httponly=True,
|
||||
samesite="lax",
|
||||
)
|
||||
return sso_redirect
|
||||
elif ui_username is not None:
|
||||
# No Google, Microsoft SSO
|
||||
# Use UI Credentials set in .env
|
||||
|
|
@ -1312,12 +1324,17 @@ async def auth_callback(request: Request, state: Optional[str] = None): # noqa:
|
|||
request=request, key=key_id, existing_key=existing_key, result=result
|
||||
)
|
||||
|
||||
# Control-plane cross-origin: read return_to from cookie.
|
||||
# Starlette's cookie_parser already handles RFC 2109 unquoting.
|
||||
cp_return_to: Optional[str] = request.cookies.get("litellm_cp_return_to")
|
||||
|
||||
return await SSOAuthenticationHandler.get_redirect_response_from_openid(
|
||||
result=result,
|
||||
request=request,
|
||||
received_response=received_response,
|
||||
generic_client_id=generic_client_id,
|
||||
ui_access_mode=ui_access_mode,
|
||||
return_to=cp_return_to,
|
||||
)
|
||||
|
||||
|
||||
|
|
@ -1760,6 +1777,38 @@ class SSOAuthenticationHandler:
|
|||
Handler for SSO Authentication across all SSO providers
|
||||
"""
|
||||
|
||||
@staticmethod
|
||||
def _validate_return_to(return_to: str) -> None:
|
||||
"""
|
||||
Validate that return_to matches the configured control_plane_url origin.
|
||||
|
||||
Raises HTTPException(400) if:
|
||||
- control_plane_url is not configured in general_settings
|
||||
- return_to origin does not match control_plane_url origin
|
||||
"""
|
||||
from litellm.proxy.proxy_server import general_settings
|
||||
|
||||
control_plane_url = general_settings.get("control_plane_url")
|
||||
if control_plane_url is None:
|
||||
raise HTTPException(
|
||||
status_code=400,
|
||||
detail="return_to is not allowed: control_plane_url is not configured",
|
||||
)
|
||||
|
||||
def _origin(url: str) -> tuple:
|
||||
parsed = urlparse(url)
|
||||
scheme = (parsed.scheme or "").lower()
|
||||
hostname = (parsed.hostname or "").lower()
|
||||
default_port = 443 if scheme == "https" else 80
|
||||
port = parsed.port if parsed.port is not None else default_port
|
||||
return (scheme, hostname, port)
|
||||
|
||||
if _origin(return_to) != _origin(control_plane_url):
|
||||
raise HTTPException(
|
||||
status_code=400,
|
||||
detail="return_to does not match the configured control_plane_url",
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
async def get_sso_login_redirect(
|
||||
redirect_url: str,
|
||||
|
|
@ -2358,6 +2407,7 @@ class SSOAuthenticationHandler:
|
|||
received_response: Optional[dict] = None,
|
||||
generic_client_id: Optional[str] = None,
|
||||
ui_access_mode: Optional[Dict] = None,
|
||||
return_to: Optional[str] = None,
|
||||
) -> RedirectResponse:
|
||||
import jwt
|
||||
|
||||
|
|
@ -2367,6 +2417,7 @@ class SSOAuthenticationHandler:
|
|||
master_key,
|
||||
premium_user,
|
||||
proxy_logging_obj,
|
||||
redis_usage_cache,
|
||||
user_api_key_cache,
|
||||
user_custom_sso,
|
||||
)
|
||||
|
|
@ -2534,6 +2585,36 @@ class SSOAuthenticationHandler:
|
|||
master_key or "",
|
||||
algorithm="HS256",
|
||||
)
|
||||
|
||||
# Control-plane cross-origin: store JWT behind a single-use opaque
|
||||
# code (60s TTL) so the token never appears in browser history / logs.
|
||||
# The control plane redeems it via POST /v3/login/exchange.
|
||||
if return_to is not None:
|
||||
SSOAuthenticationHandler._validate_return_to(return_to)
|
||||
|
||||
code = secrets.token_urlsafe(32)
|
||||
cache_key = f"login_code:{code}"
|
||||
cache_value = {"token": jwt_token, "redirect_url": return_to}
|
||||
if redis_usage_cache is not None:
|
||||
await redis_usage_cache.async_set_cache(
|
||||
key=cache_key, value=cache_value, ttl=60
|
||||
)
|
||||
else:
|
||||
await user_api_key_cache.async_set_cache(
|
||||
key=cache_key, value=cache_value, ttl=60
|
||||
)
|
||||
|
||||
separator = "&" if "?" in return_to else "?"
|
||||
redirect_url = (
|
||||
return_to + separator + urlencode({"login": "success", "code": code})
|
||||
)
|
||||
verbose_proxy_logger.info(
|
||||
"Cross-origin SSO: redirecting to control plane with login code"
|
||||
)
|
||||
redirect_response = RedirectResponse(url=redirect_url, status_code=303)
|
||||
redirect_response.delete_cookie("litellm_cp_return_to")
|
||||
return redirect_response
|
||||
|
||||
if user_id is not None and isinstance(user_id, str):
|
||||
litellm_dashboard_ui += "?login=success"
|
||||
verbose_proxy_logger.info(f"Redirecting to {litellm_dashboard_ui}")
|
||||
|
|
|
|||
|
|
@ -541,6 +541,7 @@ from litellm.types.llms.anthropic import (
|
|||
AnthropicResponseUsageBlock,
|
||||
)
|
||||
from litellm.types.llms.openai import HttpxBinaryResponseContent
|
||||
from litellm.types.proxy.control_plane_endpoints import WorkerRegistryEntry
|
||||
from litellm.types.proxy.management_endpoints.model_management_endpoints import (
|
||||
ModelGroupInfoProxy,
|
||||
)
|
||||
|
|
@ -1546,6 +1547,7 @@ user_custom_key_generate = None
|
|||
# Sentinel: prevents PKCE-no-Redis advisory from re-logging on config hot-reload.
|
||||
# Tests that need to reset it can patch 'litellm.proxy.proxy_server._pkce_no_redis_warning_emitted'.
|
||||
_pkce_no_redis_warning_emitted: bool = False
|
||||
_cp_no_redis_warning_emitted: bool = False
|
||||
user_custom_sso = None
|
||||
user_custom_ui_sso_sign_in_handler = None
|
||||
use_background_health_checks = None
|
||||
|
|
@ -2295,6 +2297,7 @@ class ProxyConfig:
|
|||
self.config: Dict[str, Any] = {}
|
||||
self._last_semantic_filter_config: Optional[Dict[str, Any]] = None
|
||||
self._last_hashicorp_vault_config: Optional[Dict[str, Any]] = None
|
||||
self.worker_registry: List["WorkerRegistryEntry"] = []
|
||||
|
||||
def is_yaml(self, config_file_path: str) -> bool:
|
||||
if not os.path.isfile(config_file_path):
|
||||
|
|
@ -3095,6 +3098,21 @@ class ProxyConfig:
|
|||
"Set PKCE_STRICT_CACHE_MISS=true to fail fast with a 401 on cache misses "
|
||||
"instead of continuing without a code_verifier."
|
||||
)
|
||||
|
||||
### CONTROL PLANE CODE-EXCHANGE PREREQUISITE CHECK ###
|
||||
cp_url = general_settings.get("control_plane_url")
|
||||
if cp_url and redis_usage_cache is None:
|
||||
global _cp_no_redis_warning_emitted
|
||||
if not _cp_no_redis_warning_emitted:
|
||||
_cp_no_redis_warning_emitted = True
|
||||
verbose_proxy_logger.warning(
|
||||
"control_plane_url is configured but Redis is not configured for LiteLLM caching. "
|
||||
"Login codes (SSO and /v3/login) will not be shared across instances — "
|
||||
"the /v3/login/exchange call may land on a different pod and fail with 401. "
|
||||
"Configure Redis via the 'cache' section in your proxy config, "
|
||||
"or ensure sticky sessions for single-instance deployments."
|
||||
)
|
||||
|
||||
### STORE MODEL IN DB ### feature flag for `/model/new`
|
||||
store_model_in_db = general_settings.get("store_model_in_db", False)
|
||||
if store_model_in_db is None:
|
||||
|
|
@ -3385,7 +3403,15 @@ class ProxyConfig:
|
|||
litellm.vector_store_registry.load_vector_stores_from_config(
|
||||
vector_store_registry_config
|
||||
)
|
||||
pass
|
||||
|
||||
## WORKER REGISTRY (Control Plane)
|
||||
worker_registry_config = config.get("worker_registry", None)
|
||||
if worker_registry_config:
|
||||
self.worker_registry = [
|
||||
WorkerRegistryEntry(**e) for e in worker_registry_config
|
||||
]
|
||||
else:
|
||||
self.worker_registry = []
|
||||
|
||||
async def _init_policy_engine(
|
||||
self,
|
||||
|
|
@ -11095,6 +11121,165 @@ async def login_v2(request: Request): # noqa: PLR0915
|
|||
)
|
||||
|
||||
|
||||
@router.post(
|
||||
"/v3/login", include_in_schema=False
|
||||
) # control-plane login — always returns token in body for cross-origin use
|
||||
async def login_v3(request: Request): # noqa: PLR0915
|
||||
global premium_user, general_settings, master_key
|
||||
from litellm.proxy.auth.login_utils import authenticate_user, create_ui_token_object
|
||||
from litellm.proxy.utils import get_custom_url
|
||||
|
||||
try:
|
||||
if not general_settings.get("control_plane_url"):
|
||||
raise ProxyException(
|
||||
message="/v3/login is only available on workers with control_plane_url configured",
|
||||
type=ProxyErrorTypes.not_found_error,
|
||||
param="control_plane_url",
|
||||
code=status.HTTP_404_NOT_FOUND,
|
||||
)
|
||||
|
||||
body = await request.json()
|
||||
username = str(body.get("username"))
|
||||
password = str(body.get("password"))
|
||||
|
||||
login_result = await authenticate_user(
|
||||
username=username,
|
||||
password=password,
|
||||
master_key=master_key,
|
||||
prisma_client=prisma_client,
|
||||
)
|
||||
|
||||
returned_ui_token_object = create_ui_token_object(
|
||||
login_result=login_result,
|
||||
general_settings=general_settings,
|
||||
premium_user=premium_user,
|
||||
)
|
||||
|
||||
import jwt
|
||||
|
||||
jwt_token = jwt.encode(
|
||||
cast(dict, returned_ui_token_object),
|
||||
cast(str, master_key),
|
||||
algorithm="HS256",
|
||||
)
|
||||
|
||||
litellm_dashboard_ui = get_custom_url(str(request.base_url))
|
||||
if litellm_dashboard_ui.endswith("/"):
|
||||
litellm_dashboard_ui += "ui/"
|
||||
else:
|
||||
litellm_dashboard_ui += "/ui/"
|
||||
litellm_dashboard_ui += "?login=success"
|
||||
|
||||
# Store JWT behind a single-use opaque code (60s TTL)
|
||||
code = secrets.token_urlsafe(32)
|
||||
cache_key = f"login_code:{code}"
|
||||
cache_value = {"token": jwt_token, "redirect_url": litellm_dashboard_ui}
|
||||
if redis_usage_cache is not None:
|
||||
await redis_usage_cache.async_set_cache(
|
||||
key=cache_key, value=cache_value, ttl=60
|
||||
)
|
||||
else:
|
||||
await user_api_key_cache.async_set_cache(
|
||||
key=cache_key, value=cache_value, ttl=60
|
||||
)
|
||||
|
||||
return JSONResponse(
|
||||
content={"code": code, "expires_in": 60},
|
||||
status_code=status.HTTP_200_OK,
|
||||
)
|
||||
except Exception as e:
|
||||
verbose_proxy_logger.exception(
|
||||
"litellm.proxy.proxy_server.login_v3(): Exception occurred - {}".format(
|
||||
str(e)
|
||||
)
|
||||
)
|
||||
if isinstance(e, ProxyException):
|
||||
raise e
|
||||
elif isinstance(e, HTTPException):
|
||||
raise ProxyException(
|
||||
message=getattr(e, "detail", str(e)),
|
||||
type=ProxyErrorTypes.auth_error,
|
||||
param=getattr(e, "param", "None"),
|
||||
code=getattr(e, "status_code", status.HTTP_500_INTERNAL_SERVER_ERROR),
|
||||
)
|
||||
else:
|
||||
error_msg = f"{str(e)}"
|
||||
raise ProxyException(
|
||||
message=error_msg,
|
||||
type=ProxyErrorTypes.auth_error,
|
||||
param="None",
|
||||
code=status.HTTP_500_INTERNAL_SERVER_ERROR,
|
||||
)
|
||||
|
||||
|
||||
@router.post(
|
||||
"/v3/login/exchange", include_in_schema=False
|
||||
) # exchange single-use opaque code for JWT
|
||||
async def login_v3_exchange(request: Request):
|
||||
try:
|
||||
if not general_settings.get("control_plane_url"):
|
||||
raise ProxyException(
|
||||
message="/v3/login/exchange is only available on workers with control_plane_url configured",
|
||||
type=ProxyErrorTypes.not_found_error,
|
||||
param="control_plane_url",
|
||||
code=status.HTTP_404_NOT_FOUND,
|
||||
)
|
||||
|
||||
body = await request.json()
|
||||
code = body.get("code")
|
||||
if not code:
|
||||
raise ProxyException(
|
||||
message="Missing 'code' parameter",
|
||||
type=ProxyErrorTypes.auth_error,
|
||||
param="code",
|
||||
code=status.HTTP_400_BAD_REQUEST,
|
||||
)
|
||||
|
||||
cache_key = f"login_code:{code}"
|
||||
if redis_usage_cache is not None:
|
||||
cached_data = await redis_usage_cache.async_get_cache(key=cache_key)
|
||||
else:
|
||||
cached_data = await user_api_key_cache.async_get_cache(key=cache_key)
|
||||
|
||||
if not cached_data or not isinstance(cached_data, dict):
|
||||
raise ProxyException(
|
||||
message="Invalid or expired login code",
|
||||
type=ProxyErrorTypes.auth_error,
|
||||
param="code",
|
||||
code=status.HTTP_401_UNAUTHORIZED,
|
||||
)
|
||||
|
||||
# Single-use: delete immediately
|
||||
if redis_usage_cache is not None:
|
||||
await redis_usage_cache.async_delete_cache(key=cache_key)
|
||||
else:
|
||||
await user_api_key_cache.async_delete_cache(key=cache_key)
|
||||
|
||||
json_response = JSONResponse(
|
||||
content={
|
||||
"token": cached_data["token"],
|
||||
"redirect_url": cached_data["redirect_url"],
|
||||
},
|
||||
status_code=status.HTTP_200_OK,
|
||||
)
|
||||
json_response.set_cookie(key="token", value=cached_data["token"])
|
||||
return json_response
|
||||
except ProxyException:
|
||||
raise
|
||||
except Exception as e:
|
||||
verbose_proxy_logger.exception(
|
||||
"litellm.proxy.proxy_server.login_v3_exchange(): Exception occurred - {}".format(
|
||||
str(e)
|
||||
)
|
||||
)
|
||||
raise ProxyException(
|
||||
message=str(e),
|
||||
type=ProxyErrorTypes.auth_error,
|
||||
param="None",
|
||||
code=status.HTTP_500_INTERNAL_SERVER_ERROR,
|
||||
)
|
||||
|
||||
|
||||
@app.get("/onboarding/get_token", include_in_schema=False)
|
||||
async def onboarding(invite_link: str, request: Request):
|
||||
"""
|
||||
|
|
|
|||
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"] == []
|
||||
|
|
|
|||
|
|
@ -6441,3 +6441,53 @@ async def test_list_team_v1_batches_key_queries():
|
|||
assert result[0].keys == [key1, key2]
|
||||
assert result[1].team_id == "team-2"
|
||||
assert result[1].keys == [key3]
|
||||
|
||||
|
||||
def test_new_team_request_accepts_team_member_budget_duration():
|
||||
"""Test that NewTeamRequest does not silently drop team_member_budget_duration."""
|
||||
from litellm.proxy._types import NewTeamRequest
|
||||
|
||||
request = NewTeamRequest(
|
||||
team_member_budget=20.0,
|
||||
team_member_budget_duration="30d",
|
||||
)
|
||||
assert request.team_member_budget == 20.0
|
||||
assert request.team_member_budget_duration == "30d"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_create_team_member_budget_table_with_duration():
|
||||
"""Verify that create_team_member_budget_table passes budget_duration
|
||||
through to the new_budget call when team_member_budget_duration is provided."""
|
||||
from litellm.proxy._types import NewTeamRequest, UserAPIKeyAuth, LitellmUserRoles
|
||||
from litellm.proxy.management_endpoints.team_endpoints import TeamMemberBudgetHandler
|
||||
|
||||
mock_budget_response = MagicMock(budget_id="budget-abc")
|
||||
mock_admin = UserAPIKeyAuth(
|
||||
user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN
|
||||
)
|
||||
|
||||
data = NewTeamRequest(
|
||||
team_alias="test-team",
|
||||
team_member_budget=20.0,
|
||||
team_member_budget_duration="30d",
|
||||
)
|
||||
|
||||
with patch(
|
||||
"litellm.proxy.management_endpoints.budget_management_endpoints.new_budget",
|
||||
new_callable=AsyncMock,
|
||||
return_value=mock_budget_response,
|
||||
) as mock_new_budget:
|
||||
result = await TeamMemberBudgetHandler.create_team_member_budget_table(
|
||||
data=data,
|
||||
new_team_data_json={"metadata": None},
|
||||
user_api_key_dict=mock_admin,
|
||||
team_member_budget=20.0,
|
||||
team_member_budget_duration="30d",
|
||||
)
|
||||
|
||||
mock_new_budget.assert_awaited_once()
|
||||
budget_request = mock_new_budget.call_args.kwargs["budget_obj"]
|
||||
assert budget_request.budget_duration == "30d"
|
||||
assert budget_request.max_budget == 20.0
|
||||
assert result["metadata"]["team_member_budget_id"] == "budget-abc"
|
||||
|
|
|
|||
|
|
@ -6,7 +6,7 @@ from unittest.mock import AsyncMock, MagicMock, patch
|
|||
|
||||
import httpx
|
||||
import pytest
|
||||
from fastapi import Request
|
||||
from fastapi import HTTPException, Request
|
||||
|
||||
from litellm._uuid import uuid
|
||||
|
||||
|
|
@ -5160,3 +5160,99 @@ def test_generic_response_convertor_extra_attributes_missing_field(monkeypatch):
|
|||
assert result.extra_fields["missing_field"] is None
|
||||
assert result.extra_fields["another_missing"] is None
|
||||
|
||||
|
||||
class TestValidateReturnTo:
|
||||
"""Tests for SSOAuthenticationHandler._validate_return_to"""
|
||||
|
||||
def test_rejects_when_no_control_plane_url_configured(self, monkeypatch):
|
||||
"""return_to should be rejected if control_plane_url is not in general_settings."""
|
||||
monkeypatch.setattr(
|
||||
"litellm.proxy.proxy_server.general_settings", {}
|
||||
)
|
||||
with pytest.raises(HTTPException) as exc_info:
|
||||
SSOAuthenticationHandler._validate_return_to("https://cp.example.com/ui")
|
||||
assert exc_info.value.status_code == 400
|
||||
assert "not configured" in exc_info.value.detail
|
||||
|
||||
def test_allows_matching_origin(self, monkeypatch):
|
||||
"""return_to matching the configured control_plane_url origin should pass."""
|
||||
monkeypatch.setattr(
|
||||
"litellm.proxy.proxy_server.general_settings",
|
||||
{"control_plane_url": "https://cp.example.com"},
|
||||
)
|
||||
# Should not raise
|
||||
SSOAuthenticationHandler._validate_return_to("https://cp.example.com/ui?page=models")
|
||||
|
||||
def test_allows_matching_origin_with_trailing_slash(self, monkeypatch):
|
||||
"""Trailing slash on control_plane_url should not affect origin comparison."""
|
||||
monkeypatch.setattr(
|
||||
"litellm.proxy.proxy_server.general_settings",
|
||||
{"control_plane_url": "https://cp.example.com/"},
|
||||
)
|
||||
SSOAuthenticationHandler._validate_return_to("https://cp.example.com/ui")
|
||||
|
||||
def test_rejects_prefix_attack(self, monkeypatch):
|
||||
"""return_to like cp.example.com.evil.com must be rejected (not just prefix match)."""
|
||||
monkeypatch.setattr(
|
||||
"litellm.proxy.proxy_server.general_settings",
|
||||
{"control_plane_url": "https://cp.example.com"},
|
||||
)
|
||||
with pytest.raises(HTTPException) as exc_info:
|
||||
SSOAuthenticationHandler._validate_return_to("https://cp.example.com.evil.com/steal")
|
||||
assert exc_info.value.status_code == 400
|
||||
|
||||
def test_rejects_different_origin(self, monkeypatch):
|
||||
"""return_to pointing to a completely different domain should be rejected."""
|
||||
monkeypatch.setattr(
|
||||
"litellm.proxy.proxy_server.general_settings",
|
||||
{"control_plane_url": "https://cp.example.com"},
|
||||
)
|
||||
with pytest.raises(HTTPException) as exc_info:
|
||||
SSOAuthenticationHandler._validate_return_to("https://evil.com/phish")
|
||||
assert exc_info.value.status_code == 400
|
||||
|
||||
def test_case_insensitive_hostname(self, monkeypatch):
|
||||
"""Hostname comparison should be case-insensitive per RFC 3986."""
|
||||
monkeypatch.setattr(
|
||||
"litellm.proxy.proxy_server.general_settings",
|
||||
{"control_plane_url": "https://CP.Example.COM"},
|
||||
)
|
||||
# Should not raise
|
||||
SSOAuthenticationHandler._validate_return_to("https://cp.example.com/ui")
|
||||
|
||||
def test_rejects_scheme_mismatch(self, monkeypatch):
|
||||
"""http:// must be rejected when control_plane_url uses https://."""
|
||||
monkeypatch.setattr(
|
||||
"litellm.proxy.proxy_server.general_settings",
|
||||
{"control_plane_url": "https://cp.example.com"},
|
||||
)
|
||||
with pytest.raises(HTTPException) as exc_info:
|
||||
SSOAuthenticationHandler._validate_return_to("http://cp.example.com/ui")
|
||||
assert exc_info.value.status_code == 400
|
||||
|
||||
def test_rejects_port_mismatch(self, monkeypatch):
|
||||
"""Non-default port must be rejected."""
|
||||
monkeypatch.setattr(
|
||||
"litellm.proxy.proxy_server.general_settings",
|
||||
{"control_plane_url": "https://cp.example.com"},
|
||||
)
|
||||
with pytest.raises(HTTPException) as exc_info:
|
||||
SSOAuthenticationHandler._validate_return_to("https://cp.example.com:8443/ui")
|
||||
assert exc_info.value.status_code == 400
|
||||
|
||||
def test_allows_explicit_default_port(self, monkeypatch):
|
||||
"""https://host:443 should match https://host (default port normalisation)."""
|
||||
monkeypatch.setattr(
|
||||
"litellm.proxy.proxy_server.general_settings",
|
||||
{"control_plane_url": "https://cp.example.com"},
|
||||
)
|
||||
SSOAuthenticationHandler._validate_return_to("https://cp.example.com:443/ui")
|
||||
|
||||
def test_allows_matching_custom_port(self, monkeypatch):
|
||||
"""Both sides on the same custom port should match."""
|
||||
monkeypatch.setattr(
|
||||
"litellm.proxy.proxy_server.general_settings",
|
||||
{"control_plane_url": "https://cp.example.com:3000"},
|
||||
)
|
||||
SSOAuthenticationHandler._validate_return_to("https://cp.example.com:3000/ui")
|
||||
|
||||
|
|
|
|||
|
|
@ -236,6 +236,217 @@ def test_login_v2_returns_json_on_invalid_json_body(monkeypatch):
|
|||
assert isinstance(data["error"], dict)
|
||||
|
||||
|
||||
def test_login_v3_rejected_without_control_plane_url(monkeypatch):
|
||||
"""v3/login returns 404 when control_plane_url is not configured."""
|
||||
mock_prisma_client = MagicMock()
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.master_key", "test-master-key")
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", {})
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client)
|
||||
|
||||
client = TestClient(app)
|
||||
response = client.post(
|
||||
"/v3/login",
|
||||
json={"username": "alice", "password": "secret"},
|
||||
)
|
||||
|
||||
assert response.status_code == 404
|
||||
assert "control_plane_url" in response.json()["error"]["message"]
|
||||
|
||||
|
||||
def test_login_v3_returns_code(monkeypatch):
|
||||
"""v3/login returns an opaque code, not the JWT directly."""
|
||||
mock_prisma_client = MagicMock()
|
||||
monkeypatch.setattr(
|
||||
"litellm.proxy.auth.login_utils.authenticate_user",
|
||||
AsyncMock(return_value={"user_id": "test-user"}),
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
"litellm.proxy.auth.login_utils.create_ui_token_object",
|
||||
MagicMock(return_value={"user_id": "test-user"}),
|
||||
)
|
||||
monkeypatch.setattr("jwt.encode", MagicMock(return_value="signed-token"))
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.master_key", "test-master-key")
|
||||
monkeypatch.setattr(
|
||||
"litellm.proxy.proxy_server.general_settings",
|
||||
{"control_plane_url": "https://cp.example.com"},
|
||||
)
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.premium_user", False)
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client)
|
||||
mock_config = MagicMock()
|
||||
mock_config.worker_registry = []
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.proxy_config", mock_config)
|
||||
monkeypatch.setattr("litellm.proxy.utils.get_server_root_path", lambda: "")
|
||||
monkeypatch.setattr("litellm.proxy.utils.get_proxy_base_url", lambda: None)
|
||||
|
||||
client = TestClient(app)
|
||||
response = client.post(
|
||||
"/v3/login",
|
||||
json={"username": "alice", "password": "secret"},
|
||||
)
|
||||
|
||||
assert response.status_code == 200
|
||||
data = response.json()
|
||||
assert "code" in data
|
||||
assert data["expires_in"] == 60
|
||||
assert "token" not in data
|
||||
|
||||
|
||||
def test_login_v3_exchange_happy_path(monkeypatch):
|
||||
"""Full flow: v3/login returns code, v3/login/exchange redeems it for JWT."""
|
||||
mock_prisma_client = MagicMock()
|
||||
monkeypatch.setattr(
|
||||
"litellm.proxy.auth.login_utils.authenticate_user",
|
||||
AsyncMock(return_value={"user_id": "test-user"}),
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
"litellm.proxy.auth.login_utils.create_ui_token_object",
|
||||
MagicMock(return_value={"user_id": "test-user"}),
|
||||
)
|
||||
monkeypatch.setattr("jwt.encode", MagicMock(return_value="signed-token"))
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.master_key", "test-master-key")
|
||||
monkeypatch.setattr(
|
||||
"litellm.proxy.proxy_server.general_settings",
|
||||
{"control_plane_url": "https://cp.example.com"},
|
||||
)
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.premium_user", False)
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client)
|
||||
mock_config = MagicMock()
|
||||
mock_config.worker_registry = []
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.proxy_config", mock_config)
|
||||
monkeypatch.setattr("litellm.proxy.utils.get_server_root_path", lambda: "")
|
||||
monkeypatch.setattr("litellm.proxy.utils.get_proxy_base_url", lambda: None)
|
||||
|
||||
client = TestClient(app)
|
||||
|
||||
# Step 1: login — get code
|
||||
login_response = client.post(
|
||||
"/v3/login",
|
||||
json={"username": "alice", "password": "secret"},
|
||||
)
|
||||
assert login_response.status_code == 200
|
||||
code = login_response.json()["code"]
|
||||
|
||||
# Step 2: exchange — get JWT
|
||||
exchange_response = client.post(
|
||||
"/v3/login/exchange",
|
||||
json={"code": code},
|
||||
)
|
||||
assert exchange_response.status_code == 200
|
||||
exchange_data = exchange_response.json()
|
||||
assert exchange_data["token"] == "signed-token"
|
||||
assert "redirect_url" in exchange_data
|
||||
assert exchange_response.cookies.get("token") == "signed-token"
|
||||
|
||||
|
||||
def test_login_v3_exchange_single_use(monkeypatch):
|
||||
"""Code can only be redeemed once."""
|
||||
mock_prisma_client = MagicMock()
|
||||
monkeypatch.setattr(
|
||||
"litellm.proxy.auth.login_utils.authenticate_user",
|
||||
AsyncMock(return_value={"user_id": "test-user"}),
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
"litellm.proxy.auth.login_utils.create_ui_token_object",
|
||||
MagicMock(return_value={"user_id": "test-user"}),
|
||||
)
|
||||
monkeypatch.setattr("jwt.encode", MagicMock(return_value="signed-token"))
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.master_key", "test-master-key")
|
||||
monkeypatch.setattr(
|
||||
"litellm.proxy.proxy_server.general_settings",
|
||||
{"control_plane_url": "https://cp.example.com"},
|
||||
)
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.premium_user", False)
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client)
|
||||
mock_config = MagicMock()
|
||||
mock_config.worker_registry = []
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.proxy_config", mock_config)
|
||||
monkeypatch.setattr("litellm.proxy.utils.get_server_root_path", lambda: "")
|
||||
monkeypatch.setattr("litellm.proxy.utils.get_proxy_base_url", lambda: None)
|
||||
|
||||
client = TestClient(app)
|
||||
|
||||
login_response = client.post(
|
||||
"/v3/login",
|
||||
json={"username": "alice", "password": "secret"},
|
||||
)
|
||||
code = login_response.json()["code"]
|
||||
|
||||
# First exchange succeeds
|
||||
first = client.post("/v3/login/exchange", json={"code": code})
|
||||
assert first.status_code == 200
|
||||
|
||||
# Second exchange fails
|
||||
second = client.post("/v3/login/exchange", json={"code": code})
|
||||
assert second.status_code == 401
|
||||
|
||||
|
||||
def test_login_v3_exchange_invalid_code(monkeypatch):
|
||||
"""Random code returns 401."""
|
||||
monkeypatch.setattr(
|
||||
"litellm.proxy.proxy_server.general_settings",
|
||||
{"control_plane_url": "https://cp.example.com"},
|
||||
)
|
||||
client = TestClient(app)
|
||||
response = client.post(
|
||||
"/v3/login/exchange",
|
||||
json={"code": "nonexistent-code"},
|
||||
)
|
||||
assert response.status_code == 401
|
||||
|
||||
|
||||
def test_login_v3_exchange_rejected_without_control_plane_url(monkeypatch):
|
||||
"""v3/login/exchange returns 404 when control_plane_url is not configured."""
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", {})
|
||||
|
||||
client = TestClient(app)
|
||||
response = client.post(
|
||||
"/v3/login/exchange",
|
||||
json={"code": "some-code"},
|
||||
)
|
||||
|
||||
assert response.status_code == 404
|
||||
assert "control_plane_url" in response.json()["error"]["message"]
|
||||
|
||||
|
||||
def test_login_v3_returns_json_on_proxy_exception(monkeypatch):
|
||||
"""Test that /v3/login returns JSON error when ProxyException is raised"""
|
||||
from litellm.proxy._types import ProxyErrorTypes, ProxyException
|
||||
|
||||
mock_prisma_client = MagicMock()
|
||||
mock_authenticate_user = AsyncMock(
|
||||
side_effect=ProxyException(
|
||||
message="Invalid credentials",
|
||||
type=ProxyErrorTypes.auth_error,
|
||||
param="password",
|
||||
code=401,
|
||||
)
|
||||
)
|
||||
|
||||
monkeypatch.setattr(
|
||||
"litellm.proxy.auth.login_utils.authenticate_user",
|
||||
mock_authenticate_user,
|
||||
)
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.master_key", "test-master-key")
|
||||
monkeypatch.setattr(
|
||||
"litellm.proxy.proxy_server.general_settings",
|
||||
{"control_plane_url": "https://cp.example.com"},
|
||||
)
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client)
|
||||
|
||||
client = TestClient(app)
|
||||
response = client.post(
|
||||
"/v3/login",
|
||||
json={"username": "alice", "password": "wrong"},
|
||||
)
|
||||
|
||||
assert response.status_code == 401
|
||||
assert response.headers["content-type"] == "application/json"
|
||||
data = response.json()
|
||||
assert "error" in data
|
||||
assert data["error"]["message"] == "Invalid credentials"
|
||||
assert data["error"]["type"] == "auth_error"
|
||||
|
||||
|
||||
def test_fallback_login_has_no_deprecation_banner(client_no_auth):
|
||||
response = client_no_auth.get("/fallback/login")
|
||||
|
||||
|
|
|
|||
|
|
@ -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 { exchangeLoginCode, getProxyBaseUrl, switchToWorkerUrl } from "@/components/networking";
|
||||
import { clearTokenCookies, getCookie } from "@/utils/cookieUtils";
|
||||
import { isJwtExpired } from "@/utils/jwtUtils";
|
||||
import { consumeReturnUrl, getReturnUrl, isValidReturnUrl } from "@/utils/returnUrlUtils";
|
||||
import { InfoCircleOutlined } from "@ant-design/icons";
|
||||
import { Alert, Button, Card, Form, Input, Popover, Space, Typography } from "antd";
|
||||
import { InfoCircleOutlined, CloudServerOutlined } from "@ant-design/icons";
|
||||
import { Alert, Button, Card, Form, Input, Popover, Select, Space, Typography } from "antd";
|
||||
import { useRouter } from "next/navigation";
|
||||
import { useEffect, useState } from "react";
|
||||
import { useWorker } from "@/hooks/useWorker";
|
||||
|
||||
function LoginPageContent() {
|
||||
const [username, setUsername] = useState("");
|
||||
|
|
@ -19,6 +20,17 @@ function LoginPageContent() {
|
|||
const { data: uiConfig, isLoading: isConfigLoading } = useUIConfig();
|
||||
const loginMutation = useLogin();
|
||||
const router = useRouter();
|
||||
const { workers, selectWorker } = useWorker();
|
||||
const [selectedWorkerId, setSelectedWorkerId] = useState<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,44 @@ function LoginPageContent() {
|
|||
return;
|
||||
}
|
||||
|
||||
// Cross-origin SSO: worker redirected back with a single-use code.
|
||||
// Exchange it for the JWT via the worker's /v3/login/exchange endpoint.
|
||||
const params = new URLSearchParams(window.location.search);
|
||||
const ssoCode = params.get("code");
|
||||
if (ssoCode) {
|
||||
const workerUrl = localStorage.getItem("litellm_worker_url");
|
||||
exchangeLoginCode(ssoCode, workerUrl).then(() => {
|
||||
params.delete("code");
|
||||
const cleanSearch = params.toString();
|
||||
window.history.replaceState(null, "", window.location.pathname + (cleanSearch ? `?${cleanSearch}` : ""));
|
||||
router.replace("/ui/?login=success");
|
||||
});
|
||||
return;
|
||||
}
|
||||
|
||||
// Backwards compat: handle direct token in URL (legacy flow)
|
||||
const urlToken = params.get("token");
|
||||
if (urlToken && !isJwtExpired(urlToken)) {
|
||||
document.cookie = `token=${urlToken}; path=/; SameSite=Lax`;
|
||||
params.delete("token");
|
||||
const cleanSearch = params.toString();
|
||||
window.history.replaceState(
|
||||
null,
|
||||
"",
|
||||
window.location.pathname + (cleanSearch ? `?${cleanSearch}` : ""),
|
||||
);
|
||||
router.replace("/ui/?login=success");
|
||||
return;
|
||||
}
|
||||
|
||||
// If switching workers on a control plane, clear the old token and show login
|
||||
const switchingWorker = params.has("worker");
|
||||
if (switchingWorker && uiConfig?.is_control_plane) {
|
||||
clearTokenCookies();
|
||||
setIsLoading(false);
|
||||
return;
|
||||
}
|
||||
|
||||
const rawToken = getCookie("token");
|
||||
if (rawToken && !isJwtExpired(rawToken)) {
|
||||
// User already logged in - redirect to return URL or default
|
||||
|
|
@ -38,7 +88,7 @@ function LoginPageContent() {
|
|||
if (returnUrl) {
|
||||
router.replace(returnUrl);
|
||||
} else {
|
||||
router.replace(`${getProxyBaseUrl()}/ui`);
|
||||
router.replace("/ui");
|
||||
}
|
||||
return;
|
||||
}
|
||||
|
|
@ -58,16 +108,35 @@ function LoginPageContent() {
|
|||
}, [isConfigLoading, router, uiConfig]);
|
||||
|
||||
const handleSubmit = () => {
|
||||
// If a worker is selected, point proxyBaseUrl at it before login
|
||||
const selectedWorker = workers.find((w) => w.worker_id === selectedWorkerId);
|
||||
if (selectedWorker) {
|
||||
switchToWorkerUrl(selectedWorker.url);
|
||||
}
|
||||
|
||||
loginMutation.mutate(
|
||||
{ username, password },
|
||||
{ username, password, useV3: !!selectedWorker },
|
||||
{
|
||||
onSuccess: (data) => {
|
||||
// Check if we have a return URL to use instead of the default redirect
|
||||
const returnUrl = consumeReturnUrl();
|
||||
if (returnUrl) {
|
||||
router.push(returnUrl);
|
||||
// Update the worker context with the selected worker
|
||||
if (selectedWorker) {
|
||||
selectWorker(selectedWorker.worker_id);
|
||||
// Stay on the CP's UI — proxyBaseUrl already points at the worker
|
||||
router.push("/ui/?login=success");
|
||||
} else {
|
||||
router.push(data.redirect_url);
|
||||
// Normal (non-control-plane) login — follow the server's redirect
|
||||
const returnUrl = consumeReturnUrl();
|
||||
if (returnUrl) {
|
||||
router.push(returnUrl);
|
||||
} else {
|
||||
router.push(data.redirect_url);
|
||||
}
|
||||
}
|
||||
},
|
||||
onError: () => {
|
||||
// Reset proxyBaseUrl on login failure
|
||||
if (selectedWorker) {
|
||||
switchToWorkerUrl(null);
|
||||
}
|
||||
},
|
||||
},
|
||||
|
|
@ -154,6 +223,22 @@ function LoginPageContent() {
|
|||
{error && <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 +294,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,19 @@ const Navbar: React.FC<NavbarProps> = ({
|
|||
|
||||
const handleLogout = () => {
|
||||
clearTokenCookies();
|
||||
localStorage.removeItem("litellm_selected_worker_id");
|
||||
localStorage.removeItem("litellm_worker_url");
|
||||
window.location.href = logoutUrl;
|
||||
};
|
||||
|
||||
const handleWorkerSwitch = (workerId: string) => {
|
||||
clearTokenCookies();
|
||||
clearStoredReturnUrl();
|
||||
localStorage.removeItem("litellm_selected_worker_id");
|
||||
localStorage.removeItem("litellm_worker_url");
|
||||
window.location.href = `/ui/login?worker=${encodeURIComponent(workerId)}`;
|
||||
};
|
||||
|
||||
return (
|
||||
<nav className="bg-white border-b border-gray-200 sticky top-0 z-10">
|
||||
<div className="w-full">
|
||||
|
|
@ -169,6 +181,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 {
|
||||
|
|
@ -9030,15 +9088,20 @@ export const deriveErrorMessage = (errorData: any): string => {
|
|||
export interface LoginRequest {
|
||||
username: string;
|
||||
password: string;
|
||||
useV3?: boolean;
|
||||
}
|
||||
|
||||
interface LoginResponse {
|
||||
redirect_url: string;
|
||||
token?: string;
|
||||
code?: string;
|
||||
expires_in?: number;
|
||||
}
|
||||
|
||||
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,
|
||||
|
|
@ -9060,10 +9123,65 @@ 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();
|
||||
|
||||
// v3 returns an opaque code — exchange it for the real JWT
|
||||
if (useV3 && data.code) {
|
||||
const exchangeUrl = proxyBaseUrl
|
||||
? `${proxyBaseUrl}/v3/login/exchange`
|
||||
: "/v3/login/exchange";
|
||||
|
||||
const exchangeResponse = await fetch(exchangeUrl, {
|
||||
method: "POST",
|
||||
body: JSON.stringify({ code: data.code }),
|
||||
credentials: "include",
|
||||
headers: { "Content-Type": "application/json" },
|
||||
});
|
||||
|
||||
if (!exchangeResponse.ok) {
|
||||
const errorData = await exchangeResponse.json();
|
||||
throw new Error(deriveErrorMessage(errorData));
|
||||
}
|
||||
|
||||
const exchangeData: LoginResponse = await exchangeResponse.json();
|
||||
if (exchangeData.token) {
|
||||
document.cookie = `token=${exchangeData.token}; path=/; SameSite=Lax`;
|
||||
}
|
||||
return exchangeData;
|
||||
}
|
||||
|
||||
// Backwards compatibility: v2 or old v3 returns token directly
|
||||
if (data.token) {
|
||||
document.cookie = `token=${data.token}; path=/; SameSite=Lax`;
|
||||
}
|
||||
|
||||
return data;
|
||||
};
|
||||
|
||||
/**
|
||||
* Exchange a single-use login code for a JWT token.
|
||||
* Used by the SSO callback when the worker redirects back with ?code=.
|
||||
*/
|
||||
export const exchangeLoginCode = async (code: string, workerBaseUrl?: string | null): Promise<string> => {
|
||||
const base = workerBaseUrl || getProxyBaseUrl();
|
||||
const response = await fetch(`${base}/v3/login/exchange`, {
|
||||
method: "POST",
|
||||
body: JSON.stringify({ code }),
|
||||
headers: { "Content-Type": "application/json" },
|
||||
});
|
||||
|
||||
if (!response.ok) {
|
||||
const errorData = await response.json();
|
||||
throw new Error(deriveErrorMessage(errorData));
|
||||
}
|
||||
|
||||
const data = await response.json();
|
||||
if (data.token) {
|
||||
document.cookie = `token=${data.token}; path=/; SameSite=Lax`;
|
||||
}
|
||||
return data.token;
|
||||
};
|
||||
|
||||
export const getUiSettings = async () => {
|
||||
const proxyBaseUrl = getProxyBaseUrl();
|
||||
const url = proxyBaseUrl ? `${proxyBaseUrl}/get/ui_settings` : `/get/ui_settings`;
|
||||
|
|
|
|||
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,
|
||||
};
|
||||
};
|
||||
Loading…
Add table
Reference in a new issue