From 32bcdc9739d4b806c209a584c6c81a7a22366482 Mon Sep 17 00:00:00 2001 From: Mario Zechner Date: Tue, 19 May 2026 09:52:58 +0200 Subject: [PATCH] fix(coding-agent): simplify agent session settlement --- packages/coding-agent/CHANGELOG.md | 1 + .../coding-agent/src/core/agent-session.ts | 236 ++++++------------ .../agent-session-bash-persistence.test.ts | 4 +- .../suite/agent-session-compaction.test.ts | 24 +- ...113-agent-session-event-settlement.test.ts | 95 +++++++ 5 files changed, 187 insertions(+), 173 deletions(-) create mode 100644 packages/coding-agent/test/suite/regressions/1717-2113-agent-session-event-settlement.test.ts diff --git a/packages/coding-agent/CHANGELOG.md b/packages/coding-agent/CHANGELOG.md index 7ff38a6a..147b3b19 100644 --- a/packages/coding-agent/CHANGELOG.md +++ b/packages/coding-agent/CHANGELOG.md @@ -4,6 +4,7 @@ ### Fixed +- Fixed AgentSession retry, compaction, and event settlement to use the awaited agent lifecycle instead of a separate event queue. - Fixed the subagent extension's parallel mode to return useful per-task output and failed-task diagnostics to the parent model instead of 100-character previews ([#4710](https://github.com/earendil-works/pi/issues/4710)). - Fixed Windows local bash execution to hide helper console windows when launched from background SDK processes ([#4699](https://github.com/earendil-works/pi/issues/4699)). diff --git a/packages/coding-agent/src/core/agent-session.ts b/packages/coding-agent/src/core/agent-session.ts index 1d51163b..2f5b805d 100644 --- a/packages/coding-agent/src/core/agent-session.ts +++ b/packages/coding-agent/src/core/agent-session.ts @@ -252,7 +252,6 @@ export class AgentSession { // Event subscription state private _unsubscribeAgent?: () => void; private _eventListeners: AgentSessionEventListener[] = []; - private _agentEventQueue: Promise = Promise.resolve(); /** Tracks pending steering messages for UI display. Removed when delivered. */ private _steeringMessages: string[] = []; @@ -272,8 +271,6 @@ export class AgentSession { // Retry state private _retryAbortController: AbortController | undefined = undefined; private _retryAttempt = 0; - private _retryPromise: Promise | undefined = undefined; - private _retryResolve: (() => void) | undefined = undefined; // Bash execution state private _bashAbortController: AbortController | undefined = undefined; @@ -395,8 +392,6 @@ export class AgentSession { return undefined; } - await this._agentEventQueue; - try { return await runner.emitToolCall({ type: "tool_call", @@ -463,54 +458,7 @@ export class AgentSession { private _lastAssistantMessage: AssistantMessage | undefined = undefined; /** Internal handler for agent events - shared by subscribe and reconnect */ - private _handleAgentEvent = (event: AgentEvent): void => { - // Create retry promise synchronously before queueing async processing. - // Agent.emit() calls this handler synchronously, and prompt() calls waitForRetry() - // as soon as agent.prompt() resolves. If _retryPromise is created only inside - // _processAgentEvent, slow earlier queued events can delay agent_end processing - // and waitForRetry() can miss the in-flight retry. - this._createRetryPromiseForAgentEnd(event); - - this._agentEventQueue = this._agentEventQueue.then( - () => this._processAgentEvent(event), - () => this._processAgentEvent(event), - ); - - // Keep queue alive if an event handler fails - this._agentEventQueue.catch(() => {}); - }; - - private _createRetryPromiseForAgentEnd(event: AgentEvent): void { - if (event.type !== "agent_end" || this._retryPromise) { - return; - } - - const settings = this.settingsManager.getRetrySettings(); - if (!settings.enabled) { - return; - } - - const lastAssistant = this._findLastAssistantInMessages(event.messages); - if (!lastAssistant || !this._isRetryableError(lastAssistant)) { - return; - } - - this._retryPromise = new Promise((resolve) => { - this._retryResolve = resolve; - }); - } - - private _findLastAssistantInMessages(messages: AgentMessage[]): AssistantMessage | undefined { - for (let i = messages.length - 1; i >= 0; i--) { - const message = messages[i]; - if (message.role === "assistant") { - return message as AssistantMessage; - } - } - return undefined; - } - - private async _processAgentEvent(event: AgentEvent): Promise { + private _handleAgentEvent = async (event: AgentEvent): Promise => { // When a user message starts, check if it's from either queue and remove it BEFORE emitting // This ensures the UI sees the updated queue state if (event.type === "message_start" && event.message.role === "user") { @@ -581,31 +529,7 @@ export class AgentSession { } } } - - // Check auto-retry and auto-compaction after agent completes - if (event.type === "agent_end" && this._lastAssistantMessage) { - const msg = this._lastAssistantMessage; - this._lastAssistantMessage = undefined; - - // Check for retryable errors first (overloaded, rate limit, server errors) - if (this._isRetryableError(msg)) { - const didRetry = await this._handleRetryableError(msg); - if (didRetry) return; // Retry was initiated, don't proceed to compaction - } - - this._resolveRetry(); - await this._checkCompaction(msg); - } - } - - /** Resolve the pending retry promise */ - private _resolveRetry(): void { - if (this._retryResolve) { - this._retryResolve(); - this._retryResolve = undefined; - this._retryPromise = undefined; - } - } + }; /** Extract text content from a message */ private _getUserMessageText(message: Message): string { @@ -630,7 +554,7 @@ export class AgentSession { private _replaceMessageInPlace(target: AgentMessage, replacement: AgentMessage): void { // Agent-core stores the finalized message object in its state before emitting message_end. - // SessionManager persistence happens later in _processAgentEvent() with event.message. + // SessionManager persistence happens later in _handleAgentEvent() with event.message. // Mutating this object in place keeps agent state, later turn/agent events, listeners, // and the eventual SessionManager.appendMessage(event.message) persistence in sync. if (target === replacement) { @@ -968,6 +892,41 @@ export class AgentSession { // Prompting // ========================================================================= + private async _runAgentPrompt(messages: AgentMessage | AgentMessage[]): Promise { + try { + await this.agent.prompt(messages); + while (await this._handlePostAgentRun()) { + await this.agent.continue(); + } + } finally { + this._flushPendingBashMessages(); + } + } + + private async _handlePostAgentRun(): Promise { + const msg = this._lastAssistantMessage; + this._lastAssistantMessage = undefined; + if (!msg) { + return false; + } + + if (this._isRetryableError(msg) && (await this._prepareRetry(msg))) { + return true; + } + + if (msg.stopReason === "error" && this._retryAttempt > 0) { + this._emit({ + type: "auto_retry_end", + success: false, + attempt: this._retryAttempt, + finalError: msg.errorMessage, + }); + this._retryAttempt = 0; + } + + return await this._checkCompaction(msg); + } + /** * Send a prompt to the agent. * - Handles extension commands (registered via pi.registerCommand) immediately, even during streaming @@ -1058,8 +1017,15 @@ export class AgentSession { // Check if we need to compact before sending (catches aborted responses) const lastAssistant = this._findLastAssistantMessage(); - if (lastAssistant) { - await this._checkCompaction(lastAssistant, false); + if (lastAssistant && (await this._checkCompaction(lastAssistant, false))) { + try { + await this.agent.continue(); + while (await this._handlePostAgentRun()) { + await this.agent.continue(); + } + } finally { + this._flushPendingBashMessages(); + } } // Build messages array (custom message if any, then user message) @@ -1119,8 +1085,7 @@ export class AgentSession { } preflightResult?.(true); - await this.agent.prompt(messages); - await this.waitForRetry(); + await this._runAgentPrompt(messages); } /** @@ -1306,7 +1271,7 @@ export class AgentSession { this.agent.steer(appMessage); } } else if (options?.triggerTurn) { - await this.agent.prompt(appMessage); + await this._runAgentPrompt(appMessage); } else { this.agent.state.messages.push(appMessage); this.sessionManager.appendCustomMessageEntry( @@ -1777,12 +1742,12 @@ export class AgentSession { * @param assistantMessage The assistant message to check * @param skipAbortedCheck If false, include aborted messages (for pre-prompt check). Default: true */ - private async _checkCompaction(assistantMessage: AssistantMessage, skipAbortedCheck = true): Promise { + private async _checkCompaction(assistantMessage: AssistantMessage, skipAbortedCheck = true): Promise { const settings = this.settingsManager.getCompactionSettings(); - if (!settings.enabled) return; + if (!settings.enabled) return false; // Skip if message was aborted (user cancelled) - unless skipAbortedCheck is false - if (skipAbortedCheck && assistantMessage.stopReason === "aborted") return; + if (skipAbortedCheck && assistantMessage.stopReason === "aborted") return false; const contextWindow = this.model?.contextWindow ?? 0; @@ -1800,7 +1765,7 @@ export class AgentSession { const assistantIsFromBeforeCompaction = compactionEntry !== null && assistantMessage.timestamp <= new Date(compactionEntry.timestamp).getTime(); if (assistantIsFromBeforeCompaction) { - return; + return false; } // Case 1: Overflow - LLM returned context overflow error @@ -1815,7 +1780,7 @@ export class AgentSession { errorMessage: "Context overflow recovery failed after one compact-and-retry attempt. Try reducing context or switching to a larger-context model.", }); - return; + return false; } this._overflowRecoveryAttempted = true; @@ -1825,8 +1790,7 @@ export class AgentSession { if (messages.length > 0 && messages[messages.length - 1].role === "assistant") { this.agent.state.messages = messages.slice(0, -1); } - await this._runAutoCompaction("overflow", true); - return; + return await this._runAutoCompaction("overflow", true); } // Case 2: Threshold - context is getting large @@ -1836,7 +1800,7 @@ export class AgentSession { if (assistantMessage.stopReason === "error") { const messages = this.agent.state.messages; const estimate = estimateContextTokens(messages); - if (estimate.lastUsageIndex === null) return; // No usage data at all + if (estimate.lastUsageIndex === null) return false; // No usage data at all // Verify the usage source is post-compaction. Kept pre-compaction messages // have stale usage reflecting the old (larger) context and would falsely // trigger compaction right after one just finished. @@ -1846,21 +1810,22 @@ export class AgentSession { usageMsg.role === "assistant" && (usageMsg as AssistantMessage).timestamp <= new Date(compactionEntry.timestamp).getTime() ) { - return; + return false; } contextTokens = estimate.tokens; } else { contextTokens = calculateContextTokens(assistantMessage.usage); } if (shouldCompact(contextTokens, contextWindow, settings)) { - await this._runAutoCompaction("threshold", false); + return await this._runAutoCompaction("threshold", false); } + return false; } /** * Internal: Run auto-compaction with events. */ - private async _runAutoCompaction(reason: "overflow" | "threshold", willRetry: boolean): Promise { + private async _runAutoCompaction(reason: "overflow" | "threshold", willRetry: boolean): Promise { const settings = this.settingsManager.getCompactionSettings(); this._emit({ type: "compaction_start", reason }); @@ -1875,7 +1840,7 @@ export class AgentSession { aborted: false, willRetry: false, }); - return; + return false; } let apiKey: string | undefined; @@ -1890,7 +1855,7 @@ export class AgentSession { aborted: false, willRetry: false, }); - return; + return false; } apiKey = authResult.apiKey; headers = authResult.headers; @@ -1909,7 +1874,7 @@ export class AgentSession { aborted: false, willRetry: false, }); - return; + return false; } let extensionCompaction: CompactionResult | undefined; @@ -1932,7 +1897,7 @@ export class AgentSession { aborted: true, willRetry: false, }); - return; + return false; } if (extensionResult?.compaction) { @@ -1978,7 +1943,7 @@ export class AgentSession { aborted: true, willRetry: false, }); - return; + return false; } this.sessionManager.appendCompaction(summary, firstKeptEntryId, tokensBefore, details, fromExtension); @@ -2013,17 +1978,12 @@ export class AgentSession { if (lastMsg?.role === "assistant" && (lastMsg as AssistantMessage).stopReason === "error") { this.agent.state.messages = messages.slice(0, -1); } - - setTimeout(() => { - this.agent.continue().catch(() => {}); - }, 100); - } else if (this.agent.hasQueuedMessages()) { - // Auto-compaction can complete while follow-up/steering/custom messages are waiting. - // Kick the loop so queued messages are actually delivered. - setTimeout(() => { - this.agent.continue().catch(() => {}); - }, 100); + return true; } + + // Auto-compaction can complete while follow-up/steering/custom messages are waiting. + // Continue once so queued messages are delivered. + return this.agent.hasQueuedMessages(); } catch (error) { const errorMessage = error instanceof Error ? error.message : "compaction failed"; this._emit({ @@ -2037,6 +1997,7 @@ export class AgentSession { ? `Context overflow recovery failed: ${errorMessage}` : `Auto-compaction failed: ${errorMessage}`, }); + return false; } finally { this._autoCompactionAbortController = undefined; } @@ -2448,36 +2409,20 @@ export class AgentSession { } /** - * Handle retryable errors with exponential backoff. - * @returns true if retry was initiated, false if max retries exceeded or disabled + * Prepare a retryable error for continuation with exponential backoff. + * @returns true if the caller should continue the agent, false otherwise */ - private async _handleRetryableError(message: AssistantMessage): Promise { + private async _prepareRetry(message: AssistantMessage): Promise { const settings = this.settingsManager.getRetrySettings(); if (!settings.enabled) { - this._resolveRetry(); return false; } - // Retry promise is created synchronously in _handleAgentEvent for agent_end. - // Keep a defensive fallback here in case a future refactor bypasses that path. - if (!this._retryPromise) { - this._retryPromise = new Promise((resolve) => { - this._retryResolve = resolve; - }); - } - this._retryAttempt++; if (this._retryAttempt > settings.maxRetries) { - // Max retries exceeded, emit final failure and reset - this._emit({ - type: "auto_retry_end", - success: false, - attempt: this._retryAttempt - 1, - finalError: message.errorMessage, - }); - this._retryAttempt = 0; - this._resolveRetry(); // Resolve so waitForRetry() completes + // Preserve the completed attempt count so post-run handling can emit the final failure. + this._retryAttempt--; return false; } @@ -2505,24 +2450,16 @@ export class AgentSession { // Aborted during sleep - emit end event so UI can clean up const attempt = this._retryAttempt; this._retryAttempt = 0; - this._retryAbortController = undefined; this._emit({ type: "auto_retry_end", success: false, attempt, finalError: "Retry cancelled", }); - this._resolveRetry(); return false; + } finally { + this._retryAbortController = undefined; } - this._retryAbortController = undefined; - - // Retry via continue() - use setTimeout to break out of event handler chain - setTimeout(() => { - this.agent.continue().catch(() => { - // Retry failed - will be caught by next agent_end - }); - }, 0); return true; } @@ -2532,26 +2469,11 @@ export class AgentSession { */ abortRetry(): void { this._retryAbortController?.abort(); - // Note: _retryAttempt is reset in the catch block of _autoRetry - this._resolveRetry(); - } - - /** - * Wait for any in-progress retry to complete. - * Returns immediately if no retry is in progress. - */ - private async waitForRetry(): Promise { - if (!this._retryPromise) { - return; - } - - await this._retryPromise; - await this.agent.waitForIdle(); } /** Whether auto-retry is currently in progress */ get isRetrying(): boolean { - return this._retryPromise !== undefined; + return this._retryAbortController !== undefined; } /** Whether auto-retry is enabled */ diff --git a/packages/coding-agent/test/suite/agent-session-bash-persistence.test.ts b/packages/coding-agent/test/suite/agent-session-bash-persistence.test.ts index a597903c..f59146b4 100644 --- a/packages/coding-agent/test/suite/agent-session-bash-persistence.test.ts +++ b/packages/coding-agent/test/suite/agent-session-bash-persistence.test.ts @@ -85,8 +85,8 @@ describe("AgentSession bash and persistence characterization", () => { releaseToolExecution?.(); await firstPrompt; - expect(harness.session.hasPendingBashMessages).toBe(true); - expect(harness.session.messages.some((message) => message.role === "bashExecution")).toBe(false); + expect(harness.session.hasPendingBashMessages).toBe(false); + expect(harness.session.messages.some((message) => message.role === "bashExecution")).toBe(true); await harness.session.prompt("next turn"); diff --git a/packages/coding-agent/test/suite/agent-session-compaction.test.ts b/packages/coding-agent/test/suite/agent-session-compaction.test.ts index e2d28911..8ca5d0e5 100644 --- a/packages/coding-agent/test/suite/agent-session-compaction.test.ts +++ b/packages/coding-agent/test/suite/agent-session-compaction.test.ts @@ -8,8 +8,8 @@ import { afterEach, describe, expect, it, vi } from "vitest"; import { createHarness, type Harness } from "./harness.js"; type SessionWithCompactionInternals = { - _checkCompaction: (assistantMessage: AssistantMessage, skipAbortedCheck?: boolean) => Promise; - _runAutoCompaction: (reason: "overflow" | "threshold", willRetry: boolean) => Promise; + _checkCompaction: (assistantMessage: AssistantMessage, skipAbortedCheck?: boolean) => Promise; + _runAutoCompaction: (reason: "overflow" | "threshold", willRetry: boolean) => Promise; }; function createUsage(totalTokens: number) { @@ -217,13 +217,9 @@ describe("AgentSession compaction characterization", () => { timestamp: Date.now(), }); - const continueSpy = vi.spyOn(harness.session.agent, "continue").mockResolvedValue(); const sessionInternals = harness.session as unknown as SessionWithCompactionInternals; - await sessionInternals._runAutoCompaction("threshold", false); - await vi.advanceTimersByTimeAsync(100); - - expect(continueSpy).toHaveBeenCalledTimes(1); + await expect(sessionInternals._runAutoCompaction("threshold", false)).resolves.toBe(true); }); it("does not retry overflow recovery more than once", async () => { @@ -235,7 +231,7 @@ describe("AgentSession compaction characterization", () => { errorMessage: "prompt is too long", timestamp: Date.now(), }); - const runAutoCompactionSpy = vi.spyOn(sessionInternals, "_runAutoCompaction").mockResolvedValue(); + const runAutoCompactionSpy = vi.spyOn(sessionInternals, "_runAutoCompaction").mockResolvedValue(false); const compactionErrors: string[] = []; harness.session.subscribe((event) => { if (event.type === "compaction_end" && event.errorMessage) { @@ -283,7 +279,7 @@ describe("AgentSession compaction characterization", () => { timestamp: Date.now(), }); - const runAutoCompactionSpy = vi.spyOn(sessionInternals, "_runAutoCompaction").mockResolvedValue(); + const runAutoCompactionSpy = vi.spyOn(sessionInternals, "_runAutoCompaction").mockResolvedValue(false); await sessionInternals._checkCompaction(staleAssistant, false); @@ -311,7 +307,7 @@ describe("AgentSession compaction characterization", () => { errorAssistant, ]; - const runAutoCompactionSpy = vi.spyOn(sessionInternals, "_runAutoCompaction").mockResolvedValue(); + const runAutoCompactionSpy = vi.spyOn(sessionInternals, "_runAutoCompaction").mockResolvedValue(false); await sessionInternals._checkCompaction(errorAssistant); @@ -332,7 +328,7 @@ describe("AgentSession compaction characterization", () => { errorAssistant, ]; - const runAutoCompactionSpy = vi.spyOn(sessionInternals, "_runAutoCompaction").mockResolvedValue(); + const runAutoCompactionSpy = vi.spyOn(sessionInternals, "_runAutoCompaction").mockResolvedValue(false); await sessionInternals._checkCompaction(errorAssistant); @@ -377,7 +373,7 @@ describe("AgentSession compaction characterization", () => { errorAssistant, ]; - const runAutoCompactionSpy = vi.spyOn(sessionInternals, "_runAutoCompaction").mockResolvedValue(); + const runAutoCompactionSpy = vi.spyOn(sessionInternals, "_runAutoCompaction").mockResolvedValue(false); await sessionInternals._checkCompaction(errorAssistant); @@ -395,8 +391,8 @@ describe("AgentSession compaction characterization", () => { const belowThresholdInternals = belowThresholdHarness.session as unknown as SessionWithCompactionInternals; const disabledInternals = disabledHarness.session as unknown as SessionWithCompactionInternals; - const belowThresholdSpy = vi.spyOn(belowThresholdInternals, "_runAutoCompaction").mockResolvedValue(); - const disabledSpy = vi.spyOn(disabledInternals, "_runAutoCompaction").mockResolvedValue(); + const belowThresholdSpy = vi.spyOn(belowThresholdInternals, "_runAutoCompaction").mockResolvedValue(false); + const disabledSpy = vi.spyOn(disabledInternals, "_runAutoCompaction").mockResolvedValue(false); await belowThresholdInternals._checkCompaction( createAssistant(belowThresholdHarness, { stopReason: "stop", totalTokens: 1_000, timestamp: Date.now() }), diff --git a/packages/coding-agent/test/suite/regressions/1717-2113-agent-session-event-settlement.test.ts b/packages/coding-agent/test/suite/regressions/1717-2113-agent-session-event-settlement.test.ts new file mode 100644 index 00000000..bea38c9e --- /dev/null +++ b/packages/coding-agent/test/suite/regressions/1717-2113-agent-session-event-settlement.test.ts @@ -0,0 +1,95 @@ +import type { AgentTool } 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 { createHarness, type Harness } from "../harness.js"; + +function createEchoTool(): AgentTool { + return { + name: "echo", + label: "Echo", + description: "Echo text back", + parameters: Type.Object({ text: Type.String() }), + execute: async (_toolCallId, params) => { + const text = typeof params === "object" && params !== null && "text" in params ? String(params.text) : ""; + return { content: [{ type: "text", text }], details: { text } }; + }, + }; +} + +describe("regressions #1717/#2113: agent session event settlement", () => { + const harnesses: Harness[] = []; + + afterEach(() => { + while (harnesses.length > 0) { + harnesses.pop()?.cleanup(); + } + }); + + it("keeps persisted assistant/toolResult message order when extension message_end handlers yield", async () => { + const harness = await createHarness({ + tools: [createEchoTool()], + extensionFactories: [ + (pi) => { + pi.on("message_end", async (event) => { + if (event.message.role === "assistant") { + await new Promise((resolve) => setTimeout(resolve, 20)); + } + }); + }, + ], + }); + harnesses.push(harness); + harness.setResponses([ + fauxAssistantMessage([fauxToolCall("echo", { text: "one" }), fauxToolCall("echo", { text: "two" })], { + stopReason: "toolUse", + }), + fauxAssistantMessage("done"), + ]); + await harness.session.prompt("run tools"); + + const branchMessages = harness.sessionManager + .getBranch() + .filter((entry) => entry.type === "message") + .map((entry) => entry.message); + expect(branchMessages.map((message) => message.role)).toEqual([ + "user", + "assistant", + "toolResult", + "toolResult", + "assistant", + ]); + const firstToolResultIndex = branchMessages.findIndex((message) => message.role === "toolResult"); + expect(firstToolResultIndex).toBeGreaterThan(0); + expect(branchMessages[firstToolResultIndex - 1]?.role).toBe("assistant"); + }); + + it("runs tool_call handlers after the assistant tool-use message is settled in the session", async () => { + let harness: Harness; + const branchRolesAtToolCall: string[][] = []; + harness = await createHarness({ + tools: [createEchoTool()], + extensionFactories: [ + (pi) => { + pi.on("tool_call", () => { + branchRolesAtToolCall.push( + harness.sessionManager + .getBranch() + .filter((entry) => entry.type === "message") + .map((entry) => entry.message.role), + ); + }); + }, + ], + }); + harnesses.push(harness); + harness.setResponses([ + fauxAssistantMessage([fauxToolCall("echo", { text: "hello" })], { stopReason: "toolUse" }), + fauxAssistantMessage("done"), + ]); + + await harness.session.prompt("run tool"); + + expect(branchRolesAtToolCall).toEqual([["user", "assistant"]]); + }); +});