For canonicalization, there is a workaround in IREE, but it is a bit heavy. Essentially, it creates a custom version of the canonicalization driver that can notify when an op is replaced with something else in a pattern. That notification is then used to identify, in op-specific way, the op that should replace the erased op in the transform-handle<->payload-op mapping.
On a more general compiler construction side of things, I think we should use canonicalization less often and more localized than what we currently do.
Have you seen [RFC] Type System for the Transform Dialect ? It has been implemented and is available, although not all transform dialect extensions migrated to use it. It is not as strongly typed as it could be, and I wanted to avoid a proliferation of casts, but you may be able to harden it a bit in the dialect. For example, add a custom cast op that statically verifies whether types are compatible (e.g., !transform.op<"scf.for"> should not be casted to !transform.op<"anything-else">.