Working with discriminated unions inside of collection causes TS troubles

[pawelblaszczyk5]

Hello 👋

I'm building some small app with tanstack db and one things that bugs me is working with discriminated unions inside of collections. The app is AI chat, so for example I have MessagePart shape, which can be TextMessagePart, ReasoningMessagePart etc. Unfortunately, it's not properly narrowed in many places when using it.

I made a small reproduction, where you can see these issues, that's much easier then to describe it:

https://github.com/pawelblaszczyk5/tanstack-db-react-compiler-reproduction/blob/c2470dd6a18e869625c531cc9ccbc0bb997c84cb/src/lib/collections.ts

As an example I chosen a schema where discriminator is also TKey to showcase also the 2. issue, but it also doesn't work in cases where it's just one of the fields. In my app I have shape with 3 fields basically id, type, data and data is different based on type, I can't properly query this without typecasting/asserting atm

But for completeness:

1. Can't discriminate union in either .where or .fn.where. It always result in the full model
2. Collection TKey is only inferred properly for few collection types, doesn't work e.g. with local storage or electric ones. They don't pass necessary generics
3. get and update in cases where TKey is discriminator doesn't properly discriminate

[joernroeder]

I'm running into the same issue and would have expected this to be narrowed down, similar to a .select in tanstack/query

[joernroeder]

Workaround: Zod .transform() to denormalize variant-specific fields

We have a collection of events typed as a Zod z.discriminatedUnion('kind', [...]) with variants like:

const schema = z.discriminatedUnion('kind', [
  z.object({ kind: z.literal('payment'), data: z.object({ payment: z.object({ merchantId: z.string(), amount: z.number() }) }) }),
  z.object({ kind: z.literal('refund'),  data: z.object({ refund:  z.object({ merchantId: z.string(), reason: z.string() }) }) }),
  z.object({ kind: z.literal('payout'),  data: z.object({ payout:  z.object({ bankAccountId: z.string(), amount: z.number() }) }) }),
])

To filter payments by merchant in an expression-based .where(), we'd need eq(a.data.payment.merchantId, id) — but Ref<T> only exposes keys shared by all union members, so variant-specific nested properties like data.payment are inaccessible.

Our workaround is a dedicated "enriched for collection" schema that wraps the discriminated union with a Zod .transform(), hoisting variant-specific fields to the top level:

const enrichedSchema = schema.transform((item) => ({
  ...item,
  merchantId: item.kind === 'payment' ? item.data.payment.merchantId
            : item.kind === 'refund'  ? item.data.refund.merchantId
            : undefined,
}))

This schema is passed to queryCollectionOptions({ schema }), so the collection's item type includes the denormalized fields and Ref<EnrichedItem> exposes them for expression-based .where() queries — enabling eq(a.merchantId, id) with full predicate push-down via loadSubsetOptions. The enrichment is purely client-side — the API schema is untouched. Each denormalized field is documented with a reference back to this issue so they can be removed when Ref<T> supports discriminated union narrowing natively.

— written by claude