diff --git a/.changeset/cleanup-drafts-on-session-delete.md b/.changeset/cleanup-drafts-on-session-delete.md new file mode 100644 index 0000000000..67d1bbc9e3 --- /dev/null +++ b/.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. \ No newline at end of file diff --git a/packages/kilo-vscode/src/KiloProvider.ts b/packages/kilo-vscode/src/KiloProvider.ts index d8fbaa2810..e559376f43 100644 --- a/packages/kilo-vscode/src/KiloProvider.ts +++ b/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( diff --git a/packages/kilo-vscode/src/kilo-provider/session-stream-scheduler.ts b/packages/kilo-vscode/src/kilo-provider/session-stream-scheduler.ts index 3209d3c6e9..46e2afba74 100644 --- a/packages/kilo-vscode/src/kilo-provider/session-stream-scheduler.ts +++ b/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 diff --git a/packages/kilo-vscode/src/services/cli-backend/connection-service.ts b/packages/kilo-vscode/src/services/cli-backend/connection-service.ts index c2f0f5b6f7..c91beacdfe 100644 --- a/packages/kilo-vscode/src/services/cli-backend/connection-service.ts +++ b/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() } /** diff --git a/packages/kilo-vscode/tests/unit/prompt-drafts.test.ts b/packages/kilo-vscode/tests/unit/prompt-drafts.test.ts index a10e100ed4..d23bcb9762 100644 --- a/packages/kilo-vscode/tests/unit/prompt-drafts.test.ts +++ b/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:" 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:" 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("race") + const [draftSessionID, setDraftSessionID] = createSignal("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:". + 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") diff --git a/packages/kilo-vscode/tests/unit/prompt-send-contract.test.ts b/packages/kilo-vscode/tests/unit/prompt-send-contract.test.ts index 5c2c53dcdb..0aee37e70b 100644 --- a/packages/kilo-vscode/tests/unit/prompt-send-contract.test.ts +++ b/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:" to ":session:". 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/) + }) + + 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:" 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:" and restoreFailed can no longer + // match :session: 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[] 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:" 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:" 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\(\)/) + }) +}) diff --git a/packages/kilo-vscode/tests/unit/session-stream-scheduler.test.ts b/packages/kilo-vscode/tests/unit/session-stream-scheduler.test.ts index 6248b89067..108aed9aea 100644 --- a/packages/kilo-vscode/tests/unit/session-stream-scheduler.test.ts +++ b/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), { 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 7050ddb12b..cad89ad618 100644 --- a/packages/kilo-vscode/webview-ui/src/components/chat/PromptInput.tsx +++ b/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() -const reviewDrafts = new Map() -const imageDrafts = new Map() const scrolls = new Map() - 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 = (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() + 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 diff --git a/packages/kilo-vscode/webview-ui/src/context/session-cloud-prune.ts b/packages/kilo-vscode/webview-ui/src/context/session-cloud-prune.ts new file mode 100644 index 0000000000..2f497b6feb --- /dev/null +++ b/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[""]). + * 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) => void) => void, + stash: PruneStash, +) => { + const pendingCloudPrune = new Map>() + + 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(get: () => T, clear: () => void, key: T) { + if (get() === key) clear() +} diff --git a/packages/kilo-vscode/webview-ui/src/context/session.tsx b/packages/kilo-vscode/webview-ui/src/context/session.tsx index 9dc87f2180..4336c85a2c 100644 --- a/packages/kilo-vscode/webview-ui/src/context/session.tsx +++ b/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, ids: Iterable): Set { + const next = new Set(prev) + for (const id of ids) next.delete(id) + return next.size === prev.size ? prev : next +} + type MessageMutation = Exclude | "append" | "update" interface MessagePageState { @@ -294,6 +303,7 @@ interface SessionContextValue { selectCloudSession: (cloudSessionId: string) => void draftSessionID: Accessor setDraftSessionID: (id: string | undefined) => void + userClearedSession: Accessor } export const SessionContext = createContext() @@ -308,6 +318,7 @@ export const SessionProvider: ParentComponent = (props) => { // Current session ID const [currentSessionID, setCurrentSessionID] = createSignal() const [draftSessionID, setDraftSessionID] = createSignal() + const [userClearedSession, setUserClearedSession] = createSignal(false) // Per-session status map — keyed by sessionID const [statusMap, setStatusMap] = createStore>({}) @@ -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:" 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[""] = [...]). 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("") — 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 {props.children} diff --git a/packages/kilo-vscode/webview-ui/src/stories/StoryProviders.tsx b/packages/kilo-vscode/webview-ui/src/stories/StoryProviders.tsx index 38fe24e3d1..71bd4db4f3 100644 --- a/packages/kilo-vscode/webview-ui/src/stories/StoryProviders.tsx +++ b/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: () => [], diff --git a/packages/kilo-vscode/webview-ui/src/utils/draft-store.ts b/packages/kilo-vscode/webview-ui/src/utils/draft-store.ts new file mode 100644 index 0000000000..25d93883a7 --- /dev/null +++ b/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() +export const reviewDrafts = new Map() +export const imageDrafts = new Map() + +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) + } + } + } +}