@cloudflare/codemode 0.1.0 → 0.1.1

package/src/tests/types.test.ts

@@ -4,8 +4,16 @@
 import { describe, it, expect } from "vitest";
 import { generateTypes, sanitizeToolName } from "../types";
 import { z } from "zod";
+import { fromJSONSchema } from "zod/v4";
+import { jsonSchema } from "ai";
+import type { ToolSet } from "ai";
 import type { ToolDescriptors } from "../types";
 
+// Helper: cast loosely-typed tool objects for generateTypes
+function genTypes(tools: Record<string, unknown>): string {
+  return generateTypes(tools as unknown as ToolSet);
+}
+
 describe("sanitizeToolName", () => {
   it("should replace hyphens with underscores", () => {
     expect(sanitizeToolName("get-weather")).toBe("get_weather");
@@ -168,4 +176,271 @@ describe("generateTypes", () => {
     // toCamelCase("get_weather") → "GetWeather"
     expect(result).toContain("GetWeatherInput");
   });
+
+  it("should handle MCP tools with input and output schemas (fromJSONSchema)", () => {
+    // MCP tools use JSON Schema format for both input and output
+    const inputSchema = {
+      type: "object" as const,
+      properties: {
+        city: { type: "string" as const, description: "City name" },
+        units: {
+          type: "string" as const,
+          enum: ["celsius", "fahrenheit"],
+          description: "Temperature units"
+        },
+        includeForecast: { type: "boolean" as const }
+      },
+      required: ["city"]
+    };
+
+    const outputSchema = {
+      type: "object" as const,
+      properties: {
+        temperature: { type: "number" as const, description: "Current temp" },
+        humidity: { type: "number" as const },
+        conditions: { type: "string" as const },
+        forecast: {
+          type: "array" as const,
+          items: {
+            type: "object" as const,
+            properties: {
+              day: { type: "string" as const },
+              high: { type: "number" as const },
+              low: { type: "number" as const }
+            }
+          }
+        }
+      },
+      required: ["temperature", "conditions"]
+    };
+
+    const tools: ToolDescriptors = {
+      getWeather: {
+        description: "Get weather for a city",
+        inputSchema: fromJSONSchema(inputSchema),
+        outputSchema: fromJSONSchema(outputSchema)
+      }
+    };
+
+    const result = generateTypes(tools);
+
+    // Input schema types
+    expect(result).toContain("type GetWeatherInput");
+    expect(result).toContain("city: string");
+    expect(result).toContain("units?:");
+    expect(result).toContain("includeForecast?: boolean");
+
+    // Output schema types (not unknown)
+    expect(result).toContain("type GetWeatherOutput");
+    expect(result).not.toContain("GetWeatherOutput = unknown");
+    expect(result).toContain("temperature: number");
+    expect(result).toContain("humidity?: number");
+    expect(result).toContain("conditions: string");
+    expect(result).toContain("forecast?:");
+    expect(result).toContain("day?: string");
+    expect(result).toContain("high?: number");
+    expect(result).toContain("low?: number");
+
+    // JSDoc
+    expect(result).toContain("@param input.city - City name");
+    expect(result).toContain("/** Current temp */");
+  });
+
+  it("should handle Zod schemas with input and output schemas", () => {
+    // Direct ToolDescriptors with Zod schemas (what generateTypes operates on)
+    const tools: ToolDescriptors = {
+      getWeather: {
+        description: "Get weather for a city",
+        inputSchema: z.object({
+          city: z.string().describe("City name"),
+          units: z.enum(["celsius", "fahrenheit"]).optional()
+        }),
+        outputSchema: z.object({
+          temperature: z.number().describe("Current temperature"),
+          humidity: z.number().describe("Humidity percentage"),
+          conditions: z.string().describe("Weather conditions"),
+          forecast: z.array(
+            z.object({
+              day: z.string(),
+              high: z.number(),
+              low: z.number()
+            })
+          )
+        })
+      }
+    };
+
+    const result = generateTypes(tools);
+
+    // Verify input schema
+    expect(result).toContain("type GetWeatherInput");
+    expect(result).toContain("city: string");
+    expect(result).toContain("units?:");
+    expect(result).toContain('"celsius"');
+    expect(result).toContain('"fahrenheit"');
+
+    // Verify output schema is properly typed (not unknown)
+    expect(result).toContain("type GetWeatherOutput");
+    expect(result).not.toContain("GetWeatherOutput = unknown");
+    expect(result).toContain("temperature: number");
+    expect(result).toContain("humidity: number");
+    expect(result).toContain("conditions: string");
+    expect(result).toContain("forecast:");
+    expect(result).toContain("day: string");
+    expect(result).toContain("high: number");
+    expect(result).toContain("low: number");
+
+    // Verify JSDoc comments from .describe()
+    expect(result).toContain("/** City name */");
+    expect(result).toContain("/** Current temperature */");
+    expect(result).toContain("@param input.city - City name");
+  });
+
+  it("should handle null inputSchema gracefully", () => {
+    const tools = {
+      broken: {
+        description: "Broken tool",
+        inputSchema: null
+      }
+    };
+
+    const result = genTypes(tools);
+
+    expect(result).toContain("type BrokenInput = unknown");
+    expect(result).toContain("type BrokenOutput = unknown");
+    expect(result).toContain("broken:");
+  });
+
+  it("should handle undefined inputSchema gracefully", () => {
+    const tools = {
+      broken: {
+        description: "Broken tool",
+        inputSchema: undefined
+      }
+    };
+
+    const result = genTypes(tools);
+
+    expect(result).toContain("type BrokenInput = unknown");
+    expect(result).toContain("type BrokenOutput = unknown");
+    expect(result).toContain("broken:");
+  });
+
+  it("should handle string inputSchema gracefully", () => {
+    const tools = {
+      broken: {
+        description: "Broken tool",
+        inputSchema: "not a schema"
+      }
+    };
+
+    const result = genTypes(tools);
+
+    expect(result).toContain("type BrokenInput = unknown");
+    expect(result).toContain("broken:");
+  });
+
+  it("should isolate errors: one throwing tool does not break others", () => {
+    // Create a tool with a getter that throws
+    const throwingSchema = {
+      get jsonSchema(): never {
+        throw new Error("Schema explosion");
+      }
+    };
+
+    const tools = {
+      good1: {
+        description: "Good first",
+        inputSchema: jsonSchema({
+          type: "object" as const,
+          properties: { a: { type: "string" as const } }
+        })
+      },
+      bad: {
+        description: "Bad tool",
+        inputSchema: throwingSchema
+      },
+      good2: {
+        description: "Good second",
+        inputSchema: jsonSchema({
+          type: "object" as const,
+          properties: { b: { type: "number" as const } }
+        })
+      }
+    };
+
+    const result = genTypes(tools);
+
+    // Good tools should work fine
+    expect(result).toContain("type Good1Input");
+    expect(result).toContain("a?: string;");
+    expect(result).toContain("type Good2Input");
+    expect(result).toContain("b?: number;");
+
+    // Bad tool should degrade to unknown
+    expect(result).toContain("type BadInput = unknown");
+    expect(result).toContain("type BadOutput = unknown");
+
+    // All three tools should appear in the codemode declaration
+    expect(result).toContain("good1:");
+    expect(result).toContain("bad:");
+    expect(result).toContain("good2:");
+  });
+
+  it("should handle AI SDK jsonSchema wrapper (MCP tools)", () => {
+    // This is what MCP tools look like when using the AI SDK jsonSchema wrapper
+    const inputJsonSchema = {
+      type: "object" as const,
+      properties: {
+        query: { type: "string" as const, description: "Search query" },
+        limit: { type: "number" as const, description: "Max results" }
+      },
+      required: ["query"]
+    };
+
+    const outputJsonSchema = {
+      type: "object" as const,
+      properties: {
+        results: {
+          type: "array" as const,
+          items: {
+            type: "object" as const,
+            properties: {
+              title: { type: "string" as const },
+              url: { type: "string" as const }
+            }
+          }
+        },
+        total: { type: "number" as const }
+      }
+    };
+
+    // Use AI SDK jsonSchema wrapper (what MCP client returns)
+    const tools = {
+      search: {
+        description: "Search the web",
+        inputSchema: jsonSchema(inputJsonSchema),
+        outputSchema: jsonSchema(outputJsonSchema)
+      }
+    };
+
+    const result = generateTypes(tools as unknown as ToolDescriptors);
+
+    // Input schema types
+    expect(result).toContain("type SearchInput");
+    expect(result).toContain("query: string");
+    expect(result).toContain("limit?: number");
+
+    // Output schema types (not unknown)
+    expect(result).toContain("type SearchOutput");
+    expect(result).not.toContain("SearchOutput = unknown");
+    expect(result).toContain("results?:");
+    expect(result).toContain("title?: string");
+    expect(result).toContain("url?: string");
+    expect(result).toContain("total?: number");
+
+    // JSDoc from JSON Schema descriptions
+    expect(result).toContain("@param input.query - Search query");
+    expect(result).toContain("@param input.limit - Max results");
+  });
 });

package/src/tool.ts

@@ -1,4 +1,4 @@
-import { tool, type Tool } from "ai";
+import { tool, type Tool, asSchema } from "ai";
 import { z } from "zod";
 import type { ToolSet } from "ai";
 import * as acorn from "acorn";
@@ -10,7 +10,8 @@ const DEFAULT_DESCRIPTION = `Execute code to achieve a goal.
 Available:
 {{types}}
 
-Write an async arrow function that returns the result.
+Write an async arrow function in JavaScript that returns the result.
+Do NOT use TypeScript syntax — no type annotations, interfaces, or generics.
 Do NOT define named functions then call them — just write the arrow function body directly.
 
 Example: async () => { const r = await codemode.searchWeb({ query: "test" }); return r; }`;
@@ -97,7 +98,9 @@ export function createCodeTool(
     description,
     inputSchema: codeSchema,
     execute: async ({ code }) => {
-      // Extract execute functions from tools, keyed by name
+      // Extract execute functions from tools, keyed by name.
+      // Wrap each with its schema so arguments from the sandbox
+      // are validated before reaching the tool function.
       const fns: Record<string, (...args: unknown[]) => Promise<unknown>> = {};
 
       for (const [name, t] of Object.entries(tools)) {
@@ -106,7 +109,25 @@ export function createCodeTool(
             ? (t.execute as (args: unknown) => Promise<unknown>)
             : undefined;
         if (execute) {
-          fns[sanitizeToolName(name)] = execute;
+          const rawSchema =
+            "inputSchema" in t
+              ? t.inputSchema
+              : "parameters" in t
+                ? (t as Record<string, unknown>).parameters
+                : undefined;
+
+          // Use AI SDK's asSchema() to normalize any schema type
+          // (Zod v3/v4, Standard Schema, JSON Schema) into a unified
+          // Schema with an optional .validate() method.
+          const schema = rawSchema != null ? asSchema(rawSchema) : undefined;
+
+          fns[sanitizeToolName(name)] = schema?.validate
+            ? async (args: unknown) => {
+                const result = await schema.validate!(args);
+                if (!result.success) throw result.error;
+                return execute(result.value);
+              }
+            : execute;
         }
       }