fix(desktop): preserve inline image placement (#1240)

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(
       <AssistantMessage
         sessionId="s1"
@@ -194,6 +194,7 @@ describe('AssistantMessage output-target cards', () => {
       '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()
   })
 

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 && (
             <InlineImageGallery

desktop/src/components/chat/InlineImageGallery.tsx

@@ -1,7 +1,10 @@
 import { useMemo, useState } from 'react'
 import { ImageGalleryModal } from './ImageGalleryModal'
-import { localImageFileUrl } from '../../lib/attachmentImages'
-import { extractAssistantOutputTargets } from '../../lib/assistantOutputTargets'
+import { isManagedGeneratedImagePath, localImageFileUrl } from '../../lib/attachmentImages'
+import {
+  extractAssistantOutputTargets,
+  extractMarkdownImageSources,
+} from '../../lib/assistantOutputTargets'
 import { isAbsoluteLocalPath, previewFsUrl } from '../../lib/handlePreviewLink'
 import { getServerBaseUrl } from '../../lib/desktopRuntime'
 
@@ -34,6 +37,17 @@ function fileName(filePath: string): string {
   return filePath.split('/').pop() || filePath
 }
 
+function normalizeImageReference(value: string): string {
+  const withoutSuffix = value.trim().split('#')[0]!.split('?')[0]!
+  let decoded = withoutSuffix
+  try {
+    decoded = decodeURIComponent(withoutSuffix)
+  } catch {
+    // Keep malformed escapes comparable without turning them into a URL.
+  }
+  return decoded.replaceAll('\\', '/').replace(/^\.\//, '')
+}
+
 type GalleryImage = {
   src: string
   name: string
@@ -56,6 +70,11 @@ type Props = {
 export function InlineImageGallery({ text, sessionId, workDir, changedFiles, suppressManagedGeneratedImages = false }: Props) {
   const [activeIndex, setActiveIndex] = useState<number | null>(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/')
-}

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(<MessageList sessionId={ACTIVE_TAB} />)
+    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<{

desktop/src/components/markdown/MarkdownRenderer.tsx

@@ -28,11 +28,11 @@ type Props = {
   streaming?: boolean
   onLinkClick?: (href: string, event: ReactMouseEvent<HTMLDivElement>) => 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<HTMLImageElement | HTMLSourceElement>('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 <img> 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 <img> tags get
+      // the same treatment as Markdown images.
       const resolved = src ? resolveImageSrc(src) : null
       if (resolved) image.setAttribute('src', resolved)
       else {

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.
  *

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)}`

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()
+  })
+})

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 —