@namzu/sdk 34.2.0 → 35.0.0

package/src/store/run/conformance.ts

@@ -33,7 +33,7 @@
  *
  * ## Consuming it
  *
- * See `docs/sdk/runtime/checkpoint-store-conformance.md`. In short:
+ * In short:
  *
  * ```typescript
  * import { describe, expect, it } from 'your-runner'
@@ -62,7 +62,13 @@ import type { CostInfo } from '../../types/common/index.js'
 import type { IterationCheckpoint } from '../../types/hitl/index.js'
 import type { ProjectId, SessionId, TenantId } from '../../types/ids/index.js'
 import type { CheckpointStore } from '../../types/run/checkpoint-store.js'
-import { asCheckpointId, asProjectId, asRunId, asSessionId, asTenantId } from '../../utils/id.js'
+import {
+	asProjectId,
+	asRunId,
+	asSessionId,
+	asTenantId,
+	generateCheckpointId,
+} from '../../utils/id.js'
 import { claimRun, listDurableRuns, releaseRun } from './listing.js'
 
 /**
@@ -167,13 +173,13 @@ export interface CheckpointStoreConformanceOptions {
 
 /** The tenant/project/session every case addresses unless it needs a second. */
 const BINDING: CheckpointStoreBinding = {
-	tenantId: asTenantId('tnt_conformance'),
-	projectId: asProjectId('prj_conformance'),
-	sessionId: asSessionId('ses_conformance'),
+	tenantId: asTenantId('528cd907-2b3e-4524-b67d-bcf63bfd4b60'),
+	projectId: asProjectId('e74a687e-883b-43e9-a1e4-0c254b9d9eb1'),
+	sessionId: asSessionId('5abc3e06-d1a0-4036-84f4-903b11b1d781'),
 }
 
 /** A tenant no backend under test holds runs for. */
-const FOREIGN_TENANT = asTenantId('tnt_conformance_other')
+const FOREIGN_TENANT = asTenantId('4ddf9f0b-a8c3-4f53-853b-1c604fc1ba52')
 
 /** Fixed instant, so every expiry in a case is judged against one clock. */
 const NOW = 5_000_000
@@ -191,12 +197,9 @@ const NO_COST: CostInfo = {
 	unpricedTokens: 0,
 }
 
-let checkpointSeq = 0
-
 function checkpoint(runId: string): IterationCheckpoint {
-	checkpointSeq += 1
 	return {
-		id: asCheckpointId(`cp_conformance_${checkpointSeq}`),
+		id: generateCheckpointId(),
 		runId: asRunId(runId),
 		iteration: 1,
 		messages: [],
@@ -213,7 +216,7 @@ function checkpoint(runId: string): IterationCheckpoint {
 	}
 }
 
-function runScope(binding: CheckpointStoreBinding, runId = 'run_conformance_a') {
+function runScope(binding: CheckpointStoreBinding, runId = '5c84eacc-d287-4155-ac3b-336deb7b2d63') {
 	return { ...binding, runId: asRunId(runId) }
 }
 
@@ -427,7 +430,7 @@ export function defineCheckpointStoreConformance(options: CheckpointStoreConform
 						})
 						await store.writeCheckpoint(
 							runScope(BINDING),
-							checkpoint('run_conformance_a'),
+							checkpoint('5c84eacc-d287-4155-ac3b-336deb7b2d63'),
 							claim?.fence,
 						)
 						const written = await store.listCheckpoints(runScope(BINDING))
@@ -446,7 +449,10 @@ export function defineCheckpointStoreConformance(options: CheckpointStoreConform
 						// A host that adopts claims on one worker must not break the
 						// workers that have not adopted them yet. Refusing here would
 						// make the capability impossible to roll out incrementally.
-						await store.writeCheckpoint(runScope(BINDING), checkpoint('run_conformance_a'))
+						await store.writeCheckpoint(
+							runScope(BINDING),
+							checkpoint('5c84eacc-d287-4155-ac3b-336deb7b2d63'),
+						)
 						const written = await store.listCheckpoints(runScope(BINDING))
 						expect(written.length).toBe(1)
 					}),
@@ -476,7 +482,7 @@ export function defineCheckpointStoreConformance(options: CheckpointStoreConform
 							() =>
 								store.writeCheckpoint(
 									runScope(BINDING),
-									checkpoint('run_conformance_a'),
+									checkpoint('5c84eacc-d287-4155-ac3b-336deb7b2d63'),
 									first?.fence,
 								),
 							/no longer holds it/,
@@ -515,7 +521,7 @@ export function defineCheckpointStoreConformance(options: CheckpointStoreConform
 							() =>
 								store.writeCheckpoint(
 									runScope(BINDING),
-									checkpoint('run_conformance_a'),
+									checkpoint('5c84eacc-d287-4155-ac3b-336deb7b2d63'),
 									stale?.fence,
 								),
 							/no longer holds it/,
@@ -543,7 +549,7 @@ export function defineCheckpointStoreConformance(options: CheckpointStoreConform
 							() =>
 								store.writeCheckpoint(
 									runScope(BINDING),
-									checkpoint('run_conformance_a'),
+									checkpoint('5c84eacc-d287-4155-ac3b-336deb7b2d63'),
 									claim?.fence,
 								),
 							/no longer holds it/,
@@ -559,8 +565,8 @@ export function defineCheckpointStoreConformance(options: CheckpointStoreConform
 					'answers nothing for a tenant it does not hold',
 					withStore(async (store) => {
 						await store.writeCheckpoint(
-							runScope(BINDING, 'run_conformance_a'),
-							checkpoint('run_conformance_a'),
+							runScope(BINDING, '5c84eacc-d287-4155-ac3b-336deb7b2d63'),
+							checkpoint('5c84eacc-d287-4155-ac3b-336deb7b2d63'),
 						)
 
 						const foreign = await listDurableRuns(store, { tenantId: FOREIGN_TENANT })
@@ -575,13 +581,13 @@ export function defineCheckpointStoreConformance(options: CheckpointStoreConform
 					'answers nothing for a project it does not hold',
 					withStore(async (store) => {
 						await store.writeCheckpoint(
-							runScope(BINDING, 'run_conformance_a'),
-							checkpoint('run_conformance_a'),
+							runScope(BINDING, '5c84eacc-d287-4155-ac3b-336deb7b2d63'),
+							checkpoint('5c84eacc-d287-4155-ac3b-336deb7b2d63'),
 						)
 
 						const foreign = await listDurableRuns(store, {
 							tenantId: BINDING.tenantId,
-							projectId: asProjectId('prj_conformance_other'),
+							projectId: asProjectId('f89abf78-869a-40b1-8463-9a3051be514f'),
 						})
 						expect(foreign.entries.map((e) => e.runId)).toEqual([])
 					}),
@@ -610,18 +616,18 @@ export function defineCheckpointStoreConformance(options: CheckpointStoreConform
 					'rebuilds an addressable row for every run it holds',
 					withStore(async (store) => {
 						await store.writeCheckpoint(
-							runScope(BINDING, 'run_conformance_a'),
-							checkpoint('run_conformance_a'),
+							runScope(BINDING, '5c84eacc-d287-4155-ac3b-336deb7b2d63'),
+							checkpoint('5c84eacc-d287-4155-ac3b-336deb7b2d63'),
 						)
 						await store.writeCheckpoint(
-							runScope(BINDING, 'run_conformance_b'),
-							checkpoint('run_conformance_b'),
+							runScope(BINDING, '31e7abcd-79c7-4720-ad81-f57f89d1ec63'),
+							checkpoint('31e7abcd-79c7-4720-ad81-f57f89d1ec63'),
 						)
 
 						const page = await listDurableRuns(store, { tenantId: BINDING.tenantId })
 						expect(page.entries.map((e) => e.runId)).toEqual([
-							'run_conformance_a',
-							'run_conformance_b',
+							'31e7abcd-79c7-4720-ad81-f57f89d1ec63',
+							'5c84eacc-d287-4155-ac3b-336deb7b2d63',
 						])
 						// A row that cannot be turned back into a scope is a report,
 						// not a work queue.
@@ -641,24 +647,28 @@ export function defineCheckpointStoreConformance(options: CheckpointStoreConform
 						'keeps one tenant’s runs out of another tenant’s listing',
 						withStore(async (store) => {
 							await store.writeCheckpoint(
-								runScope(BINDING, 'run_conformance_a'),
-								checkpoint('run_conformance_a'),
+								runScope(BINDING, '5c84eacc-d287-4155-ac3b-336deb7b2d63'),
+								checkpoint('5c84eacc-d287-4155-ac3b-336deb7b2d63'),
 							)
 							await store.writeCheckpoint(
 								{
 									tenantId: FOREIGN_TENANT,
 									projectId: BINDING.projectId,
 									sessionId: BINDING.sessionId,
-									runId: asRunId('run_conformance_z'),
+									runId: asRunId('a8085004-143d-49b4-9beb-8c396f723803'),
 								},
-								checkpoint('run_conformance_z'),
+								checkpoint('a8085004-143d-49b4-9beb-8c396f723803'),
 							)
 
 							const mine = await listDurableRuns(store, { tenantId: BINDING.tenantId })
-							expect(mine.entries.map((e) => e.runId)).toEqual(['run_conformance_a'])
+							expect(mine.entries.map((e) => e.runId)).toEqual([
+								'5c84eacc-d287-4155-ac3b-336deb7b2d63',
+							])
 
 							const theirs = await listDurableRuns(store, { tenantId: FOREIGN_TENANT })
-							expect(theirs.entries.map((e) => e.runId)).toEqual(['run_conformance_z'])
+							expect(theirs.entries.map((e) => e.runId)).toEqual([
+								'a8085004-143d-49b4-9beb-8c396f723803',
+							])
 						}),
 					)
 				}
@@ -667,15 +677,19 @@ export function defineCheckpointStoreConformance(options: CheckpointStoreConform
 					it(
 						'shows the queue reader which runs nobody holds',
 						withStore(async (store) => {
-							for (const id of ['run_conformance_a', 'run_conformance_b', 'run_conformance_c'])
+							for (const id of [
+								'5c84eacc-d287-4155-ac3b-336deb7b2d63',
+								'31e7abcd-79c7-4720-ad81-f57f89d1ec63',
+								'727c095d-acad-4264-89f6-66cadc0bb644',
+							])
 								await store.writeCheckpoint(runScope(BINDING, id), checkpoint(id))
 
-							await claimRun(store, runScope(BINDING, 'run_conformance_a'), {
+							await claimRun(store, runScope(BINDING, '5c84eacc-d287-4155-ac3b-336deb7b2d63'), {
 								holder: 'w1',
 								ttlMs: 60_000,
 								now: NOW,
 							})
-							await claimRun(store, runScope(BINDING, 'run_conformance_b'), {
+							await claimRun(store, runScope(BINDING, '31e7abcd-79c7-4720-ad81-f57f89d1ec63'), {
 								holder: 'w1',
 								ttlMs: 1_000,
 								now: NOW,
@@ -691,8 +705,8 @@ export function defineCheckpointStoreConformance(options: CheckpointStoreConform
 							// forever — the failure the lease exists to prevent,
 							// reintroduced by the filter that reads it.
 							expect(free.entries.map((e) => e.runId)).toEqual([
-								'run_conformance_b',
-								'run_conformance_c',
+								'31e7abcd-79c7-4720-ad81-f57f89d1ec63',
+								'727c095d-acad-4264-89f6-66cadc0bb644',
 							])
 
 							const taken = await listDurableRuns(
@@ -700,7 +714,9 @@ export function defineCheckpointStoreConformance(options: CheckpointStoreConform
 								{ tenantId: BINDING.tenantId },
 								{ claimed: true, now: NOW + 2_000 },
 							)
-							expect(taken.entries.map((e) => e.runId)).toEqual(['run_conformance_a'])
+							expect(taken.entries.map((e) => e.runId)).toEqual([
+								'5c84eacc-d287-4155-ac3b-336deb7b2d63',
+							])
 							expect(taken.entries[0]?.claim?.holder).toBe('w1')
 							expect(taken.entries[0]?.claim?.expired).toBe(false)
 						}),

package/src/store/run/disk.ts

@@ -16,6 +16,7 @@ import type {
 	RunStore,
 } from '../../types/run/store.js'
 import { atomicWriteFile } from '../../utils/atomic-write.js'
+import { asCheckpointId, asRunId } from '../../utils/id.js'
 import { SCOPE_ATTRIBUTE } from '../../utils/log/types.js'
 import { type Logger, resolveLogger } from '../../utils/logger.js'
 import { defineSchema, migrate, stamp } from '../schema.js'
@@ -29,6 +30,13 @@ import { defineSchema, migrate, stamp } from '../schema.js'
  * the shape changes.
  */
 const SCHEMA = defineSchema({ kind: 'run-store', current: 1, migrations: {} })
+// Older readers must refuse ledger-bound checkpoints instead of dropping their
+// budget authority. Other run-store records retain their existing schema.
+const CHECKPOINT_SCHEMA = defineSchema({
+	kind: 'run-checkpoint',
+	current: 2,
+	migrations: { 1: (record) => record },
+})
 
 /**
  * One finished tool call, recovered from the transcript.
@@ -58,6 +66,8 @@ export class RunDiskStore implements RunStore {
 	}
 
 	async initRun(runId: string, parentRunId?: string): Promise<string> {
+		asRunId(runId)
+		if (parentRunId !== undefined) asRunId(parentRunId)
 		if (parentRunId) {
 			this.runDir = join(this.baseDir, parentRunId, 'children', runId)
 		} else {
@@ -195,6 +205,8 @@ export class RunDiskStore implements RunStore {
 			status: run.status,
 			metadata: run.metadata,
 			tokenUsage: run.tokenUsage,
+			budget: run.budget,
+			budgetBinding: run.budgetBinding,
 			currentIteration: run.currentIteration,
 			startedAt: run.startedAt,
 			endedAt: run.endedAt,
@@ -237,6 +249,7 @@ export class RunDiskStore implements RunStore {
 	}
 
 	async writeCheckpoint(checkpoint: IterationCheckpoint): Promise<void> {
+		asCheckpointId(checkpoint.id)
 		const dir = this.requireInit()
 		const cpDir = join(dir, 'checkpoints')
 		await mkdir(cpDir, { recursive: true })
@@ -247,10 +260,11 @@ export class RunDiskStore implements RunStore {
 		// as if it were the older shape, and the refusal that exists to
 		// prevent exactly that never fires. The stamp is what gives the
 		// migration chain something to hang on.
-		await atomicWriteJson(join(cpDir, `${checkpoint.id}.json`), stamp(SCHEMA, checkpoint))
+		await atomicWriteJson(join(cpDir, `${checkpoint.id}.json`), checkpoint, CHECKPOINT_SCHEMA)
 	}
 
 	async readCheckpoint(checkpointId: CheckpointId): Promise<IterationCheckpoint | null> {
+		asCheckpointId(checkpointId)
 		const dir = this.requireInit()
 		try {
 			const content = await readFile(join(dir, 'checkpoints', `${checkpointId}.json`), 'utf-8')
@@ -266,6 +280,7 @@ export class RunDiskStore implements RunStore {
 	}
 
 	async deleteCheckpoint(checkpointId: CheckpointId): Promise<void> {
+		asCheckpointId(checkpointId)
 		const dir = this.requireInit()
 		try {
 			await unlink(join(dir, 'checkpoints', `${checkpointId}.json`))
@@ -611,8 +626,8 @@ export async function readCheckpointsIn(runDir: string): Promise<IterationCheckp
 	return checkpoints.sort((a, b) => a.createdAt - b.createdAt)
 }
 
-async function atomicWriteJson(filePath: string, value: unknown): Promise<void> {
-	await atomicWriteFile(filePath, JSON.stringify(stamp(SCHEMA, value), null, 2))
+async function atomicWriteJson(filePath: string, value: unknown, schema = SCHEMA): Promise<void> {
+	await atomicWriteFile(filePath, JSON.stringify(stamp(schema, value), null, 2))
 }
 
 /**
@@ -656,7 +671,7 @@ function hasUsableBudgets(record: Partial<IterationCheckpoint>): boolean {
 }
 
 function parseCheckpoint(content: string, file: string): IterationCheckpoint {
-	const parsed = migrate<unknown>(SCHEMA, JSON.parse(content))
+	const parsed = migrate<unknown>(CHECKPOINT_SCHEMA, JSON.parse(content))
 	const record = parsed as Partial<IterationCheckpoint> | null
 
 	if (

package/src/store/run/token-budget-disk.ts

@@ -0,0 +1,247 @@
+import { mkdir, readFile, rename, unlink, writeFile } from 'node:fs/promises'
+import { dirname, join } from 'node:path'
+import {
+	TokenBudget,
+	type TokenBudgetSnapshot,
+	validateTokenBudgetSnapshot,
+} from '../../run/token-budget.js'
+import { DefaultPathBuilder, type PathBuilder } from '../../session/workspace/path-builder.js'
+import type { TokenUsage } from '../../types/common/index.js'
+import type { TokenBudgetScope, TokenBudgetStore } from '../../types/run/token-budget-store.js'
+import { temporaryPathFor } from '../../utils/atomic-write.js'
+import { asProjectId, asRunId, asSessionId, asTenantId } from '../../utils/id.js'
+
+interface BudgetRecord {
+	readonly version: 1
+	readonly scope: TokenBudgetScope
+	readonly snapshot: TokenBudgetSnapshot
+}
+
+export interface DiskTokenBudgetStoreConfig {
+	/** Root session's runs directory, shared by every account in a tree. */
+	readonly baseDir: string
+}
+
+/**
+ * Canonical ledger beside the root run, independent of checkpoint pruning.
+ *
+ * Writes from this instance are serialized and use private atomic sidecars.
+ * New directories are mode 0700 and files are mode 0600. This is a single
+ * writer store; atomic replacement does not provide a distributed lease.
+ */
+export class DiskTokenBudgetStore implements TokenBudgetStore {
+	private readonly writes = new Map<string, Promise<void>>()
+
+	constructor(private readonly config: DiskTokenBudgetStoreConfig) {}
+
+	private path(scope: TokenBudgetScope): string {
+		validateScope(scope)
+		return join(this.config.baseDir, scope.runId, 'token-budget.json')
+	}
+
+	async load(scope: TokenBudgetScope): Promise<TokenBudgetSnapshot | null> {
+		const path = this.path(scope)
+		await this.writes.get(path)
+		return this.read(path, scope)
+	}
+
+	async save(scope: TokenBudgetScope, snapshot: TokenBudgetSnapshot): Promise<void> {
+		const path = this.path(scope)
+		const checkedScope = validateScope(scope)
+		const checkedSnapshot = validateSnapshotScope(snapshot, checkedScope)
+		const previous = this.writes.get(path) ?? Promise.resolve()
+		const writing = previous
+			.catch(() => undefined)
+			.then(async () => {
+				const existing = await this.read(path, checkedScope)
+				if (existing) assertSameRoot(existing, checkedSnapshot)
+				await mkdir(dirname(path), { recursive: true, mode: 0o700 })
+				const temporary = temporaryPathFor(path)
+				const record: BudgetRecord = {
+					version: 1,
+					scope: checkedScope,
+					snapshot: checkedSnapshot,
+				}
+				try {
+					await writeFile(temporary, `${JSON.stringify(record)}\n`, {
+						encoding: 'utf8',
+						flag: 'wx',
+						mode: 0o600,
+					})
+					await rename(temporary, path)
+				} catch (error) {
+					await unlink(temporary).catch(() => undefined)
+					throw error
+				}
+			})
+		this.writes.set(path, writing)
+		try {
+			await writing
+		} finally {
+			if (this.writes.get(path) === writing) this.writes.delete(path)
+		}
+	}
+
+	private async read(path: string, scope: TokenBudgetScope): Promise<TokenBudgetSnapshot | null> {
+		let text: string
+		try {
+			text = await readFile(path, 'utf8')
+		} catch (error) {
+			if ((error as NodeJS.ErrnoException).code === 'ENOENT') return null
+			throw error
+		}
+		const raw: unknown = JSON.parse(text)
+		if (typeof raw !== 'object' || raw === null || Array.isArray(raw)) {
+			throw new Error('Invalid token budget record')
+		}
+		const record = raw as Partial<BudgetRecord>
+		if (record.version !== 1) throw new Error('Unsupported token budget record version')
+		if (!record.scope) throw new Error('Token budget record is missing its root scope')
+		const storedScope = validateScope(record.scope)
+		if (!sameScope(scope, storedScope)) throw new Error('Token budget root scope mismatch')
+		return validateSnapshotScope(record.snapshot, scope)
+	}
+}
+
+export interface OpenTokenBudgetOptions {
+	readonly store?: TokenBudgetStore
+	readonly scope: TokenBudgetScope
+	/** Required for a new root. Existing roots can reopen without changing their cap. */
+	readonly limit?: number
+	readonly pathBuilder?: PathBuilder
+	readonly workingDirectory?: string
+	readonly accountId?: string
+	/** A checkpoint binding must never create a fresh ledger if its record is gone. */
+	readonly requireExisting?: boolean
+}
+
+/** Open the latest authority; a checkpoint only selects an existing account. */
+export async function openTokenBudget(options: OpenTokenBudgetOptions): Promise<TokenBudget> {
+	const scope = validateScope(options.scope)
+	const pathBuilder =
+		options.pathBuilder ??
+		new DefaultPathBuilder(join(options.workingDirectory ?? process.cwd(), '.namzu'))
+	const store =
+		options.store ??
+		new DiskTokenBudgetStore({
+			baseDir: join(pathBuilder.sessionDir(scope.projectId, scope.sessionId), 'runs'),
+		})
+	const persistence = {
+		scope,
+		save: (snapshot: TokenBudgetSnapshot) => store.save(scope, snapshot),
+	}
+	const saved = await store.load(scope)
+	let budget: TokenBudget
+	if (saved === null) {
+		if (options.requireExisting || options.accountId !== undefined) {
+			throw new Error('The token budget ledger required by this run is missing')
+		}
+		if (options.limit === undefined)
+			throw new Error('A new token budget requires an explicit limit')
+		budget = TokenBudget.create(options.limit, scope.runId, persistence)
+	} else {
+		const snapshot = validateSnapshotScope(saved, scope)
+		if (options.limit !== undefined && rootLimit(snapshot) !== options.limit) {
+			throw new Error('Token budget root limit mismatch')
+		}
+		budget = TokenBudget.restore(snapshot, persistence)
+	}
+	const account = options.accountId === undefined ? budget : budget.account(options.accountId)
+	await budget.flush()
+	return account
+}
+
+function validateScope(scope: TokenBudgetScope): TokenBudgetScope {
+	return {
+		tenantId: asTenantId(scope.tenantId),
+		projectId: asProjectId(scope.projectId),
+		sessionId: asSessionId(scope.sessionId),
+		runId: asRunId(scope.runId),
+	}
+}
+
+function sameScope(left: TokenBudgetScope, right: TokenBudgetScope): boolean {
+	return (
+		left.tenantId === right.tenantId &&
+		left.projectId === right.projectId &&
+		left.sessionId === right.sessionId &&
+		left.runId === right.runId
+	)
+}
+
+function validateSnapshotScope(snapshot: unknown, scope: TokenBudgetScope): TokenBudgetSnapshot {
+	const checked = validateTokenBudgetSnapshot(snapshot)
+	if (checked.rootRunId !== scope.runId) throw new Error('Token budget root run mismatch')
+	return checked
+}
+
+function rootLimit(snapshot: TokenBudgetSnapshot): number {
+	const account = snapshot.accounts.find((account) => account.id === snapshot.rootAccountId)
+	if (!account) throw new Error('Token budget root account is missing')
+	return account.limit
+}
+
+function assertSameRoot(before: TokenBudgetSnapshot, after: TokenBudgetSnapshot): void {
+	if (before.rootAccountId !== after.rootAccountId) {
+		throw new Error('Cannot replace an existing token budget with a different root')
+	}
+	const reconciled = before.poisoned && !after.poisoned
+	if (reconciled && after.requests.length > 0) {
+		throw new Error('Cannot reopen a token budget with unresolved requests')
+	}
+	const accounts = new Map(after.accounts.map((account) => [account.id, account]))
+	for (const previous of before.accounts) {
+		const current = accounts.get(previous.id)
+		if (
+			!current ||
+			(previous.limit > 0 && (current.limit === 0 || current.limit > previous.limit)) ||
+			current.parentId !== previous.parentId ||
+			(previous.runId !== undefined && current.runId !== previous.runId) ||
+			(previous.settled && !current.settled)
+		) {
+			throw new Error('Cannot replace or remove an existing token budget account')
+		}
+		for (const [runId, spent] of Object.entries(previous.usage)) {
+			const recorded = current.usage[runId]
+			if (!usageAtLeast(recorded, spent)) {
+				throw new Error('Cannot reduce recorded token budget usage')
+			}
+		}
+	}
+	const receipts = new Map(after.completedRequests.map((request) => [request.id, request]))
+	for (const previous of before.completedRequests) {
+		const current = receipts.get(previous.id)
+		if (
+			!current ||
+			current.accountId !== previous.accountId ||
+			current.runId !== previous.runId ||
+			!usageAtLeast(current.usage, previous.usage)
+		) {
+			throw new Error('Cannot remove or replace a token budget completion receipt')
+		}
+	}
+	const active = new Map(after.requests.map((request) => [request.id, request]))
+	for (const previous of before.requests) {
+		const current = active.get(previous.id) ?? receipts.get(previous.id)
+		if (
+			!current ||
+			current.accountId !== previous.accountId ||
+			current.runId !== previous.runId ||
+			(reconciled && current.usage === undefined) ||
+			!usageAtLeast(current.usage, previous.usage)
+		) {
+			throw new Error(
+				'Cannot forget an unresolved token budget request without a completion receipt',
+			)
+		}
+	}
+}
+
+function usageAtLeast(current: TokenUsage | undefined, previous: TokenUsage | undefined): boolean {
+	if (previous === undefined) return true
+	if (current === undefined) return false
+	return Object.entries(previous).every(
+		([field, amount]) =>
+			typeof amount !== 'number' || (current[field as keyof TokenUsage] ?? 0) >= amount,
+	)
+}