From f34fe4758a88cae5124348f7863a24e93febfc1a Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Thu, 19 Mar 2026 21:19:51 -0700 Subject: [PATCH] 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