Thanks to visit codestin.com
Credit goes to github.com

Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
7 changes: 7 additions & 0 deletions .changeset/code-mode-tool-validation.md
Original file line number Diff line number Diff line change
@@ -0,0 +1,7 @@
---
'@tanstack/ai-code-mode': patch
---

fix(code-mode): validate tool input/output against schemas

Code mode converted a tool's input/output schema to JSON Schema for the prompt but never validated against it β€” unlike the normal agent-loop path. A provided `inputSchema`/`outputSchema` was effectively just documentation: inner (`external_*`) tools received raw, un-coerced sandbox args and outputs went unchecked. `toolToBinding` now runs the same Standard Schema validation, so defaults/coercions apply and invalid input/output throws an agent-readable error.
38 changes: 34 additions & 4 deletions packages/ai-code-mode/src/bindings/tool-to-binding.ts
Original file line number Diff line number Diff line change
@@ -1,4 +1,8 @@
import { convertSchemaToJsonSchema } from '@tanstack/ai'
import {
convertSchemaToJsonSchema,
isStandardSchema,
parseWithStandardSchema,
} from '@tanstack/ai'
import type { ToolExecutionContext } from '@tanstack/ai'
import type { CodeModeTool, ToolBinding } from '../types'

Expand Down Expand Up @@ -51,9 +55,35 @@ export function toolToBinding(
}

const toolExecute = tool.execute
const execute = (args: unknown, context?: ToolExecutionContext) => {
// Pass context to the underlying tool so it can emit custom events
return Promise.resolve(toolExecute(args, context))
const execute = async (args: unknown, context?: ToolExecutionContext) => {
let input = args
if (tool.inputSchema && isStandardSchema(tool.inputSchema)) {
try {
input = parseWithStandardSchema(tool.inputSchema, args)
} catch (error) {
const message =
error instanceof Error ? error.message : 'Validation failed'
throw new Error(
`Input validation failed for tool ${tool.name}: ${message}`,
)
}
}

let result = await Promise.resolve(toolExecute(input, context))

if (tool.outputSchema && isStandardSchema(tool.outputSchema)) {
try {
result = parseWithStandardSchema(tool.outputSchema, result)
} catch (error) {
const message =
error instanceof Error ? error.message : 'Validation failed'
throw new Error(
`Output validation failed for tool ${tool.name}: ${message}`,
)
}
}

return result
}

return {
Expand Down
39 changes: 39 additions & 0 deletions packages/ai-code-mode/tests/tool-to-binding.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -59,6 +59,45 @@ describe('toolToBinding', () => {
const result = await binding.execute({ query: 'test' })
expect(result).toEqual({ result: 'response to test' })
})

it('validates input against the tool schema and throws an agent-readable error', async () => {
const tool = createMockServerTool('fetchData')
const binding = toolToBinding(tool)

await expect(binding.execute({ query: 123 })).rejects.toThrow(
/Input validation failed for tool fetchData/,
)
})

it('coerces/defaults input via the schema before calling execute', async () => {
const def = toolDefinition({
name: 'withDefault',
description: 'has a defaulted field',
inputSchema: z.object({ n: z.number().default(7) }),
outputSchema: z.object({ n: z.number() }),
})
const execute = vi.fn(async (input: { n?: number }) => ({ n: input.n! }))
const binding = toolToBinding(def.server(execute))

const result = await binding.execute({})
expect(execute).toHaveBeenCalledWith({ n: 7 }, undefined)
expect(result).toEqual({ n: 7 })
})

it('validates output against the tool schema', async () => {
const def = toolDefinition({
name: 'badOutput',
description: 'returns the wrong shape',
inputSchema: z.object({}),
outputSchema: z.object({ result: z.string() }),
})
// @ts-expect-error - deliberately returns a non-conforming value
const binding = toolToBinding(def.server(async () => ({ result: 123 })))

await expect(binding.execute({})).rejects.toThrow(
/Output validation failed for tool badOutput/,
)
})
})

describe('toolsToBindings', () => {
Expand Down
Loading