From 1ebe40bbc352b22977974451c8529ea2b314c180 Mon Sep 17 00:00:00 2001 From: Omer Aplak Date: Thu, 15 Jan 2026 11:14:32 -0800 Subject: [PATCH 1/2] feat: support streaming tool outputs by returning an AsyncIterable from `execute` --- .changeset/upset-shrimps-guess.md | 30 +++ examples/with-tools/src/tools/weather.ts | 44 ++-- packages/core/src/agent/agent.ts | 191 +++++++++++++----- packages/core/src/agent/test-utils.ts | 2 +- packages/core/src/tool/index.ts | 8 +- packages/core/src/tool/manager/ToolManager.ts | 6 +- packages/core/src/tool/manager/index.spec.ts | 13 +- website/docs/agents/tools.md | 52 +++++ 8 files changed, 273 insertions(+), 73 deletions(-) create mode 100644 .changeset/upset-shrimps-guess.md diff --git a/.changeset/upset-shrimps-guess.md b/.changeset/upset-shrimps-guess.md new file mode 100644 index 000000000..0def4f0a1 --- /dev/null +++ b/.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, + }; + }, +}); +``` diff --git a/examples/with-tools/src/tools/weather.ts b/examples/with-tools/src/tools/weather.ts index b4b62d0e7..f03b8aae4 100644 --- a/examples/with-tools/src/tools/weather.ts +++ b/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.`, }; }, }); diff --git a/packages/core/src/agent/agent.ts b/packages/core/src/agent/agent.ts index 66fbfc7f4..b98b3fed7 100644 --- a/packages/core/src/agent/agent.ts +++ b/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(): Deferred { return { promise, resolve }; } +const asyncGeneratorFunction = Object.getPrototypeOf(async function* () {}).constructor; + +function isAsyncGeneratorFunction( + value: unknown, +): value is (...args: any[]) => AsyncIterable { + return typeof value === "function" && value.constructor === asyncGeneratorFunction; +} + +function isAsyncIterable(value: unknown): value is AsyncIterable { + if (typeof value !== "object" || value === null) { + return false; + } + + return typeof (value as AsyncIterable)[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 { - return (tool: BaseTool) => async (args: any, options?: ToolExecutionOptions) => { + ): (tool: BaseTool) => (args: any, options?: ToolExecutionOptions) => ToolExecutionResult { + 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,116 @@ 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 { + 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; + } finally { + oc.traceContext.endChildSpan(toolSpan, "completed", {}); + } + }.call(this); + } + + return oc.traceContext.withSpan(toolSpan, async () => { try { // Call tool start hook - can throw ToolDeniedError await hooks.onToolStart?.({ @@ -3637,56 +3761,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", {}); diff --git a/packages/core/src/agent/test-utils.ts b/packages/core/src/agent/test-utils.ts index 39b7bda49..2fc0cabd9 100644 --- a/packages/core/src/agent/test-utils.ts +++ b/packages/core/src/agent/test-utils.ts @@ -362,7 +362,7 @@ export function createTestAgent(options?: Partial): Agent { */ export function createMockTool( name: string, - execute?: (params: any, options?: any) => Promise | any, + execute?: (params: any, options?: any) => PromiseLike | AsyncIterable | any, options?: { description?: string; parameters?: ToolSchema; diff --git a/packages/core/src/tool/index.ts b/packages/core/src/tool/index.ts index a5c697e77..e5acb93f5 100644 --- a/packages/core/src/tool/index.ts +++ b/packages/core/src/tool/index.ts @@ -9,6 +9,8 @@ import { LoggerProxy } from "../logger"; */ type JSONValue = string | number | boolean | null | { [key: string]: JSONValue } | Array; +export type ToolExecutionResult = PromiseLike | AsyncIterable | 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, options?: ToolExecuteOptions, - ) => Promise : unknown>; + ) => ToolExecutionResult : unknown>; }; /** @@ -205,11 +208,12 @@ export class Tool, options?: ToolExecuteOptions, - ) => Promise : unknown>; + ) => ToolExecutionResult : unknown>; /** * Whether this tool should be executed on the client side. diff --git a/packages/core/src/tool/manager/ToolManager.ts b/packages/core/src/tool/manager/ToolManager.ts index c02b0b7b6..ce74d37bf 100644 --- a/packages/core/src/tool/manager/ToolManager.ts +++ b/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 (args: any, options?: ToolExecutionOptions) => Promise, + ) => (args: any, options?: ToolExecutionOptions) => ToolExecutionResult, ): Record { type ManagedTool = { description: string; inputSchema: AgentTool["parameters"]; - execute?: (args: any, options?: ToolExecutionOptions) => Promise; + execute?: (args: any, options?: ToolExecutionOptions) => ToolExecutionResult; needsApproval?: AgentTool["needsApproval"]; providerOptions?: AgentTool["providerOptions"]; toModelOutput?: AgentTool["toModelOutput"]; diff --git a/packages/core/src/tool/manager/index.spec.ts b/packages/core/src/tool/manager/index.spec.ts index 08fb07023..73e37c25c 100644 --- a/packages/core/src/tool/manager/index.spec.ts +++ b/packages/core/src/tool/manager/index.spec.ts @@ -1,6 +1,13 @@ import { z } from "zod"; import type { ToolExecuteOptions } from "../../agent/providers/base/types"; -import { type AgentTool, type ProviderTool, Tool, ToolManager, createTool } from "../index"; +import { + type AgentTool, + type ProviderTool, + Tool, + type ToolExecutionResult, + ToolManager, + createTool, +} from "../index"; import { createToolkit } from "../toolkit"; import type { Toolkit } from "../toolkit"; @@ -234,7 +241,7 @@ describe("ToolManager", () => { const serverToolEntry = preparedTools[mockTool1.name] as { description: string; inputSchema: unknown; - execute?: (args: any, options?: ToolExecuteOptions) => Promise; + execute?: (args: any, options?: ToolExecuteOptions) => ToolExecutionResult; }; expect(createToolExecuteFunction).toHaveBeenCalledTimes(1); @@ -262,7 +269,7 @@ describe("ToolManager", () => { const clientToolEntry = preparedTools[clientTool.name] as { description: string; inputSchema: unknown; - execute?: (args: any, options?: ToolExecuteOptions) => Promise; + execute?: (args: any, options?: ToolExecuteOptions) => ToolExecutionResult; }; expect(createToolExecuteFunction).not.toHaveBeenCalled(); diff --git a/website/docs/agents/tools.md b/website/docs/agents/tools.md index cd37750f0..24b30b9a3 100644 --- a/website/docs/agents/tools.md +++ b/website/docs/agents/tools.md @@ -46,6 +46,58 @@ Each tool has: The `execute` function's parameter types are automatically inferred from the Zod schema, providing full IntelliSense support. +## Streaming Tool Results (Preliminary) + +If your tool can provide progress or intermediate status, return an `AsyncIterable` from `execute`. +Each `yield` is emitted as a preliminary tool result, and the **last yielded value** is treated as the final result. + +```ts +import { createTool } from "@voltagent/core"; +import { z } from "zod"; + +const weatherOutputSchema = z.discriminatedUnion("status", [ + z.object({ + status: z.literal("loading"), + text: z.string(), + weather: z.undefined().optional(), + }), + z.object({ + status: z.literal("success"), + text: z.string(), + weather: z.object({ + location: z.string(), + temperature: z.number(), + }), + }), +]); + +const weatherTool = createTool({ + name: "get_weather", + description: "Get the current weather for a location", + parameters: z.object({ + location: z.string().describe("The city name"), + }), + outputSchema: weatherOutputSchema, + async *execute({ location }) { + yield { + status: "loading" as const, + text: `Getting weather for ${location}`, + weather: undefined, + }; + + await new Promise((resolve) => setTimeout(resolve, 3000)); + + const temperature = 72 + Math.floor(Math.random() * 21) - 10; + + yield { + status: "success" as const, + text: `The weather in ${location} is ${temperature}F`, + weather: { location, temperature }, + }; + }, +}); +``` + ### Tool Tags Tags are optional string labels that help organize and categorize tools. From 0319ac7647c20c96e502800bd99e098307ccfc2e Mon Sep 17 00:00:00 2001 From: Omer Aplak Date: Thu, 15 Jan 2026 11:24:19 -0800 Subject: [PATCH 2/2] fix: duplicate endSpan --- packages/core/src/agent/agent.ts | 2 -- 1 file changed, 2 deletions(-) diff --git a/packages/core/src/agent/agent.ts b/packages/core/src/agent/agent.ts index b98b3fed7..b3657d20f 100644 --- a/packages/core/src/agent/agent.ts +++ b/packages/core/src/agent/agent.ts @@ -3740,8 +3740,6 @@ export class Agent { return await handleToolError(e); }); yield errorResult; - } finally { - oc.traceContext.endChildSpan(toolSpan, "completed", {}); } }.call(this); }