Add skip_setup_flow flag to oauth

app/controllers/application_controller.rb

@@ -53,6 +53,20 @@ def safe_return_url(url)
     url
   end
 
+  # Build a return_data hash from a continue URL, extracting known query
+  # params (like skip_setup_flow) so they survive auth redirects.
+  def build_return_data(continue_url)
+    url = safe_return_url(continue_url)
+    return {} if url.blank?
+
+    data = { "url" => url }
+    query = Rack::Utils.parse_query(URI.parse(url).query.to_s)
+    data["skip_setup_flow"] = true if query.key?("skip_setup_flow")
+    data
+  rescue URI::InvalidURIError
+    { "url" => url }
+  end
+
   def authenticate_user!
     unless user_signed_in?
       redirect_to signin_path(continue: request.fullpath), alert: "Please sign in first!"

app/controllers/sessions_controller.rb

@@ -1,6 +1,6 @@
 class SessionsController < ApplicationController
   def hca_new
-    session[:return_data] = { "url" => safe_return_url(params[:continue].presence) } if params[:continue].present?
+    session[:return_data] = build_return_data(params[:continue]) if params[:continue].present?
     Rails.logger.info("Sessions return data: #{session[:return_data]}")
     redirect_uri = url_for(action: :hca_create, only_path: false)
 
@@ -29,6 +29,7 @@ def hca_create
       else
         redirect_to root_path, notice: notice
       end
+
     else
       redirect_to root_path, alert: "Failed to authenticate with Hack Club Auth!"
     end
@@ -71,7 +72,7 @@ def slack_create
       if slack_state&.dig("close_window")
         redirect_to close_window_path
       elsif @user.previously_new_record?
-        session[:return_data] = { "url" => continue_url }
+        session[:return_data] = build_return_data(continue_url)
         redirect_to my_wakatime_setup_path, notice: notice
       elsif continue_url.present?
         redirect_to continue_url, notice: notice # codeql[rb/url-redirection]
@@ -234,9 +235,8 @@ def token
       valid_token.mark_used!
       reset_session
       session[:user_id] = valid_token.user_id
-      session[:return_data] = valid_token.return_data || {}
-
       continue_url = safe_return_url(valid_token.continue_param)
+      session[:return_data] = (valid_token.return_data || {}).merge(build_return_data(continue_url))
       if continue_url.present?
         redirect_to continue_url, notice: "Successfully signed in!" # codeql[rb/url-redirection]
       else

app/controllers/users_controller.rb

@@ -7,13 +7,20 @@ class UsersController < InertiaController
 
   def wakatime_setup
     api_key = ensure_api_key
+    skipping = session.dig(:return_data, "skip_setup_flow") || params[:skip_setup_flow].present?
+
+    # Clear so it doesn't persist across future visits
+    session[:return_data]&.delete("skip_setup_flow") if skipping
 
     render inertia: "WakatimeSetup/Index", props: {
       current_user_api_key: api_key.token,
       setup_os: detect_setup_os(request.user_agent).to_s,
       # Full URL (with host) is shown to users in their config file, so we
       # build it server-side rather than via js_from_routes.
-      api_url: api_hackatime_v1_url
+      api_url: api_hackatime_v1_url,
+      skip_setup_flow: skipping,
+      return_url: skipping ? session.dig(:return_data, "url") : nil,
+      return_button_text: skipping ? (session.dig(:return_data, "button_text") || "Done") : nil
     }
   end