From 537e5ea9d8d5220d4af56eecd040b86c47179abb Mon Sep 17 00:00:00 2001 From: Dinh Le Date: Fri, 4 Sep 2026 10:57:37 +0700 Subject: [PATCH] feat(ai-sdk): stream tool outputs for async generator handlers Add isAsyncGeneratorFunction to @orpc/shared and use it in createToolFactory as a fallback when the output schema is not an asyncIteratorObject, so async function* handlers stream preliminary results even without a schema. --- apps/content/docs/integrations/ai-sdk.mdx | 2 +- packages/ai-sdk/src/tool.test-d.ts | 13 ++++++ packages/ai-sdk/src/tool.test.ts | 54 +++++++++++++++++++++++ packages/ai-sdk/src/tool.ts | 9 +++- packages/shared/src/function.test.ts | 24 +++++++++- packages/shared/src/function.ts | 9 ++++ 6 files changed, 107 insertions(+), 4 deletions(-) diff --git a/apps/content/docs/integrations/ai-sdk.mdx b/apps/content/docs/integrations/ai-sdk.mdx index 7f6214efe..a4c61ea07 100644 --- a/apps/content/docs/integrations/ai-sdk.mdx +++ b/apps/content/docs/integrations/ai-sdk.mdx @@ -185,7 +185,7 @@ const getWeatherTool = createTool(getWeatherProcedure, { ### Streaming Tool Outputs -When a procedure outputs an [AsyncIteratorObject](/docs/async-iterator-object) validated with `asyncIteratorObject`, the resulting tool streams every event as a [preliminary tool result](https://ai-sdk.dev/docs/ai-sdk-core/tools-and-tool-calling#preliminary-tool-results): each event replaces the tool output in the UI, and the last event becomes the final tool result sent to the model. +When a procedure outputs an [AsyncIteratorObject](/docs/async-iterator-object), either validated with `asyncIteratorObject` or produced by an `async function*` handler, the resulting tool streams every event as a [preliminary tool result](https://ai-sdk.dev/docs/ai-sdk-core/tools-and-tool-calling#preliminary-tool-results): each event replaces the tool output in the UI, and the last event becomes the final tool result sent to the model. ```ts import { asyncIteratorObject, os } from '@orpc/server' diff --git a/packages/ai-sdk/src/tool.test-d.ts b/packages/ai-sdk/src/tool.test-d.ts index 4511888dd..507c37001 100644 --- a/packages/ai-sdk/src/tool.test-d.ts +++ b/packages/ai-sdk/src/tool.test-d.ts @@ -158,4 +158,17 @@ describe('createToolFactory', () => { expectTypeOf>().toEqualTypeOf<{ status: string }>() }) + + it('infer output as yield type for async generator handlers without an output schema', () => { + const procedure = os + .input(z.object({ location: z.string() })) + .handler(async function* () { + yield { status: 'building' } + return { url: 'https://example.com' } + }) + + const tool = createToolFactory()(procedure) + + expectTypeOf>().toEqualTypeOf<{ status: string }>() + }) }) diff --git a/packages/ai-sdk/src/tool.test.ts b/packages/ai-sdk/src/tool.test.ts index c6d7d0b56..75bb1e91e 100644 --- a/packages/ai-sdk/src/tool.test.ts +++ b/packages/ai-sdk/src/tool.test.ts @@ -448,5 +448,59 @@ describe('createToolFactory', () => { await expect(iterator.next()).resolves.toEqual({ done: false, value: { message: 'one' } }) await expect(iterator.next()).rejects.toThrow('AsyncIteratorObject validation failed') }) + + it('streams events when the handler is an async generator without an output schema', async () => { + const procedure = os + .input(inputSchema) + .handler(async function* ({ input }) { + yield { message: `one ${input.name}` } + yield { message: `two ${input.name}` } + return { count: 2 } + }) + + const tool = createToolFactory()(procedure) + + expect(tool.outputSchema).toBeUndefined() + + const outputs: unknown[] = [] + for await (const output of (tool as any).execute({ name: 'Alice' }, { abortSignal })) { + outputs.push(output) + } + + expect(outputs).toEqual([{ message: 'one Alice' }, { message: 'two Alice' }]) + }) + + it('streams events when the handler is an async generator and the output schema is not asyncIteratorObject', async () => { + const procedure = os + .input(inputSchema) + .output(type>()) + .handler(async function* () { + yield { message: 'one' } + yield { message: 'two' } + }) + + const tool = createToolFactory()(procedure) + + const outputs: unknown[] = [] + for await (const output of (tool as any).execute({ name: 'Alice' }, { abortSignal })) { + outputs.push(output) + } + + expect(outputs).toEqual([{ message: 'one' }, { message: 'two' }]) + }) + + it('does not stream when a non-generator handler returns an async iterator without an asyncIteratorObject schema', async () => { + const procedure = os + .input(inputSchema) + .handler(async () => (async function* () { + yield { message: 'one' } + })()) + + const tool = createToolFactory()(procedure) + + const output = await (tool as any).execute({ name: 'Alice' }, { abortSignal }) + + expect(output[Symbol.asyncIterator]).toBeTypeOf('function') + }) }) }) diff --git a/packages/ai-sdk/src/tool.ts b/packages/ai-sdk/src/tool.ts index 2053de9d2..c405fbd73 100644 --- a/packages/ai-sdk/src/tool.ts +++ b/packages/ai-sdk/src/tool.ts @@ -8,7 +8,7 @@ import type { FunctionTool } from './tool-meta' import { getAsyncIteratorObjectSchemaDetails } from '@orpc/contract' import { combineJsonSchemasWithComposition } from '@orpc/json-schema' import { call, Procedure } from '@orpc/server' -import { isPlainObject, mergeTwoLevels, ORPC_NAME, resolveMaybeOptionalOptions, toArray } from '@orpc/shared' +import { isAsyncGeneratorFunction, isPlainObject, mergeTwoLevels, ORPC_NAME, resolveMaybeOptionalOptions, toArray } from '@orpc/shared' import { tool } from 'ai' import { getAiSdkToolMeta } from './tool-meta' @@ -275,12 +275,17 @@ export function createToolFactory( disableInputValidation: true, }) + /** + * Output schemas are the source of truth, but an `async function*` handler always + * returns an async iterator, so it streams even without an `asyncIteratorObject` schema. + */ const isIteratorOutput = getIteratorYieldSchemas(toArray(procedure['~orpc'].outputSchemas)) !== undefined + || isAsyncGeneratorFunction(procedure['~orpc'].handler) return implementTool(procedure, { ...toolOptions as any, /** - * For `asyncIteratorObject` outputs, the tool streams each event as a + * For async iterator outputs, the tool streams each event as a * [preliminary result](https://ai-sdk.dev/docs/ai-sdk-core/tools-and-tool-calling#preliminary-tool-results), * and the last event becomes the final result. */ diff --git a/packages/shared/src/function.test.ts b/packages/shared/src/function.test.ts index ec2f34378..4bb9ef29b 100644 --- a/packages/shared/src/function.test.ts +++ b/packages/shared/src/function.test.ts @@ -1,4 +1,4 @@ -import { defer, once, tryOrUndefined } from './function' +import { defer, isAsyncGeneratorFunction, once, tryOrUndefined } from './function' it('once', () => { const fn = vi.fn(() => ({})) @@ -66,3 +66,25 @@ describe('tryOrUndefined', () => { expect(tryOrUndefined(() => null)).toBeNull() }) }) + +describe('isAsyncGeneratorFunction', () => { + it('returns true for async generator functions', () => { + async function* gen() {} + + expect(isAsyncGeneratorFunction(gen)).toBe(true) + expect(isAsyncGeneratorFunction(gen.bind(null))).toBe(true) + expect(isAsyncGeneratorFunction(async function* () {})).toBe(true) + expect(isAsyncGeneratorFunction({ async* method() {} }.method)).toBe(true) + }) + + it('returns false for anything else', () => { + expect(isAsyncGeneratorFunction(function* () {})).toBe(false) + expect(isAsyncGeneratorFunction(async () => {})).toBe(false) + expect(isAsyncGeneratorFunction(() => {})).toBe(false) + expect(isAsyncGeneratorFunction(class {})).toBe(false) + expect(isAsyncGeneratorFunction((async function* () {})())).toBe(false) + expect(isAsyncGeneratorFunction({})).toBe(false) + expect(isAsyncGeneratorFunction(null)).toBe(false) + expect(isAsyncGeneratorFunction(undefined)).toBe(false) + }) +}) diff --git a/packages/shared/src/function.ts b/packages/shared/src/function.ts index 9d37db83e..841a13c28 100644 --- a/packages/shared/src/function.ts +++ b/packages/shared/src/function.ts @@ -1,5 +1,14 @@ export type AnyFunction = (...args: any[]) => any +const AsyncGeneratorFunction = Object.getPrototypeOf(async function* () {}).constructor + +/** + * Checks whether a value is an async generator function (`async function*`). + */ +export function isAsyncGeneratorFunction(value: unknown): value is (...args: any[]) => AsyncGenerator { + return value instanceof AsyncGeneratorFunction +} + export function once(fn: () => T): () => T { let cached: { result: T } | undefined