copilotkit--copilotkit/packages/runtime/src/agent/__tests__/zod-regression.test.ts

/**
 * Regression tests proving that the agent package's defineTool and
 * convertToolDefinitionsToVercelAITools still work identically with
 * Zod schemas after the Standard Schema migration.
 *
 * Covers:
 * 1. defineTool with complex Zod schemas
 * 2. convertToolDefinitionsToVercelAITools passes Zod schemas directly
 *    (not through JSON Schema conversion — critical behavioral regression)
 * 3. BuiltInAgent with Zod tools still works end-to-end
 * 4. execute callback receives correctly shaped args
 */
import { describe, it, expect, vi, beforeEach, afterEach } from "vitest";
import { z } from "zod";
import {
  defineTool,
  convertToolDefinitionsToVercelAITools,
  BuiltInAgent,
} from "../index";
import type { RunAgentInput } from "@ag-ui/client";
import { streamText } from "ai";
import { mockStreamTextResponse, finish, collectEvents } from "./test-helpers";

vi.mock("ai", async () => {
  const actual = await vi.importActual("ai");
  return {
    ...actual,
    streamText: vi.fn(),
    tool: vi.fn((config) => config),
    stepCountIs: vi.fn((count: number) => ({ type: "stepCount", count })),
  };
});

vi.mock("@ai-sdk/openai", () => ({
  createOpenAI: vi.fn(() => (modelId: string) => ({
    modelId,
    provider: "openai",
  })),
}));

describe("defineTool Zod regression", () => {
  it("complex Zod schema with nested objects, arrays, unions", () => {
    const tool = defineTool({
      name: "complexTool",
      description: "A complex tool",
      parameters: z.object({
        query: z.string().describe("Search query"),
        filters: z
          .object({
            category: z.enum(["books", "movies", "music"]).optional(),
            minRating: z.number().min(0).max(5).optional(),
            tags: z.array(z.string()).optional(),
          })
          .optional(),
        pagination: z
          .object({
            page: z.number().int().positive().default(1),
            perPage: z.number().int().positive().default(20),
          })
          .optional(),
      }),
      execute: async (args) => {
        return { query: args.query };
      },
    });

    expect(tool.name).toBe("complexTool");
    expect(tool.parameters["~standard"].vendor).toBe("zod");
    expect(tool.parameters["~standard"].version).toBe(1);
  });

  it("Zod schema with discriminated union", () => {
    const tool = defineTool({
      name: "actionTool",
      description: "An action tool",
      parameters: z.object({
        action: z.discriminatedUnion("type", [
          z.object({ type: z.literal("search"), query: z.string() }),
          z.object({ type: z.literal("navigate"), url: z.string().url() }),
          z.object({
            type: z.literal("execute"),
            code: z.string(),
            language: z.enum(["javascript", "python"]),
          }),
        ]),
      }),
      execute: async (args) => {
        return { action: args.action.type };
      },
    });

    expect(tool.name).toBe("actionTool");
    expect(tool.parameters["~standard"].vendor).toBe("zod");
  });

  it("Zod schema with nullable and record fields", () => {
    const tool = defineTool({
      name: "flexTool",
      description: "Flexible tool",
      parameters: z.object({
        title: z.string(),
        description: z.string().nullable(),
        metadata: z.record(z.string(), z.unknown()).optional(),
      }),
      execute: async (args) => {
        return { title: args.title };
      },
    });

    expect(tool.name).toBe("flexTool");
    expect(tool.parameters["~standard"].vendor).toBe("zod");
  });

  it("execute callback receives correctly typed args", async () => {
    const receivedArgs: unknown[] = [];

    const tool = defineTool({
      name: "argCapture",
      description: "Captures args",
      parameters: z.object({
        city: z.string(),
        units: z.enum(["celsius", "fahrenheit"]),
      }),
      execute: async (args) => {
        receivedArgs.push(args);
        return { temp: 22 };
      },
    });

    await tool.execute({ city: "Berlin", units: "celsius" });

    expect(receivedArgs).toHaveLength(1);
    expect(receivedArgs[0]).toEqual({ city: "Berlin", units: "celsius" });
  });
});

describe("convertToolDefinitionsToVercelAITools Zod regression", () => {
  it("Zod schemas are passed directly to AI SDK (not converted via JSON Schema)", () => {
    const zodSchema = z.object({
      city: z.string(),
      units: z.enum(["celsius", "fahrenheit"]).optional(),
    });

    const tools = [
      defineTool({
        name: "weather",
        description: "Get weather",
        parameters: zodSchema,
        execute: async () => ({ temp: 22 }),
      }),
    ];

    const aiTools = convertToolDefinitionsToVercelAITools(tools);

    expect(aiTools).toHaveProperty("weather");
    // The inputSchema should be the original Zod schema (passed directly),
    // NOT a jsonSchema() wrapper. We verify by checking the Zod-specific
    // ~standard.vendor property is preserved on the schema.
    const inputSchema = aiTools.weather.inputSchema;
    expect(inputSchema).toBeDefined();
    // For Zod schemas, the AI SDK receives the raw Zod object (which has ~standard.vendor === "zod")
    // If it went through jsonSchema(), it would lose the ~standard property
    expect(inputSchema["~standard"]?.vendor).toBe("zod");
  });

  it("multiple Zod tools all get direct pass-through", () => {
    const tools = [
      defineTool({
        name: "tool1",
        description: "Tool 1",
        parameters: z.object({ a: z.string() }),
        execute: async () => ({}),
      }),
      defineTool({
        name: "tool2",
        description: "Tool 2",
        parameters: z.object({ b: z.number(), c: z.boolean().optional() }),
        execute: async () => ({}),
      }),
      defineTool({
        name: "tool3",
        description: "Tool 3",
        parameters: z.object({
          nested: z.object({ x: z.number(), y: z.number() }),
        }),
        execute: async () => ({}),
      }),
    ];

    const aiTools = convertToolDefinitionsToVercelAITools(tools);

    for (const name of ["tool1", "tool2", "tool3"]) {
      expect(aiTools).toHaveProperty(name);
      expect(aiTools[name].inputSchema["~standard"]?.vendor).toBe("zod");
      expect(typeof aiTools[name].execute).toBe("function");
    }
  });

  it("preserves tool description and execute function", () => {
    const executeFn = vi.fn().mockResolvedValue({ result: "ok" });

    const tools = [
      defineTool({
        name: "myTool",
        description: "My custom tool description",
        parameters: z.object({ input: z.string() }),
        execute: executeFn,
      }),
    ];

    const aiTools = convertToolDefinitionsToVercelAITools(tools);

    expect(aiTools.myTool.description).toBe("My custom tool description");
    expect(aiTools.myTool.execute).toBeDefined();
  });
});

describe("BuiltInAgent Zod regression", () => {
  const originalEnv = process.env;

  beforeEach(() => {
    vi.clearAllMocks();
    process.env = { ...originalEnv };
    process.env.OPENAI_API_KEY = "test-key";
  });

  afterEach(() => {
    process.env = originalEnv;
  });

  it("Zod tools are included in streamText call", async () => {
    const agent = new BuiltInAgent({
      model: "openai/gpt-4o",
      tools: [
        defineTool({
          name: "weather",
          description: "Get weather",
          parameters: z.object({
            city: z.string(),
            units: z.enum(["celsius", "fahrenheit"]).optional(),
          }),
          execute: async () => ({ temp: 22 }),
        }),
      ],
    });

    vi.mocked(streamText).mockReturnValue(
      mockStreamTextResponse([finish()]) as any,
    );

    const input: RunAgentInput = {
      threadId: "thread1",
      runId: "run1",
      messages: [],
      tools: [],
      context: [],
      state: {},
    };

    await collectEvents(agent["run"](input));

    const callArgs = vi.mocked(streamText).mock.calls[0][0];
    expect(callArgs.tools).toHaveProperty("weather");
    expect(typeof callArgs.tools?.weather.execute).toBe("function");
  });

  it("multiple Zod tools coexist with built-in state tools", async () => {
    const agent = new BuiltInAgent({
      model: "openai/gpt-4o",
      tools: [
        defineTool({
          name: "search",
          description: "Search",
          parameters: z.object({ query: z.string() }),
          execute: async () => ({}),
        }),
        defineTool({
          name: "calculate",
          description: "Calculate",
          parameters: z.object({
            expression: z.string(),
            precision: z.number().int().optional(),
          }),
          execute: async () => ({}),
        }),
      ],
    });

    vi.mocked(streamText).mockReturnValue(
      mockStreamTextResponse([finish()]) as any,
    );

    const input: RunAgentInput = {
      threadId: "thread1",
      runId: "run1",
      messages: [],
      tools: [],
      context: [],
      state: {},
    };

    await collectEvents(agent["run"](input));

    const callArgs = vi.mocked(streamText).mock.calls[0][0];
    // User-defined Zod tools
    expect(callArgs.tools).toHaveProperty("search");
    expect(callArgs.tools).toHaveProperty("calculate");
    // Built-in state tools
    expect(callArgs.tools).toHaveProperty("AGUISendStateSnapshot");
    expect(callArgs.tools).toHaveProperty("AGUISendStateDelta");
  });

  it("Zod tool execute is invocable after conversion", async () => {
    const executeFn = vi.fn().mockResolvedValue({ temp: 22 });

    const agent = new BuiltInAgent({
      model: "openai/gpt-4o",
      tools: [
        defineTool({
          name: "weather",
          description: "Get weather",
          parameters: z.object({ city: z.string() }),
          execute: executeFn,
        }),
      ],
    });

    vi.mocked(streamText).mockReturnValue(
      mockStreamTextResponse([finish()]) as any,
    );

    const input: RunAgentInput = {
      threadId: "thread1",
      runId: "run1",
      messages: [],
      tools: [],
      context: [],
      state: {},
    };

    await collectEvents(agent["run"](input));

    const callArgs = vi.mocked(streamText).mock.calls[0][0];
    const weatherTool = callArgs.tools?.weather;

    // Invoke the execute function to verify it's wired correctly
    await weatherTool.execute(
      { city: "Berlin" },
      { toolCallId: "tool-call-1", messages: [] },
    );
    expect(executeFn).toHaveBeenCalledWith(
      { city: "Berlin" },
      { toolCallId: "tool-call-1", messages: [] },
    );
  });
});