From a2b22d07c2b773dcbcc71b8e6999a9bd1c118b91 Mon Sep 17 00:00:00 2001 From: marius-kilocode Date: Thu, 30 Jul 2026 15:04:42 +0200 Subject: [PATCH 1/2] fix(agent-manager): propagate base branch override to active diff source --- .../agent-manager-diff-base-override.md | 5 ++ .../src/agent-manager/AgentManagerProvider.ts | 7 +- .../agent-manager/worktree-diff-controller.ts | 9 +- .../unit/worktree-diff-controller.test.ts | 89 +++++++++++++++++++ 4 files changed, 105 insertions(+), 5 deletions(-) create mode 100644 .changeset/agent-manager-diff-base-override.md create mode 100644 packages/kilo-vscode/tests/unit/worktree-diff-controller.test.ts diff --git a/.changeset/agent-manager-diff-base-override.md b/.changeset/agent-manager-diff-base-override.md new file mode 100644 index 00000000000..5d38f0e6cbc --- /dev/null +++ b/.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. diff --git a/packages/kilo-vscode/src/agent-manager/AgentManagerProvider.ts b/packages/kilo-vscode/src/agent-manager/AgentManagerProvider.ts index a7cdd033cdc..5dd9b1987cd 100644 --- a/packages/kilo-vscode/src/agent-manager/AgentManagerProvider.ts +++ b/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") { diff --git a/packages/kilo-vscode/src/agent-manager/worktree-diff-controller.ts b/packages/kilo-vscode/src/agent-manager/worktree-diff-controller.ts index dff233ca5cc..2883ae77afc 100644 --- a/packages/kilo-vscode/src/agent-manager/worktree-diff-controller.ts +++ b/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. */ diff --git a/packages/kilo-vscode/tests/unit/worktree-diff-controller.test.ts b/packages/kilo-vscode/tests/unit/worktree-diff-controller.test.ts new file mode 100644 index 00000000000..b60f5b983f0 --- /dev/null +++ b/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 { + 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() + }) +}) From ff7fdf65b2567f2c8559abf20596ab192ea2f6c0 Mon Sep 17 00:00:00 2001 From: marius-kilocode Date: Thu, 30 Jul 2026 15:42:46 +0200 Subject: [PATCH 2/2] fix(agent-manager): preserve diff watch mode across base changes --- .../agent-manager/worktree-diff-controller.ts | 10 ++++-- .../unit/worktree-diff-controller.test.ts | 32 ++++++++++++++++++- 2 files changed, 39 insertions(+), 3 deletions(-) diff --git a/packages/kilo-vscode/src/agent-manager/worktree-diff-controller.ts b/packages/kilo-vscode/src/agent-manager/worktree-diff-controller.ts index 2883ae77afc..9a95ae5f96a 100644 --- a/packages/kilo-vscode/src/agent-manager/worktree-diff-controller.ts +++ b/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() @@ -184,6 +186,7 @@ export class WorktreeDiffController { public stop(): void { this.controller.stop() this.target = undefined + this.poll = false } /** @@ -200,8 +203,10 @@ export class WorktreeDiffController { 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) + // 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. */ @@ -215,6 +220,7 @@ export class WorktreeDiffController { private async activate(id: string, poll: boolean, fetch: boolean): Promise { this.target = undefined + this.poll = poll await this.ready("stateReady rejected, continuing diff activate:") const { ctx } = parseDiffId(id) const resolved = await this.resolve(ctx) diff --git a/packages/kilo-vscode/tests/unit/worktree-diff-controller.test.ts b/packages/kilo-vscode/tests/unit/worktree-diff-controller.test.ts index b60f5b983f0..7c3078f7a02 100644 --- a/packages/kilo-vscode/tests/unit/worktree-diff-controller.test.ts +++ b/packages/kilo-vscode/tests/unit/worktree-diff-controller.test.ts @@ -9,14 +9,16 @@ import type { WorktreeStateManager } from "../../src/agent-manager/WorktreeState // 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() { +function make(onFetch?: (n: number) => Promise) { 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: [] } }, } @@ -86,4 +88,32 @@ describe("WorktreeDiffController.setBase", () => { 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((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() + }) })