refactoring: introduce FlexibleSchema (#6818)

## Background

A generalized schema type is needed for output schema support.

## Summary

Replace ToolInputSchema with FlexibleSchema.

Author: lgrammel

.changeset/calm-eels-obey.md

@@ -0,0 +1,5 @@
+---
+'@ai-sdk/provider-utils': patch
+---
+
+refactoring: introduce FlexibleSchema

packages/ai/core/tool/index.ts

@@ -14,7 +14,6 @@ export type {
   Tool,
   ToolCallOptions,
   ToolExecuteFunction,
-  ToolInputSchema,
   InferToolInput,
   InferToolOutput,
 } from '@ai-sdk/provider-utils';

packages/ai/core/tool/mcp/types.ts

@@ -1,6 +1,6 @@
 import { z } from 'zod';
 import { JSONObject } from '@ai-sdk/provider';
-import { Tool, ToolInputSchema } from '@ai-sdk/provider-utils';
+import { FlexibleSchema, Tool } from '@ai-sdk/provider-utils';
 
 export const LATEST_PROTOCOL_VERSION = '2024-11-05';
 export const SUPPORTED_PROTOCOL_VERSIONS = [
@@ -9,7 +9,7 @@ export const SUPPORTED_PROTOCOL_VERSIONS = [
 ];
 
 export type ToolSchemas =
-  | Record<string, { inputSchema: ToolInputSchema<JSONObject | unknown> }>
+  | Record<string, { inputSchema: FlexibleSchema<JSONObject | unknown> }>
   | 'automatic'
   | undefined;
 
@@ -21,14 +21,14 @@ type MappedTool<T extends Tool | JSONObject, OUTPUT extends any> =
       : never;
 
 export type McpToolSet<TOOL_SCHEMAS extends ToolSchemas = 'automatic'> =
-  TOOL_SCHEMAS extends Record<string, { inputSchema: ToolInputSchema<unknown> }>
+  TOOL_SCHEMAS extends Record<string, { inputSchema: FlexibleSchema<unknown> }>
     ? {
         [K in keyof TOOL_SCHEMAS]: MappedTool<TOOL_SCHEMAS[K], CallToolResult> &
           Required<
             Pick<MappedTool<TOOL_SCHEMAS[K], CallToolResult>, 'execute'>
           >;
       }
-    : McpToolSet<Record<string, { inputSchema: ToolInputSchema<unknown> }>>;
+    : McpToolSet<Record<string, { inputSchema: FlexibleSchema<unknown> }>>;
 
 const ClientOrServerImplementationSchema = z
   .object({

packages/ai/core/tool/tool.test-d.ts

@@ -3,7 +3,7 @@ import {
   Tool,
   ToolCallOptions,
   ToolExecuteFunction,
-  ToolInputSchema,
+  FlexibleSchema,
 } from '@ai-sdk/provider-utils';
 import { tool } from './tool';
 
@@ -26,7 +26,7 @@ describe('tool helper', () => {
     expectTypeOf(toolType.execute).toEqualTypeOf<undefined>();
     expectTypeOf(toolType.execute).not.toEqualTypeOf<Function>();
     expectTypeOf(toolType.inputSchema).toEqualTypeOf<
-      ToolInputSchema<{ number: number }>
+      FlexibleSchema<{ number: number }>
     >();
   });
 
@@ -61,7 +61,7 @@ describe('tool helper', () => {
     >();
     expectTypeOf(toolType.execute).not.toEqualTypeOf<undefined>();
     expectTypeOf(toolType.inputSchema).toEqualTypeOf<
-      ToolInputSchema<{ number: number }>
+      FlexibleSchema<{ number: number }>
     >();
   });
 });

packages/provider-utils/src/index.ts

@@ -22,7 +22,13 @@ export {
 export * from './remove-undefined-entries';
 export * from './resolve';
 export * from './response-handler';
-export { asSchema, jsonSchema, type InferSchema, type Schema } from './schema';
+export {
+  asSchema,
+  jsonSchema,
+  type FlexibleSchema,
+  type InferSchema,
+  type Schema,
+} from './schema';
 export * from './uint8-utils';
 export * from './validate-types';
 export * from './validator';

packages/provider-utils/src/parse-json.ts

@@ -29,9 +29,9 @@ export async function parseJSON(options: {
  * @returns {Promise<T>} - The parsed object.
  */
 export async function parseJSON<
-  SCHEMA extends z4.$ZodType | z3.Schema | Validator,
-  VALUE = InferSchema<SCHEMA>,
->(options: { text: string; schema: SCHEMA }): Promise<VALUE>;
+  VALIDATOR extends z4.$ZodType | z3.Schema | Validator,
+  VALUE = InferSchema<VALIDATOR>,
+>(options: { text: string; schema: VALIDATOR }): Promise<VALUE>;
 export async function parseJSON<
   SCHEMA extends z4.$ZodType | z3.Schema | Validator,
   VALUE = InferSchema<SCHEMA>,

packages/provider-utils/src/provider-defined-client-tool-factory.ts

@@ -1,4 +1,5 @@
-import { Tool, ToolExecuteFunction, ToolInputSchema } from './types/tool';
+import { Tool, ToolExecuteFunction } from './types/tool';
+import { FlexibleSchema } from './schema';
 
 export type ProviderDefinedClientToolFactory<INPUT, ARGS extends object> = <
   OUTPUT,
@@ -20,7 +21,7 @@ export function createProviderDefinedClientToolFactory<
   inputSchema,
 }: {
   id: string;
-  inputSchema: ToolInputSchema<INPUT>;
+  inputSchema: FlexibleSchema<INPUT>;
 }): ProviderDefinedClientToolFactory<INPUT, ARGS> {
   return <OUTPUT>({
     execute,

packages/provider-utils/src/schema.ts

@@ -26,6 +26,8 @@ export type Schema<OBJECT = unknown> = Validator<OBJECT> & {
   readonly jsonSchema: JSONSchema7;
 };
 
+export type FlexibleSchema<T> = z4.$ZodType<T> | z3.Schema<T> | Schema<T>;
+
 export type InferSchema<SCHEMA> = SCHEMA extends z3.Schema
   ? z3.infer<SCHEMA>
   : SCHEMA extends z4.$ZodType

packages/provider-utils/src/types/index.ts

@@ -18,7 +18,6 @@ export type {
   Tool,
   ToolCallOptions,
   ToolExecuteFunction,
-  ToolInputSchema,
   InferToolInput,
   InferToolOutput,
 } from './tool';

packages/provider-utils/src/types/tool.ts

@@ -1,18 +1,7 @@
-import {
-  JSONObject,
-  JSONValue,
-  LanguageModelV2ToolResultPart,
-} from '@ai-sdk/provider';
-import * as z3 from 'zod/v3';
-import * as z4 from 'zod/v4/core';
-import { Schema } from '../schema';
+import { JSONValue, LanguageModelV2ToolResultPart } from '@ai-sdk/provider';
+import { FlexibleSchema } from '../schema';
 import { ModelMessage } from './model-message';
 
-export type ToolInputSchema<T = JSONObject> =
-  | z4.$ZodType<T>
-  | z3.Schema<T>
-  | Schema<T>;
-
 export interface ToolCallOptions {
   /**
    * The ID of the tool call. You can use it e.g. when sending tool-call related information with stream data.
@@ -66,7 +55,7 @@ The schema of the input that the tool expects. The language model will use this
 It is also used to validate the output of the language model.
 Use descriptions to make the input understandable for the language model.
    */
-    inputSchema: ToolInputSchema<INPUT>;
+    inputSchema: FlexibleSchema<INPUT>;
   }
 > &
   NeverOptional<