feat(billing): checkout carries the GPU wanted and the room is made when payment lands

apps/api/app/routes/billing.ts

@@ -1,5 +1,6 @@
 import { Actor } from '@nestri/core/actor';
 import { Billing } from '@nestri/core/billing/index';
+import { CheckoutRoom } from '@nestri/core/billing/checkout-room';
 import { Polar } from '@nestri/core/billing/polar';
 import { ErrorCodes, VisibleError } from '@nestri/core/error';
 import { Team } from '@nestri/core/team/index';
@@ -134,6 +135,44 @@ export namespace BillingApi {
 					...ids
 				});
 
+				// The room they paid for. Never allowed to fail the delivery:
+				// the plan is already applied, and a retry loop cannot make a
+				// GPU free.
+				if (
+					updated &&
+					delivery.gpu &&
+					delivery.subscriptionId &&
+					delivery.standing.plan === 'paid' &&
+					delivery.standing.status === 'active'
+				) {
+					try {
+						const room = await CheckoutRoom.forSubscription({
+							teamId: delivery.teamId,
+							subscriptionId: delivery.subscriptionId,
+							gpu: delivery.gpu
+						});
+						console.log('billing checkout room', {
+							team: delivery.teamId,
+							requested: delivery.gpu,
+							...room
+						});
+						if (room.reason === 'no gpu free') {
+							console.error('a paid checkout asked for a GPU and none was free', {
+								team: delivery.teamId,
+								requested: delivery.gpu,
+								...ids
+							});
+						}
+					} catch (err) {
+						console.error('a paid checkout room failed to be created', {
+							team: delivery.teamId,
+							requested: delivery.gpu,
+							...ids,
+							err
+						});
+					}
+				}
+
 				return c.json({
 					data: { applied: Boolean(updated), type: delivery.type }
 				});
@@ -197,19 +236,26 @@ export namespace BillingApi {
 					.object({
 						successUrl: z.url().optional().meta({
 							description: 'Where to return to once they have paid'
+						}),
+						gpu: z.string().min(1).max(128).optional().meta({
+							description:
+								'A fleet GPU model, as `GET /location` lists it. When payment lands, a room on it is created (or on the closest free model if none is free)'
 						})
 					})
 					.strict()
 			),
 			async (c) => {
 				mustBeConfigured();
+				const { successUrl, gpu } = c.req.valid('json');
+				if (gpu) await CheckoutRoom.assertKnown(gpu);
 				const teamId = await payingTeam();
 				const user = await User.fromID(Actor.userID);
 				return c.json({
 					data: await Polar.checkout({
 						teamId,
 						email: user?.email ?? undefined,
-						successUrl: c.req.valid('json').successUrl
+						successUrl,
+						gpu
 					})
 				});
 			}

apps/api/test/billing-checkout-gpu.test.ts

@@ -0,0 +1,232 @@
+import { createHmac } from 'node:crypto';
+
+import { afterAll, afterEach, describe, expect, test } from 'bun:test';
+
+import { AccessToken } from '@nestri/core/access-token/index';
+import { Polar } from '@nestri/core/billing/polar';
+import { Fixtures } from '@nestri/core/db/fixtures';
+import { testDb } from '@nestri/core/db/test';
+import { Env } from '@nestri/core/env';
+import { Identifier } from '@nestri/core/id';
+import { Machine } from '@nestri/core/machine/index';
+
+import { app } from '../app/index';
+import { TEST_FRONTEND_URL } from './setup';
+
+const sql = testDb();
+const createdUserIds: string[] = [];
+const createdOrgIds: string[] = [];
+const realFetch = globalThis.fetch;
+
+const PAID = 'prod_paid';
+const V710 = { model: 'AMD Radeon Pro V710', vramMib: 28672, pciId: '1002:7461' };
+const RX = { model: 'AMD Radeon RX 9070', vramMib: 16384, pciId: '1002:7550' };
+const H100 = { model: 'NVIDIA H100', vramMib: 81920, pciId: '10de:2330' };
+const RTX = { model: 'NVIDIA RTX 4090', vramMib: 24576, pciId: '10de:2684' };
+
+const secret = `whsec_${Buffer.from('d'.repeat(32)).toString('base64')}`;
+
+function configure(organisationId?: string) {
+	Env.init({
+		NODE_ENV: 'test',
+		FRONTEND_URL: TEST_FRONTEND_URL,
+		POLAR_ACCESS_TOKEN: 'polar_test',
+		POLAR_PRODUCT_ID: PAID,
+		POLAR_WEBHOOK_SECRET: secret,
+		...(organisationId ? { FLEET_ORGANISATION_ID: organisationId } : {})
+	});
+	Polar.reset();
+}
+
+async function person(label: string) {
+	const owner = await Fixtures.owner(label);
+	createdUserIds.push(owner.userId);
+	const pat = await AccessToken.create({
+		id: Identifier.ascending('accessToken'),
+		ownerUserId: owner.userId,
+		teamId: null,
+		name: label
+	});
+	return {
+		owner,
+		headers: {
+			authorization: `Bearer ${pat.token}`,
+			'content-type': 'application/json'
+		} as Record<string, string>
+	};
+}
+
+async function fleet(gpus: (typeof V710)[][]) {
+	const operator = await Fixtures.owner('operator');
+	createdUserIds.push(operator.userId);
+	const f = await Fixtures.fleet(operator, gpus.length);
+	createdOrgIds.push(f.organisationId);
+	for (const [i, g] of gpus.entries()) await Machine.setGpus({ id: f.machines[i]!.id, gpus: g });
+	configure(f.organisationId);
+	return f;
+}
+
+/** Stand in for Polar, remembering what was sent to create a checkout. */
+function polarCheckout() {
+	const sent: any[] = [];
+	globalThis.fetch = (async (input: Request | string | URL, init?: RequestInit) => {
+		const req = input instanceof Request ? input : new Request(String(input), init);
+		if (req.url.includes('polar.sh')) {
+			sent.push(await req.json());
+			return new Response(JSON.stringify({ id: 'chk_1', url: 'https://polar.sh/pay/chk_1' }), {
+				status: 201,
+				headers: { 'content-type': 'application/json' }
+			});
+		}
+		return realFetch(input as never, init);
+	}) as typeof fetch;
+	return sent;
+}
+
+let n = 0;
+async function deliver(type: string, data: Record<string, unknown>) {
+	const body = JSON.stringify({ type, data });
+	const id = `msg_cg_${++n}`;
+	const ts = Math.floor(Date.now() / 1000);
+	// Standard Webhooks: HMAC-SHA256 over `id.timestamp.body`, keyed by the decoded secret.
+	const signature = `v1,${createHmac('sha256', Buffer.from(secret.slice(6), 'base64'))
+		.update(`${id}.${ts}.${body}`)
+		.digest('base64')}`;
+	return app.request('/billing/webhook', {
+		method: 'POST',
+		headers: {
+			'webhook-id': id,
+			'webhook-timestamp': String(ts),
+			'webhook-signature': signature
+		},
+		body
+	});
+}
+
+const paid = (teamId: string, sub: string, gpu?: string) => ({
+	id: sub,
+	productId: PAID,
+	customer: { externalId: teamId },
+	metadata: gpu ? { teamId, gpu } : { teamId }
+});
+
+const rooms = (userId: string) =>
+	sql`select b.gpu_pci_id, b.label from box b where b.user_id = ${userId} and b.time_deleted is null`;
+
+afterEach(() => {
+	globalThis.fetch = realFetch;
+	Polar.reset();
+});
+
+afterAll(async () => {
+	await sql`delete from checkout_room where team_id in (select id from team where owner_id = any(${createdUserIds}))`;
+	await sql`delete from box where user_id = any(${createdUserIds})`;
+	await sql`delete from install_token where created_by_user_id = any(${createdUserIds})`;
+	await sql`delete from machine where owner_user_id = any(${createdUserIds})`;
+	await sql`delete from organisation where id = any(${createdOrgIds})`;
+	await sql`delete from "user" where id = any(${createdUserIds})`;
+});
+
+describe('POST /billing/checkout with a GPU', () => {
+	test('the model travels to the provider as checkout metadata', async () => {
+		await fleet([[V710], [H100]]);
+		const p = await person('co-meta');
+		const sent = polarCheckout();
+		const res = await app.request('/billing/checkout', {
+			method: 'POST',
+			headers: p.headers,
+			body: JSON.stringify({ gpu: H100.model })
+		});
+		// The stand-in's reply is not a full checkout, so the SDK may refuse
+		// to parse it; what matters here is what was sent.
+		void res;
+		expect(sent[0].metadata).toEqual({ teamId: p.owner.teamId, gpu: H100.model });
+	});
+
+	test('no gpu leaves the metadata as it was', async () => {
+		await fleet([[V710]]);
+		const p = await person('co-plain');
+		const sent = polarCheckout();
+		const res = await app.request('/billing/checkout', {
+			method: 'POST',
+			headers: p.headers,
+			body: JSON.stringify({})
+		});
+		void res;
+		expect(sent[0].metadata).toEqual({ teamId: p.owner.teamId });
+	});
+
+	test('a model the fleet does not have is refused, and nothing reaches the provider', async () => {
+		await fleet([[V710]]);
+		const p = await person('co-unknown');
+		const sent = polarCheckout();
+		const res = await app.request('/billing/checkout', {
+			method: 'POST',
+			headers: p.headers,
+			body: JSON.stringify({ gpu: 'Imaginary 9000' })
+		});
+		expect(res.status).toBe(400);
+		expect(sent).toHaveLength(0);
+	});
+});
+
+describe('the webhook makes the room that was paid for', () => {
+	test('a paid subscription with a gpu creates a room on that model', async () => {
+		await fleet([[V710], [H100]]);
+		const p = await person('wh-room');
+		const res = await deliver('subscription.active', paid(p.owner.teamId, 'sub_wh_room', H100.model));
+		expect(res.status).toBe(200);
+		const made = await rooms(p.owner.userId);
+		expect(made.map((r) => [r.gpu_pci_id, r.label])).toEqual([[H100.pciId, H100.model]]);
+		const [row] = await sql`select substituted, placed_model from checkout_room where subscription_id = 'sub_wh_room'`;
+		expect(row).toMatchObject({ substituted: false, placed_model: H100.model });
+	});
+
+	test('a redelivery, and the created event after the active one, add no second room', async () => {
+		await fleet([[H100], [H100]]);
+		const p = await person('wh-twice');
+		const data = paid(p.owner.teamId, 'sub_wh_twice', H100.model);
+		await deliver('subscription.created', data);
+		await deliver('subscription.active', data);
+		await deliver('subscription.active', data);
+		expect(await rooms(p.owner.userId)).toHaveLength(1);
+	});
+
+	test('with the model taken, the closest free one of the same vendor is used and recorded', async () => {
+		await fleet([[H100], [V710], [RX], [RTX]]);
+		const taker = await person('wh-taker');
+		await deliver('subscription.active', paid(taker.owner.teamId, 'sub_wh_taker', H100.model));
+		const p = await person('wh-sub');
+		// H100 is gone. Nvidia first, so the 4090 beats the closer-in-memory AMD cards.
+		await deliver('subscription.active', paid(p.owner.teamId, 'sub_wh_sub', H100.model));
+		const made = await rooms(p.owner.userId);
+		expect(made.map((r) => r.gpu_pci_id)).toEqual([RTX.pciId]);
+		const [row] = await sql`select substituted, requested_model, placed_model from checkout_room where subscription_id = 'sub_wh_sub'`;
+		expect(row).toMatchObject({
+			substituted: true,
+			requested_model: H100.model,
+			placed_model: RTX.model
+		});
+	});
+
+	test('with nothing free, no room is made, the plan is still applied, and a retry stays quiet', async () => {
+		await fleet([[V710]]);
+		const taker = await person('wh-none-taker');
+		await deliver('subscription.active', paid(taker.owner.teamId, 'sub_wh_nt', V710.model));
+		const p = await person('wh-none');
+		const data = paid(p.owner.teamId, 'sub_wh_none', V710.model);
+		expect((await deliver('subscription.active', data)).status).toBe(200);
+		expect(await rooms(p.owner.userId)).toHaveLength(0);
+		const [team] = await sql`select plan from team where id = ${p.owner.teamId}`;
+		expect(team!.plan).toBe('paid');
+		await deliver('subscription.active', data);
+		expect(await rooms(p.owner.userId)).toHaveLength(0);
+	});
+
+	test('a paid subscription with no gpu creates nothing', async () => {
+		await fleet([[V710]]);
+		const p = await person('wh-nogpu');
+		await deliver('subscription.active', paid(p.owner.teamId, 'sub_wh_nogpu'));
+		expect(await rooms(p.owner.userId)).toHaveLength(0);
+	});
+});

packages/core/migrations/0031_checkout_room.sql

@@ -0,0 +1,16 @@
+CREATE TABLE "checkout_room" (
+	"id" char(30) PRIMARY KEY NOT NULL,
+	"time_created" timestamp with time zone DEFAULT now() NOT NULL,
+	"time_updated" timestamp with time zone DEFAULT now() NOT NULL,
+	"time_deleted" timestamp with time zone,
+	"team_id" char(30) NOT NULL,
+	"subscription_id" text NOT NULL,
+	"requested_model" text NOT NULL,
+	"placed_model" text,
+	"substituted" boolean DEFAULT false NOT NULL,
+	"box_id" char(30)
+);
+--> statement-breakpoint
+ALTER TABLE "checkout_room" ADD CONSTRAINT "checkout_room_team_id_team_id_fk" FOREIGN KEY ("team_id") REFERENCES "public"."team"("id") ON DELETE cascade ON UPDATE no action;--> statement-breakpoint
+ALTER TABLE "checkout_room" ADD CONSTRAINT "checkout_room_box_id_box_id_fk" FOREIGN KEY ("box_id") REFERENCES "public"."box"("id") ON DELETE set null ON UPDATE no action;--> statement-breakpoint
+CREATE UNIQUE INDEX "checkout_room_subscription_idx" ON "checkout_room" USING btree ("subscription_id");

packages/core/migrations/meta/_journal.json

@@ -218,6 +218,13 @@
       "when": 1792436540630,
       "tag": "0030_box_gpu_slot",
       "breakpoints": true
+    },
+    {
+      "idx": 31,
+      "version": "7",
+      "when": 1792438603353,
+      "tag": "0031_checkout_room",
+      "breakpoints": true
     }
   ]
 }

packages/core/src/billing/checkout-room.sql.ts

@@ -0,0 +1,32 @@
+import { boolean, pgTable, text, uniqueIndex } from 'drizzle-orm/pg-core';
+
+import { BoxTable } from '../box/box.sql.js';
+import { id, timestamps, ulid } from '../db/types.js';
+import { TeamTable } from '../team/team.sql.js';
+
+/**
+ * The room a paid checkout asked for, one row per subscription.
+ *
+ * The unique subscription id is what makes a redelivered webhook harmless: the
+ * first delivery claims the row, every later one finds it taken and does
+ * nothing. The row is kept even when no room could be made, so a retry cannot
+ * hand out a card the person was already told was not there.
+ */
+export const CheckoutRoomTable = pgTable(
+	'checkout_room',
+	{
+		...id,
+		...timestamps,
+		teamId: ulid('team_id')
+			.notNull()
+			.references(() => TeamTable.id, { onDelete: 'cascade' }),
+		subscriptionId: text('subscription_id').notNull(),
+		/** The fleet model the person picked before paying. */
+		requestedModel: text('requested_model').notNull(),
+		/** The model the room went on; null when nothing was free. */
+		placedModel: text('placed_model'),
+		substituted: boolean('substituted').notNull().default(false),
+		boxId: ulid('box_id').references(() => BoxTable.id, { onDelete: 'set null' })
+	},
+	(table) => [uniqueIndex('checkout_room_subscription_idx').on(table.subscriptionId)]
+);

packages/core/src/billing/checkout-room.ts

@@ -0,0 +1,187 @@
+import { and, eq, isNull } from 'drizzle-orm';
+import z from 'zod';
+
+import { BoxTable } from '../box/box.sql.js';
+import { Box } from '../box/index.js';
+import { Placement } from '../box/placement.js';
+import { Database } from '../db/index.js';
+import { Env } from '../env.js';
+import { ErrorCodes, VisibleError } from '../error.js';
+import { fn } from '../fn.js';
+import { Identifier } from '../id.js';
+import { Team } from '../team/index.js';
+import { Billing } from './index.js';
+import { CheckoutRoomTable } from './checkout-room.sql.js';
+import { Rooms } from './rooms.js';
+
+/**
+ * The room somebody paid to get.
+ *
+ * Pro checkout can carry the fleet GPU the person was looking at. When the
+ * payment lands, that room is made for them, so the money and the thing it was
+ * for arrive together rather than leaving them to find the button again.
+ */
+export namespace CheckoutRoom {
+	/** Every model the fleet has a card of, free or not. */
+	export async function knownModels(): Promise<string[]> {
+		const organisationId = Env.get().FLEET_ORGANISATION_ID;
+		if (!organisationId) return [];
+		const cards = await Placement.fleetCards(organisationId);
+		return [...new Set(cards.map((c) => c.model))];
+	}
+
+	/** Refuse a model the fleet does not have. */
+	export async function assertKnown(model: string): Promise<void> {
+		if (!(await knownModels()).includes(model)) {
+			throw new VisibleError(
+				'validation',
+				ErrorCodes.Validation.INVALID_PARAMETER,
+				`${model} is not a Nestri GPU`,
+				'gpu'
+			);
+		}
+	}
+
+	const vendor = (pciId: string) => pciId.split(':')[0]?.toLowerCase() ?? '';
+
+	/**
+	 * Which model to use: the one asked for when a card of it is free,
+	 * otherwise the closest free one — same vendor first, then nearest memory.
+	 * Null when no fleet card is free.
+	 */
+	export function choose(
+		requested: string,
+		cards: readonly Placement.Card[]
+	): { model: string; substituted: boolean } | null {
+		const free = cards.filter((c) => c.free);
+		if (free.some((c) => c.model === requested)) return { model: requested, substituted: false };
+		if (free.length === 0) return null;
+		const wanted = cards.find((c) => c.model === requested);
+		const wantedVendor = wanted ? vendor(wanted.pciId) : '';
+		const wantedVram = wanted?.vramMib ?? 0;
+		const rank = (c: Placement.Card) => [
+			vendor(c.pciId) === wantedVendor ? 0 : 1,
+			Math.abs((c.vramMib ?? 0) - wantedVram)
+		];
+		const best = [...free].sort((a, b) => {
+			const [av, ad] = rank(a)!;
+			const [bv, bd] = rank(b)!;
+			return av! - bv! || ad! - bd! || a.model.localeCompare(b.model);
+		})[0]!;
+		return { model: best.model, substituted: true };
+	}
+
+	export const Outcome = z.object({
+		created: z.boolean(),
+		reason: z.enum(['created', 'duplicate', 'no gpu free', 'no fleet', 'no team', 'over limit']),
+		model: z.string().nullable(),
+		substituted: z.boolean(),
+		boxId: z.string().nullable()
+	});
+	export type Outcome = z.infer<typeof Outcome>;
+
+	/**
+	 * Make the room a paid subscription asked for, once.
+	 *
+	 * Keyed on the subscription id: a redelivery finds the row already claimed
+	 * and does nothing, including when the first delivery found no GPU free.
+	 */
+	export const forSubscription = fn(
+		z.object({ teamId: z.string(), subscriptionId: z.string(), gpu: z.string() }),
+		async (input): Promise<Outcome> => {
+			const none = (reason: Outcome['reason']): Outcome => ({
+				created: false,
+				reason,
+				model: null,
+				substituted: false,
+				boxId: null
+			});
+			const organisationId = Env.get().FLEET_ORGANISATION_ID;
+			const team = await Team.fromID(input.teamId);
+			if (!team) return none('no team');
+
+			return Database.transaction(async () => {
+				const claimed = await Database.use((tx) =>
+					tx
+						.insert(CheckoutRoomTable)
+						.values({
+							id: Identifier.ascending('checkoutRoom'),
+							teamId: input.teamId,
+							subscriptionId: input.subscriptionId,
+							requestedModel: input.gpu
+						})
+						.onConflictDoNothing()
+						.returning({ id: CheckoutRoomTable.id })
+						.then((r) => r[0])
+				);
+				if (!claimed) return none('duplicate');
+				if (!organisationId) return none('no fleet');
+
+				let picked: ReturnType<typeof choose> = null;
+				const placer: Placement.Placer = async (request) => {
+					picked = choose(input.gpu, await Placement.fleetCards(organisationId));
+					if (!picked) {
+						throw new VisibleError(
+							'rate_limit',
+							ErrorCodes.Server.SERVICE_UNAVAILABLE,
+							'No Nestri GPU is free right now'
+						);
+					}
+					return Placement.onFleetModel(picked.model)(request);
+				};
+				let overLimit = false;
+				let box: Awaited<ReturnType<typeof Box.createPlaced>> | null = null;
+				try {
+					box = await Box.createPlaced(
+						{
+							id: Identifier.ascending('box'),
+							userId: team.ownerId,
+							label: input.gpu.slice(0, 64),
+							tier: 'sm'
+						},
+						placer,
+						async () => {
+							if ((await Billing.cloudRooms(team.ownerId)) >= Rooms.allowed(team.plan)) {
+								overLimit = true;
+								throw Rooms.atLimit(input.gpu, team.plan);
+							}
+						}
+					);
+				} catch (err) {
+					if (!(err instanceof VisibleError)) throw err;
+				}
+				if (!box) {
+					return none(overLimit ? 'over limit' : 'no gpu free');
+				}
+				const final = picked as { model: string; substituted: boolean } | null;
+				// A room is named after its GPU, which is not the one asked for
+				// when another was put in its place.
+				if (final?.substituted) {
+					await Database.use((tx) =>
+						tx
+							.update(BoxTable)
+							.set({ label: final.model.slice(0, 64) })
+							.where(eq(BoxTable.id, box!.id))
+					);
+				}
+				await Database.use((tx) =>
+					tx
+						.update(CheckoutRoomTable)
+						.set({
+							boxId: box!.id,
+							placedModel: final?.model ?? input.gpu,
+							substituted: final?.substituted ?? false
+						})
+						.where(and(eq(CheckoutRoomTable.id, claimed.id), isNull(CheckoutRoomTable.timeDeleted)))
+				);
+				return {
+					created: true,
+					reason: 'created',
+					model: final?.model ?? input.gpu,
+					substituted: final?.substituted ?? false,
+					boxId: box.id
+				};
+			});
+		}
+	);
+}

packages/core/src/billing/polar.ts

@@ -173,7 +173,9 @@ export namespace Polar {
 		z.object({
 			teamId: z.string(),
 			email: z.email().optional(),
-			successUrl: z.url().optional()
+			successUrl: z.url().optional(),
+			/** The fleet GPU the person wanted a room on; made when payment lands. */
+			gpu: z.string().min(1).max(128).optional()
 		}),
 		async (input) => {
 			const { productId } = settings();
@@ -191,7 +193,7 @@ export namespace Polar {
 				// subscription. A customer who already existed under this email
 				// keeps the customer record they had, without our id on it, and
 				// this is then the only place the team survives.
-				metadata: { teamId: input.teamId },
+				metadata: input.gpu ? { teamId: input.teamId, gpu: input.gpu } : { teamId: input.teamId },
 				customerEmail: input.email,
 				successUrl: input.successUrl
 			});
@@ -309,6 +311,8 @@ export namespace Polar {
 		 */
 		subscriptionId: string | null;
 		customerId: string | null;
+		/** The fleet GPU the checkout asked a room on, if it named one. */
+		gpu: string | null;
 	}
 
 	/**
@@ -382,7 +386,7 @@ export namespace Polar {
 			// existed before checkout keeps the record they had, without it, so
 			// the team then comes from the checkout's metadata, which the
 			// subscription carries. A delivery with neither was not started here.
-			const metadata = data.metadata as { teamId?: unknown } | undefined;
+			const metadata = data.metadata as { teamId?: unknown; gpu?: unknown } | undefined;
 			const teamId =
 				customer?.externalId ??
 				customer?.external_id ??
@@ -411,7 +415,8 @@ export namespace Polar {
 				teamId,
 				standing: standingFor(event.type, productId),
 				subscriptionId,
-				customerId
+				customerId,
+				gpu: typeof metadata?.gpu === 'string' && metadata.gpu ? metadata.gpu : null
 			};
 		}
 	);

packages/core/src/id.ts

@@ -17,6 +17,7 @@ export namespace Identifier {
 		trialClaim: 'trc',
 		steamLinkRequest: 'slr',
 		box: 'box',
+		checkoutRoom: 'cor',
 		session: 'ses',
 		launch: 'lch',
 		accessToken: 'pat',