fix(coding-agent): resolve waitForRetry() race when auto-retry produces tool calls (#2440)

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.

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);
 		}
 	}

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);
+	});
 });