Merge PR #7475: fix(session): prevent stale /btw branch promotion (@roboomp)

packages/coding-agent/CHANGELOG.md

@@ -10,6 +10,7 @@
 - Fixed the Windows bash tool silently taking down the whole omp process when a command blocked until its timeout: cancelling a timed-out run walked the spawned child's descendant tree from raw `th32ParentProcessID` links, and a recycled pid matching the harness's stale recorded parent pid could enumerate omp as a false descendant and `TerminateProcess` it, killing the session with no `session_exit` record. Run-cancellation sweeps now refuse to signal the harness or any process collected beneath it, while still reaping the timed-out target when it owns a recycled ancestor pid ([#7452](https://github.com/can1357/oh-my-pi/issues/7452)).
 - Fixed a supervised process reaching a terminal state without telling the session that launched it. The broker recorded the final snapshot for `hub ps` to read but emitted no event, so an idle owner only learned that a process had exited by polling. The broker now sends one `daemon-completed` notification per terminal exit to the owner socket recorded at launch, and the session turns it into a model-visible message. Restart transitions stay live, and a stale or disposed session drops the notification without changing explicit `hub wait` or persistence behavior.
 - Fixed PCRE2-only grep patterns terminating Bun on macOS by using the interpreted PCRE2 engine instead of its unsafe same-process JIT in native grep, embedded `grep -P`, and embedded `rg --pcre2`/`--engine=auto` ([#7399](https://github.com/can1357/oh-my-pi/issues/7399)).
+- 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
 

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":

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;
@@ -50,17 +52,39 @@ export class BtwController {
 	}
 
 	canBranch(): boolean {
+		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.#branchInFlight &&
-			this.#activeRequest?.component.isBranchable() === true &&
 			this.#lastQuestion !== undefined &&
 			this.#lastReplyText !== undefined &&
 			this.#lastAssistantMessage !== undefined &&
-			this.#lastLeafId !== null &&
-			this.#lastLeafId === this.ctx.sessionManager.getLeafId()
+			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";
+		if (!this.#lastQuestion || !this.#lastReplyText || !this.#lastAssistantMessage) {
+			return "the answer is unavailable";
+		}
+		if (!this.#lastLeafId) return "the session has no branch point";
+		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";
+		return undefined;
+	}
+
 	canCopy(): boolean {
 		return (
 			!this.#copyInFlight && this.#activeRequest?.component.isCopyable() === true && this.#lastCopyText !== undefined
@@ -83,17 +107,34 @@ export class BtwController {
 	}
 
 	async handleBranch(): Promise<boolean> {
-		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;
+		const sessionId = this.#lastSessionId;
+		if (!request || !question || !assistantMessage || !leafId || !sessionId) return false;
+
 		this.#branchInFlight = true;
+		request.component.markBranching();
 		try {
-			await this.ctx.handleBtwBranch(this.#lastQuestion, this.#lastAssistantMessage);
+			await this.ctx.handleBtwBranch(question, assistantMessage, leafId, sessionId);
 			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,10 +160,15 @@ 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(),
+			sessionId: this.ctx.sessionManager.getSessionId(),
 		};
 		this.ctx.btwContainer.clear();
 		this.ctx.btwContainer.addChild(request.component);
@@ -156,6 +202,7 @@ export class BtwController {
 				this.#lastCopyText = copyText;
 				this.#lastAssistantMessage = assistantMessageWithReplyText(assistantMessage, replyText);
 				this.#lastLeafId = request.leafId;
+				this.#lastSessionId = request.sessionId;
 			} else {
 				this.#clearCompletedState();
 			}
@@ -190,6 +237,7 @@ export class BtwController {
 		this.#lastAssistantMessage = undefined;
 		this.#lastCopyText = undefined;
 		this.#lastLeafId = undefined;
+		this.#lastSessionId = undefined;
 	}
 
 	#isActiveRequest(request: BtwRequest): boolean {

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.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();

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<boolean> {
 		return this.#btwController.handleBranch();
 	}
@@ -4856,9 +4861,14 @@ export class InteractiveMode implements InteractiveModeContext {
 		return this.#btwController.handleCopy();
 	}
 
-	async handleBtwBranch(question: string, assistantMessage: AssistantMessage): Promise<void> {
+	async handleBtwBranch(
+		question: string,
+		assistantMessage: AssistantMessage,
+		leafId: string,
+		sessionId: string,
+	): Promise<void> {
 		try {
-			const result = await this.session.branchFromBtw(question, assistantMessage);
+			const result = await this.session.branchFromBtw(question, assistantMessage, leafId, sessionId);
 			if (result.cancelled) {
 				this.showStatus("/btw branch cancelled", { dim: true });
 				return;

packages/coding-agent/src/modes/types.ts

@@ -418,9 +418,15 @@ export interface InteractiveModeContext {
 	handleBtwEscape(): boolean;
 	handleBtwBranchKey(): Promise<boolean>;
 	canBranchBtw(): boolean;
+	handlesBtwBranchKey(): boolean;
 	canCopyBtw(): boolean;
 	handleBtwCopyKey(): Promise<boolean>;
-	handleBtwBranch(question: string, assistantMessage: AssistantMessage): Promise<void>;
+	handleBtwBranch(
+		question: string,
+		assistantMessage: AssistantMessage,
+		leafId: string,
+		sessionId: string,
+	): Promise<void>;
 	handleOmfgCommand(complaint: string): Promise<void>;
 	hasActiveOmfg(): boolean;
 	handleOmfgEscape(): boolean;

packages/coding-agent/src/session/agent-session.ts

@@ -347,6 +347,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 */
 // ============================================================================
@@ -3724,7 +3725,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) });
 		}
@@ -7618,21 +7623,24 @@ export class AgentSession {
 		}
 	}
 
+	/** 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");
 		}
 
-		const leafId = this.sessionManager.getLeafId();
-		if (!leafId) {
-			throw new Error("Cannot branch /btw: current session has no leaf");
+		if (!leafId || this.sessionManager.getSessionId() !== sessionId || this.sessionManager.getLeafId() !== leafId) {
+			throw new Error("Cannot branch /btw: session changed since /btw started");
 		}
 
 		if (
+			this.isStreaming ||
 			this.isBashRunning ||
 			this.isEvalRunning ||
 			this.isCompacting ||
@@ -7653,8 +7661,17 @@ export class AgentSession {
 			}
 		}
 
-		await this.#cancelPostPromptTasks();
+		if (this.sessionManager.getSessionId() !== sessionId || 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 ||
@@ -7667,10 +7684,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();
@@ -7684,6 +7697,9 @@ export class AgentSession {
 			advisorRecordersDetached = true;
 			await this.#advisors.drainAndDetachRecorders();
 			try {
+				if (this.sessionManager.getSessionId() !== sessionId || 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();

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,12 @@ 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),
+			activeSession.sessionManager.getSessionId(),
+		);
 
 		expect(result.cancelled).toBe(false);
 		expect(result.sessionFile).toBe(activeSession.sessionFile);
@@ -155,7 +166,12 @@ describe("AgentSession.branchFromBtw", () => {
 			return result;
 		});
 
-		const result = await activeSession.branchFromBtw("question", createBtwAssistant());
+		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");
@@ -176,7 +192,12 @@ 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),
+			activeSession.sessionManager.getSessionId(),
+		);
 
 		expect(result).toEqual({ cancelled: true, sessionFile: originalFile });
 		expect(activeSession.sessionFile).toBe(originalFile);
@@ -186,6 +207,52 @@ describe("AgentSession.branchFromBtw", () => {
 		});
 	});
 
+	it("refuses when the session leaf advances while a branch hook is pending", async () => {
+		const hookStarted = Promise.withResolvers<void>();
+		const hookRelease = Promise.withResolvers<void>();
+		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),
+			activeSession.sessionManager.getSessionId(),
+		);
+		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("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"),
@@ -197,7 +264,12 @@ 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),
+			activeSession.sessionManager.getSessionId(),
+		);
 
 		expect(result.cancelled).toBe(false);
 		const messages = activeSession.messages;
@@ -208,47 +280,35 @@ 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<void>();
 		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),
+				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);
+
+		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 +321,14 @@ 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),
+				activeSession.sessionManager.getSessionId(),
+			),
+		).rejects.toThrow("Cannot branch /btw while session maintenance or user work is still running");
 
 		activeSession.abortBash();
 		await bashPromise.catch(() => undefined);
@@ -278,9 +343,14 @@ 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),
+				activeSession.sessionManager.getSessionId(),
+			),
+		).rejects.toThrow("Cannot branch /btw while session maintenance or user work is still running");
 
 		abortController.abort();
 	});
@@ -295,12 +365,17 @@ 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),
+				activeSession.sessionManager.getSessionId(),
+			),
+		).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<void>();
 		const extensionRunner = {
 			hasHandlers: vi.fn((eventType: string) => eventType === "session_before_branch"),
@@ -312,6 +387,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 +397,32 @@ describe("AgentSession.branchFromBtw", () => {
 		});
 		expect(activeSession.hasPostPromptWork).toBe(true);
 
-		const branchPromise = activeSession.branchFromBtw("question", createBtwAssistant());
+		const branchPromise = activeSession.branchFromBtw(
+			"question",
+			createBtwAssistant(),
+			requiredLeafId(activeSession),
+			activeSession.sessionManager.getSessionId(),
+		);
 		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),
+				activeSession.sessionManager.getSessionId(),
+			),
+		).rejects.toThrow("Cannot branch /btw: session is not persisted");
 	});
 });

packages/coding-agent/test/input-controller-keybindings.test.ts

@@ -109,6 +109,8 @@ 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 handlesBtwBranchKey = vi.fn(() => false);
 	const editor: FakeEditor = {
 		setText(text: string) {
 			editorText = text;
@@ -210,7 +212,8 @@ async function createContext() {
 		toggleThinkingBlockVisibility: vi.fn(),
 		showModelSelector,
 		updateEditorBorderColor: vi.fn(),
-		hasActiveBtw: vi.fn(() => false),
+		hasActiveBtw,
+		handlesBtwBranchKey,
 		handleBtwBranchKey,
 		canBranchBtw,
 		canCopyBtw,
@@ -242,6 +245,8 @@ async function createContext() {
 			handleBtwBranchKey,
 			addInputListener,
 			canBranchBtw,
+			hasActiveBtw,
+			handlesBtwBranchKey,
 			handleBtwCopyKey,
 			canCopyBtw,
 			showError,
@@ -392,7 +397,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.handlesBtwBranchKey.mockReturnValue(true);
 		const controller = new InputController(ctx);
 
 		controller.setupKeyHandlers();
@@ -406,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();
-		(ctx.canBranchBtw as unknown as { mockReturnValue(value: boolean): void }).mockReturnValue(true);
+		spies.handlesBtwBranchKey.mockReturnValue(true);
 		editor.setText("build a branch");
 		const controller = new InputController(ctx);
 
@@ -419,8 +424,23 @@ describe("InputController keybinding setup", () => {
 		expect(spies.handleBtwBranchKey).not.toHaveBeenCalled();
 	});
 
-	it("lets b fall through when /btw is not branchable", async () => {
+	it("consumes b while a completed /btw branch is unavailable", async () => {
 		const { InputController, ctx, spies } = await createContext();
+		spies.handlesBtwBranchKey.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 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();
@@ -434,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();
-		(ctx.canBranchBtw as unknown as { mockReturnValue(value: boolean): void }).mockReturnValue(true);
+		spies.handlesBtwBranchKey.mockReturnValue(true);
 		setFocused({ pasteText: vi.fn() });
 		const controller = new InputController(ctx);
 

packages/coding-agent/test/modes/controllers/btw-controller.test.ts

@@ -45,24 +45,35 @@ function makeFakeSession(
 ): InteractiveModeContext["session"] {
 	return {
 		model: { provider: "anthropic", id: "claude-sonnet-4-5" },
+		isStreaming: false,
 		runEphemeralTurn,
 	} as unknown as InteractiveModeContext["session"];
 }
 
 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();
@@ -101,6 +112,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", () => {
@@ -226,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 () => {
@@ -241,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 () => {
@@ -253,6 +278,54 @@ 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 () => {
+		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(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", {
+			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 }));
+		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);
+		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);
+		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 () => {
@@ -267,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 () => {
@@ -307,7 +381,34 @@ describe("BtwController", () => {
 		await drainBtwRequest();
 
 		expect(await controller.handleBranch()).toBe(true);
-		expect(ctx.handleBtwBranch).toHaveBeenCalledWith("Question?", assistantMessage);
+		expect(ctx.handleBtwBranch).toHaveBeenCalledWith("Question?", assistantMessage, "leaf-1", "session-1");
+	});
+
+	it("keeps a pending branch visible and refuses to dismiss it", async () => {
+		const branch = Promise.withResolvers<void>();
+		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();
+		expect(controller.handlesBranchKey()).toBe(true);
+
+		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 +434,18 @@ 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",
+			"session-1",
+		);
 	});
 
 	it("copies the sanitized visible reply text after a complete non-empty reply", async () => {
@@ -418,14 +524,19 @@ 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",
+			"session-1",
+		);
 	});
 
 	it("ignores duplicate branch requests while branch promotion is in flight", async () => {