From 8576d1a9e97294acdd578a1505275be14dc3ec8f Mon Sep 17 00:00:00 2001
From: Danielle Maywood <danielle@themaywoods.com>
Date: Wed, 25 Mar 2026 17:35:57 +0000
Subject: [PATCH] fix(site): persist file attachments across navigations on
 create form (#23609)

---
 site/src/pages/AgentsPage/AgentsPage.test.ts  | 223 +++++++++++++++++-
 .../AgentsPage/components/AgentChatInput.tsx  |   2 +-
 .../components/AgentCreateForm.stories.tsx    |  57 +++++
 .../AgentsPage/components/AgentCreateForm.tsx |   9 +-
 .../AgentsPage/hooks/useFileAttachments.ts    | 164 ++++++++++++-
 5 files changed, 441 insertions(+), 14 deletions(-)

diff --git a/site/src/pages/AgentsPage/AgentsPage.test.ts b/site/src/pages/AgentsPage/AgentsPage.test.ts
index 33c4fc1d56..2c365bf717 100644
--- a/site/src/pages/AgentsPage/AgentsPage.test.ts
+++ b/site/src/pages/AgentsPage/AgentsPage.test.ts
@@ -1,9 +1,13 @@
 import { act, renderHook } from "@testing-library/react";
-import { beforeEach, describe, expect, it } from "vitest";
+import { afterEach, beforeEach, describe, expect, it, vi } from "vitest";
 import {
 	emptyInputStorageKey,
 	useEmptyStateDraft,
 } from "./components/AgentCreateForm";
+import {
+	persistedAttachmentsStorageKey,
+	useFileAttachments,
+} from "./hooks/useFileAttachments";
 
 describe("useEmptyStateDraft", () => {
 	beforeEach(() => {
@@ -166,3 +170,220 @@ describe("useEmptyStateDraft", () => {
 		unmount();
 	});
 });
+
+describe("useFileAttachments persistence", () => {
+	beforeEach(() => {
+		localStorage.clear();
+	});
+
+	afterEach(() => {
+		vi.restoreAllMocks();
+	});
+
+	const renderFileAttachments = () =>
+		renderHook(() => useFileAttachments("org-1", { persist: true }));
+
+	const makePersistedEntry = (
+		overrides: Partial<{
+			fileId: string;
+			fileName: string;
+			fileType: string;
+			lastModified: number;
+		}> = {},
+	) => ({
+		fileId: "file-1",
+		fileName: "photo.png",
+		fileType: "image/png",
+		lastModified: 1000,
+		...overrides,
+	});
+
+	it("restores uploaded attachments from localStorage on mount", () => {
+		const entry = makePersistedEntry();
+		localStorage.setItem(
+			persistedAttachmentsStorageKey,
+			JSON.stringify([entry]),
+		);
+
+		const { result, unmount } = renderFileAttachments();
+
+		expect(result.current.attachments).toHaveLength(1);
+		expect(result.current.attachments[0].name).toBe("photo.png");
+		expect(result.current.attachments[0].type).toBe("image/png");
+
+		const file = result.current.attachments[0];
+		const state = result.current.uploadStates.get(file);
+		expect(state).toEqual({ status: "uploaded", fileId: "file-1" });
+
+		const previewUrl = result.current.previewUrls.get(file);
+		expect(previewUrl).toBe("/api/experimental/chats/files/file-1");
+		unmount();
+	});
+
+	it("does not create preview URLs for non-image attachments", () => {
+		const entry = makePersistedEntry({
+			fileType: "text/plain",
+			fileName: "notes.txt",
+		});
+		localStorage.setItem(
+			persistedAttachmentsStorageKey,
+			JSON.stringify([entry]),
+		);
+
+		const { result, unmount } = renderFileAttachments();
+
+		expect(result.current.attachments).toHaveLength(1);
+		const file = result.current.attachments[0];
+		expect(result.current.previewUrls.has(file)).toBe(false);
+		expect(result.current.uploadStates.get(file)).toEqual({
+			status: "uploaded",
+			fileId: "file-1",
+		});
+		unmount();
+	});
+
+	it("returns empty state when nothing is persisted", () => {
+		const { result, unmount } = renderFileAttachments();
+
+		expect(result.current.attachments).toHaveLength(0);
+		expect(result.current.uploadStates.size).toBe(0);
+		expect(result.current.previewUrls.size).toBe(0);
+		unmount();
+	});
+
+	it("does not restore when persist option is false", () => {
+		const entry = makePersistedEntry();
+		localStorage.setItem(
+			persistedAttachmentsStorageKey,
+			JSON.stringify([entry]),
+		);
+
+		const { result, unmount } = renderHook(() =>
+			useFileAttachments("org-1", { persist: false }),
+		);
+
+		expect(result.current.attachments).toHaveLength(0);
+		unmount();
+	});
+
+	it("does not restore when no options argument is passed", () => {
+		const entry = makePersistedEntry();
+		localStorage.setItem(
+			persistedAttachmentsStorageKey,
+			JSON.stringify([entry]),
+		);
+
+		const { result, unmount } = renderHook(() => useFileAttachments("org-1"));
+
+		expect(result.current.attachments).toHaveLength(0);
+		unmount();
+	});
+
+	it("clears persisted attachments on resetAttachments", () => {
+		const entry = makePersistedEntry();
+		localStorage.setItem(
+			persistedAttachmentsStorageKey,
+			JSON.stringify([entry]),
+		);
+
+		const { result, unmount } = renderFileAttachments();
+
+		act(() => {
+			result.current.resetAttachments();
+		});
+
+		expect(localStorage.getItem(persistedAttachmentsStorageKey)).toBeNull();
+		expect(result.current.attachments).toHaveLength(0);
+		unmount();
+	});
+
+	it("removes the correct entry when an attachment is removed", () => {
+		const entries = [
+			makePersistedEntry({ fileId: "file-1", fileName: "a.png" }),
+			makePersistedEntry({ fileId: "file-2", fileName: "b.png" }),
+		];
+		localStorage.setItem(
+			persistedAttachmentsStorageKey,
+			JSON.stringify(entries),
+		);
+
+		const { result, unmount } = renderFileAttachments();
+		expect(result.current.attachments).toHaveLength(2);
+
+		act(() => {
+			result.current.handleRemoveAttachment(0);
+		});
+
+		expect(result.current.attachments).toHaveLength(1);
+		expect(result.current.attachments[0].name).toBe("b.png");
+
+		const stored = JSON.parse(
+			localStorage.getItem(persistedAttachmentsStorageKey)!,
+		);
+		expect(stored).toHaveLength(1);
+		expect(stored[0].fileId).toBe("file-2");
+		unmount();
+	});
+
+	it("handles corrupt localStorage gracefully", () => {
+		localStorage.setItem(persistedAttachmentsStorageKey, "not-valid-json");
+
+		const { result, unmount } = renderFileAttachments();
+
+		expect(result.current.attachments).toHaveLength(0);
+		unmount();
+	});
+
+	it("persists attachment metadata after successful upload", async () => {
+		const { API } = await import("#/api/api");
+		vi.spyOn(API.experimental, "uploadChatFile").mockResolvedValue({
+			id: "new-file-id",
+		});
+		vi.spyOn(globalThis, "fetch").mockResolvedValue(new Response());
+
+		const { result, unmount } = renderFileAttachments();
+
+		const file = new File(["hello"], "test.png", { type: "image/png" });
+
+		act(() => {
+			result.current.handleAttach([file]);
+		});
+
+		// Wait for the async upload to complete and state to update.
+		await vi.waitFor(() => {
+			const state = result.current.uploadStates.get(file);
+			expect(state?.status).toBe("uploaded");
+		});
+
+		const stored = JSON.parse(
+			localStorage.getItem(persistedAttachmentsStorageKey)!,
+		);
+		expect(stored).toHaveLength(1);
+		expect(stored[0].fileId).toBe("new-file-id");
+		expect(stored[0].fileName).toBe("test.png");
+		unmount();
+	});
+
+	it("does not persist attachment metadata when upload fails", async () => {
+		const { API } = await import("#/api/api");
+		vi.spyOn(API.experimental, "uploadChatFile").mockRejectedValue(
+			new Error("server error"),
+		);
+
+		const { result, unmount } = renderFileAttachments();
+
+		const file = new File(["hello"], "test.png", { type: "image/png" });
+
+		act(() => {
+			result.current.handleAttach([file]);
+		});
+
+		await vi.waitFor(() => {
+			const state = result.current.uploadStates.get(file);
+			expect(state?.status).toBe("error");
+		});
+
+		expect(localStorage.getItem(persistedAttachmentsStorageKey)).toBeNull();
+		unmount();
+	});
+});
diff --git a/site/src/pages/AgentsPage/components/AgentChatInput.tsx b/site/src/pages/AgentsPage/components/AgentChatInput.tsx
index 202dc591db..6130a88db8 100644
--- a/site/src/pages/AgentsPage/components/AgentChatInput.tsx
+++ b/site/src/pages/AgentsPage/components/AgentChatInput.tsx
@@ -443,7 +443,7 @@ export const AttachmentPreview: FC<{
 						)}
 						<button
 							type="button"
-							onClick={() => onRemove(index)}
+							onClick={() => onRemove(file)}
 							className="absolute -right-2 -top-2 flex h-6 w-6 cursor-pointer items-center justify-center rounded-full border-0 bg-surface-primary text-content-secondary shadow-sm opacity-0 transition-opacity hover:bg-surface-secondary hover:text-content-primary group-hover:opacity-100 group-focus-within:opacity-100 focus:opacity-100"
 							aria-label={`Remove ${file.name}`}
 						>
diff --git a/site/src/pages/AgentsPage/components/AgentCreateForm.stories.tsx b/site/src/pages/AgentsPage/components/AgentCreateForm.stories.tsx
index 9683186f3d..52ed1bf56a 100644
--- a/site/src/pages/AgentsPage/components/AgentCreateForm.stories.tsx
+++ b/site/src/pages/AgentsPage/components/AgentCreateForm.stories.tsx
@@ -189,6 +189,63 @@ export const NoModelsConfigured: Story = {
 	},
 };
 
+export const PreservesAttachmentsOnFailedSend: Story = {
+	args: {
+		...defaultArgs,
+		onCreateChat: fn().mockRejectedValue(new Error("server error")),
+	},
+	beforeEach: () => {
+		localStorage.clear();
+		// Pre-persist an uploaded attachment so it is restored on mount.
+		localStorage.setItem(
+			"agents.persisted-attachments",
+			JSON.stringify([
+				{
+					fileId: "persisted-file-1",
+					fileName: "photo.png",
+					fileType: "image/png",
+					lastModified: 1000,
+				},
+			]),
+		);
+		spyOn(API, "getWorkspaces").mockResolvedValue({
+			workspaces: [],
+			count: 0,
+		});
+	},
+	play: async ({ canvasElement, args }) => {
+		const canvas = within(canvasElement);
+
+		// The restored attachment should appear on mount.
+		await waitFor(() => {
+			expect(canvas.getByLabelText("Remove photo.png")).toBeInTheDocument();
+		});
+
+		// Type a message and submit.
+		const input = canvas.getByTestId("chat-message-input");
+		await userEvent.click(input);
+		await userEvent.keyboard("test message");
+		await userEvent.click(canvas.getByRole("button", { name: "Send" }));
+
+		// Wait for onCreateChat to have been called (and rejected).
+		await waitFor(() => {
+			expect(args.onCreateChat).toHaveBeenCalled();
+		});
+
+		// The attachment must still be visible after the failed send.
+		await waitFor(() => {
+			expect(canvas.getByLabelText("Remove photo.png")).toBeInTheDocument();
+		});
+
+		// localStorage must still have the persisted attachment.
+		const stored = localStorage.getItem("agents.persisted-attachments");
+		expect(stored).not.toBeNull();
+		const parsed = JSON.parse(stored!);
+		expect(parsed).toHaveLength(1);
+		expect(parsed[0].fileId).toBe("persisted-file-1");
+	},
+};
+
 export const UsageLimitExceeded: Story = {
 	args: {
 		...defaultArgs,
diff --git a/site/src/pages/AgentsPage/components/AgentCreateForm.tsx b/site/src/pages/AgentsPage/components/AgentCreateForm.tsx
index 2b1f71e3ef..32c0fad085 100644
--- a/site/src/pages/AgentsPage/components/AgentCreateForm.tsx
+++ b/site/src/pages/AgentsPage/components/AgentCreateForm.tsx
@@ -255,10 +255,13 @@ export const AgentCreateForm: FC<AgentCreateFormProps> = ({
 				selectedMCPServerIdsRef.current.length > 0
 					? [...selectedMCPServerIdsRef.current]
 					: undefined,
-		}).catch(() => {
+		}).catch((err) => {
 			// Re-enable draft persistence so the user can edit
-			// and retry after a failed send attempt.
+			// and retry after a failed send attempt, then rethrow
+			// so callers (handleSendWithAttachments) can preserve
+			// attachments on failure.
 			resetDraft();
+			throw err;
 		});
 	};
 
@@ -270,7 +273,7 @@ export const AgentCreateForm: FC<AgentCreateFormProps> = ({
 		handleAttach,
 		handleRemoveAttachment,
 		resetAttachments,
-	} = useFileAttachments(organizations[0]?.id);
+	} = useFileAttachments(organizations[0]?.id, { persist: true });
 
 	const handleSendWithAttachments = async (message: string) => {
 		const fileIds: string[] = [];
diff --git a/site/src/pages/AgentsPage/hooks/useFileAttachments.ts b/site/src/pages/AgentsPage/hooks/useFileAttachments.ts
index faeb01733d..731d3d286d 100644
--- a/site/src/pages/AgentsPage/hooks/useFileAttachments.ts
+++ b/site/src/pages/AgentsPage/hooks/useFileAttachments.ts
@@ -9,6 +9,114 @@ import {
 } from "react";
 import type { UploadState } from "../components/AgentChatInput";
 
+/** @internal Exported for testing. */
+export const persistedAttachmentsStorageKey = "agents.persisted-attachments";
+
+/**
+ * Serializable metadata stored in localStorage so that already-uploaded
+ * attachments survive page navigations on the create form.
+ */
+interface PersistedAttachment {
+	fileId: string;
+	fileName: string;
+	fileType: string;
+	lastModified: number;
+}
+
+/**
+ * Restore previously persisted attachments from localStorage.
+ * Creates synthetic File objects (empty blobs with correct metadata)
+ * and populates the corresponding Maps so the UI can render them.
+ */
+function restorePersistedAttachments(): {
+	attachments: File[];
+	uploadStates: Map<File, UploadState>;
+	previewUrls: Map<File, string>;
+} {
+	const stored = localStorage.getItem(persistedAttachmentsStorageKey);
+	if (!stored) {
+		return {
+			attachments: [],
+			uploadStates: new Map(),
+			previewUrls: new Map(),
+		};
+	}
+	try {
+		const persisted: PersistedAttachment[] = JSON.parse(stored);
+		const attachments: File[] = [];
+		const uploadStates = new Map<File, UploadState>();
+		const previewUrls = new Map<File, string>();
+
+		for (const p of persisted) {
+			if (!p.fileId || !p.fileName) continue;
+			// Synthetic File used as a Map key only. Its content is
+			// never read because the existing file_id is reused at
+			// send time.
+			const file = new File([], p.fileName, {
+				type: p.fileType,
+				lastModified: p.lastModified,
+			});
+			attachments.push(file);
+			uploadStates.set(file, { status: "uploaded", fileId: p.fileId });
+			if (p.fileType.startsWith("image/")) {
+				previewUrls.set(file, `/api/experimental/chats/files/${p.fileId}`);
+			}
+		}
+		return { attachments, uploadStates, previewUrls };
+	} catch {
+		return {
+			attachments: [],
+			uploadStates: new Map(),
+			previewUrls: new Map(),
+		};
+	}
+}
+
+function addPersistedAttachment(file: File, fileId: string) {
+	const stored = localStorage.getItem(persistedAttachmentsStorageKey);
+	let persisted: PersistedAttachment[];
+	try {
+		persisted = stored ? JSON.parse(stored) : [];
+	} catch {
+		persisted = [];
+	}
+	persisted.push({
+		fileId,
+		fileName: file.name,
+		fileType: file.type,
+		lastModified: file.lastModified,
+	});
+	localStorage.setItem(
+		persistedAttachmentsStorageKey,
+		JSON.stringify(persisted),
+	);
+}
+
+function removePersistedAttachment(fileId: string) {
+	const stored = localStorage.getItem(persistedAttachmentsStorageKey);
+	if (!stored) {
+		return;
+	}
+	try {
+		const persisted: PersistedAttachment[] = JSON.parse(stored);
+		const filtered = persisted.filter((p) => p.fileId !== fileId);
+		if (filtered.length > 0) {
+			localStorage.setItem(
+				persistedAttachmentsStorageKey,
+				JSON.stringify(filtered),
+			);
+		} else {
+			localStorage.removeItem(persistedAttachmentsStorageKey);
+		}
+	} catch {
+		localStorage.removeItem(persistedAttachmentsStorageKey);
+	}
+}
+
+function clearPersistedAttachments() {
+	localStorage.removeItem(persistedAttachmentsStorageKey);
+}
+
 interface UseFileAttachmentsReturn {
 	attachments: File[];
 	textContents: Map<File, string>;
@@ -25,12 +133,25 @@ interface UseFileAttachmentsReturn {
 
 export function useFileAttachments(
 	organizationId: string | undefined,
+	options?: { persist?: boolean },
 ): UseFileAttachmentsReturn {
-	const [attachments, setAttachments] = useState<File[]>([]);
-	const [uploadStates, setUploadStates] = useState(
-		() => new Map<File, UploadState>(),
+	const persist = options?.persist ?? false;
+
+	// Restore previously uploaded attachments from localStorage
+	// when persistence is enabled. Computed once on first render.
+	const [restored] = useState(() =>
+		persist
+			? restorePersistedAttachments()
+			: {
+					attachments: [] as File[],
+					uploadStates: new Map<File, UploadState>(),
+					previewUrls: new Map<File, string>(),
+				},
 	);
-	const [previewUrls, setPreviewUrls] = useState(() => new Map<File, string>());
+
+	const [attachments, setAttachments] = useState<File[]>(restored.attachments);
+	const [uploadStates, setUploadStates] = useState(restored.uploadStates);
+	const [previewUrls, setPreviewUrls] = useState(restored.previewUrls);
 	const [textContents, setTextContents] = useState(
 		() => new Map<File, string>(),
 	);
@@ -71,6 +192,9 @@ export function useFileAttachments(
 						fileId: result.id,
 					}),
 				);
+				if (persist) {
+					addPersistedAttachment(file, result.id);
+				}
 				// Pre-warm the browser HTTP cache for images so the
 				// timeline can render them instantly after send. We
 				// intentionally skip text attachments because the
@@ -140,6 +264,22 @@ export function useFileAttachments(
 	};
 
 	const handleRemoveAttachment = (attachment: number | File) => {
+		// Resolve the file to remove and perform localStorage side
+		// effects before entering state updaters. React may call
+		// updaters more than once (StrictMode, React Compiler), so
+		// they must stay pure.
+		const idx =
+			typeof attachment === "number"
+				? attachment
+				: attachments.indexOf(attachment);
+		const removed = idx >= 0 ? attachments[idx] : undefined;
+		if (persist && removed) {
+			const state = uploadStates.get(removed);
+			if (state?.status === "uploaded" && state.fileId) {
+				removePersistedAttachment(state.fileId);
+			}
+		}
+
 		setAttachments((prev) => {
 			const index =
 				typeof attachment === "number" ? attachment : prev.indexOf(attachment);
@@ -147,22 +287,22 @@ export function useFileAttachments(
 				return prev;
 			}
 
-			const removed = prev[index];
+			const removedFile = prev[index];
 			setUploadStates((prevStates) => {
 				const next = new Map(prevStates);
-				next.delete(removed);
+				next.delete(removedFile);
 				return next;
 			});
 			setPreviewUrls((prevUrls) => {
-				const url = prevUrls.get(removed);
+				const url = prevUrls.get(removedFile);
 				if (url?.startsWith("blob:")) URL.revokeObjectURL(url);
 				const next = new Map(prevUrls);
-				next.delete(removed);
+				next.delete(removedFile);
 				return next;
 			});
 			setTextContents((prevContents) => {
 				const next = new Map(prevContents);
-				next.delete(removed);
+				next.delete(removedFile);
 				return next;
 			});
 			return prev.filter((_, i) => i !== index);
@@ -177,6 +317,9 @@ export function useFileAttachments(
 		setTextContents(new Map());
 		setUploadStates(new Map());
 		setAttachments([]);
+		if (persist) {
+			clearPersistedAttachments();
+		}
 	};
 
 	return {
@@ -188,6 +331,9 @@ export function useFileAttachments(
 		handleRemoveAttachment,
 		startUpload,
 		resetAttachments,
+		// Raw setters exposed for AgentDetailContent to pre-populate
+		// attachments from existing chat messages. These bypass
+		// localStorage persistence. Only use when persist is false.
 		setAttachments,
 		setPreviewUrls,
 		setUploadStates,