From cefa40ed8f8dbcd6ffd67b5e7a0eaa75a5050d1a Mon Sep 17 00:00:00 2001 From: Thomas Mustier <6326440+tmustier@users.noreply.github.com> Date: Tue, 28 Jul 2026 13:46:30 +0100 Subject: [PATCH] ***WIP it's a PoC *** fix(coding-agent): guard tree navigation during responses (#7022) * fix(coding-agent): guard tree navigation during responses * fix(coding-agent): simplify tree navigation regression * fix(coding-agent): abort active run on session replacement and tree navigation Session replacement (resume, new, fork, import) now settles the active response in teardownCurrent before disposing the outgoing session, so the aborted turn and its tool results persist instead of leaving a dangling tool call. Tree navigation aborts the active response directly once the user commits, without a confirmation gate. --- .../src/core/agent-session-runtime.ts | 3 ++ .../coding-agent/src/core/agent-session.ts | 4 ++ .../src/modes/interactive/interactive-mode.ts | 8 ++- .../test/suite/agent-session-runtime.test.ts | 53 ++++++++++++++++++- .../regressions/tree-during-streaming.test.ts | 35 ++++++++++++ 5 files changed, 101 insertions(+), 2 deletions(-) create mode 100644 packages/coding-agent/test/suite/regressions/tree-during-streaming.test.ts 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(); + } + }); +});