fix(ai-sdk): keep output validation enabled for tool procedures (#1791)

dinwwwh · web-flow · commit f36f85e019c3 · 2026-08-05T14:52:44.000+07:00
`createToolFactory` disabled both input and output validation on the assumption that the AI SDK re-validates against the tool schemas. That is only true for input: the AI SDK treats `outputSchema` as type metadata and never validates the value `execute` returns against it (confirmed in [vercel/ai#10222](vercel/ai#10222), docs corrected in [vercel/ai#11016](vercel/ai#11016)). Procedures run as tools therefore skipped their `.output()` schemas entirely, including transforms and defaults. ## Fixes - Output validation stays enabled; only input validation remains disabled, since that half is genuinely redundant. - A handler returning invalid output now rejects with an output validation error instead of passing through silently. - Streamed `asyncIteratorObject` outputs validate every yielded event, and a handler returning a non-iterator now errors instead of being yielded once as the final result (the old fallback branch became unreachable and is removed). ## Testing - A new `generateText` test with a mock model proves the AI SDK returns schema-violating `execute` output untouched, so it will start failing if a future `ai` release adds output validation and makes ours redundant. - Remaining tests cover input validation being skipped, invalid output rejection, per-event stream validation, and non-iterator rejection.
diff --git a/packages/ai-sdk/src/tool.test.ts b/packages/ai-sdk/src/tool.test.ts
@@ -1,5 +1,7 @@
 import { asyncIteratorObject, oc, type } from '@orpc/contract'
 import { os } from '@orpc/server'
+import { generateText } from 'ai'
+import { MockLanguageModelV4 } from 'ai/test'
 import z from 'zod'
 import { createToolFactory, implementToolFactory } from './tool'
 import { aiSdkTool } from './tool-meta'
@@ -201,6 +203,39 @@ describe('implementToolFactory', () => {
     })
   })
 
+  it('the AI SDK does not validate execute results against outputSchema, so oRPC must validate output itself', async () => {
+    const contract = oc.input(inputSchema).output(outputSchema)
+
+    const greet = implementToolFactory()(contract, {
+      execute: async () => ({ greeting: 123 }) as any,
+    })
+
+    const result = await generateText({
+      model: new MockLanguageModelV4({
+        doGenerate: async () => ({
+          content: [{
+            type: 'tool-call',
+            toolCallId: 'call-1',
+            toolName: 'greet',
+            input: JSON.stringify({ name: 'Alice' }),
+          }],
+          finishReason: { unified: 'tool-calls', raw: undefined },
+          usage: {
+            inputTokens: { total: 10, noCache: 10, cacheRead: undefined, cacheWrite: undefined },
+            outputTokens: { total: 20, text: 20, reasoning: undefined },
+          },
+          warnings: [],
+        }),
+      }),
+      tools: { greet },
+      prompt: 'Greet Alice',
+    })
+
+    expect(result.toolResults).toEqual([
+      expect.objectContaining({ output: { greeting: 123 } }),
+    ])
+  })
+
   describe('async iterator output schema', () => {
     const yieldSchema = z.object({ message: z.string() })
     const returnSchema = z.object({ count: z.number() })
@@ -275,15 +310,30 @@ describe('createToolFactory', () => {
     expect(tool.metadata).toEqual({ source: 'weather-service' })
   })
 
-  it('disable validation at oRPC level to avoid twice times validation', async () => {
+  it('disable input validation at oRPC level to avoid validating twice', async () => {
+    const handler = vi.fn(() => ({ greeting: 'Hello!' }))
+
+    const procedure = os
+      .input(inputSchema)
+      .output(outputSchema)
+      .handler(handler)
+
+    const tool = createToolFactory()(procedure)
+
+    await expect(tool.execute?.('invalid' as any, { abortSignal } as any)).resolves.toEqual({ greeting: 'Hello!' })
+
+    expect(handler).toHaveBeenCalledWith(expect.objectContaining({ input: 'invalid' }), 'invalid')
+  })
+
+  it('keeps output validation enabled because the AI SDK does not validate execute results', async () => {
     const procedure = os
       .input(inputSchema)
       .output(outputSchema)
-      .handler(({ input }) => input as any)
+      .handler(() => ({ greeting: 123 }) as any)
 
     const tool = createToolFactory()(procedure)
 
-    await expect(tool.execute?.('invalid' as any, { abortSignal } as any)).resolves.toEqual('invalid')
+    await expect(tool.execute?.({ name: 'Alice' }, { abortSignal } as any)).rejects.toThrow('Output validation failed')
   })
 
   describe('async iterator output', () => {
@@ -335,20 +385,32 @@ describe('createToolFactory', () => {
       expect(finallyCalled).toBe(true)
     })
 
-    it('yields non-iterator output once when handler ignores the declared iterator schema', async () => {
+    it('rejects when handler ignores the declared iterator schema', async () => {
       const procedure = os
         .input(inputSchema)
         .output(asyncIteratorObject(yieldSchema))
         .handler(async () => ({ message: 'not an iterator' }) as any)
 
       const tool = createToolFactory()(procedure)
 
-      const outputs: unknown[] = []
-      for await (const output of (tool as any).execute({ name: 'Alice' }, { abortSignal })) {
-        outputs.push(output)
-      }
+      const iterator = (tool as any).execute({ name: 'Alice' }, { abortSignal })
+      await expect(iterator.next()).rejects.toThrow('Output validation failed')
+    })
+
+    it('validates each streamed event against the yield schema', async () => {
+      const procedure = os
+        .input(inputSchema)
+        .output(asyncIteratorObject(yieldSchema))
+        .handler(async function* () {
+          yield { message: 'one' }
+          yield { message: 123 } as any
+        })
 
-      expect(outputs).toEqual([{ message: 'not an iterator' }])
+      const tool = createToolFactory()(procedure)
+
+      const iterator = (tool as any).execute({ name: 'Alice' }, { abortSignal })
+      await expect(iterator.next()).resolves.toEqual({ done: false, value: { message: 'one' } })
+      await expect(iterator.next()).rejects.toThrow('AsyncIteratorObject validation failed')
     })
   })
 })
diff --git 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 { isAsyncIteratorObject, ORPC_NAME, resolveMaybeOptionalOptions, toArray } from '@orpc/shared'
+import { ORPC_NAME, resolveMaybeOptionalOptions, toArray } from '@orpc/shared'
 import { tool } from 'ai'
 import { getAiSdkToolMeta } from './tool-meta'
 
@@ -262,12 +262,13 @@ export function createToolFactory<TInitialContext extends Context = object>(
 
     /**
      * The AI SDK already validates input against the tool's `inputSchema`,
-     * so validation is disabled at the oRPC level to avoid validating twice.
+     * so input validation is disabled at the oRPC level to avoid validating twice.
+     * Output validation stays enabled because the AI SDK does not validate
+     * the value returned from `execute` against the tool's `outputSchema`.
      */
     const disabledValidation = new Procedure({
       ...procedure['~orpc'],
       disableInputValidation: true,
-      disableOutputValidation: true,
     })
 
     const isIteratorOutput = getIteratorYieldSchemas(toArray(procedure['~orpc'].outputSchemas)) !== undefined
@@ -281,14 +282,7 @@ export function createToolFactory<TInitialContext extends Context = object>(
        */
       execute: isIteratorOutput
         ? async function* (input, callingOptions) {
-          const output = await call(disabledValidation, input as any, { signal: callingOptions.abortSignal, ...options })
-
-          if (!isAsyncIteratorObject(output)) {
-            yield output
-            return
-          }
-
-          yield* output
+          yield* await call(disabledValidation, input as any, { signal: callingOptions.abortSignal, ...options }) as AsyncIterable<any>
         }
         : (input, callingOptions) => {
             return call(disabledValidation, input as any, { signal: callingOptions.abortSignal, ...options })
