Add type inference to `loadBeforeMutationEntities` and `loadAfterMutationEntities`

[genu]

Is there a reliable way to infer the proper return type of loadAfterMutationEntities and loadBeforeMutationEntities?

https://github.com/zenstackhq/zenstack-v3/blob/35a9e0d6575d7baf7d0aa3e29484d7aedb56b4d0/packages/runtime/src/client/plugin.ts#L170-L174

https://github.com/zenstackhq/zenstack-v3/blob/35a9e0d6575d7baf7d0aa3e29484d7aedb56b4d0/packages/runtime/src/client/plugin.ts#L153-L158

If the return data is typed, it becomes easier working with any follow up queries:

export const onAfterPostDelete = definePlugin<SchemaType>({
  id: "onAfterPostDelete",
  onEntityMutation: {
    async beforeEntityMutation({ model, action, loadBeforeMutationEntities, client, queryId }) {
      const deletedPost = await loadBeforeMutationEntities() // <--- typed as Record<string, unknown>[] | undefined but its really a Post[]

      if (action === "delete" && model === "Post") {
        // Do some logic
        const tags = await client.tag.findMany({
          where: {
            postId: deletedPost[0].id, // <-- This doesn't work unless we cast deletedPost to proper type like Post[]
          },
        })
      }
    },
  },
})

[genu]

After investigating the codebase, here is a potential strategy:

Why this is tricky

The hook callback receives model as a runtime value — a union of all model name strings (GetModels<Schema>). TypeScript can't use a runtime value to conditionally type another property in the same object. So when you write if (model === "Post"), TS narrows model to "Post" but has no way to simultaneously narrow the return type of loadBeforeMutationEntities().

All the building blocks for typed entities already exist internally (DefaultModelResult<Schema, Model> computes the plain entity shape for any given model), they just aren't wired up to the plugin hook types.

Possible strategy

Per-model hook registration (medium effort, best DX, this may be more in line with current patterns anyway)

Allow registering callbacks per model instead of (or in addition to) a single catch-all:

onEntityMutation: {
  Post: {
    async beforeEntityMutation({ action, loadBeforeMutationEntities }) {
      const posts = await loadBeforeMutationEntities(); // ← typed as Post[]
    }
  }
}

Personally I think this is the best balance of effort and DX — the type machinery already exists, it just needs to be wired into a per-model registration shape.

[genu]

@ymc9 Any thoughts on the proposed solution above?

we can keep the global hooks as is (for backward compatibility), and add the per-model hook registration.

[ymc9]

@ymc9 Any thoughts on the proposed solution above?

we can keep the global hooks as is (for backward compatibility), and add the per-model hook registration.

Hi @genu , I agree per-model registration is the solution to this issue. Probably in most cases people would be interested in hooking into a few models instead of all. If we make that change, it'll look nicer if the "all-model" reg is put under a special "$all" key, to be consistent with conventions in other places. But it'll incur a breaking change.

I'm wondering if this is a high priority. Are you implementing a complex hook?

[genu]

Hi @ymc9

You're right, it would technically, be a breaking change. It isn't a high priority for me since I'm able to work around it.

I can create a PR for this that can be merged in at some point in the future along with other breaking changes if you'd like.

I just wanted to bump this, so it isn't forgotten.