Open History to see the last commit
1import {
2 CommandId,
3 type EnvironmentId,
4 type ModelSelection,
5 type ScopedThreadRef,
6 type ThreadId,
7 PROVIDER_SEND_TURN_MAX_ATTACHMENTS,
8} from "@t3tools/contracts";
9import {
10 wasBootstrapThreadDeleted,
11 wasBootstrapThreadNotCreated,
12} from "@t3tools/client-runtime/errors";
13import { Atom } from "effect/reactivity";
14import {
15 squashAtomCommandFailure,
16 type AtomCommandResult,
17} from "@t3tools/client-runtime/state/runtime";
18import { readLocalApi } from "../../localApi";
19import { awaitAttachmentUploads, getUploadedAttachments } from "../../lib/attachmentUploadQueue";
20import { findRevertCheckpoint } from "../../session-logic";
21import { type ChatMessage } from "../../types";
22import { randomUUID } from "~/lib/utils";
23import { stackedThreadToast, toastManager } from "../ui/toast";
24import { newThreadId } from "~/lib/utils";
25import {
26 type ComposerFileAttachment,
27 type ComposerImageAttachment,
28 useComposerDraftStore,
29 DraftId,
30} from "../../composerDraftStore";
31import { ATTACHMENT_ONLY_BOOTSTRAP_PROMPT } from "@t3tools/shared/composerContextReferences";
32import { createModelSelection } from "@t3tools/shared/model";
33import type { ProviderInstanceEntry } from "../../providerInstances";
34import { getComposerPromptInjectionState, getComposerProviderState } from "./composerProviderState";
35import {
36 formatOutgoingPrompt,
37 getAntigravitySendBlockReason,
38 cloneComposerImageForRetry,
39 type ComposerSendSnapshot,
40 readFileAsDataUrl,
41 prepareRevertedMessageAttachments,
42 waitForRevertedMessage,
43} from "../ChatView.logic";
44import { readPreparedConnection } from "../../state/session";
45import { recallableComposerPrompt } from "./composerPromptHistory";
46
47// Kept per route so a send in flight, and the multiple-model selections, survive
48// ChatView remounting while the user navigates.
49export const draftFanoutSelectionsAtom = Atom.family((_routeKey: string) =>
50 Atom.make<ReadonlyArray<ModelSelection> | null>(null).pipe(Atom.keepAlive),
51);
52const draftSendStateByRouteKey = new Map<
53 string,
54 {
55 readonly sendInFlight: { current: boolean };
56 readonly uncertainSubmissions: { current: Map<string, ThreadId> };
57 }
58>();
59export function draftSendState(routeKey: string) {
60 let state = draftSendStateByRouteKey.get(routeKey);
61 if (!state) {
62 state = { sendInFlight: { current: false }, uncertainSubmissions: { current: new Map() } };
63 draftSendStateByRouteKey.set(routeKey, state);
64 }
65 return state;
66}
67
68/**
69 * Starts one thread of a multiple-model send. A request that may have
70 * reached the server before failing is remembered by `retryKey`, so sending
71 * again cannot silently create a duplicate thread; its failure toast offers
72 * to open that thread or to allow the retry.
73 */
74export async function startFanoutThread(input: {
75 readonly retryKey: string;
76 readonly model: string;
77 readonly uncertainSubmissions: Map<string, ThreadId>;
78 readonly start: (threadId: ThreadId) => Promise<AtomCommandResult<unknown, unknown>>;
79 readonly openThread: (threadId: ThreadId) => void;
80}): Promise<boolean> {
81 const { retryKey, uncertainSubmissions } = input;
82 const uncertainThreadId = uncertainSubmissions.get(retryKey);
83 const targetThreadId = uncertainThreadId ?? newThreadId();
84 let requestMayHaveStarted = false;
85 try {
86 if (uncertainThreadId) {
87 throw new Error(
88 "The previous request may have started. Open its thread to check before sending again.",
89 );
90 }
91 requestMayHaveStarted = true;
92 const result = await input.start(targetThreadId);
93 if (result._tag === "Failure") {
94 const error = squashAtomCommandFailure(result);
95 if (wasBootstrapThreadDeleted(error) || wasBootstrapThreadNotCreated(error)) {
96 requestMayHaveStarted = false;
97 }
98 throw error;
99 }
100 return true;
101 } catch (error) {
102 if (requestMayHaveStarted && !uncertainSubmissions.has(retryKey)) {
103 uncertainSubmissions.set(retryKey, targetThreadId);
104 }
105 const retainedThreadId = uncertainSubmissions.get(retryKey);
106 const failureToastId = toastManager.add(
107 stackedThreadToast({
108 type: "error",
109 title: `Could not start ${input.model}`,
110 description: error instanceof Error ? error.message : "Failed to send message.",
111 ...(retainedThreadId
112 ? {
113 timeout: 0,
114 data: {
115 secondaryActionProps: {
116 children: "Allow retry",
117 onClick: () => {
118 void readLocalApi()
119 ?.dialogs.confirm(
120 "The previous request may already be running. Check its thread first. Allow another send that could create a duplicate thread?",
121 )
122 .then(
123 (confirmed) => {
124 if (
125 confirmed &&
126 uncertainSubmissions.get(retryKey) === retainedThreadId
127 ) {
128 uncertainSubmissions.delete(retryKey);
129 toastManager.close(failureToastId);
130 }
131 },
132 () => undefined,
133 );
134 },
135 },
136 },
137 actionProps: {
138 children: "Open thread",
139 onClick: () => input.openThread(retainedThreadId),
140 },
141 }
142 : {}),
143 }),
144 );
145 return false;
146 }
147}
148
149/**
150 * Reverts the thread to before `message` and puts that message's text and
151 * attachments back in the composer draft. Returns the restored prompt, or the
152 * error to show; on error the draft is untouched.
153 */
154export async function rewindToMessage(input: {
155 readonly message: ChatMessage;
156 readonly environmentId: EnvironmentId;
157 readonly routeThreadRef: ScopedThreadRef;
158 readonly threadId: ThreadId;
159 readonly turnCount: number;
160 readonly restoreFiles: boolean;
161 readonly checkpoint: ReturnType<typeof findRevertCheckpoint>;
162 readonly draftTarget: ScopedThreadRef | DraftId;
163 readonly composerHasPendingAttachments: boolean;
164 readonly createAssetUrl: Parameters<
165 typeof prepareRevertedMessageAttachments
166 >[0]["createAssetUrl"];
167 readonly revertThreadCheckpoint: (value: {
168 environmentId: EnvironmentId;
169 input: {
170 commandId: CommandId;
171 threadId: ThreadId;
172 turnCount: number;
173 restoreFiles: boolean;
174 checkpointId?: string;
175 scopeId?: string;
176 };
177 }) => Promise<AtomCommandResult<unknown, unknown>>;
178}): Promise<
179 | { readonly error: string; readonly prompt?: never }
180 | { readonly error: null; readonly prompt: string }
181> {
182 const { message, environmentId, draftTarget } = input;
183 try {
184 if (input.composerHasPendingAttachments) {
185 throw new Error("Wait for attachments to finish preparing before rewinding.");
186 }
187 const connection = readPreparedConnection(environmentId);
188 if (!connection) throw new Error("The environment is not connected.");
189 const files = await prepareRevertedMessageAttachments({
190 message,
191 environmentId,
192 httpBaseUrl: connection.httpBaseUrl,
193 createAssetUrl: input.createAssetUrl,
194 });
195 const store = useComposerDraftStore.getState();
196 const draft = store.getComposerDraft(draftTarget);
197 if (
198 (draft?.images.length ?? 0) + (draft?.files.length ?? 0) + files.length >
199 PROVIDER_SEND_TURN_MAX_ATTACHMENTS
200 ) {
201 throw new Error("Make room for this message's attachments in the composer before rewinding.");
202 }
203 const commandId = CommandId.make(randomUUID());
204 await waitForRevertedMessage(input.routeThreadRef, input.turnCount, commandId, async () => {
205 const result = await input.revertThreadCheckpoint({
206 environmentId,
207 input: {
208 commandId,
209 threadId: input.threadId,
210 turnCount: input.turnCount,
211 restoreFiles: input.restoreFiles,
212 ...(input.checkpoint === undefined
213 ? {}
214 : { checkpointId: input.checkpoint.id, scopeId: input.checkpoint.scopeId }),
215 },
216 });
217 if (result._tag === "Failure") throw squashAtomCommandFailure(result);
218 });
219 const currentPrompt = store.getComposerDraft(draftTarget)?.prompt ?? "";
220 const restoredPrompt = recallableComposerPrompt(message.text);
221 const nextPrompt =
222 restoredPrompt.length === 0
223 ? currentPrompt
224 : currentPrompt.length > 0
225 ? `${currentPrompt}\n\n${restoredPrompt}`
226 : restoredPrompt;
227 store.setPrompt(draftTarget, nextPrompt);
228 const images: ComposerImageAttachment[] = [];
229 const restoredFiles: ComposerFileAttachment[] = [];
230 files.forEach((file, index) => {
231 const attachment = {
232 id: randomUUID(),
233 name: file.name,
234 mimeType: file.type,
235 sizeBytes: file.size,
236 file,
237 };
238 if (message.attachments?.[index]?.type === "image") {
239 images.push({ ...attachment, type: "image", previewUrl: URL.createObjectURL(file) });
240 } else {
241 restoredFiles.push({ ...attachment, type: "file" });
242 }
243 });
244 store.addImages(draftTarget, images, { allowDuplicates: true });
245 store.addFiles(draftTarget, restoredFiles, { allowDuplicates: true });
246 return { error: null, prompt: nextPrompt };
247 } catch (error) {
248 return { error: error instanceof Error ? error.message : "Failed to revert thread state." };
249 }
250}
251
252/**
253 * Writes a failed send's composer contents back into its draft. Returns the
254 * restored images, which are fresh copies so their preview URLs outlive the
255 * optimistic row that held the originals.
256 */
257export function restoreComposerDraft(
258 draftTarget: ScopedThreadRef | DraftId,
259 snapshot: ComposerSendSnapshot,
260): ComposerImageAttachment[] {
261 const store = useComposerDraftStore.getState();
262 const images = snapshot.images.map(cloneComposerImageForRetry);
263 store.setPrompt(draftTarget, snapshot.draftText);
264 store.setReplyTo(draftTarget, snapshot.replyTo);
265 store.addImages(draftTarget, images);
266 store.addFiles(draftTarget, snapshot.files);
267 store.setTerminalContexts(draftTarget, snapshot.terminalContexts);
268 store.setPreviewAnnotations(draftTarget, snapshot.previewAnnotations);
269 store.setReviewComments(draftTarget, snapshot.reviewComments);
270 store.setThreadContexts(draftTarget, snapshot.threadContexts);
271 store.setEnvironmentContexts(draftTarget, snapshot.environmentContexts);
272 return images;
273}
274
275export interface MultipleModelTarget {
276 readonly selection: ModelSelection;
277 readonly text: string;
278}
279
280/**
281 * The thread each selected model starts for a multiple-model send, in
282 * selection order, or why that model cannot take the message.
283 */
284export function resolveMultipleModelTargets(input: {
285 readonly selections: ReadonlyArray<ModelSelection>;
286 readonly providerInstanceEntries: ReadonlyArray<ProviderInstanceEntry>;
287 readonly messageText: string;
288}): Array<MultipleModelTarget | { readonly error: string }> {
289 const { messageText } = input;
290 return input.selections.map((selection) => {
291 const provider = input.providerInstanceEntries.find(
292 (entry) => entry.instanceId === selection.instanceId,
293 );
294 if (!provider?.enabled || !provider.isAvailable || provider.status !== "ready") {
295 return { error: `Provider for ${selection.model} is unavailable.` };
296 }
297 const providerBlockReason = getAntigravitySendBlockReason(provider.snapshot, selection.model);
298 if (providerBlockReason) return { error: providerBlockReason };
299 const providerState = getComposerProviderState({
300 provider: provider.driverKind,
301 model: selection.model,
302 models: provider.models,
303 modelOptions: selection.options,
304 promptInjectionState: getComposerPromptInjectionState(messageText),
305 });
306 return {
307 selection: createModelSelection(
308 selection.instanceId,
309 selection.model,
310 providerState.modelOptionsForDispatch,
311 ),
312 text: formatOutgoingPrompt({
313 provider: provider.driverKind,
314 model: selection.model,
315 models: provider.models,
316 effort: providerState.promptEffort,
317 text: messageText || ATTACHMENT_ONLY_BOOTSTRAP_PROMPT,
318 }),
319 };
320 });
321}
322
323/**
324 * The attachments a turn sends: the server's upload records when it takes uploads, otherwise
325 * images inlined as data URLs. Waits for in-flight uploads; rejects when one cannot go out.
326 */
327export function prepareTurnAttachments(input: {
328 readonly environmentId: EnvironmentId;
329 readonly attachments: ReadonlyArray<ComposerImageAttachment | ComposerFileAttachment>;
330 readonly usesAttachmentUploads: boolean;
331}) {
332 const { environmentId, usesAttachmentUploads } = input;
333 return Promise.all(
334 input.attachments.map(async (attachment) => {
335 if (usesAttachmentUploads) {
336 await awaitAttachmentUploads([attachment.id]);
337 const uploaded = getUploadedAttachments({ environmentId, images: [attachment] })?.[0];
338 if (!uploaded) {
339 throw new Error(`Attachment '${attachment.name}' failed to upload.`);
340 }
341 return uploaded;
342 }
343 if (attachment.type !== "image") {
344 throw new Error("This server does not support file attachments.");
345 }
346 return {
347 type: "image" as const,
348 id: attachment.id,
349 name: attachment.name,
350 mimeType: attachment.mimeType,
351 sizeBytes: attachment.sizeBytes,
352 dataUrl: await readFileAsDataUrl(attachment.file),
353 ...(attachment.source ? { source: attachment.source } : {}),
354 };
355 }),
356 );
357}