1import { ApprovalRequestId, EnvironmentId, ThreadId } from "@t3tools/contracts";
2import { beforeEach, expect, it, vi } from "vite-plus/test";
3import { useComposerDraftStore } from "./composerDraftStore";
4import {
5 questionAttachmentDraftId,
6 countQuestionAttachments,
7 questionAttachmentDraftPrefix,
8 changeQuestionAttachmentPreparation,
9 clearQuestionAttachmentDraft,
10 useQuestionAttachmentPreparation,
11} from "./questionAttachments";
12
13const release = vi.hoisted(() => vi.fn());
14vi.mock("./lib/attachmentUploadQueue", () => ({ releaseDraftAttachments: release }));
15const environmentId = EnvironmentId.make("environment-1");
16const threadId = ThreadId.make("thread-1");
17const requestId = ApprovalRequestId.make("request-1");
18beforeEach(() => {
19 useComposerDraftStore.setState({ draftsByThreadKey: {}, draftThreadsByThreadKey: {} });
20 useQuestionAttachmentPreparation.setState({ counts: {} });
21 release.mockClear();
22});
23it("keeps question files separate from the normal draft and other questions", () => {
24 const revoke = vi.spyOn(URL, "revokeObjectURL");
25 const store = useComposerDraftStore.getState();
26 const first = questionAttachmentDraftId(environmentId, threadId, requestId, "first");
27 const second = questionAttachmentDraftId(environmentId, threadId, requestId, "second");
28 store.setPrompt({ environmentId, threadId }, "Unsent prompt");
29 const image = {
30 type: "image" as const,
31 id: "image-1",
32 name: "image.png",
33 mimeType: "image/png",
34 sizeBytes: 5,
35 previewUrl: "blob:question-image",
36 file: new File(["image"], "image.png", { type: "image/png" }),
37 };
38 store.addImages(first, [image]);
39 store.addFiles(second, [
40 {
41 type: "file",
42 id: "file-1",
43 name: "spec.txt",
44 mimeType: "text/plain",
45 sizeBytes: 4,
46 file: new File(["spec"], "spec.txt", { type: "text/plain" }),
47 },
48 ]);
49 changeQuestionAttachmentPreparation(first, 1);
50 expect(store.getComposerDraft(first)?.images).toHaveLength(1);
51 expect(store.getComposerDraft(second)?.files).toHaveLength(1);
52 clearQuestionAttachmentDraft(first);
53 expect(store.getComposerDraft(first)).toBeNull();
54 expect(store.getComposerDraft(second)?.files).toHaveLength(1);
55 expect(store.getComposerDraft({ environmentId, threadId })?.prompt).toBe("Unsent prompt");
56 expect(useQuestionAttachmentPreparation.getState().counts[first]).toBeUndefined();
57 expect(release).toHaveBeenCalledWith([image]);
58 expect(revoke).toHaveBeenCalledWith(image.previewUrl);
59 revoke.mockRestore();
60});
61it("scopes provider request ids to their environment and thread", () => {
62 const keys = [
63 questionAttachmentDraftId(environmentId, threadId, requestId, "q"),
64 questionAttachmentDraftId(EnvironmentId.make("environment-2"), threadId, requestId, "q"),
65 questionAttachmentDraftId(environmentId, ThreadId.make("thread-2"), requestId, "q"),
66 ];
67 expect(new Set(keys).size).toBe(3);
68});
69it("does not match drafts belonging to threads with a shared id prefix", () => {
70 const prefix = questionAttachmentDraftPrefix(environmentId, threadId);
71 expect(
72 questionAttachmentDraftId(environmentId, threadId, requestId, "q").startsWith(prefix),
73 ).toBe(true);
74 for (const suffix of ["-extra", ":extra", "/extra", "%extra"]) {
75 expect(
76 questionAttachmentDraftId(
77 environmentId,
78 ThreadId.make(`${threadId}${suffix}`),
79 requestId,
80 "q",
81 ).startsWith(prefix),
82 ).toBe(false);
83 }
84});
85
86it("does not clear another environment whose id contains a question prefix", () => {
87 const otherEnvironment = EnvironmentId.make(
88 `${environmentId}:question-${encodeURIComponent(JSON.stringify(threadId))}-nested`,
89 );
90 const otherKey = questionAttachmentDraftId(otherEnvironment, threadId, requestId, "q");
91 const ownKey = questionAttachmentDraftId(environmentId, threadId, requestId, "q");
92 const store = useComposerDraftStore.getState();
93 store.setPrompt(otherKey, "Keep this answer");
94 store.setPrompt(ownKey, "Discard this answer");
95 const prefix = questionAttachmentDraftPrefix(environmentId, threadId);
96 for (const key of [ownKey, otherKey]) {
97 if (key.startsWith(prefix)) clearQuestionAttachmentDraft(key);
98 }
99 expect(store.getComposerDraft(ownKey)).toBeNull();
100 expect(store.getComposerDraft(otherKey)?.prompt).toBe("Keep this answer");
101});
102
103it("counts files, images, and pending preparation across the shared question budget", () => {
104 const first = questionAttachmentDraftId(environmentId, threadId, requestId, "first");
105 const second = questionAttachmentDraftId(environmentId, threadId, requestId, "second");
106 const other = questionAttachmentDraftId(
107 EnvironmentId.make("other"),
108 threadId,
109 requestId,
110 "first",
111 );
112 const store = useComposerDraftStore.getState();
113 store.addFiles(
114 second,
115 Array.from({ length: 6 }, (_, index) => ({
116 type: "file" as const,
117 id: `file-${index}`,
118 name: `spec-${index}.txt`,
119 mimeType: "text/plain",
120 sizeBytes: 4,
121 file: new File(["spec"], `spec-${index}.txt`),
122 })),
123 );
124 store.addImages(first, [
125 {
126 type: "image",
127 id: "image",
128 name: "image.png",
129 mimeType: "image/png",
130 sizeBytes: 5,
131 previewUrl: "blob:count-test",
132 file: new File(["image"], "image.png"),
133 },
134 ]);
135 changeQuestionAttachmentPreparation(second, 1);
136 changeQuestionAttachmentPreparation(other, 8);
137 expect(countQuestionAttachments([first, second])).toBe(8);
138 expect(countQuestionAttachments([second])).toBe(7);
139 changeQuestionAttachmentPreparation(second, -1);
140 expect(countQuestionAttachments([first, second])).toBe(7);
141});