diff --git a/site/src/pages/AgentsPage/components/ChatConversation/AttachmentBlocks.tsx b/site/src/pages/AgentsPage/components/ChatConversation/AttachmentBlocks.tsx index 7b94b9a28f..ce8c3e22d9 100644 --- a/site/src/pages/AgentsPage/components/ChatConversation/AttachmentBlocks.tsx +++ b/site/src/pages/AgentsPage/components/ChatConversation/AttachmentBlocks.tsx @@ -26,7 +26,7 @@ import { formatTextAttachmentPreview, } from "../../utils/fetchTextAttachment"; import { ImageThumbnail } from "../AgentChatInput"; -import { useExpiredFileIds } from "./ExpiredFileIdsContext"; +import { useFileProbes } from "./FileProbeContext"; import type { RenderBlock } from "./types"; export type PreviewTextAttachment = { @@ -293,7 +293,7 @@ const RemoteTextAttachmentButton: FC<{ onPreview, showStatus = false, }) => { - const { hasExpired, markExpired } = useExpiredFileIds(); + const { hasExpired, markExpired } = useFileProbes(); const isKnownExpired = hasExpired(fileId); const [content, setContent] = useState(null); const [isLoading, setIsLoading] = useState(false); @@ -414,7 +414,15 @@ const RemoteImageBlock: FC<{ displayName: string; onImageClick?: (src: string) => void; }> = ({ fileId, href, displayName, onImageClick }) => { - const { hasExpired, markExpired } = useExpiredFileIds(); + const { + hasExpired, + markExpired, + isPending, + markPending, + clearPending, + getProbeResult, + setProbeResult, + } = useFileProbes(); const isKnownExpired = fileId !== undefined && hasExpired(fileId); const [failureState, setFailureState] = useState( () => (isKnownExpired ? { kind: "expired" } : { kind: "idle" }), @@ -429,10 +437,12 @@ const RemoteImageBlock: FC<{ /> ); } - if (failureState.kind !== "idle") { + const sharedResult = fileId ? getProbeResult(fileId) : undefined; + const effectiveFailure = sharedResult ?? failureState; + if (effectiveFailure.kind !== "idle") { return ( ); @@ -464,7 +474,13 @@ const RemoteImageBlock: FC<{ setFailureState({ kind: "expired" }); return; } + // Dedup: skip probe, context will propagate the result. + if (isPending(fileId)) { + setFailureState({ kind: "failed" }); + return; + } + markPending(fileId); const controller = probeRequest.start(); // Optimistically swap to the generic failure tile. The // probe will either upgrade it to "expired" or fill in @@ -474,22 +490,27 @@ const RemoteImageBlock: FC<{ void probeAttachmentFailure(href, controller.signal) .then((reason) => { - if (!probeRequest.clear(controller)) { - return; - } + clearPending(fileId); + // Context writes stay above the clear() guard so + // siblings get the result even if this block unmounted. if (reason.kind === "expired") { markExpired(fileId); } - setFailureState(reason); + setProbeResult(fileId, reason); + if (probeRequest.clear(controller)) { + setFailureState(reason); + } }) .catch((error) => { - if (!probeRequest.clear(controller)) { - return; - } + clearPending(fileId); if (isAbortError(error)) { return; } - setFailureState(attachmentFailureFromError(error)); + const failure = attachmentFailureFromError(error); + setProbeResult(fileId, failure); + if (probeRequest.clear(controller)) { + setFailureState(failure); + } }); }} /> diff --git a/site/src/pages/AgentsPage/components/ChatConversation/ConversationTimeline.stories.tsx b/site/src/pages/AgentsPage/components/ChatConversation/ConversationTimeline.stories.tsx index 150e5e7f63..bb599a01e5 100644 --- a/site/src/pages/AgentsPage/components/ChatConversation/ConversationTimeline.stories.tsx +++ b/site/src/pages/AgentsPage/components/ChatConversation/ConversationTimeline.stories.tsx @@ -432,6 +432,7 @@ export const UserMessageWithRepeatedExpiredImage: Story = { const images = canvas.getAllByRole("img", { name: "Attached image" }); expect(images).toHaveLength(2); fireEvent.error(images[0]); + fireEvent.error(images[1]); await waitFor(() => expect( canvas.getAllByRole("img", { name: "Image expired" }), @@ -444,6 +445,49 @@ export const UserMessageWithRepeatedExpiredImage: Story = { }, }; +/** Duplicate file IDs with a non-expired probe reuse the first result. */ +export const UserMessageWithRepeatedFailedImage: Story = { + args: buildStoryArgs( + buildUserMessage({ + id: 1, + text: "First reference to the failed upload", + files: [buildImageAttachmentPart("storybook-failed-image")], + }), + buildUserMessage({ + id: 2, + text: "Second reference to the same failed upload", + files: [buildImageAttachmentPart("storybook-failed-image")], + }), + ), + play: async ({ canvasElement }) => { + const canvas = within(canvasElement); + const images = canvas.getAllByRole("img", { name: "Attached image" }); + expect(images).toHaveLength(2); + fireEvent.error(images[0]); + fireEvent.error(images[1]); + await waitFor(() => + expect( + canvas.getAllByRole("img", { name: "Image failed to load" }), + ).toHaveLength(2), + ); + expect(getAttachmentFetchCount("storybook-failed-image")).toBe(1); + expect( + canvas.queryByRole("button", { name: "View Attached image" }), + ).not.toBeInTheDocument(); + + const tiles = await waitFor(() => { + const t = canvas.getAllByRole("img", { name: "Image failed to load" }); + for (const tile of t) { + expect(tile).toHaveAttribute("data-state"); + } + return t; + }); + for (const tile of tiles) { + await hoverAndExpectTooltip(tile, FAILED_ATTACHMENT_API_MESSAGE); + } + }, +}; + /** File-id images that fail with a non-404 status render a generic failure tile. */ export const UserMessageWithFailedRemoteImage: Story = { args: buildStoryArgs( diff --git a/site/src/pages/AgentsPage/components/ChatConversation/ConversationTimeline.tsx b/site/src/pages/AgentsPage/components/ChatConversation/ConversationTimeline.tsx index 313333243a..e157092f3e 100644 --- a/site/src/pages/AgentsPage/components/ChatConversation/ConversationTimeline.tsx +++ b/site/src/pages/AgentsPage/components/ChatConversation/ConversationTimeline.tsx @@ -44,7 +44,7 @@ import { AttachmentBlock, type PreviewTextAttachment, } from "./AttachmentBlocks"; -import { ExpiredFileIdsProvider } from "./ExpiredFileIdsContext"; +import { FileProbeProvider } from "./FileProbeContext"; import { deriveMessageDisplayState } from "./messageHelpers"; import { getEditableUserMessagePayload } from "./messageParsing"; import { useSmoothStreamingText } from "./SmoothText"; @@ -1056,7 +1056,7 @@ export const ConversationTimeline = memo( : undefined; return ( - +
( ); })}
- + ); }, ); diff --git a/site/src/pages/AgentsPage/components/ChatConversation/ExpiredFileIdsContext.tsx b/site/src/pages/AgentsPage/components/ChatConversation/ExpiredFileIdsContext.tsx deleted file mode 100644 index cd98c70632..0000000000 --- a/site/src/pages/AgentsPage/components/ChatConversation/ExpiredFileIdsContext.tsx +++ /dev/null @@ -1,45 +0,0 @@ -import { - createContext, - type FC, - type PropsWithChildren, - useContext, - useState, -} from "react"; - -type ExpiredFileIdsContextValue = { - hasExpired: (fileId: string) => boolean; - markExpired: (fileId: string) => void; -}; - -const ExpiredFileIdsContext = createContext({ - hasExpired: () => false, - markExpired: () => {}, -}); - -export const ExpiredFileIdsProvider: FC = ({ children }) => { - const [expiredFileIds, setExpiredFileIds] = useState>( - () => new Set(), - ); - - return ( - expiredFileIds.has(fileId), - markExpired: (fileId) => { - setExpiredFileIds((previous) => { - if (previous.has(fileId)) { - return previous; - } - const next = new Set(previous); - next.add(fileId); - return next; - }); - }, - }} - > - {children} - - ); -}; - -export const useExpiredFileIds = () => useContext(ExpiredFileIdsContext); diff --git a/site/src/pages/AgentsPage/components/ChatConversation/FileProbeContext.tsx b/site/src/pages/AgentsPage/components/ChatConversation/FileProbeContext.tsx new file mode 100644 index 0000000000..560fd31bd9 --- /dev/null +++ b/site/src/pages/AgentsPage/components/ChatConversation/FileProbeContext.tsx @@ -0,0 +1,78 @@ +import { + createContext, + type FC, + type PropsWithChildren, + useContext, + useRef, + useState, +} from "react"; +import type { AttachmentFailure } from "../../utils/chatAttachments"; + +type FileProbeContextValue = { + hasExpired: (fileId: string) => boolean; + markExpired: (fileId: string) => void; + isPending: (fileId: string) => boolean; + markPending: (fileId: string) => void; + clearPending: (fileId: string) => void; + getProbeResult: (fileId: string) => AttachmentFailure | undefined; + setProbeResult: (fileId: string, result: AttachmentFailure) => void; +}; + +const FileProbeContext = createContext({ + hasExpired: () => false, + markExpired: () => {}, + isPending: () => false, + markPending: () => {}, + clearPending: () => {}, + getProbeResult: () => undefined, + setProbeResult: () => {}, +}); + +export const FileProbeProvider: FC = ({ children }) => { + const [expiredFileIds, setExpiredFileIds] = useState>( + () => new Set(), + ); + // Ref, not state: must be readable synchronously by the second + // onError handler before React re-renders. + const pendingProbeFileIds = useRef>(new Set()); + const [probeResults, setProbeResults] = useState< + Map + >(() => new Map()); + + return ( + expiredFileIds.has(fileId), + markExpired: (fileId) => { + setExpiredFileIds((previous) => { + if (previous.has(fileId)) { + return previous; + } + const next = new Set(previous); + next.add(fileId); + return next; + }); + }, + isPending: (fileId) => pendingProbeFileIds.current.has(fileId), + markPending: (fileId) => { + pendingProbeFileIds.current.add(fileId); + }, + clearPending: (fileId) => { + pendingProbeFileIds.current.delete(fileId); + }, + getProbeResult: (fileId) => probeResults.get(fileId), + setProbeResult: (fileId, result) => { + setProbeResults((prev) => { + const next = new Map(prev); + next.set(fileId, result); + return next; + }); + }, + }} + > + {children} + + ); +}; + +export const useFileProbes = () => useContext(FileProbeContext);