From da29a25e54dd3aa79e0331805e2fb4a39227dd09 Mon Sep 17 00:00:00 2001
From: KamijoToma <admin@misakacloud.net>
Date: Thu, 18 Jun 2026 03:10:09 +0800
Subject: [PATCH] fix(coding-agent): cancel btw post-prompt work

---
 .../coding-agent/src/session/agent-session.ts | 11 +++++++
 .../test/agent-session-btw-branch.test.ts     | 32 +++++++++++++++++++
 2 files changed, 43 insertions(+)

diff --git a/packages/coding-agent/src/session/agent-session.ts b/packages/coding-agent/src/session/agent-session.ts
index 23af22cae..5b1be048e 100644
--- a/packages/coding-agent/src/session/agent-session.ts
+++ b/packages/coding-agent/src/session/agent-session.ts
@@ -11192,6 +11192,17 @@ export class AgentSession {
 			}
 		}
 
+		await this.#cancelPostPromptTasks();
+		if (
+			this.isBashRunning ||
+			this.isEvalRunning ||
+			this.isCompacting ||
+			this.isGeneratingHandoff ||
+			this.isRetrying
+		) {
+			throw new Error("Cannot branch /btw while session maintenance or user work is still running");
+		}
+
 		this.#pendingNextTurnMessages = [];
 		this.#scheduledHiddenNextTurnGeneration = undefined;
 		this.agent.replaceQueues([], []);
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 6f4d0e6eb..62fdb04fb 100644
--- a/packages/coding-agent/test/agent-session-btw-branch.test.ts
+++ b/packages/coding-agent/test/agent-session-btw-branch.test.ts
@@ -252,6 +252,38 @@ describe("AgentSession.branchFromBtw", () => {
 		);
 	});
 
+	it("cancels post-prompt work after branch hooks before switching sessions", async () => {
+		const hookRelease = Promise.withResolvers<void>();
+		const extensionRunner = {
+			hasHandlers: vi.fn((eventType: string) => eventType === "session_before_branch"),
+			emit: vi.fn(async () => {
+				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();
+		activeSession.queueDeferredMessage({
+			role: "custom",
+			customType: "test-hidden-message",
+			content: "hidden",
+			display: false,
+			timestamp: Date.now(),
+		});
+		expect(activeSession.hasPostPromptWork).toBe(true);
+
+		const branchPromise = activeSession.branchFromBtw("question", createBtwAssistant());
+		await Promise.resolve();
+		expect(activeSession.hasPostPromptWork).toBe(true);
+
+		hookRelease.resolve();
+		const result = await branchPromise;
+
+		expect(result.cancelled).toBe(false);
+		expect(activeSession.hasPostPromptWork).toBe(false);
+	});
+
 	it("throws for in-memory sessions", async () => {
 		const activeSession = await createSession({ persisted: false });
 		activeSession.sessionManager.appendMessage({ role: "user", content: "seed", timestamp: Date.now() });