diff --git a/CHANGELOG.md b/CHANGELOG.md index ef73d3f..584e6a2 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -2,6 +2,12 @@ ## Unreleased +## 0.9.1 + +### Fixed + +- The usage ledger lost counts when several processes ran at once. Each process read `usage.json` once and then wrote its own copy over the file, so the last writer won: in a test, 2 processes making 200 updates each left 200, not 400. The spend that `/typesafe status` reported fell short of the real use, and the `usdPerDay`, `requestsPerDay`, and `inputTokensPerDay` caps did not stop at their limit. Each update now adds its count to the file's current content under a lock file next to it, and `today()`, `blocked()`, and `describe()` read the file, so the caps count every process. An update waits at most 200 milliseconds for the lock, takes over a lock older than 5 seconds, and keeps a count it could not write until its next update. `UsageLedger.flush()` adds any such count at once. Older versions still replace the whole file on each write and can erase the counts of processes that run this one. + ## 0.9.0 ### Added diff --git a/docs/api.md b/docs/api.md index 39b1d8f..ce6d756 100644 --- a/docs/api.md +++ b/docs/api.md @@ -85,7 +85,7 @@ Every item comes back as `{ ok: true, index, value }` or `{ ok: false, index, er `getUsage()` returns this client's session counters (`requestsStarted`, `requestsSucceeded`, `requestsFailed`, `inputTokens`, `outputTokens`, `estimatedUsd`). `getSpend()` adds today's persisted totals, the caps in force, and the cap currently reached. -Day caps live in `~/.pi/agent/pi-typesafe/usage.json` (owner-only, atomic, best-effort: an unwritable ledger never fails a request) and roll over at local midnight. +Day caps live in `~/.pi/agent/pi-typesafe/usage.json` (owner-only, atomic, best-effort: an unwritable ledger never fails a request) and roll over at local midnight. Every process that writes the ledger adds its counts to the file's current content under a lock file next to it (`usage.json.lock`), and `today()`, `blocked()`, and `describe()` read the file, so a day cap counts every process. An update waits at most 200 milliseconds for the lock; a lock older than 5 seconds belongs to a dead process and is taken over. A count that could not be written in time stays in the process and is added with its next update, or by `flush()`, which you can call before a process exits. Versions before 0.9.1 replace the whole file on each write and can still erase the counts of other processes. | Option | Environment | Bounds | | --- | --- | --- | diff --git a/package-lock.json b/package-lock.json index 6f9d9fd..8a46258 100644 --- a/package-lock.json +++ b/package-lock.json @@ -1,12 +1,12 @@ { "name": "pi-typesafe", - "version": "0.9.0", + "version": "0.9.1", "lockfileVersion": 3, "requires": true, "packages": { "": { "name": "pi-typesafe", - "version": "0.9.0", + "version": "0.9.1", "license": "MIT", "dependencies": { "@typesafe-ai/sdk": "^0.6.0" diff --git a/package.json b/package.json index 2342feb..7fa9bf5 100644 --- a/package.json +++ b/package.json @@ -1,6 +1,6 @@ { "name": "pi-typesafe", - "version": "0.9.0", + "version": "0.9.1", "description": "TypeSafe AI (Jev) decisions for Pi: batched Choice/Score/Noul evaluation tool, terminal playground, and a typed API other extensions build on.", "type": "module", "license": "MIT", diff --git a/src/usage.ts b/src/usage.ts index 8cf4d8b..106598c 100644 --- a/src/usage.ts +++ b/src/usage.ts @@ -1,4 +1,4 @@ -import { chmodSync, mkdirSync, readFileSync, renameSync, rmSync, writeFileSync } from "node:fs"; +import { chmodSync, closeSync, mkdirSync, openSync, readFileSync, renameSync, rmSync, statSync, unlinkSync, writeFileSync } from "node:fs"; import { dirname, join } from "node:path"; import { piTypesafeDir } from "./credentials.js"; @@ -104,9 +104,25 @@ function dayOf(value: unknown): DayRecord { return costed ? { ...totalsOf(value), costUsd, costedInputTokens } : totalsOf(value); } -function readDays(path: string): Record { +/** Longest an update waits for the ledger lock. A request that cannot get it keeps its delta and adds it with the next update. */ +export const LOCK_WAIT_MS = 200; +/** A lock this old belongs to a process that died while holding it; the next update takes it over. */ +export const LOCK_STALE_MS = 5_000; +const LOCK_RETRY_MS = 2; + +/** + * The days stored in the file: `{}` for a missing or corrupt file, `undefined` when the file exists but cannot be read, + * so the caller never replaces a ledger it could not see. + */ +function readDays(path: string): Record | undefined { + let text: string; try { - const parsed: unknown = JSON.parse(readFileSync(path, "utf8")); + text = readFileSync(path, "utf8"); + } catch (error) { + return (error as NodeJS.ErrnoException).code === "ENOENT" ? {} : undefined; + } + try { + const parsed: unknown = JSON.parse(text); const days = parsed && typeof parsed === "object" ? (parsed as { days?: unknown }).days : undefined; if (!days || typeof days !== "object" || Array.isArray(days)) return {}; const result: Record = {}; @@ -115,7 +131,7 @@ function readDays(path: string): Record { } return result; } catch { - // A missing, unreadable, or corrupt ledger restarts today's count; it never blocks a request. + // A corrupt ledger restarts the count; it never blocks a request. return {}; } } @@ -128,8 +144,23 @@ function keepRecent(days: Record, today: string): Record): void { +/** `base` plus `delta`. Cost fields appear once either side has them, so a day from an older file stays as it was until a cost is recorded. */ +function addDay(base: DayRecord, delta: DayRecord): DayRecord { + return { + requestsStarted: base.requestsStarted + delta.requestsStarted, + requestsSucceeded: base.requestsSucceeded + delta.requestsSucceeded, + requestsFailed: base.requestsFailed + delta.requestsFailed, + inputTokens: base.inputTokens + delta.inputTokens, + outputTokens: base.outputTokens + delta.outputTokens, + ...(delta.costUsd === undefined && base.costUsd === undefined ? {} : { + costUsd: Math.round(((base.costUsd ?? 0) + (delta.costUsd ?? 0)) * 1e12) / 1e12, + costedInputTokens: (base.costedInputTokens ?? 0) + (delta.costedInputTokens ?? 0), + }), + }; +} + +/** Owner-only and atomic. Returns false when the file could not be written: a ledger this process cannot write never fails a request. */ +function writeDays(path: string, days: Record): boolean { const file: LedgerFile = { version: USAGE_VERSION, days }; const temporary = `${path}.${process.pid}.tmp`; try { @@ -137,8 +168,47 @@ function writeDays(path: string, days: Record): void { writeFileSync(temporary, `${JSON.stringify(file, null, 2)}\n`, { mode: 0o600, flag: "w" }); chmodSync(temporary, 0o600); renameSync(temporary, path); + return true; } catch { try { rmSync(temporary, { force: true }); } catch { /* best-effort cleanup only */ } + return false; + } +} + +function pause(ms: number): void { + Atomics.wait(new Int32Array(new SharedArrayBuffer(4)), 0, 0, ms); +} + +/** + * Take the lock file next to the ledger, created exclusively. Waits at most LOCK_WAIT_MS; a lock older than + * LOCK_STALE_MS is taken over. Returns the release function, or undefined when the lock could not be taken. + */ +function takeLock(lock: string): (() => void) | undefined { + const deadline = Date.now() + LOCK_WAIT_MS; + try { + mkdirSync(dirname(lock), { recursive: true, mode: 0o700 }); + for (;;) { + try { + closeSync(openSync(lock, "wx", 0o600)); + return () => { try { unlinkSync(lock); } catch { /* already gone */ } }; + } catch (error) { + if ((error as NodeJS.ErrnoException).code !== "EEXIST") return undefined; + } + try { + if (Date.now() - statSync(lock).mtimeMs > LOCK_STALE_MS) { + unlinkSync(lock); + continue; + } + } catch (error) { + // The holder released it, or another process took it over, between our open and our stat. + if ((error as NodeJS.ErrnoException).code !== "ENOENT") return undefined; + continue; + } + if (Date.now() >= deadline) return undefined; + pause(LOCK_RETRY_MS); + } + } catch { + return undefined; } } @@ -185,9 +255,11 @@ export interface UsageLedgerOptions { } /** - * One day of persisted request, token, and cost totals, plus the caps that stop the next request. The in-memory copy is - * authoritative for this process; the file is the cross-process, across-restart record. Reads are defensive, writes are - * atomic and best-effort, and a day rolls over on the local date, so a long eval cannot accumulate forever unnoticed. + * One day of persisted request, token, and cost totals, plus the caps that stop the next request. The file is the + * record for every process: each update adds its delta to the file's current content under a lock, and `today()`, + * `blocked()`, and `describe()` read the file, so a cap counts every process that writes it. A delta that cannot be + * written in time (the lock is held, or the file is unwritable) stays in this process and is added with its next + * update. Reads are defensive, writes are atomic and best-effort, and a day rolls over on the local date. */ export interface UsageLedger { readonly path: string; @@ -203,15 +275,17 @@ export interface UsageLedger { blocked(caps: SpendCaps): BlockedCap | undefined; /** One line for status output: today's requests, tokens, and cost, with the caps that apply. */ describe(caps?: SpendCaps): string; + /** Add to the file any delta this process still holds. Updates do this themselves; call it before the process exits. */ + flush?(): void; } export function openUsageLedger(options: UsageLedgerOptions = {}): UsageLedger { const path = options.path ?? usagePath(); + const lock = `${path}.lock`; const now = options.now ?? (() => new Date()); const usdPerMTok = options.usdPerMTok !== undefined && Number.isFinite(options.usdPerMTok) && options.usdPerMTok >= 0 ? options.usdPerMTok : DEFAULT_USD_PER_MTOK; - let day = localDay(now()); - let days = keepRecent(readDays(path), day); - let totals = days[day] as DayRecord; + /** Deltas this process has recorded and the file does not hold yet, by day. */ + let pending: Record = {}; /** Recorded costs plus an estimate for the input tokens that have none. */ const costOf = (value: DayRecord): number => @@ -227,39 +301,38 @@ export function openUsageLedger(options: UsageLedgerOptions = {}): UsageLedger { estimatedUsd: costOf(value), }); - const save = () => { - days = keepRecent({ ...days, [day]: totals }, day); - writeDays(path, days); + /** The file's current content for today, plus whatever this process still holds. */ + const current = (): { day: string; totals: DayRecord } => { + const day = localDay(now()); + const stored = readDays(path)?.[day] ?? emptyTotals(); + return { day, totals: pending[day] === undefined ? stored : addDay(stored, pending[day] as DayRecord) }; }; - const roll = () => { - const current = localDay(now()); - if (current === day) return; - day = current; - totals = days[day] ?? emptyTotals(); - days = keepRecent(days, day); + const flush = (): void => { + const owed = Object.entries(pending); + if (owed.length === 0) return; + const release = takeLock(lock); + if (!release) return; + try { + const stored = readDays(path); + if (!stored) return; + for (const [name, delta] of owed) stored[name] = addDay(stored[name] ?? emptyTotals(), delta); + if (writeDays(path, keepRecent(stored, localDay(now())))) pending = {}; + } finally { + release(); + } }; const add = (delta: Partial & { costUsd?: number; costedInputTokens?: number }) => { - roll(); - totals = { - requestsStarted: totals.requestsStarted + (delta.requestsStarted ?? 0), - requestsSucceeded: totals.requestsSucceeded + (delta.requestsSucceeded ?? 0), - requestsFailed: totals.requestsFailed + (delta.requestsFailed ?? 0), - inputTokens: totals.inputTokens + (delta.inputTokens ?? 0), - outputTokens: totals.outputTokens + (delta.outputTokens ?? 0), - ...(delta.costUsd === undefined && totals.costUsd === undefined ? {} : { - costUsd: Math.round(((totals.costUsd ?? 0) + (delta.costUsd ?? 0)) * 1e12) / 1e12, - costedInputTokens: (totals.costedInputTokens ?? 0) + (delta.costedInputTokens ?? 0), - }), - }; - save(); + const day = localDay(now()); + pending[day] = addDay(pending[day] ?? emptyTotals(), { ...emptyTotals(), ...delta }); + flush(); }; return { path, usdPerMTok, - today: () => { roll(); return report(totals, day); }, + today: () => { const { day, totals } = current(); return report(totals, day); }, recordStart: () => add({ requestsStarted: 1 }), recordSuccess: (inputTokens, outputTokens, usd) => { const input = Number.isSafeInteger(inputTokens) && inputTokens > 0 ? inputTokens : 0; @@ -271,8 +344,9 @@ export function openUsageLedger(options: UsageLedgerOptions = {}): UsageLedger { }); }, recordFailure: () => add({ requestsFailed: 1 }), + flush, blocked: (caps: SpendCaps) => { - roll(); + const { day, totals } = current(); const checks: Array<[BlockedCap["cap"], number | undefined, number]> = [ ["requestsPerDay", caps.maxRequestsPerDay, totals.requestsStarted], ["inputTokensPerDay", caps.maxInputTokensPerDay, totals.inputTokens], @@ -284,14 +358,14 @@ export function openUsageLedger(options: UsageLedgerOptions = {}): UsageLedger { return undefined; }, describe: (caps: SpendCaps = {}) => { - roll(); - const current = report(totals, day); + const { day, totals } = current(); + const today = report(totals, day); const limits = [ - caps.maxRequestsPerDay === undefined ? undefined : `${current.requestsStarted}/${caps.maxRequestsPerDay} requests`, - caps.maxInputTokensPerDay === undefined ? undefined : `${current.inputTokens}/${caps.maxInputTokensPerDay} input tokens`, - caps.maxUsdPerDay === undefined ? undefined : `$${current.estimatedUsd.toFixed(4)}/$${caps.maxUsdPerDay.toFixed(2)}`, + caps.maxRequestsPerDay === undefined ? undefined : `${today.requestsStarted}/${caps.maxRequestsPerDay} requests`, + caps.maxInputTokensPerDay === undefined ? undefined : `${today.inputTokens}/${caps.maxInputTokensPerDay} input tokens`, + caps.maxUsdPerDay === undefined ? undefined : `$${today.estimatedUsd.toFixed(4)}/$${caps.maxUsdPerDay.toFixed(2)}`, ].filter((part): part is string => part !== undefined); - return `${current.requestsStarted} requests today (${current.requestsSucceeded} ok, ${current.requestsFailed} failed), ${current.inputTokens} input / ${current.outputTokens} output tokens, ~$${current.estimatedUsd.toFixed(4)}${limits.length ? `; caps ${limits.join(", ")}` : "; no daily cap"}`; + return `${today.requestsStarted} requests today (${today.requestsSucceeded} ok, ${today.requestsFailed} failed), ${today.inputTokens} input / ${today.outputTokens} output tokens, ~$${today.estimatedUsd.toFixed(4)}${limits.length ? `; caps ${limits.join(", ")}` : "; no daily cap"}`; }, }; } diff --git a/tests/helpers/usage-writer.ts b/tests/helpers/usage-writer.ts new file mode 100644 index 0000000..f1efd3b --- /dev/null +++ b/tests/helpers/usage-writer.ts @@ -0,0 +1,12 @@ +// Child process for the usage ledger tests: writes to one ledger file, then flushes anything the lock made it hold back. +import { openUsageLedger } from "../../src/usage.js"; + +const [path, iterations, startAt, usd] = process.argv.slice(2) as [string, string, string, string]; +const ledger = openUsageLedger({ path }); +// Start together, so the writers overlap. +while (Date.now() < Number(startAt)) { /* wait for the start time */ } +for (let index = 0; index < Number(iterations); index++) { + ledger.recordStart(); + ledger.recordSuccess(10, 3, Number(usd)); +} +ledger.flush?.(); diff --git a/tests/usage.test.ts b/tests/usage.test.ts index 7d54dc5..392473e 100644 --- a/tests/usage.test.ts +++ b/tests/usage.test.ts @@ -1,10 +1,11 @@ import assert from "node:assert/strict"; -import { mkdtempSync, readFileSync, statSync, writeFileSync } from "node:fs"; +import { spawn } from "node:child_process"; +import { existsSync, mkdtempSync, readFileSync, rmSync, statSync, utimesSync, writeFileSync } from "node:fs"; import { tmpdir } from "node:os"; import { join } from "node:path"; import { test } from "node:test"; import { - DEFAULT_USD_PER_MTOK, capsFromEnvironment, estimateUsd, localDay, mergeCaps, openUsageLedger, usagePath, + DEFAULT_USD_PER_MTOK, LOCK_STALE_MS, LOCK_WAIT_MS, capsFromEnvironment, estimateUsd, localDay, mergeCaps, openUsageLedger, usagePath, } from "../src/usage.js"; const workspace = mkdtempSync(join(tmpdir(), "pi-typesafe-usage-")); @@ -104,3 +105,119 @@ test("the default usage path sits with the key store", () => { if (saved === undefined) delete process.env.PI_CODING_AGENT_DIR; else process.env.PI_CODING_AGENT_DIR = saved; } }); + +const writer = join(import.meta.dirname, "helpers", "usage-writer.ts"); + +/** Run one writer process to completion; they all start at `startAt`. */ +const runWriter = (path: string, iterations: number, startAt: number, usd: number) => + new Promise((resolve, reject) => { + const child = spawn(process.execPath, ["--import", "tsx", writer, path, String(iterations), String(startAt), String(usd)], { stdio: "inherit" }); + child.on("error", reject); + child.on("exit", (code) => (code === 0 ? resolve() : reject(new Error(`writer exited with ${code}`)))); + }); + +test("processes writing one ledger at the same time lose no count", async () => { + const path = ledgerPath("concurrent"); + const processes = 4; + const iterations = 200; // Each iteration records a start and a success: 2 updates, so 1,600 updates in all. + const startAt = Date.now() + 1_500; + await Promise.all(Array.from({ length: processes }, () => runWriter(path, iterations, startAt, 0.000125))); + const today = openUsageLedger({ path }).today(); + const total = processes * iterations; + assert.equal(today.requestsStarted, total); + assert.equal(today.requestsSucceeded, total); + assert.equal(today.inputTokens, total * 10); + assert.equal(today.outputTokens, total * 3); + assert.ok(Math.abs(today.estimatedUsd - total * 0.000125) < 1e-9, `costUsd ${today.estimatedUsd}`); + const file = JSON.parse(readFileSync(path, "utf8")); + assert.equal(file.version, 1); + assert.equal(file.days[localDay()].costedInputTokens, total * 10); + assert.ok(Math.abs(file.days[localDay()].costUsd - total * 0.000125) < 1e-9); + assert.equal(statSync(path).mode & 0o777, 0o600); + assert.equal(existsSync(`${path}.lock`), false); +}); + +test("a cap sees the usage another process recorded", async () => { + const path = ledgerPath("cross-process-cap"); + const ledger = openUsageLedger({ path }); + assert.equal(ledger.blocked({ maxRequestsPerDay: 3 }), undefined); + await runWriter(path, 3, 0, 0.0001); + // The ledger was opened before the other process wrote: it still reads the file, not a copy. + assert.equal(ledger.today().requestsStarted, 3); + assert.equal(ledger.blocked({ maxRequestsPerDay: 3 })?.cap, "requestsPerDay"); + assert.equal(ledger.blocked({ maxInputTokensPerDay: 30 })?.cap, "inputTokensPerDay"); + assert.equal(ledger.blocked({ maxUsdPerDay: 0.0003 })?.cap, "usdPerDay"); + // This process adds to what the other one wrote instead of replacing it. + ledger.recordStart(); + assert.equal(openUsageLedger({ path }).today().requestsStarted, 4); +}); + +test("a lock left by a dead process is taken over", () => { + const path = ledgerPath("stale-lock"); + writeFileSync(`${path}.lock`, ""); + const old = new Date(Date.now() - LOCK_STALE_MS - 5_000); + utimesSync(`${path}.lock`, old, old); + const ledger = openUsageLedger({ path }); + const started = Date.now(); + ledger.recordStart(); + assert.ok(Date.now() - started < LOCK_WAIT_MS, "a stale lock must not delay the update"); + assert.equal(JSON.parse(readFileSync(path, "utf8")).days[localDay()].requestsStarted, 1); + assert.equal(existsSync(`${path}.lock`), false); +}); + +test("a held lock delays an update at most the bound, and the delta is added later", () => { + const path = ledgerPath("held-lock"); + writeFileSync(`${path}.lock`, ""); + const ledger = openUsageLedger({ path }); + const started = Date.now(); + ledger.recordStart(); + const waited = Date.now() - started; + assert.ok(waited >= LOCK_WAIT_MS - 20, `waited ${waited} ms`); + assert.ok(waited < LOCK_WAIT_MS + 300, `waited ${waited} ms`); + // The holder's lock is untouched, the file has not been written, and this process still counts the request. + assert.equal(existsSync(`${path}.lock`), true); + assert.equal(existsSync(path), false); + assert.equal(ledger.today().requestsStarted, 1); + assert.equal(ledger.blocked({ maxRequestsPerDay: 1 })?.cap, "requestsPerDay"); + rmSync(`${path}.lock`); + ledger.recordSuccess(40, 5, 0.001); + const day = JSON.parse(readFileSync(path, "utf8")).days[localDay()]; + assert.equal(day.requestsStarted, 1); + assert.equal(day.requestsSucceeded, 1); + assert.equal(day.inputTokens, 40); + assert.equal(day.costUsd, 0.001); + assert.equal(existsSync(`${path}.lock`), false); +}); + +test("a held-back delta lands on the day it was recorded", () => { + let day = 8; + const path = ledgerPath("held-rollover"); + const ledger = openUsageLedger({ path, now: () => at(day) }); + writeFileSync(`${path}.lock`, ""); + ledger.recordStart(); + rmSync(`${path}.lock`); + day = 9; + ledger.recordFailure(); + const file = JSON.parse(readFileSync(path, "utf8")); + assert.equal(file.days["2026-01-08"].requestsStarted, 1); + assert.equal(file.days["2026-01-09"].requestsFailed, 1); + assert.equal(file.days["2026-01-09"].requestsStarted, 0); +}); + +test("a ledger file written before costs were recorded is still read and extended", () => { + const path = ledgerPath("old-file"); + writeFileSync(path, JSON.stringify({ version: 1, days: { "2026-01-10": { requestsStarted: 2, requestsSucceeded: 2, requestsFailed: 0, inputTokens: 1_000_000, outputTokens: 5 } } })); + const ledger = openUsageLedger({ path, now: at.bind(null, 10), usdPerMTok: 1 }); + assert.equal(ledger.today().inputTokens, 1_000_000); + assert.equal(ledger.today().estimatedUsd, 1); + ledger.recordStart(); + assert.equal(JSON.parse(readFileSync(path, "utf8")).days["2026-01-10"].costUsd, undefined); + ledger.recordSuccess(500_000, 0, 0.25); + const day = JSON.parse(readFileSync(path, "utf8")).days["2026-01-10"]; + assert.equal(day.requestsStarted, 3); + assert.equal(day.inputTokens, 1_500_000); + assert.equal(day.costUsd, 0.25); + assert.equal(day.costedInputTokens, 500_000); + // 1,000,000 old tokens at the ledger's rate plus the recorded 0.25. + assert.equal(ledger.today().estimatedUsd, 1.25); +});