From 8a0529ed9d63f929df0a3b5958ead779ca8c0f50 Mon Sep 17 00:00:00 2001 From: Petr Baudis Date: Fri, 20 Mar 2026 00:33:17 +0100 Subject: [PATCH] fix(coding-agent): resolve waitForRetry() race when auto-retry produces tool calls (#2440) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit When auto-retry fires after a retryable error (e.g. overloaded_error) and the retry response includes tool_use, session.prompt() returned prematurely because _resolveRetry() was called on the first successful message_end — while the agent loop was still executing tools via the fire-and-forget agent.continue(). This caused callers to observe isStreaming=true after prompt() returned, and follow-up session.prompt() calls threw 'Agent is already processing'. The tool execution results were silently lost. Fix: move _resolveRetry() from the message_end handler to the agent_end handler. The _retryAttempt counter reset stays on message_end (preventing accumulation across LLM calls within a turn), but the promise that unblocks waitForRetry() now only resolves when the full agent loop completes. --- .../coding-agent/src/core/agent-session.ts | 2 +- .../test/agent-session-retry.test.ts | 91 ++++++++++++++++++- 2 files changed, 91 insertions(+), 2 deletions(-) diff --git a/packages/coding-agent/src/core/agent-session.ts b/packages/coding-agent/src/core/agent-session.ts index ddca14c9..4f44c22a 100644 --- a/packages/coding-agent/src/core/agent-session.ts +++ b/packages/coding-agent/src/core/agent-session.ts @@ -490,7 +490,6 @@ export class AgentSession { attempt: this._retryAttempt, }); this._retryAttempt = 0; - this._resolveRetry(); } } } @@ -506,6 +505,7 @@ export class AgentSession { if (didRetry) return; // Retry was initiated, don't proceed to compaction } + this._resolveRetry(); await this._checkCompaction(msg); } } diff --git a/packages/coding-agent/test/agent-session-retry.test.ts b/packages/coding-agent/test/agent-session-retry.test.ts index 962e5005..1697f2ae 100644 --- a/packages/coding-agent/test/agent-session-retry.test.ts +++ b/packages/coding-agent/test/agent-session-retry.test.ts @@ -1,8 +1,9 @@ import { existsSync, mkdirSync, rmSync } from "node:fs"; import { tmpdir } from "node:os"; import { join } from "node:path"; -import { Agent, type AgentEvent } from "@mariozechner/pi-agent-core"; +import { Agent, type AgentEvent, type AgentTool } from "@mariozechner/pi-agent-core"; import { type AssistantMessage, type AssistantMessageEvent, EventStream, getModel } from "@mariozechner/pi-ai"; +import { Type } from "@sinclair/typebox"; import { afterEach, beforeEach, describe, expect, it } from "vitest"; import { AgentSession } from "../src/core/agent-session.js"; import { AuthStorage } from "../src/core/auth-storage.js"; @@ -226,4 +227,92 @@ describe("AgentSession retry", () => { expect(callCount).toBe(2); expect(events).toEqual(["start:1", "end:success=true"]); }); + + it("prompt waits for full agent loop when retry produces tool calls", async () => { + // Regression: when auto-retry fires and the retry response includes tool_use, + // session.prompt() must wait for the entire tool loop to finish before returning. + // Previously, _resolveRetry() on the first successful message_end would unblock + // waitForRetry() while the agent was still executing tools. + let callCount = 0; + const toolExecuted = { value: false }; + + const echoTool: AgentTool = { + name: "echo", + label: "Echo", + description: "Echo text back", + parameters: Type.Object({ text: Type.String() }), + execute: async () => { + toolExecuted.value = true; + return { content: [{ type: "text", text: "echoed" }], details: undefined }; + }, + }; + + const model = getModel("anthropic", "claude-sonnet-4-5")!; + const agent = new Agent({ + getApiKey: () => "test-key", + initialState: { model, systemPrompt: "Test", tools: [] }, + streamFn: () => { + callCount++; + const stream = new MockAssistantStream(); + queueMicrotask(() => { + if (callCount === 1) { + // First call: overloaded error + const msg = createAssistantMessage("", { + stopReason: "error", + errorMessage: "overloaded_error", + }); + stream.push({ type: "start", partial: msg }); + stream.push({ type: "error", reason: "error", error: msg }); + } else if (callCount === 2) { + // Second call (retry): text + tool_use + const msg: AssistantMessage = { + ...createAssistantMessage("Looking that up now."), + stopReason: "toolUse", + content: [ + { type: "text", text: "Looking that up now." }, + { type: "toolCall", id: "call_1", name: "echo", arguments: { text: "hello" } }, + ], + }; + stream.push({ type: "start", partial: msg }); + stream.push({ type: "done", reason: "toolUse", message: msg }); + } else { + // Third call (after tool result): final response + const msg = createAssistantMessage("Final answer."); + stream.push({ type: "start", partial: msg }); + stream.push({ type: "done", reason: "stop", message: msg }); + } + }); + return stream; + }, + }); + + const sessionManager = SessionManager.inMemory(); + const settingsManager = SettingsManager.create(tempDir, tempDir); + const authStorage = AuthStorage.create(join(tempDir, "auth.json")); + const modelRegistry = new ModelRegistry(authStorage, tempDir); + authStorage.setRuntimeApiKey("anthropic", "test-key"); + settingsManager.applyOverrides({ retry: { enabled: true, maxRetries: 3, baseDelayMs: 1 } }); + + session = new AgentSession({ + agent, + sessionManager, + settingsManager, + cwd: tempDir, + modelRegistry, + resourceLoader: createTestResourceLoader(), + baseToolsOverride: { echo: echoTool }, + }); + + await session.prompt("Test"); + + // All three LLM calls must have completed + expect(callCount).toBe(3); + // Tool must have been executed + expect(toolExecuted.value).toBe(true); + // Agent must not be streaming after prompt returns + expect(session.isStreaming).toBe(false); + // A follow-up prompt must work (no "Agent is already processing" error) + await session.prompt("Follow-up"); + expect(callCount).toBe(4); + }); });