From 8da145d623a9d43ad3f0f543e5a0da754c3357f7 Mon Sep 17 00:00:00 2001
From: nora <163450896+24c02@users.noreply.github.com>
Date: Fri, 13 Jun 2025 13:48:21 -0400
Subject: [PATCH] add program return button (#314)

---
 app/controllers/api/internal/magic_links_controller.rb |  2 +-
 app/controllers/sessions_controller.rb                 |  1 +
 app/views/users/wakatime_setup_step_4.html.erb         | 10 ++++++++--
 ...20250613163710_add_return_data_to_sign_in_tokens.rb |  5 +++++
 db/schema.rb                                           |  3 ++-
 5 files changed, 17 insertions(+), 4 deletions(-)
 create mode 100644 db/migrate/20250613163710_add_return_data_to_sign_in_tokens.rb

diff --git a/app/controllers/api/internal/magic_links_controller.rb b/app/controllers/api/internal/magic_links_controller.rb
index d431023..3ff9e2d 100644
--- a/app/controllers/api/internal/magic_links_controller.rb
+++ b/app/controllers/api/internal/magic_links_controller.rb
@@ -38,7 +38,7 @@ module Api
       end
 
       def magic_link_params
-        params.permit(:continue_param)
+        params.permit(:continue_param, return_data: {})
       end
     end
   end
diff --git a/app/controllers/sessions_controller.rb b/app/controllers/sessions_controller.rb
index b056b34..d5fd757 100644
--- a/app/controllers/sessions_controller.rb
+++ b/app/controllers/sessions_controller.rb
@@ -156,6 +156,7 @@ class SessionsController < ApplicationController
     if valid_token
       valid_token.mark_used!
       session[:user_id] = valid_token.user_id
+      session[:return_data] = valid_token.return_data || {}
 
       if valid_token.continue_param.present?
         redirect_to valid_token.continue_param, notice: "Successfully signed in!"
diff --git a/app/views/users/wakatime_setup_step_4.html.erb b/app/views/users/wakatime_setup_step_4.html.erb
index 13032db..138447c 100644
--- a/app/views/users/wakatime_setup_step_4.html.erb
+++ b/app/views/users/wakatime_setup_step_4.html.erb
@@ -14,6 +14,12 @@
   <button>Set up another editor</button>
 <% end %>
 
-<%= link_to root_path do %>
-  <button>Done</button>
+<% if (url = session.dig(:return_data, "url")) %>
+  <%= link_to url do %>
+    <button><%= session.dig(:return_data, "button_text") || "Done" %></button>
+  <% end %>
+  <% else %>
+  <%= link_to root_path do %>
+    <button>Done</button>
+  <% end %>
 <% end %>
diff --git a/db/migrate/20250613163710_add_return_data_to_sign_in_tokens.rb b/db/migrate/20250613163710_add_return_data_to_sign_in_tokens.rb
new file mode 100644
index 0000000..a2675c5
--- /dev/null
+++ b/db/migrate/20250613163710_add_return_data_to_sign_in_tokens.rb
@@ -0,0 +1,5 @@
+class AddReturnDataToSignInTokens < ActiveRecord::Migration[8.0]
+  def change
+    add_column :sign_in_tokens, :return_data, :jsonb
+  end
+end
diff --git a/db/schema.rb b/db/schema.rb
index bba0cfd..4439986 100644
--- a/db/schema.rb
+++ b/db/schema.rb
@@ -10,7 +10,7 @@
 #
 # It's strongly recommended that you check this file into your version control system.
 
-ActiveRecord::Schema[8.0].define(version: 2025_06_11_071124) do
+ActiveRecord::Schema[8.0].define(version: 2025_06_13_163710) do
   create_schema "pganalyze"
 
   # These are extensions that must be enabled in order to support this database
@@ -464,6 +464,7 @@ ActiveRecord::Schema[8.0].define(version: 2025_06_11_071124) do
     t.datetime "created_at", null: false
     t.datetime "updated_at", null: false
     t.string "continue_param"
+    t.jsonb "return_data"
     t.index ["token"], name: "index_sign_in_tokens_on_token"
     t.index ["user_id"], name: "index_sign_in_tokens_on_user_id"
   end