Merge pull request #11442 from IamCoder18/vscode/cleanup-drafts-on-session-delete

fix(vscode): clear unsent drafts on session deletion

.changeset/cleanup-drafts-on-session-delete.md

@@ -0,0 +1,7 @@
+---
+"kilo-code": patch
+---
+
+Free webview memory for deleted VS Code sessions by clearing unsent prompt text, review comments, and pending image attachments that were retained in the per-session draft cache after `sessionDeleted`.
+
+Also restores an in-flight failed draft into the live prompt after a session is deleted mid-send (whether user-initiated or via external CLI/TUI/cascade delete), while never rehydrating it into a prompt the user explicitly cleared.

packages/kilo-vscode/src/KiloProvider.ts

@@ -1505,6 +1505,16 @@ export class KiloProvider implements vscode.WebviewViewProvider, TelemetryProper
           // for the busy-session warning on Save.
           if (event.type === "session.status") return true
 
+          // session.deleted must always pass through so the webview can run its cleanup
+          // (messages, parts, stash, todos, permissions, drafts, etc.) — including for
+          // sessions that were never explicitly tracked here (e.g. child sessions
+          // cascade-deleted with the parent, or external CLI deletions). We deliberately
+          // do NOT re-track the deleted id: handleLoadMessages intentionally drops late
+          // responses for sessions that have been pruned, and re-tracking would let an
+          // in-flight messagesLoaded response resurrect transcript state for a session
+          // the webview just cleaned up.
+          if (event.type === "session.deleted") return true
+
           return this.trackedSessionIds.has(sessionId)
         },
         (payload, directory) => {
@@ -1949,6 +1959,43 @@ export class KiloProvider implements vscode.WebviewViewProvider, TelemetryProper
     }
   }
 
+  /**
+   * Drops every per-session cache entry we hold for the given id. Shared between
+   * the user-initiated delete path (handleDeleteSession, after the backend
+   * confirms) and the SSE session.deleted path (cascaded child deletes and
+   * external CLI/TUI deletes that arrive via the event stream), so both paths
+   * leave trackedSessionIds, sessionDirectories, and the related Maps in the
+   * same state — including currentSession / contextSessionID / focused-session
+   * registration. Without clearing those three, resolveSession() would still
+   * see the deleted id via this.currentSession and the next send would target
+   * a session the backend has already deleted.
+   */
+  private pruneDeletedSession(sessionID: string): void {
+    this.trackedSessionIds.delete(sessionID)
+    this.streams.drop(sessionID)
+    this.visibleTaskStreams.delete(sessionID)
+    this.syncedChildSessions.delete(sessionID)
+    this.sessionDirectories.delete(sessionID)
+    this.aborts.delete(sessionID)
+    this.lastReconciledAt.delete(sessionID)
+    this.checkpoints.delete(sessionID)
+    this.revisions.delete(sessionID)
+    this.refreshes.delete(sessionID)
+    this.sessionStatusMap.delete(sessionID)
+    this.costs.onSessionDeleted(sessionID)
+    const deletedAlertLimit = this.activeAlerts.get(sessionID)
+    if (deletedAlertLimit !== undefined) {
+      this.activeAlerts.delete(sessionID)
+      this.postMessage({ type: "sessionCostAlertResolved", sessionID: sessionID, limit: deletedAlertLimit })
+    }
+    this.connectionService.pruneSession(sessionID)
+    if (this.currentSession?.id === sessionID) {
+      this.contextSessionID = undefined
+      this.setCurrentSession(null)
+    }
+    if (this.streams.focused === sessionID) this.focusSession(undefined)
+  }
+
   /**
    * Handle deleting a session.
    */
@@ -1965,23 +2012,7 @@ export class KiloProvider implements vscode.WebviewViewProvider, TelemetryProper
       )
       await stopSessionProcesses(this.client, sessionID, workspaceDir)
       await this.client.session.delete({ sessionID, directory: workspaceDir }, { throwOnError: true })
-      this.trackedSessionIds.delete(sessionID)
-      this.streams.drop(sessionID)
-      this.visibleTaskStreams.delete(sessionID)
-      this.syncedChildSessions.delete(sessionID)
-      this.costs.onSessionDeleted(sessionID)
-      const deletedAlertLimit = this.activeAlerts.get(sessionID)
-      if (deletedAlertLimit !== undefined) {
-        this.activeAlerts.delete(sessionID)
-        this.postMessage({ type: "sessionCostAlertResolved", sessionID: sessionID, limit: deletedAlertLimit })
-      }
-      this.sessionDirectories.delete(sessionID)
-      this.aborts.delete(sessionID)
-      this.lastReconciledAt.delete(sessionID)
-      this.checkpoints.delete(sessionID)
-      this.revisions.delete(sessionID)
-      this.refreshes.delete(sessionID)
-      this.connectionService.pruneSession(sessionID)
+      this.pruneDeletedSession(sessionID)
       if (this.currentSession?.id === sessionID) {
         this.contextSessionID = undefined
         this.setCurrentSession(null)
@@ -3727,9 +3758,13 @@ export class KiloProvider implements vscode.WebviewViewProvider, TelemetryProper
     // message.part.* events are always session-scoped; drop if session unknown.
     if (!sessionID && isSessionScopedPartEvent(event.type)) return
     if (this.postModelUsageChanged(event, sessionID)) return
-    if (event.type !== "indexing.status" && sessionID && !this.trackedSessionIds.has(sessionID)) {
+    if (
+      event.type !== "indexing.status" &&
+      event.type !== "session.deleted" &&
+      sessionID &&
+      !this.trackedSessionIds.has(sessionID)
+    )
       return
-    }
 
     if (event.type === "session.updated" && typeof event.properties.info.cost === "number") {
       const cost = this.costs.setSessionCost(event.properties.sessionID, event.properties.info.cost)
@@ -3818,6 +3853,15 @@ export class KiloProvider implements vscode.WebviewViewProvider, TelemetryProper
       }
     }
 
+    // Drop the per-session caches for deleted sessions so a late
+    // handleLoadMessages response (or any other guarded read) can't resurrect
+    // transcript state for a session the webview just cleaned up. The
+    // prefilter lets session.deleted through without re-tracking, and the
+    // handleEvent guard does the same — this is the matching prune.
+    if (event.type === "session.deleted" && sessionID) {
+      this.pruneDeletedSession(sessionID)
+    }
+
     if (!isLegacySyncEvent(event)) {
       const props = event.properties
       handleNetworkEvent(

packages/kilo-vscode/src/kilo-provider/session-stream-scheduler.ts

@@ -157,6 +157,11 @@ export class SessionStreamScheduler {
     if (sessionID) this.flush(sessionID)
   }
 
+  /** Currently focused (active-lane) session ID, if any. */
+  get focused(): string | undefined {
+    return this.active
+  }
+
   setVisible(sessionID: string, visible: boolean): void {
     const changed = visible ? !this.visible.has(sessionID) : this.visible.has(sessionID)
     if (!changed) return

packages/kilo-vscode/src/services/cli-backend/connection-service.ts

@@ -257,11 +257,25 @@ export class KiloConnectionService {
    * Remove all messageID → sessionID entries for a given session.
    * Called when a session is deleted or otherwise pruned so the map
    * does not grow unbounded over the extension lifetime.
+   *
+   * Also drops the session from any provider's focused or opened set
+   * so the server's `viewed` notification stops advertising a deleted
+   * id after external (CLI/TUI/cascade) deletes arrive via SSE.
    */
   pruneSession(sessionId: string): void {
     for (const [mid, sid] of this.messageSessionIdsByMessageId) {
       if (sid === sessionId) this.messageSessionIdsByMessageId.delete(mid)
     }
+    for (const [key, sid] of this.focused) {
+      if (sid === sessionId) this.focused.delete(key)
+    }
+    for (const [key, ids] of this.opened) {
+      if (!ids.includes(sessionId)) continue
+      const next = ids.filter((id) => id !== sessionId)
+      if (next.length === 0) this.opened.delete(key)
+      else this.opened.set(key, next)
+    }
+    this.flushViewed()
   }
 
   /**

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

@@ -1,4 +1,6 @@
-import { describe, it, expect } from "bun:test"
+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 {
   createdDraftKey,
   movePromptDraft,
@@ -7,6 +9,102 @@ import {
   sessionDraftKey,
 } from "../../webview-ui/src/utils/prompt-drafts"
 
+beforeEach(() => {
+  drafts.clear()
+  reviewDrafts.clear()
+  imageDrafts.clear()
+})
+
+describe("deleteDraftsForSession", () => {
+  it("clears deleted-session drafts without touching other sessions", () => {
+    drafts.set("prompt:default:session:a", "draft a")
+    drafts.set("prompt:default:pending:a", "pending a")
+    drafts.set("prompt:default:session:b", "draft b")
+    reviewDrafts.set("prompt:default:session:a", [])
+    imageDrafts.set("prompt:default:session:a", [])
+
+    deleteDraftsForSession("a")
+
+    expect(drafts.has("prompt:default:session:a")).toBe(false)
+    expect(drafts.has("prompt:default:pending:a")).toBe(false)
+    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)
+  })
+
+  it("is a no-op when given an empty id", () => {
+    drafts.set("prompt:default:session:a", "draft a")
+    deleteDraftsForSession("")
+    expect(drafts.get("prompt:default:session:a")).toBe("draft a")
+  })
+
+  it("clears drafts that PromptInput's draftKey effect recreates after the batch", () => {
+    // Production race that motivated the post-batch deleteDraftsForSession call:
+    //   1. handleSessionDeleted batches setCurrentSessionID(undefined) +
+    //      setDraftSessionID(undefined). PromptInput's draftKey memo transitions from
+    //      ":session:<id>" to the "new" bucket.
+    //   2. PromptInput's createEffect(on(draftKey, ...)) runs after the batch and calls
+    //      saveDraft(prev, currentText, currentImages), writing the live prompt and any
+    //      attached image data URLs back into the just-cleared ":session:<id>" key.
+    //   3. deleteDraftsForSession runs after the effect and clears the re-added entry.
+    //
+    // The test wires the same reactive plumbing — real Solid createSignal/createEffect/on
+    // against the same scopeDraftKey/sessionDraftKey/pendingDraftKey helpers PromptInput
+    // uses — so a regression that moves the cleanup back inside the batch (or drops it
+    // entirely) leaks the recreated draft and the final assertion fails.
+    const img = {
+      id: "i1",
+      filename: "x.png",
+      mime: "image/png",
+      dataUrl: "data:image/png;base64,AAAA",
+    }
+    const draftKey = "prompt:default:session:race"
+
+    createRoot((dispose) => {
+      // Live prompt state, the way PromptInput tracks it.
+      const [text, setText] = createSignal("draft a")
+      const [images] = createSignal([img])
+      const [currentSessionID, setCurrentSessionID] = createSignal<string | undefined>("race")
+      const [draftSessionID, setDraftSessionID] = createSignal<string | undefined>("race")
+
+      const boxKey = "prompt:default"
+      const rawKey = () =>
+        sessionDraftKey(currentSessionID()) ?? pendingDraftKey(draftSessionID() ?? undefined) ?? "new"
+      const key = () => scopeDraftKey(boxKey, rawKey())
+
+      // Pre-deletion: the user has unsent text and an attached image for this session.
+      drafts.set(draftKey, text())
+      imageDrafts.set(draftKey, images())
+
+      // Mirror the saveDraft behavior PromptInput's effect runs when draftKey transitions.
+      createEffect(
+        on(key, (k, prev) => {
+          if (prev !== undefined && prev !== k) {
+            drafts.set(prev, text())
+            imageDrafts.set(prev, images())
+          }
+        }),
+      )
+
+      // Production batch: clear the ids so draftKey transitions off ":session:<id>".
+      setCurrentSessionID(undefined)
+      setDraftSessionID(undefined)
+      // Solid has now run the effect; the recreate happened. Sanity-check before cleanup.
+      expect(drafts.has(draftKey)).toBe(true)
+      expect(imageDrafts.has(draftKey)).toBe(true)
+
+      // The post-batch cleanup. A single in-batch call (run before the effect) would
+      // have been wiped by the recreate above and not catch this — the post-batch
+      // call is what actually frees the entry.
+      deleteDraftsForSession("race")
+      dispose()
+    })
+
+    expect(drafts.has(draftKey)).toBe(false)
+    expect(imageDrafts.has(draftKey)).toBe(false)
+  })
+})
+
 describe("sessionDraftKey", () => {
   it("prefixes session ids", () => {
     expect(sessionDraftKey("abc")).toBe("session:abc")

packages/kilo-vscode/tests/unit/prompt-send-contract.test.ts

@@ -13,11 +13,14 @@
 import { describe, it, expect } from "bun:test"
 import fs from "node:fs"
 import path from "node:path"
+import { clearIfOn } from "../../webview-ui/src/context/session-cloud-prune"
 
 const ROOT = path.resolve(import.meta.dir, "../..")
 const SESSION_FILE = path.join(ROOT, "webview-ui/src/context/session.tsx")
 const CHATVIEW_FILE = path.join(ROOT, "webview-ui/src/components/chat/ChatView.tsx")
 const PROMPT_UTILS_FILE = path.join(ROOT, "webview-ui/src/components/chat/prompt-input-utils.ts")
+const KILOPROVIDER_FILE = path.join(ROOT, "src/KiloProvider.ts")
+const CONNECTION_SERVICE_FILE = path.join(ROOT, "src/services/cli-backend/connection-service.ts")
 
 function readFile(filePath: string): string {
   return fs.readFileSync(filePath, "utf-8")
@@ -116,3 +119,380 @@ describe("isPromptBlocked signature contract", () => {
     expect(params).toHaveLength(1)
   })
 })
+
+describe("handleSessionDeleted draft cleanup contract", () => {
+  const source = readFile(SESSION_FILE)
+
+  it("clears draftSessionID independently of currentSessionID when it equals the deleted id", () => {
+    const body = extractFunctionBody(source, "handleSessionDeleted")
+    const draftBlock = body.match(/if \(draftSessionID\(\) === sessionID\) \{([\s\S]*?)\}/)
+    expect(draftBlock).not.toBeNull()
+    expect(draftBlock![1]).toContain("setDraftSessionID(undefined)")
+    // Must be a sibling check, not nested inside the currentSessionID branch —
+    // otherwise a deleted but non-active session leaves draftSessionID stale.
+    const activeBlock = body.match(/if \(currentSessionID\(\) === sessionID\) \{([\s\S]*?)\}/)
+    expect(activeBlock![1]).not.toContain("setDraftSessionID")
+  })
+
+  it("calls deleteDraftsForSession outside the cleanup batch so PromptInput's recreate is also cleaned up", () => {
+    const body = extractFunctionBody(source, "handleSessionDeleted")
+    const batchMatch = body.match(/batch\(\(\) => \{([\s\S]*?)\}\)/)
+    expect(batchMatch).not.toBeNull()
+    expect(batchMatch![1]).not.toContain("deleteDraftsForSession(sessionID)")
+    const postBatch = body.slice((batchMatch!.index ?? 0) + batchMatch![0].length)
+    expect(postBatch).toContain("deleteDraftsForSession(sessionID)")
+  })
+
+  it("removes the deleted id from the loaded Set so cascade/external deletes free the marker", () => {
+    // The user-initiated deleteSession() path prunes loaded optimistically, but
+    // cascade deletes and external CLI/TUI deletes only come through
+    // handleSessionDeleted. Without this, those ids stay in loaded until reload.
+    const body = extractFunctionBody(source, "handleSessionDeleted")
+    expect(body).toMatch(
+      /setLoaded\(\s*\(prev\)\s*=>\s*\{[\s\S]*?prev\.has\(sessionID\)[\s\S]*?next\.delete\(sessionID\)[\s\S]*?\}\)/,
+    )
+  })
+
+  it("drops respondingPermissions entries that belong to the deleted session", () => {
+    // setPermissions is cleared by removeSessionPermissions, but respondingPermissions
+    // (the Set of in-flight permission ids) is a separate accessor that doesn't know
+    // which ids belong to which session. Without an explicit prune here, a permission
+    // request that the user was responding to when the session was deleted would keep
+    // its id resident and block future requests with the same id.
+    const body = extractFunctionBody(source, "handleSessionDeleted")
+    expect(body).toContain("setRespondingPermissions")
+  })
+})
+
+describe("KiloProvider pruneDeletedSession contract", () => {
+  const source = readFile(KILOPROVIDER_FILE)
+
+  it("drops sessionStatusMap entries alongside the other per-session caches", () => {
+    // sessionStatusMap is the source of truth for the destructive-config busy-session
+    // warning (sessionStatusMap.size === 0 short-circuit, the allStatusMap fed to the
+    // Settings panel). Without this prune, deleted sessions stay marked as
+    // busy/retry/etc. until provider dispose, suppressing the "you have a busy session"
+    // warning for the new current session.
+    const match = source.match(/pruneDeletedSession\(sessionID: string\): void \{([\s\S]*?)\n  \}/)
+    expect(match).not.toBeNull()
+    expect(match![1]).toContain("this.sessionStatusMap.delete(sessionID)")
+  })
+
+  it("clears currentSession and contextSessionID when the deleted id matches", () => {
+    // The SSE session.deleted path runs pruneDeletedSession; if it leaves
+    // currentSession pointing at the deleted session, resolveSession() in the
+    // next sendMessage falls back to currentSession.id and targets a session
+    // the backend has already deleted. The user-initiated delete path
+    // (handleDeleteSession) does this clearing after the prune; pruneDeletedSession
+    // itself must do the same so the SSE path is symmetric.
+    const match = source.match(/pruneDeletedSession\(sessionID: string\): void \{([\s\S]*?)\n  \}/)
+    expect(match).not.toBeNull()
+    expect(match![1]).toMatch(
+      /if \(this\.currentSession\?\.id === sessionID\)\s*\{[\s\S]*?this\.contextSessionID = undefined[\s\S]*?this\.setCurrentSession\(null\)/,
+    )
+  })
+
+  it("unfocuses the streams when the deleted id matches the focused session", () => {
+    // Without this, connectionService.focused still reports the deleted id to
+    // the backend (viewed.focused), and focusSession() never calls
+    // unregisterFocused for this instance.
+    const match = source.match(/pruneDeletedSession\(sessionID: string\): void \{([\s\S]*?)\n  \}/)
+    expect(match).not.toBeNull()
+    expect(match![1]).toMatch(/if \(this\.streams\.focused === sessionID\) this\.focusSession\(undefined\)/)
+  })
+})
+
+describe("sendMessage / sendCommand draft id contract", () => {
+  const source = readFile(SESSION_FILE)
+
+  it("sendMessage mints a draftID when there is no current session and none was supplied", () => {
+    // External session deletions leave currentSessionID() undefined and clear
+    // draftSessionID(). Without minting a draftID here, the webview posts
+    // {type: "sendMessage", sessionID: undefined, draftID: undefined} and the
+    // extension's sessionCreated echo has no key to migrate the in-flight draft
+    // from ":pending:<id>" to ":session:<newSessionId>". The user loses the
+    // typed message and the new session starts empty.
+    const body = extractFunctionBody(source, "sendMessage")
+    expect(body).toMatch(/!sid && !draftID \? crypto\.randomUUID\(\) : draftID/)
+  })
+
+  it("sendCommand mints a draftID when there is no current session and none was supplied", () => {
+    const body = extractFunctionBody(source, "sendCommand")
+    expect(body).toMatch(/!sid && !draftID \? crypto\.randomUUID\(\) : draftID/)
+  })
+})
+
+describe("PromptInput restoreFailed fallback contract", () => {
+  const PROMPT_FILE = path.join(ROOT, "webview-ui/src/components/chat/PromptInput.tsx")
+  const source = readFile(PROMPT_FILE)
+
+  it("targets draftKey() instead of computing a key from failed.sessionID", () => {
+    // The contract: restoreFailed early-returns when userClearedSession is
+    // true (covers BOTH "user clicked New Task" and the Delete-current-session
+    // race window where currentSessionID/draftSessionID haven't been cleared
+    // yet but userClearedSession is already true). When the user did NOT
+    // explicitly clear, candidates come from the failure's sessionID/draftID
+    // (the keys the send was actually scoped to), plus :new ONLY when the
+    // user has effectively returned to the empty state via an external
+    // session.deleted.
+    const match = source.match(/const restoreFailed = \(failed: SendMessageFailedMessage\) => \{([\s\S]*?)\n  \}/)
+    expect(match).not.toBeNull()
+    expect(match![1]).not.toMatch(/const effectiveSessionID/)
+    expect(match![1]).toMatch(/if \(session\.userClearedSession\(\)\) return/)
+    expect(match![1]).toMatch(
+      /if \(failed\.sessionID\) candidates\.add\(scopeDraftKey\(boxKey\(\),\s*sessionDraftKey\(failed\.sessionID\)\)\)/,
+    )
+    expect(match![1]).toMatch(
+      /if \(failed\.draftID\) candidates\.add\(scopeDraftKey\(boxKey\(\),\s*pendingDraftKey\(failed\.draftID\)\)\)/,
+    )
+    expect(match![1]).toMatch(
+      /if \(!session\.currentSessionID\(\) && !session\.draftSessionID\(\)\) candidates\.add\(scopeDraftKey\(boxKey\(\),\s*"new"\)\)/,
+    )
+    expect(match![1]).toMatch(/const target = draftKey\(\)/)
+    expect(match![1]).toMatch(/candidates\.has\(target\)/)
+  })
+
+  it("does NOT add :new when the user is on a different live session or pending draft", () => {
+    // Guard against the unconditional-:new regression: if the user has
+    // navigated to a different session/pending draft, the failed draft
+    // must NOT be rehydrated into that unrelated prompt even if the
+    // failure carries scope IDs that no longer match the live state.
+    const match = source.match(/const restoreFailed = \(failed: SendMessageFailedMessage\) => \{([\s\S]*?)\n  \}/)
+    expect(match).not.toBeNull()
+    expect(match![1]).not.toMatch(
+      /if \(!failed\.sessionID && !failed\.draftID\) candidates\.add\(scopeDraftKey\(boxKey\(\),\s*"new"\)\)/,
+    )
+  })
+})
+
+describe("SessionContext userClearedSession contract", () => {
+  const source = readFile(SESSION_FILE)
+
+  it("declares userClearedSession on the context interface", () => {
+    // restoreFailed uses session.userClearedSession() to decide whether :new
+    // is a legitimate restore target after the user clicks New Task or
+    // deletes their current/draft session. The accessor must be exposed.
+    expect(source).toMatch(/userClearedSession:\s*Accessor<boolean>/)
+  })
+
+  it("clearCurrentSession sets the flag", () => {
+    // User clicking New Task while a failure is pending must NOT restore
+    // the failed draft into the new prompt.
+    const body = extractFunctionBody(source, "clearCurrentSession")
+    expect(body).toMatch(/setUserClearedSession\(true\)/)
+  })
+
+  it("deleteSession sets the flag when deleting the current or draft session", () => {
+    // User clicking Delete on their current/draft session is morally the
+    // same as New Task — both land in :new without wanting a stale restore.
+    const body = extractFunctionBody(source, "deleteSession")
+    expect(body).toMatch(
+      /if \(id === currentSessionID\(\) \|\| id === draftSessionID\(\)\) setUserClearedSession\(true\)/,
+    )
+  })
+
+  it("handleSessionCreated resets the flag when adopting the new session", () => {
+    // After the user creates a new session, the flag is stale and must be
+    // cleared so a later external delete of that new session can restore
+    // into :new again.
+    const body = extractFunctionBody(source, "handleSessionCreated")
+    expect(body).toMatch(/setUserClearedSession\(false\)/)
+  })
+
+  it("selectSession resets the flag when picking an existing session", () => {
+    const body = extractFunctionBody(source, "selectSession")
+    expect(body).toMatch(/setUserClearedSession\(false\)/)
+  })
+
+  it("exposes userClearedSession in the SessionContext value", () => {
+    expect(source).toMatch(/userClearedSession,?\s*\n\s*\}/m)
+  })
+
+  it("sendMessage resets userClearedSession when starting a fresh draft from :new", () => {
+    // Race: user on session A, sends, failure pending; clicks New Task
+    // (userClearedSession=true), then types new text and clicks Send. We mint
+    // a draftID and adopt it as draftSessionID. If a failure for the new
+    // send returns BEFORE sessionCreated lands (so currentSessionID is
+    // still undefined and userClearedSession is still true), the failure's
+    // draftID matches draftSessionID() but the flag would suppress restore.
+    // Resetting the flag at the moment the user starts the new draft closes
+    // that window: the failure is for the current in-progress draft and must
+    // be restorable.
+    const body = extractFunctionBody(source, "sendMessage")
+    const block = body.match(/if \(!sid\) \{([\s\S]*?)\}/)
+    expect(block).not.toBeNull()
+    expect(block![1]).toMatch(/setUserClearedSession\(false\)/)
+    expect(block![1]).toMatch(/setDraftSessionID\(scope\)/)
+  })
+
+  it("sendCommand resets userClearedSession when starting a fresh draft from :new", () => {
+    const body = extractFunctionBody(source, "sendCommand")
+    const block = body.match(/if \(!sid\) \{([\s\S]*?)\}/)
+    expect(block).not.toBeNull()
+    expect(block![1]).toMatch(/setUserClearedSession\(false\)/)
+    expect(block![1]).toMatch(/setDraftSessionID\(scope\)/)
+  })
+
+  it("selectCloudSession resets userClearedSession when picking a cloud session", () => {
+    // After clearCurrentSession set the flag, selecting a cloud session
+    // must clear it (mirrors selectSession's reset). Without this, any
+    // post-import failure exits restoration early and loses the cleared
+    // text, review comments, and images.
+    const body = extractFunctionBody(source, "selectCloudSession")
+    expect(body).toMatch(/setUserClearedSession\(false\)/)
+  })
+
+  it("handleCloudSessionImported resets userClearedSession after the import completes", () => {
+    // Defense in depth: even if selectCloudSession's reset was missed
+    // (e.g. deleteSession set the flag against the synthetic cloud key
+    // between select and import), the import confirmation must clear the
+    // flag so a later post-import send failure is not suppressed.
+    const body = extractFunctionBody(source, "handleCloudSessionImported")
+    expect(body).toMatch(/setUserClearedSession\(false\)/)
+  })
+
+  it("handleCloudSessionImported migrates draftSessionID from the cloud key to the real session id", () => {
+    // Without this, draftSessionID stays on the synthetic "cloud:<id>" key.
+    // After a later external delete of the imported session,
+    // handleSessionDeleted only clears draftSessionID when it equals the
+    // deleted id; the synthetic cloud key never matches, so draftKey()
+    // falls back to ":pending:cloud:<id>" and restoreFailed can no longer
+    // match :session:<id> or :new — silently losing the failed draft.
+    const body = extractFunctionBody(source, "handleCloudSessionImported")
+    expect(body).toMatch(/setDraftSessionID\(session\.id\)/)
+  })
+})
+
+describe("Cloud import parts cleanup contract", () => {
+  const source = readFile(SESSION_FILE)
+
+  it("declares a pendingCloudPrune tracker for cloud message IDs", () => {
+    // Without a tracker, repeated preview -> import cycles accumulate full
+    // cloud transcripts in store.parts because handleMessagesLoaded never
+    // knows which keys belong to the carried-over cloud messages.
+    expect(source).toMatch(/pendingCloudPrune/)
+  })
+
+  it("handleCloudSessionDataLoaded registers the cloud message IDs", () => {
+    const body = extractFunctionBody(source, "handleCloudSessionDataLoaded")
+    expect(body).toMatch(/pendingCloudPrune\.set\(/)
+  })
+
+  it("handleCloudSessionImported transfers the prune set to the new session id", () => {
+    const body = extractFunctionBody(source, "handleCloudSessionImported")
+    expect(body).toMatch(/pendingCloudPrune\.set\(session\.id,/)
+    expect(body).toMatch(/pendingCloudPrune\.delete\(cloudKey\)/)
+  })
+
+  it("handleMessagesLoaded prunes cloud-import orphans from store.parts and stash", () => {
+    // The carried-over cloud messages are gone from store.messages after
+    // this call, so any store.parts[<cloud-msg-id>] entry is unreachable.
+    const body = extractFunctionBody(source, "handleMessagesLoaded")
+    expect(body).toMatch(/pendingCloudPrune\.get\(sessionID\)/)
+    expect(body).toMatch(/pendingCloudPrune\.delete\(sessionID\)/)
+  })
+
+  it("handleSessionDeleted prunes cloud-import orphans if the imported session is deleted before loadMessages returns", () => {
+    const body = extractFunctionBody(source, "handleSessionDeleted")
+    expect(body).toMatch(/pruneCloudOrphans\(sessionID\)/)
+  })
+
+  it("handleCloudSessionImportFailed prunes cloud parts and the synthetic session entries", () => {
+    // Implemented as a switch case inside handleExtensionMessage, not a
+    // standalone function, so search the source for the case body directly.
+    const idx = source.indexOf('case "cloudSessionImportFailed"')
+    expect(idx).toBeGreaterThan(-1)
+    const after = source.slice(idx, idx + 4000)
+    expect(after).toMatch(/pruneCloudOrphans\(failedKey\)/)
+    expect(after).toMatch(/delete sessions\[failedKey\]/)
+    expect(after).toMatch(/delete messages\[failedKey\]/)
+  })
+
+  it("handleCloudSessionImportFailed clears cloudPreviewId, currentSessionID, draftSessionID, and loading only when still on the failed cloud session", () => {
+    // The failure arrives asynchronously. selectCloudSession sets the
+    // preview id to the RAW cloud session id, both session/draft ids to
+    // the synthetic "cloud:<id>" key, and the loading spinner, but the
+    // user can start previewing a different cloud session, switch
+    // sessions, or start a new task before the failure comes back.
+    // Unconditionally resetting any of them would clobber that newer
+    // scope: cloudPreviewId blanking drops a later preview response and
+    // disables import-mode sends; currentSessionID blanking blanks
+    // the active session; draftSessionID blanking leaves draftKey()
+    // at ":new"; and unguarded setLoading(false) drops the spinner
+    // for a newer preview before its data arrives, leaving the UI
+    // looking idle while still loading. Clear only if still on the
+    // dead preview's scope: cloudPreviewId is compared against the raw
+    // message.cloudSessionId, while currentSessionID/draftSessionID are
+    // compared against the "cloud:<id>" failedKey. The guard is
+    // extracted into a clearIfOn helper to keep the switch-case
+    // complexity under the lint cap.
+    //
+    // The loading check MUST run before cloudPreviewId is nulled,
+    // otherwise `cloudPreviewId() === message.cloudSessionId` would be
+    // false even on the failing preview and the spinner would stick
+    // until later navigation clears it.
+    const idx = source.indexOf('case "cloudSessionImportFailed"')
+    expect(idx).toBeGreaterThan(-1)
+    const after = source.slice(idx, idx + 4000)
+    expect(after).toMatch(/clearIfOn\(cloudPreviewId, \(\) => setLoading\(false\), message\.cloudSessionId\)/)
+    expect(after).toMatch(/clearIfOn\(cloudPreviewId, \(\) => setCloudPreviewId\(null\), message\.cloudSessionId\)/)
+    expect(after).toMatch(/clearIfOn\(currentSessionID, \(\) => setCurrentSessionID\(undefined\), failedKey\)/)
+    expect(after).toMatch(/clearIfOn\(draftSessionID, \(\) => setDraftSessionID\(undefined\), failedKey\)/)
+    expect(after).not.toMatch(/^\s*setLoading\(false\)\s*$/m)
+    // Loading check must come before cloudPreviewId null in the case body.
+    const loadIdx = after.indexOf("setLoading(false)")
+    const nullIdx = after.indexOf("setCloudPreviewId(null)")
+    expect(loadIdx).toBeGreaterThan(-1)
+    expect(nullIdx).toBeGreaterThan(-1)
+    expect(loadIdx).toBeLessThan(nullIdx)
+  })
+
+  it("clearIfOn runs the clear callback only while the scope still matches the key", () => {
+    // Used by cloudSessionImportFailed so the switch case stays under the
+    // complexity cap. The helper must compare get() to the key before
+    // calling the clear callback: a stale async failure must not clobber
+    // a newer scope the user has navigated to. Takes a clear callback
+    // rather than a setter so the same helper works for both
+    // undefined-cleared signals (currentSessionID / draftSessionID) and
+    // null-cleared signals (cloudPreviewId) without changing their setter
+    // signatures.
+    let cleared = 0
+    let value = "pending"
+    clearIfOn(
+      () => value,
+      () => {
+        cleared++
+      },
+      "pending",
+    )
+    expect(cleared).toBe(1)
+
+    // Scope has moved on (user navigated to a different preview / session)
+    // — the clear callback must NOT run.
+    value = "other"
+    clearIfOn(
+      () => value,
+      () => {
+        cleared++
+      },
+      "pending",
+    )
+    expect(cleared).toBe(1)
+  })
+})
+
+describe("KiloConnectionService pruneSession contract", () => {
+  const source = readFile(CONNECTION_SERVICE_FILE)
+
+  it("drops the deleted session from focused and opened Maps", () => {
+    // KiloProvider's pruneDeletedSession calls connectionService.pruneSession.
+    // Without clearing focused/opened entries whose value is the deleted id,
+    // the backend keeps receiving viewed.focused with the dead session id and
+    // any background tab opener stays registered for it.
+    const match = source.match(/pruneSession\(sessionId: string\): void \{([\s\S]*?)\n  \}/)
+    expect(match).not.toBeNull()
+    expect(match![1]).toMatch(/this\.focused\.delete\(key\)/)
+    expect(match![1]).toMatch(/this\.opened\.(?:set|delete)/)
+    expect(match![1]).toMatch(/this\.flushViewed\(\)/)
+  })
+})

packages/kilo-vscode/tests/unit/session-stream-scheduler.test.ts

@@ -255,6 +255,17 @@ describe("SessionStreamScheduler / focus and lifecycle", () => {
     expect(stats.active).toBe(0)
   })
 
+  it("exposes the focused session id via the getter so pruneDeletedSession can match it", () => {
+    const queue = new SessionStreamScheduler(() => {})
+    expect(queue.focused).toBeUndefined()
+    queue.focus("sess-1")
+    expect(queue.focused).toBe("sess-1")
+    queue.focus("sess-2")
+    expect(queue.focused).toBe("sess-2")
+    queue.focus(undefined)
+    expect(queue.focused).toBeUndefined()
+  })
+
   it("dispose() stops further emissions from queued work", async () => {
     const sent: Sent[] = []
     const queue = new SessionStreamScheduler((msg) => sent.push(msg), {

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

@@ -57,16 +57,12 @@ import {
   scopeDraftKey,
   sessionDraftKey,
 } from "../../utils/prompt-drafts"
+import { drafts, imageDrafts, reviewDrafts } from "../../utils/draft-store"
 import { ReviewComments } from "./ReviewComments"
 import { partReview, reviewBody } from "../../../../src/shared/review-comments"
 import { isEnterKeyCommitNotIme } from "../../utils/ime-enter"
 
-// Per-session input text storage (module-level so it survives remounts)
-const drafts = new Map<string, string>()
-const reviewDrafts = new Map<string, ReviewComment[]>()
-const imageDrafts = new Map<string, ImageAttachment[]>()
 const scrolls = new Map<string, number>()
-
 function mergeReviewComments(current: ReviewComment[], incoming: ReviewComment[]): ReviewComment[] {
   if (incoming.length === 0) return current
   const map = new Map(current.map((item) => [item.id, item]))
@@ -472,11 +468,33 @@ export const PromptInput: Component<PromptInputProps> = (props) => {
 
   const restoreFailed = (failed: SendMessageFailedMessage) => {
     // Only restore a failed draft when the user has not started another one.
-    const target = scopeDraftKey(
-      boxKey(),
-      sessionDraftKey(failed.sessionID) ?? pendingDraftKey(failed.draftID) ?? "new",
-    )
-    if (target !== draftKey() || text().trim() || reviewComments().length > 0 || imageAttach.images().length > 0) return
+    if (text().trim() || reviewComments().length > 0 || imageAttach.images().length > 0) return
+
+    // If the user explicitly transitioned out of the original send's scope
+    // (clearCurrentSession() or Delete on the current/draft session), don't
+    // restore anywhere. This covers BOTH the obvious "user clicked New Task
+    // and we land in :new" case AND the tighter race window where the user
+    // clicked Delete on the current session: the backend's sessionDeleted
+    // round-trip hasn't completed yet so currentSessionID/draftSessionID
+    // still point at the dead session, but userClearedSession is true. Without
+    // this guard, the session-scoped candidate on the previous lines would
+    // match the still-current draftKey and rehydrate the failed draft into
+    // the session the user explicitly chose to delete.
+    if (session.userClearedSession()) return
+
+    // Build candidates from the keys the original send was actually scoped
+    // under. :new is only added when the user has effectively returned to the
+    // empty state — i.e. no current session and no pending draft. Combined
+    // with the userClearedSession early return above, this catches both
+    // "send from session -> session deleted mid-round-trip" and "send from
+    // :new (mints draftID) -> session created mid-round-trip -> session
+    // deleted externally" without rehydrating into any user-explicit clear.
+    const candidates = new Set<string>()
+    if (failed.sessionID) candidates.add(scopeDraftKey(boxKey(), sessionDraftKey(failed.sessionID)))
+    if (failed.draftID) candidates.add(scopeDraftKey(boxKey(), pendingDraftKey(failed.draftID)))
+    if (!session.currentSessionID() && !session.draftSessionID()) candidates.add(scopeDraftKey(boxKey(), "new"))
+    const target = draftKey()
+    if (!candidates.has(target)) return
 
     const draft = failed.review ? reviewBody(failed.review, failed.text) : failed.text
     if (draft === undefined) return

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

@@ -0,0 +1,39 @@
+/**
+ * Tracks cloud-import message IDs awaiting parts cleanup. During preview,
+ * parts are stored keyed by original cloud message IDs (store.parts["<cloud-msg-id>"]).
+ * On import, the carried-over messages keep those IDs until handleMessagesLoaded
+ * replaces them with server-assigned IDs. Once the carried-over messages are gone
+ * from store.messages, their parts are orphans — pruneCloudOrphans drops them so
+ * preview -> import cycles don't accumulate full transcripts in the reactive store.
+ */
+import type { Part } from "../types/messages"
+
+export interface PruneStash {
+  remove: (id: string) => void
+}
+
+export const createCloudPrune = (
+  setParts: (mutator: (parts: Record<string, Part[]>) => void) => void,
+  stash: PruneStash,
+) => {
+  const pendingCloudPrune = new Map<string, Set<string>>()
+
+  const prune = (key: string) => {
+    const ids = pendingCloudPrune.get(key)
+    if (!ids) return
+    setParts((parts) => {
+      for (const id of ids) delete parts[id]
+    })
+    for (const id of ids) stash.remove(id)
+    pendingCloudPrune.delete(key)
+  }
+
+  return { pendingCloudPrune, prune }
+}
+
+/** Clear a scope only if it still points at the given key. Async failure paths
+ * must not clobber scopes the user has navigated to since the operation was
+ * started. */
+export function clearIfOn<T>(get: () => T, clear: () => void, key: T) {
+  if (get() === key) clear()
+}

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

@@ -76,12 +76,21 @@ import { getVariant, sessionVariantKeys, transferVariants, variantKey } from "./
 import { KILO_AUTO, KILO_PROVIDER_ID, parseModelString } from "../../../src/shared/provider-model"
 import { reviewMetadata, type ReviewMessageData } from "../../../src/shared/review-comments"
 import { visibleMessages as filterVisibleMessages } from "./session-queue"
+import { deleteDraftsForSession } from "../utils/draft-store"
 import { createAbortState } from "./abort-state"
+import { clearIfOn, createCloudPrune } from "./session-cloud-prune"
 import { isSameSessionTree } from "./model-usage"
 
 const RECENT_LIMIT = 5
 const MESSAGE_PAGE_LIMIT = 80
 
+/** Remove ids from a Set immutably, returning the original when nothing changed. */
+function dropSet(prev: Set<string>, ids: Iterable<string>): Set<string> {
+  const next = new Set(prev)
+  for (const id of ids) next.delete(id)
+  return next.size === prev.size ? prev : next
+}
+
 type MessageMutation = Exclude<MessageLoadMode, "focus"> | "append" | "update"
 
 interface MessagePageState {
@@ -294,6 +303,7 @@ interface SessionContextValue {
   selectCloudSession: (cloudSessionId: string) => void
   draftSessionID: Accessor<string | undefined>
   setDraftSessionID: (id: string | undefined) => void
+  userClearedSession: Accessor<boolean>
 }
 
 export const SessionContext = createContext<SessionContextValue>()
@@ -308,6 +318,7 @@ export const SessionProvider: ParentComponent = (props) => {
   // Current session ID
   const [currentSessionID, setCurrentSessionID] = createSignal<string | undefined>()
   const [draftSessionID, setDraftSessionID] = createSignal<string | undefined>()
+  const [userClearedSession, setUserClearedSession] = createSignal(false)
 
   // Per-session status map — keyed by sessionID
   const [statusMap, setStatusMap] = createStore<Record<string, SessionStatusInfo>>({})
@@ -547,6 +558,8 @@ export const SessionProvider: ParentComponent = (props) => {
 
   const agentNames = createMemo(() => new Set(agents().map((agent) => agent.name)))
 
+  const { pendingCloudPrune, prune: pruneCloudOrphans } = createCloudPrune((m) => setStore("parts", produce(m)), stash)
+
   /** Per-mode model from config (e.g. config.agent.code.model). */
   function getModeModel(agentName: string): ModelSelection | null {
     return parseModelString(config().agent?.[agentName]?.model)
@@ -1156,10 +1169,33 @@ export const SessionProvider: ParentComponent = (props) => {
         handleCloudSessionImported(message.cloudSessionId, message.session)
         break
 
-      case "cloudSessionImportFailed":
-        setCloudPreviewId(null)
-        setCurrentSessionID(undefined)
-        setLoading(false)
+      case "cloudSessionImportFailed": {
+        const failedKey = `cloud:${message.cloudSessionId}`
+        pruneCloudOrphans(failedKey)
+        setStore(
+          "sessions",
+          produce((sessions) => {
+            delete sessions[failedKey]
+          }),
+        )
+        setStore(
+          "messages",
+          produce((messages) => {
+            delete messages[failedKey]
+          }),
+        )
+        setStore(
+          "toolParts",
+          produce((toolParts) => {
+            delete toolParts[failedKey]
+          }),
+        )
+        // cloudPreviewId stores the raw cloud session id (see selectCloudSession),
+        // not the synthetic "cloud:<id>" key used for session/draft ids.
+        clearIfOn(cloudPreviewId, () => setLoading(false), message.cloudSessionId)
+        clearIfOn(cloudPreviewId, () => setCloudPreviewId(null), message.cloudSessionId)
+        clearIfOn(currentSessionID, () => setCurrentSessionID(undefined), failedKey)
+        clearIfOn(draftSessionID, () => setDraftSessionID(undefined), failedKey)
         showToast({
           variant: "error",
           title: language.t("session.cloud.import.failed") ?? "Failed to import cloud session",
@@ -1167,6 +1203,7 @@ export const SessionProvider: ParentComponent = (props) => {
         })
         console.error("[Kilo New] Cloud session import failed:", message.error)
         break
+      }
 
       case "worktreeStatsLoaded":
         setWorktreeStats({ files: message.files, additions: message.additions, deletions: message.deletions })
@@ -1292,6 +1329,7 @@ export const SessionProvider: ParentComponent = (props) => {
       if (!draftID || draft === draftID || active === draftID) {
         setCurrentSessionID(session.id)
         setDraftSessionID(session.id)
+        setUserClearedSession(false)
       }
     })
   }
@@ -1484,6 +1522,19 @@ export const SessionProvider: ParentComponent = (props) => {
       const revert = store.sessions[sessionID]?.revert ?? undefined
       if (revert) resetTodos(sessionID, revert)
       recoverPrefs(sessionID, merged)
+
+      const cloudIDs = pendingCloudPrune.get(sessionID)
+      if (cloudIDs?.size) {
+        const live = new Set(messages.map((m) => m.id))
+        setStore(
+          "parts",
+          produce((p) => {
+            for (const id of cloudIDs) if (!live.has(id)) delete p[id]
+          }),
+        )
+        for (const id of cloudIDs) stash.remove(id)
+        pendingCloudPrune.delete(sessionID)
+      }
     })
     if (reset) requestAnimationFrame(() => patchPage(sessionID, { lastMutation: undefined }))
   }
@@ -1924,104 +1975,61 @@ export const SessionProvider: ParentComponent = (props) => {
       clearHiddenErrors(msgIds)
 
       setStore(
-        "sessions",
-        produce((sessions) => {
-          delete sessions[sessionID]
-        }),
-      )
-      setStore(
-        "messages",
-        produce((messages) => {
-          delete messages[sessionID]
-        }),
-      )
-      setStore(
-        "parts",
-        produce((parts) => {
-          for (const id of msgIds) {
-            delete parts[id]
+        produce((s) => {
+          delete s.sessions[sessionID]
+          delete s.messages[sessionID]
+          for (const id of msgIds) delete s.parts[id]
+          delete s.toolParts[sessionID]
+          delete s.todos[sessionID]
+          for (const [id, state] of Object.entries(s.modelUsage)) {
+            if (id === sessionID || state.data?.sessionIDs.includes(sessionID)) delete s.modelUsage[id]
           }
+          delete s.agentSelections[sessionID]
+          delete s.sessionOverrides[sessionID]
+          for (const key of sessionVariantKeys(s.variantSelections, sessionID)) delete s.variantSelections[key]
         }),
       )
-      setStore(
-        "toolParts",
-        produce((parts) => {
-          delete parts[sessionID]
-        }),
-      )
-      setStore(
-        "todos",
-        produce((todos) => {
-          delete todos[sessionID]
-        }),
-      )
-      setStore(
-        "modelUsage",
-        produce((usage) => {
-          for (const [id, state] of Object.entries(usage)) {
-            if (id === sessionID || state.data?.sessionIDs.includes(sessionID)) delete usage[id]
-          }
-        }),
-      )
-      setPages(
-        produce((map) => {
-          delete map[sessionID]
-        }),
-      )
-      setStore(
-        "agentSelections",
-        produce((selections) => {
-          delete selections[sessionID]
-        }),
-      )
+      // prettier-ignore
+      setPages(produce((map) => { delete map[sessionID] }))
       // Clean up pending questions/errors for the deleted session
       const deleted = questions()
         .filter((q) => q.sessionID === sessionID)
         .map((q) => q.id)
       if (deleted.length > 0) {
         setQuestions((prev) => prev.filter((q) => q.sessionID !== sessionID))
-        setQuestionErrors((prev) => {
-          const next = new Set(prev)
-          for (const id of deleted) next.delete(id)
-          if (next.size === prev.size) return prev
-          return next
-        })
+        setQuestionErrors((prev) => dropSet(prev, deleted))
       }
       const gone = suggestions()
         .filter((item) => item.sessionID === sessionID)
         .map((item) => item.id)
       if (gone.length > 0) {
         setSuggestions((prev) => prev.filter((item) => item.sessionID !== sessionID))
-        setSuggestionErrors((prev) => {
-          const next = new Set(prev)
-          for (const id of gone) next.delete(id)
-          if (next.size === prev.size) return prev
-          return next
-        })
-        setRespondingSuggestions((prev) => {
-          const next = new Set(prev)
-          for (const id of gone) next.delete(id)
-          if (next.size === prev.size) return prev
-          return next
-        })
+        setSuggestionErrors((prev) => dropSet(prev, gone))
+        setRespondingSuggestions((prev) => dropSet(prev, gone))
       }
+      const staleResponding = permissions()
+        .filter((p) => p.sessionID === sessionID)
+        .map((p) => p.id)
       setPermissions((prev) => removeSessionPermissions(prev, sessionID))
-      setStatusMap(
-        produce((map) => {
-          delete map[sessionID]
-        }),
-      )
+      if (staleResponding.length > 0) {
+        setRespondingPermissions((prev) => dropSet(prev, staleResponding))
+      }
+      // prettier-ignore
+      setLoaded((prev) => { if (!prev.has(sessionID)) return prev; const next = new Set(prev); next.delete(sessionID); return next })
+      // prettier-ignore
+      setStatusMap(produce((map) => { delete map[sessionID] }))
       clearClose(sessionID)
-      setBusySinceMap(
-        produce((map) => {
-          delete map[sessionID]
-        }),
-      )
+      // prettier-ignore
+      setBusySinceMap(produce((map) => { delete map[sessionID] }))
       if (currentSessionID() === sessionID) {
         setCurrentSessionID(undefined)
         setLoading(false)
       }
+      // prettier-ignore
+      if (draftSessionID() === sessionID) { setDraftSessionID(undefined) }
     })
+    deleteDraftsForSession(sessionID)
+    pruneCloudOrphans(sessionID)
   }
 
   // Splices the message from the store and deletes its parts.
@@ -2044,6 +2052,7 @@ export const SessionProvider: ParentComponent = (props) => {
   function handleCloudSessionDataLoaded(cloudSessionId: string, title: string, messages: Message[]) {
     if (cloudPreviewId() !== cloudSessionId) return
     const key = `cloud:${cloudSessionId}`
+    pendingCloudPrune.set(key, new Set(messages.map((m) => m.id)))
     batch(() => {
       setLoaded((prev) => {
         if (prev.has(key)) return prev
@@ -2093,31 +2102,10 @@ export const SessionProvider: ParentComponent = (props) => {
 
       setCloudPreviewId(null)
       setCurrentSessionID(session.id)
+      setDraftSessionID(session.id)
+
+      setUserClearedSession(false)
 
-      // Clean up synthetic cloud: entries from sessions/messages stores.
-      //
-      // Why we do NOT delete cloud parts here:
-      //
-      // During preview, parts are stored keyed by the original cloud message IDs
-      // (e.g. store.parts["<cloud-msg-id>"] = [...]). When the import completes
-      // we carry cloudMessages into the new local session (above) so the UI
-      // renders immediately without a loading flash. Those carried-over message
-      // objects still hold their original cloud IDs, so every SessionTurn
-      // calls getParts("<cloud-msg-id>") — which means the parts must remain in
-      // the store for now.
-      //
-      // If we deleted them here, every message would temporarily render with no
-      // parts (parts().length === 0), showing only a loading shimmer until the
-      // real data arrives.
-      //
-      // Instead, right after this batch we dispatch a "loadMessages" request
-      // (below). When the extension responds with the "messagesLoaded" event,
-      // handleMessagesLoaded() replaces the messages array with server-assigned
-      // IDs and writes new parts keyed by those IDs. The old cloud-keyed part
-      // entries become orphans — no message in the store references them anymore.
-      // They remain in store.parts until the webview reloads or the store is
-      // reset, which is a bounded, one-session-worth amount of data that does
-      // not accumulate over time.
       setStore(
         "sessions",
         produce((sessions) => {
@@ -2137,6 +2125,11 @@ export const SessionProvider: ParentComponent = (props) => {
         }),
       )
     })
+    const cloudPruneIDs = pendingCloudPrune.get(cloudKey)
+    if (cloudPruneIDs) {
+      pendingCloudPrune.set(session.id, cloudPruneIDs)
+      pendingCloudPrune.delete(cloudKey)
+    }
     // Load real messages in the background (picks up server-assigned IDs
     // and the new user message once the send completes via SSE)
     patchPage(session.id, { loadingInitial: true, before: undefined, hasMore: false })
@@ -2263,12 +2256,16 @@ export const SessionProvider: ParentComponent = (props) => {
       dismissQuestion(q.id)
     }
 
-    const scope = draftID ?? sid
+    const effectiveDraftID = !sid && !draftID ? crypto.randomUUID() : draftID
+    const scope = effectiveDraftID ?? sid
     if (scope) {
       clearClose(scope)
       addOptimistic(scope, messageID, text, files, review)
       startSubmission(scope, messageID)
-      if (!sid) setDraftSessionID(scope)
+      if (!sid) {
+        setUserClearedSession(false)
+        setDraftSessionID(scope)
+      }
     }
     const agent = promptAgent(scope)
 
@@ -2277,7 +2274,7 @@ export const SessionProvider: ParentComponent = (props) => {
       text,
       messageID,
       sessionID: sid,
-      draftID,
+      draftID: effectiveDraftID,
       providerID,
       modelID,
       agent,
@@ -2331,12 +2328,16 @@ export const SessionProvider: ParentComponent = (props) => {
       dismissQuestion(q.id)
     }
 
-    const scope = draftID ?? sid
+    const effectiveDraftID = !sid && !draftID ? crypto.randomUUID() : draftID
+    const scope = effectiveDraftID ?? sid
     if (scope) {
       clearClose(scope)
       addOptimistic(scope, messageID, `/${command} ${args}`.trim(), files)
       startSubmission(scope, messageID)
-      if (!sid) setDraftSessionID(scope)
+      if (!sid) {
+        setUserClearedSession(false)
+        setDraftSessionID(scope)
+      }
     }
     const agent = promptAgent(scope)
 
@@ -2346,7 +2347,7 @@ export const SessionProvider: ParentComponent = (props) => {
       arguments: args,
       messageID,
       sessionID: sid,
-      draftID,
+      draftID: effectiveDraftID,
       providerID,
       modelID,
       agent,
@@ -2504,6 +2505,7 @@ export const SessionProvider: ParentComponent = (props) => {
   }
 
   function clearCurrentSession() {
+    setUserClearedSession(true)
     setCurrentSessionID(undefined)
     setDraftSessionID(undefined)
     setCloudPreviewId(null)
@@ -2553,6 +2555,7 @@ export const SessionProvider: ParentComponent = (props) => {
     // the worktree selection) still moved (the reported "only the diff changes").
     setCurrentSessionID(id)
     setDraftSessionID(id)
+    setUserClearedSession(false)
     setLoading(!ready)
     if (!ready) patchPage(id, { loadingInitial: true, loadingOlder: false, before: undefined, hasMore: false })
     // Only the message fetch needs the backend. Defer it while offline and let
@@ -2599,6 +2602,7 @@ export const SessionProvider: ParentComponent = (props) => {
     setCloudPreviewId(cloudSessionId)
     setCurrentSessionID(key)
     setDraftSessionID(key)
+    setUserClearedSession(false)
     setLoading(true)
     vscode.postMessage({ type: "requestCloudSessionData", sessionId: cloudSessionId })
   }
@@ -2621,6 +2625,7 @@ export const SessionProvider: ParentComponent = (props) => {
       next.delete(id)
       return next
     })
+    if (id === currentSessionID() || id === draftSessionID()) setUserClearedSession(true)
     vscode.postMessage({ type: "deleteSession", sessionID: id })
   }
 
@@ -2960,6 +2965,7 @@ export const SessionProvider: ParentComponent = (props) => {
     selectCloudSession,
     draftSessionID,
     setDraftSessionID,
+    userClearedSession,
   }
 
   return <SessionContext.Provider value={value}>{props.children}</SessionContext.Provider>

packages/kilo-vscode/webview-ui/src/stories/StoryProviders.tsx

@@ -193,6 +193,7 @@ export function mockSessionValue(overrides?: {
     submitting: () => false,
     draftSessionID: () => undefined,
     setDraftSessionID: noop,
+    userClearedSession: () => false,
     messageMutation: () => undefined,
     messages: () => [],
     visibleMessages: () => [],

packages/kilo-vscode/webview-ui/src/utils/draft-store.ts

@@ -0,0 +1,21 @@
+import type { ReviewComment } from "../types/messages"
+import type { ImageAttachment } from "../hooks/useImageAttachments"
+
+export const drafts = new Map<string, string>()
+export const reviewDrafts = new Map<string, ReviewComment[]>()
+export const imageDrafts = new Map<string, ImageAttachment[]>()
+
+export function deleteDraftsForSession(id: string) {
+  if (!id) return
+  const sessionSuffix = `:session:${id}`
+  const pendingSuffix = `:pending:${id}`
+  const maps = [drafts, reviewDrafts, imageDrafts]
+  for (const map of maps) {
+    for (const key of map.keys()) {
+      if (typeof key !== "string") continue
+      if (key.endsWith(sessionSuffix) || key.endsWith(pendingSuffix)) {
+        map.delete(key)
+      }
+    }
+  }
+}