fix(mcp): persist generated images - #2094
Conversation
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
📝 WalkthroughWalkthroughThe change persists temporary MCP image URLs as ChangesMCP image persistence
Estimated code review effort: 4 (Complex) | ~60 minutes Sequence Diagram(s)sequenceDiagram
participant MCPTool
participant McpService
participant ImageCache
participant ToolManager
participant MarkdownRenderer
MCPTool->>McpService: return temporary image URLs
McpService->>ImageCache: cache unique image URLs
ImageCache-->>McpService: return imgcache:// references
McpService-->>ToolManager: persist normalized tool content
ToolManager->>ImageCache: resolve exact cached arguments
ImageCache-->>ToolManager: return data URLs
McpService->>MarkdownRenderer: pass promoted local image sources
MarkdownRenderer-->>McpService: remove matching Markdown image nodes
Possibly related PRs
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches 💡 1📝 Generate docstrings 💡
🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
There was a problem hiding this comment.
🧹 Nitpick comments (1)
test/main/platform/imageCache.test.ts (1)
48-77: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAdd a symbolic-link escape test.
The resolver rejects symbolic links and validates the real path, but this suite does not test that contract. Add a platform-compatible test for an image-cache symlink that targets a file outside the cache root.
As per coding guidelines, “Add the smallest regression test for user-visible behavior or a documented contract.”
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@test/main/platform/imageCache.test.ts` around lines 48 - 77, Add a focused test alongside the existing resolveCachedImageDataUrl rejection cases that creates an image-cache symbolic link pointing to a file outside electronMock.userDataPath/images, then asserts the resolver rejects it with the established invalid cached image path error. Use the platform-compatible symlink setup already supported by the test environment and clean up any created files or links.Source: Coding guidelines
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Nitpick comments:
In `@test/main/platform/imageCache.test.ts`:
- Around line 48-77: Add a focused test alongside the existing
resolveCachedImageDataUrl rejection cases that creates an image-cache symbolic
link pointing to a file outside electronMock.userDataPath/images, then asserts
the resolver rejects it with the established invalid cached image path error.
Use the platform-compatible symlink setup already supported by the test
environment and clean up any created files or links.
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro Plus
Run ID: 13df7559-8958-40c3-b3bb-2e487af1d296
📒 Files selected for processing (16)
docs/issues/mcp-generated-image-persistence/spec.mdsrc/main/lib/toolCallImagePreviews.tssrc/main/mcp/index.tssrc/main/mcp/toolManager.tssrc/main/platform/imageCache.tssrc/renderer/src/components/markdown/MarkdownRenderer.vuesrc/renderer/src/components/message/MessageBlockContent.vuesrc/renderer/src/components/message/MessageItemAssistant.vuetest/main/agent/deepchat/runtime/contextBuilder.test.tstest/main/lib/toolCallImagePreviews.test.tstest/main/mcp/mcpService.test.tstest/main/mcp/toolManager.test.tstest/main/platform/imageCache.test.tstest/renderer/components/MarkdownRenderer.test.tstest/renderer/components/message/MessageBlockContent.test.tstest/renderer/components/message/MessageItemAssistant.test.ts
|
Note GitHub couldn't provide a complete incremental comparison for this pull request, so CodeRabbit is performing a full review instead. This review may take a little longer. |
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@src/main/platform/imageCache.ts`:
- Around line 115-119: Update the URI validation in the resolver around
normalizedSource to match cachePreviewData’s case-insensitive imgcache://
contract, accepting persisted references with uppercase scheme characters while
preserving rejection of unsupported schemes. Add one regression test covering an
uppercase-scheme cached image reference.
🪄 Autofix
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro Plus
Run ID: cfbc2f19-dc91-44d4-adce-6b86ca205628
📒 Files selected for processing (16)
docs/issues/mcp-generated-image-persistence/spec.mdsrc/main/lib/toolCallImagePreviews.tssrc/main/mcp/index.tssrc/main/mcp/toolManager.tssrc/main/platform/imageCache.tssrc/renderer/src/components/markdown/MarkdownRenderer.vuesrc/renderer/src/components/message/MessageBlockContent.vuesrc/renderer/src/components/message/MessageItemAssistant.vuetest/main/agent/deepchat/runtime/contextBuilder.test.tstest/main/lib/toolCallImagePreviews.test.tstest/main/mcp/mcpService.test.tstest/main/mcp/toolManager.test.tstest/main/platform/imageCache.test.tstest/renderer/components/MarkdownRenderer.test.tstest/renderer/components/message/MessageBlockContent.test.tstest/renderer/components/message/MessageItemAssistant.test.ts
🚧 Files skipped from review as they are similar to previous changes (13)
- test/renderer/components/MarkdownRenderer.test.ts
- src/renderer/src/components/message/MessageItemAssistant.vue
- src/main/mcp/index.ts
- src/renderer/src/components/message/MessageBlockContent.vue
- test/main/mcp/mcpService.test.ts
- src/renderer/src/components/markdown/MarkdownRenderer.vue
- test/main/mcp/toolManager.test.ts
- test/main/lib/toolCallImagePreviews.test.ts
- src/main/mcp/toolManager.ts
- test/renderer/components/message/MessageItemAssistant.test.ts
- test/renderer/components/message/MessageBlockContent.test.ts
- test/main/agent/deepchat/runtime/contextBuilder.test.ts
- src/main/lib/toolCallImagePreviews.ts
|
Addressed the cache URI compatibility and safety follow-up in 565380d:
Validated locally:
|
yyhhyyyyyy
left a comment
There was a problem hiding this comment.
The overall approach is coherent: normalize temporary URLs before they first enter provider
history, resolve imgcache:// only in cloned execution arguments, and deduplicate the
presentation without rewriting persisted Markdown.
I did not find a P0 issue, but I think the following two P1 issues should be addressed before
merge. There are also three P2 correctness/resource issues closely related to the new behavior.
-
[P1] The automatic download path needs a trust-aware SSRF policy.
extractStringImagePreviews
now turns any matching URL in MCP-controlled text into a main-process download.
cacheImageFromUrl
accepts HTTP, follows redirects by default, and does not reject loopback, link-local, or
private-network destinations.This gives an untrusted or compromised remote MCP server a client-side SSRF primitive.
The extension check does not close the boundary because a public.pngURL can redirect to
an arbitrary private endpoint. The resulting cache URI can subsequently be resolved to a
data URL and sent to another MCP call, creating a possible read-and-return path for the
fetched bytes.Please make the policy depend on the MCP server's trust/source, validate every redirect
target, and reject private destinations for remote servers unless explicitly allowed.
A public-to-private redirect regression test would be important here. -
[P1] Automatic caching is unbounded in both count and response size.
The loop in
prepareToolCallImageContent
downloads every distinct matching URL sequentially. Each failed URL may consume the full
10-second timeout, while axios buffers the complete response as anarraybufferwithout a
content-length limit. The abort signal only stops waiting at the caller; it is not passed to
axios, so the active request may continue and write a file after cancellation.A normal search/browser MCP result containing many image links can therefore block a tool
round for a long time, while a hostile result can cause excessive memory or disk use.
Please add a maximum image count, per-image and aggregate byte limits, a total time budget,
and propagate cancellation to the underlying request. -
[P2] The cache producer does not preserve the actual MIME type.
getImageExtensionFromMimeType
stores AVIF, BMP, ICO,application/octet-stream, and every unknown content type as.jpg.
The resolver later trusts that extension and emitsdata:image/jpeg, even when the bytes are
another format. A200HTML or JSON response is also treated as a durable image.This breaks the stated MIME-correct data URL contract and may make follow-up image tools
reject otherwise valid cached images. Please validate that the response is a supported image
and keep the producer/resolver MIME mapping consistent. -
[P2] Raw
replaceAllcan rewrite unrelated URLs.replaceImageReferences
replaces arbitrary string occurrences rather than the URL tokens identified during
extraction. For example, after cachinghttps://host/a.png, an unrelated
https://host/a.png.jsonin the same result is changed to
imgcache://cached.png.json.Please replace only the exact token spans collected by the extractor and add a regression
test for URLs where one value is a prefix of another. -
[P2] Expanded MCP arguments have no aggregate size limit.
resolveCachedImageArguments
enforces an 8 MiB limit per file, but it does not cap the number of references, deduplicate
repeated reads, or validate the final serialized request size. Repeating a small
imgcache://value several times can expand a small model-generated argument into a
request of tens or hundreds of MiB.Please add per-call memoization and an aggregate encoded-size/reference-count limit before
invoking the MCP client.
The renderer-side deduplication and the execution-only argument cloning otherwise look sound.
The two P1 items above are merge blockers from my perspective.
There was a problem hiding this comment.
Actionable comments posted: 1
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (2)
src/main/platform/imageCache.ts (2)
296-351: 🔒 Security & Privacy | 🟠 Major | ⚡ Quick winPath Traversal (CWE-367): Time-of-check Time-of-use (TOCTOU) Race Condition
Reachability: Internal
Read the validated resolved path atomically.
fs.promises.readFile(realFilePath, { signal })still opens the resolved pathname, so the real-path containment check can race with a local symlink replacement. Resolve once after validation, or open a resolved no-follow descriptor and read from it, then add a race-oriented regression test.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@src/main/platform/imageCache.ts` around lines 296 - 351, Update resolveCachedImageDataUrl to read the validated cache file atomically after containment checks, avoiding a second pathname-based lookup that can race with symlink replacement; use a resolved no-follow file descriptor or equivalent single-resolution approach and preserve abort handling and MIME validation. Add a regression test that replaces the path with a symlink during resolution and verifies the outside file is not read.
208-212: 🗄️ Data Integrity & Integration | 🟠 Major | ⚡ Quick winMake cancelled image writes atomic.
fs.promises.writeFilewith anAbortSignalis best-effort and can leave a truncated file after opening a new cache filename. IncacheImageFromUrl, the internal timeout catch returns the original URL but leavessaveFileName; incacheImageFromBase64, the same path can rethrow but may leave the final filename too. Write to a unique temporary file and rename after success, or delete the target on every failed write. Add cancellation coverage for both the URL and Base64 paths.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@src/main/platform/imageCache.ts` around lines 208 - 212, Make image writes atomic in cacheImageFromUrl and cacheImageFromBase64 by writing data to a unique temporary file in cacheDir, then renaming it to saveFileName only after the write completes successfully; clean up the temporary file on cancellation or any write failure, and ensure no partially written final cache file remains. Add cancellation coverage for both URL and Base64 caching paths.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@src/main/platform/imageCache.ts`:
- Around line 104-109: Update the SVG detection logic around the bounded `text`
preprocessing and root-element regex to skip valid XML prolog content, including
`<!DOCTYPE svg>`, comments, and processing instructions, before matching the
`<svg>` root. Preserve the 4096-byte bound and existing BOM/whitespace handling,
and add regression cases covering these prologs for both URL and Base64 cache
paths.
---
Outside diff comments:
In `@src/main/platform/imageCache.ts`:
- Around line 296-351: Update resolveCachedImageDataUrl to read the validated
cache file atomically after containment checks, avoiding a second pathname-based
lookup that can race with symlink replacement; use a resolved no-follow file
descriptor or equivalent single-resolution approach and preserve abort handling
and MIME validation. Add a regression test that replaces the path with a symlink
during resolution and verifies the outside file is not read.
- Around line 208-212: Make image writes atomic in cacheImageFromUrl and
cacheImageFromBase64 by writing data to a unique temporary file in cacheDir,
then renaming it to saveFileName only after the write completes successfully;
clean up the temporary file on cancellation or any write failure, and ensure no
partially written final cache file remains. Add cancellation coverage for both
URL and Base64 caching paths.
🪄 Autofix
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro Plus
Run ID: e8168e62-d979-4959-9430-98ed606b647a
📒 Files selected for processing (3)
docs/issues/mcp-generated-image-persistence/spec.mdsrc/main/platform/imageCache.tstest/main/platform/imageCache.test.ts
🚧 Files skipped from review as they are similar to previous changes (2)
- test/main/platform/imageCache.test.ts
- docs/issues/mcp-generated-image-persistence/spec.md
| const text = data | ||
| .subarray(0, 4096) | ||
| .toString('utf8') | ||
| .replace(/^\uFEFF/, '') | ||
| .trimStart() | ||
| if (/^(?:<\?xml[^>]*>\s*)?<svg(?:\s|>)/i.test(text)) return 'svg' |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
Handle valid SVG prologs.
This detector accepts only an optional XML declaration followed immediately by <svg>. Valid SVG files can include a <!DOCTYPE svg> declaration, comments, or processing instructions before the root element. The URL and Base64 cache paths then return the original source instead of a stable imgcache:// reference. Add regression cases and handle the bounded XML prolog before checking the root element.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@src/main/platform/imageCache.ts` around lines 104 - 109, Update the SVG
detection logic around the bounded `text` preprocessing and root-element regex
to skip valid XML prolog content, including `<!DOCTYPE svg>`, comments, and
processing instructions, before matching the `<svg>` root. Preserve the
4096-byte bound and existing BOM/whitespace handling, and add regression cases
covering these prologs for both URL and Base64 cache paths.
Summary
Root cause
MCP image generators can return temporary signed image URLs inside ordinary text. Those URLs were neither detected nor cached, so the provider and persisted transcript retained an expiring remote URL. After expiry or restart, the renderer correctly reloaded that URL and received an authorization failure.
UI behavior
BEFORE
AFTER
Unrelated Markdown images and links remain visible.
Validation
The SDD records the persistence contract, execution-only image resolution, security boundaries, and release smoke-test procedure.
Summary by CodeRabbit