TypeScript: Implementing Type-Safe Tool Calling for AI Agents

Introduction to Type-Safe Agent Tools

Building autonomous AI agents requires invoking deterministic code functions from non-deterministic language models. When an agent decides to execute a search query, query a database, or trigger a deployment pipeline, it outputs structured text—usually JSON—representing the tool name and its parameters. In a production environment, trusting the model's output without strict structural validation leads to runtime failures, malformed payloads, and silent execution bugs.

Type safety in this domain spans two distinct boundaries: static verification during compilation and runtime validation before execution. By leveraging advanced TypeScript features like mapped types, template literal types, and runtime schema validators, we can bridge the gap between static types and dynamic LLM generations.

Defining the Tool Schema Contract

To allow a Large Language Model to call our functions, we must expose a schema that describes the function signature, argument names, data types, and required fields. Standard JSON Schema is the universal format accepted by major model providers. Writing these schemas by hand alongside TypeScript types introduces drift. Instead, we should derive JSON schemas directly from TypeScript types or use a validation library like Zod to define both.

Here is how we define a structured tool definition interface in TypeScript:

import { z } from 'zod';
 
export interface ToolDefinition<TParameters extends z.ZodTypeAny = z.ZodTypeAny> {
  name: string;
  description: string;
  parameters: TParameters;
  execute: (args: z.infer<TParameters>) => Promise<unknown>;
}

Using Zod allows us to extract the inferred TypeScript type using z.infer<TParameters>, ensuring that the execution function receives strongly typed arguments without manual casting.

Creating a Type-Safe Tool Registry

An agentic workflow often exposes dozens of tools. Managing these tools requires a registry that can aggregate their JSON schemas for the model provider while retaining precise argument types for the execution phase. If an agent selects a tool, the dispatcher must extract the correct parameter type and pass it safely to the implementation.

export class ToolRegistry<TTools extends Record<string, ToolDefinition>> {
  constructor(private tools: TTools) {}
 
  public getProviderSchemas() {
    return Object.values(this.tools).map(tool => ({
      type: 'function',
      function: {
        name: tool.name,
        description: tool.description,
        parameters: zodToJsonSchema(tool.parameters),
      },
    }));
  }
 
  public async execute<TName extends keyof TTools>(
    name: TName,
    rawArgs: unknown
  ): Promise<Awaited<ReturnType<TTools[TName]['execute']>>> {
    const tool = this.tools[name];
    if (!tool) {
      throw new Error(`Tool ${String(name)} not found`);
    }
 
    const parsedArgs = await tool.parameters.parseAsync(rawArgs);
    return tool.execute(parsedArgs);
  }
}

By constraining the name parameter with keyof TTools, the registry ensures that callers cannot invoke non-existent tools, and the return type is automatically inferred based on the specific tool being executed.

Handling Validation Failures and Self-Correction

LLMs frequently generate malformed JSON or omit required arguments, especially when handling complex nested objects. A naive implementation throws an unhandled exception, abruptly halting the agent loop. A robust implementation captures validation errors and feeds them back to the model as a system or tool-error message, allowing the agent to self-correct.

async function safeExecuteStep(registry: ToolRegistry<any>, toolCall: ToolCallRecord) {
  try {
    const result = await registry.execute(toolCall.name, toolCall.arguments);
    return { success: true, data: result };
  } catch (error) {
    if (error instanceof z.ZodError) {
      return {
        success: false,
        error: `Invalid arguments provided: ${error.message}. Please correct the payload and retry.`,
      };
    }
    return {
      success: false,
      error: error instanceof Error ? error.message : 'Unknown execution error',
    };
  }
}

This feedback loop transforms structural failures from fatal application crashes into manageable conversational states.

Architectural Tradeoffs and Constraints

Implementing strict type safety introduces overhead. Deriving JSON schemas at runtime adds CPU cycles on initialization, though this is negligible compared to network latency from model inference. Furthermore, strictly typed tools reduce the flexibility of open-ended agents. If an agent needs to pass arbitrary key-value pairs, overly rigid Zod schemas will reject valid creative attempts.

Security Implications

Allowing an LLM to invoke code paths dynamically creates an attack surface akin to Remote Code Execution (RCE). Even with strict schemas, an attacker who prompts the model maliciously (prompt injection) could force a valid tool to execute destructive actions, such as dropping a database table or deleting files.

Mitigate this by:

• Enforcing user-in-the-loop confirmation steps for destructive mutations.
• Scoping database and file-system tools to isolated sandboxes or read-only roles.
• Implementing rate limits on resource-intensive tools.

Operational Verification and Testing

Testing tool-calling loops requires mocking model responses to ensure deterministic behavior. Because the execution layer relies on standard TypeScript functions, you can unit-test each tool in isolation without spinning up an LLM mock.

import { describe, it, expect } from 'vitest';
 
describe('DatabaseQueryTool', () => {
  it('parses valid SQL parameters and executes', async () => {
    const result = await databaseTool.execute({ query: 'SELECT 1' });
    expect(result).toBeDefined();
  });
 
  it('rejects invalid parameters at runtime', async () => {
    await expect(databaseTool.execute({ query: 123 })).rejects.toThrow();
  });
});

Integration tests should verify that the registry correctly translates Zod schemas into the format expected by the upstream model provider API, ensuring that schema modifications do not silently break tool serialization.