From e3de43ef87eb7d841e88067d52439b7b00a0059f Mon Sep 17 00:00:00 2001 From: Alem Tuzlak Date: Tue, 13 Oct 2026 20:27:26 +0200 Subject: [PATCH 01/16] feat(ai): check tool input at the last moment, after approval edits and middleware --- .changeset/tool-input-final-check.md | 21 + docs/advanced/middleware.md | 20 +- docs/tools/server-tools.md | 6 + packages/ai/src/activities/chat/index.ts | 88 +-- .../src/activities/chat/middleware/types.ts | 4 +- .../activities/chat/tools/input-validation.ts | 33 ++ .../src/activities/chat/tools/tool-calls.ts | 257 ++++---- packages/ai/src/middlewares/otel.ts | 3 + .../ai/tests/chat-client-tool-wait.test.ts | 6 +- .../ai/tests/tool-input-validation.test.ts | 552 ++++++++++++++++++ testing/e2e/src/routes/api.middleware-test.ts | 12 + testing/e2e/src/routes/middleware-test.tsx | 1 + testing/e2e/tests/middleware.spec.ts | 35 ++ 13 files changed, 846 insertions(+), 192 deletions(-) create mode 100644 .changeset/tool-input-final-check.md create mode 100644 packages/ai/src/activities/chat/tools/input-validation.ts create mode 100644 packages/ai/tests/tool-input-validation.test.ts diff --git a/.changeset/tool-input-final-check.md b/.changeset/tool-input-final-check.md new file mode 100644 index 0000000000..e5ff50ec7f --- /dev/null +++ b/.changeset/tool-input-final-check.md @@ -0,0 +1,21 @@ +--- +'@tanstack/ai': minor +--- + +Check a tool input at the last moment, just before the tool runs or the client gets the call. + +Behavior change: + +- The tool's Standard Schema `inputSchema` now checks the input after approval `editedArgs` and after middleware `onBeforeToolCall`. Before, edits and `transformArgs` values reached `execute` or the client without a check. An input that does not fit the schema now becomes a tool error (state `output-error`), and the tool does not run. +- `onBeforeToolCall` gets the parsed input before the schema check. It has no schema defaults or transforms yet, and it can be invalid. +- Async Standard Schema input validators are awaited. Before, they threw. +- The tool gets the schema's output, so a transform applies once. +- `onBeforeToolCall` now also runs for a client tool before the server sends the call to the client. This includes a client request at an interrupt boundary. It does not run while a call waits for approval, for a denied call, or when the client result comes back. `otelMiddleware` records no `execute_tool` span for a client tool, because the server does not run it. +- A validation error after the hook runs `onAfterToolCall` with `ok: false`. +- The error text now also shows the received arguments. + +Migration: + +- If a middleware `transformArgs` or an approval edit sends a value that the schema rejects, make the value fit the schema, or widen the schema. +- If `onBeforeToolCall` reads `hookCtx.args`, check its shape first. Do not expect schema defaults there. +- If a middleware counts or blocks calls in `onBeforeToolCall` and must see only server tools, return early when `hookCtx.tool?.execute` is missing. diff --git a/docs/advanced/middleware.md b/docs/advanced/middleware.md index e0800eb5e5..1371fa62e8 100644 --- a/docs/advanced/middleware.md +++ b/docs/advanced/middleware.md @@ -668,7 +668,25 @@ system prompt, is in [Apply Answers](../interrupts/apply-answers). ### onBeforeToolCall -Called before each tool executes. The first middleware that returns a non-void decision short-circuits — remaining middleware are skipped for that tool call. +Called before a server tool runs, and before a client tool goes to the client. The first middleware that returns a non-void decision skips the remaining middleware for that call. + +Each tool call goes through these steps: + +1. The engine parses the JSON arguments from the model. +2. `onBeforeToolCall` gets the parsed input. +3. The tool's `inputSchema` checks the final input. An async validator is awaited. A schema transform applies once, here. +4. The server tool runs, or the client gets the call, with the checked input. + +The schema checks the input after your hook: + +- `hookCtx.args` has no schema defaults or transforms yet, and it can be invalid. Check its shape before you read it. +- A `transformArgs` value or an in-place edit must fit the schema. If it does not fit, the call gets a tool error and the tool does not run. + +When the hook runs: + +- A client tool runs the hook before the server sends the call to the client. The hook does not run when the client result comes back. +- A call that waits for approval does not run the hook. A denied call does not run it either. +- After approval, the hook gets the approved input, including any `editedArgs`. See [Tool Approval](../interrupts/tool-approval). ```typescript import { type ChatMiddleware } from "@tanstack/ai"; diff --git a/docs/tools/server-tools.md b/docs/tools/server-tools.md index 12aaf118a7..4a5104ad64 100644 --- a/docs/tools/server-tools.md +++ b/docs/tools/server-tools.md @@ -357,6 +357,12 @@ const getUserData = getUserDataDef.server(async ({ userId }) => { Malformed JSON arguments and Standard Schema input validation failures are also surfaced as tool-result errors. The tool implementation is not called, and the agent loop can return the error to the model so it can repair the tool call. This also applies to Chat Completions tools with no required input fields. +The `inputSchema` checks the input just before the tool runs: + +- The check runs after approval edits and after middleware `onBeforeToolCall`. An edited input that does not fit the schema becomes a tool error. +- An async validator is awaited. +- The tool gets the schema's output, so a transform such as `z.string().transform(Number)` applies once. + ## Using JSON Schema If you have existing JSON Schema definitions or prefer not to use Zod, you can define tool schemas using raw JSON Schema objects: diff --git a/packages/ai/src/activities/chat/index.ts b/packages/ai/src/activities/chat/index.ts index a4155fef00..26afde8a9e 100644 --- a/packages/ai/src/activities/chat/index.ts +++ b/packages/ai/src/activities/chat/index.ts @@ -3928,7 +3928,15 @@ class TextEngine< this.addTerminalAssistantMessages() } } - const actionable = this.getBoundaryActionableToolRequests(toolCalls) + const actionable = yield* this.runWhileYielding( + this.getBoundaryActionableToolRequests(toolCalls), + ) + for (const chunk of this.buildToolResultChunks( + actionable.results, + finishEvent, + )) { + yield* this.pipeThroughMiddleware(chunk) + } yield* this.emitActionableInterruptBoundary( finishEvent, actionable.approvals, @@ -3939,51 +3947,65 @@ class TextEngine< return true } - private getBoundaryActionableToolRequests( + private async getBoundaryActionableToolRequests( toolCalls: ReadonlyArray, - ): { + ): Promise<{ approvals: Array clientRequests: Array - } { + results: Array + }> { const { approvals, clientToolResults } = this.collectClientState() - const approvalRequests: Array = [] - const clientRequests: Array = [] - for (const toolCall of toolCalls) { + const pendingCalls = toolCalls.filter((toolCall) => { // Provider-executed calls are complete; never surface them as client work. - if (isProviderExecutedToolCall(toolCall)) continue + if (isProviderExecutedToolCall(toolCall)) return false const tool = this.resolveExecutableTools([toolCall]).find( (candidate) => candidate.name === toolCall.function.name, - ) as RuntimeToolWithApproval | undefined - if (!tool) continue - let input: unknown = {} - try { - const parsed = JSON.parse(toolCall.function.arguments.trim() || '{}') - input = parsed && typeof parsed === 'object' ? parsed : {} - } catch { - input = {} - } - const approvalId = `approval_${toolCall.id}` - if (tool.needsApproval && !approvals.has(approvalId)) { - approvalRequests.push({ - toolCallId: toolCall.id, - toolName: toolCall.function.name, - input, - approvalId, - }) - } else if ( + ) + if (!tool) return false + const resolution = + approvals.get(toolCall.id) ?? approvals.get(`approval_${toolCall.id}`) + if (tool.needsApproval && resolution === undefined) return true + return ( !tool.execute && !clientToolResults.has(toolCall.id) && !this.resumeClientToolErrors.has(toolCall.id) && !this.resumeCancelledToolCallIds.has(toolCall.id) - ) { - clientRequests.push({ - toolCallId: toolCall.id, - toolName: toolCall.function.name, - input, - }) + ) + }) + // Approval previews and client requests get the same final input check + // as the tool phase. No server tool runs here: a server tool is in the + // list only while it waits for approval. + const generator = executeToolCalls( + pendingCalls, + this.resolveExecutableTools(pendingCalls), + approvals, + clientToolResults, + undefined, + { + onBeforeToolCall: (toolCall, tool, args) => + this.middlewareRunner.runOnBeforeToolCall(this.middlewareCtx, { + toolCall, + tool, + args, + toolName: toolCall.function.name, + toolCallId: toolCall.id, + }), + onAfterToolCall: (info) => + this.middlewareRunner.runOnAfterToolCall(this.middlewareCtx, info), + }, + this.middlewareCtx.context, + this.toolAbortSignal, + ) + while (true) { + const next = await generator.next() + if (next.done) { + return { + approvals: next.value.needsApproval, + clientRequests: next.value.needsClientExecution, + results: next.value.results, + } } } - return { approvals: approvalRequests, clientRequests } } private completeEphemeralInterruptBindings(chunk: StreamChunk): StreamChunk { diff --git a/packages/ai/src/activities/chat/middleware/types.ts b/packages/ai/src/activities/chat/middleware/types.ts index 1b5a603446..8ec7ed505c 100644 --- a/packages/ai/src/activities/chat/middleware/types.ts +++ b/packages/ai/src/activities/chat/middleware/types.ts @@ -841,8 +841,10 @@ export interface ChatMiddleware< | Promise | null> /** - * Called before a tool is executed. + * Called before a server tool runs and before a client tool is sent to the + * client. Not called while the call waits for approval. * Can observe, transform args, skip execution, or abort the run. + * The tool's `inputSchema` checks the args after this hook. */ onBeforeToolCall?: ( ctx: ChatMiddlewareContext, diff --git a/packages/ai/src/activities/chat/tools/input-validation.ts b/packages/ai/src/activities/chat/tools/input-validation.ts new file mode 100644 index 0000000000..64c6c06ebc --- /dev/null +++ b/packages/ai/src/activities/chat/tools/input-validation.ts @@ -0,0 +1,33 @@ +import { + StandardSchemaValidationError, + isStandardSchema, + validateWithStandardSchema, +} from './schema-converter' +import type { SchemaInput } from '../../../types' + +function receivedArgumentsText(received: unknown) { + try { + return JSON.stringify(received, null, 2) ?? String(received) + } catch { + return '[Arguments cannot be serialized as JSON]' + } +} + +/** + * The final check of a tool input, just before the tool runs or the client + * gets the call. A Standard Schema runs here (an async validator too), and + * its output is the input, so transforms apply once. The error goes back to + * the model as the tool result. + */ +export async function validateToolInput( + schema: SchemaInput | undefined, + received: unknown, + toolName: string, +): Promise { + if (!isStandardSchema(schema)) return received + const validation = await validateWithStandardSchema(schema, received) + if (validation.success) return validation.data + const error = new StandardSchemaValidationError(validation.issues) + error.message = `Input validation failed for tool ${toolName}: ${error.message}\nReceived arguments:\n${receivedArgumentsText(received)}` + throw error +} diff --git a/packages/ai/src/activities/chat/tools/tool-calls.ts b/packages/ai/src/activities/chat/tools/tool-calls.ts index 3bc608fee5..359241134d 100644 --- a/packages/ai/src/activities/chat/tools/tool-calls.ts +++ b/packages/ai/src/activities/chat/tools/tool-calls.ts @@ -2,6 +2,7 @@ import { normalizeToolResult } from '../../../utilities/tool-result' import { tanstackMetadata } from '../../../utilities/merge-metadata' import { isProviderExecutedToolCall } from '../../../utilities/provider-executed' import { mergeStreams } from '../../../utilities/merge-streams' +import { validateToolInput } from './input-validation' import type { AdapterYieldChunk } from '../../../utilities/adapter-yield-chunk' import { StandardSchemaValidationError, @@ -373,20 +374,7 @@ export class ToolCallManager< ) } - // Validate input against inputSchema (for Standard Schema compliant schemas) - if (tool.inputSchema && isStandardSchema(tool.inputSchema)) { - try { - args = parseWithStandardSchema(tool.inputSchema, args) - } catch (validationError: unknown) { - const message = - validationError instanceof Error - ? validationError.message - : 'Validation failed' - throw new Error( - `Input validation failed for tool ${tool.name}: ${message}`, - ) - } - } + args = await validateToolInput(tool.inputSchema, args, tool.name) // Execute the tool const executionContext = { @@ -1104,27 +1092,95 @@ export async function* executeToolCalls( } } - // Validate input against inputSchema (for Standard Schema compliant schemas) - if (tool.inputSchema && isStandardSchema(tool.inputSchema)) { + const resolution = tool.needsApproval + ? approvalResolution(approvals, toolCall.id) + : undefined + // A call that waits for approval shows its checked input. Middleware + // runs only once the call is approved. + const needsPreview = tool.needsApproval && resolution === undefined + if (resolution !== undefined && !isApproved(resolution)) { + // User declined + results.push({ + toolCallId: toolCall.id, + toolName, + result: + resumeState?.deniedToolResults?.get(toolCall.id) ?? + deniedApprovalResult(resolution), + input, + state: 'output-error', + outcome: 'denied', + }) + continue + } + if (resolution !== undefined) { + input = editedApprovalArgs(resolution) ?? input + } + // A client tool whose result came back already went through the hook + // before dispatch. Its input is checked again, with no second hook run. + const clientError = resumeState?.clientToolErrors?.get(toolCall.id) + const clientDone = + !tool.execute && + (clientResults.has(toolCall.id) || clientError !== undefined) + + if (!needsPreview && !clientDone && middlewareHooks) { + let decision: Awaited> try { - input = parseWithStandardSchema(tool.inputSchema, input) - } catch (validationError: unknown) { - const message = - validationError instanceof Error - ? validationError.message - : 'Validation failed' - results.push({ + decision = await applyBeforeToolCallDecision( + toolCall, + tool, + input, + toolName, + middlewareHooks, + results, + ) + } catch (error) { + // The calls prepared so far still run, as they would one at a time. + failures.push(error) + break + } + if (!decision.proceed) continue + input = decision.input + } + + // The final check, after approval edits and middleware: invalid input + // becomes a tool error, and the checked value is what runs. + try { + input = await validateToolInput(tool.inputSchema, input, toolName) + } catch (error) { + const message = + error instanceof Error ? error.message : 'Validation failed' + await pushResultAndRunAfterHook( + results, + { toolCallId: toolCall.id, toolName, - result: { - error: `Input validation failed for tool ${tool.name}: ${message}`, - }, - // raw parse may have failed validation — still attach best-effort input + result: { error: message }, input, state: 'output-error', - }) - continue - } + duration: 0, + }, + { + toolCall, + tool, + toolName, + toolCallId: toolCall.id, + ok: false, + duration: 0, + error, + }, + // A preview call and a returned client call run no hooks. + needsPreview || clientDone ? undefined : middlewareHooks, + ) + continue + } + if (needsPreview) { + needsApproval.push({ + toolCallId: toolCall.id, + toolName, + input, + approvalId: `approval_${toolCall.id}`, + }) + continue } // Create a ToolExecutionContext for this tool call with event emission @@ -1155,142 +1211,31 @@ export async function* executeToolCalls( }, } as ToolExecutionContext - // CASE 1: Client-side tool (no execute function) + // Client-side tool (no execute function), approved or with no approval if (!tool.execute) { - // Check if tool needs approval - if (tool.needsApproval) { - const approvalId = `approval_${toolCall.id}` - const resolution = approvalResolution(approvals, toolCall.id) - - // Check if approval decision exists - if (resolution !== undefined) { - const approved = isApproved(resolution) - - if (approved) { - input = editedApprovalArgs(resolution) ?? input - // Approved - check if client has executed - const clientError = resumeState?.clientToolErrors?.get(toolCall.id) - if (clientResults.has(toolCall.id) || clientError !== undefined) { - results.push( - await buildClientToolResult( - toolCall.id, - toolName, - tool, - clientResults.get(toolCall.id), - input, - clientError, - ), - ) - } else { - // Approved but not executed yet - request client execution - needsClientExecution.push({ - toolCallId: toolCall.id, - toolName, - input, - }) - } - } else { - // User declined - results.push({ - toolCallId: toolCall.id, - toolName, - result: - resumeState?.deniedToolResults?.get(toolCall.id) ?? - deniedApprovalResult(resolution), - input, - state: 'output-error', - outcome: 'denied', - }) - } - } else { - // Need approval first - needsApproval.push({ - toolCallId: toolCall.id, - toolName: toolCall.function.name, - input, - approvalId, - }) - } - } else { - // No approval needed - check if client has executed - const clientError = resumeState?.clientToolErrors?.get(toolCall.id) - if (clientResults.has(toolCall.id) || clientError !== undefined) { - results.push( - await buildClientToolResult( - toolCall.id, - toolName, - tool, - clientResults.get(toolCall.id), - input, - clientError, - ), - ) - } else { - // Request client execution - needsClientExecution.push({ - toolCallId: toolCall.id, + if (clientDone) { + results.push( + await buildClientToolResult( + toolCall.id, toolName, + tool, + clientResults.get(toolCall.id), input, - }) - } - } - continue - } - - // CASE 2: Server tool with approval required - if (tool.needsApproval) { - const approvalId = `approval_${toolCall.id}` - const resolution = approvalResolution(approvals, toolCall.id) - - if (resolution === undefined) { - // Need approval - needsApproval.push({ - toolCallId: toolCall.id, - toolName, - input, - approvalId, - }) - continue - } - if (!isApproved(resolution)) { - // User declined - results.push({ + clientError, + ), + ) + } else { + // Request client execution + needsClientExecution.push({ toolCallId: toolCall.id, toolName, - result: - resumeState?.deniedToolResults?.get(toolCall.id) ?? - deniedApprovalResult(resolution), input, - state: 'output-error', - outcome: 'denied', }) - continue } - // Approved: run it like any other server tool below. - input = editedApprovalArgs(resolution) ?? input - } - - // CASE 3: Server tool, approved or with no approval - if (middlewareHooks) { - let decision: Awaited> - try { - decision = await applyBeforeToolCallDecision( - toolCall, - tool, - input, - toolName, - middlewareHooks, - results, - ) - } catch (error) { - // The calls prepared so far still run, as they would one at a time. - failures.push(error) - break - } - if (!decision.proceed) continue - input = decision.input + continue } + // Server tool, approved or with no approval const run = runServerTool( toolCall, tool, diff --git a/packages/ai/src/middlewares/otel.ts b/packages/ai/src/middlewares/otel.ts index cef80e2e8d..00994530f7 100644 --- a/packages/ai/src/middlewares/otel.ts +++ b/packages/ai/src/middlewares/otel.ts @@ -899,6 +899,9 @@ export function otelMiddleware( onBeforeToolCall(ctx, hookCtx) { safeCall('otel.onBeforeToolCall', () => { + // A client tool runs in the client, not here, so it gets no + // execute_tool span. Middleware sees its input before dispatch. + if (hookCtx.tool && !hookCtx.tool.execute) return const state = stateByCtx.get(ctx) if (!state) return const parent = state.currentIterationSpan ?? state.rootSpan diff --git a/packages/ai/tests/chat-client-tool-wait.test.ts b/packages/ai/tests/chat-client-tool-wait.test.ts index a49217ae20..24ee54158d 100644 --- a/packages/ai/tests/chat-client-tool-wait.test.ts +++ b/packages/ai/tests/chat-client-tool-wait.test.ts @@ -214,7 +214,11 @@ describe('client-tool wait lifecycle', () => { }), ) - expect(fakeTracer.spans).toHaveLength(2) + // The client runs the tool, so the server records no execute_tool span. + expect(fakeTracer.spans.map((span) => span.name)).toEqual([ + 'chat test-model', + 'chat test-model #0', + ]) expect(fakeTracer.spans.every((span) => span.ended)).toBe(true) expect( fakeMeter.records.filter( diff --git a/packages/ai/tests/tool-input-validation.test.ts b/packages/ai/tests/tool-input-validation.test.ts new file mode 100644 index 0000000000..1e8106f45b --- /dev/null +++ b/packages/ai/tests/tool-input-validation.test.ts @@ -0,0 +1,552 @@ +import { describe, expect, it } from 'vitest' +import { z } from 'zod' +import { chat } from '../src/activities/chat/index' +import { defineChatMiddleware } from '../src/activities/chat/middleware/define' +import { defineInterrupt } from '../src/interrupt-definition' +import { executeToolCalls } from '../src/activities/chat/tools/tool-calls' +import { otelMiddleware } from '../src/middlewares/otel' +import { EventType } from '../src/types' +import { collectChunks, createMockAdapter, ev } from './test-utils' +import { createFakeMeter, createFakeTracer } from './middlewares/fake-otel' +import type { StandardSchemaV1 } from '@standard-schema/spec' +import type { ChatMiddleware } from '../src/activities/chat/middleware/types' +import type { ToolExecutionMiddlewareHooks } from '../src/activities/chat/tools/tool-calls' +import type { ToolApprovalResolution } from '../src/interrupts' +import type { + ModelMessage, + SchemaInput, + StreamChunk, + Tool, + ToolCall, +} from '../src/types' + +const schema = z.object({ count: z.number().int().min(1) }) + +async function runTool( + rawArguments: string, + inputSchema: SchemaInput, + options: { + client?: boolean + needsApproval?: boolean + approvals?: Map + middleware?: ToolExecutionMiddlewareHooks + clientResults?: Map + } = {}, +) { + const executed: Array = [] + const tool: Tool = { + name: 'check', + description: 'Check input', + inputSchema, + needsApproval: options.needsApproval, + ...(options.client + ? {} + : { + execute: (input: unknown) => { + executed.push(input) + return input + }, + }), + } + const call: ToolCall = { + id: 'call-1', + type: 'function', + function: { name: 'check', arguments: rawArguments }, + } + const generator = executeToolCalls( + [call], + [tool], + options.approvals, + options.clientResults, + undefined, + options.middleware, + ) + while (true) { + const next = await generator.next() + if (next.done) return { ...next.value, executed } + } +} + +describe('final tool input validation', () => { + it('awaits an async Standard Schema validator', async () => { + const asyncSchema: StandardSchemaV1 = { + '~standard': { + version: 1, + vendor: 'test', + validate: async (value) => + typeof value === 'object' && + value !== null && + 'count' in value && + typeof value.count === 'number' + ? { value: { count: value.count * 10 } } + : { issues: [{ message: 'count must be a number' }] }, + }, + } + const ok = await runTool('{"count":2}', asyncSchema) + expect(ok.executed).toEqual([{ count: 20 }]) + const bad = await runTool('{"count":"2"}', asyncSchema) + expect(bad.executed).toEqual([]) + expect(bad.results[0]?.state).toBe('output-error') + expect(bad.results[0]?.result).toEqual({ + error: expect.stringContaining('count must be a number'), + }) + }) + + it('includes the received arguments in the error', async () => { + const result = await runTool('{"count":0}', schema) + expect(result.results[0]?.result).toEqual({ + error: expect.stringContaining('Input validation failed for tool check'), + }) + expect(result.results[0]?.result).toEqual({ + error: expect.stringContaining('"count": 0'), + }) + expect(result.executed).toEqual([]) + }) + + it('keeps empty and null arguments as {}', async () => { + for (const raw of ['', ' ', 'null']) { + const result = await runTool(raw, z.object({})) + expect(result.executed).toEqual([{}]) + } + }) + + it.each([false, true])( + 'rejects a middleware replacement before server or client execution: client=%s', + async (client) => { + const result = await runTool('{"count":2}', schema, { + client, + middleware: { + onBeforeToolCall: async () => ({ + type: 'transformArgs', + args: { count: 'wrong' }, + }), + }, + }) + expect(result.results[0]?.state).toBe('output-error') + expect(result.executed).toEqual([]) + expect(result.needsClientExecution).toEqual([]) + }, + ) + + it('rejects a middleware edit made in place', async () => { + const result = await runTool('{"count":2}', schema, { + middleware: { + onBeforeToolCall: async (_call, _tool, input) => { + if (typeof input === 'object' && input !== null) + Reflect.set(input, 'count', 'wrong') + }, + }, + }) + expect(result.results[0]?.state).toBe('output-error') + expect(result.executed).toEqual([]) + }) + + it('runs onAfterToolCall with the validation error', async () => { + const after: Array = [] + await runTool('{"count":2}', schema, { + middleware: { + onBeforeToolCall: async () => ({ + type: 'transformArgs', + args: { count: 'wrong' }, + }), + onAfterToolCall: async (info) => { + after.push(info.ok) + }, + }, + }) + expect(after).toEqual([false]) + }) + + it('lets middleware repair invalid input', async () => { + const result = await runTool('{"count":"wrong"}', schema, { + middleware: { + onBeforeToolCall: async () => ({ + type: 'transformArgs', + args: { count: 4 }, + }), + }, + }) + expect(result.executed).toEqual([{ count: 4 }]) + }) + + it('shows middleware the raw input and runs a transform once', async () => { + const seen: Array = [] + let transforms = 0 + const inputSchema = z.object({ + count: z.string().transform((value) => { + transforms++ + return Number(value) + }), + }) + const result = await runTool('{"count":"2"}', inputSchema, { + middleware: { + onBeforeToolCall: async (_call, _tool, input) => { + seen.push(structuredClone(input)) + }, + }, + }) + expect(seen).toEqual([{ count: '2' }]) + expect(result.executed).toEqual([{ count: 2 }]) + expect(transforms).toBe(1) + }) + + it('sends the checked middleware input to the client', async () => { + const result = await runTool('{"count":2}', schema, { + client: true, + middleware: { + onBeforeToolCall: async () => ({ + type: 'transformArgs', + args: { count: 5 }, + }), + }, + }) + expect(result.needsClientExecution).toEqual([ + { toolCallId: 'call-1', toolName: 'check', input: { count: 5 } }, + ]) + }) + + it.each([false, true])( + 'checks an approval edit for a client tool: resumed=%s', + async (resumed) => { + const result = await runTool('{"count":2}', schema, { + client: true, + needsApproval: true, + approvals: new Map([ + ['call-1', { approved: true, editedArgs: { count: 'wrong' } }], + ]), + clientResults: resumed + ? new Map([['call-1', 'client result']]) + : new Map(), + }) + expect(result.results[0]?.state).toBe('output-error') + expect(result.needsClientExecution).toEqual([]) + }, + ) + + it('rejects an invalid approval edit for a server tool', async () => { + const result = await runTool('{"count":2}', schema, { + needsApproval: true, + approvals: new Map([ + ['call-1', { approved: true, editedArgs: { count: 0 } }], + ]), + }) + expect(result.results[0]?.state).toBe('output-error') + expect(result.executed).toEqual([]) + }) + + it.each([false, true])( + 'keeps the transformed approval edit: client=%s', + async (client) => { + const result = await runTool('"2"', z.string().transform(Number), { + client, + needsApproval: true, + approvals: new Map([['call-1', { approved: true, editedArgs: '4' }]]), + }) + if (client) { + expect(result.needsClientExecution).toEqual([ + { toolCallId: 'call-1', toolName: 'check', input: 4 }, + ]) + } else { + expect(result.executed).toEqual([4]) + } + }, + ) + + it.each([false, true])( + 'runs no middleware and dispatches nothing while approval is pending: client=%s', + async (client) => { + let middlewareCalls = 0 + const result = await runTool('{"count":2}', schema, { + client, + needsApproval: true, + middleware: { + onBeforeToolCall: async () => { + middlewareCalls++ + }, + }, + }) + expect(result.needsApproval).toEqual([ + { + toolCallId: 'call-1', + toolName: 'check', + input: { count: 2 }, + approvalId: 'approval_call-1', + }, + ]) + expect(result.needsClientExecution).toEqual([]) + expect(middlewareCalls).toBe(0) + expect(result.executed).toEqual([]) + }, + ) + + it.each([false, true])( + 'runs no middleware for a denied tool: client=%s', + async (client) => { + let middlewareCalls = 0 + const result = await runTool('{"count":2}', schema, { + client, + needsApproval: true, + approvals: new Map([['call-1', false]]), + middleware: { + onBeforeToolCall: async () => { + middlewareCalls++ + }, + }, + }) + expect(result.results[0]?.outcome).toBe('denied') + expect(result.needsClientExecution).toEqual([]) + expect(middlewareCalls).toBe(0) + expect(result.executed).toEqual([]) + }, + ) +}) + +const clientCheck = ( + inputSchema: SchemaInput, + needsApproval = false, +): Tool => ({ + name: 'check', + description: 'Check input', + inputSchema, + needsApproval, +}) + +const history = (input: unknown): Array => [ + { + role: 'assistant', + content: null, + toolCalls: [ + { + id: 'call', + type: 'function', + function: { name: 'check', arguments: JSON.stringify(input) }, + }, + ], + }, +] + +function descriptors(chunks: Array) { + return chunks.flatMap((chunk) => + chunk.type === EventType.RUN_FINISHED && chunk.outcome?.type === 'interrupt' + ? chunk.outcome.interrupts + : [], + ) +} + +function clientInput(chunks: Array) { + return descriptors(chunks).find((entry) => entry.id === 'client_tool_call') + ?.metadata?.input +} + +const review = defineInterrupt({ + id: 'boundary-review', + responseSchema: z.boolean(), +}) + +// Pauses after the model call, so client requests come from the boundary. +const reviewAfterModel = defineChatMiddleware({ + name: 'boundary', + onInterruptBoundary(ctx) { + if (ctx.phase !== 'afterModel') return undefined + return { + interrupts: [ + review.interrupt({ + key: 'review', + reason: 'review', + message: 'Review', + }), + ], + } + }, +}) + +async function initial( + input: unknown, + tool: Tool, + middleware: Array = [], + boundary = false, +) { + const { adapter } = createMockAdapter({ + iterations: [ + [ + ev.runStarted(), + ev.toolStart('call', 'check'), + ev.toolArgs('call', JSON.stringify(input)), + ev.runFinished('tool_calls'), + ], + ], + }) + return collectChunks( + chat({ + adapter, + messages: [{ role: 'user', content: 'check' }], + tools: [tool], + ...(boundary + ? { + interrupts: [review], + middleware: [reviewAfterModel, ...middleware], + } + : { middleware }), + }), + ) +} + +describe('client tool input in chat()', () => { + it.each([false, true])( + 'shows middleware the raw input, then sends the checked input: boundary=%s', + async (boundary) => { + let transforms = 0 + const seen: Array = [] + const inputSchema = z.object({ + count: z.string().transform((value) => { + transforms++ + return Number(value) + }), + }) + const chunks = await initial( + { count: '2' }, + clientCheck(inputSchema), + [ + { + onBeforeToolCall(_ctx, hook) { + seen.push(structuredClone(hook.args)) + }, + }, + ], + boundary, + ) + expect(seen).toEqual([{ count: '2' }]) + expect(transforms).toBe(1) + expect(clientInput(chunks)).toEqual({ count: 2 }) + }, + ) + + it.each([false, true])( + 'sends no client request when middleware makes the input invalid: boundary=%s', + async (boundary) => { + const chunks = await initial( + { count: 2 }, + clientCheck(schema), + [ + { + onBeforeToolCall: () => ({ + type: 'transformArgs', + args: { count: 'wrong' }, + }), + }, + ], + boundary, + ) + expect(clientInput(chunks)).toBeUndefined() + expect( + chunks.some( + (chunk) => + chunk.type === EventType.TOOL_CALL_RESULT && + chunk.content.includes('Input validation failed'), + ), + ).toBe(true) + }, + ) + + it('runs no middleware for the pending approval preview', async () => { + let hooks = 0 + const chunks = await initial( + { count: '2' }, + clientCheck(z.object({ count: z.string().transform(Number) }), true), + [ + { + onBeforeToolCall() { + hooks++ + }, + }, + ], + ) + expect(hooks).toBe(0) + expect( + descriptors(chunks).find((entry) => entry.id === 'approval_call') + ?.metadata?.input, + ).toEqual({ count: 2 }) + expect(clientInput(chunks)).toBeUndefined() + }) + + it('checks an approval edit after middleware on resume', async () => { + const { adapter, calls } = createMockAdapter({ iterations: [] }) + const seen: Array = [] + const inputSchema = z.object({ + count: z.number().transform((value) => value + 1), + }) + const chunks = await collectChunks( + chat({ + adapter, + parentRunId: 'old', + messages: history({ count: 1 }), + tools: [clientCheck(inputSchema, true)], + resume: [ + { + interruptId: 'approval_call', + status: 'resolved', + payload: { approved: true, editedArgs: { count: 2 } }, + }, + ], + middleware: [ + { + onBeforeToolCall(_ctx, hook) { + seen.push(structuredClone(hook.args)) + }, + }, + ], + }), + ) + expect(calls).toHaveLength(0) + expect(seen).toEqual([{ count: 2 }]) + expect(clientInput(chunks)).toEqual({ count: 3 }) + }) + + it('checks the client input, with no second hook run, when the result comes back', async () => { + const { adapter, calls } = createMockAdapter({ + iterations: [[ev.runStarted(), ev.runFinished()]], + }) + let hooks = 0 + await collectChunks( + chat({ + adapter, + parentRunId: 'old', + messages: history({ count: 'wrong' }), + tools: [clientCheck(schema)], + resume: [ + { + interruptId: 'client_tool_call', + status: 'resolved', + payload: { ok: true }, + }, + ], + middleware: [ + { + onBeforeToolCall() { + hooks++ + }, + }, + ], + }), + ) + expect(hooks).toBe(0) + const results = calls[0]?.messages.filter( + (message) => message.role === 'tool', + ) + expect(results).toHaveLength(1) + expect(JSON.stringify(results?.[0])).toContain('Input validation failed') + }) + + it('records no execute_tool span for a client tool', async () => { + const fakeTracer = createFakeTracer() + await initial({ count: 2 }, clientCheck(schema), [ + otelMiddleware({ + tracer: fakeTracer.tracer, + meter: createFakeMeter().meter, + }), + ]) + expect( + fakeTracer.spans.some((span) => span.name.startsWith('execute_tool')), + ).toBe(false) + }) +}) diff --git a/testing/e2e/src/routes/api.middleware-test.ts b/testing/e2e/src/routes/api.middleware-test.ts index 8c7213fe5a..989a91f817 100644 --- a/testing/e2e/src/routes/api.middleware-test.ts +++ b/testing/e2e/src/routes/api.middleware-test.ts @@ -177,6 +177,16 @@ const toolReplaceMiddleware: ChatMiddleware = { }, } +// Rewrites the args to a value the tool's inputSchema rejects. The final +// input check runs after this hook, so the tool must not run. +const toolBadArgsMiddleware: ChatMiddleware = { + name: 'tool-bad-args', + onBeforeToolCall(_ctx, hookCtx) { + if (hookCtx.toolName !== 'get_weather') return undefined + return { type: 'transformArgs' as const, args: { city: 42 } } + }, +} + /** * Capability provide/consume flow (`capability` mode). * @@ -579,6 +589,8 @@ export const Route = createFileRoute('/api/middleware-test')({ middleware.push(toolSkipMiddleware) if (middlewareMode === 'tool-replace') middleware.push(toolReplaceMiddleware) + if (middlewareMode === 'tool-bad-args') + middleware.push(toolBadArgsMiddleware) if (middlewareMode === 'capability') { // Order matters: the provider's setup() must run before the // consumer reads the capability. Array order is preserved. diff --git a/testing/e2e/src/routes/middleware-test.tsx b/testing/e2e/src/routes/middleware-test.tsx index 6bde171f32..59d3643469 100644 --- a/testing/e2e/src/routes/middleware-test.tsx +++ b/testing/e2e/src/routes/middleware-test.tsx @@ -19,6 +19,7 @@ const MIDDLEWARE_MODES = [ { id: 'chunk-transform', label: 'Chunk Transform (prefix text)' }, { id: 'tool-skip', label: 'Tool Skip (skip with custom result)' }, { id: 'tool-replace', label: 'Tool Replace (replace result after call)' }, + { id: 'tool-bad-args', label: 'Tool Bad Args (invalid args rewrite)' }, { id: 'capability', label: 'Capability (provide/consume prefix)' }, { id: 'phase-recorder', label: 'Phase Recorder (capture phase + chunks)' }, { id: 'otel', label: 'OpenTelemetry (capture spans/metrics)' }, diff --git a/testing/e2e/tests/middleware.spec.ts b/testing/e2e/tests/middleware.spec.ts index dd1d3dcd72..f527393a81 100644 --- a/testing/e2e/tests/middleware.spec.ts +++ b/testing/e2e/tests/middleware.spec.ts @@ -132,6 +132,41 @@ test.describe('Middleware Lifecycle', () => { expect(toolResults[0].content).toContain('skipped') }) + test('onBeforeToolCall args that fail the inputSchema become a tool error', async ({ + page, + testId, + aimockPort, + }) => { + const params = new URLSearchParams() + if (testId) params.set('testId', testId) + if (aimockPort) params.set('aimockPort', String(aimockPort)) + const qs = params.toString() + await page.goto(`/middleware-test${qs ? '?' + qs : ''}`) + await page.waitForTimeout(2000) + await page.locator('#mw-scenario-select').selectOption('with-tool') + await page.locator('#mw-mode-select').selectOption('tool-bad-args') + await page.locator('#mw-run-button').click() + + await page.waitForFunction( + () => + document + .querySelector('#mw-metadata') + ?.getAttribute('data-test-complete') === 'true', + { timeout: 10000 }, + ) + + const messagesJson = await page.locator('#mw-messages-json').textContent() + const messages = JSON.parse(messagesJson || '[]') + const toolResults = messages.flatMap((m: any) => + m.parts.filter((p: any) => p.type === 'tool-result'), + ) + // The middleware rewrote { city: 'NYC' } to { city: 42 }. The schema + // checks the rewritten args, so get_weather never runs. + expect(toolResults).toHaveLength(1) + expect(toolResults[0].content).toContain('Input validation failed') + expect(toolResults[0].content).not.toContain('temperature') + }) + test('onAfterToolCall replaceResult changes what the client and the model see', async ({ page, request, From 3ff07fd28ac4c92cbf3bb08de28921758cb67e72 Mon Sep 17 00:00:00 2001 From: Alem Tuzlak Date: Tue, 13 Oct 2026 20:49:46 +0200 Subject: [PATCH 02/16] test(ai): fix the boundary client tool tests --- .../ai/tests/tool-input-validation.test.ts | 21 ++++++++++--------- 1 file changed, 11 insertions(+), 10 deletions(-) diff --git a/packages/ai/tests/tool-input-validation.test.ts b/packages/ai/tests/tool-input-validation.test.ts index 1e8106f45b..b39b63f92a 100644 --- a/packages/ai/tests/tool-input-validation.test.ts +++ b/packages/ai/tests/tool-input-validation.test.ts @@ -343,11 +343,11 @@ const review = defineInterrupt({ responseSchema: z.boolean(), }) -// Pauses after the model call, so client requests come from the boundary. -const reviewAfterModel = defineChatMiddleware({ +// Pauses before the tools run, so client requests come from the boundary. +const reviewBeforeTools = defineChatMiddleware({ name: 'boundary', onInterruptBoundary(ctx) { - if (ctx.phase !== 'afterModel') return undefined + if (ctx.phase !== 'beforeTools') return undefined return { interrupts: [ review.interrupt({ @@ -376,17 +376,17 @@ async function initial( ], ], }) + const messages: Array = [{ role: 'user', content: 'check' }] + if (!boundary) { + return collectChunks(chat({ adapter, messages, tools: [tool], middleware })) + } return collectChunks( chat({ adapter, - messages: [{ role: 'user', content: 'check' }], + messages, tools: [tool], - ...(boundary - ? { - interrupts: [review], - middleware: [reviewAfterModel, ...middleware], - } - : { middleware }), + interrupts: [review], + middleware: [reviewBeforeTools, ...middleware], }), ) } @@ -442,6 +442,7 @@ describe('client tool input in chat()', () => { chunks.some( (chunk) => chunk.type === EventType.TOOL_CALL_RESULT && + typeof chunk.content === 'string' && chunk.content.includes('Input validation failed'), ), ).toBe(true) From 76a687df8627ed5752bd1ada82729b18b9b780f8 Mon Sep 17 00:00:00 2001 From: Alem Tuzlak Date: Tue, 13 Oct 2026 20:45:37 +0200 Subject: [PATCH 03/16] feat(ai): validate JSON Schema tool input with @cfworker/json-schema --- .changeset/json-schema-tool-input.md | 16 ++ docs/tools/server-tools.md | 15 +- packages/ai/package.json | 1 + .../activities/chat/tools/input-validation.ts | 94 ++++++++- .../ai/tests/tool-input-validation.test.ts | 186 +++++++++++++++++- pnpm-lock.yaml | 64 +++--- testing/e2e/src/routeTree.gen.ts | 21 ++ .../src/routes/api.json-schema-tool-input.ts | 46 +++++ .../e2e/tests/json-schema-tool-input.spec.ts | 13 ++ 9 files changed, 423 insertions(+), 33 deletions(-) create mode 100644 .changeset/json-schema-tool-input.md create mode 100644 testing/e2e/src/routes/api.json-schema-tool-input.ts create mode 100644 testing/e2e/tests/json-schema-tool-input.spec.ts diff --git a/.changeset/json-schema-tool-input.md b/.changeset/json-schema-tool-input.md new file mode 100644 index 0000000000..fa6c635970 --- /dev/null +++ b/.changeset/json-schema-tool-input.md @@ -0,0 +1,16 @@ +--- +'@tanstack/ai': minor +--- + +Check a raw JSON Schema tool input before the tool runs. This covers plain JSON Schema tools and every MCP tool. + +Behavior change: + +- A tool with a JSON Schema `inputSchema` now gets a tool error (state `output-error`) when the input does not fit the schema. The tool does not run. Before, any input ran. +- The check runs at the same place as the Standard Schema check: after approval edits and middleware `onBeforeToolCall`. +- The draft comes from `$schema` (draft-04, draft-07, 2019-09, 2020-12). With no `$schema`, the check uses 2020-12, or draft-04 when the schema has a boolean `exclusiveMinimum` or `exclusiveMaximum`. `format` is not checked. +- The error names each failed location and shows the received arguments: `Input validation failed for tool : #/count: ...`. +- A schema that the check cannot use, such as one with a remote `$ref`, gives a tool error on each call. +- The check uses `@cfworker/json-schema`. It does not compile code, so it runs on Cloudflare Workers. + +Migration: if a JSON Schema tool must accept input that its schema rejects, widen the schema. diff --git a/docs/tools/server-tools.md b/docs/tools/server-tools.md index 4a5104ad64..4866bec070 100644 --- a/docs/tools/server-tools.md +++ b/docs/tools/server-tools.md @@ -355,7 +355,7 @@ const getUserData = getUserDataDef.server(async ({ userId }) => { **Throwing vs. returning an error:** if your `.server()` function throws, the SDK catches it and surfaces it as a tool-result *error* (the model sees the failure but you lose control over the message). Returning a structured `{ error }` shape keeps the model in control of how to recover and is usually preferable. Either way, when an `outputSchema` is defined the returned value is validated against it (Zod) before being added to the conversation — so include the `error` field in your `outputSchema` if you return it. -Malformed JSON arguments and Standard Schema input validation failures are also surfaced as tool-result errors. The tool implementation is not called, and the agent loop can return the error to the model so it can repair the tool call. This also applies to Chat Completions tools with no required input fields. +Malformed JSON arguments and input validation failures (Standard Schema or JSON Schema) are also surfaced as tool-result errors. The tool implementation is not called, and the agent loop can return the error to the model so it can repair the tool call. This also applies to Chat Completions tools with no required input fields. The `inputSchema` checks the input just before the tool runs: @@ -409,14 +409,23 @@ const getUserData = getUserDataDef.server(async (args) => { }); ``` -> **Note:** JSON Schema tools skip runtime validation. Zod schemas are recommended for full type safety and validation. +The JSON Schema checks each call before the tool runs. MCP tools get the same check. If the input does not fit, the call gets a tool error, and the error names each failed field: + +```text +Input validation failed for tool get_user_data: #: Instance does not have required property "userId". +``` + +- The check uses the draft that `$schema` names: draft-04, draft-07, 2019-09, or 2020-12. With no `$schema`, it uses 2020-12. +- The check does not test `format` (such as `uri` or `email`). A `format` only describes the value. +- The check does not compile code, so it also runs on Cloudflare Workers. +- `args` is still typed as `unknown`. A Zod schema gives you the type too. > **Tip:** When you pass `.client()` tools to `useChat`, a check on `part.name` narrows `part.input` and `part.output`. See [Type-safe tool call events](../chat/stream-events#type-safe-tool-call-events). ## Best Practices 1. **Keep tools focused** - Each tool should do one thing well -2. **Validate inputs** - Use Zod schemas to ensure type safety (JSON Schema skips validation) +2. **Validate inputs** - The `inputSchema` checks every call. A Zod schema also gives you types 3. **Handle errors** - Return meaningful error messages 4. **Use descriptions** - Clear descriptions help the model use tools correctly 5. **Secure sensitive operations** - Never expose API keys or secrets to the client diff --git a/packages/ai/package.json b/packages/ai/package.json index 9f2eca7f77..81fd734828 100644 --- a/packages/ai/package.json +++ b/packages/ai/package.json @@ -102,6 +102,7 @@ ], "dependencies": { "@ag-ui/core": "1.0.0", + "@cfworker/json-schema": "^4.1.1", "@standard-schema/spec": "^1.1.0", "@tanstack/ai-event-client": "workspace:^", "@tanstack/ai-utils": "workspace:^", diff --git a/packages/ai/src/activities/chat/tools/input-validation.ts b/packages/ai/src/activities/chat/tools/input-validation.ts index 64c6c06ebc..f1657e77fd 100644 --- a/packages/ai/src/activities/chat/tools/input-validation.ts +++ b/packages/ai/src/activities/chat/tools/input-validation.ts @@ -1,9 +1,54 @@ +import { Validator } from '@cfworker/json-schema' import { StandardSchemaValidationError, + convertSchemaToJsonSchema, isStandardSchema, validateWithStandardSchema, } from './schema-converter' -import type { SchemaInput } from '../../../types' +import type { Schema, SchemaDraft } from '@cfworker/json-schema' +import type { JSONSchema, SchemaInput } from '../../../types' + +const validators = new WeakMap() + +/** + * The draft named by `$schema`. Default: 2020-12. A schema with no `$schema` + * that writes `exclusiveMinimum: true` (draft-04, OpenAPI 3.0) is draft-04. + */ +function draftOf(schema: JSONSchema, json: string): SchemaDraft { + const uri = typeof schema.$schema === 'string' ? schema.$schema : '' + if (uri.includes('draft-04')) return '4' + if (uri.includes('draft-06') || uri.includes('draft-07')) return '7' + if (uri.includes('2019-09')) return '2019-09' + if (!uri && /"exclusiveM(?:in|ax)imum":(?:true|false)/.test(json)) return '4' + return '2020-12' +} + +function validatorFor(schema: JSONSchema | boolean) { + if (typeof schema === 'boolean') return new Validator(schema) + const cached = validators.get(schema) + if (cached) return cached + // The validator writes lookup data onto the schema it gets. A JSON copy + // keeps the tool's own schema, which also goes to the provider, unchanged. + // `format` is an annotation, as in 2019-09 and later: a tool that declares + // `format: 'uri'` still gets a plain path. + const json = JSON.stringify(schema) + const copy: Schema = JSON.parse(json, (key, value) => + key === 'format' && typeof value === 'string' ? undefined : value, + ) + const validator = new Validator(copy, draftOf(schema, json), false) + validators.set(schema, validator) + return validator +} + +/** Errors that only say that a nested check failed. The nested error says why. */ +const WRAPPER_KEYWORDS = new Set([ + 'properties', + 'patternProperties', + 'items', + 'prefixItems', + '$ref', + 'allOf', +]) function receivedArgumentsText(received: unknown) { try { @@ -13,18 +58,57 @@ function receivedArgumentsText(received: unknown) { } } +function validateJsonInput( + schema: JSONSchema | boolean, + received: unknown, + toolName: string, +) { + let errors: string + try { + // Check the JSON form: a key set to `undefined` is the same as no key. + const result = validatorFor(schema).validate( + JSON.parse(JSON.stringify(received)), + ) + if (result.valid) return received + const leaves = result.errors.filter( + (error) => !WRAPPER_KEYWORDS.has(error.keyword), + ) + errors = (leaves.length > 0 ? leaves : result.errors) + .map((error) => `${error.instanceLocation}: ${error.error}`) + .join(', ') + } catch (error) { + // A schema the validator cannot use, such as an unresolved `$ref`, or a + // value with no JSON form. The first line is enough for the model. + const message = error instanceof Error ? error.message : String(error) + errors = message.split('\n', 1).join('') + } + throw new Error( + `Input validation failed for tool ${toolName}: ${errors}\nReceived arguments:\n${receivedArgumentsText(received)}`, + ) +} + /** * The final check of a tool input, just before the tool runs or the client * gets the call. A Standard Schema runs here (an async validator too), and - * its output is the input, so transforms apply once. The error goes back to - * the model as the tool result. + * its output is the input, so transforms apply once. A JSON Schema (an MCP + * tool, a plain JSON Schema tool) is checked with `@cfworker/json-schema`. + * The error goes back to the model as the tool result. */ export async function validateToolInput( - schema: SchemaInput | undefined, + schema: SchemaInput | boolean | undefined, received: unknown, toolName: string, ): Promise { - if (!isStandardSchema(schema)) return received + if (schema === undefined) return received + if (typeof schema === 'boolean') { + return validateJsonInput(schema, received, toolName) + } + if (!isStandardSchema(schema)) { + const jsonSchema = convertSchemaToJsonSchema(schema, { io: 'input' }) + return jsonSchema + ? validateJsonInput(jsonSchema, received, toolName) + : received + } const validation = await validateWithStandardSchema(schema, received) if (validation.success) return validation.data const error = new StandardSchemaValidationError(validation.issues) diff --git a/packages/ai/tests/tool-input-validation.test.ts b/packages/ai/tests/tool-input-validation.test.ts index b39b63f92a..938186fa2c 100644 --- a/packages/ai/tests/tool-input-validation.test.ts +++ b/packages/ai/tests/tool-input-validation.test.ts @@ -1,4 +1,4 @@ -import { describe, expect, it } from 'vitest' +import { describe, expect, it, vi } from 'vitest' import { z } from 'zod' import { chat } from '../src/activities/chat/index' import { defineChatMiddleware } from '../src/activities/chat/middleware/define' @@ -13,6 +13,7 @@ import type { ChatMiddleware } from '../src/activities/chat/middleware/types' import type { ToolExecutionMiddlewareHooks } from '../src/activities/chat/tools/tool-calls' import type { ToolApprovalResolution } from '../src/interrupts' import type { + JSONSchema, ModelMessage, SchemaInput, StreamChunk, @@ -301,6 +302,189 @@ describe('final tool input validation', () => { ) }) +describe('JSON Schema tool input', () => { + const jsonSchema: JSONSchema = { + type: 'object', + properties: { count: { type: 'integer', minimum: 1 } }, + required: ['count'], + additionalProperties: false, + } + + it('runs valid input', async () => { + const result = await runTool('{"count":2}', jsonSchema) + expect(result.executed).toEqual([{ count: 2 }]) + }) + + it.each(['{}', '{"count":"2"}', '{"count":0}', '{"count":2,"extra":true}'])( + 'rejects invalid input without executing: %s', + async (raw) => { + const result = await runTool(raw, jsonSchema) + expect(result.executed).toEqual([]) + expect(result.results[0]?.state).toBe('output-error') + expect(result.results[0]?.result).toEqual({ + error: expect.stringContaining( + 'Input validation failed for tool check', + ), + }) + }, + ) + + it('names the failed location and shows the received arguments', async () => { + const result = await runTool('{"count":0}', jsonSchema) + expect(result.results[0]?.result).toEqual({ + error: expect.stringMatching(/#\/count: .*\n.*Received arguments:/s), + }) + expect(result.results[0]?.result).toEqual({ + error: expect.stringContaining('"count": 0'), + }) + }) + + it('checks local references and conditional fields', async () => { + const schema: JSONSchema = { + type: 'object', + $defs: { count: { type: 'integer' } }, + properties: { + mode: { type: 'string' }, + count: { $ref: '#/$defs/count' }, + }, + if: { properties: { mode: { const: 'count' } }, required: ['mode'] }, + then: { required: ['count'] }, + } + expect( + (await runTool('{"mode":"count","count":3}', schema)).executed, + ).toEqual([{ mode: 'count', count: 3 }]) + expect((await runTool('{"mode":"count"}', schema)).executed).toEqual([]) + expect((await runTool('{"count":"3"}', schema)).executed).toEqual([]) + }) + + it('uses the draft named by $schema', async () => { + // Draft 4 spells an exclusive bound as a boolean next to `maximum`. + const schema: JSONSchema = JSON.parse( + JSON.stringify({ + $schema: 'http://json-schema.org/draft-04/schema#', + type: 'object', + properties: { + count: { type: 'number', maximum: 5, exclusiveMaximum: true }, + }, + }), + ) + expect((await runTool('{"count":4}', schema)).executed).toEqual([ + { count: 4 }, + ]) + expect((await runTool('{"count":5}', schema)).executed).toEqual([]) + }) + + it('runs empty and null arguments as {} when nothing is required', async () => { + for (const raw of ['', 'null']) { + const result = await runTool(raw, { type: 'object', properties: {} }) + expect(result.executed).toEqual([{}]) + } + }) + + it('rejects empty arguments when a field is required', async () => { + const result = await runTool('', jsonSchema) + expect(result.executed).toEqual([]) + expect(result.results[0]?.state).toBe('output-error') + }) + + it('checks a middleware rewrite and an approval edit', async () => { + const rewritten = await runTool('{"count":2}', jsonSchema, { + middleware: { + onBeforeToolCall: async () => ({ + type: 'transformArgs', + args: { count: 'wrong' }, + }), + }, + }) + expect(rewritten.executed).toEqual([]) + const edited = await runTool('{"count":2}', jsonSchema, { + needsApproval: true, + approvals: new Map([ + ['call-1', { approved: true, editedArgs: { count: -1 } }], + ]), + }) + expect(edited.executed).toEqual([]) + expect(edited.results[0]?.state).toBe('output-error') + }) + + it('turns an unusable schema into a tool error', async () => { + const result = await runTool('{"count":2}', { + type: 'object', + properties: { count: { $ref: 'https://example.com/count.json' } }, + }) + expect(result.executed).toEqual([]) + expect(result.results[0]?.result).toEqual({ + error: expect.stringContaining('Unresolved $ref'), + }) + expect(JSON.stringify(result.results[0]?.result)).not.toContain( + 'Known schemas', + ) + }) + + it('reads a boolean exclusive bound with no $schema as draft-04', async () => { + const schema: JSONSchema = JSON.parse( + '{"type":"object","properties":{"count":{"type":"number","maximum":5,"exclusiveMaximum":true}}}', + ) + expect((await runTool('{"count":4}', schema)).executed).toEqual([ + { count: 4 }, + ]) + expect((await runTool('{"count":5}', schema)).executed).toEqual([]) + }) + + it('treats format as an annotation', async () => { + const result = await runTool('{"path":"/tmp/report.txt"}', { + type: 'object', + properties: { path: { type: 'string', format: 'uri' } }, + required: ['path'], + }) + expect(result.executed).toEqual([{ path: '/tmp/report.txt' }]) + }) + + it('treats a key set to undefined as a missing key', async () => { + const result = await runTool('{"count":2}', jsonSchema, { + middleware: { + onBeforeToolCall: async () => ({ + type: 'transformArgs', + args: { count: 2, note: undefined }, + }), + }, + }) + expect(result.executed).toEqual([{ count: 2, note: undefined }]) + }) + + it('does not change the tool schema', async () => { + const schema: JSONSchema = { + type: 'object', + $defs: { count: { type: 'integer' } }, + properties: { count: { $ref: '#/$defs/count' } }, + } + const before = JSON.stringify(schema) + await runTool('{"count":2}', schema) + expect(JSON.stringify(schema)).toBe(before) + expect(Object.getOwnPropertyNames(schema.properties?.count)).toEqual([ + '$ref', + ]) + }) + + it('works when runtime code generation is blocked', async () => { + vi.stubGlobal('Function', function blockedCodeGeneration() { + throw new Error('Code generation is blocked') + }) + try { + const schema: JSONSchema = { + type: 'object', + properties: { count: { type: 'integer' } }, + } + expect((await runTool('{"count":2}', schema)).executed).toEqual([ + { count: 2 }, + ]) + expect((await runTool('{"count":"2"}', schema)).executed).toEqual([]) + } finally { + vi.unstubAllGlobals() + } + }) +}) + const clientCheck = ( inputSchema: SchemaInput, needsApproval = false, diff --git a/pnpm-lock.yaml b/pnpm-lock.yaml index aa9e45eb08..cc63ffd52a 100644 --- a/pnpm-lock.yaml +++ b/pnpm-lock.yaml @@ -1072,10 +1072,10 @@ importers: version: 2.68.0 '@mcp-ui/client': specifier: ^7.1.1 - version: 7.1.1(react-dom@19.2.3(react@19.2.3))(react@19.2.3) + version: 7.1.1(@cfworker/json-schema@4.1.1)(react-dom@19.2.3(react@19.2.3))(react@19.2.3) '@modelcontextprotocol/sdk': specifier: ^1.29.0 - version: 1.29.0(supports-color@7.2.0)(zod@4.2.1) + version: 1.29.0(@cfworker/json-schema@4.1.1)(supports-color@7.2.0)(zod@4.2.1) '@ngrok/ngrok': specifier: ^1.7.0 version: 1.7.0 @@ -2011,6 +2011,9 @@ importers: '@ag-ui/core': specifier: 1.0.0 version: 1.0.0(zod@4.2.1) + '@cfworker/json-schema': + specifier: ^4.1.1 + version: 4.1.1 '@standard-schema/spec': specifier: ^1.1.0 version: 1.1.0 @@ -2497,7 +2500,7 @@ importers: dependencies: '@google/genai': specifier: ^2.10.0 - version: 2.10.0(@modelcontextprotocol/sdk@1.29.0(zod@4.3.6)) + version: 2.10.0(@modelcontextprotocol/sdk@1.29.0(@cfworker/json-schema@4.1.1)(zod@4.3.6)) '@tanstack/ai-utils': specifier: workspace:^ version: link:../ai-utils @@ -2976,7 +2979,7 @@ importers: devDependencies: '@mcp-ui/client': specifier: ^7.1.1 - version: 7.1.1(react-dom@19.2.3(react@19.2.3))(react@19.2.3) + version: 7.1.1(@cfworker/json-schema@4.1.1)(react-dom@19.2.3(react@19.2.3))(react@19.2.3) '@tanstack/ai': specifier: workspace:* version: link:../ai @@ -3007,7 +3010,7 @@ importers: devDependencies: '@mcp-ui/client': specifier: ^7.1.1 - version: 7.1.1(react-dom@19.2.3(react@19.2.3))(react@19.2.3) + version: 7.1.1(@cfworker/json-schema@4.1.1)(react-dom@19.2.3(react@19.2.3))(react@19.2.3) '@standard-schema/spec': specifier: ^1.1.0 version: 1.1.0 @@ -3100,7 +3103,7 @@ importers: dependencies: '@modelcontextprotocol/sdk': specifier: ^1.29.0 - version: 1.29.0(zod@4.3.6) + version: 1.29.0(@cfworker/json-schema@4.1.1)(zod@4.3.6) '@tanstack/ai-skills': specifier: workspace:^ version: link:../ai-skills @@ -3125,7 +3128,7 @@ importers: dependencies: '@blaxel/core': specifier: ^0.3.10 - version: 0.3.11(@hey-api/openapi-ts@0.99.0(magicast@0.5.2)(typescript@7.0.2)) + version: 0.3.11(@cfworker/json-schema@4.1.1)(@hey-api/openapi-ts@0.99.0(magicast@0.5.2)(typescript@7.0.2)) devDependencies: '@tanstack/ai-sandbox': specifier: workspace:* @@ -3450,7 +3453,7 @@ importers: devDependencies: '@google/genai': specifier: ^2.10.0 - version: 2.10.0(@modelcontextprotocol/sdk@1.29.0(zod@4.3.6)) + version: 2.10.0(@modelcontextprotocol/sdk@1.29.0(@cfworker/json-schema@4.1.1)(zod@4.6.5)) '@tanstack/ai': specifier: workspace:* version: link:../ai @@ -3641,7 +3644,7 @@ importers: version: 1.34.0(vitest@4.1.11(@opentelemetry/api@1.9.1)(@types/node@24.10.3)(@vitest/coverage-v8@4.1.10(vitest@4.1.10))(happy-dom@20.11.2)(jsdom@27.4.0(@noble/hashes@2.3.0))(vite@8.2.1(@types/node@24.10.3)(esbuild@0.28.1)(jiti@2.7.0)(less@4.6.6)(sass@1.101.0)(terser@5.44.1)(tsx@4.23.12)(yaml@2.9.0))) '@modelcontextprotocol/sdk': specifier: ^1.29.0 - version: 1.29.0(zod@4.3.6) + version: 1.29.0(@cfworker/json-schema@4.1.1)(zod@4.3.6) '@openrouter/sdk': specifier: 0.13.20 version: 0.13.20 @@ -5106,6 +5109,9 @@ packages: '@bufbuild/protobuf@2.15.0': resolution: {integrity: sha512-DAheWUkVr/SJTWCc+lg9dhY0eN4SaWlf4+bG1KzHeXbnqt0AfB/NX0Z+VunGlM1ki1B4zVvye27MpKh/svySUA==} + '@cfworker/json-schema@4.1.1': + resolution: {integrity: sha512-gAmrUZSGtKc3AiBL71iNWxDsyUC5uMaKKGdvzYsBoTW/xi42JQHl7eKV2OYzCUqvc+D2RCcf7EXY2iCyFIk6og==} + '@changesets/apply-release-plan@7.1.0': resolution: {integrity: sha512-yq8ML3YS7koKQ/9bk1PqO0HMzApIFNwjlwCnwFEXMzNe8NpzeeYYKCmnhWJGkN8g7E51MnWaSbqRcTcdIxUgnQ==} @@ -20988,10 +20994,10 @@ snapshots: '@bcoe/v8-coverage@1.0.2': {} - '@blaxel/core@0.3.11(@hey-api/openapi-ts@0.99.0(magicast@0.5.2)(typescript@7.0.2))': + '@blaxel/core@0.3.11(@cfworker/json-schema@4.1.1)(@hey-api/openapi-ts@0.99.0(magicast@0.5.2)(typescript@7.0.2))': dependencies: '@hey-api/client-fetch': 0.10.2(@hey-api/openapi-ts@0.99.0(magicast@0.5.2)(typescript@7.0.2)) - '@modelcontextprotocol/sdk': 1.29.0(zod@4.3.6) + '@modelcontextprotocol/sdk': 1.29.0(@cfworker/json-schema@4.1.1)(zod@4.3.6) archiver: 7.0.1 axios: 1.18.1(debug@4.4.3(supports-color@7.2.0))(supports-color@7.2.0) dockerfile-ast: 0.7.1 @@ -21022,6 +21028,8 @@ snapshots: '@bufbuild/protobuf@2.15.0': {} + '@cfworker/json-schema@4.1.1': {} + '@changesets/apply-release-plan@7.1.0': dependencies: '@changesets/config': 3.1.3 @@ -22223,27 +22231,27 @@ snapshots: '@shikijs/types': 3.23.0 '@shikijs/vscode-textmate': 10.0.2 - '@google/genai@2.10.0(@modelcontextprotocol/sdk@1.29.0(zod@4.3.6))': + '@google/genai@2.10.0(@modelcontextprotocol/sdk@1.29.0(@cfworker/json-schema@4.1.1)(zod@4.3.6))': dependencies: google-auth-library: 10.5.0 p-retry: 4.6.2 protobufjs: 7.6.4 ws: 8.21.0 optionalDependencies: - '@modelcontextprotocol/sdk': 1.29.0(zod@4.3.6) + '@modelcontextprotocol/sdk': 1.29.0(@cfworker/json-schema@4.1.1)(zod@4.3.6) transitivePeerDependencies: - bufferutil - supports-color - utf-8-validate - '@google/genai@2.10.0(@modelcontextprotocol/sdk@1.29.0(zod@4.6.5))': + '@google/genai@2.10.0(@modelcontextprotocol/sdk@1.29.0(@cfworker/json-schema@4.1.1)(zod@4.6.5))': dependencies: google-auth-library: 10.5.0 p-retry: 4.6.2 protobufjs: 7.6.4 ws: 8.21.0 optionalDependencies: - '@modelcontextprotocol/sdk': 1.29.0(zod@4.6.5) + '@modelcontextprotocol/sdk': 1.29.0(@cfworker/json-schema@4.1.1)(zod@4.6.5) transitivePeerDependencies: - bufferutil - supports-color @@ -22637,10 +22645,10 @@ snapshots: - encoding - supports-color - '@mcp-ui/client@7.1.1(react-dom@19.2.3(react@19.2.3))(react@19.2.3)': + '@mcp-ui/client@7.1.1(@cfworker/json-schema@4.1.1)(react-dom@19.2.3(react@19.2.3))(react@19.2.3)': dependencies: - '@modelcontextprotocol/ext-apps': 1.7.4(@modelcontextprotocol/sdk@1.29.0(supports-color@7.2.0)(zod@4.2.1))(react-dom@19.2.3(react@19.2.3))(react@19.2.3)(zod@3.25.76) - '@modelcontextprotocol/sdk': 1.29.0(zod@3.25.76) + '@modelcontextprotocol/ext-apps': 1.7.4(@modelcontextprotocol/sdk@1.29.0(@cfworker/json-schema@4.1.1)(supports-color@7.2.0)(zod@4.2.1))(react-dom@19.2.3(react@19.2.3))(react@19.2.3)(zod@3.25.76) + '@modelcontextprotocol/sdk': 1.29.0(@cfworker/json-schema@4.1.1)(zod@3.25.76) react: 19.2.3 react-dom: 19.2.3(react@19.2.3) zod: 3.25.76 @@ -22738,16 +22746,16 @@ snapshots: dependencies: zod: 4.3.6 - '@modelcontextprotocol/ext-apps@1.7.4(@modelcontextprotocol/sdk@1.29.0(supports-color@7.2.0)(zod@4.2.1))(react-dom@19.2.3(react@19.2.3))(react@19.2.3)(zod@3.25.76)': + '@modelcontextprotocol/ext-apps@1.7.4(@modelcontextprotocol/sdk@1.29.0(@cfworker/json-schema@4.1.1)(supports-color@7.2.0)(zod@4.2.1))(react-dom@19.2.3(react@19.2.3))(react@19.2.3)(zod@3.25.76)': dependencies: - '@modelcontextprotocol/sdk': 1.29.0(supports-color@7.2.0)(zod@4.2.1) + '@modelcontextprotocol/sdk': 1.29.0(@cfworker/json-schema@4.1.1)(supports-color@7.2.0)(zod@4.2.1) '@standard-schema/spec': 1.1.0 zod: 3.25.76 optionalDependencies: react: 19.2.3 react-dom: 19.2.3(react@19.2.3) - '@modelcontextprotocol/sdk@1.29.0(supports-color@7.2.0)(zod@4.2.1)': + '@modelcontextprotocol/sdk@1.29.0(@cfworker/json-schema@4.1.1)(supports-color@7.2.0)(zod@4.2.1)': dependencies: '@hono/node-server': 1.19.17(hono@4.13.8) ajv: 8.20.0 @@ -22766,10 +22774,12 @@ snapshots: raw-body: 3.0.2 zod: 4.2.1 zod-to-json-schema: 3.25.2(zod@4.2.1) + optionalDependencies: + '@cfworker/json-schema': 4.1.1 transitivePeerDependencies: - supports-color - '@modelcontextprotocol/sdk@1.29.0(zod@3.25.76)': + '@modelcontextprotocol/sdk@1.29.0(@cfworker/json-schema@4.1.1)(zod@3.25.76)': dependencies: '@hono/node-server': 1.19.17(hono@4.13.8) ajv: 8.20.0 @@ -22788,10 +22798,12 @@ snapshots: raw-body: 3.0.2 zod: 3.25.76 zod-to-json-schema: 3.25.2(zod@3.25.76) + optionalDependencies: + '@cfworker/json-schema': 4.1.1 transitivePeerDependencies: - supports-color - '@modelcontextprotocol/sdk@1.29.0(zod@4.3.6)': + '@modelcontextprotocol/sdk@1.29.0(@cfworker/json-schema@4.1.1)(zod@4.3.6)': dependencies: '@hono/node-server': 1.19.17(hono@4.13.8) ajv: 8.20.0 @@ -22810,10 +22822,12 @@ snapshots: raw-body: 3.0.2 zod: 4.3.6 zod-to-json-schema: 3.25.2(zod@4.3.6) + optionalDependencies: + '@cfworker/json-schema': 4.1.1 transitivePeerDependencies: - supports-color - '@modelcontextprotocol/sdk@1.29.0(zod@4.6.5)': + '@modelcontextprotocol/sdk@1.29.0(@cfworker/json-schema@4.1.1)(zod@4.6.5)': dependencies: '@hono/node-server': 1.19.17(hono@4.13.8) ajv: 8.20.0 @@ -22832,6 +22846,8 @@ snapshots: raw-body: 3.0.2 zod: 4.6.5 zod-to-json-schema: 3.25.2(zod@4.6.5) + optionalDependencies: + '@cfworker/json-schema': 4.1.1 transitivePeerDependencies: - supports-color optional: true diff --git a/testing/e2e/src/routeTree.gen.ts b/testing/e2e/src/routeTree.gen.ts index aa9537f6db..f964927f01 100644 --- a/testing/e2e/src/routeTree.gen.ts +++ b/testing/e2e/src/routeTree.gen.ts @@ -95,6 +95,7 @@ import { Route as ApiGenerationPersistenceServerRouteImport } from './routes/api import { Route as ApiImageRouteImport } from './routes/api.image' import { Route as ApiInterruptsTestRouteImport } from './routes/api.interrupts-test' import { Route as ApiJoinRunClientToolRouteImport } from './routes/api.join-run-client-tool' +import { Route as ApiJsonSchemaToolInputRouteImport } from './routes/api.json-schema-tool-input' import { Route as ApiJsonTransportRouteImport } from './routes/api.json-transport' import { Route as ApiKeyedAdapterRouteImport } from './routes/api.keyed-adapter' import { Route as ApiLazyToolsWireRouteImport } from './routes/api.lazy-tools-wire' @@ -628,6 +629,11 @@ const ApiJoinRunClientToolRoute = ApiJoinRunClientToolRouteImport.update({ path: '/api/join-run-client-tool', getParentRoute: () => rootRouteImport, } as any) +const ApiJsonSchemaToolInputRoute = ApiJsonSchemaToolInputRouteImport.update({ + id: '/api/json-schema-tool-input', + path: '/api/json-schema-tool-input', + getParentRoute: () => rootRouteImport, +} as any) const ApiJsonTransportRoute = ApiJsonTransportRouteImport.update({ id: '/api/json-transport', path: '/api/json-transport', @@ -1146,6 +1152,7 @@ export interface FileRoutesByFullPath { '/api/image': typeof ApiImageRouteWithChildren '/api/interrupts-test': typeof ApiInterruptsTestRoute '/api/join-run-client-tool': typeof ApiJoinRunClientToolRoute + '/api/json-schema-tool-input': typeof ApiJsonSchemaToolInputRoute '/api/json-transport': typeof ApiJsonTransportRoute '/api/keyed-adapter': typeof ApiKeyedAdapterRoute '/api/lazy-tools-wire': typeof ApiLazyToolsWireRoute @@ -1316,6 +1323,7 @@ export interface FileRoutesByTo { '/api/image': typeof ApiImageRouteWithChildren '/api/interrupts-test': typeof ApiInterruptsTestRoute '/api/join-run-client-tool': typeof ApiJoinRunClientToolRoute + '/api/json-schema-tool-input': typeof ApiJsonSchemaToolInputRoute '/api/json-transport': typeof ApiJsonTransportRoute '/api/keyed-adapter': typeof ApiKeyedAdapterRoute '/api/lazy-tools-wire': typeof ApiLazyToolsWireRoute @@ -1487,6 +1495,7 @@ export interface FileRoutesById { '/api/image': typeof ApiImageRouteWithChildren '/api/interrupts-test': typeof ApiInterruptsTestRoute '/api/join-run-client-tool': typeof ApiJoinRunClientToolRoute + '/api/json-schema-tool-input': typeof ApiJsonSchemaToolInputRoute '/api/json-transport': typeof ApiJsonTransportRoute '/api/keyed-adapter': typeof ApiKeyedAdapterRoute '/api/lazy-tools-wire': typeof ApiLazyToolsWireRoute @@ -1659,6 +1668,7 @@ export interface FileRouteTypes { | '/api/image' | '/api/interrupts-test' | '/api/join-run-client-tool' + | '/api/json-schema-tool-input' | '/api/json-transport' | '/api/keyed-adapter' | '/api/lazy-tools-wire' @@ -1829,6 +1839,7 @@ export interface FileRouteTypes { | '/api/image' | '/api/interrupts-test' | '/api/join-run-client-tool' + | '/api/json-schema-tool-input' | '/api/json-transport' | '/api/keyed-adapter' | '/api/lazy-tools-wire' @@ -1999,6 +2010,7 @@ export interface FileRouteTypes { | '/api/image' | '/api/interrupts-test' | '/api/join-run-client-tool' + | '/api/json-schema-tool-input' | '/api/json-transport' | '/api/keyed-adapter' | '/api/lazy-tools-wire' @@ -2170,6 +2182,7 @@ export interface RootRouteChildren { ApiImageRoute: typeof ApiImageRouteWithChildren ApiInterruptsTestRoute: typeof ApiInterruptsTestRoute ApiJoinRunClientToolRoute: typeof ApiJoinRunClientToolRoute + ApiJsonSchemaToolInputRoute: typeof ApiJsonSchemaToolInputRoute ApiJsonTransportRoute: typeof ApiJsonTransportRoute ApiKeyedAdapterRoute: typeof ApiKeyedAdapterRoute ApiLazyToolsWireRoute: typeof ApiLazyToolsWireRoute @@ -2854,6 +2867,13 @@ declare module '@tanstack/react-router' { preLoaderRoute: typeof ApiJoinRunClientToolRouteImport parentRoute: typeof rootRouteImport } + '/api/json-schema-tool-input': { + id: '/api/json-schema-tool-input' + path: '/api/json-schema-tool-input' + fullPath: '/api/json-schema-tool-input' + preLoaderRoute: typeof ApiJsonSchemaToolInputRouteImport + parentRoute: typeof rootRouteImport + } '/api/json-transport': { id: '/api/json-transport' path: '/api/json-transport' @@ -3580,6 +3600,7 @@ const rootRouteChildren: RootRouteChildren = { ApiImageRoute: ApiImageRouteWithChildren, ApiInterruptsTestRoute: ApiInterruptsTestRoute, ApiJoinRunClientToolRoute: ApiJoinRunClientToolRoute, + ApiJsonSchemaToolInputRoute: ApiJsonSchemaToolInputRoute, ApiJsonTransportRoute: ApiJsonTransportRoute, ApiKeyedAdapterRoute: ApiKeyedAdapterRoute, ApiLazyToolsWireRoute: ApiLazyToolsWireRoute, diff --git a/testing/e2e/src/routes/api.json-schema-tool-input.ts b/testing/e2e/src/routes/api.json-schema-tool-input.ts new file mode 100644 index 0000000000..fe5a21b4b4 --- /dev/null +++ b/testing/e2e/src/routes/api.json-schema-tool-input.ts @@ -0,0 +1,46 @@ +import { createFileRoute } from '@tanstack/react-router' +import { chat, toolDefinition } from '@tanstack/ai' +import { fakeText } from '@tanstack/ai/testing' + +// No network: `fakeText` answers from a script. The first answer calls a +// JSON Schema tool with a wrong type, the second answer repeats the tool +// result. MCP tools have the same raw JSON Schema input. +export const Route = createFileRoute('/api/json-schema-tool-input')({ + server: { + handlers: { + POST: async () => { + const fake = fakeText() + fake.setResponses([ + { toolCalls: [{ name: 'count_items', input: { count: 'many' } }] }, + ({ request }) => ({ + text: `Tool said: ${String(request.messages.at(-1)?.content)}`, + }), + ]) + + const runs: Array = [] + const countItems = toolDefinition({ + name: 'count_items', + description: 'Count the items', + inputSchema: { + $schema: 'http://json-schema.org/draft-07/schema#', + type: 'object', + properties: { count: { type: 'integer' } }, + required: ['count'], + }, + }).server((input) => { + runs.push(input) + return 'counted' + }) + + const { text } = await chat({ + adapter: fake, + messages: [{ role: 'user', content: 'Count the items' }], + tools: [countItems], + stream: false, + }) + + return Response.json({ text, runs }) + }, + }, + }, +}) diff --git a/testing/e2e/tests/json-schema-tool-input.spec.ts b/testing/e2e/tests/json-schema-tool-input.spec.ts new file mode 100644 index 0000000000..101a1fdbe8 --- /dev/null +++ b/testing/e2e/tests/json-schema-tool-input.spec.ts @@ -0,0 +1,13 @@ +import { test, expect } from './fixtures' + +test('a JSON Schema tool with invalid input returns a tool error and does not run', async ({ + request, +}) => { + const response = await request.post('/api/json-schema-tool-input') + + expect(response.ok()).toBe(true) + const body = await response.json() + expect(body.runs).toEqual([]) + expect(body.text).toContain('Input validation failed for tool count_items') + expect(body.text).toContain('#/count') +}) From fff8dafadd890bca8f68971d8f3307cde3ae5f65 Mon Sep 17 00:00:00 2001 From: Alem Tuzlak Date: Tue, 13 Oct 2026 21:20:11 +0200 Subject: [PATCH 04/16] chore: keep one @google/genai copy for ai-vertex and ai-gemini --- pnpm-lock.yaml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pnpm-lock.yaml b/pnpm-lock.yaml index cc63ffd52a..9c11f88a32 100644 --- a/pnpm-lock.yaml +++ b/pnpm-lock.yaml @@ -3453,7 +3453,7 @@ importers: devDependencies: '@google/genai': specifier: ^2.10.0 - version: 2.10.0(@modelcontextprotocol/sdk@1.29.0(@cfworker/json-schema@4.1.1)(zod@4.6.5)) + version: 2.10.0(@modelcontextprotocol/sdk@1.29.0(@cfworker/json-schema@4.1.1)(zod@4.3.6)) '@tanstack/ai': specifier: workspace:* version: link:../ai From e065ccb1698f78eb962d34c994cc229f61a62063 Mon Sep 17 00:00:00 2001 From: Alem Tuzlak Date: Tue, 13 Oct 2026 20:59:25 +0200 Subject: [PATCH 05/16] feat(ai): repair common model mistakes in tool input before the final check --- .changeset/tool-input-coercion.md | 11 + docs/tools/server-tools.md | 14 + .../activities/chat/tools/input-validation.ts | 314 ++++++++++++++++-- .../ai/tests/tool-input-validation.test.ts | 219 +++++++++++- .../src/routes/api.json-schema-tool-input.ts | 17 +- testing/e2e/src/routes/api.middleware-test.ts | 2 +- testing/e2e/src/routes/api.tools-test.ts | 5 +- .../e2e/tests/json-schema-tool-input.spec.ts | 6 +- testing/e2e/tests/middleware.spec.ts | 2 +- 9 files changed, 539 insertions(+), 51 deletions(-) create mode 100644 .changeset/tool-input-coercion.md diff --git a/.changeset/tool-input-coercion.md b/.changeset/tool-input-coercion.md new file mode 100644 index 0000000000..95825e8005 --- /dev/null +++ b/.changeset/tool-input-coercion.md @@ -0,0 +1,11 @@ +--- +'@tanstack/ai': minor +--- + +Repair common model mistakes in a tool input before the final input check. + +- When a JSON Schema tool's input (and every MCP tool's input) fails the check, it is repaired and checked again: `"5"` as `5` (decimal text only), `"true"` as `true`, `5` as `"5"` for a string, and `null` as `""`, `0`, or `false` for a required field. A `null` for an optional field that does not allow `null` is removed. The repair follows unions, `$ref`, `if`/`then`/`else`, `prefixItems`, `patternProperties`, and `additionalProperties`. +- A Standard Schema runs first as written. When it fails and the schema can export JSON Schema (Zod, ArkType, Valibot), the input is repaired and the schema runs again. A failure reports the issues of the first run. +- The received arguments do not change. The tool gets the repaired copy. + +Behavior change: a call that failed the check only because of these mistakes now runs with the repaired input. diff --git a/docs/tools/server-tools.md b/docs/tools/server-tools.md index 4866bec070..0271174d18 100644 --- a/docs/tools/server-tools.md +++ b/docs/tools/server-tools.md @@ -363,6 +363,20 @@ The `inputSchema` checks the input just before the tool runs: - An async validator is awaited. - The tool gets the schema's output, so a transform such as `z.string().transform(Number)` applies once. +Models often send a value with the wrong type, such as `"5"` for a number. The check repairs these against the schema before it fails the call: + +| The model sends | The schema wants | The tool gets | +| --- | --- | --- | +| `"5"` | `number` or `integer` | `5` | +| `"true"` | `boolean` | `true` | +| `5` | `string` | `"5"` | +| `null` for an optional field | a type with no `null` | no field | +| `null` for a required field | `string`, `number`, or `boolean` | `""`, `0`, or `false` | + +- A JSON Schema tool, and an MCP tool, gets the repair when its input fails the check. A valid input runs as it is. +- A Standard Schema that can export JSON Schema (Zod, ArkType, Valibot) runs first as written. If it fails, the input is repaired and the schema runs again. If it still fails, the error shows the issues of the first run. +- An input that cannot be repaired, such as `"many"` for a number, becomes a tool error. + ## Using JSON Schema If you have existing JSON Schema definitions or prefer not to use Zod, you can define tool schemas using raw JSON Schema objects: diff --git a/packages/ai/src/activities/chat/tools/input-validation.ts b/packages/ai/src/activities/chat/tools/input-validation.ts index f1657e77fd..95d859eb9d 100644 --- a/packages/ai/src/activities/chat/tools/input-validation.ts +++ b/packages/ai/src/activities/chat/tools/input-validation.ts @@ -1,14 +1,23 @@ -import { Validator } from '@cfworker/json-schema' +import { dereference, validate } from '@cfworker/json-schema' import { StandardSchemaValidationError, convertSchemaToJsonSchema, + isStandardJSONSchema, isStandardSchema, validateWithStandardSchema, } from './schema-converter' +import type { StandardSchemaV1 } from '@standard-schema/spec' import type { Schema, SchemaDraft } from '@cfworker/json-schema' import type { JSONSchema, SchemaInput } from '../../../types' -const validators = new WeakMap() +/** A JSON Schema ready for `validate`: the copy, its draft, its `$ref` lookup. */ +interface CompiledSchema { + schema: Schema | boolean + draft: SchemaDraft + lookup: Record +} + +const compiled = new WeakMap() /** * The draft named by `$schema`. Default: 2020-12. A schema with no `$schema` @@ -23,11 +32,13 @@ function draftOf(schema: JSONSchema, json: string): SchemaDraft { return '2020-12' } -function validatorFor(schema: JSONSchema | boolean) { - if (typeof schema === 'boolean') return new Validator(schema) - const cached = validators.get(schema) +function compile(schema: JSONSchema | boolean): CompiledSchema { + if (typeof schema === 'boolean') { + return { schema, draft: '2020-12', lookup: dereference(schema) } + } + const cached = compiled.get(schema) if (cached) return cached - // The validator writes lookup data onto the schema it gets. A JSON copy + // `dereference` writes lookup data onto the schema it gets. A JSON copy // keeps the tool's own schema, which also goes to the provider, unchanged. // `format` is an annotation, as in 2019-09 and later: a tool that declares // `format: 'uri'` still gets a plain path. @@ -35,9 +46,229 @@ function validatorFor(schema: JSONSchema | boolean) { const copy: Schema = JSON.parse(json, (key, value) => key === 'format' && typeof value === 'string' ? undefined : value, ) - const validator = new Validator(copy, draftOf(schema, json), false) - validators.set(schema, validator) - return validator + const result = { + schema: copy, + draft: draftOf(schema, json), + lookup: dereference(copy), + } + compiled.set(schema, result) + return result +} + +const DECIMAL = /^-?\d+(?:\.\d+)?(?:[eE][+-]?\d+)?$/ + +function convertPrimitive(value: unknown, type: string) { + switch (type) { + case 'number': + case 'integer': { + if (value === null) return 0 + if (typeof value === 'boolean') return value ? 1 : 0 + // Decimal text only: no hex, and no integer past the safe range. + if (typeof value !== 'string' || !DECIMAL.test(value)) return value + const number = Number(value) + const valid = + type === 'integer' + ? Number.isSafeInteger(number) + : Number.isFinite(number) + return valid ? number : value + } + case 'boolean': + if (value === null || value === 'false' || value === 0) return false + if (value === 'true' || value === 1) return true + return value + case 'string': + if (value === null) return '' + return typeof value === 'number' || typeof value === 'boolean' + ? String(value) + : value + case 'null': + return value === '' || value === 0 || value === false ? null : value + default: + return value + } +} + +function matchesType(value: unknown, type: string) { + switch (type) { + case 'null': + return value === null + case 'array': + return Array.isArray(value) + case 'object': + return ( + typeof value === 'object' && value !== null && !Array.isArray(value) + ) + case 'integer': + return typeof value === 'number' && Number.isInteger(value) + default: + return typeof value === type + } +} + +/** + * Repair common model mistakes against the schema before the check: `"5"` + * for 5, `"true"` for true, `null` for an optional field. It walks unions, + * `$ref`, `if`/`then`/`else`, tuples, and pattern and additional properties. + * A value that cannot be repaired stays as it is, so the check reports it. + */ +function coerceArguments(value: unknown, root: CompiledSchema) { + const { draft, lookup } = root + + function check(schema: Schema | boolean, candidate: unknown) { + if (typeof schema === 'boolean') return schema + return validate(candidate, schema, draft, lookup).valid + } + + function refTarget(schema: Schema) { + const uri = schema.__absolute_ref__ + const target = uri === undefined ? undefined : lookup[uri] + return typeof target === 'object' ? target : undefined + } + + function coerceUnion( + candidate: unknown, + schemas: Array, + seen: Set, + required: Set, + ) { + if (schemas.some((schema) => check(schema, candidate))) return candidate + for (const schema of schemas) { + const converted = coerce( + structuredClone(candidate), + schema, + seen, + required, + ) + if (check(schema, converted)) return converted + } + return candidate + } + + function requiredKeys( + schema: Schema | boolean, + candidate: unknown, + seen = new Set(), + ): Set { + if (typeof schema === 'boolean' || seen.has(schema)) return new Set() + const nextSeen = new Set(seen).add(schema) + const required = new Set( + Array.isArray(schema.required) ? schema.required : [], + ) + const children: Array = [...(schema.allOf ?? [])] + const target = refTarget(schema) + if (target) children.push(target) + if (schema.if !== undefined) { + const branch = check(schema.if, candidate) ? schema.then : schema.else + if (branch) children.push(branch) + } + for (const child of children) { + for (const key of requiredKeys(child, candidate, nextSeen)) { + required.add(key) + } + } + return required + } + + function coerce( + candidate: unknown, + schema: Schema | boolean, + seen = new Set(), + inheritedRequired = new Set(), + ): unknown { + if (typeof schema === 'boolean') return candidate + // A reference cycle at the same value cannot add another conversion. + if (seen.has(schema)) return candidate + const nextSeen = new Set(seen).add(schema) + const required = requiredKeys(schema, candidate) + for (const key of inheritedRequired) required.add(key) + let next = candidate + const target = refTarget(schema) + if (target) next = coerce(next, target, nextSeen, required) + for (const nested of schema.allOf ?? []) { + next = coerce(next, nested, nextSeen, required) + } + if (schema.anyOf) next = coerceUnion(next, schema.anyOf, nextSeen, required) + if (schema.oneOf) next = coerceUnion(next, schema.oneOf, nextSeen, required) + + const types = + typeof schema.type === 'string' ? [schema.type] : (schema.type ?? []) + if (!types.some((type) => matchesType(next, type))) { + // `null` last: `0` for `['null', 'string']` is `"0"`, not `null`. + const order = types.filter((type) => type !== 'null') + if (types.includes('null')) order.push('null') + for (const type of order) { + const converted = convertPrimitive(next, type) + if (converted !== next) { + next = converted + break + } + } + } + + if (typeof next === 'object' && next !== null && !Array.isArray(next)) { + // Own keys only, `__proto__` included. + const object: Record = Object.fromEntries( + Object.entries(next), + ) + const properties = Object.entries(schema.properties ?? {}) + const knownKeys = new Set(properties.map(([key]) => key)) + for (const [key, property] of properties) { + if (!Object.hasOwn(object, key)) continue + if ( + object[key] === null && + !required.has(key) && + !check(property, null) + ) { + // A strict-mode provider sends `null` for an optional field. + delete object[key] + } else { + object[key] = coerce(object[key], property) + } + } + for (const [pattern, property] of Object.entries( + schema.patternProperties ?? {}, + )) { + const expression = new RegExp(pattern, 'u') + for (const key of Object.keys(object)) { + if (!expression.test(key)) continue + knownKeys.add(key) + object[key] = coerce(object[key], property) + } + } + const additional = schema.additionalProperties + if (typeof additional === 'object') { + for (const [key, entry] of Object.entries(object)) { + if (!knownKeys.has(key)) object[key] = coerce(entry, additional) + } + } + next = object + } + if (Array.isArray(next)) { + const tuple = + schema.prefixItems ?? + (Array.isArray(schema.items) ? schema.items : undefined) + const rest = Array.isArray(schema.items) ? undefined : schema.items + next = next.map((item: unknown, index) => { + const itemSchema = tuple?.[index] ?? rest + return itemSchema === undefined ? item : coerce(item, itemSchema) + }) + } + if (schema.if !== undefined) { + const branch = check(schema.if, next) ? schema.then : schema.else + if (branch) next = coerce(next, branch, nextSeen, required) + } + return next + } + + return coerce(value, root.schema) +} + +function receivedArgumentsText(received: unknown) { + try { + return JSON.stringify(received, null, 2) ?? String(received) + } catch { + return '[Arguments cannot be serialized as JSON]' + } } /** Errors that only say that a nested check failed. The nested error says why. */ @@ -50,14 +281,6 @@ const WRAPPER_KEYWORDS = new Set([ 'allOf', ]) -function receivedArgumentsText(received: unknown) { - try { - return JSON.stringify(received, null, 2) ?? String(received) - } catch { - return '[Arguments cannot be serialized as JSON]' - } -} - function validateJsonInput( schema: JSONSchema | boolean, received: unknown, @@ -66,10 +289,23 @@ function validateJsonInput( let errors: string try { // Check the JSON form: a key set to `undefined` is the same as no key. - const result = validatorFor(schema).validate( - JSON.parse(JSON.stringify(received)), - ) + const json = JSON.stringify(received) + const value: unknown = JSON.parse(json) + const root = compile(schema) + const check = (candidate: unknown) => + validate(candidate, root.schema, root.draft, root.lookup, false) + const result = check(value) if (result.valid) return received + // Repair only an input that fails. The error reports the input as sent. + let converted = value + try { + converted = coerceArguments(value, root) + } catch { + // A schema that cannot be used: the first check reports it. + } + if (JSON.stringify(converted) !== json && check(converted).valid) { + return converted + } const leaves = result.errors.filter( (error) => !WRAPPER_KEYWORDS.has(error.keyword), ) @@ -87,12 +323,23 @@ function validateJsonInput( ) } +function standardInputError( + issues: ReadonlyArray, + received: unknown, + toolName: string, +) { + const error = new StandardSchemaValidationError(issues) + error.message = `Input validation failed for tool ${toolName}: ${error.message}\nReceived arguments:\n${receivedArgumentsText(received)}` + return error +} + /** * The final check of a tool input, just before the tool runs or the client * gets the call. A Standard Schema runs here (an async validator too), and * its output is the input, so transforms apply once. A JSON Schema (an MCP * tool, a plain JSON Schema tool) is checked with `@cfworker/json-schema`. - * The error goes back to the model as the tool result. + * Common model mistakes are repaired against the schema first. The error + * goes back to the model as the tool result. */ export async function validateToolInput( schema: SchemaInput | boolean | undefined, @@ -109,9 +356,24 @@ export async function validateToolInput( ? validateJsonInput(jsonSchema, received, toolName) : received } - const validation = await validateWithStandardSchema(schema, received) - if (validation.success) return validation.data - const error = new StandardSchemaValidationError(validation.issues) - error.message = `Input validation failed for tool ${toolName}: ${error.message}\nReceived arguments:\n${receivedArgumentsText(received)}` - throw error + const first = await validateWithStandardSchema(schema, received) + if (first.success) return first.data + // The authored schema runs first. When it fails and it can export JSON + // Schema, repair the input and run the authored schema again. A failure + // reports the issues of the first run. + const original = standardInputError(first.issues, received, toolName) + if (!isStandardJSONSchema(schema)) throw original + let converted: unknown + try { + const jsonSchema = convertSchemaToJsonSchema(schema, { io: 'input' }) + if (!jsonSchema) throw original + converted = validateJsonInput(jsonSchema, received, toolName) + } catch { + throw original + } + // Nothing was repaired: a second run gives the same issues. + if (converted === received) throw original + const retried = await validateWithStandardSchema(schema, converted) + if (!retried.success) throw original + return retried.data } diff --git a/packages/ai/tests/tool-input-validation.test.ts b/packages/ai/tests/tool-input-validation.test.ts index 938186fa2c..b5e7346def 100644 --- a/packages/ai/tests/tool-input-validation.test.ts +++ b/packages/ai/tests/tool-input-validation.test.ts @@ -4,6 +4,7 @@ import { chat } from '../src/activities/chat/index' import { defineChatMiddleware } from '../src/activities/chat/middleware/define' import { defineInterrupt } from '../src/interrupt-definition' import { executeToolCalls } from '../src/activities/chat/tools/tool-calls' +import { validateToolInput } from '../src/activities/chat/tools/input-validation' import { otelMiddleware } from '../src/middlewares/otel' import { EventType } from '../src/types' import { collectChunks, createMockAdapter, ev } from './test-utils' @@ -315,19 +316,19 @@ describe('JSON Schema tool input', () => { expect(result.executed).toEqual([{ count: 2 }]) }) - it.each(['{}', '{"count":"2"}', '{"count":0}', '{"count":2,"extra":true}'])( - 'rejects invalid input without executing: %s', - async (raw) => { - const result = await runTool(raw, jsonSchema) - expect(result.executed).toEqual([]) - expect(result.results[0]?.state).toBe('output-error') - expect(result.results[0]?.result).toEqual({ - error: expect.stringContaining( - 'Input validation failed for tool check', - ), - }) - }, - ) + it.each([ + '{}', + '{"count":"wrong"}', + '{"count":0}', + '{"count":2,"extra":true}', + ])('rejects invalid input without executing: %s', async (raw) => { + const result = await runTool(raw, jsonSchema) + expect(result.executed).toEqual([]) + expect(result.results[0]?.state).toBe('output-error') + expect(result.results[0]?.result).toEqual({ + error: expect.stringContaining('Input validation failed for tool check'), + }) + }) it('names the failed location and shows the received arguments', async () => { const result = await runTool('{"count":0}', jsonSchema) @@ -354,7 +355,7 @@ describe('JSON Schema tool input', () => { (await runTool('{"mode":"count","count":3}', schema)).executed, ).toEqual([{ mode: 'count', count: 3 }]) expect((await runTool('{"mode":"count"}', schema)).executed).toEqual([]) - expect((await runTool('{"count":"3"}', schema)).executed).toEqual([]) + expect((await runTool('{"count":"x"}', schema)).executed).toEqual([]) }) it('uses the draft named by $schema', async () => { @@ -478,13 +479,201 @@ describe('JSON Schema tool input', () => { expect((await runTool('{"count":2}', schema)).executed).toEqual([ { count: 2 }, ]) - expect((await runTool('{"count":"2"}', schema)).executed).toEqual([]) + expect((await runTool('{"count":"2"}', schema)).executed).toEqual([ + { count: 2 }, + ]) + expect((await runTool('{"count":"x"}', schema)).executed).toEqual([]) } finally { vi.unstubAllGlobals() } }) }) +describe('tool input repair', () => { + const schema: JSONSchema = { + type: 'object', + properties: { count: { type: 'integer', minimum: 1 } }, + required: ['count'], + additionalProperties: false, + } + + it('coerces a numeric field before execution', async () => { + const result = await runTool('{"count":"2"}', schema) + expect(result.executed).toEqual([{ count: 2 }]) + }) + + it('does not change the received arguments', async () => { + const received = { count: '2' } + expect(await validateToolInput(schema, received, 'check')).toEqual({ + count: 2, + }) + expect(received).toEqual({ count: '2' }) + }) + + it('shows the original arguments when the repaired value still fails', async () => { + const result = await runTool('{"count":"0"}', schema) + expect(result.executed).toEqual([]) + expect(result.results[0]?.result).toEqual({ + error: expect.stringContaining('"count": "0"'), + }) + }) + + it('follows local references and drops a null optional field', async () => { + const result = await runTool('{"count":"3","label":null,"omitted":null}', { + type: 'object', + $defs: { count: { type: 'integer', minimum: 1 } }, + properties: { + count: { $ref: '#/$defs/count' }, + label: { type: ['string', 'null'] }, + omitted: { type: 'string' }, + }, + required: ['count'], + }) + expect(result.executed).toEqual([{ count: 3, label: null }]) + }) + + it.each<[unknown, JSONSchema, unknown]>([ + ['12', { anyOf: [{ type: 'integer' }, { type: 'string' }] }, '12'], + ['12', { type: ['integer', 'string'] }, '12'], + ['3', { type: 'integer' }, 3], + [['1', '2'], { type: 'array', items: { type: 'integer' } }, [1, 2]], + [ + ['1', 'x'], + { prefixItems: [{ type: 'integer' }, { type: 'string' }] }, + [1, 'x'], + ], + ['true', { type: 'boolean' }, true], + [null, { type: 'string' }, ''], + ])('repairs %j against %j', async (received, inputSchema, expected) => { + expect(await validateToolInput(inputSchema, received, 'check')).toEqual( + expected, + ) + }) + + it('rejects a primitive that cannot be repaired', async () => { + await expect( + validateToolInput({ type: 'integer' }, 'wrong', 'check'), + ).rejects.toThrow('Input validation failed for tool check') + }) + + it('keeps a valid referenced union branch', async () => { + const result = await runTool('{"count":"12"}', { + type: 'object', + $defs: { numeric: { type: 'integer' }, text: { type: 'string' } }, + properties: { + count: { + anyOf: [{ $ref: '#/$defs/numeric' }, { $ref: '#/$defs/text' }], + }, + }, + required: ['count'], + }) + expect(result.executed).toEqual([{ count: '12' }]) + }) + + it('coerces conditional branch fields', async () => { + const result = await runTool('{"mode":"count","count":"4"}', { + type: 'object', + if: { properties: { mode: { const: 'count' } }, required: ['mode'] }, + then: { + properties: { count: { type: 'integer' } }, + required: ['count'], + }, + }) + expect(result.executed).toEqual([{ mode: 'count', count: 4 }]) + }) + + it('keeps a null field that a conditional branch requires', async () => { + const result = await runTool('{"mode":"count","count":null}', { + type: 'object', + properties: { mode: { type: 'string' }, count: { type: 'integer' } }, + if: { properties: { mode: { const: 'count' } }, required: ['mode'] }, + then: { required: ['count'] }, + }) + expect(result.executed).toEqual([{ mode: 'count', count: 0 }]) + }) + + it('keeps a null field that an intersection requires', async () => { + const result = await runTool('{"count":null}', { + type: 'object', + allOf: [ + { properties: { count: { type: 'integer' } } }, + { required: ['count'] }, + ], + }) + expect(result.executed).toEqual([{ count: 0 }]) + }) + + it('keeps a null field required outside a union branch', async () => { + const result = await runTool('{"count":null}', { + type: 'object', + required: ['count'], + anyOf: [{ properties: { count: { type: 'integer' } } }], + }) + expect(result.executed).toEqual([{ count: 0 }]) + }) + + it('coerces pattern and additional properties', async () => { + const result = await runTool('{"n_1":"4","extra":"5"}', { + type: 'object', + patternProperties: { '^n_': { type: 'integer' } }, + additionalProperties: { type: 'number' }, + }) + expect(result.executed).toEqual([{ n_1: 4, extra: 5 }]) + }) + + it('coerces an approval edit before execution', async () => { + const result = await runTool('{"count":2}', schema, { + needsApproval: true, + approvals: new Map([ + ['call-1', { approved: true, editedArgs: { count: '4' } }], + ]), + }) + expect(result.executed).toEqual([{ count: 4 }]) + }) + + it('coerces a middleware rewrite before a client request', async () => { + const result = await runTool('{"count":2}', schema, { + client: true, + middleware: { + onBeforeToolCall: async () => ({ + type: 'transformArgs', + args: { count: '5' }, + }), + }, + }) + expect(result.needsClientExecution).toEqual([ + { toolCallId: 'call-1', toolName: 'check', input: { count: 5 } }, + ]) + }) + + it('repairs input for a Standard Schema that exports JSON Schema', async () => { + const result = await runTool( + '{"count":"2","label":null}', + z.object({ count: z.number(), label: z.string().optional() }), + ) + expect(result.executed).toEqual([{ count: 2 }]) + }) + + it('runs the authored transform on the repaired input', async () => { + const result = await runTool( + '{"count":"2"}', + z.object({ count: z.number().transform((value) => value * 10) }), + ) + expect(result.executed).toEqual([{ count: 20 }]) + }) + + it('reports the issues of the first run when repair does not help', async () => { + const result = await runTool( + '{"count":"wrong"}', + z.object({ count: z.number() }), + ) + expect(result.executed).toEqual([]) + expect(result.results[0]?.result).toEqual({ + error: expect.stringContaining('"count": "wrong"'), + }) + }) +}) + const clientCheck = ( inputSchema: SchemaInput, needsApproval = false, diff --git a/testing/e2e/src/routes/api.json-schema-tool-input.ts b/testing/e2e/src/routes/api.json-schema-tool-input.ts index fe5a21b4b4..044e547174 100644 --- a/testing/e2e/src/routes/api.json-schema-tool-input.ts +++ b/testing/e2e/src/routes/api.json-schema-tool-input.ts @@ -3,17 +3,26 @@ import { chat, toolDefinition } from '@tanstack/ai' import { fakeText } from '@tanstack/ai/testing' // No network: `fakeText` answers from a script. The first answer calls a -// JSON Schema tool with a wrong type, the second answer repeats the tool -// result. MCP tools have the same raw JSON Schema input. +// JSON Schema tool twice: once with a value that cannot be repaired, once +// with a number sent as a string. The second answer repeats the tool +// results. MCP tools have the same raw JSON Schema input. export const Route = createFileRoute('/api/json-schema-tool-input')({ server: { handlers: { POST: async () => { const fake = fakeText() fake.setResponses([ - { toolCalls: [{ name: 'count_items', input: { count: 'many' } }] }, + { + toolCalls: [ + { name: 'count_items', input: { count: 'many' } }, + { name: 'count_items', input: { count: '3' } }, + ], + }, ({ request }) => ({ - text: `Tool said: ${String(request.messages.at(-1)?.content)}`, + text: `Tool said: ${request.messages + .filter((message) => message.role === 'tool') + .map((message) => String(message.content)) + .join(' | ')}`, }), ]) diff --git a/testing/e2e/src/routes/api.middleware-test.ts b/testing/e2e/src/routes/api.middleware-test.ts index 989a91f817..d52d183451 100644 --- a/testing/e2e/src/routes/api.middleware-test.ts +++ b/testing/e2e/src/routes/api.middleware-test.ts @@ -183,7 +183,7 @@ const toolBadArgsMiddleware: ChatMiddleware = { name: 'tool-bad-args', onBeforeToolCall(_ctx, hookCtx) { if (hookCtx.toolName !== 'get_weather') return undefined - return { type: 'transformArgs' as const, args: { city: 42 } } + return { type: 'transformArgs' as const, args: { city: { name: 'NYC' } } } }, } diff --git a/testing/e2e/src/routes/api.tools-test.ts b/testing/e2e/src/routes/api.tools-test.ts index 3e40bc8957..43f5815264 100644 --- a/testing/e2e/src/routes/api.tools-test.ts +++ b/testing/e2e/src/routes/api.tools-test.ts @@ -65,9 +65,10 @@ function createProviderFreeAdapter(scenario: string): AnyTextAdapter { : scenario === 'client-tool-input-error' || scenario === 'invalid-client-tool-retry' ? { - arguments: '{"message":42,"type":"info"}', + // An object cannot be repaired into the string `message`. + arguments: '{"message":{"text":"hi"},"type":"info"}', initialText: 'Showing a notification.', - input: { message: 42, type: 'info' }, + input: { message: { text: 'hi' }, type: 'info' }, name: scenario === 'invalid-client-tool-retry' ? 'invalid-client-tool-retry-test' diff --git a/testing/e2e/tests/json-schema-tool-input.spec.ts b/testing/e2e/tests/json-schema-tool-input.spec.ts index 101a1fdbe8..6eda7f40fe 100644 --- a/testing/e2e/tests/json-schema-tool-input.spec.ts +++ b/testing/e2e/tests/json-schema-tool-input.spec.ts @@ -1,13 +1,15 @@ import { test, expect } from './fixtures' -test('a JSON Schema tool with invalid input returns a tool error and does not run', async ({ +test('a JSON Schema tool repairs "3" to 3, and returns a tool error for input it cannot repair', async ({ request, }) => { const response = await request.post('/api/json-schema-tool-input') expect(response.ok()).toBe(true) const body = await response.json() - expect(body.runs).toEqual([]) + // Only the repaired call ran. + expect(body.runs).toEqual([{ count: 3 }]) expect(body.text).toContain('Input validation failed for tool count_items') expect(body.text).toContain('#/count') + expect(body.text).toContain('counted') }) diff --git a/testing/e2e/tests/middleware.spec.ts b/testing/e2e/tests/middleware.spec.ts index f527393a81..faa976e1e4 100644 --- a/testing/e2e/tests/middleware.spec.ts +++ b/testing/e2e/tests/middleware.spec.ts @@ -160,7 +160,7 @@ test.describe('Middleware Lifecycle', () => { const toolResults = messages.flatMap((m: any) => m.parts.filter((p: any) => p.type === 'tool-result'), ) - // The middleware rewrote { city: 'NYC' } to { city: 42 }. The schema + // The middleware rewrote { city: 'NYC' } to { city: { name: 'NYC' } }. The schema // checks the rewritten args, so get_weather never runs. expect(toolResults).toHaveLength(1) expect(toolResults[0].content).toContain('Input validation failed') From 3dcb866a244459dbef52f632599cb6cccbec5c32 Mon Sep 17 00:00:00 2001 From: Alem Tuzlak Date: Tue, 13 Oct 2026 21:20:25 +0200 Subject: [PATCH 06/16] fix(ai): type the repair order list --- packages/ai/src/activities/chat/tools/input-validation.ts | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/packages/ai/src/activities/chat/tools/input-validation.ts b/packages/ai/src/activities/chat/tools/input-validation.ts index 95d859eb9d..7964e812ce 100644 --- a/packages/ai/src/activities/chat/tools/input-validation.ts +++ b/packages/ai/src/activities/chat/tools/input-validation.ts @@ -194,7 +194,7 @@ function coerceArguments(value: unknown, root: CompiledSchema) { typeof schema.type === 'string' ? [schema.type] : (schema.type ?? []) if (!types.some((type) => matchesType(next, type))) { // `null` last: `0` for `['null', 'string']` is `"0"`, not `null`. - const order = types.filter((type) => type !== 'null') + const order: Array = types.filter((type) => type !== 'null') if (types.includes('null')) order.push('null') for (const type of order) { const converted = convertPrimitive(next, type) From 68de395e00c8a8f1179956cfd7284e6fefd1624f Mon Sep 17 00:00:00 2001 From: Alem Tuzlak Date: Tue, 13 Oct 2026 21:20:06 +0200 Subject: [PATCH 07/16] feat(ai, adapters): keep the raw tool argument string for replay --- .changeset/raw-tool-arguments.md | 17 + docs/chat/stream-events.md | 2 +- packages/ai-anthropic/src/adapters/text.ts | 27 +- .../tests/anthropic-adapter.test.ts | 94 ++++++ .../src/converse/stream-processor.ts | 8 + .../tests/converse/stream-processor.test.ts | 63 ++++ packages/ai-gemini/src/adapters/text.ts | 9 +- .../experimental/text-interactions/adapter.ts | 39 ++- .../ai-gemini/tests/gemini-adapter.test.ts | 77 +++++ .../tests/text-interactions-adapter.test.ts | 97 +++++- packages/ai-mistral/src/adapters/text.ts | 21 +- .../ai-mistral/tests/mistral-adapter.test.ts | 103 ++++++ packages/ai-ollama/src/adapters/text.ts | 8 +- packages/ai-ollama/tests/text-adapter.test.ts | 47 +++ .../src/adapters/responses-text.ts | 15 +- packages/ai-openrouter/src/adapters/text.ts | 10 +- .../tests/openrouter-adapter.test.ts | 48 +++ .../openrouter-responses-adapter.test.ts | 102 ++++++ packages/ai/src/activities/chat/index.ts | 20 ++ .../src/activities/chat/stream/processor.ts | 17 +- .../src/activities/chat/tools/tool-calls.ts | 31 +- packages/ai/src/types.ts | 2 + .../ai/src/utilities/tool-call-arguments.ts | 309 ++++++++++++++++++ packages/ai/tests/tool-call-arguments.test.ts | 173 ++++++++++ .../src/adapters/chat-completions-text.ts | 7 +- .../src/adapters/responses-text.ts | 15 +- .../tests/chat-completions-text.test.ts | 61 +++- .../openai-base/tests/responses-text.test.ts | 102 ++++++ .../middleware-test/raw-tool-args.json | 27 ++ testing/e2e/src/routes/api.middleware-test.ts | 2 +- testing/e2e/src/routes/middleware-test.tsx | 1 + testing/e2e/tests/middleware.spec.ts | 57 ++++ 32 files changed, 1543 insertions(+), 68 deletions(-) create mode 100644 .changeset/raw-tool-arguments.md create mode 100644 packages/ai/src/utilities/tool-call-arguments.ts create mode 100644 packages/ai/tests/tool-call-arguments.test.ts create mode 100644 testing/e2e/fixtures/middleware-test/raw-tool-args.json diff --git a/.changeset/raw-tool-arguments.md b/.changeset/raw-tool-arguments.md new file mode 100644 index 0000000000..9911077c5d --- /dev/null +++ b/.changeset/raw-tool-arguments.md @@ -0,0 +1,17 @@ +--- +'@tanstack/ai': minor +'@tanstack/openai-base': patch +'@tanstack/ai-openrouter': patch +'@tanstack/ai-bedrock': patch +'@tanstack/ai-mistral': patch +'@tanstack/ai-gemini': patch +'@tanstack/ai-ollama': patch +'@tanstack/ai-anthropic': patch +--- + +Keep the provider's raw tool argument string for replay. + +- Adapters now send the raw argument string as `args` on `TOOL_CALL_END`. +- The tool call in the history keeps that exact string when it is the same JSON value as the parsed input, so the next request sends the provider's own bytes back (for example `1.50`, or an id past `Number.MAX_SAFE_INTEGER`). When the parsed input differs, only the changed values are written again. +- The `TOOL_CALL_END` that reaches the client carries the same string in `metadata.tanstack.args`, so the client history matches the server history. `input` is sent only when JSON keeps it exactly. +- Malformed tool arguments are no longer turned into `{}` by an adapter. The tool gets a "Failed to parse tool arguments as JSON" error that the model can repair. Mistral still ends the run with an error on malformed arguments, as before. diff --git a/docs/chat/stream-events.md b/docs/chat/stream-events.md index a88110295d..432a4c8ac0 100644 --- a/docs/chat/stream-events.md +++ b/docs/chat/stream-events.md @@ -238,7 +238,7 @@ The media generation hooks take a `threadId` too. There it names a slot, not a c ## Tool input and output -SSE and HTTP `TOOL_CALL_END` does not carry parsed `input`. In-process `chat()` still has `input`. Tool input and output also live on `UIMessage` parts. +SSE and HTTP `TOOL_CALL_END` does not carry parsed `input` at the top level. In-process `chat()` still has `input`. `metadata.tanstack.args` carries the tool argument string as the provider sent it. Tool input and output also live on `UIMessage` parts. On the server, feed chunks into `StreamProcessor`. On the client, read `useChat` `messages`. diff --git a/packages/ai-anthropic/src/adapters/text.ts b/packages/ai-anthropic/src/adapters/text.ts index d850bbe1f3..63e08efd6b 100644 --- a/packages/ai-anthropic/src/adapters/text.ts +++ b/packages/ai-anthropic/src/adapters/text.ts @@ -1512,7 +1512,13 @@ export class AnthropicTextAdapter< let accumulatedSignature = '' const toolCallsMap = new Map< number, - { id: string; name: string; input: string; started: boolean } + { + id: string + name: string + input: string + started: boolean + hasInputDelta: boolean + } >() let currentToolIndex = -1 // Server-side tools share the `input_json_delta` wire format with client @@ -1577,8 +1583,11 @@ export class AnthropicTextAdapter< toolCallsMap.set(currentToolIndex, { id: event.content_block.id, name: event.content_block.name, - input: '', + // The initial input snapshot. The first `input_json_delta` + // replaces it. + input: JSON.stringify(event.content_block.input) ?? '', started: false, + hasInputDelta: false, }) } else if (event.content_block.type === 'server_tool_use') { currentServerTool = { @@ -1851,6 +1860,10 @@ export class AnthropicTextAdapter< } } + if (!existing.hasInputDelta) { + existing.input = '' + existing.hasInputDelta = true + } existing.input += event.delta.partial_json yield { @@ -1932,13 +1945,14 @@ export class AnthropicTextAdapter< } } - // Emit TOOL_CALL_END - let parsedInput: unknown = {} + // Emit TOOL_CALL_END. Malformed JSON sends no `input`, so core + // reports the error from the raw `args`. + let parsedInput: unknown try { const parsed = existing.input ? JSON.parse(existing.input) : {} parsedInput = parsed && typeof parsed === 'object' ? parsed : {} } catch { - parsedInput = {} + // Keep `parsedInput` undefined for the tool error. } yield { @@ -1948,7 +1962,8 @@ export class AnthropicTextAdapter< toolName: existing.name, model, timestamp: Date.now(), - input: parsedInput, + ...(existing.input !== '' && { args: existing.input }), + ...(parsedInput !== undefined && { input: parsedInput }), } // Reset so a new TEXT_MESSAGE_START is emitted if text follows tool calls diff --git a/packages/ai-anthropic/tests/anthropic-adapter.test.ts b/packages/ai-anthropic/tests/anthropic-adapter.test.ts index 86e950aa5a..8b26d21658 100644 --- a/packages/ai-anthropic/tests/anthropic-adapter.test.ts +++ b/packages/ai-anthropic/tests/anthropic-adapter.test.ts @@ -12,6 +12,7 @@ import { AnthropicTextAdapter } from '../src/adapters/text' import type { AnthropicTextProviderOptions } from '../src/adapters/text' import type { AnthropicDocumentMetadata } from '../src/message-types' import { ANTHROPIC_MAX_NONSTREAMING_TOKENS } from '../src/model-meta' +import { createSilentLogger } from './utils/logger' import { z } from 'zod' const mocks = vi.hoisted(() => { @@ -2520,6 +2521,99 @@ describe('Anthropic stream processing', () => { 'text B', ]) }) + + /** The TOOL_CALL_END of one streamed `tool_use` block. */ + async function toolCallEnd( + initialInput: Record | undefined, + deltas: Array, + ) { + mocks.betaMessagesCreate.mockResolvedValueOnce( + (async function* () { + yield { + type: 'content_block_start', + index: 0, + content_block: { + type: 'tool_use', + id: 'toolu_raw', + name: 'inspect', + ...(initialInput !== undefined && { input: initialInput }), + }, + } + for (const partial_json of deltas) { + yield { + type: 'content_block_delta', + index: 0, + delta: { type: 'input_json_delta', partial_json }, + } + } + yield { type: 'content_block_stop', index: 0 } + yield { + type: 'message_delta', + delta: { stop_reason: 'tool_use' }, + usage: { output_tokens: 1 }, + } + yield { type: 'message_stop' } + })(), + ) + + const chunks: Array = [] + for await (const chunk of createAdapter('claude-opus-4-1').chatStream({ + model: 'claude-opus-4-1', + messages: [{ role: 'user', content: 'Go' }], + logger: createSilentLogger(), + })) { + chunks.push(chunk) + } + return chunks.find((chunk) => chunk.type === 'TOOL_CALL_END') + } + + it.each([ + { + name: 'the streamed input', + initial: {}, + deltas: ['{"count":2}'], + args: '{"count":2}', + input: { count: 2 }, + }, + { + name: 'the streamed bytes as sent', + initial: {}, + deltas: ['{"n": ', '1.50}'], + args: '{"n": 1.50}', + input: { n: 1.5 }, + }, + { + name: 'the initial input snapshot when no delta arrives', + initial: { value: 'start' }, + deltas: [], + args: '{"value":"start"}', + input: { value: 'start' }, + }, + { + name: 'the streamed input over the initial snapshot', + initial: { value: 'start' }, + deltas: ['{"count":2}'], + args: '{"count":2}', + input: { count: 2 }, + }, + ])( + 'sends $name as the raw TOOL_CALL_END args', + async ({ initial, deltas, args, input }) => { + expect(await toolCallEnd(initial, deltas)).toMatchObject({ args, input }) + }, + ) + + it('sends malformed tool input raw on TOOL_CALL_END with no input', async () => { + const end = await toolCallEnd({}, ['{"a":']) + expect(end).toMatchObject({ args: '{"a":' }) + expect(end).not.toHaveProperty('input') + }) + + it('sends empty input as {} with no raw args', async () => { + const end = await toolCallEnd(undefined, []) + expect(end).toMatchObject({ input: {} }) + expect(end).not.toHaveProperty('args') + }) }) describe('Anthropic adapter error handling', () => { diff --git a/packages/ai-bedrock/src/converse/stream-processor.ts b/packages/ai-bedrock/src/converse/stream-processor.ts index 911d90d7e3..875d864a5c 100644 --- a/packages/ai-bedrock/src/converse/stream-processor.ts +++ b/packages/ai-bedrock/src/converse/stream-processor.ts @@ -72,6 +72,8 @@ export async function* processConverseStream( id: string name: string started: boolean + arguments: string + hasArguments: boolean } >() @@ -167,6 +169,8 @@ export async function* processConverseStream( id, name, started: true, + arguments: '', + hasArguments: false, }) yield { type: EventType.TOOL_CALL_START, @@ -189,6 +193,8 @@ export async function* processConverseStream( if (delta && 'toolUse' in delta && delta.toolUse?.input !== undefined) { const toolCall = toolCallsByIndex.get(index) if (toolCall?.started) { + toolCall.arguments += delta.toolUse.input + toolCall.hasArguments = true yield { type: EventType.TOOL_CALL_ARGS, toolCallId: toolCall.id, @@ -295,6 +301,7 @@ export async function* processConverseStream( yield { type: EventType.TOOL_CALL_END, toolCallId: toolCall.id, + ...(toolCall.hasArguments && { args: toolCall.arguments }), toolCallName: toolCall.name, toolName: toolCall.name, timestamp: Date.now(), @@ -337,6 +344,7 @@ export async function* processConverseStream( yield { type: EventType.TOOL_CALL_END, toolCallId: toolCall.id, + ...(toolCall.hasArguments && { args: toolCall.arguments }), toolCallName: toolCall.name, toolName: toolCall.name, timestamp: Date.now(), diff --git a/packages/ai-bedrock/tests/converse/stream-processor.test.ts b/packages/ai-bedrock/tests/converse/stream-processor.test.ts index 620b304d45..b6e216bbdf 100644 --- a/packages/ai-bedrock/tests/converse/stream-processor.test.ts +++ b/packages/ai-bedrock/tests/converse/stream-processor.test.ts @@ -377,4 +377,67 @@ describe('processConverseStream', () => { events.filter((e) => e.type === EventType.TOOL_CALL_END), ).toHaveLength(2) }) + + it('puts the raw streamed argument string on TOOL_CALL_END', async () => { + const events = await collect( + { + contentBlockStart: { + start: { toolUse: { toolUseId: 't1', name: 'getX' } }, + contentBlockIndex: 0, + }, + }, + { + contentBlockDelta: { + delta: { toolUse: { input: '{"n": ' } }, + contentBlockIndex: 0, + }, + }, + { + contentBlockDelta: { + delta: { toolUse: { input: '1.50}' } }, + contentBlockIndex: 0, + }, + }, + { contentBlockStop: { contentBlockIndex: 0 } }, + { messageStop: { stopReason: 'tool_use' } }, + ) + const end = events.find((e) => e.type === EventType.TOOL_CALL_END) + expect(end).toMatchObject({ toolCallId: 't1', args: '{"n": 1.50}' }) + }) + + it('puts malformed raw arguments on TOOL_CALL_END without an input', async () => { + const events = await collect( + { + contentBlockStart: { + start: { toolUse: { toolUseId: 't1', name: 'getX' } }, + contentBlockIndex: 0, + }, + }, + { + contentBlockDelta: { + delta: { toolUse: { input: '{"a":' } }, + contentBlockIndex: 0, + }, + }, + // Truncated: the drained END still carries the raw string. + ) + const end = events.find((e) => e.type === EventType.TOOL_CALL_END) + expect(end).toMatchObject({ toolCallId: 't1', args: '{"a":' }) + expect(end).not.toHaveProperty('input') + }) + + it('omits args on TOOL_CALL_END when no argument fragment arrived', async () => { + const events = await collect( + { + contentBlockStart: { + start: { toolUse: { toolUseId: 't1', name: 'getX' } }, + contentBlockIndex: 0, + }, + }, + { contentBlockStop: { contentBlockIndex: 0 } }, + ) + const end = events.find((e) => e.type === EventType.TOOL_CALL_END) + expect(end).toBeDefined() + expect(end).not.toHaveProperty('args') + }) }) diff --git a/packages/ai-gemini/src/adapters/text.ts b/packages/ai-gemini/src/adapters/text.ts index 8286eb1624..546af7386e 100644 --- a/packages/ai-gemini/src/adapters/text.ts +++ b/packages/ai-gemini/src/adapters/text.ts @@ -1158,12 +1158,14 @@ export class GeminiTextAdapter< // this chunk (including UNEXPECTED_TOOL_CALL finishes) were already // registered and started by the per-part loop above. for (const [toolCallId, toolCallData] of toolCallMap.entries()) { - let parsedInput: unknown = {} + // Malformed JSON sends no `input`, so core reports the error from + // the raw `args`. + let parsedInput: unknown try { const parsed = JSON.parse(toolCallData.args) parsedInput = parsed && typeof parsed === 'object' ? parsed : {} } catch { - parsedInput = {} + // Keep `parsedInput` undefined for the tool error. } yield { @@ -1173,7 +1175,8 @@ export class GeminiTextAdapter< toolName: toolCallData.name, model, timestamp: Date.now(), - input: parsedInput, + args: toolCallData.args, + ...(parsedInput !== undefined && { input: parsedInput }), } } diff --git a/packages/ai-gemini/src/experimental/text-interactions/adapter.ts b/packages/ai-gemini/src/experimental/text-interactions/adapter.ts index 40dc9a26ce..5886afe491 100644 --- a/packages/ai-gemini/src/experimental/text-interactions/adapter.ts +++ b/packages/ai-gemini/src/experimental/text-interactions/adapter.ts @@ -920,6 +920,22 @@ function parsePartialToolArguments( } } +// The raw `args` and parsed `input` of a TOOL_CALL_END. The streamed argument +// text wins over the initial `step.arguments` snapshot. Malformed JSON sends +// no `input`, so core reports the error from the raw `args`. +function toolCallEndArguments( + streamed: string | undefined, + state: ToolCallState, +): { args: string; input?: unknown } { + const args = streamed?.trim() ? streamed : JSON.stringify(state.args) + try { + const parsed: unknown = JSON.parse(args) + return { args, input: parsed && typeof parsed === 'object' ? parsed : {} } + } catch { + return { args } + } +} + // `satisfies` pins these arrays to the SDK's narrow mime-type unions: if // Google removes a format the build breaks, and if they add one ours keeps // working (we just won't accept the new one until added here). @@ -1251,7 +1267,7 @@ async function* translateInteractionEvents( toolName: state.name, model, timestamp, - input: state.args, + ...toolCallEndArguments(argStringByToolCallId.get(toolCallId), state), } } if (hasEmittedTextMessageStart) { @@ -1307,8 +1323,9 @@ async function* translateInteractionEvents( indexToToolCallId.set(index, toolCallId) // `step.arguments` is required on FunctionCallStep but may // be an empty `{}` placeholder when streaming, where the - // real args arrive as `arguments_delta` events. Treat both - // uniformly: stash whatever we got, stringify once. + // real args arrive as `arguments_delta` events. Keep this + // initial snapshot apart from the streamed text: the first + // argument delta starts the raw argument string. const initialArgs = step.arguments const state: ToolCallState = { name: step.name, @@ -1318,12 +1335,6 @@ async function* translateInteractionEvents( ended: false, } toolCalls.set(toolCallId, state) - argStringByToolCallId.set( - toolCallId, - Object.keys(initialArgs).length > 0 - ? JSON.stringify(initialArgs) - : '', - ) yield { type: EventType.TOOL_CALL_START, toolCallId, @@ -1579,8 +1590,9 @@ async function* translateInteractionEvents( break } const state = toolCalls.get(toolCallId) - if (!state) break - const fragment = delta.arguments ?? '' + // A delta with no arguments keeps the initial snapshot. + if (!state || typeof delta.arguments !== 'string') break + const fragment = delta.arguments const buffer = (argStringByToolCallId.get(toolCallId) ?? '') + fragment argStringByToolCallId.set(toolCallId, buffer) @@ -1699,7 +1711,10 @@ async function* translateInteractionEvents( toolName: state.name, model, timestamp, - input: state.args, + ...toolCallEndArguments( + argStringByToolCallId.get(toolCallId), + state, + ), } } indexToToolCallId.delete(event.index) diff --git a/packages/ai-gemini/tests/gemini-adapter.test.ts b/packages/ai-gemini/tests/gemini-adapter.test.ts index 86fee86930..f3c4dc586a 100644 --- a/packages/ai-gemini/tests/gemini-adapter.test.ts +++ b/packages/ai-gemini/tests/gemini-adapter.test.ts @@ -570,6 +570,83 @@ describe('GeminiAdapter through AI', () => { expect(endIds.sort()).toEqual(startIds.sort()) }) + it.each([ + { args: { count: 2 }, raw: '{"count":2}', input: { count: 2 } }, + { args: '{"n": 1.50}', raw: '{"n": 1.50}', input: { n: 1.5 } }, + ])( + 'sends the raw function arguments $raw on TOOL_CALL_END', + async ({ args, raw, input }) => { + mocks.generateContentStreamSpy.mockResolvedValue( + createStream([ + { + candidates: [ + { + content: { + parts: [ + { functionCall: { id: 'call_raw', name: 'inspect', args } }, + ], + }, + finishReason: 'STOP', + }, + ], + }, + ]), + ) + + const chunks: Array = [] + for await (const chunk of createTextAdapter().chatStream({ + model: 'gemini-2.5-pro', + logger: testLogger, + messages: [{ role: 'user', content: 'Go' }], + })) { + chunks.push(chunk) + } + + expect(chunks.find((c) => c.type === 'TOOL_CALL_END')).toMatchObject({ + args: raw, + input, + }) + }, + ) + + it('sends malformed function arguments raw on TOOL_CALL_END with no input', async () => { + mocks.generateContentStreamSpy.mockResolvedValue( + createStream([ + { + candidates: [ + { + content: { + parts: [ + { + functionCall: { + id: 'call_bad', + name: 'inspect', + args: '{"a":', + }, + }, + ], + }, + finishReason: 'STOP', + }, + ], + }, + ]), + ) + + const chunks: Array = [] + for await (const chunk of createTextAdapter().chatStream({ + model: 'gemini-2.5-pro', + logger: testLogger, + messages: [{ role: 'user', content: 'Go' }], + })) { + chunks.push(chunk) + } + + const end = chunks.find((c) => c.type === 'TOOL_CALL_END') + expect(end).toMatchObject({ args: '{"a":' }) + expect(end).not.toHaveProperty('input') + }) + it('emits grounded web search sources and omits the provider call on the next turn', async () => { const groundingMetadata = { webSearchQueries: ['latest release'], diff --git a/packages/ai-gemini/tests/text-interactions-adapter.test.ts b/packages/ai-gemini/tests/text-interactions-adapter.test.ts index 30dffa518c..ea40b04d44 100644 --- a/packages/ai-gemini/tests/text-interactions-adapter.test.ts +++ b/packages/ai-gemini/tests/text-interactions-adapter.test.ts @@ -563,11 +563,9 @@ describe('GeminiTextInteractionsAdapter', () => { expect(endEvent.input).toEqual({ location: 'Berlin', unit: 'celsius' }) }) - it('preserves the last good args when the arguments stream truncates mid-fragment', async () => { - // If the stream ends after an incomplete fragment, the tool input must - // retain whatever was parsed so far rather than being reset to `{}` — a - // strict JSON.parse would throw on the partial buffer and lose the - // already-merged keys. + it('keeps truncated arguments raw on TOOL_CALL_END with no input', async () => { + // A partial value can appear while the arguments stream. It is not the + // final input: core reports the malformed raw arguments as a tool error. mocks.interactionsCreateSpy.mockResolvedValue( mkStream([ { @@ -602,21 +600,98 @@ describe('GeminiTextInteractionsAdapter', () => { ]), ) - const adapter = createAdapter() const chunks = await collectChunks( - chat({ - adapter, + createAdapter().chatStream({ + model: 'gemini-2.5-flash', + logger: testLogger, messages: [{ role: 'user', content: 'Weather in London?' }], tools: [{ name: 'lookup_weather', description: 'Return the weather' }], }), ) - // No parse-failure RUN_ERROR, and the completed key survives truncation. expect(chunks.find((c) => c.type === 'RUN_ERROR')).toBeUndefined() - const endEvent = chunks.find((c) => c.type === 'TOOL_CALL_END') as any - expect(endEvent.input).toEqual({ city: 'London' }) + const endEvent = chunks.find((c) => c.type === 'TOOL_CALL_END') + expect(endEvent).toMatchObject({ args: '{"city":"London","unit":' }) + expect(endEvent).not.toHaveProperty('input') }) + const rawArgumentCases: Array<{ + name: string + deltas: Array + args: string + input: Record + }> = [ + { + name: 'the initial snapshot when no argument delta arrives', + deltas: [], + args: '{"initial":1}', + input: { initial: 1 }, + }, + { + name: 'the streamed text over the initial snapshot', + deltas: ['{"n": ', '1.50}'], + args: '{"n": 1.50}', + input: { n: 1.5 }, + }, + { + name: 'the initial snapshot when a delta has no arguments', + deltas: [undefined], + args: '{"initial":1}', + input: { initial: 1 }, + }, + ] + + it.each(rawArgumentCases)( + 'sends $name as the raw TOOL_CALL_END args', + async ({ deltas, args, input }) => { + mocks.interactionsCreateSpy.mockResolvedValue( + mkStream([ + { + event_type: 'interaction.created', + interaction: { id: 'int_raw', status: 'in_progress' }, + }, + { + event_type: 'step.start', + index: 0, + step: { + type: 'function_call', + id: 'call_raw', + name: 'inspect', + arguments: { initial: 1 }, + }, + }, + ...deltas.map((fragment) => ({ + event_type: 'step.delta', + index: 0, + delta: { + type: 'arguments_delta', + ...(fragment !== undefined && { arguments: fragment }), + }, + })), + { event_type: 'step.stop', index: 0 }, + { + event_type: 'interaction.completed', + interaction: { id: 'int_raw', status: 'completed' }, + }, + ]), + ) + + const chunks = await collectChunks( + createAdapter().chatStream({ + model: 'gemini-2.5-flash', + logger: testLogger, + messages: [{ role: 'user', content: 'Go' }], + tools: [{ name: 'inspect', description: 'Inspect' }], + }), + ) + + expect(chunks.find((c) => c.type === 'TOOL_CALL_END')).toMatchObject({ + args, + input, + }) + }, + ) + it('translates thought_summary deltas into REASONING events', async () => { mocks.interactionsCreateSpy.mockResolvedValue( mkStream([ diff --git a/packages/ai-mistral/src/adapters/text.ts b/packages/ai-mistral/src/adapters/text.ts index 40845ff0b5..17be990d81 100644 --- a/packages/ai-mistral/src/adapters/text.ts +++ b/packages/ai-mistral/src/adapters/text.ts @@ -542,6 +542,7 @@ export class MistralTextAdapter< id: string name: string arguments: string + hasArguments: boolean started: boolean ended: boolean } @@ -664,6 +665,7 @@ export class MistralTextAdapter< id: toolCallDelta.id || '', name: toolCallDelta.function?.name || '', arguments: '', + hasArguments: false, started: false, ended: false, } @@ -684,6 +686,7 @@ export class MistralTextAdapter< : JSON.stringify(rawArgs) if (argsDelta !== undefined) { + toolCall.hasArguments = true toolCall.arguments += argsDelta } @@ -750,6 +753,7 @@ export class MistralTextAdapter< toolName: toolCall.name, model: chunkModel, timestamp, + ...(toolCall.hasArguments && { args: toolCall.arguments }), input: parsedInput, }) } @@ -834,6 +838,9 @@ export class MistralTextAdapter< } for (const [, toolCall] of toolCallsInProgress) { if (toolCall.started && !toolCall.ended) { + // Parse before marking ended, so a parse error still lets the + // catch below send TOOL_CALL_END with the raw args. + const parsedInput = parseToolCallInput(toolCall, normalizeToolInput) toolCall.ended = true hasEmittedToolCall = true yield asChunk({ @@ -843,7 +850,8 @@ export class MistralTextAdapter< toolName: toolCall.name, model: lastChunkModel, timestamp, - input: parseToolCallInput(toolCall, normalizeToolInput), + ...(toolCall.hasArguments && { args: toolCall.arguments }), + input: parsedInput, }) } } @@ -900,9 +908,9 @@ export class MistralTextAdapter< for (const [, toolCall] of toolCallsInProgress) { if (toolCall.started && !toolCall.ended) { toolCall.ended = true - // Best-effort parse for the partial args; if invalid, surface - // empty input rather than throwing inside the cleanup path. - let partialInput: unknown = {} + // Best-effort parse for the partial args. If invalid, leave input + // out and send the raw args, so core reports the parse error. + let partialInput: unknown try { partialInput = toolCall.arguments ? normalizeToolInput( @@ -911,7 +919,7 @@ export class MistralTextAdapter< ) : {} } catch { - partialInput = {} + partialInput = undefined } yield asChunk({ type: 'TOOL_CALL_END', @@ -920,7 +928,8 @@ export class MistralTextAdapter< toolName: toolCall.name, model: lastChunkModel, timestamp, - input: partialInput, + ...(toolCall.hasArguments && { args: toolCall.arguments }), + ...(partialInput !== undefined && { input: partialInput }), }) } } diff --git a/packages/ai-mistral/tests/mistral-adapter.test.ts b/packages/ai-mistral/tests/mistral-adapter.test.ts index 5212814b9a..0d0fea0483 100644 --- a/packages/ai-mistral/tests/mistral-adapter.test.ts +++ b/packages/ai-mistral/tests/mistral-adapter.test.ts @@ -1087,6 +1087,109 @@ describe('Mistral AG-UI event emission', () => { expect(caught?.message).toMatch(/lookup_weather/) }) + it('puts the raw argument string on TOOL_CALL_END', async () => { + setupMockStream([ + { + id: 'cmpl-raw', + model: 'mistral-large-latest', + choices: [ + { + index: 0, + delta: { + toolCalls: [ + { + index: 0, + id: 'tc_raw', + function: { + name: 'lookup_weather', + arguments: '{"location": "Berlin", ', + }, + }, + ], + }, + finishReason: null, + }, + ], + }, + { + id: 'cmpl-raw', + model: 'mistral-large-latest', + choices: [ + { + index: 0, + delta: { + toolCalls: [{ index: 0, function: { arguments: '"n": 1.50}' } }], + }, + finishReason: 'tool_calls', + }, + ], + }, + ]) + + const adapter = createMistralText('mistral-large-latest', 'test-api-key') + const chunks: Array = [] + for await (const chunk of adapter.chatStream( + chatOpts({ + model: 'mistral-large-latest', + messages: [{ role: 'user', content: 'Weather in Berlin?' }], + tools: [weatherTool], + }), + )) { + chunks.push(chunk) + } + + const end = chunks.find((c) => c.type === 'TOOL_CALL_END') + expect(end).toMatchObject({ + toolCallId: 'tc_raw', + args: '{"location": "Berlin", "n": 1.50}', + input: { location: 'Berlin' }, + }) + }) + + it('sends malformed raw arguments on TOOL_CALL_END without an input', async () => { + setupMockStream([ + { + id: 'cmpl-bad', + model: 'mistral-large-latest', + choices: [ + { + index: 0, + delta: { + toolCalls: [ + { + index: 0, + id: 'tc_bad', + function: { name: 'lookup_weather', arguments: '{"a":' }, + }, + ], + }, + finishReason: 'tool_calls', + }, + ], + }, + ]) + + const adapter = createMistralText('mistral-large-latest', 'test-api-key') + const chunks: Array = [] + await expect( + (async () => { + for await (const chunk of adapter.chatStream( + chatOpts({ + model: 'mistral-large-latest', + messages: [{ role: 'user', content: 'x' }], + tools: [weatherTool], + }), + )) { + chunks.push(chunk) + } + })(), + ).rejects.toThrow(/Failed to parse tool call arguments/) + + const end = chunks.find((c) => c.type === 'TOOL_CALL_END') + expect(end).toMatchObject({ toolCallId: 'tc_bad', args: '{"a":' }) + expect(end).not.toHaveProperty('input') + }) + it('sends stream_options.include_usage so Mistral returns usage on streaming', async () => { let capturedBody: unknown vi.stubGlobal( diff --git a/packages/ai-ollama/src/adapters/text.ts b/packages/ai-ollama/src/adapters/text.ts index bc87d6cfee..ca1b59df65 100644 --- a/packages/ai-ollama/src/adapters/text.ts +++ b/packages/ai-ollama/src/adapters/text.ts @@ -355,7 +355,7 @@ export class OllamaTextAdapter extends BaseTextAdapter< } // Serialize arguments to a string for the TOOL_CALL_ARGS event - let parsedInput: unknown = {} + let parsedInput: unknown const argsStr = typeof actualToolCall.function.arguments === 'string' ? actualToolCall.function.arguments @@ -364,7 +364,8 @@ export class OllamaTextAdapter extends BaseTextAdapter< const parsed = JSON.parse(argsStr) parsedInput = parsed && typeof parsed === 'object' ? parsed : {} } catch { - parsedInput = actualToolCall.function.arguments + // Leave input out. Core reports the malformed raw args. + parsedInput = undefined } // Emit TOOL_CALL_ARGS with full args (Ollama doesn't stream args incrementally) @@ -385,7 +386,8 @@ export class OllamaTextAdapter extends BaseTextAdapter< toolName: actualToolCall.function.name || '', model: chunk.model, timestamp: Date.now(), - input: parsedInput, + args: argsStr, + ...(parsedInput !== undefined && { input: parsedInput }), }) return events diff --git a/packages/ai-ollama/tests/text-adapter.test.ts b/packages/ai-ollama/tests/text-adapter.test.ts index dd759c34b7..14dde3a044 100644 --- a/packages/ai-ollama/tests/text-adapter.test.ts +++ b/packages/ai-ollama/tests/text-adapter.test.ts @@ -276,6 +276,53 @@ describe('OllamaTextAdapter.chatStream (tool calls)', () => { expect(startChunk!.toolCallId).toMatch(/^search_\d+/) }) + async function toolCallEnd(args: unknown) { + chatMock.mockResolvedValueOnce( + asyncIterable([ + { + message: { + role: 'assistant', + content: '', + tool_calls: [ + { id: 'tc-raw', function: { name: 'search', arguments: args } }, + ], + }, + done: true, + done_reason: 'stop', + }, + ]), + ) + const chunks = await collectStream( + createOllamaChat('llama3.2').chatStream({ + logger: testLogger, + model: 'llama3.2', + messages: [{ role: 'user', content: 'find cats' }], + tools: [searchTool], + }), + ) + return chunks.find((c) => c.type === 'TOOL_CALL_END') + } + + it('puts the raw argument string on TOOL_CALL_END', async () => { + const end = await toolCallEnd('{"q": "cats", "n": 1.50}') + expect(end).toMatchObject({ + toolCallId: 'tc-raw', + args: '{"q": "cats", "n": 1.50}', + input: { q: 'cats', n: 1.5 }, + }) + }) + + it('puts the serialized object arguments on TOOL_CALL_END', async () => { + const end = await toolCallEnd({ q: 'cats' }) + expect(end).toMatchObject({ args: '{"q":"cats"}', input: { q: 'cats' } }) + }) + + it('sends malformed raw arguments on TOOL_CALL_END without an input', async () => { + const end = await toolCallEnd('{"a":') + expect(end).toMatchObject({ toolCallId: 'tc-raw', args: '{"a":' }) + expect(end).not.toHaveProperty('input') + }) + it('forwards tools to the ollama client in provider format', async () => { chatMock.mockResolvedValueOnce( asyncIterable([ diff --git a/packages/ai-openrouter/src/adapters/responses-text.ts b/packages/ai-openrouter/src/adapters/responses-text.ts index 89a9356964..b2f288ec53 100644 --- a/packages/ai-openrouter/src/adapters/responses-text.ts +++ b/packages/ai-openrouter/src/adapters/responses-text.ts @@ -1471,7 +1471,7 @@ export class OpenRouterResponsesTextAdapter< rawArguments: chunk.arguments, }, ) - parsedInput = {} + parsedInput = undefined } } @@ -1482,7 +1482,8 @@ export class OpenRouterResponsesTextAdapter< toolName: name, model: model || options.model, timestamp: Date.now(), - input: parsedInput, + args: chunk.arguments, + ...(parsedInput !== undefined && { input: parsedInput }), } } @@ -1552,7 +1553,7 @@ export class OpenRouterResponsesTextAdapter< rawArguments: rawArgs, }, ) - parsedInput = {} + parsedInput = undefined } } yield { @@ -1562,7 +1563,8 @@ export class OpenRouterResponsesTextAdapter< toolName: name, model: model || options.model, timestamp: Date.now(), - input: parsedInput, + ...(rawArgs !== undefined && { args: rawArgs }), + ...(parsedInput !== undefined && { input: parsedInput }), } metadata.ended = true metadata.pendingArguments = undefined @@ -1679,7 +1681,7 @@ export class OpenRouterResponsesTextAdapter< rawArguments: rawArgs, }, ) - parsedInput = {} + parsedInput = undefined } } yield { @@ -1689,7 +1691,8 @@ export class OpenRouterResponsesTextAdapter< toolName: name, model: model || options.model, timestamp: Date.now(), - input: parsedInput, + ...(rawArgs !== undefined && { args: rawArgs }), + ...(parsedInput !== undefined && { input: parsedInput }), } metadata.ended = true metadata.pendingArguments = undefined diff --git a/packages/ai-openrouter/src/adapters/text.ts b/packages/ai-openrouter/src/adapters/text.ts index 677d8dc636..884272434f 100644 --- a/packages/ai-openrouter/src/adapters/text.ts +++ b/packages/ai-openrouter/src/adapters/text.ts @@ -1167,7 +1167,7 @@ export class OpenRouterTextAdapter< rawArguments: toolCall.arguments, }, ) - parsedInput = {} + parsedInput = undefined } } @@ -1179,7 +1179,8 @@ export class OpenRouterTextAdapter< toolName: toolCall.name, model: chunk.model || options.model, timestamp: Date.now(), - input: parsedInput, + args: toolCall.arguments, + ...(parsedInput !== undefined && { input: parsedInput }), } emittedAnyToolCallEnd = true } @@ -1230,7 +1231,7 @@ export class OpenRouterTextAdapter< rawArguments: toolCall.arguments, }, ) - parsedInput = {} + parsedInput = undefined } } yield { @@ -1240,7 +1241,8 @@ export class OpenRouterTextAdapter< toolName: toolCall.name, model: lastModel || options.model, timestamp: Date.now(), - input: parsedInput, + args: toolCall.arguments, + ...(parsedInput !== undefined && { input: parsedInput }), } emittedAnyToolCallEnd = true } diff --git a/packages/ai-openrouter/tests/openrouter-adapter.test.ts b/packages/ai-openrouter/tests/openrouter-adapter.test.ts index 00b6235e78..ad9bf2bc3f 100644 --- a/packages/ai-openrouter/tests/openrouter-adapter.test.ts +++ b/packages/ai-openrouter/tests/openrouter-adapter.test.ts @@ -522,6 +522,54 @@ describe('OpenRouter adapter option mapping', () => { } }) + // finishReason=null closes the call in the post-stream drain. + it.each([ + { raw: '{"n": 1.50}', finishReason: 'tool_calls' }, + { raw: '{"n": 1.50}', finishReason: null }, + { raw: '{"a":', finishReason: 'tool_calls' }, + { raw: '{"a":', finishReason: null }, + ])( + 'sends the raw argument string $raw as args on TOOL_CALL_END (finishReason=$finishReason)', + async ({ raw, finishReason }) => { + setupMockSdkClient([ + { + id: 'chatcmpl-raw', + model: 'openai/gpt-4o-mini', + choices: [ + { + delta: { + toolCalls: [ + { + index: 0, + id: 'call_raw', + type: 'function', + function: { name: 'lookup_weather', arguments: raw }, + }, + ], + }, + finishReason, + }, + ], + }, + ]) + + const chunks: Array = [] + for await (const chunk of createAdapter().chatStream({ + model: 'openai/gpt-4o-mini', + messages: [{ role: 'user', content: 'Weather?' }], + tools: [weatherTool], + logger: testLogger, + })) { + chunks.push(chunk) + } + + const toolEnd = chunks.find((c) => c.type === 'TOOL_CALL_END') + expect(toolEnd).toMatchObject({ toolCallId: 'call_raw', args: raw }) + if (raw === '{"a":') expect(toolEnd).not.toHaveProperty('input') + else expect(toolEnd).toMatchObject({ input: { n: 1.5 } }) + }, + ) + it('emits parentMessageId on tool-first tool calls matching the assistant message id', async () => { // Tool call streams before any text content. parentMessageId must bind the // tool call to the same assistant message id the eventual TEXT_MESSAGE_START diff --git a/packages/ai-openrouter/tests/openrouter-responses-adapter.test.ts b/packages/ai-openrouter/tests/openrouter-responses-adapter.test.ts index e02554a996..caeaf5bc9d 100644 --- a/packages/ai-openrouter/tests/openrouter-responses-adapter.test.ts +++ b/packages/ai-openrouter/tests/openrouter-responses-adapter.test.ts @@ -1353,6 +1353,108 @@ describe('OpenRouter responses adapter — stream event bridge', () => { }) }) + describe('raw tool arguments', () => { + type EndSite = 'arguments.done' | 'output_item.done' | 'response.completed' + + // Each site emits TOOL_CALL_END from a different event. Without a name on + // output_item.added, args.done is deferred to a later backfill. + function rawArgsStream( + site: EndSite, + raw: string, + ): Array> { + const named = { + type: 'function_call', + id: 'item_raw', + callId: 'call_raw', + name: 'lookup_weather', + } + return [ + { + type: 'response.output_item.added', + sequenceNumber: 0, + outputIndex: 0, + item: + site === 'arguments.done' + ? { ...named, arguments: '' } + : { + type: 'function_call', + id: 'item_raw', + callId: 'call_raw', + arguments: '', + }, + }, + { + type: 'response.function_call_arguments.done', + sequenceNumber: 1, + itemId: 'item_raw', + outputIndex: 0, + arguments: raw, + }, + ...(site === 'output_item.done' + ? [ + { + type: 'response.output_item.done', + sequenceNumber: 2, + outputIndex: 0, + item: { ...named, arguments: raw }, + }, + ] + : []), + { + type: 'response.completed', + sequenceNumber: 3, + response: { + model: 'm', + output: [{ ...named, arguments: raw }], + usage: { inputTokens: 0, outputTokens: 0, totalTokens: 0 }, + }, + }, + ] + } + + async function toolEnds(site: EndSite, raw: string) { + setupMockSdkClient(rawArgsStream(site, raw)) + const chunks: Array = [] + for await (const chunk of createAdapter().chatStream({ + model: 'openai/gpt-4o-mini', + messages: [{ role: 'user', content: 'hi' }], + tools: [weatherTool], + logger: testLogger, + })) { + chunks.push(chunk) + } + return chunks.filter((chunk) => chunk.type === 'TOOL_CALL_END') + } + + it.each([ + 'arguments.done', + 'output_item.done', + 'response.completed', + ])( + 'sends the raw argument string as args on TOOL_CALL_END (%s)', + async (site) => { + const ends = await toolEnds(site, '{"n": 1.50}') + expect(ends).toHaveLength(1) + expect(ends[0]).toMatchObject({ + toolCallId: 'call_raw', + args: '{"n": 1.50}', + input: { n: 1.5 }, + }) + }, + ) + + it.each([ + 'arguments.done', + 'output_item.done', + 'response.completed', + ])('sends malformed arguments as args with no input (%s)', async (site) => { + const ends = await toolEnds(site, '{"a":') + expect(ends).toHaveLength(1) + expect(ends[0]).toMatchObject({ toolCallId: 'call_raw', args: '{"a":' }) + expect(ends[0]).not.toHaveProperty('input') + }) + }) + it('round-trips distinct Responses item and call IDs through a server tool', async () => { const firstTurn = [ { diff --git a/packages/ai/src/activities/chat/index.ts b/packages/ai/src/activities/chat/index.ts index 26afde8a9e..600aba0775 100644 --- a/packages/ai/src/activities/chat/index.ts +++ b/packages/ai/src/activities/chat/index.ts @@ -39,6 +39,7 @@ import { tanstackMetadata, withTanstackMetadata, } from '../../utilities/merge-metadata' +import { isToolInputJsonLossless } from '../../utilities/tool-call-arguments' import { subagentHostMessageId } from '../../utilities/subagent-wire' import { withDurabilityBatchHint } from '../../utilities/durability-batch' import { normalizeStreamChunk } from '../../utilities/normalize-stream-chunk' @@ -2345,6 +2346,25 @@ class TextEngine< private handleToolCallEndEvent(chunk: ToolCallEndEvent): void { this.toolCallManager.completeToolCall(chunk) const end = chunk as AdapterYieldChunk + const call = this.toolCallManager + .getToolCalls() + .find((candidate) => candidate.id === chunk.toolCallId) + if (call) { + // The client stores the same argument string as the server history. + // `input` rides along only when JSON keeps it exactly. + const args = call.function.arguments + end.args = args + end.metadata = withTanstackMetadata(end, { args }).metadata + try { + const projection: unknown = args.trim() === '' ? {} : JSON.parse(args) + if (isToolInputJsonLossless(projection)) end.input = projection + else delete end.input + } catch { + delete end.input + } + const metadata = tanstackMetadata(end) + if (metadata) delete metadata.input + } const state = end.state ?? tanstackMetadata(end)?.state if (state !== 'output-error' || end.result === undefined) return this.handleToolCallResultEvent({ diff --git a/packages/ai/src/activities/chat/stream/processor.ts b/packages/ai/src/activities/chat/stream/processor.ts index 12c960686d..770018c99c 100644 --- a/packages/ai/src/activities/chat/stream/processor.ts +++ b/packages/ai/src/activities/chat/stream/processor.ts @@ -35,6 +35,7 @@ import { withTanstackMetadata, } from '../../../utilities/merge-metadata' import { getChunkRunId } from '../../../utilities/chunk-ids' +import { reconcileToolCallArguments } from '../../../utilities/tool-call-arguments' import type { AdapterYieldChunk } from '../../../utilities/adapter-yield-chunk' import { isContentPartArray, @@ -2516,8 +2517,20 @@ export class StreamProcessor { // `arguments` and `input` then stay with the streamed value, so the // two never disagree. const serializedInput = serializeToolInput(input) - if (serializedInput !== undefined) { - existingToolCall.arguments = serializedInput + // The provider's raw argument string wins when it is the same JSON + // value as the input, so the history keeps the exact bytes. + const rawArgs = + (chunk as AdapterYieldChunk).args ?? tanstackMetadata(chunk)?.args + let args = serializedInput + if (rawArgs !== undefined) { + try { + args = reconcileToolCallArguments(rawArgs, input) + } catch { + // An input with no JSON form keeps the canonical string. + } + } + if (args !== undefined) { + existingToolCall.arguments = args } const index = msgState.toolCallOrder.indexOf(chunk.toolCallId) diff --git a/packages/ai/src/activities/chat/tools/tool-calls.ts b/packages/ai/src/activities/chat/tools/tool-calls.ts index 359241134d..0db2c37f3d 100644 --- a/packages/ai/src/activities/chat/tools/tool-calls.ts +++ b/packages/ai/src/activities/chat/tools/tool-calls.ts @@ -1,3 +1,4 @@ +import { reconcileToolCallArguments } from '../../../utilities/tool-call-arguments' import { normalizeToolResult } from '../../../utilities/tool-result' import { tanstackMetadata } from '../../../utilities/merge-metadata' import { isProviderExecutedToolCall } from '../../../utilities/provider-executed' @@ -313,10 +314,32 @@ export class ToolCallManager< completeToolCall(event: ToolCallEndEvent): void { for (const toolCall of this.toolCallsMap.values()) { if (toolCall.id !== event.toolCallId) continue - if (event.input === undefined) return - const normalized = - event.input && typeof event.input === 'object' ? event.input : {} - toolCall.function.arguments = JSON.stringify(normalized) + // The provider's raw argument string wins when the parsed input is the + // same JSON value, so replay sends the exact bytes back. Malformed raw + // JSON stays as it is, for the tool error. + const extra = event as AdapterYieldChunk + const snapshot = extra.args ?? tanstackMetadata(extra)?.args + // A non-object input (an empty tool_use block) is `{}`. + const input = + event.input === undefined + ? undefined + : event.input && typeof event.input === 'object' + ? event.input + : {} + let raw: string | undefined = snapshot ?? toolCall.function.arguments + if (snapshot === undefined && input !== undefined) { + try { + JSON.parse(raw) + } catch { + raw = undefined + } + } + try { + toolCall.function.arguments = reconcileToolCallArguments(raw, input) + } catch { + // An input with no JSON form (a cycle): keep the raw arguments. + if (raw !== undefined) toolCall.function.arguments = raw + } return } } diff --git a/packages/ai/src/types.ts b/packages/ai/src/types.ts index 45adefd676..eaf111b0f4 100644 --- a/packages/ai/src/types.ts +++ b/packages/ai/src/types.ts @@ -773,6 +773,8 @@ export interface TanStackRunMetadata { state?: ToolOutputState /** Parsed `TOOL_CALL_END` input. Spec `TOOL_CALL_END` has no top-level `input`. */ input?: unknown + /** The raw tool argument string on `TOOL_CALL_END`, as the provider sent it. */ + args?: string } /** diff --git a/packages/ai/src/utilities/tool-call-arguments.ts b/packages/ai/src/utilities/tool-call-arguments.ts new file mode 100644 index 0000000000..3f0355690e --- /dev/null +++ b/packages/ai/src/utilities/tool-call-arguments.ts @@ -0,0 +1,309 @@ +type JsonNode = { + start: number + end: number + value: unknown + keyToken?: string + properties?: Map + items?: Array +} + +function jsonInput( + value: unknown, + key = '', + seen = new Set(), +): unknown { + if ( + (value !== null && + (typeof value === 'object' || typeof value === 'function')) || + typeof value === 'bigint' + ) { + const toJSON: unknown = Reflect.get(Object(value), 'toJSON') + if (typeof toJSON === 'function') + value = Reflect.apply(toJSON, value, [key]) + } + if (value !== null && typeof value === 'object') { + let boxedType: 'number' | 'string' | 'boolean' | 'bigint' | undefined + try { + Number.prototype.valueOf.call(value) + boxedType = 'number' + } catch {} + if (boxedType === undefined) { + try { + String.prototype.valueOf.call(value) + boxedType = 'string' + } catch {} + } + if (boxedType === undefined) { + try { + Boolean.prototype.valueOf.call(value) + boxedType = 'boolean' + } catch {} + } + if (boxedType === undefined) { + try { + BigInt.prototype.valueOf.call(value) + boxedType = 'bigint' + } catch {} + } + if (boxedType === 'number') value = +Object(value) + else if (boxedType === 'string') value = String(value) + else if (boxedType === 'boolean') + value = Boolean.prototype.valueOf.call(value) + else if (boxedType === 'bigint') + value = BigInt.prototype.valueOf.call(value) + } + if (value === null || typeof value !== 'object') return value + if (seen.has(value)) + throw new TypeError('Cannot serialize circular tool input') + seen.add(value) + try { + if (Array.isArray(value)) { + const items: Array = [] + const length = Math.min( + Math.max(Math.trunc(+value.length) || 0, 0), + Number.MAX_SAFE_INTEGER, + ) + for (let index = 0; index < length; index++) + items.push(jsonInput(value[index], String(index), seen)) + return items + } + const result: Record = {} + for (const property of Object.keys(value)) { + const child = jsonInput(Reflect.get(value, property), property, seen) + if ( + child !== undefined && + typeof child !== 'function' && + typeof child !== 'symbol' + ) + Object.defineProperty(result, property, { + value: child, + enumerable: true, + configurable: true, + writable: true, + }) + } + return result + } finally { + seen.delete(value) + } +} + +function sameJsonValue(left: unknown, right: unknown): boolean { + if (Object.is(left, right)) return true + if (Array.isArray(left) && Array.isArray(right)) + return ( + left.length === right.length && + left.every((value, index) => sameJsonValue(value, right[index])) + ) + if ( + left === null || + right === null || + typeof left !== 'object' || + typeof right !== 'object' || + Array.isArray(left) || + Array.isArray(right) + ) + return false + const leftKeys = Object.keys(left) + const rightKeys = Object.keys(right) + return ( + leftKeys.length === rightKeys.length && + leftKeys.every( + (key) => + Object.prototype.hasOwnProperty.call(right, key) && + sameJsonValue(Reflect.get(left, key), Reflect.get(right, key)), + ) + ) +} + +function scanJson(raw: string, parsed: unknown): JsonNode { + let cursor = 0 + const whitespace = () => { + while ( + raw[cursor] === ' ' || + raw[cursor] === '\t' || + raw[cursor] === '\r' || + raw[cursor] === '\n' + ) + cursor++ + } + const stringToken = () => { + const start = cursor++ + while (cursor < raw.length) { + if (raw[cursor] === '\\') cursor += 2 + else if (raw[cursor++] === '"') break + } + return raw.slice(start, cursor) + } + const value = (parsedValue: unknown): JsonNode => { + whitespace() + const start = cursor + if (raw[cursor] === '{') { + cursor++ + whitespace() + const properties = new Map() + while (raw[cursor] !== '}') { + const keyToken = stringToken() + const key: unknown = JSON.parse(keyToken) + whitespace() + cursor++ + const child = + parsedValue !== null && + typeof parsedValue === 'object' && + typeof key === 'string' + ? Reflect.get(parsedValue, key) + : undefined + const childNode = value(child) + if (typeof key === 'string') + properties.set(key, { ...childNode, keyToken }) + whitespace() + if (raw[cursor] !== ',') break + cursor++ + whitespace() + } + cursor++ + return { start, end: cursor, value: parsedValue, properties } + } + if (raw[cursor] === '[') { + cursor++ + whitespace() + const items: Array = [] + while (raw[cursor] !== ']') { + items.push( + value( + Array.isArray(parsedValue) ? parsedValue[items.length] : undefined, + ), + ) + whitespace() + if (raw[cursor] !== ',') break + cursor++ + } + cursor++ + return { start, end: cursor, value: parsedValue, items } + } + if (raw[cursor] === '"') stringToken() + else + while (cursor < raw.length && !/[\s,}\]]/.test(raw[cursor] ?? '')) + cursor++ + return { start, end: cursor, value: parsedValue } + } + return value(parsed) +} + +function serializeValue( + value: unknown, + raw: string, + node: JsonNode | undefined, + seen: Set, +): string | undefined { + if (node && sameJsonValue(node.value, value)) + return raw.slice(node.start, node.end) + if (typeof value === 'number') { + if (Object.is(value, -0)) return '-0' + if (value === Infinity) return '1e999' + if (value === -Infinity) return '-1e999' + return JSON.stringify(value) + } + if (value === null || typeof value !== 'object') return JSON.stringify(value) + if (seen.has(value)) + throw new TypeError('Cannot serialize circular tool input') + seen.add(value) + try { + if (Array.isArray(value)) + return ( + '[' + + Array.from( + value, + (item, index) => + serializeValue(item, raw, node?.items?.[index], seen) ?? 'null', + ).join(',') + + ']' + ) + const properties: Array = [] + for (const key of Object.keys(value)) { + const serialized = serializeValue( + Reflect.get(value, key), + raw, + node?.properties?.get(key), + seen, + ) + if (serialized !== undefined) + properties.push( + (node?.properties?.get(key)?.keyToken ?? JSON.stringify(key)) + + ':' + + serialized, + ) + } + return '{' + properties.join(',') + '}' + } finally { + seen.delete(value) + } +} + +export function reconcileToolCallArguments( + raw: string | undefined, + input: unknown, +): string { + if (input === undefined) return raw ?? '' + if (raw === undefined) + return serializeValue(jsonInput(input), '', undefined, new Set()) ?? '' + let parsed: unknown + try { + parsed = JSON.parse(raw) + } catch { + // Empty raw arguments are no arguments: the input is the value to keep. + if (raw.trim() === '') + return serializeValue(jsonInput(input), '', undefined, new Set()) ?? raw + return raw + } + const value = jsonInput(input) + if (sameJsonValue(parsed, value)) return raw + return serializeValue(value, raw, scanJson(raw, parsed), new Set()) ?? raw +} + +export function isToolInputJsonLossless( + input: unknown, + seen = new Set(), +): boolean { + if (input === null || typeof input === 'string' || typeof input === 'boolean') + return true + if (typeof input === 'number') + return Number.isFinite(input) && !Object.is(input, -0) + if (typeof input !== 'object' || seen.has(input)) return false + const prototype: unknown = Object.getPrototypeOf(input) + if ( + !Array.isArray(input) && + prototype !== Object.prototype && + prototype !== null + ) + return false + const toJSON = Object.getOwnPropertyDescriptor(input, 'toJSON') + if ( + toJSON && + (toJSON.get !== undefined || typeof toJSON.value === 'function') + ) + return false + seen.add(input) + try { + if (Array.isArray(input)) { + for (let index = 0; index < input.length; index++) { + const descriptor = Object.getOwnPropertyDescriptor(input, String(index)) + if ( + descriptor?.get !== undefined || + !isToolInputJsonLossless(descriptor?.value, seen) + ) + return false + } + return true + } + return Object.keys(input).every((key) => { + const descriptor = Object.getOwnPropertyDescriptor(input, key) + return ( + descriptor?.get === undefined && + isToolInputJsonLossless(descriptor?.value, seen) + ) + }) + } finally { + seen.delete(input) + } +} diff --git a/packages/ai/tests/tool-call-arguments.test.ts b/packages/ai/tests/tool-call-arguments.test.ts new file mode 100644 index 0000000000..1a6e876d9c --- /dev/null +++ b/packages/ai/tests/tool-call-arguments.test.ts @@ -0,0 +1,173 @@ +import { describe, expect, it } from 'vitest' +import { chat } from '../src/activities/chat/index' +import { EventType } from '../src/types' +import { tanstackMetadata } from '../src/utilities/merge-metadata' +import { reconcileToolCallArguments } from '../src/utilities/tool-call-arguments' +import { collectChunks, createMockAdapter, ev } from './test-utils' + +describe('lexical tool argument reconciliation', () => { + it.each([ + ' 1e999 ', + '-0', + '9007199254740993', + '1.20e+3', + '"\\u0061"', + '{"a":1e999,"b":[-0,9007199254740993,"\\u0061"]}', + ])('keeps unchanged raw JSON exactly: %s', (raw) => { + expect(reconcileToolCallArguments(raw, JSON.parse(raw))).toBe(raw) + }) + it('reuses unchanged nested values after removing an optional null', () => { + const raw = + '{"removed":null,"nested":{"number":1e999,"text":"\\u0061"},"array":[-0,9007199254740993]}' + expect( + reconcileToolCallArguments(raw, { + nested: { number: Infinity, text: 'a' }, + array: [-0, 9007199254740992], + }), + ).toBe( + '{"nested":{"number":1e999,"text":"\\u0061"},"array":[-0,9007199254740993]}', + ) + }) + it('retains unchanged array positions while correcting another position', () => { + expect(reconcileToolCallArguments('[1e999,2,-0]', [Infinity, 3, -0])).toBe( + '[1e999,3,-0]', + ) + }) + it('uses valid JSON spelling for corrected numbers', () => { + expect( + reconcileToolCallArguments('{}', { + positive: Infinity, + negative: -Infinity, + zero: -0, + }), + ).toBe('{"positive":1e999,"negative":-1e999,"zero":-0}') + }) + it('does not repair malformed raw JSON with a parsed projection', () => { + expect(reconcileToolCallArguments('{"broken":', {})).toBe('{"broken":') + }) + it('serializes END-only input when no raw arguments exist', () => { + expect(reconcileToolCallArguments(undefined, null)).toBe('null') + expect(reconcileToolCallArguments(undefined, [1, true])).toBe('[1,true]') + }) + it('leaves raw arguments alone when input is undefined', () => { + expect(reconcileToolCallArguments('1e999', undefined)).toBe('1e999') + }) +}) + +describe('raw JSON key spelling', () => { + it('keeps an escaped key when removing an unrelated null', () => { + expect( + reconcileToolCallArguments('{"\\u0061":9007199254740993,"drop":null}', { + a: 9007199254740992, + }), + ).toBe('{"\\u0061":9007199254740993}') + }) + it('uses the last duplicate key value and its original spelling', () => { + expect( + reconcileToolCallArguments('{"a":1,"\\u0061":1e999,"drop":null}', { + a: Infinity, + }), + ).toBe('{"\\u0061":1e999}') + }) + it('serializes sparse custom array positions as null slots', () => { + const values: Array = new Array(2) + values[1] = -0 + expect(reconcileToolCallArguments(undefined, values)).toBe('[null,-0]') + }) +}) + +describe('ordinary custom input JSON semantics', () => { + it('uses Date and boxed primitive JSON values', () => { + expect( + reconcileToolCallArguments('{}', new Date('2026-10-04T00:00:00Z')), + ).toBe('"2026-10-04T00:00:00.000Z"') + expect(reconcileToolCallArguments(undefined, Object(7))).toBe('7') + }) + it('passes the property key to custom toJSON', () => { + const input = { + value: { + toJSON(key: string) { + return key + }, + }, + } + expect(reconcileToolCallArguments(undefined, input)).toBe( + '{"value":"value"}', + ) + }) +}) + +describe('raw tool arguments through chat()', () => { + async function run(raw: string, input?: unknown) { + const executed: Array = [] + const { adapter, calls } = createMockAdapter({ + iterations: [ + [ + ev.runStarted(), + ev.toolStart('call-1', 'check'), + ev.toolArgs('call-1', raw), + { + type: EventType.TOOL_CALL_END, + toolCallId: 'call-1', + timestamp: Date.now(), + args: raw, + ...(input !== undefined ? { input } : {}), + }, + ev.runFinished('tool_calls'), + ], + [ev.runStarted('run-2'), ev.runFinished('stop', 'run-2')], + ], + }) + const chunks = await collectChunks( + chat({ + adapter, + messages: [{ role: 'user', content: 'check' }], + tools: [ + { + name: 'check', + description: 'Check', + execute: (value: unknown) => { + executed.push(value) + return 'ok' + }, + }, + ], + }), + ) + return { chunks, calls, executed } + } + + function replayedArguments( + calls: ReturnType['calls'], + ) { + return calls[1]?.messages + .flatMap((message) => message.toolCalls ?? []) + .map((call) => call.function.arguments) + } + + it('replays the exact provider bytes', async () => { + const raw = '{"count": 1.50, "id": 9007199254740993}' + const { calls, executed } = await run(raw, JSON.parse(raw)) + expect(replayedArguments(calls)).toEqual([raw]) + expect(executed).toEqual([{ count: 1.5, id: 9007199254740992 }]) + }) + + it('sends the raw arguments to the client in TOOL_CALL_END metadata', async () => { + const raw = '{"count": 1.50}' + const { chunks } = await run(raw, JSON.parse(raw)) + const end = chunks.find((chunk) => chunk.type === EventType.TOOL_CALL_END) + const metadata = end ? tanstackMetadata(end) : undefined + expect(metadata && 'args' in metadata ? metadata.args : undefined).toBe(raw) + }) + + it('reports malformed raw arguments instead of running with {}', async () => { + const { calls, executed } = await run('{"count":') + expect(executed).toEqual([]) + const toolMessage = calls[1]?.messages.find( + (message) => message.role === 'tool', + ) + expect(JSON.stringify(toolMessage)).toContain( + 'Failed to parse tool arguments as JSON', + ) + }) +}) diff --git a/packages/openai-base/src/adapters/chat-completions-text.ts b/packages/openai-base/src/adapters/chat-completions-text.ts index ae0e6172f8..43048659ba 100644 --- a/packages/openai-base/src/adapters/chat-completions-text.ts +++ b/packages/openai-base/src/adapters/chat-completions-text.ts @@ -220,6 +220,7 @@ export abstract class OpenAIBaseChatCompletionsTextAdapter< toolCallId, toolCallName: rejectedToolCall.toolName, toolName: rejectedToolCall.toolName, + args: rejectedToolCall.arguments, ...(rejectedToolCall.input !== undefined && { input: rejectedToolCall.input, }), @@ -1166,7 +1167,8 @@ export abstract class OpenAIBaseChatCompletionsTextAdapter< toolName: toolCall.name, model: chunk.model || options.model, timestamp: Date.now(), - input: parsedInput, + args: toolCall.arguments, + ...(parsedInput !== undefined && { input: parsedInput }), } emittedAnyToolCallEnd = true } @@ -1236,7 +1238,8 @@ export abstract class OpenAIBaseChatCompletionsTextAdapter< toolName: toolCall.name, model: lastModel || options.model, timestamp: Date.now(), - input: parsedInput, + args: toolCall.arguments, + ...(parsedInput !== undefined && { input: parsedInput }), } pendingToolCount += 1 emittedAnyToolCallEnd = true diff --git a/packages/openai-base/src/adapters/responses-text.ts b/packages/openai-base/src/adapters/responses-text.ts index cb8c05b562..03f567ba84 100644 --- a/packages/openai-base/src/adapters/responses-text.ts +++ b/packages/openai-base/src/adapters/responses-text.ts @@ -2074,7 +2074,7 @@ export abstract class OpenAIBaseResponsesTextAdapter< rawArguments: chunk.arguments, }, ) - parsedInput = {} + parsedInput = undefined } } @@ -2085,7 +2085,8 @@ export abstract class OpenAIBaseResponsesTextAdapter< toolName: name, model: model || options.model, timestamp: Date.now(), - input: parsedInput, + args: chunk.arguments, + ...(parsedInput !== undefined && { input: parsedInput }), } } @@ -2171,7 +2172,7 @@ export abstract class OpenAIBaseResponsesTextAdapter< rawArguments: rawArgs, }, ) - parsedInput = {} + parsedInput = undefined } } yield { @@ -2181,7 +2182,8 @@ export abstract class OpenAIBaseResponsesTextAdapter< toolName: name, model: model || options.model, timestamp: Date.now(), - input: parsedInput, + ...(rawArgs !== undefined && { args: rawArgs }), + ...(parsedInput !== undefined && { input: parsedInput }), } metadata.ended = true metadata.pendingArguments = undefined @@ -2342,7 +2344,7 @@ export abstract class OpenAIBaseResponsesTextAdapter< rawArguments: rawArgs, }, ) - parsedInput = {} + parsedInput = undefined } } yield { @@ -2352,7 +2354,8 @@ export abstract class OpenAIBaseResponsesTextAdapter< toolName: name, model: model || options.model, timestamp: Date.now(), - input: parsedInput, + ...(rawArgs !== undefined && { args: rawArgs }), + ...(parsedInput !== undefined && { input: parsedInput }), } metadata.ended = true metadata.pendingArguments = undefined diff --git a/packages/openai-base/tests/chat-completions-text.test.ts b/packages/openai-base/tests/chat-completions-text.test.ts index 1fe050309a..8b16063ece 100644 --- a/packages/openai-base/tests/chat-completions-text.test.ts +++ b/packages/openai-base/tests/chat-completions-text.test.ts @@ -1240,7 +1240,8 @@ describe('OpenAIBaseChatCompletionsTextAdapter', () => { const toolEnd = chunks.find((chunk) => chunk.type === 'TOOL_CALL_END') expect(toolEnd).toBeDefined() - expect(toolEnd?.input).toBeUndefined() + expect(toolEnd).toMatchObject({ args: '{"location":' }) + expect(toolEnd).not.toHaveProperty('input') expect( chunks.find((chunk) => chunk.type === 'TOOL_CALL_ARGS'), ).toMatchObject({ @@ -1263,6 +1264,64 @@ describe('OpenAIBaseChatCompletionsTextAdapter', () => { ) }) + describe('raw tool arguments', () => { + it.each([true, false])( + 'sends the accumulated argument string as args on TOOL_CALL_END (finish_reason=%s)', + async (withFinishReason) => { + setupMockSdkClient([ + { + id: 'chatcmpl-raw', + model: 'test-model', + choices: [ + { + delta: { + tool_calls: [ + { + index: 0, + id: 'call_raw', + type: 'function', + function: { name: 'lookup_weather', arguments: '{"n": ' }, + }, + ], + }, + finish_reason: null, + }, + ], + }, + { + id: 'chatcmpl-raw', + model: 'test-model', + choices: [ + { + delta: { + tool_calls: [{ index: 0, function: { arguments: '1.50}' } }], + }, + finish_reason: withFinishReason ? 'tool_calls' : null, + }, + ], + }, + ]) + const adapter = new TestChatCompletionsAdapter(testConfig, 'test-model') + const chunks: Array = [] + for await (const chunk of adapter.chatStream({ + logger: testLogger, + model: 'test-model', + messages: [{ role: 'user', content: 'Weather?' }], + tools: [weatherTool], + })) { + chunks.push(chunk) + } + + const toolEnd = chunks.find((chunk) => chunk.type === 'TOOL_CALL_END') + expect(toolEnd).toMatchObject({ + toolCallId: 'call_raw', + args: '{"n": 1.50}', + input: { n: 1.5 }, + }) + }, + ) + }) + describe('subclassing', () => { it('allows subclassing with custom name', () => { class MyProviderAdapter extends OpenAIBaseChatCompletionsTextAdapter { diff --git a/packages/openai-base/tests/responses-text.test.ts b/packages/openai-base/tests/responses-text.test.ts index 404d41239a..6dc49e6bcd 100644 --- a/packages/openai-base/tests/responses-text.test.ts +++ b/packages/openai-base/tests/responses-text.test.ts @@ -1869,6 +1869,108 @@ describe('OpenAIBaseResponsesTextAdapter', () => { }) }) + describe('raw tool arguments', () => { + type EndSite = 'arguments.done' | 'output_item.done' | 'response.completed' + + // Each site emits TOOL_CALL_END from a different event. Without a name on + // output_item.added, args.done is deferred to a later backfill. + function rawArgsStream( + site: EndSite, + raw: string, + ): Array> { + const named = { + type: 'function_call', + id: 'fc_raw', + name: 'lookup_weather', + } + return [ + { + type: 'response.created', + response: { + id: 'resp-raw', + model: 'test-model', + status: 'in_progress', + }, + }, + { + type: 'response.output_item.added', + output_index: 0, + item: + site === 'arguments.done' + ? named + : { type: 'function_call', id: 'fc_raw' }, + }, + { + type: 'response.function_call_arguments.done', + item_id: 'fc_raw', + arguments: raw, + }, + ...(site === 'output_item.done' + ? [ + { + type: 'response.output_item.done', + output_index: 0, + item: { ...named, arguments: raw }, + }, + ] + : []), + { + type: 'response.completed', + response: { + id: 'resp-raw', + model: 'test-model', + status: 'completed', + output: [{ ...named, arguments: raw }], + usage: { input_tokens: 1, output_tokens: 1, total_tokens: 2 }, + }, + }, + ] + } + + async function toolEnds(site: EndSite, raw: string) { + setupMockResponsesClient(rawArgsStream(site, raw)) + const adapter = new TestResponsesAdapter(testConfig, 'test-model') + const chunks: Array = [] + for await (const chunk of adapter.chatStream({ + logger: testLogger, + model: 'test-model', + messages: [{ role: 'user', content: 'q' }], + tools: [weatherTool], + })) { + chunks.push(chunk) + } + return chunks.filter((chunk) => chunk.type === 'TOOL_CALL_END') + } + + it.each([ + 'arguments.done', + 'output_item.done', + 'response.completed', + ])( + 'sends the raw argument string as args on TOOL_CALL_END (%s)', + async (site) => { + const ends = await toolEnds(site, '{"n": 1.50}') + expect(ends).toHaveLength(1) + expect(ends[0]).toMatchObject({ + toolCallId: 'fc_raw', + args: '{"n": 1.50}', + input: { n: 1.5 }, + }) + }, + ) + + it.each([ + 'arguments.done', + 'output_item.done', + 'response.completed', + ])('sends malformed arguments as args with no input (%s)', async (site) => { + const ends = await toolEnds(site, '{"a":') + expect(ends).toHaveLength(1) + expect(ends[0]).toMatchObject({ toolCallId: 'fc_raw', args: '{"a":' }) + expect(ends[0]).not.toHaveProperty('input') + }) + }) + describe('content_part events', () => { it('emits TEXT_MESSAGE_START on content_part.added with output_text', async () => { const streamChunks = [ diff --git a/testing/e2e/fixtures/middleware-test/raw-tool-args.json b/testing/e2e/fixtures/middleware-test/raw-tool-args.json new file mode 100644 index 0000000000..97e159b49d --- /dev/null +++ b/testing/e2e/fixtures/middleware-test/raw-tool-args.json @@ -0,0 +1,27 @@ +{ + "fixtures": [ + { + "match": { + "userMessage": "[raw-tool-args] run test", + "sequenceIndex": 0 + }, + "response": { + "toolCalls": [ + { + "name": "get_weather", + "arguments": "{\"city\": \"NYC\", \"days\": 1.50}" + } + ] + } + }, + { + "match": { + "userMessage": "[raw-tool-args] run test", + "sequenceIndex": 1 + }, + "response": { + "content": "The weather is sunny." + } + } + ] +} diff --git a/testing/e2e/src/routes/api.middleware-test.ts b/testing/e2e/src/routes/api.middleware-test.ts index d52d183451..f5d7cb672a 100644 --- a/testing/e2e/src/routes/api.middleware-test.ts +++ b/testing/e2e/src/routes/api.middleware-test.ts @@ -675,7 +675,7 @@ export const Route = createFileRoute('/api/middleware-test')({ const tools = genericScenario && testId ? genericTools(testId, genericScenario) - : scenario === 'with-tool' + : scenario === 'with-tool' || scenario === 'raw-tool-args' ? [weatherTool] : scenario === 'structured-client-tool-wait' ? [clientContextToolDefinition] diff --git a/testing/e2e/src/routes/middleware-test.tsx b/testing/e2e/src/routes/middleware-test.tsx index 59d3643469..fc326cc99b 100644 --- a/testing/e2e/src/routes/middleware-test.tsx +++ b/testing/e2e/src/routes/middleware-test.tsx @@ -254,6 +254,7 @@ function MiddlewareTestPage() { +