From 7a1183da91225eec2dd4e46dfa037e66567d8980 Mon Sep 17 00:00:00 2001 From: roboomp Date: Mon, 3 Aug 2026 09:07:12 +0000 Subject: [PATCH 1/3] fix(session): prevented stale /btw branch promotion - Passed the authorized leaf through the branch executor and revalidated it before rewriting history. - Refused promotion during active turns and bounded post-prompt drains. - Consumed unavailable branch keys while showing pending and refusal state in the panel. Fixes #7474 --- packages/coding-agent/CHANGELOG.md | 4 + .../src/modes/components/btw-panel.ts | 45 ++++++- .../src/modes/controllers/btw-controller.ts | 50 +++++-- .../src/modes/controllers/input-controller.ts | 2 +- .../src/modes/interactive-mode.ts | 4 +- packages/coding-agent/src/modes/types.ts | 2 +- .../coding-agent/src/session/agent-session.ts | 33 +++-- .../test/agent-session-btw-branch.test.ts | 127 +++++++++++------- .../test/input-controller-keybindings.test.ts | 26 +++- .../modes/controllers/btw-controller.test.ts | 100 +++++++++++--- 10 files changed, 293 insertions(+), 100 deletions(-) diff --git a/packages/coding-agent/CHANGELOG.md b/packages/coding-agent/CHANGELOG.md index b20591af8..890c853d4 100644 --- a/packages/coding-agent/CHANGELOG.md +++ b/packages/coding-agent/CHANGELOG.md @@ -2,6 +2,10 @@ ## [Unreleased] +### Fixed + +- Fixed `/btw` branch promotion parking behind active turns, cutting from an outdated session leaf, and leaving rejected branch keys indistinguishable from composer input ([#7474](https://github.com/can1357/oh-my-pi/issues/7474)). + ## [17.2.5] - 2026-08-03 ### Breaking Changes diff --git a/packages/coding-agent/src/modes/components/btw-panel.ts b/packages/coding-agent/src/modes/components/btw-panel.ts index 0c8ba38bf..36209a4ed 100644 --- a/packages/coding-agent/src/modes/components/btw-panel.ts +++ b/packages/coding-agent/src/modes/components/btw-panel.ts @@ -3,16 +3,37 @@ import { replaceTabs } from "../../tools/render-utils"; import { getMarkdownTheme, theme } from "../theme/theme"; import { DynamicBorder } from "./dynamic-border"; -type BtwPanelState = "running" | "complete" | "aborted" | "error"; +type BtwPanelState = "running" | "complete" | "branching" | "aborted" | "error"; interface BtwPanelComponentOptions { question: string; tui: TUI; + canBranch?: () => boolean; +} + +class BtwFooter implements Component { + #getLine: () => string; + #line: string | undefined; + #text: Text | undefined; + + constructor(getLine: () => string) { + this.#getLine = getLine; + } + + render(width: number): readonly string[] { + const line = this.#getLine(); + if (line !== this.#line || !this.#text) { + this.#line = line; + this.#text = new Text(line, 1, 0); + } + return this.#text.render(width); + } } export class BtwPanelComponent extends Container { #question: string; #tui: TUI; + #canBranch: (() => boolean) | undefined; #state: BtwPanelState = "running"; #answer = ""; #errorMessage: string | undefined; @@ -23,6 +44,7 @@ export class BtwPanelComponent extends Container { super(); this.#question = options.question; this.#tui = options.tui; + this.#canBranch = options.canBranch; this.#rebuild(); } @@ -47,6 +69,14 @@ export class BtwPanelComponent extends Container { this.#rebuild(); } + /** Shows that the completed answer is being promoted into the chat session. */ + markBranching(): void { + if (this.#closed) return; + this.#state = "branching"; + this.#errorMessage = undefined; + this.#rebuild(); + } + markAborted(): void { if (this.#closed) return; this.#state = "aborted"; @@ -86,7 +116,7 @@ export class BtwPanelComponent extends Container { this.addChild(new Spacer(1)); this.addChild(this.#contentComponent()); this.addChild(new Spacer(1)); - this.addChild(new Text(this.#footerLine(), 1, 0)); + this.addChild(new BtwFooter(() => this.#footerLine())); this.addChild(new Spacer(1)); this.addChild(new DynamicBorder(str => theme.fg("dim", str))); // Component-scoped: a rebuild replaces only this panel's own children @@ -100,8 +130,15 @@ export class BtwPanelComponent extends Container { switch (this.#state) { case "running": return theme.fg("muted", "Esc cancel /btw"); - case "complete": - return theme.fg("muted", this.isCopyable() ? "c copy · b branch to chat · Esc dismiss" : "Esc dismiss"); + case "complete": { + if (!this.isCopyable()) return theme.fg("muted", "Esc dismiss"); + const actions = ["c copy"]; + if (this.#canBranch?.() ?? this.isBranchable()) actions.push("b branch to chat"); + actions.push("Esc dismiss"); + return theme.fg("muted", actions.join(" · ")); + } + case "branching": + return theme.fg("muted", `${theme.status.pending} Branching to chat…`); case "aborted": return theme.fg("warning", `${theme.status.warning} Cancelled · Esc dismiss`); case "error": diff --git a/packages/coding-agent/src/modes/controllers/btw-controller.ts b/packages/coding-agent/src/modes/controllers/btw-controller.ts index 5013a5f99..6e3c79d6f 100644 --- a/packages/coding-agent/src/modes/controllers/btw-controller.ts +++ b/packages/coding-agent/src/modes/controllers/btw-controller.ts @@ -50,15 +50,21 @@ export class BtwController { } canBranch(): boolean { - return ( - !this.#branchInFlight && - this.#activeRequest?.component.isBranchable() === true && - this.#lastQuestion !== undefined && - this.#lastReplyText !== undefined && - this.#lastAssistantMessage !== undefined && - this.#lastLeafId !== null && - this.#lastLeafId === this.ctx.sessionManager.getLeafId() - ); + return this.#branchUnavailableReason() === undefined; + } + + #branchUnavailableReason(): string | undefined { + if (this.#branchInFlight) return "a branch is already in progress"; + if (this.#activeRequest?.component.isBranchable() !== true) return "the answer is not ready"; + if (!this.#lastQuestion || !this.#lastReplyText || !this.#lastAssistantMessage) { + return "the answer is unavailable"; + } + if (!this.#lastLeafId) return "the session has no branch point"; + if (this.#lastLeafId !== this.ctx.sessionManager.getLeafId()) { + return "the session changed since /btw started"; + } + if (this.ctx.session.isStreaming) return "a turn is still running"; + return undefined; } canCopy(): boolean { @@ -83,17 +89,33 @@ export class BtwController { } async handleBranch(): Promise { - if (!this.canBranch() || !this.#lastQuestion || !this.#lastAssistantMessage) return false; + const unavailableReason = this.#branchUnavailableReason(); + if (unavailableReason) { + this.ctx.showStatus(`/btw branch unavailable: ${unavailableReason}`, { dim: true }); + return false; + } + const request = this.#activeRequest; + const question = this.#lastQuestion; + const assistantMessage = this.#lastAssistantMessage; + const leafId = this.#lastLeafId; + if (!request || !question || !assistantMessage || !leafId) return false; + this.#branchInFlight = true; + request.component.markBranching(); try { - await this.ctx.handleBtwBranch(this.#lastQuestion, this.#lastAssistantMessage); + await this.ctx.handleBtwBranch(question, assistantMessage, leafId); return true; } finally { this.#branchInFlight = false; + if (this.#activeRequest === request) request.component.markComplete(); } } handleEscape(): boolean { + if (this.#branchInFlight) { + this.ctx.showStatus("/btw branch is in progress", { dim: true }); + return true; + } if (!this.#activeRequest) return false; this.#closeActiveRequest({ abort: this.#activeRequest.abortController.signal.aborted === false }); return true; @@ -119,7 +141,11 @@ export class BtwController { this.#closeActiveRequest({ abort: true }); const request: BtwRequest = { - component: new BtwPanelComponent({ question: trimmedQuestion, tui: this.ctx.ui }), + component: new BtwPanelComponent({ + question: trimmedQuestion, + tui: this.ctx.ui, + canBranch: () => this.canBranch(), + }), abortController: new AbortController(), question: trimmedQuestion, leafId: this.ctx.sessionManager.getLeafId(), diff --git a/packages/coding-agent/src/modes/controllers/input-controller.ts b/packages/coding-agent/src/modes/controllers/input-controller.ts index d29d2b049..dc265a984 100644 --- a/packages/coding-agent/src/modes/controllers/input-controller.ts +++ b/packages/coding-agent/src/modes/controllers/input-controller.ts @@ -250,7 +250,7 @@ export class InputController { this.#btwBranchListenerInstalled = true; this.ctx.ui.addInputListener(data => { if (!matchesKey(data, "b")) return undefined; - if (!this.ctx.canBranchBtw()) return undefined; + if (!this.ctx.hasActiveBtw()) return undefined; if (this.ctx.ui.getFocused() !== this.ctx.editor) return undefined; if (this.ctx.editor.getText().trim()) return undefined; void this.ctx.handleBtwBranchKey(); diff --git a/packages/coding-agent/src/modes/interactive-mode.ts b/packages/coding-agent/src/modes/interactive-mode.ts index 122046703..cc5be0312 100644 --- a/packages/coding-agent/src/modes/interactive-mode.ts +++ b/packages/coding-agent/src/modes/interactive-mode.ts @@ -4856,9 +4856,9 @@ export class InteractiveMode implements InteractiveModeContext { return this.#btwController.handleCopy(); } - async handleBtwBranch(question: string, assistantMessage: AssistantMessage): Promise { + async handleBtwBranch(question: string, assistantMessage: AssistantMessage, leafId: string): Promise { try { - const result = await this.session.branchFromBtw(question, assistantMessage); + const result = await this.session.branchFromBtw(question, assistantMessage, leafId); if (result.cancelled) { this.showStatus("/btw branch cancelled", { dim: true }); return; diff --git a/packages/coding-agent/src/modes/types.ts b/packages/coding-agent/src/modes/types.ts index c204b8386..ef2b3a4c7 100644 --- a/packages/coding-agent/src/modes/types.ts +++ b/packages/coding-agent/src/modes/types.ts @@ -420,7 +420,7 @@ export interface InteractiveModeContext { canBranchBtw(): boolean; canCopyBtw(): boolean; handleBtwCopyKey(): Promise; - handleBtwBranch(question: string, assistantMessage: AssistantMessage): Promise; + handleBtwBranch(question: string, assistantMessage: AssistantMessage, leafId: string): Promise; handleOmfgCommand(complaint: string): Promise; hasActiveOmfg(): boolean; handleOmfgEscape(): boolean; diff --git a/packages/coding-agent/src/session/agent-session.ts b/packages/coding-agent/src/session/agent-session.ts index 73b02eafd..95e93876f 100644 --- a/packages/coding-agent/src/session/agent-session.ts +++ b/packages/coding-agent/src/session/agent-session.ts @@ -340,6 +340,7 @@ import { TodoTracker, type TodoTrackerHost } from "./todo-tracker"; import { TtsrCoordinator, type TtsrCoordinatorHost } from "./ttsr-coordinator"; const PLAN_MODE_REMINDER_MAX = 3; +const POST_PROMPT_DRAIN_TIMEOUT_MS = 5_000; /** Internal marker for hook messages queued through the agent loop */ // ============================================================================ @@ -3676,7 +3677,11 @@ export class AgentSession { const postPromptDrain = this.#cancelPostPromptTasks(); this.agent.abort(); try { - await withTimeout(postPromptDrain, 5_000, "Timed out draining post-prompt tasks during dispose"); + await withTimeout( + postPromptDrain, + POST_PROMPT_DRAIN_TIMEOUT_MS, + "Timed out draining post-prompt tasks during dispose", + ); } catch (error) { logger.warn("Post-prompt tasks still draining at dispose deadline", { error: String(error) }); } @@ -7553,21 +7558,23 @@ export class AgentSession { } } + /** Promotes a completed /btw answer from the explicitly authorized session leaf. */ async branchFromBtw( question: string, assistantMessage: AssistantMessage, + leafId: string, ): Promise<{ cancelled: boolean; sessionFile: string | undefined }> { const previousSessionFile = this.sessionFile; if (!this.sessionManager.getSessionFile()) { throw new Error("Cannot branch /btw: session is not persisted"); } - const leafId = this.sessionManager.getLeafId(); - if (!leafId) { - throw new Error("Cannot branch /btw: current session has no leaf"); + if (!leafId || this.sessionManager.getLeafId() !== leafId) { + throw new Error("Cannot branch /btw: session changed since /btw started"); } if ( + this.isStreaming || this.isBashRunning || this.isEvalRunning || this.isCompacting || @@ -7588,8 +7595,17 @@ export class AgentSession { } } - await this.#cancelPostPromptTasks(); + if (this.sessionManager.getLeafId() !== leafId) { + throw new Error("Cannot branch /btw: session changed since /btw started"); + } + + await withTimeout( + this.#cancelPostPromptTasks(), + POST_PROMPT_DRAIN_TIMEOUT_MS, + "Timed out draining post-prompt tasks before /btw branch", + ); if ( + this.isStreaming || this.isBashRunning || this.isEvalRunning || this.isCompacting || @@ -7602,10 +7618,6 @@ export class AgentSession { this.#pendingNextTurnMessages = []; this.#scheduledHiddenNextTurnGeneration = undefined; this.agent.replaceQueues([], []); - if (this.isStreaming) { - await this.abort({ goalReason: "internal", reason: "branching /btw" }); - this.agent.replaceQueues([], []); - } await this.#bash.flushPending(); await this.sessionManager.flush(); const bashTransition = this.#bash.beginSessionTransition(); @@ -7619,6 +7631,9 @@ export class AgentSession { advisorRecordersDetached = true; await this.#advisors.drainAndDetachRecorders(); try { + if (this.sessionManager.getLeafId() !== leafId) { + throw new Error("Cannot branch /btw: session changed since /btw started"); + } this.sessionManager.createBranchedSession(leafId); this.#bash.markSessionTransition(bashTransition); this.#advisors.clearCost(); diff --git a/packages/coding-agent/test/agent-session-btw-branch.test.ts b/packages/coding-agent/test/agent-session-btw-branch.test.ts index 052584b44..07ff57620 100644 --- a/packages/coding-agent/test/agent-session-btw-branch.test.ts +++ b/packages/coding-agent/test/agent-session-btw-branch.test.ts @@ -48,6 +48,12 @@ function expectSanitizedBtwAssistant(message: AssistantMessage): void { ]); } +function requiredLeafId(session: AgentSession): string { + const leafId = session.sessionManager.getLeafId(); + if (!leafId) throw new Error("Expected session leaf"); + return leafId; +} + describe("AgentSession.branchFromBtw", () => { let tempDir: string; let session: AgentSession | undefined; @@ -122,7 +128,11 @@ describe("AgentSession.branchFromBtw", () => { const originalRaw = fs.readFileSync(originalFile!, "utf8"); const assistantMessage = createBtwAssistant(); - const result = await activeSession.branchFromBtw("why did this fail?", assistantMessage); + const result = await activeSession.branchFromBtw( + "why did this fail?", + assistantMessage, + requiredLeafId(activeSession), + ); expect(result.cancelled).toBe(false); expect(result.sessionFile).toBe(activeSession.sessionFile); @@ -155,7 +165,7 @@ describe("AgentSession.branchFromBtw", () => { return result; }); - const result = await activeSession.branchFromBtw("question", createBtwAssistant()); + const result = await activeSession.branchFromBtw("question", createBtwAssistant(), requiredLeafId(activeSession)); expect(result.cancelled).toBe(false); const replacementSessionFile = activeSession.sessionFile; if (!replacementSessionFile) throw new Error("Expected the replacement session to be persisted"); @@ -176,7 +186,7 @@ describe("AgentSession.branchFromBtw", () => { await activeSession.sessionManager.flush(); const originalFile = activeSession.sessionFile; - const result = await activeSession.branchFromBtw("question", createBtwAssistant()); + const result = await activeSession.branchFromBtw("question", createBtwAssistant(), requiredLeafId(activeSession)); expect(result).toEqual({ cancelled: true, sessionFile: originalFile }); expect(activeSession.sessionFile).toBe(originalFile); @@ -186,6 +196,36 @@ describe("AgentSession.branchFromBtw", () => { }); }); + it("refuses when the session leaf advances while a branch hook is pending", async () => { + const hookStarted = Promise.withResolvers(); + const hookRelease = Promise.withResolvers(); + const extensionRunner = { + hasHandlers: vi.fn((eventType: string) => eventType === "session_before_branch"), + emit: vi.fn(async () => { + hookStarted.resolve(); + await hookRelease.promise; + return undefined; + }), + } as unknown as ExtensionRunner; + const activeSession = await createSession({ extensionRunner }); + activeSession.sessionManager.appendMessage({ role: "user", content: "seed", timestamp: Date.now() }); + await activeSession.sessionManager.flush(); + const originalFile = activeSession.sessionFile; + + const branchPromise = activeSession.branchFromBtw( + "question", + createBtwAssistant(), + requiredLeafId(activeSession), + ); + await hookStarted.promise; + activeSession.sessionManager.appendMessage({ role: "user", content: "late work", timestamp: Date.now() }); + await activeSession.sessionManager.flush(); + hookRelease.resolve(); + + await expect(branchPromise).rejects.toThrow("Cannot branch /btw: session changed since /btw started"); + expect(activeSession.sessionFile).toBe(originalFile); + }); + it("syncs promoted /btw messages into live context even when hooks skip conversation restore", async () => { const extensionRunner = { hasHandlers: vi.fn((eventType: string) => eventType === "session_before_branch"), @@ -197,7 +237,7 @@ describe("AgentSession.branchFromBtw", () => { await activeSession.sessionManager.flush(); const assistantMessage = createBtwAssistant(); - const result = await activeSession.branchFromBtw("question", assistantMessage); + const result = await activeSession.branchFromBtw("question", assistantMessage, requiredLeafId(activeSession)); expect(result.cancelled).toBe(false); const messages = activeSession.messages; @@ -208,47 +248,30 @@ describe("AgentSession.branchFromBtw", () => { expectSanitizedBtwAssistant(promoted); }); - it("aborts an in-flight main stream before switching to the /btw branch", async () => { + it("refuses to defer a /btw branch while the main turn is streaming", async () => { const providerStarted = Promise.withResolvers(); const activeSession = await createSession({ handler: () => { providerStarted.resolve(); - return { content: ["main response should not move"], delayMs: 60_000 }; + return { content: ["main response"], delayMs: 60_000 }; }, }); activeSession.sessionManager.appendMessage({ role: "user", content: "seed", timestamp: Date.now() }); await activeSession.sessionManager.flush(); + const originalFile = activeSession.sessionFile; const promptPromise = activeSession.prompt("main prompt"); await providerStarted.promise; expect(activeSession.isStreaming).toBe(true); - await activeSession.followUp("queued follow-up should not move"); - expect(activeSession.queuedMessageCount).toBe(1); - const assistantMessage = createBtwAssistant(); - const result = await activeSession.branchFromBtw("question", assistantMessage); + await expect( + activeSession.branchFromBtw("question", createBtwAssistant(), requiredLeafId(activeSession)), + ).rejects.toThrow("Cannot branch /btw while session maintenance or user work is still running"); + expect(activeSession.isStreaming).toBe(true); + expect(activeSession.sessionFile).toBe(originalFile); + + await activeSession.abort({ goalReason: "internal", reason: "test cleanup" }); await promptPromise; - - expect(result.cancelled).toBe(false); - const messages = activeSession.messages; - expect(messages.at(-2)).toMatchObject({ role: "user", content: [{ type: "text", text: "question" }] }); - const promoted = messages.at(-1); - expect(promoted?.role).toBe("assistant"); - if (promoted?.role !== "assistant") throw new Error("Expected promoted assistant message"); - expectSanitizedBtwAssistant(promoted); - expect(messages).not.toContainEqual( - expect.objectContaining({ - role: "assistant", - content: [{ type: "text", text: "main response should not move" }], - }), - ); - expect(activeSession.queuedMessageCount).toBe(0); - expect(messages).not.toContainEqual( - expect.objectContaining({ - role: "user", - content: [{ type: "text", text: "queued follow-up should not move" }], - }), - ); }); it("refuses to branch /btw while user bash work is still running", async () => { @@ -261,9 +284,9 @@ describe("AgentSession.branchFromBtw", () => { }); while (!activeSession.isBashRunning) await Bun.sleep(1); - await expect(activeSession.branchFromBtw("question", createBtwAssistant())).rejects.toThrow( - "Cannot branch /btw while session maintenance or user work is still running", - ); + await expect( + activeSession.branchFromBtw("question", createBtwAssistant(), requiredLeafId(activeSession)), + ).rejects.toThrow("Cannot branch /btw while session maintenance or user work is still running"); activeSession.abortBash(); await bashPromise.catch(() => undefined); @@ -278,9 +301,9 @@ describe("AgentSession.branchFromBtw", () => { activeSession.trackEvalExecution(execution, abortController).catch(() => undefined); expect(activeSession.isEvalRunning).toBe(true); - await expect(activeSession.branchFromBtw("question", createBtwAssistant())).rejects.toThrow( - "Cannot branch /btw while session maintenance or user work is still running", - ); + await expect( + activeSession.branchFromBtw("question", createBtwAssistant(), requiredLeafId(activeSession)), + ).rejects.toThrow("Cannot branch /btw while session maintenance or user work is still running"); abortController.abort(); }); @@ -295,12 +318,12 @@ describe("AgentSession.branchFromBtw", () => { }); sessionWithMaintenance._maintenanceForTest = true; - await expect(activeSession.branchFromBtw("question", createBtwAssistant())).rejects.toThrow( - "Cannot branch /btw while session maintenance or user work is still running", - ); + await expect( + activeSession.branchFromBtw("question", createBtwAssistant(), requiredLeafId(activeSession)), + ).rejects.toThrow("Cannot branch /btw while session maintenance or user work is still running"); }); - it("cancels post-prompt work after branch hooks before switching sessions", async () => { + it("refuses when post-prompt work starts a turn while a branch hook is pending", async () => { const hookRelease = Promise.withResolvers(); const extensionRunner = { hasHandlers: vi.fn((eventType: string) => eventType === "session_before_branch"), @@ -312,6 +335,7 @@ describe("AgentSession.branchFromBtw", () => { const activeSession = await createSession({ extensionRunner }); activeSession.sessionManager.appendMessage({ role: "user", content: "seed", timestamp: Date.now() }); await activeSession.sessionManager.flush(); + const originalFile = activeSession.sessionFile; activeSession.queueDeferredMessage({ role: "custom", customType: "test-hidden-message", @@ -321,23 +345,26 @@ describe("AgentSession.branchFromBtw", () => { }); expect(activeSession.hasPostPromptWork).toBe(true); - const branchPromise = activeSession.branchFromBtw("question", createBtwAssistant()); + const branchPromise = activeSession.branchFromBtw( + "question", + createBtwAssistant(), + requiredLeafId(activeSession), + ); await Promise.resolve(); - expect(activeSession.hasPostPromptWork).toBe(true); - hookRelease.resolve(); - const result = await branchPromise; - expect(result.cancelled).toBe(false); - expect(activeSession.hasPostPromptWork).toBe(false); + await expect(branchPromise).rejects.toThrow( + "Cannot branch /btw while session maintenance or user work is still running", + ); + expect(activeSession.sessionFile).toBe(originalFile); }); it("throws for in-memory sessions", async () => { const activeSession = await createSession({ persisted: false }); activeSession.sessionManager.appendMessage({ role: "user", content: "seed", timestamp: Date.now() }); - await expect(activeSession.branchFromBtw("question", createBtwAssistant())).rejects.toThrow( - "Cannot branch /btw: session is not persisted", - ); + await expect( + activeSession.branchFromBtw("question", createBtwAssistant(), requiredLeafId(activeSession)), + ).rejects.toThrow("Cannot branch /btw: session is not persisted"); }); }); diff --git a/packages/coding-agent/test/input-controller-keybindings.test.ts b/packages/coding-agent/test/input-controller-keybindings.test.ts index 1a4b63893..bb54f68fc 100644 --- a/packages/coding-agent/test/input-controller-keybindings.test.ts +++ b/packages/coding-agent/test/input-controller-keybindings.test.ts @@ -109,6 +109,7 @@ async function createContext() { const handleBtwCopyKey = vi.fn(async () => true); const canBranchBtw = vi.fn(() => false); const canCopyBtw = vi.fn(() => false); + const hasActiveBtw = vi.fn(() => false); const editor: FakeEditor = { setText(text: string) { editorText = text; @@ -210,7 +211,7 @@ async function createContext() { toggleThinkingBlockVisibility: vi.fn(), showModelSelector, updateEditorBorderColor: vi.fn(), - hasActiveBtw: vi.fn(() => false), + hasActiveBtw, handleBtwBranchKey, canBranchBtw, canCopyBtw, @@ -242,6 +243,7 @@ async function createContext() { handleBtwBranchKey, addInputListener, canBranchBtw, + hasActiveBtw, handleBtwCopyKey, canCopyBtw, showError, @@ -392,7 +394,7 @@ describe("InputController keybinding setup", () => { it("routes b to branch a branchable /btw panel", async () => { const { InputController, ctx, spies } = await createContext(); - (ctx.canBranchBtw as unknown as { mockReturnValue(value: boolean): void }).mockReturnValue(true); + spies.hasActiveBtw.mockReturnValue(true); const controller = new InputController(ctx); controller.setupKeyHandlers(); @@ -406,7 +408,7 @@ describe("InputController keybinding setup", () => { it("lets b fall through while the editor has draft text", async () => { const { InputController, ctx, editor, spies } = await createContext(); - (ctx.canBranchBtw as unknown as { mockReturnValue(value: boolean): void }).mockReturnValue(true); + spies.hasActiveBtw.mockReturnValue(true); editor.setText("build a branch"); const controller = new InputController(ctx); @@ -419,7 +421,21 @@ describe("InputController keybinding setup", () => { expect(spies.handleBtwBranchKey).not.toHaveBeenCalled(); }); - it("lets b fall through when /btw is not branchable", async () => { + it("consumes b while an active /btw branch is unavailable", async () => { + const { InputController, ctx, spies } = await createContext(); + spies.hasActiveBtw.mockReturnValue(true); + const controller = new InputController(ctx); + + controller.setupKeyHandlers(); + const listener = spies.addInputListener.mock.calls[1]?.[0]; + expect(listener).toBeDefined(); + const result = listener?.("b"); + + expect(result).toEqual({ consume: true }); + expect(spies.handleBtwBranchKey).toHaveBeenCalledTimes(1); + }); + + it("lets b fall through when no /btw panel is active", async () => { const { InputController, ctx, spies } = await createContext(); const controller = new InputController(ctx); @@ -434,7 +450,7 @@ describe("InputController keybinding setup", () => { it("lets b fall through while another input is focused", async () => { const { InputController, ctx, setFocused, spies } = await createContext(); - (ctx.canBranchBtw as unknown as { mockReturnValue(value: boolean): void }).mockReturnValue(true); + spies.hasActiveBtw.mockReturnValue(true); setFocused({ pasteText: vi.fn() }); const controller = new InputController(ctx); diff --git a/packages/coding-agent/test/modes/controllers/btw-controller.test.ts b/packages/coding-agent/test/modes/controllers/btw-controller.test.ts index 4585aafea..3bf60b8ea 100644 --- a/packages/coding-agent/test/modes/controllers/btw-controller.test.ts +++ b/packages/coding-agent/test/modes/controllers/btw-controller.test.ts @@ -45,6 +45,7 @@ function makeFakeSession( ): InteractiveModeContext["session"] { return { model: { provider: "anthropic", id: "claude-sonnet-4-5" }, + isStreaming: false, runEphemeralTurn, } as unknown as InteractiveModeContext["session"]; } @@ -101,6 +102,18 @@ describe("BtwPanelComponent", () => { expect(rendered).toContain("b branch to chat"); expect(rendered).toContain("Esc dismiss"); }); + + it("hides the branch action when the controller rejects the current leaf", () => { + const ui = { requestRender: vi.fn(), requestComponentRender: vi.fn() } as unknown as TUI; + const panel = new BtwPanelComponent({ question: "Question?", tui: ui, canBranch: () => false }); + + panel.setAnswer("Answer"); + panel.markComplete(); + + const rendered = Bun.stripANSI(panel.render(120).join("\n")); + expect(rendered).toContain("c copy"); + expect(rendered).not.toContain("b branch to chat"); + }); }); describe("BtwController", () => { @@ -255,6 +268,27 @@ describe("BtwController", () => { expect(controller.canBranch()).toBe(true); }); + it("refuses a completed branch while the main turn is streaming", async () => { + const assistantMessage = createAssistantMessage("Answer"); + const runEphemeralTurn = vi.fn(async () => ({ replyText: "Answer", assistantMessage })); + const session = makeFakeSession(runEphemeralTurn); + Object.defineProperty(session, "isStreaming", { value: true }); + const btwContainer = new Container(); + const ctx = makeCtx(session, btwContainer); + const controller = new BtwController(ctx); + + await controller.start("Question?"); + await drainBtwRequest(); + + expect(controller.canBranch()).toBe(false); + const panel = btwContainer.children[0]; + expect(Bun.stripANSI(panel?.render(120).join("\n") ?? "")).not.toContain("b branch to chat"); + expect(await controller.handleBranch()).toBe(false); + expect(ctx.showStatus).toHaveBeenCalledWith("/btw branch unavailable: a turn is still running", { + dim: true, + }); + }); + it("does not allow branch after a complete empty reply", async () => { const runEphemeralTurn = vi.fn(async () => ({ replyText: " ", @@ -307,7 +341,33 @@ describe("BtwController", () => { await drainBtwRequest(); expect(await controller.handleBranch()).toBe(true); - expect(ctx.handleBtwBranch).toHaveBeenCalledWith("Question?", assistantMessage); + expect(ctx.handleBtwBranch).toHaveBeenCalledWith("Question?", assistantMessage, "leaf-1"); + }); + + it("keeps a pending branch visible and refuses to dismiss it", async () => { + const branch = Promise.withResolvers(); + const assistantMessage = createAssistantMessage("Answer"); + const runEphemeralTurn = vi.fn(async () => ({ replyText: "Answer", assistantMessage })); + const btwContainer = new Container(); + const ctx = makeCtx(makeFakeSession(runEphemeralTurn), btwContainer); + ctx.handleBtwBranch = vi.fn(async () => { + await branch.promise; + }); + const controller = new BtwController(ctx); + + await controller.start("Question?"); + await drainBtwRequest(); + const branchPromise = controller.handleBranch(); + await Promise.resolve(); + + const panel = btwContainer.children[0]; + expect(Bun.stripANSI(panel?.render(120).join("\n") ?? "")).toContain("Branching to chat"); + expect(controller.handleEscape()).toBe(true); + expect(btwContainer.children).toHaveLength(1); + expect(ctx.showStatus).toHaveBeenCalledWith("/btw branch is in progress", { dim: true }); + + branch.resolve(); + await branchPromise; }); it("branches the sanitized reply text while preserving non-text assistant content", async () => { @@ -333,13 +393,17 @@ describe("BtwController", () => { await drainBtwRequest(); expect(await controller.handleBranch()).toBe(true); - expect(ctx.handleBtwBranch).toHaveBeenCalledWith("Question?", { - ...assistantMessage, - content: [ - { type: "thinking", thinking: "Keep this reasoning." }, - { type: "text", text: "sanitized" }, - ], - }); + expect(ctx.handleBtwBranch).toHaveBeenCalledWith( + "Question?", + { + ...assistantMessage, + content: [ + { type: "thinking", thinking: "Keep this reasoning." }, + { type: "text", text: "sanitized" }, + ], + }, + "leaf-1", + ); }); it("copies the sanitized visible reply text after a complete non-empty reply", async () => { @@ -418,14 +482,18 @@ describe("BtwController", () => { await drainBtwRequest(); expect(await controller.handleBranch()).toBe(true); - expect(ctx.handleBtwBranch).toHaveBeenCalledWith("Question?", { - ...assistantMessage, - content: [ - { type: "thinking", thinking: "reasoning" }, - { type: "text", text: "sanitized" }, - ], - providerPayload: undefined, - }); + expect(ctx.handleBtwBranch).toHaveBeenCalledWith( + "Question?", + { + ...assistantMessage, + content: [ + { type: "thinking", thinking: "reasoning" }, + { type: "text", text: "sanitized" }, + ], + providerPayload: undefined, + }, + "leaf-1", + ); }); it("ignores duplicate branch requests while branch promotion is in flight", async () => { From 98f484bd2f9a11571f807b8b0b1441cc7a7a6671 Mon Sep 17 00:00:00 2001 From: roboomp Date: Mon, 3 Aug 2026 09:22:32 +0000 Subject: [PATCH 2/3] fix(session): validated session identity for /btw branch Branched session files preserve entry ids, so leaf-id equality alone let a stale /btw answer promote into a different loaded session. Capture the originating session id at /btw start and require it to match at both the controller gate and every branchFromBtw checkpoint. Fixes #7474 --- .../src/modes/controllers/btw-controller.ts | 15 +++- .../src/modes/interactive-mode.ts | 9 ++- packages/coding-agent/src/modes/types.ts | 7 +- .../coding-agent/src/session/agent-session.ts | 9 ++- .../test/agent-session-btw-branch.test.ts | 74 +++++++++++++++++-- .../modes/controllers/btw-controller.test.ts | 42 ++++++++++- 6 files changed, 135 insertions(+), 21 deletions(-) diff --git a/packages/coding-agent/src/modes/controllers/btw-controller.ts b/packages/coding-agent/src/modes/controllers/btw-controller.ts index 6e3c79d6f..1f3a2b315 100644 --- a/packages/coding-agent/src/modes/controllers/btw-controller.ts +++ b/packages/coding-agent/src/modes/controllers/btw-controller.ts @@ -10,6 +10,7 @@ interface BtwRequest { abortController: AbortController; question: string; leafId: string | null; + sessionId: string; } function assistantMessageWithReplyText(assistantMessage: AssistantMessage, replyText: string): AssistantMessage { @@ -39,6 +40,7 @@ export class BtwController { #lastReplyText: string | undefined; #lastAssistantMessage: AssistantMessage | undefined; #lastLeafId: string | null | undefined; + #lastSessionId: string | undefined; #branchInFlight = false; #lastCopyText: string | undefined; #copyInFlight = false; @@ -60,7 +62,10 @@ export class BtwController { return "the answer is unavailable"; } if (!this.#lastLeafId) return "the session has no branch point"; - if (this.#lastLeafId !== this.ctx.sessionManager.getLeafId()) { + if ( + this.#lastSessionId !== this.ctx.sessionManager.getSessionId() || + this.#lastLeafId !== this.ctx.sessionManager.getLeafId() + ) { return "the session changed since /btw started"; } if (this.ctx.session.isStreaming) return "a turn is still running"; @@ -98,12 +103,13 @@ export class BtwController { const question = this.#lastQuestion; const assistantMessage = this.#lastAssistantMessage; const leafId = this.#lastLeafId; - if (!request || !question || !assistantMessage || !leafId) return false; + const sessionId = this.#lastSessionId; + if (!request || !question || !assistantMessage || !leafId || !sessionId) return false; this.#branchInFlight = true; request.component.markBranching(); try { - await this.ctx.handleBtwBranch(question, assistantMessage, leafId); + await this.ctx.handleBtwBranch(question, assistantMessage, leafId, sessionId); return true; } finally { this.#branchInFlight = false; @@ -149,6 +155,7 @@ export class BtwController { abortController: new AbortController(), question: trimmedQuestion, leafId: this.ctx.sessionManager.getLeafId(), + sessionId: this.ctx.sessionManager.getSessionId(), }; this.ctx.btwContainer.clear(); this.ctx.btwContainer.addChild(request.component); @@ -182,6 +189,7 @@ export class BtwController { this.#lastCopyText = copyText; this.#lastAssistantMessage = assistantMessageWithReplyText(assistantMessage, replyText); this.#lastLeafId = request.leafId; + this.#lastSessionId = request.sessionId; } else { this.#clearCompletedState(); } @@ -216,6 +224,7 @@ export class BtwController { this.#lastAssistantMessage = undefined; this.#lastCopyText = undefined; this.#lastLeafId = undefined; + this.#lastSessionId = undefined; } #isActiveRequest(request: BtwRequest): boolean { diff --git a/packages/coding-agent/src/modes/interactive-mode.ts b/packages/coding-agent/src/modes/interactive-mode.ts index cc5be0312..c4f809a49 100644 --- a/packages/coding-agent/src/modes/interactive-mode.ts +++ b/packages/coding-agent/src/modes/interactive-mode.ts @@ -4856,9 +4856,14 @@ export class InteractiveMode implements InteractiveModeContext { return this.#btwController.handleCopy(); } - async handleBtwBranch(question: string, assistantMessage: AssistantMessage, leafId: string): Promise { + async handleBtwBranch( + question: string, + assistantMessage: AssistantMessage, + leafId: string, + sessionId: string, + ): Promise { try { - const result = await this.session.branchFromBtw(question, assistantMessage, leafId); + const result = await this.session.branchFromBtw(question, assistantMessage, leafId, sessionId); if (result.cancelled) { this.showStatus("/btw branch cancelled", { dim: true }); return; diff --git a/packages/coding-agent/src/modes/types.ts b/packages/coding-agent/src/modes/types.ts index ef2b3a4c7..f0167e098 100644 --- a/packages/coding-agent/src/modes/types.ts +++ b/packages/coding-agent/src/modes/types.ts @@ -420,7 +420,12 @@ export interface InteractiveModeContext { canBranchBtw(): boolean; canCopyBtw(): boolean; handleBtwCopyKey(): Promise; - handleBtwBranch(question: string, assistantMessage: AssistantMessage, leafId: string): Promise; + handleBtwBranch( + question: string, + assistantMessage: AssistantMessage, + leafId: string, + sessionId: string, + ): Promise; handleOmfgCommand(complaint: string): Promise; hasActiveOmfg(): boolean; handleOmfgEscape(): boolean; diff --git a/packages/coding-agent/src/session/agent-session.ts b/packages/coding-agent/src/session/agent-session.ts index 95e93876f..3bcba7650 100644 --- a/packages/coding-agent/src/session/agent-session.ts +++ b/packages/coding-agent/src/session/agent-session.ts @@ -7558,18 +7558,19 @@ export class AgentSession { } } - /** Promotes a completed /btw answer from the explicitly authorized session leaf. */ + /** Promotes a completed /btw answer from the explicitly authorized session and leaf. */ async branchFromBtw( question: string, assistantMessage: AssistantMessage, leafId: string, + sessionId: string, ): Promise<{ cancelled: boolean; sessionFile: string | undefined }> { const previousSessionFile = this.sessionFile; if (!this.sessionManager.getSessionFile()) { throw new Error("Cannot branch /btw: session is not persisted"); } - if (!leafId || this.sessionManager.getLeafId() !== leafId) { + if (!leafId || this.sessionManager.getSessionId() !== sessionId || this.sessionManager.getLeafId() !== leafId) { throw new Error("Cannot branch /btw: session changed since /btw started"); } @@ -7595,7 +7596,7 @@ export class AgentSession { } } - if (this.sessionManager.getLeafId() !== leafId) { + if (this.sessionManager.getSessionId() !== sessionId || this.sessionManager.getLeafId() !== leafId) { throw new Error("Cannot branch /btw: session changed since /btw started"); } @@ -7631,7 +7632,7 @@ export class AgentSession { advisorRecordersDetached = true; await this.#advisors.drainAndDetachRecorders(); try { - if (this.sessionManager.getLeafId() !== leafId) { + if (this.sessionManager.getSessionId() !== sessionId || this.sessionManager.getLeafId() !== leafId) { throw new Error("Cannot branch /btw: session changed since /btw started"); } this.sessionManager.createBranchedSession(leafId); diff --git a/packages/coding-agent/test/agent-session-btw-branch.test.ts b/packages/coding-agent/test/agent-session-btw-branch.test.ts index 07ff57620..47703e216 100644 --- a/packages/coding-agent/test/agent-session-btw-branch.test.ts +++ b/packages/coding-agent/test/agent-session-btw-branch.test.ts @@ -132,6 +132,7 @@ describe("AgentSession.branchFromBtw", () => { "why did this fail?", assistantMessage, requiredLeafId(activeSession), + activeSession.sessionManager.getSessionId(), ); expect(result.cancelled).toBe(false); @@ -165,7 +166,12 @@ describe("AgentSession.branchFromBtw", () => { return result; }); - const result = await activeSession.branchFromBtw("question", createBtwAssistant(), requiredLeafId(activeSession)); + const result = await activeSession.branchFromBtw( + "question", + createBtwAssistant(), + requiredLeafId(activeSession), + activeSession.sessionManager.getSessionId(), + ); expect(result.cancelled).toBe(false); const replacementSessionFile = activeSession.sessionFile; if (!replacementSessionFile) throw new Error("Expected the replacement session to be persisted"); @@ -186,7 +192,12 @@ describe("AgentSession.branchFromBtw", () => { await activeSession.sessionManager.flush(); const originalFile = activeSession.sessionFile; - const result = await activeSession.branchFromBtw("question", createBtwAssistant(), requiredLeafId(activeSession)); + const result = await activeSession.branchFromBtw( + "question", + createBtwAssistant(), + requiredLeafId(activeSession), + activeSession.sessionManager.getSessionId(), + ); expect(result).toEqual({ cancelled: true, sessionFile: originalFile }); expect(activeSession.sessionFile).toBe(originalFile); @@ -216,6 +227,7 @@ describe("AgentSession.branchFromBtw", () => { "question", createBtwAssistant(), requiredLeafId(activeSession), + activeSession.sessionManager.getSessionId(), ); await hookStarted.promise; activeSession.sessionManager.appendMessage({ role: "user", content: "late work", timestamp: Date.now() }); @@ -226,6 +238,21 @@ describe("AgentSession.branchFromBtw", () => { expect(activeSession.sessionFile).toBe(originalFile); }); + it("refuses when the authorized session id no longer matches the loaded session", async () => { + const activeSession = await createSession(); + activeSession.sessionManager.appendMessage({ role: "user", content: "seed", timestamp: Date.now() }); + await activeSession.sessionManager.flush(); + const originalFile = activeSession.sessionFile; + const leafId = requiredLeafId(activeSession); + + // A resumed/branched session preserves the entry id, so the leaf still matches + // while the loaded session is different. + await expect( + activeSession.branchFromBtw("question", createBtwAssistant(), leafId, "some-other-session"), + ).rejects.toThrow("Cannot branch /btw: session changed since /btw started"); + expect(activeSession.sessionFile).toBe(originalFile); + }); + it("syncs promoted /btw messages into live context even when hooks skip conversation restore", async () => { const extensionRunner = { hasHandlers: vi.fn((eventType: string) => eventType === "session_before_branch"), @@ -237,7 +264,12 @@ describe("AgentSession.branchFromBtw", () => { await activeSession.sessionManager.flush(); const assistantMessage = createBtwAssistant(); - const result = await activeSession.branchFromBtw("question", assistantMessage, requiredLeafId(activeSession)); + const result = await activeSession.branchFromBtw( + "question", + assistantMessage, + requiredLeafId(activeSession), + activeSession.sessionManager.getSessionId(), + ); expect(result.cancelled).toBe(false); const messages = activeSession.messages; @@ -265,7 +297,12 @@ describe("AgentSession.branchFromBtw", () => { expect(activeSession.isStreaming).toBe(true); await expect( - activeSession.branchFromBtw("question", createBtwAssistant(), requiredLeafId(activeSession)), + activeSession.branchFromBtw( + "question", + createBtwAssistant(), + requiredLeafId(activeSession), + activeSession.sessionManager.getSessionId(), + ), ).rejects.toThrow("Cannot branch /btw while session maintenance or user work is still running"); expect(activeSession.isStreaming).toBe(true); expect(activeSession.sessionFile).toBe(originalFile); @@ -285,7 +322,12 @@ describe("AgentSession.branchFromBtw", () => { while (!activeSession.isBashRunning) await Bun.sleep(1); await expect( - activeSession.branchFromBtw("question", createBtwAssistant(), requiredLeafId(activeSession)), + activeSession.branchFromBtw( + "question", + createBtwAssistant(), + requiredLeafId(activeSession), + activeSession.sessionManager.getSessionId(), + ), ).rejects.toThrow("Cannot branch /btw while session maintenance or user work is still running"); activeSession.abortBash(); @@ -302,7 +344,12 @@ describe("AgentSession.branchFromBtw", () => { expect(activeSession.isEvalRunning).toBe(true); await expect( - activeSession.branchFromBtw("question", createBtwAssistant(), requiredLeafId(activeSession)), + activeSession.branchFromBtw( + "question", + createBtwAssistant(), + requiredLeafId(activeSession), + activeSession.sessionManager.getSessionId(), + ), ).rejects.toThrow("Cannot branch /btw while session maintenance or user work is still running"); abortController.abort(); @@ -319,7 +366,12 @@ describe("AgentSession.branchFromBtw", () => { sessionWithMaintenance._maintenanceForTest = true; await expect( - activeSession.branchFromBtw("question", createBtwAssistant(), requiredLeafId(activeSession)), + activeSession.branchFromBtw( + "question", + createBtwAssistant(), + requiredLeafId(activeSession), + activeSession.sessionManager.getSessionId(), + ), ).rejects.toThrow("Cannot branch /btw while session maintenance or user work is still running"); }); @@ -349,6 +401,7 @@ describe("AgentSession.branchFromBtw", () => { "question", createBtwAssistant(), requiredLeafId(activeSession), + activeSession.sessionManager.getSessionId(), ); await Promise.resolve(); hookRelease.resolve(); @@ -364,7 +417,12 @@ describe("AgentSession.branchFromBtw", () => { activeSession.sessionManager.appendMessage({ role: "user", content: "seed", timestamp: Date.now() }); await expect( - activeSession.branchFromBtw("question", createBtwAssistant(), requiredLeafId(activeSession)), + activeSession.branchFromBtw( + "question", + createBtwAssistant(), + requiredLeafId(activeSession), + activeSession.sessionManager.getSessionId(), + ), ).rejects.toThrow("Cannot branch /btw: session is not persisted"); }); }); diff --git a/packages/coding-agent/test/modes/controllers/btw-controller.test.ts b/packages/coding-agent/test/modes/controllers/btw-controller.test.ts index 3bf60b8ea..88f7be15d 100644 --- a/packages/coding-agent/test/modes/controllers/btw-controller.test.ts +++ b/packages/coding-agent/test/modes/controllers/btw-controller.test.ts @@ -52,18 +52,28 @@ function makeFakeSession( function makeCtx(session: InteractiveModeContext["session"], btwContainer = new Container()): InteractiveModeContext { let leafId: string | null = "leaf-1"; + let sessionId = "session-1"; return { ui: { requestRender: vi.fn(), requestComponentRender: vi.fn() } as unknown as TUI, btwContainer, session, - sessionManager: { getLeafId: () => leafId } as unknown as InteractiveModeContext["sessionManager"], + sessionManager: { + getLeafId: () => leafId, + getSessionId: () => sessionId, + } as unknown as InteractiveModeContext["sessionManager"], showStatus: vi.fn(), showError: vi.fn(), handleBtwBranch: vi.fn(async () => {}), setTestLeafId(nextLeafId: string | null) { leafId = nextLeafId; }, - } as unknown as InteractiveModeContext & { setTestLeafId(nextLeafId: string | null): void }; + setTestSessionId(nextSessionId: string) { + sessionId = nextSessionId; + }, + } as unknown as InteractiveModeContext & { + setTestLeafId(nextLeafId: string | null): void; + setTestSessionId(nextSessionId: string): void; + }; } afterEach(() => { vi.restoreAllMocks(); @@ -268,6 +278,30 @@ describe("BtwController", () => { expect(controller.canBranch()).toBe(true); }); + it("refuses branch when the loaded session changed but the leaf id still matches", async () => { + const assistantMessage = createAssistantMessage("Answer"); + const runEphemeralTurn = vi.fn(async () => ({ replyText: "Answer", assistantMessage })); + const ctx = makeCtx(makeFakeSession(runEphemeralTurn)) as InteractiveModeContext & { + setTestSessionId(nextSessionId: string): void; + }; + const controller = new BtwController(ctx); + + await controller.start("Question?"); + await drainBtwRequest(); + expect(controller.canBranch()).toBe(true); + + // A resumed/branched session preserves the entry id, so the leaf still matches; + // the session id must still gate the promotion. + ctx.setTestSessionId("session-2"); + + expect(controller.canBranch()).toBe(false); + expect(await controller.handleBranch()).toBe(false); + expect(ctx.handleBtwBranch).not.toHaveBeenCalled(); + expect(ctx.showStatus).toHaveBeenCalledWith("/btw branch unavailable: the session changed since /btw started", { + dim: true, + }); + }); + it("refuses a completed branch while the main turn is streaming", async () => { const assistantMessage = createAssistantMessage("Answer"); const runEphemeralTurn = vi.fn(async () => ({ replyText: "Answer", assistantMessage })); @@ -341,7 +375,7 @@ describe("BtwController", () => { await drainBtwRequest(); expect(await controller.handleBranch()).toBe(true); - expect(ctx.handleBtwBranch).toHaveBeenCalledWith("Question?", assistantMessage, "leaf-1"); + expect(ctx.handleBtwBranch).toHaveBeenCalledWith("Question?", assistantMessage, "leaf-1", "session-1"); }); it("keeps a pending branch visible and refuses to dismiss it", async () => { @@ -403,6 +437,7 @@ describe("BtwController", () => { ], }, "leaf-1", + "session-1", ); }); @@ -493,6 +528,7 @@ describe("BtwController", () => { providerPayload: undefined, }, "leaf-1", + "session-1", ); }); From a54a2f79cbfd9b3a3e7b136cea83533064ca9d0a Mon Sep 17 00:00:00 2001 From: roboomp Date: Mon, 3 Aug 2026 09:33:07 +0000 Subject: [PATCH 3/3] fix(tui): preserved b before /btw completion Reserve the plain b shortcut only after /btw has a completed answer or a branch is already pending. Running, empty, aborted, and failed panels now leave the key for the composer, while completed-but-refused branches still consume it with an explanation. Fixes #7474 --- .../src/modes/controllers/btw-controller.ts | 13 +++++++++++++ .../src/modes/controllers/input-controller.ts | 2 +- .../coding-agent/src/modes/interactive-mode.ts | 5 +++++ packages/coding-agent/src/modes/types.ts | 1 + .../test/input-controller-keybindings.test.ts | 16 ++++++++++------ .../modes/controllers/btw-controller.test.ts | 7 +++++++ 6 files changed, 37 insertions(+), 7 deletions(-) diff --git a/packages/coding-agent/src/modes/controllers/btw-controller.ts b/packages/coding-agent/src/modes/controllers/btw-controller.ts index 1f3a2b315..1a45a614d 100644 --- a/packages/coding-agent/src/modes/controllers/btw-controller.ts +++ b/packages/coding-agent/src/modes/controllers/btw-controller.ts @@ -55,6 +55,19 @@ export class BtwController { return this.#branchUnavailableReason() === undefined; } + /** Whether plain `b` is currently reserved for a completed or pending branch action. */ + handlesBranchKey(): boolean { + if (this.#branchInFlight) return true; + if (this.#activeRequest?.component.isBranchable() !== true) return false; + return ( + this.#lastQuestion !== undefined && + this.#lastReplyText !== undefined && + this.#lastAssistantMessage !== undefined && + this.#lastLeafId !== undefined && + this.#lastSessionId !== undefined + ); + } + #branchUnavailableReason(): string | undefined { if (this.#branchInFlight) return "a branch is already in progress"; if (this.#activeRequest?.component.isBranchable() !== true) return "the answer is not ready"; diff --git a/packages/coding-agent/src/modes/controllers/input-controller.ts b/packages/coding-agent/src/modes/controllers/input-controller.ts index dc265a984..711a2d8b8 100644 --- a/packages/coding-agent/src/modes/controllers/input-controller.ts +++ b/packages/coding-agent/src/modes/controllers/input-controller.ts @@ -250,7 +250,7 @@ export class InputController { this.#btwBranchListenerInstalled = true; this.ctx.ui.addInputListener(data => { if (!matchesKey(data, "b")) return undefined; - if (!this.ctx.hasActiveBtw()) return undefined; + if (!this.ctx.handlesBtwBranchKey()) return undefined; if (this.ctx.ui.getFocused() !== this.ctx.editor) return undefined; if (this.ctx.editor.getText().trim()) return undefined; void this.ctx.handleBtwBranchKey(); diff --git a/packages/coding-agent/src/modes/interactive-mode.ts b/packages/coding-agent/src/modes/interactive-mode.ts index c4f809a49..87dd71003 100644 --- a/packages/coding-agent/src/modes/interactive-mode.ts +++ b/packages/coding-agent/src/modes/interactive-mode.ts @@ -4844,6 +4844,11 @@ export class InteractiveMode implements InteractiveModeContext { return this.#btwController.canBranch(); } + /** Reserves plain `b` only after /btw has a completed branch action to handle. */ + handlesBtwBranchKey(): boolean { + return this.#btwController.handlesBranchKey(); + } + handleBtwBranchKey(): Promise { return this.#btwController.handleBranch(); } diff --git a/packages/coding-agent/src/modes/types.ts b/packages/coding-agent/src/modes/types.ts index f0167e098..d62fd9fda 100644 --- a/packages/coding-agent/src/modes/types.ts +++ b/packages/coding-agent/src/modes/types.ts @@ -418,6 +418,7 @@ export interface InteractiveModeContext { handleBtwEscape(): boolean; handleBtwBranchKey(): Promise; canBranchBtw(): boolean; + handlesBtwBranchKey(): boolean; canCopyBtw(): boolean; handleBtwCopyKey(): Promise; handleBtwBranch( diff --git a/packages/coding-agent/test/input-controller-keybindings.test.ts b/packages/coding-agent/test/input-controller-keybindings.test.ts index bb54f68fc..6e0b387c0 100644 --- a/packages/coding-agent/test/input-controller-keybindings.test.ts +++ b/packages/coding-agent/test/input-controller-keybindings.test.ts @@ -110,6 +110,7 @@ async function createContext() { const canBranchBtw = vi.fn(() => false); const canCopyBtw = vi.fn(() => false); const hasActiveBtw = vi.fn(() => false); + const handlesBtwBranchKey = vi.fn(() => false); const editor: FakeEditor = { setText(text: string) { editorText = text; @@ -212,6 +213,7 @@ async function createContext() { showModelSelector, updateEditorBorderColor: vi.fn(), hasActiveBtw, + handlesBtwBranchKey, handleBtwBranchKey, canBranchBtw, canCopyBtw, @@ -244,6 +246,7 @@ async function createContext() { addInputListener, canBranchBtw, hasActiveBtw, + handlesBtwBranchKey, handleBtwCopyKey, canCopyBtw, showError, @@ -394,7 +397,7 @@ describe("InputController keybinding setup", () => { it("routes b to branch a branchable /btw panel", async () => { const { InputController, ctx, spies } = await createContext(); - spies.hasActiveBtw.mockReturnValue(true); + spies.handlesBtwBranchKey.mockReturnValue(true); const controller = new InputController(ctx); controller.setupKeyHandlers(); @@ -408,7 +411,7 @@ describe("InputController keybinding setup", () => { it("lets b fall through while the editor has draft text", async () => { const { InputController, ctx, editor, spies } = await createContext(); - spies.hasActiveBtw.mockReturnValue(true); + spies.handlesBtwBranchKey.mockReturnValue(true); editor.setText("build a branch"); const controller = new InputController(ctx); @@ -421,9 +424,9 @@ describe("InputController keybinding setup", () => { expect(spies.handleBtwBranchKey).not.toHaveBeenCalled(); }); - it("consumes b while an active /btw branch is unavailable", async () => { + it("consumes b while a completed /btw branch is unavailable", async () => { const { InputController, ctx, spies } = await createContext(); - spies.hasActiveBtw.mockReturnValue(true); + spies.handlesBtwBranchKey.mockReturnValue(true); const controller = new InputController(ctx); controller.setupKeyHandlers(); @@ -435,8 +438,9 @@ describe("InputController keybinding setup", () => { expect(spies.handleBtwBranchKey).toHaveBeenCalledTimes(1); }); - it("lets b fall through when no /btw panel is active", async () => { + it("lets b reach the composer before an active /btw answer is branchable", async () => { const { InputController, ctx, spies } = await createContext(); + spies.hasActiveBtw.mockReturnValue(true); const controller = new InputController(ctx); controller.setupKeyHandlers(); @@ -450,7 +454,7 @@ describe("InputController keybinding setup", () => { it("lets b fall through while another input is focused", async () => { const { InputController, ctx, setFocused, spies } = await createContext(); - spies.hasActiveBtw.mockReturnValue(true); + spies.handlesBtwBranchKey.mockReturnValue(true); setFocused({ pasteText: vi.fn() }); const controller = new InputController(ctx); diff --git a/packages/coding-agent/test/modes/controllers/btw-controller.test.ts b/packages/coding-agent/test/modes/controllers/btw-controller.test.ts index 88f7be15d..e0d55017f 100644 --- a/packages/coding-agent/test/modes/controllers/btw-controller.test.ts +++ b/packages/coding-agent/test/modes/controllers/btw-controller.test.ts @@ -249,6 +249,7 @@ describe("BtwController", () => { await controller.start("Question?"); expect(controller.canBranch()).toBe(false); + expect(controller.handlesBranchKey()).toBe(false); }); it("does not allow branch when the completed answer has no originating leaf", async () => { @@ -264,6 +265,7 @@ describe("BtwController", () => { await drainBtwRequest(); expect(controller.canBranch()).toBe(false); + expect(controller.handlesBranchKey()).toBe(true); }); it("allows branch after a complete non-empty reply", async () => { @@ -276,6 +278,7 @@ describe("BtwController", () => { await drainBtwRequest(); expect(controller.canBranch()).toBe(true); + expect(controller.handlesBranchKey()).toBe(true); }); it("refuses branch when the loaded session changed but the leaf id still matches", async () => { @@ -295,6 +298,7 @@ describe("BtwController", () => { ctx.setTestSessionId("session-2"); expect(controller.canBranch()).toBe(false); + expect(controller.handlesBranchKey()).toBe(true); expect(await controller.handleBranch()).toBe(false); expect(ctx.handleBtwBranch).not.toHaveBeenCalled(); expect(ctx.showStatus).toHaveBeenCalledWith("/btw branch unavailable: the session changed since /btw started", { @@ -315,6 +319,7 @@ describe("BtwController", () => { await drainBtwRequest(); expect(controller.canBranch()).toBe(false); + expect(controller.handlesBranchKey()).toBe(true); const panel = btwContainer.children[0]; expect(Bun.stripANSI(panel?.render(120).join("\n") ?? "")).not.toContain("b branch to chat"); expect(await controller.handleBranch()).toBe(false); @@ -335,6 +340,7 @@ describe("BtwController", () => { await drainBtwRequest(); expect(controller.canBranch()).toBe(false); + expect(controller.handlesBranchKey()).toBe(false); }); it("does not allow branch after aborted or errored requests", async () => { @@ -393,6 +399,7 @@ describe("BtwController", () => { await drainBtwRequest(); const branchPromise = controller.handleBranch(); await Promise.resolve(); + expect(controller.handlesBranchKey()).toBe(true); const panel = btwContainer.children[0]; expect(Bun.stripANSI(panel?.render(120).join("\n") ?? "")).toContain("Branching to chat");