update

pi-mono-reference/packages/agent/test/agent.test.ts

@@ -0,0 +1,463 @@
+import { type AssistantMessage, type AssistantMessageEvent, EventStream, getModel } from "@mariozechner/pi-ai";
+import { describe, expect, it } from "vitest";
+import { Agent } from "../src/index.js";
+
+// Mock stream that mimics AssistantMessageEventStream
+class MockAssistantStream extends EventStream<AssistantMessageEvent, AssistantMessage> {
+	constructor() {
+		super(
+			(event) => event.type === "done" || event.type === "error",
+			(event) => {
+				if (event.type === "done") return event.message;
+				if (event.type === "error") return event.error;
+				throw new Error("Unexpected event type");
+			},
+		);
+	}
+}
+
+function createAssistantMessage(text: string): AssistantMessage {
+	return {
+		role: "assistant",
+		content: [{ type: "text", text }],
+		api: "openai-responses",
+		provider: "openai",
+		model: "mock",
+		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(),
+	};
+}
+
+function createDeferred(): {
+	promise: Promise<void>;
+	resolve: () => void;
+} {
+	let resolve = () => {};
+	const promise = new Promise<void>((resolvePromise) => {
+		resolve = resolvePromise;
+	});
+	return { promise, resolve };
+}
+
+describe("Agent", () => {
+	it("should create an agent instance with default state", () => {
+		const agent = new Agent();
+
+		expect(agent.state).toBeDefined();
+		expect(agent.state.systemPrompt).toBe("");
+		expect(agent.state.model).toBeDefined();
+		expect(agent.state.thinkingLevel).toBe("off");
+		expect(agent.state.tools).toEqual([]);
+		expect(agent.state.messages).toEqual([]);
+		expect(agent.state.isStreaming).toBe(false);
+		expect(agent.state.streamingMessage).toBe(undefined);
+		expect(agent.state.pendingToolCalls).toEqual(new Set());
+		expect(agent.state.errorMessage).toBeUndefined();
+	});
+
+	it("should create an agent instance with custom initial state", () => {
+		const customModel = getModel("openai", "gpt-4o-mini");
+		const agent = new Agent({
+			initialState: {
+				systemPrompt: "You are a helpful assistant.",
+				model: customModel,
+				thinkingLevel: "low",
+			},
+		});
+
+		expect(agent.state.systemPrompt).toBe("You are a helpful assistant.");
+		expect(agent.state.model).toBe(customModel);
+		expect(agent.state.thinkingLevel).toBe("low");
+	});
+
+	it("should subscribe to events", () => {
+		const agent = new Agent();
+
+		let eventCount = 0;
+		const unsubscribe = agent.subscribe((_event) => {
+			eventCount++;
+		});
+
+		// No initial event on subscribe
+		expect(eventCount).toBe(0);
+
+		// State mutators don't emit events
+		agent.state.systemPrompt = "Test prompt";
+		expect(eventCount).toBe(0);
+		expect(agent.state.systemPrompt).toBe("Test prompt");
+
+		// Unsubscribe should work
+		unsubscribe();
+		agent.state.systemPrompt = "Another prompt";
+		expect(eventCount).toBe(0); // Should not increase
+	});
+
+	it("should await async subscribers before prompt resolves", async () => {
+		const barrier = createDeferred();
+		const agent = new Agent({
+			streamFn: () => {
+				const stream = new MockAssistantStream();
+				queueMicrotask(() => {
+					stream.push({ type: "done", reason: "stop", message: createAssistantMessage("ok") });
+				});
+				return stream;
+			},
+		});
+
+		let listenerFinished = false;
+		agent.subscribe(async (event) => {
+			if (event.type === "agent_end") {
+				await barrier.promise;
+				listenerFinished = true;
+			}
+		});
+
+		let promptResolved = false;
+		const promptPromise = agent.prompt("hello").then(() => {
+			promptResolved = true;
+		});
+
+		await new Promise((resolve) => setTimeout(resolve, 10));
+		expect(promptResolved).toBe(false);
+		expect(listenerFinished).toBe(false);
+		expect(agent.state.isStreaming).toBe(true);
+
+		barrier.resolve();
+		await promptPromise;
+
+		expect(listenerFinished).toBe(true);
+		expect(promptResolved).toBe(true);
+		expect(agent.state.isStreaming).toBe(false);
+	});
+
+	it("waitForIdle should wait for async subscribers", async () => {
+		const barrier = createDeferred();
+		const agent = new Agent({
+			streamFn: () => {
+				const stream = new MockAssistantStream();
+				queueMicrotask(() => {
+					stream.push({ type: "done", reason: "stop", message: createAssistantMessage("ok") });
+				});
+				return stream;
+			},
+		});
+
+		agent.subscribe(async (event) => {
+			if (event.type === "message_end" && event.message.role === "assistant") {
+				await barrier.promise;
+			}
+		});
+
+		const promptPromise = agent.prompt("hello");
+		let idleResolved = false;
+		const idlePromise = agent.waitForIdle().then(() => {
+			idleResolved = true;
+		});
+
+		await new Promise((resolve) => setTimeout(resolve, 10));
+		expect(idleResolved).toBe(false);
+		expect(agent.state.isStreaming).toBe(true);
+
+		barrier.resolve();
+		await Promise.all([promptPromise, idlePromise]);
+
+		expect(idleResolved).toBe(true);
+		expect(agent.state.isStreaming).toBe(false);
+	});
+
+	it("should pass the active abort signal to subscribers", async () => {
+		let receivedSignal: AbortSignal | undefined;
+		const agent = new Agent({
+			streamFn: (_model, _context, options) => {
+				const stream = new MockAssistantStream();
+				queueMicrotask(() => {
+					stream.push({ type: "start", partial: createAssistantMessage("") });
+					const checkAbort = () => {
+						if (options?.signal?.aborted) {
+							stream.push({ type: "error", reason: "aborted", error: createAssistantMessage("Aborted") });
+						} else {
+							setTimeout(checkAbort, 5);
+						}
+					};
+					checkAbort();
+				});
+				return stream;
+			},
+		});
+
+		agent.subscribe((event, signal) => {
+			if (event.type === "agent_start") {
+				receivedSignal = signal;
+			}
+		});
+
+		const promptPromise = agent.prompt("hello");
+		await new Promise((resolve) => setTimeout(resolve, 10));
+
+		expect(receivedSignal).toBeDefined();
+		expect(receivedSignal?.aborted).toBe(false);
+
+		agent.abort();
+		await promptPromise;
+
+		expect(receivedSignal?.aborted).toBe(true);
+	});
+
+	it("should update state with mutators", () => {
+		const agent = new Agent();
+
+		// Test setSystemPrompt
+		agent.state.systemPrompt = "Custom prompt";
+		expect(agent.state.systemPrompt).toBe("Custom prompt");
+
+		// Test setModel
+		const newModel = getModel("google", "gemini-2.5-flash");
+		agent.state.model = newModel;
+		expect(agent.state.model).toBe(newModel);
+
+		// Test setThinkingLevel
+		agent.state.thinkingLevel = "high";
+		expect(agent.state.thinkingLevel).toBe("high");
+
+		// Test setTools
+		const tools = [{ name: "test", description: "test tool" } as any];
+		agent.state.tools = tools;
+		expect(agent.state.tools).toEqual(tools);
+		expect(agent.state.tools).not.toBe(tools); // Should be a copy
+
+		// Test replaceMessages
+		const messages = [{ role: "user" as const, content: "Hello", timestamp: Date.now() }];
+		agent.state.messages = messages;
+		expect(agent.state.messages).toEqual(messages);
+		expect(agent.state.messages).not.toBe(messages); // Should be a copy
+
+		// Test appendMessage
+		const newMessage = { role: "assistant" as const, content: [{ type: "text" as const, text: "Hi" }] };
+		agent.state.messages.push(newMessage as any);
+		expect(agent.state.messages).toHaveLength(2);
+		expect(agent.state.messages[1]).toBe(newMessage);
+
+		// Test clearMessages
+		agent.state.messages = [];
+		expect(agent.state.messages).toEqual([]);
+	});
+
+	it("should support steering message queue", async () => {
+		const agent = new Agent();
+
+		const message = { role: "user" as const, content: "Steering message", timestamp: Date.now() };
+		agent.steer(message);
+
+		// The message is queued but not yet in state.messages
+		expect(agent.state.messages).not.toContainEqual(message);
+	});
+
+	it("should support follow-up message queue", async () => {
+		const agent = new Agent();
+
+		const message = { role: "user" as const, content: "Follow-up message", timestamp: Date.now() };
+		agent.followUp(message);
+
+		// The message is queued but not yet in state.messages
+		expect(agent.state.messages).not.toContainEqual(message);
+	});
+
+	it("should handle abort controller", () => {
+		const agent = new Agent();
+
+		// Should not throw even if nothing is running
+		expect(() => agent.abort()).not.toThrow();
+	});
+
+	it("should throw when prompt() called while streaming", async () => {
+		let abortSignal: AbortSignal | undefined;
+		const agent = new Agent({
+			// Use a stream function that responds to abort
+			streamFn: (_model, _context, options) => {
+				abortSignal = options?.signal;
+				const stream = new MockAssistantStream();
+				queueMicrotask(() => {
+					stream.push({ type: "start", partial: createAssistantMessage("") });
+					// Check abort signal periodically
+					const checkAbort = () => {
+						if (abortSignal?.aborted) {
+							stream.push({ type: "error", reason: "aborted", error: createAssistantMessage("Aborted") });
+						} else {
+							setTimeout(checkAbort, 5);
+						}
+					};
+					checkAbort();
+				});
+				return stream;
+			},
+		});
+
+		// Start first prompt (don't await, it will block until abort)
+		const firstPrompt = agent.prompt("First message");
+
+		// Wait a tick for isStreaming to be set
+		await new Promise((resolve) => setTimeout(resolve, 10));
+		expect(agent.state.isStreaming).toBe(true);
+
+		// Second prompt should reject
+		await expect(agent.prompt("Second message")).rejects.toThrow(
+			"Agent is already processing a prompt. Use steer() or followUp() to queue messages, or wait for completion.",
+		);
+
+		// Cleanup - abort to stop the stream
+		agent.abort();
+		await firstPrompt.catch(() => {}); // Ignore abort error
+	});
+
+	it("should throw when continue() called while streaming", async () => {
+		let abortSignal: AbortSignal | undefined;
+		const agent = new Agent({
+			streamFn: (_model, _context, options) => {
+				abortSignal = options?.signal;
+				const stream = new MockAssistantStream();
+				queueMicrotask(() => {
+					stream.push({ type: "start", partial: createAssistantMessage("") });
+					const checkAbort = () => {
+						if (abortSignal?.aborted) {
+							stream.push({ type: "error", reason: "aborted", error: createAssistantMessage("Aborted") });
+						} else {
+							setTimeout(checkAbort, 5);
+						}
+					};
+					checkAbort();
+				});
+				return stream;
+			},
+		});
+
+		// Start first prompt
+		const firstPrompt = agent.prompt("First message");
+		await new Promise((resolve) => setTimeout(resolve, 10));
+		expect(agent.state.isStreaming).toBe(true);
+
+		// continue() should reject
+		await expect(agent.continue()).rejects.toThrow(
+			"Agent is already processing. Wait for completion before continuing.",
+		);
+
+		// Cleanup
+		agent.abort();
+		await firstPrompt.catch(() => {});
+	});
+
+	it("continue() should process queued follow-up messages after an assistant turn", async () => {
+		const agent = new Agent({
+			streamFn: () => {
+				const stream = new MockAssistantStream();
+				queueMicrotask(() => {
+					stream.push({ type: "done", reason: "stop", message: createAssistantMessage("Processed") });
+				});
+				return stream;
+			},
+		});
+
+		agent.state.messages = [
+			{
+				role: "user",
+				content: [{ type: "text", text: "Initial" }],
+				timestamp: Date.now() - 10,
+			},
+			createAssistantMessage("Initial response"),
+		];
+
+		agent.followUp({
+			role: "user",
+			content: [{ type: "text", text: "Queued follow-up" }],
+			timestamp: Date.now(),
+		});
+
+		await expect(agent.continue()).resolves.toBeUndefined();
+
+		const hasQueuedFollowUp = agent.state.messages.some((message) => {
+			if (message.role !== "user") return false;
+			if (typeof message.content === "string") return message.content === "Queued follow-up";
+			return message.content.some((part) => part.type === "text" && part.text === "Queued follow-up");
+		});
+
+		expect(hasQueuedFollowUp).toBe(true);
+		expect(agent.state.messages[agent.state.messages.length - 1].role).toBe("assistant");
+	});
+
+	it("continue() should keep one-at-a-time steering semantics from assistant tail", async () => {
+		let responseCount = 0;
+		const agent = new Agent({
+			streamFn: () => {
+				const stream = new MockAssistantStream();
+				responseCount++;
+				queueMicrotask(() => {
+					stream.push({
+						type: "done",
+						reason: "stop",
+						message: createAssistantMessage(`Processed ${responseCount}`),
+					});
+				});
+				return stream;
+			},
+		});
+
+		agent.state.messages = [
+			{
+				role: "user",
+				content: [{ type: "text", text: "Initial" }],
+				timestamp: Date.now() - 10,
+			},
+			createAssistantMessage("Initial response"),
+		];
+
+		agent.steer({
+			role: "user",
+			content: [{ type: "text", text: "Steering 1" }],
+			timestamp: Date.now(),
+		});
+		agent.steer({
+			role: "user",
+			content: [{ type: "text", text: "Steering 2" }],
+			timestamp: Date.now() + 1,
+		});
+
+		await expect(agent.continue()).resolves.toBeUndefined();
+
+		const recentMessages = agent.state.messages.slice(-4);
+		expect(recentMessages.map((m) => m.role)).toEqual(["user", "assistant", "user", "assistant"]);
+		expect(responseCount).toBe(2);
+	});
+
+	it("forwards sessionId to streamFn options", async () => {
+		let receivedSessionId: string | undefined;
+		const agent = new Agent({
+			sessionId: "session-abc",
+			streamFn: (_model, _context, options) => {
+				receivedSessionId = options?.sessionId;
+				const stream = new MockAssistantStream();
+				queueMicrotask(() => {
+					const message = createAssistantMessage("ok");
+					stream.push({ type: "done", reason: "stop", message });
+				});
+				return stream;
+			},
+		});
+
+		await agent.prompt("hello");
+		expect(receivedSessionId).toBe("session-abc");
+
+		// Test setter
+		agent.sessionId = "session-def";
+		expect(agent.sessionId).toBe("session-def");
+
+		await agent.prompt("hello again");
+		expect(receivedSessionId).toBe("session-def");
+	});
+});

pi-mono-reference/packages/agent/test/e2e.test.ts

@@ -0,0 +1,404 @@
+import {
+	type AssistantMessage,
+	type FauxProviderRegistration,
+	fauxAssistantMessage,
+	fauxText,
+	fauxThinking,
+	fauxToolCall,
+	type Model,
+	registerFauxProvider,
+	type ToolResultMessage,
+	type UserMessage,
+} from "@mariozechner/pi-ai";
+import { afterEach, describe, expect, it } from "vitest";
+import { Agent, type AgentEvent } from "../src/index.js";
+import { calculateTool } from "./utils/calculate.js";
+
+const registrations: FauxProviderRegistration[] = [];
+
+function createFauxRegistration(options: Parameters<typeof registerFauxProvider>[0] = {}): FauxProviderRegistration {
+	const registration = registerFauxProvider(options);
+	registrations.push(registration);
+	return registration;
+}
+
+function getTextContent(message: AssistantMessage | ToolResultMessage): string {
+	return message.content
+		.filter((block) => block.type === "text")
+		.map((block) => block.text)
+		.join("\n");
+}
+
+afterEach(() => {
+	while (registrations.length > 0) {
+		registrations.pop()?.unregister();
+	}
+});
+
+async function basicPrompt(model: Model<string>) {
+	const agent = new Agent({
+		initialState: {
+			systemPrompt: "You are a helpful assistant. Keep your responses concise.",
+			model,
+			thinkingLevel: "off",
+			tools: [],
+		},
+	});
+
+	await agent.prompt("What is 2+2? Answer with just the number.");
+
+	expect(agent.state.isStreaming).toBe(false);
+	expect(agent.state.messages.length).toBe(2);
+	expect(agent.state.messages[0].role).toBe("user");
+	expect(agent.state.messages[1].role).toBe("assistant");
+
+	const assistantMessage = agent.state.messages[1];
+	if (assistantMessage.role !== "assistant") throw new Error("Expected assistant message");
+	expect(getTextContent(assistantMessage)).toContain("4");
+}
+
+async function toolExecution(model: Model<string>) {
+	const agent = new Agent({
+		initialState: {
+			systemPrompt: "You are a helpful assistant. Always use the calculator tool for math.",
+			model,
+			thinkingLevel: "off",
+			tools: [calculateTool],
+		},
+	});
+
+	const pendingToolCallsDuringEvents: Array<{ type: AgentEvent["type"]; ids: string[] }> = [];
+	agent.subscribe((event) => {
+		if (event.type === "tool_execution_start" || event.type === "tool_execution_end") {
+			pendingToolCallsDuringEvents.push({
+				type: event.type,
+				ids: [...agent.state.pendingToolCalls],
+			});
+		}
+	});
+
+	await agent.prompt("Calculate 123 * 456 using the calculator tool.");
+
+	expect(agent.state.isStreaming).toBe(false);
+	expect(agent.state.messages.length).toBeGreaterThanOrEqual(4);
+	const toolResultMsg = agent.state.messages.find((message) => message.role === "toolResult");
+	expect(toolResultMsg).toBeDefined();
+	if (toolResultMsg?.role !== "toolResult") throw new Error("Expected tool result message");
+	expect(getTextContent(toolResultMsg)).toContain("123 * 456 = 56088");
+
+	const finalMessage = agent.state.messages[agent.state.messages.length - 1];
+	if (finalMessage.role !== "assistant") throw new Error("Expected final assistant message");
+	expect(getTextContent(finalMessage)).toContain("56088");
+	expect(agent.state.pendingToolCalls.size).toBe(0);
+	expect(pendingToolCallsDuringEvents).toEqual([
+		{ type: "tool_execution_start", ids: ["calc-1"] },
+		{ type: "tool_execution_end", ids: [] },
+	]);
+}
+
+async function abortExecution(model: Model<string>) {
+	const agent = new Agent({
+		initialState: {
+			systemPrompt: "You are a helpful assistant.",
+			model,
+			thinkingLevel: "off",
+			tools: [],
+		},
+	});
+
+	const promptPromise = agent.prompt("Count slowly from 1 to 20.");
+	setTimeout(() => {
+		agent.abort();
+	}, 30);
+
+	await promptPromise;
+
+	expect(agent.state.isStreaming).toBe(false);
+	expect(agent.state.messages.length).toBeGreaterThanOrEqual(2);
+
+	const lastMessage = agent.state.messages[agent.state.messages.length - 1];
+	if (lastMessage.role !== "assistant") throw new Error("Expected assistant message");
+	expect(lastMessage.stopReason).toBe("aborted");
+	expect(lastMessage.errorMessage).toBeDefined();
+	expect(agent.state.errorMessage).toBe(lastMessage.errorMessage);
+}
+
+async function stateUpdates(model: Model<string>) {
+	const agent = new Agent({
+		initialState: {
+			systemPrompt: "You are a helpful assistant.",
+			model,
+			thinkingLevel: "off",
+			tools: [],
+		},
+	});
+
+	const events: AgentEvent["type"][] = [];
+	agent.subscribe((event) => {
+		events.push(event.type);
+	});
+
+	await agent.prompt("Count from 1 to 5.");
+
+	expect(events).toContain("agent_start");
+	expect(events).toContain("turn_start");
+	expect(events).toContain("message_start");
+	expect(events).toContain("message_update");
+	expect(events).toContain("message_end");
+	expect(events).toContain("turn_end");
+	expect(events).toContain("agent_end");
+	expect(events.indexOf("agent_start")).toBeLessThan(events.indexOf("message_start"));
+	expect(events.indexOf("message_start")).toBeLessThan(events.indexOf("message_end"));
+	expect(events.indexOf("message_end")).toBeLessThan(events.lastIndexOf("agent_end"));
+
+	expect(agent.state.isStreaming).toBe(false);
+	expect(agent.state.messages.length).toBe(2);
+}
+
+async function multiTurnConversation(model: Model<string>) {
+	const agent = new Agent({
+		initialState: {
+			systemPrompt: "You are a helpful assistant.",
+			model,
+			thinkingLevel: "off",
+			tools: [],
+		},
+	});
+
+	await agent.prompt("My name is Alice.");
+	expect(agent.state.messages.length).toBe(2);
+
+	await agent.prompt("What is my name?");
+	expect(agent.state.messages.length).toBe(4);
+
+	const lastMessage = agent.state.messages[3];
+	if (lastMessage.role !== "assistant") throw new Error("Expected assistant message");
+	expect(getTextContent(lastMessage).toLowerCase()).toContain("alice");
+}
+
+describe("Agent integration with faux provider", () => {
+	it("handles a basic text prompt", async () => {
+		const faux = createFauxRegistration();
+		faux.setResponses([fauxAssistantMessage("4")]);
+		await basicPrompt(faux.getModel());
+	});
+
+	it("executes tools and tracks pending tool calls", async () => {
+		const faux = createFauxRegistration();
+		faux.setResponses([
+			fauxAssistantMessage(
+				[
+					fauxText("Let me calculate that."),
+					fauxToolCall("calculate", { expression: "123 * 456" }, { id: "calc-1" }),
+				],
+				{ stopReason: "toolUse" },
+			),
+			fauxAssistantMessage("The result is 56088."),
+		]);
+		await toolExecution(faux.getModel());
+	});
+
+	it("handles abort during streaming", async () => {
+		const faux = createFauxRegistration({
+			tokensPerSecond: 20,
+			tokenSize: { min: 2, max: 2 },
+		});
+		faux.setResponses([
+			fauxAssistantMessage(
+				"one two three four five six seven eight nine ten eleven twelve thirteen fourteen fifteen",
+			),
+		]);
+		await abortExecution(faux.getModel());
+	});
+
+	it("emits lifecycle updates while streaming", async () => {
+		const faux = createFauxRegistration({ tokenSize: { min: 1, max: 1 } });
+		faux.setResponses([fauxAssistantMessage("1 2 3 4 5")]);
+		await stateUpdates(faux.getModel());
+	});
+
+	it("maintains context across multiple turns", async () => {
+		const faux = createFauxRegistration();
+		faux.setResponses([
+			fauxAssistantMessage("Nice to meet you, Alice."),
+			(context) => {
+				const hasAlice = context.messages.some((message) => {
+					if (message.role !== "user") return false;
+					if (typeof message.content === "string") return message.content.includes("Alice");
+					return message.content.some((block) => block.type === "text" && block.text.includes("Alice"));
+				});
+				return fauxAssistantMessage(hasAlice ? "Your name is Alice." : "I do not know your name.");
+			},
+		]);
+		await multiTurnConversation(faux.getModel());
+	});
+
+	it("preserves thinking content blocks", async () => {
+		const faux = createFauxRegistration({ models: [{ id: "faux-reasoning", reasoning: true }] });
+		faux.setResponses([fauxAssistantMessage([fauxThinking("step by step"), fauxText("4")])]);
+
+		const agent = new Agent({
+			initialState: {
+				systemPrompt: "You are a helpful assistant.",
+				model: faux.getModel(),
+				thinkingLevel: "low",
+				tools: [],
+			},
+		});
+
+		await agent.prompt("What is 2+2?");
+
+		const assistantMessage = agent.state.messages[1];
+		if (assistantMessage?.role !== "assistant") throw new Error("Expected assistant message");
+		expect(assistantMessage.content).toEqual([
+			{ type: "thinking", thinking: "step by step" },
+			{ type: "text", text: "4" },
+		]);
+	});
+});
+
+describe("Agent.continue() with faux provider", () => {
+	describe("validation", () => {
+		it("throws when no messages in context", async () => {
+			const faux = createFauxRegistration();
+			const agent = new Agent({
+				initialState: {
+					systemPrompt: "Test",
+					model: faux.getModel(),
+				},
+			});
+
+			await expect(agent.continue()).rejects.toThrow("No messages to continue from");
+		});
+
+		it("throws when last message is assistant", async () => {
+			const faux = createFauxRegistration();
+			const model = faux.getModel();
+			const agent = new Agent({
+				initialState: {
+					systemPrompt: "Test",
+					model,
+				},
+			});
+
+			const assistantMessage: AssistantMessage = {
+				role: "assistant",
+				content: [{ type: "text", text: "Hello" }],
+				api: model.api,
+				provider: model.provider,
+				model: model.id,
+				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(),
+			};
+			agent.state.messages = [assistantMessage];
+
+			await expect(agent.continue()).rejects.toThrow("Cannot continue from message role: assistant");
+		});
+	});
+
+	describe("continue from user message", () => {
+		it("continues and gets a response when last message is user", async () => {
+			const faux = createFauxRegistration();
+			faux.setResponses([fauxAssistantMessage("HELLO WORLD")]);
+			const agent = new Agent({
+				initialState: {
+					systemPrompt: "You are a helpful assistant. Follow instructions exactly.",
+					model: faux.getModel(),
+					thinkingLevel: "off",
+					tools: [],
+				},
+			});
+
+			const userMessage: UserMessage = {
+				role: "user",
+				content: [{ type: "text", text: "Say exactly: HELLO WORLD" }],
+				timestamp: Date.now(),
+			};
+			agent.state.messages = [userMessage];
+
+			await agent.continue();
+
+			expect(agent.state.isStreaming).toBe(false);
+			expect(agent.state.messages.length).toBe(2);
+			expect(agent.state.messages[0].role).toBe("user");
+			expect(agent.state.messages[1].role).toBe("assistant");
+
+			const assistantMsg = agent.state.messages[1];
+			if (assistantMsg.role !== "assistant") throw new Error("Expected assistant message");
+			expect(getTextContent(assistantMsg).toUpperCase()).toContain("HELLO WORLD");
+		});
+	});
+
+	describe("continue from tool result", () => {
+		it("continues and processes tool results", async () => {
+			const faux = createFauxRegistration();
+			const model = faux.getModel();
+			faux.setResponses([fauxAssistantMessage("The answer is 8.")]);
+			const agent = new Agent({
+				initialState: {
+					systemPrompt:
+						"You are a helpful assistant. After getting a calculation result, state the answer clearly.",
+					model,
+					thinkingLevel: "off",
+					tools: [calculateTool],
+				},
+			});
+
+			const userMessage: UserMessage = {
+				role: "user",
+				content: [{ type: "text", text: "What is 5 + 3?" }],
+				timestamp: Date.now(),
+			};
+
+			const assistantMessage: AssistantMessage = {
+				role: "assistant",
+				content: [
+					{ type: "text", text: "Let me calculate that." },
+					{ type: "toolCall", id: "calc-1", name: "calculate", arguments: { expression: "5 + 3" } },
+				],
+				api: model.api,
+				provider: model.provider,
+				model: model.id,
+				usage: {
+					input: 0,
+					output: 0,
+					cacheRead: 0,
+					cacheWrite: 0,
+					totalTokens: 0,
+					cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 },
+				},
+				stopReason: "toolUse",
+				timestamp: Date.now(),
+			};
+
+			const toolResult: ToolResultMessage = {
+				role: "toolResult",
+				toolCallId: "calc-1",
+				toolName: "calculate",
+				content: [{ type: "text", text: "5 + 3 = 8" }],
+				isError: false,
+				timestamp: Date.now(),
+			};
+
+			agent.state.messages = [userMessage, assistantMessage, toolResult];
+
+			await agent.continue();
+
+			expect(agent.state.isStreaming).toBe(false);
+			expect(agent.state.messages.length).toBeGreaterThanOrEqual(4);
+
+			const lastMessage = agent.state.messages[agent.state.messages.length - 1];
+			expect(lastMessage.role).toBe("assistant");
+			if (lastMessage.role !== "assistant") throw new Error("Expected assistant message");
+			expect(getTextContent(lastMessage)).toContain("8");
+		});
+	});
+});

pi-mono-reference/packages/agent/test/utils/calculate.ts

@@ -0,0 +1,32 @@
+import { type Static, Type } from "@sinclair/typebox";
+import type { AgentTool, AgentToolResult } from "../../src/types.js";
+
+export interface CalculateResult extends AgentToolResult<undefined> {
+	content: Array<{ type: "text"; text: string }>;
+	details: undefined;
+}
+
+export function calculate(expression: string): CalculateResult {
+	try {
+		const result = new Function(`return ${expression}`)();
+		return { content: [{ type: "text", text: `${expression} = ${result}` }], details: undefined };
+	} catch (e: any) {
+		throw new Error(e.message || String(e));
+	}
+}
+
+const calculateSchema = Type.Object({
+	expression: Type.String({ description: "The mathematical expression to evaluate" }),
+});
+
+type CalculateParams = Static<typeof calculateSchema>;
+
+export const calculateTool: AgentTool<typeof calculateSchema, undefined> = {
+	label: "Calculator",
+	name: "calculate",
+	description: "Evaluate mathematical expressions",
+	parameters: calculateSchema,
+	execute: async (_toolCallId: string, args: CalculateParams) => {
+		return calculate(args.expression);
+	},
+};

pi-mono-reference/packages/agent/test/utils/get-current-time.ts

@@ -0,0 +1,46 @@
+import { type Static, Type } from "@sinclair/typebox";
+import type { AgentTool, AgentToolResult } from "../../src/types.js";
+
+export interface GetCurrentTimeResult extends AgentToolResult<{ utcTimestamp: number }> {}
+
+export async function getCurrentTime(timezone?: string): Promise<GetCurrentTimeResult> {
+	const date = new Date();
+	if (timezone) {
+		try {
+			const timeStr = date.toLocaleString("en-US", {
+				timeZone: timezone,
+				dateStyle: "full",
+				timeStyle: "long",
+			});
+			return {
+				content: [{ type: "text", text: timeStr }],
+				details: { utcTimestamp: date.getTime() },
+			};
+		} catch (_e) {
+			throw new Error(`Invalid timezone: ${timezone}. Current UTC time: ${date.toISOString()}`);
+		}
+	}
+	const timeStr = date.toLocaleString("en-US", { dateStyle: "full", timeStyle: "long" });
+	return {
+		content: [{ type: "text", text: timeStr }],
+		details: { utcTimestamp: date.getTime() },
+	};
+}
+
+const getCurrentTimeSchema = Type.Object({
+	timezone: Type.Optional(
+		Type.String({ description: "Optional timezone (e.g., 'America/New_York', 'Europe/London')" }),
+	),
+});
+
+type GetCurrentTimeParams = Static<typeof getCurrentTimeSchema>;
+
+export const getCurrentTimeTool: AgentTool<typeof getCurrentTimeSchema, { utcTimestamp: number }> = {
+	label: "Current Time",
+	name: "get_current_time",
+	description: "Get the current date and time",
+	parameters: getCurrentTimeSchema,
+	execute: async (_toolCallId: string, args: GetCurrentTimeParams) => {
+		return getCurrentTime(args.timezone);
+	},
+};