From 17128054ebb494a45baefa505673e9b749063983 Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Tue, 21 Jul 2026 15:54:32 -0700 Subject: [PATCH] fix(ui): revoke code interpreter image object URLs on unmount The effect cleanup captured imageUrls from the render it closed over, which is always the initial empty object since the URLs are created later by the async fetch. Every container image blob URL was therefore never revoked and stayed pinned until full page unload. Track created URLs in an effect-local array and revoke those, and abort in-flight image fetches on unmount so late-resolving blobs cannot create URLs after cleanup has run. --- .../chat_ui/CodeInterpreterOutput.test.tsx | 94 +++++++++++++++++++ .../chat_ui/CodeInterpreterOutput.tsx | 13 ++- 2 files changed, 105 insertions(+), 2 deletions(-) diff --git a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/CodeInterpreterOutput.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/CodeInterpreterOutput.test.tsx index a42bd42ed81..0aa385b9ca7 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/CodeInterpreterOutput.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/CodeInterpreterOutput.test.tsx @@ -295,6 +295,100 @@ describe("CodeInterpreterOutput", () => { }); }); + it("revokes created object URLs on unmount", async () => { + let urlCounter = 0; + (URL.createObjectURL as ReturnType).mockImplementation(() => `blob:test-${++urlCounter}`); + + const mockBlob = new Blob(["image data"], { type: "image/png" }); + (global.fetch as any).mockResolvedValue({ + ok: true, + blob: vi.fn().mockResolvedValue(mockBlob), + }); + + const annotations = [ + { + type: "container_file_citation" as const, + container_id: "container-1", + file_id: "file-1", + filename: "chart.png", + start_index: 0, + end_index: 10, + }, + { + type: "container_file_citation" as const, + container_id: "container-1", + file_id: "file-2", + filename: "plot.jpg", + start_index: 0, + end_index: 10, + }, + ]; + + const { unmount } = render( + , + ); + + await waitFor(() => { + expect(URL.createObjectURL).toHaveBeenCalledTimes(2); + }); + + unmount(); + + expect(URL.revokeObjectURL).toHaveBeenCalledWith("blob:test-1"); + expect(URL.revokeObjectURL).toHaveBeenCalledWith("blob:test-2"); + }); + + it("does not leak object URLs from fetches that complete after unmount", async () => { + let urlCounter = 0; + (URL.createObjectURL as ReturnType).mockImplementation(() => `blob:late-${++urlCounter}`); + + const mockBlob = new Blob(["image data"], { type: "image/png" }); + let resolveBlob: (value: Blob) => void; + (global.fetch as any).mockResolvedValue({ + ok: true, + blob: vi.fn().mockReturnValue( + new Promise((resolve) => { + resolveBlob = resolve; + }), + ), + }); + + const annotations = [ + { + type: "container_file_citation" as const, + container_id: "container-1", + file_id: "file-1", + filename: "chart.png", + start_index: 0, + end_index: 10, + }, + ]; + + const { unmount } = render( + , + ); + + await waitFor(() => { + expect(global.fetch).toHaveBeenCalledTimes(1); + }); + + unmount(); + resolveBlob!(mockBlob); + await new Promise((resolve) => setTimeout(resolve, 0)); + + const created = (URL.createObjectURL as ReturnType).mock.results.map((r) => r.value); + const revoked = (URL.revokeObjectURL as ReturnType).mock.calls.map((c) => c[0]); + expect(revoked).toEqual(expect.arrayContaining(created)); + }); + it("should handle fetch errors gracefully", async () => { const consoleErrorSpy = vi.spyOn(console, "error").mockImplementation(() => {}); (global.fetch as any).mockRejectedValue(new Error("Network error")); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/CodeInterpreterOutput.tsx b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/CodeInterpreterOutput.tsx index c27273a116d..95e3e132de7 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/CodeInterpreterOutput.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/CodeInterpreterOutput.tsx @@ -39,6 +39,9 @@ const CodeInterpreterOutput: React.FC = ({ // Fetch images from container files API useEffect(() => { + const createdUrls: string[] = []; + const abortController = new AbortController(); + const fetchImages = async () => { for (const annotation of annotations) { const isImage = @@ -58,16 +61,21 @@ const CodeInterpreterOutput: React.FC = ({ headers: { [getGlobalLitellmHeaderName()]: `Bearer ${accessToken}`, }, + signal: abortController.signal, }, ); if (response.ok) { const blob = await response.blob(); + if (abortController.signal.aborted) return; const url = URL.createObjectURL(blob); + createdUrls.push(url); setImageUrls((prev) => ({ ...prev, [annotation.file_id]: url })); } } catch (error) { - console.error("Error fetching image:", error); + if (!abortController.signal.aborted) { + console.error("Error fetching image:", error); + } } finally { setLoadingImages((prev) => ({ ...prev, [annotation.file_id]: false })); } @@ -81,7 +89,8 @@ const CodeInterpreterOutput: React.FC = ({ // Cleanup URLs on unmount return () => { - Object.values(imageUrls).forEach((url) => URL.revokeObjectURL(url)); + abortController.abort(); + createdUrls.forEach((url) => URL.revokeObjectURL(url)); }; }, [annotations, accessToken, proxyBaseUrl]);