Open History to see the last commit
1import { type EnvironmentId, type MessageId, type ThreadId } from "@t3tools/contracts";
2import { useAtomValue } from "@effect/atom-react";
3import { type Dispatch, type SetStateAction, useEffect } from "react";
4import { selectHandoffImageResources } from "../../session-logic";
5import type { useQuietThreadVisibleTurnItems } from "../../state/entities";
6import { useAssetUrls } from "../../assets/assetUrls";
7import {
8 dropOptimisticHandoff,
9 loadedHandoffMessageIdsAtom,
10 type OptimisticBlobPreviews,
11 type OptimisticUserMessages,
12} from "./optimisticUserMessages";
13
14interface AttachmentPreviewHandoffInput {
15 readonly environmentId: EnvironmentId;
16 readonly threadId: ThreadId;
17 readonly serverVisibleTurnItems: ReturnType<typeof useQuietThreadVisibleTurnItems>;
18 /** Blob previews of acknowledged messages, by message id; they stand in for the server's copy. */
19 readonly handoff: OptimisticUserMessages["handoff"];
20 readonly blobPreviews: OptimisticBlobPreviews;
21 readonly setOptimisticUserMessages: Dispatch<SetStateAction<OptimisticUserMessages>>;
22}
23
24/**
25 * A sent image keeps showing its local blob preview until the server's copy has loaded, so the
26 * timeline does not flash an empty frame. The handoff ends, and its blobs are revoked, once the
27 * server previews load.
28 */
29export function useAttachmentPreviewHandoff(input: AttachmentPreviewHandoffInput) {
30 const {
31 environmentId,
32 threadId,
33 serverVisibleTurnItems,
34 blobPreviews,
35 setOptimisticUserMessages,
36 } = input;
37 const handoffByMessageId = input.handoff;
38
39 // Blob URLs are a browser resource owned per thread; leaving the thread or the view ends them.
40 // oxlint-disable-next-line t3code/no-unreasoned-state-sync -- revokes the thread's blob preview URLs when the thread changes or the view unmounts
41 useEffect(() => () => blobPreviews.revokeThread(threadId), [blobPreviews, threadId]);
42
43 const serverAttachmentResources = selectHandoffImageResources(
44 Object.keys(handoffByMessageId).length === 0
45 ? undefined
46 : serverVisibleTurnItems.flatMap(({ item }) =>
47 item.type === "user_message"
48 ? [{ id: item.messageId, role: "user" as const, attachments: item.attachments }]
49 : [],
50 ),
51 handoffByMessageId,
52 );
53 const serverAttachmentUrls = useAssetUrls(environmentId, serverAttachmentResources);
54 const serverAttachmentUrlById = new Map(
55 serverAttachmentResources.flatMap((resource, index) => {
56 const url = serverAttachmentUrls[index];
57 return url ? [[resource.attachmentId, url] as const] : [];
58 }),
59 );
60
61 // Each handoff whose server previews are all minted, as [messageId, serverUrls, blobUrls].
62 const pendingPreloads: Array<readonly [MessageId, ReadonlyArray<string>, ReadonlyArray<string>]> =
63 [];
64 for (const row of serverVisibleTurnItems) {
65 if (row.item.type !== "user_message") continue;
66 const handoffUrls = handoffByMessageId[row.item.messageId];
67 if (handoffUrls === undefined || row.item.attachments.length === 0) continue;
68 const serverPreviewUrls = row.item.attachments.flatMap((attachment) => {
69 if (attachment.type !== "image") return [];
70 const url = serverAttachmentUrlById.get(attachment.id);
71 return url === undefined ? [] : [url];
72 });
73 if (
74 serverPreviewUrls.length === 0 ||
75 serverPreviewUrls.length !== handoffUrls.length ||
76 serverPreviewUrls.some((previewUrl) => previewUrl.startsWith("blob:"))
77 ) {
78 continue;
79 }
80 pendingPreloads.push([row.item.messageId, serverPreviewUrls, handoffUrls]);
81 }
82 const loadedHandoffMessageIds = useAtomValue(
83 loadedHandoffMessageIdsAtom(JSON.stringify(pendingPreloads)),
84 );
85 // The handoff is state of the rendering ChatView, so dropping a loaded one here is React's
86 // adjust-state-while-rendering pattern; the preload atom has already revoked its blobs.
87 if (loadedHandoffMessageIds.length > 0) {
88 setOptimisticUserMessages((state) =>
89 loadedHandoffMessageIds.reduce(
90 (next, messageId) => dropOptimisticHandoff(next, threadId, messageId),
91 state,
92 ),
93 );
94 }
95
96 const timelineAttachmentUrlById = (() => {
97 const urls = new Map(serverAttachmentUrlById);
98 for (const row of serverVisibleTurnItems) {
99 if (row.item.type !== "user_message") continue;
100 const handoffUrls = handoffByMessageId[row.item.messageId];
101 if (handoffUrls === undefined) continue;
102 let imageIndex = 0;
103 for (const attachment of row.item.attachments) {
104 if (attachment.type !== "image") continue;
105 const handoffUrl = handoffUrls[imageIndex];
106 imageIndex += 1;
107 if (handoffUrl !== undefined) urls.set(attachment.id, handoffUrl);
108 }
109 }
110 return urls;
111 })();
112 return { timelineAttachmentUrlById };
113}