diff --git a/packages/coding-agent/src/core/agent-session-runtime.ts b/packages/coding-agent/src/core/agent-session-runtime.ts index c3b60569e..d2eb9a3fb 100644 --- a/packages/coding-agent/src/core/agent-session-runtime.ts +++ b/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, diff --git a/packages/coding-agent/test/suite/regressions/8724-in-memory-fork-active-tool.test.ts b/packages/coding-agent/test/suite/regressions/8724-in-memory-fork-active-tool.test.ts new file mode 100644 index 000000000..0868b49a0 --- /dev/null +++ b/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> = []; + + 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((resolve) => { + markToolStarted = resolve; + }); + const blockingTool: AgentTool = { + name: "block", + label: "Block", + description: "Wait until aborted", + parameters: Type.Object({}), + execute: (_toolCallId, _params, signal) => + new Promise>((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"]); + }); +});