Test changed-key persistence and Autumn usage rollup

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<typeof setup>, 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<typeof setup>, 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);
+  });
+});

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, {});