mirror of
https://github.com/coder/coder.git
synced 2026-09-24 15:04:27 +08:00
fix(site): persist file attachments across navigations on create form (#23609)
This commit is contained in:
@@ -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();
|
||||
});
|
||||
});
|
||||
|
||||
@@ -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}`}
|
||||
>
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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[] = [];
|
||||
|
||||
@@ -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,
|
||||
|
||||
Reference in New Issue
Block a user