lib/unleash/feature_toggle.rb in unleash-4.4.2 vs lib/unleash/feature_toggle.rb in unleash-4.4.3

@@ -95,13 +95,18 @@
     def sum_variant_defs_weights
       self.variant_definitions.map(&:weight).reduce(0, :+)
     end
 
     def variant_salt(context, stickiness = "default")
-      return context.get_by_name(stickiness) unless stickiness == "default"
-      return context.user_id unless context.user_id.to_s.empty?
-      return context.session_id unless context.session_id.to_s.empty?
-      return context.remote_address unless context.remote_address.to_s.empty?
+      begin
+        return context.get_by_name(stickiness) if !context.nil? && stickiness != "default"
+      rescue KeyError
+        Unleash.logger.warn "Custom stickiness key (#{stickiness}) not found in the provided context #{context}. " \
+          "Falling back to default behavior."
+      end
+      return context.user_id unless context&.user_id.to_s.empty?
+      return context.session_id unless context&.session_id.to_s.empty?
+      return context.remote_address unless context&.remote_address.to_s.empty?
 
       SecureRandom.random_number
     end
 
     def variant_from_override_match(context)