diff --git a/ui/litellm-dashboard/src/app/(dashboard)/layout.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/layout.test.tsx index 0fdbd047786..19b70b1d42d 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/layout.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/layout.test.tsx @@ -1,7 +1,8 @@ import { describe, it, expect, vi, beforeEach, afterEach } from "vitest"; -import { fireEvent, renderWithProviders as render, screen, waitFor, within } from "../../../tests/test-utils"; +import { act, fireEvent, renderWithProviders as render, screen, waitFor, within } from "../../../tests/test-utils"; import { usePathname } from "next/navigation"; -import { AuthProvider } from "@/contexts/AuthContext"; +import { AuthProvider, useAuth } from "@/contexts/AuthContext"; +import { clearTokenCookies, storeLoginToken } from "@/utils/cookieUtils"; import { recordUiEvent } from "@/lib/telemetry/uiEvents"; import Layout from "./layout"; @@ -241,10 +242,49 @@ describe("(dashboard) Layout", () => { }; afterEach(() => { + clearTokenCookies(); document.cookie = "token=; Max-Age=0; Path=/"; document.cookie = "litellm_return_url=; Max-Age=0; Path=/"; }); + function SessionIdentity() { + const { authLoading, userEmail, userRole } = useAuth(); + if (authLoading) return
Loading session
; + return
{userEmail ? `${userEmail} (${userRole})` : "Signed out"}
; + } + + it("updates the mounted session after login and when a different user signs in", async () => { + clearTokenCookies(); + render( + + + , + ); + pendingUiConfig.resolve(); + expect(await screen.findByText("Signed out")).toBeInTheDocument(); + + act(() => storeLoginToken(sessionCookie({ user_email: "admin@example.test", user_role: "proxy_admin" }))); + expect(await screen.findByText("admin@example.test (Admin)")).toBeInTheDocument(); + + act(() => storeLoginToken(sessionCookie({ user_email: "user@example.test", user_role: "internal_user" }))); + expect(await screen.findByText("user@example.test (Internal User)")).toBeInTheDocument(); + }); + + it.each(["", " ", "invalid-token"])("does not authenticate a stored invalid token %j", async (token) => { + clearTokenCookies(); + render( + + + , + ); + pendingUiConfig.resolve(); + expect(await screen.findByText("Signed out")).toBeInTheDocument(); + + act(() => storeLoginToken(token)); + + expect(screen.getByText("Signed out")).toBeInTheDocument(); + }); + it("routes a session flagged password_reset_required to the change-password page", async () => { const flaggedClaims = { user_id: "flagged-user", diff --git a/ui/litellm-dashboard/src/contexts/AuthContext.tsx b/ui/litellm-dashboard/src/contexts/AuthContext.tsx index c68c8b9d81a..18dd45cb56a 100644 --- a/ui/litellm-dashboard/src/contexts/AuthContext.tsx +++ b/ui/litellm-dashboard/src/contexts/AuthContext.tsx @@ -2,7 +2,7 @@ import React, { createContext, useContext, useEffect, useState } from "react"; import { jwtDecode } from "jwt-decode"; -import { clearTokenCookies, getCookie } from "@/utils/cookieUtils"; +import { clearTokenCookies, getCookie, LOGIN_TOKEN_CHANGED_EVENT } from "@/utils/cookieUtils"; import { isJwtExpired } from "@/utils/jwtUtils"; import { effectiveSessionRole } from "@/utils/roles"; import { getUiConfig, setGlobalLitellmHeaderName } from "@/components/networking"; @@ -56,6 +56,17 @@ export function AuthProvider({ children }: { children: React.ReactNode }) { useEffect(() => { let cancelled = false; + const refreshToken = () => { + const raw = getCookie("token"); + const valid = raw && !isJwtExpired(raw) ? raw : null; + if (raw && !valid) { + deleteCookie("token", "/"); + } + setToken(valid); + }; + + window.addEventListener(LOGIN_TOKEN_CHANGED_EVENT, refreshToken); + (async () => { try { await getUiConfig(); @@ -65,20 +76,13 @@ export function AuthProvider({ children }: { children: React.ReactNode }) { if (cancelled) return; - const raw = getCookie("token"); - const valid = raw && !isJwtExpired(raw) ? raw : null; - - // Clear expired/invalid token so downstream code doesn't keep trying to use it. - if (raw && !valid) { - deleteCookie("token", "/"); - } - - setToken(valid); + refreshToken(); setAuthLoading(false); })(); return () => { cancelled = true; + window.removeEventListener(LOGIN_TOKEN_CHANGED_EVENT, refreshToken); }; }, []); diff --git a/ui/litellm-dashboard/src/utils/cookieUtils.ts b/ui/litellm-dashboard/src/utils/cookieUtils.ts index 66eb807ed2d..a403775b20e 100644 --- a/ui/litellm-dashboard/src/utils/cookieUtils.ts +++ b/ui/litellm-dashboard/src/utils/cookieUtils.ts @@ -4,6 +4,8 @@ import { clearAllMcpTokens } from "./mcpTokenStore"; +export const LOGIN_TOKEN_CHANGED_EVENT = "litellm-login-token-changed"; + /** * Returns the cookie path for the UI. * Derives the path from window.location.pathname so it works when @@ -103,6 +105,8 @@ export function storeLoginToken(token: string) { } catch { // sessionStorage may be unavailable (e.g. private browsing quota exceeded) } + + window.dispatchEvent(new Event(LOGIN_TOKEN_CHANGED_EVENT)); } /**