Merge pull request #18575 from BerriAI/litellm_server_base_url_fix

[Feature] UI - Clicking on Logo Directs to Correct URL

ui/litellm-dashboard/src/app/(dashboard)/hooks/healthReadiness/useHealthReadiness.ts

@@ -0,0 +1,27 @@
+import { getProxyBaseUrl } from "@/components/networking";
+import { useQuery, UseQueryResult } from "@tanstack/react-query";
+import { createQueryKeys } from "../common/queryKeysFactory";
+
+const healthReadinessKeys = createQueryKeys("healthReadiness");
+
+interface HealthReadinessResponse {
+  litellm_version?: string;
+  [key: string]: any;
+}
+
+const fetchHealthReadiness = async (): Promise<HealthReadinessResponse> => {
+  const baseUrl = getProxyBaseUrl();
+  const response = await fetch(`${baseUrl}/health/readiness`);
+  if (!response.ok) {
+    throw new Error(`Failed to fetch health readiness: ${response.statusText}`);
+  }
+  return response.json();
+};
+
+export const useHealthReadiness = (): UseQueryResult<HealthReadinessResponse> => {
+  return useQuery<HealthReadinessResponse>({
+    queryKey: healthReadinessKeys.detail("readiness"),
+    queryFn: fetchHealthReadiness,
+    staleTime: 5 * 60 * 1000, // 5 minutes
+  });
+};

ui/litellm-dashboard/src/components/navbar.test.tsx

@@ -0,0 +1,181 @@
+import userEvent from "@testing-library/user-event";
+import { describe, expect, it, vi } from "vitest";
+import { renderWithProviders, screen, waitFor } from "../../tests/test-utils";
+import Navbar from "./navbar";
+
+// Mock the hooks and utilities
+vi.mock("@/components/networking", () => ({
+  getProxyBaseUrl: vi.fn(() => "http://localhost:4000"),
+}));
+
+vi.mock("@/utils/proxyUtils", () => ({
+  fetchProxySettings: vi.fn(),
+}));
+
+// Create mock functions that can be controlled in tests
+let mockUseThemeImpl = () => ({ logoUrl: null as string | null });
+let mockUseHealthReadinessImpl = () => ({ data: null as any });
+let mockGetLocalStorageItemImpl = () => null as string | null;
+
+vi.mock("@/contexts/ThemeContext", () => ({
+  useTheme: () => mockUseThemeImpl(),
+}));
+
+vi.mock("@/app/(dashboard)/hooks/healthReadiness/useHealthReadiness", () => ({
+  useHealthReadiness: () => mockUseHealthReadinessImpl(),
+}));
+
+vi.mock("@/utils/localStorageUtils", () => ({
+  getLocalStorageItem: () => mockGetLocalStorageItemImpl(),
+  setLocalStorageItem: vi.fn(),
+  removeLocalStorageItem: vi.fn(),
+  emitLocalStorageChange: vi.fn(),
+}));
+
+vi.mock("@/utils/cookieUtils", () => ({
+  clearTokenCookies: vi.fn(),
+}));
+
+// Mock window.location.href for logout testing
+Object.defineProperty(window, "location", {
+  value: { href: "" },
+  writable: true,
+});
+
+describe("Navbar", () => {
+  const defaultProps = {
+    userID: "test-user",
+    userEmail: "test@example.com",
+    userRole: "Admin",
+    premiumUser: false,
+    proxySettings: {},
+    setProxySettings: vi.fn(),
+    accessToken: "test-token",
+    isPublicPage: false,
+  };
+
+  it("should render without crashing", () => {
+    renderWithProviders(<Navbar {...defaultProps} />);
+
+    expect(screen.getByText("Docs")).toBeInTheDocument();
+    expect(screen.getByText("User")).toBeInTheDocument();
+  });
+
+  it("should display user information in dropdown", async () => {
+    const user = userEvent.setup();
+    renderWithProviders(<Navbar {...defaultProps} />);
+
+    await user.click(screen.getByText("User"));
+
+    await waitFor(() => {
+      expect(screen.getByText("test-user")).toBeInTheDocument();
+    });
+    expect(screen.getByText("Admin")).toBeInTheDocument();
+    expect(screen.getByText("test@example.com")).toBeInTheDocument();
+  });
+
+  it("should show sidebar toggle button when onToggleSidebar is provided", () => {
+    const mockToggle = vi.fn();
+    renderWithProviders(<Navbar {...defaultProps} onToggleSidebar={mockToggle} />);
+
+    const toggleButton = screen.getByTitle("Collapse sidebar");
+    expect(toggleButton).toBeInTheDocument();
+  });
+
+  it("should call onToggleSidebar when sidebar button is clicked", async () => {
+    const mockToggle = vi.fn();
+    const user = userEvent.setup();
+    renderWithProviders(<Navbar {...defaultProps} onToggleSidebar={mockToggle} />);
+
+    const toggleButton = screen.getByTitle("Collapse sidebar");
+    await user.click(toggleButton);
+
+    expect(mockToggle).toHaveBeenCalledTimes(1);
+  });
+
+  it("should show premium user badge when premiumUser is true", async () => {
+    const user = userEvent.setup();
+    const premiumProps = { ...defaultProps, premiumUser: true };
+    renderWithProviders(<Navbar {...premiumProps} />);
+
+    await user.click(screen.getByText("User"));
+
+    await waitFor(() => {
+      expect(screen.getByText("Premium")).toBeInTheDocument();
+    });
+  });
+
+  it("should show version badge when health data contains version", () => {
+    mockUseHealthReadinessImpl = () => ({ data: { litellm_version: "1.0.0" } });
+
+    renderWithProviders(<Navbar {...defaultProps} />);
+
+    expect(screen.getByText("v1.0.0")).toBeInTheDocument();
+
+    // Reset mock
+    mockUseHealthReadinessImpl = () => ({ data: null });
+  });
+
+  it("should use custom logo from theme context", () => {
+    mockUseThemeImpl = () => ({ logoUrl: "https://example.com/custom-logo.png" });
+
+    renderWithProviders(<Navbar {...defaultProps} />);
+
+    const logoImg = screen.getByAltText("LiteLLM Brand");
+    expect(logoImg).toHaveAttribute("src", "https://example.com/custom-logo.png");
+
+    // Reset mock
+    mockUseThemeImpl = () => ({ logoUrl: null });
+  });
+
+  it("should hide user dropdown on public pages", () => {
+    const publicPageProps = { ...defaultProps, isPublicPage: true };
+    renderWithProviders(<Navbar {...publicPageProps} />);
+
+    expect(screen.queryByText("User")).not.toBeInTheDocument();
+  });
+
+  it("should handle hide new features toggle", async () => {
+    const user = userEvent.setup();
+
+    // Initially disabled
+    mockGetLocalStorageItemImpl = () => "false";
+
+    renderWithProviders(<Navbar {...defaultProps} />);
+
+    await user.click(screen.getByText("User"));
+
+    await waitFor(() => {
+      expect(screen.getByText("test-user")).toBeInTheDocument();
+    });
+
+    // Find and click the toggle switch
+    const toggleSwitch = screen.getByLabelText("Toggle hide new feature indicators");
+    await user.click(toggleSwitch);
+
+    // The functions are mocked globally, so we can check if they were called
+    // by accessing them through the mock registry
+    const localStorageUtils = vi.mocked(await import("@/utils/localStorageUtils"));
+    expect(localStorageUtils.setLocalStorageItem).toHaveBeenCalledWith("disableShowNewBadge", "true");
+    expect(localStorageUtils.emitLocalStorageChange).toHaveBeenCalledWith("disableShowNewBadge");
+  });
+
+  it("should handle logout functionality", async () => {
+    const user = userEvent.setup();
+
+    renderWithProviders(<Navbar {...defaultProps} />);
+
+    await user.click(screen.getByText("User"));
+
+    await waitFor(() => {
+      expect(screen.getByText("test-user")).toBeInTheDocument();
+    });
+
+    // Click logout
+    await user.click(screen.getByText("Logout"));
+
+    const cookieUtils = vi.mocked(await import("@/utils/cookieUtils"));
+    expect(cookieUtils.clearTokenCookies).toHaveBeenCalled();
+    expect(window.location.href).toBe("");
+  });
+});

ui/litellm-dashboard/src/components/navbar.tsx

@@ -1,22 +1,27 @@
-import Link from "next/link";
-import React, { useState, useEffect } from "react";
-import type { MenuProps } from "antd";
-import { Dropdown, Switch, Tooltip } from "antd";
+import { useHealthReadiness } from "@/app/(dashboard)/hooks/healthReadiness/useHealthReadiness";
 import { getProxyBaseUrl } from "@/components/networking";
+import { useTheme } from "@/contexts/ThemeContext";
+import { clearTokenCookies } from "@/utils/cookieUtils";
+import {
+  emitLocalStorageChange,
+  getLocalStorageItem,
+  removeLocalStorageItem,
+  setLocalStorageItem,
+} from "@/utils/localStorageUtils";
+import { fetchProxySettings } from "@/utils/proxyUtils";
 import {
-  UserOutlined,
-  LogoutOutlined,
   CrownOutlined,
+  LogoutOutlined,
   MailOutlined,
-  SafetyOutlined,
   MenuFoldOutlined,
   MenuUnfoldOutlined,
+  SafetyOutlined,
+  UserOutlined,
 } from "@ant-design/icons";
-import { clearTokenCookies } from "@/utils/cookieUtils";
-import { getLocalStorageItem, setLocalStorageItem, removeLocalStorageItem } from "@/utils/localStorageUtils";
-import { fetchProxySettings } from "@/utils/proxyUtils";
-import { useTheme } from "@/contexts/ThemeContext";
-import { emitLocalStorageChange } from "@/utils/localStorageUtils";
+import type { MenuProps } from "antd";
+import { Dropdown, Switch, Tooltip } from "antd";
+import Link from "next/link";
+import React, { useEffect, useState } from "react";
 
 interface NavbarProps {
   userID: string | null;
@@ -44,30 +49,16 @@ const Navbar: React.FC<NavbarProps> = ({
   onToggleSidebar,
 }) => {
   const baseUrl = getProxyBaseUrl();
+  console.log("baseUrl", baseUrl);
   const [logoutUrl, setLogoutUrl] = useState("");
-  const [version, setVersion] = useState("");
   const [disableShowNewBadge, setDisableShowNewBadge] = useState(false);
   const { logoUrl } = useTheme();
+  const { data: healthData } = useHealthReadiness();
+  const version = healthData?.litellm_version;
 
   // Simple logo URL: use custom logo if available, otherwise default
   const imageUrl = logoUrl || `${baseUrl}/get_image`;
 
-  useEffect(() => {
-    const fetchVersion = async () => {
-      try {
-        const response = await fetch(`${baseUrl}/health/readiness`);
-        const data = await response.json();
-        if (data.litellm_version) {
-          setVersion(data.litellm_version);
-        }
-      } catch (error) {
-        console.error("Failed to fetch version:", error);
-      }
-    };
-
-    fetchVersion();
-  }, [baseUrl]);
-
   useEffect(() => {
     const initializeProxySettings = async () => {
       if (accessToken) {
@@ -190,7 +181,7 @@ const Navbar: React.FC<NavbarProps> = ({
             )}
 
             <div className="flex items-center">
-              <Link href="/" className="flex items-center">
+              <Link href={baseUrl ? baseUrl : "/"} className="flex items-center">
                 <div className="relative">
                   <img src={imageUrl} alt="LiteLLM Brand" className="h-10 w-auto" />
                   <span

ui/litellm-dashboard/src/components/public_model_hub.test.tsx

@@ -27,6 +27,10 @@ vi.mock("./networking", async (importOriginal) => {
   };
 });
 
+vi.mock("./navbar", () => ({
+  default: vi.fn(() => <div data-testid="navbar">Navbar Component</div>),
+}));
+
 beforeAll(() => {
   Object.defineProperty(window, "matchMedia", {
     writable: true,