From ba5c8fcfdf556e3269d98988a7a4a232071550bc Mon Sep 17 00:00:00 2001
From: Anatol Zakrividoroga <53095479+anatolzak@users.noreply.github.com>
Date: Fri, 19 Jun 2026 15:55:37 +0300
Subject: [PATCH] fix(code-mode): validate tool input/output against schemas

---
 .changeset/code-mode-tool-validation.md       |  7 ++++
 .../src/bindings/tool-to-binding.ts           | 38 ++++++++++++++++--
 .../tests/tool-to-binding.test.ts             | 39 +++++++++++++++++++
 3 files changed, 80 insertions(+), 4 deletions(-)
 create mode 100644 .changeset/code-mode-tool-validation.md

diff --git a/.changeset/code-mode-tool-validation.md b/.changeset/code-mode-tool-validation.md
new file mode 100644
index 000000000..052da081a
--- /dev/null
+++ b/.changeset/code-mode-tool-validation.md
@@ -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.
diff --git a/packages/ai-code-mode/src/bindings/tool-to-binding.ts b/packages/ai-code-mode/src/bindings/tool-to-binding.ts
index 6469fe945..b6c722f51 100644
--- a/packages/ai-code-mode/src/bindings/tool-to-binding.ts
+++ b/packages/ai-code-mode/src/bindings/tool-to-binding.ts
@@ -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'
 
@@ -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 {
diff --git a/packages/ai-code-mode/tests/tool-to-binding.test.ts b/packages/ai-code-mode/tests/tool-to-binding.test.ts
index dc19c9267..2c4b65f4c 100644
--- a/packages/ai-code-mode/tests/tool-to-binding.test.ts
+++ b/packages/ai-code-mode/tests/tool-to-binding.test.ts
@@ -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', () => {