feat(vscode): edit queued messages in the prompt input

.changeset/edit-queued-messages.md

@@ -0,0 +1,6 @@
+---
+"kilo-code": minor
+"@kilocode/cli": patch
+---
+
+Edit queued messages in VS Code before they are sent, while preserving their text and attachments.

packages/kilo-ui/src/components/message-part.tsx

@@ -752,6 +752,8 @@ export function UserMessageDisplay(props: {
   text?: string
   copyText?: string
   header?: JSX.Element
+  edit?: { label: string; onClick: () => void; disabled?: boolean }
+  queuedDisabled?: boolean
   onDelete?: () => void
   onFork?: () => void
   onRevert?: () => void
@@ -835,6 +837,7 @@ export function UserMessageDisplay(props: {
           icon="close-small"
           size="normal"
           variant="ghost"
+          disabled={props.queuedDisabled}
           onMouseDown={(e) => e.preventDefault()}
           onClick={(event) => {
             event.stopPropagation()
@@ -846,6 +849,28 @@ export function UserMessageDisplay(props: {
     </Show>
   )
 
+  const Edit = () => (
+    <Show when={props.edit}>
+      {(edit) => (
+        <Tooltip value={edit().label} placement="right" gutter={4}>
+          <IconButton
+            data-slot="user-message-edit"
+            icon="edit"
+            size="small"
+            variant="ghost"
+            disabled={props.queuedDisabled || edit().disabled}
+            onMouseDown={(event) => event.preventDefault()}
+            onClick={(event) => {
+              event.stopPropagation()
+              edit().onClick()
+            }}
+            aria-label={edit().label}
+          />
+        </Tooltip>
+      )}
+    </Show>
+  )
+
   return (
     <GrowBox animate={!!props.animate} fade class="w-full min-w-0 self-stretch max-w-full">
       <div data-component="user-message" data-interrupted={props.interrupted ? "" : undefined}>
@@ -885,6 +910,7 @@ export function UserMessageDisplay(props: {
         <Show when={!text() && !props.header && props.queued}>
           <div data-slot="user-message-queued-indicator">
             <TextShimmer text={i18n.t("ui.message.queued")} />
+            <Edit />
             <Delete />
           </div>
         </Show>
@@ -900,6 +926,7 @@ export function UserMessageDisplay(props: {
               <GrowBox animate={!!props.animate} open={!!props.queued}>
                 <div data-slot="user-message-queued-indicator">
                   <TextShimmer text={i18n.t("ui.message.queued")} />
+                  <Edit />
                   <Delete />
                 </div>
               </GrowBox>

packages/kilo-vscode/src/KiloProvider.ts

@@ -1163,7 +1163,7 @@ export class KiloProvider implements vscode.WebviewViewProvider, TelemetryProper
           this.checkpoint(message.sessionID, () => this.handleUnrevertSession(message.sessionID))
           break
         case "deleteMessage":
-          await this.handleDeleteMessage(message.sessionID, message.messageID)
+          await this.handleDeleteMessage(message.sessionID, message.messageID, message.requestID)
           break
         case "permissionResponse":
           await handlePermissionResponse(
@@ -2425,24 +2425,34 @@ export class KiloProvider implements vscode.WebviewViewProvider, TelemetryProper
     }
   }
 
-  private async handleDeleteMessage(sessionID: string, messageID: string): Promise<void> {
-    if (!this.client) {
+  private async handleDeleteMessage(sessionID: string, messageID: string, requestID?: string): Promise<void> {
+    const result = {
+      type: "deleteMessageResult" as const,
+      sessionID,
+      messageID,
+      ...(requestID !== undefined ? { requestID } : {}),
+    }
+    const client = this.client
+    if (!client) {
       this.postMessage({ type: "error", message: "Not connected to CLI backend", sessionID })
+      this.postMessage({ ...result, success: false })
       return
     }
 
     try {
-      await this.client.session.deleteMessage(
-        { sessionID, messageID, directory: this.getWorkspaceDirectory(sessionID) },
+      const response = await client.session.deleteMessage(
+        { sessionID, messageID, directory: this.getWorkspaceDirectory(sessionID), queued: true },
         { throwOnError: true },
       )
+      this.postMessage({ ...result, success: response.data === true })
     } catch (error) {
       console.error("[Kilo New] KiloProvider: Failed to delete message:", error)
       this.postMessage({
         type: "error",
         message: getErrorMessage(error) || "Failed to delete message",
         sessionID,
       })
+      this.postMessage({ ...result, success: false })
     }
   }
 

packages/kilo-vscode/tests/unit/kilo-provider-load-messages.test.ts

@@ -70,6 +70,8 @@ function createClient(options?: {
   createDeferred?: Deferred<{ data: ReturnType<typeof mkCreatedSession> }>
   abortFailures?: string[]
   abortDeferred?: Deferred<void>
+  deleteResult?: unknown
+  deleteError?: boolean
   supportDeferred?: Deferred<{ data: { available: boolean; reason?: string } }>
   sandboxDeferred?: Deferred<{ data: unknown }>
   sandboxStarted?: Deferred<void>
@@ -79,6 +81,12 @@ function createClient(options?: {
   const stopped: { sessionID: string; directory?: string }[] = []
   const aborted: { sessionID: string; directory?: string }[] = []
   const deleted: { sessionID: string; directory?: string }[] = []
+  const deletedMessages: Array<{
+    sessionID: string
+    messageID: string
+    directory?: string
+    queued?: boolean
+  }> = []
   const prompted: Array<Record<string, unknown>> = []
   const reverted: Array<Record<string, unknown>> = []
   const created: Array<Record<string, unknown>> = []
@@ -90,6 +98,7 @@ function createClient(options?: {
     stopped,
     aborted,
     deleted,
+    deletedMessages,
     prompted,
     reverted,
     created,
@@ -133,6 +142,11 @@ function createClient(options?: {
         if (options?.deleteDeferred) return options.deleteDeferred.promise
         return { data: {} }
       },
+      deleteMessage: async (params: { sessionID: string; messageID: string; directory?: string; queued?: boolean }) => {
+        deletedMessages.push(params)
+        if (options?.deleteError) throw new Error("delete failed")
+        return { data: options?.deleteResult }
+      },
     },
     sandbox: {
       support: async (params: Record<string, unknown>) => {
@@ -171,7 +185,7 @@ function createClient(options?: {
   }
 }
 
-function createConnection(client: ReturnType<typeof createClient>) {
+function createConnection(client: ReturnType<typeof createClient> | null) {
   const state = { value: undefined as boolean | undefined, revision: 0, pending: Promise.resolve() }
   return {
     sandboxPreference: {
@@ -251,13 +265,14 @@ type ProviderInternals = {
   refreshGitStatus: (directory?: string, sessionID?: string) => Promise<void>
   handleLoadMessages: (sid: string, opts?: { mode?: string; before?: string; limit?: number }) => Promise<void>
   handleDeleteSession: (sid: string) => Promise<void>
+  handleDeleteMessage: (sid: string, mid: string, rid?: string) => Promise<void>
 }
 
-function makeProvider(client: ReturnType<typeof createClient>) {
+function makeProvider(client: ReturnType<typeof createClient> | null) {
   const connection = createConnection(client)
   const provider = new KiloProvider({} as never, connection as never)
   const internal = provider as unknown as ProviderInternals
-  internal.connectionState = "connected"
+  internal.connectionState = client ? "connected" : "disconnected"
   const sent: unknown[] = []
   internal.webview = {
     postMessage: async (message: unknown) => {
@@ -1159,6 +1174,31 @@ describe("KiloProvider.handleDeleteSession / background processes", () => {
   })
 })
 
+describe("KiloProvider.handleDeleteMessage", () => {
+  const ids = { sessionID: "s1", messageID: "m1", requestID: "r1" }
+
+  it.each([true, false, undefined, "true"])("confirms only a true queued deletion result: %p", async (result) => {
+    const client = createClient({ deleteResult: result })
+    const { internal, sent } = makeProvider(client)
+    await internal.handleDeleteMessage(ids.sessionID, ids.messageID, ids.requestID)
+    expect(client.deletedMessages).toEqual([{ sessionID: "s1", messageID: "m1", directory: "/repo", queued: true }])
+    expect(sent).toContainEqual({ type: "deleteMessageResult", ...ids, success: result === true })
+  })
+
+  it.each([true, false])("confirms removal failures when connected=%p", async (connected) => {
+    const error = spyOn(console, "error").mockImplementation(() => {})
+    const { internal, sent } = makeProvider(connected ? createClient({ deleteError: true }) : null)
+    await internal.handleDeleteMessage(ids.sessionID, ids.messageID, ids.requestID)
+    expect(sent).toContainEqual({
+      type: "error",
+      message: connected ? "delete failed" : "Not connected to CLI backend",
+      sessionID: ids.sessionID,
+    })
+    expect(sent).toContainEqual({ type: "deleteMessageResult", ...ids, success: false })
+    error.mockRestore()
+  })
+})
+
 describe("KiloProvider.handleLoadMessages / slim payload", () => {
   it("shows a cost alert even when cost arrives after the session is idle", () => {
     const client = createClient()

packages/kilo-vscode/tests/unit/prompt-drafts.test.ts

@@ -1,6 +1,12 @@
 import { beforeEach, describe, it, expect } from "bun:test"
 import { createEffect, createRoot, createSignal, on } from "solid-js"
-import { deleteDraftsForSession, drafts, imageDrafts, reviewDrafts } from "../../webview-ui/src/utils/draft-store"
+import {
+  deleteDraftsForSession,
+  drafts,
+  imageDrafts,
+  mentionDrafts,
+  reviewDrafts,
+} from "../../webview-ui/src/utils/draft-store"
 import {
   createdDraftKey,
   movePromptDraft,
@@ -13,6 +19,7 @@ beforeEach(() => {
   drafts.clear()
   reviewDrafts.clear()
   imageDrafts.clear()
+  mentionDrafts.clear()
 })
 
 describe("deleteDraftsForSession", () => {
@@ -22,6 +29,7 @@ describe("deleteDraftsForSession", () => {
     drafts.set("prompt:default:session:b", "draft b")
     reviewDrafts.set("prompt:default:session:a", [])
     imageDrafts.set("prompt:default:session:a", [])
+    mentionDrafts.set("prompt:default:session:a", { paths: ["file with spaces.ts"], sessions: [] })
 
     deleteDraftsForSession("a")
 
@@ -30,6 +38,7 @@ describe("deleteDraftsForSession", () => {
     expect(drafts.get("prompt:default:session:b")).toBe("draft b")
     expect(reviewDrafts.has("prompt:default:session:a")).toBe(false)
     expect(imageDrafts.has("prompt:default:session:a")).toBe(false)
+    expect(mentionDrafts.has("prompt:default:session:a")).toBe(false)
   })
 
   it("is a no-op when given an empty id", () => {

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

@@ -4,12 +4,19 @@ import {
   messageTurns,
   partitionTurns,
   queuedUserMessageIDs,
+  removeQueuedMessage,
   stableMessageTurns,
   visibleMessages,
   visibleParts,
   type RevertBoundary,
 } from "../../webview-ui/src/context/session-queue"
-import type { Message, Part, SessionStatusInfo } from "../../webview-ui/src/types/messages"
+import type {
+  ExtensionMessage,
+  Message,
+  Part,
+  SessionStatusInfo,
+  WebviewMessage,
+} from "../../webview-ui/src/types/messages"
 
 const base = {
   sessionID: "session",
@@ -634,3 +641,57 @@ describe("activeUserMessageID", () => {
     expect(activeUserMessageID(messages, { type: "busy" })).toBe("message_3")
   })
 })
+
+describe("queued deletion cleanup", () => {
+  const setup = (timeout = 10_000) => {
+    const listeners = new Set<(message: ExtensionMessage) => void>()
+    const sent: WebviewMessage[] = []
+    const promise = removeQueuedMessage(
+      {
+        onMessage: (handler) => {
+          listeners.add(handler)
+          return () => {
+            listeners.delete(handler)
+          }
+        },
+        postMessage: (message) => {
+          sent.push(message)
+        },
+      },
+      "session",
+      "message",
+      timeout,
+    )
+    return {
+      promise,
+      listeners,
+      sent,
+      emit: (message: ExtensionMessage) => listeners.forEach((handler) => handler(message)),
+    }
+  }
+
+  it.each([true, false])("settles only the matching acknowledgment: %p", async (success) => {
+    const state = setup()
+    const request = state.sent.at(0)
+    if (request?.type !== "deleteMessage") throw new Error("Missing deletion request")
+    state.emit({ type: "connectionState", state: "connecting" })
+    state.emit({ ...request, type: "deleteMessageResult", requestID: "unrelated", success: true })
+    expect(state.listeners.size).toBe(1)
+    state.emit({ ...request, type: "deleteMessageResult", success })
+    expect(await state.promise).toBe(success)
+    expect(state.listeners.size).toBe(0)
+  })
+
+  const interruptions: Array<ExtensionMessage | undefined> = [
+    undefined,
+    { type: "connectionState", state: "disconnected" },
+    { type: "connectionState", state: "error" },
+    { type: "sessionDeleted", sessionID: "session" },
+  ]
+  it.each(interruptions)("releases the waiter and listener after timeout or interruption: %p", async (message) => {
+    const state = setup(1)
+    if (message) state.emit(message)
+    expect(await state.promise).toBe(false)
+    expect(state.listeners.size).toBe(0)
+  })
+})

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

@@ -68,7 +68,12 @@ export const ChatView: Component<ChatViewProps> = (props) => {
   // the very first send instead of growing once the message lands.
   const hasMessages = () => session.messages().length > 0 || session.submitting()
 
-  // "Continue in Worktree" state
+  const [editable, setEditable] = createSignal(false)
+  const [editing, setEditing] = createSignal<{ sessionID: string; messageID: string }>()
+  const edit = (sessionID: string, messageID: string) => {
+    if (props.readonly || !editable() || editing()) return
+    setEditing({ sessionID, messageID })
+  }
   const [transferring, setTransferring] = createSignal(false)
   const [transferDetail, setTransferDetail] = createSignal("")
   const [repoBranch, setRepoBranch] = createSignal<string>()
@@ -356,6 +361,9 @@ export const ChatView: Component<ChatViewProps> = (props) => {
               onSelectSession={props.onSelectSession}
               onShowHistory={props.onShowHistory}
               onForkMessage={props.onForkMessage}
+              onEditMessage={edit}
+              queuedDisabled={editing()?.sessionID === id() && !!editing()}
+              editDisabled={!editable() || !!editing()}
               questions={standaloneQuestions}
               suggestions={standaloneSuggestions}
               readonly={props.readonly}
@@ -388,6 +396,9 @@ export const ChatView: Component<ChatViewProps> = (props) => {
             <Show when={!props.readonly}>
               <PromptInput
                 blocked={blocked}
+                edit={editing()}
+                onEditComplete={() => setEditing(undefined)}
+                onEditReady={setEditable}
                 suggesting={suggesting}
                 questioning={questioning}
                 worktree={props.worktree}