From aed18fdfbf694c0cd45a711fcccb248aa1c29017 Mon Sep 17 00:00:00 2001 From: David Metcalfe <80915+DavidMetcalfe@users.noreply.github.com> Date: Fri, 7 Aug 2026 19:32:32 -0700 Subject: [PATCH] fix(desktop): keep MEDIA attachments and prose when a turn also generated an image MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit stripGeneratedImageEchoes stripped EVERY markdown image and #media: link from a final message whenever an image_generate tool result was present in the same turn's parts. A completed response that carried a MEDIA: directive for a different file lost its attachment link entirely, and a message whose only content was such a directive rendered blank after the empty text part was dropped by dedupeGeneratedImageEchoesInParts (#80621). Scope the strip to echoes of the actual generated source: - markdown images are removed only when the src matches a source, with a lookahead so a longer path (a.png.bak) is not eaten by a prefix source (a.png) - #media: links are removed only when the encoded path matches a source, anchored to the closing paren so prefix paths survive - bare-path stripping no longer matches inside markdown link destinations, so [label](/tmp/a.png) is not corrupted to [label]() Also wrap parseMarkdownIntoBlocksCached in a lexer-throw fallback so a single pathological message cannot unwind through MessageRenderBoundary to the root boundary and blank the whole transcript. Closes #80621 Rebased onto current main: the strip keeps main's whitespace semantics — main deliberately stopped normalizing whole-document whitespace through media extraction (3e1b7af21be), so the same-file echo test now expects the two-space gap the removed image leaves instead of the old collapsed form. The scoped strip itself is unchanged. --- apps/desktop/src/lib/generated-images.test.ts | 73 +++++++++++++++++++ apps/desktop/src/lib/generated-images.ts | 35 +++++++-- apps/desktop/src/lib/markdown-blocks.test.ts | 21 +++++- apps/desktop/src/lib/markdown-blocks.ts | 16 ++++ 4 files changed, 136 insertions(+), 9 deletions(-) diff --git a/apps/desktop/src/lib/generated-images.test.ts b/apps/desktop/src/lib/generated-images.test.ts index 7fb3210510..8c934a5e78 100644 --- a/apps/desktop/src/lib/generated-images.test.ts +++ b/apps/desktop/src/lib/generated-images.test.ts @@ -38,6 +38,50 @@ describe('stripGeneratedImageEchoes', () => { 'Saved image: ' ) }) + + it('keeps a MEDIA: attachment link whose path is NOT a generated source (#80621)', () => { + // The reporter's case A: a final response with prose + a MEDIA: directive, + // in a turn that also completed an image_generate tool call for a DIFFERENT + // file. The attachment link must survive the echo-dedupe. + const text = + 'Here is the rendered proof:\n[Image: render-proof.jpg](#media:%2Ftmp%2Frender-proof.jpg)\nIt should show the layout correctly.' + + expect(stripGeneratedImageEchoes(text, ['/tmp/other-generated.png'])).toBe(text) + }) + + it('keeps a MEDIA: directive for a different file while stripping the same-file echo (#80621)', () => { + const text = + 'Here is the rendered proof:\n[Image: render-proof.jpg](#media:%2Ftmp%2Frender-proof.jpg)\nAnd the raw output: ![Generated](/tmp/generated.png) done.' + + // Two spaces where the echo stood: main stopped normalizing whole-document + // whitespace through media extraction (3e1b7af21be — it erased hard breaks + // and fenced-code spacing), so the strip leaves the gap as-is. + expect(stripGeneratedImageEchoes(text, ['/tmp/generated.png'])).toBe( + 'Here is the rendered proof:\n[Image: render-proof.jpg](#media:%2Ftmp%2Frender-proof.jpg)\nAnd the raw output: done.' + ) + }) + + it('keeps backticked absolute paths inert (#80621 case B)', () => { + const text = 'The proof is at `/tmp/render-proof.png` and the other at `/tmp/render-proof-2.png`.' + + expect(stripGeneratedImageEchoes(text, ['/tmp/render-proof.png'])).toBe(text) + }) + + it('does not over-match a longer path that merely starts with the source', () => { + // `/tmp/a.png` must not eat a link to `/tmp/a.png.bak` — prefix over-match + // would silently drop a legitimate attachment. + const text = 'Keep this: [Image: a.png.bak](#media:%2Ftmp%2Fa.png.bak)' + + expect(stripGeneratedImageEchoes(text, ['/tmp/a.png'])).toBe(text) + }) + + it('does not corrupt a markdown link whose destination is the source path', () => { + // A bare-path strip inside `[label](/tmp/a.png)` would leave a broken + // `[label]()` link; the link must survive intact. + const text = 'Open [the file](/tmp/a.png) to view it.' + + expect(stripGeneratedImageEchoes(text, ['/tmp/a.png'])).toBe(text) + }) }) describe('generatedImageEchoSources', () => { @@ -123,4 +167,33 @@ describe('dedupeGeneratedImageEchoesInParts', () => { expect(dedupeGeneratedImageEchoesInParts(parts)).toEqual(parts) }) + + it('keeps the MEDIA: attachment link when it references a different file than the generation (#80621)', () => { + const parts = [ + { + text: 'Here is the rendered proof:\n[Image: render-proof.jpg](#media:%2Ftmp%2Frender-proof.jpg)\nIt should show the layout correctly.', + type: 'text' as const + }, + { + result: { host_image: '/tmp/other-generated.png', image: '/tmp/other-generated.png', success: true }, + toolName: 'image_generate', + type: 'tool-call' as const + } + ] + + expect(dedupeGeneratedImageEchoesInParts(parts)).toEqual(parts) + }) + + it('does not drop a message whose only content was a MEDIA: link to a different file (#80621)', () => { + const parts = [ + { text: '[Image: render-proof.jpg](#media:%2Ftmp%2Frender-proof.jpg)', type: 'text' as const }, + { + result: { host_image: '/tmp/other-generated.png', image: '/tmp/other-generated.png', success: true }, + toolName: 'image_generate', + type: 'tool-call' as const + } + ] + + expect(dedupeGeneratedImageEchoesInParts(parts)).toEqual(parts) + }) }) diff --git a/apps/desktop/src/lib/generated-images.ts b/apps/desktop/src/lib/generated-images.ts index d88f5e463a..c8989cd2d9 100644 --- a/apps/desktop/src/lib/generated-images.ts +++ b/apps/desktop/src/lib/generated-images.ts @@ -101,21 +101,40 @@ export function generatedImageEchoSources(parts: readonly ToolLike[]): string[] } /** Strip a generated image out of prose so it only ever shows in the tool slot. - * Once a generation succeeded (`sources` is non-empty) we drop every embedded - * image and media link from that message — the model frequently restates the - * remote URL while the result holds the local path, so matching the exact - * source is not enough. Bare occurrences of the known paths/URLs are removed - * too. Surrounding prose is preserved. */ + * Once a generation succeeded (`sources` is non-empty) we drop embedded images + * and media links whose src/path matches one of the known generated sources — + * the model frequently restates the image in prose. Bare occurrences of the + * known paths/URLs are removed too. Surrounding prose is preserved, and + * references to files OTHER than the generated source (e.g. a separate + * `MEDIA:` delivery directive) are never touched (#80621). */ export function stripGeneratedImageEchoes(text: string, sources: readonly string[]): string { if (!text || sources.length === 0) { return text } - // Remove only attachment spans; surrounding whitespace may be Markdown syntax. - let next = text.replace(/!\[[^\]\n]*\]\([^)\n]*\)/g, '').replace(/\[[^\]\n]*\]\(\s*#media:[^)\n]*\)/g, '') + let next = text for (const source of unique([...sources])) { - next = next.replace(new RegExp(String.raw`(^|[\s([{])?(?=$|[\s)\]},.!?])`, 'g'), '$1') + const escaped = regexEscape(source) + // Markdown image whose src is this generated source: ![alt]( [title]). + // The lookahead after the escaped source stops a prefix over-match — a + // link to `/tmp/a.png.bak` must not be eaten by source `/tmp/a.png`. + next = next.replace(new RegExp(String.raw`!\[[^\]\n]*\]\(\s*?(?=$|[\s)\\])[^)\n]*\)`, 'g'), '') + // Media link whose path is this generated source: [label](#media:). + // Anchored to the closing paren so a link to a path that merely STARTS with + // the encoded source (e.g. `#media:%2Ftmp%2Fa.png.bak` vs `/tmp/a.png`) + // survives; case-insensitive so lowercase hex (`%2f`) is caught too. Only + // whitespace is allowed before the closing paren — a raw space can never be + // part of the href (encodeURIComponent turns it into %20), so this cannot + // over-match a longer path. + next = next.replace( + new RegExp(String.raw`\[[^\]\n]*\]\(\s*#media:${regexEscape(encodeURIComponent(source))}\s*\)`, 'gi'), + '' + ) + // Bare occurrence of the source path/URL. `(`/`)` are excluded from the + // boundary classes so a path inside a markdown link destination + // (`[label](/tmp/a.png)`) is not corrupted into `[label]()`. + next = next.replace(new RegExp(String.raw`(^|[\s[{])?(?=$|[\s\]}.,!?])`, 'g'), '$1') } return next diff --git a/apps/desktop/src/lib/markdown-blocks.test.ts b/apps/desktop/src/lib/markdown-blocks.test.ts index ff8d9495c6..47de04db91 100644 --- a/apps/desktop/src/lib/markdown-blocks.test.ts +++ b/apps/desktop/src/lib/markdown-blocks.test.ts @@ -1,5 +1,5 @@ import { parseMarkdownIntoBlocks } from '@assistant-ui/react-streamdown' -import { describe, expect, it } from 'vitest' +import { describe, expect, it, vi } from 'vitest' import { parseMarkdownIntoBlocksCached } from './markdown-blocks' @@ -161,4 +161,23 @@ describe('parseMarkdownIntoBlocksCached', () => { } } }, 30_000) + + it('falls back to the raw text block when the lexer throws instead of crashing the thread (#80621)', () => { + // The cached splitter feeds the whole transcript renderer; an unhandled + // lexer throw would unwind through MessageRenderBoundary to the root + // boundary and blank every message. Spy on String.prototype.split (which + // the lexer uses) to force a throw, and assert the fallback keeps the text + // visible as one raw block. + const splitSpy = vi.spyOn(String.prototype, 'split').mockImplementation(() => { + throw new Error('boom') + }) + + try { + const text = 'The proof is at `/tmp/render-proof.png` and it looks good.' + + expect(parseMarkdownIntoBlocksCached(text)).toEqual([text]) + } finally { + splitSpy.mockRestore() + } + }) }) diff --git a/apps/desktop/src/lib/markdown-blocks.ts b/apps/desktop/src/lib/markdown-blocks.ts index 8aac6a3a61..96c1a9d2c4 100644 --- a/apps/desktop/src/lib/markdown-blocks.ts +++ b/apps/desktop/src/lib/markdown-blocks.ts @@ -115,6 +115,22 @@ function lexIncrementally(text: string): null | string[] { } export function parseMarkdownIntoBlocksCached(markdown: string): string[] { + try { + return parseMarkdownIntoBlocksCachedUnsafe(markdown) + } catch { + // A message that the markdown lexer cannot parse must never take the whole + // thread down: MessageRenderBoundary re-throws non-transient errors to the + // root boundary, which blanks the entire transcript. Fall back to the raw + // text as a single block so the message stays visible (join(blocks) === + // markdown still holds). Drop any partial cache entry the failed lex may + // have left behind so a retry is not served stale data. See #80621. + exactCache.delete(markdown) + + return [markdown] + } +} + +function parseMarkdownIntoBlocksCachedUnsafe(markdown: string): string[] { const hit = exactCache.get(markdown) if (hit) {