Redirect to setup for new users entering OAuth flow (#931)

* Fix redirects in new account flow?

* Fix redirects! And tests!

* Make tests better

* I will test this club

* Uh oh, let's redirect to WakaTime setup

* Make it actually detect new users properly?

app/controllers/sessions_controller.rb

@@ -33,8 +33,13 @@ class SessionsController < ApplicationController
         MigrateUserFromHackatimeJob.perform_later(@user.id)
       end
 
-      if @user.created_at > 5.seconds.ago
-        session[:return_data] = { "url" => safe_return_url(params[:continue].presence) }
+      if !@user.heartbeats.exists?
+        # User hasn't set up editor yet; preserve return_data already set by hca_new,
+        # only override if a new, safely sanitized continue URL is available.
+        if params[:continue].present?
+          sanitized_continue_url = safe_return_url(params[:continue].presence)
+          session[:return_data] = { "url" => sanitized_continue_url } if sanitized_continue_url.present?
+        end
         Rails.logger.info("Sessions return data: #{session[:return_data]}")
         redirect_to my_wakatime_setup_path, notice: "Successfully signed in with Hack Club Auth! Welcome!"
       elsif session[:return_data]&.dig("url").present?
@@ -84,11 +89,11 @@ class SessionsController < ApplicationController
       state = JSON.parse(params[:state]) rescue {}
       if state["close_window"]
         redirect_to close_window_path
-      elsif @user.created_at > 5.seconds.ago
+      elsif !@user.heartbeats.exists?
         session[:return_data] = { "url" => safe_return_url(state["continue"].presence) }
         redirect_to my_wakatime_setup_path, notice: "Successfully signed in with Slack! Welcome!"
-      elsif state["continue"].present? && safe_return_url(state["continue"]).present?
-        redirect_to safe_return_url(state["continue"]), notice: "Successfully signed in with Slack! Welcome!"
+      elsif (continue_url = safe_return_url(state["continue"].presence))
+        redirect_to continue_url, notice: "Successfully signed in with Slack! Welcome!"
       else
         redirect_to root_path, notice: "Successfully signed in with Slack! Welcome!"
       end
@@ -249,9 +254,15 @@ class SessionsController < ApplicationController
       valid_token.mark_used!
       session[:user_id] = valid_token.user_id
       session[:return_data] = valid_token.return_data || {}
+      user = User.find(valid_token.user_id)
+      continue_url = safe_return_url(valid_token.continue_param)
 
-      if valid_token.continue_param.present? && safe_return_url(valid_token.continue_param).present?
-        redirect_to safe_return_url(valid_token.continue_param), notice: "Successfully signed in!"
+      if !user.heartbeats.exists?
+        # User hasn't set up editor yet; send through wakatime setup first
+        session[:return_data]["url"] = continue_url if continue_url.present?
+        redirect_to my_wakatime_setup_path, notice: "Successfully signed in!"
+      elsif continue_url.present?
+        redirect_to continue_url, notice: "Successfully signed in!"
       else
         redirect_to root_path, notice: "Successfully signed in!"
       end

db/migrate/20260210094354_add_heartbeat_indexes.rb

@@ -0,0 +1,35 @@
+class AddHeartbeatIndexes < ActiveRecord::Migration[8.1]
+  disable_ddl_transaction!
+
+  def change
+    add_index :heartbeats, [ :user_id, :category, :time ],
+              name: "idx_heartbeats_user_category_time",
+              where: "(deleted_at IS NULL)",
+              algorithm: :concurrently,
+              if_not_exists: true
+
+    add_index :heartbeats, [ :user_id, :editor, :time ],
+              name: "idx_heartbeats_user_editor_time",
+              where: "(deleted_at IS NULL)",
+              algorithm: :concurrently,
+              if_not_exists: true
+
+    add_index :heartbeats, [ :user_id, :language, :time ],
+              name: "idx_heartbeats_user_language_time",
+              where: "(deleted_at IS NULL)",
+              algorithm: :concurrently,
+              if_not_exists: true
+
+    add_index :heartbeats, [ :user_id, :operating_system, :time ],
+              name: "idx_heartbeats_user_operating_system_time",
+              where: "(deleted_at IS NULL)",
+              algorithm: :concurrently,
+              if_not_exists: true
+
+    add_index :heartbeats, [ :user_id, :project ],
+              name: "index_heartbeats_on_user_id_and_project",
+              where: "(deleted_at IS NULL)",
+              algorithm: :concurrently,
+              if_not_exists: true
+  end
+end

test/controllers/sessions_controller_test.rb

@@ -5,8 +5,6 @@ class SessionsControllerTest < ActionDispatch::IntegrationTest
     ActiveRecord::FixtureSet.reset_cache
   end
 
-  # -- HCA: hca_new stores continue in session --
-
   test "hca_new stores continue path for oauth authorize" do
     continue_query = {
       client_id: "Ck47_6hihaBqZO7z3CLmJlCB-0NzHtZHGeDBwG4CqRs",
@@ -77,9 +75,12 @@ class SessionsControllerTest < ActionDispatch::IntegrationTest
 
     # LoopsMailer forces SMTP delivery even in test; temporarily override
     original_delivery_method = LoopsMailer.delivery_method
-    LoopsMailer.delivery_method = :test
-    post email_auth_path, params: { email: "test@example.com", continue: oauth_path }
-    LoopsMailer.delivery_method = original_delivery_method
+    begin
+      LoopsMailer.delivery_method = :test
+      post email_auth_path, params: { email: "test@example.com", continue: oauth_path }
+    ensure
+      LoopsMailer.delivery_method = original_delivery_method
+    end
 
     assert_response :redirect
 
@@ -88,8 +89,9 @@ class SessionsControllerTest < ActionDispatch::IntegrationTest
     assert_equal oauth_path, token.continue_param
   end
 
-  test "email token redirects to continue param after sign in" do
+  test "email token redirects user with heartbeats to continue param after sign in" do
     user = User.create!
+    user.heartbeats.create!(time: Time.current.to_f, source_type: :test_entry)
     oauth_path = "/oauth/authorize?client_id=test&response_type=code"
     sign_in_token = user.sign_in_tokens.create!(
       auth_type: :email,
@@ -103,8 +105,25 @@ class SessionsControllerTest < ActionDispatch::IntegrationTest
     assert_equal user.id, session[:user_id]
   end
 
-  test "email token falls back to root when no continue param" do
+  test "email token redirects user without heartbeats to wakatime setup with continue stored in session" do
     user = User.create!
+    oauth_path = "/oauth/authorize?client_id=test&response_type=code"
+    sign_in_token = user.sign_in_tokens.create!(
+      auth_type: :email,
+      continue_param: oauth_path
+    )
+
+    get auth_token_path(token: sign_in_token.token)
+
+    assert_response :redirect
+    assert_redirected_to my_wakatime_setup_path
+    assert_equal user.id, session[:user_id]
+    assert_equal oauth_path, session.dig(:return_data, "url")
+  end
+
+  test "email token falls back to root when no continue param for user with heartbeats" do
+    user = User.create!
+    user.heartbeats.create!(time: Time.current.to_f, source_type: :test_entry)
     sign_in_token = user.sign_in_tokens.create!(auth_type: :email)
 
     get auth_token_path(token: sign_in_token.token)
@@ -114,8 +133,20 @@ class SessionsControllerTest < ActionDispatch::IntegrationTest
     assert_equal user.id, session[:user_id]
   end
 
-  test "email token rejects external continue URL" do
+  test "email token sends user without heartbeats to wakatime setup when no continue param" do
     user = User.create!
+    sign_in_token = user.sign_in_tokens.create!(auth_type: :email)
+
+    get auth_token_path(token: sign_in_token.token)
+
+    assert_response :redirect
+    assert_redirected_to my_wakatime_setup_path
+    assert_equal user.id, session[:user_id]
+  end
+
+  test "email token rejects external continue URL for user with heartbeats" do
+    user = User.create!
+    user.heartbeats.create!(time: Time.current.to_f, source_type: :test_entry)
     sign_in_token = user.sign_in_tokens.create!(
       auth_type: :email,
       continue_param: "https://evil.example.com/phish"
@@ -128,8 +159,9 @@ class SessionsControllerTest < ActionDispatch::IntegrationTest
     assert_equal user.id, session[:user_id]
   end
 
-  test "email token rejects protocol-relative continue URL" do
+  test "email token rejects protocol-relative continue URL for user with heartbeats" do
     user = User.create!
+    user.heartbeats.create!(time: Time.current.to_f, source_type: :test_entry)
     sign_in_token = user.sign_in_tokens.create!(
       auth_type: :email,
       continue_param: "//evil.example.com/phish"
@@ -170,4 +202,80 @@ class SessionsControllerTest < ActionDispatch::IntegrationTest
     assert_redirected_to root_path
     assert_nil session[:user_id]
   end
+
+  # -- HCA callback: new user without heartbeats --
+
+  test "hca_create redirects user without heartbeats to wakatime setup" do
+    user = User.create!
+    User.define_singleton_method(:from_hca_token) { |_code, _uri| user }
+    get "/auth/hca/callback", params: { code: "fake_code" }
+
+    assert_response :redirect
+    assert_redirected_to my_wakatime_setup_path
+    assert_equal user.id, session[:user_id]
+  ensure
+    User.singleton_class.remove_method(:from_hca_token)
+  end
+
+  test "hca_create stores continue in session for user without heartbeats" do
+    user = User.create!
+    oauth_path = "/oauth/authorize?client_id=test&response_type=code"
+
+    User.define_singleton_method(:from_hca_token) { |_code, _uri| user }
+    get "/auth/hca/callback", params: { code: "fake_code", continue: oauth_path }
+
+    assert_response :redirect
+    assert_redirected_to my_wakatime_setup_path
+    assert_equal oauth_path, session.dig(:return_data, "url")
+  ensure
+    User.singleton_class.remove_method(:from_hca_token)
+  end
+
+  test "hca_create does not overwrite return_data with nil for unsafe continue URL" do
+    user = User.create!
+
+    # Simulate hca_new having set a valid return URL
+    get hca_auth_path(continue: "/oauth/authorize?client_id=test")
+
+    User.define_singleton_method(:from_hca_token) { |_code, _uri| user }
+    get "/auth/hca/callback", params: { code: "fake_code", continue: "https://evil.example.com" }
+
+    assert_response :redirect
+    assert_redirected_to my_wakatime_setup_path
+    assert_equal "/oauth/authorize?client_id=test", session.dig(:return_data, "url")
+  ensure
+    User.singleton_class.remove_method(:from_hca_token)
+  end
+
+  # -- Slack callback: new user without heartbeats --
+
+  test "slack_create redirects user without heartbeats to wakatime setup" do
+    user = User.create!
+    state = { "continue" => "/oauth/authorize?client_id=test" }.to_json
+
+    User.define_singleton_method(:from_slack_token) { |_code, _uri| user }
+    get "/auth/slack/callback", params: { code: "fake_code", state: state }
+
+    assert_response :redirect
+    assert_redirected_to my_wakatime_setup_path
+    assert_equal user.id, session[:user_id]
+    assert_equal "/oauth/authorize?client_id=test", session.dig(:return_data, "url")
+  ensure
+    User.singleton_class.remove_method(:from_slack_token)
+  end
+
+  test "slack_create redirects user with heartbeats to continue URL" do
+    user = User.create!
+    user.heartbeats.create!(time: Time.current.to_f, source_type: :test_entry)
+    oauth_path = "/oauth/authorize?client_id=test&response_type=code"
+    state = { "continue" => oauth_path }.to_json
+
+    User.define_singleton_method(:from_slack_token) { |_code, _uri| user }
+    get "/auth/slack/callback", params: { code: "fake_code", state: state }
+
+    assert_response :redirect
+    assert_redirected_to oauth_path
+  ensure
+    User.singleton_class.remove_method(:from_slack_token)
+  end
 end