diff --git a/packages/@emulators/autumn/src/__tests__/autumn-usage-rollup.test.ts b/packages/@emulators/autumn/src/__tests__/autumn-usage-rollup.test.ts new file mode 100644 index 00000000..ef0df4c3 --- /dev/null +++ b/packages/@emulators/autumn/src/__tests__/autumn-usage-rollup.test.ts @@ -0,0 +1,94 @@ +import { describe, it, expect } from "vitest"; +import { Store } from "@emulators/core"; + +import { getAutumnStore, seedFromConfig } from "../index.js"; +import { balanceForFeature, checkAndConsume, compactUsage } from "../serialize.js"; + +const DAY_MS = 86_400_000; +const BASE = "http://localhost:0"; + +// Usage events are rolled up so storage and balance reads stay bounded. Every +// balance Autumn can report must be the same before and after a rollup. +function setup() { + const store = new Store(); + seedFromConfig(store, BASE, { + plans: [ + { id: "daily", name: "Daily", items: [{ feature_id: "executions", included: 1000, reset: { interval: "day" } }] }, + { id: "flat", name: "Flat", items: [{ feature_id: "executions", included: 1000 }] }, + ], + customers: [ + { id: "org_window", subscriptions: [{ plan_id: "daily", status: "active" }] }, + { id: "org_watermark", subscriptions: [{ plan_id: "flat", status: "active" }] }, + ], + }); + return getAutumnStore(store); +} + +function track(as: ReturnType, customerId: string, count: number, at?: number) { + for (let i = 0; i < count; i++) { + const event = as.events.insert({ customer_id: customerId, feature_id: "executions", value: 1 }); + if (at !== undefined) as.events.update(event.id, { created_at: new Date(at).toISOString() }); + } +} + +function customer(as: ReturnType, id: string) { + return as.customers.findOneBy("customer_id", id)!; +} + +describe("autumn usage rollup", () => { + it("keeps the current reset window's usage when older usage is rolled up", () => { + const as = setup(); + const now = Date.now(); + const current = customer(as, "org_window"); + // The subscription started 36 hours ago, so its current daily window began 12 hours ago. + as.customers.update(current.id, { + subscriptions: current.subscriptions.map((sub) => ({ ...sub, started_at: now - 1.5 * DAY_MS, usage_epoch: 0 })), + }); + track(as, "org_window", 50, now - DAY_MS); + track(as, "org_window", 70, now - 60 * 60 * 1000); + const before = balanceForFeature(as, customer(as, "org_window"), "executions"); + expect(before?.usage).toBe(70); + + compactUsage(as, customer(as, "org_window"), "executions"); + + expect(as.events.findBy("customer_id", "org_window")).toHaveLength(2); + expect(balanceForFeature(as, customer(as, "org_window"), "executions")).toEqual(before); + }); + + it("keeps usage after a subscription's watermark separate from usage before it", () => { + const as = setup(); + track(as, "org_watermark", 40); + const watermark = Math.max(...as.events.findBy("customer_id", "org_watermark").map((event) => event.id)); + const current = customer(as, "org_watermark"); + as.customers.update(current.id, { + subscriptions: current.subscriptions.map((sub) => ({ ...sub, usage_epoch: watermark })), + }); + track(as, "org_watermark", 60); + const before = balanceForFeature(as, customer(as, "org_watermark"), "executions"); + expect(before?.usage).toBe(60); + + compactUsage(as, customer(as, "org_watermark"), "executions"); + + expect(as.events.findBy("customer_id", "org_watermark")).toHaveLength(2); + expect(balanceForFeature(as, customer(as, "org_watermark"), "executions")).toEqual(before); + }); + + it("leaves a small history untouched", () => { + const as = setup(); + track(as, "org_watermark", 10); + compactUsage(as, customer(as, "org_watermark"), "executions"); + expect(as.events.findBy("customer_id", "org_watermark")).toHaveLength(10); + }); + + it("keeps enforcing the limit while consumption is rolled up", () => { + const as = setup(); + const current = customer(as, "org_watermark"); + for (let i = 0; i < 1000; i++) { + expect(checkAndConsume(as, customer(as, "org_watermark"), "executions", 1, true).allowed).toBe(true); + } + const denied = checkAndConsume(as, current, "executions", 1, true); + expect(denied.allowed).toBe(false); + expect(denied.balance).toMatchObject({ usage: 1000, remaining: 0 }); + expect(as.events.findBy("customer_id", "org_watermark").length).toBeLessThanOrEqual(65); + }); +}); diff --git a/packages/@emulators/cloudflare/src/__tests__/worker.test.ts b/packages/@emulators/cloudflare/src/__tests__/worker.test.ts index 3266ec6e..46e37af3 100644 --- a/packages/@emulators/cloudflare/src/__tests__/worker.test.ts +++ b/packages/@emulators/cloudflare/src/__tests__/worker.test.ts @@ -569,6 +569,87 @@ describe("cloudflare durable object control plane", () => { expect([...storage.keys()].filter((key) => key.startsWith("ledger:entry:"))).toHaveLength(70); }); + it("writes only what an admission changed, never rescans storage, and keeps storage bounded", async () => { + const { state, storage, puts } = makeState(); + let lists = 0; + const list = state.storage.list.bind(state.storage); + state.storage.list = (options) => { + lists++; + return list(options); + }; + const headers = { + "x-emulator-service": "autumn", + "x-emulator-instance": "admission", + "x-emulator-base-url": "https://autumn.admission.emulators.dev", + authorization: "Bearer am_sk_test", + "content-type": "application/json", + }; + const call = (path: string, body: unknown) => + do1.fetch( + new Request(`https://autumn.admission.emulators.dev${path}`, { + method: "POST", + headers, + body: JSON.stringify(body), + }), + ); + const do1 = new EmulatorDurableObject(state, {}); + const seeded = await do1.fetch( + new Request("https://autumn.admission.emulators.dev/__seed", { + method: "POST", + headers, + body: JSON.stringify({ + strict: false, + autumn: { + plans: [{ id: "team", name: "Team", items: [{ feature_id: "executions", included: 500 }] }], + customers: [{ id: "org_synthetic", subscriptions: [{ plan_id: "team", status: "active" }] }], + }, + }), + }), + ); + expect(seeded.status).toBe(200); + const getOrCreate = () => call("/v1/customers.get_or_create", { customer_id: "org_synthetic" }); + const consume = () => + call("/v1/balances.check", { + customer_id: "org_synthetic", + feature_id: "executions", + required_balance: 1, + send_event: true, + }); + expect((await getOrCreate()).status).toBe(200); + + // A lookup that changes no billing state adds only its ledger entry and the ledger index. + puts.length = 0; + lists = 0; + expect((await getOrCreate()).status).toBe(200); + expect(puts.map((put) => put.key.replace(/req_\d+$/, "req_N")).sort()).toEqual([ + "ledger:entry:req_N", + "ledger:meta", + ]); + expect(lists).toBe(0); + + // Consumption writes the changed records only, and stored events stay bounded. + for (let i = 0; i < 300; i++) expect((await consume()).status).toBe(200); + expect(lists).toBe(0); + puts.length = 0; + const last = (await (await consume()).json()) as { balance: { usage: number; remaining: number } }; + expect(last.balance).toMatchObject({ usage: 301, remaining: 199 }); + expect(puts.length).toBeLessThan(10); + const events = [...storage.keys()].filter((key) => key.startsWith("snapshot:item:autumn.events:")); + expect(events.length).toBeLessThanOrEqual(65); + + // A fresh Durable Object over the same storage restores the same balance. + const do2 = new EmulatorDurableObject(state, {}); + const restored = await do2.fetch( + new Request("https://autumn.admission.emulators.dev/v1/customers.get_or_create", { + method: "POST", + headers, + body: JSON.stringify({ customer_id: "org_synthetic" }), + }), + ); + const customer = (await restored.json()) as { balances: { executions: { usage: number } } }; + expect(customer.balances.executions.usage).toBe(301); + }); + it("persists the request ledger across durable object eviction", async () => { const { state } = makeState(); const do1 = new EmulatorDurableObject(state, {});