From f3c6915d617d3e8c7ea44547a2ce42accf51e87f Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Thu, 19 Mar 2026 17:56:33 -0700 Subject: [PATCH 1/4] feat: add useChatHistory hook with tests (extracted from ChatUI) Co-Authored-By: Claude Sonnet 4.6 --- .../playground/chat_ui/useChatHistory.test.ts | 506 ++++++++++++++++++ .../playground/chat_ui/useChatHistory.ts | 402 ++++++++++++++ ui/litellm-dashboard/tests/setupTests.ts | 5 + 3 files changed, 913 insertions(+) create mode 100644 ui/litellm-dashboard/src/components/playground/chat_ui/useChatHistory.test.ts create mode 100644 ui/litellm-dashboard/src/components/playground/chat_ui/useChatHistory.ts 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>; + mcpEvents: MCPEvent[]; + setMCPEvents: React.Dispatch>; + messageTraceId: string | null; + setMessageTraceId: React.Dispatch>; + responsesSessionId: string | null; + setResponsesSessionId: React.Dispatch>; + useApiSessionManagement: boolean; + setUseApiSessionManagement: React.Dispatch>; + + // 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(() => { + 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([]); + + const [messageTraceId, setMessageTraceId] = useState( + () => sessionStorage.getItem("messageTraceId") || null, + ); + + const [responsesSessionId, setResponsesSessionId] = useState( + () => sessionStorage.getItem("responsesSessionId") || null, + ); + + const [useApiSessionManagement, setUseApiSessionManagement] = useState(() => { + 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 { From e88425b88113d6f7f9a67043c4265be6049876ed Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Thu, 19 Mar 2026 18:02:50 -0700 Subject: [PATCH 2/4] refactor: wire ChatUI to use useChatHistory hook Co-Authored-By: Claude Sonnet 4.6 --- .../components/playground/chat_ui/ChatUI.tsx | 362 ++---------------- .../playground/chat_ui/useChatHistory.ts | 2 +- 2 files changed, 36 insertions(+), 328 deletions(-) diff --git a/ui/litellm-dashboard/src/components/playground/chat_ui/ChatUI.tsx b/ui/litellm-dashboard/src/components/playground/chat_ui/ChatUI.tsx index bc0d52cf58..88f3c85c21 100644 --- a/ui/litellm-dashboard/src/components/playground/chat_ui/ChatUI.tsx +++ b/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"; @@ -62,7 +61,6 @@ import { generateCodeSnippet } from "./CodeSnippets"; import EndpointSelector from "./EndpointSelector"; import FilePreviewCard from "./FilePreviewCard"; import MCPEventsDisplay from "./MCPEventsDisplay"; -import type { MCPEvent } from "../../mcp_tools/types"; import { EndpointType, getEndpointType } from "./mode_endpoint_mapping"; import ReasoningContent from "./ReasoningContent"; import ResponseMetrics, { TokenUsage } from "./ResponseMetrics"; @@ -74,6 +72,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; @@ -134,6 +133,34 @@ const ChatUI: React.FC = ({ 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) { @@ -150,16 +177,6 @@ const ChatUI: React.FC = ({ () => sessionStorage.getItem("customProxyBaseUrl") || "", ); const [inputMessage, setInputMessage] = useState(""); - const [chatHistory, setChatHistory] = useState(() => { - 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(simplified ? fixedModel : undefined); const [showCustomModelInput, setShowCustomModelInput] = useState(false); const [modelInfo, setModelInfo] = useState([]); @@ -217,16 +234,6 @@ const ChatUI: React.FC = ({ return []; } }); - const [messageTraceId, setMessageTraceId] = useState( - () => sessionStorage.getItem("messageTraceId") || null, - ); - const [responsesSessionId, setResponsesSessionId] = useState( - () => sessionStorage.getItem("responsesSessionId") || null, - ); - const [useApiSessionManagement, setUseApiSessionManagement] = useState(() => { - const saved = sessionStorage.getItem("useApiSessionManagement"); - return saved ? JSON.parse(saved) : true; // Default to API session management - }); const [uploadedImages, setUploadedImages] = useState([]); const [imagePreviewUrls, setImagePreviewUrls] = useState([]); const [responsesUploadedImage, setResponsesUploadedImage] = useState(null); @@ -237,7 +244,6 @@ const ChatUI: React.FC = ({ const [isGetCodeModalVisible, setIsGetCodeModalVisible] = useState(false); const [generatedCode, setGeneratedCode] = useState(""); const [selectedSdk, setSelectedSdk] = useState<"openai" | "azure">("openai"); - const [mcpEvents, setMCPEvents] = useState([]); const [temperature, setTemperature] = useState(1.0); const [maxTokens, setMaxTokens] = useState(2048); const [useAdvancedParams, setUseAdvancedParams] = useState(false); @@ -331,17 +337,6 @@ const ChatUI: React.FC = ({ 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); @@ -362,17 +357,6 @@ const ChatUI: React.FC = ({ 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, @@ -384,9 +368,6 @@ const ChatUI: React.FC = ({ selectedVectorStores, selectedGuardrails, selectedPolicies, - messageTraceId, - responsesSessionId, - useApiSessionManagement, selectedMCPServers, mcpServerToolRestrictions, selectedVoice, @@ -478,264 +459,6 @@ const ChatUI: React.FC = ({ } }, [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) => { if (event.key === "Enter" && !event.shiftKey) { event.preventDefault(); // Prevent default to avoid newline @@ -966,7 +689,7 @@ const ChatUI: React.FC = ({ } 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); @@ -1222,26 +945,11 @@ const ChatUI: React.FC = ({ }; 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."); }; diff --git a/ui/litellm-dashboard/src/components/playground/chat_ui/useChatHistory.ts b/ui/litellm-dashboard/src/components/playground/chat_ui/useChatHistory.ts index 8830b8321c..52575ef026 100644 --- a/ui/litellm-dashboard/src/components/playground/chat_ui/useChatHistory.ts +++ b/ui/litellm-dashboard/src/components/playground/chat_ui/useChatHistory.ts @@ -1,4 +1,4 @@ -import { useState, useEffect } from "react"; +import React, { useState, useEffect } from "react"; import { MessageType, A2ATaskMetadata } from "./types"; import { TokenUsage } from "./ResponseMetrics"; import { MCPEvent } from "../../mcp_tools/types"; From 954bfcbd421429996db0f71d45908e9ee843203f Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Thu, 19 Mar 2026 20:46:51 -0700 Subject: [PATCH 3/4] chore: remove verbose console.log calls from useChatHistory hook Removes 17 console.log statements that fired on every streaming chunk in updateTextUI, updateTimingData, updateUsageData, and other hot-path functions. The console.error for sessionStorage parse failures is kept. Co-Authored-By: Claude Opus 4.6 --- .../playground/chat_ui/useChatHistory.ts | 23 ++----------------- 1 file changed, 2 insertions(+), 21 deletions(-) diff --git a/ui/litellm-dashboard/src/components/playground/chat_ui/useChatHistory.ts b/ui/litellm-dashboard/src/components/playground/chat_ui/useChatHistory.ts index 52575ef026..36e16c1701 100644 --- a/ui/litellm-dashboard/src/components/playground/chat_ui/useChatHistory.ts +++ b/ui/litellm-dashboard/src/components/playground/chat_ui/useChatHistory.ts @@ -91,7 +91,6 @@ export function useChatHistory({ simplified }: { simplified: boolean }): UseChat }, [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 @@ -149,27 +148,21 @@ export function useChatHistory({ simplified }: { simplified: boolean }): UseChat }; 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 = [ + return [ ...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, { @@ -180,24 +173,20 @@ export function useChatHistory({ simplified }: { simplified: boolean }): UseChat ]; } - 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]; } @@ -207,7 +196,6 @@ export function useChatHistory({ simplified }: { simplified: boolean }): UseChat }; const updateA2AMetadata = (a2aMetadata: A2ATaskMetadata) => { - console.log("Received A2A metadata:", a2aMetadata); setChatHistory((prevHistory) => { const lastMessage = prevHistory[prevHistory.length - 1]; @@ -242,12 +230,10 @@ export function useChatHistory({ simplified }: { simplified: boolean }): UseChat }; 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, @@ -261,7 +247,6 @@ export function useChatHistory({ simplified }: { simplified: boolean }): UseChat }; const handleResponseId = (responseId: string) => { - console.log("Received response ID for session management:", responseId); if (useApiSessionManagement) { setResponsesSessionId(responseId); } @@ -276,7 +261,6 @@ export function useChatHistory({ simplified }: { simplified: boolean }): UseChat }; 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") @@ -291,13 +275,10 @@ export function useChatHistory({ simplified }: { simplified: boolean }): UseChat : 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; + return [...prev, event]; }); }; From f34fe4758a88cae5124348f7863a24e93febfc1a Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Thu, 19 Mar 2026 21:19:51 -0700 Subject: [PATCH 4/4] fix: stale closure, simplified session isolation, debounce-race in useChatHistory - clearChatHistory: use functional setChatHistory updater so blob URL revocation operates on the latest snapshot, not a stale closure capture. - Simplified mode: skip sessionStorage hydration and persistence for messageTraceId, responsesSessionId, and useApiSessionManagement so embedded widgets don't cross-contaminate the full playground session. - Debounce race: skip re-writing empty chatHistory to sessionStorage after clearChatHistory already removed the key. - Added 5 new tests covering these fixes (39 total). Co-Authored-By: Claude Opus 4.6 --- .../playground/chat_ui/useChatHistory.test.ts | 85 +++++++++++++++++++ .../playground/chat_ui/useChatHistory.ts | 27 ++++-- 2 files changed, 103 insertions(+), 9 deletions(-) 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 index 312d22c28e..c2067fc698 100644 --- a/ui/litellm-dashboard/src/components/playground/chat_ui/useChatHistory.test.ts +++ b/ui/litellm-dashboard/src/components/playground/chat_ui/useChatHistory.test.ts @@ -436,6 +436,7 @@ describe("useChatHistory", () => { }); it("should clear sessionStorage when not simplified", () => { + vi.useFakeTimers(); const { result } = renderHook(() => useChatHistory({ simplified: false })); sessionStorage.setItem("chatHistory", "[]"); @@ -446,9 +447,16 @@ describe("useChatHistory", () => { 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", () => { @@ -463,6 +471,83 @@ describe("useChatHistory", () => { // 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", () => { diff --git a/ui/litellm-dashboard/src/components/playground/chat_ui/useChatHistory.ts b/ui/litellm-dashboard/src/components/playground/chat_ui/useChatHistory.ts index 36e16c1701..8e38191259 100644 --- a/ui/litellm-dashboard/src/components/playground/chat_ui/useChatHistory.ts +++ b/ui/litellm-dashboard/src/components/playground/chat_ui/useChatHistory.ts @@ -51,14 +51,15 @@ export function useChatHistory({ simplified }: { simplified: boolean }): UseChat const [mcpEvents, setMCPEvents] = useState([]); const [messageTraceId, setMessageTraceId] = useState( - () => sessionStorage.getItem("messageTraceId") || null, + () => (simplified ? null : sessionStorage.getItem("messageTraceId") || null), ); const [responsesSessionId, setResponsesSessionId] = useState( - () => sessionStorage.getItem("responsesSessionId") || null, + () => (simplified ? null : sessionStorage.getItem("responsesSessionId") || null), ); const [useApiSessionManagement, setUseApiSessionManagement] = useState(() => { + if (simplified) return true; const saved = sessionStorage.getItem("useApiSessionManagement"); return saved ? JSON.parse(saved) : true; // Default to API session management }); @@ -66,6 +67,9 @@ export function useChatHistory({ simplified }: { simplified: boolean }): UseChat // 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 @@ -77,6 +81,7 @@ export function useChatHistory({ simplified }: { simplified: boolean }): UseChat // messageTraceId/responsesSessionId/useApiSessionManagement persistence useEffect(() => { + if (simplified) return; if (messageTraceId) { sessionStorage.setItem("messageTraceId", messageTraceId); } else { @@ -88,7 +93,7 @@ export function useChatHistory({ simplified }: { simplified: boolean }): UseChat sessionStorage.removeItem("responsesSessionId"); } sessionStorage.setItem("useApiSessionManagement", JSON.stringify(useApiSessionManagement)); - }, [messageTraceId, responsesSessionId, useApiSessionManagement]); + }, [messageTraceId, responsesSessionId, useApiSessionManagement, simplified]); const updateTextUI = (role: string, chunk: string, model?: string) => { setChatHistory((prev) => { @@ -330,14 +335,18 @@ export function useChatHistory({ simplified }: { simplified: boolean }): UseChat }; const clearChatHistory = () => { - // Clean up audio object URLs before clearing history - chatHistory.forEach((message) => { - if (message.isAudio && typeof message.content === "string") { - URL.revokeObjectURL(message.content); - } + // 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 []; }); - setChatHistory([]); setMessageTraceId(null); setResponsesSessionId(null); // Clear responses session ID setMCPEvents([]); // Clear MCP events