From 2eff74c85428bcd5035a691a7c8767e038e5a903 Mon Sep 17 00:00:00 2001
From: Shivam Mishra <scm.mymail@gmail.com>
Date: Thu, 8 Jan 2026 21:26:14 +0530
Subject: [PATCH] feat: make response usage increment atomic

---
 .../account/plan_usage_and_limits.rb          | 14 +++++--
 spec/enterprise/models/account_spec.rb        | 40 +++++++++++++++++++
 2 files changed, 51 insertions(+), 3 deletions(-)

diff --git a/enterprise/app/models/enterprise/account/plan_usage_and_limits.rb b/enterprise/app/models/enterprise/account/plan_usage_and_limits.rb
index ce03efa41..72289c000 100644
--- a/enterprise/app/models/enterprise/account/plan_usage_and_limits.rb
+++ b/enterprise/app/models/enterprise/account/plan_usage_and_limits.rb
@@ -16,9 +16,17 @@ module Enterprise::Account::PlanUsageAndLimits
   end
 
   def increment_response_usage
-    current_usage = custom_attributes[CAPTAIN_RESPONSES_USAGE].to_i || 0
-    custom_attributes[CAPTAIN_RESPONSES_USAGE] = current_usage + 1
-    save
+    increment_sql = <<~SQL.squish
+      custom_attributes = jsonb_set(
+        custom_attributes,
+        '{#{CAPTAIN_RESPONSES_USAGE}}',
+        to_jsonb(COALESCE((custom_attributes->>'#{CAPTAIN_RESPONSES_USAGE}')::int, 0) + 1),
+        true
+      )
+    SQL
+
+    updated = self.class.where(id: id).update_all(increment_sql)
+    reload if updated.positive?
   end
 
   def reset_response_usage
diff --git a/spec/enterprise/models/account_spec.rb b/spec/enterprise/models/account_spec.rb
index 1c727328e..d58edb3b8 100644
--- a/spec/enterprise/models/account_spec.rb
+++ b/spec/enterprise/models/account_spec.rb
@@ -78,6 +78,46 @@ RSpec.describe Account, type: :model do
         expect(responses_limits[:current_available]).to eq captain_limits[:startups][:responses] - 1
       end
 
+      it 'handles concurrent increments without losing updates' do
+        # Simulate concurrent calls to increment_response_usage
+        # This tests that the atomic SQL update prevents race conditions
+        concurrent_calls = 50
+
+        # Use a barrier to ensure all threads start at the same time
+        # This maximizes the chance of overlapping reads/writes
+        ready = Concurrent::CountDownLatch.new(concurrent_calls)
+        start = Concurrent::CountDownLatch.new(1)
+
+        threads = concurrent_calls.times.map do
+          Thread.new do
+            # Reload account in each thread to simulate separate workers
+            thread_account = Account.find(account.id)
+
+            # Signal ready and wait for all threads to be ready
+            ready.count_down
+            start.wait
+
+            # Now all threads will execute simultaneously
+            thread_account.increment_response_usage
+          end
+        end
+
+        # Wait for all threads to be ready
+        ready.wait
+        # Release all threads at once
+        start.count_down
+
+        # Wait for all threads to complete
+        threads.each(&:join)
+
+        # Verify all increments were captured without lost updates
+        expect(account.reload.custom_attributes['captain_responses_usage']).to eq concurrent_calls
+
+        responses_limits = account.usage_limits[:captain][:responses]
+        expect(responses_limits[:consumed]).to eq concurrent_calls
+        expect(responses_limits[:current_available]).to eq captain_limits[:startups][:responses] - concurrent_calls
+      end
+
       it 'reseting responses limits updates usage_limits' do
         account.custom_attributes['captain_responses_usage'] = 30
         account.save!