From 4bafebbcdcdc515eca916dcc9082cbb53f2f3883 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Sat, 17 Oct 2026 14:39:33 -0700 Subject: [PATCH] fix(ui): take the dashboard's user role and identity from the server (#47025) * fix(ui): take the dashboard's user role and identity from the server Adds an internal GET /management/v1/users/me that returns the user id, role and email the proxy resolved for the calling key. The dashboard now reads identity from it instead of the session cookie's claims, so its navigation and role gating always match what the backend authorizes. The layout waits on it behind the existing loading screen, and a failed lookup clears the session and returns to login. Sessions that must change their password can also read it * fix(ui): sign out when the current user lookup fails and wait for it on chat The dashboard layout now holds the loading screen and signs out when /management/v1/users/me fails, instead of mounting the shell and looping back to loading as each consumer refetched the failed query. The chat layout also waits for the identity so its panels never query without a user id. The lookup goes through the typed client with the regenerated schema.d.ts * fix(ui): send the session's auth header on the first user lookup and keep chat closed until authorized The first lookup runs before AuthProvider registers the custom header name, so it now reads auth_header_name from the session cookie. Chat also waits for isAuthorized, so a failed lookup no longer mounts the shell with an empty user id * fix(ui): type the dashboard callers of the now typed user id useAuthorized used to hand back the cookie's user id untyped, which hid six strict type errors in its callers. The id is now string | null, so the production build failed. Callers handle the null case, and query key filters accept the boolean blocked filter * fix(ui): look up the current user only after the proxy URL is discovered Pages outside the dashboard shell, like chat and connect, could start the lookup before UI config discovery set the proxy base URL. On a split-origin deployment both attempts then hit the UI origin and a valid session was signed out * fix(ui): keep a loaded identity while UI config discovery reruns Pausing the lookup by passing a null token switched the query key, so on the dashboard the identity the layout had already loaded vanished until the shell's own UI config request finished. The key now stays the same and only the fetch waits for discovery --- litellm/proxy/_types.py | 1 + litellm/proxy/auth/login_utils.py | 2 +- litellm/proxy/management/users/endpoints.py | 20 ++ litellm/proxy/management/users/exceptions.py | 13 + litellm/proxy/management/users/schemas.py | 11 + .../management_v1/__init__.py | 2 + tests/unit/proxy/auth/test_login_utils.py | 4 +- tests/unit/proxy/auth/test_route_checks.py | 25 +- .../proxy/management/users/test_endpoints.py | 67 +++++ .../hooks/common/queryKeysFactory.ts | 2 +- .../hooks/currentUser/useCurrentUser.ts | 28 ++ .../app/(dashboard)/hooks/redirectToLogin.ts | 8 + .../(dashboard)/hooks/useAuthorized.test.ts | 273 +++++++++++++++++- .../app/(dashboard)/hooks/useAuthorized.ts | 36 ++- .../(dashboard)/layout.integration.test.tsx | 2 + .../src/app/(dashboard)/layout.test.tsx | 117 ++++++++ .../src/app/(dashboard)/layout.tsx | 20 +- .../components/AllModelsTab.tsx | 2 +- .../policies/_components/add_policy_form.tsx | 2 +- .../src/app/chat/layout.test.tsx | 30 ++ ui/litellm-dashboard/src/app/chat/layout.tsx | 14 +- .../DeploymentHealth.integration.test.tsx | 6 + .../src/components/add_model/AddModelForm.tsx | 2 +- .../liteadmin/LiteAdmin.integration.test.tsx | 9 +- .../organisms/create_key_button.tsx | 4 +- ui/litellm-dashboard/src/lib/http/schema.d.ts | 45 +++ 26 files changed, 702 insertions(+), 43 deletions(-) create mode 100644 litellm/proxy/management/users/endpoints.py create mode 100644 litellm/proxy/management/users/exceptions.py create mode 100644 litellm/proxy/management/users/schemas.py create mode 100644 tests/unit/proxy/management/users/test_endpoints.py create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/hooks/currentUser/useCurrentUser.ts create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/hooks/redirectToLogin.ts diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index dc7383569dc..a7bd16804c1 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -563,6 +563,7 @@ class LiteLLMRoutes(enum.Enum): "/team/metadata_schema", "/user/info", "/v2/user/info", + "/management/v1/users/me", "/model/info", "/v1/model/info", "/v2/model/info", diff --git a/litellm/proxy/auth/login_utils.py b/litellm/proxy/auth/login_utils.py index 629b31024e2..9783f44b41f 100644 --- a/litellm/proxy/auth/login_utils.py +++ b/litellm/proxy/auth/login_utils.py @@ -57,7 +57,7 @@ if TYPE_CHECKING: from prisma import types as prisma_types BREACH_RECHECK_INTERVAL: Final = timedelta(hours=24) -PASSWORD_RESET_ALLOWED_ROUTES: Final = ("/user/password/change", "/session/logout") +PASSWORD_RESET_ALLOWED_ROUTES: Final = ("/user/password/change", "/session/logout", "/management/v1/users/me") PASSWORD_SESSION_METADATA: Final = MappingProxyType({"login_method": "username_password"}) diff --git a/litellm/proxy/management/users/endpoints.py b/litellm/proxy/management/users/endpoints.py new file mode 100644 index 00000000000..62dc3a2a518 --- /dev/null +++ b/litellm/proxy/management/users/endpoints.py @@ -0,0 +1,20 @@ +from typing import Annotated, Final + +from fastapi import APIRouter, Depends + +from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.auth.user_api_key_auth import user_api_key_auth +from litellm.proxy.management.users.exceptions import current_user_not_found +from litellm.proxy.management.users.schemas import CurrentUserResponse +from litellm.proxy.management_endpoints.management_v1.common import MANAGEMENT_V1_PREFIX + +router: Final = APIRouter(prefix=MANAGEMENT_V1_PREFIX) + + +@router.get("/users/me", include_in_schema=False) +async def get_current_user( + caller: Annotated[UserAPIKeyAuth, Depends(user_api_key_auth)], +) -> CurrentUserResponse: + if caller.user_id is None: + raise current_user_not_found() + return CurrentUserResponse(user_id=caller.user_id, user_role=caller.user_role, user_email=caller.user_email) diff --git a/litellm/proxy/management/users/exceptions.py b/litellm/proxy/management/users/exceptions.py new file mode 100644 index 00000000000..2c1e58c9283 --- /dev/null +++ b/litellm/proxy/management/users/exceptions.py @@ -0,0 +1,13 @@ +from litellm.proxy.list_api.common import PROBLEM_TYPE_BASE, ManagementProblem +from litellm.types.proxy.management_endpoints.management_v1 import ProblemDetail + + +def current_user_not_found() -> ManagementProblem: + return ManagementProblem( + ProblemDetail( + type=f"{PROBLEM_TYPE_BASE}user-not-found", + title="User not found", + status=404, + detail="The key used for this request is not bound to a user.", + ) + ) diff --git a/litellm/proxy/management/users/schemas.py b/litellm/proxy/management/users/schemas.py new file mode 100644 index 00000000000..243e970c3a1 --- /dev/null +++ b/litellm/proxy/management/users/schemas.py @@ -0,0 +1,11 @@ +from pydantic import BaseModel, ConfigDict + +from litellm.proxy._types import LitellmUserRoles + + +class CurrentUserResponse(BaseModel): + model_config = ConfigDict(frozen=True) + + user_id: str + user_role: LitellmUserRoles | None + user_email: str | None diff --git a/litellm/proxy/management_endpoints/management_v1/__init__.py b/litellm/proxy/management_endpoints/management_v1/__init__.py index b29b7fe5dd5..a02d989cc28 100644 --- a/litellm/proxy/management_endpoints/management_v1/__init__.py +++ b/litellm/proxy/management_endpoints/management_v1/__init__.py @@ -4,6 +4,7 @@ from typing import Final from fastapi import APIRouter +from litellm.proxy.management.users.endpoints import router as current_user_router from litellm.proxy.management_endpoints.management_v1.budgets import ( router as budgets_router, ) @@ -22,5 +23,6 @@ router.include_router(budgets_router) router.include_router(spend_logs_router) router.include_router(teams_router) router.include_router(users_router) +router.include_router(current_user_router) __all__ = ["router"] diff --git a/tests/unit/proxy/auth/test_login_utils.py b/tests/unit/proxy/auth/test_login_utils.py index 1fb648fe5e6..ec9b9a04f72 100644 --- a/tests/unit/proxy/auth/test_login_utils.py +++ b/tests/unit/proxy/auth/test_login_utils.py @@ -2144,7 +2144,7 @@ class TestPasswordResetRequiredSessionMinting: row = _db_user_row(password="Str0ng!Passw0rd", password_reset_required=True) result, key_kwargs = await self._login(_prisma_with_user(row)) - assert key_kwargs["allowed_routes"] == ["/user/password/change", "/session/logout"] + assert key_kwargs["allowed_routes"] == ["/user/password/change", "/session/logout", "/management/v1/users/me"] assert key_kwargs["metadata"] == {"login_method": "username_password", "password_reset_required": True} assert result.password_reset_required is True @@ -2205,7 +2205,7 @@ class TestPasswordResetRequiredSessionMinting: result, key_kwargs, _ = await self._login_with_screen_result(mock_prisma_client, breached=True) - assert key_kwargs["allowed_routes"] == ["/user/password/change", "/session/logout"] + assert key_kwargs["allowed_routes"] == ["/user/password/change", "/session/logout", "/management/v1/users/me"] assert key_kwargs["metadata"] == {"login_method": "username_password", "password_reset_required": True} assert result.password_reset_required is True diff --git a/tests/unit/proxy/auth/test_route_checks.py b/tests/unit/proxy/auth/test_route_checks.py index 7bdcb8afd37..a519a4e576a 100644 --- a/tests/unit/proxy/auth/test_route_checks.py +++ b/tests/unit/proxy/auth/test_route_checks.py @@ -17,6 +17,7 @@ from litellm.proxy._types import ( ) from litellm.proxy.auth.auth_checks import _is_api_route_allowed from litellm.proxy.auth.auth_checks_organization import user_is_org_admin +from litellm.proxy.auth.login_utils import PASSWORD_RESET_ALLOWED_ROUTES from litellm.proxy.auth.route_checks import RouteChecks from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import router as llm_passthrough_router @@ -3992,20 +3993,38 @@ def _password_reset_session_token() -> UserAPIKeyAuth: `password_reset_required`.""" return UserAPIKeyAuth( user_id="flagged_user", - allowed_routes=["/user/password/change"], + allowed_routes=list(PASSWORD_RESET_ALLOWED_ROUTES), metadata={"password_reset_required": True}, ) -def test_password_reset_session_can_reach_change_password(): +@pytest.mark.parametrize("route", ["/user/password/change", "/management/v1/users/me"]) +def test_password_reset_session_can_reach_change_password_and_its_own_identity(route): result = RouteChecks.is_virtual_key_allowed_to_call_route( - route="/user/password/change", + route=route, valid_token=_password_reset_session_token(), ) assert result is True +@pytest.mark.parametrize( + "role", + [LitellmUserRoles.INTERNAL_USER.value, LitellmUserRoles.INTERNAL_USER_VIEW_ONLY.value], +) +def test_non_admin_user_can_read_its_own_identity(role): + allowed = RouteChecks.non_proxy_admin_allowed_routes_check( + user_obj=LiteLLM_UserTable(user_id="test_user", user_role=role), + _user_role=role, + route="/management/v1/users/me", + request=MagicMock(spec=Request, query_params={}), + valid_token=UserAPIKeyAuth(user_id="test_user", user_role=role), + request_data={}, + ) + + assert allowed is None + + @pytest.mark.parametrize( "route", [ diff --git a/tests/unit/proxy/management/users/test_endpoints.py b/tests/unit/proxy/management/users/test_endpoints.py new file mode 100644 index 00000000000..7ec0e5b7f53 --- /dev/null +++ b/tests/unit/proxy/management/users/test_endpoints.py @@ -0,0 +1,67 @@ +from collections.abc import Iterator +from typing import Final + +import pytest +from fastapi import FastAPI, Request +from fastapi.testclient import TestClient + +from litellm.proxy._types import LitellmUserRoles +from litellm.proxy.auth.user_api_key_auth import UserAPIKeyAuth, user_api_key_auth +from litellm.proxy.list_api.common import ManagementProblem, problem_response +from litellm.proxy.management_endpoints.management_v1 import router + +app: Final = FastAPI() + + +@app.exception_handler(ManagementProblem) +async def management_problem_exception_handler(request: Request, exc: ManagementProblem): + return problem_response(exc.problem) + + +app.include_router(router) +client: Final = TestClient(app) + +USERS_ME_PATH: Final = "/management/v1/users/me" + + +@pytest.fixture +def caller() -> Iterator[None]: + yield + app.dependency_overrides.clear() + + +def authenticate_as(resolved: UserAPIKeyAuth) -> None: + app.dependency_overrides[user_api_key_auth] = lambda: resolved + + +@pytest.mark.parametrize( + "role", + [LitellmUserRoles.INTERNAL_USER, LitellmUserRoles.PROXY_ADMIN, LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY], +) +def test_returns_the_identity_auth_resolved_for_the_caller(caller, role): + authenticate_as(UserAPIKeyAuth(user_id="u-1", user_role=role, user_email="u1@example.com")) + + response = client.get(USERS_ME_PATH, headers={"Authorization": "Bearer k"}) + + assert response.status_code == 200 + assert response.json() == {"user_id": "u-1", "user_role": role.value, "user_email": "u1@example.com"} + + +def test_key_without_a_user_gets_a_404_problem_document(caller): + authenticate_as(UserAPIKeyAuth(team_id="team-1", user_role=LitellmUserRoles.INTERNAL_USER)) + + response = client.get(USERS_ME_PATH, headers={"Authorization": "Bearer k"}) + + assert response.status_code == 404 + assert response.headers["content-type"] == "application/problem+json" + assert response.json()["type"] == "urn:litellm:error:user-not-found" + assert response.json()["status"] == 404 + + +def test_route_is_served_but_hidden_from_the_openapi_schema(caller): + authenticate_as(UserAPIKeyAuth(user_id="u-1", user_role=LitellmUserRoles.INTERNAL_USER)) + paths: Final = app.openapi()["paths"] + + assert client.get(USERS_ME_PATH).status_code == 200 + assert "/management/v1/users/bulk" in paths + assert USERS_ME_PATH not in paths diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/common/queryKeysFactory.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/common/queryKeysFactory.ts index eceaae9d422..a7367069014 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/hooks/common/queryKeysFactory.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/common/queryKeysFactory.ts @@ -3,7 +3,7 @@ type ListParams = { page?: number; limit?: number; - filters?: Record; + filters?: Record; }; /** diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/currentUser/useCurrentUser.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/currentUser/useCurrentUser.ts new file mode 100644 index 00000000000..6b688ec2668 --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/currentUser/useCurrentUser.ts @@ -0,0 +1,28 @@ +import { skipToken, useQuery, type UseQueryOptions } from "@tanstack/react-query"; +import { fetchClient } from "@/lib/http/api"; +import { getAuthHeaderName } from "@/lib/http/runtime"; +import type { components } from "@/lib/http/schema"; +import { createQueryKeys } from "../common/queryKeysFactory"; + +export type CurrentUser = components["schemas"]["CurrentUserResponse"]; + +const currentUserKeys = createQueryKeys("currentUser"); + +const getCurrentUser = async (accessToken: string, authHeaderName: string): Promise => { + const { data } = await fetchClient.GET("/management/v1/users/me", { + headers: { [authHeaderName]: `Bearer ${accessToken}` }, + }); + if (!data) throw new Error("Current user response is empty"); + return data; +}; + +export const useCurrentUser = (accessToken: string | null, authHeaderName?: string, canFetch = true) => { + const options: UseQueryOptions = { + queryKey: currentUserKeys.detail(accessToken ?? ""), + queryFn: + accessToken && canFetch ? () => getCurrentUser(accessToken, authHeaderName ?? getAuthHeaderName()) : skipToken, + staleTime: Infinity, + retry: 1, + }; + return useQuery(options); +}; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/redirectToLogin.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/redirectToLogin.ts new file mode 100644 index 00000000000..ec53589a8b2 --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/redirectToLogin.ts @@ -0,0 +1,8 @@ +import { getProxyBaseUrl, getUiPageUrl } from "@/components/networking"; +import { buildLoginUrlWithReturn, getLoginUrl, storeReturnUrl } from "@/utils/returnUrlUtils"; + +export const redirectToLogin = () => { + storeReturnUrl(); + const baseLoginUrl = getUiPageUrl(getLoginUrl(getProxyBaseUrl())); + window.location.replace(buildLoginUrlWithReturn(baseLoginUrl)); +}; 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 1ca0581a391..5d7825b310a 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/hooks/useAuthorized.test.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/useAuthorized.test.ts @@ -1,9 +1,11 @@ /* @vitest-environment jsdom */ import React from "react"; -import { renderHook, waitFor } from "@testing-library/react"; +import { act, renderHook, waitFor } from "@testing-library/react"; import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; import { QueryClient, QueryClientProvider } from "@tanstack/react-query"; import useAuthorized from "./useAuthorized"; +import { useUIConfig } from "./uiConfig/useUIConfig"; +import { useCurrentUser } from "./currentUser/useCurrentUser"; // Unmock useAuthorized to test the actual implementation vi.unmock("@/app/(dashboard)/hooks/useAuthorized"); @@ -60,6 +62,32 @@ vi.mock("@/utils/returnUrlUtils", async (importOriginal) => { storeReturnUrl: vi.fn(), }; }); +const fetchMock = vi.fn<(request: Request) => Promise>(); + +const CURRENT_USER_PATH = "/management/v1/users/me"; + +const currentUserRequests = () => + fetchMock.mock.calls + .map(([request]) => request) + .filter((request) => new URL(request.url).pathname === CURRENT_USER_PATH); + +type CurrentUserBody = { user_id: string; user_email: string | null; user_role: string | null }; + +const serveCurrentUser = (reply: CurrentUserBody | ((request: Request) => CurrentUserBody)) => { + fetchMock.mockImplementation(async (request) => Response.json(typeof reply === "function" ? reply(request) : reply)); + vi.stubGlobal("fetch", fetchMock); +}; + +const stallCurrentUser = () => { + fetchMock.mockImplementation(() => new Promise(() => {})); + vi.stubGlobal("fetch", fetchMock); +}; + +const failCurrentUser = (status: number) => { + fetchMock.mockImplementation(async () => Response.json({ detail: "rejected" }, { status })); + vi.stubGlobal("fetch", fetchMock); +}; + const createQueryClient = () => new QueryClient({ defaultOptions: { @@ -114,10 +142,12 @@ describe("useAuthorized", () => { decodeTokenMock.mockReset(); checkTokenValidityMock.mockReset(); buildLoginUrlWithReturnMock.mockClear(); + vi.unstubAllGlobals(); + fetchMock.mockReset(); clearCookie(); }); - it("should decode the token and expose user details", async () => { + it("should take user details from the server and the rest from the token", async () => { getUiConfigMock.mockResolvedValue({ server_root_path: "/", proxy_base_url: null, @@ -138,6 +168,7 @@ describe("useAuthorized", () => { decodeTokenMock.mockReturnValue(decodedPayload); checkTokenValidityMock.mockReturnValue(true); + serveCurrentUser({ user_id: "user-1", user_email: "user@example.com", user_role: "proxy_admin" }); const token = createJwt(decodedPayload); document.cookie = `token=${token}; path=/;`; @@ -145,9 +176,13 @@ describe("useAuthorized", () => { const { result } = renderHook(() => useAuthorized(), { wrapper }); await waitFor(() => { - expect(result.current.token).toBe(token); + expect(result.current.userId).toBe("user-1"); }); + expect(currentUserRequests().map((request) => request.headers.get("Authorization"))).toEqual([ + "Bearer api-key-123", + ]); + 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"); @@ -161,6 +196,39 @@ describe("useAuthorized", () => { expect(clearTokenCookiesMock).not.toHaveBeenCalled(); }); + it("sends the first lookup on the auth header named in the session cookie", async () => { + const uiConfig = { + server_root_path: "/", + proxy_base_url: null, + auto_redirect_to_sso: false, + admin_ui_disabled: false, + sso_configured: false, + }; + getUiConfigMock.mockResolvedValue(uiConfig); + + const decodedPayload = { + key: "api-key-custom-header", + user_id: "user-3", + user_role: "internal_user", + auth_header_name: "X-Litellm-Key", + }; + + decodeTokenMock.mockReturnValue(decodedPayload); + checkTokenValidityMock.mockReturnValue(true); + serveCurrentUser({ user_id: "user-3", user_email: null, user_role: "internal_user" }); + document.cookie = `token=${createJwt(decodedPayload)}; path=/;`; + + const { result } = renderHook(() => useAuthorized(), { wrapper }); + + await waitFor(() => { + expect(result.current.userId).toBe("user-3"); + }); + + expect(currentUserRequests().map((request) => request.headers.get("X-Litellm-Key"))).toEqual([ + "Bearer api-key-custom-header", + ]); + }); + it("should present proxy_admin_viewer as Admin while flagging it view-only", async () => { getUiConfigMock.mockResolvedValue({ server_root_path: "/", @@ -182,6 +250,11 @@ describe("useAuthorized", () => { decodeTokenMock.mockReturnValue(decodedPayload); checkTokenValidityMock.mockReturnValue(true); + serveCurrentUser({ + user_id: "user-2", + user_email: "viewer@example.com", + user_role: "proxy_admin_viewer", + }); const token = createJwt(decodedPayload); document.cookie = `token=${token}; path=/;`; @@ -189,7 +262,7 @@ describe("useAuthorized", () => { const { result } = renderHook(() => useAuthorized(), { wrapper }); await waitFor(() => { - expect(result.current.token).toBe(token); + expect(result.current.userId).toBe("user-2"); }); expect(result.current.userRole).toBe("Admin"); @@ -222,6 +295,7 @@ describe("useAuthorized", () => { expect(replaceMock).toHaveBeenCalledWith("http://proxy.example/ui/login/"); expect(result.current.accessToken).toBeNull(); expect(result.current.userRole).toBe("Undefined Role"); + expect(currentUserRequests()).toHaveLength(0); }); it("should redirect even with valid token if admin_ui_disabled is true", async () => { @@ -245,6 +319,7 @@ describe("useAuthorized", () => { decodeTokenMock.mockReturnValue(decodedPayload); checkTokenValidityMock.mockReturnValue(true); + serveCurrentUser({ user_id: "user-1", user_email: "user@example.com", user_role: "proxy_admin" }); const token = createJwt(decodedPayload); document.cookie = `token=${token}; path=/;`; @@ -256,8 +331,6 @@ describe("useAuthorized", () => { }); expect(result.current.accessToken).toBe("api-key-123"); - expect(result.current.userId).toBe("user-1"); - expect(result.current.userEmail).toBe("user@example.com"); }); it("should redirect when token is missing", async () => { @@ -313,5 +386,193 @@ describe("useAuthorized", () => { expect(replaceMock).toHaveBeenCalledWith("http://proxy.example/ui/login/"); expect(checkTokenValidityMock).toHaveBeenCalledWith(token); + expect(currentUserRequests()).toHaveLength(0); + }); + + describe("identity from the server", () => { + const uiConfig = { + server_root_path: "/", + proxy_base_url: null, + auto_redirect_to_sso: false, + admin_ui_disabled: false, + sso_configured: false, + }; + + const setSessionCookie = (claims: Record) => { + decodeTokenMock.mockReturnValue(claims); + checkTokenValidityMock.mockReturnValue(true); + const token = createJwt(claims); + document.cookie = `token=${token}; path=/;`; + return token; + }; + + it("uses the role the server resolved, not the role claimed in the cookie", async () => { + getUiConfigMock.mockResolvedValue(uiConfig); + setSessionCookie({ key: "internal-key", user_id: "internal-1", user_role: "proxy_admin" }); + serveCurrentUser({ + user_id: "internal-1", + user_email: "internal@example.com", + user_role: "internal_user", + }); + + const { result } = renderHook(() => useAuthorized(), { wrapper }); + + await waitFor(() => { + expect(result.current.userRole).toBe("Internal User"); + }); + expect(result.current.userRoleLabel).toBe("Internal User"); + expect(result.current.isViewOnly).toBe(false); + expect(result.current.userEmail).toBe("internal@example.com"); + expect(replaceMock).not.toHaveBeenCalled(); + }); + + it("stays loading with no role until the server answers", async () => { + getUiConfigMock.mockResolvedValue(uiConfig); + setSessionCookie({ key: "admin-claim-key", user_id: "u-1", user_role: "proxy_admin" }); + stallCurrentUser(); + + const { result } = renderHook(() => ({ auth: useAuthorized(), uiConfig: useUIConfig() }), { wrapper }); + + await waitFor(() => { + expect(result.current.uiConfig.isSuccess).toBe(true); + }); + expect(result.current.auth.isLoading).toBe(true); + expect(result.current.auth.userRole).toBe("Undefined Role"); + expect(replaceMock).not.toHaveBeenCalled(); + }); + + it("waits for backend discovery before looking up the current user", async () => { + let finishDiscovery!: (config: typeof uiConfig) => void; + getUiConfigMock.mockReturnValue(new Promise((resolve) => (finishDiscovery = resolve))); + setSessionCookie({ key: "discovery-key", user_id: "u-9", user_role: "internal_user" }); + serveCurrentUser({ user_id: "u-9", user_email: null, user_role: "internal_user" }); + + const { result } = renderHook(() => useAuthorized(), { wrapper }); + + await waitFor(() => expect(getUiConfigMock).toHaveBeenCalled()); + await act(() => new Promise((resolve) => setTimeout(resolve, 50))); + expect(currentUserRequests()).toHaveLength(0); + expect(result.current.isLoading).toBe(true); + + await act(async () => finishDiscovery(uiConfig)); + + await waitFor(() => expect(result.current.userId).toBe("u-9")); + expect(currentUserRequests()).toHaveLength(1); + expect(replaceMock).not.toHaveBeenCalled(); + }); + + it("keeps an identity the page already loaded while backend discovery reruns", async () => { + getUiConfigMock.mockReturnValue(new Promise(() => {})); + setSessionCookie({ key: "cached-key", user_id: "u-7", user_role: "proxy_admin" }); + serveCurrentUser({ user_id: "u-7", user_email: null, user_role: "internal_user" }); + const queryClient = createQueryClient(); + const sharedWrapper = ({ children }: { children: React.ReactNode }) => + React.createElement(QueryClientProvider, { client: queryClient }, children); + + const { result: layoutLookup } = renderHook(() => useCurrentUser("cached-key"), { wrapper: sharedWrapper }); + await waitFor(() => expect(layoutLookup.current.isSuccess).toBe(true)); + + const { result } = renderHook(() => useAuthorized(), { wrapper: sharedWrapper }); + + expect(result.current.userId).toBe("u-7"); + expect(result.current.userRole).toBe("Internal User"); + expect(currentUserRequests()).toHaveLength(1); + }); + + it("clears the cookie and redirects to login when the current user cannot be loaded", async () => { + getUiConfigMock.mockResolvedValue(uiConfig); + setSessionCookie({ key: "orphan-key", user_id: "u-1", user_role: "internal_user" }); + failCurrentUser(404); + + renderHook(() => useAuthorized(), { wrapper }); + + await waitFor( + () => { + expect(replaceMock).toHaveBeenCalledWith("http://proxy.example/ui/login/"); + }, + { timeout: 4000 }, + ); + expect(clearTokenCookiesMock).toHaveBeenCalled(); + expect(currentUserRequests()).toHaveLength(2); + }); + + it("refetches the identity when a different session key logs in", async () => { + getUiConfigMock.mockResolvedValue(uiConfig); + serveCurrentUser((request) => + request.headers.get("Authorization") === "Bearer first-key" + ? { user_id: "first", user_email: null, user_role: "proxy_admin" } + : { user_id: "second", user_email: null, user_role: "internal_user" }, + ); + const queryClient = createQueryClient(); + const sharedClientWrapper = ({ children }: { children: React.ReactNode }) => + React.createElement(QueryClientProvider, { client: queryClient }, children); + setSessionCookie({ key: "first-key", user_id: "first" }); + + const { result, rerender } = renderHook(() => useAuthorized(), { wrapper: sharedClientWrapper }); + await waitFor(() => { + expect(result.current.userId).toBe("first"); + }); + + setSessionCookie({ key: "second-key", user_id: "second" }); + act(() => rerender()); + + await waitFor(() => { + expect(result.current.userId).toBe("second"); + }); + expect(result.current.userRole).toBe("Internal User"); + }); + + it.each([ + { cookieRole: "internal_user", serverRole: "proxy_admin_viewer", isViewOnly: true }, + { cookieRole: "proxy_admin_viewer", serverRole: "proxy_admin", isViewOnly: false }, + ])("derives view-only from the server role ($serverRole), not the cookie ($cookieRole)", async (scenario) => { + getUiConfigMock.mockResolvedValue(uiConfig); + setSessionCookie({ key: `view-only-${scenario.serverRole}`, user_id: "u-1", user_role: scenario.cookieRole }); + serveCurrentUser({ user_id: "u-1", user_email: null, user_role: scenario.serverRole }); + + const { result } = renderHook(() => useAuthorized(), { wrapper }); + + await waitFor(() => { + expect(result.current.userId).toBe("u-1"); + }); + expect(result.current.isViewOnly).toBe(scenario.isViewOnly); + }); + + it("does not refetch the identity when a consumer mounts after it has loaded", async () => { + getUiConfigMock.mockResolvedValue(uiConfig); + setSessionCookie({ key: "late-key", user_id: "u-1" }); + serveCurrentUser({ user_id: "u-1", user_email: null, user_role: "internal_user" }); + const queryClient = createQueryClient(); + const sharedClientWrapper = ({ children }: { children: React.ReactNode }) => + React.createElement(QueryClientProvider, { client: queryClient }, children); + + const first = renderHook(() => useAuthorized(), { wrapper: sharedClientWrapper }); + await waitFor(() => { + expect(first.result.current.userId).toBe("u-1"); + }); + const second = renderHook(() => useAuthorized(), { wrapper: sharedClientWrapper }); + + expect(second.result.current.userId).toBe("u-1"); + await act(async () => {}); + expect(currentUserRequests()).toHaveLength(1); + }); + + it("fetches the identity once for many consumers on the page", async () => { + getUiConfigMock.mockResolvedValue(uiConfig); + setSessionCookie({ key: "shared-key", user_id: "u-1" }); + serveCurrentUser({ user_id: "u-1", user_email: null, user_role: "internal_user" }); + const queryClient = createQueryClient(); + const sharedClientWrapper = ({ children }: { children: React.ReactNode }) => + React.createElement(QueryClientProvider, { client: queryClient }, children); + + const { result } = renderHook(() => [useAuthorized(), useAuthorized(), useAuthorized()], { + wrapper: sharedClientWrapper, + }); + + await waitFor(() => { + expect(result.current.map((auth) => auth.userId)).toEqual(["u-1", "u-1", "u-1"]); + }); + expect(currentUserRequests()).toHaveLength(1); + }); }); }); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/useAuthorized.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/useAuthorized.ts index 0d240c02f45..47716fccfeb 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/hooks/useAuthorized.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/useAuthorized.ts @@ -1,11 +1,11 @@ "use client"; -import { getProxyBaseUrl, getUiPageUrl } from "@/components/networking"; import { clearTokenCookies, getCookie } from "@/utils/cookieUtils"; import { checkTokenValidity, decodeToken } from "@/utils/jwtUtils"; -import { buildLoginUrlWithReturn, getLoginUrl, storeReturnUrl } from "@/utils/returnUrlUtils"; -import { useCallback, useEffect, useMemo } from "react"; +import { useEffect, useMemo } from "react"; import { effectiveSessionRole, formatUserRole, isViewOnlySessionRole } from "@/utils/roles"; +import { useCurrentUser } from "./currentUser/useCurrentUser"; +import { redirectToLogin } from "./redirectToLogin"; import { useUIConfig } from "./uiConfig/useUIConfig"; const useAuthorized = () => { @@ -15,16 +15,14 @@ const useAuthorized = () => { const decoded = useMemo(() => decodeToken(token), [token]); const isTokenValid = useMemo(() => checkTokenValidity(token), [token]); - const isLoading = isUIConfigLoading; - const isAuthorized = isTokenValid && !uiConfig?.admin_ui_disabled; - - // Helper function to redirect to login while preserving the current URL - const redirectToLogin = useCallback(() => { - storeReturnUrl(); - const baseLoginUrl = getUiPageUrl(getLoginUrl(getProxyBaseUrl())); - const loginUrlWithReturn = buildLoginUrlWithReturn(baseLoginUrl); - window.location.replace(loginUrlWithReturn); - }, []); + const { + data: currentUser, + isLoading: isCurrentUserLoading, + isError: isCurrentUserError, + } = useCurrentUser(isTokenValid ? decoded?.key ?? null : null, decoded?.auth_header_name, !isUIConfigLoading); + const serverRole = currentUser?.user_role ?? ""; + const isLoading = isUIConfigLoading || isCurrentUserLoading; + const isAuthorized = isTokenValid && !uiConfig?.admin_ui_disabled && !isCurrentUserError; // Single useEffect for all redirect logic useEffect(() => { @@ -36,18 +34,18 @@ const useAuthorized = () => { } redirectToLogin(); } - }, [isLoading, isAuthorized, token, redirectToLogin]); + }, [isLoading, isAuthorized, token]); return { isLoading, isAuthorized, token: isAuthorized ? token : null, accessToken: decoded?.key ?? null, - userId: decoded?.user_id ?? null, - userEmail: decoded?.user_email ?? null, - userRole: effectiveSessionRole(decoded?.user_role), - userRoleLabel: formatUserRole(decoded?.user_role), - isViewOnly: isViewOnlySessionRole(decoded?.user_role), + userId: currentUser?.user_id ?? null, + userEmail: currentUser?.user_email ?? null, + userRole: effectiveSessionRole(serverRole), + userRoleLabel: formatUserRole(serverRole), + isViewOnly: isViewOnlySessionRole(serverRole), premiumUser: decoded?.premium_user ?? null, disabledPersonalKeyCreation: decoded?.disabled_non_admin_personal_key_creation ?? null, loginMethod: decoded?.login_method ?? null, diff --git a/ui/litellm-dashboard/src/app/(dashboard)/layout.integration.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/layout.integration.test.tsx index aa2069eb74e..c4c32d5f7f3 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/layout.integration.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/layout.integration.test.tsx @@ -81,6 +81,8 @@ beforeEach(() => { }; return Response.json(usage); } + case "/management/v1/users/me": + return Response.json({ user_id: "layout-admin", user_role: "proxy_admin", user_email: null }); case "/sso/get/ui_settings": return Response.json({}); case "/get/user_banner": diff --git a/ui/litellm-dashboard/src/app/(dashboard)/layout.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/layout.test.tsx index 154f72719e2..0fdbd047786 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/layout.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/layout.test.tsx @@ -7,6 +7,17 @@ import Layout from "./layout"; const { replaceMock } = vi.hoisted(() => ({ replaceMock: vi.fn() })); +const CURRENT_USER_PATH = "/management/v1/users/me"; +const fetchMock = vi.fn<(input: RequestInfo | URL) => Promise>(); +const pathOf = (input: RequestInfo | URL) => + new URL(input instanceof Request ? input.url : String(input), "http://localhost").pathname; +let currentUserReply: () => Promise; + +const currentUserRequests = () => + fetchMock.mock.calls + .map(([input]) => input) + .filter((input): input is Request => input instanceof Request && pathOf(input) === CURRENT_USER_PATH); + let searchParamsValue = new URLSearchParams(); vi.mock("@/lib/telemetry/uiEvents", () => ({ @@ -79,6 +90,17 @@ describe("(dashboard) Layout", () => { pendingUiConfig = createDeferred(); searchParamsValue = new URLSearchParams(); vi.mocked(usePathname).mockReturnValue("/ui/guardrails"); + currentUserReply = async () => + Response.json({ user_id: "normal-user", user_role: "internal_user", user_email: null }); + fetchMock.mockReset(); + fetchMock.mockImplementation((input) => + pathOf(input) === CURRENT_USER_PATH ? currentUserReply() : Promise.resolve(Response.json({}, { status: 404 })), + ); + vi.stubGlobal("fetch", fetchMock); + }); + + afterEach(() => { + vi.unstubAllGlobals(); }); it("starts mobile navigation closed, opens a modal drawer and closes it after choosing a page", async () => { @@ -220,6 +242,7 @@ describe("(dashboard) Layout", () => { afterEach(() => { document.cookie = "token=; Max-Age=0; Path=/"; + document.cookie = "litellm_return_url=; Max-Age=0; Path=/"; }); it("routes a session flagged password_reset_required to the change-password page", async () => { @@ -264,5 +287,99 @@ describe("(dashboard) Layout", () => { expect(await screen.findByTestId("page-content")).toBeInTheDocument(); expect(replaceMock).not.toHaveBeenCalled(); }); + + it("shows the loading screen instead of the shell until the current user resolves", async () => { + let resolveCurrentUser!: (response: Response) => void; + currentUserReply = () => new Promise((resolve) => (resolveCurrentUser = resolve)); + document.cookie = `token=${sessionCookie({ user_id: "normal-user", key: "sk-pending-session", user_role: "proxy_admin" })}; Path=/`; + + render( + + +
+ + , + ); + pendingUiConfig.resolve(); + + await waitFor(() => + expect(currentUserRequests().map((request) => request.headers.get("Authorization"))).toEqual([ + "Bearer sk-pending-session", + ]), + ); + expect(screen.getByTestId("loading-screen")).toBeInTheDocument(); + expect(screen.queryByTestId("sidebar")).not.toBeInTheDocument(); + expect(screen.queryByTestId("page-content")).not.toBeInTheDocument(); + + resolveCurrentUser(Response.json({ user_id: "normal-user", user_role: "internal_user", user_email: null })); + + expect(await screen.findByTestId("page-content")).toBeInTheDocument(); + expect(screen.getByTestId("sidebar")).toBeInTheDocument(); + }); + + it("signs out instead of mounting the shell when the current user cannot be loaded", async () => { + const originalLocation = window.location; + const locationReplace = vi.fn(); + Object.defineProperty(window, "location", { + value: { ...originalLocation, replace: locationReplace }, + writable: true, + }); + currentUserReply = async () => Response.json({ detail: "rejected" }, { status: 401 }); + document.cookie = `token=${sessionCookie({ user_id: "gone-user", key: "sk-rejected-session" })}; Path=/`; + + try { + render( + + +
+ + , + ); + pendingUiConfig.resolve(); + + await waitFor(() => expect(locationReplace).toHaveBeenCalledWith(expect.stringContaining("/login")), { + timeout: 4000, + }); + expect(document.cookie).not.toContain("token="); + expect(document.cookie).toContain("litellm_return_url="); + expect(screen.queryByTestId("sidebar")).not.toBeInTheDocument(); + expect(screen.queryByTestId("page-content")).not.toBeInTheDocument(); + expect(currentUserRequests()).toHaveLength(2); + } finally { + Object.defineProperty(window, "location", { value: originalLocation, writable: true }); + } + }); + + it("keeps the session of an invitation link whose current user cannot be loaded", async () => { + const originalLocation = window.location; + const locationReplace = vi.fn(); + Object.defineProperty(window, "location", { + value: { ...originalLocation, replace: locationReplace }, + writable: true, + }); + searchParamsValue = new URLSearchParams("invitation_id=abc123"); + currentUserReply = async () => Response.json({ detail: "rejected" }, { status: 401 }); + document.cookie = `token=${sessionCookie({ user_id: "invited-user", key: "sk-invited-session" })}; Path=/`; + + try { + render( + + +
+ + , + ); + pendingUiConfig.resolve(); + + await waitFor(() => expect(currentUserRequests()).toHaveLength(2), { timeout: 4000 }); + await waitFor(() => + expect(replaceMock).toHaveBeenCalledWith(expect.stringContaining("/onboarding?invitation_id=abc123")), + ); + expect(locationReplace).not.toHaveBeenCalled(); + expect(document.cookie).toContain("token="); + } finally { + Object.defineProperty(window, "location", { value: originalLocation, writable: true }); + } + }); }); }); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/layout.tsx b/ui/litellm-dashboard/src/app/(dashboard)/layout.tsx index 083118bc65d..92c8806df07 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/layout.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/layout.tsx @@ -7,6 +7,10 @@ import Navbar from "@/components/navbar"; import LoadingScreen from "@/components/common_components/LoadingScreen"; import { ThemeProvider } from "@/contexts/ThemeContext"; import { useAuth } from "@/contexts/AuthContext"; +import { useCurrentUser } from "@/app/(dashboard)/hooks/currentUser/useCurrentUser"; +import { redirectToLogin } from "@/app/(dashboard)/hooks/redirectToLogin"; +import { clearTokenCookies } from "@/utils/cookieUtils"; +import { decodeToken } from "@/utils/jwtUtils"; import SidebarProvider from "@/app/(dashboard)/components/SidebarProvider"; import { usePathname, useRouter, useSearchParams } from "next/navigation"; import { @@ -218,8 +222,14 @@ function LayoutContent({ children }: { children: React.ReactNode }) { const router = useRouter(); const searchParams = useSearchParams(); const pathname = usePathname(); - const { accessToken, authLoading, passwordResetRequired } = useAuth(); + const { accessToken, token, authLoading, passwordResetRequired } = useAuth(); + const session = decodeToken(token); + const { isLoading: currentUserLoading, isError: currentUserFailed } = useCurrentUser( + session?.key ?? null, + session?.auth_header_name, + ); const isInvitationFlow = Boolean(searchParams.get("invitation_id")); + const identityPending = currentUserLoading || currentUserFailed; // Legacy invitation links point at /ui/?invitation_id=; the onboarding form now lives at its own // /onboarding route. Redirect once ui-config has loaded so uiHref resolves the SERVER_ROOT_PATH base. @@ -237,7 +247,13 @@ function LayoutContent({ children }: { children: React.ReactNode }) { } }, [authLoading, passwordResetRequired, pathname, router]); - if (authLoading || isInvitationFlow) { + useEffect(() => { + if (!currentUserFailed || isInvitationFlow) return; + clearTokenCookies(); + redirectToLogin(); + }, [currentUserFailed, isInvitationFlow]); + + if (authLoading || identityPending || isInvitationFlow) { return ; } diff --git a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/AllModelsTab.tsx b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/AllModelsTab.tsx index 478ddcd2458..92890031ddd 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/AllModelsTab.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/AllModelsTab.tsx @@ -337,7 +337,7 @@ const AllModelsTab = ({ availableModelGroups={availableModelGroups} availableModelAccessGroups={availableModelAccessGroups} userRole={userRole} - userID={userId} + userID={userId ?? ""} isViewOnly={isViewOnly} onModelIdClick={setSelectedModelId} onTeamIdClick={setSelectedTeamId} diff --git a/ui/litellm-dashboard/src/app/(dashboard)/policies/_components/add_policy_form.tsx b/ui/litellm-dashboard/src/app/(dashboard)/policies/_components/add_policy_form.tsx index b43bcba7bd0..0d1fc4f7aba 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/policies/_components/add_policy_form.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/policies/_components/add_policy_form.tsx @@ -247,7 +247,7 @@ const AddPolicyForm: React.FC = ({ }, [visible, accessToken]); const loadAvailableModels = async () => { - if (!accessToken) return; + if (!accessToken || !userId) return; try { const response = await modelAvailableCall(accessToken, userId, userRole); if (response?.data) { diff --git a/ui/litellm-dashboard/src/app/chat/layout.test.tsx b/ui/litellm-dashboard/src/app/chat/layout.test.tsx index ac4fe2068fb..d455b84afa3 100644 --- a/ui/litellm-dashboard/src/app/chat/layout.test.tsx +++ b/ui/litellm-dashboard/src/app/chat/layout.test.tsx @@ -7,6 +7,8 @@ const { mockUseAuthorized, mockUseUISettings, mockReplace, mockUiHref, state } = const state = { enableChatUI: false, isUISettingsLoading: false, + isIdentityLoading: false, + isAuthorized: true, }; return { state, @@ -18,6 +20,8 @@ const { mockUseAuthorized, mockUseUISettings, mockReplace, mockUiHref, state } = userId: "user-1", userEmail: "user@example.com", premiumUser: false, + isLoading: state.isIdentityLoading, + isAuthorized: state.isAuthorized, })), mockUseUISettings: vi.fn(() => ({ data: { values: { enable_chat_ui: state.enableChatUI } }, @@ -47,6 +51,8 @@ describe("ChatLayout", () => { afterEach(() => { state.enableChatUI = false; state.isUISettingsLoading = false; + state.isIdentityLoading = false; + state.isAuthorized = true; mockReplace.mockClear(); mockUiHref.mockClear(); }); @@ -84,4 +90,28 @@ describe("ChatLayout", () => { expect(screen.queryByTestId("chat-shell")).not.toBeInTheDocument(); expect(mockReplace).not.toHaveBeenCalled(); }); + + it("renders nothing until the signed-in user's identity has loaded", () => { + state.enableChatUI = true; + state.isIdentityLoading = true; + render( + +
+ , + ); + expect(screen.queryByTestId("chat-shell")).not.toBeInTheDocument(); + expect(screen.queryByTestId("page-content")).not.toBeInTheDocument(); + }); + + it("renders nothing when the signed-in user's identity could not be confirmed", () => { + state.enableChatUI = true; + state.isAuthorized = false; + render( + +
+ , + ); + expect(screen.queryByTestId("chat-shell")).not.toBeInTheDocument(); + expect(screen.queryByTestId("page-content")).not.toBeInTheDocument(); + }); }); diff --git a/ui/litellm-dashboard/src/app/chat/layout.tsx b/ui/litellm-dashboard/src/app/chat/layout.tsx index 0bed5bda4a8..bfe494d2234 100644 --- a/ui/litellm-dashboard/src/app/chat/layout.tsx +++ b/ui/litellm-dashboard/src/app/chat/layout.tsx @@ -12,7 +12,15 @@ import { uiHref } from "@/utils/uiHref"; // The nuqs Next adapter uses useSearchParams, so keep the chat tree behind Suspense. function ChatLayoutContent({ children }: { children: React.ReactNode }) { - const { accessToken, userRole, userId, userEmail, premiumUser } = useAuthorized(); + const { + accessToken, + userRole, + userId, + userEmail, + premiumUser, + isLoading: isIdentityLoading, + isAuthorized, + } = useAuthorized(); const { data: uiSettings, isLoading: isUISettingsLoading } = useUISettings(); const router = useRouter(); @@ -23,7 +31,9 @@ function ChatLayoutContent({ children }: { children: React.ReactNode }) { if (blocked) router.replace(uiHref("")); }, [blocked, router]); - if (isUISettingsLoading || blocked) return null; + const identityReady = !isIdentityLoading && isAuthorized; + + if (!identityReady || isUISettingsLoading || blocked) return null; return ( diff --git a/ui/litellm-dashboard/src/components/DeploymentHealth/DeploymentHealth.integration.test.tsx b/ui/litellm-dashboard/src/components/DeploymentHealth/DeploymentHealth.integration.test.tsx index 1b4d96c352d..87c34165fd9 100644 --- a/ui/litellm-dashboard/src/components/DeploymentHealth/DeploymentHealth.integration.test.tsx +++ b/ui/litellm-dashboard/src/components/DeploymentHealth/DeploymentHealth.integration.test.tsx @@ -36,7 +36,10 @@ const snoozedBy = (user: string, until: string): DeploymentWarning["ignore"] => ignored_until: until, }); +let signedInRole = "proxy_admin"; + const signIn = (userRole: string) => { + signedInRole = userRole; const claims = { key: "sk-health-test", user_id: "admin-1", user_role: userRole, exp: NOW.getTime() / 1000 + 3600 }; const token = `${btoa(JSON.stringify({ alg: "none" }))}.${btoa(JSON.stringify(claims))}.test`; document.cookie = `token=${token}; Path=/`; @@ -47,6 +50,9 @@ const handleRequest = async (request: Request): Promise => { if (pathname === "/litellm/.well-known/litellm-ui-config") { return Response.json({ server_root_path: "", proxy_base_url: "", admin_ui_disabled: false }); } + if (pathname === "/management/v1/users/me") { + return Response.json({ user_id: "admin-1", user_role: signedInRole, user_email: null }); + } if (pathname === "/health/deployment_warnings" && request.method === "GET") { proxy.gets += 1; return Response.json({ warnings: proxy.warnings, ignore_available: proxy.ignoreAvailable }); diff --git a/ui/litellm-dashboard/src/components/add_model/AddModelForm.tsx b/ui/litellm-dashboard/src/components/add_model/AddModelForm.tsx index 7e90991523c..beab81c16a2 100644 --- a/ui/litellm-dashboard/src/components/add_model/AddModelForm.tsx +++ b/ui/litellm-dashboard/src/components/add_model/AddModelForm.tsx @@ -202,7 +202,7 @@ const AddModelForm: React.FC = ({ : null; const isAdmin = all_admin_roles.includes(userRole); - const isTeamAdmin = isUserTeamAdminForAnyTeam(teams, userId); + const isTeamAdmin = isUserTeamAdminForAnyTeam(teams, userId ?? ""); // Same owner the Auto-Routers tab uses, so the two creation forms cannot disagree about // who has to name a team. This form is only reachable when creation is allowed at all. const createScope = modelCreationScope( diff --git a/ui/litellm-dashboard/src/components/liteadmin/LiteAdmin.integration.test.tsx b/ui/litellm-dashboard/src/components/liteadmin/LiteAdmin.integration.test.tsx index 61c3dbbc1f1..b0ff9871570 100644 --- a/ui/litellm-dashboard/src/components/liteadmin/LiteAdmin.integration.test.tsx +++ b/ui/litellm-dashboard/src/components/liteadmin/LiteAdmin.integration.test.tsx @@ -64,15 +64,18 @@ function deferred() { return { promise, resolve }; } +let signedInUser = { user_id: "first-admin", user_role: "proxy_admin" }; + function session( role = "proxy_admin", user = "first-admin", license: { premium_user?: boolean | null | string } = { premium_user: true }, + key = `sk-session-${user}`, ) { const encode = (value: object) => btoa(JSON.stringify(value)).replaceAll("=", "").replaceAll("+", "-").replaceAll("/", "_"); const claims = { - key: `sk-session-${user}`, + key, user_id: user, user_role: role, ...license, @@ -81,6 +84,7 @@ function session( }; const token = `${encode({ alg: "none" })}.${encode(claims)}.test`; document.cookie = `token=${token}; Path=/`; + signedInUser = { user_id: user, user_role: role }; return token; } @@ -134,6 +138,7 @@ function gateway(replies: (ModelReply | Promise)[], options: Gateway const recorded = { url: request.url, headers: request.headers, body, signal: request.signal }; const path = new URL(request.url).pathname; if (path.endsWith("/ui/litellm-runtime-config.json")) return json({}); + if (path.endsWith("/management/v1/users/me")) return json({ ...signedInUser, user_email: null }); requests.push(recorded); if (path.endsWith("/litellm-ui-config")) return json({ proxy_base_url: MANAGEMENT, server_root_path: "", admin_ui_disabled: false }); @@ -1018,7 +1023,7 @@ describe("LiteAdmin in the gateway", () => { settings.target = `${INFERENCE}/updated`; await act(async () => view.client.invalidateQueries({ queryKey: ["proxySettings"] })); } else { - session("proxy_admin_viewer"); + session("proxy_admin_viewer", "first-admin", { premium_user: true }, "sk-session-first-admin-viewer"); view.refresh(); expect(screen.queryByRole("button", { name: "LiteAdmin" })).not.toBeInTheDocument(); expect(screen.queryByRole("region", { name: "Ask LiteAdmin", hidden: true })).not.toBeInTheDocument(); diff --git a/ui/litellm-dashboard/src/components/organisms/create_key_button.tsx b/ui/litellm-dashboard/src/components/organisms/create_key_button.tsx index a15f6c6e31c..30c7a80c456 100644 --- a/ui/litellm-dashboard/src/components/organisms/create_key_button.tsx +++ b/ui/litellm-dashboard/src/components/organisms/create_key_button.tsx @@ -479,7 +479,7 @@ const CreateKey: React.FC = ({ team, teams, data, addKey, autoOp const response = endpoint === "service_account" ? await keyCreateServiceAccountCall(accessToken, payload) - : await keyCreateCall(accessToken, userID, payload); + : await keyCreateCall(accessToken, userID ?? "", payload); // Add the data to the state in the parent component // Also directly update the keys list in VirtualKeysTable without an API call @@ -1902,7 +1902,7 @@ const CreateKey: React.FC = ({ team, teams, data, addKey, autoOp Create New User