diff --git a/apps/web/src/components/ChatView.logic.ts b/apps/web/src/components/ChatView.logic.ts index b3594ed28ef72..67c10d9dfeadf 100644 --- a/apps/web/src/components/ChatView.logic.ts +++ b/apps/web/src/components/ChatView.logic.ts @@ -667,13 +667,6 @@ export function resolveTerminalWorktreePath(input: { return input.threadWorktreePath; } -export function revokeBlobPreviewUrl(previewUrl: string | undefined): void { - if (!previewUrl || typeof URL === "undefined" || !previewUrl.startsWith("blob:")) { - return; - } - URL.revokeObjectURL(previewUrl); -} - export async function resolveFileAttachmentUrl(input: { attachment: ChatFileAttachment; environmentId: EnvironmentId; @@ -733,18 +726,6 @@ export async function prepareRevertedMessageAttachments(input: { ); } -export function revokeUserMessagePreviewUrls(message: ChatMessage): void { - if (message.role !== "user" || !message.attachments) { - return; - } - for (const attachment of message.attachments) { - if (!isImageAttachment(attachment)) { - continue; - } - revokeBlobPreviewUrl(attachment.previewUrl); - } -} - export function timelineHasEphemeralPreviewUrls( entries: ReadonlyArray & { message?: ChatMessage }>, ): boolean { diff --git a/apps/web/src/components/ChatView.tsx b/apps/web/src/components/ChatView.tsx index 8bf2ab68558ca..5a698ae9c7dae 100644 --- a/apps/web/src/components/ChatView.tsx +++ b/apps/web/src/components/ChatView.tsx @@ -13,6 +13,12 @@ import { usePullRequestDraftThread } from "./chat/usePullRequestDraftThread"; import { useComposerSelection } from "./chat/useComposerSelection"; import { useProjectCloneBanner } from "./chat/useProjectCloneBanner"; import { useAttachmentPreviewHandoff } from "./chat/useAttachmentPreviewHandoff"; +import { + applyOptimisticUserMessagesAction, + createOptimisticBlobPreviews, + emptyOptimisticUserMessages, + reconcileOptimisticUserMessages, +} from "./chat/optimisticUserMessages"; import { useTimelineRewind } from "./chat/useTimelineRewind"; import { useThreadLifecycle } from "./chat/useThreadLifecycle"; import { useRightPanelActions } from "./chat/useRightPanelActions"; @@ -322,7 +328,6 @@ import { } from "./chat/draftHeroTransition"; import { buildLocalDraftThread, - collectUserMessageBlobPreviewUrls, deriveCommittedServerUserMessageIds, dismissBranchMismatchForSession, shouldShowPlanFollowUpPrompt, @@ -333,7 +338,6 @@ import { resolveComposerProviderSelection, resolveThreadMetadataUpdateForNextTurn, resolveSendEnvMode, - revokeUserMessagePreviewUrls, startNewThreadForProject, codexArtifactTemplatePromptToAppend, shouldRefocusComposerOnWindowFocus, @@ -632,7 +636,26 @@ export default function ChatView(props: ChatViewProps) { }; const citeAssistantTextStable = useStableCallback(citeAssistantText); const [isWorkspaceFileDragActive, setIsWorkspaceFileDragActive] = useState(false); - const [optimisticUserMessages, setOptimisticUserMessages] = useState([]); + const [optimisticBlobPreviews] = useState(createOptimisticBlobPreviews); + const [heldOptimisticUserMessages, setHeldOptimisticUserMessages] = useState(() => + emptyOptimisticUserMessages(threadId), + ); + // Acknowledged messages leave the timeline in the same render the server's copy enters it. + const optimisticState = reconcileOptimisticUserMessages( + heldOptimisticUserMessages, + threadId, + serverAcknowledgedUserMessageIds, + ); + if (optimisticState !== heldOptimisticUserMessages) { + // Idempotent, and a discarded render only delays revocation to the thread's end. + optimisticBlobPreviews.handOff(optimisticState); + setHeldOptimisticUserMessages(optimisticState); + } + const optimisticUserMessages = optimisticState.messages; + const setOptimisticUserMessages = (action: SetStateAction) => + setHeldOptimisticUserMessages((state) => + applyOptimisticUserMessagesAction(state, threadId, action), + ); // Draft ids survive promotion and expire on navigation, keeping the pending resend scoped correctly. const [workLocallyResendDraftId, setWorkLocallyResendDraftId] = useState(null); const [feedbackSubmissionsByThreadKey, setFeedbackSubmissionsByThreadKey] = useState< @@ -1401,12 +1424,14 @@ export default function ChatView(props: ChatViewProps) { } }; const openFileAttachmentStable = useStableCallback(openFileAttachment); - const { handoffAttachmentPreviewsStable, timelineAttachmentUrlById } = - useAttachmentPreviewHandoff({ - environmentId, - serverVisibleTurnItems, - optimisticUserMessages, - }); + const { timelineAttachmentUrlById } = useAttachmentPreviewHandoff({ + environmentId, + threadId, + serverVisibleTurnItems, + handoff: optimisticState.handoff, + blobPreviews: optimisticBlobPreviews, + setOptimisticUserMessages: setHeldOptimisticUserMessages, + }); const anchoredTimelineMessages = feedbackSubmissions.flatMap((submission) => submission.status === "interrupted" ? [] @@ -2040,57 +2065,12 @@ export default function ChatView(props: ChatViewProps) { }; }, [activeThread?.id, focusComposerStable, isMobileViewport, terminalUiState.terminalOpen]); - useEffect(() => { - if (!activeThread?.id) return; - if (activeMessageCount === 0) { - return; - } - const removedMessages = optimisticUserMessages.filter((message) => - serverAcknowledgedUserMessageIds.has(message.id), - ); - if (removedMessages.length === 0) { - return; - } - const timer = window.setTimeout(() => { - setOptimisticUserMessages((existing) => - existing.filter((message) => !serverAcknowledgedUserMessageIds.has(message.id)), - ); - }, 0); - for (const removedMessage of removedMessages) { - const previewUrls = collectUserMessageBlobPreviewUrls(removedMessage); - if (previewUrls.length > 0) { - handoffAttachmentPreviewsStable(removedMessage.id, previewUrls); - continue; - } - revokeUserMessagePreviewUrls(removedMessage); - } - return () => { - window.clearTimeout(timer); - }; - }, [ - activeMessageCount, - activeThread?.id, - serverAcknowledgedUserMessageIds, - handoffAttachmentPreviewsStable, - optimisticUserMessages, - ]); - // Promotion drops draftId while a send may still be dispatching, so key by thread. - const [optimisticThreadId, setOptimisticThreadId] = useState(threadId); - const [retiredOptimisticMessages, setRetiredOptimisticMessages] = useState< - ReadonlyArray - >([]); - if (optimisticThreadId !== threadId) { - setOptimisticThreadId(threadId); - setRetiredOptimisticMessages(optimisticUserMessages); - setOptimisticUserMessages([]); + const [localDispatchThreadId, setLocalDispatchThreadId] = useState(threadId); + if (localDispatchThreadId !== threadId) { + setLocalDispatchThreadId(threadId); resetLocalDispatch(); } - useEffect(() => { - for (const message of retiredOptimisticMessages) { - revokeUserMessagePreviewUrls(message); - } - }, [retiredOptimisticMessages]); const activeWorktreePath = activeThread?.worktreePath ?? null; const derivedEnvMode: DraftThreadEnvMode = resolveEffectiveEnvMode({ @@ -2817,6 +2797,10 @@ export default function ChatView(props: ChatViewProps) { setDockedDraftHeroThreadKey, setSetupQueueSendThreadKey, setOptimisticUserMessages, + optimisticBlobPreviews: { + track: (message) => optimisticBlobPreviews.track(threadId, message), + release: (message) => optimisticBlobPreviews.release(threadId, message), + }, setFeedbackSubmissionsByThreadKey, setThreadError, setKeepFullHistory, diff --git a/apps/web/src/components/chat/optimisticUserMessages.test.ts b/apps/web/src/components/chat/optimisticUserMessages.test.ts new file mode 100644 index 0000000000000..0f7ce644e10bd --- /dev/null +++ b/apps/web/src/components/chat/optimisticUserMessages.test.ts @@ -0,0 +1,173 @@ +import { MessageId, ThreadId } from "@t3tools/contracts"; +import { AtomRegistry } from "effect/reactivity"; +import { describe, expect, it, vi } from "vite-plus/test"; + +import type { ChatMessage } from "../../types"; +import { + applyOptimisticUserMessagesAction, + createOptimisticBlobPreviews, + dropOptimisticHandoff, + emptyOptimisticUserMessages, + loadedHandoffMessageIdsAtom, + reconcileOptimisticUserMessages, +} from "./optimisticUserMessages"; + +const threadA = ThreadId.make("thread-a"); +const threadB = ThreadId.make("thread-b"); + +function message(id: string, previewUrl?: string): ChatMessage { + return { + id: MessageId.make(id), + role: "user", + text: id, + runId: null, + createdAt: "2026-01-01T00:00:00.000Z", + updatedAt: "2026-01-01T00:00:00.000Z", + streaming: false, + ...(previewUrl + ? { + attachments: [ + { + type: "image", + id: `${id}-image`, + name: "image.png", + mimeType: "image/png", + sizeBytes: 1, + previewUrl, + }, + ], + } + : {}), + } as ChatMessage; +} + +const ids = (...values: string[]) => new Set(values.map((value) => MessageId.make(value))); + +describe("optimistic user messages", () => { + it("an acknowledged message stays gone after the server drops its row", () => { + let state = applyOptimisticUserMessagesAction(emptyOptimisticUserMessages(threadA), threadA, [ + message("m1"), + message("m2", "blob:two"), + ]); + state = reconcileOptimisticUserMessages(state, threadA, ids("m1", "m2")); + expect(state.messages).toEqual([]); + expect(state.handoff).toEqual({ m2: ["blob:two"] }); + + // A cancelled queued run or a rewind removes the server's copy again. + state = reconcileOptimisticUserMessages(state, threadA, ids()); + expect(state.messages).toEqual([]); + + state = dropOptimisticHandoff(state, threadA, MessageId.make("m2")); + expect(state.handoff).toEqual({}); + }); + + it("releases the previous thread's messages when the thread changes", () => { + let state = applyOptimisticUserMessagesAction(emptyOptimisticUserMessages(threadA), threadA, [ + message("a1"), + ]); + state = reconcileOptimisticUserMessages(state, threadB, ids()); + state = reconcileOptimisticUserMessages(state, threadA, ids()); + expect(state).toEqual(emptyOptimisticUserMessages(threadA)); + }); + + it("drops a late update for a thread the view has left", () => { + const state = reconcileOptimisticUserMessages( + applyOptimisticUserMessagesAction(emptyOptimisticUserMessages(threadA), threadA, [ + message("a1"), + ]), + threadB, + ids(), + ); + expect(applyOptimisticUserMessagesAction(state, threadA, [message("a2")])).toBe(state); + }); + + it("revokes a thread's blob previews once, when their use ends or the thread is left", () => { + using revoke = vi.spyOn(URL, "revokeObjectURL").mockImplementation(() => {}); + const previews = createOptimisticBlobPreviews(); + previews.track(threadA, message("a1", "blob:a1")); + previews.track(threadA, message("a2", "blob:a2")); + previews.track(threadB, message("b1", "blob:b1")); + + previews.revoke(threadA, ["blob:a1"]); + previews.revokeThread(threadA); + expect(revoke.mock.calls.map(([url]) => url)).toEqual(["blob:a1", "blob:a2"]); + }); + + it("a send that fails after the server acknowledged it leaves the handed-off previews shown", () => { + using revoke = vi.spyOn(URL, "revokeObjectURL").mockImplementation(() => {}); + const previews = createOptimisticBlobPreviews(); + const sent = message("m1", "blob:m1"); + previews.track(threadA, sent); + let state = applyOptimisticUserMessagesAction(emptyOptimisticUserMessages(threadA), threadA, [ + sent, + ]); + state = reconcileOptimisticUserMessages(state, threadA, ids("m1")); + previews.handOff(state); + + previews.release(threadA, sent); + expect(revoke).not.toHaveBeenCalled(); + + previews.revokeThread(threadA); + expect(revoke.mock.calls.map(([url]) => url)).toEqual(["blob:m1"]); + }); + + it("a send that fails before acknowledgement revokes its previews", () => { + using revoke = vi.spyOn(URL, "revokeObjectURL").mockImplementation(() => {}); + const previews = createOptimisticBlobPreviews(); + const sent = message("m1", "blob:m1"); + previews.track(threadA, sent); + const state = applyOptimisticUserMessagesAction(emptyOptimisticUserMessages(threadA), threadA, [ + sent, + ]); + previews.handOff(reconcileOptimisticUserMessages(state, threadA, ids())); + + previews.release(threadA, sent); + expect(revoke.mock.calls.map(([url]) => url)).toEqual(["blob:m1"]); + }); +}); + +class FakeImage { + static readonly created: FakeImage[] = []; + src = ""; + private readonly listeners = new Map void>(); + constructor() { + FakeImage.created.push(this); + } + addEventListener(type: string, listener: () => void) { + this.listeners.set(type, listener); + } + fire(type: "load" | "error") { + this.listeners.get(type)?.(); + } +} + +describe("handoff preloads", () => { + it("reports a handoff loaded, and revokes its blobs, once every server preview loads", async () => { + vi.stubGlobal("Image", FakeImage); + try { + using revoke = vi.spyOn(URL, "revokeObjectURL").mockImplementation(() => undefined); + FakeImage.created.length = 0; + const registry = AtomRegistry.make(); + const atom = loadedHandoffMessageIdsAtom( + JSON.stringify([ + ["m-loaded", ["https://server/a.png", "https://server/b.png"], ["blob:a", "blob:b"]], + ["m-failed", ["https://server/c.png"], ["blob:c"]], + ]), + ); + const release = registry.mount(atom); + const image = (src: string) => FakeImage.created.find((candidate) => candidate.src === src); + + image("https://server/a.png")?.fire("load"); + image("https://server/c.png")?.fire("error"); + await Promise.resolve(); + expect(registry.get(atom)).toEqual([]); + + image("https://server/b.png")?.fire("load"); + await vi.waitFor(() => expect(registry.get(atom)).toEqual(["m-loaded"])); + expect(revoke.mock.calls.map(([url]) => url)).toEqual(["blob:a", "blob:b"]); + release(); + } finally { + vi.unstubAllGlobals(); + } + }); +}); diff --git a/apps/web/src/components/chat/optimisticUserMessages.ts b/apps/web/src/components/chat/optimisticUserMessages.ts new file mode 100644 index 0000000000000..eddf94168e7a7 --- /dev/null +++ b/apps/web/src/components/chat/optimisticUserMessages.ts @@ -0,0 +1,180 @@ +import type { MessageId, ThreadId } from "@t3tools/contracts"; +import * as Effect from "effect/Effect"; +import * as Schema from "effect/Schema"; +import { AsyncResult, Atom } from "effect/reactivity"; +import type { SetStateAction } from "react"; +import type { ChatMessage } from "../../types"; +import { collectUserMessageBlobPreviewUrls } from "../ChatView.logic"; + +/** + * Messages shown before the server echoes them. An acknowledged message leaves `messages` in the + * render its ack arrives, so it never reappears even if the server later drops its row. Its blob + * previews stay in `handoff` until the server's copies load (see useAttachmentPreviewHandoff). + */ +export interface OptimisticUserMessages { + readonly threadId: ThreadId; + readonly messages: ChatMessage[]; + readonly handoff: Readonly>>; +} + +export function emptyOptimisticUserMessages(threadId: ThreadId): OptimisticUserMessages { + return { threadId, messages: [], handoff: {} }; +} + +/** + * The view's state for `threadId`: the previous thread's messages are released when the thread + * changes, and acknowledged messages are moved out of `messages`. Returns `state` when nothing changes. + */ +export function reconcileOptimisticUserMessages( + state: OptimisticUserMessages, + threadId: ThreadId, + acknowledgedIds: ReadonlySet, +): OptimisticUserMessages { + if (state.threadId !== threadId) return emptyOptimisticUserMessages(threadId); + if (!state.messages.some((message) => acknowledgedIds.has(message.id))) return state; + const messages: ChatMessage[] = []; + const handoff: Record> = { ...state.handoff }; + for (const message of state.messages) { + if (!acknowledgedIds.has(message.id)) { + messages.push(message); + continue; + } + const previewUrls = collectUserMessageBlobPreviewUrls(message); + if (previewUrls.length > 0) handoff[message.id] = previewUrls; + } + return { ...state, messages, handoff }; +} + +/** Applies a send-path update. An update for a thread the view has since left is dropped. */ +export function applyOptimisticUserMessagesAction( + state: OptimisticUserMessages, + threadId: ThreadId, + action: SetStateAction, +): OptimisticUserMessages { + if (state.threadId !== threadId) return state; + const next = typeof action === "function" ? action(state.messages) : action; + const unchanged = + next.length === state.messages.length && + next.every((message, index) => message === state.messages[index]); + return unchanged ? state : { ...state, messages: next }; +} + +export function dropOptimisticHandoff( + state: OptimisticUserMessages, + threadId: ThreadId, + messageId: MessageId, +): OptimisticUserMessages { + if (state.threadId !== threadId || !(messageId in state.handoff)) return state; + const { [messageId]: _released, ...handoff } = state.handoff; + return { ...state, handoff }; +} + +/** + * Blob previews a chat view owns per thread. Revocation cannot run in a state updater, so handlers + * revoke here and leaving the thread revokes the rest. A handed-off preview still stands in for the + * server's row, so releasing its message (a send failing after the ack) leaves it to the handoff. + */ +export interface OptimisticBlobPreviews { + readonly track: (threadId: ThreadId, message: ChatMessage) => void; + readonly handOff: (state: OptimisticUserMessages) => void; + readonly revoke: (threadId: ThreadId, previewUrls: ReadonlyArray) => void; + readonly release: (threadId: ThreadId, message: ChatMessage) => void; + readonly revokeThread: (threadId: ThreadId) => void; +} + +export function createOptimisticBlobPreviews(): OptimisticBlobPreviews { + const ownedByThread = new Map>(); + const handedOff = new Set(); + const revoke = (threadId: ThreadId, previewUrls: ReadonlyArray) => { + const owned = ownedByThread.get(threadId); + for (const previewUrl of previewUrls) { + owned?.delete(previewUrl); + handedOff.delete(previewUrl); + if (typeof URL !== "undefined") URL.revokeObjectURL(previewUrl); + } + }; + return { + track: (threadId, message) => { + const previewUrls = collectUserMessageBlobPreviewUrls(message); + if (previewUrls.length === 0) return; + const owned = ownedByThread.get(threadId) ?? new Set(); + for (const previewUrl of previewUrls) owned.add(previewUrl); + ownedByThread.set(threadId, owned); + }, + handOff: (state) => { + for (const previewUrls of Object.values(state.handoff)) { + for (const previewUrl of previewUrls) handedOff.add(previewUrl); + } + }, + revoke, + release: (threadId, message) => + revoke( + threadId, + collectUserMessageBlobPreviewUrls(message).filter((url) => !handedOff.has(url)), + ), + revokeThread: (threadId) => { + revoke(threadId, [...(ownedByThread.get(threadId) ?? [])]); + ownedByThread.delete(threadId); + }, + }; +} + +// Server asset URLs carry a signed token, so the error keeps only where the image was loaded from. +class ImagePreloadError extends Schema.TaggedError()("ImagePreloadError", { + protocol: Schema.String, + host: Schema.String, +}) { + static fromUrl(url: string) { + const parsed = URL.parse(url, globalThis.location?.href); + return new ImagePreloadError({ protocol: parsed?.protocol ?? "", host: parsed?.host ?? "" }); + } +} + +const preloadImage = Effect.fn("optimisticUserMessages.preloadImage")(function* (url: string) { + yield* Effect.callback((resume) => { + const image = new Image(); + image.addEventListener("load", () => resume(Effect.void), { once: true }); + image.addEventListener("error", () => resume(Effect.fail(ImagePreloadError.fromUrl(url))), { + once: true, + }); + image.src = url; + return Effect.sync(() => { + image.src = ""; + }); + }); +}); + +/** + * Loads one handoff's server previews, keyed by `[serverUrls, blobUrls]`. Once they have all + * loaded the blobs are no longer shown, so they are revoked here, once per load. A failed load + * keeps the handoff, and its blobs, until the thread is left. + */ +const handoffPreloadAtom = Atom.family((key: string) => { + const [serverUrls, blobUrls] = JSON.parse(key) as [string[], string[]]; + return Atom.make( + typeof Image === "undefined" + ? Effect.never + : Effect.forEach(serverUrls, preloadImage, { concurrency: "unbounded", discard: true }).pipe( + Effect.tap(() => + Effect.sync(() => { + for (const blobUrl of blobUrls) URL.revokeObjectURL(blobUrl); + }), + ), + ), + ); +}); + +/** + * The message ids, among `[messageId, serverUrls, blobUrls]` entries, whose server previews have + * loaded. Each handoff preloads in its own atom, so a change to the set does not restart the rest. + */ +export const loadedHandoffMessageIdsAtom = Atom.family((key: string) => { + const preloads = JSON.parse(key) as Array<[MessageId, string[], string[]]>; + return Atom.make((get) => + preloads.flatMap(([messageId, serverUrls, blobUrls]) => + AsyncResult.isSuccess(get(handoffPreloadAtom(JSON.stringify([serverUrls, blobUrls])))) + ? [messageId] + : [], + ), + ); +}); diff --git a/apps/web/src/components/chat/useAttachmentPreviewHandoff.ts b/apps/web/src/components/chat/useAttachmentPreviewHandoff.ts index 4556a3180c286..c42047f2f4843 100644 --- a/apps/web/src/components/chat/useAttachmentPreviewHandoff.ts +++ b/apps/web/src/components/chat/useAttachmentPreviewHandoff.ts @@ -1,106 +1,53 @@ -import { useStableCallback } from "../../hooks/useStableCallback"; -import { type EnvironmentId, type MessageId } from "@t3tools/contracts"; -import { useEffect, useLayoutEffect, useRef, useState } from "react"; +import { type EnvironmentId, type MessageId, type ThreadId } from "@t3tools/contracts"; +import { useAtomValue } from "@effect/atom-react"; +import { type Dispatch, type SetStateAction, useEffect } from "react"; import { selectHandoffImageResources } from "../../session-logic"; -import { type ChatMessage } from "../../types"; import type { useQuietThreadVisibleTurnItems } from "../../state/entities"; -import { revokeBlobPreviewUrl, revokeUserMessagePreviewUrls } from "../ChatView.logic"; import { useAssetUrls } from "../../assets/assetUrls"; +import { + dropOptimisticHandoff, + loadedHandoffMessageIdsAtom, + type OptimisticBlobPreviews, + type OptimisticUserMessages, +} from "./optimisticUserMessages"; + interface AttachmentPreviewHandoffInput { readonly environmentId: EnvironmentId; + readonly threadId: ThreadId; readonly serverVisibleTurnItems: ReturnType; - readonly optimisticUserMessages: ReadonlyArray; + /** Blob previews of acknowledged messages, by message id; they stand in for the server's copy. */ + readonly handoff: OptimisticUserMessages["handoff"]; + readonly blobPreviews: OptimisticBlobPreviews; + readonly setOptimisticUserMessages: Dispatch>; } /** * A sent image keeps showing its local blob preview until the server's copy has loaded, so the - * timeline does not flash an empty frame. Returns the attachment URLs the timeline renders and - * the handoff for a message whose optimistic row was replaced by the server's. + * timeline does not flash an empty frame. The handoff ends, and its blobs are revoked, once the + * server previews load. */ export function useAttachmentPreviewHandoff(input: AttachmentPreviewHandoffInput) { - const { environmentId, serverVisibleTurnItems, optimisticUserMessages } = input; + const { + environmentId, + threadId, + serverVisibleTurnItems, + blobPreviews, + setOptimisticUserMessages, + } = input; + const handoffByMessageId = input.handoff; - const optimisticUserMessagesRef = useRef(optimisticUserMessages); - useLayoutEffect(() => { - optimisticUserMessagesRef.current = optimisticUserMessages; - }); - const [attachmentPreviewHandoffByMessageId, setAttachmentPreviewHandoffByMessageId] = useState< - Record - >({}); - const attachmentPreviewHandoffByMessageIdRef = useRef>({}); - const attachmentPreviewPromotionInFlightByMessageIdRef = useRef>({}); - useEffect(() => { - attachmentPreviewHandoffByMessageIdRef.current = attachmentPreviewHandoffByMessageId; - }, [attachmentPreviewHandoffByMessageId]); - const clearAttachmentPreviewHandoff = ( - messageId: MessageId, - previewUrls?: ReadonlyArray, - ) => { - delete attachmentPreviewPromotionInFlightByMessageIdRef.current[messageId]; - const currentPreviewUrls = - previewUrls ?? attachmentPreviewHandoffByMessageIdRef.current[messageId] ?? []; - setAttachmentPreviewHandoffByMessageId((existing) => { - if (!(messageId in existing)) { - return existing; - } - const next = { ...existing }; - delete next[messageId]; - attachmentPreviewHandoffByMessageIdRef.current = next; - return next; - }); - for (const previewUrl of currentPreviewUrls) { - revokeBlobPreviewUrl(previewUrl); - } - }; - const clearAttachmentPreviewHandoffStable = useStableCallback(clearAttachmentPreviewHandoff); - const clearAttachmentPreviewHandoffs = () => { - attachmentPreviewPromotionInFlightByMessageIdRef.current = {}; - for (const previewUrls of Object.values(attachmentPreviewHandoffByMessageIdRef.current)) { - for (const previewUrl of previewUrls) { - revokeBlobPreviewUrl(previewUrl); - } - } - attachmentPreviewHandoffByMessageIdRef.current = {}; - setAttachmentPreviewHandoffByMessageId({}); - }; - const clearAttachmentPreviewHandoffsStable = useStableCallback(clearAttachmentPreviewHandoffs); - useEffect(() => { - return () => { - clearAttachmentPreviewHandoffsStable(); - for (const message of optimisticUserMessagesRef.current) { - revokeUserMessagePreviewUrls(message); - } - }; - }, [clearAttachmentPreviewHandoffsStable]); - const handoffAttachmentPreviews = (messageId: MessageId, previewUrls: string[]) => { - if (previewUrls.length === 0) return; + // Blob URLs are a browser resource owned per thread; leaving the thread or the view ends them. + useEffect(() => () => blobPreviews.revokeThread(threadId), [blobPreviews, threadId]); - const previousPreviewUrls = attachmentPreviewHandoffByMessageIdRef.current[messageId] ?? []; - const nextPreviewUrlSet = new Set(previewUrls); - for (const previewUrl of previousPreviewUrls) { - if (!nextPreviewUrlSet.has(previewUrl)) { - revokeBlobPreviewUrl(previewUrl); - } - } - setAttachmentPreviewHandoffByMessageId((existing) => { - const next = { - ...existing, - [messageId]: previewUrls, - }; - attachmentPreviewHandoffByMessageIdRef.current = next; - return next; - }); - }; - const handoffAttachmentPreviewsStable = useStableCallback(handoffAttachmentPreviews); const serverAttachmentResources = selectHandoffImageResources( - Object.keys(attachmentPreviewHandoffByMessageId).length === 0 + Object.keys(handoffByMessageId).length === 0 ? undefined : serverVisibleTurnItems.flatMap(({ item }) => item.type === "user_message" ? [{ id: item.messageId, role: "user" as const, attachments: item.attachments }] : [], ), - attachmentPreviewHandoffByMessageId, + handoffByMessageId, ); const serverAttachmentUrls = useAssetUrls(environmentId, serverAttachmentResources); const serverAttachmentUrlById = new Map( @@ -109,104 +56,47 @@ export function useAttachmentPreviewHandoff(input: AttachmentPreviewHandoffInput return url ? [[resource.attachmentId, url] as const] : []; }), ); - useEffect(() => { - if (typeof Image === "undefined" || serverVisibleTurnItems.length === 0) { - return; - } - const cleanups: Array<() => void> = []; - const userMessagesById = new Map( - serverVisibleTurnItems.flatMap((row) => - row.item.type === "user_message" ? [[String(row.item.messageId), row.item] as const] : [], + // Each handoff whose server previews are all minted, as [messageId, serverUrls, blobUrls]. + const pendingPreloads: Array, ReadonlyArray]> = + []; + for (const row of serverVisibleTurnItems) { + if (row.item.type !== "user_message") continue; + const handoffUrls = handoffByMessageId[row.item.messageId]; + if (handoffUrls === undefined || row.item.attachments.length === 0) continue; + const serverPreviewUrls = row.item.attachments.flatMap((attachment) => { + if (attachment.type !== "image") return []; + const url = serverAttachmentUrlById.get(attachment.id); + return url === undefined ? [] : [url]; + }); + if ( + serverPreviewUrls.length === 0 || + serverPreviewUrls.length !== handoffUrls.length || + serverPreviewUrls.some((previewUrl) => previewUrl.startsWith("blob:")) + ) { + continue; + } + pendingPreloads.push([row.item.messageId, serverPreviewUrls, handoffUrls]); + } + const loadedHandoffMessageIds = useAtomValue( + loadedHandoffMessageIdsAtom(JSON.stringify(pendingPreloads)), + ); + // The handoff is state of the rendering ChatView, so dropping a loaded one here is React's + // adjust-state-while-rendering pattern; the preload atom has already revoked its blobs. + if (loadedHandoffMessageIds.length > 0) { + setOptimisticUserMessages((state) => + loadedHandoffMessageIds.reduce( + (next, messageId) => dropOptimisticHandoff(next, threadId, messageId), + state, ), ); + } - for (const [messageId, handoffPreviewUrls] of Object.entries( - attachmentPreviewHandoffByMessageId, - )) { - if (attachmentPreviewPromotionInFlightByMessageIdRef.current[messageId]) { - continue; - } - - const serverMessage = userMessagesById.get(messageId); - if (serverMessage === undefined || serverMessage.attachments.length === 0) { - continue; - } - - const serverPreviewUrls = serverMessage.attachments.flatMap((attachment) => - attachment.type === "image" - ? [serverAttachmentUrlById.get(attachment.id)].filter( - (previewUrl): previewUrl is string => previewUrl !== undefined, - ) - : [], - ); - if ( - serverPreviewUrls.length === 0 || - serverPreviewUrls.length !== handoffPreviewUrls.length || - serverPreviewUrls.some((previewUrl) => previewUrl.startsWith("blob:")) - ) { - continue; - } - - attachmentPreviewPromotionInFlightByMessageIdRef.current[messageId] = true; - - let cancelled = false; - const imageInstances: HTMLImageElement[] = []; - - const preloadServerPreviews = Promise.all( - serverPreviewUrls.map( - (previewUrl) => - new Promise((resolve, reject) => { - const image = new Image(); - imageInstances.push(image); - const handleLoad = () => resolve(); - const handleError = () => - reject(new Error(`Failed to load server preview for ${messageId}.`)); - image.addEventListener("load", handleLoad, { once: true }); - image.addEventListener("error", handleError, { once: true }); - image.src = previewUrl; - }), - ), - ); - - void preloadServerPreviews - .then(() => { - if (cancelled) { - return; - } - clearAttachmentPreviewHandoffStable(messageId as MessageId, handoffPreviewUrls); - }) - .catch(() => { - if (!cancelled) { - delete attachmentPreviewPromotionInFlightByMessageIdRef.current[messageId]; - } - }); - - cleanups.push(() => { - cancelled = true; - delete attachmentPreviewPromotionInFlightByMessageIdRef.current[messageId]; - for (const image of imageInstances) { - image.src = ""; - } - }); - } - - return () => { - for (const cleanup of cleanups) { - cleanup(); - } - }; - }, [ - attachmentPreviewHandoffByMessageId, - clearAttachmentPreviewHandoffStable, - serverAttachmentUrlById, - serverVisibleTurnItems, - ]); const timelineAttachmentUrlById = (() => { const urls = new Map(serverAttachmentUrlById); for (const row of serverVisibleTurnItems) { if (row.item.type !== "user_message") continue; - const handoffUrls = attachmentPreviewHandoffByMessageId[row.item.messageId]; + const handoffUrls = handoffByMessageId[row.item.messageId]; if (handoffUrls === undefined) continue; let imageIndex = 0; for (const attachment of row.item.attachments) { @@ -218,8 +108,5 @@ export function useAttachmentPreviewHandoff(input: AttachmentPreviewHandoffInput } return urls; })(); - return { - handoffAttachmentPreviewsStable, - timelineAttachmentUrlById, - }; + return { timelineAttachmentUrlById }; } diff --git a/apps/web/src/components/chat/useComposerSend.ts b/apps/web/src/components/chat/useComposerSend.ts index 597cafb35e92f..c63a6a3f4fce1 100644 --- a/apps/web/src/components/chat/useComposerSend.ts +++ b/apps/web/src/components/chat/useComposerSend.ts @@ -5,7 +5,6 @@ import { restorePlanFollowUpComposer, buildExpiredTerminalContextToastCopy, deriveComposerSendState, - revokeUserMessagePreviewUrls, boundAttachmentIds, formatOutgoingPrompt, optimisticMessageAttachments, @@ -224,6 +223,11 @@ interface ComposerSendInput { readonly setDockedDraftHeroThreadKey: Dispatch>; readonly setSetupQueueSendThreadKey: Dispatch>; readonly setOptimisticUserMessages: Dispatch>; + /** Takes ownership of, and later releases, an optimistic message's blob previews. */ + readonly optimisticBlobPreviews: { + readonly track: (message: ChatMessage) => void; + readonly release: (message: ChatMessage) => void; + }; readonly setFeedbackSubmissionsByThreadKey: Dispatch< SetStateAction>> >; @@ -311,6 +315,7 @@ export function useComposerSend(input: ComposerSendInput) { setDockedDraftHeroThreadKey, setSetupQueueSendThreadKey, setOptimisticUserMessages, + optimisticBlobPreviews, setFeedbackSubmissionsByThreadKey, setThreadError, setKeepFullHistory, @@ -1112,25 +1117,30 @@ export function useComposerSend(input: ComposerSendInput) { messageIdForSend, ); } - setOptimisticUserMessages((existing) => [ - ...existing, - { - id: messageIdForSend, - role: "user", - text: outgoingMessageText, - ...(optimisticAttachments.length > 0 ? { attachments: optimisticAttachments } : {}), - ...(outgoingMessageContext !== undefined ? { context: outgoingMessageContext } : {}), - runId: null, - createdAt: messageCreatedAt, - updatedAt: messageCreatedAt, - streaming: false, - ...(shouldQueueBehindActiveRun - ? { inputIntent: "queued_turn" as const } - : phase === "running" && dispatchMode === "steer" - ? { inputIntent: "steer" as const } - : {}), - }, - ]); + const optimisticMessage: ChatMessage = { + id: messageIdForSend, + role: "user", + text: outgoingMessageText, + ...(optimisticAttachments.length > 0 ? { attachments: optimisticAttachments } : {}), + ...(outgoingMessageContext !== undefined ? { context: outgoingMessageContext } : {}), + runId: null, + createdAt: messageCreatedAt, + updatedAt: messageCreatedAt, + streaming: false, + ...(shouldQueueBehindActiveRun + ? { inputIntent: "queued_turn" as const } + : phase === "running" && dispatchMode === "steer" + ? { inputIntent: "steer" as const } + : {}), + }; + // After leaving the thread during an upload, its previews were already released and the + // message is never shown, so tracking them would keep them alive past the thread. + if (currentRouteThreadKeyRef.current === routeThreadKey) { + optimisticBlobPreviews.track(optimisticMessage); + } else { + optimisticBlobPreviews.release(optimisticMessage); + } + setOptimisticUserMessages((existing) => [...existing, optimisticMessage]); setThreadError(threadIdForSend, null); if (expiredTerminalContextCount > 0) { const toastCopy = buildExpiredTerminalContextToastCopy( @@ -1369,13 +1379,11 @@ export function useComposerSend(input: ComposerSendInput) { ) ) { setOptimisticUserMessages((existing) => { - const removed = existing.filter((message) => message.id === messageIdForSend); - for (const message of removed) { - revokeUserMessagePreviewUrls(message); - } const next = existing.filter((message) => message.id !== messageIdForSend); return next.length === existing.length ? existing : next; }); + // The restored draft clones its images, so the sent copies are no longer used. + optimisticBlobPreviews.release(optimisticMessage); const retryComposerImages = restoreComposerDraft(composerDraftTarget, sendSnapshot); if (currentRouteThreadKeyRef.current === routeThreadKey) { promptRef.current = messageTextForSend;