fix(billing): preserve generated schedule edits

server/src/internal/billing/v2/actions/generateRequest/compute/buildGenerationPrompts.ts

@@ -22,6 +22,7 @@ export const buildSystemPrompt = (tool: GenerateBillingTool) => {
 		"- Never omit a required field. Always set plan_id to your best match from the context plans; when sibling variants exist (e.g. monthly vs yearly), pick the one matching the stated interval or amount, defaulting to the monthly variant.",
 		"- When the request states a price for a plan (e.g. 'at 10k/mo'), always set customize.price to it — including Enterprise/custom placeholder plans where the docs say to ask about the base price.",
 		"- customer.current_plans[].effective_plan is the subscription's live configuration in the same shape as context.plans. When changing a plan or version, explicitly preserve any current term the request says to keep by copying it into the corresponding customize override.",
+		"- customer.schedules shows the persisted phase topology; each customer_product_id references customer.current_plans. Preserve it unless the request explicitly changes it.",
 		"- Use ONLY plan ids and feature ids that appear in the context below.",
 		"- Monetary amounts are in major currency units (e.g. dollars). Never convert to cents.",
 		"- The operation always targets the customer in the context. Ignore any other customer mentioned in the request.",

server/src/internal/billing/v2/actions/generateRequest/setup/setupGenerationContext.ts

@@ -12,6 +12,8 @@ import {
 } from "@autumn/shared";
 import type { AutumnContext } from "@/honoUtils/HonoEnv";
 import { CusService } from "@/internal/customers/CusService";
+import { CusProductService } from "@/internal/customers/cusProducts/CusProductService";
+import { getCustomerSchedulesByScope } from "@/internal/customers/cusUtils/getFullCustomerSchedule";
 import { FeatureService } from "@/internal/features/FeatureService";
 import { ProductService } from "@/internal/products/ProductService";
 
@@ -61,11 +63,12 @@ export const compactCustomerProduct = ({
 	const entity = entities.find(
 		(candidate) => candidate.internal_id === customerProduct.internal_entity_id,
 	);
+	const entityId = entity?.id ?? customerProduct.entity_id;
 	return {
 		customer_product_id: customerProduct.id,
 		plan_id: customerProduct.product.id,
 		status: customerProduct.status,
-		...(entity?.id ? { entity_id: entity.id } : {}),
+		...(entityId ? { entity_id: entityId } : {}),
 		...(customerProduct.canceled ? { canceled: true } : {}),
 		...(customerProduct.trial_ends_at
 			? { trial_ends_at: customerProduct.trial_ends_at }
@@ -86,8 +89,66 @@ export const setupGenerationContext = async ({
 	const [fullProducts, features, fullCustomer] = await Promise.all([
 		ProductService.listFull({ db: ctx.db, env: ctx.env, orgId: ctx.org.id }),
 		FeatureService.list({ db: ctx.db, env: ctx.env, orgId: ctx.org.id }),
-		CusService.getFull({ ctx, idOrInternalId: customerId }),
+		CusService.getFull({ ctx, idOrInternalId: customerId, withEntities: true }),
 	]);
+	const [{ customerSchedule, entitySchedules }, allCustomerProducts] =
+		await Promise.all([
+			getCustomerSchedulesByScope({
+				ctx,
+				internalCustomerId: fullCustomer.internal_id,
+			}),
+			CusProductService.list({
+				db: ctx.db,
+				internalCustomerId: fullCustomer.internal_id,
+			}),
+		]);
+	const scheduledProductIds = new Set(
+		[customerSchedule, ...Object.values(entitySchedules)].flatMap(
+			(schedule) =>
+				schedule?.phases.flatMap((phase) => phase.customer_product_ids) ?? [],
+		),
+	);
+	const loadedProductIds = new Set(
+		fullCustomer.customer_products.map((product) => product.id),
+	);
+	const customerProducts = allCustomerProducts.filter(
+		(product) =>
+			loadedProductIds.has(product.id) || scheduledProductIds.has(product.id),
+	);
+	const entities = fullCustomer.entities ?? [];
+	const customerProductById = new Map(
+		customerProducts.map((product) => [product.id, product]),
+	);
+	const entityIdByInternalId = new Map(
+		customerProducts.flatMap((product) =>
+			product.internal_entity_id && product.entity_id
+				? [[product.internal_entity_id, product.entity_id]]
+				: [],
+		),
+	);
+	const compactSchedule = (
+		schedule: NonNullable<typeof customerSchedule>,
+		entityId?: string,
+	) => ({
+		...(entityId ? { entity_id: entityId } : {}),
+		phases: schedule.phases.map(({ starts_at, customer_product_ids }) => ({
+			starts_at,
+			customer_product_ids,
+			...(customer_product_ids.some(
+				(id) =>
+					customerProductById.get(id)?.billing_cycle_anchor_resets_at ===
+					starts_at,
+			)
+				? { billing_cycle_anchor: "phase_start" as const }
+				: {}),
+		})),
+	});
+	const schedules = [
+		...(customerSchedule ? [compactSchedule(customerSchedule)] : []),
+		...Object.entries(entitySchedules).map(([internalEntityId, schedule]) =>
+			compactSchedule(schedule, entityIdByInternalId.get(internalEntityId)),
+		),
+	];
 
 	const now = Date.now();
 
@@ -96,23 +157,22 @@ export const setupGenerationContext = async ({
 			customer: {
 				id: fullCustomer.id,
 				...(fullCustomer.name ? { name: fullCustomer.name } : {}),
-				current_plans: fullCustomer.customer_products.map(
-					(customerProduct) => ({
-						...compactCustomerProduct({
-							customerProduct,
-							entities: fullCustomer.entities ?? [],
-						}),
-						effective_plan: compactPlan({
-							features,
-							product: cusProductToProduct({ cusProduct: customerProduct }),
-						}),
+				current_plans: customerProducts.map((customerProduct) => ({
+					...compactCustomerProduct({
+						customerProduct,
+						entities,
 					}),
-				),
-				entities: (fullCustomer.entities ?? []).map((entity) => ({
+					effective_plan: compactPlan({
+						features,
+						product: cusProductToProduct({ cusProduct: customerProduct }),
+					}),
+				})),
+				entities: entities.map((entity) => ({
 					id: entity.id,
 					...(entity.name ? { name: entity.name } : {}),
 					...(entity.feature_id ? { feature_id: entity.feature_id } : {}),
 				})),
+				...(schedules.length ? { schedules } : {}),
 			},
 			features: features.map((feature) => ({
 				id: feature.id,

server/tests/evals/generateBillingRequest/fixtures.ts

@@ -104,22 +104,182 @@ export const versionedPlanContext = ({
 		})),
 	}) as unknown as GenerationContext;
 
-export const versionedScheduleContext = (): GenerationContext => {
-	const context = versionedPlanContext();
+export const complexScheduleStarts = [
+	Date.UTC(2026, 7, 24, 12),
+	Date.UTC(2027, 7, 24, 12),
+	Date.UTC(2028, 7, 24, 12),
+	Date.UTC(2029, 7, 24, 12),
+];
+
+const monthlyItem = (feature_id: string, included: number) => ({
+	feature_id,
+	included,
+	pooled: false,
+	reset: { interval: "month" },
+});
+
+const planFeatureIds: Record<string, string> = {
+	"analytics-addon": "reports",
+	core: "messages",
+	"success-addon": "sessions",
+	"support-addon": "tickets",
+};
+
+const customizedPlan = ({
+	amount,
+	included,
+	plan_id,
+	version = 1,
+}: {
+	amount: number;
+	included: number;
+	plan_id: string;
+	version?: number;
+}) => ({
+	customize: {
+		items: [monthlyItem(planFeatureIds[plan_id]!, included)],
+		price: { amount, interval: "month" },
+	},
+	plan_id,
+	version,
+});
+
+export const complexScheduleRequest = ({
+	phaseTwoVersion = 2,
+	phaseThreeSupportPrice = 275,
+	phaseFourMessages = 40_000,
+}: {
+	phaseTwoVersion?: number;
+	phaseThreeSupportPrice?: number;
+	phaseFourMessages?: number;
+} = {}): Record<string, unknown> => ({
+	billing_behavior: "none",
+	billing_cycle_anchor: "now",
+	phases: [
+		[900, 12_000, 1, 175, 5_000, 225, 25],
+		[1_100, 20_000, phaseTwoVersion, 190, 7_500, 250, 50],
+		[1_350, 30_000, 3, 200, 9_000, phaseThreeSupportPrice, 75],
+		[1_500, phaseFourMessages, 3, 225, 10_000, 300, 100],
+	].map(
+		(
+			[
+				corePrice,
+				messages,
+				version,
+				analyticsPrice,
+				reports,
+				supportPrice,
+				tickets,
+			],
+			index,
+		) => ({
+			...(index % 2 ? { billing_cycle_anchor: "phase_start" } : {}),
+			plans: [
+				customizedPlan({
+					amount: corePrice,
+					included: messages,
+					plan_id: "core",
+					version,
+				}),
+				customizedPlan({
+					amount: analyticsPrice,
+					included: reports,
+					plan_id: "analytics-addon",
+				}),
+				customizedPlan({
+					amount: supportPrice,
+					included: tickets,
+					plan_id: "support-addon",
+				}),
+			],
+			starts_at: complexScheduleStarts[index],
+		}),
+	),
+	unscheduled_plans: [
+		customizedPlan({
+			amount: 95,
+			included: 25,
+			plan_id: "success-addon",
+		}),
+	],
+});
+
+export const complexScheduleContext = (): GenerationContext => {
+	const plans = [
+		...[1, 2, 3].map((version) => ({
+			id: "core",
+			items: [monthlyItem("messages", version * 10_000)],
+			name: "Core",
+			price: { amount: version * 500, interval: "month" },
+			version,
+		})),
+		...(
+			[
+				["analytics-addon", "Analytics Add-on", 200, 5_000, "reports"],
+				["support-addon", "Support Add-on", 300, 50, "tickets"],
+				["success-addon", "Success Add-on", 100, 25, "sessions"],
+			] as const
+		).map(([id, name, amount, included, featureId]) => ({
+			id,
+			is_add_on: true,
+			items: [monthlyItem(featureId, included)],
+			name,
+			price: { amount, interval: "month" },
+			version: 1,
+		})),
+	];
+	const schedule = complexScheduleRequest() as {
+		phases: {
+			billing_cycle_anchor?: string;
+			plans: ReturnType<typeof customizedPlan>[];
+			starts_at: number;
+		}[];
+	};
+	const customerProductId = (phase: number, planId: string) =>
+		`cp_${phase + 1}_${planId}`;
+
 	return {
-		...context,
-		plans: [
-			...context.plans,
-			{
-				id: "support-addon",
-				is_add_on: true,
-				items: [],
-				name: "Support Add-on",
-				price: { amount: 10, interval: BillingInterval.Month },
-				version: 1,
-			},
+		customer: {
+			current_plans: schedule.phases.flatMap((phase, phaseIndex) =>
+				phase.plans.map((plan) => ({
+					customer_product_id: customerProductId(phaseIndex, plan.plan_id),
+					effective_plan: {
+						...plans.find(
+							(candidate) =>
+								candidate.id === plan.plan_id &&
+								candidate.version === plan.version,
+						),
+						...plan.customize,
+					},
+					plan_id: plan.plan_id,
+					status: phaseIndex ? "scheduled" : "active",
+				})),
+			),
+			id: "cus_complex_schedule",
+			name: "Example Company",
+			schedules: [
+				{
+					phases: schedule.phases.map((phase, phaseIndex) => ({
+						...(phase.billing_cycle_anchor
+							? { billing_cycle_anchor: phase.billing_cycle_anchor }
+							: {}),
+						customer_product_ids: phase.plans.map((plan) =>
+							customerProductId(phaseIndex, plan.plan_id),
+						),
+						starts_at: phase.starts_at,
+					})),
+				},
+			],
+		},
+		features: [
+			{ id: "messages", name: "Messages", type: "single_use" },
+			{ id: "reports", name: "Reports", type: "single_use" },
+			{ id: "sessions", name: "Success Sessions", type: "single_use" },
+			{ id: "tickets", name: "Support Tickets", type: "single_use" },
 		],
-	};
+		now,
+		plans,
+	} as unknown as GenerationContext;
 };
 
 /** Enterprise org with a volume-tiered prepaid credits ladder — the shape that