diff --git a/litellm/proxy/discovery_endpoints/ui_discovery_endpoints.py b/litellm/proxy/discovery_endpoints/ui_discovery_endpoints.py index 9aaa2fb838..cbe28849b1 100644 --- a/litellm/proxy/discovery_endpoints/ui_discovery_endpoints.py +++ b/litellm/proxy/discovery_endpoints/ui_discovery_endpoints.py @@ -18,9 +18,11 @@ async def get_ui_config(): from litellm.proxy.auth.auth_utils import _has_user_setup_sso auto_redirect_ui_login_to_sso = os.getenv("AUTO_REDIRECT_UI_LOGIN_TO_SSO", "true").lower() == "true" + admin_ui_disabled = os.getenv("DISABLE_ADMIN_UI", "false").lower() == "true" return UiDiscoveryEndpoints( server_root_path=get_server_root_path(), proxy_base_url=get_proxy_base_url(), auto_redirect_to_sso=_has_user_setup_sso() and auto_redirect_ui_login_to_sso, + admin_ui_disabled=admin_ui_disabled, ) diff --git a/litellm/types/proxy/discovery_endpoints/ui_discovery_endpoints.py b/litellm/types/proxy/discovery_endpoints/ui_discovery_endpoints.py index f100dd35fa..dc167667bc 100644 --- a/litellm/types/proxy/discovery_endpoints/ui_discovery_endpoints.py +++ b/litellm/types/proxy/discovery_endpoints/ui_discovery_endpoints.py @@ -7,3 +7,4 @@ class UiDiscoveryEndpoints(BaseModel): server_root_path: str proxy_base_url: Optional[str] auto_redirect_to_sso: bool + admin_ui_disabled: bool diff --git a/tests/test_litellm/proxy/discovery_endpoints/test_ui_discovery_endpoints.py b/tests/test_litellm/proxy/discovery_endpoints/test_ui_discovery_endpoints.py index 599d543758..1d31d4f7bf 100644 --- a/tests/test_litellm/proxy/discovery_endpoints/test_ui_discovery_endpoints.py +++ b/tests/test_litellm/proxy/discovery_endpoints/test_ui_discovery_endpoints.py @@ -30,6 +30,7 @@ def test_ui_discovery_endpoints_with_defaults(): assert data["server_root_path"] == "/" assert data["proxy_base_url"] is None assert data["auto_redirect_to_sso"] is False + assert data["admin_ui_disabled"] is False def test_ui_discovery_endpoints_with_custom_server_root_path(): @@ -144,3 +145,43 @@ def test_ui_discovery_endpoints_both_routes_return_same_data(): assert response2.status_code == 200 assert response1.json() == response2.json() + +def test_ui_discovery_endpoints_with_admin_ui_disabled(): + app = FastAPI() + app.include_router(router) + client = TestClient(app) + + 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.dict(os.environ, {"DISABLE_ADMIN_UI": "true"}, clear=False): + + response = client.get("/.well-known/litellm-ui-config") + + assert response.status_code == 200 + data = response.json() + assert data["server_root_path"] == "/" + assert data["proxy_base_url"] is None + assert data["auto_redirect_to_sso"] is False + assert data["admin_ui_disabled"] is True + + +def test_ui_discovery_endpoints_with_admin_ui_enabled(): + app = FastAPI() + app.include_router(router) + client = TestClient(app) + + 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.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"] == "/" + assert data["proxy_base_url"] is None + assert data["auto_redirect_to_sso"] is False + assert data["admin_ui_disabled"] is False + diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/useAuthorized.test.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/useAuthorized.test.ts index 9198450a63..2668461937 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/hooks/useAuthorized.test.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/useAuthorized.test.ts @@ -1,12 +1,15 @@ /* @vitest-environment jsdom */ -import { renderHook } from "@testing-library/react"; +import React from "react"; +import { renderHook, waitFor } from "@testing-library/react"; import { afterEach, describe, expect, it, vi } from "vitest"; +import { QueryClient, QueryClientProvider } from "@tanstack/react-query"; import useAuthorized from "./useAuthorized"; -const { replaceMock, clearTokenCookiesMock, getProxyBaseUrlMock } = vi.hoisted(() => ({ +const { replaceMock, clearTokenCookiesMock, getProxyBaseUrlMock, getUiConfigMock } = vi.hoisted(() => ({ replaceMock: vi.fn(), clearTokenCookiesMock: vi.fn(), getProxyBaseUrlMock: vi.fn(() => "http://proxy.example"), + getUiConfigMock: vi.fn(), })); vi.mock("next/navigation", () => ({ @@ -15,9 +18,14 @@ vi.mock("next/navigation", () => ({ }), })); -vi.mock("@/components/networking", () => ({ - getProxyBaseUrl: getProxyBaseUrlMock, -})); +vi.mock("@/components/networking", async (importOriginal) => { + const actual = await importOriginal(); + return { + ...actual, + getProxyBaseUrl: getProxyBaseUrlMock, + getUiConfig: getUiConfigMock, + }; +}); vi.mock("@/utils/cookieUtils", async (importOriginal) => { const actual = await importOriginal(); @@ -27,6 +35,21 @@ vi.mock("@/utils/cookieUtils", async (importOriginal) => { }; }); +const createQueryClient = () => + new QueryClient({ + defaultOptions: { + queries: { + retry: false, + gcTime: 0, + }, + }, + }); + +const wrapper = ({ children }: { children: React.ReactNode }) => { + const queryClient = createQueryClient(); + return React.createElement(QueryClientProvider, { client: queryClient }, children); +}; + const createJwt = (payload: Record) => { const base64Url = btoa(JSON.stringify(payload)).replace(/=+$/, "").replace(/\+/g, "-").replace(/\//g, "_"); return `eyJhbGciOiJub25lIn0.${base64Url}.signature`; @@ -41,10 +64,18 @@ describe("useAuthorized", () => { replaceMock.mockReset(); clearTokenCookiesMock.mockReset(); getProxyBaseUrlMock.mockClear(); + getUiConfigMock.mockReset(); clearCookie(); }); - it("should decode the token and expose user details", () => { + it("should decode the token and expose user details", async () => { + getUiConfigMock.mockResolvedValue({ + server_root_path: "/", + proxy_base_url: null, + auto_redirect_to_sso: false, + admin_ui_disabled: false, + }); + const token = createJwt({ key: "api-key-123", user_id: "user-1", @@ -56,9 +87,12 @@ describe("useAuthorized", () => { }); document.cookie = `token=${token}; path=/;`; - const { result } = renderHook(() => useAuthorized()); + const { result } = renderHook(() => useAuthorized(), { wrapper }); + + await waitFor(() => { + expect(result.current.token).toBe(token); + }); - expect(result.current.token).toBe(token); expect(result.current.accessToken).toBe("api-key-123"); expect(result.current.userId).toBe("user-1"); expect(result.current.userEmail).toBe("user@example.com"); @@ -69,14 +103,54 @@ describe("useAuthorized", () => { expect(replaceMock).not.toHaveBeenCalled(); }); - it("should clear cookies and redirect on an invalid token", () => { + it("should clear cookies and redirect on an invalid token", async () => { + getUiConfigMock.mockResolvedValue({ + server_root_path: "/", + proxy_base_url: null, + auto_redirect_to_sso: false, + admin_ui_disabled: false, + }); + document.cookie = "token=invalid-token; path=/;"; - const { result } = renderHook(() => useAuthorized()); + const { result } = renderHook(() => useAuthorized(), { wrapper }); + + await waitFor(() => { + expect(clearTokenCookiesMock).toHaveBeenCalled(); + }); - expect(clearTokenCookiesMock).toHaveBeenCalled(); expect(replaceMock).toHaveBeenCalledWith("http://proxy.example/ui/login"); expect(result.current.accessToken).toBeNull(); expect(result.current.userRole).toBe("Undefined Role"); }); + + it("should redirect even with valid token if admin_ui_disabled is true", async () => { + getUiConfigMock.mockResolvedValue({ + server_root_path: "/", + proxy_base_url: null, + auto_redirect_to_sso: false, + admin_ui_disabled: true, + }); + + const token = createJwt({ + key: "api-key-123", + user_id: "user-1", + user_email: "user@example.com", + user_role: "app_admin", + premium_user: true, + disabled_non_admin_personal_key_creation: false, + login_method: "username_password", + }); + document.cookie = `token=${token}; path=/;`; + + const { result } = renderHook(() => useAuthorized(), { wrapper }); + + await waitFor(() => { + expect(replaceMock).toHaveBeenCalledWith("http://proxy.example/ui/login"); + }); + + expect(result.current.accessToken).toBe("api-key-123"); + expect(result.current.userId).toBe("user-1"); + expect(result.current.userEmail).toBe("user@example.com"); + }); }); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/useAuthorized.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/useAuthorized.ts index 7610c6346b..62d514f066 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/hooks/useAuthorized.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/useAuthorized.ts @@ -1,10 +1,11 @@ "use client"; -import { useEffect, useMemo } from "react"; -import { useRouter } from "next/navigation"; -import { jwtDecode } from "jwt-decode"; -import { clearTokenCookies, getCookie } from "@/utils/cookieUtils"; import { getProxyBaseUrl } from "@/components/networking"; +import { clearTokenCookies, getCookie } from "@/utils/cookieUtils"; +import { jwtDecode } from "jwt-decode"; +import { useRouter } from "next/navigation"; +import { useEffect, useMemo } from "react"; +import { useUIConfig } from "./uiConfig/useUIConfig"; function formatUserRole(userRole: string) { if (!userRole) { @@ -37,15 +38,19 @@ function formatUserRole(userRole: string) { const useAuthorized = () => { const router = useRouter(); + const { data: uiConfig, isLoading: isUIConfigLoading } = useUIConfig(); const token = typeof document !== "undefined" ? getCookie("token") : null; // Redirect after mount if missing/invalid token useEffect(() => { - if (!token) { + if (isUIConfigLoading) { + return; + } + if (!token || uiConfig?.admin_ui_disabled) { router.replace(`${getProxyBaseUrl()}/ui/login`); } - }, [token, router]); + }, [token, router, isUIConfigLoading, uiConfig]); // Decode safely const decoded = useMemo(() => { diff --git a/ui/litellm-dashboard/src/app/login/LoginPage.test.tsx b/ui/litellm-dashboard/src/app/login/LoginPage.test.tsx index cce063eceb..7983451260 100644 --- a/ui/litellm-dashboard/src/app/login/LoginPage.test.tsx +++ b/ui/litellm-dashboard/src/app/login/LoginPage.test.tsx @@ -169,4 +169,27 @@ describe("LoginPage", () => { expect(mockPush).not.toHaveBeenCalled(); }); + + it("should show alert when admin_ui_disabled is true", async () => { + (useUIConfig as ReturnType).mockReturnValue({ + data: { admin_ui_disabled: true, server_root_path: "/", proxy_base_url: null }, + isLoading: false, + }); + (getCookie as ReturnType).mockReturnValue(null); + + const queryClient = createQueryClient(); + render( + + + , + ); + + await waitFor(() => { + expect(screen.getByRole("alert")).toBeInTheDocument(); + expect(screen.getByText("Admin UI Disabled")).toBeInTheDocument(); + }); + + expect(mockPush).not.toHaveBeenCalled(); + expect(mockReplace).not.toHaveBeenCalled(); + }); }); diff --git a/ui/litellm-dashboard/src/app/login/LoginPage.tsx b/ui/litellm-dashboard/src/app/login/LoginPage.tsx index 85f2c6dd87..620cb41dfe 100644 --- a/ui/litellm-dashboard/src/app/login/LoginPage.tsx +++ b/ui/litellm-dashboard/src/app/login/LoginPage.tsx @@ -25,6 +25,12 @@ function LoginPageContent() { return; } + // Check if admin UI is disabled + if (uiConfig && uiConfig.admin_ui_disabled) { + setIsLoading(false); + return; + } + const rawToken = getCookie("token"); if (rawToken && !isJwtExpired(rawToken)) { router.replace(`${getProxyBaseUrl()}/ui`); @@ -59,6 +65,38 @@ function LoginPageContent() { return ; } + // Show disabled message if admin UI is disabled + if (uiConfig && uiConfig.admin_ui_disabled) { + return ( +
+ + +
+ 🚅 LiteLLM +
+ + + + The Admin UI has been disabled by the administrator. To re-enable it, please update the following + environment variable: + + + DISABLE_ADMIN_UI=False + + + } + type="warning" + showIcon + /> +
+
+
+ ); + } + return (
diff --git a/ui/litellm-dashboard/src/components/networking.tsx b/ui/litellm-dashboard/src/components/networking.tsx index 21f1816814..aeba207db2 100644 --- a/ui/litellm-dashboard/src/components/networking.tsx +++ b/ui/litellm-dashboard/src/components/networking.tsx @@ -230,6 +230,7 @@ export interface LiteLLMWellKnownUiConfig { server_root_path: string; proxy_base_url: string | null; auto_redirect_to_sso: boolean; + admin_ui_disabled: boolean; } export interface CredentialsResponse {