1import { EnvironmentId, PROVIDER_SEND_TURN_MAX_ATTACHMENTS } from "@t3tools/contracts";
2import { act, createElement, useRef } from "react";
3import { create, type ReactTestRenderer } from "react-test-renderer";
4import { describe, expect, it, vi } from "vite-plus/test";
5
6import { DraftId, useComposerDraftStore } from "../../composerDraftStore";
7import {
8 importComposerThreadAttachment,
9 remainingComposerAttachmentSlots,
10} from "./composerThreadImport";
11import { THREAD_CONTEXT_DROP_EVENT, useThreadContextDrop } from "./threadContextDrag";
12
13function deferredFile() {
14 let resolve!: (file: File) => void;
15 let reject!: (error: Error) => void;
16 const promise = new Promise<File>((resolveFile, rejectFile) => {
17 resolve = resolveFile;
18 reject = rejectFile;
19 });
20 return { promise, resolve, reject };
21}
22
23describe("importComposerThreadAttachment", () => {
24 it.each(["expand", "draft switch", "environment switch", "unmount"])(
25 "keeps import ownership correct after %s while loading",
26 async (change) => {
27 const target = new EventTarget();
28 const transcript = deferredFile();
29 const pendingImports = new Map<string, number>();
30 const attached: string[] = [];
31 const expandedAtDrop: boolean[] = [];
32 let importing!: Promise<void>;
33 function Composer({
34 environmentId,
35 targetKey,
36 expanded,
37 }: {
38 environmentId: EnvironmentId;
39 targetKey: string;
40 expanded: boolean;
41 }) {
42 const targetRef = useRef<EventTarget | null>(target);
43 useThreadContextDrop(targetRef, environmentId, targetKey, (_refs, isActive) => {
44 expandedAtDrop.push(expanded);
45 importing = importComposerThreadAttachment({
46 targetKey,
47 pendingImports,
48 countReservedAttachments: () => pendingImports.get(targetKey) ?? 0,
49 load: () => transcript.promise,
50 isActive,
51 attach: async (file) => {
52 attached.push(file.name);
53 return true;
54 },
55 onLimitReached: () => {
56 throw new Error("The slot should be available");
57 },
58 });
59 });
60 return null;
61 }
62
63 const props = {
64 environmentId: EnvironmentId.make("destination"),
65 targetKey: "draft",
66 expanded: false,
67 };
68 vi.stubGlobal("IS_REACT_ACT_ENVIRONMENT", true);
69 let renderer: ReactTestRenderer | undefined;
70 try {
71 act(() => {
72 renderer = create(createElement(Composer, props));
73 });
74 act(() => target.dispatchEvent(new CustomEvent(THREAD_CONTEXT_DROP_EVENT, { detail: [] })));
75 expect(pendingImports.get("draft")).toBe(1);
76 if (change === "unmount") {
77 act(() => renderer?.unmount());
78 } else {
79 act(() =>
80 renderer?.update(
81 createElement(Composer, {
82 ...props,
83 expanded: true,
84 ...(change === "draft switch" ? { targetKey: "another-draft" } : {}),
85 ...(change === "environment switch"
86 ? { environmentId: EnvironmentId.make("another-environment") }
87 : {}),
88 }),
89 ),
90 );
91 if (change !== "expand") {
92 act(() => renderer?.update(createElement(Composer, props)));
93 }
94 }
95 transcript.resolve(new File(["transcript"], "thread.txt"));
96 await act(async () => importing);
97 expect(attached).toEqual(change === "expand" ? ["thread.txt"] : []);
98 expect(pendingImports.size).toBe(0);
99 if (change === "expand") {
100 act(() =>
101 target.dispatchEvent(new CustomEvent(THREAD_CONTEXT_DROP_EVENT, { detail: [] })),
102 );
103 await act(async () => importing);
104 expect(expandedAtDrop).toEqual([false, true]);
105 expect(attached).toEqual(["thread.txt", "thread.txt"]);
106 }
107 } finally {
108 act(() => renderer?.unmount());
109 vi.unstubAllGlobals();
110 }
111 },
112 );
113
114 it("rejects the next drop before loading when the draft takes the final slot", async () => {
115 const targetKey = DraftId.make("thread-import-synchronous-admission");
116 const store = useComposerDraftStore.getState();
117 const pendingImports = new Map<string, number>();
118 const files = Array.from({ length: PROVIDER_SEND_TURN_MAX_ATTACHMENTS - 1 }, (_, index) => ({
119 type: "file" as const,
120 id: `existing-${index}`,
121 name: `existing-${index}.txt`,
122 mimeType: "text/plain",
123 sizeBytes: 1,
124 file: new File(["x"], `existing-${index}.txt`, { type: "text/plain" }),
125 }));
126 store.addFiles(targetKey, files);
127 const load = vi.fn(async () => new File(["transcript"], "thread.jsonl"));
128 const onLimitReached = vi.fn();
129 const input = {
130 targetKey,
131 pendingImports,
132 countReservedAttachments: () =>
133 (store.getComposerDraft(targetKey)?.files.length ?? 0) +
134 (pendingImports.get(targetKey) ?? 0),
135 load,
136 isActive: () => true,
137 attach: async (file: File) =>
138 store.addFiles(targetKey, [
139 {
140 type: "file",
141 id: "transcript",
142 name: file.name,
143 mimeType: file.type,
144 sizeBytes: file.size,
145 file,
146 },
147 ]).length > 0,
148 onLimitReached,
149 };
150
151 await importComposerThreadAttachment(input);
152 await importComposerThreadAttachment(input);
153
154 expect(load).toHaveBeenCalledTimes(1);
155 expect(onLimitReached).toHaveBeenCalledTimes(1);
156 expect(store.getComposerDraft(targetKey)?.files.at(-1)?.name).toBe("thread.jsonl");
157 expect(pendingImports.size).toBe(0);
158 store.clearComposerContent(targetKey);
159 });
160
161 it.each(["file", "image"])(
162 "reserves the final slot while restoring a stashed %s",
163 async (kind) => {
164 const pendingImports = new Map<string, number>();
165 const targetKey = "draft";
166 const attached = Array.from(
167 { length: PROVIDER_SEND_TURN_MAX_ATTACHMENTS - 1 },
168 (_, index) => `existing-${index}`,
169 );
170 const countReservedAttachments = () => attached.length + (pendingImports.get(targetKey) ?? 0);
171 const attach = async (file: File) => {
172 if (
173 remainingComposerAttachmentSlots(attached.length, pendingImports.get(targetKey) ?? 0) ===
174 0
175 ) {
176 return false;
177 }
178 attached.push(file.name);
179 return true;
180 };
181 const transcript = deferredFile();
182 const importing = importComposerThreadAttachment({
183 targetKey,
184 pendingImports,
185 countReservedAttachments,
186 load: () => transcript.promise,
187 isActive: () => true,
188 attach,
189 onLimitReached: () => {
190 throw new Error("The final slot should be available");
191 },
192 });
193
194 expect(await attach(new File(["unrelated"], `unrelated.${kind}`))).toBe(false);
195 transcript.resolve(new File(["transcript"], "thread.txt"));
196 await importing;
197
198 expect(attached.at(-1)).toBe("thread.txt");
199 expect(attached).toHaveLength(PROVIDER_SEND_TURN_MAX_ATTACHMENTS);
200 expect(pendingImports.size).toBe(0);
201 },
202 );
203
204 it("rejects a repeated drop until the reserved slot is released", async () => {
205 const pendingImports = new Map<string, number>();
206 const transcript = deferredFile();
207 const load = vi.fn(() => transcript.promise);
208 const onLimitReached = vi.fn();
209 const attach = vi.fn(async (_file: File) => true);
210 const input = {
211 targetKey: "draft",
212 pendingImports,
213 countReservedAttachments: () =>
214 PROVIDER_SEND_TURN_MAX_ATTACHMENTS - 1 + (pendingImports.get("draft") ?? 0),
215 load,
216 isActive: () => true,
217 attach,
218 onLimitReached,
219 };
220 const first = importComposerThreadAttachment(input);
221 await importComposerThreadAttachment(input);
222 expect(load).toHaveBeenCalledTimes(1);
223 expect(onLimitReached).toHaveBeenCalledTimes(1);
224
225 transcript.resolve(new File(["transcript"], "thread.txt"));
226 await first;
227 await importComposerThreadAttachment(input);
228 expect(attach).toHaveBeenCalledTimes(2);
229 expect(pendingImports.size).toBe(0);
230 });
231
232 it("keeps the second drop's slot reserved when the first drop finishes", async () => {
233 const pendingImports = new Map<string, number>();
234 const first = deferredFile();
235 const second = deferredFile();
236 const attached: string[] = [];
237 const countReservedAttachments = () =>
238 PROVIDER_SEND_TURN_MAX_ATTACHMENTS - 2 + attached.length + (pendingImports.get("draft") ?? 0);
239 const attach = async (file: File) => {
240 if (countReservedAttachments() >= PROVIDER_SEND_TURN_MAX_ATTACHMENTS) return false;
241 attached.push(file.name);
242 return true;
243 };
244 const input = {
245 targetKey: "draft",
246 pendingImports,
247 countReservedAttachments,
248 isActive: () => true,
249 attach,
250 onLimitReached: () => {
251 throw new Error("Both slots should be available");
252 },
253 };
254 const firstImport = importComposerThreadAttachment({ ...input, load: () => first.promise });
255 const secondImport = importComposerThreadAttachment({ ...input, load: () => second.promise });
256 first.resolve(new File(["first"], "first-thread.txt"));
257 await firstImport;
258 expect(await attach(new File(["unrelated"], "unrelated.txt"))).toBe(false);
259
260 second.resolve(new File(["second"], "second-thread.txt"));
261 await secondImport;
262 expect(attached).toEqual(["first-thread.txt", "second-thread.txt"]);
263 expect(pendingImports.size).toBe(0);
264 });
265
266 it.each(["failure", "inactive"])(
267 "releases the original draft's slot after %s",
268 async (outcome) => {
269 const pendingImports = new Map<string, number>();
270 const transcript = deferredFile();
271 const attach = vi.fn(async (_file: File) => true);
272 let activeTargetKey = "original-draft";
273 const input = {
274 targetKey: "original-draft",
275 pendingImports,
276 countReservedAttachments: () =>
277 PROVIDER_SEND_TURN_MAX_ATTACHMENTS - 1 + (pendingImports.get("original-draft") ?? 0),
278 load: () => transcript.promise,
279 isActive: () => activeTargetKey === "original-draft",
280 attach,
281 onLimitReached: () => {
282 throw new Error("The slot should have been released");
283 },
284 };
285 const importing = importComposerThreadAttachment(input);
286 if (outcome === "inactive") activeTargetKey = "new-draft";
287 expect(pendingImports.get("new-draft") ?? 0).toBe(0);
288 if (outcome === "failure") {
289 const failed = expect(importing).rejects.toThrow("Failed to load");
290 transcript.reject(new Error("Failed to load"));
291 await failed;
292 } else {
293 transcript.resolve(new File(["transcript"], "stale-thread.txt"));
294 await importing;
295 }
296 expect(attach).not.toHaveBeenCalled();
297 expect(pendingImports.size).toBe(0);
298
299 await importComposerThreadAttachment({
300 ...input,
301 load: async () => new File(["transcript"], "retry-thread.txt"),
302 isActive: () => true,
303 });
304 expect(attach.mock.calls[0]?.[0]?.name).toBe("retry-thread.txt");
305 expect(pendingImports.size).toBe(0);
306 },
307 );
308});