gsd-pi 2.76.0 → 2.77.0-dev.2daa994b6

package/packages/pi-agent-core/src/agent-loop.test.ts

@@ -1,51 +1,81 @@
 // agent-loop tests
 // Covers: pauseTurn handling (#2869), schema overload retry cap (#2783)
 
-import { describe, it, mock } from "node:test";
+import { describe, it } from "node:test";
 import assert from "node:assert/strict";
-import { readFileSync } from "node:fs";
-import { join, dirname } from "node:path";
-import { fileURLToPath } from "node:url";
 import { Type } from "@sinclair/typebox";
 import { agentLoop, MAX_CONSECUTIVE_VALIDATION_FAILURES } from "./agent-loop.js";
 import type { AgentContext, AgentLoopConfig, AgentTool, AgentEvent, AgentMessage } from "./types.js";
 import { AssistantMessageEventStream, EventStream } from "@gsd/pi-ai";
 import type { AssistantMessage, AssistantMessageEvent, Model } from "@gsd/pi-ai";
 
-const __dirname = dirname(fileURLToPath(import.meta.url));
-
 describe("agent-loop — pauseTurn handling (#2869)", () => {
-	it("sets hasMoreToolCalls when stopReason is pauseTurn", () => {
-		const source = readFileSync(join(__dirname, "agent-loop.ts"), "utf-8");
-
-		// The agent loop must treat pauseTurn as a reason to continue the inner
-		// loop, just like toolUse. This prevents incomplete server_tool_use blocks
-		// from being saved to history, which would cause a 400 on the next request.
-		assert.match(
-			source,
-			/pauseTurn/,
-			"agent-loop.ts must handle the pauseTurn stop reason",
-		);
+	it("continues to a second assistant turn when stopReason is pauseTurn", async () => {
+		const pauseTurn = makeAssistantMessage({
+			content: [{ type: "text", text: "Still working..." }],
+			stopReason: "pauseTurn",
+		});
+		const finalStop = makeAssistantMessage({
+			content: [{ type: "text", text: "Done." }],
+			stopReason: "stop",
+		});
+
+		let streamCallCount = 0;
+		const seenLastRoles: string[] = [];
+		const streamFn = (_model: Model<any>, llmContext: any): AssistantMessageEventStream => {
+			streamCallCount++;
+			seenLastRoles.push(llmContext.messages.at(-1)?.role ?? "none");
+			const message = streamCallCount === 1 ? pauseTurn : finalStop;
+			const stream = new AssistantMessageEventStream();
+			queueMicrotask(() => {
+				stream.push({ type: "start", partial: message });
+				stream.push({ type: "done", message });
+				stream.end(message);
+			});
+			return stream;
+		};
 
-		// Verify it sets hasMoreToolCalls = true for pauseTurn
-		assert.match(
-			source,
-			/stopReason\s*===?\s*["']pauseTurn["']/,
-			'agent-loop.ts must check for stopReason === "pauseTurn"',
+		const context: AgentContext = {
+			systemPrompt: "You are a test agent.",
+			messages: [{ role: "user", content: [{ type: "text", text: "Keep going" }], timestamp: Date.now() }],
+			tools: [],
+		};
+
+		const config: AgentLoopConfig = {
+			model: TEST_MODEL,
+			convertToLlm: (msgs) => msgs.filter((m): m is any => m.role !== "custom"),
+			toolExecution: "sequential",
+		};
+
+		const stream = agentLoop(
+			[{ role: "user", content: [{ type: "text", text: "Keep going" }], timestamp: Date.now() }],
+			context,
+			config,
+			undefined,
+			streamFn as any,
 		);
-	});
 
-	it("pauseTurn is in the StopReason union type", () => {
-		// Read the pi-ai types to ensure pauseTurn is a valid StopReason
-		const typesPath = join(__dirname, "..", "..", "pi-ai", "src", "types.ts");
-		const typesSource = readFileSync(typesPath, "utf-8");
-		assert.match(
-			typesSource,
-			/["']pauseTurn["']/,
-			'StopReason type must include "pauseTurn"',
+		const events = await collectEvents(stream);
+		const assistantMessages = events.filter(
+			(event): event is Extract<AgentEvent, { type: "message_end" }> =>
+				event.type === "message_end" && event.message.role === "assistant",
 		);
+		const turnStarts = events.filter((event) => event.type === "turn_start");
+
+		assert.equal(streamCallCount, 2, "pauseTurn must cause a second provider call");
+		assert.equal(assistantMessages.length, 2, "expected two assistant turns");
+		assert.equal(assistantMessages[0].message.stopReason, "pauseTurn");
+		assert.equal(assistantMessages[1].message.stopReason, "stop");
+		assert.equal(turnStarts.length, 2, "expected a second turn_start after pauseTurn");
+		assert.deepEqual(seenLastRoles, ["user", "assistant"], "second turn must continue from the paused assistant state");
 	});
 
+	// The behavioural test above ("continues to a second assistant turn …")
+	// already exercises `stopReason: "pauseTurn"` at the type level (TS won't
+	// compile without the union member) and at runtime (the stream emits it
+	// and the loop acts on it). A separate source-text grep added nothing.
+	// Deleted per #4797.
+
 	it("uses provider-supplied external tool results instead of the placeholder", async () => {
 		const externalMessage = makeAssistantMessage({
 			content: [
@@ -98,6 +128,171 @@ describe("agent-loop — pauseTurn handling (#2869)", () => {
 		assert.deepEqual(toolEnd.result.details, { source: "claude-code" });
 		assert.equal(toolEnd.isError, false);
 	});
+
+	it("injects queued steering messages before the next assistant turn and skips remaining tools", async () => {
+		const tool = makeToolWithSchema();
+		const steeringMessage: AgentMessage = {
+			role: "user",
+			content: [{ type: "text", text: "Stop after the first tool." }],
+			timestamp: Date.now(),
+		};
+		const toolTurn = makeAssistantMessage({
+			content: [
+				{
+					type: "toolCall",
+					id: "write-1",
+					name: "write_file",
+					arguments: { path: "/tmp/one", content: "first" },
+				},
+				{
+					type: "toolCall",
+					id: "write-2",
+					name: "write_file",
+					arguments: { path: "/tmp/two", content: "second" },
+				},
+			],
+			stopReason: "toolUse",
+		});
+		const finalStop = makeAssistantMessage({
+			content: [{ type: "text", text: "Stopped after the first tool." }],
+			stopReason: "stop",
+		});
+
+		let steeringPollCount = 0;
+		const seenLastMessages: unknown[] = [];
+		const streamFn = (_model: Model<any>, llmContext: any): AssistantMessageEventStream => {
+			seenLastMessages.push(llmContext.messages.at(-1));
+			const message = seenLastMessages.length === 1 ? toolTurn : finalStop;
+			const stream = new AssistantMessageEventStream();
+			queueMicrotask(() => {
+				stream.push({ type: "start", partial: message });
+				stream.push({ type: "done", message });
+				stream.end(message);
+			});
+			return stream;
+		};
+
+		const context: AgentContext = {
+			systemPrompt: "You are a test agent.",
+			messages: [{ role: "user", content: [{ type: "text", text: "Write both files" }], timestamp: Date.now() }],
+			tools: [tool],
+		};
+
+		const config: AgentLoopConfig = {
+			model: TEST_MODEL,
+			convertToLlm: (msgs) => msgs.filter((m): m is any => m.role !== "custom"),
+			toolExecution: "sequential",
+			getSteeringMessages: async () => {
+				steeringPollCount++;
+				return steeringPollCount === 2 ? [steeringMessage] : [];
+			},
+		};
+
+		const stream = agentLoop(
+			[{ role: "user", content: [{ type: "text", text: "Write both files" }], timestamp: Date.now() }],
+			context,
+			config,
+			undefined,
+			streamFn as any,
+		);
+
+		const events = await collectEvents(stream);
+		const toolEnds = events.filter(
+			(event): event is Extract<AgentEvent, { type: "tool_execution_end" }> => event.type === "tool_execution_end",
+		);
+		const steeringEnd = events.find(
+			(event): event is Extract<AgentEvent, { type: "message_end" }> =>
+				event.type === "message_end" &&
+				event.message.role === "user" &&
+				(event.message.content[0] as any)?.text === "Stop after the first tool.",
+		);
+
+		assert.equal(toolEnds.length, 2, "expected one executed tool and one skipped tool");
+		assert.ok(
+			toolEnds.some((event) => JSON.stringify(event.result.content) === JSON.stringify([{ type: "text", text: "done" }])),
+			"expected one completed tool result",
+		);
+		assert.ok(
+			toolEnds.some(
+				(event) =>
+					JSON.stringify(event.result.content) ===
+					JSON.stringify([{ type: "text", text: "Skipped due to queued user message." }]),
+			),
+			"expected one skipped tool result after steering interruption",
+		);
+		assert.ok(steeringEnd, "queued steering message should be emitted before the next assistant turn");
+		assert.equal((seenLastMessages[1] as any)?.content?.[0]?.text, "Stop after the first tool.");
+	});
+
+	it("restarts the outer loop when follow-up messages arrive after a stop", async () => {
+		const initialStop = makeAssistantMessage({
+			content: [{ type: "text", text: "First answer." }],
+			stopReason: "stop",
+		});
+		const followUpStop = makeAssistantMessage({
+			content: [{ type: "text", text: "Follow-up answer." }],
+			stopReason: "stop",
+		});
+		const followUpMessage: AgentMessage = {
+			role: "user",
+			content: [{ type: "text", text: "One more thing." }],
+			timestamp: Date.now(),
+		};
+
+		let followUpPollCount = 0;
+		const seenLastMessages: unknown[] = [];
+		const streamFn = (_model: Model<any>, llmContext: any): AssistantMessageEventStream => {
+			seenLastMessages.push(llmContext.messages.at(-1));
+			const message = seenLastMessages.length === 1 ? initialStop : followUpStop;
+			const stream = new AssistantMessageEventStream();
+			queueMicrotask(() => {
+				stream.push({ type: "start", partial: message });
+				stream.push({ type: "done", message });
+				stream.end(message);
+			});
+			return stream;
+		};
+
+		const context: AgentContext = {
+			systemPrompt: "You are a test agent.",
+			messages: [{ role: "user", content: [{ type: "text", text: "Initial prompt" }], timestamp: Date.now() }],
+			tools: [],
+		};
+
+		const config: AgentLoopConfig = {
+			model: TEST_MODEL,
+			convertToLlm: (msgs) => msgs.filter((m): m is any => m.role !== "custom"),
+			toolExecution: "sequential",
+			getFollowUpMessages: async () => {
+				followUpPollCount++;
+				return followUpPollCount === 1 ? [followUpMessage] : [];
+			},
+		};
+
+		const stream = agentLoop(
+			[{ role: "user", content: [{ type: "text", text: "Initial prompt" }], timestamp: Date.now() }],
+			context,
+			config,
+			undefined,
+			streamFn as any,
+		);
+
+		const events = await collectEvents(stream);
+		const assistantEnds = events.filter(
+			(event): event is Extract<AgentEvent, { type: "message_end" }> =>
+				event.type === "message_end" && event.message.role === "assistant",
+		);
+		const followUpEnd = events.find(
+			(event): event is Extract<AgentEvent, { type: "message_end" }> =>
+				event.type === "message_end" &&
+				event.message.role === "user" &&
+				(event.message.content[0] as any)?.text === "One more thing.",
+		);
+
+		assert.equal(assistantEnds.length, 2, "expected the agent loop to restart for the follow-up turn");
+		assert.ok(followUpEnd, "follow-up message should be emitted before the restarted assistant turn");
+		assert.equal((seenLastMessages[1] as any)?.content?.[0]?.text, "One more thing.");
+	});
 });
 
 /**

package/packages/pi-agent-core/src/agent-loop.ts

@@ -222,6 +222,20 @@ async function runLoop(
 				// backoff, network unavailable). Convert to a graceful error message so the
 				// agent loop can end cleanly instead of crashing with an unhandled rejection.
 				const errorText = error instanceof Error ? error.message : String(error);
+				if (config.onStreamError) {
+					try {
+						await config.onStreamError(
+							{
+								error: error instanceof Error ? error : new Error(errorText),
+								partialText: "",
+								willRetry: false,
+							},
+							signal,
+						);
+					} catch {
+						// Hook failures must not crash the loop.
+					}
+				}
 				message = {
 					role: "assistant",
 					content: [],

package/packages/pi-agent-core/src/agent.test.ts

@@ -1,131 +1,125 @@
 // Agent activeInferenceModel regression tests
-// Verifies that activeInferenceModel is set/cleared correctly in _runLoop,
-// and that the footer reads activeInferenceModel instead of state.model.
+// Verifies that activeInferenceModel is set before streaming begins and
+// cleared after streaming completes — observed via the streamFn seam and
+// post-condition, not the source text.
 // Regression test for https://github.com/gsd-build/gsd-2/issues/1844 Bug 2
 
 import { describe, it } from "node:test";
 import assert from "node:assert/strict";
-import { readFileSync } from "node:fs";
-import { join, dirname } from "node:path";
-import { fileURLToPath } from "node:url";
 import { Agent } from "./agent.ts";
 import { getModel, type AssistantMessageEventStream } from "@gsd/pi-ai";
 
-const __dirname = dirname(fileURLToPath(import.meta.url));
+function makeDoneStream(modelId: string): AssistantMessageEventStream {
+	const usage = {
+		input: 0,
+		output: 0,
+		cacheRead: 0,
+		cacheWrite: 0,
+		totalTokens: 0,
+		cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 },
+	};
+	const message = {
+		role: "assistant" as const,
+		content: [{ type: "text" as const, text: "ok" }],
+		api: "anthropic-messages" as const,
+		provider: "anthropic",
+		model: modelId,
+		usage,
+		stopReason: "stop" as const,
+		timestamp: Date.now(),
+	};
+	return {
+		async *[Symbol.asyncIterator]() {
+			yield { type: "start", partial: message };
+			yield { type: "done", message };
+		},
+		result: async () => message,
+		[Symbol.asyncDispose]: async () => {},
+	} as AssistantMessageEventStream;
+}
 
 describe("Agent — activeInferenceModel (#1844 Bug 2)", () => {
-	it("activeInferenceModel is declared in AgentState interface", () => {
-		const typesSource = readFileSync(join(__dirname, "types.ts"), "utf-8");
-		assert.match(typesSource, /activeInferenceModel\??:\s*Model/,
-			"AgentState must declare activeInferenceModel field");
-	});
+	it("_runLoop sets activeInferenceModel = model mid-stream and clears it when finished", async () => {
+		const model = getModel("anthropic", "claude-3-5-sonnet-20241022");
+		let midStreamModel: unknown = "<not-captured>";
 
-	it("_runLoop sets activeInferenceModel before streaming and clears in finally", () => {
-		const agentSource = readFileSync(join(__dirname, "agent.ts"), "utf-8");
+		const agent = new Agent({
+			initialState: { model, systemPrompt: "test", tools: [] },
+			streamFn: (streamModel): AssistantMessageEventStream => {
+				// streamFn is invoked AFTER `activeInferenceModel = model` and
+				// BEFORE the finally block that clears it. Capture state here.
+				midStreamModel = agent.state.activeInferenceModel;
+				return makeDoneStream(streamModel.id);
+			},
+		});
 
-		// Must set activeInferenceModel = model before streaming starts
-		const setLine = agentSource.indexOf("this._state.activeInferenceModel = model");
-		assert.ok(setLine > -1, "agent.ts must set activeInferenceModel = model in _runLoop");
+		// Baseline: undefined before any inference.
+		assert.equal(
+			agent.state.activeInferenceModel,
+			undefined,
+			"activeInferenceModel must be undefined before prompt()",
+		);
 
-		// Must clear activeInferenceModel = undefined after streaming completes
-		const clearLine = agentSource.indexOf("this._state.activeInferenceModel = undefined");
-		assert.ok(clearLine > -1, "agent.ts must clear activeInferenceModel in finally block");
+		await agent.prompt("hello");
 
-		// The set must come before the clear
-		assert.ok(setLine < clearLine, "activeInferenceModel must be set before cleared");
-	});
+		// Mid-stream: set to the inference model.
+		assert.equal(
+			(midStreamModel as { id?: string } | undefined)?.id,
+			model.id,
+			"activeInferenceModel must equal the inference model while streaming",
+		);
 
-	it("footer displays activeInferenceModel instead of state.model", () => {
-		const footerPath = join(__dirname, "..", "..", "pi-coding-agent", "src",
-			"modes", "interactive", "components", "footer.ts");
-		const footerSource = readFileSync(footerPath, "utf-8");
-		assert.match(footerSource, /activeInferenceModel/,
-			"footer.ts must reference activeInferenceModel for display");
+		// Post-stream: cleared back to undefined (finally block).
+		assert.equal(
+			agent.state.activeInferenceModel,
+			undefined,
+			"activeInferenceModel must be undefined after prompt() resolves",
+		);
 	});
 
-	it("activeInferenceModel is set before AbortController creation", () => {
-		const agentSource = readFileSync(join(__dirname, "agent.ts"), "utf-8");
+	it("activeInferenceModel is also cleared when the stream throws", async () => {
+		const model = getModel("anthropic", "claude-3-5-sonnet-20241022");
+		let midStreamModel: unknown = "<not-captured>";
+
+		const agent = new Agent({
+			initialState: { model, systemPrompt: "test", tools: [] },
+			streamFn: (): AssistantMessageEventStream => {
+				midStreamModel = agent.state.activeInferenceModel;
+				return {
+					async *[Symbol.asyncIterator]() {
+						throw new Error("boom");
+					},
+					result: async () => {
+						throw new Error("boom");
+					},
+					[Symbol.asyncDispose]: async () => {},
+				} as AssistantMessageEventStream;
+			},
+		});
+
+		await agent.prompt("hello");
 
-		const setLine = agentSource.indexOf("this._state.activeInferenceModel = model");
-		const abortLine = agentSource.indexOf("this.abortController = new AbortController");
-		assert.ok(setLine > -1 && abortLine > -1);
-		assert.ok(setLine < abortLine,
-			"activeInferenceModel must be set before streaming infrastructure is created");
+		assert.equal(
+			(midStreamModel as { id?: string } | undefined)?.id,
+			model.id,
+			"activeInferenceModel must be set even when the stream later throws",
+		);
+		assert.equal(
+			agent.state.activeInferenceModel,
+			undefined,
+			"activeInferenceModel must be cleared in finally even after stream errors",
+		);
 	});
 
 	it("getProviderOptions are forwarded into the provider stream call", async () => {
 		let capturedOptions: Record<string, unknown> | undefined;
+		const model = getModel("anthropic", "claude-3-5-sonnet-20241022");
 		const agent = new Agent({
-			initialState: {
-				model: getModel("anthropic", "claude-3-5-sonnet-20241022"),
-				systemPrompt: "test",
-				tools: [],
-			},
+			initialState: { model, systemPrompt: "test", tools: [] },
 			getProviderOptions: async () => ({ customRuntimeOption: "present" }),
 			streamFn: (_model, _context, options): AssistantMessageEventStream => {
 				capturedOptions = options as Record<string, unknown> | undefined;
-				return {
-					async *[Symbol.asyncIterator]() {
-						yield {
-							type: "start",
-							partial: {
-								role: "assistant",
-								content: [],
-								api: "anthropic-messages",
-								provider: "anthropic",
-								model: "claude-3-5-sonnet-20241022",
-								usage: {
-									input: 0,
-									output: 0,
-									cacheRead: 0,
-									cacheWrite: 0,
-									totalTokens: 0,
-									cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 },
-								},
-								stopReason: "stop",
-								timestamp: Date.now(),
-							},
-						};
-						yield {
-							type: "done",
-							message: {
-								role: "assistant",
-								content: [{ type: "text", text: "ok" }],
-								api: "anthropic-messages",
-								provider: "anthropic",
-								model: "claude-3-5-sonnet-20241022",
-								usage: {
-									input: 0,
-									output: 0,
-									cacheRead: 0,
-									cacheWrite: 0,
-									totalTokens: 0,
-									cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 },
-								},
-								stopReason: "stop",
-								timestamp: Date.now(),
-							},
-						};
-					},
-					result: async () => ({
-						role: "assistant",
-						content: [{ type: "text", text: "ok" }],
-						api: "anthropic-messages",
-						provider: "anthropic",
-						model: "claude-3-5-sonnet-20241022",
-						usage: {
-							input: 0,
-							output: 0,
-							cacheRead: 0,
-							cacheWrite: 0,
-							totalTokens: 0,
-							cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 },
-						},
-						stopReason: "stop",
-						timestamp: Date.now(),
-					}),
-					[Symbol.asyncDispose]: async () => {},
-				} as AssistantMessageEventStream;
+				return makeDoneStream(model.id);
 			},
 		});
 

package/packages/pi-agent-core/src/types.ts

@@ -203,6 +203,40 @@ export interface AgentLoopConfig extends SimpleStreamOptions {
 	 * but skips tool.execute() and does not add tool results to context.
 	 */
 	externalToolExecution?: boolean;
+
+	/**
+	 * Called when the agent is about to retry after a recoverable provider error.
+	 *
+	 * Fires after each transient failure (rate limit, connection reset, etc.)
+	 * before the retry backoff sleep. Useful for surfacing retry telemetry to
+	 * the host app without instrumenting the provider layer directly.
+	 */
+	onRetry?: (context: RetryContext, signal?: AbortSignal) => Promise<void> | void;
+
+	/**
+	 * Called when the LLM stream fails mid-response.
+	 *
+	 * Fires once per failed stream — `onRetry` may still fire separately if the
+	 * loop elects to retry. Use this for logging or flushing partial state.
+	 */
+	onStreamError?: (context: StreamErrorContext, signal?: AbortSignal) => Promise<void> | void;
+}
+
+/** Context passed to `onRetry`. */
+export interface RetryContext {
+	attempt: number;
+	maxRetries: number;
+	delayMs: number;
+	error: Error;
+}
+
+/** Context passed to `onStreamError`. */
+export interface StreamErrorContext {
+	error: Error;
+	/** The text already emitted by the stream before it failed (may be empty). */
+	partialText: string;
+	/** Whether the loop will attempt to retry after this failure. */
+	willRetry: boolean;
 }
 
 /**