fix(langchain/zod): add caching to interopZodTransformInputSchema to respect references

libs/langchain-core/src/utils/types/tests/zod.test.ts

@@ -1525,6 +1525,209 @@ describe("Zod utility functions", () => {
         expect(elementShape.name).toBeInstanceOf(z4.ZodString);
         expect(elementShape.age).toBeInstanceOf(z4.ZodNumber);
       });
+
+      it("should cache and reuse sanitized sub-schemas when same schema is used in multiple properties", () => {
+        // Create a shared sub-schema that will be used in multiple places
+        const addressSchema = z4.object({
+          street: z4.string().transform((s) => s.toUpperCase()),
+          city: z4.string(),
+        });
+
+        const inputSchema = z4.object({
+          homeAddress: addressSchema,
+          workAddress: addressSchema,
+          billingAddress: addressSchema,
+        });
+
+        const result = interopZodTransformInputSchema(inputSchema, true);
+
+        expect(result).toBeInstanceOf(z4.ZodObject);
+        const resultShape = getInteropZodObjectShape(result as any);
+
+        // All three address properties should reference the exact same sanitized instance
+        expect(resultShape.homeAddress).toBe(resultShape.workAddress);
+        expect(resultShape.workAddress).toBe(resultShape.billingAddress);
+        expect(resultShape.homeAddress).toBe(resultShape.billingAddress);
+      });
+
+      it("should cache and reuse sanitized sub-schemas when same schema is used in array and as property", () => {
+        // Create a shared user schema
+        const userSchema = z4.object({
+          name: z4.string().transform((s) => s.toUpperCase()),
+          age: z4.number(),
+        });
+
+        const inputSchema = z4.object({
+          primaryUser: userSchema,
+          secondaryUser: userSchema,
+          allUsers: z4.array(userSchema),
+        });
+
+        const result = interopZodTransformInputSchema(inputSchema, true);
+
+        expect(result).toBeInstanceOf(z4.ZodObject);
+        const resultShape = getInteropZodObjectShape(result as any);
+
+        // Extract the array element schema
+        const arrayElement = (resultShape.allUsers as any)._zod.def.element;
+
+        // All user schema references should be the exact same instance
+        expect(resultShape.primaryUser).toBe(resultShape.secondaryUser);
+        expect(resultShape.primaryUser).toBe(arrayElement);
+        expect(resultShape.secondaryUser).toBe(arrayElement);
+      });
+
+      it("should cache and reuse deeply nested shared schemas", () => {
+        // Create a deeply nested shared schema
+        const tagSchema = z4.object({
+          name: z4.string().transform((s) => s.toLowerCase()),
+          color: z4.string(),
+        });
+
+        const postSchema = z4.object({
+          title: z4.string(),
+          tags: z4.array(tagSchema),
+          primaryTag: tagSchema,
+        });
+
+        const inputSchema = z4.object({
+          posts: z4.array(postSchema),
+          featuredPost: postSchema,
+        });
+
+        const result = interopZodTransformInputSchema(inputSchema, true);
+
+        expect(result).toBeInstanceOf(z4.ZodObject);
+        const resultShape = getInteropZodObjectShape(result as any);
+
+        // Extract the post schema from the array
+        const postsArrayElement = (resultShape.posts as any)._zod.def.element;
+
+        // The featuredPost and the array element should be the same instance
+        expect(resultShape.featuredPost).toBe(postsArrayElement);
+
+        // Extract tag schemas from both places
+        const featuredPostShape = getInteropZodObjectShape(
+          resultShape.featuredPost as any
+        );
+        const postsArrayElementShape = getInteropZodObjectShape(
+          postsArrayElement as any
+        );
+
+        // Tag schemas should also be reused
+        const featuredPostTagsArray = (featuredPostShape.tags as any)._zod.def
+          .element;
+        const postsArrayTagsArray = (postsArrayElementShape.tags as any)._zod
+          .def.element;
+
+        expect(featuredPostShape.primaryTag).toBe(featuredPostTagsArray);
+        expect(postsArrayElementShape.primaryTag).toBe(postsArrayTagsArray);
+        expect(featuredPostTagsArray).toBe(postsArrayTagsArray);
+      });
+
+      it("should handle caching when same schema appears at different nesting levels", () => {
+        // Create a schema that appears at different nesting levels
+        const metadataSchema = z4.object({
+          key: z4.string(),
+          value: z4.string().transform((s) => s.trim()),
+        });
+
+        const inputSchema = z4.object({
+          topLevelMetadata: metadataSchema,
+          nested: z4.object({
+            nestedMetadata: metadataSchema,
+            deepNested: z4.object({
+              deepMetadata: metadataSchema,
+            }),
+          }),
+          metadataList: z4.array(metadataSchema),
+        });
+
+        const result = interopZodTransformInputSchema(inputSchema, true);
+
+        expect(result).toBeInstanceOf(z4.ZodObject);
+        const resultShape = getInteropZodObjectShape(result as any);
+        const nestedShape = getInteropZodObjectShape(resultShape.nested as any);
+        const deepNestedShape = getInteropZodObjectShape(
+          nestedShape.deepNested as any
+        );
+        const metadataListElement = (resultShape.metadataList as any)._zod.def
+          .element;
+
+        // All metadata schemas should reference the same sanitized instance
+        expect(resultShape.topLevelMetadata).toBe(nestedShape.nestedMetadata);
+        expect(nestedShape.nestedMetadata).toBe(deepNestedShape.deepMetadata);
+        expect(deepNestedShape.deepMetadata).toBe(metadataListElement);
+        expect(resultShape.topLevelMetadata).toBe(metadataListElement);
+      });
+
+      it("should generate JSON schema with $ref for reused sub-schemas", () => {
+        // Create a shared address schema with transforms
+        const addressSchema = z4.object({
+          street: z4.string().transform((s) => s.toUpperCase()),
+          city: z4.string(),
+          zipCode: z4.string(),
+        });
+
+        const inputSchema = z4.object({
+          homeAddress: addressSchema,
+          workAddress: addressSchema,
+          billingAddress: addressSchema,
+          shippingAddresses: z4.array(addressSchema),
+        });
+
+        // Get the sanitized input schema with transforms removed
+        const sanitizedSchema = interopZodTransformInputSchema(
+          inputSchema,
+          true
+        );
+
+        // Generate JSON schema with refs for reused schemas
+        const jsonSchema = z4.toJSONSchema(sanitizedSchema as any, {
+          cycles: "ref",
+          reused: "ref",
+          override(ctx) {
+            ctx.jsonSchema.title = "TestSchema";
+          },
+        });
+
+        // Verify exact JSON schema structure with $ref for reused schemas
+        expect(jsonSchema).toEqual({
+          $schema: "https://json-schema.org/draft/2020-12/schema",
+          title: "TestSchema",
+          type: "object",
+          properties: {
+            homeAddress: { $ref: "#/$defs/__schema0" },
+            workAddress: { $ref: "#/$defs/__schema0" },
+            billingAddress: { $ref: "#/$defs/__schema0" },
+            shippingAddresses: {
+              title: "TestSchema",
+              type: "array",
+              items: { $ref: "#/$defs/__schema0" },
+            },
+          },
+          required: [
+            "homeAddress",
+            "workAddress",
+            "billingAddress",
+            "shippingAddresses",
+          ],
+          additionalProperties: false,
+          $defs: {
+            __schema0: {
+              title: "TestSchema",
+              type: "object",
+              properties: {
+                street: { title: "TestSchema", type: "string" },
+                city: { title: "TestSchema", type: "string" },
+                zipCode: { title: "TestSchema", type: "string" },
+              },
+              required: ["street", "city", "zipCode"],
+              additionalProperties: false,
+            },
+          },
+        });
+      });
     });
 
     it("should throw error for non-schema values", () => {

libs/langchain-core/src/utils/types/zod.ts

@@ -779,23 +779,24 @@ function isZodTransformV4(schema: InteropZodType): schema is z4.$ZodPipe {
   return isZodSchemaV4(schema) && schema._zod.def.type === "pipe";
 }
 
-/**
- * Returns the input type of a Zod transform schema, for both v3 and v4.
- * If the schema is not a transform, returns undefined. If `recursive` is true,
- * recursively processes nested object schemas and arrays of object schemas.
- *
- * @param schema - The Zod schema instance (v3 or v4)
- * @param {boolean} [recursive=false] - Whether to recursively process nested objects/arrays.
- * @returns The input Zod schema of the transform, or undefined if not a transform
- */
-export function interopZodTransformInputSchema(
+function interopZodTransformInputSchemaImpl(
   schema: InteropZodType,
-  recursive = false
+  recursive: boolean,
+  cache: WeakMap<InteropZodType, InteropZodType>
 ): InteropZodType {
+  const cached = cache.get(schema);
+  if (cached !== undefined) {
+    return cached;
+  }
+
   // Zod v3: ._def.schema is the input schema for ZodEffects (transform)
   if (isZodSchemaV3(schema)) {
     if (isZodTransformV3(schema)) {
-      return interopZodTransformInputSchema(schema._def.schema, recursive);
+      return interopZodTransformInputSchemaImpl(
+        schema._def.schema,
+        recursive,
+        cache
+      );
     }
     // TODO: v3 schemas aren't recursively handled here
     // (currently not necessary since zodToJsonSchema handles this)
@@ -806,9 +807,10 @@ export function interopZodTransformInputSchema(
   if (isZodSchemaV4(schema)) {
     let outputSchema: InteropZodType = schema;
     if (isZodTransformV4(schema)) {
-      outputSchema = interopZodTransformInputSchema(
+      outputSchema = interopZodTransformInputSchemaImpl(
         schema._zod.def.in,
-        recursive
+        recursive,
+        cache
       );
     }
     if (recursive) {
@@ -818,9 +820,10 @@ export function interopZodTransformInputSchema(
         for (const [key, keySchema] of Object.entries(
           outputSchema._zod.def.shape
         )) {
-          outputShape[key] = interopZodTransformInputSchema(
+          outputShape[key] = interopZodTransformInputSchemaImpl(
             keySchema,
-            recursive
+            recursive,
+            cache
           ) as z4.$ZodType;
         }
         outputSchema = clone<ZodObjectV4>(outputSchema, {
@@ -830,9 +833,10 @@ export function interopZodTransformInputSchema(
       }
       // Handle nested array schemas
       else if (isZodArrayV4(outputSchema)) {
-        const elementSchema = interopZodTransformInputSchema(
+        const elementSchema = interopZodTransformInputSchemaImpl(
           outputSchema._zod.def.element,
-          recursive
+          recursive,
+          cache
         );
         outputSchema = clone<z4.$ZodArray>(outputSchema, {
           ...outputSchema._zod.def,
@@ -842,12 +846,30 @@ export function interopZodTransformInputSchema(
     }
     const meta = globalRegistry.get(schema);
     if (meta) globalRegistry.add(outputSchema as z4.$ZodType, meta);
+    cache.set(schema, outputSchema);
     return outputSchema;
   }
 
   throw new Error("Schema must be an instance of z3.ZodType or z4.$ZodType");
 }
 
+/**
+ * Returns the input type of a Zod transform schema, for both v3 and v4.
+ * If the schema is not a transform, returns undefined. If `recursive` is true,
+ * recursively processes nested object schemas and arrays of object schemas.
+ *
+ * @param schema - The Zod schema instance (v3 or v4)
+ * @param {boolean} [recursive=false] - Whether to recursively process nested objects/arrays.
+ * @returns The input Zod schema of the transform, or undefined if not a transform
+ */
+export function interopZodTransformInputSchema(
+  schema: InteropZodType,
+  recursive = false
+): InteropZodType {
+  const cache = new WeakMap<InteropZodType, InteropZodType>();
+  return interopZodTransformInputSchemaImpl(schema, recursive, cache);
+}
+
 /**
  * Creates a modified version of a Zod object schema where fields matching a predicate are made optional.
  * Supports both Zod v3 and v4 schemas and preserves the original schema version.