From 02390f5a4f9731dca9ede01b421d5de72ba88dec Mon Sep 17 00:00:00 2001
From: qdivan <77005282+qdivan@users.noreply.github.com>
Date: Fri, 24 Jul 2026 22:33:42 +0800
Subject: [PATCH] fix(ui): revoke OCR preview URL on unmount

---
 .../components/chat_ui/ChatUI.test.tsx        | 43 +++++++++++++++++++
 .../playground/components/chat_ui/ChatUI.tsx  | 28 +++++++++---
 .../chat_ui/EndpointSelector.test.tsx         |  1 +
 3 files changed, 67 insertions(+), 5 deletions(-)

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(
+      <ChatUI
+        accessToken="1234567890"
+        token="1234567890"
+        userRole="user"
+        userID="1234567890"
+        disabledPersonalKeyCreation={false}
+      />,
+    );
+
+    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<ChatUIProps> = ({
   const codeInterpreter = useCodeInterpreter();
 
   const chatEndRef = useRef<HTMLDivElement>(null);
+  const ocrFilePreviewUrlRef = useRef<string | null>(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<ChatUIProps> = ({
   };
 
   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<ChatUIProps> = ({
   };
 
   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();