feat: add generic type support for ToolInvocation

[liby]

Description

This PR adds generic type support to the ToolInvocation type, allowing for better type safety when working with tool calls and results.

Example

// With specific types
type WeatherTool = ToolInvocation<
  { city: string },
  { temperature: number; condition: string }
>;

[lgrammel]

Thanks! How would you use this in client code? do you have an example?

[liby]

Thanks! How would you use this in client code? do you have an example?

Example

• agents/tools/types.ts

  import type { z } from 'zod';
  import type { Tool } from 'ai';
  import type { allTools, ToolName } from './tools';
  
  export type ToolCreationFunction = () => Tool<any, any>;
  
  /**
   * Infer parameter types from a tool creation function
   */
  export type InferToolArgs<T extends ToolCreationFunction> =
    z.infer<ReturnType<T>['parameters']>;
  
  /**
   * Infer result types from a tool creation function
   */
  export type InferToolResult<T extends ToolCreationFunction> =
    Awaited<ReturnType<NonNullable<ReturnType<T>['execute']>>>;
  
  /**
   * Extract complete type information from a tool creation function
   */
  export interface ToolTypes<T extends ToolCreationFunction> {
    Args: InferToolArgs<T>;
    Result: InferToolResult<T>;
  }
  
  // The ToolInvocation in ai-sdk is not generic, so we write our own
  // Note: we need to make sure the types are compatible, check this when upgrading ai-sdk
  export type GenericToolInvocation<ARGS = any, RESULT = any> =
    ({
      state: 'partial-call' | 'call';
      toolCallId: string;
      toolName: ToolName;
      args: ARGS;
    } | {
      state: 'result';
      toolCallId: string;
      toolName: ToolName;
      args: ARGS;
      result: RESULT;
    });
  
  // Infer the type of the tool render from the tool creation function
  export type ToolRender<T extends ToolCreationFunction> =
    (props: {
      invocation: GenericToolInvocation<
        ToolTypes<T>['Args'],
        ToolTypes<T>['Result']
      >
    }) => React.ReactNode;
  
  export type AllToolRenders = Record<ToolName, ToolRender<typeof allTools[ToolName]['createTool']>>

• agents/tools/tools.ts

  import "server-only"
  import { createScrapingTool } from "./scraping/server";
  import type { ToolCreationFunction } from "./types";
  
  interface ToolDefinition {
    createTool: ToolCreationFunction;
  }
  
  export const allTools = {
    scrapeUrl: {
      createTool: createScrapingTool,
    },
  } as const satisfies Record<string, ToolDefinition>;
  
  export type ToolName = keyof typeof allTools;

• agents/tools/renders.ts

  import { renderScraping } from "./scraping/client";
  import type { AllToolRenders } from "./types";
  
  export const toolRenders: AllToolRenders = {
      scrapeUrl: renderScraping,
  }

• client-component.tsx

  'use client';
  
  import { toolRenders } from '@/agents/tools/renders';
  import type { ToolName } from '@/agents/tools/tools';
  import type { GenericToolInvocation } from '@/agents/tools/types';
  
  const ClientComponent = ({ part }: { part: TextUIPart | ReasoningUIPart | ToolInvocationUIPart | SourceUIPart }) => {
    if (part.type === 'tool-invocation') {
      const Render = toolRenders[part.toolInvocation.toolName as ToolName];
  
      if (!Render) {
        return (
          <Render
            key={part.toolInvocation.toolCallId}
            invocation={
              part.toolInvocation as GenericToolInvocation<any,any>
            }
          />
        );
      }
    }
  }