From a5b928a4c0daba48cf7fe9cb3a2a7d4ba93fb65b Mon Sep 17 00:00:00 2001 From: Dinh Le Date: Wed, 5 Aug 2026 14:43:21 +0700 Subject: [PATCH] fix(ai-sdk): keep output validation enabled for tool procedures --- packages/ai-sdk/src/tool.test.ts | 80 ++++++++++++++++++++++++++++---- packages/ai-sdk/src/tool.ts | 16 ++----- 2 files changed, 76 insertions(+), 20 deletions(-) diff --git a/packages/ai-sdk/src/tool.test.ts b/packages/ai-sdk/src/tool.test.ts index ca56c325b..5e08c92e5 100644 --- 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,7 +385,7 @@ 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)) @@ -343,12 +393,24 @@ describe('createToolFactory', () => { 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 index c2e8bbfd1..e254db36c 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 { 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( /** * 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( */ 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 } : (input, callingOptions) => { return call(disabledValidation, input as any, { signal: callingOptions.abortSignal, ...options })