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.
This commit is contained in:
ryan-crabbe-berri 2026-07-24 16:56:06 -07:00
parent 1b6cee9b12
commit e8c8d9ba22
4 changed files with 55 additions and 17 deletions

View file

@ -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<ProxySettings | null> {
return queryClient.ensureQueryData({
queryKey: [...proxySettingsKeys.all, accessToken],
queryFn: () => fetchProxySettings(accessToken),
});
}

View file

@ -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<void> {
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()));
};
}

View file

@ -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(<Navbar {...defaultProps} accessToken="test-token-pending-settings" />);
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(<Navbar {...defaultProps} />);

View file

@ -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<NavbarProps> = ({
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<NavbarProps> = ({
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();