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)