1import { afterEach, beforeEach, describe, expect, it } from "vite-plus/test";
2import { EnvironmentId, MessageId, ThreadId } from "@t3tools/contracts";
3import {
4 collectAssistantCitations,
5 serializeAssistantCitation,
6} from "@t3tools/shared/assistantCitations";
7
8import { removeLocalStorageItem } from "./hooks/useLocalStorage";
9
10import {
11 MAX_STASH_ENTRIES,
12 PROMPT_STASH_STORAGE_KEY,
13 MAX_STASH_ENTRY_ATTACHMENT_CHARS,
14 partitionStashAttachments,
15 usePromptStashStore,
16 writePromptStashStorageForTest,
17 type PromptStashEntry,
18} from "./promptStashStore";
19
20function makeEntry(input: {
21 id: string;
22 prompt?: string;
23 attachmentChars?: number;
24}): PromptStashEntry {
25 return {
26 id: input.id,
27 createdAt: "2026-07-24T12:00:00.000Z",
28 prompt: input.prompt ?? `prompt ${input.id}`,
29 attachments:
30 input.attachmentChars !== undefined
31 ? [
32 {
33 id: `${input.id}-img`,
34 name: "shot.png",
35 mimeType: "image/png",
36 sizeBytes: input.attachmentChars,
37 dataUrl: "x".repeat(input.attachmentChars),
38 },
39 ]
40 : [],
41 droppedImageNames: [],
42 };
43}
44
45function resetPromptStashStore() {
46 usePromptStashStore.setState({ entries: [] });
47 writePromptStashStorageForTest("");
48 removeLocalStorageItem(PROMPT_STASH_STORAGE_KEY);
49}
50
51describe("partitionStashAttachments", () => {
52 it("keeps attachments within the budget and reports dropped names in order", () => {
53 const small = {
54 id: "a",
55 name: "small.png",
56 mimeType: "image/png",
57 sizeBytes: 10,
58 dataUrl: "x".repeat(10),
59 };
60 const huge = {
61 id: "b",
62 name: "huge.png",
63 mimeType: "image/png",
64 sizeBytes: MAX_STASH_ENTRY_ATTACHMENT_CHARS,
65 dataUrl: "x".repeat(MAX_STASH_ENTRY_ATTACHMENT_CHARS),
66 };
67 const alsoSmall = {
68 id: "c",
69 name: "also-small.png",
70 mimeType: "image/png",
71 sizeBytes: 10,
72 dataUrl: "x".repeat(10),
73 };
74 const { kept, droppedNames } = partitionStashAttachments([small, huge, alsoSmall]);
75 expect(kept.map((attachment) => attachment.id)).toEqual(["a", "c"]);
76 expect(droppedNames).toEqual(["huge.png"]);
77 });
78
79 it("admits a single attachment that exactly fits the budget", () => {
80 const exact = {
81 id: "a",
82 name: "exact.png",
83 mimeType: "image/png",
84 sizeBytes: MAX_STASH_ENTRY_ATTACHMENT_CHARS,
85 dataUrl: "x".repeat(MAX_STASH_ENTRY_ATTACHMENT_CHARS),
86 };
87 const { kept, droppedNames } = partitionStashAttachments([exact]);
88 expect(kept).toHaveLength(1);
89 expect(droppedNames).toEqual([]);
90 });
91});
92
93describe("promptStashStore", () => {
94 beforeEach(() => {
95 resetPromptStashStore();
96 });
97
98 afterEach(() => {
99 resetPromptStashStore();
100 });
101
102 it("prepends entries so the newest stash is first", () => {
103 const store = usePromptStashStore.getState();
104 store.stashEntry(makeEntry({ id: "first" }));
105 store.stashEntry(makeEntry({ id: "second" }));
106 const entries = usePromptStashStore.getState().entries;
107 expect(entries.map((entry) => entry.id)).toEqual(["second", "first"]);
108 });
109
110 it("restores citations and bound comments from a persisted stash", () => {
111 const citation = {
112 version: 1 as const,
113 environmentId: EnvironmentId.make("remote-source"),
114 threadId: ThreadId.make("source-thread"),
115 messageId: MessageId.make("source-message"),
116 text: "A quote to revisit later.",
117 comment: "Revisit this tomorrow.\nCheck whether the example still applies.",
118 start: 0,
119 end: 24,
120 prefix: "",
121 suffix: " Next.",
122 };
123 const prompt = `Revisit ${serializeAssistantCitation(citation)}`;
124 writePromptStashStorageForTest(
125 JSON.stringify({
126 version: 2,
127 state: { entries: [makeEntry({ id: "citation-stash", prompt })] },
128 }),
129 );
130 const restored = usePromptStashStore.getState().takeEntry("citation-stash").entry;
131 expect(restored?.prompt).toBe(prompt);
132 expect(
133 collectAssistantCitations(restored?.prompt ?? "").map((entry) => entry.citation),
134 ).toEqual([citation]);
135 expect(usePromptStashStore.getState().takeEntry("citation-stash").entry).toBeNull();
136 });
137
138 it("evicts the oldest entry past the cap and returns it", () => {
139 const store = usePromptStashStore.getState();
140 for (let index = 0; index < MAX_STASH_ENTRIES; index += 1) {
141 expect(store.stashEntry(makeEntry({ id: `entry-${index}` })).evicted).toBeNull();
142 }
143 const { evicted } = store.stashEntry(makeEntry({ id: "overflow" }));
144 expect(evicted?.id).toBe("entry-0");
145 const entries = usePromptStashStore.getState().entries;
146 expect(entries).toHaveLength(MAX_STASH_ENTRIES);
147 expect(entries[0]?.id).toBe("overflow");
148 });
149
150 // This test environment has no `localStorage`, so the store runs on its
151 // in-memory fallback — the exact "kept for this session, gone on reload"
152 // case the composer must distinguish from an outright write failure.
153 it("distinguishes a memory-only write (written, not durable) from a failed one", () => {
154 const store = usePromptStashStore.getState();
155 const result = store.stashEntry(makeEntry({ id: "memory-only" }));
156 expect(result.written).toBe(true);
157 expect(result.durable).toBe(false);
158 expect(usePromptStashStore.getState().entries.map((entry) => entry.id)).toEqual([
159 "memory-only",
160 ]);
161 });
162
163 it("takeEntry removes and returns the entry; second take returns null", () => {
164 const store = usePromptStashStore.getState();
165 store.stashEntry(makeEntry({ id: "keep" }));
166 store.stashEntry(makeEntry({ id: "take" }));
167 expect(store.takeEntry("take").entry?.id).toBe("take");
168 expect(store.takeEntry("take").entry).toBeNull();
169 const entries = usePromptStashStore.getState().entries;
170 expect(entries.map((entry) => entry.id)).toEqual(["keep"]);
171 });
172
173 it("finalizeEntryImages attaches images and clears the pending count", () => {
174 const store = usePromptStashStore.getState();
175 store.stashEntry({ ...makeEntry({ id: "pending" }), pendingImageCount: 2 });
176
177 const { attached } = store.finalizeEntryImages("pending", {
178 attachments: [
179 {
180 id: "img-1",
181 name: "a.webp",
182 mimeType: "image/webp",
183 sizeBytes: 10,
184 dataUrl: "data:image/webp;base64,AAAA",
185 },
186 ],
187 droppedImageNames: ["big.png"],
188 unreadableImageNames: [],
189 });
190
191 expect(attached).toBe(true);
192 const entry = usePromptStashStore.getState().entries[0];
193 expect(entry?.attachments).toHaveLength(1);
194 expect(entry?.droppedImageNames).toEqual(["big.png"]);
195 expect(entry?.pendingImageCount).toBe(0);
196 });
197
198 it("takeEntry returns images and drop metadata finalized after a menu snapshot", () => {
199 const store = usePromptStashStore.getState();
200 store.stashEntry({ ...makeEntry({ id: "pending-restore" }), pendingImageCount: 1 });
201 const menuSnapshot = usePromptStashStore.getState().entries[0];
202 expect(menuSnapshot?.attachments).toEqual([]);
203
204 store.finalizeEntryImages("pending-restore", {
205 attachments: [
206 {
207 id: "img-finalized",
208 name: "finalized.webp",
209 mimeType: "image/webp",
210 sizeBytes: 12,
211 dataUrl: "data:image/webp;base64,BBBB",
212 },
213 ],
214 droppedImageNames: ["too-large.png"],
215 unreadableImageNames: ["unreadable.png"],
216 });
217
218 const { entry: taken } = store.takeEntry("pending-restore");
219 expect(taken).not.toBe(menuSnapshot);
220 expect(taken).toMatchObject({
221 attachments: [
222 {
223 id: "img-finalized",
224 name: "finalized.webp",
225 },
226 ],
227 droppedImageNames: ["too-large.png"],
228 unreadableImageNames: ["unreadable.png"],
229 pendingImageCount: 0,
230 });
231 });
232
233 it("preserves uploaded file references without storing file contents", () => {
234 const store = usePromptStashStore.getState();
235 const file = {
236 id: "file-1",
237 name: "report.pdf",
238 mimeType: "application/pdf",
239 sizeBytes: 42,
240 attachmentId: "pending-report-pdf",
241 environmentId: EnvironmentId.make("environment-1"),
242 };
243
244 store.stashEntry({ ...makeEntry({ id: "with-file" }), files: [file] });
245 store.finalizeEntryImages("with-file", {
246 attachments: [],
247 droppedImageNames: [],
248 unreadableImageNames: [],
249 });
250
251 expect(usePromptStashStore.getState().entries[0]?.files).toEqual([file]);
252 });
253
254 it("finalizeEntryImages reports false when the entry was already taken", () => {
255 const store = usePromptStashStore.getState();
256 store.stashEntry({ ...makeEntry({ id: "racing" }), pendingImageCount: 1 });
257 // Restored (or deleted) while its images were still encoding.
258 store.takeEntry("racing");
259
260 const { attached } = store.finalizeEntryImages("racing", {
261 attachments: [],
262 droppedImageNames: [],
263 unreadableImageNames: [],
264 });
265
266 expect(attached).toBe(false);
267 });
268
269 it("settles a pending count left behind by a crashed or closed session", () => {
270 writePromptStashStorageForTest(
271 JSON.stringify({
272 version: 2,
273 state: {
274 entries: [{ ...makeEntry({ id: "orphan" }), pendingImageCount: 2 }],
275 },
276 }),
277 );
278
279 // Hydration must settle the stale count, or the entry would stay stuck
280 // showing "saving…" with images that no longer exist anywhere.
281 const entry = usePromptStashStore.getState().entries[0];
282 expect(entry?.pendingImageCount).toBe(0);
283 expect(entry?.unreadableImageNames).toHaveLength(2);
284 });
285
286 it("ignores an unreadable v1 payload seeded under the current key", () => {
287 // The v1 shape (per-provider queues) does not decode as v2; hydration
288 // must fall back to an empty stash rather than throw.
289 writePromptStashStorageForTest(
290 JSON.stringify({
291 version: 1,
292 state: { queuesByScopeKey: { "provider:claudeAgent": [] } },
293 }),
294 );
295 expect(usePromptStashStore.getState().entries).toEqual([]);
296 });
297});
298
299describe("prompt stash context records", () => {
300 it("keeps the records behind a stashed prompt's chips", () => {
301 const store = usePromptStashStore.getState();
302 const records = [
303 {
304 version: 1 as const,
305 contextId: "ctx-1" as never,
306 kind: "terminal" as const,
307 label: "Terminal 1 line 4",
308 terminalId: "default",
309 terminalLabel: "Terminal 1",
310 lineStart: 4,
311 lineEnd: 4,
312 text: "boom",
313 },
314 ];
315 store.stashEntry({
316 id: "entry-records",
317 createdAt: "2026-01-01T00:00:00.000Z",
318 prompt: "see [Terminal 1 line 4](t3-context://v1/terminal/ctx-1)",
319 attachments: [],
320 droppedImageNames: [],
321 records,
322 });
323 const taken = usePromptStashStore.getState().takeEntry("entry-records");
324 expect(taken.entry?.records).toEqual(records);
325 });
326});