zod plugin outputs invalid schema with nested discriminated unions

[hnykda]

Description

Hello,

I want to apologise for non-reproducible example, but I think it's still worth reporting. Feel free to delete this if you think it's not good enough.

I have a quite dynamic schema, but fully determined by one or two discriminated unions. Unfortunately, the generated zod schema is invalid:

export const zSubmitTaskBody = z.object({
  payload: z.union([
    z
      .object({
        task_type: z.literal("find_number").optional(),
      })
      .merge(zFindNumberRequestParams),
    z
      .object({
        task_type: z.literal("find_sources").optional(),
      })
      .merge(zFindSourcesRequestParams),
    z
      .object({
        task_type: z.literal("populate_reference_class").optional(),
      })
      .merge(zPopulateReferenceClassRequestParams),
    z
      .object({
        task_type: z.literal("chat").optional(),
      })
      .merge(zChatRequestPayload),
    z
      .object({
        task_type: z.literal("agent").optional(),
      })
      .merge(zAgentRequestPayload),
    z
      .object({
        task_type: z.literal("filter").optional(),
      })
      .merge(zFilterRequest),
    z
      .object({
        task_type: z.literal("join").optional(),
      })
      .merge(zJoinRequest),
  ]),
  session_id: z.string().uuid(),
  task_id: z.string().uuid().optional(),
});

export const zChatRequestPayload = z.union([
  z
    .object({
      processing_mode: z.literal("standalone").optional(),
    })
    .merge(
      z.object({
        task_type: z.literal("chat"),
        processing_mode: z.literal("standalone"),
        context_artifacts: z
          .union([z.array(z.string().uuid()), z.null()])
          .optional(),
        query: z.object({
          prompt: z.string(),
          sources: z.union([z.array(z.string()), z.null()]).optional(),
          response_schema_instructions: z
            .union([z.string(), z.null()])
            .optional(),
        }),
      }),
    ),
  z
    .object({
      processing_mode: z.literal("map").optional(),
    })
    .merge(
      z.object({
        task_type: z.literal("chat"),
        input_artifacts: z.array(z.string().uuid()),
        context_artifacts: z
          .union([z.array(z.string().uuid()), z.null()])
          .optional(),
        query: z.object({
          task_type: z.literal("chat"),
          prompt: z.string(),
          source_selection_instructions: z
            .union([z.string(), z.null()])
            .optional(),
          response_schema_instructions: z
            .union([z.string(), z.null()])
            .optional(),
        }),
        join_with_input: z.boolean().optional().default(false),
        processing_mode: z.literal("map"),
      }),
    ),
  z
    .object({
      processing_mode: z.literal("reduce").optional(),
    })
    .merge(
      z.object({
        task_type: z.literal("chat"),
        input_artifacts: z.array(z.string().uuid()),
        context_artifacts: z
          .union([z.array(z.string().uuid()), z.null()])
          .optional(),
        query: z.object({
          task_type: z.literal("chat"),
          prompt: z.string(),
          source_selection_instructions: z
            .union([z.string(), z.null()])
            .optional(),
          response_schema_instructions: z
            .union([z.string(), z.null()])
            .optional(),
        }),
        join_with_input: z.boolean().optional().default(false),
        processing_mode: z.literal("reduce"),
      }),
    ),
  z
    .object({
      processing_mode: z.literal("expand").optional(),
    })
    .merge(
      z.object({
        task_type: z.literal("chat"),
        processing_mode: z.literal("expand"),
        context_artifacts: z
          .union([z.array(z.string().uuid()), z.null()])
          .optional(),
        query: z.object({
          prompt: z.string(),
          sources: z.union([z.array(z.string()), z.null()]).optional(),
          items_to_list: z.string(),
          response_schema_instructions: z
            .union([z.string(), z.null()])
            .optional(),
        }),
      }),
    ),
  z
    .object({
      processing_mode: z.literal("map_expand").optional(),
    })
    .merge(
      z.object({
        task_type: z.literal("chat"),
        processing_mode: z.literal("map_expand"),
        query: z.object({
          task_type: z.literal("chat"),
          prompt: z.string(),
          source_selection_instructions: z
            .union([z.string(), z.null()])
            .optional(),
          items_to_list: z.string(),
          response_schema_instructions: z
            .union([z.string(), z.null()])
            .optional(),
        }),
        input_artifacts: z.array(z.string().uuid()),
        context_artifacts: z
          .union([z.array(z.string().uuid()), z.null()])
          .optional(),
        join_with_input: z.boolean().optional().default(false),
      }),
    ),
]);


export const zAgentRequestPayload = z.union([
  z
    .object({
      processing_mode: z.literal("standalone").optional(),
    })
    .merge(
      z.object({
        task_type: z.literal("agent"),
        processing_mode: z.literal("standalone"),
        context_artifacts: z
          .union([z.array(z.string().uuid()), z.null()])
          .optional(),
        query: z.object({
          task_type: z.literal("agent"),
          task: z.string(),
          number_of_steps: z.number().int().optional().default(5),
          response_schema_instructions: z
            .union([z.string(), z.null()])
            .optional(),
        }),
      }),
    ),
  z
    .object({
      processing_mode: z.literal("map").optional(),
    })
    .merge(
      z.object({
        task_type: z.literal("agent"),
        input_artifacts: z
          .union([z.array(z.string().uuid()), z.null()])
          .optional(),
        context_artifacts: z
          .union([z.array(z.string().uuid()), z.null()])
          .optional(),
        query: z.object({
          task_type: z.literal("agent"),
          task_creation_instructions: z.string(),
          response_schema_instructions: z
            .union([z.string(), z.null()])
            .optional(),
          number_of_steps: z.number().int().optional().default(5),
        }),
        join_with_input: z.boolean().optional().default(false),
        processing_mode: z.literal("map"),
      }),
    ),
  z
    .object({
      processing_mode: z.literal("reduce").optional(),
    })
    .merge(
      z.object({
        task_type: z.literal("agent"),
        input_artifacts: z
          .union([z.array(z.string().uuid()), z.null()])
          .optional(),
        context_artifacts: z
          .union([z.array(z.string().uuid()), z.null()])
          .optional(),
        query: z.object({
          task_type: z.literal("agent"),
          task_creation_instructions: z.string(),
          response_schema_instructions: z
            .union([z.string(), z.null()])
            .optional(),
          number_of_steps: z.number().int().optional().default(5),
        }),
        join_with_input: z.boolean().optional().default(false),
        processing_mode: z.literal("reduce"),
      }),
    ),
]);

the error is:

    z
Argument of type 'ZodUnion<[ZodObject<extendShape<{ processing_mode: ZodOptional<ZodLiteral<"standalone">>; }, { task_type: ZodLiteral<"chat">; processing_mode: ZodLiteral<"standalone">; context_artifacts: ZodOptional<...>; query: ZodObject<...>; }>, "strip", ZodTypeAny, { ...; }, { ...; }>, ZodObject<...>, ZodObject<...>, ZodObject<...' is not assignable to parameter of type 'AnyZodObject'.
  Type 'ZodUnion<[ZodObject<extendShape<{ processing_mode: ZodOptional<ZodLiteral<"standalone">>; }, { task_type: ZodLiteral<"chat">; processing_mode: ZodLiteral<"standalone">; context_artifacts: ZodOptional<...>; query: ZodObject<...>; }>, "strip", ZodTypeAny, { ...; }, { ...; }>, ZodObject<...>, ZodObject<...>, ZodObject<...' is missing the following properties from type 'ZodObject<any, any, any, { [x: string]: any; }, { [x: string]: any; }>': _cached, _getCached, shape, strict, and 14 more.ts(2345)

The way I work around it is running sed -i '' '/export const zSubmitTaskBody/,/});/d' src/lib/engine_api/generated/zod.gen.ts after the generation. Would be great to whitelist stuff on the individual plugin.

Reproducible example or configuration

Sorry, I can't share the whole schema and spec.

OpenAPI specification (optional)

No response

System information (optional)

No response

[mrlubos]

Hey, what's the correct solution to make the schema valid?

[hnykda]

I think this one: colinhacks/zod#147 (comment)

[jagregory]

I'm also hitting this issue but with a much simpler schema: https://stackblitz.com/edit/hey-api-client-fetch-example-qkcgdpvu?file=src%2Fclient%2Fzod.gen.ts (see: schema.json and the zod.gen.ts files).

It appears if an allOf has an item that's a oneOf before an object then the resulting zod code will incorrectly attempt to call .merge on an enum. If the object is before the oneOf in the allOf then the .merge is called on the object which is ok. I assume if all the items in an allOf are complex then it wouldn't work either.

FWIW, this schema is generated from a fairly trivial set of serde'd rust structs just using #[serde(flatten)] on a field with an enum type, which I wouldn't expect to be particularly uncommon.

[mrlubos]

@jagregory this will be fixed in v0.66.7. @hnykda since I can't test against your spec, you'll have to let me know if it still doesn't work!