fix(agent-manager): keep chat in sync with session selection while offline

.changeset/fix-agent-manager-session-switch-offline.md

@@ -0,0 +1,5 @@
+---
+"kilo-code": patch
+---
+
+Keep the chat in sync with the selected Agent Manager session when the backend connection is briefly unavailable, so switching sessions no longer updates only the side diff while the conversation stays on the previous session.

packages/kilo-vscode/tests/unit/session-select-connection.test.ts

@@ -0,0 +1,72 @@
+/**
+ * Source contract test for selectSession's connection handling.
+ *
+ * Static analysis — reads session.tsx and verifies that selectSession updates
+ * the current session id BEFORE (and independently of) the backend connection
+ * check, and only defers the message fetch when offline.
+ *
+ * Regression guard: previously selectSession bailed out entirely when
+ * `server.isConnected()` was false. In Agent Manager the side diff is resolved
+ * from the worktree selection independently of currentSessionID, so a switch
+ * during a transient disconnect moved the diff but left the chat frozen on the
+ * previous session ("switching only changes the sidebar diff"). The chat must
+ * always follow the selection; only the network fetch may wait for reconnect.
+ */
+
+import { describe, it, expect } from "bun:test"
+import fs from "node:fs"
+import path from "node:path"
+
+const ROOT = path.resolve(import.meta.dir, "../..")
+const SESSION_FILE = path.join(ROOT, "webview-ui/src/context/session.tsx")
+
+const source = fs.readFileSync(SESSION_FILE, "utf-8")
+
+describe("selectSession keeps the chat in sync with the selection while offline", () => {
+  const start = source.indexOf("function selectSession(")
+  const cloudGuard = source.indexOf('id.startsWith("cloud:")', start)
+  const setCurrent = source.indexOf("setCurrentSessionID(id)", start)
+  const offlineDefer = source.indexOf("if (!server.isConnected()) {", start)
+
+  it("selectSession exists", () => {
+    expect(start).toBeGreaterThan(-1)
+  })
+
+  it("returns early for cloud preview ids before touching the current session", () => {
+    expect(cloudGuard).toBeGreaterThan(start)
+    expect(cloudGuard).toBeLessThan(setCurrent)
+  })
+
+  it("sets currentSessionID before checking the connection (chat follows selection offline)", () => {
+    expect(setCurrent).toBeGreaterThan(-1)
+    expect(offlineDefer).toBeGreaterThan(-1)
+    // The whole point of the fix: the local selection update must precede the
+    // connection guard, so a disconnected switch no longer freezes the chat.
+    expect(setCurrent).toBeLessThan(offlineDefer)
+  })
+
+  it("defers the fetch for any session while offline, including cached ones", () => {
+    const body = source.slice(start, source.indexOf("\n  function loadFocusedMessages("))
+    // Queue a replay unconditionally. The earlier `deferredFetch = ready ? undefined : id`
+    // form skipped cached sessions, so a reconnect never re-sent the focus load that
+    // re-focuses the backend (focusSession/contextSessionID/SSE tracking/reconcile).
+    expect(body).toContain("deferredFetch = id")
+    expect(body).not.toMatch(/deferredFetch\s*=\s*ready\s*\?/)
+  })
+})
+
+describe("a deferred fetch is replayed on reconnect", () => {
+  it("watches the connection and replays the deferred session load", () => {
+    expect(source).toContain("on(server.isConnected")
+    const effect = source.slice(source.indexOf("on(server.isConnected"))
+    expect(effect).toContain("deferredFetch")
+    // Replays with the focus/replace choice so cached sessions still re-focus the backend.
+    expect(effect).toMatch(/loadFocusedMessages\(\s*id,\s*loaded\(\)\.has\(id\)\s*\)/)
+  })
+
+  it("the focused load helper sends focus for cached sessions and replace otherwise", () => {
+    const helper = source.slice(source.indexOf("function loadFocusedMessages("))
+    expect(helper).toMatch(/mode: "focus"/)
+    expect(helper).toMatch(/mode: "replace"/)
+  })
+})

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

@@ -10,6 +10,7 @@ import {
   createSignal,
   createMemo,
   createEffect,
+  on,
   onMount,
   onCleanup,
   batch,
@@ -2519,27 +2520,61 @@ export const SessionProvider: ParentComponent = (props) => {
     })
   }
 
+  // Session whose message fetch was deferred because the backend was offline at
+  // selection time. Replayed by the reconnect effect below.
+  let deferredFetch: string | undefined
+
   function selectSession(id: string) {
-    if (!server.isConnected()) {
-      console.warn("[Kilo New] Cannot select session: not connected")
-      return
-    }
+    // Cloud preview sessions use a separate keyed path (selectCloudSession).
     if (id.startsWith("cloud:")) {
       console.warn("[Kilo New] Cannot select cloud preview session via selectSession")
       return
     }
     const ready = loaded().has(id)
+    // Reflect the selection locally and synchronously so the chat always tracks
+    // the sidebar/tab selection. These are local signals and need no backend, so
+    // they update even while disconnected. Bailing out here when not connected
+    // froze the chat on the previous session while the side diff (resolved from
+    // the worktree selection) still moved (the reported "only the diff changes").
     setCurrentSessionID(id)
     setDraftSessionID(id)
     setLoading(!ready)
-    if (ready) {
-      vscode.postMessage({ type: "loadMessages", sessionID: id, mode: "focus" })
+    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
+    // the reconnect effect replay it. We defer even for cached sessions: the
+    // load message is what re-focuses the backend (focusSession, contextSessionID,
+    // SSE tracking, active worktree) and runs the reconcile self-heal, so skipping
+    // it would leave the extension focused on the previously selected session.
+    if (!server.isConnected()) {
+      deferredFetch = id
       return
     }
-    patchPage(id, { loadingInitial: true, loadingOlder: false, before: undefined, hasMore: false })
-    vscode.postMessage({ type: "loadMessages", sessionID: id, mode: "replace", limit: MESSAGE_PAGE_LIMIT })
+    deferredFetch = undefined
+    loadFocusedMessages(id, ready)
   }
 
+  function loadFocusedMessages(id: string, ready: boolean) {
+    vscode.postMessage(
+      ready
+        ? { type: "loadMessages", sessionID: id, mode: "focus" }
+        : { type: "loadMessages", sessionID: id, mode: "replace", limit: MESSAGE_PAGE_LIMIT },
+    )
+  }
+
+  // Replay a fetch deferred while offline once the backend reconnects. Scoped to
+  // the still-current session so the normal connected path never double-fetches.
+  // Uses the same focus/replace choice as a live selection so a reconnect after
+  // a cached-session switch still re-focuses the backend and reconciles.
+  createEffect(
+    on(server.isConnected, (connected) => {
+      if (!connected) return
+      const id = deferredFetch
+      deferredFetch = undefined
+      if (!id || id !== currentSessionID()) return
+      loadFocusedMessages(id, loaded().has(id))
+    }),
+  )
+
   function selectCloudSession(cloudSessionId: string) {
     if (!server.isConnected()) {
       console.warn("[Kilo New] Cannot select cloud session: not connected")