@mastra/schema-compat 0.10.2-alpha.2 → 0.10.2-alpha.3

package/src/schema-compatibility.ts

@@ -1,6 +1,7 @@
 import type { Schema, LanguageModelV1 } from 'ai';
 import type { JSONSchema7 } from 'json-schema';
-import { z } from 'zod';
+import { z, ZodOptional, ZodObject, ZodArray, ZodUnion, ZodString, ZodNumber, ZodDate, ZodDefault } from 'zod';
+import type { ZodTypeAny } from 'zod';
 import type { Targets } from 'zod-to-json-schema';
 import { convertZodSchemaToAISDKSchema } from './utils';
 
@@ -26,6 +27,15 @@ export const ALL_NUMBER_CHECKS = [
  */
 export const ALL_ARRAY_CHECKS = ['min', 'max', 'length'] as const;
 
+export const isOptional = (v: ZodTypeAny): v is ZodOptional<any> => v instanceof ZodOptional;
+export const isObj = (v: ZodTypeAny): v is ZodObject<any, any, any> => v instanceof ZodObject;
+export const isArr = (v: ZodTypeAny): v is ZodArray<any, any> => v instanceof ZodArray;
+export const isUnion = (v: ZodTypeAny): v is ZodUnion<[ZodTypeAny, ...ZodTypeAny[]]> => v instanceof ZodUnion;
+export const isString = (v: ZodTypeAny): v is ZodString => v instanceof ZodString;
+export const isNumber = (v: ZodTypeAny): v is ZodNumber => v instanceof ZodNumber;
+export const isDate = (v: ZodTypeAny): v is ZodDate => v instanceof ZodDate;
+export const isDefault = (v: ZodTypeAny): v is ZodDefault<any> => v instanceof ZodDefault;
+
 /**
  * Zod types that are not supported by most AI model providers and should be avoided.
  * @constant
@@ -214,30 +224,7 @@ export abstract class SchemaCompatLayer {
    * @returns The processed Zod type
    * @abstract
    */
-  abstract processZodType<T extends z.AnyZodObject>(value: z.ZodTypeAny): ShapeValue<T>;
-
-  /**
-   * Applies compatibility transformations to a Zod object schema.
-   *
-   * @param zodSchema - The Zod object schema to transform
-   * @returns Object containing the transformed schema
-   * @private
-   */
-  private applyZodSchemaCompatibility(zodSchema: z.AnyZodObject): {
-    schema: z.AnyZodObject;
-  } {
-    const newSchema = z.object(
-      Object.entries<z.ZodTypeAny>(zodSchema.shape || {}).reduce(
-        (acc, [key, value]) => ({
-          ...acc,
-          [key]: this.processZodType<any>(value),
-        }),
-        {},
-      ),
-    );
-
-    return { schema: newSchema };
-  }
+  abstract processZodType(value: ZodTypeAny): ZodTypeAny;
 
   /**
    * Default handler for Zod object types. Recursively processes all properties in the object.
@@ -245,22 +232,25 @@ export abstract class SchemaCompatLayer {
    * @param value - The Zod object to process
    * @returns The processed Zod object
    */
-  public defaultZodObjectHandler<T extends z.AnyZodObject>(value: z.ZodTypeAny): ShapeValue<T> {
-    const zodObject = value as z.ZodObject<any, any, any>;
-    const processedShape = Object.entries(zodObject.shape || {}).reduce<Record<string, z.ZodTypeAny>>(
-      (acc, [key, propValue]) => {
-        const typedPropValue = propValue as z.ZodTypeAny;
-        const processedValue = this.processZodType<T>(typedPropValue);
-        acc[key] = processedValue;
-        return acc;
-      },
-      {},
-    );
-    let result = z.object(processedShape);
+  public defaultZodObjectHandler(value: ZodObject<any, any, any>): ZodObject<any, any, any> {
+    const processedShape = Object.entries(value.shape).reduce<Record<string, ZodTypeAny>>((acc, [key, propValue]) => {
+      acc[key] = this.processZodType(propValue as ZodTypeAny);
+      return acc;
+    }, {});
+
+    let result: ZodObject<any, any, any> = z.object(processedShape);
+
+    if (value._def.unknownKeys === 'strict') {
+      result = result.strict();
+    }
+    if (value._def.catchall && !(value._def.catchall instanceof z.ZodNever)) {
+      result = result.catchall(value._def.catchall);
+    }
+
     if (value.description) {
       result = result.describe(value.description);
     }
-    return result as ShapeValue<T>;
+    return result;
   }
 
   /**
@@ -314,40 +304,46 @@ export abstract class SchemaCompatLayer {
    * @param handleChecks - Array constraints to convert to descriptions vs keep as validation
    * @returns The processed Zod array
    */
-  public defaultZodArrayHandler<T extends z.AnyZodObject>(
-    value: z.ZodTypeAny,
+  public defaultZodArrayHandler(
+    value: ZodArray<any, any>,
     handleChecks: readonly ArrayCheckType[] = ALL_ARRAY_CHECKS,
-  ): ShapeValue<T> {
-    const zodArray = (value as z.ZodArray<any>)._def;
-    const arrayType = zodArray.type;
-    const constraints: ArrayConstraints = {};
-    if (zodArray.minLength?.value !== undefined && handleChecks.includes('min')) {
-      constraints.minLength = zodArray.minLength.value;
-    }
-    if (zodArray.maxLength?.value !== undefined && handleChecks.includes('max')) {
-      constraints.maxLength = zodArray.maxLength.value;
-    }
-    if (zodArray.exactLength?.value !== undefined && handleChecks.includes('length')) {
-      constraints.exactLength = zodArray.exactLength.value;
-    }
-    const processedType =
-      arrayType._def.typeName === 'ZodObject' ? this.processZodType<T>(arrayType as z.ZodTypeAny) : arrayType;
+  ): ZodArray<any, any> {
+    const zodArrayDef = value._def;
+    const processedType = this.processZodType(zodArrayDef.type);
+
     let result = z.array(processedType);
-    if (zodArray.minLength?.value !== undefined && !handleChecks.includes('min')) {
-      result = result.min(zodArray.minLength.value);
+
+    const constraints: ArrayConstraints = {};
+
+    if (zodArrayDef.minLength?.value !== undefined) {
+      if (handleChecks.includes('min')) {
+        constraints.minLength = zodArrayDef.minLength.value;
+      } else {
+        result = result.min(zodArrayDef.minLength.value);
+      }
     }
-    if (zodArray.maxLength?.value !== undefined && !handleChecks.includes('max')) {
-      result = result.max(zodArray.maxLength.value);
+
+    if (zodArrayDef.maxLength?.value !== undefined) {
+      if (handleChecks.includes('max')) {
+        constraints.maxLength = zodArrayDef.maxLength.value;
+      } else {
+        result = result.max(zodArrayDef.maxLength.value);
+      }
     }
-    if (zodArray.exactLength?.value !== undefined && !handleChecks.includes('length')) {
-      result = result.length(zodArray.exactLength.value);
+
+    if (zodArrayDef.exactLength?.value !== undefined) {
+      if (handleChecks.includes('length')) {
+        constraints.exactLength = zodArrayDef.exactLength.value;
+      } else {
+        result = result.length(zodArrayDef.exactLength.value);
+      }
     }
 
     const description = this.mergeParameterDescription(value.description, constraints);
     if (description) {
       result = result.describe(description);
     }
-    return result as ShapeValue<T>;
+    return result;
   }
 
   /**
@@ -357,15 +353,14 @@ export abstract class SchemaCompatLayer {
    * @returns The processed Zod union
    * @throws Error if union has fewer than 2 options
    */
-  public defaultZodUnionHandler<T extends z.AnyZodObject>(value: z.ZodTypeAny): ShapeValue<T> {
-    const zodUnion = value as z.ZodUnion<[z.ZodTypeAny, ...z.ZodTypeAny[]]>;
-    const processedOptions = zodUnion._def.options.map((option: z.ZodTypeAny) => this.processZodType<T>(option));
+  public defaultZodUnionHandler(value: ZodUnion<[ZodTypeAny, ...ZodTypeAny[]]>): ZodTypeAny {
+    const processedOptions = value._def.options.map((option: ZodTypeAny) => this.processZodType(option));
     if (processedOptions.length < 2) throw new Error('Union must have at least 2 options');
-    let result = z.union(processedOptions as [z.ZodTypeAny, z.ZodTypeAny, ...z.ZodTypeAny[]]);
+    let result = z.union(processedOptions as [ZodTypeAny, ZodTypeAny, ...ZodTypeAny[]]);
     if (value.description) {
       result = result.describe(value.description);
     }
-    return result as ShapeValue<T>;
+    return result;
   }
 
   /**
@@ -375,13 +370,12 @@ export abstract class SchemaCompatLayer {
    * @param handleChecks - String constraints to convert to descriptions vs keep as validation
    * @returns The processed Zod string
    */
-  public defaultZodStringHandler<T extends z.AnyZodObject>(
-    value: z.ZodTypeAny,
+  public defaultZodStringHandler(
+    value: ZodString,
     handleChecks: readonly StringCheckType[] = ALL_STRING_CHECKS,
-  ): ShapeValue<T> {
-    const zodString = value as z.ZodString;
+  ): ZodString {
     const constraints: StringConstraints = {};
-    const checks = zodString._def.checks || [];
+    const checks = value._def.checks || [];
     type ZodStringCheck = (typeof checks)[number];
     const newChecks: ZodStringCheck[] = [];
     for (const check of checks) {
@@ -437,7 +431,7 @@ export abstract class SchemaCompatLayer {
     if (description) {
       result = result.describe(description);
     }
-    return result as ShapeValue<T>;
+    return result;
   }
 
   /**
@@ -447,13 +441,12 @@ export abstract class SchemaCompatLayer {
    * @param handleChecks - Number constraints to convert to descriptions vs keep as validation
    * @returns The processed Zod number
    */
-  public defaultZodNumberHandler<T extends z.AnyZodObject>(
-    value: z.ZodTypeAny,
+  public defaultZodNumberHandler(
+    value: ZodNumber,
     handleChecks: readonly NumberCheckType[] = ALL_NUMBER_CHECKS,
-  ): ShapeValue<T> {
-    const zodNumber = value as z.ZodNumber;
+  ): ZodNumber {
     const constraints: NumberConstraints = {};
-    const checks = zodNumber._def.checks || [];
+    const checks = value._def.checks || [];
     type ZodNumberCheck = (typeof checks)[number];
     const newChecks: ZodNumberCheck[] = [];
     for (const check of checks) {
@@ -501,7 +494,7 @@ export abstract class SchemaCompatLayer {
     if (description) {
       result = result.describe(description);
     }
-    return result as ShapeValue<T>;
+    return result;
   }
 
   /**
@@ -510,10 +503,9 @@ export abstract class SchemaCompatLayer {
    * @param value - The Zod date to process
    * @returns A Zod string schema representing the date in ISO format
    */
-  public defaultZodDateHandler<T extends z.AnyZodObject>(value: z.ZodTypeAny): ShapeValue<T> {
-    const zodDate = value as z.ZodDate;
+  public defaultZodDateHandler(value: ZodDate): ZodString {
     const constraints: DateConstraints = {};
-    const checks = zodDate._def.checks || [];
+    const checks = value._def.checks || [];
     type ZodDateCheck = (typeof checks)[number];
     const newChecks: ZodDateCheck[] = [];
     for (const check of checks) {
@@ -542,7 +534,7 @@ export abstract class SchemaCompatLayer {
     if (description) {
       result = result.describe(description);
     }
-    return result as ShapeValue<T>;
+    return result;
   }
 
   /**
@@ -552,14 +544,14 @@ export abstract class SchemaCompatLayer {
    * @param handleTypes - Types that should be processed vs passed through
    * @returns The processed Zod optional
    */
-  public defaultZodOptionalHandler<T extends z.AnyZodObject>(
-    value: z.ZodTypeAny,
+  public defaultZodOptionalHandler(
+    value: ZodOptional<any>,
     handleTypes: readonly AllZodType[] = SUPPORTED_ZOD_TYPES,
-  ): ShapeValue<T> {
+  ): ZodTypeAny {
     if (handleTypes.includes(value._def.innerType._def.typeName as AllZodType)) {
       return this.processZodType(value._def.innerType).optional();
     } else {
-      return value as ShapeValue<T>;
+      return value;
     }
   }
 
@@ -569,10 +561,10 @@ export abstract class SchemaCompatLayer {
    * @param zodSchema - The Zod object schema to process
    * @returns An AI SDK Schema with provider-specific compatibility applied
    */
-  public processToAISDKSchema(zodSchema: z.AnyZodObject): Schema {
-    const { schema } = this.applyZodSchemaCompatibility(zodSchema);
+  public processToAISDKSchema(zodSchema: z.ZodSchema): Schema {
+    const processedSchema = this.processZodType(zodSchema);
 
-    return convertZodSchemaToAISDKSchema(schema, this.getSchemaTarget());
+    return convertZodSchemaToAISDKSchema(processedSchema, this.getSchemaTarget());
   }
 
   /**
@@ -581,7 +573,7 @@ export abstract class SchemaCompatLayer {
    * @param zodSchema - The Zod object schema to process
    * @returns A JSONSchema7 object with provider-specific compatibility applied
    */
-  public processToJSONSchema(zodSchema: z.AnyZodObject): JSONSchema7 {
+  public processToJSONSchema(zodSchema: z.ZodSchema): JSONSchema7 {
     return this.processToAISDKSchema(zodSchema).jsonSchema;
   }
 }

package/src/utils.test.ts

@@ -31,6 +31,12 @@ class MockSchemaCompatibility extends SchemaCompatLayer {
     if (value._def.typeName === 'ZodString') {
       return z.string().describe('processed string');
     }
+    if (value instanceof z.ZodObject) {
+      return this.defaultZodObjectHandler(value);
+    }
+    if (value instanceof z.ZodArray) {
+      return this.defaultZodArrayHandler(value);
+    }
     return value;
   }
 }
@@ -97,6 +103,13 @@ describe('Builder Functions', () => {
       expect(result).toHaveProperty('jsonSchema');
       expect(result.jsonSchema).toHaveProperty('properties');
     });
+
+    it('should handle array schemas', () => {
+      const zodSchema = z.array(z.string());
+      const result = convertZodSchemaToAISDKSchema(zodSchema);
+      expect(result.jsonSchema.type).toBe('array');
+      expect((result.jsonSchema.items as any)?.type).toBe('string');
+    });
   });
 
   describe('convertSchemaToZod', () => {
@@ -158,6 +171,20 @@ describe('Builder Functions', () => {
       const parseResult = result.safeParse(validData);
       expect(parseResult.success).toBe(true);
     });
+
+    it('should convert AI SDK array schema to Zod', () => {
+      const aiSchema: Schema = jsonSchema({
+        type: 'array',
+        items: {
+          type: 'string',
+        },
+      });
+
+      const result = convertSchemaToZod(aiSchema);
+
+      expect(result).toBeInstanceOf(z.ZodArray);
+      expect((result as z.ZodArray<any>).element).toBeInstanceOf(z.ZodString);
+    });
   });
 
   describe('applyCompatLayer', () => {
@@ -275,29 +302,6 @@ describe('Builder Functions', () => {
       expect(result).toHaveProperty('validate');
     });
 
-    it('should convert non-object AI SDK schema correctly', () => {
-      const stringSchema: Schema = jsonSchema({ type: 'string' });
-
-      const result = applyCompatLayer({
-        schema: stringSchema,
-        compatLayers: [mockCompatibility],
-        mode: 'aiSdkSchema',
-      });
-
-      expect(result).toHaveProperty('jsonSchema');
-      expect(result).toHaveProperty('validate');
-
-      // Verify the schema structure shows the string was wrapped in a 'value' property
-      const resultSchema = (result as Schema).jsonSchema;
-      expect(resultSchema).toHaveProperty('type', 'object');
-      expect(resultSchema).toHaveProperty('properties');
-      expect(resultSchema.properties).toHaveProperty('value');
-
-      // Verify the string property has the description added by the mock compatibility layer
-      const valueProperty = resultSchema.properties!.value as any;
-      expect(valueProperty).toHaveProperty('description', 'processed string');
-    });
-
     it('should handle complex schema with multiple compatLayers', () => {
       const compat1 = new MockSchemaCompatibility(mockModel, false);
       const compat2 = new MockSchemaCompatibility(mockModel, true);
@@ -324,5 +328,107 @@ describe('Builder Functions', () => {
       expect(compat1.processZodType).not.toHaveBeenCalled();
       expect(compat2.processZodType).toHaveBeenCalled();
     });
+
+    it('should process Zod array schema with compatibility', () => {
+      const arraySchema = z.array(z.string());
+
+      const result = applyCompatLayer({
+        schema: arraySchema,
+        compatLayers: [mockCompatibility],
+        mode: 'aiSdkSchema',
+      });
+
+      expect(result.jsonSchema.type).toBe('array');
+      if (
+        result.jsonSchema.items &&
+        !Array.isArray(result.jsonSchema.items) &&
+        typeof result.jsonSchema.items === 'object'
+      ) {
+        expect(result.jsonSchema.items.type).toBe('string');
+        expect(result.jsonSchema.items.description).toBe('processed string');
+      } else {
+        expect.fail('items is not a single schema object');
+      }
+    });
+
+    it('should process AI SDK array schema with compatibility', () => {
+      const aiSchema: Schema = jsonSchema({
+        type: 'array',
+        items: {
+          type: 'string',
+        },
+      });
+
+      const result = applyCompatLayer({
+        schema: aiSchema,
+        compatLayers: [mockCompatibility],
+        mode: 'aiSdkSchema',
+      });
+
+      expect(result.jsonSchema.type).toBe('array');
+      if (
+        result.jsonSchema.items &&
+        !Array.isArray(result.jsonSchema.items) &&
+        typeof result.jsonSchema.items === 'object'
+      ) {
+        expect(result.jsonSchema.items.type).toBe('string');
+        expect(result.jsonSchema.items.description).toBe('processed string');
+      } else {
+        expect.fail('items is not a single schema object');
+      }
+    });
+
+    it('should handle a complex array of objects schema', () => {
+      const complexArraySchema = z.array(
+        z.object({
+          id: z.string(),
+          user: z.object({
+            name: z.string(),
+          }),
+        }),
+      );
+
+      const result = applyCompatLayer({
+        schema: complexArraySchema,
+        compatLayers: [mockCompatibility],
+        mode: 'aiSdkSchema',
+      });
+
+      const { jsonSchema } = result;
+      expect(jsonSchema.type).toBe('array');
+
+      const items = jsonSchema.items;
+      if (items && !Array.isArray(items) && typeof items === 'object') {
+        expect(items.type).toBe('object');
+        expect(items.properties).toHaveProperty('id');
+        expect(items.properties).toHaveProperty('user');
+
+        const idProperty = items.properties!.id as any;
+        expect(idProperty.description).toBe('processed string');
+
+        const userProperty = items.properties!.user as any;
+        expect(userProperty.type).toBe('object');
+        expect(userProperty.properties).toHaveProperty('name');
+
+        const nameProperty = userProperty.properties.name as any;
+        expect(nameProperty.description).toBe('processed string');
+      } else {
+        expect.fail('items is not a single schema object');
+      }
+    });
+
+    it('should handle a scalar zod schema', () => {
+      const scalarSchema = z.string().email();
+
+      const result = applyCompatLayer({
+        schema: scalarSchema,
+        compatLayers: [mockCompatibility],
+        mode: 'aiSdkSchema',
+      });
+
+      const { jsonSchema } = result;
+      expect(jsonSchema.type).toBe('string');
+      expect(jsonSchema.description).toBe('processed string');
+    });
   });
 });

package/src/utils.ts

@@ -1,8 +1,7 @@
 import { jsonSchema } from 'ai';
 import type { Schema } from 'ai';
 import type { JSONSchema7 } from 'json-schema';
-import type { ZodSchema } from 'zod';
-import { z } from 'zod';
+import type { z, ZodSchema } from 'zod';
 import { convertJsonSchemaToZod } from 'zod-from-json-schema';
 import type { JSONSchema as ZodFromJSONSchema_JSONSchema } from 'zod-from-json-schema';
 import type { Targets } from 'zod-to-json-schema';
@@ -118,7 +117,7 @@ export function convertSchemaToZod(schema: Schema | z.ZodSchema): z.ZodType {
  * @returns Processed schema as an AI SDK Schema
  */
 export function applyCompatLayer(options: {
-  schema: Schema | z.AnyZodObject;
+  schema: Schema | z.ZodSchema;
   compatLayers: SchemaCompatLayer[];
   mode: 'aiSdkSchema';
 }): Schema;
@@ -133,7 +132,7 @@ export function applyCompatLayer(options: {
  * @returns Processed schema as a JSONSchema7
  */
 export function applyCompatLayer(options: {
-  schema: Schema | z.AnyZodObject;
+  schema: Schema | z.ZodSchema;
   compatLayers: SchemaCompatLayer[];
   mode: 'jsonSchema';
 }): JSONSchema7;
@@ -178,29 +177,17 @@ export function applyCompatLayer({
   compatLayers,
   mode,
 }: {
-  schema: Schema | z.AnyZodObject;
+  schema: Schema | z.ZodSchema;
   compatLayers: SchemaCompatLayer[];
   mode: 'jsonSchema' | 'aiSdkSchema';
 }): JSONSchema7 | Schema {
-  let zodSchema: z.AnyZodObject;
+  let zodSchema: z.ZodSchema;
 
   if (!isZodType(schema)) {
-    // Convert Schema to ZodObject
-    const convertedSchema = convertSchemaToZod(schema);
-    if (convertedSchema instanceof z.ZodObject) {
-      zodSchema = convertedSchema;
-    } else {
-      // If it's not an object schema, wrap it in an object
-      zodSchema = z.object({ value: convertedSchema });
-    }
+    // Convert non-zod schema to Zod
+    zodSchema = convertSchemaToZod(schema);
   } else {
-    // Ensure it's a ZodObject
-    if (schema instanceof z.ZodObject) {
-      zodSchema = schema;
-    } else {
-      // Wrap non-object schemas in an object
-      zodSchema = z.object({ value: schema });
-    }
+    zodSchema = schema;
   }
 
   for (const compat of compatLayers) {