feat: optimize entities cache

2026-04-22 20:06:09 +01:00
parent 9f6ffb5409
commit 21658f1b78
11 changed files with 573 additions and 134 deletions
--- a/server/experiments/benchmarkV2CacheRedis.ts
+++ b/server/experiments/benchmarkV2CacheRedis.ts
@@ -0,0 +1,440 @@
+import { AppEnv } from "@autumn/shared";
+import type { Redis } from "ioredis";
+import {
+	initDrizzle,
+	prodTestCustomerId,
+	prodTestEntityId,
+	prodTestOrgId,
+} from "./experimentEnv";
+
+// ---------------------------------------------------------------------------
+// EDIT ME — hardcode the customer / feature you want to diagnose
+// ---------------------------------------------------------------------------
+const CONFIG = {
+	orgId: prodTestOrgId as string,
+	env: AppEnv.Live as AppEnv,
+	customerId: prodTestCustomerId as string,
+	entityId: prodTestEntityId as string | undefined, // or undefined
+	featureId: undefined as string | undefined, // set to isolate to one feature
+	iterations: 5,
+	runLua: true, // lua writes back entitlement hashes even for amount=0 — disable if paranoid
+	forceRepopulate: true, // DEL the cached subject before priming, so it's rebuilt from the DB
+};
+// ---------------------------------------------------------------------------
+
+const { resolveRedisV2 } = await import("../src/external/redis/resolveRedisV2");
+const { warmupRedisV2 } = await import("../src/external/redis/initRedisV2");
+const { buildFullSubjectKey } = await import(
+	"../src/internal/customers/cache/fullSubject/builders/buildFullSubjectKey"
+);
+const { buildSharedFullSubjectBalanceKey } = await import(
+	"../src/internal/customers/cache/fullSubject/builders/buildSharedFullSubjectBalanceKey"
+);
+const { buildDeductFromSubjectBalancesKeys } = await import(
+	"../src/internal/customers/cache/fullSubject/builders/buildDeductFromSubjectBalancesKeys"
+);
+const { getOrInitFullSubjectViewEpoch } = await import(
+	"../src/internal/customers/cache/fullSubject/actions/invalidate/getOrInitFullSubjectViewEpoch"
+);
+const { getCachedFeatureBalancesBatch } = await import(
+	"../src/internal/customers/cache/fullSubject/balances/getCachedFeatureBalances"
+);
+const { getOrSetCachedFullSubject } = await import(
+	"../src/internal/customers/cache/fullSubject/actions/getOrSetCachedFullSubject"
+);
+const { sanitizeCachedFullSubject } = await import(
+	"../src/internal/customers/cache/fullSubject/sanitize/index"
+);
+const { AGGREGATED_BALANCE_FIELD } = await import(
+	"../src/internal/customers/cache/fullSubject/config/fullSubjectCacheConfig"
+);
+
+type Step = { name: string; ms: number; extra?: string };
+type Iter = { index: number; steps: Step[] };
+
+const nowMs = () => performance.now();
+
+const time = async <T>(
+	name: string,
+	fn: () => Promise<T>,
+): Promise<{ step: Step; value: T }> => {
+	const start = nowMs();
+	const value = await fn();
+	const ms = nowMs() - start;
+	return { step: { name, ms }, value };
+};
+
+const makeCtx = ({ redisV2, db }: { redisV2: Redis; db: unknown }) =>
+	({
+		org: { id: CONFIG.orgId },
+		env: CONFIG.env,
+		redisV2,
+		db,
+		dbGeneral: db,
+		features: [],
+		skipCache: false,
+		isPublic: false,
+		extraLogs: {},
+		logger: {
+			debug: (...a: unknown[]) => console.debug(...a),
+			info: (...a: unknown[]) => console.info(...a),
+			warn: (...a: unknown[]) => console.warn(...a),
+			error: (...a: unknown[]) => console.error(...a),
+		},
+	}) as unknown as Parameters<typeof getOrInitFullSubjectViewEpoch>[0]["ctx"];
+
+const pct = (sorted: number[], p: number): number => {
+	if (sorted.length === 0) return 0;
+	const idx = Math.min(sorted.length - 1, Math.floor((p / 100) * sorted.length));
+	return sorted[idx];
+};
+
+const summarize = (all: Iter[]) => {
+	const byStep = new Map<string, number[]>();
+	for (const iter of all) {
+		for (const step of iter.steps) {
+			const arr = byStep.get(step.name) ?? [];
+			arr.push(step.ms);
+			byStep.set(step.name, arr);
+		}
+	}
+
+	console.log("\n=== Summary (ms) ===");
+	console.log(
+		"step".padEnd(48),
+		"min".padStart(8),
+		"p50".padStart(8),
+		"p95".padStart(8),
+		"max".padStart(8),
+	);
+	for (const [name, arr] of byStep) {
+		const sorted = [...arr].sort((a, b) => a - b);
+		console.log(
+			name.padEnd(48),
+			sorted[0].toFixed(2).padStart(8),
+			pct(sorted, 50).toFixed(2).padStart(8),
+			pct(sorted, 95).toFixed(2).padStart(8),
+			sorted[sorted.length - 1].toFixed(2).padStart(8),
+		);
+	}
+};
+
+const runIteration = async ({
+	i,
+	redisV2,
+	ctx,
+}: {
+	i: number;
+	redisV2: Redis;
+	// biome-ignore lint/suspicious/noExplicitAny: stub ctx
+	ctx: any;
+}): Promise<Iter> => {
+	const steps: Step[] = [];
+	const subjectKey = buildFullSubjectKey({
+		orgId: CONFIG.orgId,
+		env: CONFIG.env,
+		customerId: CONFIG.customerId,
+		entityId: CONFIG.entityId,
+	});
+
+	// A1. GET subject
+	const { step: s1, value: cachedRaw } = await time(
+		"GET full_subject",
+		() => redisV2.get(subjectKey),
+	);
+	steps.push({ ...s1, extra: cachedRaw ? `${cachedRaw.length}b` : "null" });
+
+	if (!cachedRaw) {
+		console.log(`  iter ${i}: no cached subject — stopping this iteration`);
+		return { index: i, steps };
+	}
+
+	// parse
+	const { step: sParse, value: cached } = await time(
+		"JSON.parse + sanitize",
+		async () =>
+			sanitizeCachedFullSubject({
+				cachedFullSubject: JSON.parse(cachedRaw),
+			}),
+	);
+	steps.push(sParse);
+
+	if (i === 1) {
+		const sizeOf = (v: unknown) => JSON.stringify(v ?? null).length;
+		const cachedAny = cached as unknown as Record<string, unknown>;
+		const entries = Object.keys(cachedAny)
+			.map((k) => ({ k, bytes: sizeOf(cachedAny[k]) }))
+			.sort((a, b) => b.bytes - a.bytes);
+		const total = entries.reduce((a, e) => a + e.bytes, 0);
+		console.log(`    blob anatomy (total ${(total / 1024).toFixed(1)}kb):`);
+		for (const e of entries.slice(0, 8)) {
+			const pct = ((e.bytes / total) * 100).toFixed(1);
+			console.log(
+				`      ${(e.bytes / 1024).toFixed(1).padStart(8)}kb  ${pct.padStart(5)}%  ${e.k}`,
+			);
+		}
+	}
+
+	// A2. epoch
+	const { step: s2 } = await time("getOrInitFullSubjectViewEpoch", () =>
+		getOrInitFullSubjectViewEpoch({ ctx, customerId: CONFIG.customerId }),
+	);
+	steps.push(s2);
+
+	const meteredFeatures = CONFIG.featureId
+		? cached.meteredFeatures.filter((f: string) => f === CONFIG.featureId)
+		: cached.meteredFeatures;
+
+	const includeAggregated = !CONFIG.entityId;
+
+	// A3. batch hmget pipeline
+	const { step: s3, value: batchOutcome } = await time(
+		`getCachedFeatureBalancesBatch (${meteredFeatures.length} features)`,
+		() =>
+			getCachedFeatureBalancesBatch({
+				ctx,
+				customerId: CONFIG.customerId,
+				featureIds: meteredFeatures,
+				customerEntitlementIdsByFeatureId:
+					cached.customerEntitlementIdsByFeatureId,
+				includeAggregated,
+			}),
+	);
+	steps.push({ ...s3, extra: batchOutcome.kind });
+
+	// A4. sequential per-feature HMGET with per-feature timing (slow-key hunt)
+	const perFeature: Array<{ featureId: string; ms: number; bytes: number }> =
+		[];
+	const { step: s4 } = await time(
+		`sequential HMGET per feature (${meteredFeatures.length} calls)`,
+		async () => {
+			for (const featureId of meteredFeatures) {
+				const ids =
+					cached.customerEntitlementIdsByFeatureId[featureId] ?? [];
+				const fields = includeAggregated
+					? [...ids, AGGREGATED_BALANCE_FIELD]
+					: ids;
+				const balanceKey = buildSharedFullSubjectBalanceKey({
+					orgId: CONFIG.orgId,
+					env: CONFIG.env,
+					customerId: CONFIG.customerId,
+					featureId,
+				});
+				const t = nowMs();
+				let bytes = 0;
+				if (fields.length === 0) {
+					await redisV2.exists(balanceKey);
+				} else {
+					const vals = (await redisV2.hmget(
+						balanceKey,
+						...fields,
+					)) as (string | null)[];
+					bytes = vals.reduce((a, v) => a + (v?.length ?? 0), 0);
+				}
+				perFeature.push({ featureId, ms: nowMs() - t, bytes });
+			}
+		},
+	);
+	steps.push(s4);
+
+	// Only print on first iteration to keep output tidy
+	if (i === 1) {
+		const sorted = [...perFeature].sort((a, b) => b.ms - a.ms);
+		console.log("    top-5 slowest features (seq HMGET):");
+		for (const f of sorted.slice(0, 5)) {
+			console.log(
+				`      ${f.ms.toFixed(2).padStart(7)}ms  ${(f.bytes / 1024).toFixed(1).padStart(6)}kb  ${f.featureId}`,
+			);
+		}
+	}
+
+	// C. dry lua invocation (amount_to_deduct=0)
+	if (CONFIG.runLua && meteredFeatures.length > 0) {
+		const customerEntitlementDeductions: Array<{
+			customer_entitlement_id: string;
+			credit_cost: number;
+			feature_id: string;
+			entity_feature_id: string | null;
+			usage_allowed: boolean;
+			min_balance: number;
+			max_balance: number;
+		}> = [];
+
+		for (const featureId of meteredFeatures) {
+			const ids = cached.customerEntitlementIdsByFeatureId[featureId] ?? [];
+			for (const id of ids) {
+				customerEntitlementDeductions.push({
+					customer_entitlement_id: id,
+					credit_cost: 0,
+					feature_id: featureId,
+					entity_feature_id: null,
+					usage_allowed: false,
+					min_balance: 0,
+					max_balance: 0,
+				});
+			}
+		}
+
+		if (customerEntitlementDeductions.length === 0) {
+			console.log("  (skipping lua — no customer entitlements)");
+		} else {
+			const routingKey = subjectKey;
+			const { keys, balanceKeyIndexByFeatureId } =
+				buildDeductFromSubjectBalancesKeys({
+					orgId: CONFIG.orgId,
+					env: CONFIG.env,
+					customerId: CONFIG.customerId,
+					routingKey,
+					lockReceiptKey: null,
+					customerEntitlementDeductions,
+					fallbackFeatureId:
+						CONFIG.featureId ?? meteredFeatures[0] ?? "unknown",
+				});
+
+			const luaParams = {
+				org_id: CONFIG.orgId,
+				env: CONFIG.env,
+				customer_id: CONFIG.customerId,
+				customer_entitlement_deductions: customerEntitlementDeductions,
+				balance_key_index_by_feature_id: balanceKeyIndexByFeatureId,
+				spend_limit_by_feature_id: null,
+				usage_based_cus_ent_ids_by_feature_id: null,
+				amount_to_deduct: 0,
+				target_balance: null,
+				target_entity_id: CONFIG.entityId ?? null,
+				rollovers: null,
+				skip_additional_balance: false,
+				alter_granted_balance: false,
+				overage_behaviour: "cap",
+				feature_id: CONFIG.featureId ?? meteredFeatures[0],
+				lock: null,
+				unwind_value: null,
+				debug: true,
+			};
+
+			const { step: sLua, value: luaRaw } = await time(
+				`deductFromSubjectBalances lua (${keys.length} keys, ${customerEntitlementDeductions.length} ents)`,
+				() =>
+					// biome-ignore lint/suspicious/noExplicitAny: ioredis custom command typed via module augmentation not available here
+					(redisV2 as any).deductFromSubjectBalances(
+						keys.length,
+						...keys,
+						JSON.stringify(luaParams),
+					),
+			);
+			steps.push(sLua);
+
+			try {
+				const luaResult = JSON.parse(luaRaw as string);
+				if (luaResult.error) {
+					console.log(`    lua error: ${luaResult.error}`);
+				}
+				if (luaResult.logs?.length) {
+					console.log(`    lua logs (${luaResult.logs.length} lines):`);
+					for (const line of luaResult.logs) console.log(`      ${line}`);
+				}
+			} catch {
+				console.log("    (lua result unparseable)");
+			}
+		}
+	}
+
+	return { index: i, steps };
+};
+
+const main = async () => {
+	console.log("=== V2 cache Redis benchmark ===");
+	console.log("config:", {
+		orgId: CONFIG.orgId,
+		env: CONFIG.env,
+		customerId: CONFIG.customerId,
+		entityId: CONFIG.entityId,
+		featureId: CONFIG.featureId,
+		iterations: CONFIG.iterations,
+		runLua: CONFIG.runLua,
+	});
+
+	const { db } = initDrizzle();
+	const redisV2 = resolveRedisV2();
+	await warmupRedisV2();
+
+	const ctx = makeCtx({ redisV2, db });
+
+	// Prime the cache if cold, so the iteration benchmarks always measure a hit path.
+	const subjectKey = buildFullSubjectKey({
+		orgId: CONFIG.orgId,
+		env: CONFIG.env,
+		customerId: CONFIG.customerId,
+		entityId: CONFIG.entityId,
+	});
+	if (CONFIG.forceRepopulate) {
+		// Grab meteredFeatures from the existing subject so we know which
+		// balance hashes to drop. All keys share the {customerId} hash tag,
+		// so a single DEL across them is slot-safe.
+		const staleRaw = await redisV2.get(subjectKey);
+		const balanceKeys: string[] = [];
+		if (staleRaw) {
+			try {
+				const stale = JSON.parse(staleRaw) as {
+					meteredFeatures?: string[];
+				};
+				for (const featureId of stale.meteredFeatures ?? []) {
+					balanceKeys.push(
+						buildSharedFullSubjectBalanceKey({
+							orgId: CONFIG.orgId,
+							env: CONFIG.env,
+							customerId: CONFIG.customerId,
+							featureId,
+						}),
+					);
+				}
+			} catch {
+				console.warn("  (could not parse stale subject to extract balance keys)");
+			}
+		}
+		const toDel = [subjectKey, ...balanceKeys];
+		const deleted = await redisV2.del(...toDel);
+		console.log(
+			`\n--- forceRepopulate: DEL ${toDel.length} keys (subject + ${balanceKeys.length} balance hashes) → ${deleted} removed ---`,
+		);
+	}
+
+	const existing = await redisV2.get(subjectKey);
+	if (!existing) {
+		console.log("\n--- priming cache (cold) ---");
+		const t0 = nowMs();
+		try {
+			await getOrSetCachedFullSubject({
+				ctx,
+				customerId: CONFIG.customerId,
+				entityId: CONFIG.entityId,
+				source: "benchmarkV2CacheRedis",
+			});
+			console.log(`  primed in ${(nowMs() - t0).toFixed(2)}ms`);
+		} catch (err) {
+			console.error("  prime failed:", err);
+			console.error(
+				"  (cache is cold and cannot be populated from this script — aborting)",
+			);
+			process.exit(1);
+		}
+	} else {
+		console.log(`\ncache already warm (${existing.length}b)`);
+	}
+
+	const all: Iter[] = [];
+	for (let i = 1; i <= CONFIG.iterations; i++) {
+		console.log(`\n--- iteration ${i} ---`);
+		const iter = await runIteration({ i, redisV2, ctx });
+		for (const step of iter.steps) {
+			const extra = step.extra ? `  (${step.extra})` : "";
+			console.log(`  ${step.ms.toFixed(2).padStart(8)}ms  ${step.name}${extra}`);
+		}
+		all.push(iter);
+	}
+
+	summarize(all);
+	process.exit(0);
+};
+
+await main();
--- a/server/experiments/explainEntityAggregate.ts
+++ b/server/experiments/explainEntityAggregate.ts
@@ -31,7 +31,7 @@ const main = async () => {

 	console.log("--- Running entity aggregate query ---");
 	const start = performance.now();
-	const result = await getEntityAggregateForSync({
+	const parsedResult = await getEntityAggregateForSync({
 		db,
 		orgId,
 		env,
@@ -39,23 +39,19 @@ const main = async () => {
 		internalFeatureIds: features?.map((feature) => feature.internal_id),
 	});
 	const elapsed = performance.now() - start;
-	console.log(`Rows returned: ${result.length}`);
-	console.log(`Wall-clock time: ${elapsed.toFixed(2)}ms`);
-	console.log("Result:", JSON.stringify(result, null, 2));
-	console.log();
+	console.log(`Rows (after schema parse): ${parsedResult.length}`);
+	console.log(`Wall-clock time: ${elapsed.toFixed(2)}ms\n`);

-	// Build the same query inline for EXPLAIN ANALYZE
+	// Also fetch raw rows so we can see the output even if schema parse fails
 	const { getEntityAggregateFragments } = await import(
 		"../src/internal/customers/repos/getFullSubject/getEntityAggregateFragments"
 	);
-
-	const statusFilter = sql`AND cp.status = ANY(ARRAY['active', 'past_due', 'scheduled'])`;
-	const entityFragments = getEntityAggregateFragments({
-		statusFilter,
+	const statusFilterRaw = sql`AND cp.status = ANY(ARRAY['active', 'past_due', 'scheduled'])`;
+	const entityFragmentsRaw = getEntityAggregateFragments({
+		statusFilter: statusFilterRaw,
 		internalFeatureIds: features?.map((feature) => feature.internal_id),
 	});
-
-	const query = sql`
+	const rawQuery = sql`
 		WITH subject_customer_records AS (
 			SELECT *
 			FROM customers c
@@ -65,15 +61,36 @@ const main = async () => {
 			ORDER BY (c.id = ${customerId}) DESC
 			LIMIT 1
 		)
-
-		${entityFragments.ctes}
-
+		${entityFragmentsRaw.ctes}
 		SELECT *
 		FROM entity_aggregated_cus_entitlements
 	`;
+	const rawResult = await db.execute(rawQuery);
+	const result = rawResult as unknown as Record<string, unknown>[];
+	console.log(`Raw rows: ${result.length}\n`);
+
+	console.log("--- Entity aggregate output (per feature) ---\n");
+	for (const row of result) {
+		const entities = (row as unknown as { entities?: Record<string, unknown> })
+			.entities;
+		const entityCount = (row as unknown as { entity_count?: number })
+			.entity_count;
+		const entityKeys = entities ? Object.keys(entities) : [];
+
+		console.log(`feature_id: ${row.feature_id}`);
+		console.log(`  internal_feature_id: ${row.internal_feature_id}`);
+		console.log(`  balance: ${row.balance}`);
+		console.log(`  adjustment: ${row.adjustment}`);
+		console.log(`  additional_balance: ${row.additional_balance}`);
+		console.log(`  rollover_balance: ${row.rollover_balance}`);
+		console.log(`  entity_count: ${entityCount}`);
+		console.log(`  entities keys (${entityKeys.length}): ${JSON.stringify(entityKeys)}`);
+		console.log(`  entities: ${JSON.stringify(entities, null, 2)}`);
+		console.log();
+	}

 	console.log("--- EXPLAIN (ANALYZE, BUFFERS) ---\n");
-	const explainQuery = sql`EXPLAIN (ANALYZE, BUFFERS, FORMAT TEXT) ${query}`;
+	const explainQuery = sql`EXPLAIN (ANALYZE, BUFFERS, FORMAT TEXT) ${rawQuery}`;
 	const explainResult = await db.execute(explainQuery);

 	for (const row of explainResult) {
--- a/server/experiments/explainGetFullSubject.ts
+++ b/server/experiments/explainGetFullSubject.ts
@@ -1,3 +1,5 @@
+import { writeFileSync } from "node:fs";
+import { resolve } from "node:path";
 import { AppEnv } from "@autumn/shared";
 import { sql } from "drizzle-orm";
 import {
@@ -43,17 +45,67 @@ const main = async () => {
 	const elapsed = performance.now() - start;
 	console.log(`Rows returned: ${result.length}`);
 	console.log(`Wall-clock time: ${elapsed.toFixed(2)}ms`);
-	console.log("Result:", JSON.stringify(result, null, 2));
+
+	// The full row payload is enormous — dump it to a file and print a summary.
+	const outDir = resolve(import.meta.dir, "out");
+	try {
+		const { mkdirSync } = await import("node:fs");
+		mkdirSync(outDir, { recursive: true });
+	} catch {
+		// ignore
+	}
+	const resultPath = resolve(outDir, "getFullSubject-result.json");
+	writeFileSync(resultPath, JSON.stringify(result, null, 2));
+	console.log(`Result written to: ${resultPath}`);
+
+	const row0 = (result[0] ?? {}) as Record<string, unknown>;
+	const summarize = (value: unknown): unknown => {
+		if (Array.isArray(value)) return `Array(len=${value.length})`;
+		if (value && typeof value === "object")
+			return `Object(keys=${Object.keys(value).length})`;
+		return value;
+	};
+	console.log("\n--- Top-level fields (summary) ---");
+	for (const [key, value] of Object.entries(row0)) {
+		console.log(`  ${key}: ${summarize(value)}`);
+	}
+
+	const aggregated = (row0.aggregated_customer_entitlements ?? []) as Array<
+		Record<string, unknown>
+	>;
+	if (aggregated.length > 0) {
+		console.log(
+			`\n--- aggregated_customer_entitlements (${aggregated.length}) ---`,
+		);
+		for (const ae of aggregated) {
+			const entities = (ae.entities ?? {}) as Record<string, unknown>;
+			console.log(
+				`  ${ae.feature_id}: balance=${ae.balance}, adj=${ae.adjustment}, add=${ae.additional_balance}, rollover=${ae.rollover_balance}, entities.keys=${Object.keys(entities).length}`,
+			);
+		}
+	}
 	console.log();

 	console.log("--- EXPLAIN (ANALYZE, BUFFERS) ---\n");
 	const explainQuery = sql`EXPLAIN (ANALYZE, BUFFERS, FORMAT TEXT) ${query}`;
 	const explainResult = await db.execute(explainQuery);

+	const explainPath = resolve(outDir, "getFullSubject-explain.txt");
+	const explainLines: string[] = [];
 	for (const row of explainResult) {
 		const line = (row as Record<string, unknown>)["QUERY PLAN"];
-		console.log(line);
+		if (typeof line === "string") explainLines.push(line);
 	}
+	writeFileSync(explainPath, explainLines.join("\n"));
+	console.log(`EXPLAIN written to: ${explainPath}`);
+
+	// Print just the top of the plan (total time, first ~20 lines) inline
+	const head = explainLines.slice(0, 20);
+	for (const line of head) console.log(line);
+	const execLine = explainLines.find((line) => /Execution Time:/.test(line));
+	const planLine = explainLines.find((line) => /Planning Time:/.test(line));
+	if (planLine) console.log(planLine);
+	if (execLine) console.log(execLine);

 	process.exit(0);
 };
--- a/server/src/internal/customers/repos/getFullSubject/getEntityAggregateFragments.ts
+++ b/server/src/internal/customers/repos/getFullSubject/getEntityAggregateFragments.ts
@@ -2,60 +2,39 @@ import { type SQL, sql } from "drizzle-orm";
 import { getEntityOptionsAggregateFragments } from "./getEntityOptionsAggregateFragments.js";

 /**
- * Rollover CTEs driven from the shared `entity_product_cus_ents` and
- * `entity_loose_cus_ents` base CTEs — avoids re-scanning `customer_entitlements`
- * and `customer_products` per branch.
+ * Rollover CTEs:
+ *   - entity_rollover_keys: per-entity rollover balances, built ONLY from
+ *     jsonb_each(r.entities). Keys only exist in the map when a rollover row
+ *     explicitly carries a per-entity breakdown.
+ *   - entity_rollover_feature: feature-level rollover totals, summed directly
+ *     off r.balance / r.usage (one add per active rollover).
 */
 const buildEntityRolloverCtes = () => sql`
-		entity_rollover_rows AS (
-			-- Top-level rollover: one row per (rollover × cus_ent).
-			-- entity_key = cp.internal_entity_id for product-attached, ce.internal_entity_id for loose.
-			SELECT
-				ce.internal_feature_id,
-				ce.internal_customer_id,
-				COALESCE(ce.cp_entity_key, ce.internal_entity_id) AS entity_key,
-				r.balance::numeric AS rollover_balance,
-				COALESCE(r.usage, 0)::numeric AS rollover_usage
-			FROM rollovers r
-			JOIN entity_level_cus_ents ce ON r.cus_ent_id = ce.id
-			WHERE (r.expires_at IS NULL OR r.expires_at > EXTRACT(EPOCH FROM now()) * 1000)
-
-			UNION ALL
-
-			-- Per-entity rollover from jsonb_each(r.entities).
+		entity_rollover_keys AS (
 			SELECT
 				ce.internal_feature_id,
 				ce.internal_customer_id,
 				kv.entity_key AS entity_key,
-				COALESCE((kv.entity_value->>'balance')::numeric, 0) AS rollover_balance,
-				COALESCE((kv.entity_value->>'usage')::numeric, 0) AS rollover_usage
+				SUM(COALESCE((kv.entity_value->>'balance')::numeric, 0)) AS rollover_balance,
+				SUM(COALESCE((kv.entity_value->>'usage')::numeric, 0)) AS rollover_usage
 			FROM rollovers r
 			JOIN entity_level_cus_ents ce ON r.cus_ent_id = ce.id
 			CROSS JOIN LATERAL jsonb_each(r.entities) AS kv(entity_key, entity_value)
 			WHERE jsonb_typeof(r.entities) = 'object'
 				AND (r.expires_at IS NULL OR r.expires_at > EXTRACT(EPOCH FROM now()) * 1000)
-		),
-
-		entity_rollover_keys AS (
-			SELECT
-				internal_feature_id,
-				internal_customer_id,
-				entity_key,
-				SUM(rollover_balance) AS rollover_balance,
-				SUM(rollover_usage) AS rollover_usage
-			FROM entity_rollover_rows
-			WHERE entity_key IS NOT NULL
-			GROUP BY internal_feature_id, internal_customer_id, entity_key
+			GROUP BY ce.internal_feature_id, ce.internal_customer_id, kv.entity_key
 		),

 		entity_rollover_feature AS (
 			SELECT
-				internal_feature_id,
-				internal_customer_id,
-				SUM(rollover_balance) AS rollover_balance,
-				SUM(rollover_usage) AS rollover_usage
-			FROM entity_rollover_rows
-			GROUP BY internal_feature_id, internal_customer_id
+				ce.internal_feature_id,
+				ce.internal_customer_id,
+				SUM(r.balance::numeric) AS rollover_balance,
+				SUM(COALESCE(r.usage, 0)::numeric) AS rollover_usage
+			FROM rollovers r
+			JOIN entity_level_cus_ents ce ON r.cus_ent_id = ce.id
+			WHERE (r.expires_at IS NULL OR r.expires_at > EXTRACT(EPOCH FROM now()) * 1000)
+			GROUP BY ce.internal_feature_id, ce.internal_customer_id
 		)
 `;

@@ -145,65 +124,23 @@ export const getEntityAggregateFragments = ({

 		${entityOptionsAggregateFragments.ctes},

-		entity_balance_rows AS (
-			-- Top-level: one row per cus_ent (product-attached or loose).
-			-- entity_key = cp.internal_entity_id for product-attached, ce.internal_entity_id for loose.
-			SELECT
-				COALESCE(ce.external_id, ce.id) AS api_id,
-				ce.internal_feature_id,
-				ce.internal_customer_id,
-				ce.feature_id,
-				COALESCE(ent.allowance, 0)::numeric AS allowance,
-				ce.balance::numeric AS balance,
-				COALESCE(ce.adjustment, 0)::numeric AS adjustment,
-				COALESCE(ce.additional_balance, 0)::numeric AS additional_balance,
-				ce.unlimited,
-				ce.usage_allowed,
-				COALESCE(ce.cp_entity_key, ce.internal_entity_id) AS entity_key,
-				ce.balance::numeric AS entity_balance,
-				COALESCE(ce.adjustment, 0)::numeric AS entity_adjustment,
-				COALESCE(ce.additional_balance, 0)::numeric AS entity_additional_balance
-			FROM entity_level_cus_ents ce
-			JOIN entitlements ent ON ce.entitlement_id = ent.id
-
-			UNION ALL
-
-			-- Per-entity: N rows per cus_ent from jsonb_each(ce.entities).
-			-- balance/adj/additional = 0 to avoid double-counting at the aggregate level.
-			SELECT
-				COALESCE(ce.external_id, ce.id) AS api_id,
-				ce.internal_feature_id,
-				ce.internal_customer_id,
-				ce.feature_id,
-				COALESCE(ent.allowance, 0)::numeric AS allowance,
-				0::numeric AS balance,
-				0::numeric AS adjustment,
-				0::numeric AS additional_balance,
-				ce.unlimited,
-				ce.usage_allowed,
-				kv.entity_key AS entity_key,
-				(kv.entity_value->>'balance')::numeric AS entity_balance,
-				COALESCE((kv.entity_value->>'adjustment')::numeric, 0) AS entity_adjustment,
-				COALESCE((kv.entity_value->>'additional_balance')::numeric, 0) AS entity_additional_balance
-			FROM entity_level_cus_ents ce
-			JOIN entitlements ent ON ce.entitlement_id = ent.id
-			CROSS JOIN LATERAL jsonb_each(ce.entities) AS kv(entity_key, entity_value)
-			WHERE jsonb_typeof(ce.entities) = 'object'
-		),
-
 		${buildEntityRolloverCtes()},

+		-- Per-entity balance map: built ONLY from jsonb_each(ce.entities) so
+		-- that entities keys reflect real per-entity breakdowns, not every
+		-- product-attached cus_ent. Summed across cus_ents for the same key.
 		entity_balance_keys AS (
 			SELECT
-				internal_feature_id,
-				internal_customer_id,
-				entity_key,
-				SUM(entity_balance) AS balance,
-				SUM(entity_adjustment) AS adjustment,
-				SUM(entity_additional_balance) AS additional_balance
-			FROM entity_balance_rows
-			WHERE entity_key IS NOT NULL
-			GROUP BY internal_feature_id, internal_customer_id, entity_key
+				ce.internal_feature_id,
+				ce.internal_customer_id,
+				kv.entity_key,
+				SUM((kv.entity_value->>'balance')::numeric) AS balance,
+				SUM(COALESCE((kv.entity_value->>'adjustment')::numeric, 0)) AS adjustment,
+				SUM(COALESCE((kv.entity_value->>'additional_balance')::numeric, 0)) AS additional_balance
+			FROM entity_level_cus_ents ce
+			CROSS JOIN LATERAL jsonb_each(ce.entities) AS kv(entity_key, entity_value)
+			WHERE jsonb_typeof(ce.entities) = 'object'
+			GROUP BY ce.internal_feature_id, ce.internal_customer_id, kv.entity_key
 		),

 		entity_aggregate_map AS (
@@ -250,21 +187,23 @@ export const getEntityAggregateFragments = ({
 			GROUP BY ejk.internal_feature_id, ejk.internal_customer_id
 		),

+		-- Feature-level totals: summed directly off entity_level_cus_ents
+		-- (one add per cus_ent).
 		entity_aggregate_totals AS (
 			SELECT
-				MIN(ebr.api_id) AS api_id,
-				ebr.internal_feature_id,
-				ebr.internal_customer_id,
-				MIN(ebr.feature_id) AS feature_id,
-				SUM(ebr.allowance) AS allowance_total,
-				SUM(ebr.balance) AS balance,
-				SUM(ebr.adjustment) AS adjustment,
-				SUM(ebr.additional_balance) AS additional_balance,
-				BOOL_OR(ebr.unlimited) AS unlimited,
-				BOOL_OR(ebr.usage_allowed) AS usage_allowed,
-				COUNT(DISTINCT ebr.entity_key) FILTER (WHERE ebr.entity_key IS NOT NULL) AS entity_count
-			FROM entity_balance_rows ebr
-			GROUP BY ebr.internal_feature_id, ebr.internal_customer_id
+				MIN(COALESCE(ce.external_id, ce.id)) AS api_id,
+				ce.internal_feature_id,
+				ce.internal_customer_id,
+				MIN(ce.feature_id) AS feature_id,
+				SUM(COALESCE(ent.allowance, 0)::numeric) AS allowance_total,
+				SUM(ce.balance::numeric) AS balance,
+				SUM(COALESCE(ce.adjustment, 0)::numeric) AS adjustment,
+				SUM(COALESCE(ce.additional_balance, 0)::numeric) AS additional_balance,
+				BOOL_OR(ce.unlimited) AS unlimited,
+				BOOL_OR(ce.usage_allowed) AS usage_allowed
+			FROM entity_level_cus_ents ce
+			JOIN entitlements ent ON ce.entitlement_id = ent.id
+			GROUP BY ce.internal_feature_id, ce.internal_customer_id
 		),

 		entity_aggregated_cus_entitlements AS (
@@ -282,7 +221,6 @@ export const getEntityAggregateFragments = ({
 				COALESCE(erf.rollover_usage, 0) AS rollover_usage,
 				eat.unlimited,
 				eat.usage_allowed,
-				eat.entity_count,
 				eam.entities
 			FROM entity_aggregate_totals eat
 			LEFT JOIN entity_aggregate_map eam
--- a/server/tests/integration/db/full-subject/full-subject-aggregate-options.test.ts
+++ b/server/tests/integration/db/full-subject/full-subject-aggregate-options.test.ts
@@ -236,7 +236,6 @@ describe(`${chalk.yellowBright("fullSubject aggregate options")}`, () => {
 				expect(aggregateMessagesBalance?.allowance_total).toBe(600);
 				expect(aggregateMessagesBalance?.prepaid_grant_from_options).toBe(1000);
 				expect(aggregateMessagesBalance?.balance).toBe(1600);
-				expect(aggregateMessagesBalance?.entity_count).toBe(2);
 			},
 		});
 	});
--- a/server/tests/integration/db/full-subject/full-subject-aggregates.test.ts
+++ b/server/tests/integration/db/full-subject/full-subject-aggregates.test.ts
@@ -106,7 +106,6 @@ describe(`${chalk.yellowBright("fullSubject aggregates")}`, () => {
 				expect(aggregate.balance).toBe(30);
 				expect(aggregate.adjustment).toBe(0);
 				expect(aggregate.additional_balance).toBe(0);
-				expect(aggregate.entity_count).toBe(2);
 				expect(JSON.stringify(aggregate.feature.id)).toBe(
 					JSON.stringify(scenario.customerEntitlements[1]!.feature_id),
 				);
@@ -157,7 +156,6 @@ describe(`${chalk.yellowBright("fullSubject aggregates")}`, () => {
 				const aggregate = comparable.aggregated_customer_entitlements[0]!;
 				expect(aggregate.allowance_total).toBe(200);
 				expect(aggregate.balance).toBe(30);
-				expect(aggregate.entity_count).toBe(2);
 				expect(aggregate.entities).toMatchObject({
 					[scenario.ids.internalEntityIds[0]!]: {
 						id: scenario.ids.internalEntityIds[0],
--- a/server/tests/integration/db/full-subject/utils/buildComparableFullSubject.ts
+++ b/server/tests/integration/db/full-subject/utils/buildComparableFullSubject.ts
@@ -148,7 +148,6 @@ const comparableAggregatedEntitlement = (
 	additional_balance: customerEntitlement.additional_balance,
 	unlimited: customerEntitlement.unlimited,
 	usage_allowed: customerEntitlement.usage_allowed,
-	entity_count: customerEntitlement.entity_count,
 	entities: customerEntitlement.entities,
 	feature: {
 		id: customerEntitlement.feature.id,
--- a/server/tests/unit/full-subject-cache/full-subject-aggregate-balance.test.ts
+++ b/server/tests/unit/full-subject-cache/full-subject-aggregate-balance.test.ts
@@ -22,7 +22,6 @@ describe("fullSubject aggregate balance", () => {
 			rollover_usage: 0,
 			unlimited: false,
 			usage_allowed: false,
-			entity_count: 2,
 			entities: {
 				ent1: {
 					id: "ent1",
--- a/shared/models/cusProductModels/cusEntModels/aggregatedCusEnt.ts
+++ b/shared/models/cusProductModels/cusEntModels/aggregatedCusEnt.ts
@@ -24,7 +24,6 @@ export const AggregatedFeatureBalanceSchema = z.object({
 	rollover_usage: z.number().default(0),
 	unlimited: z.boolean(),
 	usage_allowed: z.boolean(),
-	entity_count: z.number(),
 	entities: z.record(z.string(), AggregatedEntityBalanceSchema).nullish(),
 });

--- a/shared/utils/cusUtils/fullCusUtils/logFullCustomer.ts
+++ b/shared/utils/cusUtils/fullCusUtils/logFullCustomer.ts
@@ -42,7 +42,6 @@ export const logFullCustomer = ({
 		balance: ae.balance,
 		adjustment: ae.adjustment,
 		unlimited: ae.unlimited,
-		entity_count: ae.entity_count,
 		entitlement:
 			"entitlement" in ae
 				? (ae as Record<string, unknown>).entitlement
--- a/shared/utils/fullSubjectUtils/logFullSubject.ts
+++ b/shared/utils/fullSubjectUtils/logFullSubject.ts
@@ -62,7 +62,6 @@ export const logFullSubject = ({
 				feature_id: ae.feature_id,
 				balance: ae.balance,
 				unlimited: ae.unlimited,
-				entity_count: ae.entity_count,
 			})) ?? "N/A",
 	};