fix: lazy load chat diffs

packages/kilo-vscode/src/KiloProvider.ts

@@ -622,6 +622,9 @@ export class KiloProvider implements vscode.WebviewViewProvider, TelemetryProper
           // isn't blocked by slow responses for earlier sessions.
           void this.handleLoadMessages(message.sessionID)
           break
+        case "requestSessionMessageDiff":
+          void this.handleLoadSessionMessageDiff(message.sessionID, message.messageID)
+          break
         case "syncSession":
           this.handleSyncSession(message.sessionID, message.parentSessionID).catch((e) =>
             console.error("[Kilo New] handleSyncSession failed:", e),
@@ -1369,6 +1372,38 @@ export class KiloProvider implements vscode.WebviewViewProvider, TelemetryProper
     }
   }
 
+  private async handleLoadSessionMessageDiff(sessionID: string, messageID: string): Promise<void> {
+    if (!this.client) {
+      this.postMessage({
+        type: "sessionMessageDiffError",
+        sessionID,
+        messageID,
+        error: "Not connected to CLI backend",
+      })
+      return
+    }
+
+    try {
+      const workspaceDir = this.getWorkspaceDirectory(sessionID)
+      const { data } = await retry(() =>
+        this.client!.session.diff({ sessionID, messageID, directory: workspaceDir }, { throwOnError: true }),
+      )
+      this.postMessage({
+        type: "sessionMessageDiffLoaded",
+        sessionID,
+        messageID,
+        diffs: data ?? [],
+      })
+    } catch (error) {
+      this.postMessage({
+        type: "sessionMessageDiffError",
+        sessionID,
+        messageID,
+        error: getErrorMessage(error) || "Failed to load message diff",
+      })
+    }
+  }
+
   /**
    * Handle syncing a child session (e.g. spawned by the task tool).
    * Tracks the session for SSE events and fetches its messages.

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

@@ -18,6 +18,7 @@ export function resolveEventSessionId(
     case "session.status":
     case "session.idle":
     case "session.error":
+    case "session.diff":
     case "todo.updated":
       return event.properties.sessionID
     case "message.updated":

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

@@ -46,6 +46,14 @@ describe("resolveEventSessionId", () => {
     expect(resolveEventSessionId(e, noLookup)).toBe("s4")
   })
 
+  it("returns sessionID from session.diff", () => {
+    const e = event({
+      type: "session.diff",
+      properties: { sessionID: "s5", diff: [] },
+    })
+    expect(resolveEventSessionId(e, noLookup)).toBe("s5")
+  })
+
   it("returns sessionID from message.updated and calls onMessageUpdated", () => {
     const e = event({
       type: "message.updated",

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

@@ -9,7 +9,7 @@
  * - Simpler flat structure without overflow containers
  */
 
-import { Component, createMemo, For, Show, createSignal, createEffect, on } from "solid-js"
+import { Component, createMemo, For, Show, createSignal, createEffect, on, onCleanup } from "solid-js"
 import { Dynamic } from "solid-js/web"
 import { UserMessageDisplay } from "@kilocode/kilo-ui/message-part"
 import { Collapsible } from "@kilocode/kilo-ui/collapsible"
@@ -121,12 +121,51 @@ export const VscodeSessionTurn: Component<VscodeSessionTurnProps> = (props) => {
 
   const [open, setOpen] = createSignal(false)
   const [expanded, setExpanded] = createSignal<string[]>([])
+  const [full, setFull] = createSignal<FileDiff[] | undefined>()
+  const [diffLoading, setDiffLoading] = createSignal(false)
+  const [diffError, setDiffError] = createSignal<string | undefined>()
+  let seq = 0
+
+  const fullByFile = createMemo(() => new Map((full() ?? []).map((diff) => [diff.file, diff])))
+
+  function loadDiffs() {
+    if (full() || diffLoading()) return
+    const token = ++seq
+    setDiffLoading(true)
+    setDiffError(undefined)
+    session
+      .loadMessageDiff(props.sessionID, props.messageID)
+      .then((diffs) => {
+        if (token !== seq || !open()) return
+        setFull(diffs as FileDiff[])
+      })
+      .catch((err: unknown) => {
+        if (token !== seq || !open()) return
+        setDiffError(err instanceof Error ? err.message : String(err))
+      })
+      .finally(() => {
+        if (token === seq) setDiffLoading(false)
+      })
+  }
+
+  onCleanup(() => {
+    seq++
+  })
 
   createEffect(
     on(
       open,
       (value, prev) => {
-        if (!value && prev) setExpanded([])
+        if (value) {
+          loadDiffs()
+          return
+        }
+        if (!prev) return
+        seq++
+        setExpanded([])
+        setFull(undefined)
+        setDiffError(undefined)
+        setDiffLoading(false)
       },
       { defer: true },
     ),
@@ -219,6 +258,7 @@ export const VscodeSessionTurn: Component<VscodeSessionTurnProps> = (props) => {
                         <For each={diffs()}>
                           {(diff) => {
                             const active = createMemo(() => expanded().includes(diff.file))
+                            const loaded = createMemo(() => fullByFile().get(diff.file))
                             const [visible, setVisible] = createSignal(false)
 
                             createEffect(
@@ -263,13 +303,26 @@ export const VscodeSessionTurn: Component<VscodeSessionTurnProps> = (props) => {
                                 </StickyAccordionHeader>
                                 <Accordion.Content>
                                   <Show when={visible()}>
-                                    <div data-slot="session-turn-diff-view" data-scrollable>
-                                      <Dynamic
-                                        component={diffComponent}
-                                        before={{ name: diff.file, contents: diff.before }}
-                                        after={{ name: diff.file, contents: diff.after }}
-                                      />
-                                    </div>
+                                    <Show
+                                      when={loaded()}
+                                      fallback={
+                                        <div data-slot="session-turn-diff-placeholder">
+                                          {diffLoading()
+                                            ? language.t("session.review.loadingChanges")
+                                            : diffError() || language.t("session.review.noChanges")}
+                                        </div>
+                                      }
+                                    >
+                                      {(item) => (
+                                        <div data-slot="session-turn-diff-view" data-scrollable>
+                                          <Dynamic
+                                            component={diffComponent}
+                                            before={{ name: item().file, contents: item().before }}
+                                            after={{ name: item().file, contents: item().after }}
+                                          />
+                                        </div>
+                                      )}
+                                    </Show>
                                   </Show>
                                 </Accordion.Content>
                               </Accordion.Item>

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

@@ -15,6 +15,7 @@ import { useLanguage } from "./language"
 import { showToast } from "@kilocode/kilo-ui/toast"
 import type {
   SessionInfo,
+  SessionFileDiff,
   Message,
   Part,
   PartDelta,
@@ -62,6 +63,21 @@ interface SessionStore {
   favoriteModels: ModelSelection[]
 }
 
+interface DiffRequest {
+  promise: Promise<SessionFileDiff[]>
+  resolve: (diffs: SessionFileDiff[]) => void
+  reject: (err: Error) => void
+}
+
+function deferred(): DiffRequest {
+  const req = {} as DiffRequest
+  req.promise = new Promise<SessionFileDiff[]>((resolve, reject) => {
+    req.resolve = resolve
+    req.reject = reject
+  })
+  return req
+}
+
 interface SessionContextValue {
   // Current session
   currentSessionID: Accessor<string | undefined>
@@ -194,6 +210,7 @@ interface SessionContextValue {
   createSession: () => void
   clearCurrentSession: () => void
   loadSessions: () => void
+  loadMessageDiff: (sessionID: string, messageID: string) => Promise<SessionFileDiff[]>
   selectSession: (id: string) => void
   deleteSession: (id: string) => void
   renameSession: (id: string, title: string) => void
@@ -323,6 +340,27 @@ export const SessionProvider: ParentComponent = (props) => {
   // Prevents handleMessagesLoaded from wiping them when it replaces the array.
   const pendingOptimistic = new Map<string, Set<string>>()
 
+  const diffRequests = new Map<string, DiffRequest>()
+  const diffKey = (sessionID: string, messageID: string) => `${sessionID}\u0000${messageID}`
+
+  function loadMessageDiff(sessionID: string, messageID: string): Promise<SessionFileDiff[]> {
+    const key = diffKey(sessionID, messageID)
+    const hit = diffRequests.get(key)
+    if (hit) return hit.promise
+
+    const req = deferred()
+    diffRequests.set(key, req)
+    vscode.postMessage({ type: "requestSessionMessageDiff", sessionID, messageID })
+    return req.promise
+  }
+
+  onCleanup(() => {
+    for (const req of diffRequests.values()) {
+      req.reject(new Error("Session provider disposed"))
+    }
+    diffRequests.clear()
+  })
+
   // Store for sessions, messages, parts, todos, modelSelections, agentSelections
   const [store, setStore] = createStore<SessionStore>({
     sessions: {},
@@ -640,6 +678,26 @@ export const SessionProvider: ParentComponent = (props) => {
     vscode.postMessage({ type: "toggleFavorite", action, providerID, modelID })
   }
 
+  function handleDiffMessage(message: ExtensionMessage): void {
+    if (message.type === "sessionMessageDiffLoaded") {
+      const key = diffKey(message.sessionID, message.messageID)
+      const req = diffRequests.get(key)
+      diffRequests.delete(key)
+      req?.resolve(message.diffs)
+      return
+    }
+
+    if (message.type === "sessionMessageDiffError") {
+      const key = diffKey(message.sessionID, message.messageID)
+      const req = diffRequests.get(key)
+      diffRequests.delete(key)
+      req?.reject(new Error(message.error))
+    }
+  }
+
+  const unsubDiffs = vscode.onMessage(handleDiffMessage)
+  onCleanup(unsubDiffs)
+
   // Handle messages from extension
   onMount(() => {
     const unsubscribe = vscode.onMessage((message: ExtensionMessage) => {
@@ -1876,6 +1934,7 @@ export const SessionProvider: ParentComponent = (props) => {
     createSession,
     clearCurrentSession,
     loadSessions,
+    loadMessageDiff,
     selectSession,
     deleteSession,
     renameSession,

packages/kilo-vscode/webview-ui/src/types/messages.ts

@@ -560,6 +560,20 @@ export interface MessageCreatedMessage {
   message: Message
 }
 
+export interface SessionMessageDiffLoadedMessage {
+  type: "sessionMessageDiffLoaded"
+  sessionID: string
+  messageID: string
+  diffs: SessionFileDiff[]
+}
+
+export interface SessionMessageDiffErrorMessage {
+  type: "sessionMessageDiffError"
+  sessionID: string
+  messageID: string
+  error: string
+}
+
 export interface SessionsLoadedMessage {
   type: "sessionsLoaded"
   sessions: SessionInfo[]
@@ -1419,6 +1433,8 @@ export type ExtensionMessage =
   | MessageRemovedMessage
   | MessagesLoadedMessage
   | MessageCreatedMessage
+  | SessionMessageDiffLoadedMessage
+  | SessionMessageDiffErrorMessage
   | SessionsLoadedMessage
   | CloudSessionsLoadedMessage
   | GitRemoteUrlLoadedMessage
@@ -1569,6 +1585,12 @@ export interface LoadMessagesRequest {
   sessionID: string
 }
 
+export interface RequestSessionMessageDiffRequest {
+  type: "requestSessionMessageDiff"
+  sessionID: string
+  messageID: string
+}
+
 export interface LoadSessionsRequest {
   type: "loadSessions"
 }
@@ -2334,6 +2356,7 @@ export type WebviewMessage =
   | CreateSessionRequest
   | ClearSessionRequest
   | LoadMessagesRequest
+  | RequestSessionMessageDiffRequest
   | LoadSessionsRequest
   | RequestCloudSessionsMessage
   | RequestGitRemoteUrlMessage

packages/opencode/src/session/index.ts

@@ -496,32 +496,35 @@ export namespace Session {
       const updateMessage = <T extends MessageV2.Info>(msg: T): Effect.Effect<T> =>
         Effect.gen(function* () {
           // kilocode_change start - ignore FK errors when session was deleted while processor was still running
+          const info = MessageV2.stripMessageMetadata(msg) as T
           yield* Effect.sync(() =>
-            KiloSession.runSyncSafe(
-              () => SyncEvent.run(MessageV2.Event.Updated, { sessionID: msg.sessionID, info: msg }),
-              { type: "message update", id: msg.id, sessionID: msg.sessionID },
-            ),
+            KiloSession.runSyncSafe(() => SyncEvent.run(MessageV2.Event.Updated, { sessionID: info.sessionID, info }), {
+              type: "message update",
+              id: info.id,
+              sessionID: info.sessionID,
+            }),
           )
+          return info
           // kilocode_change end
-          return msg
         }).pipe(Effect.withSpan("Session.updateMessage"))
 
       const updatePart = <T extends MessageV2.Part>(part: T): Effect.Effect<T> =>
         Effect.gen(function* () {
-          // kilocode_change start - ignore FK errors when session was deleted while processor was still running
+          // kilocode_change start - strip bulky metadata and ignore FK errors when session was deleted while processor was still running
+          const info = MessageV2.stripPartMetadata(part) as T
           yield* Effect.sync(() =>
             KiloSession.runSyncSafe(
               () =>
                 SyncEvent.run(MessageV2.Event.PartUpdated, {
-                  sessionID: part.sessionID,
-                  part: structuredClone(part),
+                  sessionID: info.sessionID,
+                  part: structuredClone(info),
                   time: Date.now(),
                 }),
-              { type: "part update", id: part.id, sessionID: part.sessionID },
+              { type: "part update", id: info.id, sessionID: info.sessionID },
             ),
           )
+          return info
           // kilocode_change end
-          return part
         }).pipe(Effect.withSpan("Session.updatePart"))
 
       const create = Effect.fn("Session.create")(function* (input?: {

packages/opencode/src/session/revert.ts

@@ -67,7 +67,7 @@ export namespace SessionRevert {
       await Storage.write(["session_diff", input.sessionID], diffs)
       Bus.publish(Session.Event.Diff, {
         sessionID: input.sessionID,
-        diff: diffs,
+        diff: SessionSummary.slim(diffs), // kilocode_change
       })
       // kilocode_change start - strip full file contents before persisting to DB
       const summaryDiffs = diffs.map((d) => ({