From 062d11f996704007c0cc5b098eb8059640865ab3 Mon Sep 17 00:00:00 2001 From: marius-kilocode Date: Mon, 24 Aug 2026 15:19:37 +0200 Subject: [PATCH] fix(vscode): restore images when undoing prompts --- .changeset/restore-undo-attachments.md | 5 ++ .../tests/unit/session-utils.test.ts | 50 +++++++++++++++++++ .../src/components/chat/PromptInput.tsx | 12 +++++ .../webview-ui/src/context/session-utils.ts | 34 +++++++++++++ .../webview-ui/src/context/session.tsx | 38 +++++--------- .../src/types/messages/extension-messages.ts | 13 +++++ 6 files changed, 126 insertions(+), 26 deletions(-) create mode 100644 .changeset/restore-undo-attachments.md diff --git a/.changeset/restore-undo-attachments.md b/.changeset/restore-undo-attachments.md new file mode 100644 index 0000000000..e50268ff3f --- /dev/null +++ b/.changeset/restore-undo-attachments.md @@ -0,0 +1,5 @@ +--- +"kilo-code": patch +--- + +Restore attached images when undoing the last prompt. diff --git a/packages/kilo-vscode/tests/unit/session-utils.test.ts b/packages/kilo-vscode/tests/unit/session-utils.test.ts index 4343037949..a401bb2ab1 100644 --- a/packages/kilo-vscode/tests/unit/session-utils.test.ts +++ b/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>) => + ({ type: "file", id: "f", mime: "text/plain", url: "", ...overrides }) as Extract + + 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: [] }) + }) +}) diff --git a/packages/kilo-vscode/webview-ui/src/components/chat/PromptInput.tsx b/packages/kilo-vscode/webview-ui/src/components/chat/PromptInput.tsx index 9b1983fae3..0958cd3eb8 100644 --- a/packages/kilo-vscode/webview-ui/src/components/chat/PromptInput.tsx +++ b/packages/kilo-vscode/webview-ui/src/components/chat/PromptInput.tsx @@ -724,6 +724,18 @@ export const PromptInput: Component = (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") { diff --git a/packages/kilo-vscode/webview-ui/src/context/session-utils.ts b/packages/kilo-vscode/webview-ui/src/context/session-utils.ts index ec3ef9de42..b783496f7a 100644 --- a/packages/kilo-vscode/webview-ui/src/context/session-utils.ts +++ b/packages/kilo-vscode/webview-ui/src/context/session-utils.ts @@ -34,6 +34,40 @@ export function messageParts(messages: Message[]): Record { 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 => 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 diff --git a/packages/kilo-vscode/webview-ui/src/context/session.tsx b/packages/kilo-vscode/webview-ui/src/context/session.tsx index bf6019f83c..6bf6779bc8 100644 --- a/packages/kilo-vscode/webview-ui/src/context/session.tsx +++ b/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 => 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 => 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 }) } diff --git a/packages/kilo-vscode/webview-ui/src/types/messages/extension-messages.ts b/packages/kilo-vscode/webview-ui/src/types/messages/extension-messages.ts index 34252b71c2..89b52abe1a 100644 --- a/packages/kilo-vscode/webview-ui/src/types/messages/extension-messages.ts +++ b/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 {