feat(ai): Implement Zod-based tool validation and improve Agent API

- Replace JSON Schema with Zod schemas for tool parameter definitions
- Add runtime validation for all tool calls at provider level
- Create shared validation module with detailed error formatting
- Update Agent API with comprehensive event system
- Add agent tests with calculator tool for multi-turn execution
- Add abort test to verify proper handling of aborted requests
- Update documentation with detailed event flow examples
- Rename generate.ts to stream.ts for clarity

packages/ai/src/agent/agent.ts

@@ -1,97 +1,68 @@
 import { EventStream } from "../event-stream";
-import { streamSimple } from "../generate.js";
-import type {
-	AssistantMessage,
-	Context,
-	Message,
-	Model,
-	SimpleGenerateOptions,
-	ToolResultMessage,
-	UserMessage,
-} from "../types.js";
-import type { AgentContext, AgentTool, AgentToolResult } from "./types";
-
-// Event types
-export type AgentEvent =
-	| { type: "message_start"; message: Message }
-	| { type: "message_update"; message: AssistantMessage }
-	| { type: "message_complete"; message: Message }
-	| { type: "tool_execution_start"; toolCallId: string; toolName: string; args: any }
-	| {
-			type: "tool_execution_complete";
-			toolCallId: string;
-			toolName: string;
-			result: AgentToolResult<any> | string;
-			isError: boolean;
-	  }
-	| { type: "turn_complete"; messages: AgentContext["messages"] };
-
-// Configuration for prompt execution
-export interface PromptConfig {
-	model: Model<any>;
-	apiKey: string;
-	enableThinking?: boolean;
-	preprocessor?: (messages: AgentContext["messages"], abortSignal?: AbortSignal) => Promise<AgentContext["messages"]>;
-}
+import { streamSimple } from "../stream.js";
+import type { AssistantMessage, Context, Message, ToolResultMessage, UserMessage } from "../types.js";
+import { validateToolArguments } from "../validation.js";
+import type { AgentContext, AgentEvent, AgentTool, AgentToolResult, PromptConfig } from "./types";
 
 // Main prompt function - returns a stream of events
 export function prompt(
+	prompt: UserMessage,
 	context: AgentContext,
 	config: PromptConfig,
-	prompt: UserMessage,
 	signal?: AbortSignal,
 ): EventStream<AgentEvent, AgentContext["messages"]> {
 	const stream = new EventStream<AgentEvent, AgentContext["messages"]>(
-		(event) => event.type === "turn_complete",
-		(event) => (event.type === "turn_complete" ? event.messages : []),
+		(event) => event.type === "agent_end",
+		(event) => (event.type === "agent_end" ? event.messages : []),
 	);
 
 	// Run the prompt async
 	(async () => {
-		try {
-			// Track new messages generated during this prompt
-			const newMessages: AgentContext["messages"] = [];
+		// Track new messages generated during this prompt
+		const newMessages: AgentContext["messages"] = [];
+		// Create user message for the prompt
+		const messages = [...context.messages, prompt];
+		newMessages.push(prompt);
 
-			// Create user message
-			const messages = [...context.messages, prompt];
-			newMessages.push(prompt);
+		stream.push({ type: "agent_start" });
+		stream.push({ type: "turn_start" });
+		stream.push({ type: "message_start", message: prompt });
+		stream.push({ type: "message_end", message: prompt });
 
-			stream.push({ type: "message_start", message: prompt });
-			stream.push({ type: "message_complete", message: prompt });
+		// Update context with new messages
+		const currentContext: AgentContext = {
+			...context,
+			messages,
+		};
 
-			// Update context with new messages
-			const currentContext: AgentContext = {
-				...context,
-				messages,
-			};
-
-			// Keep looping while we have tool calls
-			let hasMoreToolCalls = true;
-			while (hasMoreToolCalls) {
-				// Stream assistant response
-				const assistantMessage = await streamAssistantResponse(currentContext, config, signal, stream);
-				newMessages.push(assistantMessage);
-
-				// Check for tool calls
-				const toolCalls = assistantMessage.content.filter((c) => c.type === "toolCall");
-				hasMoreToolCalls = toolCalls.length > 0;
-
-				if (hasMoreToolCalls) {
-					// Execute tool calls
-					const toolResults = await executeToolCalls(currentContext.tools, assistantMessage, signal, stream);
-					newMessages.push(...toolResults);
-
-					// Add tool results to context
-					currentContext.messages = [...currentContext.messages, ...toolResults];
-				}
+		// Keep looping while we have tool calls
+		let hasMoreToolCalls = true;
+		let firstTurn = true;
+		while (hasMoreToolCalls) {
+			if (!firstTurn) {
+				stream.push({ type: "turn_start" });
+			} else {
+				firstTurn = false;
 			}
+			// Stream assistant response
+			const assistantMessage = await streamAssistantResponse(currentContext, config, signal, stream);
+			newMessages.push(assistantMessage);
 
-			stream.push({ type: "turn_complete", messages: newMessages });
-		} catch (error) {
-			// End stream on error
-			stream.end([]);
-			throw error;
+			// Check for tool calls
+			const toolCalls = assistantMessage.content.filter((c) => c.type === "toolCall");
+			hasMoreToolCalls = toolCalls.length > 0;
+
+			const toolResults: ToolResultMessage[] = [];
+			if (hasMoreToolCalls) {
+				// Execute tool calls
+				toolResults.push(...(await executeToolCalls(currentContext.tools, assistantMessage, signal, stream)));
+				currentContext.messages.push(...toolResults);
+				newMessages.push(...toolResults);
+			}
+			stream.push({ type: "turn_end", assistantMessage, toolResults: toolResults });
 		}
+		stream.push({ type: "agent_end", messages: newMessages });
+		stream.end(newMessages);
 	})();
 
 	return stream;
@@ -122,16 +93,7 @@ async function streamAssistantResponse(
 		tools: context.tools, // AgentTool extends Tool, so this works
 	};
 
-	const options: SimpleGenerateOptions = {
-		apiKey: config.apiKey,
-		signal,
-	};
-
-	if (config.model.reasoning && config.enableThinking) {
-		options.reasoning = "medium";
-	}
-
-	const response = await streamSimple(config.model, processedContext, options);
+	const response = await streamSimple(config.model, processedContext, { ...config, signal });
 
 	let partialMessage: AssistantMessage | null = null;
 	let addedPartial = false;
@@ -147,14 +109,17 @@ async function streamAssistantResponse(
 
 			case "text_start":
 			case "text_delta":
+			case "text_end":
 			case "thinking_start":
 			case "thinking_delta":
+			case "thinking_end":
 			case "toolcall_start":
 			case "toolcall_delta":
+			case "toolcall_end":
 				if (partialMessage) {
 					partialMessage = event.partial;
 					context.messages[context.messages.length - 1] = partialMessage;
-					stream.push({ type: "message_update", message: { ...partialMessage } });
+					stream.push({ type: "message_update", assistantMessageEvent: event, message: { ...partialMessage } });
 				}
 				break;
 
@@ -166,7 +131,7 @@ async function streamAssistantResponse(
 				} else {
 					context.messages.push(finalMessage);
 				}
-				stream.push({ type: "message_complete", message: finalMessage });
+				stream.push({ type: "message_end", message: finalMessage });
 				return finalMessage;
 			}
 		}
@@ -176,7 +141,7 @@ async function streamAssistantResponse(
 }
 
 async function executeToolCalls<T>(
-	tools: AgentTool<T>[] | undefined,
+	tools: AgentTool<any, T>[] | undefined,
 	assistantMessage: AssistantMessage,
 	signal: AbortSignal | undefined,
 	stream: EventStream<AgentEvent, Message[]>,
@@ -199,14 +164,19 @@ async function executeToolCalls<T>(
 
 		try {
 			if (!tool) throw new Error(`Tool ${toolCall.name} not found`);
-			resultOrError = await tool.execute(toolCall.arguments, toolCall.id, signal);
+
+			// Validate arguments using shared validation function
+			const validatedArgs = validateToolArguments(tool, toolCall);
+
+			// Execute with validated, typed arguments
+			resultOrError = await tool.execute(toolCall.id, validatedArgs, signal);
 		} catch (e) {
-			resultOrError = `Error: ${e instanceof Error ? e.message : String(e)}`;
+			resultOrError = e instanceof Error ? e.message : String(e);
 			isError = true;
 		}
 
 		stream.push({
-			type: "tool_execution_complete",
+			type: "tool_execution_end",
 			toolCallId: toolCall.id,
 			toolName: toolCall.name,
 			result: resultOrError,
@@ -224,7 +194,7 @@ async function executeToolCalls<T>(
 
 		results.push(toolResultMessage);
 		stream.push({ type: "message_start", message: toolResultMessage });
-		stream.push({ type: "message_complete", message: toolResultMessage });
+		stream.push({ type: "message_end", message: toolResultMessage });
 	}
 
 	return results;

packages/ai/src/agent/index.ts

@@ -1,3 +1,3 @@
-export { type AgentEvent, type PromptConfig, prompt } from "./agent";
+export { prompt } from "./agent";
 export * from "./tools";
-export type { AgentContext, AgentTool } from "./types";
+export type { AgentContext, AgentEvent, AgentTool, PromptConfig } from "./types";

packages/ai/src/agent/tools/calculate.ts

@@ -1,3 +1,4 @@
+import { z } from "zod";
 import type { AgentTool } from "../../agent";
 
 export interface CalculateResult {
@@ -14,21 +15,16 @@ export function calculate(expression: string): CalculateResult {
 	}
 }
 
-export const calculateTool: AgentTool<undefined> = {
+const calculateSchema = z.object({
+	expression: z.string().describe("The mathematical expression to evaluate"),
+});
+
+export const calculateTool: AgentTool<typeof calculateSchema, undefined> = {
 	label: "Calculator",
 	name: "calculate",
 	description: "Evaluate mathematical expressions",
-	parameters: {
-		type: "object",
-		properties: {
-			expression: {
-				type: "string",
-				description: "The mathematical expression to evaluate",
-			},
-		},
-		required: ["expression"],
-	},
-	execute: async (args: { expression: string }) => {
+	parameters: calculateSchema,
+	execute: async (_toolCallId, args) => {
 		return calculate(args.expression);
 	},
 };

packages/ai/src/agent/tools/get-current-time.ts

@@ -1,3 +1,4 @@
+import { z } from "zod";
 import type { AgentTool } from "../../agent";
 import type { AgentToolResult } from "../types";
 
@@ -25,20 +26,16 @@ export async function getCurrentTime(timezone?: string): Promise<GetCurrentTimeR
 	};
 }
 
-export const getCurrentTimeTool: AgentTool<{ utcTimestamp: number }> = {
+const getCurrentTimeSchema = z.object({
+	timezone: z.string().optional().describe("Optional timezone (e.g., 'America/New_York', 'Europe/London')"),
+});
+
+export const getCurrentTimeTool: AgentTool<typeof getCurrentTimeSchema, { utcTimestamp: number }> = {
 	label: "Current Time",
 	name: "get_current_time",
 	description: "Get the current date and time",
-	parameters: {
-		type: "object",
-		properties: {
-			timezone: {
-				type: "string",
-				description: "Optional timezone (e.g., 'America/New_York', 'Europe/London')",
-			},
-		},
-	},
-	execute: async (args: { timezone?: string }) => {
+	parameters: getCurrentTimeSchema,
+	execute: async (_toolCallId, args) => {
 		return getCurrentTime(args.timezone);
 	},
 };

packages/ai/src/agent/types.ts

@@ -1,4 +1,13 @@
-import type { Message, Tool } from "../types.js";
+import type { ZodSchema, z } from "zod";
+import type {
+	AssistantMessage,
+	AssistantMessageEvent,
+	Message,
+	Model,
+	SimpleStreamOptions,
+	Tool,
+	ToolResultMessage,
+} from "../types.js";
 
 export interface AgentToolResult<T> {
 	// Output of the tool to be given to the LLM in ToolResultMessage.content
@@ -8,10 +17,14 @@ export interface AgentToolResult<T> {
 }
 
 // AgentTool extends Tool but adds the execute function
-export interface AgentTool<TDetails> extends Tool {
+export interface AgentTool<TParameters extends ZodSchema = ZodSchema, TDetails = any> extends Tool<TParameters> {
 	// A human-readable label for the tool to be displayed in UI
 	label: string;
-	execute: (params: any, toolCallId: string, signal?: AbortSignal) => Promise<AgentToolResult<TDetails>>;
+	execute: (
+		toolCallId: string,
+		params: z.infer<TParameters>,
+		signal?: AbortSignal,
+	) => Promise<AgentToolResult<TDetails>>;
 }
 
 // AgentContext is like Context but uses AgentTool
@@ -20,3 +33,37 @@ export interface AgentContext {
 	messages: Message[];
 	tools?: AgentTool<any>[];
 }
+
+// Event types
+export type AgentEvent =
+	// Emitted when the agent starts. An agent can emit multiple turns
+	| { type: "agent_start" }
+	// Emitted when a turn starts. A turn can emit an optional user message (initial prompt), an assistant message (response) and multiple tool result messages
+	| { type: "turn_start" }
+	// Emitted when a user, assistant or tool result message starts
+	| { type: "message_start"; message: Message }
+	// Emitted when an asssitant messages is updated due to streaming
+	| { type: "message_update"; assistantMessageEvent: AssistantMessageEvent; message: AssistantMessage }
+	// Emitted when a user, assistant or tool result message is complete
+	| { type: "message_end"; message: Message }
+	// Emitted when a tool execution starts
+	| { type: "tool_execution_start"; toolCallId: string; toolName: string; args: any }
+	// Emitted when a tool execution completes
+	| {
+			type: "tool_execution_end";
+			toolCallId: string;
+			toolName: string;
+			result: AgentToolResult<any> | string;
+			isError: boolean;
+	  }
+	// Emitted when a full turn completes
+	| { type: "turn_end"; assistantMessage: AssistantMessage; toolResults: ToolResultMessage[] }
+	// Emitted when the agent has completed all its turns. All messages from every turn are
+	// contained in messages, which can be appended to the context
+	| { type: "agent_end"; messages: AgentContext["messages"] };
+
+// Configuration for prompt execution
+export interface PromptConfig extends SimpleStreamOptions {
+	model: Model<any>;
+	preprocessor?: (messages: AgentContext["messages"], abortSignal?: AbortSignal) => Promise<AgentContext["messages"]>;
+}