From e2848db0e8d890619e242d562c94df05361fe502 Mon Sep 17 00:00:00 2001
From: johnyeo <johnyeocx@gmail.com>
Date: Thu, 11 Jun 2026 18:26:52 +0100
Subject: [PATCH] clear entity cache on usage limit update

---
 .../fullSubject/updateEntityDataV2.lua        | 28 ++++++++
 .../deductionV2/executeRedisDeductionV2.ts    |  4 +-
 .../actions/updateCachedEntityData.ts         |  5 +-
 .../entity-usage-window-enforcement.test.ts   | 65 +++++++++++++++++++
 4 files changed, 100 insertions(+), 2 deletions(-)

diff --git a/server/src/_luaScriptsV2/fullSubject/updateEntityDataV2.lua b/server/src/_luaScriptsV2/fullSubject/updateEntityDataV2.lua
index 8bc42eebd..b93a8d3f6 100644
--- a/server/src/_luaScriptsV2/fullSubject/updateEntityDataV2.lua
+++ b/server/src/_luaScriptsV2/fullSubject/updateEntityDataV2.lua
@@ -50,6 +50,34 @@ for field_name, field_value in pairs(updates) do
   table.insert(updated_fields, field_name)
 end
 
+if updates.usage_limits ~= nil then
+  local seen_feature_ids = {}
+  local usage_window_feature_ids = {}
+
+  local function append_usage_limit_feature_ids(usage_limits)
+    if type(usage_limits) ~= 'table' then
+      return
+    end
+
+    for _, usage_limit in ipairs(usage_limits) do
+      if type(usage_limit) == 'table' and usage_limit.feature_id ~= nil then
+        local feature_id = usage_limit.feature_id
+        if seen_feature_ids[feature_id] == nil then
+          seen_feature_ids[feature_id] = true
+          table.insert(usage_window_feature_ids, feature_id)
+        end
+      end
+    end
+  end
+
+  if cached.customer ~= nil then
+    append_usage_limit_feature_ids(cached.customer.usage_limits)
+  end
+  append_usage_limit_feature_ids(cached.entity.usage_limits)
+
+  cached.usageWindowFeatureIds = usage_window_feature_ids
+end
+
 redis.call("SET", subject_key, cjson.encode(cached), "EX", cache_ttl)
 
 return cjson.encode({ success = true, updated_fields = updated_fields })
diff --git a/server/src/internal/balances/utils/deductionV2/executeRedisDeductionV2.ts b/server/src/internal/balances/utils/deductionV2/executeRedisDeductionV2.ts
index dbc24afca..acbf41009 100644
--- a/server/src/internal/balances/utils/deductionV2/executeRedisDeductionV2.ts
+++ b/server/src/internal/balances/utils/deductionV2/executeRedisDeductionV2.ts
@@ -278,6 +278,8 @@ export const executeRedisDeductionV2 = async ({
 		)
 			? resultJson.usage_window_mutations
 			: [];
+		const usageWindowsByFeatureId =
+			resultJson.usage_windows_by_feature_id ?? {};
 		const modifiedCustomerEntitlementIds = Array.isArray(
 			resultJson.modified_customer_entitlement_ids,
 		)
@@ -301,7 +303,7 @@ export const executeRedisDeductionV2 = async ({
 		// Typed handoff for the PG mirror; empty arrays kept (prune-to-empty
 		// must still full-replace).
 		for (const [featureId, usageWindows] of Object.entries(
-			resultJson.usage_windows_by_feature_id ?? {},
+			usageWindowsByFeatureId,
 		)) {
 			allUsageWindowUpdates[featureId] = {
 				internal_customer_id: fullSubject.internalCustomerId,
diff --git a/server/src/internal/customers/cache/fullSubject/actions/updateCachedEntityData.ts b/server/src/internal/customers/cache/fullSubject/actions/updateCachedEntityData.ts
index 028aa165e..009e86261 100644
--- a/server/src/internal/customers/cache/fullSubject/actions/updateCachedEntityData.ts
+++ b/server/src/internal/customers/cache/fullSubject/actions/updateCachedEntityData.ts
@@ -20,7 +20,10 @@ export const updateCachedEntityData = async ({
 	customerId: string;
 	entityId: string;
 	updates: Partial<
-		Pick<Entity, "spend_limits" | "usage_alerts" | "overage_allowed">
+		Pick<
+			Entity,
+			"spend_limits" | "usage_limits" | "usage_alerts" | "overage_allowed"
+		>
 	>;
 }): Promise<void> => {
 	if (Object.keys(updates).length === 0) return;
diff --git a/server/tests/integration/balances/usage-windows/entities/entity-usage-window-enforcement.test.ts b/server/tests/integration/balances/usage-windows/entities/entity-usage-window-enforcement.test.ts
index b05e6950e..9d3f4c7b0 100644
--- a/server/tests/integration/balances/usage-windows/entities/entity-usage-window-enforcement.test.ts
+++ b/server/tests/integration/balances/usage-windows/entities/entity-usage-window-enforcement.test.ts
@@ -131,6 +131,71 @@ test.concurrent(
 	},
 );
 
+test.concurrent(
+	`${chalk.yellowBright("ent-uw-enforce1b: entities.update usage_limits patches a pre-existing cached subject")}`,
+	async () => {
+		const perEntityProduct = products.base({
+			id: "ent-uw-enforce-cached-update",
+			items: [
+				items.monthlyMessages({
+					includedUsage: 100,
+					entityFeatureId: TestFeature.Users,
+				}),
+			],
+		});
+
+		const customerId = "ent-uw-enforce-cached-1";
+		const { entities } = await initScenario({
+			customerId,
+			setup: [
+				s.customer({ testClock: false }),
+				s.products({ list: [perEntityProduct] }),
+				s.entities({ count: 1, featureId: TestFeature.Users }),
+			],
+			actions: [s.billing.attach({ productId: perEntityProduct.id })],
+		});
+
+		await autumnV2_3.check({
+			customer_id: customerId,
+			entity_id: entities[0].id,
+			feature_id: TestFeature.Messages,
+		});
+
+		await setEntityUsageLimit({
+			autumn: autumnV2_3,
+			customerId,
+			entityId: entities[0].id,
+			featureId: TestFeature.Messages,
+			limit: 5,
+		});
+
+		await autumnV2_3.track({
+			customer_id: customerId,
+			entity_id: entities[0].id,
+			feature_id: TestFeature.Messages,
+			value: 7,
+		});
+
+		await expectEntityFeatureBalance({
+			autumn: autumnV2_3,
+			customerId,
+			entityId: entities[0].id,
+			featureId: TestFeature.Messages,
+			granted: 100,
+			remaining: 95,
+			usage: 5,
+		});
+		await expectEntityUsageLimit({
+			autumn: autumnV2_3,
+			customerId,
+			entityId: entities[0].id,
+			featureId: TestFeature.Messages,
+			usage: 5,
+			limit: 5,
+		});
+	},
+);
+
 test.concurrent(
 	`${chalk.yellowBright("ent-uw-enforce2: two entities with different caps stay isolated while customer balance aggregates")}`,
 	async () => {