diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/proxySettings/useProxySettings.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/proxySettings/useProxySettings.ts index 82cefd800f4..038de1242e0 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/hooks/proxySettings/useProxySettings.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/proxySettings/useProxySettings.ts @@ -1,5 +1,5 @@ import { fetchProxySettings } from "@/utils/proxyUtils"; -import { useQuery } from "@tanstack/react-query"; +import { QueryClient, useQuery } from "@tanstack/react-query"; import { createQueryKeys } from "../common/queryKeysFactory"; export const proxySettingsKeys = createQueryKeys("proxySettings"); @@ -24,3 +24,13 @@ export default function useProxySettings(accessToken: string | null): ProxySetti }); return data ?? EMPTY_PROXY_SETTINGS; } + +export function ensureProxySettings( + queryClient: QueryClient, + accessToken: string | null, +): Promise { + return queryClient.ensureQueryData({ + queryKey: [...proxySettingsKeys.all, accessToken], + queryFn: () => fetchProxySettings(accessToken), + }); +} diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/useLogout.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/useLogout.ts index 8da057ef9be..3bf0a0228cc 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/hooks/useLogout.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/useLogout.ts @@ -1,19 +1,23 @@ +import { getProxyBaseUrl } from "@/components/networking"; import { clearTokenCookies } from "@/utils/cookieUtils"; -import { clearStoredReturnUrl } from "@/utils/returnUrlUtils"; -import useProxySettings from "@/app/(dashboard)/hooks/proxySettings/useProxySettings"; +import { clearStoredReturnUrl, getLoginUrl } from "@/utils/returnUrlUtils"; +import { useQueryClient } from "@tanstack/react-query"; +import useProxySettings, { ensureProxySettings } from "@/app/(dashboard)/hooks/proxySettings/useProxySettings"; /** * Shared sign-out handler. Used by both the top navbar and the sidebar footer so * the two entry points can never drift on which client state gets cleared. */ -export function useLogout(accessToken: string | null): () => void { - const proxySettings = useProxySettings(accessToken); +export function useLogout(accessToken: string | null): () => Promise { + const queryClient = useQueryClient(); + useProxySettings(accessToken); - return () => { + return async () => { + const settings = await ensureProxySettings(queryClient, accessToken).catch(() => null); clearTokenCookies(); clearStoredReturnUrl(); localStorage.removeItem("litellm_selected_worker_id"); localStorage.removeItem("litellm_worker_url"); - window.location.href = proxySettings.PROXY_LOGOUT_URL || ""; + window.location.replace(settings?.PROXY_LOGOUT_URL || getLoginUrl(getProxyBaseUrl())); }; } diff --git a/ui/litellm-dashboard/src/components/navbar.test.tsx b/ui/litellm-dashboard/src/components/navbar.test.tsx index bbf78aa9bc7..8df278d894a 100644 --- a/ui/litellm-dashboard/src/components/navbar.test.tsx +++ b/ui/litellm-dashboard/src/components/navbar.test.tsx @@ -340,6 +340,38 @@ describe("Navbar", () => { }); }); + it("should wait for an in-flight proxy settings fetch so a configured logout URL is not skipped", async () => { + const user = userEvent.setup(); + vi.mocked(window.location.replace).mockClear(); + + const proxyUtils = vi.mocked(await import("@/utils/proxyUtils")); + let resolveSettings!: (value: { PROXY_BASE_URL: string; PROXY_LOGOUT_URL: string }) => void; + proxyUtils.fetchProxySettings.mockImplementationOnce( + () => + new Promise((resolve) => { + resolveSettings = resolve; + }), + ); + + renderWithProviders(); + + await user.click(screen.getByRole("button", { name: /open account menu/i })); + + await waitFor(() => { + expect(screen.getByText("test-user")).toBeInTheDocument(); + }); + + await user.click(screen.getByText("Logout")); + + expect(window.location.replace).not.toHaveBeenCalled(); + + resolveSettings({ PROXY_BASE_URL: "", PROXY_LOGOUT_URL: "https://sso.example.com/logout" }); + + await waitFor(() => { + expect(window.location.replace).toHaveBeenCalledWith("https://sso.example.com/logout"); + }); + }); + it("should not render dark mode toggle slider", () => { renderWithProviders(); diff --git a/ui/litellm-dashboard/src/components/navbar.tsx b/ui/litellm-dashboard/src/components/navbar.tsx index 5376d1490f7..82e2ed6cbba 100644 --- a/ui/litellm-dashboard/src/components/navbar.tsx +++ b/ui/litellm-dashboard/src/components/navbar.tsx @@ -6,7 +6,7 @@ import { getProxyBaseUrl } from "@/components/networking"; import { useTheme } from "@/contexts/ThemeContext"; import { clearTokenCookies } from "@/utils/cookieUtils"; import { clearStoredReturnUrl, getLoginUrl } from "@/utils/returnUrlUtils"; -import useProxySettings from "@/app/(dashboard)/hooks/proxySettings/useProxySettings"; +import { useLogout } from "@/app/(dashboard)/hooks/useLogout"; import { DownOutlined, MenuFoldOutlined, MenuUnfoldOutlined } from "@ant-design/icons"; import { Tag } from "antd"; import Link from "next/link"; @@ -33,7 +33,7 @@ const Navbar: React.FC = ({ onToggleSidebar, }) => { const baseUrl = getProxyBaseUrl(); - const proxySettings = useProxySettings(accessToken); + const handleLogout = useLogout(accessToken); const { logoUrl } = useTheme(); const { data: healthData } = useHealthReadinessDetails(accessToken); const version = healthData?.litellm_version; @@ -44,14 +44,6 @@ const Navbar: React.FC = ({ const imageUrl = logoUrl || `${baseUrl}/get_image`; - const handleLogout = () => { - clearTokenCookies(); - clearStoredReturnUrl(); - localStorage.removeItem("litellm_selected_worker_id"); - localStorage.removeItem("litellm_worker_url"); - window.location.replace(proxySettings.PROXY_LOGOUT_URL || getLoginUrl(baseUrl)); - }; - const handleWorkerSwitch = (workerId: string) => { clearTokenCookies(); clearStoredReturnUrl();