diff --git a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/ChatUI.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/ChatUI.test.tsx index 9d9c5dc86ab..56bc9503a79 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/ChatUI.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/ChatUI.test.tsx @@ -441,6 +441,49 @@ describe("ChatUI", () => { expect(customProxyInput).toHaveValue(testProxyUrl); }); + it("revokes the active OCR preview URL on unmount", async () => { + const createObjectURL = vi.fn(() => "blob:http://localhost/ocr-preview"); + const revokeObjectURL = vi.fn(); + URL.createObjectURL = createObjectURL; + URL.revokeObjectURL = revokeObjectURL; + + const { getByText, unmount } = render( + , + ); + + await waitFor(() => { + expect(getByText("Test Key")).toBeInTheDocument(); + }); + + await selectComboboxOption("Select an endpoint", "/v1/ocr"); + + await waitFor(() => { + expect(screen.getByText("Click or drag a document or image to upload")).toBeInTheDocument(); + }); + + const uploadInput = screen.getByLabelText(/Click or drag a document or image to upload/); + + const imageFile = new File(["ocr"], "ocr.png", { type: "image/png" }); + await act(async () => { + fireEvent.change(uploadInput, { target: { files: [imageFile] } }); + }); + + await waitFor(() => { + expect(createObjectURL).toHaveBeenCalledWith(imageFile); + expect(screen.getByAltText("Upload preview")).toBeInTheDocument(); + }); + + unmount(); + + expect(revokeObjectURL).toHaveBeenCalledWith("blob:http://localhost/ocr-preview"); + }); + it("should enable search functionality for MCP server selector", async () => { const user = userEvent.setup(); render( diff --git a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/ChatUI.tsx b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/ChatUI.tsx index 5920fb57126..def4e60e27a 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/ChatUI.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/ChatUI.tsx @@ -275,6 +275,24 @@ const ChatUI: React.FC = ({ const codeInterpreter = useCodeInterpreter(); const chatEndRef = useRef(null); + const ocrFilePreviewUrlRef = useRef(null); + + const revokeOcrFilePreviewUrl = () => { + if (ocrFilePreviewUrlRef.current) { + URL.revokeObjectURL(ocrFilePreviewUrlRef.current); + } + ocrFilePreviewUrlRef.current = null; + setOcrFilePreviewUrl(null); + }; + + useEffect(() => { + return () => { + if (ocrFilePreviewUrlRef.current) { + URL.revokeObjectURL(ocrFilePreviewUrlRef.current); + ocrFilePreviewUrlRef.current = null; + } + }; + }, []); // Fetch MCP servers and toolsets const loadMCPServers = async () => { @@ -704,8 +722,11 @@ const ChatUI: React.FC = ({ }; const handleOcrFileUpload = (file: File): false => { + revokeOcrFilePreviewUrl(); setUploadedOcrFile(file); - setOcrFilePreviewUrl(file.type.startsWith("image/") ? URL.createObjectURL(file) : null); + const previewUrl = file.type.startsWith("image/") ? URL.createObjectURL(file) : null; + ocrFilePreviewUrlRef.current = previewUrl; + setOcrFilePreviewUrl(previewUrl); return false; }; @@ -718,11 +739,8 @@ const ChatUI: React.FC = ({ }; const handleRemoveOcrFile = () => { - if (ocrFilePreviewUrl) { - URL.revokeObjectURL(ocrFilePreviewUrl); - } setUploadedOcrFile(null); - setOcrFilePreviewUrl(null); + revokeOcrFilePreviewUrl(); }; const handleSendMessage = async () => { diff --git a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/EndpointSelector.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/EndpointSelector.test.tsx index 7b6d71bfb2e..b330b2c8fb5 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/EndpointSelector.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/EndpointSelector.test.tsx @@ -35,6 +35,7 @@ describe("EndpointSelector", () => { await user.click(combobox); const input = await screen.findByRole("combobox"); + await user.clear(input); await user.type(input, "ocr"); expect(await screen.findByText("/v1/ocr")).toBeInTheDocument();