fix(coding-agent): settle active turn before in-memory fork (#8937)

Fixes #8724

Co-authored-by: yaogangqiang <yaogangqiang@kanyun.com>

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

@@ -333,12 +333,12 @@ export class AgentSessionRuntime {
 		}
 
 		const sessionManager = this.session.sessionManager;
+		await this.teardownCurrent("fork", sessionManager.getSessionFile());
 		if (!targetLeafId) {
-			sessionManager.newSession({ parentSession: this.session.sessionFile });
+			sessionManager.newSession({ parentSession: previousSessionFile });
 		} else {
 			sessionManager.createBranchedSession(targetLeafId);
 		}
-		await this.teardownCurrent("fork", sessionManager.getSessionFile());
 		this.apply(
 			await this.createRuntime({
 				cwd: this.cwd,

packages/coding-agent/test/suite/regressions/8724-in-memory-fork-active-tool.test.ts

@@ -0,0 +1,100 @@
+import type { AgentTool, AgentToolResult } from "@earendil-works/pi-agent-core";
+import { fauxAssistantMessage, fauxToolCall } from "@earendil-works/pi-ai";
+import { Type } from "typebox";
+import { afterEach, describe, expect, it } from "vitest";
+import {
+	AgentSessionRuntime,
+	type AgentSessionServices,
+	type CreateAgentSessionRuntimeFactory,
+	createAgentSessionFromServices,
+} from "../../../src/core/agent-session-runtime.ts";
+import { createHarness } from "../harness.ts";
+
+describe("regression #8724: in-memory fork during an active tool turn", () => {
+	const cleanups: Array<() => Promise<void> | void> = [];
+
+	afterEach(async () => {
+		while (cleanups.length > 0) {
+			await cleanups.pop()?.();
+		}
+	});
+
+	it("does not append the aborted turn to the replacement session", async () => {
+		let markToolStarted = () => {};
+		const toolStarted = new Promise<void>((resolve) => {
+			markToolStarted = resolve;
+		});
+		const blockingTool: AgentTool = {
+			name: "block",
+			label: "Block",
+			description: "Wait until aborted",
+			parameters: Type.Object({}),
+			execute: (_toolCallId, _params, signal) =>
+				new Promise<AgentToolResult<unknown>>((resolve) => {
+					markToolStarted();
+					signal?.addEventListener(
+						"abort",
+						() => resolve({ content: [{ type: "text", text: "tool aborted" }], details: {} }),
+						{ once: true },
+					);
+				}),
+		};
+		const harness = await createHarness({ tools: [blockingTool] });
+		const services: AgentSessionServices = {
+			cwd: harness.tempDir,
+			agentDir: harness.tempDir,
+			modelRuntime: harness.session.modelRuntime,
+			settingsManager: harness.settingsManager,
+			resourceLoader: harness.session.resourceLoader,
+			diagnostics: [],
+		};
+		const createRuntime: CreateAgentSessionRuntimeFactory = async ({ sessionManager, sessionStartEvent }) => ({
+			...(await createAgentSessionFromServices({
+				services,
+				sessionManager,
+				sessionStartEvent,
+				model: harness.getModel(),
+				noTools: "all",
+			})),
+			services,
+			diagnostics: [],
+		});
+		const runtime = new AgentSessionRuntime(harness.session, services, createRuntime);
+		cleanups.push(async () => {
+			if (runtime.session !== harness.session) {
+				await runtime.dispose();
+			}
+			harness.cleanup();
+		});
+
+		harness.setResponses([
+			fauxAssistantMessage("first response"),
+			fauxAssistantMessage(fauxToolCall("block", {}), { stopReason: "toolUse" }),
+			fauxAssistantMessage("unused after abort"),
+		]);
+		await runtime.session.prompt("first prompt");
+		const firstUserEntryId = runtime.session.getUserMessagesForForking()[0]?.entryId;
+		expect(firstUserEntryId).toBeDefined();
+
+		const outgoingPrompt = runtime.session.prompt("start blocking tool");
+		await toolStarted;
+		const forkResult = await runtime.fork(firstUserEntryId!);
+		await outgoingPrompt;
+		await runtime.session.bindExtensions({});
+
+		expect(forkResult).toEqual({ cancelled: false, selectedText: "first prompt" });
+		expect(runtime.session.messages).toEqual([]);
+		expect(runtime.session.sessionManager.getEntries().filter((entry) => entry.type === "message")).toEqual([]);
+
+		let capturedRoles: string[] = [];
+		harness.setResponses([
+			(context) => {
+				capturedRoles = context.messages.map((message) => message.role);
+				return fauxAssistantMessage("next response");
+			},
+		]);
+		await runtime.session.prompt("next prompt");
+
+		expect(capturedRoles).toEqual(["user"]);
+	});
+});