feat: added spend limits to entity billing controls

bun.lock

@@ -101,7 +101,7 @@
     },
     "packages/autumn-js": {
       "name": "autumn-js",
-      "version": "1.0.0-beta.6",
+      "version": "1.0.0-beta.8",
       "dependencies": {
         "query-string": "^9.2.2",
         "rou3": "^0.6.1",

server/src/_luaScriptsV2/customers/updateEntityInCustomer.lua

@@ -0,0 +1,73 @@
+--[[
+  Lua Script: Update Entity Fields in Cached FullCustomer
+
+  Atomically updates specific fields on an entity inside the cached
+  FullCustomer JSON object. Matches by entity id or internal_id.
+
+  KEYS[1] = FullCustomer cache key
+
+  ARGV[1] = JSON:
+    {
+      id_or_internal_id: string,
+      updates: { field: value, ... }
+    }
+
+  Returns JSON:
+    { "ok": true, "updated_count": number }
+    { "ok": false, "error": string }
+]]
+
+local cache_key = KEYS[1]
+local params = cjson.decode(ARGV[1])
+
+local id_or_internal_id = params.id_or_internal_id
+local updates = params.updates
+
+if not id_or_internal_id then
+  return cjson.encode({ ok = false, error = "missing_id_or_internal_id" })
+end
+
+if not updates then
+  return cjson.encode({ ok = false, error = "missing_updates" })
+end
+
+local raw = redis.call('JSON.GET', cache_key, '.')
+if not raw then
+  return cjson.encode({ ok = false, error = "cache_miss" })
+end
+
+local full_customer = cjson.decode(raw)
+
+if full_customer.entities == nil then
+  return cjson.encode({ ok = false, error = "no_entities" })
+end
+
+if #full_customer.entities == 0 then
+  return cjson.encode({ ok = false, error = "empty_entities" })
+end
+
+local entity_idx = nil
+for idx, entity in ipairs(full_customer.entities) do
+  if entity.id == id_or_internal_id or entity.internal_id == id_or_internal_id then
+    entity_idx = idx - 1
+    break
+  end
+end
+
+if entity_idx == nil then
+  return cjson.encode({ ok = false, error = "entity_not_found" })
+end
+
+local base_path = '$.entities[' .. entity_idx .. '].'
+local updated_count = 0
+
+for field, value in pairs(updates) do
+  if value == cjson.null then
+    redis.call('JSON.SET', cache_key, base_path .. field, 'null')
+  else
+    redis.call('JSON.SET', cache_key, base_path .. field, cjson.encode(value))
+  end
+  updated_count = updated_count + 1
+end
+
+return cjson.encode({ ok = true, updated_count = updated_count })

server/src/_luaScriptsV2/luaScriptsV2.ts

@@ -214,6 +214,14 @@ export const APPEND_ENTITY_TO_CUSTOMER_SCRIPT = readFileSync(
 	"utf-8",
 );
 
+/**
+ * Atomically update specific fields on an entity inside the cached FullCustomer.
+ */
+export const UPDATE_ENTITY_IN_CUSTOMER_SCRIPT = readFileSync(
+	join(CUSTOMER_DIR, "updateEntityInCustomer.lua"),
+	"utf-8",
+);
+
 /**
  * Atomically upsert an invoice in the customer's invoices array.
  * Matches by stripe_id — replaces if found, appends if not.

server/src/external/autumn/autumnCli.ts

@@ -8,6 +8,7 @@ import {
 	type ApiCusFeatureV3,
 	type ApiCusProductV3,
 	type ApiCustomerV3,
+	type ApiEntityBillingControlsInput,
 	type ApiEntityV0,
 	type AttachBodyV0,
 	type AttachParamsV0Input,
@@ -593,6 +594,21 @@ export class AutumnInt {
 			);
 			return data;
 		},
+
+		update: async (
+			customerId: string,
+			entityId: string,
+			updates: {
+				billing_controls?: ApiEntityBillingControlsInput;
+			},
+		) => {
+			const data = await this.post(`/entities.update`, {
+				customer_id: customerId,
+				entity_id: entityId,
+				...updates,
+			});
+			return data;
+		},
 	};
 
 	products = {

server/src/external/autumn/autumnCliV2.ts

@@ -4,6 +4,8 @@ import dotenv from "dotenv";
 dotenv.config();
 
 import {
+	type ApiCustomerV3,
+	type ApiEntityBillingControlsInput,
 	type AttachBodyV0,
 	type CancelBody,
 	type CheckoutParams,
@@ -13,7 +15,6 @@ import {
 	type CheckResponseV1,
 	type CreateEntityParams,
 	type CreateRewardProgram,
-	type ApiCustomerV3,
 	CustomerExpand,
 	EntityExpand,
 	ErrCode,
@@ -143,6 +144,33 @@ export class AutumnCliV2 {
 		return response.json();
 	}
 
+	async patch(path: string, body: any) {
+		const response = await fetch(`${this.baseUrl}${path}`, {
+			method: "PATCH",
+			headers: this.headers,
+			body: JSON.stringify(body),
+		});
+
+		if (response.status !== 200) {
+			let error: any;
+			try {
+				error = await response.json();
+			} catch (_e) {
+				throw new AutumnError({
+					message: `PATCH ${path} failed with status ${response.status}`,
+					code: ErrCode.InternalError,
+				});
+			}
+
+			throw new AutumnError({
+				message: error.message || `PATCH ${path} failed`,
+				code: error.code || ErrCode.InternalError,
+			});
+		}
+
+		return response.json();
+	}
+
 	async delete(
 		path: string,
 		{
@@ -335,6 +363,20 @@ export class AutumnCliV2 {
 		delete: async (customerId: string, entityId: string) => {
 			return await this.delete(`/customers/${customerId}/entities/${entityId}`);
 		},
+
+		update: async (
+			customerId: string,
+			entityId: string,
+			updates: {
+				billing_controls?: ApiEntityBillingControlsInput;
+			},
+		) => {
+			return await this.post(`/entities.update`, {
+				customer_id: customerId,
+				entity_id: entityId,
+				...updates,
+			});
+		},
 	};
 
 	products = {

server/src/external/redis/initRedis.ts

@@ -25,6 +25,7 @@ import {
 	UPDATE_CUSTOMER_DATA_SCRIPT,
 	UPDATE_CUSTOMER_ENTITLEMENTS_SCRIPT,
 	UPDATE_CUSTOMER_PRODUCT_SCRIPT,
+	UPDATE_ENTITY_IN_CUSTOMER_SCRIPT,
 	UPSERT_INVOICE_IN_CUSTOMER_SCRIPT,
 } from "../../_luaScriptsV2/luaScriptsV2.js";
 
@@ -207,6 +208,11 @@ const configureRedisInstance = (redisInstance: Redis): Redis => {
 		lua: APPEND_ENTITY_TO_CUSTOMER_SCRIPT,
 	});
 
+	redisInstance.defineCommand("updateEntityInCustomer", {
+		numberOfKeys: 1,
+		lua: UPDATE_ENTITY_IN_CUSTOMER_SCRIPT,
+	});
+
 	redisInstance.defineCommand("upsertInvoiceInCustomer", {
 		numberOfKeys: 1,
 		lua: UPSERT_INVOICE_IN_CUSTOMER_SCRIPT,
@@ -418,6 +424,10 @@ declare module "ioredis" {
 			cacheKey: string,
 			entityJson: string,
 		): Promise<string>;
+		updateEntityInCustomer(
+			cacheKey: string,
+			paramsJson: string,
+		): Promise<string>;
 		upsertInvoiceInCustomer(
 			cacheKey: string,
 			invoiceJson: string,

server/src/internal/api/entities/EntityService.ts

@@ -70,15 +70,12 @@ export class EntityService {
 				),
 		});
 	}
-	static async insert({ db, data }: { db: DrizzleCli; data: any }) {
+	static async insert({ db, data }: { db: DrizzleCli; data: Entity[] }) {
 		if (data.length === 0) {
 			return [];
 		}
 
-		const results = await db
-			.insert(entities)
-			.values(data as any)
-			.returning();
+		const results = await db.insert(entities).values(data).returning();
 
 		return results as Entity[];
 	}
@@ -135,7 +132,7 @@ export class EntityService {
 	}: {
 		db: DrizzleCli;
 		internalId: string;
-		update: any;
+		update: Partial<Entity>;
 	}) {
 		const results = await db
 			.update(entities)

server/src/internal/balances/track/utils/runRedisTrack.ts

@@ -90,7 +90,7 @@ export const runRedisTrack = async ({
 		executeRedisDeduction({
 			ctx,
 			fullCustomer,
-			entityId: fullCustomer.entity?.id,
+			entityId: fullCustomer.entity?.id ?? undefined,
 			deductions: featureDeductions,
 			deductionOptions: {
 				overageBehaviour: overageBehavior || "cap",

server/src/internal/balances/updateBalance/runRedisUpdateBalanceV2.ts

@@ -28,7 +28,7 @@ export const runRedisUpdateBalanceV2 = async ({
 }) => {
 	const { org, env } = ctx;
 	const customerId = fullCustomer.id || fullCustomer.internal_id;
-	const entityId = fullCustomer.entity?.id;
+	const entityId = fullCustomer.entity?.id ?? undefined;
 
 	const deductionOptions: DeductionOptions = {
 		overageBehaviour: "allow", // Allow bypasses granted_balance cap for balance updates

server/src/internal/balances/updateBalance/updateGrantedBalance.ts

@@ -42,7 +42,7 @@ export const updateGrantedBalance = async ({
 
 	const currentAllowance = cusEntsToAllowance({
 		cusEnts,
-		entityId: fullCustomer.entity?.id,
+		entityId: fullCustomer.entity?.id ?? undefined,
 		withRollovers: false,
 	});
 

server/src/internal/balances/utils/allocatedInvoice/setupAllocatedInvoiceContext.ts

@@ -36,7 +36,7 @@ export const setupAllocatedInvoiceContext = async ({
 		ctx,
 		params: {
 			customer_id: oldFullCustomer.id ?? oldFullCustomer.internal_id,
-			entity_id: oldFullCustomer.entity?.id,
+			entity_id: oldFullCustomer.entity?.id ?? undefined,
 		},
 	});
 

server/src/internal/billing/v2/actions/legacy/legacyAttach.ts

@@ -54,7 +54,7 @@ export const legacyAttach = async ({
 
 	const params: AttachParamsV1 = {
 		customer_id: fullCustomer.id || fullCustomer.internal_id,
-		entity_id: fullCustomer.entity?.id,
+		entity_id: fullCustomer.entity?.id ?? undefined,
 		plan_id: fullProduct.id,
 
 		invoice_mode: attachParamsToInvoiceModeParams({ attachParams }),

server/src/internal/billing/v2/actions/legacy/renew.ts

@@ -65,7 +65,7 @@ export const renew = async ({
 
 	const params: UpdateSubscriptionV1Params = {
 		customer_id: fullCustomer.id || fullCustomer.internal_id,
-		entity_id: fullCustomer.entity?.id,
+		entity_id: fullCustomer.entity?.id ?? undefined,
 		plan_id: fullProduct.id,
 
 		invoice_mode: body.invoice

server/src/internal/billing/v2/actions/legacy/updateQuantity.ts

@@ -67,7 +67,7 @@ export const updateQuantity = async ({
 
 	const params: UpdateSubscriptionV1Params = {
 		customer_id: fullCustomer.id || fullCustomer.internal_id,
-		entity_id: fullCustomer.entity?.id,
+		entity_id: fullCustomer.entity?.id ?? undefined,
 		plan_id: fullProduct.id,
 
 		invoice_mode: body.invoice

server/src/internal/billing/v2/actions/migrate/migrate.ts

@@ -47,7 +47,7 @@ export async function migrate({
 	const updateSubscriptionParams: UpdateSubscriptionV1Params = {
 		customer_id: fullCustomer.id || fullCustomer.internal_id,
 		customer_product_id: currentCustomerProduct.id,
-		entity_id: entity?.id,
+		entity_id: entity?.id ?? undefined,
 
 		proration_behavior: "none",
 		version: newProduct.version, // to trigger update custom plan intent

server/src/internal/billing/v2/utils/billingResult/billingResultToResponse.ts

@@ -31,7 +31,7 @@ export const billingResultToResponse = ({
 
 	return {
 		customer_id: customerId,
-		entity_id: fullCustomer.entity?.id,
+		entity_id: fullCustomer.entity?.id ?? undefined,
 		invoice: stripeInvoice
 			? {
 					status: stripeInvoice.status,

server/src/internal/billing/v2/utils/initFullCustomerProduct/applyExisting/applyExistingStatesToCustomerProduct.ts

@@ -32,7 +32,7 @@ export const applyExistingStatesToCustomerProduct = ({
 
 		existingUsages = cusProductToExistingUsages({
 			cusProduct: fromCustomerProduct,
-			entityId: fullCustomer.entity?.id,
+			entityId: fullCustomer.entity?.id ?? undefined,
 			...existingUsagesConfig,
 		});
 	}

server/src/internal/billing/v2/utils/initFullCustomerProduct/initCustomerEntitlement/initCustomerEntitlementEntities.ts

@@ -26,6 +26,7 @@ export const initCustomerEntitlementEntities = ({
 		});
 
 		if (!featureMatches) continue;
+		if (!entity.id) continue;
 
 		entities[entity.id] = {
 			id: entity.id,

server/src/internal/checkouts/handlers/handleGetCheckout.ts

@@ -78,7 +78,7 @@ export const handleGetCheckout = createRoute({
 			},
 			entity: fullCustomer.entity
 				? {
-						id: fullCustomer.entity.id,
+						id: fullCustomer.entity.id ?? fullCustomer.entity.internal_id,
 						name: fullCustomer.entity.name || null,
 					}
 				: null,

server/src/internal/checkouts/handlers/handlePreviewCheckout.ts

@@ -97,7 +97,7 @@ export const handlePreviewCheckout = createRoute({
 			},
 			entity: fullCustomer.entity
 				? {
-						id: fullCustomer.entity.id,
+						id: fullCustomer.entity.id ?? fullCustomer.entity.internal_id,
 						name: fullCustomer.entity.name || null,
 					}
 				: null,

server/src/internal/customers/add-product/initCusEnt.ts

@@ -41,7 +41,7 @@ const initCusEntEntities = ({
 
 	for (const entity of entities) {
 		if (!entitlementHasEntityFeature({ entitlement, entity })) continue;
-
+		if (!entity.id) continue;
 		if (existingCusEnt?.entities?.[entity.id]) {
 			continue;
 		}

server/src/internal/customers/cusUtils/apiCusUtils/getApiBalance/getApiBalance.ts

@@ -46,7 +46,7 @@ const getApiBalanceBreakdownItem = ({
 	fullCus: FullCustomer;
 	customerEntitlement: FullCusEntWithFullCusProduct;
 }): ApiBalanceBreakdownV1 => {
-	const entityId = fullCus.entity?.id;
+	const entityId = fullCus.entity?.id ?? fullCus.entity?.internal_id;
 
 	const planId = cusEntsToPlanId({ cusEnts: [customerEntitlement] });
 
@@ -123,7 +123,7 @@ export const getApiBalance = ({
 	cusEnts: FullCusEntWithFullCusProduct[];
 	feature: Feature;
 }): { data: ApiBalanceV1 } => {
-	const entityId = fullCus.entity?.id;
+	const entityId = fullCus.entity?.id ?? fullCus.entity?.internal_id;
 
 	const apiFeature = expandIncludes({
 		expand: ctx.expand,

server/src/internal/customers/cusUtils/fullCustomerCacheUtils/updateEntityInCache.ts

@@ -0,0 +1,105 @@
+import type { Entity } from "@autumn/shared";
+import { redis } from "@/external/redis/initRedis.js";
+import type { AutumnContext } from "@/honoUtils/HonoEnv.js";
+import { tryRedisWrite } from "@/utils/cacheUtils/cacheUtils.js";
+import { deleteCachedFullCustomer } from "./deleteCachedFullCustomer.js";
+import { buildFullCustomerCacheKey } from "./fullCustomerCacheConfig.js";
+
+type UpdateEntityInCacheResult = {
+	ok: boolean;
+	updatedCount?: number;
+	error?: string;
+};
+
+const shouldDeleteCache = ({ error }: { error?: string }) => {
+	return (
+		error === "no_entities" ||
+		error === "empty_entities" ||
+		error === "entity_not_found"
+	);
+};
+
+export const updateEntityInCache = async ({
+	ctx,
+	customerId,
+	idOrInternalId,
+	updates,
+}: {
+	ctx: AutumnContext;
+	customerId: string;
+	idOrInternalId: string;
+	updates: Partial<Pick<Entity, "spend_limits">>;
+}): Promise<UpdateEntityInCacheResult | null> => {
+	try {
+		if (Object.keys(updates).length === 0) {
+			return { ok: true, updatedCount: 0 };
+		}
+
+		const { org, env, logger } = ctx;
+
+		const cacheKey = buildFullCustomerCacheKey({
+			orgId: org.id,
+			env,
+			customerId,
+		});
+
+		const result = await tryRedisWrite(async () => {
+			return await redis.updateEntityInCustomer(
+				cacheKey,
+				JSON.stringify({
+					id_or_internal_id: idOrInternalId,
+					updates,
+				}),
+			);
+		});
+
+		if (result === null) {
+			logger.warn(
+				`[updateEntityInCache] Redis write failed for entity ${idOrInternalId}`,
+			);
+			return null;
+		}
+
+		const parsed = JSON.parse(result) as {
+			ok: boolean;
+			updated_count?: number;
+			error?: string;
+		};
+
+		if (parsed.ok) {
+			return {
+				ok: true,
+				updatedCount: parsed.updated_count,
+			};
+		}
+
+		if (parsed.error === "cache_miss") {
+			return {
+				ok: false,
+				error: parsed.error,
+			};
+		}
+
+		if (shouldDeleteCache({ error: parsed.error })) {
+			await deleteCachedFullCustomer({
+				ctx,
+				customerId,
+				source: "updateEntityInCache",
+			});
+		}
+
+		logger.warn(
+			`[updateEntityInCache] entity ${idOrInternalId}: ${parsed.error ?? "unknown_error"}`,
+		);
+
+		return {
+			ok: false,
+			error: parsed.error,
+		};
+	} catch (error) {
+		ctx.logger.error(
+			`[updateEntityInCache] entity ${idOrInternalId}: error, ${error}`,
+		);
+		return null;
+	}
+};

server/src/internal/customers/handlers/handleTransferProductV2.ts

@@ -95,7 +95,7 @@ export const handleTransferProductV2 = createRoute({
 		if (toCusProduct) {
 			throw new CusProductAlreadyExistsError({
 				productId: toCusProduct.product?.id,
-				entityId: toEntity?.id,
+				entityId: toEntity?.id ?? toEntity?.internal_id,
 				customerId: from_entity_id && !to_entity_id ? customer_id : undefined,
 			});
 		}

server/src/internal/entities/actions/batchCreateEntities.ts

@@ -73,6 +73,9 @@ export const batchCreateEntities = async ({
 			update: {
 				id: inputEntities[0].id,
 				name: inputEntities[0].name,
+				...(inputEntities[0].billing_controls && {
+					spend_limits: inputEntities[0].billing_controls.spend_limits,
+				}),
 			},
 		});
 
@@ -103,7 +106,7 @@ export const batchCreateEntities = async ({
 		const apiEntity = await getApiEntity({
 			ctx,
 			customerId,
-			entityId: entity.id,
+			entityId: entity.id ?? entity.internal_id,
 			fullCus: clonedFullCus,
 			withAutumnId,
 		});

server/src/internal/entities/actions/deleteEntity.ts

@@ -94,6 +94,8 @@ export const deleteEntity = async ({
 			let newEntities: {
 				[key: string]: EntityBalance;
 			};
+
+			if (entity.id === null) continue;
 			if (replaceable) {
 				const { newEntities: newEntities_ } = replaceEntityInCusEnt({
 					cusEnt: linkedCusEnt,

server/src/internal/entities/actions/index.ts

@@ -1,7 +1,9 @@
 import { batchCreateEntities } from "./batchCreateEntities";
 import { deleteEntity } from "./deleteEntity";
+import { updateEntity } from "./updateEntity";
 
 export const entityActions = {
 	batchCreate: batchCreateEntities,
 	delete: deleteEntity,
+	update: updateEntity,
 } as const;

server/src/internal/entities/actions/updateEntity.ts

@@ -0,0 +1,51 @@
+import {
+	CustomerExpand,
+	CustomerNotFoundError,
+	EntityNotFoundError,
+	type UpdateEntityParams,
+} from "@autumn/shared";
+import type { AutumnContext } from "@/honoUtils/HonoEnv.js";
+import { CusService } from "@/internal/customers/CusService.js";
+import { updateEntityDbAndCache } from "./updateEntityDbAndCache.js";
+
+export const updateEntity = async ({
+	ctx,
+	params,
+}: {
+	ctx: AutumnContext;
+	params: UpdateEntityParams;
+}) => {
+	const {
+		customer_id: customerId,
+		entity_id: entityId,
+		billing_controls,
+	} = params;
+	if (!customerId) {
+		throw new CustomerNotFoundError({ customerId: "" });
+	}
+
+	const fullCustomer = await CusService.getFull({
+		ctx,
+		idOrInternalId: customerId,
+		// withEntities: true,
+		entityId: entityId,
+		expand: [CustomerExpand.Invoices],
+	});
+
+	const entity = fullCustomer.entity;
+
+	if (!entity) {
+		throw new EntityNotFoundError({ entityId });
+	}
+
+	await updateEntityDbAndCache({
+		ctx,
+		customerId,
+		entity,
+		updates: {
+			spend_limits: billing_controls?.spend_limits,
+		},
+	});
+
+	return entity.id ?? entity.internal_id;
+};

server/src/internal/entities/actions/updateEntityDbAndCache.ts

@@ -0,0 +1,39 @@
+import type { Entity } from "@autumn/shared";
+import type { AutumnContext } from "@/honoUtils/HonoEnv.js";
+import { EntityService } from "@/internal/api/entities/EntityService.js";
+import { updateEntityInCache } from "@/internal/customers/cusUtils/fullCustomerCacheUtils/updateEntityInCache.js";
+
+export const updateEntityDbAndCache = async ({
+	ctx,
+	customerId,
+	entity,
+	updates,
+}: {
+	ctx: AutumnContext;
+	customerId: string;
+	entity: Entity;
+	updates: Partial<Pick<Entity, "spend_limits">>;
+}) => {
+	const filteredUpdates = Object.fromEntries(
+		Object.entries(updates).filter(([, value]) => value !== undefined),
+	) as Partial<Pick<Entity, "spend_limits">>;
+
+	if (Object.keys(filteredUpdates).length === 0) {
+		return entity;
+	}
+
+	const updatedEntity = await EntityService.update({
+		db: ctx.db,
+		internalId: entity.internal_id,
+		update: filteredUpdates,
+	});
+
+	await updateEntityInCache({
+		ctx,
+		customerId,
+		idOrInternalId: entity.id ?? entity.internal_id,
+		updates: filteredUpdates,
+	});
+
+	return updatedEntity;
+};

server/src/internal/entities/entityRouter.ts

@@ -7,6 +7,7 @@ import { handleDeleteEntityV2 } from "./handlers/handleDeleteEntity/handleDelete
 import { handleGetEntity } from "./handlers/handleGetEntity/handleGetEntity.js";
 import { handleGetEntityV2 } from "./handlers/handleGetEntity/handleGetEntityV2.js";
 import { handleListEntities } from "./handlers/handleListEntities.js";
+import { handleUpdateEntity } from "./handlers/handleUpdateEntity/handleUpdateEntity.js";
 
 export const entityRouter = new Hono<HonoEnv>();
 
@@ -30,3 +31,4 @@ export const entityRpcRouter = new Hono<HonoEnv>();
 entityRpcRouter.post("/entities.create", ...handleCreateEntityV2);
 entityRpcRouter.post("/entities.get", ...handleGetEntityV2);
 entityRpcRouter.post("/entities.delete", ...handleDeleteEntityV2);
+entityRpcRouter.post("/entities.update", ...handleUpdateEntity);

server/src/internal/entities/entityUtils/apiEntityUtils/getApiEntityBase.ts

@@ -71,6 +71,7 @@ export const getApiEntityBase = async ({
 		subscriptions: apiSubscriptions,
 		purchases: apiPurchases,
 		balances: apiBalances,
+		billing_controls: { spend_limits: entity.spend_limits ?? undefined },
 	} satisfies ApiEntityV2);
 
 	return {

server/src/internal/entities/entityUtils/entityUtils.ts

@@ -1,10 +1,10 @@
-import {
-	type AppEnv,
-	type Entity,
-	FeatureType,
-	FullCusEntWithFullCusProduct,
+import type {
+	AppEnv,
+	CreateEntityParams,
+	Entity,
+	Feature,
 } from "@autumn/shared";
-import { generateId, notNullish, nullish } from "@/utils/genUtils.js";
+import { generateId } from "@/utils/genUtils.js";
 
 export const constructEntity = ({
 	inputEntity,
@@ -14,8 +14,8 @@ export const constructEntity = ({
 	env,
 	deleted = false,
 }: {
-	inputEntity: any;
-	feature: any;
+	inputEntity: CreateEntityParams;
+	feature: Feature;
 	internalCustomerId: string;
 	orgId: string;
 	env: AppEnv;
@@ -24,7 +24,7 @@ export const constructEntity = ({
 	const entity: Entity = {
 		internal_id: generateId("ety"),
 		id: inputEntity.id,
-		name: inputEntity.name,
+		name: inputEntity.name ?? null,
 		internal_customer_id: internalCustomerId,
 		feature_id: feature.id,
 		internal_feature_id: feature.internal_id,
@@ -32,6 +32,7 @@ export const constructEntity = ({
 		env,
 		deleted,
 		created_at: Date.now(),
+		spend_limits: inputEntity.billing_controls?.spend_limits,
 	};
 
 	return entity;

server/src/internal/entities/handlers/handleCreateEntity/getInputEntities.ts

@@ -65,7 +65,9 @@ export const validateAndGetInputEntities = async ({
 
 	for (const entity of existingEntities) {
 		if (inputEntities.some((e: any) => e.id === entity.id) && !entity.deleted) {
-			throw new EntityAlreadyExistsError({ entityId: entity.id });
+			throw new EntityAlreadyExistsError({
+				entityId: entity.id ?? entity.internal_id,
+			});
 		}
 	}
 

server/src/internal/entities/handlers/handleCreateEntity/handleCreateEntityV2.ts

@@ -21,6 +21,7 @@ export const handleCreateEntityV2 = createRoute({
 					id: body.entity_id,
 					name: body.name,
 					feature_id: body.feature_id,
+					billing_controls: body.billing_controls,
 				},
 			],
 			customerData: customer_data,

server/src/internal/entities/handlers/handleDeleteEntity/handleDeleteEntity.ts

@@ -89,6 +89,8 @@ export const handleDeleteEntity = createRoute({
 				let newEntities: {
 					[key: string]: EntityBalance;
 				};
+				if (entity.id === null) continue;
+
 				if (replaceable) {
 					const { newEntities: newEntities_ } = replaceEntityInCusEnt({
 						cusEnt: linkedCusEnt,
@@ -104,34 +106,34 @@ export const handleDeleteEntity = createRoute({
 					newEntities = newEntities_;
 				}
 
-			await CusEntService.update({
-				ctx: {
-					db,
-					logger,
-					org,
-					env,
-					customerId: customer_id,
-				},
-				id: linkedCusEnt.id,
-				updates: {
-					entities: newEntities,
-				},
-			});
+				await CusEntService.update({
+					ctx: {
+						db,
+						logger,
+						org,
+						env,
+						customerId: customer_id,
+					},
+					id: linkedCusEnt.id,
+					updates: {
+						entities: newEntities,
+					},
+				});
 			}
 
-		if (!replaceable) {
-			await CusEntService.increment({
-				ctx: {
-					db,
-					logger,
-					org,
-					env,
-					customerId: customer_id,
-				},
-				id: mainCusEnt.id,
-				amount: 1,
-			});
-		}
+			if (!replaceable) {
+				await CusEntService.increment({
+					ctx: {
+						db,
+						logger,
+						org,
+						env,
+						customerId: customer_id,
+					},
+					id: mainCusEnt.id,
+					amount: 1,
+				});
+			}
 		}
 
 		// Cancel any subs

server/src/internal/entities/handlers/handleListEntities.ts

@@ -9,10 +9,14 @@ export const handleListEntities = createRoute({
 		const fullCus = await CusService.getFull({
 			ctx,
 			idOrInternalId: customer_id,
+			withEntities: true,
 		});
 
 		return c.json({
-			list: fullCus.entities,
+			list: fullCus.entities.map(({ spend_limits, ...entity }) => ({
+				...entity,
+				billing_controls: { spend_limits: spend_limits ?? undefined },
+			})),
 		});
 	},
 });

server/src/internal/entities/handlers/handleUpdateEntity/handleUpdateEntity.ts

@@ -0,0 +1,48 @@
+import {
+	AffectedResource,
+	CustomerNotFoundError,
+	UpdateEntityParamsSchema,
+} from "@autumn/shared";
+import { createRoute } from "@/honoMiddlewares/routeHandler.js";
+import { findCustomerForEntity } from "../../actions/findCustomer.js";
+import { entityActions } from "../../actions/index.js";
+import { getApiEntity } from "../../entityUtils/apiEntityUtils/getApiEntity.js";
+
+export const handleUpdateEntity = createRoute({
+	body: UpdateEntityParamsSchema,
+	resource: AffectedResource.Entity,
+	handler: async (c) => {
+		const ctx = c.get("ctx");
+		const body = c.req.valid("json");
+
+		let customerId = body.customer_id;
+		if (!customerId) {
+			const customer = await findCustomerForEntity({
+				ctx,
+				entityId: body.entity_id,
+			});
+
+			customerId = customer?.id ?? undefined;
+		}
+
+		if (!customerId) {
+			throw new CustomerNotFoundError({ customerId: body.customer_id ?? "" });
+		}
+
+		await entityActions.update({
+			ctx,
+			params: {
+				...body,
+				customer_id: customerId,
+			},
+		});
+
+		const apiEntity = await getApiEntity({
+			ctx,
+			customerId,
+			entityId: body.entity_id,
+		});
+
+		return c.json(apiEntity);
+	},
+});

server/src/internal/products/productUtils.ts

@@ -310,7 +310,7 @@ export const attachToInsertParams = (
 
 		if (entity) {
 			internalEntityId = entity.internal_id;
-			attachEntityId = entity.id;
+			attachEntityId = entity.id ?? entity.internal_id;
 		}
 	}
 

server/src/utils/importUtils/addProductFromSubs.ts

@@ -101,7 +101,7 @@ export const addProductFromSubs = async ({
 			cusProducts: cusProducts,
 			features: [],
 			internalEntityId: entity?.internal_id,
-			entityId: entity?.id,
+			entityId: entity?.id ?? entity?.internal_id,
 			isCustom: isCustom,
 		},
 

server/tests/integration/crud/entities/update-entity-billing-controls.test.ts

@@ -0,0 +1,116 @@
+import { expect, test } from "bun:test";
+import type { ApiEntityV2, EntityBillingControls } from "@autumn/shared";
+import { TestFeature } from "@tests/setup/v2Features.js";
+import { expectAutumnError } from "@tests/utils/expectUtils/expectErrUtils.js";
+import ctx from "@tests/utils/testInitUtils/createTestContext.js";
+import { initScenario, s } from "@tests/utils/testInitUtils/initScenario.js";
+import chalk from "chalk";
+import { CusService } from "@/internal/customers/CusService.js";
+
+const initialBillingControls: EntityBillingControls = {
+	spend_limits: [
+		{
+			feature_id: TestFeature.Messages,
+			enabled: true,
+			overage_limit: 25,
+		},
+	],
+};
+
+test.concurrent(`${chalk.yellowBright("entity billing controls: create and update entity spend limits")}`, async () => {
+	const { customerId, autumnV2_1 } = await initScenario({
+		customerId: "entity-billing-controls-1",
+		setup: [s.customer({})],
+		actions: [],
+	});
+
+	const created = await autumnV2_1.entities.create(customerId, {
+		id: "entity-1",
+		name: "Entity 1",
+		feature_id: TestFeature.Users,
+		billing_controls: initialBillingControls,
+	});
+
+	expect((created as ApiEntityV2).billing_controls?.spend_limits).toEqual(
+		initialBillingControls.spend_limits,
+	);
+
+	const fetched = await autumnV2_1.entities.get<ApiEntityV2>(
+		customerId,
+		"entity-1",
+	);
+
+	expect(fetched.billing_controls?.spend_limits).toEqual(
+		initialBillingControls.spend_limits,
+	);
+
+	const updatedBillingControls: EntityBillingControls = {
+		spend_limits: [
+			{
+				feature_id: TestFeature.Credits,
+				enabled: false,
+				overage_limit: 100,
+			},
+		],
+	};
+
+	await autumnV2_1.entities.update(customerId, "entity-1", {
+		billing_controls: updatedBillingControls,
+	});
+
+	const updated = await autumnV2_1.entities.get<ApiEntityV2>(
+		customerId,
+		"entity-1",
+	);
+
+	expect(updated.billing_controls?.spend_limits).toEqual(
+		updatedBillingControls.spend_limits,
+	);
+
+	const fromDb = await CusService.getFull({
+		ctx,
+		idOrInternalId: customerId,
+		withEntities: true,
+	});
+
+	const entity = fromDb.entities.find(
+		(candidate) => candidate.id === "entity-1",
+	);
+	expect(entity?.spend_limits).toEqual(updatedBillingControls.spend_limits);
+
+	await autumnV2_1.entities.update(customerId, "entity-1", {
+		billing_controls: { spend_limits: [] },
+	});
+
+	const cleared = await autumnV2_1.entities.get<ApiEntityV2>(
+		customerId,
+		"entity-1",
+	);
+
+	expect(cleared.billing_controls?.spend_limits).toEqual([]);
+});
+
+test.concurrent(`${chalk.yellowBright("entity billing controls: require feature_id when overage_limit is set")}`, async () => {
+	const { customerId, autumnV2_1 } = await initScenario({
+		customerId: "entity-billing-controls-2",
+		setup: [s.customer({})],
+		actions: [],
+	});
+
+	await expectAutumnError({
+		func: async () =>
+			await autumnV2_1.entities.create(customerId, {
+				id: "entity-2",
+				name: "Entity 2",
+				feature_id: TestFeature.Users,
+				billing_controls: {
+					spend_limits: [
+						// @ts-expect-error
+						{
+							overage_limit: 10,
+						},
+					],
+				},
+			}),
+	});
+});

shared/api/common/entityData.ts

@@ -1,4 +1,5 @@
 import { z } from "zod/v4";
+import { ApiEntityBillingControlsSchema } from "../entities/billingControls/entityBillingControls.js";
 
 export const EntityDataSchema = z
 	.object({
@@ -8,6 +9,9 @@ export const EntityDataSchema = z
 		name: z.string().optional().meta({
 			description: "Name of the entity",
 		}),
+		billing_controls: ApiEntityBillingControlsSchema.optional().meta({
+			description: "Billing controls for the entity.",
+		}),
 	})
 	.meta({
 		title: "EntityData",

shared/api/entities/apiEntityV2.ts

@@ -6,12 +6,16 @@ import {
 } from "../customers/cusPlans/apiSubscriptionV1.js";
 import { ApiInvoiceV1Schema } from "../others/apiInvoice/apiInvoiceV1.js";
 import { ApiBaseEntitySchema } from "./apiBaseEntity.js";
+import { ApiEntityBillingControlsSchema } from "./billingControls/entityBillingControls.js";
 
 // V2 base entity - uses V1 subscriptions (single array with status field)
 export const BaseApiEntityV2Schema = ApiBaseEntitySchema.extend({
 	subscriptions: z.array(ApiSubscriptionV1Schema),
 	purchases: z.array(ApiPurchaseV0Schema),
 	balances: z.record(z.string(), ApiBalanceV1Schema),
+	billing_controls: ApiEntityBillingControlsSchema.optional().meta({
+		description: "Billing controls for the entity.",
+	}),
 });
 
 export const ApiEntityExpandSchema = z.object({

shared/api/entities/billingControls/entityBillingControls.ts

@@ -0,0 +1,15 @@
+import { z } from "zod/v4";
+import { ApiEntitySpendLimitSchema } from "./entitySpendLimit.js";
+
+export const ApiEntityBillingControlsSchema = z.object({
+	spend_limits: z.array(ApiEntitySpendLimitSchema).optional().meta({
+		description: "List of overage spend limits per feature.",
+	}),
+});
+
+export type ApiEntityBillingControls = z.infer<
+	typeof ApiEntityBillingControlsSchema
+>;
+export type ApiEntityBillingControlsInput = z.input<
+	typeof ApiEntityBillingControlsSchema
+>;

shared/api/entities/billingControls/entitySpendLimit.ts

@@ -0,0 +1,6 @@
+import type { z } from "zod/v4";
+import { EntitySpendLimitSchema } from "../../../models/cusModels/billingControls/entitySpendLimit.js";
+
+export const ApiEntitySpendLimitSchema = EntitySpendLimitSchema;
+
+export type ApiEntitySpendLimit = z.infer<typeof ApiEntitySpendLimitSchema>;

shared/api/entities/billingControls/index.ts

@@ -0,0 +1,2 @@
+export * from "./entityBillingControls.js";
+export * from "./entitySpendLimit.js";

shared/api/entities/crud/createEntityParams.ts

@@ -1,5 +1,6 @@
 import { z } from "zod/v4";
 import { CustomerDataSchema } from "../../common/customerData.js";
+import { ApiEntityBillingControlsSchema } from "../billingControls/entityBillingControls.js";
 
 export const CreateEntityParamsV0Schema = z.object({
 	id: z
@@ -17,6 +18,9 @@ export const CreateEntityParamsV0Schema = z.object({
 	feature_id: z.string().meta({
 		description: "The ID of the feature this entity is associated with",
 	}),
+	billing_controls: ApiEntityBillingControlsSchema.optional().meta({
+		description: "Billing controls for the entity.",
+	}),
 	customer_data: CustomerDataSchema.optional().meta({
 		description:
 			"Customer attributes used to resolve the customer when customer_id is not provided.",

shared/api/entities/crud/index.ts

@@ -1,3 +1,4 @@
 export * from "./createEntityParams.js";
 export * from "./deleteEntityParams.js";
 export * from "./getEntityParams.js";
+export * from "./updateEntityParams.js";

shared/api/entities/crud/updateEntityParams.ts

@@ -0,0 +1,16 @@
+import { z } from "zod/v4";
+import { ApiEntityBillingControlsSchema } from "../billingControls/entityBillingControls.js";
+
+export const UpdateEntityParamsSchema = z.object({
+	customer_id: z.string().optional().meta({
+		description: "The ID of the customer that owns the entity.",
+	}),
+	entity_id: z.string().meta({
+		description: "The ID of the entity.",
+	}),
+	billing_controls: ApiEntityBillingControlsSchema.optional().meta({
+		description: "Billing controls to replace on the entity.",
+	}),
+});
+
+export type UpdateEntityParams = z.infer<typeof UpdateEntityParamsSchema>;

shared/api/entities/index.ts

@@ -1 +1,2 @@
+export * from "./billingControls/index.js";
 export * from "./crud/index.js";

shared/models/cusModels/billingControlModels.ts

@@ -1,4 +1,13 @@
 import { z } from "zod/v4";
+import {
+	type EntityBillingControls,
+	type EntityBillingControlsInput,
+	EntityBillingControlsSchema,
+} from "./billingControls/entityBillingControls.js";
+import {
+	type EntitySpendLimit,
+	EntitySpendLimitSchema,
+} from "./billingControls/entitySpendLimit.js";
 import { PurchaseLimitIntervalEnum } from "./billingControls/purchaseLimitInterval.js";
 
 export const AutoTopupPurchaseLimitSchema = z.object({
@@ -49,3 +58,10 @@ export type CustomerBillingControls = z.infer<
 export type CustomerBillingControlsInput = z.input<
 	typeof CustomerBillingControlsSchema
 >;
+
+export { EntityBillingControlsSchema, EntitySpendLimitSchema };
+export type {
+	EntityBillingControls,
+	EntityBillingControlsInput,
+	EntitySpendLimit,
+};

shared/models/cusModels/billingControls/entityBillingControls.ts

@@ -0,0 +1,13 @@
+import { z } from "zod/v4";
+import { EntitySpendLimitSchema } from "./entitySpendLimit.js";
+
+export const EntityBillingControlsSchema = z.object({
+	spend_limits: z.array(EntitySpendLimitSchema).optional().meta({
+		description: "List of overage spend limits per feature.",
+	}),
+});
+
+export type EntityBillingControls = z.infer<typeof EntityBillingControlsSchema>;
+export type EntityBillingControlsInput = z.input<
+	typeof EntityBillingControlsSchema
+>;

shared/models/cusModels/billingControls/entitySpendLimit.ts

@@ -0,0 +1,29 @@
+import { z } from "zod/v4";
+
+export const EntitySpendLimitSchema = z
+	.object({
+		feature_id: z.string().optional().meta({
+			description: "Optional feature ID this spend limit applies to.",
+		}),
+		enabled: z.boolean().default(false).meta({
+			description: "Whether this spend limit is enabled.",
+		}),
+		overage_limit: z.number().min(0).optional().meta({
+			description: "Maximum allowed overage spend for the target feature.",
+		}),
+	})
+	.refine(
+		(data) => {
+			if (data.overage_limit === undefined) {
+				return true;
+			}
+
+			return data.feature_id !== undefined;
+		},
+		{
+			message: "feature_id is required when overage_limit is provided",
+			path: ["feature_id"],
+		},
+	);
+
+export type EntitySpendLimit = z.infer<typeof EntitySpendLimitSchema>;

shared/models/cusModels/entityModels/entityModels.ts

@@ -1,17 +1,19 @@
 import { z } from "zod/v4";
 import type { Feature } from "../../featureModels/featureModels.js";
+import { EntitySpendLimitSchema } from "../billingControlModels.js";
 
 export const EntitySchema = z.object({
-	id: z.string(),
+	id: z.string().nullable(),
 	org_id: z.string(),
 	created_at: z.number(),
 	internal_id: z.string(),
 	internal_customer_id: z.string(),
 	env: z.string(),
-	name: z.string(),
+	name: z.string().nullable(),
 	deleted: z.boolean(),
 	feature_id: z.string(),
 	internal_feature_id: z.string(),
+	spend_limits: z.array(EntitySpendLimitSchema).nullish(),
 });
 
 // export const CreateEntitySchema = z.object({

shared/models/cusModels/entityModels/entityTable.ts

@@ -2,6 +2,7 @@ import {
 	boolean,
 	foreignKey,
 	index,
+	jsonb,
 	numeric,
 	pgTable,
 	text,
@@ -9,6 +10,7 @@ import {
 } from "drizzle-orm/pg-core";
 import { features } from "../../featureModels/featureTable.js";
 import { organizations } from "../../orgModels/orgTable.js";
+import type { EntitySpendLimit } from "../billingControlModels.js";
 import { customers } from "../cusTable.js";
 
 export const entities = pgTable(
@@ -23,6 +25,7 @@ export const entities = pgTable(
 		name: text(),
 		deleted: boolean().default(false).notNull(),
 		internal_feature_id: text("internal_feature_id"),
+		spend_limits: jsonb().$type<EntitySpendLimit[]>(),
 
 		// Optional...
 		feature_id: text("feature_id"),

shared/models/cusModels/index.ts

@@ -1,2 +1,4 @@
 export * from "./billingControlModels.js";
+export * from "./billingControls/entityBillingControls.js";
+export * from "./billingControls/entitySpendLimit.js";
 export * from "./billingControls/purchaseLimitInterval.js";

shared/utils/cusUtils/fullCusUtils/fullCustomerToCustomerEntitlements.ts

@@ -78,7 +78,7 @@ export const fullCustomerToCustomerEntitlements = ({
 	sortCusEntsForDeduction({
 		cusEnts,
 		reverseOrder,
-		entityId: entity?.id,
+		entityId: entity?.id ?? undefined,
 		customerEntitlementFilters,
 	});