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

@@ -34,6 +34,8 @@ export class WorktreeDiffController {
   private readonly controller: SourceController
   private target: Target | undefined
   private applying: string | undefined
+  /** Intended watch mode for the active context; isPolling lags the initial fetch. */
+  private poll = false
   /** Ephemeral per-context base override, keyed by context id. */
   private baseOverrides = new Map<string, string>()
 
@@ -184,6 +186,7 @@ export class WorktreeDiffController {
   public stop(): void {
     this.controller.stop()
     this.target = undefined
+    this.poll = false
   }
 
   /**
@@ -195,8 +198,15 @@ 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(). The
+    // recorded poll intent preserves watch mode even when the initial fetch
+    // is still in flight (isPolling only turns true once it resolves).
+    await this.activate(id, this.poll, true)
   }
 
   /** Branch picker data for a context's directory, using any active override. */
@@ -210,6 +220,7 @@ export class WorktreeDiffController {
 
   private async activate(id: string, poll: boolean, fetch: boolean): Promise<void> {
     this.target = undefined
+    this.poll = poll
     await this.ready("stateReady rejected, continuing diff activate:")
     const { ctx } = parseDiffId(id)
     const resolved = await this.resolve(ctx)

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

@@ -0,0 +1,119 @@
+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(onFetch?: (n: number) => Promise<void>) {
+  const builds: { id: string; ctx: PanelContext }[] = []
+  let fetches = 0
+  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() {
+          await onFetch?.(++fetches)
+          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()
+  })
+
+  it("keeps watching when the base changes during the initial fetch", async () => {
+    // Hold the first activation's fetch in flight, simulating a slow worktree
+    // diff. isPolling is still false in this window, but the watch intent must
+    // survive the base change rather than downgrading the panel to one-shot.
+    let release: () => void = () => {}
+    const gate = new Promise<void>((resolve) => (release = resolve))
+    const { controller, builds } = make(async (n) => {
+      if (n === 1) await gate
+    })
+
+    controller.start("s1#branch")
+    await waitFor(() => builds.length === 1)
+
+    const change = controller.setBase("s1#branch", "feature-x")
+    release()
+    await change
+    expect(builds.length).toBe(2)
+    expect(builds[1]!.ctx.baseBranch).toBe("feature-x")
+
+    // Polling survives: start() early-returns for an id that is already
+    // watched. A downgraded one-shot panel would re-activate and rebuild here.
+    controller.start("s1#branch")
+    await tick()
+    expect(builds.length).toBe(2)
+
+    controller.stop()
+  })
+})