shumengya/sproutclaw
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fix(coding-agent): simplify agent session settlement

Browse Source
This commit is contained in:
Mario Zechner
2026-05-19 09:52:58 +02:00
parent 6a21a52d91
commit 32bcdc9739
5 changed files with 187 additions and 173 deletions
Download Patch File Download Diff File Expand all files Collapse all files

1
packages/coding-agent/CHANGELOG.md
View File

@@ -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)).

236
packages/coding-agent/src/core/agent-session.ts
View File

@@ -252,7 +252,6 @@ export class AgentSession {
// Event subscription state
private _unsubscribeAgent?: () => void;
private _eventListeners: AgentSessionEventListener[] = [];
private _agentEventQueue: Promise<void> = 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<void> | 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<void> {
private _handleAgentEvent = async (event: AgentEvent): Promise<void> => {
// 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<void> {
try {
await this.agent.prompt(messages);
while (await this._handlePostAgentRun()) {
await this.agent.continue();
}
} finally {
this._flushPendingBashMessages();
}
}
private async _handlePostAgentRun(): Promise<boolean> {
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<void> {
private async _checkCompaction(assistantMessage: AssistantMessage, skipAbortedCheck = true): Promise<boolean> {
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<void> {
private async _runAutoCompaction(reason: "overflow" | "threshold", willRetry: boolean): Promise<boolean> {
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<boolean> {
private async _prepareRetry(message: AssistantMessage): Promise<boolean> {
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<void> {
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 */

4
packages/coding-agent/test/suite/agent-session-bash-persistence.test.ts
View File

@@ -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");

24
packages/coding-agent/test/suite/agent-session-compaction.test.ts
View File

@@ -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<void>;
_runAutoCompaction: (reason: "overflow" | "threshold", willRetry: boolean) => Promise<void>;
_checkCompaction: (assistantMessage: AssistantMessage, skipAbortedCheck?: boolean) => Promise<boolean>;
_runAutoCompaction: (reason: "overflow" | "threshold", willRetry: boolean) => Promise<boolean>;
};
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() }),

95
packages/coding-agent/test/suite/regressions/1717-2113-agent-session-event-settlement.test.ts Normal file
View File

@@ -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"]]);
});
});