diff --git a/packages/app/src/attachments/local-file-attachment-store.test.ts b/packages/app/src/attachments/local-file-attachment-store.test.ts new file mode 100644 index 000000000..d6aa4fdb1 --- /dev/null +++ b/packages/app/src/attachments/local-file-attachment-store.test.ts @@ -0,0 +1,70 @@ +import { beforeEach, describe, expect, it, vi } from "vitest"; +import { createLocalFileAttachmentStore } from "./local-file-attachment-store"; + +const fileSystemMock = vi.hoisted(() => ({ + getInfoAsync: vi.fn(async (uri: string) => { + if (uri.endsWith("/preview-assets/")) { + return { exists: true, isDirectory: true }; + } + return { exists: true, isDirectory: false, size: 4 }; + }), + makeDirectoryAsync: vi.fn(async () => {}), + writeAsStringAsync: vi.fn(async () => {}), + copyAsync: vi.fn(async () => {}), + readAsStringAsync: vi.fn(async () => "AAECAw=="), + deleteAsync: vi.fn(async () => {}), + readDirectoryAsync: vi.fn(async () => []), +})); + +vi.mock("expo-file-system/legacy", () => ({ + cacheDirectory: "file:///cache/", + EncodingType: { Base64: "base64" }, + getInfoAsync: fileSystemMock.getInfoAsync, + makeDirectoryAsync: fileSystemMock.makeDirectoryAsync, + writeAsStringAsync: fileSystemMock.writeAsStringAsync, + copyAsync: fileSystemMock.copyAsync, + readAsStringAsync: fileSystemMock.readAsStringAsync, + deleteAsync: fileSystemMock.deleteAsync, + readDirectoryAsync: fileSystemMock.readDirectoryAsync, +})); + +describe("local file attachment store", () => { + beforeEach(() => { + fileSystemMock.getInfoAsync.mockClear(); + fileSystemMock.makeDirectoryAsync.mockClear(); + fileSystemMock.writeAsStringAsync.mockClear(); + fileSystemMock.copyAsync.mockClear(); + fileSystemMock.readAsStringAsync.mockClear(); + fileSystemMock.deleteAsync.mockClear(); + fileSystemMock.readDirectoryAsync.mockClear(); + }); + + it("writes raw base64 sources directly to the managed file path", async () => { + const store = createLocalFileAttachmentStore({ + storageType: "native-file", + baseDirectoryName: "preview-assets", + resolvePreviewUrl: async (attachment) => `file://${attachment.storageKey}`, + }); + + const attachment = await store.save({ + id: "preview_8_test", + mimeType: "image/png", + fileName: "result.png", + source: { kind: "base64", base64: "AAECAw==" }, + }); + + expect(fileSystemMock.writeAsStringAsync).toHaveBeenCalledWith( + "file:///cache/preview-assets/preview_8_test.png", + "AAECAw==", + { encoding: "base64" }, + ); + expect(attachment).toMatchObject({ + id: "preview_8_test", + mimeType: "image/png", + storageType: "native-file", + storageKey: "/cache/preview-assets/preview_8_test.png", + fileName: "result.png", + byteSize: 4, + }); + }); +}); diff --git a/packages/app/src/attachments/local-file-attachment-store.ts b/packages/app/src/attachments/local-file-attachment-store.ts index f170a4ecc..181771eaf 100644 --- a/packages/app/src/attachments/local-file-attachment-store.ts +++ b/packages/app/src/attachments/local-file-attachment-store.ts @@ -72,6 +72,13 @@ async function writeFromSource(input: { return; } + if (input.source.kind === "base64") { + await FileSystem.writeAsStringAsync(input.targetUri, input.source.base64, { + encoding: FileSystem.EncodingType.Base64, + }); + return; + } + const base64 = await blobToBase64(input.source.blob); await FileSystem.writeAsStringAsync(input.targetUri, base64, { encoding: FileSystem.EncodingType.Base64, diff --git a/packages/app/src/attachments/service.test.ts b/packages/app/src/attachments/service.test.ts new file mode 100644 index 000000000..81e3002a5 --- /dev/null +++ b/packages/app/src/attachments/service.test.ts @@ -0,0 +1,96 @@ +import { afterEach, describe, expect, it } from "vitest"; +import type { AttachmentMetadata, AttachmentStore, SaveAttachmentInput } from "@/attachments/types"; +import { __setAttachmentStoreForTests } from "./store"; +import { encodeAttachmentsForSend, persistAttachmentFromBase64 } from "./service"; + +function createAttachment(input: Partial = {}): AttachmentMetadata { + return { + id: input.id ?? "att_1", + mimeType: input.mimeType ?? "image/png", + storageType: input.storageType ?? "web-indexeddb", + storageKey: input.storageKey ?? "att_1", + fileName: input.fileName, + byteSize: input.byteSize, + createdAt: input.createdAt ?? 1700000000000, + }; +} + +function createRecordingStore(): AttachmentStore & { + savedSources: SaveAttachmentInput[]; + releasedUrls: string[]; +} { + const savedSources: SaveAttachmentInput[] = []; + const releasedUrls: string[] = []; + + return { + storageType: "web-indexeddb", + savedSources, + releasedUrls, + async save(input) { + savedSources.push(input); + return createAttachment({ + id: input.id, + mimeType: input.mimeType, + fileName: input.fileName, + byteSize: 4, + }); + }, + async encodeBase64({ attachment }) { + return `${attachment.id}:base64`; + }, + async resolvePreviewUrl({ attachment }) { + return `blob:${attachment.id}`; + }, + async releasePreviewUrl({ url }) { + releasedUrls.push(url); + }, + async delete() {}, + async garbageCollect() {}, + }; +} + +describe("attachment service", () => { + afterEach(() => { + __setAttachmentStoreForTests(null); + }); + + it("persists raw base64 without requiring a data URL wrapper", async () => { + const store = createRecordingStore(); + __setAttachmentStoreForTests(store); + + const attachment = await persistAttachmentFromBase64({ + id: "att_base64", + base64: "AAECAw==", + mimeType: "image/png", + fileName: "image.png", + }); + + expect(attachment).toEqual({ + id: "att_base64", + mimeType: "image/png", + storageType: "web-indexeddb", + storageKey: "att_1", + fileName: "image.png", + byteSize: 4, + createdAt: 1700000000000, + }); + expect(store.savedSources).toEqual([ + { + id: "att_base64", + mimeType: "image/png", + fileName: "image.png", + source: { kind: "base64", base64: "AAECAw==" }, + }, + ]); + }); + + it("keeps provider send output byte-compatible", async () => { + const store = createRecordingStore(); + __setAttachmentStoreForTests(store); + const attachment = createAttachment({ id: "att_send", mimeType: "image/jpeg" }); + + await expect(encodeAttachmentsForSend([attachment])).resolves.toEqual([ + { data: "att_send:base64", mimeType: "image/jpeg" }, + ]); + }); +}); diff --git a/packages/app/src/attachments/service.ts b/packages/app/src/attachments/service.ts index 51a608966..c44b64339 100644 --- a/packages/app/src/attachments/service.ts +++ b/packages/app/src/attachments/service.ts @@ -31,6 +31,21 @@ export async function persistAttachmentFromDataUrl(input: { }); } +export async function persistAttachmentFromBase64(input: { + base64: string; + mimeType?: string; + fileName?: string | null; + id?: string; +}): Promise { + const store = await getAttachmentStore(); + return await store.save({ + id: input.id, + mimeType: input.mimeType, + fileName: input.fileName, + source: { kind: "base64", base64: input.base64 }, + }); +} + export async function persistAttachmentFromFileUri(input: { uri: string; mimeType?: string; diff --git a/packages/app/src/attachments/types.ts b/packages/app/src/attachments/types.ts index a859212f3..72de571d9 100644 --- a/packages/app/src/attachments/types.ts +++ b/packages/app/src/attachments/types.ts @@ -23,6 +23,7 @@ export type ComposerAttachment = | { kind: "github_pr"; item: GitHubSearchItem }; export type AttachmentDataSource = + | { kind: "base64"; base64: string } | { kind: "blob"; blob: Blob } | { kind: "data_url"; dataUrl: string } | { kind: "file_uri"; uri: string }; diff --git a/packages/app/src/attachments/utils.test.ts b/packages/app/src/attachments/utils.test.ts index d840964aa..d8e8a6a37 100644 --- a/packages/app/src/attachments/utils.test.ts +++ b/packages/app/src/attachments/utils.test.ts @@ -1,5 +1,5 @@ import { describe, expect, it } from "vitest"; -import { pathToFileUri } from "./utils"; +import { createImageSourceCacheKey, parseDataUrl, parseImageDataUrl, pathToFileUri } from "./utils"; describe("pathToFileUri", () => { it("converts POSIX absolute paths to file URIs", () => { @@ -22,3 +22,35 @@ describe("pathToFileUri", () => { expect(pathToFileUri("relative/path")).toBe("relative/path"); }); }); + +describe("parseDataUrl", () => { + it("accepts base64 data URLs with media-type parameters", () => { + expect(parseDataUrl("data:image/png;charset=utf-8;name=preview;base64,AAECAw==")).toEqual({ + mimeType: "image/png", + base64: "AAECAw==", + }); + }); + + it("rejects non-base64 data URLs", () => { + expect(() => parseDataUrl("data:image/png,not-base64")).toThrow( + "Attachment data URL is not base64 encoded.", + ); + }); +}); + +describe("parseImageDataUrl", () => { + it("returns a compact cache key for image data URLs", () => { + const dataUrl = `data:image/png;base64,${"a".repeat(512)}`; + + expect(parseImageDataUrl(dataUrl)).toMatchObject({ + mimeType: "image/png", + base64: "a".repeat(512), + }); + expect(createImageSourceCacheKey(dataUrl)).toMatch(/^data-image:image\/png:512:/); + expect(createImageSourceCacheKey(dataUrl)).not.toContain("a".repeat(128)); + }); + + it("ignores non-image data URLs", () => { + expect(parseImageDataUrl("data:text/plain;base64,SGVsbG8=")).toBeNull(); + }); +}); diff --git a/packages/app/src/attachments/utils.ts b/packages/app/src/attachments/utils.ts index 1a24b8fc4..448b706d2 100644 --- a/packages/app/src/attachments/utils.ts +++ b/packages/app/src/attachments/utils.ts @@ -14,11 +14,17 @@ export function normalizeMimeType(input: string | undefined | null): string { } export function parseDataUrl(dataUrl: string): { mimeType: string; base64: string } { - const match = /^data:([^;,]+)?;base64,(.+)$/i.exec(dataUrl); + const match = /^data:([^,]*),([\s\S]+)$/i.exec(dataUrl.trim()); if (!match) { throw new Error("Malformed data URL for attachment."); } - const [, mimeTypeRaw, base64] = match; + const metadata = match[1] ?? ""; + const base64 = match[2]?.replace(/\s/g, ""); + const [mimeTypeRaw, ...parameters] = metadata.split(";").map((part) => part.trim()); + const isBase64 = parameters.some((part) => part.toLowerCase() === "base64"); + if (!isBase64) { + throw new Error("Attachment data URL is not base64 encoded."); + } if (!base64) { throw new Error("Attachment data URL is missing base64 payload."); } @@ -28,6 +34,61 @@ export function parseDataUrl(dataUrl: string): { mimeType: string; base64: strin }; } +function hashString(value: string): string { + let hash = 2166136261; + for (let index = 0; index < value.length; index += 1) { + hash ^= value.charCodeAt(index); + hash = Math.imul(hash, 16777619); + } + return (hash >>> 0).toString(36); +} + +export function parseImageDataUrl( + uri: string, +): { mimeType: string; base64: string; cacheKey: string } | null { + if (!uri.trim().toLowerCase().startsWith("data:image/")) { + return null; + } + + try { + const parsed = parseDataUrl(uri); + if (!parsed.mimeType.toLowerCase().startsWith("image/")) { + return null; + } + return { + ...parsed, + cacheKey: `data-image:${parsed.mimeType}:${parsed.base64.length}:${hashString(parsed.base64)}`, + }; + } catch { + return null; + } +} + +export function createImageSourceCacheKey(source: string): string { + return parseImageDataUrl(source)?.cacheKey ?? source; +} + +export function getFileNameFromPath(path: string | null | undefined): string | null { + const trimmed = path?.trim(); + if (!trimmed) { + return null; + } + + const normalized = trimmed.replace(/\\/g, "/").replace(/\/+$/, ""); + const fileName = normalized.split("/").pop()?.trim(); + return fileName || null; +} + +export function createPreviewAttachmentId(input: { + base64: string; + mimeType: string; + path?: string | null; +}): string { + const path = input.path?.trim() ?? ""; + const hash = hashString(`${input.mimeType}\0${path}\0${input.base64}`); + return `preview_${input.base64.length}_${hash}`; +} + export async function blobToBase64(blob: Blob): Promise { return await new Promise((resolve, reject) => { const reader = new FileReader(); diff --git a/packages/app/src/attachments/web/indexeddb-attachment-store.ts b/packages/app/src/attachments/web/indexeddb-attachment-store.ts index 38c3b2368..c807b7bb8 100644 --- a/packages/app/src/attachments/web/indexeddb-attachment-store.ts +++ b/packages/app/src/attachments/web/indexeddb-attachment-store.ts @@ -74,6 +74,15 @@ function runTx( }); } +function base64ToBlob(input: { base64: string; mimeType: string }): Blob { + const binary = atob(input.base64); + const bytes = new Uint8Array(binary.length); + for (let index = 0; index < binary.length; index += 1) { + bytes[index] = binary.charCodeAt(index); + } + return new Blob([bytes], { type: input.mimeType }); +} + async function sourceToBlob(input: SaveAttachmentInput): Promise<{ blob: Blob; mimeType: string }> { const source = input.source; if (source.kind === "blob") { @@ -96,6 +105,14 @@ async function sourceToBlob(input: SaveAttachmentInput): Promise<{ blob: Blob; m }; } + if (source.kind === "base64") { + const mimeType = normalizeMimeType(input.mimeType); + return { + blob: base64ToBlob({ base64: source.base64, mimeType }), + mimeType, + }; + } + const response = await fetch(source.uri); const blob = await response.blob(); const mimeType = normalizeMimeType(input.mimeType ?? blob.type); diff --git a/packages/app/src/components/file-pane.tsx b/packages/app/src/components/file-pane.tsx index 5f9965349..d0bff6a59 100644 --- a/packages/app/src/components/file-pane.tsx +++ b/packages/app/src/components/file-pane.tsx @@ -25,6 +25,10 @@ import { lineNumberGutterWidth } from "@/components/code-insets"; import { isRenderedMarkdownFile } from "@/components/file-pane-render-mode"; import { isWeb } from "@/constants/platform"; import { createMarkdownStyles } from "@/styles/markdown-styles"; +import type { AttachmentMetadata } from "@/attachments/types"; +import { useAttachmentPreviewUrl } from "@/attachments/use-attachment-preview-url"; +import { persistAttachmentFromBase64 } from "@/attachments/service"; +import { createPreviewAttachmentId, getFileNameFromPath } from "@/attachments/utils"; interface CodeLineProps { tokens: HighlightToken[]; @@ -40,6 +44,7 @@ interface FilePreviewBodyProps { showDesktopWebScrollbar: boolean; isMobile: boolean; filePath: string; + imagePreviewUri: string | null; } function trimNonEmpty(value: string | null | undefined): string | null { @@ -60,6 +65,32 @@ function formatFileSize({ size }: { size: number }): string { return `${(size / (1024 * 1024)).toFixed(1)} MB`; } +async function createFilePanePreview(file: ExplorerFile | null): Promise<{ + file: ExplorerFile | null; + imageAttachment: AttachmentMetadata | null; +}> { + if (!file || file.kind !== "image" || !file.content) { + return { file, imageAttachment: null }; + } + + const { content: _content, ...imageFile } = file; + const imageAttachment = await persistAttachmentFromBase64({ + id: createPreviewAttachmentId({ + base64: file.content, + mimeType: file.mimeType ?? "image/png", + path: file.path, + }), + base64: file.content, + mimeType: file.mimeType, + fileName: getFileNameFromPath(file.path), + }); + + return { + file: imageFile, + imageAttachment, + }; +} + const CodeLine = React.memo(function CodeLine({ tokens, lineNumber, @@ -116,6 +147,7 @@ function FilePreviewBody({ showDesktopWebScrollbar, isMobile, filePath, + imagePreviewUri, }: FilePreviewBodyProps) { const { theme } = useUnistyles(); const isDark = theme.colorScheme === "dark"; @@ -230,7 +262,16 @@ function FilePreviewBody({ ); } - if (preview.kind === "image" && preview.content) { + if (preview.kind === "image") { + if (!imagePreviewUri) { + return ( + + + Loading file… + + ); + } + return ( @@ -312,6 +359,7 @@ export function FilePane({ showDesktopWebScrollbar={showDesktopWebScrollbar} isMobile={isMobile} filePath={filePath} + imagePreviewUri={imagePreviewUri} /> ); diff --git a/packages/app/src/components/message.tsx b/packages/app/src/components/message.tsx index 269e273de..b137b8254 100644 --- a/packages/app/src/components/message.tsx +++ b/packages/app/src/components/message.tsx @@ -78,11 +78,17 @@ import { setAssistantImageMetadata, } from "@/utils/assistant-image-metadata"; import { resolveAssistantImageSource } from "@/utils/assistant-image-source"; +import { + createPreviewAttachmentId, + getFileNameFromPath, + parseImageDataUrl, +} from "@/attachments/utils"; export type { InlinePathTarget } from "@/utils/inline-path"; import { PlanCard } from "./plan-card"; import { useToolCallSheet } from "./tool-call-sheet"; import { ToolCallDetailsContent } from "./tool-call-details"; import { useAttachmentPreviewUrl } from "@/attachments/use-attachment-preview-url"; +import { persistAttachmentFromBase64, persistAttachmentFromDataUrl } from "@/attachments/service"; import type { DaemonClient } from "@server/client/daemon-client"; import { isWeb, isNative } from "@/constants/platform"; @@ -595,6 +601,7 @@ function AssistantMarkdownImage({ () => resolveAssistantImageSource({ source, workspaceRoot }), [source, workspaceRoot], ); + const dataImage = useMemo(() => parseImageDataUrl(source), [source]); const containerStyle = useMemo>( () => ({ marginTop: hasLeadingContent ? theme.spacing[4] : 0, @@ -621,16 +628,44 @@ function AssistantMarkdownImage({ if (payload.error) { throw new Error(payload.error); } - if (!payload.file || payload.file.kind !== "image" || !payload.file.content) { + const file = payload.file; + if (!file || file.kind !== "image" || !file.content) { throw new Error("Image preview unavailable."); } - return `data:${payload.file.mimeType ?? "image/png"};base64,${payload.file.content}`; + return await persistAttachmentFromBase64({ + id: createPreviewAttachmentId({ + base64: file.content, + mimeType: file.mimeType ?? "image/png", + path: file.path || resolution.path, + }), + base64: file.content, + mimeType: file.mimeType, + fileName: getFileNameFromPath(file.path || resolution.path), + }); + }, + }); + const dataImageQuery = useQuery({ + queryKey: ["assistantMarkdownDataImage", dataImage?.cacheKey ?? null], + enabled: dataImage !== null, + staleTime: 30_000, + queryFn: async () => { + if (!dataImage) { + return null; + } + + return await persistAttachmentFromDataUrl({ + id: createPreviewAttachmentId(dataImage), + dataUrl: source, + mimeType: dataImage.mimeType, + }); }, }); - const directUri = resolution?.kind === "direct" ? resolution.uri : null; - const resolvedUri = directUri ?? query.data ?? null; + const fileAssetUri = useAttachmentPreviewUrl(query.data); + const dataImageAssetUri = useAttachmentPreviewUrl(dataImageQuery.data); + const directUri = resolution?.kind === "direct" && !dataImage ? resolution.uri : null; + const resolvedUri = directUri ?? dataImageAssetUri ?? fileAssetUri ?? null; if (resolvedUri) { return ( @@ -645,7 +680,7 @@ function AssistantMarkdownImage({ ); } - if (query.isLoading) { + if (query.isLoading || dataImageQuery.isLoading) { return ( - {query.error instanceof Error ? query.error.message : "Unable to load image preview."} + {query.error instanceof Error + ? query.error.message + : dataImageQuery.error instanceof Error + ? dataImageQuery.error.message + : "Unable to load image preview."} ); diff --git a/packages/app/src/desktop/attachments/desktop-attachment-store.test.ts b/packages/app/src/desktop/attachments/desktop-attachment-store.test.ts index 1eaa8b3c7..1bc2ceb39 100644 --- a/packages/app/src/desktop/attachments/desktop-attachment-store.test.ts +++ b/packages/app/src/desktop/attachments/desktop-attachment-store.test.ts @@ -86,6 +86,33 @@ describe("desktop attachment store", () => { }); }); + it("saves raw base64 sources via desktop filesystem writes", async () => { + const store = createDesktopAttachmentStore(); + const attachment = await store.save({ + id: "att_base64", + mimeType: "image/png", + fileName: "inline.png", + source: { + kind: "base64", + base64: "AAECAw==", + }, + }); + + expect(writeDesktopAttachmentBase64Mock).toHaveBeenCalledWith({ + attachmentId: "att_base64", + base64: "AAECAw==", + extension: ".png", + }); + expect(attachment).toMatchObject({ + id: "att_base64", + mimeType: "image/png", + storageType: "desktop-file", + storageKey: "/managed/att_2.png", + fileName: "inline.png", + byteSize: 4, + }); + }); + it("delegates encode/preview/delete/gc to desktop command path", async () => { const store = createDesktopAttachmentStore(); const attachment = { diff --git a/packages/app/src/desktop/attachments/desktop-attachment-store.ts b/packages/app/src/desktop/attachments/desktop-attachment-store.ts index c6e7f9989..779c09f64 100644 --- a/packages/app/src/desktop/attachments/desktop-attachment-store.ts +++ b/packages/app/src/desktop/attachments/desktop-attachment-store.ts @@ -166,6 +166,16 @@ export function createDesktopAttachmentStore(): AttachmentStore { }); } + if (input.source.kind === "base64") { + const mimeType = normalizeMimeType(input.mimeType); + return await saveDesktopAttachmentFromBase64({ + id, + base64: input.source.base64, + mimeType, + fileName, + }); + } + const mimeType = normalizeMimeType(input.mimeType ?? input.source.blob.type); const base64 = await blobToBase64(input.source.blob); return await saveDesktopAttachmentFromBase64({ diff --git a/packages/app/src/utils/assistant-image-metadata.test.ts b/packages/app/src/utils/assistant-image-metadata.test.ts index 4c7054983..65e2cde8c 100644 --- a/packages/app/src/utils/assistant-image-metadata.test.ts +++ b/packages/app/src/utils/assistant-image-metadata.test.ts @@ -55,4 +55,15 @@ describe("assistant image metadata", () => { ), ).toBeGreaterThan(220); }); + + it("estimates image-only data-image markdown without caching the full payload as text", () => { + const source = `data:image/png;base64,${"a".repeat(512)}`; + setAssistantImageMetadata({ source }, { width: 1200, height: 800 }); + + const imageOnlyHeight = estimateAssistantMessageHeightFromCache(`![Screenshot](${source})`); + const mixedHeight = estimateAssistantMessageHeightFromCache(`Text\n\n![Screenshot](${source})`); + + expect(imageOnlyHeight).toBeGreaterThan(220); + expect(mixedHeight).toBeGreaterThan(imageOnlyHeight ?? 0); + }); }); diff --git a/packages/app/src/utils/assistant-image-metadata.ts b/packages/app/src/utils/assistant-image-metadata.ts index 64e8180c6..9f6e159b0 100644 --- a/packages/app/src/utils/assistant-image-metadata.ts +++ b/packages/app/src/utils/assistant-image-metadata.ts @@ -1,5 +1,6 @@ import { MAX_CONTENT_WIDTH } from "@/constants/layout"; import { resolveAssistantImageSource } from "@/utils/assistant-image-source"; +import { createImageSourceCacheKey } from "@/attachments/utils"; export interface AssistantImageMetadata { width: number; @@ -48,8 +49,25 @@ function normalizeAssistantImageSourceToken(value: string): string | null { return source || null; } +function parseAssistantImageMarkdown(markdown: string): { + sources: string[]; + hasNonImageText: boolean; +} { + const sources: string[] = []; + for (const match of markdown.matchAll(MARKDOWN_IMAGE_PATTERN)) { + const normalized = normalizeAssistantImageSourceToken(match[1] ?? ""); + if (normalized) { + sources.push(normalized); + } + } + return { + sources, + hasNonImageText: markdown.replace(MARKDOWN_IMAGE_PATTERN, "").trim().length > 0, + }; +} + function createSourceAliasKey(source: string): string { - return `source:${source}`; + return `source:${createImageSourceCacheKey(source)}`; } function createResolutionKey(input: { @@ -65,7 +83,7 @@ function createResolutionKey(input: { return null; } if (resolution.kind === "direct") { - return `direct:${resolution.uri}`; + return `direct:${createImageSourceCacheKey(resolution.uri)}`; } return `file:${input.serverId ?? "unknown-server"}:${resolution.cwd}:${resolution.path}`; } @@ -140,7 +158,8 @@ export function setAssistantImageMetadata( } export function extractAssistantImageSources(markdown: string): string[] { - const cachedParse = assistantImageParseCache.get(markdown); + const shouldCacheParse = !/data:image\//i.test(markdown); + const cachedParse = shouldCacheParse ? assistantImageParseCache.get(markdown) : undefined; if (cachedParse) { touchCacheEntry( assistantImageParseCache, @@ -151,32 +170,15 @@ export function extractAssistantImageSources(markdown: string): string[] { return cachedParse.sources; } - const sources: string[] = []; - for (const match of markdown.matchAll(MARKDOWN_IMAGE_PATTERN)) { - const normalized = normalizeAssistantImageSourceToken(match[1] ?? ""); - if (normalized) { - sources.push(normalized); - } + const parsed = parseAssistantImageMarkdown(markdown); + if (shouldCacheParse) { + touchCacheEntry(assistantImageParseCache, markdown, parsed, ASSISTANT_IMAGE_PARSE_CACHE_LIMIT); } - const hasNonImageText = markdown.replace(MARKDOWN_IMAGE_PATTERN, "").trim().length > 0; - touchCacheEntry( - assistantImageParseCache, - markdown, - { sources, hasNonImageText }, - ASSISTANT_IMAGE_PARSE_CACHE_LIMIT, - ); - return sources; + return parsed.sources; } export function estimateAssistantMessageHeightFromCache(markdown: string): number | null { - const cachedParse = assistantImageParseCache.get(markdown); - const parsed = - cachedParse ?? - (() => { - const sources = extractAssistantImageSources(markdown); - const nextParsed = assistantImageParseCache.get(markdown); - return nextParsed ?? { sources, hasNonImageText: true }; - })(); + const parsed = assistantImageParseCache.get(markdown) ?? parseAssistantImageMarkdown(markdown); if (parsed.sources.length === 0) { return null; } diff --git a/packages/app/src/utils/image-attachments-from-files.test.ts b/packages/app/src/utils/image-attachments-from-files.test.ts index 557cbb97d..5f64ee93d 100644 --- a/packages/app/src/utils/image-attachments-from-files.test.ts +++ b/packages/app/src/utils/image-attachments-from-files.test.ts @@ -28,8 +28,10 @@ function createTestStore(): AttachmentStore { byteSize = input.source.blob.size; } else if (input.source.kind === "data_url") { byteSize = input.source.dataUrl.length; - } else { + } else if (input.source.kind === "file_uri") { byteSize = input.source.uri.length; + } else { + byteSize = input.source.base64.length; } return { id,