discriminatedUnion for schema of a JSON column, based on type column?

[canadaduane]

We have a table that looks roughly like so:

/// @zod.import(["import { chatMessageSchema } from '@sai/primitives';"])
model Chat {
  id              String    @id @default(cuid())
  type            ChatType
  messages        Json      @default("[]") /// @zod.custom.use(z.array(chatMessageSchema))
}

But we would like to transition to a schema on the messages column that is a union depending on the type column. Is there a way to make a discriminated union of a JSON schema depend on a model's column?

For example, making up some kind of notation for this:

model Chat {
  id              String    @id @default(cuid())
  type            ChatType
  messages        Json      @default("[]") /// @zod.custom.use(z.array(z.discriminatedUnion(@Chat.type, [z.object({ type: "v1", ...}), z.object({ type: "v2", ...})])))
}