1import { useAtomValue } from "@effect/atom-react";
2import {
3 makeQueuedRunEditStore,
4 makeQueuedRunsAtom,
5} from "@t3tools/client-runtime/queued-run-edit";
6import type { ThreadStructure } from "@t3tools/client-runtime/state/shell";
7import {
8 PROVIDER_SEND_TURN_MAX_ATTACHMENTS,
9 type ChatAttachment,
10 type ComposerContextRecord,
11 type EnvironmentContextRecord,
12 type ResponseContextRecord,
13 type EnvironmentId,
14 type MessageId,
15 type OrchestrationMessageContext,
16 type PreviewAnnotationPayload,
17 type RunId,
18 type ScopedThreadRef,
19 type ThreadContextRecord,
20 type ThreadId,
21} from "@t3tools/contracts";
22import { scopedThreadKey } from "@t3tools/client-runtime/environment";
23import {
24 restoreAttachmentOnlyBootstrapPrompt,
25 stripAttachmentOnlyBootstrapPrompt,
26} from "@t3tools/shared/composerContextReferences";
27import { Atom } from "effect/reactivity";
28import type { RefObject } from "react";
29import { useAssetUrls } from "../../assets/assetUrls";
30import {
31 DraftId,
32 useComposerDraftStore,
33 type ComposerFileAttachment,
34 type ComposerImageAttachment,
35 type ComposerThreadTarget,
36} from "../../composerDraftStore";
37import { useStableCallback } from "../../hooks/useStableCallback";
38import {
39 awaitAttachmentUploads,
40 getUploadedAttachments,
41 releaseDraftAttachments,
42 startAttachmentUpload,
43} from "../../lib/attachmentUploadQueue";
44import { buildMessageContext, isResponseContextRecord } from "../../lib/composerContextRecords";
45import {
46 collectInlineContextIds,
47 removeInlineContextReference,
48} from "../../lib/composerContextReferences";
49import { readFileAsDataUrl } from "../../lib/imageCompression";
50import type { TerminalContextDraft } from "../../lib/terminalContext";
51import type { ReviewCommentContext } from "../../reviewCommentContext";
52import { appAtomRegistry } from "../../rpc/atomRegistry";
53import { environmentServerConfigsAtom } from "../../state/server";
54import {
55 environmentThreadDetails,
56 environmentThreadShells,
57 threadEnvironment,
58} from "../../state/threads";
59import { useAtomCommand } from "../../state/use-atom-command";
60import { stackedThreadToast, toastManager } from "../ui/toast";
61import type { ChatComposerHandle } from "./ChatComposer";
62import { fileAttachmentCapabilityBlockReason } from "./composerAttachmentFiles";
63import { prepareQueuedEditAttachments, recoverQueuedMessageEdit } from "./queuedMessageEdit";
64import type { EditQueuedRunRequest } from "./QueuedRunsControl";
65
66interface EditingQueuedRun {
67 readonly threadId: ThreadId;
68 readonly runId: RunId;
69 readonly messageId: MessageId;
70 readonly originalText: string;
71 readonly existingAttachments: ReadonlyArray<ChatAttachment>;
72 readonly context?: OrchestrationMessageContext | undefined;
73 /** Removing a server attachment changes the message even when the draft looks untouched. */
74 readonly removedExistingAttachment?: true;
75}
76
77// Keyed by the route thread, so an edit survives navigating away and back.
78const queuedRunEdits = makeQueuedRunEditStore<EditingQueuedRun>({
79 registry: appAtomRegistry,
80 label: "web-queued-run-edits",
81});
82
83interface QueuedRunEditContent {
84 readonly prompt: string;
85 readonly images: ReadonlyArray<ComposerImageAttachment>;
86 readonly files: ReadonlyArray<ComposerFileAttachment>;
87 readonly terminalContexts: ReadonlyArray<TerminalContextDraft>;
88 readonly reviewComments: ReadonlyArray<ReviewCommentContext>;
89 readonly previewAnnotations: ReadonlyArray<PreviewAnnotationPayload>;
90 readonly threadContexts: ReadonlyArray<ThreadContextRecord>;
91 readonly environmentContexts: ReadonlyArray<EnvironmentContextRecord>;
92 readonly replyTo?: ResponseContextRecord | null;
93}
94
95const isDraftOwnedRecord = (record: ComposerContextRecord) =>
96 record.kind === "thread" || record.kind === "environment";
97
98function queuedEditDraftTarget(routeThreadRef: ScopedThreadRef, runId: RunId) {
99 return DraftId.make(`queued-edit:${scopedThreadKey(routeThreadRef)}:${runId}`);
100}
101
102async function submitQueuedRunEdit(input: {
103 readonly environmentId: EnvironmentId;
104 readonly draftTarget: ComposerThreadTarget;
105 readonly editing: EditingQueuedRun;
106 readonly content: QueuedRunEditContent;
107 readonly editQueuedRun: ReturnType<typeof useEditQueuedRunCommand>;
108}): Promise<string | null> {
109 const { environmentId, draftTarget, editing, content } = input;
110 const editText = content.prompt.trim();
111 const newAttachments = [...content.images, ...content.files];
112 const uploadFiles = async (files: ReadonlyArray<ComposerFileAttachment>) => {
113 const validateFiles = () => {
114 const config = appAtomRegistry.get(environmentServerConfigsAtom).get(environmentId) ?? null;
115 const reason = fileAttachmentCapabilityBlockReason({
116 files,
117 attachmentUploadsCapabilityKnown: config !== null,
118 supportsAttachmentUploads: config !== null,
119 maxFileAttachmentBytes:
120 config?.environment.capabilities.fileAttachments?.maxUploadBytes ?? null,
121 });
122 if (reason !== null) throw new Error(reason);
123 };
124 validateFiles();
125 for (const file of files) startAttachmentUpload({ environmentId, image: file, draftTarget });
126 await awaitAttachmentUploads(files.map((file) => file.id));
127 validateFiles();
128 const uploaded = getUploadedAttachments({ environmentId, images: files });
129 if (uploaded === null) throw new Error("Retry or remove failed uploads before saving.");
130 return uploaded;
131 };
132 try {
133 const uploads = await prepareQueuedEditAttachments({
134 existingAttachments: editing.existingAttachments,
135 images: content.images,
136 files: content.files,
137 readImage: readFileAsDataUrl,
138 uploadFiles,
139 });
140 // The edit draft owns thread and machine mentions, so a chip removed while editing is dropped.
141 const keptRecords = (editing.context?.records ?? []).filter(
142 (record) =>
143 !isDraftOwnedRecord(record) &&
144 (!("attachmentId" in record) ||
145 editing.existingAttachments.some((attachment) => attachment.id === record.attachmentId)),
146 );
147 const newRecords =
148 buildMessageContext({
149 terminalContexts: content.terminalContexts.filter(
150 (context) => context.text.trim().length > 0,
151 ),
152 reviewComments: content.reviewComments,
153 previewAnnotations: content.previewAnnotations,
154 threadContexts: content.threadContexts,
155 environmentContexts: content.environmentContexts,
156 replyTo: content.replyTo ?? null,
157 attachments: newAttachments.map((attachment, index) => ({
158 attachment,
159 attachmentId: uploads[editing.existingAttachments.length + index]?.id ?? attachment.id,
160 })),
161 })?.records ?? [];
162 const result = await input.editQueuedRun({
163 environmentId,
164 input: {
165 threadId: editing.threadId,
166 runId: editing.runId,
167 text: restoreAttachmentOnlyBootstrapPrompt(editText, uploads.length > 0),
168 edit: {
169 messageId: editing.messageId,
170 attachments: uploads,
171 context: {
172 version: 1,
173 records: [
174 ...new Map(
175 [...keptRecords, ...newRecords].map(
176 (record) => [record.contextId, record] as const,
177 ),
178 ).values(),
179 ],
180 },
181 },
182 },
183 });
184 return result._tag === "Failure" ? "Could not save the edited queued message." : null;
185 } catch (error) {
186 return error instanceof Error ? error.message : "Could not save the edited queued message.";
187 }
188}
189
190function useEditQueuedRunCommand() {
191 return useAtomCommand(threadEnvironment.editQueuedRun, { reportFailure: false });
192}
193
194/** Edits preserve the thread draft; attachment removals stay local until save. */
195export function useQueuedRunEdit(input: {
196 readonly environmentId: EnvironmentId;
197 readonly routeThreadRef: ScopedThreadRef;
198 readonly threadDraftTarget: ComposerThreadTarget;
199 readonly activeThreadId: ThreadId | null;
200 readonly serverProjection: ThreadStructure | null;
201 readonly promptRef: { current: string };
202 readonly composerRef: RefObject<ChatComposerHandle | null>;
203}) {
204 const {
205 environmentId,
206 routeThreadRef,
207 threadDraftTarget,
208 activeThreadId,
209 serverProjection,
210 promptRef,
211 composerRef,
212 } = input;
213 const editQueuedRunCommand = useEditQueuedRunCommand();
214 const editKey = scopedThreadKey(routeThreadRef);
215 const storedEdit = useAtomValue(queuedRunEdits.atom)[editKey] ?? null;
216 const editing = activeThreadId === null ? null : storedEdit;
217 const isSaving = useAtomValue(queuedRunEdits.savingAtom)[editKey] === true;
218 const draftTarget =
219 editing === null ? threadDraftTarget : queuedEditDraftTarget(routeThreadRef, editing.runId);
220
221 const imageResources = (editing?.existingAttachments ?? [])
222 .filter((attachment) => attachment.type === "image")
223 .map((attachment) => ({ _tag: "attachment" as const, attachmentId: attachment.id }));
224 const imageUrls = useAssetUrls(environmentId, imageResources);
225 const urlByAttachmentId = new Map(
226 imageResources.map((resource, index) => [resource.attachmentId, imageUrls[index] ?? null]),
227 );
228 const editingAttachments =
229 editing === null
230 ? null
231 : editing.existingAttachments.map((attachment) => ({
232 attachment,
233 url: urlByAttachmentId.get(attachment.id) ?? null,
234 }));
235
236 const focusComposerSoon = () => {
237 window.requestAnimationFrame(() => composerRef.current?.focusAtEnd());
238 };
239 const clearEditDraft = (runId: RunId) =>
240 useComposerDraftStore
241 .getState()
242 .clearComposerContent(queuedEditDraftTarget(routeThreadRef, runId));
243
244 const beginEdit = useStableCallback((request: EditQueuedRunRequest) => {
245 if (activeThreadId === null) return;
246 if (editing !== null && editing.runId !== request.runId) clearEditDraft(editing.runId);
247 const target = queuedEditDraftTarget(routeThreadRef, request.runId);
248 const sentText = stripAttachmentOnlyBootstrapPrompt(request.text);
249 const context = serverProjection?.messages.find(
250 (message) => message.id === request.messageId,
251 )?.context;
252 const referenced = new Set(collectInlineContextIds(sentText));
253 const records = (context?.records ?? []).filter((record) => referenced.has(record.contextId));
254 // A reply shows above the composer while its message is edited, not as a chip in the text.
255 // Only the reply the text still points at: a removed or replaced one stays behind as a record.
256 const reply = records.find(isResponseContextRecord);
257 const text = reply ? removeInlineContextReference(sentText, reply.contextId).prompt : sentText;
258 const store = useComposerDraftStore.getState();
259 store.clearComposerContent(target);
260 store.setPrompt(target, text);
261 store.setThreadContexts(
262 target,
263 records.filter((record): record is ThreadContextRecord => record.kind === "thread"),
264 );
265 store.setEnvironmentContexts(
266 target,
267 records.filter((record): record is EnvironmentContextRecord => record.kind === "environment"),
268 );
269 if (reply) store.setReplyTo(target, reply);
270 const threadRef = { environmentId, threadId: activeThreadId };
271 queuedRunEdits.begin(
272 editKey,
273 {
274 threadId: activeThreadId,
275 runId: request.runId,
276 messageId: request.messageId,
277 originalText: text,
278 existingAttachments: request.attachments,
279 context,
280 },
281 {
282 runs: makeQueuedRunsAtom({
283 threadId: activeThreadId,
284 snapshot: threadEnvironment.snapshotAtom(environmentId),
285 detailStatus: environmentThreadDetails.statusAtom(threadRef),
286 runs: Atom.make(
287 (get) =>
288 get(environmentThreadDetails.quietThreadAtom(threadRef))?.projection.runs ?? null,
289 ),
290 }),
291 // The run started or was cancelled from another client. A changed
292 // edit moves into the thread's own draft when that draft is empty;
293 // otherwise it is dropped. Edits outlive navigation, so the toast is
294 // unscoped and names the thread: the user may be anywhere by now.
295 onRunLeftQueue: (_, edit) => {
296 const recovery = recoverQueuedMessageEdit({
297 editTarget: queuedEditDraftTarget(routeThreadRef, edit.runId),
298 threadTarget: threadDraftTarget,
299 originalText: edit.originalText,
300 removedExistingAttachment: edit.removedExistingAttachment === true,
301 });
302 if (recovery === "clean") return;
303 const title = appAtomRegistry.get(
304 environmentThreadShells.threadShellAtom(threadRef),
305 )?.title;
306 const where = title === undefined ? "" : ` in “${title}”`;
307 toastManager.add(
308 stackedThreadToast({
309 type: recovery === "kept" ? "info" : "warning",
310 title: "Queued message is no longer queued",
311 description:
312 recovery === "kept"
313 ? `Your unsaved edit${where} was kept in the composer.`
314 : `Your unsaved edit${where} was discarded.`,
315 }),
316 );
317 },
318 },
319 );
320 focusComposerSoon();
321 });
322 const cancelEdit = useStableCallback(() => {
323 if (editing === null) return;
324 clearEditDraft(editing.runId);
325 queuedRunEdits.end(editKey);
326 focusComposerSoon();
327 });
328 const removeEditingAttachment = useStableCallback((attachmentId: string) => {
329 queuedRunEdits.update(editKey, (current) => ({
330 ...current,
331 removedExistingAttachment: true,
332 existingAttachments: current.existingAttachments.filter(
333 (attachment) => attachment.id !== attachmentId,
334 ),
335 }));
336 });
337
338 const saveEdit = async (
339 content: QueuedRunEditContent,
340 setThreadError: (threadId: ThreadId, error: string | null) => void,
341 ) => {
342 if (editing === null || isSaving) return;
343 const newAttachments = [...content.images, ...content.files];
344 if (
345 editing.existingAttachments.length + newAttachments.length >
346 PROVIDER_SEND_TURN_MAX_ATTACHMENTS
347 ) {
348 setThreadError(
349 editing.threadId,
350 `A message can have at most ${PROVIDER_SEND_TURN_MAX_ATTACHMENTS} attachments.`,
351 );
352 return;
353 }
354 if (
355 content.prompt.trim().length === 0 &&
356 editing.existingAttachments.length === 0 &&
357 newAttachments.length === 0
358 ) {
359 return;
360 }
361 // The save's result decides the edit's fate, even if the run starts meanwhile.
362 const finishSave = queuedRunEdits.startSave(editKey);
363 if (finishSave === null) return;
364 const error = await submitQueuedRunEdit({
365 environmentId,
366 draftTarget,
367 editing,
368 content,
369 editQueuedRun: editQueuedRunCommand,
370 });
371 setThreadError(editing.threadId, error);
372 const current = finishSave(error === null);
373 if (error !== null) return;
374 releaseDraftAttachments(newAttachments);
375 // The edit was cancelled or begun again while this save ran; its drafts are the newer edit's.
376 if (!current) return;
377 promptRef.current = "";
378 clearEditDraft(editing.runId);
379 composerRef.current?.resetCursorState();
380 focusComposerSoon();
381 };
382
383 return {
384 editingRunId: editing?.runId ?? null,
385 isEditing: editing !== null,
386 isSaving,
387 draftTarget,
388 editingAttachments,
389 beginEdit,
390 cancelEdit,
391 removeEditingAttachment,
392 saveEdit,
393 };
394}