Better Devise Password Reset UX

I'm currently redesigning some devise views when I noted the following in PasswordController#edit:

Basically, upon getting the clicking on the reset link, Devise doesn't actually check whether it exists or whether whether the token is still valid. From my perspective that is bad UX, as the user will have to retype the password, only to realize that it was for nothing. Thus, I'd say the following is better:

It will throw a 404 if not found, but more importantly will immediately tell you if the password reset link has expired so that you can immediately display it to the user.

As far as I can see, the security risk is the same, as existence of a token can be deducted from update as well.