feat: support streaming tool outputs by returning an AsyncIterable fr…

.changeset/upset-shrimps-guess.md

@@ -0,0 +1,30 @@
+---
+"@voltagent/core": patch
+---
+
+feat: support streaming tool outputs by returning an AsyncIterable from `execute`, emitting preliminary results before the final output.
+
+```ts
+import { createTool } from "@voltagent/core";
+import { z } from "zod";
+
+const weatherTool = createTool({
+  name: "get_weather",
+  description: "Get the current weather for a location",
+  parameters: z.object({
+    location: z.string(),
+  }),
+  async *execute({ location }) {
+    yield { status: "loading" as const, text: `Getting weather for ${location}` };
+
+    await new Promise((resolve) => setTimeout(resolve, 3000));
+
+    const temperature = 72;
+    yield {
+      status: "success" as const,
+      text: `The weather in ${location} is ${temperature}F`,
+      temperature,
+    };
+  },
+});
+```

examples/with-tools/src/tools/weather.ts

@@ -1,17 +1,25 @@
 import { createTool } from "@voltagent/core";
 import { z } from "zod";
 
-// Define the output schema for weather data
-const weatherOutputSchema = z.object({
-  weather: z.object({
-    location: z.string(),
-    temperature: z.number(),
-    condition: z.string(),
-    humidity: z.number(),
-    windSpeed: z.number(),
+// Define the output schema for weather data (supports preliminary updates)
+const weatherOutputSchema = z.discriminatedUnion("status", [
+  z.object({
+    status: z.literal("loading"),
+    text: z.string(),
+    weather: z.undefined().optional(),
   }),
-  message: z.string(),
-});
+  z.object({
+    status: z.literal("success"),
+    text: z.string(),
+    weather: z.object({
+      location: z.string(),
+      temperature: z.number(),
+      condition: z.string(),
+      humidity: z.number(),
+      windSpeed: z.number(),
+    }),
+  }),
+]);
 
 /**
  * A tool for fetching weather information for a given location
@@ -24,11 +32,18 @@ export const weatherTool = createTool({
     location: z.string().describe("The city or location to get weather for"),
   }),
   outputSchema: weatherOutputSchema,
-  needsApproval: true,
+  /* needsApproval: true, */
+
+  async *execute({ location }) {
+    yield {
+      status: "loading" as const,
+      text: `Getting weather for ${location}`,
+      weather: undefined,
+    };
 
-  execute: async ({ location }) => {
     // In a real implementation, this would call a weather API
     // This is a mock implementation for demonstration purposes
+    await new Promise((resolve) => setTimeout(resolve, 3000));
 
     const mockWeatherData = {
       location,
@@ -40,9 +55,10 @@ export const weatherTool = createTool({
       windSpeed: Math.floor(Math.random() * 30), // Random wind speed between 0-30 km/h
     };
 
-    return {
+    yield {
+      status: "success" as const,
+      text: `Current weather in ${location}: ${mockWeatherData.temperature}°C and ${mockWeatherData.condition.toLowerCase()} with ${mockWeatherData.humidity}% humidity and wind speed of ${mockWeatherData.windSpeed} km/h.`,
       weather: mockWeatherData,
-      message: `Current weather in ${location}: ${mockWeatherData.temperature}°C and ${mockWeatherData.condition.toLowerCase()} with ${mockWeatherData.humidity}% humidity and wind speed of ${mockWeatherData.windSpeed} km/h.`,
     };
   },
 });

packages/core/src/agent/agent.ts

@@ -50,7 +50,7 @@ import { TRIGGER_CONTEXT_KEY } from "../observability/context-keys";
 import { type ObservabilityFlushState, flushObservability } from "../observability/utils";
 import { AgentRegistry } from "../registries/agent-registry";
 import type { BaseRetriever } from "../retriever/retriever";
-import type { Tool, Toolkit, VercelTool } from "../tool";
+import type { Tool, ToolExecutionResult, Toolkit, VercelTool } from "../tool";
 import { createTool } from "../tool";
 import { ToolManager } from "../tool/manager";
 import { randomUUID } from "../utils/id";
@@ -320,6 +320,22 @@ function createDeferred<T>(): Deferred<T> {
   return { promise, resolve };
 }
 
+const asyncGeneratorFunction = Object.getPrototypeOf(async function* () {}).constructor;
+
+function isAsyncGeneratorFunction(
+  value: unknown,
+): value is (...args: any[]) => AsyncIterable<unknown> {
+  return typeof value === "function" && value.constructor === asyncGeneratorFunction;
+}
+
+function isAsyncIterable(value: unknown): value is AsyncIterable<unknown> {
+  if (typeof value !== "object" || value === null) {
+    return false;
+  }
+
+  return typeof (value as AsyncIterable<unknown>)[Symbol.asyncIterator] === "function";
+}
+
 /**
  * Base options for all generation methods
  * Extends AI SDK's CallSettings for full compatibility
@@ -3583,8 +3599,8 @@ export class Agent {
   private createToolExecutionFactory(
     oc: OperationContext,
     hooks: AgentHooks,
-  ): (tool: BaseTool) => (args: any, options?: ToolExecutionOptions) => Promise<any> {
-    return (tool: BaseTool) => async (args: any, options?: ToolExecutionOptions) => {
+  ): (tool: BaseTool) => (args: any, options?: ToolExecutionOptions) => ToolExecutionResult<any> {
+    return (tool: BaseTool) => (args: any, options?: ToolExecutionOptions) => {
       // AI SDK passes ToolExecutionOptions with fields: toolCallId, messages, abortSignal
       const toolCallId = options?.toolCallId ?? randomUUID();
       const messages = options?.messages ?? [];
@@ -3621,8 +3637,114 @@ export class Agent {
       oc.systemContext.set("historyEntryId", oc.operationId);
       oc.systemContext.set("parentToolSpan", toolSpan);
 
-      // Execute tool and handle span lifecycle
-      return await oc.traceContext.withSpan(toolSpan, async () => {
+      const handleToolSuccess = async (result: any, validatedResult: any) => {
+        toolSpan.setAttribute("output", safeStringify(result));
+        toolSpan.setStatus({ code: SpanStatusCode.OK });
+        toolSpan.end();
+
+        await hooks.onToolEnd?.({
+          agent: this,
+          tool,
+          output: validatedResult,
+          error: undefined,
+          context: oc,
+          options: executionOptions,
+        });
+      };
+
+      const handleToolError = async (errorValue: unknown) => {
+        const error = errorValue instanceof Error ? errorValue : new Error(String(errorValue));
+        const voltAgentError = createVoltAgentError(error, {
+          stage: "tool_execution",
+          toolError: {
+            toolCallId,
+            toolName: tool.name,
+            toolExecutionError: error,
+            toolArguments: args,
+          },
+        });
+        const errorResult = buildToolErrorResult(error, toolCallId, tool.name);
+
+        toolSpan.setStatus({ code: SpanStatusCode.ERROR, message: error.message });
+        toolSpan.recordException(error);
+        toolSpan.end();
+
+        await hooks.onToolEnd?.({
+          agent: this,
+          tool,
+          output: undefined,
+          error: voltAgentError,
+          context: oc,
+          options: executionOptions,
+        });
+
+        if (isToolDeniedError(errorValue)) {
+          oc.abortController.abort(errorValue);
+        }
+
+        return errorResult;
+      };
+
+      const execute = tool.execute;
+      if (execute && isAsyncGeneratorFunction(execute)) {
+        return async function* (this: Agent): AsyncGenerator<any, void, void> {
+          try {
+            await oc.traceContext.withSpan(toolSpan, async () => {
+              await hooks.onToolStart?.({
+                agent: this,
+                tool,
+                context: oc,
+                args,
+                options: executionOptions,
+              });
+            });
+
+            const result = execute(args, executionOptions);
+
+            if (!isAsyncIterable(result)) {
+              const resolved = await result;
+              const validatedResult = await this.validateToolOutput(resolved, tool);
+              await oc.traceContext.withSpan(toolSpan, async () => {
+                await handleToolSuccess(resolved, validatedResult);
+              });
+              yield resolved;
+              return;
+            }
+
+            const iterator = result[Symbol.asyncIterator]();
+            let finalOutput: any = undefined;
+            let validatedResult: any = undefined;
+            let hasOutput = false;
+
+            while (true) {
+              const next = await oc.traceContext.withSpan(toolSpan, () => iterator.next());
+              if (next.done) {
+                break;
+              }
+
+              finalOutput = next.value;
+              hasOutput = true;
+              validatedResult = await this.validateToolOutput(finalOutput, tool);
+              yield finalOutput;
+            }
+
+            if (!hasOutput) {
+              validatedResult = await this.validateToolOutput(finalOutput, tool);
+            }
+
+            await oc.traceContext.withSpan(toolSpan, async () => {
+              await handleToolSuccess(finalOutput, validatedResult);
+            });
+          } catch (e) {
+            const errorResult = await oc.traceContext.withSpan(toolSpan, async () => {
+              return await handleToolError(e);
+            });
+            yield errorResult;
+          }
+        }.call(this);
+      }
+
+      return oc.traceContext.withSpan(toolSpan, async () => {
         try {
           // Call tool start hook - can throw ToolDeniedError
           await hooks.onToolStart?.({
@@ -3637,56 +3759,23 @@ export class Agent {
           if (!tool.execute) {
             throw new Error(`Tool ${tool.name} does not have "execute" method`);
           }
-          const result = await tool.execute(args, executionOptions);
-          const validatedResult = await this.validateToolOutput(result, tool);
+          let result = await tool.execute(args, executionOptions);
 
-          // End OTEL span
-          toolSpan.setAttribute("output", safeStringify(result));
-          toolSpan.setStatus({ code: SpanStatusCode.OK });
-          toolSpan.end();
+          if (isAsyncIterable(result)) {
+            let lastOutput: any = undefined;
+            for await (const output of result) {
+              lastOutput = output;
+            }
+            result = lastOutput;
+          }
 
-          // Call tool end hook
-          await hooks.onToolEnd?.({
-            agent: this,
-            tool,
-            output: validatedResult,
-            error: undefined,
-            context: oc,
-            options: executionOptions,
-          });
+          const validatedResult = await this.validateToolOutput(result, tool);
+
+          await handleToolSuccess(result, validatedResult);
 
           return result;
         } catch (e) {
-          const error = e instanceof Error ? e : new Error(String(e));
-          const voltAgentError = createVoltAgentError(error, {
-            stage: "tool_execution",
-            toolError: {
-              toolCallId,
-              toolName: tool.name,
-              toolExecutionError: error,
-              toolArguments: args,
-            },
-          });
-          const errorResult = buildToolErrorResult(error, toolCallId, tool.name);
-
-          toolSpan.setStatus({ code: SpanStatusCode.ERROR, message: error.message });
-          toolSpan.recordException(error);
-          toolSpan.end();
-
-          await hooks.onToolEnd?.({
-            agent: this,
-            tool,
-            output: undefined,
-            error: voltAgentError,
-            context: oc,
-            options: executionOptions,
-          });
-
-          if (isToolDeniedError(e)) {
-            oc.abortController.abort(e);
-          }
-
-          return errorResult;
+          return await handleToolError(e);
         } finally {
           // End the span if it was created
           oc.traceContext.endChildSpan(toolSpan, "completed", {});

packages/core/src/agent/test-utils.ts

@@ -362,7 +362,7 @@ export function createTestAgent(options?: Partial<AgentOptions>): Agent {
  */
 export function createMockTool(
   name: string,
-  execute?: (params: any, options?: any) => Promise<any> | any,
+  execute?: (params: any, options?: any) => PromiseLike<any> | AsyncIterable<any> | any,
   options?: {
     description?: string;
     parameters?: ToolSchema;

packages/core/src/tool/index.ts

@@ -9,6 +9,8 @@ import { LoggerProxy } from "../logger";
  */
 type JSONValue = string | number | boolean | null | { [key: string]: JSONValue } | Array<JSONValue>;
 
+export type ToolExecutionResult<T> = PromiseLike<T> | AsyncIterable<T> | T;
+
 /**
  * Tool result output format for multi-modal content.
  * Matches AI SDK's LanguageModelV2ToolResultOutput type.
@@ -134,11 +136,12 @@ export type ToolOptions<
    * Function to execute when the tool is called.
    * @param args - The arguments passed to the tool
    * @param options - Optional execution options including context, abort signals, etc.
+   * @returns A result or an AsyncIterable of results (last value is final).
    */
   execute?: (
     args: z.infer<T>,
     options?: ToolExecuteOptions,
-  ) => Promise<O extends ToolSchema ? z.infer<O> : unknown>;
+  ) => ToolExecutionResult<O extends ToolSchema ? z.infer<O> : unknown>;
 };
 
 /**
@@ -205,11 +208,12 @@ export class Tool<T extends ToolSchema = ToolSchema, O extends ToolSchema | unde
    * Function to execute when the tool is called.
    * @param args - The arguments passed to the tool
    * @param options - Optional execution options including context, abort signals, etc.
+   * @returns A result or an AsyncIterable of results (last value is final).
    */
   readonly execute?: (
     args: z.infer<T>,
     options?: ToolExecuteOptions,
-  ) => Promise<O extends ToolSchema ? z.infer<O> : unknown>;
+  ) => ToolExecutionResult<O extends ToolSchema ? z.infer<O> : unknown>;
 
   /**
    * Whether this tool should be executed on the client side.

packages/core/src/tool/manager/ToolManager.ts

@@ -2,7 +2,7 @@ import type { ToolExecutionOptions } from "@ai-sdk/provider-utils";
 import type { Logger } from "@voltagent/internal";
 import type { ApiToolInfo } from "../../agent/types";
 import { zodSchemaToJsonUI } from "../../utils/toolParser";
-import type { AgentTool, ProviderTool, VercelTool } from "../index";
+import type { AgentTool, ProviderTool, ToolExecutionResult, VercelTool } from "../index";
 import type { Toolkit } from "../toolkit";
 import { BaseToolManager } from "./BaseToolManager";
 import { ToolkitManager } from "./ToolkitManager";
@@ -48,12 +48,12 @@ export class ToolManager extends BaseToolManager<AgentTool | VercelTool | Toolki
   public prepareToolsForExecution(
     createToolExecuteFunction: (
       tool: AgentTool,
-    ) => (args: any, options?: ToolExecutionOptions) => Promise<any>,
+    ) => (args: any, options?: ToolExecutionOptions) => ToolExecutionResult<any>,
   ): Record<string, any> {
     type ManagedTool = {
       description: string;
       inputSchema: AgentTool["parameters"];
-      execute?: (args: any, options?: ToolExecutionOptions) => Promise<any>;
+      execute?: (args: any, options?: ToolExecutionOptions) => ToolExecutionResult<any>;
       needsApproval?: AgentTool["needsApproval"];
       providerOptions?: AgentTool["providerOptions"];
       toModelOutput?: AgentTool["toModelOutput"];