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

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 */