Merge pull request #24260 from BerriAI/litellm_yj_march_19_2026

[Infra] Merge internal dev branch with main

litellm/proxy/management_endpoints/mcp_management_endpoints.py

@@ -1214,17 +1214,9 @@ if MCP_AVAILABLE:
                     "error": "User does not have permission to create mcp servers. You can only create mcp servers if you are a PROXY_ADMIN."
                 },
             )
-        elif payload.server_id is not None:
-            # fail if the mcp server with id already exists
-            mcp_server = await get_mcp_server(prisma_client, payload.server_id)
-            if mcp_server is not None:
-                raise HTTPException(
-                    status_code=status.HTTP_400_BAD_REQUEST,
-                    detail={
-                        "error": f"MCP Server with id {payload.server_id} already exists. Cannot create another."
-                    },
-                )
-        elif (
+
+        # Block reserved special server IDs
+        if (
             SpecialMCPServerName.all_team_servers == payload.server_id
             or SpecialMCPServerName.all_proxy_servers == payload.server_id
         ):
@@ -1235,6 +1227,17 @@ if MCP_AVAILABLE:
                 },
             )
 
+        if payload.server_id is not None:
+            # fail if the mcp server with id already exists
+            mcp_server = await get_mcp_server(prisma_client, payload.server_id)
+            if mcp_server is not None:
+                raise HTTPException(
+                    status_code=status.HTTP_400_BAD_REQUEST,
+                    detail={
+                        "error": f"MCP Server with id {payload.server_id} already exists. Cannot create another."
+                    },
+                )
+
         # TODO: audit log for create
 
         # Admin-created servers are always active — clear any submission lifecycle

ui/litellm-dashboard/src/app/(dashboard)/api-reference/page.tsx

@@ -1,16 +1,10 @@
 "use client";
 
 import APIReferenceView from "@/app/(dashboard)/api-reference/APIReferenceView";
-import { useState } from "react";
-
-interface ProxySettings {
-  PROXY_BASE_URL: string;
-  PROXY_LOGOUT_URL: string;
-  LITELLM_UI_API_DOC_BASE_URL?: string | null;
-}
+import useProxySettings from "@/app/(dashboard)/hooks/proxySettings/useProxySettings";
 
 const APIReferencePage = () => {
-  const [proxySettings, setProxySettings] = useState<ProxySettings>({ PROXY_BASE_URL: "", PROXY_LOGOUT_URL: "" });
+  const proxySettings = useProxySettings();
 
   return <APIReferenceView proxySettings={proxySettings} />;
 };

ui/litellm-dashboard/src/app/(dashboard)/components/Sidebar2.tsx

@@ -195,7 +195,7 @@ const menuItems: MenuItemCfg[] = [
     icon: <UserOutlined style={{ fontSize: 18 }} />,
     roles: all_admin_roles,
   },
-  { key: "14", page: "api_ref", label: "API Reference", icon: <ApiOutlined style={{ fontSize: 18 }} /> },
+  { key: "14", page: "api-reference", label: "API Reference", icon: <ApiOutlined style={{ fontSize: 18 }} /> },
   {
     key: "16",
     page: "model-hub-table",

ui/litellm-dashboard/src/app/(dashboard)/hooks/common/queryKeysFactory.test.ts

@@ -0,0 +1,34 @@
+import { describe, it, expect } from "vitest";
+import { createQueryKeys } from "./queryKeysFactory";
+
+describe("createQueryKeys", () => {
+  const keys = createQueryKeys("books");
+
+  it("should return the resource name as the base key", () => {
+    expect(keys.all).toEqual(["books"]);
+  });
+
+  it("should generate a lists key", () => {
+    expect(keys.lists()).toEqual(["books", "list"]);
+  });
+
+  it("should generate a list key with params", () => {
+    expect(keys.list({ page: 1, limit: 10 })).toEqual([
+      "books",
+      "list",
+      { params: { page: 1, limit: 10 } },
+    ]);
+  });
+
+  it("should generate a list key with undefined params when none provided", () => {
+    expect(keys.list()).toEqual(["books", "list", { params: undefined }]);
+  });
+
+  it("should generate a details key", () => {
+    expect(keys.details()).toEqual(["books", "detail"]);
+  });
+
+  it("should generate a detail key for a specific ID", () => {
+    expect(keys.detail("123")).toEqual(["books", "detail", "123"]);
+  });
+});

ui/litellm-dashboard/src/app/(dashboard)/hooks/proxySettings/useProxySettings.ts

@@ -0,0 +1,21 @@
+import { useState, useEffect } from "react";
+import { fetchProxySettings } from "@/utils/proxyUtils";
+import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized";
+
+export default function useProxySettings() {
+  const { accessToken } = useAuthorized();
+  const [proxySettings, setProxySettings] = useState({
+    PROXY_BASE_URL: "",
+    PROXY_LOGOUT_URL: "",
+    LITELLM_UI_API_DOC_BASE_URL: null as string | null,
+  });
+
+  useEffect(() => {
+    if (!accessToken) return;
+    fetchProxySettings(accessToken).then((settings) => {
+      if (settings) setProxySettings(settings);
+    });
+  }, [accessToken]);
+
+  return proxySettings;
+}

ui/litellm-dashboard/src/app/(dashboard)/teams/components/TeamsHeaderTabs.test.tsx

@@ -0,0 +1,54 @@
+import { render, screen } from "@testing-library/react";
+import React from "react";
+import { describe, expect, it, vi } from "vitest";
+import TeamsHeaderTabs from "./TeamsHeaderTabs";
+
+vi.mock("@tremor/react", () => ({
+  TabGroup: ({ children, ...props }: any) => <div data-testid="tab-group" {...props}>{children}</div>,
+  TabList: ({ children, ...props }: any) => <div data-testid="tab-list" {...props}>{children}</div>,
+  Tab: ({ children, ...props }: any) => <button {...props}>{children}</button>,
+  TabPanels: ({ children, ...props }: any) => <div data-testid="tab-panels" {...props}>{children}</div>,
+  Text: ({ children, ...props }: any) => <span {...props}>{children}</span>,
+  Icon: ({ onClick, ...props }: any) => <button data-testid="refresh-icon" onClick={onClick} />,
+}));
+
+vi.mock("@heroicons/react/outline", () => ({
+  RefreshIcon: () => <svg data-testid="refresh-svg" />,
+}));
+
+const renderTabs = (props: Partial<Parameters<typeof TeamsHeaderTabs>[0]> = {}) => {
+  const defaults = {
+    lastRefreshed: "",
+    onRefresh: vi.fn(),
+    userRole: "Internal User",
+    children: <div data-testid="panel-content">Panel</div>,
+  };
+  return render(<TeamsHeaderTabs {...defaults} {...props} />);
+};
+
+describe("TeamsHeaderTabs", () => {
+  it("should render 'Your Teams' and 'Available Teams' tabs", () => {
+    renderTabs();
+
+    expect(screen.getByText("Your Teams")).toBeInTheDocument();
+    expect(screen.getByText("Available Teams")).toBeInTheDocument();
+  });
+
+  it("should render 'Default Team Settings' tab when user is Admin", () => {
+    renderTabs({ userRole: "Admin" });
+
+    expect(screen.getByText("Default Team Settings")).toBeInTheDocument();
+  });
+
+  it("should not render 'Default Team Settings' tab for non-admin users", () => {
+    renderTabs({ userRole: "Internal User" });
+
+    expect(screen.queryByText("Default Team Settings")).not.toBeInTheDocument();
+  });
+
+  it("should display last refreshed time when provided", () => {
+    renderTabs({ lastRefreshed: "2024-06-01 12:00:00" });
+
+    expect(screen.getByText("Last Refreshed: 2024-06-01 12:00:00")).toBeInTheDocument();
+  });
+});

ui/litellm-dashboard/src/app/(dashboard)/teams/components/TeamsTable/TeamsTable.test.tsx

@@ -0,0 +1,129 @@
+import { render, screen } from "@testing-library/react";
+import userEvent from "@testing-library/user-event";
+import React from "react";
+import { describe, expect, it, vi } from "vitest";
+import { Team } from "@/components/key_team_helpers/key_list";
+import TeamsTable from "./TeamsTable";
+
+vi.mock("@tremor/react", () => ({
+  Button: React.forwardRef<HTMLButtonElement, any>(({ children, ...props }, ref) =>
+    React.createElement("button", { ...props, ref }, children),
+  ),
+  Icon: ({ onClick, ...props }: any) => <button data-testid={props["data-testid"] || "icon-btn"} onClick={onClick} aria-label={props["aria-label"]} />,
+  Table: ({ children }: any) => <table>{children}</table>,
+  TableHead: ({ children }: any) => <thead>{children}</thead>,
+  TableBody: ({ children }: any) => <tbody>{children}</tbody>,
+  TableRow: ({ children }: any) => <tr>{children}</tr>,
+  TableHeaderCell: ({ children }: any) => <th>{children}</th>,
+  TableCell: ({ children, ...props }: any) => <td {...props}>{children}</td>,
+  Text: ({ children }: any) => <span>{children}</span>,
+}));
+
+vi.mock("antd", () => ({
+  Tooltip: ({ children }: any) => <>{children}</>,
+}));
+
+vi.mock("@heroicons/react/outline", () => ({
+  PencilAltIcon: () => <svg data-testid="pencil-icon" />,
+  TrashIcon: () => <svg data-testid="trash-icon" />,
+}));
+
+vi.mock("@/utils/dataUtils", () => ({
+  formatNumberWithCommas: (val: number, decimals: number) =>
+    val != null ? val.toFixed(decimals) : "N/A",
+}));
+
+vi.mock("@/app/(dashboard)/teams/components/TeamsTable/ModelsCell", () => ({
+  default: ({ team }: any) => <td data-testid="models-cell">{team.models.join(",")}</td>,
+}));
+
+vi.mock("@/app/(dashboard)/teams/components/TeamsTable/YourRoleCell/YourRoleCell", () => ({
+  default: ({ team }: any) => <td data-testid="role-cell">{team.team_id}</td>,
+}));
+
+const makeTeam = (overrides: Partial<Team> = {}): Team => ({
+  team_id: "team-abc1234",
+  team_alias: "Platform",
+  models: ["gpt-4"],
+  max_budget: 500,
+  budget_duration: null,
+  tpm_limit: null,
+  rpm_limit: null,
+  organization_id: "org-1",
+  created_at: "2024-06-01T00:00:00Z",
+  keys: [],
+  members_with_roles: [],
+  spend: 123.4567,
+  ...overrides,
+});
+
+const defaultPerTeamInfo = {
+  "team-abc1234": {
+    keys: [{ token: "tok-1" } as any, { token: "tok-2" } as any],
+    team_info: {
+      members_with_roles: [{ user_id: "u1", role: "admin" } as any],
+    },
+  },
+};
+
+const renderTable = (overrides: Partial<Parameters<typeof TeamsTable>[0]> = {}) => {
+  const defaults = {
+    teams: [makeTeam()],
+    currentOrg: null,
+    perTeamInfo: defaultPerTeamInfo,
+    userRole: "Admin",
+    userId: "user-1",
+    setSelectedTeamId: vi.fn(),
+    setEditTeam: vi.fn(),
+    onDeleteTeam: vi.fn(),
+  };
+  return render(<TeamsTable {...defaults} {...overrides} />);
+};
+
+describe("TeamsTable", () => {
+  it("should render table headers", () => {
+    renderTable();
+
+    expect(screen.getByText("Team Name")).toBeInTheDocument();
+    expect(screen.getByText("Team ID")).toBeInTheDocument();
+    expect(screen.getByText("Created")).toBeInTheDocument();
+    expect(screen.getByText("Spend (USD)")).toBeInTheDocument();
+    expect(screen.getByText("Budget (USD)")).toBeInTheDocument();
+    expect(screen.getByText("Models")).toBeInTheDocument();
+    expect(screen.getByText("Organization")).toBeInTheDocument();
+    expect(screen.getByText("Your Role")).toBeInTheDocument();
+    expect(screen.getByText("Info")).toBeInTheDocument();
+  });
+
+  it("should render team rows with team data", () => {
+    renderTable();
+
+    expect(screen.getByText("Platform")).toBeInTheDocument();
+    expect(screen.getByText("team-ab...")).toBeInTheDocument();
+    expect(screen.getByText("org-1")).toBeInTheDocument();
+  });
+
+  it("should show edit and delete icons for Admin users", () => {
+    renderTable({ userRole: "Admin" });
+
+    expect(screen.getAllByTestId("icon-btn").length).toBeGreaterThanOrEqual(2);
+  });
+
+  it("should not show edit and delete icons for non-Admin users", () => {
+    renderTable({ userRole: "Internal User" });
+
+    // Only the team ID button should be present, no icon-btn for edit/delete
+    const iconBtns = screen.queryAllByTestId("icon-btn");
+    expect(iconBtns).toHaveLength(0);
+  });
+
+  it("should call setSelectedTeamId when team ID button is clicked", async () => {
+    const user = userEvent.setup();
+    const setSelectedTeamId = vi.fn();
+    renderTable({ setSelectedTeamId });
+
+    await user.click(screen.getByText("team-ab..."));
+
+    expect(setSelectedTeamId).toHaveBeenCalledWith("team-abc1234");
+  });
+});

ui/litellm-dashboard/src/app/onboarding/OnboardingForm.test.tsx

@@ -0,0 +1,95 @@
+import { render, screen } from "@testing-library/react";
+import React from "react";
+import { describe, expect, it, vi } from "vitest";
+import { OnboardingForm } from "./OnboardingForm";
+
+const mockUseOnboardingCredentials = vi.fn();
+const mockClaimToken = vi.fn();
+
+vi.mock("next/navigation", () => ({
+  useSearchParams: () => new URLSearchParams("invitation_id=inv-123"),
+}));
+
+vi.mock("jwt-decode", () => ({
+  jwtDecode: vi.fn(() => ({
+    user_email: "alice@example.com",
+    user_id: "user-1",
+    key: "access-tok",
+  })),
+}));
+
+vi.mock("@/app/(dashboard)/hooks/onboarding/useOnboarding", () => ({
+  useOnboardingCredentials: (...args: unknown[]) => mockUseOnboardingCredentials(...args),
+  useClaimOnboardingToken: () => ({ mutate: mockClaimToken, isPending: false }),
+}));
+
+vi.mock("@/components/networking", () => ({
+  getProxyBaseUrl: vi.fn(() => ""),
+}));
+
+vi.mock("./OnboardingLoadingView", () => ({
+  OnboardingLoadingView: () => <div data-testid="loading-view">Loading</div>,
+}));
+
+vi.mock("./OnboardingErrorView", () => ({
+  OnboardingErrorView: () => <div data-testid="error-view">Error</div>,
+}));
+
+vi.mock("./OnboardingFormBody", () => ({
+  OnboardingFormBody: ({ variant, userEmail }: { variant: string; userEmail: string }) => (
+    <div data-testid="form-body" data-variant={variant} data-email={userEmail}>
+      Form Body
+    </div>
+  ),
+}));
+
+describe("OnboardingForm", () => {
+  it("should render loading view when credentials are loading", () => {
+    mockUseOnboardingCredentials.mockReturnValue({
+      data: undefined,
+      isLoading: true,
+      isError: false,
+    });
+
+    render(<OnboardingForm variant="signup" />);
+
+    expect(screen.getByTestId("loading-view")).toBeInTheDocument();
+  });
+
+  it("should render error view when credentials fail to load", () => {
+    mockUseOnboardingCredentials.mockReturnValue({
+      data: undefined,
+      isLoading: false,
+      isError: true,
+    });
+
+    render(<OnboardingForm variant="signup" />);
+
+    expect(screen.getByTestId("error-view")).toBeInTheDocument();
+  });
+
+  it("should render form body with decoded email when credentials are loaded", () => {
+    mockUseOnboardingCredentials.mockReturnValue({
+      data: { token: "fake-jwt-token" },
+      isLoading: false,
+      isError: false,
+    });
+
+    render(<OnboardingForm variant="signup" />);
+
+    expect(screen.getByTestId("form-body")).toBeInTheDocument();
+    expect(screen.getByTestId("form-body")).toHaveAttribute("data-email", "alice@example.com");
+  });
+
+  it("should pass variant prop to OnboardingFormBody", () => {
+    mockUseOnboardingCredentials.mockReturnValue({
+      data: { token: "fake-jwt-token" },
+      isLoading: false,
+      isError: false,
+    });
+
+    render(<OnboardingForm variant="reset_password" />);
+
+    expect(screen.getByTestId("form-body")).toHaveAttribute("data-variant", "reset_password");
+  });
+});

ui/litellm-dashboard/src/app/page.tsx

@@ -1,6 +1,5 @@
 "use client";
 
-import APIReferenceView from "@/app/(dashboard)/api-reference/APIReferenceView";
 import SidebarProvider from "@/app/(dashboard)/components/SidebarProvider";
 import OldModelDashboard from "@/app/(dashboard)/models-and-endpoints/ModelsAndEndpointsView";
 import PlaygroundPage from "@/app/(dashboard)/playground/page";
@@ -9,7 +8,7 @@ import AgentsPanel from "@/components/agents";
 import BudgetPanel from "@/components/budgets/budget_panel";
 import CacheDashboard from "@/components/cache_dashboard";
 import ClaudeCodePluginsPanel from "@/components/claude_code_plugins";
-import { fetchTeams } from "@/components/common_components/fetch_teams";
+import { teamListCall as v2TeamListCall } from "@/app/(dashboard)/hooks/teams/useTeams";
 import LoadingScreen from "@/components/common_components/LoadingScreen";
 import { CostTrackingSettings } from "@/components/CostTrackingSettings";
 import GeneralSettings from "@/components/general_settings";
@@ -48,7 +47,7 @@ import { buildLoginUrlWithReturn, consumeReturnUrl, normalizeUrlForCompare, stor
 import { formatUserRole, isAdminRole } from "@/utils/roles";
 import { QueryClient, QueryClientProvider } from "@tanstack/react-query";
 import { jwtDecode } from "jwt-decode";
-import { useSearchParams } from "next/navigation";
+import { useRouter, useSearchParams } from "next/navigation";
 import { Suspense, useEffect, useMemo, useRef, useState } from "react";
 import { ConfigProvider, theme } from "antd";
 
@@ -75,6 +74,16 @@ interface ProxySettings {
   LITELLM_UI_API_DOC_BASE_URL?: string | null;
 }
 
+/**
+ * Map of legacy query-param page keys → new path-based route segments.
+ * When a user visits ?page=<key>, they are redirected to /ui/<value>.
+ * Add entries here as pages are migrated from the if/else chain to path-based routes.
+ */
+const LEGACY_REDIRECTS: Record<string, string> = {
+  api_ref: "api-reference",
+  "api-reference": "api-reference",
+};
+
 function CreateKeyPageContent() {
   const [userRole, setUserRole] = useState("");
   const [premiumUser, setPremiumUser] = useState(false);
@@ -90,6 +99,7 @@ function CreateKeyPageContent() {
   });
 
   const [showSSOBanner, setShowSSOBanner] = useState<boolean>(true);
+  const router = useRouter();
   const searchParams = useSearchParams()!;
   const [modelData, setModelData] = useState<any>({ data: [] });
   const [token, setToken] = useState<string | null>(null);
@@ -243,6 +253,15 @@ function CreateKeyPageContent() {
     }
   }, [redirectToLogin]);
 
+  // Redirect legacy query-param pages to their new path-based routes
+  const isLegacyRedirect = page in LEGACY_REDIRECTS;
+  useEffect(() => {
+    if (!authLoading && isLegacyRedirect) {
+      const base = (proxyBaseUrl || "") + "/ui";
+      router.replace(`${base}/${LEGACY_REDIRECTS[page]}`);
+    }
+  }, [authLoading, isLegacyRedirect, page, router]);
+
   // Check for a stored return URL after successful authentication
   // This handles the case where user comes back from SSO and we need to redirect to the original URL
   useEffect(() => {
@@ -339,7 +358,9 @@ function CreateKeyPageContent() {
       fetchUserModels(userID, userRole, accessToken, setUserModels);
     }
     if (accessToken && userID && userRole) {
-      fetchTeams(accessToken, userID, userRole, null, setTeams);
+      v2TeamListCall(accessToken, 1, 100, {
+        userID: userRole !== "Admin" && userRole !== "Admin Viewer" ? userID : null,
+      }).then((response) => setTeams(response.teams ?? [])).catch(console.error);
     }
     if (accessToken) {
       fetchOrganizations(accessToken, setOrganizations);
@@ -427,7 +448,7 @@ function CreateKeyPageContent() {
     setShowClaudeCodePrompt(true);
   };
 
-  if (authLoading || redirectToLogin) {
+  if (authLoading || redirectToLogin || isLegacyRedirect) {
     return <LoadingScreen />;
   }
 
@@ -536,8 +557,6 @@ function CreateKeyPageContent() {
                     <AdminPanel
                       proxySettings={proxySettings}
                     />
-                  ) : page == "api_ref" ? (
-                    <APIReferenceView proxySettings={proxySettings} />
                   ) : page == "logging-and-alerts" ? (
                     <Settings userID={userID} userRole={userRole} accessToken={accessToken} premiumUser={premiumUser} />
                   ) : page == "budgets" ? (

ui/litellm-dashboard/src/components/AccessGroups/AccessGroupsModal/AccessGroupCreateModal.tsx

@@ -1,5 +1,6 @@
 import React from "react";
-import { Modal, Form, message } from "antd";
+import { Modal, Form } from "antd";
+import MessageManager from "@/components/molecules/message_manager";
 import {
   AccessGroupBaseForm,
   AccessGroupFormValues,
@@ -37,7 +38,7 @@ export function AccessGroupCreateModal({
 
         createMutation.mutate(params, {
           onSuccess: () => {
-            message.success("Access group created successfully");
+            MessageManager.success("Access group created successfully");
             form.resetFields();
             onSuccess?.();
             onCancel();

ui/litellm-dashboard/src/components/AccessGroups/AccessGroupsModal/AccessGroupEditModal.tsx

@@ -1,5 +1,6 @@
 import React, { useEffect } from "react";
-import { Modal, Form, message } from "antd";
+import { Modal, Form } from "antd";
+import MessageManager from "@/components/molecules/message_manager";
 import {
   AccessGroupBaseForm,
   AccessGroupFormValues,
@@ -55,7 +56,7 @@ export function AccessGroupEditModal({
           { accessGroupId: accessGroup.access_group_id, params },
           {
             onSuccess: () => {
-              message.success("Access group updated successfully");
+              MessageManager.success("Access group updated successfully");
               onSuccess?.();
               onCancel();
             },

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

@@ -3,7 +3,6 @@ import {
   Modal,
   Typography,
   Divider,
-  message,
   Table,
   Select,
   InputNumber,
@@ -14,6 +13,7 @@ import {
 import { userBulkUpdateUserCall, teamBulkMemberAddCall, Member } from "./networking";
 import { UserEditView } from "./user_edit_view";
 import NotificationsManager from "./molecules/notifications_manager";
+import MessageManager from "@/components/molecules/message_manager";
 
 const { Text, Title } = Typography;
 
@@ -188,7 +188,7 @@ const BulkEditUserModal: React.FC<BulkEditUserModalProps> = ({
         }
 
         if (failedTeams.length > 0) {
-          message.warning(`Failed to add users to ${failedTeams.length} team(s)`);
+          MessageManager.warning(`Failed to add users to ${failedTeams.length} team(s)`);
         }
       }
 

ui/litellm-dashboard/src/components/CloudZeroCostTracking/CloudZeroCreateModal.tsx

@@ -1,4 +1,5 @@
-import { Form, Modal, Input, message } from "antd";
+import { Form, Modal, Input } from "antd";
+import MessageManager from "@/components/molecules/message_manager";
 import { useEffect } from "react";
 import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized";
 import { useCloudZeroCreate } from "@/app/(dashboard)/hooks/cloudzero/useCloudZeroCreate";
@@ -31,7 +32,7 @@ export default function CloudZeroCreationModal({ open, onOk, onCancel }: CloudZe
         },
         {
           onSuccess: () => {
-            message.success("CloudZero integration created successfully");
+            MessageManager.success("CloudZero integration created successfully");
             form.resetFields();
             onOk();
           },
@@ -39,7 +40,7 @@ export default function CloudZeroCreationModal({ open, onOk, onCancel }: CloudZe
             if (error?.errorFields) {
               return;
             }
-            message.error(error?.message || "Failed to create CloudZero integration");
+            MessageManager.error(error?.message || "Failed to create CloudZero integration");
           },
         },
       );
@@ -47,7 +48,7 @@ export default function CloudZeroCreationModal({ open, onOk, onCancel }: CloudZe
       if (error?.errorFields) {
         return;
       }
-      message.error(error?.message || "Failed to create CloudZero integration");
+      MessageManager.error(error?.message || "Failed to create CloudZero integration");
     }
   };
 

ui/litellm-dashboard/src/components/CloudZeroCostTracking/CloudZeroIntegrationSettings.tsx

@@ -3,7 +3,8 @@ import { useCloudZeroExport } from "@/app/(dashboard)/hooks/cloudzero/useCloudZe
 import { useCloudZeroDeleteSettings } from "@/app/(dashboard)/hooks/cloudzero/useCloudZeroSettings";
 import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized";
 import DeleteResourceModal from "@/components/common_components/DeleteResourceModal";
-import { Alert, Button, Card, Descriptions, Divider, message, Popconfirm, Tag } from "antd";
+import { Alert, Button, Card, Descriptions, Divider, Popconfirm, Tag } from "antd";
+import MessageManager from "@/components/molecules/message_manager";
 import { CheckCircle, Edit, Play, Trash2, Upload } from "lucide-react";
 import { useState } from "react";
 import CloudZeroUpdateModal from "./CloudZeroUpdateModal";
@@ -30,10 +31,10 @@ export function CloudZeroIntegrationSettings({ settings, onSettingsUpdated }: Cl
       { limit: 10 },
       {
         onSuccess: (data) => {
-          message.success("Dry run completed successfully");
+          MessageManager.success("Dry run completed successfully");
         },
         onError: (error) => {
-          message.error(error?.message || "Failed to perform dry run");
+          MessageManager.error(error?.message || "Failed to perform dry run");
         },
       },
     );
@@ -48,10 +49,10 @@ export function CloudZeroIntegrationSettings({ settings, onSettingsUpdated }: Cl
       { operation: "replace_hourly" },
       {
         onSuccess: () => {
-          message.success("Data successfully exported to CloudZero");
+          MessageManager.success("Data successfully exported to CloudZero");
         },
         onError: (error) => {
-          message.error(error?.message || "Failed to export data");
+          MessageManager.error(error?.message || "Failed to export data");
         },
       },
     );
@@ -79,12 +80,12 @@ export function CloudZeroIntegrationSettings({ settings, onSettingsUpdated }: Cl
 
     deleteMutation.mutate(undefined, {
       onSuccess: () => {
-        message.success("CloudZero integration deleted successfully");
+        MessageManager.success("CloudZero integration deleted successfully");
         setIsDeleteModalOpen(false);
         onSettingsUpdated();
       },
       onError: (error) => {
-        message.error(error?.message || "Failed to delete CloudZero integration");
+        MessageManager.error(error?.message || "Failed to delete CloudZero integration");
       },
     });
   };

ui/litellm-dashboard/src/components/CloudZeroCostTracking/CloudZeroUpdateModal.tsx

@@ -1,6 +1,7 @@
 import { useCloudZeroUpdateSettings } from "@/app/(dashboard)/hooks/cloudzero/useCloudZeroSettings";
 import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized";
-import { Form, Input, message, Modal } from "antd";
+import { Form, Input, Modal } from "antd";
+import MessageManager from "@/components/molecules/message_manager";
 import { useEffect } from "react";
 import { CloudZeroSettings } from "./types";
 
@@ -39,7 +40,7 @@ export default function CloudZeroUpdateModal({ open, onOk, onCancel, settings }:
         },
         {
           onSuccess: () => {
-            message.success("CloudZero integration updated successfully");
+            MessageManager.success("CloudZero integration updated successfully");
             form.resetFields();
             onOk();
           },
@@ -47,7 +48,7 @@ export default function CloudZeroUpdateModal({ open, onOk, onCancel, settings }:
             if (error?.errorFields) {
               return;
             }
-            message.error(error?.message || "Failed to update CloudZero integration");
+            MessageManager.error(error?.message || "Failed to update CloudZero integration");
           },
         },
       );
@@ -55,7 +56,7 @@ export default function CloudZeroUpdateModal({ open, onOk, onCancel, settings }:
       if (error?.errorFields) {
         return;
       }
-      message.error(error?.message || "Failed to update CloudZero integration");
+      MessageManager.error(error?.message || "Failed to update CloudZero integration");
     }
   };
 

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

@@ -18,8 +18,8 @@ vi.mock("./networking", () => ({
   getPoliciesList: vi.fn().mockResolvedValue({ policies: [] }),
 }));
 
-vi.mock("./common_components/fetch_teams", () => ({
-  fetchTeams: vi.fn(),
+vi.mock("@/app/(dashboard)/hooks/teams/useTeams", () => ({
+  teamListCall: vi.fn().mockResolvedValue({ teams: [], total: 0, page: 1, page_size: 100, total_pages: 0 }),
 }));
 
 vi.mock("./molecules/notifications_manager", () => ({
@@ -375,6 +375,9 @@ describe("OldTeams - handleCreate organization handling", () => {
         organizations={[]}
       />,
     );
+    await waitFor(() => {
+      expect(screen.getByTestId("delete-team-button")).toBeInTheDocument();
+    });
     const deleteTeamButton = screen.getByTestId("delete-team-button");
     act(() => {
       fireEvent.click(deleteTeamButton);
@@ -389,7 +392,7 @@ describe("OldTeams - empty state", () => {
     mockUseOrganizations.mockReturnValue({ data: [] });
   });
 
-  it("should display empty state message when teams array is empty", () => {
+  it("should display empty state message when teams array is empty", async () => {
     renderWithQueryClient(
       <OldTeams
         teams={[]}
@@ -402,11 +405,13 @@ describe("OldTeams - empty state", () => {
       />,
     );
 
-    expect(screen.getByText("No teams found")).toBeInTheDocument();
-    expect(screen.getByText("Adjust your filters or create a new team")).toBeInTheDocument();
+    await waitFor(() => {
+      expect(screen.getByText("No teams yet")).toBeInTheDocument();
+    });
+    expect(screen.getByText("Create your first team to organize members and manage access to models.")).toBeInTheDocument();
   });
 
-  it("should display empty state message when teams is null", () => {
+  it("should display empty state message when teams is null", async () => {
     renderWithQueryClient(
       <OldTeams
         teams={null}
@@ -419,11 +424,13 @@ describe("OldTeams - empty state", () => {
       />,
     );
 
-    expect(screen.getByText("No teams found")).toBeInTheDocument();
-    expect(screen.getByText("Adjust your filters or create a new team")).toBeInTheDocument();
+    await waitFor(() => {
+      expect(screen.getByText("No teams yet")).toBeInTheDocument();
+    });
+    expect(screen.getByText("Create your first team to organize members and manage access to models.")).toBeInTheDocument();
   });
 
-  it("should not display empty state when teams array has items", () => {
+  it("should not display empty state when teams array has items", async () => {
     renderWithQueryClient(
       <OldTeams
         teams={[
@@ -451,9 +458,11 @@ describe("OldTeams - empty state", () => {
       />,
     );
 
-    expect(screen.queryByText("No teams found")).not.toBeInTheDocument();
-    expect(screen.queryByText("Adjust your filters or create a new team")).not.toBeInTheDocument();
-    expect(screen.getByText("Test Team")).toBeInTheDocument();
+    await waitFor(() => {
+      expect(screen.getByText("Test Team")).toBeInTheDocument();
+    });
+    expect(screen.queryByText("No teams yet")).not.toBeInTheDocument();
+    expect(screen.queryByText("Create your first team to organize members and manage access to models.")).not.toBeInTheDocument();
   });
 });
 
@@ -621,12 +630,9 @@ describe("OldTeams - premium props", () => {
       />,
     );
 
-    const truncatedTeamId = "team-123456789".slice(0, 7);
-    const teamButton = await screen.findByRole("button", {
-      name: new RegExp(`${truncatedTeamId}\\.\\.\\.`),
-    });
+    const teamIdElement = await screen.findByText("team-123456789");
     act(() => {
-      fireEvent.click(teamButton);
+      fireEvent.click(teamIdElement);
     });
 
     await waitFor(() => expect(mockTeamInfoView).toHaveBeenCalled());
@@ -798,7 +804,7 @@ describe("OldTeams - access_group_ids in team create", () => {
       />,
     );
 
-    const createButton = screen.getByRole("button", { name: /create new team/i });
+    const createButton = screen.getAllByRole("button", { name: /create team/i })[0];
     act(() => {
       fireEvent.click(createButton);
     });
@@ -823,7 +829,8 @@ describe("OldTeams - access_group_ids in team create", () => {
     const accessGroupInput = screen.getByTestId("access-group-selector");
     fireEvent.change(accessGroupInput, { target: { value: "ag-1,ag-2" } });
 
-    const createTeamSubmitButton = screen.getByRole("button", { name: /create team/i });
+    const createTeamSubmitButtons = screen.getAllByRole("button", { name: /create team/i });
+    const createTeamSubmitButton = createTeamSubmitButtons[createTeamSubmitButtons.length - 1];
     fireEvent.click(createTeamSubmitButton);
 
     await waitFor(() => {
@@ -865,7 +872,7 @@ describe("OldTeams - models dropdown options", () => {
       expect(fetchAvailableModelsForTeamOrKey).toHaveBeenCalled();
     });
 
-    const createButton = screen.getByRole("button", { name: /create new team/i });
+    const createButton = screen.getAllByRole("button", { name: /create team/i })[0];
     act(() => {
       fireEvent.click(createButton);
     });
@@ -884,7 +891,7 @@ describe("OldTeams - organization alias display", () => {
     mockUseOrganizations.mockReturnValue({ data: [] });
   });
 
-  it("should display organization alias instead of organization id", () => {
+  it("should display organization alias instead of organization id", async () => {
     const mockOrganizations = [
       {
         organization_id: "org-123",
@@ -934,11 +941,13 @@ describe("OldTeams - organization alias display", () => {
       />,
     );
 
-    expect(screen.getByText("Test Organization")).toBeInTheDocument();
+    await waitFor(() => {
+      expect(screen.getByText("Test Organization")).toBeInTheDocument();
+    });
     expect(screen.queryByText("org-123")).not.toBeInTheDocument();
   });
 
-  it("should display organization id when alias is not found", () => {
+  it("should display organization id when alias is not found", async () => {
     mockUseOrganizations.mockReturnValue({ data: [] });
 
     renderWithQueryClient(
@@ -968,10 +977,12 @@ describe("OldTeams - organization alias display", () => {
       />,
     );
 
-    expect(screen.getByText("org-unknown")).toBeInTheDocument();
+    await waitFor(() => {
+      expect(screen.getByText("org-unknown")).toBeInTheDocument();
+    });
   });
 
-  it("should display N/A when organization_id is null", () => {
+  it("should display N/A when organization_id is null", async () => {
     mockUseOrganizations.mockReturnValue({ data: [] });
 
     renderWithQueryClient(
@@ -1001,6 +1012,9 @@ describe("OldTeams - organization alias display", () => {
       />,
     );
 
-    expect(screen.getByText("N/A")).toBeInTheDocument();
+    await waitFor(() => {
+      // When organization_id is null, the table shows "—" in the Organization column
+      expect(screen.getAllByText("—").length).toBeGreaterThan(0);
+    });
   });
 });

ui/litellm-dashboard/src/components/Projects/ProjectModals/CreateProjectModal.tsx

@@ -1,5 +1,6 @@
-import { Modal, Form, Button, Typography, message } from "antd";
+import { Modal, Form, Button, Typography } from "antd";
 import { FolderAddOutlined } from "@ant-design/icons";
+import MessageManager from "@/components/molecules/message_manager";
 import {
   useCreateProject,
   ProjectCreateParams,
@@ -32,12 +33,12 @@ export function CreateProjectModal({
 
       createMutation.mutate(params, {
         onSuccess: () => {
-          message.success("Project created successfully");
+          MessageManager.success("Project created successfully");
           form.resetFields();
           onClose();
         },
         onError: (error) => {
-          message.error(error.message || "Failed to create project");
+          MessageManager.error(error.message || "Failed to create project");
         },
       });
     } catch (error) {

ui/litellm-dashboard/src/components/Projects/ProjectModals/EditProjectModal.tsx

@@ -1,6 +1,7 @@
 import { useEffect } from "react";
-import { Modal, Form, Button, Typography, message } from "antd";
+import { Modal, Form, Button, Typography } from "antd";
 import { SaveOutlined } from "@ant-design/icons";
+import MessageManager from "@/components/molecules/message_manager";
 import { ProjectResponse } from "@/app/(dashboard)/hooks/projects/useProjects";
 import {
   useUpdateProject,
@@ -80,12 +81,12 @@ export function EditProjectModal({
         { projectId: project.project_id, params },
         {
           onSuccess: () => {
-            message.success("Project updated successfully");
+            MessageManager.success("Project updated successfully");
             onSuccess?.();
             onClose();
           },
           onError: (error) => {
-            message.error(error.message || "Failed to update project");
+            MessageManager.error(error.message || "Failed to update project");
           },
         },
       );

ui/litellm-dashboard/src/components/SearchTools/SearchToolTester.tsx

@@ -1,5 +1,6 @@
 import React, { useState } from "react";
-import { Button, Input, Typography, Spin, message } from "antd";
+import { Button, Input, Typography, Spin } from "antd";
+import MessageManager from "@/components/molecules/message_manager";
 import { SearchOutlined, LoadingOutlined } from "@ant-design/icons";
 import { searchToolQueryCall } from "../networking";
 import NotificationsManager from "../molecules/notifications_manager";
@@ -39,7 +40,7 @@ export const SearchToolTester: React.FC<SearchToolTesterProps> = ({ searchToolNa
 
   const handleSearch = async () => {
     if (!query.trim()) {
-      message.warning("Please enter a search query");
+      MessageManager.warning("Please enter a search query");
       return;
     }
 

ui/litellm-dashboard/src/components/Settings/RouterSettings/Fallbacks/AddFallbacks.tsx

@@ -5,8 +5,9 @@
  */
 
 import { Button as TremorButton } from "@tremor/react";
-import { Button, message } from "antd";
+import { Button } from "antd";
 import React, { useEffect, useState } from "react";
+import MessageManager from "@/components/molecules/message_manager";
 import NotificationManager from "../../../molecules/notifications_manager";
 import { fetchAvailableModels, ModelGroup } from "../../../playground/llm_calls/fetch_models";
 import { AddFallbacksModal } from "./AddFallbacksModal";
@@ -90,7 +91,7 @@ export default function AddFallbacks({
       (g) => !g.primaryModel || g.fallbackModels.length === 0,
     );
     if (invalidGroups.length > 0) {
-      message.error(
+      MessageManager.error(
         `Please complete configuration for all groups. ${invalidGroups.length} group(s) incomplete.`,
       );
       return;

ui/litellm-dashboard/src/components/Settings/RouterSettings/Fallbacks/FallbackSelectionForm.tsx

@@ -5,9 +5,10 @@
  */
 
 import { Button } from "@tremor/react";
-import { message, Tabs } from "antd";
+import { Tabs } from "antd";
 import { Plus } from "lucide-react";
 import React, { useEffect, useState } from "react";
+import MessageManager from "@/components/molecules/message_manager";
 import { FallbackGroup, FallbackGroupConfig } from "./FallbackGroupConfig";
 
 interface FallbackSelectionFormProps {
@@ -60,7 +61,7 @@ export function FallbackSelectionForm({
 
   const handleRemoveGroup = (targetId: string) => {
     if (groups.length === 1) {
-      message.warning("At least one group is required");
+      MessageManager.warning("At least one group is required");
       return;
     }
     const newGroups = groups.filter((g) => g.id !== targetId);

ui/litellm-dashboard/src/components/agents/add_agent_form.tsx

@@ -1,5 +1,6 @@
 import React, { useState, useEffect } from "react";
-import { Modal, Form, message, Select, Input, Steps, Radio, Tag, Divider, Switch, InputNumber, Collapse } from "antd";
+import { Modal, Form, Select, Input, Steps, Radio, Tag, Divider, Switch, InputNumber, Collapse } from "antd";
+import MessageManager from "@/components/molecules/message_manager";
 import { Button } from "@tremor/react";
 import { CheckCircleFilled, KeyOutlined, RobotOutlined, AppstoreOutlined, InfoCircleOutlined } from "@ant-design/icons";
 import CreatedKeyDisplay from "../shared/CreatedKeyDisplay";
@@ -216,7 +217,7 @@ const AddAgentForm: React.FC<AddAgentFormProps> = ({
 
   const handleCreateAgent = async () => {
     if (!accessToken) {
-      message.error("No access token available");
+      MessageManager.error("No access token available");
       return;
     }
 
@@ -226,7 +227,7 @@ const AddAgentForm: React.FC<AddAgentFormProps> = ({
       const values = { ...form.getFieldsValue(true) };
       const agentData = buildAgentData(values);
       if (!agentData) {
-        message.error("Failed to build agent data");
+        MessageManager.error("Failed to build agent data");
         setIsSubmitting(false);
         return;
       }
@@ -301,7 +302,7 @@ const AddAgentForm: React.FC<AddAgentFormProps> = ({
         setCreatedKeyValue(keyResponse.key || null);
       } else if (keyAssignOption === "existing_key") {
         if (!selectedExistingKey) {
-          message.error("Please select an existing key to assign");
+          MessageManager.error("Please select an existing key to assign");
           setIsSubmitting(false);
           return;
         }
@@ -318,7 +319,7 @@ const AddAgentForm: React.FC<AddAgentFormProps> = ({
     } catch (error) {
       console.error("Error creating agent:", error);
       const errorMessage = error instanceof Error ? error.message : String(error);
-      message.error(errorMessage ? `Failed to create agent: ${errorMessage}` : "Failed to create agent");
+      MessageManager.error(errorMessage ? `Failed to create agent: ${errorMessage}` : "Failed to create agent");
     } finally {
       setIsSubmitting(false);
     }

ui/litellm-dashboard/src/components/agents/agent_info.tsx

@@ -1,6 +1,7 @@
 import React, { useState, useEffect } from "react";
 import { Card, Title, Text, Button as TremorButton, Tab, TabGroup, TabList, TabPanel, TabPanels} from "@tremor/react";
-import { Form, Input, InputNumber, Button as AntButton, message, Spin, Descriptions, Divider } from "antd";
+import { Form, Input, InputNumber, Button as AntButton, Spin, Descriptions, Divider } from "antd";
+import MessageManager from "@/components/molecules/message_manager";
 import { ArrowLeftIcon } from "@heroicons/react/outline";
 import { getAgentInfo, patchAgentCall, getAgentCreateMetadata, AgentCreateInfo } from "../networking";
 import { Agent } from "./types";
@@ -72,7 +73,7 @@ const AgentInfoView: React.FC<AgentInfoViewProps> = ({
       }
     } catch (error) {
       console.error("Error fetching agent info:", error);
-      message.error("Failed to load agent information");
+      MessageManager.error("Failed to load agent information");
     } finally {
       setIsLoading(false);
     }
@@ -111,12 +112,12 @@ const AgentInfoView: React.FC<AgentInfoViewProps> = ({
       }
       
       await patchAgentCall(accessToken, agentId, updateData);
-      message.success("Agent updated successfully");
+      MessageManager.success("Agent updated successfully");
       setIsEditing(false);
       fetchAgentInfo();
     } catch (error) {
       console.error("Error updating agent:", error);
-      message.error("Failed to update agent");
+      MessageManager.error("Failed to update agent");
     } finally {
       setIsSaving(false);
     }

ui/litellm-dashboard/src/components/chat/ChatPage.tsx

@@ -1,7 +1,8 @@
 "use client";
 
 import React, { useCallback, useEffect, useRef, useState, useLayoutEffect } from "react";
-import { Tooltip, Skeleton, Popover, message } from "antd";
+import { Tooltip, Skeleton, Popover } from "antd";
+import MessageManager from "@/components/molecules/message_manager";
 import {
   SettingOutlined,
   PlusOutlined,
@@ -212,7 +213,7 @@ const ChatPage: React.FC<ChatPageProps> = ({ accessToken, userRole, userId, user
           localStorage.setItem(LOCALSTORAGE_MODEL_KEY, JSON.stringify([names[0]]));
         }
       })
-      .catch(() => message.error("Could not load models"))
+      .catch(() => MessageManager.error("Could not load models"))
       .finally(() => setIsLoadingModels(false));
   }, [accessToken]);
 

ui/litellm-dashboard/src/components/chat/MCPAppsPanel.tsx

@@ -5,7 +5,7 @@ import { Spin, Input, Button, Skeleton } from "antd";
 import { SearchOutlined, ArrowLeftOutlined, RightOutlined, ToolOutlined, CheckCircleOutlined } from "@ant-design/icons";
 import { deleteMCPOAuthUserCredential, fetchMCPServers, getMCPOAuthUserCredentialStatus, listMCPTools } from "../networking";
 import { AUTH_TYPE, MCPServer, MCPTool, handleTransport } from "../mcp_tools/types";
-import { message } from "antd";
+import MessageManager from "@/components/molecules/message_manager";
 import { useUserMcpOAuthFlow } from "@/hooks/useUserMcpOAuthFlow";
 
 // ── OAuth2 connect button ─────────────────────────────────────────────────────
@@ -198,7 +198,7 @@ const MCPAppsPanel: React.FC<Props> = ({ accessToken, selectedServers, onChange
       const idToFetch = serverId ?? serverName;
       const result = await listMCPTools(accessToken, idToFetch);
       if (result?.error) {
-        message.warning(`Could not load tools for ${serverName}`);
+        MessageManager.warning(`Could not load tools for ${serverName}`);
         return;
       }
       // Use the ref so we read the most up-to-date list; guard against duplicates
@@ -207,7 +207,7 @@ const MCPAppsPanel: React.FC<Props> = ({ accessToken, selectedServers, onChange
         onChange([...selectedServersRef.current, serverName]);
       }
     } catch {
-      message.warning(`Could not load tools for ${serverName}`);
+      MessageManager.warning(`Could not load tools for ${serverName}`);
     } finally {
       setTogglingOn((prev) => {
         const next = new Set(prev);

ui/litellm-dashboard/src/components/chat/MCPConnectPicker.tsx

@@ -1,5 +1,6 @@
 import React, { useEffect, useState } from "react";
-import { Switch, Spin, message } from "antd";
+import { Switch, Spin } from "antd";
+import MessageManager from "@/components/molecules/message_manager";
 import { fetchMCPServers, listMCPTools } from "../networking";
 import { MCPServer } from "../mcp_tools/types";
 
@@ -57,7 +58,7 @@ const MCPConnectPicker: React.FC<Props> = ({ accessToken, selectedServers, onCha
       const result = await listMCPTools(accessToken, serverName);
       // listMCPTools never throws; it returns { tools, error, message } on failure
       if (result?.error) {
-        message.warning(
+        MessageManager.warning(
           `Could not load tools for ${serverName} — it will be excluded from this message.`
         );
         // Do not add to selectedServers
@@ -65,7 +66,7 @@ const MCPConnectPicker: React.FC<Props> = ({ accessToken, selectedServers, onCha
       }
       onChange([...selectedServers, serverName]);
     } catch {
-      message.warning(
+      MessageManager.warning(
         `Could not load tools for ${serverName} — it will be excluded from this message.`
       );
       // Do not add to selectedServers

ui/litellm-dashboard/src/components/chat/MCPCredentialsTab.tsx

@@ -8,7 +8,8 @@
  */
 
 import React, { useCallback, useEffect, useState } from "react";
-import { Spin, message } from "antd";
+import { Spin } from "antd";
+import MessageManager from "@/components/molecules/message_manager";
 import { DeleteOutlined, LinkOutlined } from "@ant-design/icons";
 import { Badge, Table, TableBody, TableCell, TableHead, TableHeaderCell, TableRow } from "@tremor/react";
 import {
@@ -77,7 +78,7 @@ const MCPCredentialsTab: React.FC<Props> = ({ accessToken }) => {
       await deleteMCPOAuthUserCredential(accessToken, serverId);
       setCredentials((prev) => prev.filter((c) => c.server_id !== serverId));
     } catch {
-      message.error("Failed to revoke connection. Please try again.");
+      MessageManager.error("Failed to revoke connection. Please try again.");
     } finally {
       setRevoking((prev) => { const n = new Set(prev); n.delete(serverId); return n; });
     }

ui/litellm-dashboard/src/components/claude_code_plugins/add_plugin_form.tsx

@@ -1,5 +1,6 @@
 import React, { useState } from "react";
-import { Modal, Form, Input, Select, message } from "antd";
+import { Modal, Form, Input, Select } from "antd";
+import MessageManager from "@/components/molecules/message_manager";
 import { Button } from "@tremor/react";
 import { registerClaudeCodePlugin } from "../networking";
 import {
@@ -43,13 +44,13 @@ const AddPluginForm: React.FC<AddPluginFormProps> = ({
 
   const handleSubmit = async (values: any) => {
     if (!accessToken) {
-      message.error("No access token available");
+      MessageManager.error("No access token available");
       return;
     }
 
     // Validate plugin name
     if (!validatePluginName(values.name)) {
-      message.error(
+      MessageManager.error(
         "Plugin name must be kebab-case (lowercase letters, numbers, and hyphens only)"
       );
       return;
@@ -57,7 +58,7 @@ const AddPluginForm: React.FC<AddPluginFormProps> = ({
 
     // Validate semantic version if provided
     if (values.version && !isValidSemanticVersion(values.version)) {
-      message.error(
+      MessageManager.error(
         "Version must be in semantic versioning format (e.g., 1.0.0)"
       );
       return;
@@ -65,13 +66,13 @@ const AddPluginForm: React.FC<AddPluginFormProps> = ({
 
     // Validate email if provided
     if (values.authorEmail && !isValidEmail(values.authorEmail)) {
-      message.error("Invalid email format");
+      MessageManager.error("Invalid email format");
       return;
     }
 
     // Validate homepage URL if provided
     if (values.homepage && !isValidUrl(values.homepage)) {
-      message.error("Invalid homepage URL format");
+      MessageManager.error("Invalid homepage URL format");
       return;
     }
 
@@ -119,14 +120,14 @@ const AddPluginForm: React.FC<AddPluginFormProps> = ({
       }
 
       await registerClaudeCodePlugin(accessToken, pluginData);
-      message.success("Plugin registered successfully");
+      MessageManager.success("Plugin registered successfully");
       form.resetFields();
       setSourceType("github");
       onSuccess();
       onClose();
     } catch (error) {
       console.error("Error registering plugin:", error);
-      message.error("Failed to register plugin");
+      MessageManager.error("Failed to register plugin");
     } finally {
       setIsSubmitting(false);
     }

ui/litellm-dashboard/src/components/common_components/IconActionButton/TableIconActionButtons/TableIconActionButton.tsx

@@ -6,6 +6,7 @@ import {
   ChevronUpIcon,
   ChevronDownIcon,
   ExternalLinkIcon,
+  ClipboardCopyIcon,
 } from "@heroicons/react/outline";
 import { Tooltip } from "antd";
 import BaseActionButton from "../BaseActionButton";
@@ -32,6 +33,7 @@ export const TableIconActionButtonMap: Record<string, TableIconActionButtonBaseP
   Up: { icon: ChevronUpIcon, className: "hover:text-blue-600" },
   Down: { icon: ChevronDownIcon, className: "hover:text-blue-600" },
   Open: { icon: ExternalLinkIcon, className: "hover:text-green-600" },
+  Copy: { icon: ClipboardCopyIcon, className: "hover:text-blue-600" },
 };
 
 export default function TableIconActionButton({

ui/litellm-dashboard/src/components/common_components/OrganizationDropdown.tsx

@@ -1,13 +1,16 @@
 import React from "react";
-import { Select } from "antd";
+import { Select, Typography } from "antd";
 import { Organization } from "../networking";
 
+const { Text } = Typography;
+
 interface OrganizationDropdownProps {
   organizations?: Organization[] | null;
   value?: string;
   onChange?: (value: string) => void;
   disabled?: boolean;
   loading?: boolean;
+  style?: React.CSSProperties;
 }
 
 const OrganizationDropdown: React.FC<OrganizationDropdownProps> = ({
@@ -16,16 +19,18 @@ const OrganizationDropdown: React.FC<OrganizationDropdownProps> = ({
   onChange,
   disabled,
   loading,
+  style,
 }) => {
   return (
     <Select
       showSearch
-      placeholder="Search or select an organization"
+      placeholder="All Organizations"
       value={value}
       onChange={onChange}
       disabled={disabled}
       loading={loading}
       allowClear
+      style={{ minWidth: 280, ...style }}
       filterOption={(input, option) => {
         if (!option) return false;
         const org = organizations?.find((o) => o.organization_id === option.key);
@@ -37,12 +42,11 @@ const OrganizationDropdown: React.FC<OrganizationDropdownProps> = ({
 
         return orgAlias.includes(searchTerm) || orgId.includes(searchTerm);
       }}
-      optionFilterProp="children"
     >
       {organizations?.map((org) => (
         <Select.Option key={org.organization_id} value={org.organization_id}>
           <span className="font-medium">{org.organization_alias}</span>{" "}
-          <span className="text-gray-500">({org.organization_id})</span>
+          <Text type="secondary">({org.organization_id})</Text>
         </Select.Option>
       ))}
     </Select>

ui/litellm-dashboard/src/components/key_team_helpers/filter_helpers.test.ts

@@ -0,0 +1,90 @@
+import { describe, expect, it, vi } from "vitest";
+import { fetchTeamFilterOptions } from "./filter_helpers";
+
+const mockKeyListCall = vi.fn();
+
+vi.mock("@/components/networking", () => ({
+  keyListCall: (...args: unknown[]) => mockKeyListCall(...args),
+  teamListCall: vi.fn(),
+  organizationListCall: vi.fn(),
+}));
+
+describe("fetchTeamFilterOptions", () => {
+  it("should return empty arrays when accessToken is null", async () => {
+    const result = await fetchTeamFilterOptions(null, "team-1");
+
+    expect(result).toEqual({ keyAliases: [], organizationIds: [], userIds: [] });
+    expect(mockKeyListCall).not.toHaveBeenCalled();
+  });
+
+  it("should return empty arrays when teamId is empty", async () => {
+    const result = await fetchTeamFilterOptions("tok-123", "");
+
+    expect(result).toEqual({ keyAliases: [], organizationIds: [], userIds: [] });
+    expect(mockKeyListCall).not.toHaveBeenCalled();
+  });
+
+  it("should return sorted key aliases from fetched keys", async () => {
+    mockKeyListCall.mockResolvedValue({
+      keys: [
+        { key_alias: "zeta-key" },
+        { key_alias: "alpha-key" },
+        { key_alias: "mid-key" },
+      ],
+      total_pages: 1,
+    });
+
+    const result = await fetchTeamFilterOptions("tok-123", "team-1");
+
+    expect(result.keyAliases).toEqual(["alpha-key", "mid-key", "zeta-key"]);
+  });
+
+  it("should deduplicate organization IDs across pages", async () => {
+    mockKeyListCall
+      .mockResolvedValueOnce({
+        keys: [
+          { organization_id: "org-b" },
+          { organization_id: "org-a" },
+        ],
+        total_pages: 2,
+      })
+      .mockResolvedValueOnce({
+        keys: [
+          { organization_id: "org-a" },
+          { organization_id: "org-c" },
+        ],
+        total_pages: 2,
+      });
+
+    const result = await fetchTeamFilterOptions("tok-123", "team-1");
+
+    expect(result.organizationIds).toEqual(["org-a", "org-b", "org-c"]);
+  });
+
+  it("should map user IDs with email addresses", async () => {
+    mockKeyListCall.mockResolvedValue({
+      keys: [
+        { user_id: "u1", user: { user_email: "alice@example.com" } },
+        { user_id: "u2", user: { user_email: "bob@example.com" } },
+      ],
+      total_pages: 1,
+    });
+
+    const result = await fetchTeamFilterOptions("tok-123", "team-1");
+
+    expect(result.userIds).toEqual(
+      expect.arrayContaining([
+        { id: "u1", email: "alice@example.com" },
+        { id: "u2", email: "bob@example.com" },
+      ]),
+    );
+  });
+
+  it("should handle API errors gracefully and return empty arrays", async () => {
+    mockKeyListCall.mockRejectedValue(new Error("Network error"));
+
+    const result = await fetchTeamFilterOptions("tok-123", "team-1");
+
+    expect(result).toEqual({ keyAliases: [], organizationIds: [], userIds: [] });
+  });
+});

ui/litellm-dashboard/src/components/key_team_helpers/transform_key_info.test.ts

@@ -0,0 +1,62 @@
+import { describe, it, expect } from "vitest";
+import { transformKeyInfo } from "./transform_key_info";
+
+describe("transformKeyInfo", () => {
+  it("should combine key and info fields into a single object", () => {
+    const apiResponse = {
+      key: "sk-abc123",
+      info: {
+        token_id: "tok_1",
+        key_name: "my-key",
+        spend: 10.5,
+      },
+    };
+    const result = transformKeyInfo(apiResponse);
+    expect(result).toEqual({
+      token: "sk-abc123",
+      token_id: "tok_1",
+      key_name: "my-key",
+      spend: 10.5,
+    });
+  });
+
+  it("should set the token field from the key property", () => {
+    const apiResponse = {
+      key: "sk-xyz789",
+      info: { key_name: "test" },
+    };
+    const result = transformKeyInfo(apiResponse);
+    expect(result.token).toBe("sk-xyz789");
+  });
+
+  it("should preserve all info fields in the result", () => {
+    const apiResponse = {
+      key: "sk-abc",
+      info: {
+        token_id: "tok_2",
+        key_name: "prod-key",
+        spend: 42,
+        models: ["gpt-4"],
+        team_id: "team-1",
+        metadata: { env: "production" },
+      },
+    };
+    const result = transformKeyInfo(apiResponse);
+    expect(result.token_id).toBe("tok_2");
+    expect(result.key_name).toBe("prod-key");
+    expect(result.spend).toBe(42);
+    expect(result.models).toEqual(["gpt-4"]);
+    expect(result.team_id).toBe("team-1");
+    expect(result.metadata).toEqual({ env: "production" });
+  });
+
+  it("should handle empty info object", () => {
+    const apiResponse = {
+      key: "sk-empty",
+      info: {},
+    };
+    const result = transformKeyInfo(apiResponse);
+    expect(result.token).toBe("sk-empty");
+    expect(Object.keys(result)).toContain("token");
+  });
+});

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

@@ -232,8 +232,8 @@ const menuGroups: MenuGroup[] = [
     groupLabel: "DEVELOPER TOOLS",
     items: [
       {
-        key: "api_ref",
-        page: "api_ref",
+        key: "api-reference",
+        page: "api-reference",
         label: "API Reference",
         icon: <ApiOutlined />,
       },

ui/litellm-dashboard/src/components/mcp_tools/ByokCredentialModal.tsx

@@ -1,7 +1,8 @@
 "use client";
 
 import React, { useState } from "react";
-import { Modal, Input, Switch, message } from "antd";
+import { Modal, Input, Switch } from "antd";
+import MessageManager from "@/components/molecules/message_manager";
 import {
   KeyOutlined,
   LockOutlined,
@@ -46,7 +47,7 @@ export const ByokCredentialModal: React.FC<ByokCredentialModalProps> = ({
 
   const handleAuthorize = async () => {
     if (!apiKey.trim()) {
-      message.error("Please enter your API key");
+      MessageManager.error("Please enter your API key");
       return;
     }
     setLoading(true);
@@ -63,11 +64,11 @@ export const ByokCredentialModal: React.FC<ByokCredentialModalProps> = ({
         const err = await response.json();
         throw new Error(err?.detail?.error || "Failed to save credential");
       }
-      message.success(`Connected to ${serverDisplayName}`);
+      MessageManager.success(`Connected to ${serverDisplayName}`);
       onSuccess(server.server_id);
       handleClose();
     } catch (e: any) {
-      message.error(e.message || "Failed to connect");
+      MessageManager.error(e.message || "Failed to connect");
     } finally {
       setLoading(false);
     }

ui/litellm-dashboard/src/components/molecules/message_manager.tsx

@@ -0,0 +1,38 @@
+import { message as staticMessage } from "antd";
+import type { MessageInstance } from "antd/es/message/interface";
+
+let messageInstance: MessageInstance | null = null;
+
+export const setMessageInstance = (instance: MessageInstance) => {
+  messageInstance = instance;
+};
+
+const getMessageApi = () => messageInstance || staticMessage;
+
+const MessageManager = {
+  success(content: string, duration?: number) {
+    getMessageApi().success(content, duration);
+  },
+
+  error(content: string, duration?: number) {
+    getMessageApi().error(content, duration);
+  },
+
+  warning(content: string, duration?: number) {
+    getMessageApi().warning(content, duration);
+  },
+
+  info(content: string, duration?: number) {
+    getMessageApi().info(content, duration);
+  },
+
+  loading(content: string, duration?: number) {
+    return getMessageApi().loading(content, duration);
+  },
+
+  destroy() {
+    getMessageApi().destroy();
+  },
+};
+
+export default MessageManager;

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

@@ -68,7 +68,7 @@ export const getInProductNudgesCall = async (accessToken: string) => {
 /**
  * Helper file for calls being made to proxy
  */
-import { message } from "antd";
+import MessageManager from "@/components/molecules/message_manager";
 import { clearTokenCookies } from "@/utils/cookieUtils";
 import { TagNewRequest, TagUpdateRequest, TagListResponse, TagInfoResponse } from "./tag_management/types";
 import { Team } from "./key_team_helpers/key_list";
@@ -613,7 +613,7 @@ export const modelCreateCall = async (accessToken: string, formValues: Model) =>
     console.log("API Response:", data);
 
     // Close any existing messages before showing new ones
-    message.destroy();
+    MessageManager.destroy();
 
     // Sequential success messages
     NotificationsManager.success(`Model ${formValues.model_name} created successfully`);

ui/litellm-dashboard/src/components/organisms/create_key_button.tsx

@@ -8,7 +8,7 @@ import { formatNumberWithCommas } from "@/utils/dataUtils";
 import { InfoCircleOutlined } from "@ant-design/icons";
 import { useQueryClient } from "@tanstack/react-query";
 import { Accordion, AccordionBody, AccordionHeader, Button, Col, Grid, Text, TextInput, Title } from "@tremor/react";
-import { Button as Button2, Form, Input, message, Modal, Radio, Select, Switch, Tag, Tooltip } from "antd";
+import { Button as Button2, Form, Input, Modal, Radio, Select, Switch, Tag, Tooltip } from "antd";
 import debounce from "lodash/debounce";
 import React, { useCallback, useEffect, useState } from "react";
 import { rolesWithWriteAccess } from "../../utils/roles";

ui/litellm-dashboard/src/components/page_metadata.ts

@@ -24,7 +24,7 @@ export const pageDescriptions: Record<string, string> = {
   projects: "Manage projects within teams",
   "access-groups": "Manage access groups for role-based permissions",
   budgets: "Set and monitor spending budgets",
-  api_ref: "Browse API documentation and endpoints",
+  "api-reference": "Browse API documentation and endpoints",
   "model-hub-table": "Explore available AI models and providers",
   "learning-resources": "Access tutorials and documentation",
   caching: "Configure response caching settings",

ui/litellm-dashboard/src/components/playground/chat_ui/ChatUI.tsx

@@ -28,7 +28,6 @@ import ReactMarkdown from "react-markdown";
 import { Prism as SyntaxHighlighter } from "react-syntax-highlighter";
 import { coy } from "react-syntax-highlighter/dist/esm/styles/prism";
 import { v4 as uuidv4 } from "uuid";
-import { truncateString } from "../../../utils/textUtils";
 import GuardrailSelector from "../../guardrails/GuardrailSelector";
 import PolicySelector from "../../policies/PolicySelector";
 import MCPToolArgumentsForm, { MCPToolArgumentsFormRef } from "../../mcp_tools/MCPToolArgumentsForm";
@@ -61,9 +60,8 @@ import CodeInterpreterTool from "./CodeInterpreterTool";
 import { generateCodeSnippet } from "./CodeSnippets";
 import EndpointSelector from "./EndpointSelector";
 import FilePreviewCard from "./FilePreviewCard";
-import MCPEventsDisplay from "./MCPEventsDisplay";
-import type { MCPEvent } from "../../mcp_tools/types";
 import ChatMessageBubble from "./ChatMessageBubble";
+import MCPEventsDisplay from "./MCPEventsDisplay";
 import { EndpointType, getEndpointType } from "./mode_endpoint_mapping";
 import ReasoningContent from "./ReasoningContent";
 import ResponseMetrics, { TokenUsage } from "./ResponseMetrics";
@@ -75,6 +73,7 @@ import SessionManagement from "./SessionManagement";
 import RealtimePlayground from "./RealtimePlayground";
 import { A2ATaskMetadata, MessageType } from "./types";
 import { useCodeInterpreter } from "./useCodeInterpreter";
+import { useChatHistory } from "./useChatHistory";
 
 const { TextArea } = Input;
 const { Dragger } = Upload;
@@ -135,6 +134,34 @@ const ChatUI: React.FC<ChatUIProps> = ({
       return {};
     }
   });
+  const {
+    chatHistory,
+    setChatHistory,
+    mcpEvents,
+    setMCPEvents,
+    messageTraceId,
+    setMessageTraceId,
+    responsesSessionId,
+    setResponsesSessionId,
+    useApiSessionManagement,
+    setUseApiSessionManagement,
+    updateTextUI,
+    updateReasoningContent,
+    updateTimingData,
+    updateUsageData,
+    updateA2AMetadata,
+    updateTotalLatency,
+    updateSearchResults,
+    handleResponseId,
+    handleToggleSessionManagement,
+    handleMCPEvent,
+    updateImageUI,
+    updateEmbeddingsUI,
+    updateAudioUI,
+    updateChatImageUI,
+    clearChatHistory: clearChatHistoryHook,
+    clearMCPEvents,
+  } = useChatHistory({ simplified });
   const [apiKeySource, setApiKeySource] = useState<"session" | "custom">(() => {
     const saved = sessionStorage.getItem("apiKeySource");
     if (saved) {
@@ -151,16 +178,6 @@ const ChatUI: React.FC<ChatUIProps> = ({
     () => sessionStorage.getItem("customProxyBaseUrl") || "",
   );
   const [inputMessage, setInputMessage] = useState("");
-  const [chatHistory, setChatHistory] = useState<MessageType[]>(() => {
-    if (simplified) return [];
-    try {
-      const saved = sessionStorage.getItem("chatHistory");
-      return saved ? JSON.parse(saved) : [];
-    } catch (error) {
-      console.error("Error parsing chatHistory from sessionStorage", error);
-      return [];
-    }
-  });
   const [selectedModel, setSelectedModel] = useState<string | undefined>(simplified ? fixedModel : undefined);
   const [showCustomModelInput, setShowCustomModelInput] = useState<boolean>(false);
   const [modelInfo, setModelInfo] = useState<ModelGroup[]>([]);
@@ -218,16 +235,6 @@ const ChatUI: React.FC<ChatUIProps> = ({
       return [];
     }
   });
-  const [messageTraceId, setMessageTraceId] = useState<string | null>(
-    () => sessionStorage.getItem("messageTraceId") || null,
-  );
-  const [responsesSessionId, setResponsesSessionId] = useState<string | null>(
-    () => sessionStorage.getItem("responsesSessionId") || null,
-  );
-  const [useApiSessionManagement, setUseApiSessionManagement] = useState<boolean>(() => {
-    const saved = sessionStorage.getItem("useApiSessionManagement");
-    return saved ? JSON.parse(saved) : true; // Default to API session management
-  });
   const [uploadedImages, setUploadedImages] = useState<File[]>([]);
   const [imagePreviewUrls, setImagePreviewUrls] = useState<string[]>([]);
   const [responsesUploadedImage, setResponsesUploadedImage] = useState<File | null>(null);
@@ -238,7 +245,6 @@ const ChatUI: React.FC<ChatUIProps> = ({
   const [isGetCodeModalVisible, setIsGetCodeModalVisible] = useState(false);
   const [generatedCode, setGeneratedCode] = useState("");
   const [selectedSdk, setSelectedSdk] = useState<"openai" | "azure">("openai");
-  const [mcpEvents, setMCPEvents] = useState<MCPEvent[]>([]);
   const [temperature, setTemperature] = useState<number>(1.0);
   const [maxTokens, setMaxTokens] = useState<number>(2048);
   const [useAdvancedParams, setUseAdvancedParams] = useState<boolean>(false);
@@ -332,17 +338,6 @@ const ChatUI: React.FC<ChatUIProps> = ({
     proxySettings,
   ]);
 
-  useEffect(() => {
-    if (simplified) return; // Do not persist chat history in simplified (embedded) mode
-    const handler = setTimeout(() => {
-      sessionStorage.setItem("chatHistory", JSON.stringify(chatHistory));
-    }, 500); // Debounce by 500ms
-
-    return () => {
-      clearTimeout(handler);
-    };
-  }, [chatHistory, simplified]);
-
   useEffect(() => {
     sessionStorage.setItem("apiKeySource", JSON.stringify(apiKeySource));
     sessionStorage.setItem("apiKey", apiKey);
@@ -363,17 +358,6 @@ const ChatUI: React.FC<ChatUIProps> = ({
         sessionStorage.removeItem("selectedModel");
       }
     }
-    if (messageTraceId) {
-      sessionStorage.setItem("messageTraceId", messageTraceId);
-    } else {
-      sessionStorage.removeItem("messageTraceId");
-    }
-    if (responsesSessionId) {
-      sessionStorage.setItem("responsesSessionId", responsesSessionId);
-    } else {
-      sessionStorage.removeItem("responsesSessionId");
-    }
-    sessionStorage.setItem("useApiSessionManagement", JSON.stringify(useApiSessionManagement));
     // Note: codeInterpreterEnabled and selectedContainerId are persisted by useCodeInterpreter hook
   }, [
     simplified,
@@ -385,9 +369,6 @@ const ChatUI: React.FC<ChatUIProps> = ({
     selectedVectorStores,
     selectedGuardrails,
     selectedPolicies,
-    messageTraceId,
-    responsesSessionId,
-    useApiSessionManagement,
     selectedMCPServers,
     mcpServerToolRestrictions,
     selectedVoice,
@@ -479,264 +460,6 @@ const ChatUI: React.FC<ChatUIProps> = ({
     }
   }, [chatHistory]);
 
-  const updateTextUI = (role: string, chunk: string, model?: string) => {
-    console.log("updateTextUI called with:", role, chunk, model);
-    setChatHistory((prev) => {
-      const last = prev[prev.length - 1];
-      // if the last message is already from this same role, append
-      if (last && last.role === role && !last.isImage && !last.isAudio) {
-        // build a new object, but only set `model` if it wasn't there already
-        const updated: MessageType = {
-          ...last,
-          content: last.content + chunk,
-          model: last.model ?? model, // ← only use the passed‐in model on the first chunk
-        };
-        return [...prev.slice(0, -1), updated];
-      } else {
-        // otherwise start a brand new assistant bubble
-        return [
-          ...prev,
-          {
-            role,
-            content: chunk,
-            model, // model set exactly once here
-          },
-        ];
-      }
-    });
-  };
-
-  const updateReasoningContent = (chunk: string) => {
-    setChatHistory((prevHistory) => {
-      const lastMessage = prevHistory[prevHistory.length - 1];
-
-      if (lastMessage && lastMessage.role === "assistant" && !lastMessage.isImage && !lastMessage.isAudio) {
-        return [
-          ...prevHistory.slice(0, prevHistory.length - 1),
-          {
-            ...lastMessage,
-            reasoningContent: (lastMessage.reasoningContent || "") + chunk,
-          },
-        ];
-      } else {
-        // If there's no assistant message yet, we'll create one with empty content
-        // but with reasoning content
-        if (prevHistory.length > 0 && prevHistory[prevHistory.length - 1].role === "user") {
-          return [
-            ...prevHistory,
-            {
-              role: "assistant",
-              content: "",
-              reasoningContent: chunk,
-            },
-          ];
-        }
-
-        return prevHistory;
-      }
-    });
-  };
-
-  const updateTimingData = (timeToFirstToken: number) => {
-    console.log("updateTimingData called with:", timeToFirstToken);
-    setChatHistory((prevHistory) => {
-      const lastMessage = prevHistory[prevHistory.length - 1];
-      console.log("Current last message:", lastMessage);
-
-      if (lastMessage && lastMessage.role === "assistant") {
-        console.log("Updating assistant message with timeToFirstToken:", timeToFirstToken);
-        const updatedHistory = [
-          ...prevHistory.slice(0, prevHistory.length - 1),
-          {
-            ...lastMessage,
-            timeToFirstToken,
-          },
-        ];
-        console.log("Updated chat history:", updatedHistory);
-        return updatedHistory;
-      }
-      // If the last message is a user message and no assistant message exists yet,
-      // create a new assistant message with empty content
-      else if (lastMessage && lastMessage.role === "user") {
-        console.log("Creating new assistant message with timeToFirstToken:", timeToFirstToken);
-        return [
-          ...prevHistory,
-          {
-            role: "assistant",
-            content: "",
-            timeToFirstToken,
-          },
-        ];
-      }
-
-      console.log("No appropriate message found to update timing");
-      return prevHistory;
-    });
-  };
-
-  const updateUsageData = (usage: TokenUsage, toolName?: string) => {
-    console.log("Received usage data:", usage);
-    setChatHistory((prevHistory) => {
-      const lastMessage = prevHistory[prevHistory.length - 1];
-
-      if (lastMessage && lastMessage.role === "assistant") {
-        console.log("Updating message with usage data:", usage);
-        const updatedMessage = {
-          ...lastMessage,
-          usage,
-          toolName,
-        };
-        console.log("Updated message:", updatedMessage);
-
-        return [...prevHistory.slice(0, prevHistory.length - 1), updatedMessage];
-      }
-
-      return prevHistory;
-    });
-  };
-
-  const updateA2AMetadata = (a2aMetadata: A2ATaskMetadata) => {
-    console.log("Received A2A metadata:", a2aMetadata);
-    setChatHistory((prevHistory) => {
-      const lastMessage = prevHistory[prevHistory.length - 1];
-
-      if (lastMessage && lastMessage.role === "assistant") {
-        const updatedMessage = {
-          ...lastMessage,
-          a2aMetadata,
-        };
-        return [...prevHistory.slice(0, prevHistory.length - 1), updatedMessage];
-      }
-
-      return prevHistory;
-    });
-  };
-
-  const updateTotalLatency = (totalLatency: number) => {
-    setChatHistory((prevHistory) => {
-      const lastMessage = prevHistory[prevHistory.length - 1];
-
-      if (lastMessage && lastMessage.role === "assistant") {
-        return [
-          ...prevHistory.slice(0, prevHistory.length - 1),
-          {
-            ...lastMessage,
-            totalLatency,
-          },
-        ];
-      }
-
-      return prevHistory;
-    });
-  };
-
-  const updateSearchResults = (searchResults: any[]) => {
-    console.log("Received search results:", searchResults);
-    setChatHistory((prevHistory) => {
-      const lastMessage = prevHistory[prevHistory.length - 1];
-
-      if (lastMessage && lastMessage.role === "assistant") {
-        console.log("Updating message with search results");
-        const updatedMessage = {
-          ...lastMessage,
-          searchResults,
-        };
-
-        return [...prevHistory.slice(0, prevHistory.length - 1), updatedMessage];
-      }
-
-      return prevHistory;
-    });
-  };
-
-  const handleResponseId = (responseId: string) => {
-    console.log("Received response ID for session management:", responseId);
-    if (useApiSessionManagement) {
-      setResponsesSessionId(responseId);
-    }
-  };
-
-  const handleToggleSessionManagement = (useApi: boolean) => {
-    setUseApiSessionManagement(useApi);
-    if (!useApi) {
-      // Clear API session when switching to UI mode
-      setResponsesSessionId(null);
-    }
-  };
-
-  const handleMCPEvent = (event: MCPEvent) => {
-    console.log("ChatUI: Received MCP event:", event);
-    setMCPEvents((prev) => {
-      // Check if this is a duplicate event (same item_id and type)
-      // Only check for duplicates if item_id is defined (for mcp_list_tools, item_id is "mcp_list_tools")
-      const isDuplicate = event.item_id
-        ? prev.some(
-            (existingEvent) =>
-              existingEvent.item_id === event.item_id &&
-              existingEvent.type === event.type &&
-              (existingEvent.sequence_number === event.sequence_number ||
-                (existingEvent.sequence_number === undefined && event.sequence_number === undefined)),
-          )
-        : false;
-
-      if (isDuplicate) {
-        console.log("ChatUI: Duplicate MCP event, skipping");
-        return prev;
-      }
-
-      const newEvents = [...prev, event];
-      console.log("ChatUI: Updated MCP events:", newEvents);
-      return newEvents;
-    });
-  };
-
-  const updateImageUI = (imageUrl: string, model: string) => {
-    setChatHistory((prevHistory) => [...prevHistory, { role: "assistant", content: imageUrl, model, isImage: true }]);
-  };
-
-  const updateEmbeddingsUI = (embeddings: string, model?: string) => {
-    setChatHistory((prevHistory) => [
-      ...prevHistory,
-      { role: "assistant", content: truncateString(embeddings, 100), model, isEmbeddings: true },
-    ]);
-  };
-
-  const updateAudioUI = (audioUrl: string, model: string) => {
-    setChatHistory((prevHistory) => [...prevHistory, { role: "assistant", content: audioUrl, model, isAudio: true }]);
-  };
-
-  const updateChatImageUI = (imageUrl: string, model?: string) => {
-    setChatHistory((prev) => {
-      const last = prev[prev.length - 1];
-      // If the last message is from assistant and has content, add image to it
-      if (last && last.role === "assistant" && !last.isImage && !last.isAudio) {
-        const updated = {
-          ...last,
-          image: {
-            url: imageUrl,
-            detail: "auto",
-          },
-          model: last.model ?? model,
-        };
-        return [...prev.slice(0, -1), updated];
-      } else {
-        // Otherwise create a new assistant message with just the image
-        return [
-          ...prev,
-          {
-            role: "assistant",
-            content: "",
-            model,
-            image: {
-              url: imageUrl,
-              detail: "auto",
-            },
-          },
-        ];
-      }
-    });
-  };
-
   const handleKeyDown = (event: React.KeyboardEvent<HTMLTextAreaElement>) => {
     if (event.key === "Enter" && !event.shiftKey) {
       event.preventDefault(); // Prevent default to avoid newline
@@ -967,7 +690,7 @@ const ChatUI: React.FC<ChatUIProps> = ({
     }
 
     setChatHistory([...chatHistory, displayMessage]);
-    setMCPEvents([]); // Clear previous MCP events for new conversation turn
+    clearMCPEvents(); // Clear previous MCP events for new conversation turn
     codeInterpreter.clearResult(); // Clear previous code interpreter results
     setIsLoading(true);
 
@@ -1223,26 +946,11 @@ const ChatUI: React.FC<ChatUIProps> = ({
   };
 
   const clearChatHistory = () => {
-    // Clean up audio object URLs before clearing history
-    chatHistory.forEach((message) => {
-      if (message.isAudio && typeof message.content === "string") {
-        URL.revokeObjectURL(message.content);
-      }
-    });
-
-    setChatHistory([]);
-    setMessageTraceId(null);
-    setResponsesSessionId(null); // Clear responses session ID
-    setMCPEvents([]); // Clear MCP events
-    handleRemoveAllImages(); // Clear any uploaded images for image edits
-    handleRemoveResponsesImage(); // Clear any uploaded images for responses
-    handleRemoveChatImage(); // Clear any uploaded images for chat completions
-    handleRemoveAudio(); // Clear any uploaded audio for transcription
-    if (!simplified) {
-      sessionStorage.removeItem("chatHistory");
-      sessionStorage.removeItem("messageTraceId");
-      sessionStorage.removeItem("responsesSessionId");
-    }
+    clearChatHistoryHook();
+    handleRemoveAllImages();
+    handleRemoveResponsesImage();
+    handleRemoveChatImage();
+    handleRemoveAudio();
     NotificationsManager.success("Chat history cleared.");
   };
 

ui/litellm-dashboard/src/components/playground/chat_ui/CodeInterpreterTool.tsx

@@ -1,5 +1,6 @@
 import React from "react";
-import { Switch, Tooltip, message } from "antd";
+import { Switch, Tooltip } from "antd";
+import MessageManager from "@/components/molecules/message_manager";
 import { CodeOutlined, InfoCircleOutlined, ExclamationCircleOutlined } from "@ant-design/icons";
 import { Text } from "@tremor/react";
 
@@ -38,7 +39,7 @@ const CodeInterpreterTool: React.FC<CodeInterpreterToolProps> = ({
 
   const handleToggle = (checked: boolean) => {
     if (checked && !isOpenAI) {
-      message.warning("Code Interpreter is only available for OpenAI models");
+      MessageManager.warning("Code Interpreter is only available for OpenAI models");
       return;
     }
     onEnabledChange(checked);

ui/litellm-dashboard/src/components/playground/chat_ui/useChatHistory.test.ts

@@ -0,0 +1,591 @@
+import { renderHook, act } from "@testing-library/react";
+import { describe, it, expect, beforeEach, vi } from "vitest";
+import { useChatHistory } from "./useChatHistory";
+
+describe("useChatHistory", () => {
+  beforeEach(() => {
+    sessionStorage.clear();
+  });
+
+  describe("updateTextUI", () => {
+    it("should create a new assistant message when chat is empty", () => {
+      const { result } = renderHook(() => useChatHistory({ simplified: false }));
+
+      act(() => {
+        result.current.updateTextUI("assistant", "Hello", "gpt-4");
+      });
+
+      expect(result.current.chatHistory).toEqual([
+        { role: "assistant", content: "Hello", model: "gpt-4" },
+      ]);
+    });
+
+    it("should append to the last assistant message", () => {
+      const { result } = renderHook(() => useChatHistory({ simplified: false }));
+
+      act(() => {
+        result.current.updateTextUI("assistant", "Hello", "gpt-4");
+      });
+      act(() => {
+        result.current.updateTextUI("assistant", " world");
+      });
+
+      expect(result.current.chatHistory).toEqual([
+        { role: "assistant", content: "Hello world", model: "gpt-4" },
+      ]);
+    });
+
+    it("should not overwrite model on subsequent chunks", () => {
+      const { result } = renderHook(() => useChatHistory({ simplified: false }));
+
+      act(() => {
+        result.current.updateTextUI("assistant", "Hello", "gpt-4");
+      });
+      act(() => {
+        result.current.updateTextUI("assistant", " world", "gpt-3.5");
+      });
+
+      expect(result.current.chatHistory[0].model).toBe("gpt-4");
+    });
+
+    it("should create a new message when role changes", () => {
+      const { result } = renderHook(() => useChatHistory({ simplified: false }));
+
+      act(() => {
+        result.current.updateTextUI("user", "Hi");
+      });
+      act(() => {
+        result.current.updateTextUI("assistant", "Hello", "gpt-4");
+      });
+
+      expect(result.current.chatHistory).toHaveLength(2);
+      expect(result.current.chatHistory[0].role).toBe("user");
+      expect(result.current.chatHistory[1].role).toBe("assistant");
+    });
+
+    it("should not append to image messages", () => {
+      const { result } = renderHook(() => useChatHistory({ simplified: false }));
+
+      act(() => {
+        result.current.updateImageUI("http://img.png", "dall-e");
+      });
+      act(() => {
+        result.current.updateTextUI("assistant", "description", "gpt-4");
+      });
+
+      expect(result.current.chatHistory).toHaveLength(2);
+    });
+
+    it("should not append to audio messages", () => {
+      const { result } = renderHook(() => useChatHistory({ simplified: false }));
+
+      act(() => {
+        result.current.updateAudioUI("http://audio.mp3", "tts-1");
+      });
+      act(() => {
+        result.current.updateTextUI("assistant", "text", "gpt-4");
+      });
+
+      expect(result.current.chatHistory).toHaveLength(2);
+    });
+  });
+
+  describe("updateReasoningContent", () => {
+    it("should add reasoning content to existing assistant message", () => {
+      const { result } = renderHook(() => useChatHistory({ simplified: false }));
+
+      act(() => {
+        result.current.updateTextUI("assistant", "Answer", "gpt-4");
+      });
+      act(() => {
+        result.current.updateReasoningContent("thinking...");
+      });
+
+      expect(result.current.chatHistory[0].reasoningContent).toBe("thinking...");
+    });
+
+    it("should append reasoning content across chunks", () => {
+      const { result } = renderHook(() => useChatHistory({ simplified: false }));
+
+      act(() => {
+        result.current.updateTextUI("assistant", "", "gpt-4");
+      });
+      act(() => {
+        result.current.updateReasoningContent("step 1");
+      });
+      act(() => {
+        result.current.updateReasoningContent(" step 2");
+      });
+
+      expect(result.current.chatHistory[0].reasoningContent).toBe("step 1 step 2");
+    });
+
+    it("should create assistant message with reasoning when last message is user", () => {
+      const { result } = renderHook(() => useChatHistory({ simplified: false }));
+
+      act(() => {
+        result.current.setChatHistory([{ role: "user", content: "question" }]);
+      });
+      act(() => {
+        result.current.updateReasoningContent("thinking...");
+      });
+
+      expect(result.current.chatHistory).toHaveLength(2);
+      expect(result.current.chatHistory[1]).toEqual({
+        role: "assistant",
+        content: "",
+        reasoningContent: "thinking...",
+      });
+    });
+
+    it("should not update when chat is empty", () => {
+      const { result } = renderHook(() => useChatHistory({ simplified: false }));
+
+      act(() => {
+        result.current.updateReasoningContent("thinking...");
+      });
+
+      expect(result.current.chatHistory).toHaveLength(0);
+    });
+  });
+
+  describe("updateTimingData", () => {
+    it("should add timeToFirstToken to existing assistant message", () => {
+      const { result } = renderHook(() => useChatHistory({ simplified: false }));
+
+      act(() => {
+        result.current.updateTextUI("assistant", "Hello", "gpt-4");
+      });
+      act(() => {
+        result.current.updateTimingData(150);
+      });
+
+      expect(result.current.chatHistory[0].timeToFirstToken).toBe(150);
+    });
+
+    it("should create assistant message when last is user", () => {
+      const { result } = renderHook(() => useChatHistory({ simplified: false }));
+
+      act(() => {
+        result.current.setChatHistory([{ role: "user", content: "hi" }]);
+      });
+      act(() => {
+        result.current.updateTimingData(200);
+      });
+
+      expect(result.current.chatHistory).toHaveLength(2);
+      expect(result.current.chatHistory[1].timeToFirstToken).toBe(200);
+    });
+  });
+
+  describe("updateUsageData", () => {
+    it("should add usage data to assistant message", () => {
+      const { result } = renderHook(() => useChatHistory({ simplified: false }));
+
+      act(() => {
+        result.current.updateTextUI("assistant", "Hello", "gpt-4");
+      });
+
+      const usage = { completionTokens: 10, promptTokens: 5, totalTokens: 15 };
+      act(() => {
+        result.current.updateUsageData(usage);
+      });
+
+      expect(result.current.chatHistory[0].usage).toEqual(usage);
+    });
+
+    it("should add toolName when provided", () => {
+      const { result } = renderHook(() => useChatHistory({ simplified: false }));
+
+      act(() => {
+        result.current.updateTextUI("assistant", "Hello", "gpt-4");
+      });
+
+      const usage = { completionTokens: 10, promptTokens: 5, totalTokens: 15 };
+      act(() => {
+        result.current.updateUsageData(usage, "search_tool");
+      });
+
+      expect(result.current.chatHistory[0].toolName).toBe("search_tool");
+    });
+  });
+
+  describe("updateTotalLatency", () => {
+    it("should add totalLatency to assistant message", () => {
+      const { result } = renderHook(() => useChatHistory({ simplified: false }));
+
+      act(() => {
+        result.current.updateTextUI("assistant", "Hello", "gpt-4");
+      });
+      act(() => {
+        result.current.updateTotalLatency(500);
+      });
+
+      expect(result.current.chatHistory[0].totalLatency).toBe(500);
+    });
+  });
+
+  describe("updateA2AMetadata", () => {
+    it("should add A2A metadata to assistant message", () => {
+      const { result } = renderHook(() => useChatHistory({ simplified: false }));
+
+      act(() => {
+        result.current.updateTextUI("assistant", "Hello", "gpt-4");
+      });
+
+      const metadata = { taskId: "task-1", contextId: "ctx-1" };
+      act(() => {
+        result.current.updateA2AMetadata(metadata);
+      });
+
+      expect(result.current.chatHistory[0].a2aMetadata).toEqual(metadata);
+    });
+  });
+
+  describe("updateSearchResults", () => {
+    it("should add search results to assistant message", () => {
+      const { result } = renderHook(() => useChatHistory({ simplified: false }));
+
+      act(() => {
+        result.current.updateTextUI("assistant", "Hello", "gpt-4");
+      });
+
+      const searchResults = [{ object: "search", search_query: "test", data: [] }];
+      act(() => {
+        result.current.updateSearchResults(searchResults);
+      });
+
+      expect(result.current.chatHistory[0].searchResults).toEqual(searchResults);
+    });
+  });
+
+  describe("updateImageUI", () => {
+    it("should add image message to history", () => {
+      const { result } = renderHook(() => useChatHistory({ simplified: false }));
+
+      act(() => {
+        result.current.updateImageUI("http://img.png", "dall-e-3");
+      });
+
+      expect(result.current.chatHistory).toEqual([
+        { role: "assistant", content: "http://img.png", model: "dall-e-3", isImage: true },
+      ]);
+    });
+  });
+
+  describe("updateEmbeddingsUI", () => {
+    it("should add truncated embeddings message", () => {
+      const { result } = renderHook(() => useChatHistory({ simplified: false }));
+
+      act(() => {
+        result.current.updateEmbeddingsUI("[0.1, 0.2, 0.3]", "text-embedding-ada");
+      });
+
+      expect(result.current.chatHistory[0].isEmbeddings).toBe(true);
+      expect(result.current.chatHistory[0].model).toBe("text-embedding-ada");
+    });
+  });
+
+  describe("updateAudioUI", () => {
+    it("should add audio message to history", () => {
+      const { result } = renderHook(() => useChatHistory({ simplified: false }));
+
+      act(() => {
+        result.current.updateAudioUI("http://audio.mp3", "tts-1");
+      });
+
+      expect(result.current.chatHistory).toEqual([
+        { role: "assistant", content: "http://audio.mp3", model: "tts-1", isAudio: true },
+      ]);
+    });
+  });
+
+  describe("updateChatImageUI", () => {
+    it("should add image to existing assistant message", () => {
+      const { result } = renderHook(() => useChatHistory({ simplified: false }));
+
+      act(() => {
+        result.current.updateTextUI("assistant", "Here is the image", "gpt-4");
+      });
+      act(() => {
+        result.current.updateChatImageUI("http://img.png", "gpt-4");
+      });
+
+      expect(result.current.chatHistory[0].image).toEqual({
+        url: "http://img.png",
+        detail: "auto",
+      });
+    });
+
+    it("should create new assistant message with image when no assistant message exists", () => {
+      const { result } = renderHook(() => useChatHistory({ simplified: false }));
+
+      act(() => {
+        result.current.updateChatImageUI("http://img.png", "gpt-4");
+      });
+
+      expect(result.current.chatHistory[0]).toEqual({
+        role: "assistant",
+        content: "",
+        model: "gpt-4",
+        image: { url: "http://img.png", detail: "auto" },
+      });
+    });
+  });
+
+  describe("handleMCPEvent", () => {
+    it("should add MCP event", () => {
+      const { result } = renderHook(() => useChatHistory({ simplified: false }));
+
+      act(() => {
+        result.current.handleMCPEvent({ type: "tool_call", item_id: "1" });
+      });
+
+      expect(result.current.mcpEvents).toHaveLength(1);
+    });
+
+    it("should deduplicate events by item_id and type", () => {
+      const { result } = renderHook(() => useChatHistory({ simplified: false }));
+
+      const event = { type: "tool_call", item_id: "1" };
+      act(() => {
+        result.current.handleMCPEvent(event);
+      });
+      act(() => {
+        result.current.handleMCPEvent(event);
+      });
+
+      expect(result.current.mcpEvents).toHaveLength(1);
+    });
+
+    it("should allow events without item_id (no dedup)", () => {
+      const { result } = renderHook(() => useChatHistory({ simplified: false }));
+
+      act(() => {
+        result.current.handleMCPEvent({ type: "tool_call" });
+      });
+      act(() => {
+        result.current.handleMCPEvent({ type: "tool_call" });
+      });
+
+      expect(result.current.mcpEvents).toHaveLength(2);
+    });
+
+    it("should allow events with same item_id/type but different sequence_number", () => {
+      const { result } = renderHook(() => useChatHistory({ simplified: false }));
+
+      act(() => {
+        result.current.handleMCPEvent({ type: "tool_call", item_id: "1", sequence_number: 1 });
+      });
+      act(() => {
+        result.current.handleMCPEvent({ type: "tool_call", item_id: "1", sequence_number: 2 });
+      });
+
+      expect(result.current.mcpEvents).toHaveLength(2);
+    });
+  });
+
+  describe("clearMCPEvents", () => {
+    it("should clear MCP events without affecting chat history", () => {
+      const { result } = renderHook(() => useChatHistory({ simplified: false }));
+
+      act(() => {
+        result.current.updateTextUI("assistant", "Hello", "gpt-4");
+        result.current.handleMCPEvent({ type: "tool_call", item_id: "1" });
+      });
+      act(() => {
+        result.current.clearMCPEvents();
+      });
+
+      expect(result.current.mcpEvents).toEqual([]);
+      expect(result.current.chatHistory).toHaveLength(1);
+    });
+  });
+
+  describe("clearChatHistory", () => {
+    it("should clear all state", () => {
+      const { result } = renderHook(() => useChatHistory({ simplified: false }));
+
+      act(() => {
+        result.current.updateTextUI("assistant", "Hello", "gpt-4");
+        result.current.handleMCPEvent({ type: "tool_call", item_id: "1" });
+      });
+      act(() => {
+        result.current.clearChatHistory();
+      });
+
+      expect(result.current.chatHistory).toEqual([]);
+      expect(result.current.mcpEvents).toEqual([]);
+      expect(result.current.messageTraceId).toBeNull();
+      expect(result.current.responsesSessionId).toBeNull();
+    });
+
+    it("should revoke audio object URLs when clearing", () => {
+      const revokeSpy = vi.spyOn(URL, "revokeObjectURL").mockImplementation(() => {});
+      const { result } = renderHook(() => useChatHistory({ simplified: false }));
+
+      act(() => {
+        result.current.updateAudioUI("blob:http://localhost/audio-1", "tts-1");
+      });
+      act(() => {
+        result.current.clearChatHistory();
+      });
+
+      expect(revokeSpy).toHaveBeenCalledWith("blob:http://localhost/audio-1");
+      revokeSpy.mockRestore();
+    });
+
+    it("should clear sessionStorage when not simplified", () => {
+      vi.useFakeTimers();
+      const { result } = renderHook(() => useChatHistory({ simplified: false }));
+
+      sessionStorage.setItem("chatHistory", "[]");
+      sessionStorage.setItem("messageTraceId", "trace-1");
+      sessionStorage.setItem("responsesSessionId", "resp-1");
+
+      act(() => {
+        result.current.clearChatHistory();
+      });
+
+      // Advance past the 500ms debounce to verify it does not re-write the key
+      act(() => {
+        vi.advanceTimersByTime(600);
+      });
+
+      expect(sessionStorage.getItem("chatHistory")).toBeNull();
+      expect(sessionStorage.getItem("messageTraceId")).toBeNull();
+      expect(sessionStorage.getItem("responsesSessionId")).toBeNull();
+
+      vi.useRealTimers();
+    });
+
+    it("should NOT clear sessionStorage when simplified", () => {
+      sessionStorage.setItem("chatHistory", '[{"role":"user","content":"hi"}]');
+
+      const { result } = renderHook(() => useChatHistory({ simplified: true }));
+
+      act(() => {
+        result.current.clearChatHistory();
+      });
+
+      // simplified mode should not touch sessionStorage
+      expect(sessionStorage.getItem("chatHistory")).toBe('[{"role":"user","content":"hi"}]');
+    });
+
+    it("should not re-write chatHistory to sessionStorage after clear via debounce", () => {
+      vi.useFakeTimers();
+      const { result } = renderHook(() => useChatHistory({ simplified: false }));
+
+      // Add a message so the debounce has something to persist
+      act(() => {
+        result.current.updateTextUI("assistant", "Hello", "gpt-4");
+      });
+
+      // Let the debounce fire so the message is persisted
+      act(() => {
+        vi.advanceTimersByTime(600);
+      });
+      expect(sessionStorage.getItem("chatHistory")).not.toBeNull();
+
+      // Now clear
+      act(() => {
+        result.current.clearChatHistory();
+      });
+
+      // Advance past the debounce — the key should stay removed
+      act(() => {
+        vi.advanceTimersByTime(600);
+      });
+
+      expect(sessionStorage.getItem("chatHistory")).toBeNull();
+
+      vi.useRealTimers();
+    });
+  });
+
+  describe("simplified mode session isolation", () => {
+    it("should not hydrate messageTraceId from sessionStorage in simplified mode", () => {
+      sessionStorage.setItem("messageTraceId", "trace-from-playground");
+
+      const { result } = renderHook(() => useChatHistory({ simplified: true }));
+
+      expect(result.current.messageTraceId).toBeNull();
+    });
+
+    it("should not hydrate responsesSessionId from sessionStorage in simplified mode", () => {
+      sessionStorage.setItem("responsesSessionId", "resp-from-playground");
+
+      const { result } = renderHook(() => useChatHistory({ simplified: true }));
+
+      expect(result.current.responsesSessionId).toBeNull();
+    });
+
+    it("should not hydrate useApiSessionManagement from sessionStorage in simplified mode", () => {
+      sessionStorage.setItem("useApiSessionManagement", "false");
+
+      const { result } = renderHook(() => useChatHistory({ simplified: true }));
+
+      // Should get the default (true), not the stored value
+      expect(result.current.useApiSessionManagement).toBe(true);
+    });
+
+    it("should not persist session state to sessionStorage in simplified mode", () => {
+      vi.useFakeTimers();
+      const { result } = renderHook(() => useChatHistory({ simplified: true }));
+
+      act(() => {
+        result.current.setMessageTraceId("trace-embedded");
+        result.current.setResponsesSessionId("resp-embedded");
+      });
+
+      // Flush effects
+      act(() => {
+        vi.advanceTimersByTime(0);
+      });
+
+      expect(sessionStorage.getItem("messageTraceId")).toBeNull();
+      expect(sessionStorage.getItem("responsesSessionId")).toBeNull();
+
+      vi.useRealTimers();
+    });
+  });
+
+  describe("session management", () => {
+    it("handleResponseId should set responsesSessionId when useApiSessionManagement is true", () => {
+      const { result } = renderHook(() => useChatHistory({ simplified: false }));
+
+      act(() => {
+        result.current.handleResponseId("resp-123");
+      });
+
+      expect(result.current.responsesSessionId).toBe("resp-123");
+    });
+
+    it("handleResponseId should NOT set responsesSessionId when useApiSessionManagement is false", () => {
+      const { result } = renderHook(() => useChatHistory({ simplified: false }));
+
+      act(() => {
+        result.current.handleToggleSessionManagement(false);
+      });
+      act(() => {
+        result.current.handleResponseId("resp-123");
+      });
+
+      expect(result.current.responsesSessionId).toBeNull();
+    });
+
+    it("handleToggleSessionManagement should clear session when switching to UI mode", () => {
+      const { result } = renderHook(() => useChatHistory({ simplified: false }));
+
+      act(() => {
+        result.current.handleResponseId("resp-123");
+      });
+      act(() => {
+        result.current.handleToggleSessionManagement(false);
+      });
+
+      expect(result.current.useApiSessionManagement).toBe(false);
+      expect(result.current.responsesSessionId).toBeNull();
+    });
+  });
+});

ui/litellm-dashboard/src/components/playground/chat_ui/useChatHistory.ts

@@ -0,0 +1,392 @@
+import React, { useState, useEffect } from "react";
+import { MessageType, A2ATaskMetadata } from "./types";
+import { TokenUsage } from "./ResponseMetrics";
+import { MCPEvent } from "../../mcp_tools/types";
+import { truncateString } from "../../../utils/textUtils";
+
+export interface UseChatHistoryReturn {
+  // State
+  chatHistory: MessageType[];
+  setChatHistory: React.Dispatch<React.SetStateAction<MessageType[]>>;
+  mcpEvents: MCPEvent[];
+  setMCPEvents: React.Dispatch<React.SetStateAction<MCPEvent[]>>;
+  messageTraceId: string | null;
+  setMessageTraceId: React.Dispatch<React.SetStateAction<string | null>>;
+  responsesSessionId: string | null;
+  setResponsesSessionId: React.Dispatch<React.SetStateAction<string | null>>;
+  useApiSessionManagement: boolean;
+  setUseApiSessionManagement: React.Dispatch<React.SetStateAction<boolean>>;
+
+  // Actions
+  updateTextUI: (role: string, chunk: string, model?: string) => void;
+  updateReasoningContent: (chunk: string) => void;
+  updateTimingData: (timeToFirstToken: number) => void;
+  updateUsageData: (usage: TokenUsage, toolName?: string) => void;
+  updateA2AMetadata: (a2aMetadata: A2ATaskMetadata) => void;
+  updateTotalLatency: (totalLatency: number) => void;
+  updateSearchResults: (searchResults: any[]) => void;
+  handleResponseId: (responseId: string) => void;
+  handleToggleSessionManagement: (useApi: boolean) => void;
+  handleMCPEvent: (event: MCPEvent) => void;
+  updateImageUI: (imageUrl: string, model: string) => void;
+  updateEmbeddingsUI: (embeddings: string, model?: string) => void;
+  updateAudioUI: (audioUrl: string, model: string) => void;
+  updateChatImageUI: (imageUrl: string, model?: string) => void;
+  clearChatHistory: () => void;
+  clearMCPEvents: () => void;
+}
+
+export function useChatHistory({ simplified }: { simplified: boolean }): UseChatHistoryReturn {
+  const [chatHistory, setChatHistory] = useState<MessageType[]>(() => {
+    if (simplified) return [];
+    try {
+      const saved = sessionStorage.getItem("chatHistory");
+      return saved ? JSON.parse(saved) : [];
+    } catch (error) {
+      console.error("Error parsing chatHistory from sessionStorage", error);
+      return [];
+    }
+  });
+
+  const [mcpEvents, setMCPEvents] = useState<MCPEvent[]>([]);
+
+  const [messageTraceId, setMessageTraceId] = useState<string | null>(
+    () => (simplified ? null : sessionStorage.getItem("messageTraceId") || null),
+  );
+
+  const [responsesSessionId, setResponsesSessionId] = useState<string | null>(
+    () => (simplified ? null : sessionStorage.getItem("responsesSessionId") || null),
+  );
+
+  const [useApiSessionManagement, setUseApiSessionManagement] = useState<boolean>(() => {
+    if (simplified) return true;
+    const saved = sessionStorage.getItem("useApiSessionManagement");
+    return saved ? JSON.parse(saved) : true; // Default to API session management
+  });
+
+  // Debounced chatHistory persistence
+  useEffect(() => {
+    if (simplified) return; // Do not persist chat history in simplified (embedded) mode
+    // When chatHistory is empty (e.g. after clearChatHistory removed the key),
+    // don't re-write an empty array back into sessionStorage.
+    if (chatHistory.length === 0) return;
+    const handler = setTimeout(() => {
+      sessionStorage.setItem("chatHistory", JSON.stringify(chatHistory));
+    }, 500); // Debounce by 500ms
+
+    return () => {
+      clearTimeout(handler);
+    };
+  }, [chatHistory, simplified]);
+
+  // messageTraceId/responsesSessionId/useApiSessionManagement persistence
+  useEffect(() => {
+    if (simplified) return;
+    if (messageTraceId) {
+      sessionStorage.setItem("messageTraceId", messageTraceId);
+    } else {
+      sessionStorage.removeItem("messageTraceId");
+    }
+    if (responsesSessionId) {
+      sessionStorage.setItem("responsesSessionId", responsesSessionId);
+    } else {
+      sessionStorage.removeItem("responsesSessionId");
+    }
+    sessionStorage.setItem("useApiSessionManagement", JSON.stringify(useApiSessionManagement));
+  }, [messageTraceId, responsesSessionId, useApiSessionManagement, simplified]);
+
+  const updateTextUI = (role: string, chunk: string, model?: string) => {
+    setChatHistory((prev) => {
+      const last = prev[prev.length - 1];
+      // if the last message is already from this same role, append
+      if (last && last.role === role && !last.isImage && !last.isAudio) {
+        // build a new object, but only set `model` if it wasn't there already
+        const updated: MessageType = {
+          ...last,
+          content: last.content + chunk,
+          model: last.model ?? model, // ← only use the passed‐in model on the first chunk
+        };
+        return [...prev.slice(0, -1), updated];
+      } else {
+        // otherwise start a brand new assistant bubble
+        return [
+          ...prev,
+          {
+            role,
+            content: chunk,
+            model, // model set exactly once here
+          },
+        ];
+      }
+    });
+  };
+
+  const updateReasoningContent = (chunk: string) => {
+    setChatHistory((prevHistory) => {
+      const lastMessage = prevHistory[prevHistory.length - 1];
+
+      if (lastMessage && lastMessage.role === "assistant" && !lastMessage.isImage && !lastMessage.isAudio) {
+        return [
+          ...prevHistory.slice(0, prevHistory.length - 1),
+          {
+            ...lastMessage,
+            reasoningContent: (lastMessage.reasoningContent || "") + chunk,
+          },
+        ];
+      } else {
+        // If there's no assistant message yet, we'll create one with empty content
+        // but with reasoning content
+        if (prevHistory.length > 0 && prevHistory[prevHistory.length - 1].role === "user") {
+          return [
+            ...prevHistory,
+            {
+              role: "assistant",
+              content: "",
+              reasoningContent: chunk,
+            },
+          ];
+        }
+
+        return prevHistory;
+      }
+    });
+  };
+
+  const updateTimingData = (timeToFirstToken: number) => {
+    setChatHistory((prevHistory) => {
+      const lastMessage = prevHistory[prevHistory.length - 1];
+
+      if (lastMessage && lastMessage.role === "assistant") {
+        return [
+          ...prevHistory.slice(0, prevHistory.length - 1),
+          {
+            ...lastMessage,
+            timeToFirstToken,
+          },
+        ];
+      }
+      // If the last message is a user message and no assistant message exists yet,
+      // create a new assistant message with empty content
+      else if (lastMessage && lastMessage.role === "user") {
+        return [
+          ...prevHistory,
+          {
+            role: "assistant",
+            content: "",
+            timeToFirstToken,
+          },
+        ];
+      }
+
+      return prevHistory;
+    });
+  };
+
+  const updateUsageData = (usage: TokenUsage, toolName?: string) => {
+    setChatHistory((prevHistory) => {
+      const lastMessage = prevHistory[prevHistory.length - 1];
+
+      if (lastMessage && lastMessage.role === "assistant") {
+        const updatedMessage = {
+          ...lastMessage,
+          usage,
+          toolName,
+        };
+
+        return [...prevHistory.slice(0, prevHistory.length - 1), updatedMessage];
+      }
+
+      return prevHistory;
+    });
+  };
+
+  const updateA2AMetadata = (a2aMetadata: A2ATaskMetadata) => {
+    setChatHistory((prevHistory) => {
+      const lastMessage = prevHistory[prevHistory.length - 1];
+
+      if (lastMessage && lastMessage.role === "assistant") {
+        const updatedMessage = {
+          ...lastMessage,
+          a2aMetadata,
+        };
+        return [...prevHistory.slice(0, prevHistory.length - 1), updatedMessage];
+      }
+
+      return prevHistory;
+    });
+  };
+
+  const updateTotalLatency = (totalLatency: number) => {
+    setChatHistory((prevHistory) => {
+      const lastMessage = prevHistory[prevHistory.length - 1];
+
+      if (lastMessage && lastMessage.role === "assistant") {
+        return [
+          ...prevHistory.slice(0, prevHistory.length - 1),
+          {
+            ...lastMessage,
+            totalLatency,
+          },
+        ];
+      }
+
+      return prevHistory;
+    });
+  };
+
+  const updateSearchResults = (searchResults: any[]) => {
+    setChatHistory((prevHistory) => {
+      const lastMessage = prevHistory[prevHistory.length - 1];
+
+      if (lastMessage && lastMessage.role === "assistant") {
+        const updatedMessage = {
+          ...lastMessage,
+          searchResults,
+        };
+
+        return [...prevHistory.slice(0, prevHistory.length - 1), updatedMessage];
+      }
+
+      return prevHistory;
+    });
+  };
+
+  const handleResponseId = (responseId: string) => {
+    if (useApiSessionManagement) {
+      setResponsesSessionId(responseId);
+    }
+  };
+
+  const handleToggleSessionManagement = (useApi: boolean) => {
+    setUseApiSessionManagement(useApi);
+    if (!useApi) {
+      // Clear API session when switching to UI mode
+      setResponsesSessionId(null);
+    }
+  };
+
+  const handleMCPEvent = (event: MCPEvent) => {
+    setMCPEvents((prev) => {
+      // Check if this is a duplicate event (same item_id and type)
+      // Only check for duplicates if item_id is defined (for mcp_list_tools, item_id is "mcp_list_tools")
+      const isDuplicate = event.item_id
+        ? prev.some(
+            (existingEvent) =>
+              existingEvent.item_id === event.item_id &&
+              existingEvent.type === event.type &&
+              (existingEvent.sequence_number === event.sequence_number ||
+                (existingEvent.sequence_number === undefined && event.sequence_number === undefined)),
+          )
+        : false;
+
+      if (isDuplicate) {
+        return prev;
+      }
+
+      return [...prev, event];
+    });
+  };
+
+  const updateImageUI = (imageUrl: string, model: string) => {
+    setChatHistory((prevHistory) => [...prevHistory, { role: "assistant", content: imageUrl, model, isImage: true }]);
+  };
+
+  const updateEmbeddingsUI = (embeddings: string, model?: string) => {
+    setChatHistory((prevHistory) => [
+      ...prevHistory,
+      { role: "assistant", content: truncateString(embeddings, 100), model, isEmbeddings: true },
+    ]);
+  };
+
+  const updateAudioUI = (audioUrl: string, model: string) => {
+    setChatHistory((prevHistory) => [...prevHistory, { role: "assistant", content: audioUrl, model, isAudio: true }]);
+  };
+
+  const updateChatImageUI = (imageUrl: string, model?: string) => {
+    setChatHistory((prev) => {
+      const last = prev[prev.length - 1];
+      // If the last message is from assistant and has content, add image to it
+      if (last && last.role === "assistant" && !last.isImage && !last.isAudio) {
+        const updated = {
+          ...last,
+          image: {
+            url: imageUrl,
+            detail: "auto",
+          },
+          model: last.model ?? model,
+        };
+        return [...prev.slice(0, -1), updated];
+      } else {
+        // Otherwise create a new assistant message with just the image
+        return [
+          ...prev,
+          {
+            role: "assistant",
+            content: "",
+            model,
+            image: {
+              url: imageUrl,
+              detail: "auto",
+            },
+          },
+        ];
+      }
+    });
+  };
+
+  const clearChatHistory = () => {
+    // Use functional updater to get the latest snapshot — avoids stale-closure
+    // bugs where audio messages added between the last render and the click
+    // would leak their blob URLs.
+    setChatHistory((prev) => {
+      prev.forEach((message) => {
+        if (message.isAudio && typeof message.content === "string") {
+          URL.revokeObjectURL(message.content);
+        }
+      });
+      return [];
+    });
+
+    setMessageTraceId(null);
+    setResponsesSessionId(null); // Clear responses session ID
+    setMCPEvents([]); // Clear MCP events
+    if (!simplified) {
+      sessionStorage.removeItem("chatHistory");
+      sessionStorage.removeItem("messageTraceId");
+      sessionStorage.removeItem("responsesSessionId");
+    }
+  };
+
+  const clearMCPEvents = () => {
+    setMCPEvents([]);
+  };
+
+  return {
+    chatHistory,
+    setChatHistory,
+    mcpEvents,
+    setMCPEvents,
+    messageTraceId,
+    setMessageTraceId,
+    responsesSessionId,
+    setResponsesSessionId,
+    useApiSessionManagement,
+    setUseApiSessionManagement,
+    updateTextUI,
+    updateReasoningContent,
+    updateTimingData,
+    updateUsageData,
+    updateA2AMetadata,
+    updateTotalLatency,
+    updateSearchResults,
+    handleResponseId,
+    handleToggleSessionManagement,
+    handleMCPEvent,
+    updateImageUI,
+    updateEmbeddingsUI,
+    updateAudioUI,
+    updateChatImageUI,
+    clearChatHistory,
+    clearMCPEvents,
+  };
+}

ui/litellm-dashboard/src/components/policies/index.tsx

@@ -1,6 +1,7 @@
 import React, { useState, useEffect, useCallback } from "react";
 import { Button, TabGroup, TabList, Tab, TabPanels, TabPanel } from "@tremor/react";
-import { Modal, message, Alert } from "antd";
+import { Modal, Alert } from "antd";
+import MessageManager from "@/components/molecules/message_manager";
 import { ExclamationCircleOutlined, InfoCircleOutlined } from "@ant-design/icons";
 import { isAdminRole } from "@/utils/roles";
 import PolicyTable from "./policy_table";
@@ -80,7 +81,7 @@ const PoliciesPanel: React.FC<PoliciesPanelProps> = ({
       setPoliciesList(response.policies || []);
     } catch (error) {
       console.error("Error fetching policies:", error);
-      message.error("Failed to fetch policies");
+      MessageManager.error("Failed to fetch policies");
     } finally {
       setIsLoading(false);
     }
@@ -95,7 +96,7 @@ const PoliciesPanel: React.FC<PoliciesPanelProps> = ({
       setAttachmentsList(response.attachments || []);
     } catch (error) {
       console.error("Error fetching attachments:", error);
-      message.error("Failed to fetch attachments");
+      MessageManager.error("Failed to fetch attachments");
     } finally {
       setIsAttachmentsLoading(false);
     }
@@ -148,11 +149,11 @@ const PoliciesPanel: React.FC<PoliciesPanelProps> = ({
     setIsDeleting(true);
     try {
       await deletePolicyCall(accessToken, policyToDelete.policy_id);
-      message.success(`Policy "${policyToDelete.policy_name}" deleted successfully`);
+      MessageManager.success(`Policy "${policyToDelete.policy_name}" deleted successfully`);
       await fetchPolicies();
     } catch (error) {
       console.error("Error deleting policy:", error);
-      message.error("Failed to delete policy");
+      MessageManager.error("Failed to delete policy");
     } finally {
       setIsDeleting(false);
       setIsDeleteModalOpen(false);
@@ -177,11 +178,11 @@ const PoliciesPanel: React.FC<PoliciesPanelProps> = ({
         if (!accessToken) return;
         try {
           await deletePolicyAttachmentCall(accessToken, attachmentId);
-          message.success("Attachment deleted successfully");
+          MessageManager.success("Attachment deleted successfully");
           fetchAttachments();
         } catch (error) {
           console.error("Error deleting attachment:", error);
-          message.error("Failed to delete attachment");
+          MessageManager.error("Failed to delete attachment");
         }
       },
     });
@@ -193,7 +194,7 @@ const PoliciesPanel: React.FC<PoliciesPanelProps> = ({
 
   const handleUseTemplate = async (template: any) => {
     if (!accessToken) {
-      message.error("Authentication required");
+      MessageManager.error("Authentication required");
       return;
     }
 
@@ -221,7 +222,7 @@ const PoliciesPanel: React.FC<PoliciesPanelProps> = ({
       setIsGuardrailSelectionModalOpen(true);
     } catch (error) {
       console.error("Error fetching guardrails:", error);
-      message.error("Failed to load guardrails. Please try again.");
+      MessageManager.error("Failed to load guardrails. Please try again.");
     }
   };
 
@@ -271,7 +272,7 @@ const PoliciesPanel: React.FC<PoliciesPanelProps> = ({
       await proceedWithTemplate(enrichedTemplate);
     } catch (error) {
       console.error("Error enriching template:", error);
-      message.error("Failed to configure template. Please try again.");
+      MessageManager.error("Failed to configure template. Please try again.");
       setIsEnrichingTemplate(false);
     }
   };
@@ -318,15 +319,15 @@ const PoliciesPanel: React.FC<PoliciesPanelProps> = ({
 
       // Show success message
       if (createdGuardrails.length > 0) {
-        message.success(
+        MessageManager.success(
           `Created ${createdGuardrails.length} guardrail${createdGuardrails.length > 1 ? "s" : ""}! Complete the policy form to save.`
         );
       } else {
-        message.success("Template ready! Complete the policy form to save.");
+        MessageManager.success("Template ready! Complete the policy form to save.");
       }
 
       if (failedGuardrails.length > 0) {
-        message.warning(
+        MessageManager.warning(
           `Failed to create ${failedGuardrails.length} guardrail(s): ${failedGuardrails.join(", ")}. You may need to create them manually.`
         );
       }
@@ -348,7 +349,7 @@ const PoliciesPanel: React.FC<PoliciesPanelProps> = ({
       setTemplateQueue([]);
       setTemplateQueueProgress(null);
       console.error("Error creating guardrails:", error);
-      message.error("Failed to create guardrails. Please try again.");
+      MessageManager.error("Failed to create guardrails. Please try again.");
     }
   };
 

ui/litellm-dashboard/src/components/policies/pipeline_flow_builder.tsx

@@ -1,5 +1,6 @@
 import React, { useState } from "react";
-import { Select, Typography, message, Spin } from "antd";
+import { Select, Typography, Spin } from "antd";
+import MessageManager from "@/components/molecules/message_manager";
 import { Button, TextInput } from "@tremor/react";
 import { ArrowLeftIcon, PlusIcon } from "@heroicons/react/outline";
 import { DotsVerticalIcon } from "@heroicons/react/solid";
@@ -1385,17 +1386,17 @@ export const FlowBuilderPage: React.FC<FlowBuilderPageProps> = ({
 
   const handleSave = async () => {
     if (!policyName.trim()) {
-      message.error("Please enter a policy name");
+      MessageManager.error("Please enter a policy name");
       return;
     }
     if (!accessToken) {
-      message.error("No access token available");
+      MessageManager.error("No access token available");
       return;
     }
 
     const emptySteps = pipeline.steps.filter((s) => !s.guardrail);
     if (emptySteps.length > 0) {
-      message.error("Please select a guardrail for all steps");
+      MessageManager.error("Please select a guardrail for all steps");
       return;
     }
 

ui/litellm-dashboard/src/components/policies/policy_templates.tsx

@@ -1,5 +1,6 @@
 import React, { useState, useEffect, useMemo } from "react";
-import { Card, Button, Spin, message, Checkbox } from "antd";
+import { Card, Button, Spin, Checkbox } from "antd";
+import MessageManager from "@/components/molecules/message_manager";
 import {
   ShieldCheckIcon,
   ShieldExclamationIcon,
@@ -184,7 +185,7 @@ const PolicyTemplates: React.FC<PolicyTemplatesProps> = ({ onUseTemplate, onOpen
         onTemplatesLoaded?.(data);
       } catch (error) {
         console.error("Error fetching policy templates:", error);
-        message.error("Failed to fetch policy templates");
+        MessageManager.error("Failed to fetch policy templates");
       } finally {
         setIsLoading(false);
       }

ui/litellm-dashboard/src/components/shared/CreatedKeyDisplay.test.tsx

@@ -3,15 +3,11 @@ import { render, screen, act } from "@testing-library/react";
 import userEvent from "@testing-library/user-event";
 import CreatedKeyDisplay from "./CreatedKeyDisplay";
 
-vi.mock("antd", async (importOriginal) => {
-  const actual = await importOriginal<typeof import("antd")>();
-  return {
-    ...actual,
-    message: { success: vi.fn() },
-  };
-});
+vi.mock("@/components/molecules/message_manager", () => ({
+  default: { success: vi.fn(), error: vi.fn(), warning: vi.fn(), info: vi.fn(), loading: vi.fn(), destroy: vi.fn() },
+}));
 
-import { message } from "antd";
+import MessageManager from "@/components/molecules/message_manager";
 
 describe("CreatedKeyDisplay", () => {
   beforeEach(() => {
@@ -52,7 +48,7 @@ describe("CreatedKeyDisplay", () => {
 
     await user.click(screen.getByRole("button", { name: /copy virtual key/i }));
 
-    expect(message.success).toHaveBeenCalledWith("Key copied to clipboard");
+    expect(MessageManager.success).toHaveBeenCalledWith("Key copied to clipboard");
   });
 
   it("should revert button text back after 2 seconds", async () => {

ui/litellm-dashboard/src/components/shared/CreatedKeyDisplay.tsx

@@ -1,6 +1,7 @@
 import React, { useState } from "react";
 import { CopyToClipboard } from "react-copy-to-clipboard";
-import { Button, message } from "antd";
+import { Button } from "antd";
+import MessageManager from "@/components/molecules/message_manager";
 
 interface CreatedKeyDisplayProps {
   apiKey: string;
@@ -15,7 +16,7 @@ const CreatedKeyDisplay: React.FC<CreatedKeyDisplayProps> = ({ apiKey }) => {
 
   const handleCopy = () => {
     setCopied(true);
-    message.success("Key copied to clipboard");
+    MessageManager.success("Key copied to clipboard");
     setTimeout(() => setCopied(false), 2000);
   };
 

ui/litellm-dashboard/src/components/team/TeamInfo.tsx

@@ -20,7 +20,8 @@ import { isProxyAdminRole } from "@/utils/roles";
 import { EditOutlined, InfoCircleOutlined, SaveOutlined } from "@ant-design/icons";
 import { ArrowLeftIcon } from "@heroicons/react/outline";
 import { Badge, Card, Grid, Text, TextInput, Title } from "@tremor/react";
-import { Button, Form, Input, message, Select, Switch, Tabs, Tooltip } from "antd";
+import { Button, Form, Input, Select, Switch, Tabs, Tooltip } from "antd";
+import MessageManager from "@/components/molecules/message_manager";
 import { CheckIcon, CopyIcon } from "lucide-react";
 import React, { useEffect, useMemo, useState } from "react";
 import { copyToClipboard as utilCopyToClipboard } from "../../utils/dataUtils";
@@ -366,7 +367,7 @@ const TeamInfoView: React.FC<TeamInfoProps> = ({
         tpm_limit: values.tpm_limit,
         rpm_limit: values.rpm_limit,
       };
-      message.destroy(); // Remove all existing toasts
+      MessageManager.destroy(); // Remove all existing toasts
 
       await teamMemberUpdateCall(accessToken, teamId, member);
 
@@ -388,7 +389,7 @@ const TeamInfoView: React.FC<TeamInfoProps> = ({
       }
       setIsEditMemberModalVisible(false);
 
-      message.destroy(); // Remove all existing toasts
+      MessageManager.destroy(); // Remove all existing toasts
 
       NotificationsManager.fromBackend(errMsg);
       console.error("Error updating team member:", error);

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

@@ -0,0 +1,12 @@
+import { Spin } from "antd";
+import { LoadingOutlined } from "@ant-design/icons";
+
+interface AntDLoadingSpinnerProps {
+  size?: "small" | "default" | "large";
+  fontSize?: number;
+}
+
+export function AntDLoadingSpinner({ size, fontSize }: AntDLoadingSpinnerProps) {
+  const indicator = <LoadingOutlined style={fontSize ? { fontSize } : undefined} spin />;
+  return <Spin indicator={indicator} size={size} />;
+}

ui/litellm-dashboard/src/components/vector_store_management/CreateVectorStore.tsx

@@ -1,6 +1,7 @@
 import React, { useState } from "react";
 import { Card, Title, Text } from "@tremor/react";
-import { Upload, Button, Select, Form, message, Alert, Tooltip, Input } from "antd";
+import { Upload, Button, Select, Form, Alert, Tooltip, Input } from "antd";
+import MessageManager from "@/components/molecules/message_manager";
 import { InboxOutlined, InfoCircleOutlined } from "@ant-design/icons";
 import type { UploadProps } from "antd";
 import { ragIngestCall } from "../networking";
@@ -47,13 +48,13 @@ const CreateVectorStore: React.FC<CreateVectorStoreProps> = ({ accessToken, onSu
       ].includes(file.type);
 
       if (!isValidType) {
-        message.error(`${file.name} is not a supported file type. Please upload PDF, TXT, DOCX, or MD files.`);
+        MessageManager.error(`${file.name} is not a supported file type. Please upload PDF, TXT, DOCX, or MD files.`);
         return Upload.LIST_IGNORE;
       }
 
       const isLt50M = file.size / 1024 / 1024 < 50;
       if (!isLt50M) {
-        message.error(`${file.name} must be smaller than 50MB!`);
+        MessageManager.error(`${file.name} must be smaller than 50MB!`);
         return Upload.LIST_IGNORE;
       }
 
@@ -87,12 +88,12 @@ const CreateVectorStore: React.FC<CreateVectorStoreProps> = ({ accessToken, onSu
 
   const handleCreateVectorStore = async () => {
     if (documents.length === 0) {
-      message.warning("Please upload at least one document");
+      MessageManager.warning("Please upload at least one document");
       return;
     }
 
     if (!selectedProvider) {
-      message.warning("Please select a provider");
+      MessageManager.warning("Please select a provider");
       return;
     }
 
@@ -100,7 +101,7 @@ const CreateVectorStore: React.FC<CreateVectorStoreProps> = ({ accessToken, onSu
     const requiredFields = getProviderSpecificFields(selectedProvider).filter((field) => field.required);
     for (const field of requiredFields) {
       if (!providerParams[field.name]) {
-        message.warning(`Please provide ${field.label}`);
+        MessageManager.warning(`Please provide ${field.label}`);
         return;
       }
     }
@@ -108,17 +109,17 @@ const CreateVectorStore: React.FC<CreateVectorStoreProps> = ({ accessToken, onSu
     // S3 Vectors specific validation
     if (selectedProvider === "s3_vectors") {
       if (providerParams.vector_bucket_name && providerParams.vector_bucket_name.length < 3) {
-        message.warning("Vector bucket name must be at least 3 characters");
+        MessageManager.warning("Vector bucket name must be at least 3 characters");
         return;
       }
       if (providerParams.index_name && providerParams.index_name.length > 0 && providerParams.index_name.length < 3) {
-        message.warning("Index name must be at least 3 characters if provided");
+        MessageManager.warning("Index name must be at least 3 characters if provided");
         return;
       }
     }
 
     if (!accessToken) {
-      message.error("No access token available");
+      MessageManager.error("No access token available");
       return;
     }
 

ui/litellm-dashboard/src/components/vector_store_management/DocumentsTable.tsx

@@ -1,5 +1,6 @@
 import React from "react";
-import { Table, Badge, Tooltip, message } from "antd";
+import { Table, Badge, Tooltip } from "antd";
+import MessageManager from "@/components/molecules/message_manager";
 import { EyeOutlined, CopyOutlined, DeleteOutlined } from "@ant-design/icons";
 import { DocumentUpload } from "./types";
 
@@ -11,7 +12,7 @@ interface DocumentsTableProps {
 const DocumentsTable: React.FC<DocumentsTableProps> = ({ documents, onRemove }) => {
   const handleCopyId = (uid: string) => {
     navigator.clipboard.writeText(uid);
-    message.success("Document ID copied to clipboard");
+    MessageManager.success("Document ID copied to clipboard");
   };
 
   const getStatusBadge = (status: DocumentUpload["status"]) => {

ui/litellm-dashboard/src/components/vector_store_management/VectorStoreTester.tsx

@@ -1,5 +1,6 @@
 import React, { useState } from "react";
-import { Button, Input, Card, Typography, Spin, message, Divider } from "antd";
+import { Button, Input, Card, Typography, Spin, Divider } from "antd";
+import MessageManager from "@/components/molecules/message_manager";
 import { SendOutlined, DatabaseOutlined, LoadingOutlined, DownOutlined, RightOutlined } from "@ant-design/icons";
 import { vectorStoreSearchCall } from "../networking";
 import NotificationsManager from "../molecules/notifications_manager";
@@ -46,7 +47,7 @@ export const VectorStoreTester: React.FC<VectorStoreTesterProps> = ({ vectorStor
 
   const handleSearch = async () => {
     if (!query.trim()) {
-      message.warning("Please enter a search query");
+      MessageManager.warning("Please enter a search query");
       return;
     }
 

ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/CollapsibleMessage.test.tsx

@@ -0,0 +1,54 @@
+import React from "react";
+import { render, screen } from "@testing-library/react";
+import userEvent from "@testing-library/user-event";
+import { describe, it, expect } from "vitest";
+import { CollapsibleMessage } from "./CollapsibleMessage";
+
+describe("CollapsibleMessage", () => {
+  it("should return null when content is empty", () => {
+    const { container } = render(
+      <CollapsibleMessage label="SYSTEM" content="" />
+    );
+    expect(container.innerHTML).toBe("");
+  });
+
+  it("should return null when content is undefined", () => {
+    const { container } = render(<CollapsibleMessage label="SYSTEM" />);
+    expect(container.innerHTML).toBe("");
+  });
+
+  it("should render the label and char count", () => {
+    render(<CollapsibleMessage label="SYSTEM" content="Hello" />);
+    expect(screen.getByText("SYSTEM")).toBeInTheDocument();
+    expect(screen.getByText("(5 chars)")).toBeInTheDocument();
+  });
+
+  it("should show content when defaultExpanded is true", () => {
+    render(
+      <CollapsibleMessage
+        label="SYSTEM"
+        content="Visible text"
+        defaultExpanded={true}
+      />
+    );
+    expect(screen.getByText("Visible text")).toBeInTheDocument();
+  });
+
+  it("should toggle expanded state when header is clicked", async () => {
+    const user = userEvent.setup();
+    render(
+      <CollapsibleMessage
+        label="SYSTEM"
+        content="Toggle me"
+        defaultExpanded={false}
+      />
+    );
+
+    // Content is rendered in DOM but collapsed by default
+    expect(screen.getByText("Toggle me")).toBeInTheDocument();
+
+    // Click the header to expand - should still show content
+    await user.click(screen.getByText("SYSTEM"));
+    expect(screen.getByText("Toggle me")).toBeInTheDocument();
+  });
+});

ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/HistoryTree.test.tsx

@@ -0,0 +1,50 @@
+import React from "react";
+import { render, screen } from "@testing-library/react";
+import userEvent from "@testing-library/user-event";
+import { describe, it, expect } from "vitest";
+import { HistoryTree } from "./HistoryTree";
+import { ParsedMessage } from "./prettyMessagesTypes";
+
+describe("HistoryTree", () => {
+  it("should return null when messages array is empty", () => {
+    const { container } = render(<HistoryTree messages={[]} />);
+    expect(container.innerHTML).toBe("");
+  });
+
+  it('should render message count with plural "messages" for multiple messages', () => {
+    const messages: ParsedMessage[] = [
+      { role: "user", content: "Hello" },
+      { role: "assistant", content: "Hi there" },
+      { role: "user", content: "How are you?" },
+    ];
+    render(<HistoryTree messages={messages} />);
+    expect(
+      screen.getByText("HISTORY (3 messages)")
+    ).toBeInTheDocument();
+  });
+
+  it('should render message count with singular "message" for one message', () => {
+    const messages: ParsedMessage[] = [
+      { role: "user", content: "Hello" },
+    ];
+    render(<HistoryTree messages={messages} />);
+    expect(
+      screen.getByText("HISTORY (1 message)")
+    ).toBeInTheDocument();
+  });
+
+  it("should expand and show messages when header is clicked", async () => {
+    const user = userEvent.setup();
+    const messages: ParsedMessage[] = [
+      { role: "user", content: "Hello" },
+      { role: "assistant", content: "Hi there" },
+    ];
+    render(<HistoryTree messages={messages} />);
+
+    // Click to expand
+    await user.click(screen.getByText("HISTORY (2 messages)"));
+
+    expect(screen.getByText("Hello")).toBeInTheDocument();
+    expect(screen.getByText("Hi there")).toBeInTheDocument();
+  });
+});

ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/InputCard.tsx

@@ -4,7 +4,7 @@
  */
 
 import { useState } from 'react';
-import { message } from 'antd';
+import MessageManager from "@/components/molecules/message_manager";
 import { ParsedMessage } from './prettyMessagesTypes';
 import { SectionHeader } from './SectionHeader';
 import { CollapsibleMessage } from './CollapsibleMessage';
@@ -33,7 +33,7 @@ export function InputCard({ messages, promptTokens, inputCost }: InputCardProps)
   const handleCopy = () => {
     const content = lastMessage?.content || '';
     navigator.clipboard.writeText(content);
-    message.success('Input copied');
+    MessageManager.success('Input copied');
   };
 
   return (

ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/OutputCard.tsx

@@ -4,7 +4,8 @@
  */
 
 import { useState } from 'react';
-import { Typography, message as antdMessage } from 'antd';
+import { Typography } from 'antd';
+import MessageManager from "@/components/molecules/message_manager";
 import { ParsedMessage } from './prettyMessagesTypes';
 import { SectionHeader } from './SectionHeader';
 import { SimpleMessageBlock } from './SimpleMessageBlock';
@@ -25,7 +26,7 @@ export function OutputCard({ message, completionTokens, outputCost }: OutputCard
     
     const content = message.content || '';
     navigator.clipboard.writeText(content);
-    antdMessage.success('Output copied');
+    MessageManager.success('Output copied');
   };
 
   if (!message) {

ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/SimpleMessageBlock.test.tsx

@@ -0,0 +1,55 @@
+import React from "react";
+import { render, screen } from "@testing-library/react";
+import { describe, it, expect } from "vitest";
+import { SimpleMessageBlock } from "./SimpleMessageBlock";
+
+describe("SimpleMessageBlock", () => {
+  it("should render the label and content", () => {
+    render(<SimpleMessageBlock label="USER" content="Hello world" />);
+    expect(screen.getByText("USER")).toBeInTheDocument();
+    expect(screen.getByText("Hello world")).toBeInTheDocument();
+  });
+
+  it("should return null when content is empty and no tool calls", () => {
+    const { container } = render(
+      <SimpleMessageBlock label="USER" content="" />
+    );
+    expect(container.innerHTML).toBe("");
+  });
+
+  it('should return null when content is "null" string and no tool calls', () => {
+    const { container } = render(
+      <SimpleMessageBlock label="USER" content="null" />
+    );
+    expect(container.innerHTML).toBe("");
+  });
+
+  it("should render tool calls when present", () => {
+    render(
+      <SimpleMessageBlock
+        label="ASSISTANT"
+        toolCalls={[
+          { id: "tc1", name: "get_weather", arguments: { city: "Paris" } },
+        ]}
+      />
+    );
+    expect(screen.getByText("ASSISTANT")).toBeInTheDocument();
+    expect(screen.getByText("get_weather")).toBeInTheDocument();
+  });
+
+  it("should render content and tool calls together", () => {
+    render(
+      <SimpleMessageBlock
+        label="ASSISTANT"
+        content="Let me check the weather."
+        toolCalls={[
+          { id: "tc1", name: "get_weather", arguments: { city: "Paris" } },
+        ]}
+      />
+    );
+    expect(
+      screen.getByText("Let me check the weather.")
+    ).toBeInTheDocument();
+    expect(screen.getByText("get_weather")).toBeInTheDocument();
+  });
+});

ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/SimpleToolCallBlock.test.tsx

@@ -0,0 +1,51 @@
+import React from "react";
+import { render, screen } from "@testing-library/react";
+import { describe, it, expect } from "vitest";
+import { SimpleToolCallBlock } from "./SimpleToolCallBlock";
+
+describe("SimpleToolCallBlock", () => {
+  it("should render the tool name", () => {
+    render(
+      <SimpleToolCallBlock
+        tool={{ id: "1", name: "get_weather", arguments: {} }}
+      />
+    );
+    expect(screen.getByText("get_weather")).toBeInTheDocument();
+  });
+
+  it('should display "function" badge', () => {
+    render(
+      <SimpleToolCallBlock
+        tool={{ id: "1", name: "get_weather", arguments: {} }}
+      />
+    );
+    expect(screen.getByText("function")).toBeInTheDocument();
+  });
+
+  it("should render arguments when present", () => {
+    render(
+      <SimpleToolCallBlock
+        tool={{
+          id: "1",
+          name: "get_weather",
+          arguments: { city: "London", units: "metric" },
+        }}
+      />
+    );
+    expect(screen.getByText("city:")).toBeInTheDocument();
+    expect(screen.getByText('"London"')).toBeInTheDocument();
+    expect(screen.getByText("units:")).toBeInTheDocument();
+    expect(screen.getByText('"metric"')).toBeInTheDocument();
+  });
+
+  it("should not render arguments section when arguments are empty", () => {
+    const { container } = render(
+      <SimpleToolCallBlock
+        tool={{ id: "1", name: "get_weather", arguments: {} }}
+      />
+    );
+    // The tool name and "function" badge should be there, but no key: value pairs
+    expect(screen.getByText("get_weather")).toBeInTheDocument();
+    expect(screen.queryByText(/:$/)).not.toBeInTheDocument();
+  });
+});

ui/litellm-dashboard/src/contexts/AntdGlobalProvider.tsx

@@ -1,23 +1,27 @@
 "use client";
 
 import React, { useEffect, useRef } from "react";
-import { notification } from "antd";
+import { notification, message } from "antd";
 import { setNotificationInstance } from "@/components/molecules/notifications_manager";
+import { setMessageInstance } from "@/components/molecules/message_manager";
 
 export default function AntdGlobalProvider({ children }: { children: React.ReactNode }) {
-  const [api, contextHolder] = notification.useNotification();
+  const [notificationApi, notificationContextHolder] = notification.useNotification();
+  const [messageApi, messageContextHolder] = message.useMessage();
   const initialized = useRef(false);
 
   useEffect(() => {
     if (!initialized.current) {
-      setNotificationInstance(api);
+      setNotificationInstance(notificationApi);
+      setMessageInstance(messageApi);
       initialized.current = true;
     }
-  }, [api]);
+  }, [notificationApi, messageApi]);
 
   return (
     <>
-      {contextHolder}
+      {notificationContextHolder}
+      {messageContextHolder}
       {children}
     </>
   );

ui/litellm-dashboard/tests/setupTests.ts

@@ -100,6 +100,11 @@ if (!document.getAnimations) {
   document.getAnimations = () => [];
 }
 
+// Stub URL.revokeObjectURL so vi.spyOn can intercept it in tests
+if (!URL.revokeObjectURL) {
+  URL.revokeObjectURL = () => {};
+}
+
 // Mock ResizeObserver for components that use it (e.g., Tremor UI components)
 // This prevents "ResizeObserver is not defined" errors in JSDOM
 global.ResizeObserver = class ResizeObserver {