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]);