From d23e118645524d9162922b673a5a459c61754c51 Mon Sep 17 00:00:00 2001
From: Eugene Burmakin <iamfrey@gmail.com>
Date: Wed, 22 Oct 2025 21:36:51 +0200
Subject: [PATCH] Make sure family invitations are handled after sign-in

---
 app/controllers/application_controller.rb    |  8 +++
 app/controllers/users/sessions_controller.rb | 16 +----
 spec/requests/authentication_spec.rb         | 75 ++++++++++++++++++++
 3 files changed, 85 insertions(+), 14 deletions(-)

diff --git a/app/controllers/application_controller.rb b/app/controllers/application_controller.rb
index a18d9f5b..bdf00702 100644
--- a/app/controllers/application_controller.rb
+++ b/app/controllers/application_controller.rb
@@ -40,6 +40,14 @@ class ApplicationController < ActionController::Base
   end
 
   def after_sign_in_path_for(resource)
+    # Check for family invitation first
+    invitation_token = params[:invitation_token] || session[:invitation_token]
+    if invitation_token.present?
+      invitation = Family::Invitation.find_by(token: invitation_token)
+      return family_invitation_path(invitation.token) if invitation&.can_be_accepted?
+    end
+
+    # Handle iOS client flow
     client_type = request.headers['X-Dawarich-Client'] || session[:dawarich_client]
 
     case client_type
diff --git a/app/controllers/users/sessions_controller.rb b/app/controllers/users/sessions_controller.rb
index 151bddc5..e1760817 100644
--- a/app/controllers/users/sessions_controller.rb
+++ b/app/controllers/users/sessions_controller.rb
@@ -7,26 +7,14 @@ class Users::SessionsController < Devise::SessionsController
     super
   end
 
-  protected
-
-  def after_sign_in_path_for(resource)
-    if invitation_token.present?
-      invitation = Family::Invitation.find_by(token: invitation_token)
-
-      if invitation&.can_be_accepted?
-        return family_invitation_path(invitation.token)
-      end
-    end
-
-    super(resource)
-  end
-
   private
 
   def load_invitation_context
     return unless invitation_token.present?
 
     @invitation = Family::Invitation.find_by(token: invitation_token)
+    # Store token in session so it persists through the sign-in process
+    session[:invitation_token] = invitation_token if invitation_token.present?
   end
 
   def invitation_token
diff --git a/spec/requests/authentication_spec.rb b/spec/requests/authentication_spec.rb
index 621a86cc..75620f5c 100644
--- a/spec/requests/authentication_spec.rb
+++ b/spec/requests/authentication_spec.rb
@@ -166,4 +166,79 @@ RSpec.describe 'Authentication', type: :request do
       expect(response.location).not_to include('auth/ios/success')
     end
   end
+
+  describe 'Family Invitation with Authentication' do
+    let(:family) { create(:family, creator: user) }
+    let!(:membership) { create(:family_membership, user: user, family: family, role: :owner) }
+    let(:invitee) { create(:user, email: 'invitee@example.com', password: 'password123') }
+    let(:invitation) { create(:family_invitation, family: family, invited_by: user, email: invitee.email) }
+
+    it 'redirects to invitation page when signing in with invitation token in params' do
+      post user_session_path, params: {
+        user: { email: invitee.email, password: 'password123' },
+        invitation_token: invitation.token
+      }
+
+      expect(response).to redirect_to(family_invitation_path(invitation.token))
+    end
+
+    it 'redirects to invitation page when signing in with invitation token in session' do
+      # The invitation token is stored in session by Users::SessionsController#load_invitation_context
+      # when accessing the sign-in page with invitation_token param
+      get new_user_session_path, params: { invitation_token: invitation.token }
+
+      # Then sign in without the invitation_token in params (should use session value)
+      post user_session_path, params: {
+        user: { email: invitee.email, password: 'password123' }
+      }
+
+      expect(response).to redirect_to(family_invitation_path(invitation.token))
+    end
+
+    it 'prioritizes invitation over iOS flow when both are present' do
+      # Sign in with both iOS header AND invitation token
+      post user_session_path, params: {
+        user: { email: invitee.email, password: 'password123' },
+        invitation_token: invitation.token
+      }, headers: {
+        'X-Dawarich-Client' => 'ios'
+      }
+
+      # Should redirect to invitation page, NOT iOS success
+      expect(response).to redirect_to(family_invitation_path(invitation.token))
+      expect(response.location).not_to include('auth/ios/success')
+    end
+
+    it 'redirects to iOS success when invitation is expired' do
+      # Create an expired invitation
+      expired_invitation = create(:family_invitation,
+                                   family: family,
+                                   invited_by: user,
+                                   email: invitee.email,
+                                   expires_at: 1.day.ago)
+
+      # Sign in with iOS header and expired invitation token
+      post user_session_path, params: {
+        user: { email: invitee.email, password: 'password123' },
+        invitation_token: expired_invitation.token
+      }, headers: {
+        'X-Dawarich-Client' => 'ios'
+      }
+
+      # Should redirect to iOS success since invitation can't be accepted
+      expect(response).to redirect_to(%r{auth/ios/success\?token=})
+    end
+
+    it 'uses default path when invitation token is invalid' do
+      # Sign in with invalid invitation token
+      post user_session_path, params: {
+        user: { email: invitee.email, password: 'password123' },
+        invitation_token: 'invalid-token-123'
+      }
+
+      # Should use default redirect path
+      expect(response).not_to redirect_to(%r{/invitations/})
+      expect(response).to redirect_to(root_path)
+    end
+  end
 end