fix(zod): correctly add $ref definitions for transformed schemas (#1065)

Author: RobertCraigie

src/_vendor/zod-to-json-schema/parseDef.ts

@@ -99,7 +99,7 @@ export function parseDef(
 
   refs.seen.set(def, newItem);
 
-  const jsonSchema = selectParser(def, (def as any).typeName, refs);
+  const jsonSchema = selectParser(def, (def as any).typeName, refs, forceResolution);
 
   if (jsonSchema) {
     addMeta(def, refs, jsonSchema);
@@ -166,7 +166,12 @@ const getRelativePath = (pathA: string[], pathB: string[]) => {
   return [(pathA.length - i).toString(), ...pathB.slice(i)].join('/');
 };
 
-const selectParser = (def: any, typeName: ZodFirstPartyTypeKind, refs: Refs): JsonSchema7Type | undefined => {
+const selectParser = (
+  def: any,
+  typeName: ZodFirstPartyTypeKind,
+  refs: Refs,
+  forceResolution: boolean,
+): JsonSchema7Type | undefined => {
   switch (typeName) {
     case ZodFirstPartyTypeKind.ZodString:
       return parseStringDef(def, refs);
@@ -217,7 +222,7 @@ const selectParser = (def: any, typeName: ZodFirstPartyTypeKind, refs: Refs): Js
     case ZodFirstPartyTypeKind.ZodNever:
       return parseNeverDef();
     case ZodFirstPartyTypeKind.ZodEffects:
-      return parseEffectsDef(def, refs);
+      return parseEffectsDef(def, refs, forceResolution);
     case ZodFirstPartyTypeKind.ZodAny:
       return parseAnyDef();
     case ZodFirstPartyTypeKind.ZodUnknown:

src/_vendor/zod-to-json-schema/parsers/effects.ts

@@ -2,6 +2,10 @@ import { ZodEffectsDef } from 'zod';
 import { JsonSchema7Type, parseDef } from '../parseDef';
 import { Refs } from '../Refs';
 
-export function parseEffectsDef(_def: ZodEffectsDef, refs: Refs): JsonSchema7Type | undefined {
-  return refs.effectStrategy === 'input' ? parseDef(_def.schema._def, refs) : {};
+export function parseEffectsDef(
+  _def: ZodEffectsDef,
+  refs: Refs,
+  forceResolution: boolean,
+): JsonSchema7Type | undefined {
+  return refs.effectStrategy === 'input' ? parseDef(_def.schema._def, refs, forceResolution) : {};
 }

tests/lib/__snapshots__/parser.test.ts.snap

@@ -112,6 +112,37 @@ exports[`.parse() zod recursive schema extraction 2`] = `
 "
 `;
 
+exports[`.parse() zod ref schemas with \`.transform()\` 2`] = `
+"{
+  "id": "chatcmpl-A6zyLEtubMlUvGplOmr92S0mK0kiG",
+  "object": "chat.completion",
+  "created": 1726231553,
+  "model": "gpt-4o-2024-08-06",
+  "choices": [
+    {
+      "index": 0,
+      "message": {
+        "role": "assistant",
+        "content": "{\\"first\\":{\\"baz\\":true},\\"second\\":{\\"baz\\":false}}",
+        "refusal": null
+      },
+      "logprobs": null,
+      "finish_reason": "stop"
+    }
+  ],
+  "usage": {
+    "prompt_tokens": 167,
+    "completion_tokens": 13,
+    "total_tokens": 180,
+    "completion_tokens_details": {
+      "reasoning_tokens": 0
+    }
+  },
+  "system_fingerprint": "fp_143bb8492c"
+}
+"
+`;
+
 exports[`.parse() zod top-level recursive schemas 1`] = `
 "{
   "id": "chatcmpl-9uLhw79ArBF4KsQQOlsoE68m6vh6v",

tests/lib/parser.test.ts

@@ -951,5 +951,119 @@ describe('.parse()', () => {
         }
       `);
     });
+
+    test('ref schemas with `.transform()`', async () => {
+      const Inner = z.object({
+        baz: z.boolean().transform((v) => v ?? true),
+      });
+
+      const Outer = z.object({
+        first: Inner,
+        second: Inner,
+      });
+
+      expect(zodResponseFormat(Outer, 'data').json_schema.schema).toMatchInlineSnapshot(`
+        {
+          "$schema": "http://json-schema.org/draft-07/schema#",
+          "additionalProperties": false,
+          "definitions": {
+            "data": {
+              "additionalProperties": false,
+              "properties": {
+                "first": {
+                  "additionalProperties": false,
+                  "properties": {
+                    "baz": {
+                      "type": "boolean",
+                    },
+                  },
+                  "required": [
+                    "baz",
+                  ],
+                  "type": "object",
+                },
+                "second": {
+                  "$ref": "#/definitions/data_properties_first",
+                },
+              },
+              "required": [
+                "first",
+                "second",
+              ],
+              "type": "object",
+            },
+            "data_properties_first": {
+              "additionalProperties": false,
+              "properties": {
+                "baz": {
+                  "$ref": "#/definitions/data_properties_first_properties_baz",
+                },
+              },
+              "required": [
+                "baz",
+              ],
+              "type": "object",
+            },
+            "data_properties_first_properties_baz": {
+              "type": "boolean",
+            },
+          },
+          "properties": {
+            "first": {
+              "additionalProperties": false,
+              "properties": {
+                "baz": {
+                  "type": "boolean",
+                },
+              },
+              "required": [
+                "baz",
+              ],
+              "type": "object",
+            },
+            "second": {
+              "$ref": "#/definitions/data_properties_first",
+            },
+          },
+          "required": [
+            "first",
+            "second",
+          ],
+          "type": "object",
+        }
+      `);
+
+      const completion = await makeSnapshotRequest(
+        (openai) =>
+          openai.beta.chat.completions.parse({
+            model: 'gpt-4o-2024-08-06',
+            messages: [
+              {
+                role: 'user',
+                content: 'can you generate fake data matching the given response format?',
+              },
+            ],
+            response_format: zodResponseFormat(Outer, 'fakeData'),
+          }),
+        2,
+      );
+
+      expect(completion.choices[0]?.message).toMatchInlineSnapshot(`
+        {
+          "content": "{"first":{"baz":true},"second":{"baz":false}}",
+          "parsed": {
+            "first": {
+              "baz": true,
+            },
+            "second": {
+              "baz": false,
+            },
+          },
+          "refusal": null,
+          "role": "assistant",
+          "tool_calls": [],
+        }
+      `);
+    });
   });
 });