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