diff --git a/desktop/src/components/chat/AssistantMessage.linkrouting.test.tsx b/desktop/src/components/chat/AssistantMessage.linkrouting.test.tsx index 9cf1c31a..d82a64d3 100644 --- a/desktop/src/components/chat/AssistantMessage.linkrouting.test.tsx +++ b/desktop/src/components/chat/AssistantMessage.linkrouting.test.tsx @@ -169,7 +169,7 @@ describe('AssistantMessage output-target cards', () => { expect(screen.queryByText('assistantOutputs.kind.image')).toBeNull() }) - it('hands a Markdown image to the inline gallery without leaving a source-less duplicate', () => { + it('keeps a finalized Markdown image in the prose without a gallery duplicate', () => { const { container, rerender } = render( { 'src', 'http://127.0.0.1:4321/preview-fs/s1/outputs/foo/preview_frame.png', ) + expect(container.querySelector('.markdown-prose')).toContainElement(images[0]!) expect(container.querySelector('img:not([src])')).toBeNull() }) diff --git a/desktop/src/components/chat/AssistantMessage.tsx b/desktop/src/components/chat/AssistantMessage.tsx index 972d5197..bd7e83bd 100644 --- a/desktop/src/components/chat/AssistantMessage.tsx +++ b/desktop/src/components/chat/AssistantMessage.tsx @@ -13,6 +13,9 @@ import { InlineVideoGallery } from './InlineVideoGallery' import { AssistantOutputTargetCard } from './AssistantOutputTargetCard' import { openPreviewLink } from '../../lib/openPreviewLink' import { extractAssistantOutputTargets } from '../../lib/assistantOutputTargets' +import { createAssistantMarkdownImageResolver } from '../../lib/markdownImages' +import { getServerBaseUrl } from '../../lib/desktopRuntime' +import { isManagedGeneratedImagePath } from '../../lib/attachmentImages' import { useWorkspacePanelStore } from '../../stores/workspacePanelStore' import { useTranslation, type TranslationKey } from '../../i18n' @@ -97,6 +100,19 @@ export const AssistantMessage = memo(function AssistantMessage({ ), [content, isStreaming, isTurnOutputOwner, sessionId, workDir, turnChangedFiles], ) + const resolveAssistantImageSrc = useMemo( + () => { + if (isStreaming || !sessionId) return undefined + const resolveLocalImage = createAssistantMarkdownImageResolver({ + baseUrl: getServerBaseUrl(), + sessionId, + }) + return (src: string) => isManagedGeneratedImagePath(src) + ? null + : resolveLocalImage(src) + }, + [isStreaming, sessionId], + ) if (!content.trim()) return null @@ -129,6 +145,7 @@ export const AssistantMessage = memo(function AssistantMessage({ variant={documentLayout ? 'document' : 'default'} streaming={isStreaming} onLinkClick={sessionId ? handleLinkClick : undefined} + resolveImageSrc={resolveAssistantImageSrc} /> {!isStreaming && ( (null) + const markdownImageSources = useMemo( + () => new Set(extractMarkdownImageSources(text).map(normalizeImageReference)), + [text], + ) + // Absolute paths are explicitly written out in the prose (not guessed), and the // turn checkpoint can't see files written via Bash or outside its tracking scope // — so they keep the legacy behavior and render unconditionally. changedFiles @@ -63,9 +82,12 @@ export function InlineImageGallery({ text, sessionId, workDir, changedFiles, sup // need to be reconciled against what the turn actually wrote. const imagePaths = useMemo( () => extractImagePaths(text).filter( - (imagePath) => !suppressManagedGeneratedImages || !isManagedGeneratedImagePath(imagePath), + (imagePath) => ( + !markdownImageSources.has(normalizeImageReference(imagePath)) && + (!suppressManagedGeneratedImages || !isManagedGeneratedImagePath(imagePath)) + ), ), - [suppressManagedGeneratedImages, text], + [markdownImageSources, suppressManagedGeneratedImages, text], ) // An empty changedFiles only means "no TRACKED file changed" (Bash writes are @@ -86,7 +108,12 @@ export function InlineImageGallery({ text, sessionId, workDir, changedFiles, sup // a bespoke relative-path regex. const base = getServerBaseUrl() const relativeTargets = extractAssistantOutputTargets(text, { workDir, changedFiles: changedFileEvidence }).filter( - (target) => target.kind === 'image', + (target) => ( + target.kind === 'image' && + target.source !== 'markdown-link' && + !markdownImageSources.has(normalizeImageReference(target.href)) && + !markdownImageSources.has(normalizeImageReference(target.normalizedPath ?? '')) + ), ) // Dedup: an absolute path inside the workspace can be caught by BOTH sources. @@ -113,7 +140,7 @@ export function InlineImageGallery({ text, sessionId, workDir, changedFiles, sup } return [...absolute, ...relative] - }, [changedFileEvidence, imagePaths, sessionId, text, workDir]) + }, [changedFileEvidence, imagePaths, markdownImageSources, sessionId, text, workDir]) if (images.length === 0) return null @@ -170,7 +197,3 @@ export function InlineImageGallery({ text, sessionId, workDir, changedFiles, sup ) } - -function isManagedGeneratedImagePath(imagePath: string): boolean { - return imagePath.replaceAll('\\', '/').includes('/.claude/cc-haha/generated-images/') -} diff --git a/desktop/src/components/chat/MessageList.test.tsx b/desktop/src/components/chat/MessageList.test.tsx index a1fc8953..62627391 100644 --- a/desktop/src/components/chat/MessageList.test.tsx +++ b/desktop/src/components/chat/MessageList.test.tsx @@ -432,6 +432,78 @@ describe('MessageList nested tool calls', () => { expect(document.querySelector('img:not([src])')).toBeNull() }) + it('preserves local Markdown image placement when a streamed reply becomes final', () => { + const firstChunk = [ + '文字A', + '', + '![图A](01', + ].join('\n') + const secondChunk = [ + '.png)', + '', + '文字B', + '', + '![图B](nested/02.png)', + '', + '文字C', + '', + '![图C](03.png)', + '', + '裸路径仍需兜底:outputs/fallback.png', + '', + '![remote](https://attacker.example/track.png)', + '', + '![loopback](http://127.0.0.1:3456/status.png)', + ].join('\n') + + const { container } = render() + const store = useChatStore.getState() + + act(() => { + store.handleServerMessage(ACTIVE_TAB, { type: 'content_start', blockType: 'text' }) + store.handleServerMessage(ACTIVE_TAB, { type: 'content_delta', text: firstChunk }) + }) + + expect(container.querySelectorAll('img')).toHaveLength(0) + + act(() => { + store.handleServerMessage(ACTIVE_TAB, { type: 'content_delta', text: secondChunk }) + store.handleServerMessage(ACTIVE_TAB, { type: 'status', state: 'idle' }) + }) + + const assistant = container.querySelector('[data-message-shell="assistant"]') + const prose = assistant?.querySelector('.markdown-prose') + expect(prose).not.toBeNull() + + const inlineImages = Array.from(prose!.querySelectorAll('img')) + expect(inlineImages.map((image) => image.getAttribute('alt'))).toEqual(['图A', '图B', '图C']) + expect(inlineImages.map((image) => image.getAttribute('src'))).toEqual([ + 'http://127.0.0.1:3456/preview-fs/active-tab/01.png', + 'http://127.0.0.1:3456/preview-fs/active-tab/nested/02.png', + 'http://127.0.0.1:3456/preview-fs/active-tab/03.png', + ]) + + const orderedNodes = [ + screen.getByText('文字A'), + inlineImages[0]!, + screen.getByText('文字B'), + inlineImages[1]!, + screen.getByText('文字C'), + inlineImages[2]!, + ] + for (let index = 0; index < orderedNodes.length - 1; index += 1) { + expect(orderedNodes[index]!.compareDocumentPosition(orderedNodes[index + 1]!) & Node.DOCUMENT_POSITION_FOLLOWING) + .toBeTruthy() + } + + const galleryImages = Array.from(assistant!.querySelectorAll('img')).filter((image) => !prose!.contains(image)) + expect(galleryImages.map((image) => image.getAttribute('alt'))).toEqual(['fallback.png']) + expect(assistant!.querySelectorAll('img[alt="图A"]')).toHaveLength(1) + expect(assistant!.querySelectorAll('img[alt="图B"]')).toHaveLength(1) + expect(assistant!.querySelectorAll('img[alt="图C"]')).toHaveLength(1) + expect(assistant!.querySelector('img[alt="remote"], img[alt="loopback"]')).toBeNull() + }) + it('keeps fractional border-box jitter from invalidating a settled virtual row', async () => { const sessionId = 'virtual-row-measurement-jitter' const observers: Array<{ diff --git a/desktop/src/components/markdown/MarkdownRenderer.tsx b/desktop/src/components/markdown/MarkdownRenderer.tsx index 666edda3..6793f142 100644 --- a/desktop/src/components/markdown/MarkdownRenderer.tsx +++ b/desktop/src/components/markdown/MarkdownRenderer.tsx @@ -28,11 +28,11 @@ type Props = { streaming?: boolean onLinkClick?: (href: string, event: ReactMouseEvent) => boolean | void /** - * Trusted surfaces (the workspace Markdown file preview) resolve every image - * source through this callback — relative paths against the document, - * remote URLs left to CSP. Returning null strips the image. When absent the - * surface is treated as untrusted assistant output and only blob:/data: - * sources survive, which blocks tracking pixels and loopback probes. + * Resolve every image source through a policy owned by the calling surface. + * The workspace preview permits document-relative and remote sources, while + * finalized assistant prose permits only session-sandboxed local paths. + * Returning null strips the image. Without a resolver, only blob:/data: + * sources survive. */ resolveImageSrc?: (src: string) => string | null } @@ -355,9 +355,9 @@ function enhanceMarkdownHtml( container.querySelectorAll('img, source').forEach((image) => { const src = image.getAttribute('src') if (resolveImageSrc) { - // Trusted surface: the caller maps relative/remote sources to loadable - // URLs (or null to strip). Resolution happens here, after sanitization, - // so raw HTML tags get the same treatment as Markdown images. + // The caller maps permitted sources to loadable URLs (or null to strip). + // Resolution happens here, after sanitization, so raw HTML tags get + // the same treatment as Markdown images. const resolved = src ? resolveImageSrc(src) : null if (resolved) image.setAttribute('src', resolved) else { diff --git a/desktop/src/lib/assistantOutputTargets.ts b/desktop/src/lib/assistantOutputTargets.ts index be485d94..d4652d4a 100644 --- a/desktop/src/lib/assistantOutputTargets.ts +++ b/desktop/src/lib/assistantOutputTargets.ts @@ -274,6 +274,14 @@ export function extractAssistantOutputTargets( return results } +/** Markdown-authored images already have a position inside the prose. */ +export function extractMarkdownImageSources(content: string): string[] { + const codeBlocks = extractFencedCodeBlocks(content) + return extractMarkdownLinks(content) + .filter((match) => match.isImage && !isInCodeBlock(match.start, codeBlocks)) + .map((match) => match.href) +} + /** * Place a bare filename in the directory this turn actually wrote into. * diff --git a/desktop/src/lib/attachmentImages.ts b/desktop/src/lib/attachmentImages.ts index c09da91b..6b22d7dc 100644 --- a/desktop/src/lib/attachmentImages.ts +++ b/desktop/src/lib/attachmentImages.ts @@ -16,6 +16,11 @@ export function isInlineImagePath(pathOrName: string | undefined): boolean { return !!pathOrName && INLINE_IMAGE_EXTENSION_RE.test(pathOrName) } +/** Host-managed ImageGen results already render through their dedicated result card. */ +export function isManagedGeneratedImagePath(filePath: string): boolean { + return filePath.replaceAll('\\', '/').includes('/.claude/cc-haha/generated-images/') +} + /** Serves a local absolute image path through the local server. */ export function localImageFileUrl(filePath: string): string { return `${getBaseUrl()}/api/filesystem/file?path=${encodeURIComponent(filePath)}` diff --git a/desktop/src/lib/markdownImages.test.ts b/desktop/src/lib/markdownImages.test.ts index 6b2ce47a..d175c673 100644 --- a/desktop/src/lib/markdownImages.test.ts +++ b/desktop/src/lib/markdownImages.test.ts @@ -1,6 +1,7 @@ import { describe, expect, it } from 'vitest' import { + createAssistantMarkdownImageResolver, createWorkspaceMarkdownImageResolver, isSafeMarkdownImageSource, } from './markdownImages' @@ -102,3 +103,33 @@ describe('createWorkspaceMarkdownImageResolver', () => { expect(resolve('data:text/html;base64,AAAA')).toBeNull() }) }) + +describe('createAssistantMarkdownImageResolver', () => { + const resolve = createAssistantMarkdownImageResolver({ + baseUrl: 'http://127.0.0.1:3456', + sessionId: 'session-1', + }) + + it('routes relative and absolute local paths through the session workspace sandbox', () => { + expect(resolve('images/../01.png')).toBe( + 'http://127.0.0.1:3456/preview-fs/session-1/01.png', + ) + expect(resolve('/repo/output/02.png')).toBe( + 'http://127.0.0.1:3456/preview-fs/session-1//repo/output/02.png', + ) + }) + + it('keeps safe in-memory images without opening a network path', () => { + expect(resolve('data:image/png;base64,AAAA')).toBe('data:image/png;base64,AAAA') + expect(resolve('blob:https://desktop.invalid/1234')).toBe('blob:https://desktop.invalid/1234') + }) + + it('rejects remote, loopback, schemes, and workspace escapes', () => { + expect(resolve('https://attacker.example/track.png')).toBeNull() + expect(resolve('http://127.0.0.1:3456/status.png')).toBeNull() + expect(resolve('//attacker.example/track.png')).toBeNull() + expect(resolve('file:///repo/01.png')).toBeNull() + expect(resolve('../outside.png')).toBeNull() + expect(resolve('#fragment')).toBeNull() + }) +}) diff --git a/desktop/src/lib/markdownImages.ts b/desktop/src/lib/markdownImages.ts index f30c74c7..b671f879 100644 --- a/desktop/src/lib/markdownImages.ts +++ b/desktop/src/lib/markdownImages.ts @@ -22,10 +22,88 @@ export type WorkspaceMarkdownImageContext = { workDir?: string | null } +export type AssistantMarkdownImageContext = { + /** Base URL of the local server (see `getServerBaseUrl`). */ + baseUrl: string + sessionId: string +} + function splitPathSegments(value: string): string[] { return value.replace(/\\/g, '/').split('/') } +function normalizeAssistantLocalImagePath(value: string): string | null { + const slashed = value.replace(/\\/g, '/') + const isPosixAbsolute = slashed.startsWith('/') + const driveMatch = /^([A-Za-z]:)\/?/.exec(slashed) + const isAbsolute = isPosixAbsolute || Boolean(driveMatch) + const prefix = driveMatch?.[1] + ? `${driveMatch[1]}/` + : isPosixAbsolute + ? '/' + : '' + const body = driveMatch?.[0] + ? slashed.slice(driveMatch[0].length) + : slashed.replace(/^\/+/, '') + const stack: string[] = [] + + for (const segment of splitPathSegments(body)) { + if (!segment || segment === '.') continue + if (segment === '..') { + if (stack.length > 0) { + stack.pop() + continue + } + // Absolute paths remain subject to the server's canonical workdir check. + // A relative path must never make the browser normalize out of /preview-fs. + if (!isAbsolute) return null + continue + } + stack.push(segment) + } + + if (stack.length === 0) return null + return `${prefix}${stack.join('/')}` +} + +/** + * Create an image resolver for finalized assistant prose. + * + * Unlike the trusted workspace-document resolver below, this accepts no network + * URL at all. Local relative paths and absolute paths are both sent through the + * session-scoped `/preview-fs` route; the server canonicalizes the target and + * rejects anything outside that session's workdir. The caller must only attach + * this resolver after streaming finishes so an unfinished path cannot trigger a + * request. + */ +export function createAssistantMarkdownImageResolver( + context: AssistantMarkdownImageContext, +): (src: string) => string | null { + return (src: string): string | null => { + const trimmed = src.trim() + if (!trimmed) return null + if (isSafeMarkdownImageSource(trimmed)) return trimmed + if (trimmed.startsWith('#') || trimmed.startsWith('//')) return null + if (/^[a-z][a-z0-9+.-]*:/i.test(trimmed)) return null + + const withoutSuffix = trimmed.split('#')[0]!.split('?')[0]! + let decoded = withoutSuffix + try { + decoded = decodeURIComponent(withoutSuffix) + } catch { + // A malformed escape cannot name a stable local file. + return null + } + + const localPath = normalizeAssistantLocalImagePath( + decoded.split('#')[0]!.split('?')[0]!, + ) + return localPath + ? previewFsUrl(context.baseUrl, context.sessionId, localPath) + : null + } +} + /** * Create an `img src` resolver for trusted, user-owned Markdown documents (the * workspace file preview). Untrusted assistant output must NOT get a resolver —