Merge pull request #13370 from Kilo-Org/fix-undo-prompt-attachment-restoration

fix(vscode): restore images when undoing prompts

.changeset/restore-undo-attachments.md

@@ -0,0 +1,5 @@
+---
+"kilo-code": patch
+---
+
+Restore attached images when undoing the last prompt.

packages/kilo-vscode/tests/unit/session-utils.test.ts

@@ -23,6 +23,7 @@ import {
   removeSessionToolPartsForMessage,
   upsertSessionToolPart,
   recentSessions,
+  revertPromptState,
 } from "../../webview-ui/src/context/session-utils"
 import type { Message, Part, ToolPart } from "../../webview-ui/src/types/messages"
 
@@ -983,3 +984,52 @@ describe("sessionThroughput", () => {
     expect(sessionThroughput([])).toBeUndefined()
   })
 })
+
+describe("revertPromptState", () => {
+  const text = (value: string, synthetic = false) =>
+    ({ type: "text", id: `t-${value}`, text: value, synthetic }) as Part
+  const file = (overrides: Partial<Extract<Part, { type: "file" }>>) =>
+    ({ type: "file", id: "f", mime: "text/plain", url: "", ...overrides }) as Extract<Part, { type: "file" }>
+
+  it("joins non-synthetic text parts and drops synthetic ones", () => {
+    const state = revertPromptState([text("Hello "), { ...text("hidden", true) } as Part, text("world")])
+    expect(state.text).toBe("Hello world")
+  })
+
+  it("restores inline image attachments from data URLs only", () => {
+    const state = revertPromptState([
+      file({ mime: "image/png", url: "data:image/png;base64,abc", filename: "shot.png" }),
+      file({ mime: "image/jpeg", url: "https://example.com/x.jpg" }),
+      file({ mime: "application/pdf", url: "data:application/pdf;base64,def" }),
+    ])
+    expect(state.images).toEqual([{ dataUrl: "data:image/png;base64,abc", mime: "image/png", filename: "shot.png" }])
+  })
+
+  it("collects mention paths but excludes session references from paths", () => {
+    const state = revertPromptState([
+      file({ source: { type: "file", path: "a b.txt", text: { value: "@a b.txt", start: 0, end: 8 } } }),
+      file({ url: "session:ses_1", filename: "Old chat" }),
+    ])
+    expect(state.paths).toEqual(["a b.txt"])
+  })
+
+  it("maps past-chat references to session items with title fallbacks", () => {
+    const state = revertPromptState([
+      file({
+        url: "session:ses_1",
+        source: { type: "file", path: "", text: { value: "@Renamed chat", start: 0, end: 13 } },
+      }),
+      file({ url: "session:ses_2", filename: "Fallback title" }),
+    ])
+    expect(state.sessions).toEqual([
+      { id: "ses_1", title: "Renamed chat", updated: 0 },
+      { id: "ses_2", title: "Fallback title", updated: 0 },
+    ])
+  })
+
+  it("returns empty collections for tool-only messages", () => {
+    const part: Part = { type: "tool", id: "p1", tool: "bash", state: { status: "running", input: {} } }
+    const state = revertPromptState([part])
+    expect(state).toEqual({ text: "", paths: [], sessions: [], images: [] })
+  })
+})

packages/kilo-vscode/webview-ui/src/components/chat/PromptInput.tsx

@@ -724,6 +724,18 @@ export const PromptInput: Component<PromptInputProps> = (props) => {
         textareaRef.value = message.text
         adjustHeight()
       }
+      // When present, images are authoritative: replace current attachments
+      // (an empty array clears them, e.g. on redo). Absent leaves them alone.
+      if (message.images) {
+        const imgs = message.images.map((img) => ({
+          id: crypto.randomUUID(),
+          filename: img.filename ?? "image",
+          mime: img.mime,
+          dataUrl: img.dataUrl,
+        }))
+        imageAttach.replace(imgs)
+        imageDrafts.set(draftKey(), imgs)
+      }
     }
 
     if (message.type === "appendChatBoxMessage") {

packages/kilo-vscode/webview-ui/src/context/session-utils.ts

@@ -34,6 +34,40 @@ export function messageParts(messages: Message[]): Record<string, Part[]> {
   return parts
 }
 
+/** Prompt input state rebuilt from a reverted user message's parts. */
+export interface RevertPromptState {
+  text: string
+  paths: string[]
+  sessions: Array<{ id: string; title: string; updated: number }>
+  images: Array<{ dataUrl: string; mime: string; filename?: string }>
+}
+
+/**
+ * Extract the prompt content of a user message for restoration into the input
+ * box after a revert. Inline images are returned as data URLs so PromptInput
+ * can re-attach them without re-uploading.
+ */
+export function revertPromptState(parts: readonly Part[]): RevertPromptState {
+  const files = parts.filter((p): p is Extract<Part, { type: "file" }> => p.type === "file")
+  return {
+    text: parts
+      .filter((p) => p.type === "text" && !(p as { synthetic?: boolean }).synthetic)
+      .map((p) => (p as { text: string }).text ?? "")
+      .join(""),
+    paths: files.map((p) => p.source?.path).filter((p): p is string => !!p && !p.startsWith("session:")),
+    sessions: files
+      .filter((p) => p.url.startsWith("session:"))
+      .map((p) => ({
+        id: p.url.slice("session:".length),
+        title: p.source?.text?.value.replace(/^@/, "") ?? p.filename ?? p.url,
+        updated: 0,
+      })),
+    images: files
+      .filter((p) => p.mime.startsWith("image/") && p.url.startsWith("data:"))
+      .map((p) => ({ dataUrl: p.url, mime: p.mime, filename: p.filename })),
+  }
+}
+
 type SnapshotPart = {
   type?: string
   text?: string

packages/kilo-vscode/webview-ui/src/context/session.tsx

@@ -68,6 +68,7 @@ import {
   reconcileSessionToolParts,
   removeSessionToolPart,
   removeSessionToolPartsForMessage,
+  revertPromptState,
   upsertSessionToolPart,
   type MessageMutation,
   type MessagePageState,
@@ -2821,31 +2822,16 @@ export const SessionProvider: ParentComponent = (props) => {
     const id = currentSessionID()
     if (!id) return
     clearClose(id)
-    // Restore the reverted user message's prompt text into the input.
-    // Dispatch as a window message so PromptInput picks it up via onMessage.
-    const parts = store.parts[messageID]
-    if (parts) {
-      const text = parts
-        .filter((p) => p.type === "text" && !(p as { synthetic?: boolean }).synthetic)
-        .map((p) => (p as { text: string }).text ?? "")
-        .join("")
-      // Pass the original attachments' exact paths alongside the restored text
-      // so PromptInput can seed them directly rather than re-deriving mentions
-      // from the text via regex, which truncates at the first space in a
-      // filename (see PromptInput's setChatBoxMessage handler).
-      const paths = parts
-        .filter((p): p is Extract<Part, { type: "file" }> => p.type === "file")
-        .map((p) => p.source?.path)
-        .filter((p): p is string => !!p && !p.startsWith("session:"))
-      const sessions = parts
-        .filter((p): p is Extract<Part, { type: "file" }> => p.type === "file")
-        .filter((p) => p.url.startsWith("session:"))
-        .map((p) => ({
-          id: p.url.slice("session:".length),
-          title: p.source?.text?.value.replace(/^@/, "") ?? p.filename ?? p.url,
-          updated: 0,
-        }))
-      if (text) window.postMessage({ type: "setChatBoxMessage", text, paths, sessions }, "*")
+    // Restore the reverted user message's prompt text and attachments into the
+    // input. Dispatch as a window message so PromptInput picks it up via onMessage.
+    const state = revertPromptState(getParts(messageID))
+    const { text, paths, sessions, images } = state
+    // Paths carry the attachments' exact locations so PromptInput can seed them
+    // directly rather than re-deriving mentions from the text via regex, which
+    // truncates at the first space in a filename (see PromptInput's
+    // setChatBoxMessage handler).
+    if (text || paths.length > 0 || sessions.length > 0 || images.length > 0) {
+      window.postMessage({ type: "setChatBoxMessage", text, paths, sessions, images }, "*")
     }
     vscode.postMessage({ type: "revertSession", sessionID: id, messageID, partID })
   }
@@ -2854,7 +2840,7 @@ export const SessionProvider: ParentComponent = (props) => {
     const id = currentSessionID()
     if (!id) return
     // Clear the prompt input on full redo (matching TUI/desktop behavior)
-    window.postMessage({ type: "setChatBoxMessage", text: "" }, "*")
+    window.postMessage({ type: "setChatBoxMessage", text: "", images: [] }, "*")
     vscode.postMessage({ type: "unrevertSession", sessionID: id })
   }
 

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

@@ -309,6 +309,13 @@ export interface ActionMessage {
   action: string
 }
 
+/** Image attachment carried back into the prompt input when restoring a message. */
+export interface RestoredImage {
+  dataUrl: string
+  mime: string
+  filename?: string
+}
+
 export interface SetChatBoxMessage {
   type: "setChatBoxMessage"
   text: string
@@ -322,6 +329,12 @@ export interface SetChatBoxMessage {
   paths?: string[]
   /** Past chats referenced by the restored message, seeded the same way as paths. */
   sessions?: SessionSearchItem[]
+  /**
+   * Images attached to the restored message. Present means authoritative:
+   * PromptInput replaces its current attachments with this list (an empty
+   * array clears them); absent leaves current attachments untouched.
+   */
+  images?: RestoredImage[]
 }
 
 export interface AppendChatBoxMessage {