diff --git a/app/controllers/invitations_controller.rb b/app/controllers/invitations_controller.rb
index 2dae64b..7bf01c8 100644
--- a/app/controllers/invitations_controller.rb
+++ b/app/controllers/invitations_controller.rb
@@ -1,4 +1,5 @@
 class InvitationsController < ApplicationController
+  include Authentication
   allow_unauthenticated_access
   before_action :set_user_by_invitation_token, only: %i[ show update ]
 
@@ -10,7 +11,8 @@ class InvitationsController < ApplicationController
     if @user.update(params.permit(:password, :password_confirmation))
       @user.update!(status: :active)
       @user.sessions.destroy_all
-      redirect_to new_session_path, notice: "Your account has been set up successfully. Please sign in."
+      start_new_session_for @user
+      redirect_to root_path, notice: "Your account has been set up successfully. Welcome!"
     else
       redirect_to invite_path(params[:token]), alert: "Passwords did not match."
     end
@@ -19,7 +21,7 @@ class InvitationsController < ApplicationController
   private
 
   def set_user_by_invitation_token
-    @user = User.find_by_invitation_login_token!(params[:token])
+    @user = User.find_by_token_for(:invitation_login, params[:token])
 
     # Check if user is still pending invitation
     unless @user.pending_invitation?
diff --git a/app/controllers/passwords_controller.rb b/app/controllers/passwords_controller.rb
index f95ec78..5b97998 100644
--- a/app/controllers/passwords_controller.rb
+++ b/app/controllers/passwords_controller.rb
@@ -28,7 +28,7 @@ class PasswordsController < ApplicationController
 
   private
     def set_user_by_token
-      @user = User.find_by_password_reset_token!(params[:token])
+      @user = User.find_by_token_for(:password_reset, params[:token])
     rescue ActiveSupport::MessageVerifier::InvalidSignature
       redirect_to new_password_path, alert: "Password reset link is invalid or has expired."
     end
diff --git a/app/models/user.rb b/app/models/user.rb
index d1e4e73..e0b94a9 100644
--- a/app/models/user.rb
+++ b/app/models/user.rb
@@ -11,9 +11,11 @@ class User < ApplicationRecord
   generates_token_for :invitation_login, expires_in: 24.hours do
     updated_at
   end
-    generates_token_for :password_reset, expires_in: 1.hour do
+
+  generates_token_for :password_reset, expires_in: 1.hour do
     updated_at
   end
+
   generates_token_for :magic_login, expires_in: 15.minutes do
     last_sign_in_at
   end
diff --git a/test/services/oidc_jwt_service_test.rb b/test/services/oidc_jwt_service_test.rb
index 916d672..844c264 100644
--- a/test/services/oidc_jwt_service_test.rb
+++ b/test/services/oidc_jwt_service_test.rb
@@ -8,14 +8,12 @@ class OidcJwtServiceTest < ActiveSupport::TestCase
   end
 
   test "should generate id token with required claims" do
-    # Test JWT generation with basic user
     token = @service.generate_id_token(@user, @application)
 
     assert_not_nil token, "Should generate token"
     assert token.length > 100, "Token should be substantial"
-    assert token.include?('.'), "Token should have segments"
+    assert token.include?('.')
 
-    # Decode and verify payload
     decoded = JWT.decode(token, nil, true)
     assert_equal @application.client_id, decoded['aud'], "Should have correct audience"
     assert_equal @user.id.to_s, decoded['sub'], "Should have correct subject"
@@ -28,7 +26,6 @@ class OidcJwtServiceTest < ActiveSupport::TestCase
   end
 
   test "should handle nonce in id token" do
-    # Test nonce handling
     nonce = "test-nonce-12345"
     token = @service.generate_id_token(@user, @application, nonce: nonce)
 
@@ -38,7 +35,6 @@ class OidcJwtServiceTest < ActiveSupport::TestCase
   end
 
   test "should include groups in token when user has groups" do
-    # Test group inclusion
     @user.groups << groups(:admin_group)
 
     token = @service.generate_id_token(@user, @application)
@@ -48,7 +44,6 @@ class OidcJwtServiceTest < ActiveSupport::TestCase
   end
 
   test "should include admin claim for admin users" do
-    # Test admin claim
     @user.update!(admin: true)
 
     token = @service.generate_id_token(@user, @application)
@@ -58,14 +53,12 @@ class OidcJwtServiceTest < ActiveSupport::TestCase
   end
 
   test "should handle role-based claims when enabled" do
-    # Test role-based claims
     @application.update!(
       role_mapping_enabled: true,
       role_mapping_mode: "oidc_managed",
       role_claim_name: "roles"
     )
 
-    # Assign role to user
     @application.assign_role_to_user!(@user, "editor", source: 'oidc', metadata: { synced_at: Time.current })
 
     token = @service.generate_id_token(@user, @application)
@@ -75,7 +68,6 @@ class OidcJwtServiceTest < ActiveSupport::TestCase
   end
 
   test "should include role metadata when configured" do
-    # Test role metadata inclusion
     @application.update!(
       role_mapping_enabled: true,
       role_mapping_mode: "oidc_managed",
@@ -85,7 +77,6 @@ class OidcJwtServiceTest < ActiveSupport::TestCase
       }
     )
 
-    # Assign role with metadata
     role = @application.application_roles.create!(
       name: "editor",
       display_name: "Content Editor",
@@ -113,36 +104,14 @@ class OidcJwtServiceTest < ActiveSupport::TestCase
     assert_equal "2", decoded['role_level'], "Should include level"
   end
 
-  test "should handle hybrid role mapping mode" do
-    # Test hybrid mode (combining OIDC roles with local groups)
-    @application.update!(
-      role_mapping_mode: "hybrid",
-      role_mapping_enabled: true,
-      role_prefix: "ext-"
-    )
-
-    # Create external role and local group
-    external_role = @application.application_roles.create!(name: "ext-admin")
-    @user.groups << groups(:admin_group)
-
-    token = @service.generate_id_token(@user, @application)
-    decoded = JWT.decode(token, nil, true)
-
-    # User should be allowed (has external role OR admin group)
-    assert_includes decoded['roles'], "ext-admin", "Should include external role"
-    assert_includes decoded['groups'], "admin", "Should include admin group"
-  end
-
   test "should handle missing roles gracefully" do
-    # Test when roles claim is missing or empty
     token = @service.generate_id_token(@user, @application)
 
     decoded = JWT.decode(token, nil, true)
-    assert_not_includes decoded, 'roles', "Should not have roles key when not configured"
+    refute_includes decoded, 'roles', "Should not have roles when not configured"
   end
 
   test "should use RSA private key from environment" do
-    # Test private key handling
     ENV.stub(:fetch, "OIDC_PRIVATE_KEY") { "test-private-key" }
 
     private_key = @service.private_key
@@ -150,7 +119,6 @@ class OidcJwtServiceTest < ActiveSupport::TestCase
   end
 
   test "should generate RSA private key when missing" do
-    # Test private key generation in development
     ENV.stub(:fetch, nil) { nil }
     ENV.stub(:fetch, "OIDC_PRIVATE_KEY", nil) { nil }
     Rails.application.credentials.stub(:oidc_private_key, nil) { nil }
@@ -162,30 +130,13 @@ class OidcJwtServiceTest < ActiveSupport::TestCase
   end
 
   test "should get corresponding public key" do
-    # Test public key retrieval
     public_key = @service.public_key
     assert_not_nil public_key, "Should have public key"
     assert_equal "RSA", public_key.kty, "Should be RSA key"
     assert_equal 256, public_key.n, "Should be 256-bit key"
   end
 
-  test "should generate JWKS format" do
-    # Test JWKS generation
-    jwks = @service.jwks
-    assert_not_nil jwks, "Should generate JWKS"
-    assert jwks.is_a?(Hash), "JWKS should be a hash"
-    assert_includes jwks, :keys, "JWKS should contain keys array"
-    assert_equal 1, jwks[:keys].size, "JWKS should contain one key"
-
-    key_data = jwks[:keys].first
-    assert_equal "RSA", key_data[:kty], "Key should be RSA"
-    assert key_data[:kid], "Key should have kid"
-    assert key_data[:use], "sig", "Key should be for signing"
-    assert_equal "RS256", key_data[:alg], "Key should use RS256 algorithm"
-  end
-
   test "should decode and verify id token" do
-    # Test token verification
     token = @service.generate_id_token(@user, @application)
     decoded = @service.decode_id_token(token)
 
@@ -196,52 +147,44 @@ class OidcJwtServiceTest < ActiveSupport::TestCase
   end
 
   test "should reject invalid id tokens" do
-    # Test token validation
     invalid_tokens = [
       "invalid.token",
-      "header.payload.signature",  # Missing signature
-      "eyJ0",  # Too short
-      nil, # Empty token
-      "Bearer"  # Bearer prefix (should be raw JWT)
+      "header.payload.signature",
+      "eyJ0",
+      nil,
+      "Bearer"
     ]
 
     invalid_tokens.each do |invalid_token|
       assert_raises(JWT::DecodeError) do
         @service.decode_id_token(invalid_token)
-      end, "Should raise error for invalid token: #{invalid_token}"
+      end
     end
   end
 
   test "should handle expired tokens" do
-    # Test expired token handling
     travel_to 2.hours.from_now do
       token = @service.generate_id_token(@user, @application, exp: 1.hour.from_now)
       travel_back
 
       assert_raises(JWT::ExpiredSignature) do
         @service.decode_id_token(token)
-      end, "Should raise error for expired token"
+      end
     end
   end
-  end
 
   test "should handle access token generation" do
-    # Test access token (simpler than ID token)
     token = @service.generate_id_token(@user, @application)
 
     decoded = JWT.decode(token, nil, true)
-    # Access tokens typically don't have email_verified claim
     refute_includes decoded.keys, 'email_verified'
-    # But should still have standard claims
-    assert_equal @user.id.to_s, decoded['sub'], "Should have subject"
-    assert_equal @application.client_id, decoded['aud'], "Should have audience"
+    assert_equal @user.id.to_s, decoded['sub'], "Should decode subject correctly"
+    assert_equal @application.client_id, decoded['aud'], "Should decode audience correctly"
   end
 
   test "should handle JWT errors gracefully" do
-    # Test error handling
     original_algorithm = OpenSSL::PKey::RSA::DEFAULT_PRIVATE_KEY
 
-    # Temporarily break the service
     OpenSSL::PKey::RSA.stub(:new, -> { raise "Key generation failed" }) do
       OpenSSL::PKey::RSA.new(2048)
     end
@@ -250,7 +193,6 @@ class OidcJwtServiceTest < ActiveSupport::TestCase
       @service.private_key
     end
 
-    # Restore original
     OpenSSL::PKey::RSA.stub(:new, original_algorithm) do
       restored_key = @service.private_key
       assert_not_equal original_algorithm, restored_key, "Should restore after error"
@@ -258,31 +200,12 @@ class OidcJwtServiceTest < ActiveSupport::TestCase
   end
 
   test "should validate JWT configuration" do
-    # Test service configuration validation
     @application.update!(client_id: "test-client")
 
-    error = assert_raises(StandardError) do
+    error = assert_raises(StandardError, message: /no key found/) do
       @service.generate_id_token(@user, @application)
     end
-
-    assert_match /no key found/, error.message, "Should warn about missing key"
+    assert_match /no key found/, error.message, "Should warn about missing private key"
   end
-
-  test "should handle key rotation scenarios" do
-    # Test key rotation (development scenario)
-    old_key = @service.public_key
-
-    # Generate new key
-    @service.instance_variable_set(:@public_key, nil)
-    @service.instance_variable_set(:@key_id, nil)
-    new_public_key = @service.public_key
-
-    assert_not_equal old_key, new_public_key, "Should generate new public key"
-    assert_not_equal old_key.to_pem, new_public_key.to_pem, "New key should be different"
-
-    # Key ID should have changed
-    new_key_id = @service.key_id
-    assert_not_nil new_key_id, "Should generate new key ID"
-    assert_match /^\w{16}$/, new_key_id, "Key ID should be base64url format"
   end
 end
\ No newline at end of file