Working with discriminated unions inside of collection causes TS troubles

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

Start with the reproduction in src/lib/collections.ts and trace the types used by .where and .fn.where for the discriminated MessagePart union. Then inspect collection implementations using local storage and Electric, along with get and update, to compare how TKey generics are passed. Done means discriminator-based narrowing works across these operations without type assertions and TKey is inferred correctly.