72 lines
2.4 KiB
TypeScript
Executable File
72 lines
2.4 KiB
TypeScript
Executable File
// Decision-log item 9 closure: request budgets survive a gateway restart.
|
|
// Usage is an atomic KV counter hydrated into the VirtualKeyManager at boot.
|
|
|
|
import { assertEquals } from "@std/assert";
|
|
import { ConfigService } from "../../packages/config/src/service.ts";
|
|
import {
|
|
type VirtualKey,
|
|
VirtualKeyManager,
|
|
} from "../../packages/governance/src/virtual_keys.ts";
|
|
|
|
function budgetKey(): VirtualKey {
|
|
return {
|
|
id: "k-budget",
|
|
name: "budgeted",
|
|
token: "vk-durable-budget-token",
|
|
enabled: true,
|
|
budget: { maxRequests: 2 },
|
|
usedRequests: 0,
|
|
usedCostMicroUsd: 0,
|
|
};
|
|
}
|
|
|
|
/** Mirrors the createDefaultContext wiring, with awaitable persistence. */
|
|
function wire(config: ConfigService, keys: VirtualKey[]) {
|
|
const manager = new VirtualKeyManager(keys);
|
|
let lastWrite: Promise<void> = Promise.resolve();
|
|
manager.onUsage((id) => {
|
|
lastWrite = config.addUsage(id);
|
|
});
|
|
return { manager, flush: () => lastWrite };
|
|
}
|
|
|
|
Deno.test("budget usage persists across a service restart", async () => {
|
|
const dir = await Deno.makeTempDir();
|
|
const kvPath = `${dir}/frosty-test.kv`;
|
|
|
|
// --- first process lifetime -------------------------------------------
|
|
let config = await ConfigService.open(kvPath);
|
|
await config.upsertVirtualKey(budgetKey());
|
|
let { manager, flush } = wire(config, await config.listVirtualKeys());
|
|
manager.hydrateUsage(await config.loadUsage());
|
|
|
|
const first = manager.check("vk-durable-budget-token");
|
|
assertEquals(first.ok, true);
|
|
manager.recordUsage("k-budget");
|
|
await flush();
|
|
const second = manager.check("vk-durable-budget-token");
|
|
assertEquals(second.ok, true);
|
|
manager.recordUsage("k-budget");
|
|
await flush();
|
|
config.close();
|
|
|
|
// --- restart ------------------------------------------------------------
|
|
config = await ConfigService.open(kvPath);
|
|
assertEquals(await config.loadUsage(), { "k-budget": 2 });
|
|
({ manager, flush } = wire(config, await config.listVirtualKeys()));
|
|
manager.hydrateUsage(await config.loadUsage());
|
|
|
|
// The 2-request budget is exhausted even though the process restarted.
|
|
const third = manager.check("vk-durable-budget-token");
|
|
assertEquals(third.ok, false);
|
|
if (!third.ok) {
|
|
assertEquals(third.status, 402);
|
|
assertEquals(third.code, "budget_exhausted");
|
|
}
|
|
|
|
// Deleting the key clears its usage counter too.
|
|
await config.deleteVirtualKey("k-budget");
|
|
assertEquals(await config.loadUsage(), {});
|
|
config.close();
|
|
});
|