From e8c8d9ba2223dbe576fa2dac0c0fa7fda7df19e6 Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Fri, 24 Jul 2026 16:56:06 -0700 Subject: [PATCH] fix(ui): wait for proxy settings on logout and share one logout handler Greptile flagged that clicking Logout before the proxy settings query resolved treated the loading fallback as an unconfigured PROXY_LOGOUT_URL, skipping a configured SSO logout endpoint and leaving the upstream session active. Logout now awaits the settings query (instant when cached, waits when in flight) before deciding where to navigate. The navbar had also drifted from the shared useLogout hook the sidebar footer uses, which still had the old reload-in-place bug. Both entry points now go through useLogout, which carries the full fix. --- .../hooks/proxySettings/useProxySettings.ts | 12 ++++++- .../src/app/(dashboard)/hooks/useLogout.ts | 16 ++++++---- .../src/components/navbar.test.tsx | 32 +++++++++++++++++++ .../src/components/navbar.tsx | 12 ++----- 4 files changed, 55 insertions(+), 17 deletions(-) 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();