Fix stale grace state after plan upgrades and correct threshold semantics

README.md

@@ -180,10 +180,14 @@ Enforcing pricing plans is one of those boring plumbing problems that look easy
 
 - Safe under load: we use row locks and retries when setting grace/blocked/warning state, and we avoid firing the same event twice. See [grace_manager.rb](lib/pricing_plans/grace_manager.rb).
 
+- Self-healing state: when usage drops below the limit (e.g., user deletes resources, upgrades plan, or reduces usage), stale exceeded/blocked flags are automatically cleared. Methods like `grace_active?` and `should_block?` will clear outdated enforcement state as a side effect. This prevents users from remaining incorrectly flagged after remediation.
+
 - Accurate counting: persistent limits count live current rows (using `COUNT(*)`, make sure to index your foreign keys to make it fast at scale); per‑period limits record usage for the current window only. You can filter what counts with `count_scope` (Symbol/Hash/Proc/Array), and plan settings override model defaults. See [limitable.rb](lib/pricing_plans/limitable.rb) and [limit_checker.rb](lib/pricing_plans/limit_checker.rb).
 
 - Clear rules: default is to block when you hit the cap; grace periods are opt‑in. In status/UI, 0 of 0 isn’t shown as blocked. See [plan.rb](lib/pricing_plans/plan.rb), [grace_manager.rb](lib/pricing_plans/grace_manager.rb), and [view_helpers.rb](lib/pricing_plans/view_helpers.rb).
 
+- Semantic enforcement: for `grace_then_block`, grace periods start when usage goes *over* the limit (e.g., 6/5), not when it *reaches* the limit (5/5). This allows users to use their full allocation before grace begins. For `block_usage`, blocking occurs at or over the limit (e.g., at 5/5, the next creation is blocked).
+
 - Simple controllers: one‑liners to guard actions, predictable redirect order (per‑call → per‑controller → global → pricing_path), and an optional central handler. See [controller_guards.rb](lib/pricing_plans/controller_guards.rb).
 
 - Billing‑aware periods: supports billing cycle (when Pay is present), calendar month/week/day, custom time windows, and durations. See [period_calculator.rb](lib/pricing_plans/period_calculator.rb).

lib/pricing_plans.rb

@@ -27,6 +27,7 @@ class InvalidOperation < Error; end
   autoload :PaySupport, "pricing_plans/pay_support"
   autoload :LimitChecker, "pricing_plans/limit_checker"
   autoload :LimitableRegistry, "pricing_plans/limit_checker"
+  autoload :ExceededStateUtils, "pricing_plans/exceeded_state_utils"
   autoload :GraceManager, "pricing_plans/grace_manager"
   autoload :Callbacks, "pricing_plans/callbacks"
   autoload :PeriodCalculator, "pricing_plans/period_calculator"

lib/pricing_plans/callbacks.rb

@@ -106,17 +106,24 @@ def check_and_emit_limit_exceeded!(plan_owner, limit_key, current_usage, limit_c
       return if limit_config[:to] == :unlimited
 
       limit_amount = limit_config[:to].to_i
-      return unless current_usage >= limit_amount
+      after_limit = limit_config[:after_limit]
 
-      case limit_config[:after_limit]
+      case after_limit
       when :just_warn
+        return unless current_usage >= limit_amount
+
         # Just emit warning, don't track grace/block
         check_and_emit_warnings!(plan_owner, limit_key, current_usage, limit_amount)
       when :block_usage
+        return unless current_usage >= limit_amount
+
         # Do NOT mark as blocked here - this callback runs after SUCCESSFUL creation.
         # Block events are emitted from validation when creation is actually blocked.
         nil
       when :grace_then_block
+        # Grace semantics are for over-limit usage, not exact-at-limit.
+        return unless current_usage > limit_amount
+
         # Start grace period if not already in grace/blocked
         unless GraceManager.grace_active?(plan_owner, limit_key) || GraceManager.should_block?(plan_owner, limit_key)
           GraceManager.mark_exceeded!(plan_owner, limit_key, grace_period: limit_config[:grace])

lib/pricing_plans/exceeded_state_utils.rb

@@ -0,0 +1,61 @@
+# frozen_string_literal: true
+
+module PricingPlans
+  # Shared utilities for checking and managing exceeded state.
+  #
+  # This module provides common logic for determining whether usage has exceeded
+  # limits and for clearing stale exceeded flags. It is included by both
+  # GraceManager (class methods) and StatusContext (instance methods) to ensure
+  # consistent behavior.
+  #
+  # NOTE: Methods that modify state (`clear_exceeded_flags!`) are intentionally
+  # included here. The design decision is that grace/block checks should be
+  # "self-healing" - if usage drops below the limit, stale exceeded flags are
+  # automatically cleared. This prevents situations where users remain incorrectly
+  # flagged as exceeded after deleting resources or after plan upgrades.
+  module ExceededStateUtils
+    # Determine if usage has exceeded the limit based on the after_limit policy.
+    #
+    # For :grace_then_block, exceeded means strictly OVER the limit (>).
+    # For :block_usage and :just_warn, exceeded means AT or over the limit (>=).
+    #
+    # This distinction exists because:
+    # - :block_usage blocks creation of the Nth item (at limit = blocked)
+    # - :grace_then_block allows the Nth item, only starting grace when OVER
+    #
+    # @param current_usage [Integer] Current usage count
+    # @param limit_amount [Integer, Symbol] The configured limit (may be :unlimited)
+    # @param after_limit [Symbol] The enforcement policy (:block_usage, :grace_then_block, :just_warn)
+    # @return [Boolean] true if usage is considered exceeded for this policy
+    def exceeded_now?(current_usage, limit_amount, after_limit:)
+      # 0-of-0 is a special case: not considered exceeded for UX purposes
+      return false if limit_amount.to_i.zero? && current_usage.to_i.zero?
+
+      if after_limit == :grace_then_block
+        current_usage > limit_amount.to_i
+      else
+        current_usage >= limit_amount.to_i
+      end
+    end
+
+    # Clear exceeded and blocked flags from an enforcement state record.
+    #
+    # This is called when usage drops below the limit to "heal" stale state.
+    # Uses update_columns for efficiency (skips validations/callbacks).
+    #
+    # @param state [EnforcementState] The state record to clear
+    # @return [EnforcementState, nil] The updated state, or nil if no updates needed
+    def clear_exceeded_flags!(state)
+      return unless state
+
+      updates = {}
+      updates[:exceeded_at] = nil if state.exceeded_at.present?
+      updates[:blocked_at] = nil if state.blocked_at.present?
+      return state if updates.empty?
+
+      updates[:updated_at] = Time.current
+      state.update_columns(updates)
+      state
+    end
+  end
+end

lib/pricing_plans/grace_manager.rb

@@ -3,6 +3,7 @@
 module PricingPlans
   class GraceManager
     class << self
+      include ExceededStateUtils
       def mark_exceeded!(plan_owner, limit_key, grace_period: nil)
         with_lock(plan_owner, limit_key) do |state|
           # Ensure state is for the current window for per-period limits
@@ -36,6 +37,14 @@ def mark_exceeded!(plan_owner, limit_key, grace_period: nil)
       def grace_active?(plan_owner, limit_key)
         state = fresh_state_or_nil(plan_owner, limit_key)
         return false unless state&.exceeded?
+
+        plan = PlanResolver.effective_plan_for(plan_owner)
+        limit_config = plan&.limit_for(limit_key)
+        unless currently_exceeded?(plan_owner, limit_key, limit_config)
+          clear_exceeded_flags!(state)
+          return false
+        end
+
         !state.grace_expired?
       end
 
@@ -51,11 +60,16 @@ def should_block?(plan_owner, limit_key)
         limit_amount = limit_config[:to]
         return false if limit_amount == :unlimited
         current_usage = LimitChecker.current_usage_for(plan_owner, limit_key, limit_config)
-        exceeded = current_usage >= limit_amount.to_i
-        # Treat 0-of-0 as not blocked for UX/severity/status purposes
-        exceeded = false if limit_amount.to_i.zero? && current_usage.to_i.zero?
+        exceeded = exceeded_now?(current_usage, limit_amount, after_limit: after_limit)
 
-        return exceeded if after_limit == :block_usage
+        unless exceeded
+          if (state = fresh_state_or_nil(plan_owner, limit_key))
+            clear_exceeded_flags!(state)
+          end
+          return false
+        end
+
+        return true if after_limit == :block_usage
 
         # For :grace_then_block, check if grace period expired
         state = fresh_state_or_nil(plan_owner, limit_key)
@@ -119,7 +133,7 @@ def reset_state!(plan_owner, limit_key)
       end
 
       def grace_ends_at(plan_owner, limit_key)
-        state = find_state(plan_owner, limit_key)
+        state = fresh_state_or_nil(plan_owner, limit_key)
         state&.grace_ends_at
       end
 
@@ -158,6 +172,17 @@ def find_state(plan_owner, limit_key)
         EnforcementState.find_by(plan_owner: plan_owner, limit_key: limit_key.to_s)
       end
 
+      def currently_exceeded?(plan_owner, limit_key, limit_config = nil)
+        limit_config ||= PlanResolver.effective_plan_for(plan_owner)&.limit_for(limit_key)
+        return false unless limit_config
+
+        limit_amount = limit_config[:to]
+        return false if limit_amount == :unlimited
+
+        current_usage = LimitChecker.current_usage_for(plan_owner, limit_key, limit_config)
+        exceeded_now?(current_usage, limit_amount, after_limit: limit_config[:after_limit])
+      end
+
       # Returns nil if state is stale for the current period window for per-period limits
       def fresh_state_or_nil(plan_owner, limit_key)
         state = find_state(plan_owner, limit_key)

lib/pricing_plans/status_context.rb

@@ -7,6 +7,8 @@ module PricingPlans
   #
   # Thread-safe by design: each call to status() gets its own context instance.
   class StatusContext
+    include ExceededStateUtils
+
     attr_reader :plan_owner
 
     def initialize(plan_owner)
@@ -67,7 +69,26 @@ def grace_active?(limit_key)
       return @grace_active_cache[key] if @grace_active_cache.key?(key)
 
       state = fresh_enforcement_state(limit_key)
-      return @grace_active_cache[key] = false unless state&.exceeded?
+
+      # If no state exists but we're over limit for grace_then_block, lazily create grace
+      unless state&.exceeded?
+        if should_lazily_start_grace?(limit_key)
+          limit_config = limit_config_for(limit_key)
+          GraceManager.mark_exceeded!(@plan_owner, limit_key, grace_period: limit_config[:grace])
+          # Clear caches to get fresh state
+          @fresh_enforcement_state_cache&.delete(key)
+          @enforcement_state_cache&.delete(key)
+          state = fresh_enforcement_state(limit_key)
+        else
+          return @grace_active_cache[key] = false
+        end
+      end
+
+      unless currently_exceeded?(limit_key)
+        clear_exceeded_flags!(state)
+        return @grace_active_cache[key] = false
+      end
+
       @grace_active_cache[key] = !state.grace_expired?
     end
 
@@ -77,7 +98,10 @@ def grace_ends_at(limit_key)
       return @grace_ends_at_cache[key] if @grace_ends_at_cache.key?(key)
 
       state = fresh_enforcement_state(limit_key)
-      @grace_ends_at_cache[key] = state&.grace_ends_at
+      return @grace_ends_at_cache[key] = nil unless state&.exceeded?
+      return @grace_ends_at_cache[key] = nil unless grace_active?(limit_key)
+
+      @grace_ends_at_cache[key] = state.grace_ends_at
     end
 
     # Cached should block check - implemented directly to avoid GraceManager's plan resolution
@@ -95,13 +119,16 @@ def should_block?(limit_key)
       return @should_block_cache[key] = false if limit_amount == :unlimited
 
       current_usage = current_usage_for(limit_key)
-      exceeded = current_usage >= limit_amount.to_i
-      exceeded = false if limit_amount.to_i.zero? && current_usage.to_i.zero?
+      exceeded = exceeded_now?(current_usage, limit_amount, after_limit: after_limit)
 
       return @should_block_cache[key] = exceeded if after_limit == :block_usage
 
       # For :grace_then_block, check if grace period expired
-      return @should_block_cache[key] = false unless exceeded
+      unless exceeded
+        state = fresh_enforcement_state(limit_key)
+        clear_exceeded_flags!(state) if state
+        return @should_block_cache[key] = false
+      end
 
       state = fresh_enforcement_state(limit_key)
       return @should_block_cache[key] = false unless state&.exceeded?
@@ -339,5 +366,32 @@ def compute_severity(limit_key)
       highest_warn = warn_thresholds.max.to_f * 100.0
       (percent >= highest_warn && highest_warn.positive?) ? :warning : :ok
     end
+
+    def currently_exceeded?(limit_key)
+      limit_config = limit_config_for(limit_key)
+      return false unless limit_config
+
+      limit_amount = limit_config[:to]
+      return false if limit_amount == :unlimited
+
+      current_usage = current_usage_for(limit_key)
+      exceeded_now?(current_usage, limit_amount, after_limit: limit_config[:after_limit])
+    end
+
+    # Check if we should lazily start grace for this limit.
+    # This handles edge cases where usage increased without triggering callbacks
+    # (e.g., status changes, bulk imports, manual DB updates).
+    def should_lazily_start_grace?(limit_key)
+      limit_config = limit_config_for(limit_key)
+      return false unless limit_config
+      return false unless limit_config[:after_limit] == :grace_then_block
+
+      limit_amount = limit_config[:to]
+      return false if limit_amount == :unlimited
+
+      current_usage = current_usage_for(limit_key)
+      current_usage > limit_amount.to_i
+    end
+
   end
 end

test/automatic_callbacks_test.rb

@@ -81,6 +81,24 @@ def test_on_grace_start_fires_automatically_when_limit_exceeded
     assert grace_events.any?, "Expected on_grace_start to fire when exceeding limit"
   end
 
+  def test_on_grace_start_does_not_fire_at_exact_limit
+    setup_plans_with_grace
+
+    org = create_organization
+    org.assign_pricing_plan!(:pro_with_grace)
+
+    track_events_via_callbacks!(:projects)
+
+    # Reach the limit exactly; grace should only start once usage goes over limit.
+    5.times { |i| org.projects.create!(name: "Project #{i + 1}") }
+
+    grace_events = @emitted_events.select { |e| e[:type] == :grace_start && e[:key] == :projects }
+    assert_equal 0, grace_events.count
+
+    state = PricingPlans::EnforcementState.find_by(plan_owner: org, limit_key: "projects")
+    assert_nil state&.exceeded_at
+  end
+
   def test_on_block_fires_automatically_when_grace_expires
     setup_plans_with_grace
 

test/controller_guards_test.rb

@@ -186,12 +186,14 @@ def test_require_plan_limit_exceeded_with_grace_then_block_expired_grace
       PricingPlans::PlanResolver.stub(:effective_plan_for, plan) do
         result = require_plan_limit!(:projects, plan_owner: @org, by: 1)
 
-        assert result.blocked?
-        assert_match(/reached your limit/i, result.message)
+        assert result.grace?
+        assert_match(/grace period/i, result.message)
 
-        # Should have marked as blocked
+        # Should still be in exceeded state, not blocked yet, because usage is
+        # at limit and this check models the next create attempt.
         state = PricingPlans::EnforcementState.find_by(plan_owner: @org, limit_key: "projects")
-        assert state&.blocked?
+        assert state&.exceeded?
+        refute state&.blocked?
       end
     end
   end

test/integration/complete_workflow_test.rb

@@ -282,11 +282,12 @@ def test_grace_period_expiration_workflow
         PricingPlans::PlanResolver.stub(:effective_plan_for, plan) do
           result = PricingPlans::ControllerGuards.require_plan_limit!(:projects, plan_owner: org)
         end
-        assert result.blocked?
+        assert result.grace?
 
-        # Should have emitted block event
+        # Block event is not emitted here because current usage remains at the
+        # limit; this check models the next create attempt.
         block_events = events.select { |e| e[:type] == :block }
-        assert_equal 1, block_events.size
+        assert_equal 0, block_events.size
       end
     end
   end

test/plan_owner_limits_helpers_test.rb

@@ -133,7 +133,7 @@ def test_limits_severity_and_message
     assert result.grace? || result.blocked? || result.warning?
 
     sev = @org.limits_severity(:projects, :custom_models)
-    assert_includes [:warning, :grace, :blocked], sev
+    assert_includes [:warning, :at_limit, :grace, :blocked], sev
 
     msg = @org.limits_message(:projects, :custom_models)
     assert msg.nil? || msg.is_a?(String)
@@ -283,6 +283,8 @@ def test_message_for_phrasing_per_severity
     end
     org3 = create_organization
     Project.send(:limited_by_pricing_plans, :projects, plan_owner: :organization)
+    org3.projects.create!(name: "P1")
+    org3.projects.create!(name: "P2")
     PricingPlans::GraceManager.mark_exceeded!(org3, :projects)
     msg = PricingPlans.message_for(org3, :projects)
     assert_includes msg, "You’re currently over your limit"