Merge pull request #11672 from Kilo-Org/fix-sandbox-toggle-prompt

fix(vscode): preserve draft when toggling sandbox

.changeset/keep-sandbox-toggle-drafts.md

@@ -0,0 +1,5 @@
+---
+"kilo-code": patch
+---
+
+Preserve new-chat prompts and attachments when toggling sandbox before sending the first message.

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

@@ -1,5 +1,11 @@
 import { describe, it, expect } from "bun:test"
-import { pendingDraftKey, scopeDraftKey, sessionDraftKey } from "../../webview-ui/src/utils/prompt-drafts"
+import {
+  createdDraftKey,
+  movePromptDraft,
+  pendingDraftKey,
+  scopeDraftKey,
+  sessionDraftKey,
+} from "../../webview-ui/src/utils/prompt-drafts"
 
 describe("sessionDraftKey", () => {
   it("prefixes session ids", () => {
@@ -30,3 +36,41 @@ describe("scopeDraftKey", () => {
     expect(scopeDraftKey("prompt:1")).toBe("prompt:1:empty")
   })
 })
+
+describe("createdDraftKey", () => {
+  it("uses the pending key when a draft id exists", () => {
+    expect(createdDraftKey("draft-1", true)).toBe("pending:draft-1")
+  })
+
+  it("uses the new-chat key for sandbox-triggered session creation", () => {
+    expect(createdDraftKey(undefined, true)).toBe("new")
+  })
+
+  it("ignores unrelated session creation without a draft id", () => {
+    expect(createdDraftKey()).toBeUndefined()
+  })
+})
+
+describe("movePromptDraft", () => {
+  it("moves text, review comments, and images to the created session", () => {
+    const source = scopeDraftKey("prompt:default", createdDraftKey(undefined, true))
+    const target = scopeDraftKey("prompt:default", sessionDraftKey("session-1"))
+    const comment = { id: "comment-1", body: "Keep this review note" }
+    const image = { id: "image-1", dataUrl: "data:image/png;base64,abc" }
+    const text = new Map([[source, "Keep this prompt"]])
+    const comments = new Map([[source, [comment]]])
+    const images = new Map([[source, [image]]])
+
+    expect(movePromptDraft({ text, comments, images }, source, target)).toEqual({
+      text: "Keep this prompt",
+      comments: [comment],
+      images: [image],
+    })
+    expect(text.get(target)).toBe("Keep this prompt")
+    expect(comments.get(target)).toEqual([comment])
+    expect(images.get(target)).toEqual([image])
+    expect(text.has(source)).toBe(false)
+    expect(comments.has(source)).toBe(false)
+    expect(images.has(source)).toBe(false)
+  })
+})

packages/kilo-vscode/tests/unit/prompt-input-connection-guard.test.ts

@@ -28,6 +28,7 @@ describe("PromptInput sandbox toggle", () => {
     expect(start).toBeGreaterThan(-1)
     expect(end).toBeGreaterThan(start)
     expect(toggle).toContain("const sessionID = sandboxID()")
+    expect(toggle).toContain("if (!sessionID) saveDraft(draftKey(), text(), reviewComments(), imageAttach.images())")
     expect(toggle).toContain('type: "toggleSandbox"')
     expect(toggle).toContain("sessionID,")
     expect(toggle).toContain("draftID: props.pendingSessionID ?? session.draftSessionID()")
@@ -36,15 +37,19 @@ describe("PromptInput sandbox toggle", () => {
     expect(toggle).not.toContain('type: "updateConfig"')
   })
 
-  it("keeps success feedback out of the webview toast region", () => {
-    const start = src.indexOf("const handleSandboxMessage =")
-    const end = src.indexOf("const unsubscribe =", start)
-    const handler = src.slice(start, end)
+  it("captures edits made while sandbox session creation is pending", () => {
+    const start = src.indexOf('if (message.type === "sessionCreated")')
+    const end = src.indexOf('if (message.type === "action"', start)
+    const created = src.slice(start, end)
+    const save = created.indexOf(
+      "if (source === draftKey()) saveDraft(source, text(), reviewComments(), imageAttach.images())",
+    )
+    const move = created.indexOf("movePromptDraft(")
 
     expect(start).toBeGreaterThan(-1)
     expect(end).toBeGreaterThan(start)
-    expect(handler).toContain('variant: "error"')
-    expect(handler).not.toContain('variant: "success"')
+    expect(save).toBeGreaterThan(-1)
+    expect(move).toBeGreaterThan(save)
   })
 
   it("uses the internal flag for visibility and effective runtime state for the button", () => {

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

@@ -48,7 +48,13 @@ import {
 } from "./prompt-input-utils"
 import type { ExtensionMessage, ReviewComment, SendMessageFailedMessage, TextPart } from "../../types/messages"
 import { formatReviewCommentsMarkdown } from "../../utils/review-comment-markdown"
-import { pendingDraftKey, scopeDraftKey, sessionDraftKey } from "../../utils/prompt-drafts"
+import {
+  createdDraftKey,
+  movePromptDraft,
+  pendingDraftKey,
+  scopeDraftKey,
+  sessionDraftKey,
+} from "../../utils/prompt-drafts"
 import { ReviewComments } from "./ReviewComments"
 import { partReview, reviewBody } from "../../../../src/shared/review-comments"
 import { isEnterKeyCommitNotIme } from "../../utils/ime-enter"
@@ -174,6 +180,7 @@ export const PromptInput: Component<PromptInputProps> = (props) => {
     const state = sandbox()
     if ((sessionID && !state) || state?.available === false || sandboxRequest() || !server.isConnected()) return
     const requestID = crypto.randomUUID()
+    if (!sessionID) saveDraft(draftKey(), text(), reviewComments(), imageAttach.images())
     setSandboxRequest(requestID)
     setSandboxTarget(sessionID ?? null)
     vscode.postMessage({
@@ -556,19 +563,19 @@ export const PromptInput: Component<PromptInputProps> = (props) => {
       restoreFailed(message as SendMessageFailedMessage)
     }
 
-    if (message.type === "sessionCreated" && message.draftID) {
-      const target = scopeDraftKey(boxKey(), pendingDraftKey(message.draftID) ?? "new")
-      const next = scopeDraftKey(boxKey(), sessionDraftKey(message.session.id) ?? "new")
-      const draft = drafts.get(target)
-      const pending = reviewDrafts.get(target)
-      const imgs = imageDrafts.get(target)
-      if (draft !== undefined) drafts.set(next, draft)
-      if (pending) reviewDrafts.set(next, pending)
-      if (imgs) imageDrafts.set(next, imgs)
-      drafts.delete(target)
-      reviewDrafts.delete(target)
-      imageDrafts.delete(target)
-      if (!session.currentSessionID() && (props.pendingSessionID ?? session.draftSessionID()) === message.draftID) {
+    if (message.type === "sessionCreated") {
+      const raw = createdDraftKey(message.draftID, sandboxRequest() !== undefined && sandboxTarget() === null)
+      if (raw) {
+        const source = scopeDraftKey(boxKey(), raw)
+        const target = scopeDraftKey(boxKey(), sessionDraftKey(message.session.id))
+        if (source === draftKey()) saveDraft(source, text(), reviewComments(), imageAttach.images())
+        movePromptDraft({ text: drafts, comments: reviewDrafts, images: imageDrafts }, source, target)
+      }
+      if (
+        message.draftID &&
+        !session.currentSessionID() &&
+        (props.pendingSessionID ?? session.draftSessionID()) === message.draftID
+      ) {
         session.setDraftSessionID(message.session.id)
       }
     }

packages/kilo-vscode/webview-ui/src/utils/prompt-drafts.ts

@@ -13,3 +13,26 @@ export function scopeDraftKey(box: string, raw?: string): string {
   if (!raw) return `${box}:empty`
   return `${box}:${raw}`
 }
+
+export function createdDraftKey(draftID?: string, sandbox = false): string | undefined {
+  return pendingDraftKey(draftID) ?? (sandbox ? "new" : undefined)
+}
+
+export function movePromptDraft<T, C, I>(
+  stores: { text: Map<string, T>; comments: Map<string, C>; images: Map<string, I> },
+  source: string,
+  target: string,
+): { text?: T; comments?: C; images?: I } {
+  const draft = {
+    text: stores.text.get(source),
+    comments: stores.comments.get(source),
+    images: stores.images.get(source),
+  }
+  if (draft.text !== undefined) stores.text.set(target, draft.text)
+  if (draft.comments !== undefined) stores.comments.set(target, draft.comments)
+  if (draft.images !== undefined) stores.images.set(target, draft.images)
+  stores.text.delete(source)
+  stores.comments.delete(source)
+  stores.images.delete(source)
+  return draft
+}