diff --git a/packages/coding-agent/src/core/agent-session-runtime.ts b/packages/coding-agent/src/core/agent-session-runtime.ts index 93b7ee2d8..c3b60569e 100644 --- a/packages/coding-agent/src/core/agent-session-runtime.ts +++ b/packages/coding-agent/src/core/agent-session-runtime.ts @@ -165,6 +165,9 @@ export class AgentSessionRuntime { } private async teardownCurrent(reason: SessionShutdownEvent["reason"], targetSessionFile?: string): Promise { + // Settle any active response first so the aborted turn (including tool + // results) is persisted to the outgoing session before it is replaced. + await this.session.abort(); await emitSessionShutdownEvent(this.session.extensionRunner, { type: "session_shutdown", reason, diff --git a/packages/coding-agent/src/core/agent-session.ts b/packages/coding-agent/src/core/agent-session.ts index e88049855..312404d9f 100644 --- a/packages/coding-agent/src/core/agent-session.ts +++ b/packages/coding-agent/src/core/agent-session.ts @@ -2896,6 +2896,10 @@ export class AgentSession { targetId: string, options: { summarize?: boolean; customInstructions?: string; replaceInstructions?: boolean; label?: string } = {}, ): Promise<{ editorText?: string; cancelled: boolean; aborted?: boolean; summaryEntry?: BranchSummaryEntry }> { + if (this.isStreaming) { + throw new Error("Wait for the current response to finish before navigating the session tree."); + } + const oldLeafId = this.sessionManager.getLeafId(); // No-op if already at target diff --git a/packages/coding-agent/src/modes/interactive/interactive-mode.ts b/packages/coding-agent/src/modes/interactive/interactive-mode.ts index 058271937..ac9e2f80f 100644 --- a/packages/coding-agent/src/modes/interactive/interactive-mode.ts +++ b/packages/coding-agent/src/modes/interactive/interactive-mode.ts @@ -4650,7 +4650,7 @@ export class InteractiveMode { this.ui.terminal.rows, async (entryId) => { // Selecting the current leaf is a no-op (already there) - if (entryId === realLeafId) { + if (entryId === this.sessionManager.getLeafId()) { done(); this.showStatus("Already at this point"); return; @@ -4693,6 +4693,12 @@ export class InteractiveMode { } } + // The user committed to navigating: stop the active response first. + if (this.session.isStreaming) { + this.restoreQueuedMessagesToEditor(); + await this.session.abort(); + } + // Set up escape handler and status indicator if summarizing let showingSummaryIndicator = false; const originalOnEscape = this.defaultEditor.onEscape; diff --git a/packages/coding-agent/test/suite/agent-session-runtime.test.ts b/packages/coding-agent/test/suite/agent-session-runtime.test.ts index 79855e9c1..1a2a9cf67 100644 --- a/packages/coding-agent/test/suite/agent-session-runtime.test.ts +++ b/packages/coding-agent/test/suite/agent-session-runtime.test.ts @@ -1,7 +1,8 @@ import { existsSync, mkdirSync, realpathSync, rmSync } from "node:fs"; import { tmpdir } from "node:os"; import { join, parse } from "node:path"; -import { fauxAssistantMessage, registerFauxProvider } from "@earendil-works/pi-ai/compat"; +import { fauxAssistantMessage, fauxToolCall, registerFauxProvider } from "@earendil-works/pi-ai/compat"; +import { Type } from "typebox"; import { afterEach, describe, expect, it } from "vitest"; import { type CreateAgentSessionRuntimeFactory, @@ -12,6 +13,7 @@ import { import { AuthStorage } from "../../src/core/auth-storage.ts"; import { SessionManager } from "../../src/core/session-manager.ts"; import type { + AgentToolResult, ExtensionAPI, ExtensionFactory, SessionBeforeForkEvent, @@ -161,6 +163,55 @@ describe("AgentSessionRuntime characterization", () => { expect(persistedAssistant.usage.cost.total).toBe(0.123); }); + it("settles the active response before session replacement", async () => { + let toolStarted!: () => void; + const toolStartedPromise = new Promise((resolve) => { + toolStarted = resolve; + }); + const { runtime, faux } = await createRuntimeForTest((pi: ExtensionAPI) => { + pi.registerTool({ + name: "block", + label: "Block", + description: "Blocks until aborted", + parameters: Type.Object({}), + execute: (_toolCallId, _params, signal) => + new Promise>((resolve) => { + toolStarted(); + signal?.addEventListener("abort", () => + resolve({ content: [{ type: "text", text: "tool aborted" }], details: {} }), + ); + }), + }); + }); + + await runtime.session.prompt("hello"); + const firstSessionFile = runtime.session.sessionFile!; + await runtime.newSession(); + await runtime.session.bindExtensions({}); + + faux.setResponses([fauxAssistantMessage(fauxToolCall("block", {}), { stopReason: "toolUse" })]); + const outgoingSession = runtime.session; + const promptPromise = outgoingSession.prompt("start blocking tool"); + await toolStartedPromise; + + const switchResult = await runtime.switchSession(firstSessionFile); + await promptPromise; + + expect(switchResult.cancelled).toBe(false); + expect(runtime.session.sessionFile).toBe(firstSessionFile); + // The outgoing session settled before replacement: the interrupted tool + // call has a persisted tool result instead of dangling forever. + const outgoingEntries = SessionManager.open(outgoingSession.sessionFile!) + .getEntries() + .filter((entry) => entry.type === "message"); + expect(outgoingEntries.map((entry) => entry.message.role)).toEqual([ + "user", + "assistant", + "toolResult", + "assistant", + ]); + }); + it("emits session_before_switch and session_start for new and resume flows", async () => { const events: RecordedSessionEvent[] = []; const { runtime } = await createRuntimeForTest((pi: ExtensionAPI) => { diff --git a/packages/coding-agent/test/suite/regressions/tree-during-streaming.test.ts b/packages/coding-agent/test/suite/regressions/tree-during-streaming.test.ts new file mode 100644 index 000000000..e38fa1344 --- /dev/null +++ b/packages/coding-agent/test/suite/regressions/tree-during-streaming.test.ts @@ -0,0 +1,35 @@ +import { fauxAssistantMessage } from "@earendil-works/pi-ai"; +import { describe, expect, it } from "vitest"; +import { userMsg } from "../../utilities.ts"; +import { createHarness } from "../harness.ts"; + +describe("tree navigation during an active response", () => { + it("rejects navigation without changing the active leaf", async () => { + const harness = await createHarness(); + const targetId = harness.sessionManager.appendMessage(userMsg("first")); + let navigationResult: unknown; + let leafUnchanged = false; + + try { + // Navigate from inside the response factory, while the run is active. + harness.setResponses([ + async () => { + const activeLeafId = harness.sessionManager.getLeafId(); + navigationResult = await harness.session + .navigateTree(targetId, { summarize: false }) + .catch((error) => error); + leafUnchanged = activeLeafId !== targetId && harness.sessionManager.getLeafId() === activeLeafId; + return fauxAssistantMessage("response"); + }, + ]); + await harness.session.prompt("second"); + + expect(navigationResult).toEqual( + new Error("Wait for the current response to finish before navigating the session tree."), + ); + expect(leafUnchanged).toBe(true); + } finally { + harness.cleanup(); + } + }); +});