From a363b668ac930f11412de09450c847c1351c0d03 Mon Sep 17 00:00:00 2001 From: Mario Zechner Date: Sat, 25 Apr 2026 17:03:09 +0200 Subject: [PATCH] fix(coding-agent): clear tree navigation compaction state closes #3688 --- packages/coding-agent/CHANGELOG.md | 1 + .../coding-agent/src/core/agent-session.ts | 248 +++++++++--------- .../3688-tree-cancel-compacting.test.ts | 36 +++ 3 files changed, 165 insertions(+), 120 deletions(-) create mode 100644 packages/coding-agent/test/suite/regressions/3688-tree-cancel-compacting.test.ts diff --git a/packages/coding-agent/CHANGELOG.md b/packages/coding-agent/CHANGELOG.md index ee6dd6a0..33ed9b01 100644 --- a/packages/coding-agent/CHANGELOG.md +++ b/packages/coding-agent/CHANGELOG.md @@ -8,6 +8,7 @@ ### Fixed +- Fixed `/tree` cancellation via `session_before_tree` leaving the session stuck in compaction state ([#3688](https://github.com/badlogic/pi-mono/issues/3688)) - Fixed Escape interrupt handling when extensions hide the built-in working loader row ([#3674](https://github.com/badlogic/pi-mono/issues/3674)) - Fixed coding-agent test expectations for current default models and missing-auth guidance. diff --git a/packages/coding-agent/src/core/agent-session.ts b/packages/coding-agent/src/core/agent-session.ts index 098dbab7..9f9c4e6e 100644 --- a/packages/coding-agent/src/core/agent-session.ts +++ b/packages/coding-agent/src/core/agent-session.ts @@ -2711,136 +2711,144 @@ export class AgentSession { // Set up abort controller for summarization this._branchSummaryAbortController = new AbortController(); - let extensionSummary: { summary: string; details?: unknown } | undefined; - let fromExtension = false; - // Emit session_before_tree event - if (this._extensionRunner.hasHandlers("session_before_tree")) { - const result = (await this._extensionRunner.emit({ - type: "session_before_tree", - preparation, - signal: this._branchSummaryAbortController.signal, - })) as SessionBeforeTreeResult | undefined; + try { + let extensionSummary: { summary: string; details?: unknown } | undefined; + let fromExtension = false; - if (result?.cancel) { - return { cancelled: true }; + // Emit session_before_tree event + if (this._extensionRunner.hasHandlers("session_before_tree")) { + const result = (await this._extensionRunner.emit({ + type: "session_before_tree", + preparation, + signal: this._branchSummaryAbortController.signal, + })) as SessionBeforeTreeResult | undefined; + + if (result?.cancel) { + return { cancelled: true }; + } + + if (result?.summary && options.summarize) { + extensionSummary = result.summary; + fromExtension = true; + } + + // Allow extensions to override instructions and label + if (result?.customInstructions !== undefined) { + customInstructions = result.customInstructions; + } + if (result?.replaceInstructions !== undefined) { + replaceInstructions = result.replaceInstructions; + } + if (result?.label !== undefined) { + label = result.label; + } } - if (result?.summary && options.summarize) { - extensionSummary = result.summary; - fromExtension = true; + // Run default summarizer if needed + let summaryText: string | undefined; + let summaryDetails: unknown; + if (options.summarize && entriesToSummarize.length > 0 && !extensionSummary) { + const model = this.model!; + const { apiKey, headers } = await this._getRequiredRequestAuth(model); + const branchSummarySettings = this.settingsManager.getBranchSummarySettings(); + const result = await generateBranchSummary(entriesToSummarize, { + model, + apiKey, + headers, + signal: this._branchSummaryAbortController.signal, + customInstructions, + replaceInstructions, + reserveTokens: branchSummarySettings.reserveTokens, + }); + if (result.aborted) { + return { cancelled: true, aborted: true }; + } + if (result.error) { + throw new Error(result.error); + } + summaryText = result.summary; + summaryDetails = { + readFiles: result.readFiles || [], + modifiedFiles: result.modifiedFiles || [], + }; + } else if (extensionSummary) { + summaryText = extensionSummary.summary; + summaryDetails = extensionSummary.details; } - // Allow extensions to override instructions and label - if (result?.customInstructions !== undefined) { - customInstructions = result.customInstructions; - } - if (result?.replaceInstructions !== undefined) { - replaceInstructions = result.replaceInstructions; - } - if (result?.label !== undefined) { - label = result.label; - } - } + // Determine the new leaf position based on target type + let newLeafId: string | null; + let editorText: string | undefined; - // Run default summarizer if needed - let summaryText: string | undefined; - let summaryDetails: unknown; - if (options.summarize && entriesToSummarize.length > 0 && !extensionSummary) { - const model = this.model!; - const { apiKey, headers } = await this._getRequiredRequestAuth(model); - const branchSummarySettings = this.settingsManager.getBranchSummarySettings(); - const result = await generateBranchSummary(entriesToSummarize, { - model, - apiKey, - headers, - signal: this._branchSummaryAbortController.signal, - customInstructions, - replaceInstructions, - reserveTokens: branchSummarySettings.reserveTokens, + if (targetEntry.type === "message" && targetEntry.message.role === "user") { + // User message: leaf = parent (null if root), text goes to editor + newLeafId = targetEntry.parentId; + editorText = this._extractUserMessageText(targetEntry.message.content); + } else if (targetEntry.type === "custom_message") { + // Custom message: leaf = parent (null if root), text goes to editor + newLeafId = targetEntry.parentId; + editorText = + typeof targetEntry.content === "string" + ? targetEntry.content + : targetEntry.content + .filter((c): c is { type: "text"; text: string } => c.type === "text") + .map((c) => c.text) + .join(""); + } else { + // Non-user message: leaf = selected node + newLeafId = targetId; + } + + // Switch leaf (with or without summary) + // Summary is attached at the navigation target position (newLeafId), not the old branch + let summaryEntry: BranchSummaryEntry | undefined; + if (summaryText) { + // Create summary at target position (can be null for root) + const summaryId = this.sessionManager.branchWithSummary( + newLeafId, + summaryText, + summaryDetails, + fromExtension, + ); + summaryEntry = this.sessionManager.getEntry(summaryId) as BranchSummaryEntry; + + // Attach label to the summary entry + if (label) { + this.sessionManager.appendLabelChange(summaryId, label); + } + } else if (newLeafId === null) { + // No summary, navigating to root - reset leaf + this.sessionManager.resetLeaf(); + } else { + // No summary, navigating to non-root + this.sessionManager.branch(newLeafId); + } + + // Attach label to target entry when not summarizing (no summary entry to label) + if (label && !summaryText) { + this.sessionManager.appendLabelChange(targetId, label); + } + + // Update agent state + const sessionContext = this.sessionManager.buildSessionContext(); + this.agent.state.messages = sessionContext.messages; + + // Emit session_tree event + await this._extensionRunner.emit({ + type: "session_tree", + newLeafId: this.sessionManager.getLeafId(), + oldLeafId, + summaryEntry, + fromExtension: summaryText ? fromExtension : undefined, }); + + // Emit to custom tools + + return { editorText, cancelled: false, summaryEntry }; + } finally { this._branchSummaryAbortController = undefined; - if (result.aborted) { - return { cancelled: true, aborted: true }; - } - if (result.error) { - throw new Error(result.error); - } - summaryText = result.summary; - summaryDetails = { - readFiles: result.readFiles || [], - modifiedFiles: result.modifiedFiles || [], - }; - } else if (extensionSummary) { - summaryText = extensionSummary.summary; - summaryDetails = extensionSummary.details; } - - // Determine the new leaf position based on target type - let newLeafId: string | null; - let editorText: string | undefined; - - if (targetEntry.type === "message" && targetEntry.message.role === "user") { - // User message: leaf = parent (null if root), text goes to editor - newLeafId = targetEntry.parentId; - editorText = this._extractUserMessageText(targetEntry.message.content); - } else if (targetEntry.type === "custom_message") { - // Custom message: leaf = parent (null if root), text goes to editor - newLeafId = targetEntry.parentId; - editorText = - typeof targetEntry.content === "string" - ? targetEntry.content - : targetEntry.content - .filter((c): c is { type: "text"; text: string } => c.type === "text") - .map((c) => c.text) - .join(""); - } else { - // Non-user message: leaf = selected node - newLeafId = targetId; - } - - // Switch leaf (with or without summary) - // Summary is attached at the navigation target position (newLeafId), not the old branch - let summaryEntry: BranchSummaryEntry | undefined; - if (summaryText) { - // Create summary at target position (can be null for root) - const summaryId = this.sessionManager.branchWithSummary(newLeafId, summaryText, summaryDetails, fromExtension); - summaryEntry = this.sessionManager.getEntry(summaryId) as BranchSummaryEntry; - - // Attach label to the summary entry - if (label) { - this.sessionManager.appendLabelChange(summaryId, label); - } - } else if (newLeafId === null) { - // No summary, navigating to root - reset leaf - this.sessionManager.resetLeaf(); - } else { - // No summary, navigating to non-root - this.sessionManager.branch(newLeafId); - } - - // Attach label to target entry when not summarizing (no summary entry to label) - if (label && !summaryText) { - this.sessionManager.appendLabelChange(targetId, label); - } - - // Update agent state - const sessionContext = this.sessionManager.buildSessionContext(); - this.agent.state.messages = sessionContext.messages; - - // Emit session_tree event - await this._extensionRunner.emit({ - type: "session_tree", - newLeafId: this.sessionManager.getLeafId(), - oldLeafId, - summaryEntry, - fromExtension: summaryText ? fromExtension : undefined, - }); - - // Emit to custom tools - - this._branchSummaryAbortController = undefined; - return { editorText, cancelled: false, summaryEntry }; } /** diff --git a/packages/coding-agent/test/suite/regressions/3688-tree-cancel-compacting.test.ts b/packages/coding-agent/test/suite/regressions/3688-tree-cancel-compacting.test.ts new file mode 100644 index 00000000..67262e2c --- /dev/null +++ b/packages/coding-agent/test/suite/regressions/3688-tree-cancel-compacting.test.ts @@ -0,0 +1,36 @@ +import { afterEach, describe, expect, it } from "vitest"; +import { assistantMsg, userMsg } from "../../utilities.js"; +import { createHarness, type Harness } from "../harness.js"; + +describe("issue #3688 tree cancellation compaction state", () => { + const harnesses: Harness[] = []; + + afterEach(() => { + while (harnesses.length > 0) { + harnesses.pop()?.cleanup(); + } + }); + + it("clears branch summary state when session_before_tree cancels navigation", async () => { + const harness = await createHarness({ + extensionFactories: [ + (pi) => { + pi.on("session_before_tree", () => ({ cancel: true })); + }, + ], + }); + harnesses.push(harness); + + const targetId = harness.sessionManager.appendMessage(userMsg("first")); + harness.sessionManager.appendMessage(assistantMsg("reply")); + const currentLeafId = harness.sessionManager.appendMessage(userMsg("second")); + + expect(harness.sessionManager.getLeafId()).toBe(currentLeafId); + + const result = await harness.session.navigateTree(targetId, { summarize: false }); + + expect(result).toEqual({ cancelled: true }); + expect(harness.session.isCompacting).toBe(false); + expect(harness.sessionManager.getLeafId()).toBe(currentLeafId); + }); +});