fix(agent-manager): propagate base branch override to active diff source

.changeset/agent-manager-diff-base-override.md

@@ -0,0 +1,5 @@
+---
+"kilo-code": patch
+---
+
+Apply the Agent Manager base branch picker selection to the active diff immediately. Changing the base branch now refreshes the diff against the new base instead of keeping the previous comparison until the scope or session changed.

packages/kilo-vscode/src/agent-manager/AgentManagerProvider.ts

@@ -724,9 +724,10 @@ export class AgentManagerProvider implements Disposable {
       return null
     }
     if (m.type === "agentManager.setDiffBaseBranch") {
-      void this.diffs.setBase(composeDiffId(m.sessionId, normalizeScope(m.scope)), m.branch).then(() => {
-        void this.sendDiffBranches(m.sessionId, m.scope)
-      })
+      void this.diffs
+        .setBase(composeDiffId(m.sessionId, normalizeScope(m.scope)), m.branch)
+        .catch((err) => this.log("Failed to set diff base:", err instanceof Error ? err.message : String(err)))
+        .then(() => void this.sendDiffBranches(m.sessionId, m.scope))
       return null
     }
     if (m.type === "agentManager.openFile") {

packages/kilo-vscode/src/agent-manager/worktree-diff-controller.ts

@@ -195,8 +195,13 @@ export class WorktreeDiffController {
     const { ctx } = parseDiffId(id)
     if (branch) this.baseOverrides.set(ctx, branch)
     else this.baseOverrides.delete(ctx)
-    this.target = undefined
-    await this.controller.reactivate()
+    // Nothing to rebuild when the context isn't active; the override is
+    // picked up the next time start()/request() resolves it.
+    if (this.controller.currentId !== id) return
+    // Route through activate() so the base is re-resolved and pushed via
+    // setContext() — SourceController.reactivate() alone would rebuild the
+    // source against the stale context captured by the last activate().
+    await this.activate(id, this.controller.isPolling, true)
   }
 
   /** Branch picker data for a context's directory, using any active override. */

packages/kilo-vscode/tests/unit/worktree-diff-controller.test.ts

@@ -0,0 +1,89 @@
+import { describe, it, expect } from "bun:test"
+import { WorktreeDiffController } from "../../src/agent-manager/worktree-diff-controller"
+import type { DiffSourceCatalog } from "../../src/diff/sources/catalog"
+import type { DiffSource } from "../../src/diff/sources/types"
+import type { PanelContext } from "../../src/diff/types"
+import type { GitOps } from "../../src/agent-manager/GitOps"
+import type { WorktreeStateManager } from "../../src/agent-manager/WorktreeStateManager"
+
+// Records every PanelContext handed to catalog.build so tests can assert which
+// base branch the active source was (re)built with. The controller, scope
+// resolution, and SourceController lifecycle under test are all real.
+function make() {
+  const builds: { id: string; ctx: PanelContext }[] = []
+  const catalog = {
+    build: (id: string, ctx: PanelContext): DiffSource => {
+      builds.push({ id, ctx })
+      return {
+        descriptor: { id, type: "workspace", group: "Git", capabilities: { revert: true, comments: true } },
+        async fetch() {
+          return { diffs: [] }
+        },
+      }
+    },
+  } as unknown as DiffSourceCatalog
+
+  const state = {
+    getSession: (id: string) => (id === "s1" ? { id: "s1", worktreeId: "w1", createdAt: "" } : undefined),
+    getWorktree: (id: string) =>
+      id === "w1" ? { id: "w1", path: "/wt", parentBranch: "main", remote: "origin" } : undefined,
+  } as unknown as WorktreeStateManager
+
+  const controller = new WorktreeDiffController({
+    getState: () => state,
+    getRoot: () => "/repo",
+    getStateReady: () => undefined,
+    catalog,
+    git: {} as GitOps,
+    localDiffFile: async () => null,
+    post: () => {},
+    log: () => {},
+  })
+  return { controller, builds }
+}
+
+const tick = () => new Promise((resolve) => setTimeout(resolve, 0))
+
+async function waitFor(cond: () => boolean): Promise<void> {
+  for (let i = 0; i < 50; i++) {
+    if (cond()) return
+    await tick()
+  }
+  throw new Error("waitFor timed out")
+}
+
+describe("WorktreeDiffController.setBase", () => {
+  it("rebuilds the active source against the overridden base branch", async () => {
+    const { controller, builds } = make()
+    controller.start("s1#branch")
+    await waitFor(() => builds.length === 1)
+    expect(builds[0]!.ctx.dir).toBe("/wt")
+    expect(builds[0]!.ctx.baseBranch).toBe("origin/main")
+
+    await controller.setBase("s1#branch", "feature-x")
+    expect(builds.length).toBe(2)
+    expect(builds[1]!.ctx.dir).toBe("/wt")
+    expect(builds[1]!.ctx.baseBranch).toBe("feature-x")
+
+    // Clearing the override falls back to the recorded parent ref.
+    await controller.setBase("s1#branch", undefined)
+    expect(builds.length).toBe(3)
+    expect(builds[2]!.ctx.baseBranch).toBe("origin/main")
+
+    controller.stop()
+  })
+
+  it("stores the override without rebuilding when the context isn't active", async () => {
+    const { controller, builds } = make()
+
+    await controller.setBase("s1#branch", "feature-x")
+    expect(builds.length).toBe(0)
+
+    // The next activation of that context resolves the stored override.
+    controller.start("s1#branch")
+    await waitFor(() => builds.length === 1)
+    expect(builds[0]!.ctx.baseBranch).toBe("feature-x")
+
+    controller.stop()
+  })
+})