Extending zod schema with branded type produces typescript error

[guillaume-docquier-vention]

I have a brand in a base schema, which I extend to add additional data, like so:

import { z } from "zod";

const agentBaseSchema = z.object({
  id: z.string().uuid().brand("Uuid"),
})

enum AgentType {
  CODE_FREE = "CODE_FREE",
  MACHINE_CODE = "MACHINE_CODE",
}

type AgentBase = z.infer<typeof agentBaseSchema> & {
  type: AgentType 
}

const codeFreeAgentSchema = agentBaseSchema.extend({
  type: z.literal(AgentType.CODE_FREE),
  // Other specific things
})

const machineCodeAgentSchema = agentBaseSchema.extend({
  type: z.literal(AgentType.MACHINE_CODE),
  // Other specific things
})

// ERROR! Type 'string' is not assignable to type string & 'BRAND<"Uuid">'
const agentSchema = z.union([codeFreeAgentSchema, machineCodeAgentSchema]) satisfies z.ZodSchema<AgentBase>

type Agent = z.infer<typeof agentSchema>

TS playground

The agents are consumed as a discriminated union on the type, as you can see. I'm doing the satisfies bit to ensure that new schemas are created properly by other people in the team.

Is there a way to do this? Should I not be doing this?

It seems to be the same idea as https://github.com/colinhacks/zod/issues/2076, but I don't really understand how to apply the solution to my use case since I don't have input/output schemas.

[sunnylost]

type AgentBaseInput = z.input<typeof agentBaseSchema> & {
  type: AgentType 
}

type AgentBaseOutput = z.output<typeof agentBaseSchema> & {
  type: AgentType 
}

// ...

const agentSchema = z.union([codeFreeAgentSchema, machineCodeAgentSchema]) satisfies z.ZodSchema<AgentBaseOutput, any, AgentBaseInput>