From fab29b33443406d745e304f99653c4a964c1777c Mon Sep 17 00:00:00 2001 From: Joris Gallot Date: Fri, 14 Aug 2026 18:18:33 +0200 Subject: [PATCH 1/2] feat(billing): cap artifact storage per plan --- AGENTS.md | 2 +- ...260814180000_add_artifact_project_index.ts | 14 ++++ packages/server/src/billing/entitlements.ts | 33 ++++++++- packages/server/src/billing/usage.ts | 15 +++- .../server/src/db/schemas/kinora-schemas.ts | 5 +- packages/server/src/public-api/index.ts | 18 ++++- packages/server/src/router/billing.ts | 7 +- packages/server/test/ingest-cap.test.ts | 36 ++++++++-- packages/server/test/storage-cap.test.ts | 72 +++++++++++++++++++ packages/server/test/usage.test.ts | 37 +++++++++- .../web/src/pages/WorkspaceSettingsPage.vue | 44 +++++++++++- 11 files changed, 265 insertions(+), 18 deletions(-) create mode 100644 packages/server/migrations/20260814180000_add_artifact_project_index.ts create mode 100644 packages/server/test/storage-cap.test.ts diff --git a/AGENTS.md b/AGENTS.md index 124dae5..6343f40 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -89,7 +89,7 @@ The desktop app authenticates via the **OAuth 2.0 device authorization grant** ( One codebase, two deployment modes gated by `KINORA_CLOUD` (env). Self-host (`false`) unlocks every feature; cloud (`true`) enables **Polar** billing. -- **Billing** (`src/billing/`): `polar.ts` (Polar SDK + better-auth plugin), `entitlements.ts` / `usage.ts` (plan limits), `retention.ts` (per-plan run-retention windows, swept by `purge-expired-runs`). +- **Billing** (`src/billing/`): `polar.ts` (Polar SDK + better-auth plugin), `entitlements.ts` / `usage.ts` (plan limits: monthly test results, projects, and artifact bytes - the storage cap rejects an over-quota upload with a 402 and deletes the blob it just streamed), `retention.ts` (per-plan run-retention windows, swept by `purge-expired-runs`). - **Retention** (`src/billing/retention.ts`): cloud derives windows from the plan tier and is swept by an external cron calling `purge-expired-runs`. Self-host instead reads `retentionPolicy` from env (`KINORA_ARTIFACT_RETENTION_DAYS` drops blobs but keeps runs; `KINORA_RETENTION_DAYS` / `KINORA_KEEP_LAST_RUNS` delete runs), and `src/index.ts` runs a daily in-process sweep gated on that policy being non-null, so cloud never double-sweeps. - **Alerts** (`src/alerts/`): per-project notifications on new failures / regressions. Channels are `slack.ts`, `email.ts` (nodemailer/SMTP), `webhook.ts`, dispatched by `notify.ts` with an every-run / on-failure / on-regression policy (`core.ts`). - **Feedback** (`src/feedback/`, `feedback` tRPC router): in-app "Send feedback" posts bug/feature reports to the private cloud task tracker. Cloud-only: `resolveFeedbackTracker` in `env.ts` returns null unless `KINORA_CLOUD=true` and all `FEEDBACK_TRACKER_*` vars are set; `config.get.feedbackEnabled` gates the web UI. diff --git a/packages/server/migrations/20260814180000_add_artifact_project_index.ts b/packages/server/migrations/20260814180000_add_artifact_project_index.ts new file mode 100644 index 0000000..a713105 --- /dev/null +++ b/packages/server/migrations/20260814180000_add_artifact_project_index.ts @@ -0,0 +1,14 @@ +import type { Knex } from 'knex' + +// The storage quota sums an org's artifact sizes on every upload; without this it seq-scans. +export async function up(knex: Knex): Promise { + await knex.schema.alterTable('artifact', (t) => { + t.index(['project_id'], 'artifact_projectId_idx') + }) +} + +export async function down(knex: Knex): Promise { + await knex.schema.alterTable('artifact', (t) => { + t.dropIndex(['project_id'], 'artifact_projectId_idx') + }) +} diff --git a/packages/server/src/billing/entitlements.ts b/packages/server/src/billing/entitlements.ts index 9eae0c1..920a6ba 100644 --- a/packages/server/src/billing/entitlements.ts +++ b/packages/server/src/billing/entitlements.ts @@ -11,20 +11,24 @@ export interface Entitlements { maxProjects: number retentionDays: number includedResults: number + storageBytes: number alerts: boolean } +const GB = 1024 ** 3 + const UNLIMITED: Omit = { maxProjects: Number.POSITIVE_INFINITY, retentionDays: Number.POSITIVE_INFINITY, includedResults: Number.POSITIVE_INFINITY, + storageBytes: Number.POSITIVE_INFINITY, alerts: true, } const LIMITS: Record> = { - free: { maxProjects: 1, retentionDays: 7, includedResults: 2_500, alerts: false }, - team: { maxProjects: Number.POSITIVE_INFINITY, retentionDays: 90, includedResults: 10_000, alerts: true }, - pro: { maxProjects: Number.POSITIVE_INFINITY, retentionDays: 365, includedResults: 50_000, alerts: true }, + free: { maxProjects: 1, retentionDays: 7, includedResults: 2_500, storageBytes: 2 * GB, alerts: false }, + team: { maxProjects: Number.POSITIVE_INFINITY, retentionDays: 90, includedResults: 10_000, storageBytes: 50 * GB, alerts: true }, + pro: { maxProjects: Number.POSITIVE_INFINITY, retentionDays: 365, includedResults: 50_000, storageBytes: 250 * GB, alerts: true }, enterprise: UNLIMITED, selfhost: UNLIMITED, } @@ -225,3 +229,26 @@ export function ingestCapError(e: Entitlements, usedResults: number, isNewProjec return null } + +// Artifact storage cap. `incoming` is the artifact about to be stored (0 for a pre-flight check). +export function storageCapError(e: Entitlements, usedBytes: number, incoming = 0): IngestCap | null { + if (!Number.isFinite(e.storageBytes) || usedBytes + incoming <= e.storageBytes) + return null + return { + error: `Plan artifact storage limit reached (${formatBytes(e.storageBytes)}). Older runs free up space as they expire, or upgrade for more.`, + limit: e.storageBytes, + } +} + +export function formatBytes(bytes: number): string { + if (!Number.isFinite(bytes)) + return 'unlimited' + const units = ['B', 'KB', 'MB', 'GB', 'TB'] + let value = bytes + let unit = 0 + while (value >= 1024 && unit < units.length - 1) { + value /= 1024 + unit++ + } + return `${Math.round(value * 10) / 10} ${units[unit]}` +} diff --git a/packages/server/src/billing/usage.ts b/packages/server/src/billing/usage.ts index 16c0028..9fb00a7 100644 --- a/packages/server/src/billing/usage.ts +++ b/packages/server/src/billing/usage.ts @@ -1,6 +1,6 @@ -import { and, count, eq, gte } from 'drizzle-orm' +import { and, count, eq, gte, sum } from 'drizzle-orm' import { db } from '../db' -import { project, run, test } from '../db/schemas/index' +import { artifact, project, run, test } from '../db/schemas/index' export function startOfMonthUtc(): Date { const now = new Date() @@ -20,6 +20,17 @@ export async function currentPeriodResults(organizationId: string): Promise { + const [row] = await db + .select({ total: sum(artifact.size) }) + .from(artifact) + .innerJoin(project, eq(artifact.projectId, project.id)) + .where(eq(project.organizationId, organizationId)) + + return Number(row?.total ?? 0) +} + export async function projectCount(organizationId: string): Promise { const [row] = await db .select({ total: count() }) diff --git a/packages/server/src/db/schemas/kinora-schemas.ts b/packages/server/src/db/schemas/kinora-schemas.ts index 3c823e8..6f7760b 100644 --- a/packages/server/src/db/schemas/kinora-schemas.ts +++ b/packages/server/src/db/schemas/kinora-schemas.ts @@ -69,7 +69,10 @@ export const artifact = pgTable('artifact', { sha1: text('sha1'), size: integer('size'), createdAt: timestamp('created_at').defaultNow().notNull(), -}, table => [index('artifact_runId_idx').on(table.runId)]) +}, table => [ + index('artifact_runId_idx').on(table.runId), + index('artifact_projectId_idx').on(table.projectId), +]) // Cached Polar billing state, synced from the customer.state_changed webhook. export const subscription = pgTable('subscription', { diff --git a/packages/server/src/public-api/index.ts b/packages/server/src/public-api/index.ts index 1ba3d07..9e83b31 100644 --- a/packages/server/src/public-api/index.ts +++ b/packages/server/src/public-api/index.ts @@ -8,9 +8,9 @@ import { and, eq } from 'drizzle-orm' import { Hono } from 'hono' import { bodyLimit } from 'hono/body-limit' import { notifyRun } from '../alerts/notify' -import { getEntitlements, ingestCapError, quotaCrossing, quotaWarningText } from '../billing/entitlements' +import { getEntitlements, ingestCapError, quotaCrossing, quotaWarningText, storageCapError } from '../billing/entitlements' import { meterTestResults, polarClient } from '../billing/polar' -import { currentPeriodResults, projectCount, startOfMonthUtc } from '../billing/usage' +import { currentPeriodResults, projectCount, startOfMonthUtc, storageBytes } from '../billing/usage' import { db } from '../db' import { artifact, member, project, run, test, user } from '../db/schemas/index' import { auth } from '../lib/auth' @@ -320,12 +320,26 @@ publicApi.post('/runs/:runId/artifacts', async (c) => { if (!owner) return c.json({ error: 'Run not found' }, 404) + // One usage read per upload, reused after streaming to judge the artifact that just landed. + const entitlements = await getEntitlements(orgId) + const usedBytes = Number.isFinite(entitlements.storageBytes) ? await storageBytes(orgId) : 0 + // Probe with a single byte: no room for that means no room for anything, so skip reading the body. + const full = storageCapError(entitlements, usedBytes, 1) + if (full) + return c.json(full, 402) + const uploaded = await streamArtifact(c, r.projectId, runId) if (!uploaded) return c.json({ error: 'file is required' }, 400) if ('tooLarge' in uploaded) return c.json({ error: 'Artifact too large' }, 413) + const overCap = storageCapError(entitlements, usedBytes, uploaded.size) + if (overCap) { + await storage.delete(uploaded.key).catch(error => logger.error({ error, key: uploaded.key }, 'over-cap artifact cleanup failed')) + return c.json(overCap, 402) + } + const t = uploaded.testKey ? await db.query.test.findFirst({ where: and(eq(test.runId, runId), eq(test.testKey, uploaded.testKey)), diff --git a/packages/server/src/router/billing.ts b/packages/server/src/router/billing.ts index 1fb3e84..2ee688d 100644 --- a/packages/server/src/router/billing.ts +++ b/packages/server/src/router/billing.ts @@ -1,5 +1,5 @@ import { getEntitlements, getSubscription } from '../billing/entitlements' -import { currentPeriodResults } from '../billing/usage' +import { currentPeriodResults, storageBytes } from '../billing/usage' import { orgProcedure, router } from '../trpc/index' // Infinity doesn't survive JSON: unlimited limits go over the wire as null. @@ -10,10 +10,11 @@ function finiteOrNull(n: number): number | null { export const billingRouter = router({ summary: orgProcedure.query(async ({ ctx }) => { const organizationId = ctx.organizationId - const [entitlements, sub, usedResults] = await Promise.all([ + const [entitlements, sub, usedResults, usedStorageBytes] = await Promise.all([ getEntitlements(organizationId), getSubscription(organizationId), currentPeriodResults(organizationId), + storageBytes(organizationId), ]) return { @@ -23,6 +24,8 @@ export const billingRouter = router({ retentionDays: finiteOrNull(entitlements.retentionDays), includedResults: finiteOrNull(entitlements.includedResults), usedResults, + storageBytes: finiteOrNull(entitlements.storageBytes), + usedStorageBytes, status: sub?.status ?? null, currentPeriodEnd: sub?.currentPeriodEnd?.toISOString() ?? null, cancelAtPeriodEnd: sub?.cancelAtPeriodEnd ?? false, diff --git a/packages/server/test/ingest-cap.test.ts b/packages/server/test/ingest-cap.test.ts index 9b47185..1c03f69 100644 --- a/packages/server/test/ingest-cap.test.ts +++ b/packages/server/test/ingest-cap.test.ts @@ -1,9 +1,9 @@ import type { Entitlements } from '../src/billing/entitlements' import { describe, expect, it } from 'vitest' -import { ingestCapError, quotaCrossing, quotaWarningText } from '../src/billing/entitlements' +import { formatBytes, ingestCapError, quotaCrossing, quotaWarningText, storageCapError } from '../src/billing/entitlements' -const free: Entitlements = { tier: 'free', includedResults: 2500, maxProjects: 1, retentionDays: 7, alerts: false } -const selfhost: Entitlements = { tier: 'selfhost', includedResults: Infinity, maxProjects: Infinity, retentionDays: Infinity, alerts: true } +const free: Entitlements = { tier: 'free', includedResults: 2500, maxProjects: 1, retentionDays: 7, storageBytes: 2 * 1024 ** 3, alerts: false } +const selfhost: Entitlements = { tier: 'selfhost', includedResults: Infinity, maxProjects: Infinity, retentionDays: Infinity, storageBytes: Infinity, alerts: true } describe('ingestCapError', () => { it('blocks free tier once the monthly result limit is reached', () => { @@ -23,12 +23,40 @@ describe('ingestCapError', () => { }) it('the result cap only applies to free, not paid tiers', () => { - const team: Entitlements = { tier: 'team', includedResults: 10_000, maxProjects: Infinity, retentionDays: 90, alerts: true } + const team: Entitlements = { tier: 'team', includedResults: 10_000, maxProjects: Infinity, retentionDays: 90, storageBytes: 50 * 1024 ** 3, alerts: true } // Team is metered (overage billed), never hard-capped on results. expect(ingestCapError(team, 999_999, false, 0)).toBeNull() }) }) +describe('storageCapError', () => { + const small: Entitlements = { ...free, storageBytes: 1000 } + + it('allows an artifact that still fits', () => { + expect(storageCapError(small, 900, 100)).toBeNull() + }) + + it('blocks the artifact that would cross the limit', () => { + expect(storageCapError(small, 900, 101)).toEqual({ error: expect.stringContaining('storage'), limit: 1000 }) + }) + + it('blocks everything once full, including a single byte', () => { + expect(storageCapError(small, 1000, 1)).not.toBeNull() + }) + + it('never caps an unlimited tier', () => { + expect(storageCapError(selfhost, 1e15, 1e9)).toBeNull() + }) +}) + +describe('formatBytes', () => { + it('scales to the largest fitting unit', () => { + expect(formatBytes(512)).toBe('512 B') + expect(formatBytes(2 * 1024 ** 3)).toBe('2 GB') + expect(formatBytes(1536)).toBe('1.5 KB') + }) +}) + describe('quotaCrossing', () => { const LIMIT = 2500 // near = 2000 diff --git a/packages/server/test/storage-cap.test.ts b/packages/server/test/storage-cap.test.ts new file mode 100644 index 0000000..e86d8df --- /dev/null +++ b/packages/server/test/storage-cap.test.ts @@ -0,0 +1,72 @@ +import type { Entitlements } from '../src/billing/entitlements' +import { existsSync } from 'node:fs' +import { readdir } from 'node:fs/promises' +import { resolve } from 'node:path' +import { beforeEach, describe, expect, it, vi } from 'vitest' + +// Tests run self-host (unlimited), so pin a tiny cap and re-import the graph with it. +const CAP = 5 +const capped: Entitlements = { tier: 'free', maxProjects: 1, retentionDays: 7, includedResults: 2500, storageBytes: CAP, alerts: false } + +vi.mock('../src/billing/entitlements', async importOriginal => ({ + ...(await importOriginal()), + getEntitlements: async () => capped, +})) +vi.resetModules() + +const { app } = await import('../src/app') +const { db } = await import('../src/db') +const { env } = await import('../src/lib/env') +const { createApiKey, createUser, ingest, resetDb } = await import('./helpers') + +beforeEach(resetDb) + +async function upload(runId: string, apiKey: string, bytes: number) { + const form = new FormData() + form.set('file', new File([new Uint8Array(bytes)], 'trace.zip')) + form.set('name', 'trace') + return app.request(`/api/v1/runs/${runId}/artifacts`, { + method: 'POST', + headers: { Authorization: `Bearer ${apiKey}` }, + body: form, + }) +} + +async function seedRun(): Promise<{ runId: string, apiKey: string }> { + const user = await createUser() + const apiKey = await createApiKey(user.id) + await ingest(apiKey) + return { runId: (await db.query.run.findMany())[0].id, apiKey } +} + +describe('artifact storage cap', () => { + it('accepts uploads up to the plan limit', async () => { + const { runId, apiKey } = await seedRun() + expect((await upload(runId, apiKey, 4)).status).toBe(201) + expect(await db.query.artifact.findMany()).toHaveLength(1) + }) + + it('rejects the upload that would cross the limit and keeps no orphan blob', async () => { + const { runId, apiKey } = await seedRun() + expect((await upload(runId, apiKey, 4)).status).toBe(201) + + const res = await upload(runId, apiKey, 4) // 4 + 4 > 5 + expect(res.status).toBe(402) + expect(await res.json()).toMatchObject({ limit: CAP, error: expect.stringContaining('storage') }) + + const stored = await db.query.artifact.findMany() + expect(stored).toHaveLength(1) // only the first upload survived + expect(existsSync(resolve(env.STORAGE_DIR, stored[0].storageKey))).toBe(true) + // The rejected bytes left nothing behind: the run prefix holds just the accepted blob. + const runDir = resolve(env.STORAGE_DIR, stored[0].storageKey, '..') + expect(await readdir(runDir)).toHaveLength(1) + }) + + it('rejects without reading the body once the workspace is full', async () => { + const { runId, apiKey } = await seedRun() + expect((await upload(runId, apiKey, 5)).status).toBe(201) + + expect((await upload(runId, apiKey, 1)).status).toBe(402) + expect(await db.query.artifact.findMany()).toHaveLength(1) + }) +}) diff --git a/packages/server/test/usage.test.ts b/packages/server/test/usage.test.ts index 95934c5..3cb265c 100644 --- a/packages/server/test/usage.test.ts +++ b/packages/server/test/usage.test.ts @@ -1,8 +1,8 @@ import { randomUUID } from 'node:crypto' import { beforeEach, describe, expect, it } from 'vitest' -import { currentPeriodResults, projectCount } from '../src/billing/usage' +import { currentPeriodResults, projectCount, storageBytes } from '../src/billing/usage' import { db } from '../src/db' -import { project, run, test as testRow } from '../src/db/schemas/index' +import { artifact, project, run, test as testRow } from '../src/db/schemas/index' import { createUser, ownedOrgId, resetDb } from './helpers' beforeEach(resetDb) @@ -89,3 +89,36 @@ describe('projectCount', () => { expect(await projectCount(await ownedOrgId(b.id))).toBe(0) }) }) + +describe('storageBytes', () => { + async function seedArtifact(userId: string, size: number): Promise { + await seedTests(userId, 1) + const r = (await db.query.run.findMany()).at(-1)! + await db.insert(artifact).values({ + id: randomUUID(), + projectId: r.projectId, + runId: r.id, + name: 'trace', + contentType: 'application/zip', + storageKey: `${r.projectId}/${r.id}/trace.zip`, + size, + }) + } + + it('sums the artifact sizes of the org', async () => { + const user = await createUser() + const org = await ownedOrgId(user.id) + expect(await storageBytes(org)).toBe(0) + + await seedArtifact(user.id, 300) + await seedArtifact(user.id, 700) + expect(await storageBytes(org)).toBe(1000) + }) + + it('ignores another org artifacts', async () => { + const a = await createUser('a@test.dev') + const b = await createUser('b@test.dev') + await seedArtifact(a.id, 500) + expect(await storageBytes(await ownedOrgId(b.id))).toBe(0) + }) +}) diff --git a/packages/web/src/pages/WorkspaceSettingsPage.vue b/packages/web/src/pages/WorkspaceSettingsPage.vue index c471ed5..fd33bd0 100644 --- a/packages/web/src/pages/WorkspaceSettingsPage.vue +++ b/packages/web/src/pages/WorkspaceSettingsPage.vue @@ -97,6 +97,29 @@ const overCap = computed(() => { return !!b && b.includedResults != null && b.usedResults >= b.includedResults }) +const storagePct = computed(() => { + const b = billing.value + if (!b || b.storageBytes == null) + return 0 + return Math.min(100, Math.round((b.usedStorageBytes / b.storageBytes) * 100)) +}) + +const overStorageCap = computed(() => { + const b = billing.value + return !!b && b.storageBytes != null && b.usedStorageBytes >= b.storageBytes +}) + +function formatBytes(bytes: number): string { + const units = ['B', 'KB', 'MB', 'GB', 'TB'] + let value = bytes + let unit = 0 + while (value >= 1024 && unit < units.length - 1) { + value /= 1024 + unit++ + } + return `${Math.round(value * 10) / 10} ${units[unit]}` +} + interface UpgradeOption { slug: 'team' | 'pro', label: string, featured: boolean, action: 'checkout' | 'portal' } const upgradeOptions = computed(() => { @@ -200,7 +223,7 @@ function fmtDate(d: Date | string | null | undefined): string { Plan - Your subscription and monthly test-result usage. + Your subscription, monthly test-result usage, and artifact storage.
@@ -239,6 +262,25 @@ function fmtDate(d: Date | string | null | undefined): string {

+
+
+ Artifact storage + + {{ formatBytes(billing.usedStorageBytes) }} + +
+
+
+
+

+ Storage limit reached - traces and videos are rejected until older runs expire. +

+
+