fix(vscode): restore session state on first redo

.changeset/fast-redos-return.md

@@ -0,0 +1,5 @@
+---
+"kilo-code": patch
+---
+
+Restore reverted sessions on the first Redo click.

packages/kilo-vscode/src/KiloProvider.ts

@@ -196,6 +196,7 @@ type LegacySyncEvent =
       properties: Extract<SyncPayload, { name: "session.created.1" }>["data"]
     }
   | {
+      source: "sync"
       id: string
       type: "session.updated"
       properties: Extract<SyncPayload, { name: "session.updated.1" }>["data"]
@@ -206,20 +207,30 @@ type LegacySyncEvent =
       properties: Extract<SyncPayload, { name: "session.deleted.1" }>["data"]
     }
 
-type ProviderEvent = Event | LegacySyncEvent
+type FullSessionUpdatedEvent = {
+  id: string
+  type: "session.updated"
+  properties: { sessionID: string; info: Session }
+}
+
+type ProviderEvent = Event | LegacySyncEvent | FullSessionUpdatedEvent
 
 function isLegacySyncEvent(event: ProviderEvent): event is LegacySyncEvent {
+  if (event.type === "session.updated") return "source" in event && event.source === "sync"
   return (
     event.type === "message.updated" ||
     event.type === "message.removed" ||
     event.type === "message.part.updated" ||
     event.type === "message.part.removed" ||
     event.type === "session.created" ||
-    event.type === "session.updated" ||
     event.type === "session.deleted"
   )
 }
 
+function isFullSessionUpdatedEvent(event: ProviderEvent): event is FullSessionUpdatedEvent {
+  return event.type === "session.updated" && !isLegacySyncEvent(event)
+}
+
 function unwrapSyncEvent(event: GlobalEvent["payload"]): ProviderEvent | undefined {
   if (event.type !== "sync") return event
 
@@ -235,7 +246,7 @@ function unwrapSyncEvent(event: GlobalEvent["payload"]): ProviderEvent | undefin
     case "session.created.1":
       return { id: event.id, type: "session.created", properties: event.data }
     case "session.updated.1":
-      return { id: event.id, type: "session.updated", properties: event.data }
+      return { source: "sync", id: event.id, type: "session.updated", properties: event.data }
     case "session.deleted.1":
       return { id: event.id, type: "session.deleted", properties: event.data }
     default:
@@ -3102,7 +3113,8 @@ export class KiloProvider implements vscode.WebviewViewProvider, TelemetryProper
     // Drop session events from other projects before any tracking logic.
     // This must come first: the trackedSessionIds guard below would otherwise
     // let a foreign session through if it was accidentally tracked.
-    if (!isLegacySyncEvent(event) && isEventFromForeignProject(event, this.projectID)) return
+    if (!isLegacySyncEvent(event) && !isFullSessionUpdatedEvent(event) && isEventFromForeignProject(event, this.projectID))
+      return
     if (
       this.projectID &&
       (event.type === "session.created" || event.type === "session.updated") &&
@@ -3183,7 +3195,10 @@ export class KiloProvider implements vscode.WebviewViewProvider, TelemetryProper
       this.trackedSessionIds.add(event.properties.info.id)
     }
     if (event.type === "session.updated" && this.currentSession?.id === event.properties.sessionID) {
-      this.setCurrentSession(applySessionPatch(this.currentSession, event.properties.info))
+      const session = isLegacySyncEvent(event)
+        ? applySessionPatch(this.currentSession, event.properties.info)
+        : event.properties.info
+      this.setCurrentSession(session)
       this.contextSessionID = event.properties.sessionID
     }
 
@@ -3226,7 +3241,9 @@ export class KiloProvider implements vscode.WebviewViewProvider, TelemetryProper
 
     const msg = isLegacySyncEvent(event)
       ? this.mapSyncEventToWebviewMessage(event)
-      : mapSSEEventToWebviewMessage(event, sessionID)
+      : isFullSessionUpdatedEvent(event)
+        ? { type: "sessionUpdated" as const, session: this.sessionToWebview(event.properties.info) }
+        : mapSSEEventToWebviewMessage(event, sessionID)
     if (!msg) return
     if (msg.type === "partUpdated") {
       this.streams.push({ ...msg, part: this.slimPart(msg.part) })

packages/kilo-vscode/tests/unit/kilo-provider-utils.test.ts

@@ -158,6 +158,12 @@ describe("sessionToWebview", () => {
     expect(() => new Date(result.createdAt)).not.toThrow()
     expect(new Date(result.createdAt).getTime()).toBe(1700000000000)
   })
+
+  it("clears optional state omitted from a full session snapshot", () => {
+    const result = sessionToWebview(makeSession())
+    expect(result.revert).toBeNull()
+    expect(result.summary).toBeNull()
+  })
 })
 
 describe("applySessionPatch", () => {

packages/kilo-vscode/tests/unit/revert-checkpoints.test.ts

@@ -4,8 +4,18 @@ import path from "node:path"
 
 const ROOT = path.resolve(import.meta.dir, "../..")
 const TURN_FILE = path.join(ROOT, "webview-ui/src/components/chat/VscodeSessionTurn.tsx")
+const PROVIDER_FILE = path.join(ROOT, "src/KiloProvider.ts")
 
 const src = fs.readFileSync(TURN_FILE, "utf-8")
+const provider = fs.readFileSync(PROVIDER_FILE, "utf-8")
+
+function method(name: string, next: string) {
+  const start = provider.indexOf(`  private async ${name}`)
+  const end = provider.indexOf(`  private async ${next}`, start)
+  expect(start).toBeGreaterThan(-1)
+  expect(end).toBeGreaterThan(start)
+  return provider.slice(start, end)
+}
 
 describe("message revert checkpoints", () => {
   it("keeps revert actions available after a session is already reverted", () => {
@@ -18,3 +28,28 @@ describe("message revert checkpoints", () => {
     expect(src).not.toMatch(/data-revert-disabled=\{[\s\S]*?!session\.revert\(\)/)
   })
 })
+
+describe("revert session synchronization", () => {
+  it("keeps REST responses as the mutation result", () => {
+    const revert = method("handleRevertSession", "handleUnrevertSession")
+    const unrevert = method("handleUnrevertSession", "handleCompact")
+
+    expect(revert).toContain("await this.client.session.revert")
+    expect(unrevert).toContain("await this.client.session.unrevert")
+    expect(revert).toContain('type: "sessionUpdated"')
+    expect(unrevert).toContain('type: "sessionUpdated"')
+  })
+
+  it("distinguishes partial sync patches from full bus snapshots", () => {
+    expect(provider).toMatch(/source: "sync"/)
+    expect(provider).toMatch(
+      /if \(event\.type === "session\.updated"\) return "source" in event && event\.source === "sync"/,
+    )
+    expect(provider).toMatch(
+      /isLegacySyncEvent\(event\)\s*\? applySessionPatch\(this\.currentSession, event\.properties\.info\)\s*:\s*event\.properties\.info/,
+    )
+    expect(provider).toMatch(
+      /isFullSessionUpdatedEvent\(event\)\s*\? \{ type: "sessionUpdated" as const, session: this\.sessionToWebview\(event\.properties\.info\) \}/,
+    )
+  })
+})