feat: 🎸 max entities edge config

server/src/internal/customers/CusService.ts

@@ -34,7 +34,10 @@ import { executeWithHealthTracking } from "@/db/pgHealthMonitor.js";
 import type { RepoContext } from "@/db/repoContext.js";
 import type { AutumnContext } from "@/honoUtils/HonoEnv.js";
 import { withSpan } from "../analytics/tracer/spanUtils.js";
-import { getOrgCusProductLimit } from "../misc/edgeConfig/orgLimitsStore.js";
+import {
+	getOrgCusProductLimit,
+	getOrgEntitiesLimit,
+} from "../misc/edgeConfig/orgLimitsStore.js";
 import { resetCustomerEntitlements } from "./actions/resetCustomerEntitlements/resetCustomerEntitlements.js";
 import {
 	ACTIVE_STATUSES,
@@ -90,6 +93,10 @@ export class CusService {
 					orgId,
 					orgSlug: org.slug,
 				});
+				const entitiesLimit = getOrgEntitiesLimit({
+					orgId,
+					orgSlug: org.slug,
+				});
 
 				const query = getFullCusQuery({
 					idOrInternalId,
@@ -103,6 +110,7 @@ export class CusService {
 					withEvents,
 					entityId,
 					cusProductLimit,
+					entitiesLimit,
 				});
 
 				if (explain) {

server/src/internal/customers/getFullCusQuery.ts

@@ -98,7 +98,13 @@ const buildOptimizedCusProductsCTE = ({
   `;
 };
 
-const buildEntitiesCTE = (withEntities: boolean) => {
+const buildEntitiesCTE = ({
+	withEntities,
+	entitiesLimit,
+}: {
+	withEntities: boolean;
+	entitiesLimit: number;
+}) => {
 	if (!withEntities) {
 		return sql``;
 	}
@@ -114,7 +120,7 @@ const buildEntitiesCTE = (withEntities: boolean) => {
         SELECT * FROM entities e
         WHERE e.internal_customer_id = (SELECT internal_id FROM customer_record)
         ORDER BY e.internal_id DESC
-        LIMIT 300
+        LIMIT ${entitiesLimit}
       ) e
     )
   `;
@@ -284,6 +290,7 @@ export const getFullCusQuery = ({
 	withEvents,
 	entityId,
 	cusProductLimit,
+	entitiesLimit = 300,
 }: {
 	idOrInternalId: string;
 	orgId: string;
@@ -296,6 +303,7 @@ export const getFullCusQuery = ({
 	withEvents: boolean;
 	entityId?: string;
 	cusProductLimit: number;
+	entitiesLimit?: number;
 }) => {
 	const sqlChunks: SQL[] = [];
 
@@ -316,7 +324,7 @@ export const getFullCusQuery = ({
 	// Step 2: Get entities
 	if (withEntities) {
 		sqlChunks.push(sql`, `);
-		sqlChunks.push(buildEntitiesCTE(withEntities));
+		sqlChunks.push(buildEntitiesCTE({ withEntities, entitiesLimit }));
 	}
 
 	// Step 3: Get entity

server/src/internal/misc/edgeConfig/orgLimitsSchemas.ts

@@ -6,6 +6,7 @@ export const OrgLimitsConfigSchema = z.object({
 			z.string(),
 			z.object({
 				maxCusProducts: z.number().min(1).optional(),
+				maxEntities: z.number().min(1).optional(),
 			}),
 		)
 		.default({}),

server/src/internal/misc/edgeConfig/orgLimitsStore.ts

@@ -7,6 +7,7 @@ import {
 } from "./orgLimitsSchemas.js";
 
 export const DEFAULT_CUS_PRODUCT_LIMIT = 15;
+export const DEFAULT_ENTITIES_LIMIT = 300;
 
 const store = createEdgeConfigStore<OrgLimitsConfig>({
 	s3Key: ADMIN_ORG_LIMITS_CONFIG_KEY,
@@ -28,11 +29,23 @@ export const getOrgCusProductLimit = ({
 }): number => {
 	const orgs = store.get().orgs;
 	const orgConfig =
-		(orgId ? orgs[orgId] : undefined) ??
-		(orgSlug ? orgs[orgSlug] : undefined);
+		(orgId ? orgs[orgId] : undefined) ?? (orgSlug ? orgs[orgSlug] : undefined);
 	return orgConfig?.maxCusProducts ?? DEFAULT_CUS_PRODUCT_LIMIT;
 };
 
+export const getOrgEntitiesLimit = ({
+	orgId,
+	orgSlug,
+}: {
+	orgId?: string;
+	orgSlug?: string;
+}): number => {
+	const orgs = store.get().orgs;
+	const orgConfig =
+		(orgId ? orgs[orgId] : undefined) ?? (orgSlug ? orgs[orgSlug] : undefined);
+	return orgConfig?.maxEntities ?? DEFAULT_ENTITIES_LIMIT;
+};
+
 export const getOrgLimitsConfigFromSource = async () => {
 	return await store.readFromSource();
 };
@@ -63,3 +76,15 @@ export const updateFullOrgLimitsConfig = async ({
 }) => {
 	await store.writeToSource({ config });
 };
+
+/**
+ * Test-only helper: override the in-memory org limits config without touching S3.
+ * Use inside tests to simulate admin-configured org limits deterministically.
+ */
+export const _setOrgLimitsConfigForTesting = ({
+	config,
+}: {
+	config: OrgLimitsConfig;
+}) => {
+	store._setRuntimeConfigForTesting(config);
+};

server/tests/integration/crud/entities/get-entities-max-entities-limit.test.ts

@@ -0,0 +1,145 @@
+import { afterAll, expect, test } from "bun:test";
+import { TestFeature } from "@tests/setup/v2Features.js";
+import { initScenario, s } from "@tests/utils/testInitUtils/initScenario.js";
+import chalk from "chalk";
+import { CusService } from "@/internal/customers/CusService.js";
+import type { OrgLimitsConfig } from "@/internal/misc/edgeConfig/orgLimitsSchemas.js";
+import {
+	_setOrgLimitsConfigForTesting,
+	DEFAULT_ENTITIES_LIMIT,
+	getOrgEntitiesLimit,
+} from "@/internal/misc/edgeConfig/orgLimitsStore.js";
+
+/**
+ * These tests mutate the in-memory org limits store (a module-level singleton)
+ * so they must run serially, not concurrently, to avoid leaking state into
+ * unrelated tests.
+ *
+ * They call CusService.getFull directly (in-process) rather than going through
+ * the HTTP server, because the running server has its own copy of the
+ * orgLimitsStore that only refreshes from S3 every 30s. Calling in-process
+ * exercises the same SQL path (getFullCusQuery -> buildEntitiesCTE) while
+ * picking up the in-memory override immediately.
+ */
+
+const resetOrgLimits = () => {
+	_setOrgLimitsConfigForTesting({ config: { orgs: {} } });
+};
+
+afterAll(() => {
+	resetOrgLimits();
+});
+
+test(`${chalk.yellowBright("maxEntities: caps entities returned by CusService.getFull to configured limit")}`, async () => {
+	const customerId = "max-entities-cap";
+	const entityCount = 5;
+	const maxEntities = 3;
+
+	const { ctx, entities } = await initScenario({
+		customerId,
+		setup: [
+			s.customer({}),
+			s.entities({ count: entityCount, featureId: TestFeature.Users }),
+		],
+		actions: [],
+	});
+
+	expect(entities.length).toBe(entityCount);
+
+	try {
+		const config: OrgLimitsConfig = {
+			orgs: { [ctx.org.id]: { maxEntities } },
+		};
+		_setOrgLimitsConfigForTesting({ config });
+
+		// Verify the accessor resolves to the override
+		expect(getOrgEntitiesLimit({ orgId: ctx.org.id })).toBe(maxEntities);
+
+		const fullCus = await CusService.getFull({
+			ctx,
+			idOrInternalId: customerId,
+			withEntities: true,
+		});
+
+		expect(fullCus.entities).toBeDefined();
+		expect(fullCus.entities.length).toBe(maxEntities);
+	} finally {
+		resetOrgLimits();
+	}
+});
+
+test(`${chalk.yellowBright("maxEntities: returns all entities when override is absent (falls back to default)")}`, async () => {
+	const customerId = "max-entities-default";
+	const entityCount = 4;
+
+	const { ctx, entities } = await initScenario({
+		customerId,
+		setup: [
+			s.customer({}),
+			s.entities({ count: entityCount, featureId: TestFeature.Users }),
+		],
+		actions: [],
+	});
+
+	expect(entities.length).toBe(entityCount);
+
+	resetOrgLimits();
+
+	// No override for this org -> default (300) applies
+	expect(getOrgEntitiesLimit({ orgId: ctx.org.id })).toBe(
+		DEFAULT_ENTITIES_LIMIT,
+	);
+
+	const fullCus = await CusService.getFull({
+		ctx,
+		idOrInternalId: customerId,
+		withEntities: true,
+	});
+
+	expect(fullCus.entities).toBeDefined();
+	expect(fullCus.entities.length).toBe(entityCount);
+});
+
+test(`${chalk.yellowBright("maxEntities: raising the limit exposes previously-hidden entities")}`, async () => {
+	const customerId = "max-entities-raise";
+	const entityCount = 4;
+
+	const { ctx, entities } = await initScenario({
+		customerId,
+		setup: [
+			s.customer({}),
+			s.entities({ count: entityCount, featureId: TestFeature.Users }),
+		],
+		actions: [],
+	});
+
+	expect(entities.length).toBe(entityCount);
+
+	try {
+		// Start with a tight cap
+		_setOrgLimitsConfigForTesting({
+			config: { orgs: { [ctx.org.id]: { maxEntities: 2 } } },
+		});
+
+		const capped = await CusService.getFull({
+			ctx,
+			idOrInternalId: customerId,
+			withEntities: true,
+		});
+		expect(capped.entities.length).toBe(2);
+
+		// Raise the cap above the entity count
+		_setOrgLimitsConfigForTesting({
+			config: { orgs: { [ctx.org.id]: { maxEntities: 50 } } },
+		});
+
+		const uncapped = await CusService.getFull({
+			ctx,
+			idOrInternalId: customerId,
+			withEntities: true,
+		});
+		expect(uncapped.entities.length).toBe(entityCount);
+	} finally {
+		resetOrgLimits();
+	}
+});