Add Mermaid diagram export actions

.changeset/clever-mermaids-export.md

@@ -0,0 +1,5 @@
+---
+"kilo-code": minor
+---
+
+Support copying, previewing, and exporting rendered Mermaid diagrams.

packages/kilo-vscode/docs/chat-ui-features/mermaid-diagram-features.md

@@ -6,11 +6,12 @@ Chat Markdown renders fenced `mermaid` code blocks as diagrams after a response
 
 - Valid `mermaid` fences render inline as SVG diagrams.
 - The original Mermaid source remains available through the existing code-block copy button.
+- Rendered diagrams include actions to copy the Mermaid source, copy SVG, copy PNG, save SVG, save PNG, and open an image preview.
 - Invalid Mermaid syntax shows a contained error state and keeps the source visible.
 - Diagrams are not rendered while a message is streaming, which avoids repeated parse/render work on every token.
 - Diagram colors are derived from the active VS Code/Kilo CSS variables so light, dark, and high-contrast themes can render with matching backgrounds, text, borders, and link colors.
 
 ## Limitations
 
 - Mermaid is bundled by the current webview build, so bundle splitting remains a future optimization.
-- Advanced legacy actions are not restored yet: AI syntax fixing, PNG open/save, export, and zoom modal.
+- Advanced legacy actions are not restored yet: AI syntax fixing and zoom modal.

packages/kilo-vscode/src/KiloProvider.ts

@@ -16,6 +16,7 @@ import type { EditorContext, IndexingStatus } from "./services/cli-backend/types
 import { FileIgnoreController } from "./services/autocomplete/shims/FileIgnoreController"
 import { ChatTextAreaAutocomplete } from "./services/autocomplete/chat-autocomplete/ChatTextAreaAutocomplete"
 import { buildWebviewHtml } from "./utils"
+import { saveImage } from "./kilo-provider/save-image"
 import { TelemetryProxy, type TelemetryPropertiesProvider } from "./services/telemetry"
 import {
   sessionToWebview,
@@ -742,7 +743,6 @@ export class KiloProvider implements vscode.WebviewViewProvider, TelemetryProper
             console.error("[Kilo New] handleForkSession failed:", e),
           )
           break
-
         case "retryConnection":
           console.log("[Kilo New] KiloProvider: 🔄 Retrying connection...")
           this.initializeConnection().catch((e) =>
@@ -755,6 +755,8 @@ export class KiloProvider implements vscode.WebviewViewProvider, TelemetryProper
         case "previewImage":
           this.handlePreviewImage(message.dataUrl, message.filename)
           break
+        case "saveImage":
+          return void saveImage(this.getWorkspaceDirectory(this.currentSession?.id), message.dataUrl, message.filename).catch((err) => console.error("[Kilo New] KiloProvider: Failed to save image:", err))
         case "openFile":
           if (message.filePath) {
             this.handleOpenFile(message.filePath, message.line, message.column)

packages/kilo-vscode/src/agent-manager/AgentManagerProvider.ts

@@ -455,7 +455,7 @@ export class AgentManagerProvider implements Disposable {
       this.host.copyToClipboard(m.text)
       return null
     }
-    if (m.type === "previewImage") return msg
+    if (m.type === "previewImage" || m.type === "saveImage") return msg
     if (m.type === "agentManager.showExistingLocalTerminal") {
       this.terminalManager.syncLocalOnSessionSwitch()
       return null

packages/kilo-vscode/src/agent-manager/types.ts

@@ -564,6 +564,12 @@ interface PreviewImageIn {
   filename: string
 }
 
+interface SaveImageIn {
+  type: "saveImage"
+  dataUrl: string
+  filename: string
+}
+
 interface LoadMessagesIn {
   type: "loadMessages"
   sessionID: string
@@ -750,6 +756,7 @@ export type AgentManagerInMessage =
   | OpenFileIn
   | GenericOpenFileIn
   | PreviewImageIn
+  | SaveImageIn
   | LoadMessagesIn
   | SendMessageIn
   | SendCommandIn

packages/kilo-vscode/src/kilo-provider/save-image.ts

@@ -0,0 +1,16 @@
+import * as path from "path"
+import * as vscode from "vscode"
+import { parseImage } from "../image-preview"
+
+export async function saveImage(dir: string, dataUrl: string, filename: string) {
+  const img = parseImage(dataUrl, filename)
+  if (!img) return
+
+  const uri = await vscode.window.showSaveDialog({
+    defaultUri: vscode.Uri.file(path.join(dir, img.name)),
+    filters: { Images: [img.ext] },
+    saveLabel: "Save",
+  })
+  if (!uri) return
+  await vscode.workspace.fs.writeFile(uri, img.data)
+}

packages/kilo-vscode/webview-ui/src/App.tsx

@@ -149,19 +149,22 @@ export const DataBridge: Component<{ children: any }> = (props) => {
   }
 
   return (
-    <DataProvider
-      data={data}
-      directory={directory()}
-      // @ts-expect-error — onPermissionRespond/onQuestion* are extension-specific props not yet in kilo-ui's DataProvider types
-      onPermissionRespond={respond}
-      onQuestionReply={reply}
-      onQuestionReject={reject}
-      onOpenFile={open}
-      onOpenDiff={openDiff}
-      onOpenUrl={openUrl}
-    >
-      {props.children}
-    </DataProvider>
+    <>
+      <MermaidImageBridge />
+      <DataProvider
+        data={data}
+        directory={directory()}
+        // @ts-expect-error — onPermissionRespond/onQuestion* are extension-specific props not yet in kilo-ui's DataProvider types
+        onPermissionRespond={respond}
+        onQuestionReply={reply}
+        onQuestionReject={reject}
+        onOpenFile={open}
+        onOpenDiff={openDiff}
+        onOpenUrl={openUrl}
+      >
+        {props.children}
+      </DataProvider>
+    </>
   )
 }
 
@@ -178,6 +181,35 @@ export const LanguageBridge: Component<{ children: any }> = (props) => {
   )
 }
 
+type MermaidImageEvent = CustomEvent<{ dataUrl: string; filename: string }>
+
+export const MermaidImageBridge: Component = () => {
+  const vscode = useVSCode()
+
+  onMount(() => {
+    const preview = (event: Event) => {
+      const detail = (event as MermaidImageEvent).detail
+      if (!detail?.dataUrl || !detail.filename) return
+      event.preventDefault()
+      vscode.postMessage({ type: "previewImage", dataUrl: detail.dataUrl, filename: detail.filename })
+    }
+    const save = (event: Event) => {
+      const detail = (event as MermaidImageEvent).detail
+      if (!detail?.dataUrl || !detail.filename) return
+      event.preventDefault()
+      vscode.postMessage({ type: "saveImage", dataUrl: detail.dataUrl, filename: detail.filename })
+    }
+    window.addEventListener("kilo:preview-image", preview)
+    window.addEventListener("kilo:save-image", save)
+    onCleanup(() => {
+      window.removeEventListener("kilo:preview-image", preview)
+      window.removeEventListener("kilo:save-image", save)
+    })
+  })
+
+  return null
+}
+
 // Inner app component that uses the contexts
 const AppContent: Component = () => {
   const [currentView, setCurrentView] = createSignal<ViewType>("newTask")

packages/kilo-vscode/webview-ui/src/types/messages/webview-messages.ts

@@ -798,6 +798,12 @@ export interface PreviewImageRequest {
   filename: string
 }
 
+export interface SaveImageRequest {
+  type: "saveImage"
+  dataUrl: string
+  filename: string
+}
+
 // Set default base branch (webview → extension)
 export interface SetDefaultBaseBranchRequest {
   type: "agentManager.setDefaultBaseBranch"
@@ -1119,6 +1125,7 @@ export type WebviewMessage =
   | RetryConnectionRequest
   | OpenSubAgentViewerRequest
   | PreviewImageRequest
+  | SaveImageRequest
   | SetDefaultBaseBranchRequest
   | AgentManagerOpenSessionsMessage
   | RequestAutoApproveStateMessage

packages/ui/src/components/markdown.tsx

@@ -7,7 +7,7 @@ import { ComponentProps, createEffect, createResource, createSignal, onCleanup,
 import { isServer } from "solid-js/web"
 import { stream } from "./markdown-stream"
 import { tryFastRender } from "../kilocode/markdown-fast-path" // kilocode_change
-import { hasMermaid, preserveMermaid, renderMermaid } from "../kilocode/markdown-mermaid" // kilocode_change
+import { hasMermaid, preserveMermaid, renderMermaid, type MermaidLabels } from "../kilocode/markdown-mermaid" // kilocode_change
 
 type Entry = {
   hash: string
@@ -340,6 +340,22 @@ export function Markdown(
       copied: i18n.t("ui.message.copied"),
     }
 
+    // kilocode_change start: Mermaid diagram rendering
+    const mermaid = {
+      rendering: i18n.t("ui.mermaid.rendering"),
+      renderError: (message: string) => i18n.t("ui.mermaid.renderError", { message }),
+      errorDefault: i18n.t("ui.mermaid.errorDefault"),
+      errorEmpty: i18n.t("ui.mermaid.errorEmpty"),
+      copied: i18n.t("ui.message.copied"),
+      copySource: i18n.t("ui.mermaid.copySource"),
+      copySvg: i18n.t("ui.mermaid.copySvg"),
+      copyPng: i18n.t("ui.mermaid.copyPng"),
+      downloadSvg: i18n.t("ui.mermaid.downloadSvg"),
+      downloadPng: i18n.t("ui.mermaid.downloadPng"),
+      openPreview: i18n.t("ui.mermaid.openPreview"),
+    }
+    // kilocode_change end
+
     // kilocode_change start
     const fast = tryFastRender(container, content, local.streaming, decorate, setupCodeCopy, () => labels, copyCleanup)
     if (fast.handled) {
@@ -352,7 +368,7 @@ export function Markdown(
         pendingLabels = undefined
       }
       copyCleanup = fast.copyCleanup
-      kickMermaid(container, local.streaming ?? false)
+      kickMermaid(container, local.streaming ?? false, mermaid)
       kickHighlight(container, labels)
       return
     }
@@ -426,7 +442,7 @@ export function Markdown(
       })
       // kilocode_change end
 
-      kickMermaid(container, local.streaming ?? false) // kilocode_change
+      kickMermaid(container, local.streaming ?? false, mermaid) // kilocode_change
       kickHighlight(container, nextLabels)
     })
     // kilocode_change end
@@ -456,7 +472,7 @@ export function Markdown(
   // kilocode_change end
 
   // kilocode_change start: Mermaid diagram rendering
-  function kickMermaid(container: HTMLDivElement, streaming: boolean) {
+  function kickMermaid(container: HTMLDivElement, streaming: boolean, labels: MermaidLabels) {
     mermaidState.signal.aborted = true
     mermaidState.gen++
     if (!hasMermaid(container)) return
@@ -465,7 +481,7 @@ export function Markdown(
     const gen = mermaidState.gen
     const signal = { aborted: false }
     mermaidState.signal = signal
-    void renderMermaid(container, signal).catch((err) => {
+    void renderMermaid(container, signal, labels).catch((err) => {
       if (gen !== mermaidState.gen || signal.aborted) return
       console.warn("Mermaid render failed", err)
     })

packages/ui/src/i18n/en.ts

@@ -102,6 +102,16 @@ export const dict: Record<string, string> = {
   "ui.textField.copied": "Copied",
 
   "ui.imagePreview.alt": "Image preview",
+  "ui.mermaid.rendering": "Rendering Mermaid diagram...",
+  "ui.mermaid.renderError": "Mermaid render failed: {{message}}",
+  "ui.mermaid.errorDefault": "Unable to render Mermaid diagram.",
+  "ui.mermaid.errorEmpty": "Mermaid rendered an empty diagram.",
+  "ui.mermaid.copySource": "Copy Mermaid source",
+  "ui.mermaid.copySvg": "Copy SVG",
+  "ui.mermaid.copyPng": "Copy PNG",
+  "ui.mermaid.downloadSvg": "Download SVG",
+  "ui.mermaid.downloadPng": "Download PNG",
+  "ui.mermaid.openPreview": "Open image preview",
   "ui.scrollView.ariaLabel": "scrollable content",
 
   "ui.tool.read": "Read",