feat(ai-sdk): stream tool outputs for async generator handlers

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'

packages/ai-sdk/src/tool.test-d.ts

@@ -158,4 +158,17 @@ describe('createToolFactory', () => {
 
     expectTypeOf<InferToolOutput<typeof tool>>().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<InferToolOutput<typeof tool>>().toEqualTypeOf<{ status: string }>()
+  })
 })

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<AsyncIteratorObject<{ message: string }>>())
+        .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')
+    })
   })
 })

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<TInitialContext extends Context = object>(
       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.
        */

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)
+  })
+})

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<T>(fn: () => T): () => T {
   let cached: { result: T } | undefined