diff --git a/ui/litellm-dashboard/src/components/playground/chat_ui/useChatHistory.test.ts b/ui/litellm-dashboard/src/components/playground/chat_ui/useChatHistory.test.ts
new file mode 100644
index 0000000000..312d22c28e
--- /dev/null
+++ b/ui/litellm-dashboard/src/components/playground/chat_ui/useChatHistory.test.ts
@@ -0,0 +1,506 @@
+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", () => {
+      const { result } = renderHook(() => useChatHistory({ simplified: false }));
+
+      sessionStorage.setItem("chatHistory", "[]");
+      sessionStorage.setItem("messageTraceId", "trace-1");
+      sessionStorage.setItem("responsesSessionId", "resp-1");
+
+      act(() => {
+        result.current.clearChatHistory();
+      });
+
+      expect(sessionStorage.getItem("chatHistory")).toBeNull();
+      expect(sessionStorage.getItem("messageTraceId")).toBeNull();
+      expect(sessionStorage.getItem("responsesSessionId")).toBeNull();
+    });
+
+    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"}]');
+    });
+  });
+
+  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();
+    });
+  });
+});
diff --git a/ui/litellm-dashboard/src/components/playground/chat_ui/useChatHistory.ts b/ui/litellm-dashboard/src/components/playground/chat_ui/useChatHistory.ts
new file mode 100644
index 0000000000..8830b8321c
--- /dev/null
+++ b/ui/litellm-dashboard/src/components/playground/chat_ui/useChatHistory.ts
@@ -0,0 +1,402 @@
+import { 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>(
+    () => 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
+  });
+
+  // Debounced chatHistory persistence
+  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]);
+
+  // messageTraceId/responsesSessionId/useApiSessionManagement persistence
+  useEffect(() => {
+    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]);
+
+  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 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
+    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,
+  };
+}
diff --git a/ui/litellm-dashboard/tests/setupTests.ts b/ui/litellm-dashboard/tests/setupTests.ts
index e8de9f7248..5f85a80e58 100644
--- a/ui/litellm-dashboard/tests/setupTests.ts
+++ b/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 {