lib/ahoy_matey.rb in ahoy_matey-0.1.2 vs lib/ahoy_matey.rb in ahoy_matey-0.1.3

@@ -29,13 +29,14 @@
 ActionController::Base.send :include, Ahoy::Controller
 ActiveRecord::Base.send(:extend, Ahoy::Model) if defined?(ActiveRecord)
 
 if defined?(Warden)
   Warden::Manager.after_set_user except: :fetch do |user, auth, opts|
-    request = Rack::Request.new(auth.env)
-    if request.cookies["ahoy_visit"]
-      visit = Ahoy.visit_model.where(visit_token: request.cookies["ahoy_visit"]).first
-      if visit
+    request = ActionDispatch::Request.new(auth.env)
+    visit_token = request.cookies["ahoy_visit"] || request.headers["Ahoy-Visit"]
+    if visit_token
+      visit = Ahoy.visit_model.where(visit_token: visit_token).first
+      if visit and !visit.user
         visit.user = user
         visit.save!
       end
     end
   end