Password reset token lost on redirect

Problem/Motivation

This a follow-up from ✨ Support password reset UX Fixed which is no longer working, possibly due to a core update.

Steps to reproduce

Same as the previous ticket:

When a user login with one-time login token, then clicks around, this module redirects the user back to the profile page. However, because the one-time login token is not in the URL parameter anymore, the core AccountForm would require the user to type in the "Current password", which the user might not have. Therefore, the user cannot save the form, but request another one-time login token.

Proposed resolution

I'm currently using Drupal Core 9.5.2.

I've taken a quick look into why it doesn't work anymore and it looks like the pass_reset_UID is not in the $_SESSION.

Would it be better to check the URL query? This would have the "pass-reset-token". I can create a patch if this is the right approach.