diff --git a/.changeset/turn-budgets-off-event-loop.md b/.changeset/turn-budgets-off-event-loop.md new file mode 100644 index 0000000000..596b8ade8b --- /dev/null +++ b/.changeset/turn-budgets-off-event-loop.md @@ -0,0 +1,7 @@ +--- +"@schlessera/brain-ui-server": patch +"@schlessera/brain-ui-sdk": patch +--- + +Changed: the hosted server commits every turn-budget ledger write (admission, session binding, completed-response and relayed-provider usage, warnings and the SDK backstop) on its SQLite write worker at `synchronous=FULL`, in submission order, instead of on the event loop (#2850). A relayed provider request is admitted only once every earlier budget write has committed. +Changed: the provider relay awaits a promise returned from its `admit` hook. diff --git a/docs/decisions/governed-writable-workers.md b/docs/decisions/governed-writable-workers.md index e0d283493f..112e845840 100644 --- a/docs/decisions/governed-writable-workers.md +++ b/docs/decisions/governed-writable-workers.md @@ -569,7 +569,7 @@ the key and forwards it as the turn's own egress (G5). A key goes only to its built-in first-party origin, or to an exact origin the operator lists for that key in `BRAIN_UI_CLI_PROVIDER_ORIGINS`, under the 2026-10-23 ruling on #1673; the backends' origin rulings grant the CLI nothing by analogy -(`startProviderRelay`, `packages/ui-sdk/src/server/provider-relay.ts:252-257`). +(`startProviderRelay`, `packages/ui-sdk/src/server/provider-relay.ts:255-260`). The inference relay's placeholder is not a key: the launcher names a key that holds it as withheld, so a relayed Claude API-key turn's `anthropic-haiku` reaches the host's key instead of failing authentication. Video and the `jev` @@ -1443,7 +1443,7 @@ unmet rather than relaxed. **Why.** Every governed turn queues an automatic index refresh after it (`host.governedTurns?.refresh(host.workerLock)`, -`packages/ui-server/src/ws/run-session.ts:909`). The refresh usually finishes +`packages/ui-server/src/ws/run-session.ts:930`). The refresh usually finishes while the next turn runs, so that turn's finish observation found a rewritten `brain.db` and copied a whole new version of it into retained objects, inside the observation's worker hold. #2831 measured 36 versions of a ~120 MB index in @@ -1568,7 +1568,7 @@ covers everything. non-video completions, Anthropic completions, and OpenAI and Gemini images. Each upstream attempt is recorded once, with the units the provider reported, and counts toward the turn, conversation and UTC-day scopes -(`observeProvider`, `packages/ui-server/src/brain/turn-budgets.ts:236-253`). +(`observeProvider`, `packages/ui-server/src/brain/turn-budgets.ts:195-209`). The receipt's coverage names the relayed providers. Pricing follows the [cost-tracking rules](cost-tracking.md#relayed-cli-provider-requests-1673): unpriced tokens count toward the token stops, and missing usage is unknown and @@ -1576,7 +1576,7 @@ never blocks by itself. An attempt is recorded as unknown spend just before it leaves and completed when it settles, so a crash leaves it disclosed. Before each request, and again just before dispatch, the relay asks whether any scope is blocked and refuses with `budget_exceeded` if one is -(`providerAdmission`, `packages/ui-server/src/brain/turn-budgets.ts:224-228`). +(`providerAdmission`, `packages/ui-server/src/brain/turn-budgets.ts:326-330`). A request that crosses a block does not stop the turn by itself, under the ruling: each further relayed request is refused, and the agent's next completed response, which sees the same scopes, stops the turn as any usage @@ -1616,7 +1616,7 @@ rows remain historical inputs with explicit legacy attribution. Completed responses contribute once to their server-observed UTC day; legacy history uses its recorded start day. A terminal rollup never adds the live ledger a second time. The concrete ledger is (`createTurnBudgets`, -`packages/ui-server/src/brain/turn-budgets.ts:64-252`). +`packages/ui-server/src/brain/turn-budgets.ts:228-350`). Known priced lower bounds survive an unpriced model joining; unknown usage alone never blocks. A token fallback remains enforceable. diff --git a/docs/decisions/http-api-boundary.md b/docs/decisions/http-api-boundary.md index 8d96b28f2a..dd8ae50547 100644 --- a/docs/decisions/http-api-boundary.md +++ b/docs/decisions/http-api-boundary.md @@ -46,11 +46,11 @@ explicit supported rows, not silently excluded by this distinction. The source audit on 2026-09-30 used `e0875269b21b707482545be9e94901acb897893b`. The public health mount precedes auth -(`createHealthRoutes`, `packages/ui-server/src/app.ts:856`), the API guard follows -public ceremonies (`authGuard`, `packages/ui-server/src/app.ts:877`), and the -WebSocket admission is separate (`"/ws"`, `packages/ui-server/src/app.ts:1066`). +(`createHealthRoutes`, `packages/ui-server/src/app.ts:857`), the API guard follows +public ceremonies (`authGuard`, `packages/ui-server/src/app.ts:878`), and the +WebSocket admission is separate (`"/ws"`, `packages/ui-server/src/app.ts:1067`). Static serving is conditional (`if (options.staticRoot)`, -`packages/ui-server/src/app.ts:1086-1107`). That order establishes access; +`packages/ui-server/src/app.ts:1087-1108`). That order establishes access; it does not select compatibility status. A hermetic real-app launch, with temporary brain root, in-memory database, diff --git a/docs/decisions/managed-worktrees.md b/docs/decisions/managed-worktrees.md index 589db92a76..f845fc36fc 100644 --- a/docs/decisions/managed-worktrees.md +++ b/docs/decisions/managed-worktrees.md @@ -97,7 +97,7 @@ unknown; what happens next is in [hygiene and cleanup](#hygiene-and-cleanup). Each governed or strict turn whose journal is recording opens one worker generation, and every worker of that turn binds the generation's root-owned directory whole at `/tmp/brain-worktrees`, inside the turn's own `/tmp` -(`// Managed worktrees (#2430): a journaled turn's workers bind one empty generation.`, `packages/ui-server/src/ws/run-session.ts:566-578`; +(`// Managed worktrees (#2430): a journaled turn's workers bind one empty generation.`, `packages/ui-server/src/ws/run-session.ts:567-579`; `// Managed worktrees (#2430): the whole access generation`, `packages/ui-sdk/src/server/worker-launcher.ts:848-859`). A granted worktree appears there as `/tmp/brain-worktrees//` on the worker's next path lookup; the worker cannot create entries beside it, and a @@ -199,7 +199,7 @@ The host side is `createReferenceTools` start, the host records its principal and what those workers bind: the public view, the principal's private view and the turn's worker generation, if one opened -(`// Reference tools (#2432): the turn's workers bind the principal's source`, `packages/ui-server/src/ws/run-session.ts:579-585`). +(`// Reference tools (#2432): the turn's workers bind the principal's source`, `packages/ui-server/src/ws/run-session.ts:580-586`). A call carries only a repository, a revision or a worktree identity. The bridge adds the turn's own principal, and a call for a turn the host did not prepare, or under another principal, refuses diff --git a/docs/decisions/session-principals.md b/docs/decisions/session-principals.md index f3b4fd3bc2..acec1943d1 100644 --- a/docs/decisions/session-principals.md +++ b/docs/decisions/session-principals.md @@ -96,13 +96,13 @@ Both reviewers verified every row below against the source. | Cookie verify | `auth.ts:265-296` | strict `^(\d+)\.(\d+)$`, safe integers, future/expired rejection, `cookieEpoch === sessionsEpoch(db)` | | Global epoch | `auth.ts:302-318` | `settings` row; missing = 0; non-integer **throws** | | Global revoke | `auth.ts:321-330` | `setSetting` (overflow-guarded) + `clients.closeAll(1008, …)` | -| Auth routes | `authRoutes`, `app.ts:862` | mounted **before** the guard at `authGuard(authMode`, `app.ts:877`; logout guards itself (`resolveCookiePrincipal(c, auth, deps.db)`, `auth.ts:808`) | -| Passkey management | `passkeyManagementRoutes`, `app.ts:887` | mounted **after** the guard ("Mount AFTER the auth guard") | -| WS guard | `"/ws"`, `app.ts:1066-1083` | origin → `isWsAuthorized` (returns a boolean, `auth.ts:221-237`) → capacity → upgrade | +| Auth routes | `authRoutes`, `app.ts:863` | mounted **before** the guard at `authGuard(authMode`, `app.ts:878`; logout guards itself (`resolveCookiePrincipal(c, auth, deps.db)`, `auth.ts:808`) | +| Passkey management | `passkeyManagementRoutes`, `app.ts:888` | mounted **after** the guard ("Mount AFTER the auth guard") | +| WS guard | `"/ws"`, `app.ts:1067-1084` | origin → `isWsAuthorized` (returns a boolean, `auth.ts:221-237`) → capacity → upgrade | | WS upgrade | `ws/connection.ts:262-264` | `createWsUpgrade` ignores the request context | | WS admission | `ws/connection.ts:84` | `clients.add(ws)` — the socket's identity is unknown | | Turn record | `ws/run-session.ts:136-144` | recorder built from `{turnId, sessionId, billing}`; `RunningTurn` (`ws/turns.ts:55-77`) has no actor | -| Follow-up queue | `QueuedFollowUp`, `ws/turns.ts:61-103`; `slot.queue.push(entry)`, `run-session.ts:1100` | queued entries re-mint `turnId` in the same slot and may come from another socket | +| Follow-up queue | `QueuedFollowUp`, `ws/turns.ts:61-103`; `slot.queue.push(entry)`, `run-session.ts:1121` | queued entries re-mint `turnId` in the same slot and may come from another socket | | Password login | `acquirePasswordVerification(key)`, `auth.ts:749` | argon2id verify, failure counting, in-flight reservation, then `issueSessionCookie` | | Passkey login | `issueLoginSession`, `packages/ui-server/src/middleware/passkeys.ts:448` | assertion verified, then `issueSessionCookie`; `row.id` is in scope | | Passkey **registration** | `INSERT INTO passkey_credentials`, `passkeys.ts:547` | inserts a credential; calls **neither** helper | diff --git a/docs/decisions/sqlite-write-worker.md b/docs/decisions/sqlite-write-worker.md index 19d575ad90..9cc40d36ec 100644 --- a/docs/decisions/sqlite-write-worker.md +++ b/docs/decisions/sqlite-write-worker.md @@ -112,8 +112,17 @@ thread waits for it. summary the client hears of is sent once its row commits. A process without a write worker, such as the cron runner, runs the same procedures on its own connection. +- **Turn budgets followed** (#2850). Every ledger write (admission, + session binding, completed-response and relayed-provider observations, + warnings, the SDK backstop) commits on the worker, with the receipt it + computes read inside the same transaction + (`budgetLedger`, `packages/ui-server/src/brain/turn-budgets.ts`). A report + is priced where it arrives, so the backend still gets the turn's list price + at once; its warnings and any block follow the commit. A relayed provider + request is admitted only after every budget write submitted before it has + committed, and a turn's outcome waits for its budget writes. - **Other app-connection writers stay on the loop for now**: the writer gate, - turn budgets, the session catalog and message sources, the Activity + the session catalog and message sources, the Activity notifier and push ticks, auth sessions, inbox and settings. Each commits at FULL on the loop, and while a worker procedure holds its turn, such a write also waits there for the worker's commit to finish, including a journal diff --git a/docs/hosting/paid-cli-provider-qualification.md b/docs/hosting/paid-cli-provider-qualification.md index cc4b125f89..1d21bb4a95 100644 --- a/docs/hosting/paid-cli-provider-qualification.md +++ b/docs/hosting/paid-cli-provider-qualification.md @@ -152,7 +152,7 @@ restart/retry controls, rather than relying on a worker-supplied counter. Budget coverage excluded provider/tool spend until #1673; it now names the relayed CLI providers and nothing else -(`TURN_BUDGET_COVERAGE`, `packages/ui-server/src/brain/turn-budgets.ts:45`). +(`TURN_BUDGET_COVERAGE`, `packages/ui-server/src/brain/turn-budgets.ts:46`). Only qualified recorded units may advance that disclosure. Before activation, prove refusal before the next dispatch once the applicable known budget is exhausted, alongside native cancellation/cleanup and deduplication. Exact advance diff --git a/docs/plans/async-collaboration-requirements.md b/docs/plans/async-collaboration-requirements.md index 88c68bf211..4c9d09c85b 100644 --- a/docs/plans/async-collaboration-requirements.md +++ b/docs/plans/async-collaboration-requirements.md @@ -226,7 +226,7 @@ Stated before the requirements because five of them derive from it. that is the shape to copy, including its `close()` lifecycle. - R19. **The cron backstop has independent authorization before the general guard.** Mount it on the existing listener before - (`app.use("/api/*", authGuard(`, `packages/ui-server/src/app.ts:877`). + (`app.use("/api/*", authGuard(`, `packages/ui-server/src/app.ts:878`). Authorize a boot-minted ephemeral token, rotated each boot and stored in a 0600 runtime file, with the actual socket address as an additional check. Proxy headers cannot authorize it. The poke succeeds in every auth mode @@ -256,7 +256,7 @@ Stated before the requirements because five of them derive from it. (`export function createTurnRecorder(`, `packages/ui-server/src/activity/recorder.ts:125-203`). - R25. **Autonomous work gets its own pool** (`MAX_AUTONOMOUS_RUNS`, default 2) — but a second counter alone does not deliver "interactive always wins". The host cap applies only when - starting WS sessions (`const cap = host.maxConcurrentSessions();`, `packages/ui-server/src/ws/run-session.ts:1310-1322`), and an autonomous + starting WS sessions (`const cap = host.maxConcurrentSessions();`, `packages/ui-server/src/ws/run-session.ts:1331-1343`), and an autonomous turn can hold a path write lock while an interactive turn waits or is denied at 30 seconds (`export function createTurnLockBinding(`, `packages/ui-backend-claude/src/turn-lock.ts:27-119`). Required: an admission controller with reserved interactive capacity and hybrid yield at an explicit denial-risk diff --git a/docs/plans/async-collaboration.md b/docs/plans/async-collaboration.md index 7370bc8419..5e287e2108 100644 --- a/docs/plans/async-collaboration.md +++ b/docs/plans/async-collaboration.md @@ -131,7 +131,7 @@ system gate. It is not a prerequisite to rebuild or a new independent package. owns manual triggers/history; it does not tick the inbox. - **Two connections:** (`export function createUiDb(`, `packages/ui-server/src/db/client.ts:27-44`) sets WAL, foreign keys and a 5-second busy timeout. Claims are immediate transactions. -- **Auth mounting:** (`app.use("/api/*", authGuard(`, `packages/ui-server/src/app.ts:877`) follows public routes; +- **Auth mounting:** (`app.use("/api/*", authGuard(`, `packages/ui-server/src/app.ts:878`) follows public routes; (`export function authGuard(`, `packages/ui-server/src/middleware/auth.ts:189-249`) binds principals in each auth mode. An internal poke needs independent token authorization before this guard. - **Permission parking:** (`requestPermission: (req) => {`, `packages/ui-server/src/ws/bridge.ts:354-448`). Timeout unwind is @@ -156,7 +156,7 @@ system gate. It is not a prerequisite to rebuild or a new independent package. - **Share provenance/limits:** (`const result = await stageShareAt(`, `packages/ui-server/src/inbox/intake.ts:106`) assigns the source in server code. (`export const SHARE_MAX_TEXT_BYTES =`, `packages/ui-sdk/src/protocol.ts:2348`) bounds text, not binary uploads; extracted T1 context needs its own byte/token bound. -- **Interactive locks:** (`const cap = host.maxConcurrentSessions();`, `packages/ui-server/src/ws/run-session.ts:1310-1322`) gates WS starts; +- **Interactive locks:** (`const cap = host.maxConcurrentSessions();`, `packages/ui-server/src/ws/run-session.ts:1331-1343`) gates WS starts; (`export function createTurnLockBinding(`, `packages/ui-backend-claude/src/turn-lock.ts:27-119`) owns tool locks. Shared-target priority and cooperative yield belong to the keyed lock and backend lifecycle. @@ -1274,7 +1274,7 @@ full-v1 enablement; predicates and schema tests cannot substitute for it. **Approach:** **hybrid — reserve capacity normally, yield only at denial risk.** `MAX_AUTONOMOUS_RUNS` (default 2) is necessary but not sufficient: the host cap applies only at -WS session start (`const cap = host.maxConcurrentSessions();`, `packages/ui-server/src/ws/run-session.ts:1310-1322`) and an autonomous turn can hold a path write lock +WS session start (`const cap = host.maxConcurrentSessions();`, `packages/ui-server/src/ws/run-session.ts:1331-1343`) and an autonomous turn can hold a path write lock while an interactive turn waits or is denied at 30s (`export function createTurnLockBinding(`, `packages/ui-backend-claude/src/turn-lock.ts:27-119`). diff --git a/docs/plans/navigation-recovery.md b/docs/plans/navigation-recovery.md index 77b9dbdaa2..b3123ef825 100644 --- a/docs/plans/navigation-recovery.md +++ b/docs/plans/navigation-recovery.md @@ -397,9 +397,9 @@ implementation. Any PR moving these lines owns their citation repairs. - Activity resolution: (`.get("/activity/runs", (c) => {`, `packages/ui-server/src/routes/activity.ts:143-281`). -- API authorization: (`app.use("/api/*", authGuard(authMode, auth, db));`, `packages/ui-server/src/app.ts:877`). +- API authorization: (`app.use("/api/*", authGuard(authMode, auth, db));`, `packages/ui-server/src/app.ts:878`). -- Queued acceptance: (`// Queue it as the session's next turn; report queued immediately.`, `packages/ui-server/src/ws/run-session.ts:1077-1122`). +- Queued acceptance: (`// Queue it as the session's next turn; report queued immediately.`, `packages/ui-server/src/ws/run-session.ts:1098-1143`). - Approval target: (`function mutateLastAssistant(`, `packages/ui-react/src/stores/chat-state.ts:1224-1254`). diff --git a/packages/ui-sdk/src/server/provider-relay.ts b/packages/ui-sdk/src/server/provider-relay.ts index e9394be5f5..b96f6fab87 100644 --- a/packages/ui-sdk/src/server/provider-relay.ts +++ b/packages/ui-sdk/src/server/provider-relay.ts @@ -131,8 +131,11 @@ export interface ProviderRelayOptions { credential?(name: string): string | undefined; /** The concrete turn proxy: share its destination admission, byte ceiling and connection records. */ egress?: Pick>, "relayed">; - /** Called before each dispatch. A refusal is answered with `budget_exceeded` and nothing is sent. */ - admit?(): { allowed: true } | { allowed: false; message: string }; + /** + * Called before each dispatch, and awaited. A refusal is answered with + * `budget_exceeded` and nothing is sent. + */ + admit?(): { allowed: true } | { allowed: false; message: string } | Promise<{ allowed: true } | { allowed: false; message: string }>; /** One receipt per dispatched attempt. A throw is reported through `failed`, never to the worker. */ observe?(receipt: ProviderRelayReceipt): void; refused?(refusal: ProviderRelayRefusal): void; @@ -315,7 +318,7 @@ export function startProviderRelay(options: ProviderRelayOptions = {}): Provider // Header placement trims surrounding whitespace: redact and send the same value. const secret = credential(key)?.trim(); if (!secret || /[\r\n]/.test(secret)) return refuse(403, "credential_unavailable", `The host has no ${key} to relay`); - const budget = options.admit?.() ?? { allowed: true as const }; + const budget = await options.admit?.() ?? { allowed: true as const }; if (!budget.allowed) return refuse(403, "budget_exceeded", budget.message); const target = new URL(path, origin); const operation = options.egress?.relayed(request.method, target, request.signal); @@ -393,7 +396,7 @@ export function startProviderRelay(options: ProviderRelayOptions = {}): Provider }; if (operation && !operation.charge(headerBytes(`POST ${path} HTTP/1.1`, headers))) return denied("max_bytes_per_turn"); // Admission again, after every wait: a scope may have blocked meanwhile. - const now = options.admit?.() ?? { allowed: true as const }; + const now = await options.admit?.() ?? { allowed: true as const }; if (!now.allowed) { operation?.finish("denied", "budget_exceeded"); return refuse(403, "budget_exceeded", now.message); diff --git a/packages/ui-server/src/app.ts b/packages/ui-server/src/app.ts index d500eb3197..265ed1bb58 100644 --- a/packages/ui-server/src/app.ts +++ b/packages/ui-server/src/app.ts @@ -608,7 +608,8 @@ export async function createApp(options: CreateAppOptions = {}): Promise brain.index({ lock }).catch(ignoreShutdown), report: message => brainLog.emit({ severityText: "WARN", body: message }), exposures: createProviderExposures(db), ...(config.providerRecipientOrigins ? { recipientOrigins: config.providerRecipientOrigins } : {}), indexBatch }); bootCleanup.push(() => governedTurns.close()); - const budgets = createTurnBudgets(db, pricing); + // Budget writes commit on the write worker (#2850). + const budgets = createTurnBudgets(db, pricing, Date.now, writer); relayBudgets = budgets; const sessionProfiles = createSessionProfiles(db, governedTurns, budgets, { hardTimeoutMs: () => host.turnTimeoutMs }); const journal = createTurnJournal(db, { nativeManager, brainPath: config.brainPath, dbPath: config.dbPath, objectsRoot: journalRoot, writer, diff --git a/packages/ui-server/src/brain/host-boundary.ts b/packages/ui-server/src/brain/host-boundary.ts index 6598785fbd..bd82bd047a 100644 --- a/packages/ui-server/src/brain/host-boundary.ts +++ b/packages/ui-server/src/brain/host-boundary.ts @@ -230,7 +230,7 @@ export interface HostBoundaryOptions { */ providerRelay?: { origins?: ReadonlyMap; - admit(turnId: string | null): { allowed: true } | { allowed: false; message: string }; + admit(turnId: string | null): { allowed: true } | { allowed: false; message: string } | Promise<{ allowed: true } | { allowed: false; message: string }>; observe(turnId: string | null, attempt: ProviderRelayReceipt): void; }; profile?: "governed" | "strict"; diff --git a/packages/ui-server/src/brain/provider-meter.ts b/packages/ui-server/src/brain/provider-meter.ts index 3a4f858c41..29b81438d6 100644 --- a/packages/ui-server/src/brain/provider-meter.ts +++ b/packages/ui-server/src/brain/provider-meter.ts @@ -19,17 +19,18 @@ export function maintenanceProviderMeter(options: { const policy = createGovernancePolicyReader(options.brainPath); return { ...(options.origins ? { origins: options.origins } : {}), - admit(turnId) { + async admit(turnId) { const budgets = options.budgets(); if (!budgets) return { allowed: false, message: new BudgetUnavailable().message }; try { - const blocked = budgets.providerAdmission(turnId, policy.load().policy); + const blocked = await budgets.providerAdmission(turnId, policy.load().policy); return blocked ? { allowed: false, message: new BudgetExceeded(blocked, "mid_turn").message } : { allowed: true }; } catch { return { allowed: false, message: new BudgetUnavailable().message }; } }, observe(turnId: string | null, attempt: ProviderRelayReceipt) { - try { options.budgets()?.observeProvider(turnId, attempt); } - catch (error) { options.report(`a relayed provider request could not be recorded: ${error instanceof Error ? error.message : String(error)}`); } + const failed = (error: unknown) => options.report(`a relayed provider request could not be recorded: ${error instanceof Error ? error.message : String(error)}`); + try { options.budgets()?.observeProvider(turnId, attempt).catch(failed); } + catch (error) { failed(error); } }, }; } diff --git a/packages/ui-server/src/brain/turn-budgets.ts b/packages/ui-server/src/brain/turn-budgets.ts index aead033e56..f40f579c31 100644 --- a/packages/ui-server/src/brain/turn-budgets.ts +++ b/packages/ui-server/src/brain/turn-budgets.ts @@ -6,6 +6,7 @@ import { priceUsage, type RollupPricing } from "../activity/store.js"; import { priceProviderAttempt } from "../pricing/provider-pricing.js"; import type { ProviderRelayReceipt } from "@schlessera/brain-ui-sdk/internal"; import type { BudgetObservation, TurnBudgetReceipt } from "@schlessera/brain-ui-sdk/protocol"; +import { procedureCaller, type SqliteWriter, type WriteProcedures } from "../db/write-worker.js"; export class BudgetUnavailable extends Error { constructor() { super("Inference budget accounting became unavailable. The turn stopped; review any partial work before continuing."); } } @@ -60,13 +61,19 @@ function observation(scope: BudgetObservation["scope"], selected: Row[], provide warning: costUsd !== null && costUsd >= limits.warn_usd || unpricedMeasured && limits.warn_tokens !== undefined && tokens !== null && tokens >= limits.warn_tokens, blocked: costUsd !== null && costUsd >= limits.block_usd || unpricedMeasured && tokens !== null && tokens >= limits.block_tokens }; } -/** Concrete host service; policy is the already frozen admission fact, never reparsed here. */ -export function createTurnBudgets(db: Database, pricing: RollupPricing, clock: () => number = Date.now) { - const active = new Map(); +/** The usage one completed-response report measured, priced where it was reported (#2850). */ +export interface MeasuredUsage { tokens: number | null; cost: number | null; estimate: number | null; partialCost: number | null } +/** One relayed attempt as priced where it was reported. */ +export interface PricedAttempt { costUsd: number | null; estimate: boolean | null; tokens: number | null } +export type BudgetObservationResult = { receipt: TurnBudgetReceipt; warnings: BudgetObservation[]; blocked: BudgetObservation | undefined }; + +/** + * The ledger's reads and transaction bodies, on whichever connection commits + * them: the hosted server's write worker (`turn-budget-writes.ts`) or, without + * one, the caller's own. `policy` is the turn's frozen admission fact and + * `now` is taken where the call was made. + */ +export function budgetLedger(db: Database) { const retained = (id: string, sessionId: string | null, day: number): Row[] => db.query(`SELECT 1 AS legacy, pricing_estimate, run_id AS turn_id, session_id, started_at, cost_usd, cost_usd AS cost_lower_bound, CASE WHEN input_tokens IS NULL AND output_tokens IS NULL AND cache_read_tokens IS NULL AND cache_creation_tokens IS NULL THEN NULL ELSE COALESCE(input_tokens,0)+COALESCE(output_tokens,0)+COALESCE(cache_read_tokens,0)+COALESCE(cache_creation_tokens,0) END AS tokens @@ -103,17 +110,20 @@ export function createTurnBudgets(db: Database, pricing: RollupPricing, clock: ( tokens: slice.tokens === null ? null : slice.tokens + (utcBudgetDay(row.started_at) === day ? Math.max(0, (row.tokens ?? 0) - (slices.tokens ?? 0)) : 0) }; }); }; - const receipt = (id: string): TurnBudgetReceipt => { - const current = active.get(id); - if (!current) { - const saved = db.query("SELECT receipt FROM turn_budget_usage WHERE turn_id = ?").get(id) as { - receipt: string | null; - } | null; - if (!saved?.receipt) - throw new Error("Turn has no frozen budget admission"); - return JSON.parse(saved.receipt) as TurnBudgetReceipt; - } - const now = clock(), day = utcBudgetDay(now); + /** Relayed provider attempts in one scope (#1673): known prices, reported tokens and unknown spend. */ + const providerTotals = (where: string, value: string | number): ProviderTotals => db.query(`SELECT COUNT(*) AS attempts, + COALESCE(SUM(cost_usd), 0) AS cost, COUNT(cost_usd) AS priced, COALESCE(SUM(tokens), 0) AS tokens, COUNT(tokens) AS measured, + COALESCE(SUM(CASE WHEN cost_usd IS NULL AND tokens IS NOT NULL THEN 1 ELSE 0 END), 0) AS unpriced_measured + FROM provider_relay_usage WHERE ${where}`).get(value) as ProviderTotals; + const saved = (id: string): TurnBudgetReceipt => { + const row = db.query("SELECT receipt FROM turn_budget_usage WHERE turn_id = ?").get(id) as { receipt: string | null } | null; + if (!row?.receipt) + throw new Error("Turn has no frozen budget admission"); + return JSON.parse(row.receipt) as TurnBudgetReceipt; + }; + /** An active turn's receipt under its frozen policy, as of `now`. */ + const receipt = (id: string, policy: EffectiveGovernancePolicy, now: number): TurnBudgetReceipt => { + const day = utcBudgetDay(now); const ownRow = db.query("SELECT session_id FROM turn_budget_usage WHERE turn_id=?").get(id) as { session_id: string | null; }; @@ -123,19 +133,14 @@ export function createTurnBudgets(db: Database, pricing: RollupPricing, clock: ( const provider = scope === "turn" ? providerTotals("turn_id = ?", id) : scope === "session" ? own.session_id === null ? NO_PROVIDER_ATTEMPTS : providerTotals("turn_id IN (SELECT turn_id FROM turn_budget_usage WHERE session_id = ?)", own.session_id) : providerTotals("day_start = ?", day); - return observation(scope, selected, provider, current.policy, day); + return observation(scope, selected, provider, policy, day); }); return { listPriceCostUsd: own.cost_usd, listPriceIsEstimate: own.pricing_estimate === null ? null : Boolean(own.pricing_estimate), tokens: own.tokens, observedAt: now, coverage: TURN_BUDGET_COVERAGE, scopes }; }; - /** Relayed provider attempts in one scope (#1673): known prices, reported tokens and unknown spend. */ - const providerTotals = (where: string, value: string | number): ProviderTotals => db.query(`SELECT COUNT(*) AS attempts, - COALESCE(SUM(cost_usd), 0) AS cost, COUNT(cost_usd) AS priced, COALESCE(SUM(tokens), 0) AS tokens, COUNT(tokens) AS measured, - COALESCE(SUM(CASE WHEN cost_usd IS NULL AND tokens IS NOT NULL THEN 1 ELSE 0 END), 0) AS unpriced_measured - FROM provider_relay_usage WHERE ${where}`).get(value) as ProviderTotals; /** The day scope alone, for a relayed maintenance run that belongs to no turn. */ - const dayObservation = (policy: EffectiveGovernancePolicy): BudgetObservation => { - const day = utcBudgetDay(clock()); + const dayObservation = (policy: EffectiveGovernancePolicy, now: number): BudgetObservation => { + const day = utcBudgetDay(now); const selected = daySelection(rows("", null, day), day); return observation("day", selected, providerTotals("day_start = ?", day), policy, day); }; @@ -151,107 +156,195 @@ export function createTurnBudgets(db: Database, pricing: RollupPricing, clock: ( }); }; return { - begin(id: string, sessionId: string | undefined, policy: EffectiveGovernancePolicy, route?: PricingRoute) { - const rates = new Map>(); - active.set(id, { policy, route, pricing: { resolve(model, route) { - const key = JSON.stringify([model, route]); - if (!rates.has(key)) - rates.set(key, structuredClone(pricing.resolve(model, route))); - return rates.get(key)!; - } } }); - const result = db.transaction(() => { - db.query("INSERT INTO turn_budget_usage (turn_id,session_id,started_at,cost_usd,tokens,observed_at) VALUES (?,?,?,NULL,NULL,?) ON CONFLICT(turn_id) DO NOTHING").run(id, sessionId ?? null, clock(), clock()); - const result = receipt(id); + receipt, saved, dayObservation, + begin(id: string, sessionId: string | null, policy: EffectiveGovernancePolicy, now: number): TurnBudgetReceipt { + return db.transaction(() => { + db.query("INSERT INTO turn_budget_usage (turn_id,session_id,started_at,cost_usd,tokens,observed_at) VALUES (?,?,?,NULL,NULL,?) ON CONFLICT(turn_id) DO NOTHING").run(id, sessionId, now, now); + const result = receipt(id, policy, now); db.query("UPDATE turn_budget_usage SET receipt = ? WHERE turn_id = ?").run(JSON.stringify(result), id); return result; }).immediate(); - const blocked = result.scopes.find(item => item.blocked && item.scope !== "turn"); - if (blocked) - throw new BudgetExceeded(blocked, "admission"); - return result; }, - bind(id: string, sessionId: string) { db.query("UPDATE turn_budget_usage SET session_id = ? WHERE turn_id = ?").run(sessionId, id); }, - observe(id: string, usage: TurnUsage) { - const current = active.get(id); - if (!current) - return undefined; + bind(id: string, sessionId: string): void { db.query("UPDATE turn_budget_usage SET session_id = ? WHERE turn_id = ?").run(sessionId, id); }, + observe(id: string, policy: EffectiveGovernancePolicy, measured: MeasuredUsage, now: number): BudgetObservationResult { return db.transaction(() => { const previous = db.query("SELECT * FROM turn_budget_usage WHERE turn_id = ?").get(id) as Row; - const tokens = tokensOf(usage); - const priced = usage.perModel && Object.keys(usage.perModel).length ? priceUsage(usage.perModel, current.pricing, current.route) : null; - const cost = priced?.costUsd ?? null; - const knownPrices = Object.entries(usage.perModel ?? {}).filter(([, counts]) => tokensOf(counts) !== null && tokensOf(counts)! > 0).map(([model, counts]) => priceUsage({ [model]: counts }, current.pricing, current.route)).filter(value => value !== null); - const partialCost = knownPrices.length ? knownPrices.reduce((sum, value) => sum + value.costUsd, 0) : null; + const { tokens, cost, partialCost } = measured; const lower = cost ?? (partialCost === null ? previous.cost_lower_bound : Math.max(partialCost, previous.cost_lower_bound ?? 0)); if (tokens !== null && previous.tokens !== null && tokens < previous.tokens) throw new Error("Completed inference receipts must be cumulative"); - const now = clock(), day = utcBudgetDay(now); + const day = utcBudgetDay(now); const deltaCost = lower === null ? null : Math.max(0, lower - (previous.cost_lower_bound ?? 0)); const deltaTokens = tokens === null ? null : tokens - (previous.tokens ?? 0); db.query("INSERT INTO turn_budget_days(turn_id,day_start,cost_usd,cost_lower_bound,tokens) VALUES(?,?,?,?,?) ON CONFLICT(turn_id,day_start) DO UPDATE SET cost_usd = CASE WHEN turn_budget_days.cost_usd IS NULL OR excluded.cost_usd IS NULL THEN NULL ELSE turn_budget_days.cost_usd+excluded.cost_usd END, cost_lower_bound = CASE WHEN turn_budget_days.cost_lower_bound IS NULL AND excluded.cost_lower_bound IS NULL THEN NULL ELSE COALESCE(turn_budget_days.cost_lower_bound,0)+COALESCE(excluded.cost_lower_bound,0) END, tokens = CASE WHEN turn_budget_days.tokens IS NULL OR excluded.tokens IS NULL THEN NULL ELSE turn_budget_days.tokens+excluded.tokens END").run(id, day, cost === null ? null : deltaCost, deltaCost, deltaTokens); - db.query("UPDATE turn_budget_usage SET cost_usd=?,cost_lower_bound=?,pricing_estimate=?,tokens=?,observed_at=? WHERE turn_id=?").run(cost, lower, priced ? Number(priced.estimate) : null, tokens, now, id); - const result = receipt(id); + db.query("UPDATE turn_budget_usage SET cost_usd=?,cost_lower_bound=?,pricing_estimate=?,tokens=?,observed_at=? WHERE turn_id=?").run(cost, lower, measured.estimate, tokens, now, id); + const result = receipt(id, policy, now); db.query("UPDATE turn_budget_usage SET receipt = ? WHERE turn_id = ?").run(JSON.stringify(result), id); const warnings = warningCrossings(id, result); return { receipt: result, warnings, blocked: result.scopes.find(item => item.blocked) }; }).immediate(); }, - receipt, - warnings(id: string) { return db.transaction(() => warningCrossings(id, receipt(id))).immediate(); }, - backstop(id: string) { - const result = receipt(id); + warnings(id: string, policy: EffectiveGovernancePolicy, now: number) { return db.transaction(() => warningCrossings(id, receipt(id, policy, now))).immediate(); }, + backstop(id: string, policy: EffectiveGovernancePolicy, now: number): TurnBudgetReceipt { + const result = receipt(id, policy, now); result.scopes[0] = { ...result.scopes[0]!, blocked: true }; db.query("UPDATE turn_budget_usage SET receipt=? WHERE turn_id=?").run(JSON.stringify(result), id); return result; }, + /** `policy` is the active turn's, or null for a run that belongs to no active turn. */ + observeProvider(turnId: string | null, policy: EffectiveGovernancePolicy | null, attempt: ProviderRelayReceipt, priced: PricedAttempt, now: number): BudgetObservationResult | undefined { + return db.transaction(() => { + const inserted = db.query(`INSERT INTO provider_relay_usage (attempt_id, turn_id, day_start, observed_at, family, operation, origin, model, + outcome, status, usage, cost_usd, pricing_estimate, tokens) VALUES (?,?,?,?,?,?,?,?,?,?,?,?,?,?) + ON CONFLICT(attempt_id) DO UPDATE SET day_start=excluded.day_start, observed_at=excluded.observed_at, outcome=excluded.outcome, status=excluded.status, usage=excluded.usage, + cost_usd=excluded.cost_usd, pricing_estimate=excluded.pricing_estimate, tokens=excluded.tokens + WHERE provider_relay_usage.outcome = 'dispatched' AND excluded.outcome <> 'dispatched'`).run(attempt.attemptId, turnId, utcBudgetDay(now), now, + attempt.family, attempt.operation, attempt.origin, attempt.model, attempt.outcome, attempt.status, + attempt.usage ? JSON.stringify(attempt.usage) : null, priced.costUsd, priced.estimate === null ? null : Number(priced.estimate), priced.tokens).changes > 0; + if (!policy || !inserted) return undefined; + const result = receipt(turnId!, policy, now); + db.query("UPDATE turn_budget_usage SET receipt = ? WHERE turn_id = ?").run(JSON.stringify(result), turnId); + return { receipt: result, warnings: warningCrossings(turnId!, result), blocked: result.scopes.find(item => item.blocked) }; + }).immediate(); + }, + }; +} + +/** The write worker imports this module and runs the ledger's writes there (#2850). */ +export const url = import.meta.url; +export function setup(db: Database): WriteProcedures { + const ledger = budgetLedger(db); + return { begin: ledger.begin, bind: ledger.bind, observe: ledger.observe, warnings: ledger.warnings, backstop: ledger.backstop, observeProvider: ledger.observeProvider }; +} + +/** + * Concrete host service; policy is the already frozen admission fact, never + * reparsed here. With a write worker every ledger write commits there, in + * submission order (#2850); reads stay on `db`. Usage is priced where it is + * reported, so a turn's list price is known at once; its receipt, warnings + * and any block follow its commit. Admission waits for every budget write + * submitted before it. + */ +export function createTurnBudgets(db: Database, pricing: RollupPricing, clock: () => number = Date.now, writer?: SqliteWriter) { + const active = new Map(); + const ledger = budgetLedger(db); + const submit = procedureCaller(db, { url, setup }, undefined, writer); + // The latest budget write, until it settles: admission reads after it. + let tail: Promise = Promise.resolve(); + const call = (name: string, ...args: unknown[]): Promise => { + const done = submit(name, ...args); + tail = done.then(() => {}, () => {}); + return done; + }; + const settled = async () => { let waiting: Promise; do { waiting = tail; await waiting; } while (waiting !== tail); }; + const current = (id: string) => { + const turn = active.get(id); + if (turn?.failed !== undefined) throw turn.failed; + return turn; + }; + const measure = (turn: { pricing: RollupPricing; route?: PricingRoute }, usage: TurnUsage): MeasuredUsage => { + const tokens = tokensOf(usage); + const priced = usage.perModel && Object.keys(usage.perModel).length ? priceUsage(usage.perModel, turn.pricing, turn.route) : null; + const knownPrices = Object.entries(usage.perModel ?? {}).filter(([, counts]) => tokensOf(counts) !== null && tokensOf(counts)! > 0).map(([model, counts]) => priceUsage({ [model]: counts }, turn.pricing, turn.route)).filter(value => value !== null); + const partialCost = knownPrices.length ? knownPrices.reduce((sum, value) => sum + value.costUsd, 0) : null; + return { tokens, cost: priced?.costUsd ?? null, estimate: priced ? Number(priced.estimate) : null, partialCost }; + }; + const receipt = (id: string): TurnBudgetReceipt => { + const turn = active.get(id); + return turn ? ledger.receipt(id, turn.policy, clock()) : ledger.saved(id); + }; + return { + /** Settles once the turn's admission row and receipt have committed; rejects with `BudgetExceeded` when a session or day scope is blocked. */ + async begin(id: string, sessionId: string | undefined, policy: EffectiveGovernancePolicy, route?: PricingRoute): Promise { + const rates = new Map>(); + active.set(id, { policy, route, pricing: { resolve(model, route) { + const key = JSON.stringify([model, route]); + if (!rates.has(key)) + rates.set(key, structuredClone(pricing.resolve(model, route))); + return rates.get(key)!; + } } }); + const result = await call("begin", id, sessionId ?? null, policy, clock()); + const blocked = result.scopes.find(item => item.blocked && item.scope !== "turn"); + if (blocked) + throw new BudgetExceeded(blocked, "admission"); + return result; + }, + /** A failed binding makes the turn's accounting unavailable: its session scope would miss it. */ + bind(id: string, sessionId: string): Promise { + return call("bind", id, sessionId).catch((error: unknown) => { + const turn = active.get(id); + if (turn) turn.failed ??= error; + throw error; + }); + }, + /** + * One cumulative completed-response report, priced and validated now + * (an invalid one throws); settles with the receipt, any warning crossed + * and any block once it has committed. + */ + observe(id: string, usage: TurnUsage): Promise | undefined { + const turn = current(id); + if (!turn) + return undefined; + return call("observe", id, turn.policy, measure(turn, usage), clock()); + }, + /** The list price the receipt of this report will carry, known before it commits. */ + listPrice(id: string, usage: TurnUsage): number | null { + const turn = active.get(id); + return turn ? measure(turn, usage).cost : null; + }, + receipt, + warnings(id: string): Promise { + const turn = current(id); + if (!turn) return Promise.reject(new Error("Turn has no frozen budget admission")); + return call("warnings", id, turn.policy, clock()); + }, + backstop(id: string): Promise { + const turn = current(id); + if (!turn) return Promise.reject(new Error("Turn has no frozen budget admission")); + return call("backstop", id, turn.policy, clock()); + }, saved(id: string): TurnBudgetReceipt | undefined { // Presentation must not turn a missing accounting receipt into a failed journal observation. - try { - const row = db.query("SELECT receipt FROM turn_budget_usage WHERE turn_id = ?").get(id) as { - receipt: string | null; - } | null; - return row?.receipt ? JSON.parse(row.receipt) : undefined; - } + try { return ledger.saved(id); } catch { return undefined; } }, /** - * Before each relayed CLI provider request (#1673). A turn's request is - * refused once any of its scopes is blocked; a maintenance run outside any - * turn, once its UTC day is under `policy`. Nothing is reserved. + * Before each relayed CLI provider request (#1673), once every budget + * write submitted before it has committed. A turn's request is refused + * once any of its scopes is blocked; a maintenance run outside any turn, + * once its UTC day is under `policy`. Nothing is reserved. */ - providerAdmission(turnId: string | null, policy: EffectiveGovernancePolicy): BudgetObservation | undefined { - if (turnId !== null && active.has(turnId)) return receipt(turnId).scopes.find(item => item.blocked); - const day = dayObservation(policy); + async providerAdmission(turnId: string | null, policy: EffectiveGovernancePolicy): Promise { + await settled(); + if (turnId !== null && current(turnId)) return receipt(turnId).scopes.find(item => item.blocked); + const day = ledger.dayObservation(policy, clock()); return day.blocked ? day : undefined; }, /** * Record one relayed attempt once, by its attempt identity, against its * turn (or its day alone). A `dispatched` row is unknown spend until the * same attempt settles; a settled attempt never changes again, so a replay - * adds nothing. For an active turn, returns the new receipt with any warning - * crossed and the scope it blocked. + * adds nothing. For an active turn, settles with the new receipt, any + * warning crossed and the scope it blocked. */ - observeProvider(turnId: string | null, attempt: ProviderRelayReceipt) { - const current = turnId !== null ? active.get(turnId) : undefined; - const priced = priceProviderAttempt(attempt, current?.pricing ?? pricing); - const now = clock(); - return db.transaction(() => { - const inserted = db.query(`INSERT INTO provider_relay_usage (attempt_id, turn_id, day_start, observed_at, family, operation, origin, model, - outcome, status, usage, cost_usd, pricing_estimate, tokens) VALUES (?,?,?,?,?,?,?,?,?,?,?,?,?,?) - ON CONFLICT(attempt_id) DO UPDATE SET day_start=excluded.day_start, observed_at=excluded.observed_at, outcome=excluded.outcome, status=excluded.status, usage=excluded.usage, - cost_usd=excluded.cost_usd, pricing_estimate=excluded.pricing_estimate, tokens=excluded.tokens - WHERE provider_relay_usage.outcome = 'dispatched' AND excluded.outcome <> 'dispatched'`).run(attempt.attemptId, turnId, utcBudgetDay(now), now, - attempt.family, attempt.operation, attempt.origin, attempt.model, attempt.outcome, attempt.status, - attempt.usage ? JSON.stringify(attempt.usage) : null, priced.costUsd, priced.estimate === null ? null : Number(priced.estimate), priced.tokens).changes > 0; - if (!current || !inserted) return undefined; - const result = receipt(turnId!); - db.query("UPDATE turn_budget_usage SET receipt = ? WHERE turn_id = ?").run(JSON.stringify(result), turnId); - return { receipt: result, warnings: warningCrossings(turnId!, result), blocked: result.scopes.find(item => item.blocked) }; - }).immediate(); + observeProvider(turnId: string | null, attempt: ProviderRelayReceipt): Promise { + const turn = turnId !== null ? active.get(turnId) : undefined; + const priced = priceProviderAttempt(attempt, turn?.pricing ?? pricing); + return call("observeProvider", turnId, turn?.policy ?? null, attempt, + { costUsd: priced.costUsd, estimate: priced.estimate, tokens: priced.tokens }, clock()); }, remaining(id: string) { const own = receipt(id).scopes[0]!; return Math.max(0, own.blockUsd - (own.costUsd ?? 0)); }, + /** Settles once every budget write submitted before it has committed. */ + settled, end(id: string) { active.delete(id); }, }; } diff --git a/packages/ui-server/src/ws/bridge.ts b/packages/ui-server/src/ws/bridge.ts index ed19204916..bb91e5d1a2 100644 --- a/packages/ui-server/src/ws/bridge.ts +++ b/packages/ui-server/src/ws/bridge.ts @@ -318,7 +318,7 @@ export function makeBridge( host.sendToClients(scoped); // Identity must reach the client before its first scoped journal receipt. // The journal publishes its receipt once the binding commits, after this frame. - if (msg.type === "session_info" && turn.turnId === turnId) { host.budgets?.bind(turnId, msg.sessionId); void host.journal?.bindSession(turnId, msg.sessionId).catch(() => { /* The turn's record stays unbound; its observations still commit. */ }); host.scope?.bindSession(turnId, msg.sessionId); } + if (msg.type === "session_info" && turn.turnId === turnId) { void host.budgets?.bind(turnId, msg.sessionId).catch(() => { /* The turn's accounting is unavailable; its next observation stops it. */ }); void host.journal?.bindSession(turnId, msg.sessionId).catch(() => { /* The turn's record stays unbound; its observations still commit. */ }); host.scope?.bindSession(turnId, msg.sessionId); } work?.observe(scoped); if (msg.type === "result") { // Only the live turn's own result may set its disposition — a late diff --git a/packages/ui-server/src/ws/run-session.ts b/packages/ui-server/src/ws/run-session.ts index c64de4d5ec..347a9ae5bc 100644 --- a/packages/ui-server/src/ws/run-session.ts +++ b/packages/ui-server/src/ws/run-session.ts @@ -506,6 +506,7 @@ async function runRetainedSession( let referenceViews: { public: string; private: string } | null = null, referenceTurn = false; let journalFailure: Error | undefined; let budgetFailure: Error | undefined; + const budgetWrites = new Set>(); let scopeStopReason: string | undefined; turn.backendStarted = false; let enteredBackend = false; @@ -538,11 +539,11 @@ async function runRetainedSession( if (!governedHold || !host.journal) lifecycle?.skip("journal"); if (governedHold && host.budgets) { lifecycle?.mark("budgetStartedAt"); - try { host.budgets.begin(turn.turnId,resumeId,governedHold.policy.policy,billing?.pricingRoute); } + try { await host.budgets.begin(turn.turnId,resumeId,governedHold.policy.policy,billing?.pricingRoute); } catch (error) { if(error instanceof BudgetExceeded) throw error; throw new BudgetUnavailable(); } } if (governedHold && host.budgets) lifecycle?.mark("budgetAdmittedAt"); - try { if(governedHold && host.budgets) for(const item of host.budgets.warnings(turn.turnId)) host.activity?.governance?.({ + try { if(governedHold && host.budgets) for(const item of await host.budgets.warnings(turn.turnId)) host.activity?.governance?.({ kind:"budget_warning",turnId:turn.turnId,sessionId:turn.sessionId,actor:origin,model:mainModel,profile:governedHold.sandbox.profile ?? "strict", payload:{scope:item.scope,unit:item.costUsd !== null && item.costUsd >= item.warnUsd ? "usd":"tokens", limit:item.costUsd !== null && item.costUsd >= item.warnUsd ? item.warnUsd:item.warnTokens!, @@ -640,6 +641,12 @@ async function runRetainedSession( recorder?.recordBudgetStop(budgetFailure.message,"budget_unavailable"); abortController.abort(budgetFailure);coordinator.drainPendingForTurn(turn,"Inference budget accounting became unavailable."); }; + // A budget write's commit, and what the turn does with its outcome; + // the turn's outcome waits for every one (#2850). + const budgetWrite = (work: Promise) => { + const pending = work.catch(() => { if (!budgetFailure) budgetUnavailable(); }).finally(() => budgetWrites.delete(pending)); + budgetWrites.add(pending); + }; const bridge = makeBridge( host, turn, @@ -671,20 +678,28 @@ async function runRetainedSession( ...(initial.profileDraftId && !resumeId && preference ? { draftId: initial.profileDraftId, preference } : {}) } : undefined, host.budgets && governedHold ? (usage,backstop=false) => { if (budgetFailure || abortController.signal.aborted || turn.cancelled) return; + // The report is priced now; its receipt, warnings and any block + // follow its commit (#2850), judged against the turn as it is then. + const live = () => !budgetFailure && !abortController.signal.aborted && !turn.cancelled; try { - const saved = backstop ? host.budgets!.backstop(turn.turnId) : undefined; - const result = saved ? {receipt:saved,warnings:[],blocked:{...saved.scopes[0]!,blocked:true}} : host.budgets!.observe(turn.turnId,usage); - if (!result) return; - budgetWarnings(result.warnings); - if (!backstop) { + if (backstop) { + budgetWrite(host.budgets!.backstop(turn.turnId).then(saved => { if (live()) budgetBlocked({...saved.scopes[0]!,blocked:true}, true); })); + return; + } + const observed = host.budgets!.observe(turn.turnId,usage); + if (!observed) return; + const listPrice = host.budgets!.listPrice(turn.turnId,usage); + budgetWrite(observed.then(result => { + if (!live()) return; + budgetWarnings(result.warnings); const {costUsd:reportedCost,...counts}=usage; void reportedCost; host.sendToClients(withTurnScope({type:"usage",turnId:turn.turnId,usage:{...counts, ...(result.receipt.listPriceCostUsd !== null ? {costUsd:result.receipt.listPriceCostUsd} : {})}, observedAt:result.receipt.observedAt,budget:result.receipt},turn)); - } - if (result.blocked) budgetBlocked(result.blocked, backstop); - return {...usage,...(result.receipt.listPriceCostUsd !== null ? {costUsd:result.receipt.listPriceCostUsd} : {})}; + if (result.blocked) budgetBlocked(result.blocked, false); + })); + return {...usage,...(listPrice !== null ? {costUsd:listPrice} : {})}; } catch { budgetUnavailable(); } } : undefined, billing, @@ -729,23 +744,27 @@ async function runRetainedSession( // The turn this relay meters, fixed now: a late report never lands on a later turn. const egress = turnEgress, hold = governedHold, relayTurnId = turn.turnId; providerRelay = startProviderRelay({ egress, ...(host.cliProviderOrigins ? { origins: host.cliProviderOrigins } : {}), - admit: () => { + admit: async () => { if (budgetFailure) return { allowed: false, message: budgetFailure.message }; if (!host.budgets) return { allowed: true }; try { - const blocked = host.budgets.providerAdmission(relayTurnId, hold.policy.policy); + const blocked = await host.budgets.providerAdmission(relayTurnId, hold.policy.policy); + // The turn may have stopped while admission waited for earlier budget writes. + const stopped = budgetFailure as Error | undefined; + if (stopped) return { allowed: false, message: stopped.message }; return blocked ? { allowed: false, message: new BudgetExceeded(blocked, "mid_turn").message } : { allowed: true }; } catch { return { allowed: false, message: new BudgetUnavailable().message }; } }, observe: attempt => { if (!host.budgets) return; try { - const result = host.budgets.observeProvider(relayTurnId, attempt); - if (!result || budgetFailure || abortController.signal.aborted || turn.cancelled) return; - // A crossed block refuses each further relayed request (ruling D3 on - // #1673); the agent's next completed response, which sees the same - // scopes, stops the turn as any usage over a block does. - budgetWarnings(result.warnings); + budgetWrite(host.budgets.observeProvider(relayTurnId, attempt).then(result => { + if (!result || budgetFailure || abortController.signal.aborted || turn.cancelled) return; + // A crossed block refuses each further relayed request (ruling D3 on + // #1673); the agent's next completed response, which sees the same + // scopes, stops the turn as any usage over a block does. + budgetWarnings(result.warnings); + })); } catch { if (!budgetFailure) budgetUnavailable(); } }, }); @@ -813,6 +832,7 @@ async function runRetainedSession( // A worktree generation retires first: no grant can follow that observation. if (worktreeGeneration) await host.referenceWorktrees!.retireGeneration(worktreeGeneration.id); await host.journal?.finish(turn.turnId); + while (budgetWrites.size) await Promise.all([...budgetWrites]); if (budgetFailure) throw budgetFailure; if (journalFailure) throw journalFailure; // A host timeout ends the turn like a cancel: its late result is @@ -848,6 +868,7 @@ async function runRetainedSession( lifecycle?.mark("failureObservedAt"); await settleIndex?.(); await turnEgress?.settled(); + while (budgetWrites.size) await Promise.all([...budgetWrites]); const egressFailure = turnEgress?.recordingFailure(); const terminal = budgetFailure ?? journalFailure ?? egressFailure ?? err; if (terminal instanceof BudgetExceeded && terminal.stage === "admission") { diff --git a/packages/ui-server/tests/app-writes-offloop.test.ts b/packages/ui-server/tests/app-writes-offloop.test.ts index 95808be93d..b2bae87f66 100644 --- a/packages/ui-server/tests/app-writes-offloop.test.ts +++ b/packages/ui-server/tests/app-writes-offloop.test.ts @@ -10,6 +10,8 @@ import { createGovernanceWrites } from "../src/governance/event-writes.js"; import { createTurnJournal } from "../src/brain/turn-journal.js"; import { createTurnScope } from "../src/brain/turn-scope.js"; import { createTurnRecorder } from "../src/activity/recorder.js"; +import { createTurnBudgets } from "../src/brain/turn-budgets.js"; +import * as probe from "./fixtures/write-worker-probe.js"; import { createActivityStore } from "../src/activity/store.js"; import { createNativeJournalFixture } from "./helpers/native-journal-fixture.js"; import { watchWriteSites } from "./helpers/write-sites.js"; @@ -105,3 +107,32 @@ test("a turn's settled() waits for its roadblock commit", async () => { expect(settled).toBe(true); } finally { db.close(); } }); + +test("turn budget writes commit on the write worker at FULL, after earlier writes, and admission waits for them", async () => { + const f = fixture(); + const runtime = createActivityRuntime(f.db, { log: quietLog, writer: f.writer }); + try { + const sites = await watchWriteSites(f.db, f.writer, ["turn_budget_usage", "turn_budget_days", "turn_budget_warnings", "provider_relay_usage", "activity_spans"]); + const policy = resolveGovernancePolicy({ schema_version: 1, budgets: { turn: { warn_usd: 1, block_usd: 2, block_tokens: 10 } } }); + const budgets = createTurnBudgets(f.db, { resolve: () => ({ input: 1, output: 1, cacheRead: 1, cacheWrite: 1, estimate: false, source: "snapshot" }) }, Date.now, f.writer); + void runtime.store.startSpan({ spanId: "voyage:root", runId: "voyage", name: "invoke_agent", kind: "turn", origin: "session", sessionId: "ithaca" }); + await budgets.begin("voyage", undefined, policy); + await budgets.bind("voyage", "ithaca"); + const usage = (n: number) => ({ inputTokens: n, perModel: { odysseus: { inputTokens: n } } }); + expect(budgets.listPrice("voyage", usage(1)), "BUDGET_LIST_PRICE_BEFORE_COMMIT").toBe(1); + expect((await budgets.observe("voyage", usage(1))!).warnings.map(item => item.scope), "BUDGET_OBSERVE_COMMITS_ON_WORKER").toEqual(["turn"]); + // A slow commit ahead of the report that blocks the turn: admission must not read past it. + const held = f.writer.channel(probe, { table: "odysseus_log" }); await held.ready; + const slow = held.call("slow", "Calypso", 300); + const blocking = budgets.observe("voyage", usage(2))!; + expect((await budgets.providerAdmission("voyage", policy))?.scope, "BUDGET_ADMISSION_WAITS_FOR_PENDING_WRITES").toBe("turn"); + await slow; await blocking; + await budgets.warnings("voyage"); + await budgets.backstop("voyage"); + await budgets.observeProvider("voyage", { attemptId: "attempt-1", family: "gemini", operation: "generate", origin: "https://ithaca.invalid", key: "GEMINI_API_KEY", + model: "odysseus", outcome: "completed", status: 200, usage: { inputTokens: 1 } } as never); + const log = await sites.log(); + expect(log[0], "BUDGET_AFTER_EARLIER_ACTIVITY_WRITE").toBe("activity_spans"); + expect(new Set(log.slice(1)), "BUDGET_TABLES_WRITTEN_ON_WORKER").toEqual(new Set(["turn_budget_usage", "turn_budget_days", "turn_budget_warnings", "provider_relay_usage"])); + } finally { runtime.close(); await f.close(); } +}); diff --git a/packages/ui-server/tests/provider-relay-budgets.test.ts b/packages/ui-server/tests/provider-relay-budgets.test.ts index 28c7061e04..ab8fff4e25 100644 --- a/packages/ui-server/tests/provider-relay-budgets.test.ts +++ b/packages/ui-server/tests/provider-relay-budgets.test.ts @@ -5,6 +5,7 @@ * provider reported nothing. A blocked scope refuses the next request. */ import { expect, test } from "bun:test"; +import { expectRejection } from "./helpers/expect-rejection.js"; import { resolveGovernancePolicy } from "@schlessera/brain-common/internal/governance-policy"; import type { ProviderRelayReceipt } from "@schlessera/brain-ui-sdk/internal"; import { createUiDb } from "../src/db/client.js"; @@ -26,20 +27,20 @@ const attempt = (overrides: Partial = {}): ProviderRelayRe model: "gemini-embedding-2", outcome: "completed", status: 200, inlineMedia: false, usage: { inputTokens: 1_000_000, inputByModality: { TEXT: 1_000_000 } }, ...overrides }); -test("the receipt states that relayed CLI providers are metered", () => { +test("the receipt states that relayed CLI providers are metered", async () => { const f = fixture(); try { - expect(f.budgets.begin("athena", "ithaca", policy()).coverage, "COVERAGE_NAMES_RELAYED_PROVIDERS").toBe(TURN_BUDGET_COVERAGE); + expect((await f.budgets.begin("athena", "ithaca", policy())).coverage, "COVERAGE_NAMES_RELAYED_PROVIDERS").toBe(TURN_BUDGET_COVERAGE); expect(TURN_BUDGET_COVERAGE).toBe("Built-in inference and relayed brain CLI providers; other provider and tool spend is not metered"); } finally { f.close(); } }); -test("a priced embedding attempt counts once toward every scope", () => { +test("a priced embedding attempt counts once toward every scope", async () => { const f = fixture(); try { - f.budgets.begin("athena", "ithaca", policy()); + await f.budgets.begin("athena", "ithaca", policy()); const first = attempt(); - const result = f.budgets.observeProvider("athena", first)!; + const result = (await f.budgets.observeProvider("athena", first))!; const turn = result.receipt.scopes.find(scope => scope.scope === "turn")!; expect(turn.costUsd, "PROVIDER_EMBEDDING_PRICED").not.toBeNull(); expect(turn.costUsd!).toBeCloseTo(0.2, 10); @@ -48,60 +49,60 @@ test("a priced embedding attempt counts once toward every scope", () => { // The agent's own inference figures stay separate. expect(result.receipt.listPriceCostUsd).toBeNull(); // A replayed observation of the same attempt adds nothing. - expect(f.budgets.observeProvider("athena", first), "PROVIDER_ATTEMPT_RECORDED_ONCE").toBeUndefined(); + expect((await f.budgets.observeProvider("athena", first)), "PROVIDER_ATTEMPT_RECORDED_ONCE").toBeUndefined(); expect(f.budgets.receipt("athena").scopes[0]!.costUsd).toBeCloseTo(0.2, 10); } finally { f.close(); } }); -test("a blocked scope refuses the next relayed request, with nothing reserved before", () => { +test("a blocked scope refuses the next relayed request, with nothing reserved before", async () => { const f = fixture(); try { - f.budgets.begin("athena", "ithaca", policy()); - expect(f.budgets.providerAdmission("athena", policy()), "PROVIDER_ADMITTED_UNDER_LIMIT").toBeUndefined(); + await f.budgets.begin("athena", "ithaca", policy()); + expect((await f.budgets.providerAdmission("athena", policy())), "PROVIDER_ADMITTED_UNDER_LIMIT").toBeUndefined(); // $0.20 per million: 10 million text tokens is $2, the turn's block. - const result = f.budgets.observeProvider("athena", attempt({ usage: { inputTokens: 10_000_000, inputByModality: { TEXT: 10_000_000 } } }))!; + const result = (await f.budgets.observeProvider("athena", attempt({ usage: { inputTokens: 10_000_000, inputByModality: { TEXT: 10_000_000 } } })))!; expect(result.blocked?.scope, "PROVIDER_USAGE_CROSSES_BLOCK").toBe("turn"); - expect(f.budgets.providerAdmission("athena", policy())?.scope, "PROVIDER_REFUSED_ONCE_BLOCKED").toBe("turn"); + expect((await f.budgets.providerAdmission("athena", policy()))?.scope, "PROVIDER_REFUSED_ONCE_BLOCKED").toBe("turn"); } finally { f.close(); } }); -test("unpriced reported tokens count toward the token stop; missing usage never blocks", () => { +test("unpriced reported tokens count toward the token stop; missing usage never blocks", async () => { const f = fixture(); try { - f.budgets.begin("athena", "ithaca", policy()); + await f.budgets.begin("athena", "ithaca", policy()); const before = f.budgets.receipt("athena").scopes[0]!.unknownRuns; for (let index = 0; index < 5; index++) { - const unknown = f.budgets.observeProvider("athena", attempt({ usage: null, outcome: "failed", status: 500 }))!; + const unknown = (await f.budgets.observeProvider("athena", attempt({ usage: null, outcome: "failed", status: 500 })))!; expect(unknown.blocked, "UNKNOWN_SPEND_NEVER_BLOCKS").toBeUndefined(); } const turn = f.budgets.receipt("athena").scopes[0]!; expect(turn.costUsd, "UNKNOWN_SPEND_IS_NOT_ZERO").toBeNull(); expect(turn.unknownRuns, "UNKNOWN_ATTEMPTS_DISCLOSED").toBe(before + 5); - const image = f.budgets.observeProvider("athena", attempt({ family: "openai", operation: "image", model: "gpt-image-2.5-flare", - key: "OPENAI_API_KEY", usage: { inputTokens: 100, outputTokens: 900 } }))!; + const image = (await f.budgets.observeProvider("athena", attempt({ family: "openai", operation: "image", model: "gpt-image-2.5-flare", + key: "OPENAI_API_KEY", usage: { inputTokens: 100, outputTokens: 900 } })))!; expect(image.receipt.scopes[0]!.tokens).toBe(1000); expect(image.blocked?.scope, "UNPRICED_PROVIDER_TOKENS_REACH_TOKEN_STOP").toBe("turn"); } finally { f.close(); } }); -test("a maintenance run outside any turn counts toward its day only", () => { +test("a maintenance run outside any turn counts toward its day only", async () => { const f = fixture(); try { - expect(f.budgets.observeProvider(null, attempt({ usage: { inputTokens: 20_000_000, inputByModality: { TEXT: 20_000_000 } } })), + expect((await f.budgets.observeProvider(null, attempt({ usage: { inputTokens: 20_000_000, inputByModality: { TEXT: 20_000_000 } } }))), "MAINTENANCE_ATTEMPT_HAS_NO_TURN_RECEIPT").toBeUndefined(); - expect(f.budgets.providerAdmission(null, policy())?.scope, "MAINTENANCE_DAY_BLOCKS").toBe("day"); + expect((await f.budgets.providerAdmission(null, policy()))?.scope, "MAINTENANCE_DAY_BLOCKS").toBe("day"); // The day it shares refuses the next turn's admission too. - expect(() => f.budgets.begin("penelope", "loom", policy()), "MAINTENANCE_SPEND_IN_DAY_ADMISSION").toThrow("UTC day inference budget reached"); + await expectRejection(f.budgets.begin("penelope", "loom", policy()), "UTC day inference budget reached", "MAINTENANCE_SPEND_IN_DAY_ADMISSION"); f.tick(24 * 60 * 60 * 1000); // A new day: the turn's own scopes never held the maintenance run. - const fresh = f.budgets.begin("penelope", "loom", policy()).scopes; + const fresh = (await f.budgets.begin("penelope", "loom", policy())).scopes; expect(fresh.map(scope => scope.blocked), "MAINTENANCE_DAY_RESETS_FOR_TURNS").toEqual([false, false, false]); expect(fresh.find(scope => scope.scope === "day")!.costUsd ?? 0).toBe(0); - expect(f.budgets.providerAdmission(null, policy()), "MAINTENANCE_DAY_RESETS").toBeUndefined(); + expect((await f.budgets.providerAdmission(null, policy())), "MAINTENANCE_DAY_RESETS").toBeUndefined(); } finally { f.close(); } }); -test("pricing: embedding modalities, inline media without a split, and completions", () => { +test("pricing: embedding modalities, inline media without a split, and completions", async () => { expect(GEMINI_EMBEDDING_PRICES["gemini-embedding-2"]).toEqual({ TEXT: 0.2 / 1_000_000, IMAGE: 0.45 / 1_000_000 }); expect(priceProviderAttempt(attempt({ usage: { inputTokens: 2_000_000, inputByModality: { TEXT: 1_000_000, IMAGE: 1_000_000 } } }), pricing).costUsd) .toBeCloseTo(0.65, 10); @@ -116,30 +117,30 @@ test("pricing: embedding modalities, inline media without a split, and completio expect(priceProviderAttempt(attempt({ usage: null }), pricing)).toEqual({ costUsd: null, estimate: null, tokens: null }); }); -test("a dispatched attempt is unknown spend until it settles, then counts once", () => { +test("a dispatched attempt is unknown spend until it settles, then counts once", async () => { const f = fixture(); try { - f.budgets.begin("athena", "ithaca", policy()); + await f.budgets.begin("athena", "ithaca", policy()); const settled = attempt(); const before = f.budgets.receipt("athena").scopes[0]!.unknownRuns; - const pending = f.budgets.observeProvider("athena", { ...settled, outcome: "dispatched", status: 0, usage: null })!; + const pending = (await f.budgets.observeProvider("athena", { ...settled, outcome: "dispatched", status: 0, usage: null }))!; expect(pending.receipt.scopes[0]!.unknownRuns, "DISPATCHED_ATTEMPT_DISCLOSED").toBe(before + 1); - const done = f.budgets.observeProvider("athena", settled)!; + const done = (await f.budgets.observeProvider("athena", settled))!; expect(done.receipt.scopes[0]!.unknownRuns, "SETTLED_ATTEMPT_REPLACES_DISPATCH").toBe(before); expect(done.receipt.scopes[0]!.costUsd).toBeCloseTo(0.2, 10); - expect(f.budgets.observeProvider("athena", { ...settled, outcome: "dispatched", status: 0, usage: null }), "SETTLED_ATTEMPT_FINAL").toBeUndefined(); + expect((await f.budgets.observeProvider("athena", { ...settled, outcome: "dispatched", status: 0, usage: null })), "SETTLED_ATTEMPT_FINAL").toBeUndefined(); } finally { f.close(); } }); -test("an attempt that settles after midnight counts toward the day it settled in", () => { +test("an attempt that settles after midnight counts toward the day it settled in", async () => { const f = fixture(); try { f.tick(14 * 60 * 60 * 1000 + 59 * 60 * 1000); // 23:59 UTC - f.budgets.begin("athena", "ithaca", policy()); + await f.budgets.begin("athena", "ithaca", policy()); const settled = attempt({ usage: { inputTokens: 20_000_000, inputByModality: { TEXT: 20_000_000 } } }); - f.budgets.observeProvider("athena", { ...settled, outcome: "dispatched", status: 0, usage: null }); + await f.budgets.observeProvider("athena", { ...settled, outcome: "dispatched", status: 0, usage: null }); f.tick(2 * 60 * 1000); - f.budgets.observeProvider("athena", settled); - expect(f.budgets.providerAdmission(null, policy())?.scope, "SETTLED_ATTEMPT_ON_ITS_SETTLEMENT_DAY").toBe("day"); + await f.budgets.observeProvider("athena", settled); + expect((await f.budgets.providerAdmission(null, policy()))?.scope, "SETTLED_ATTEMPT_ON_ITS_SETTLEMENT_DAY").toBe("day"); } finally { f.close(); } }); diff --git a/packages/ui-server/tests/turn-budgets.test.ts b/packages/ui-server/tests/turn-budgets.test.ts index dfa662ad16..04adb75ee1 100644 --- a/packages/ui-server/tests/turn-budgets.test.ts +++ b/packages/ui-server/tests/turn-budgets.test.ts @@ -1,4 +1,5 @@ import { expect, test } from "bun:test"; +import { expectRejection } from "./helpers/expect-rejection.js"; import { createUiDb } from "../src/db/client.js"; import { createTurnBudgets, BudgetExceeded, inferenceBudgetLimits, inferenceBudgetContext } from "../src/brain/turn-budgets.js"; import { resolveGovernancePolicy } from "@schlessera/brain-common/internal/governance-policy"; @@ -10,13 +11,13 @@ function fixture() { return { db, budgets, tick(ms: number) { now += ms; }, close() { db.close(); } }; } const usage = (n: number, model = "odysseus") => ({ inputTokens: n, outputTokens: 0, cacheReadTokens: 0, cacheCreationTokens: 0, perModel: { [model]: { inputTokens: n } } }); -test("list-price completed usage warns once; resumed terminal totals and subscription effective zero cannot change it", () => { +test("list-price completed usage warns once; resumed terminal totals and subscription effective zero cannot change it", async () => { const f = fixture(); try { - f.budgets.begin("athena", "ithaca", policy()); - expect(f.budgets.observe("athena", usage(1))!.warnings.map(w => w.scope)).toEqual(["turn"]); - expect(f.budgets.observe("athena", usage(1))!.warnings).toEqual([]); - expect(f.budgets.observe("athena", usage(2))!.blocked?.scope).toBe("turn"); + await f.budgets.begin("athena", "ithaca", policy()); + expect((await f.budgets.observe("athena", usage(1)))!.warnings.map(w => w.scope)).toEqual(["turn"]); + expect((await f.budgets.observe("athena", usage(1)))!.warnings).toEqual([]); + expect((await f.budgets.observe("athena", usage(2)))!.blocked?.scope).toBe("turn"); f.db.query("INSERT INTO activity_run_rollups(run_id,origin,name,session_id,started_at,cost_usd,effective_cost_usd,input_tokens) VALUES('athena','session','invoke_agent','ithaca',?,200,0,200)").run(Date.UTC(2026, 6, 12)); expect(f.budgets.receipt("athena").scopes.find(s => s.scope === 'session')!.costUsd).toBe(2); } @@ -24,46 +25,46 @@ test("list-price completed usage warns once; resumed terminal totals and subscri f.close(); } }); -test("unpriced usage remains unknown and token fallback blocks without inventing zero", () => { +test("unpriced usage remains unknown and token fallback blocks without inventing zero", async () => { const f = fixture(); try { - f.budgets.begin("athena", "ithaca", policy()); - const first = f.budgets.observe("athena", usage(9, "unknown"))!; + await f.budgets.begin("athena", "ithaca", policy()); + const first = (await f.budgets.observe("athena", usage(9, "unknown")))!; expect(first.receipt.listPriceCostUsd).toBeNull(); expect(first.blocked).toBeUndefined(); - expect(f.budgets.observe("athena", usage(10, "unknown"))!.blocked?.scope).toBe("turn"); + expect((await f.budgets.observe("athena", usage(10, "unknown")))!.blocked?.scope).toBe("turn"); } finally { f.close(); } }); -test("retained root Activity list prices block session/day admission and live sessions share day totals", () => { +test("retained root Activity list prices block session/day admission and live sessions share day totals", async () => { const f = fixture(); try { f.db.query("INSERT INTO activity_run_rollups(run_id,origin,name,session_id,started_at,cost_usd,effective_cost_usd,input_tokens) VALUES('old','session','invoke_agent','ithaca',?,4,0,4)").run(Date.UTC(2026, 6, 12)); - expect(() => f.budgets.begin("new", "ithaca", policy())).toThrow(BudgetExceeded); + await expectRejection(f.budgets.begin("new", "ithaca", policy()), BudgetExceeded); f.tick(120000); - expect(() => f.budgets.begin("next", "other", policy())).not.toThrow(); - f.budgets.observe("next", usage(2)); - f.budgets.begin("peer", "third", policy()); - f.budgets.observe("peer", usage(2)); - expect(() => f.budgets.begin("later", "fourth", policy())).toThrow(BudgetExceeded); + await f.budgets.begin("next", "other", policy()); + await f.budgets.observe("next", usage(2)); + await f.budgets.begin("peer", "third", policy()); + await f.budgets.observe("peer", usage(2)); + await expectRejection(f.budgets.begin("later", "fourth", policy()), BudgetExceeded); expect(f.budgets.receipt("peer").scopes.find(s => s.scope === 'day')!.resetAt).toBe(Date.UTC(2026, 6, 14)); } finally { f.close(); } }); -test("UTC response slices do not recharge yesterday's cumulative usage and survive host restart", () => { +test("UTC response slices do not recharge yesterday's cumulative usage and survive host restart", async () => { const f = fixture(); try { - f.budgets.begin("athena", "ithaca", policy()); - f.budgets.observe("athena", usage(1)); + await f.budgets.begin("athena", "ithaca", policy()); + await f.budgets.observe("athena", usage(1)); f.tick(120000); - f.budgets.observe("athena", usage(2)); + await f.budgets.observe("athena", usage(2)); expect(f.budgets.receipt("athena").scopes.find(s => s.scope === 'day')!.costUsd).toBe(1); const restarted = createTurnBudgets(f.db, { resolve: () => null }, () => Date.UTC(2026, 6, 13, 0, 1)); - restarted.begin("later", "ithaca", policy()); + await restarted.begin("later", "ithaca", policy()); expect(restarted.receipt("later").scopes.find(s => s.scope === 'session')!.costUsd).toBe(2); expect(restarted.receipt("later").scopes.find(s => s.scope === 'day')!.costUsd).toBe(1); } @@ -71,45 +72,45 @@ test("UTC response slices do not recharge yesterday's cumulative usage and survi f.close(); } }); -test("the shared Activity calculator uses the admitted pricing route and frozen per-model rates for all token classes", () => { +test("the shared Activity calculator uses the admitted pricing route and frozen per-model rates for all token classes", async () => { const db = createUiDb(":memory:"); let inputRate = 0.01; const routes: Array = []; const budgets = createTurnBudgets(db, { resolve: (_model, route) => { routes.push(route); return { input: route === "openrouter" ? inputRate : 10, output: 0.1, cacheRead: 0.001, cacheWrite: 0.02, estimate: false, source: "snapshot" }; } }); try { - budgets.begin("athena", "ithaca", policy(), "openrouter"); + await budgets.begin("athena", "ithaca", policy(), "openrouter"); const sample = { inputTokens: 10, outputTokens: 2, cacheReadTokens: 100, cacheCreationTokens: 5, perModel: { odysseus: { inputTokens: 10, outputTokens: 2, cacheReadTokens: 100, cacheCreationTokens: 5 } } }; - expect(budgets.observe("athena", sample)!.receipt.listPriceCostUsd, "BUDGET_ACTIVITY_CATALOG_AND_FOUR_CLASSES").toBeCloseTo(0.5); + expect((await budgets.observe("athena", sample))!.receipt.listPriceCostUsd, "BUDGET_ACTIVITY_CATALOG_AND_FOUR_CLASSES").toBeCloseTo(0.5); inputRate = 100; - expect(budgets.observe("athena", sample)!.blocked).toBeUndefined(); + expect((await budgets.observe("athena", sample))!.blocked).toBeUndefined(); expect(routes).toEqual(["openrouter"]); } finally { db.close(); } }); -test("adding unpriced inference cannot erase a priced lower bound, warning or stop", () => { +test("adding unpriced inference cannot erase a priced lower bound, warning or stop", async () => { const f = fixture(); try { - f.budgets.begin("athena", "ithaca", policy()); - f.budgets.observe("athena", usage(1)); + await f.budgets.begin("athena", "ithaca", policy()); + await f.budgets.observe("athena", usage(1)); const mixed = (n: number) => ({ inputTokens: n + 1, outputTokens: 0, cacheReadTokens: 0, cacheCreationTokens: 0, perModel: { odysseus: { inputTokens: n }, unknown: { inputTokens: 1 } } }); - const first = f.budgets.observe("athena", mixed(1))!; + const first = (await f.budgets.observe("athena", mixed(1)))!; expect(first.receipt.listPriceCostUsd).toBeNull(); expect(first.receipt.scopes[0]).toMatchObject({ costUsd: 1, unknownRuns: 1, warning: true, blocked: false }); - expect(f.budgets.observe("athena", mixed(2))!.blocked).toMatchObject({ scope: "turn", costUsd: 2 }); + expect((await f.budgets.observe("athena", mixed(2)))!.blocked).toMatchObject({ scope: "turn", costUsd: 2 }); } finally { f.close(); } }); -test("retained warning evidence is disclosed once at admission, without inventing a response", () => { +test("retained warning evidence is disclosed once at admission, without inventing a response", async () => { const f = fixture(); try { f.db.query("INSERT INTO activity_run_rollups(run_id,origin,name,session_id,started_at,cost_usd,input_tokens) VALUES('retained','session','invoke_agent','ithaca',?,3,3)").run(Date.UTC(2026, 6, 12)); - f.budgets.begin("athena", "ithaca", policy()); - expect(f.budgets.warnings("athena").map(item => item.scope)).toEqual(["session", "day"]); - expect(f.budgets.warnings("athena")).toEqual([]); + await f.budgets.begin("athena", "ithaca", policy()); + expect((await f.budgets.warnings("athena")).map(item => item.scope)).toEqual(["session", "day"]); + expect((await f.budgets.warnings("athena"))).toEqual([]); expect(f.budgets.saved("athena")?.listPriceCostUsd).toBeNull(); } finally { @@ -117,15 +118,15 @@ test("retained warning evidence is disclosed once at admission, without inventin } }); -test("session prompt limits exclude running usage and only actionable warnings enter current context", () => { +test("session prompt limits exclude running usage and only actionable warnings enter current context", async () => { const f = fixture(); try { - f.budgets.begin("athena", "ithaca", policy()); + await f.budgets.begin("athena", "ithaca", policy()); const initial = f.budgets.receipt("athena"); const fixed = inferenceBudgetLimits(initial); expect(fixed).toContain("turn: warn $1; stop $2"); expect(inferenceBudgetContext(initial)).toBeUndefined(); - f.budgets.observe("athena", usage(1)); + await f.budgets.observe("athena", usage(1)); const changed = f.budgets.receipt("athena"); expect(inferenceBudgetLimits(changed), "FIXED_LIMITS_HAVE_NO_RUNNING_USAGE").toBe(fixed); const warning = inferenceBudgetContext(changed)!; @@ -135,3 +136,15 @@ test("session prompt limits exclude running usage and only actionable warnings e expect(fixed).not.toContain(String(changed.observedAt)); } finally { f.close(); } }); +test("a failed session binding makes the turn's accounting unavailable", async () => { + const f = fixture(); + try { + await f.budgets.begin("athena", undefined, policy()); + f.db.exec("CREATE TEMP TRIGGER odysseus_bind_failure BEFORE UPDATE OF session_id ON main.turn_budget_usage BEGIN SELECT RAISE(ABORT, 'Scylla took the binding'); END"); + await expectRejection(f.budgets.bind("athena", "ithaca"), "Scylla took the binding"); + expect(() => f.budgets.observe("athena", usage(1)), "BUDGET_FAILED_BIND_STOPS_ACCOUNTING").toThrow("Scylla took the binding"); + } + finally { + f.close(); + } +}); diff --git a/packages/ui-server/tests/ws-turn-budgets.test.ts b/packages/ui-server/tests/ws-turn-budgets.test.ts index 13734fd332..6ac0b995ec 100644 --- a/packages/ui-server/tests/ws-turn-budgets.test.ts +++ b/packages/ui-server/tests/ws-turn-budgets.test.ts @@ -22,9 +22,9 @@ import { WsHost } from "../src/ws/host.js"; import { createTurnBudgets } from "../src/brain/turn-budgets.js"; import { parseServerMessage } from "@schlessera/brain-ui-sdk/schemas"; for (const adapter of ["claude", "pi"] as const) - for (const mode of adapter === "claude" ? ["token", "usd", "sdk", "unavailable"] : ["token", "usd"]) - test(`actual ${adapter} completed response crosses the host ${mode === "unavailable" ? "unavailable accounting" : mode === "sdk" ? "SDK backstop" : mode === "usd" ? "USD budget" : "token fallback"} and stops with ${mode === "unavailable" ? "budget_unavailable" : "budget_exceeded"}, preserving partial work`, async () => { - const priced = mode !== "token", backstop = mode === "sdk", unavailable = mode === "unavailable"; + for (const mode of adapter === "claude" ? ["token", "usd", "sdk", "unavailable", "late"] : ["token", "usd"]) + test(`actual ${adapter} completed response crosses the host ${mode === "unavailable" ? "unavailable accounting" : mode === "sdk" ? "SDK backstop" : mode === "usd" ? "USD budget" : mode === "late" ? "USD budget in a commit that lands after the response" : "token fallback"} and stops with ${mode === "unavailable" ? "budget_unavailable" : "budget_exceeded"}, preserving partial work`, async () => { + const priced = mode !== "token", backstop = mode === "sdk", unavailable = mode === "unavailable", late = mode === "late"; let failAccounting: () => void = () => { }; const root = await mkdtemp(resolve(tmpdir(), "odysseus-budget-host-")), brain = resolve(root, "brain"); await createFixtureBrain(resolve(import.meta.dir, "../../.."), brain); @@ -73,7 +73,11 @@ for (const adapter of ["claude", "pi"] as const) const app = await createApp({ config: resolveServerConfig({ BRAIN_PATH: brain, DB_PATH: resolve(root, "ui.db"), HOST: "127.0.0.1", AUTH_MODE: "none", NODE_ENV: "test", BRAIN_UI_MODEL_DISCOVERY: "0", BRAIN_UI_PRICING_DISCOVERY: "0" }), registry: createStaticBackendRegistry([backend], backend.id), observability: createRecordingObservability() }); const host = priced ? new WsHost({ brainPath: brain, registry: createStaticBackendRegistry([backend], backend.id), catalog: app.wsHost.catalog, governedTurns: app.wsHost.governedTurns, journal: app.wsHost.journal, activity: app.wsHost.activity ?? undefined, - budgets: createTurnBudgets(app.db, { resolve: () => ({ input: 1, output: 1, cacheRead: 1, cacheWrite: 1, estimate: false, source: "snapshot" }) }) }) : app.wsHost; + budgets: (() => { + const budgets = createTurnBudgets(app.db, { resolve: () => ({ input: 1, output: 1, cacheRead: 1, cacheWrite: 1, estimate: false, source: "snapshot" }) }); + // The blocking report's commit lands after the backend has returned (#2850). + return late ? { ...budgets, observe: (...args: Parameters) => { const committed = budgets.observe(...args); return committed && Bun.sleep(400).then(() => committed); } } : budgets; + })() }) : app.wsHost; failAccounting = () => app.db.exec("DROP TABLE turn_budget_days"); try { const principal = createPrincipal(app.db, { authMethod: "password", label: "Odysseus device", ttlSeconds: 60 }); @@ -94,7 +98,7 @@ for (const adapter of ["claude", "pi"] as const) expect(receivedLimits, "STATIC_BUDGET_LIMITS_STILL_TRANSPORTED").toContain("turn: warn $5; stop $10"); expect(receivedContext, "NO_ROUTINE_BUDGET_TAIL_AT_ADMISSION").toBeUndefined(); expect(called, "BUDGET_ACTUAL_ADAPTER_LAUNCHED").toBe(1); - expect(aborted, "BUDGET_HOST_ABORTS_ACTUAL_ADAPTER").toBe(true); + expect(aborted, "BUDGET_HOST_ABORTS_ACTUAL_ADAPTER").toBe(!late); expect(maxBudget).toBe(10); const failure = frames().findLast(frame => frame.type === "error")!; expect(failure).toMatchObject(unavailable ? { code: "BUDGET_UNAVAILABLE", failure: { errorClass: "budget_unavailable" } } : { code: "BUDGET_EXCEEDED", failure: { errorClass: "budget_exceeded", budgetScope: "turn" } });